fix(npu): fix hybrid KV transfer with PP prefill in PD disaggregation (#38402)

python/sglang/srt/disaggregation/ascend/conn.py

@@ -174,6 +174,22 @@ class AscendKVManager(MooncakeKVManager):
         self._validate_envelope_kv_layout(
             dst_kv_ptrs, dst_kv_item_len, dst_attn_tp_size
         )
+        # Hybrid MLA prefill stages expose PP-local entries, while a PP=1
+        # decode peer registers all model layers. Pair only this layout by
+        # global layer id; every other Ascend layout keeps the legacy path.
+        if self.is_hybrid_mla_backend and self.pp_size > 1:
+            return self._send_kvcache_generic(
+                mooncake_session_id=mooncake_session_id,
+                src_data_ptrs=self.kv_args.kv_data_ptrs,
+                dst_data_ptrs=dst_kv_ptrs,
+                item_lens=self.kv_args.kv_item_lens,
+                prefill_data_indices=prefill_kv_indices,
+                dst_data_indices=dst_kv_indices,
+                executor=executor,
+                src_layer_ids=self.kv_args.kv_layer_ids,
+                dst_layer_ids=dst_layer_ids,
+            )
+
         # Group by indices
         prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
             prefill_kv_indices, dst_kv_indices

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -1501,6 +1501,7 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                             dst_indices,
                             src_state_layer_ids,
                             dst_state_layer_ids,
+                            dst_item_lens,
                         )
                         or rc
                     )
@@ -1707,6 +1708,7 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
         dst_mamba_index: list,
         src_layer_ids: Optional[List[int]] = None,
         dst_layer_ids: Optional[List[int]] = None,
+        dst_state_item_lens: Optional[list[int]] = None,
     ):
         assert len(prefill_mamba_index) == 1, "Mamba should have single state index"
 
@@ -1721,6 +1723,12 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
         for i, j in pairs:
             dst_state_ptr = dst_state_data_ptrs[j]
             length = src_state_item_lens[i]
+            if dst_state_item_lens and length != dst_state_item_lens[j]:
+                raise RuntimeError(
+                    "Prefill/Decode Mamba slot size mismatch "
+                    f"(src={length}, dst={dst_state_item_lens[j]}). "
+                    "Configure matching persistent state layouts on both peers."
+                )
             src_addr = src_state_data_ptrs[i] + length * int(prefill_mamba_index[0])
             dst_addr = dst_state_ptr + length * int(dst_mamba_index[0])
             transfer_blocks.append((src_addr, dst_addr, length))
@@ -1776,6 +1784,7 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                 dst_mamba_index,
                 src_layer_ids,
                 dst_layer_ids,
+                dst_state_item_lens,
             )
 
         local_tp_rank_in_group = self.kv_args.engine_rank % self.attn_tp_size

python/sglang/srt/mem_cache/memory_pool.py

@@ -3981,6 +3981,9 @@ class HybridLinearKVPool(KVCache):
     def get_kv_layer_ids(self):
         """Global layer ids aligned with the full-attention KV buffers."""
         layer_ids = list(self.full_attention_layer_id_mapping)
+        if self.use_mla and _is_npu and layer_ids:
+            data_ptrs, _, _ = self.get_contiguous_buf_infos()
+            return layer_ids * (len(data_ptrs) // len(layer_ids))
         return layer_ids if self.use_mla else layer_ids * 2
 
     def get_state_buf_infos(self):

python/sglang/srt/models/kimi_k3.py

@@ -575,7 +575,6 @@ class KimiK3MoE(nn.Module):
         self._shared_experts_attn_tp_comm = (
             get_parallel().enable_shared_experts_attn_tp
             and self._ep_a2a
-            and self._dp_attention
             and get_parallel().attn_tp_size > 1
         )
         shared_experts_tp_kwargs = {}

test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py

@@ -6,6 +6,7 @@ from types import SimpleNamespace
 
 import numpy as np
 
+from sglang.srt.disaggregation.ascend.conn import AscendKVManager
 from sglang.srt.disaggregation.common.conn import CommonKVManager
 from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
 from sglang.srt.disaggregation.prefill import _transfer_start_layer
@@ -142,6 +143,77 @@ class TestHybridSendUsesLayerIdPairing(CustomTestCase):
         self._run_case(model_full_ids=ids, stage_full_ids=ids[:5], start_offset=0)
 
 
+class _RecordingAscendManager:
+    def __init__(self):
+        self.is_hybrid_mla_backend = True
+        self.pp_size = 2
+        self.kv_args = SimpleNamespace(
+            kv_data_ptrs=[101, 102, 103, 104],
+            kv_item_lens=[11, 12, 13, 14],
+            kv_layer_ids=[31, 35, 31, 35],
+        )
+        self.generic_call = None
+
+    def _validate_envelope_kv_layout(self, *args):
+        pass
+
+    def _send_kvcache_generic(self, **kwargs):
+        self.generic_call = kwargs
+        return 7
+
+
+class TestAscendHybridPpDispatch(CustomTestCase):
+    def test_hybrid_pp_uses_layer_id_pairing_path(self):
+        manager = _RecordingAscendManager()
+        dst_layer_ids = [3, 7, 31, 35, 3, 7, 31, 35]
+        rc = AscendKVManager.send_kvcache(
+            manager,
+            mooncake_session_id="session",
+            prefill_kv_indices=np.array([1], dtype=np.int32),
+            dst_kv_ptrs=list(range(8)),
+            dst_kv_indices=np.array([2], dtype=np.int32),
+            executor=None,
+            dst_layer_ids=dst_layer_ids,
+        )
+
+        self.assertEqual(rc, 7)
+        self.assertEqual(manager.generic_call["src_layer_ids"], [31, 35, 31, 35])
+        self.assertEqual(manager.generic_call["dst_layer_ids"], dst_layer_ids)
+
+
+class TestMambaSlotTransfer(CustomTestCase):
+    def test_stage1_uses_paired_slot_sizes_and_offsets(self):
+        manager = _RecordingKVManager(prefill_start_layer=1, pp_size=2)
+        MooncakeKVManager._send_mamba_state(
+            manager,
+            req=SimpleNamespace(mooncake_session_id="session"),
+            prefill_mamba_index=[2],
+            src_state_data_ptrs=[1000],
+            src_state_item_lens=[16],
+            dst_state_data_ptrs=[2000, 3000],
+            dst_mamba_index=[3],
+            src_layer_ids=[7],
+            dst_layer_ids=[3, 7],
+            dst_state_item_lens=[8, 16],
+        )
+        self.assertEqual(manager.blocks, [(1032, 3048, 16)])
+
+    def test_slot_size_mismatch_rejected_before_transfer(self):
+        manager = _RecordingKVManager(prefill_start_layer=0, pp_size=1)
+        with self.assertRaisesRegex(RuntimeError, "Mamba slot size mismatch"):
+            MooncakeKVManager._send_mamba_state(
+                manager,
+                req=SimpleNamespace(mooncake_session_id="session"),
+                prefill_mamba_index=[1],
+                src_state_data_ptrs=[1000],
+                src_state_item_lens=[16],
+                dst_state_data_ptrs=[2000],
+                dst_mamba_index=[1],
+                dst_state_item_lens=[32],
+            )
+        self.assertEqual(manager.blocks, [])
+
+
 class TestGetMhaKvPtrsWithPp(CustomTestCase):
     """Derived property: the modulo heuristic in get_mha_kv_ptrs_with_pp exists
     for the decode-has-draft-KV layout [K_main, V_main, draft_K, draft_V]. Pin

test/registered/unit/mem_cache/test_mamba_unittest.py

@@ -1,5 +1,7 @@
 import unittest
 from array import array
+from types import SimpleNamespace
+from unittest.mock import patch
 
 import torch
 
@@ -80,6 +82,26 @@ class TestMamba(unittest.TestCase):
             "layer_id=1 not in full attention layers:", str(context.exception)
         )
 
+    def test_hybrid_linear_kv_pool_npu_layer_ids_match_buffer_groups(self):
+        pool = object.__new__(HybridLinearKVPool)
+        pool.full_attention_layer_id_mapping = {3: 0, 7: 1}
+        pool.use_mla = True
+
+        with patch("sglang.srt.mem_cache.memory_pool._is_npu", False):
+            self.assertEqual(pool.get_kv_layer_ids(), [3, 7])
+
+        with patch("sglang.srt.mem_cache.memory_pool._is_npu", True):
+            for group_count in (2, 3):
+                with self.subTest(group_count=group_count):
+                    pool.full_kv_pool = SimpleNamespace(
+                        get_contiguous_buf_infos=lambda: (
+                            list(range(2 * group_count)),
+                            [],
+                            [],
+                        )
+                    )
+                    self.assertEqual(pool.get_kv_layer_ids(), [3, 7] * group_count)
+
     def test_mamba_pool(self):
         max_num_reqs = 10
         mamba_cache_size = 20