Repository navigation
[None][fix] Use 128-token DSpark draft pages only where the dynamic-page kernels are verified (SM100/SM103) - #23
Closed
zheyuf wants to merge 1 commit into
Conversation
…age kernels are verified (SM100/SM103) dflash_trtllm_supports_page_size allowed power-of-two pages from 128 tokens wherever the DFlash TRTLLM ops run (SM100-SM109). On SM107 the DSpark draft pool therefore took the target's 128-token pages and paired with it for joint block reuse, which the disaggregated draft-window transfer refuses at executor creation; flashinfer's dynamic page-size kernels are also verified on SM100/SM103 only. Allow the dynamic sizes there only; elsewhere the pool keeps the static sizes (64 tokens, unpaired) as before. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Collaborator
Author
|
Superseded by #24. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
#22 lets the standalone DSpark drafter's draft pool take the target's 128-token pages when the DFlash TRTLLM ops serve them (
dflash_trtllm_supports_page_size). That rule did not depend on the GPU. The DFlash TRTLLM ops run on SM100-SM109 (is_sm_100f,tensorrt_llm/_utils.py:742-745), so on Rubin (SM107) the draft pool also became 128-token and paired with the target for joint block reuse (the default whenever block reuse is on).The SM107 MiniMax-M3.1 targets opt into the disaggregated draft-window transfer (
RuntimeHooks.disagg_draft_kv=True), which refuses a paired draft pool:tensorrt_llm/_torch/disaggregation/draft_kv.py:86-91:select_disagg_draft_kv_cache_managerraisesValueError;tensorrt_llm/_torch/disaggregation/transceiver.py:308-312: the transceiver raises the same.The selector runs at executor creation on every disaggregated worker (
tensorrt_llm/_torch/pyexecutor/_util.py:4609-4615). After #22, SM107 disaggregated serving with block reuse therefore fails to start. Separately, flashinfer's dynamic page-size trtllm-gen kernels, which serve the 128-token pages, were verified on SM100/SM103 only.This PR allows the dynamic page sizes (power-of-two pages from 128 tokens) only on SM100 and SM103. Elsewhere the draft pool keeps the static sizes, as before #22 (64 tokens, unpaired). SM107 returns exactly to its pre-#22 behavior; GB300 (SM103) does not change.
Test Coverage
test_trtllm_page_size_capability: with the SM version patched, 128-token pages without the context kernel are allowed on SM103 and refused on SM107.TestDraftPoolPageSize: 16 passed.🤖 Generated with Claude Code