diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 6bfd514dc..69a65ec4e 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1240,6 +1240,19 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): if self.req_to_token_pool.available_size() <= 0: break + # Hybrid models (e.g. K3 with KDA): guard against prealloc + # draining the mamba pool before the KV pool (would assert "Not + # enough space for mamba cache"). Evict a cached mamba slot from + # the radix tree first (only if it manages mamba states; + # ChunkCache.evict is a no-op), else stop. + mamba_allocator = getattr(self.req_to_token_pool, "mamba_allocator", None) + if mamba_allocator is not None and mamba_allocator.available_size() <= 0: + supports_mamba = self.tree_cache.supports_mamba() + if supports_mamba and hasattr(self.tree_cache, "evict"): + self.tree_cache.evict(EvictParams(num_tokens=0, mamba_num=1)) + if mamba_allocator.available_size() <= 0: + break + if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0: break diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 21902574f..2c14d17e2 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -2906,6 +2906,13 @@ def create_custom_parallel_group( Returns: The ProcessGroup if the current rank is in group_ranks, else None. + + NOTE: `group_ranks` must be the full rank list of the group, identical on + every rank of the world (e.g. obtained via get_process_group_ranks()). + Both paths below are world-collective: the general path performs a + world-size all_gather_object, and on NPU the fast path derives groups + locally from a rank-local check — a rank-local subset passed by only + some ranks would make ranks take different paths and deadlock. """ assert torch.distributed.is_initialized() @@ -2913,9 +2920,26 @@ def create_custom_parallel_group( rank = torch.distributed.get_rank() local_config = sorted(list(set(group_ranks))) - gathered_configs = [None for _ in range(world_size)] + group_size = len(local_config) - torch.distributed.all_gather_object(gathered_configs, local_config) + # Standard TP/DP partitioning: contiguous, group-aligned ranks. + is_standard_partition = ( + world_size % group_size == 0 + and local_config == list(range(local_config[0], local_config[0] + group_size)) + and local_config[0] % group_size == 0 + ) + + if not (_is_npu and is_standard_partition): + # General path: collect every rank's group via all_gather_object. + gathered_configs = [None for _ in range(world_size)] + torch.distributed.all_gather_object(gathered_configs, local_config) + else: + # NPU fast path: all_gather_object on the default HCCL PG allocates + # an HCCL buffer; instead derive the standard TP/DP groups locally. + num_groups = world_size // group_size + gathered_configs = [ + list(range(i * group_size, (i + 1) * group_size)) for i in range(num_groups) + ] unique_groups = [] seen_signatures = set() diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 5acc1cd69..e7705d9ea 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -964,12 +964,21 @@ def build_hybrid_mamba_stack( target_device_layer_num=kv_pool.layer_num, draft_layer_num=len(mtp_draft_device_pools), ) + # MambaPoolHost only supports page_first_direct; the global layout may be + # page_first_kv_split (e.g. MLA + KDA hybrid on NPU). The Mamba/KDA state + # pool has no separate K/V buffers, so kv_split does not apply; override + # to page_first_direct. + mamba_layout = ( + "page_first_direct" + if get_memory().hicache_mem_layout == "page_first_kv_split" + else get_memory().hicache_mem_layout + ) mamba_host_pool = MambaPoolHost( mamba_pool, get_memory().hicache_ratio, mamba_host_size, allocator_type=_get_allocator_type(), - layout=get_memory().hicache_mem_layout, + layout=mamba_layout, ) entries = [ build_pool_entry( @@ -1216,6 +1225,12 @@ def _build_mha_mla_host_pool( ): from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool + # The global layout is page_first_kv_split only when the target model + # uses MLA; that layout is MLA-specific, so MHA draft pools must use + # the non-MLA layout (NPU default: page_first_direct). + if isinstance(pool, MHATokenToKVPool) and layout == "page_first_kv_split": + layout = "page_first_direct" + kwargs = dict( host_to_device_ratio=host_to_device_ratio, host_size=0, diff --git a/python/sglang/srt/mem_cache/pool_host/mamba.py b/python/sglang/srt/mem_cache/pool_host/mamba.py index 51f1a400f..e8656ec15 100644 --- a/python/sglang/srt/mem_cache/pool_host/mamba.py +++ b/python/sglang/srt/mem_cache/pool_host/mamba.py @@ -26,6 +26,7 @@ _is_hip = is_hip() _is_npu = is_npu() transfer_state_per_layer_direct_pf_lf = None transfer_state_all_layer_direct_lf_pf = None +transfer_mamba_state = None if _is_cuda or _is_hip: from sgl_kernel.kvcacheio import ( transfer_kv_all_layer_direct_lf_pf, @@ -39,6 +40,12 @@ if _is_cuda or _is_hip: transfer_kv_mamba_pf_lf, ) if _is_npu: + from sgl_kernel_npu.kvcacheio import TransferDirection + + try: + from sgl_kernel_npu.kvcacheio import transfer_mamba_state + except ImportError: + transfer_mamba_state = None try: from sgl_kernel_npu.kvcacheio import ( transfer_state_all_layer_direct_lf_pf, @@ -373,6 +380,13 @@ class MambaPoolHost(HostKVCache): dst_indices=dst_indices, page_size=1, ) + elif io_backend == "kernel_ascend": + # Per-layer indexed copy: this method transfers a single layer + # (layer_first layout). The all-layer kernel path is handled by + # _copy_tensor_all_layers_lf_pf / load_to_device_per_layer. + dst[dst_indices.to(dst.device)] = src[src_indices.to(src.device)].to( + dst.device + ) else: raise ValueError(f"Unsupported io_backend: {io_backend}") @@ -474,7 +488,19 @@ class MambaPoolHost(HostKVCache): device_indices=src_indices, host_indices=dst_indices, ) + elif transfer_mamba_state is not None: + # NPU: mirror the load path — the dedicated kernel transfers all + # layers at once via a single 2D strided copy + # (device layer-first -> host page-first). + transfer_mamba_state( + device_buf=src_layers, + host_buf=dst, + device_indices=src_indices, + host_indices=dst_indices, + direction=TransferDirection.D2H, + ) else: + # Per-layer fallback when the dedicated kernel is unavailable. device_indices = src_indices.to( dtype=torch.int64, device=src_layers.device ) @@ -501,27 +527,50 @@ class MambaPoolHost(HostKVCache): is_draft: bool = False, ): if self.layout in ["page_first", "page_first_direct"]: - # no ssm state on conv-only models: nothing to transfer - if self.temporal_state_elem_size > 0: - self._copy_tensor_pf_lf( - src=self.temporal_buffer, - dst=device_pool.mamba_cache.temporal[layer_id], - src_indices=host_indices, - dst_indices=device_indices, - layer_id=layer_id, - num_layers=self.num_mamba_layers, - io_backend=io_backend, - ) - for conv_idx in range(len(self.conv_state_shapes)): - self._copy_tensor_pf_lf( - src=self.conv_buffer[conv_idx], - dst=device_pool.mamba_cache.conv[conv_idx][layer_id], - src_indices=host_indices, - dst_indices=device_indices, - layer_id=layer_id, - num_layers=self.num_mamba_layers, - io_backend=io_backend, - ) + if io_backend == "kernel_ascend" and transfer_mamba_state is not None: + # NPU: transfer all layers at once via dedicated kernel. + # layer_id == 0 covers every layer, so later calls must skip. + if layer_id == 0: + # no ssm state on conv-only models: a 0-size batched + # transfer errors, same guard as the per-layer path below + if self.temporal_state_elem_size > 0: + transfer_mamba_state( + device_buf=device_pool.mamba_cache.temporal, + host_buf=self.temporal_buffer, + device_indices=device_indices, + host_indices=host_indices, + direction=TransferDirection.H2D, + ) + for conv_idx in range(len(self.conv_state_shapes)): + transfer_mamba_state( + device_buf=device_pool.mamba_cache.conv[conv_idx], + host_buf=self.conv_buffer[conv_idx], + device_indices=device_indices, + host_indices=host_indices, + direction=TransferDirection.H2D, + ) + else: + # no ssm state on conv-only models: nothing to transfer + if self.temporal_state_elem_size > 0: + self._copy_tensor_pf_lf( + src=self.temporal_buffer, + dst=device_pool.mamba_cache.temporal[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + ) + for conv_idx in range(len(self.conv_state_shapes)): + self._copy_tensor_pf_lf( + src=self.conv_buffer[conv_idx], + dst=device_pool.mamba_cache.conv[conv_idx][layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + num_layers=self.num_mamba_layers, + io_backend=io_backend, + ) else: self._copy_tensor( self.temporal_buffer[layer_id], diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index fdac9d207..06cca265d 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -161,6 +161,10 @@ class MHATokenToKVPoolHost(HostKVCache): self.layer_num = self.target_layer_num + len(self.mtp_draft_device_pools) return self.head_dim * self.head_num * self.layer_num * self.dtype.itemsize * 2 + def get_hybrid_pool_buffer(self): + # Expose the K/V host tensors required for zero-copy I/O registration. + return [self.k_buffer, self.v_buffer] + def get_ksize_per_token(self): return self.get_size_per_token() // 2 diff --git a/python/sglang/test/ascend/test_ascend_utils.py b/python/sglang/test/ascend/test_ascend_utils.py index db19c7bca..ddbeb5b7b 100644 --- a/python/sglang/test/ascend/test_ascend_utils.py +++ b/python/sglang/test/ascend/test_ascend_utils.py @@ -116,6 +116,9 @@ KIMI_K2_5_W4A8_MODEL_PATH = os.path.join(MODEL_WEIGHTS_DIR, "Eco-Tech/Kimi-K2.5- KIMI_K2_5_EAGLE3_MODEL_PATH = os.path.join( MODEL_WEIGHTS_DIR, "lightseekorg/kimi-k2.5-eagle3" ) +KIMI_K3_W4A8_INT_MOE_WEIGHTS_PATH = os.path.join( + MODEL_WEIGHTS_DIR, "Kimi/Kimi-K3-w4a8-int-moe" +) LING_LITE_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "inclusionAI/Ling-lite") LLAMA_2_7B_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "LLM-Research/Llama-2-7B") LLAMA_3_1_8B_INSTRUCT_WEIGHTS_PATH = os.path.join( diff --git a/test/manual/hicache/test_npu_kimi_k3_hicache.py b/test/manual/hicache/test_npu_kimi_k3_hicache.py new file mode 100644 index 000000000..444a96e06 --- /dev/null +++ b/test/manual/hicache/test_npu_kimi_k3_hicache.py @@ -0,0 +1,53 @@ +import unittest + +from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin +from sglang.test.ascend.test_ascend_utils import KIMI_K3_W4A8_INT_MOE_WEIGHTS_PATH +from sglang.test.test_utils import CustomTestCase + + +class TestKimiK3MixedWithHiCacheL2(GSM8KAscendMixin, CustomTestCase): + """Testcase: Verify the inference accuracy of Kimi-K3 (MLA + KDA hybrid) on GSM8K + with mixed (non-PD) serving and HiCache L2 cache on NPU. + + [Test Category] HiCache + [Test Target] Kimi-K3 (MLA + KDA hybrid, mamba layers) + [Test Config] Mixed deployment, NPU, HiCache L2 (kernel_ascend IO backend) + """ + + model = KIMI_K3_W4A8_INT_MOE_WEIGHTS_PATH + accuracy = 0.9 + other_args = [ + "--trust-remote-code", + "--device", + "npu", + "--attention-backend", + "ascend", + "--quantization", + "modelslim", + "--dtype", + "bfloat16", + "--tp-size", + "64", + "--enable-dp-attention", + "--dp-size", + "4", + "--enable-dp-lm-head", + "--moe-a2a-backend", + "deepep", + "--deepep-mode", + "auto", + "--mem-fraction-static", + "0.75", + "--max-mamba-cache-size", + "240", + "--enable-hierarchical-cache", + "--hicache-io-backend", + "kernel_ascend", + "--enable-cache-report", + "--hicache-ratio", + "4.0", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index 2e8005c26..5dfdbc3ee 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -236,6 +236,10 @@ class TestDecodeQueueCleanup(CustomTestCase): queue._pre_alloc = MagicMock() queue.req_to_token_pool = MagicMock() queue.req_to_token_pool.available_size.return_value = 1 + # Non-hybrid pools have no mamba allocator; MagicMock would otherwise + # auto-create one and break the `available_size() <= 0` comparison in + # pop_preallocated. + queue.req_to_token_pool.mamba_allocator = None queue.req_to_metadata_buffer_idx_allocator = MagicMock() queue.req_to_metadata_buffer_idx_allocator.available_size.return_value = 1 diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 800729990..6bd4cea33 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -148,6 +148,10 @@ class TestDecodePreallocQueuePriority(unittest.TestCase): queue.req_to_token_pool = MagicMock() queue.req_to_token_pool.available_size.return_value = 100 + # Non-hybrid pools have no mamba allocator; MagicMock would otherwise + # auto-create one and break the `available_size() <= 0` comparison in + # pop_preallocated. + queue.req_to_token_pool.mamba_allocator = None queue.req_to_token_pool.req_to_token = torch.arange( 8 * 16, dtype=torch.int64 ).reshape(8, 16) diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 172d44187..af68e287b 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -442,6 +442,10 @@ class TestDecodeLockRefScenarios(CustomTestCase): queue.tree_cache.dec_lock_ref = MagicMock() queue.req_to_token_pool = MagicMock() queue.req_to_token_pool.available_size.return_value = 1 + # Non-hybrid pools have no mamba allocator; MagicMock would otherwise + # auto-create one and break the `available_size() <= 0` comparison in + # pop_preallocated. + queue.req_to_token_pool.mamba_allocator = None queue.req_to_metadata_buffer_idx_allocator = MagicMock() queue.req_to_metadata_buffer_idx_allocator.available_size.return_value = 1 queue.token_to_kv_pool = MagicMock() diff --git a/test/registered/unit/mem_cache/test_npu_mamba_async_layout.py b/test/registered/unit/mem_cache/test_npu_mamba_async_layout.py index 8bec03017..01b05bd96 100644 --- a/test/registered/unit/mem_cache/test_npu_mamba_async_layout.py +++ b/test/registered/unit/mem_cache/test_npu_mamba_async_layout.py @@ -156,6 +156,7 @@ class TestNPUMambaAsyncConfig(unittest.TestCase): MambaPoolHost.__dict__["_copy_tensor_all_layers_lf_pf"], staticmethod ) + @patch.object(mamba_pool_host, "transfer_mamba_state", None) def test_conv_only_load_skips_empty_temporal_component(self): pool = MambaPoolHost.__new__(MambaPoolHost) pool.layout = "page_first_direct"