Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ShangmingCai merged 8 commits into from
Apr 30, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
5f9eb42
Fix missing index/KV transfer for MTP layer in NSA disaggregation
zRzRzRzRzRzRzR Apr 23, 2026
59a8c5f
Merge branch 'main' into glm-pd
zRzRzRzRzRzRzR Apr 23, 2026
ddbe2c0
Merge branch 'main' into glm-pd
ShangmingCai Apr 23, 2026
35b40f6
Merge branch 'main' into glm-pd
ShangmingCai Apr 24, 2026
fcfee1c
Merge branch 'main' into glm-pd
ShangmingCai Apr 24, 2026
f1323e5
Merge branch 'main' into glm-pd
ShangmingCai Apr 27, 2026
5ae5fec
Merge branch 'main' into glm-pd
zRzRzRzRzRzRzR Apr 28, 2026
f3a4c01
Merge branch 'sgl-project:main' into glm-pd
zRzRzRzRzRzRzR Apr 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions python/sglang/srt/disaggregation/decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -369,6 +369,18 @@ def _init_kv_manager(self) -> CommonKVManager:
)
elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
kv_args.state_type = "nsa"
if self.draft_token_to_kv_pool is not None and isinstance(
self.draft_token_to_kv_pool, NSATokenToKVPool
):
Comment on lines +372 to +374

@kpham-sgl kpham-sgl Apr 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two small questions here

  1. Do we need draft_token_to_kv_pool to also be NSATokenToKVPool?
  2. Need to check if hasattr(self.draft_token_to_kv_pool,"get_state_buf_infos")

@ShangmingCai ShangmingCai Apr 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make sense. Do we need to consider the Non-MTP spec decode cases? Not sure if this is a common use-case. @zRzRzRzRzRzRzR

@zRzRzRzRzRzRzR zRzRzRzRzRzRzR Apr 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The isinstance(self.draft_token_to_kv_pool, NSATokenToKVPool) guard makes this PR a no-op for non-NSA draft pools, so non-MTP spec decode cases aren't affected. The fix kicks in only when the draft pool is also NSA, which today is the MTP-on-NSA path. If a future non-MTP spec decode also runs an NSA draft, the same logic applies and would Just Work.

(
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:
kv_args.state_type = "none"
else:
Expand Down
12 changes: 12 additions & 0 deletions python/sglang/srt/disaggregation/prefill.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,18 @@ def _init_kv_manager(self) -> CommonKVManager:
)
elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
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:
kv_args.state_type = "none"
else:
Expand Down
Loading