From dbd7281b065ecce503c92441abdf7855f727ca9b Mon Sep 17 00:00:00 2001 From: Jinyan Yi Date: Wed, 16 Sep 2026 09:10:17 +0800 Subject: [PATCH] fix(npu): fix hybrid KV transfer with PP prefill in PD disaggregation (#38402) --- .../sglang/srt/disaggregation/ascend/conn.py | 16 +++++ .../srt/disaggregation/mooncake/conn.py | 9 +++ python/sglang/srt/mem_cache/memory_pool.py | 3 + python/sglang/srt/models/kimi_k3.py | 1 - .../test_pp_hybrid_kv_transfer.py | 72 +++++++++++++++++++ .../unit/mem_cache/test_mamba_unittest.py | 22 ++++++ 6 files changed, 122 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py index 474a30fce..9c28fd871 100644 --- a/python/sglang/srt/disaggregation/ascend/conn.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 78c572369..d6bc4aed9 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 28333dffb..450b22e6f 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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): diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index b0ede48c8..1fbaf6f92 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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 = {} diff --git a/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py b/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py index a065e1cdb..0f8fd8d19 100644 --- a/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index 54b8dd21e..410fc3eea 100755 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/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