[NPU] Adapt hicache for K3 hybrid models (#39415)

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
 

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()

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,

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],

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
 

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(