Repository navigation
[None][feat] DSpark Unified KVCache standalone - #20664
Open
allisonlim-nv wants to merge 46 commits into
Open
allisonlim-nv wants to merge 46 commits into
allisonlim-nv wants to merge 46 commits into
Conversation
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Separate NVFP4 data and block scales from other equal-size buffers while keeping shared lifecycle slot allocation. Preserve matching role order and require matching scale roles. Add one regression test for mixed NVFP4 and BF16 storage, checking scale addresses at slot 1 and the final slot. Signed-off-by: Allison Lim <allim@nvidia.com>
Read the data-only view of three-dimensional pointer tables for non-FP4 layers while preserving NVFP4 scale handling and legacy two-dimensional tables. Add one native-binding regression covering mixed BF16/NVFP4 pointer tables and legacy BF16 tables. Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Preserve explicit and legacy layer configuration while applying upstream pool rebalancing settings to both paths. Signed-off-by: Allison Lim <allim@nvidia.com>
Share one KVCacheManagerV2 between compatible DFlash target and draft layers in aggregated and disaggregated serving. Preserve each layer's cache dtype and head dimension, support attention data parallelism, and retain the separate-cache fallback for unsupported layouts. Transfer unified draft layers with target pages, validate peer draft-layer counts, and add unit and B200 integration coverage. Signed-off-by: Allison Lim <allim@nvidia.com>
…lative_config in both sections In disaggregated serving, workers launched from a disaggregated config file share the draft KV cache only when both context_servers and generation_servers set speculative_config. Other workers keep the separate draft cache. Aggregated serving is unchanged. Signed-off-by: Allison Lim <allim@nvidia.com>
…parallel_size Workers launched from a disaggregated config file share the draft KV cache only when context_servers and generation_servers use the same tensor_parallel_size, in addition to both setting speculative_config. Signed-off-by: Allison Lim <allim@nvidia.com>
…y speculative_config and mismatched TP Drive the real disaggregated_mpi_worker against a config and check that the target and draft KV layers keep separate cache managers when only generation_servers sets speculative_config, and when the context and generation TP sizes differ. Signed-off-by: Allison Lim <allim@nvidia.com>
Preserve unified DFlash layer specifications alongside the standalone draft page-size override in KV-cache cost estimation. Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
1 task done
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
1 task done
Signed-off-by: Allison Lim <allim@nvidia.com>
Collaborator
|
@PerkzZheng, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~ Assignment explanation
Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once. |
Signed-off-by: Allison Lim <allim@nvidia.com>
BowenFu
reviewed
Oct 27, 2026
chienchunhung
approved these changes
Oct 27, 2026
chienchunhung
left a comment
chienchunhung
left a comment
Collaborator
There was a problem hiding this comment.
LGTM from disagg perspective after addressing the open comments.
Preserve layer offsets when using fused block-table decoding and keep unified caches on the preparation path that checks allocation bounds. Add CPU regression coverage for both integration points. Signed-off-by: Allison Lim <allim@nvidia.com>
Collaborator
|
@yuxianq, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~ @PerkzZheng requested an OOTO handoff of unstarted reviews ( Assignment explanation
Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once. |
This was referenced Oct 28, 2026
1 task
zheyuf
added a commit
to brb-nv/TensorRT-LLM
that referenced
this pull request
Oct 29, 2026
… under disaggregated serving (#25) * [None][feat] Unified kv cache manager foundation (NVIDIA#20178) Signed-off-by: Allison Lim <allim@nvidia.com> (cherry picked from commit fce28f0) Backport to this branch: the cpp/ part (BufferConfig.coalescing_group, the attentionOp 3-D pool-pointer table) is not included; _coalescing_group_kwargs passes the group only when it is nonzero, so layer specs without NVFP4 buffers run on the existing bindings, and the two mixed NVFP4/BF16 pool tests skip on bindings without coalescing_group. Hunks for main-only features (mixed-precision KV, pool rebalancing) are dropped, and layer_specs is declared on the class for managers built without __init__. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [TRTLLM-16722][feat] Dflash Unified KVCache (NVIDIA#20486) Port of NVIDIA#20486 at 219e26e (net diff against the main it merges, df95eda). Signed-off-by: Allison Lim <allim@nvidia.com> Backport to this branch: conflicts resolved against the separate draft pool's disagg window transfer (RankInfo keeps both draft fields), the drafter slot parking and seeds in DFlashSpecMetadata.prepare (the position seeding applies to the unified cache only), and the fused block-table refresh (skipped when the layer has a block offset). Hunks for code this branch does not have (spec_tail_kernels, model-level layer specs) are dropped; Gemma4 hybrid geometry stands in for _has_variable_kv_cache_geometry. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] Per View KV head transfer (NVIDIA#20709) Port of NVIDIA#20709 at 300a671 (net diff against the main it merges, 0746a23). Signed-off-by: Allison Lim <allim@nvidia.com> Backport to this branch: import and registration-check conflicts resolved by keeping both sides. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Bound DFlash's shared context by the allocation's extent DFlash took the number of valid block-table entries as the pages a request holds. On a unified cache with per-layer page tables (e.g. an FP8 target with a BF16 drafter), a sliding-window draft layer's evicted pages stay BAD_PAGE_INDEX in front of its live tail, so the count stopped at about one window (66 blocks for window 2048, page 32) whatever the context length. The decode-time allocation check then trapped once a context passed ~2.1K tokens, and prefill skipped the drafter's context store for later chunks and for prefix hits past the window. Use the extent of the allocation instead, the last valid entry + 1. Reads below the window are never issued and writes land in the live tail; for a full-attention drafter, and on the cpp backend, the extent equals the count. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit ff11364) * [None][fix] Re-seed DFlash generation slots in CUDA-graph warmup CUDA-graph warmup builds every batch from dummy requests whose ids repeat across batches (0..bs-1), so a reused id keeps its drafter slot and the previous batch's context length, which the replay after each capture advances for good. A later batch can give the same id far fewer pages, and on a unified KV cache the drafter's allocation check then trapped at server start: during the advanced-sampling capture pass the filler rows reached 31..79 positions against a 24-position allocation. Only warmup runs a CUDA-graph batch eagerly, so there re-seed each generation row's slot from this batch's positions, as prepare() already seeds a new request. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit 2b23449) * [None][feat] MiniMax-M3.1 DSpark: keep the drafter's KV in the target's manager under disagg On GB300 a disaggregated generation worker gets the target's prompt KV but not the DSpark drafter's, which lives in a separate pool: the drafter starts from an empty window at block 0 and acceptance collapses on long prompts. Hosting the draft layers in the target's KVCacheManagerV2 lets the transceiver move them with the target's pages. - MiniMaxM3KVCacheManagerV2 hosts explicit draft layer specs: index-K buffers follow the spec-built ones, dense target layers stay FP8, size-only coalescing keeps the target pools, draft layers use the base view. - KvCacheCreator unifies a standalone DSpark drafter with such a manager under disaggregated serving when the drafter reads the target's page size. Agg keeps the separate pool paired with block reuse, and a model that ships that pool's window (F0a) keeps it in disagg. TRTLLM_DISAGG_SHARED_DRAFT_KV=0 opts out. - An NVFP4 target transfers with a BF16 draft held in pools of its own. - A windowed draft lifecycle gets the hot-tier bytes of every admissible sequence's window plus one context chunk, instead of the bytes-per-token split that handed it a quarter of the KV cache. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] Reuse a unified DSpark drafter's KV only where its window is complete With the drafter's layers in the target's KV cache manager, the draft lifecycle's REQUIRED policy let the drafter's window constrain every prefix match: a request branching off a shared prefix found no committed window there and re-prefilled its whole prompt. - A windowed drafter's layer specs are OPTIONAL. The target alone sets the match, and a new cache's first resume selects the drafter's layers where their window is complete (page-aligned matches only). - Where the window is missing, prepare_context caps the claim so the prefill writes everything the drafter reads (from L + 2 - W on): at L + 1 - W, or one window before a branch point within W of the prompt end, whose window the request then publishes. A partial match first retries at its last full page. - The commit gate for unstored drafter context, its unwritten-prefix marks and branch-point publication cover a manager that hosts the drafter's layers, not only a paired draft pool. A disagg generation request needs no cap: the KV transfer requests an unselected OPTIONAL group in full. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][perf] MiniMax-M3.1 DSpark: copy only the unified drafter's block offsets per step A hybrid NVFP4 MiniMax-M3 cache is pool-less: no target layer reads its attention-op block offsets. With explicit draft layers it kept every layer's pool, so each step converted the whole per-layer table (layers x sequence slots x max_seq_len / tokens_per_block) on the device: about 3 ms per decode step at max_seq_len 1M on TP2. Keep one attention-op pool, the first draft layer's, which is the only one the unified drafter reads, and fill it from the draft lifecycle's rows with the copy a separate draft pool uses. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Allocate the whole block of a drafter that shares the target's KV cache A unified DFlash/DSpark drafter appends its block after the request's context, through the request's own pages. The manager sized every request with num_extra_kv_tokens = K-1 beyond its history, but the DSpark shift_label block is K wide (DFlash K+1). When the allocation ended on a page boundary, at the prompt end or after a step that accepts all K+1 tokens, the block's last slot had no page: the drafter then truncated the advertised context and its scratch overwrote the last context position in the draft layers, silently. When DRAFT-owner layers live in the manager, size each request beyond its history by max(num_extra_kv_tokens, drafter block width) at every allocation site and in the generation headroom. The creator passes the checkpoint-resolved width (shift_label). Adapted from NVIDIA/TensorRT-LLM Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> NVIDIA#20664 (76c4cef); num_extra_kv_tokens keeps its position-side meaning. * [None][fix] Fail at startup when TRTLLM_DISAGG_SHARED_DRAFT_KV=1 cannot be honored An explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=1 on a disaggregated worker whose target and drafter cannot share one KV cache manager silently kept the separate draft pool, whose KV is not transferred. Raise at KV cache creation instead, as NVIDIA/TensorRT-LLM NVIDIA#20664's _check_shared_draft_kv_request does. disaggregated_mpi_worker now writes "auto" rather than "1" when both sections pass its coarse checks, so its automatic choice can still fall back where a worker's finer checks refuse; the creator shares under "1" and "auto" alike. The unset default is unchanged. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Let disaggregated_mpi_worker share a DSpark drafter's KV The launcher wrote TRTLLM_DISAGG_SHARED_DRAFT_KV=0 for any drafter other than DFlash, and for any page size other than 16, 32 or 64, so a MiniMax-M3.1 DSpark deployment started through it (128-token pages) kept the separate draft pool, whose KV is not transferred. Accept a standalone DSpark drafter with attention_backend unset, AUTO or TRTLLM, as NVIDIA/TensorRT-LLM NVIDIA#20664 does, and leave its page size to each worker's KvCacheCreator: the launcher cannot know whether the drafter reads the target's page. It writes "auto"; if the two sides still disagree, the peers' draft-layer check fails loudly. An explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=0 is now kept. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Keep attention-DP dummy generation rows on DFlash's dummy slot A generation row with the attention-DP dummy id must never inherit a real slot, because the C++ offset copy maps placeholders to page 0: a slot that a context warmup stored under that id would let the dummy row write its drafter K/V past its own pages, onto another request's page. As in NVIDIA/TensorRT-LLM NVIDIA#20664, such rows now always take the dummy slot. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][perf] DFlash: size a unified MiniMax-M3.1 drafter's context stores on the host DFlash bounds each context store by the row's allocation, which it reads back from the device unless the drafter's block-table rows are known to be written in full. That was only recognized for a separate draft pool, so with the unified drafter every context chunk waited for the target's forward before the store was issued. Let the KV cache manager say whether it rewrites copied rows in full, and check the target manager's own table for a unified drafter. The MiniMax-M3 manager's drafter-only pool is copied in full. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] KVCacheManagerV2: reuse an OPTIONAL group at a partial endpoint An OPTIONAL lifecycle group was checked only at a page-aligned endpoint, so a partial prefix match always left it incomplete. A unified DSpark drafter then lost its window's partial page and the caller cut its claim back to the last full page, recomputing up to a page per request. Check an OPTIONAL group at any endpoint with the REQUIRED coverage rule: the last page needs only the matched prefix. The first resume copies a selected group's partial page into a private page, as for REQUIRED, and counts it as a partial reuse. Partial snapshots now publish OPTIONAL pages as well, so a continuation of a prompt that ended mid-page finds them. A match whose partial page is missing still falls back to the last full page. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Give the block-table refresh test's stub the extent columns The stock DFlashWorker._refresh_ctx_block_tables reads _ctx_block_cols since the allocation-extent fix; the test's object.__new__ stub never set it. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> --------- Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> Co-authored-by: allisonlim-nv <allim@nvidia.com>
Collaborator
|
[by Claude Code] @eopXD Could you review this PR? Thanks! |
1 task done
Preserve unified target/draft cache behavior alongside upstream page-table alignment and draft mapping updates. Adopt ReusePolicy and the updated cache creator signature, and retain both branches’ regression coverage. Signed-off-by: Allison Lim <allim@nvidia.com>
brb-nv
pushed a commit
to brb-nv/TensorRT-LLM
that referenced
this pull request
Oct 30, 2026
… under disaggregated serving (#25) * [None][feat] Unified kv cache manager foundation (NVIDIA#20178) Signed-off-by: Allison Lim <allim@nvidia.com> (cherry picked from commit fce28f0) Backport to this branch: the cpp/ part (BufferConfig.coalescing_group, the attentionOp 3-D pool-pointer table) is not included; _coalescing_group_kwargs passes the group only when it is nonzero, so layer specs without NVFP4 buffers run on the existing bindings, and the two mixed NVFP4/BF16 pool tests skip on bindings without coalescing_group. Hunks for main-only features (mixed-precision KV, pool rebalancing) are dropped, and layer_specs is declared on the class for managers built without __init__. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [TRTLLM-16722][feat] Dflash Unified KVCache (NVIDIA#20486) Port of NVIDIA#20486 at 219e26e (net diff against the main it merges, df95eda). Signed-off-by: Allison Lim <allim@nvidia.com> Backport to this branch: conflicts resolved against the separate draft pool's disagg window transfer (RankInfo keeps both draft fields), the drafter slot parking and seeds in DFlashSpecMetadata.prepare (the position seeding applies to the unified cache only), and the fused block-table refresh (skipped when the layer has a block offset). Hunks for code this branch does not have (spec_tail_kernels, model-level layer specs) are dropped; Gemma4 hybrid geometry stands in for _has_variable_kv_cache_geometry. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] Per View KV head transfer (NVIDIA#20709) Port of NVIDIA#20709 at 300a671 (net diff against the main it merges, 0746a23). Signed-off-by: Allison Lim <allim@nvidia.com> Backport to this branch: import and registration-check conflicts resolved by keeping both sides. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Bound DFlash's shared context by the allocation's extent DFlash took the number of valid block-table entries as the pages a request holds. On a unified cache with per-layer page tables (e.g. an FP8 target with a BF16 drafter), a sliding-window draft layer's evicted pages stay BAD_PAGE_INDEX in front of its live tail, so the count stopped at about one window (66 blocks for window 2048, page 32) whatever the context length. The decode-time allocation check then trapped once a context passed ~2.1K tokens, and prefill skipped the drafter's context store for later chunks and for prefix hits past the window. Use the extent of the allocation instead, the last valid entry + 1. Reads below the window are never issued and writes land in the live tail; for a full-attention drafter, and on the cpp backend, the extent equals the count. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit ff11364) * [None][fix] Re-seed DFlash generation slots in CUDA-graph warmup CUDA-graph warmup builds every batch from dummy requests whose ids repeat across batches (0..bs-1), so a reused id keeps its drafter slot and the previous batch's context length, which the replay after each capture advances for good. A later batch can give the same id far fewer pages, and on a unified KV cache the drafter's allocation check then trapped at server start: during the advanced-sampling capture pass the filler rows reached 31..79 positions against a 24-position allocation. Only warmup runs a CUDA-graph batch eagerly, so there re-seed each generation row's slot from this batch's positions, as prepare() already seeds a new request. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> (cherry picked from commit 2b23449) * [None][feat] MiniMax-M3.1 DSpark: keep the drafter's KV in the target's manager under disagg On GB300 a disaggregated generation worker gets the target's prompt KV but not the DSpark drafter's, which lives in a separate pool: the drafter starts from an empty window at block 0 and acceptance collapses on long prompts. Hosting the draft layers in the target's KVCacheManagerV2 lets the transceiver move them with the target's pages. - MiniMaxM3KVCacheManagerV2 hosts explicit draft layer specs: index-K buffers follow the spec-built ones, dense target layers stay FP8, size-only coalescing keeps the target pools, draft layers use the base view. - KvCacheCreator unifies a standalone DSpark drafter with such a manager under disaggregated serving when the drafter reads the target's page size. Agg keeps the separate pool paired with block reuse, and a model that ships that pool's window (F0a) keeps it in disagg. TRTLLM_DISAGG_SHARED_DRAFT_KV=0 opts out. - An NVFP4 target transfers with a BF16 draft held in pools of its own. - A windowed draft lifecycle gets the hot-tier bytes of every admissible sequence's window plus one context chunk, instead of the bytes-per-token split that handed it a quarter of the KV cache. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] Reuse a unified DSpark drafter's KV only where its window is complete With the drafter's layers in the target's KV cache manager, the draft lifecycle's REQUIRED policy let the drafter's window constrain every prefix match: a request branching off a shared prefix found no committed window there and re-prefilled its whole prompt. - A windowed drafter's layer specs are OPTIONAL. The target alone sets the match, and a new cache's first resume selects the drafter's layers where their window is complete (page-aligned matches only). - Where the window is missing, prepare_context caps the claim so the prefill writes everything the drafter reads (from L + 2 - W on): at L + 1 - W, or one window before a branch point within W of the prompt end, whose window the request then publishes. A partial match first retries at its last full page. - The commit gate for unstored drafter context, its unwritten-prefix marks and branch-point publication cover a manager that hosts the drafter's layers, not only a paired draft pool. A disagg generation request needs no cap: the KV transfer requests an unselected OPTIONAL group in full. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][perf] MiniMax-M3.1 DSpark: copy only the unified drafter's block offsets per step A hybrid NVFP4 MiniMax-M3 cache is pool-less: no target layer reads its attention-op block offsets. With explicit draft layers it kept every layer's pool, so each step converted the whole per-layer table (layers x sequence slots x max_seq_len / tokens_per_block) on the device: about 3 ms per decode step at max_seq_len 1M on TP2. Keep one attention-op pool, the first draft layer's, which is the only one the unified drafter reads, and fill it from the draft lifecycle's rows with the copy a separate draft pool uses. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Allocate the whole block of a drafter that shares the target's KV cache A unified DFlash/DSpark drafter appends its block after the request's context, through the request's own pages. The manager sized every request with num_extra_kv_tokens = K-1 beyond its history, but the DSpark shift_label block is K wide (DFlash K+1). When the allocation ended on a page boundary, at the prompt end or after a step that accepts all K+1 tokens, the block's last slot had no page: the drafter then truncated the advertised context and its scratch overwrote the last context position in the draft layers, silently. When DRAFT-owner layers live in the manager, size each request beyond its history by max(num_extra_kv_tokens, drafter block width) at every allocation site and in the generation headroom. The creator passes the checkpoint-resolved width (shift_label). Adapted from NVIDIA/TensorRT-LLM Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> NVIDIA#20664 (76c4cef); num_extra_kv_tokens keeps its position-side meaning. * [None][fix] Fail at startup when TRTLLM_DISAGG_SHARED_DRAFT_KV=1 cannot be honored An explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=1 on a disaggregated worker whose target and drafter cannot share one KV cache manager silently kept the separate draft pool, whose KV is not transferred. Raise at KV cache creation instead, as NVIDIA/TensorRT-LLM NVIDIA#20664's _check_shared_draft_kv_request does. disaggregated_mpi_worker now writes "auto" rather than "1" when both sections pass its coarse checks, so its automatic choice can still fall back where a worker's finer checks refuse; the creator shares under "1" and "auto" alike. The unset default is unchanged. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Let disaggregated_mpi_worker share a DSpark drafter's KV The launcher wrote TRTLLM_DISAGG_SHARED_DRAFT_KV=0 for any drafter other than DFlash, and for any page size other than 16, 32 or 64, so a MiniMax-M3.1 DSpark deployment started through it (128-token pages) kept the separate draft pool, whose KV is not transferred. Accept a standalone DSpark drafter with attention_backend unset, AUTO or TRTLLM, as NVIDIA/TensorRT-LLM NVIDIA#20664 does, and leave its page size to each worker's KvCacheCreator: the launcher cannot know whether the drafter reads the target's page. It writes "auto"; if the two sides still disagree, the peers' draft-layer check fails loudly. An explicit TRTLLM_DISAGG_SHARED_DRAFT_KV=0 is now kept. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Keep attention-DP dummy generation rows on DFlash's dummy slot A generation row with the attention-DP dummy id must never inherit a real slot, because the C++ offset copy maps placeholders to page 0: a slot that a context warmup stored under that id would let the dummy row write its drafter K/V past its own pages, onto another request's page. As in NVIDIA/TensorRT-LLM NVIDIA#20664, such rows now always take the dummy slot. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][perf] DFlash: size a unified MiniMax-M3.1 drafter's context stores on the host DFlash bounds each context store by the row's allocation, which it reads back from the device unless the drafter's block-table rows are known to be written in full. That was only recognized for a separate draft pool, so with the unified drafter every context chunk waited for the target's forward before the store was issued. Let the KV cache manager say whether it rewrites copied rows in full, and check the target manager's own table for a unified drafter. The MiniMax-M3 manager's drafter-only pool is copied in full. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][feat] KVCacheManagerV2: reuse an OPTIONAL group at a partial endpoint An OPTIONAL lifecycle group was checked only at a page-aligned endpoint, so a partial prefix match always left it incomplete. A unified DSpark drafter then lost its window's partial page and the caller cut its claim back to the last full page, recomputing up to a page per request. Check an OPTIONAL group at any endpoint with the REQUIRED coverage rule: the last page needs only the matched prefix. The first resume copies a selected group's partial page into a private page, as for REQUIRED, and counts it as a partial reuse. Partial snapshots now publish OPTIONAL pages as well, so a continuation of a prompt that ended mid-page finds them. A match whose partial page is missing still falls back to the last full page. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> * [None][fix] Give the block-table refresh test's stub the extent columns The stock DFlashWorker._refresh_ctx_block_tables reads _ctx_block_cols since the allocation-extent fix; the test's object.__new__ stub never set it. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> --------- Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> Co-authored-by: allisonlim-nv <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
1 task done
This branch has not been deployed
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.
Dev Engineer Review
QA Engineer Review
TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kvand a B200 acceptance-length reference.tests/integration/test_lists/test-db/l0_dgx_b200.yml. No manual-QA list change is reported.Per-File QA Perspective
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h: AddscoalescingGroup, defaulting to0. Verify that existing configurations retain compatible coalescing behavior.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp: Groups buffers by lifecycle, coalescing group, and expanded size. Verify that buffers with different groups do not share storage.cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp: Adds thecoalescing_groupconstructor argument and writable property. Verify Python constructor compatibility and property updates.cpp/tensorrt_llm/thop/attentionOp.cpp: Accepts 2D and 3D non-FP4 pointer tables. Verify that mixed-format tables select the correct data-pointer plane.docs/source/features/speculative-decoding.md: Documents shared DFlash draft KV behavior, transfer conditions, layout requirements, and separate-manager cases.tensorrt_llm/_torch/attention/backends/flashinfer.py: Separates page-index spaces using pool mapping, page-index scale, and KV factor. Verify indices for layers with different KV geometry.tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py: Acceptskv_factorbut documents that it does not affect raw slot IDs. Verify callers remain compatible.tensorrt_llm/_torch/disaggregation/native/peer.py: Rejects peers with incompatible draft-layer counts and, for shared draft KV, incompatible TP or attention-DP settings.tensorrt_llm/_torch/disaggregation/native/rank_info.py: Adds draft-layer counts to rank metadata. Verify serialization and compatibility with metadata that lacks the field.tensorrt_llm/_torch/pyexecutor/_util.py: Selects unified or separate managers and forwards explicit layer specifications. Verify sharing gates and sizing for supported and unsupported configurations.tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py: Adds per-layer cache specifications and updates sizing, allocation, indexing, and views. Verify mixed layouts, slot rounding, and the legacy path without explicit specifications.tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py: Changes a comment about disaggregated generation. The supplied change summary reports no runtime behavior change.tensorrt_llm/_torch/speculative/__init__.py: Exportsget_draft_block_width. Verify the package-level import.tensorrt_llm/_torch/speculative/dflash.py: Adds unified-cache binding, layout checks, request-slot handling, and allocation-bound checks. Verify separate-manager fallback and unified-cache error paths.tensorrt_llm/_torch/speculative/dspark.py: Adds DSpark block-width selection for shift-label checkpoints. Verify the legacy-layout fallback.tensorrt_llm/_torch/speculative/interface.py: Skips replay swapping when target and draft managers are the same object. Verify separate-manager behavior remains unchanged.tensorrt_llm/_torch/speculative/utils.py: Adds draft block-width and manager-selection helpers. Verify behavior for missing configs, embedded DSpark, and shared managers.tensorrt_llm/commands/serve.py: Validates disaggregated shared-cache settings and controlsTRTLLM_DISAGG_SHARED_DRAFT_KV. Verify startup errors for unsupported configurations.tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi: Adds thecoalescing_grouptype declaration. Verify it matches the native binding.tests/integration/defs/accuracy/references/acceptance_length.yaml: Adds a B200 acceptance-length reference for the unified DFlash test.tests/integration/defs/accuracy/test_disaggregated_serving.py: Adds a Qwen3-8B DFlash disaggregated accuracy test and acceptance-length check. The test is listed in the B200 CI test database.tests/integration/test_lists/test-db/l0_dgx_b200.yml: Adds the unified DFlash integration test to the B200 pre-merge suite.tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py: Initializeslayer_specsfor cache-cost calculations. Verify the test exercises the intended legacy path.tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py: Initializeslayer_specsin block-scale role tests. Verify role selection remains covered.tests/unittest/_torch/executor/kv_cache/test_dflash_unified_kv_gate.py: Covers unified-layout eligibility and disaggregated launcher validation, including geometry, dtype, windows, backend, and parallelism cases.tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py: Adds explanatory comments only; no executable test behavior change is reported.tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py: Covers cache sizing for sink tokens, reuse policies, draft rounding, storage-pool scaling, and model-provided layer specifications.tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py: Adds coverage for layer specifications, heterogeneous layouts and dtypes, pool sizing, block-table indexing, shared-page allocation, extra-token capacity, and coalesced-pool strides.tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py: Initializeslayer_specsin hybrid cache-manager tests. Verify quota, pool-sizing, and allocation paths.tests/unittest/_torch/speculative/hw_agnostic/test_dflash_ctx_block_tables.py: Covers DFlash block-table fallback, unified-layout errors, and allocation bounds.tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py: Covers scaled batch indices and raw, default, and truncated cache-index forms.tests/unittest/disaggregated/test_peer.py: Covers peer registration with matching and mismatched draft-layer counts.tests/unittest/disaggregated/test_rank_info.py: Covers draft-layer reporting, serialization, and the default for missing metadata.tests/unittest/_torch/attention/test_fp4_mla.py: Initializeslayer_specsbefore runtime quota sizing with pipeline slots.Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.