minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Bug Fix] missing index/KV transfer for MTP layer in NSA disaggregation (#23539)

Browse Source
This commit is contained in:
Yuxuan Zhang
2026-04-30 11:55:45 +08:00
committed by GitHub
parent 9e68c8527e
commit d040333c95
2 changed files with 24 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+12
View File
@@ -377,6 +377,18 @@ class DecodePreallocQueue:
) )
elif isinstance(self.token_to_kv_pool, NSATokenToKVPool): elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
kv_args.state_type = "nsa" kv_args.state_type = "nsa"
if self.draft_token_to_kv_pool is not None and isinstance(
self.draft_token_to_kv_pool, NSATokenToKVPool
):
(
draft_state_data_ptrs,
draft_state_data_lens,
draft_state_item_lens,
) = self.draft_token_to_kv_pool.get_state_buf_infos()
kv_args.state_data_ptrs += draft_state_data_ptrs
kv_args.state_data_lens += draft_state_data_lens
kv_args.state_item_lens += draft_state_item_lens
else: else:
kv_args.state_type = "none" kv_args.state_type = "none"
else: else:
python/sglang/srt/disaggregation/prefill.py
+12
View File
@@ -192,6 +192,18 @@ class PrefillBootstrapQueue:
) )
elif isinstance(self.token_to_kv_pool, NSATokenToKVPool): elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
kv_args.state_type = "nsa" kv_args.state_type = "nsa"
if self.draft_token_to_kv_pool is not None and isinstance(
self.draft_token_to_kv_pool, NSATokenToKVPool
):
(
draft_state_data_ptrs,
draft_state_data_lens,
draft_state_item_lens,
) = self.draft_token_to_kv_pool.get_state_buf_infos()
kv_args.state_data_ptrs += draft_state_data_ptrs
kv_args.state_data_lens += draft_state_data_lens
kv_args.state_item_lens += draft_state_item_lens
else: else:
kv_args.state_type = "none" kv_args.state_type = "none"
else: else: