执行摘要

• 一句话：NPU 后端集成 zbal 零缓冲区加速库
• 推荐动作：该 PR 值得 NPU 用户和关心内存优化的工程师精读，尤其是 lazy_init_zbal_gva_mem 的混合分配策略。建议在合并后补充 NPU 环境下的性能基准测试文档。

功能与动机

zbal 是基于 memfabric 的零缓冲区加速库，已在 Ascend NPU 上开源。PR body 中称其在 w8a8 & 4k per device in DSV3 like moe process 上实现了 30%+ 的加速和 1.2G+ 的内存效率提升。通过集成 zbal，NPU 用户可以显著提升模型推理性能并优化内存利用率。

实现拆解

1. 环境变量注册：在 python/sglang/srt/environ.py 中添加 SGLANG_ZBAL_LOCAL_MEM_SIZE 和 SGLANG_ZBAL_BOOTSTRAP_URL 两个环境变量，分别控制 zbal 分配的内存大小（MB）和分布式引导 URL。
2. 初始化入口：在 python/sglang/srt/managers/scheduler.py 的 __init__ 中新增 init_zbal_on_npu() 调用，放在 IPC 通道初始化之后、任何 PyTorch 分配之前。新建方法 init_zbal_on_npu 在 NPU 设备上调用 sglang.srt.hardware_backend.npu.utils.init_zbal。
3. 核心 init / lazy_init：在 python/sglang/srt/hardware_backend/npu/utils.py 中新增 init_zbal 和 lazy_init_zbal_gva_mem 两个函数。init_zbal 若为 mix alloc 模式则仅调用 switch_to_allocator() 并返回；否则直接通过 zbal_init 初始化全部 GVA 空间。lazy_init_zbal_gva_mem 在权重和 KV cache 加载完毕（即 model_runner 的 initialize 中）调用，计算剩余可用内存并动态分配 GVA 空间，实现混合分配。
4. 模型执行器适配：在 python/sglang/srt/model_executor/model_runner.py 的 initialize 中将 NPU 分支从 elif self.device in ["npu", "cpu"] 拆出单独处理，在 attention 后端初始化之后、设备图捕获之前插入 lazy_init_zbal_gva_mem 调用。
5. 内存查询改写：在 python/sglang/srt/utils/common.py 中修改 get_available_gpu_memory 和 get_npu_memory_capacity，当启用 zbal 时调用 zbal 自己的内存查询接口（zbal.zbal_module.mem_get_info）或直接返回预设值。
6. DeepEP 适配：在 python/sglang/srt/layers/moe/token_dispatcher/deepep.py 中根据是否启用 zbal 条件导入 zbal.zbal.deepep_adaptor 和 zbal.zbal_buffer 替代原 deep_ep。

关键文件：

• python/sglang/srt/hardware_backend/npu/utils.py（模块 NPU 工具；类别 source；类型 core-logic；符号 init_zbal, lazy_init_zbal_gva_mem）: 核心修改文件，新增 init_zbal 和 lazy_init_zbal_gva_mem 函数，实现 zbal 的初始化和延迟分配逻辑。
• python/sglang/srt/managers/scheduler.py（模块 调度器；类别 source；类型 core-logic；符号 init_zbal_on_npu）: 调度器初始化入口，新增 init_zbal_on_npu 方法，确保在 torch 分配前切换分配器。
• python/sglang/srt/model_executor/model_runner.py（模块 模型执行器；类别 source；类型 data-contract）: 模型执行器初始化，将 NPU 分支拆出并在 attention 后端初始化后、CUDA graph 捕获前调用 lazy_init_zbal_gva_mem。

关键符号：init_zbal, lazy_init_zbal_gva_mem, init_zbal_on_npu, get_available_gpu_memory, get_npu_memory_capacity

关键源码片段

python/sglang/srt/managers/scheduler.py

调度器初始化入口，新增 init_zbal_on_npu 方法，确保在 torch 分配前切换分配器。

# python/sglang/srt/managers/scheduler.py
​
class Scheduler:
    def __init__(self, ...):
        # ... 其他初始化
        self.init_ipc_channels(port_args)
​
        # Init ZBAL，切换分配器必须发生在任何 torch 分配之前
        self.init_zbal_on_npu()
​
        # Init PD-multiplexing context
        if self.enable_pdmux:
            self.init_pdmux()
        # ...
​
    def init_zbal_on_npu(self):
        if _is_npu:
            from sglang.srt.hardware_backend.npu.utils import init_zbal
​
            if self.pp_size > 1:
                logger.error(f"only zbal mix mode support pp_size > 1!")
            init_zbal(
                self.tp_size, self.gpu_id, self.tp_rank
            ) # 如果是混合模式，仅切换分配器

python/sglang/srt/model_executor/model_runner.py

模型执行器初始化，将 NPU 分支拆出并在 attention 后端初始化后、CUDA graph 捕获前调用 lazy_init_zbal_gva_mem。

# python/sglang/srt/model_executor/model_runner.py
​
class ModelRunner:
    def initialize(self, pre_model_load_memory: float):
        # ... 前面初始化
        elif self.device == "npu":
            self.init_attention_backend()
            # lazy init for zbal with mix mode ( 必须在 CUDA graph 捕获之前 )
            if envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0 and not self.is_draft_worker:
                from sglang.srt.hardware_backend.npu.utils import lazy_init_zbal_gva_mem
​
                lazy_init_zbal_gva_mem(
                    self.device,
                    self.gpu_id,
                    get_world_group().rank_in_group,
                    get_world_group().world_size,
                    get_world_group().cpu_group,
                )
            self.init_device_graphs()

评论区精华

PR 的 review 评论较少，主要由 sglang-npu-bot 触发 CI 测试并最终批准。提交历史显示作者在多次合并 main 分支后进行了代码清理和 lint 应用，整体改动较为收敛，未出现重大设计争议。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. 依赖风险：引入外部库 zbal，该库只在 Ascend NPU 环境下可用，若缺失则导入失败，但代码通过条件导入（仅在 NPU 且环境变量开启时尝试）隔离，不会阻塞其他硬件平台。
  2. 兼容性：修改了 get_available_gpu_memory 和 get_npu_memory_capacity 的行为，当启用 zbal 时返回 zbal 管理的显存量而非实际物理显存，可能影响依赖这些函数进行内存预算的其他模块（如 cache 分配）。但该修改受环境变量开关保护，默认关闭。
  3. 初始化顺序敏感：init_zbal_on_npu 必须在任何 torch 分配之前调用，在 __init__ 中已安排在 IPC 初始化之后、其他 init 之前，时序正确。lazy_init_zbal_gva_mem 在权重和 KV 加载后调用，需确保此时显存已分配完毕。
  4. 非 NPU 平台：所有改动均包含 _is_npu 或环境变量守卫，不影响其他后端。
  5. 测试覆盖不足：源码改动未包含对应测试文件，仅通过 CI 进行集成验证，缺少单元测试。
     - 影响：影响范围限于 NPU 后端用户，通过环境变量自愿开启，默认关闭。对于不设置 SGLANG_ZBAL_LOCAL_MEM_SIZE 的用户，行为完全不变。对于开启的用户，zbal 将接管显存分配，可能提升 MoE 类模型的吞吐和显存效率，但需要用户正确配置内存大小并理解 zbal 的工作模式（mix alloc 或全 GVA）。对团队而言，该 PR 是 NPU 生态建设的重要一步，为后续更深入的硬件加速铺平道路。
     - 风险标记：依赖外部库 zbal, 初始化顺序敏感, 内存查询行为变更, 缺少单元测试覆盖

关联脉络

• PR #25114 [NPU] [DOC] add performance testing and optimization docs for npu: 同属于 NPU 生态建设系列，该文档可为 zbal 的使用提供性能测试参考。