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

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