From 9e00b7ca95aac732eefeb7412cb514183f494561 Mon Sep 17 00:00:00 2001 From: litangss Date: Wed, 13 May 2026 20:52:42 +0800 Subject: [PATCH] [NPU] add zbal support for npu (#24575) --- .../sglang/srt/distributed/parallel_state.py | 2 +- python/sglang/srt/environ.py | 4 + .../sglang/srt/hardware_backend/npu/utils.py | 92 +++++++++++++++++++ .../srt/layers/moe/token_dispatcher/deepep.py | 6 +- python/sglang/srt/managers/scheduler.py | 13 +++ .../sglang/srt/model_executor/model_runner.py | 16 +++- python/sglang/srt/utils/common.py | 16 +++- 7 files changed, 144 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index c42760ce5..2877d2ed9 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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", } diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e40a24c0a..cbd81b0d3 100644 --- a/python/sglang/srt/environ.py +++ b/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) # ==================================================================== diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 538eda243..7139351fc 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index f990d7d00..790bc5b01 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/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 ( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 689cf9837..f9fac2ea5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index f64881f99..d94f7f174 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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(): diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 8043c5c08..fd9d6fd74 100644 --- a/python/sglang/srt/utils/common.py +++ b/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.")