[NPU] add zbal support for npu (#24575)

python/sglang/srt/distributed/parallel_state.py

@@ -1598,7 +1598,7 @@ _DEVICE_TO_DISTRIBUTED_BACKEND = {
     "xpu": "xccl",
     "hpu": "hccl",
     "cpu": "gloo",
-    "npu": "hccl",
+    "npu": "hccl" if not envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0 else "zbal",
     "musa": "mccl",
 }
 

python/sglang/srt/environ.py

@@ -558,6 +558,10 @@ class Envs:
     # TokenizerManager
     SGLANG_REQUEST_STATE_WAIT_TIMEOUT = EnvInt(4)
 
+    # ZBAL, zero buffer accelerate library, currently worked only in npu
+    SGLANG_ZBAL_LOCAL_MEM_SIZE = EnvInt(0)
+    SGLANG_ZBAL_BOOTSTRAP_URL = EnvStr("")
+
     SGLANG_DEFAULT_THINKING = EnvBool(False)
 
     # ====================================================================

python/sglang/srt/hardware_backend/npu/utils.py

@@ -1,5 +1,6 @@
 import functools
 import logging
+import sys
 from enum import IntEnum
 from typing import TYPE_CHECKING, Callable
 
@@ -14,6 +15,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 _is_npu = is_npu()
 indexer_weight_stream = None
+gva_is_inited = False
 
 
 class NPUACLFormat(IntEnum):
@@ -184,6 +186,96 @@ def get_indexer_weight_stream():
     return indexer_weight_stream
 
 
+def init_zbal(world_size, gpu_id, world_rank, do_check=True):
+    """
+    init zbal, if is mix alloc mode, only register for sma & comm
+    """
+    zbal_mem_size = envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get()
+    if not zbal_mem_size > 0:
+        return 1
+
+    global gva_is_inited
+    from zbal import is_mix_alloc, switch_to_allocator, zbal_init
+
+    if is_mix_alloc():
+        switch_to_allocator()
+        # use lazy init for mix alloc
+        return 1
+    else:
+        if envs.SGLANG_ZBAL_BOOTSTRAP_URL.get():
+            ret = zbal_init(
+                world_size,
+                gpu_id,
+                world_rank,
+                zbal_mem_size * (1024**2),
+                ip_port=envs.SGLANG_ZBAL_BOOTSTRAP_URL.get(),
+            )
+        else:
+            ret = zbal_init(world_size, gpu_id, world_rank, zbal_mem_size * (1024**2))
+
+        gva_is_inited = True
+
+        if do_check and not ret:
+            logger.error(f"[ZBAL] zbal init failed!")
+            sys.exit(-1)
+
+        return ret
+
+
+def lazy_init_zbal_gva_mem(
+    device, gpu_id, world_rank, world_size, cpu_group=None, do_check=True
+):
+    """
+    lazy init zbal gva mem, keep weights and kv remains alloc by dma vmm to avoid memory fragment
+    """
+    from zbal import is_mix_alloc, zbal_init
+
+    if not is_mix_alloc():
+        logger.info(
+            "lazy init is supported only in mix alloc mode, this action will be passed"
+        )
+        return 1
+
+    global gva_is_inited
+    from sglang.srt.utils.common import get_available_gpu_memory
+
+    # TODO need to use allgather if you want use total_memory stats from mem_get_info as unbalance os
+    total_memory = 61.2  # 2.5GB for other (workspace & os) outside torch
+    free_gpu_memory = get_available_gpu_memory(
+        device,
+        gpu_id,
+        distributed=world_size > 1,
+        cpu_group=cpu_group,
+        empty_cache=True,
+    )
+
+    used_memory = total_memory - free_gpu_memory
+
+    used_memory_in_mb = int(used_memory * 1024)
+    gva_in_mb = envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() - used_memory_in_mb
+    gva_in_mb = gva_in_mb - gva_in_mb % 128  # align to 128MB
+    print(f"[ZBAL] rank {world_rank} allocated {gva_in_mb} MB gva space.")
+
+    assert not gva_is_inited, "zbal gva should be inited only once"
+    # zbal_set_logger_level(0)
+    if envs.SGLANG_ZBAL_BOOTSTRAP_URL.get():
+        res = zbal_init(
+            world_size,
+            gpu_id,
+            world_rank,
+            gva_in_mb * (1024**2),
+            ip_port=envs.SGLANG_ZBAL_BOOTSTRAP_URL.get(),
+        )
+    else:
+        res = zbal_init(world_size, gpu_id, world_rank, gva_in_mb * (1024**2))
+
+    gva_is_inited = True
+    if do_check and not res:
+        logger.error(f"[ZBAL] zbal lazy init failed!")
+        sys.exit(-1)
+    return res
+
+
 share_stream = None
 routed_stream = None
 

python/sglang/srt/layers/moe/token_dispatcher/deepep.py

@@ -40,7 +40,11 @@ if TYPE_CHECKING:
     from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs
 
 try:
-    from deep_ep import Buffer, Config
+    if _is_npu and envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0:
+        from zbal.zbal.deepep_adaptor import Config
+        from zbal.zbal_buffer import Buffer
+    else:
+        from deep_ep import Buffer, Config
 
     if not _is_npu:
         from sglang.srt.layers.quantization.fp8_kernel import (

python/sglang/srt/managers/scheduler.py

@@ -417,6 +417,9 @@ class Scheduler(
         # Init inter-process communication
         self.init_ipc_channels(port_args)
 
+        # Init ZBAL, switch allocator should before any torch alloc action
+        self.init_zbal_on_npu()
+
         # Init PD-multiplexing context
         if self.enable_pdmux:
             self.init_pdmux()
@@ -496,6 +499,16 @@ class Scheduler(
 
         self.is_initializing = False
 
+    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
+            )  # only switch allocator if is mix mode
+
     def init_model_config(self):
         self.model_config = ModelConfig.from_server_args(self.server_args)
         if _is_npu:

python/sglang/srt/model_executor/model_runner.py

@@ -785,9 +785,23 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 )
             self._pre_initialize_flashinfer_allreduce_workspace()
             self.init_device_graphs()
-        elif self.device in ["npu", "cpu"]:
+        elif self.device == "cpu":
             self.init_attention_backend()
             self.init_device_graphs()
+        elif self.device == "npu":
+            self.init_attention_backend()
+            # lazy init for zbal with mix mode(before graph capture when enable_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()
         elif current_platform.is_out_of_tree():
             self.init_attention_backend()
             if current_platform.support_cuda_graph():

python/sglang/srt/utils/common.py

@@ -587,7 +587,16 @@ def get_available_gpu_memory(
             )
         if empty_cache:
             empty_device_cache(torch.npu)
-        free_gpu_memory, total_gpu_memory = torch.npu.mem_get_info()
+        if envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0:
+            import zbal
+
+            if not zbal.is_mix_alloc():
+                free_gpu_memory, total_gpu_memory = zbal.zbal_module.mem_get_info()
+            else:
+                # mix mode fall back into npu mem info since gva may not inited yet
+                free_gpu_memory, total_gpu_memory = torch.npu.mem_get_info()
+        else:
+            free_gpu_memory, total_gpu_memory = torch.npu.mem_get_info()
     elif device == "musa":
         num_gpus = torch.musa.device_count()
         assert gpu_id < num_gpus
@@ -1689,7 +1698,10 @@ def get_npu_memory_capacity():
     try:
         import torch_npu  # noqa: F401
 
-        return torch.npu.mem_get_info()[1] // 1024 // 1024  # unit: MB
+        if envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0:
+            return envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get()  # unit: MB
+        else:
+            return torch.npu.mem_get_info()[1] // 1024 // 1024  # unit: MB
     except ImportError as e:
         raise ImportError("torch_npu is required when run on npu device.")