Repository navigation
[None][feat] Unified target/draft KV cache for DFlash (port of #20178, #20486 and #20709) - #20806
Conversation
NVIDIA#20178 foundation) Squashed net change of PR NVIDIA#20486 at head 714916d relative to the main it last merged (0829689), ported onto feat/musespark-open-agentx; unrelated main commits are not included. NVIDIA#20486 carries the unified KV cache manager foundation of PR NVIDIA#20178 through eea6dcb (NVIDIA#20178's later 61248ec, 77d4eae and 636715a are not part of the NVIDIA#20709 stack and are not ported). Commits: 3bc36a5 [None][feat] add unified KV cache layout foundation dad9f1f [None][refactor] simplify unified KV cache layer specs 6ea5ad8 [None][test] remove unified KV cache foundation tests 090cbc0 [None][docs] remove unified KV cache design document 511975f [None][feat] support per-layer storage layouts in KVCacheManagerV2 bfc8bc7 fix: normalize explicit KV cache spec windows 18c71c1 fix: use explicit KV cache dtype for cold-page codec ae64dca test: cover explicit KV cache dtype for cold-page codec 701abea fix: honor physical slot strides in KV cache views 623816c test: cover mixed KV cache odd-slot addressing 7cde28e fix: restore KV cache subclass and test compatibility 4a3bb25 [None][test] fix mixed KV cache slot address arithmetic 9d0107b [None][fix] align NVFP4 data and scale page layouts e273a98 [None][fix] accept mixed KV cache pointer table layouts f6b1d08 [None][fix] preserve sink and reuse policy in KV cache costing b6db612 [TRTLLM-20434][feat] Dflash Unified KVCache 39d814d [None][fix] share the DFlash draft KV cache in disagg only with speculative_config in both sections 22cf25c [None][fix] gate the disagg shared draft KV cache on matching tensor_parallel_size feed06a [None][test] cover the disagg shared draft KV gate for generation-only speculative_config and mismatched TP 0319dc3 test: use default timeout for unified DFlash KV test 1f39128 fix: restore DFlash private arena fallback for separate KV caches 97fb214 fix: seed unified DFlash context lengths from bounded token counts caf4c45 test: cover DFlash mixed-precision block tables 875eff2 fix: require matching attention DP for shared draft KV 212fdab test: enable shared draft KV in DFlash integration test 1302d43 fix: keep mixed-window DFlash caches separate Port conflict resolutions (base = main 0829689, ours = agentx b7979ca + !10912): - kv_cache_manager_v2.py: our joint_kv_cache_reuse kwarg and the PR's layer_specs kwarg land at the same spot in _get_per_request_layer_tokens, get_cache_bytes_per_request and get_cache_quota_per_request (5 hunks): keep both. The generic layer_configs loop moves under the PR's `else:`; main's kv_pool_rebalance_config (NVIDIA#19077) is not on agentx and is dropped. - _util.py: agentx's _per_manager_cache_cost has no tokens_per_block parameter (main-only); keep the PR's unified layer_specs injection. - dflash.py: keep agentx's windowed VANILLA arena (_ctx_base) branch; ctx_len_gen (PR) equals self._ctx_len[slots]. - rank_info.py: drop main's bounded_replay_on_generation (not on agentx); keep num_draft_layers. - tests: drop main-only imports (PyTorchModelEngine, KvCacheRetentionConfig). !10912 stays intact: validate_block_attention_windows() is not re-added. Signed-off-by: Allison Lim <allim@nvidia.com> Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Squashed increment of draft PR NVIDIA#20709 at head 1703851 over NVIDIA#20486 at 714916d (both on main 0829689; the NVIDIA#20709 tree contains all of NVIDIA#20486), ported onto feat/musespark-open-agentx. Commits: 7a85a73 [None][feat] unify the standalone DSpark target and draft KV cache d492b52 [None][feat] reshard unified target/draft KV by per-view attention heads Port conflict resolutions (base = NVIDIA#20486 on main 0829689, ours = the NVIDIA#20486 port): - kv_extractor.py: agentx lacks main's view_partition hook (NVIDIA#19727); add only kv_heads_for_layer and the Callable import its annotation needs (Muse target and draft buffers differ in size, so they never share a physical pool/view). - kv_cache_manager_v2.py add_dummy_requests: keep agentx's stop_committing() for the dummy and size it with the PR's _kv_alloc_extra_tokens. - test_kvcm2_integration.py: keep agentx's DecoderRunner import and add the PR's tests only (main's generation-block reuse tests, NVIDIA#19710, are not on agentx). Signed-off-by: Allison Lim <allim@nvidia.com> Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…indings The ported NVIDIA#20178/NVIDIA#20486/NVIDIA#20709 Python assumes main's KV cache manager bindings. Two pieces are not on agentx: - Per-layer AttentionReusePolicy (main NVIDIA#19727 C++). The stack only ever uses REQUIRED, which is what agentx does for every layer, so the reuse_policy field, its AttentionLayerConfig argument and its place in the scratch and pool-group keys are dropped. The one test that needs PRIVATE is skipped. - BufferConfig.coalescing_group (the stack's own C++, which the agentx image does not carry). Group 0 is what the prebuilt bindings do implicitly, so the argument is passed only for nonzero (NVFP4) groups, which fail loudly unless the bindings are rebuilt from this branch's cpp/. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
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>
Port of main's DFlashWorker._advance_ctx_len (NVIDIA#19988, c98b976), which the unified KV cache stack relies on but agentx lacks. CUDA-graph padding, warmup and attention-DP dummy rows share one dummy slot that no request owns, and each step added the dummies' accepted tokens to its length. On a unified KV cache the drafter's allocation check covers those rows, whose pages are the padding request's, so the growing length tripped the device assert during CUDA-graph capture at server start. Empty the slot right after the add, inside the captured step; with the windowed arena its base too. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
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>
…d transfer Per-view head transfer dropped the equal-heads rule from the disaggregated layout check and the launcher's identical tensor_parallel_size and enable_attention_dp rule, but the gate tests still expected those refusals, and their speculative-config mock lacked is_dspark, which the gate now calls to name the drafter. Expect mismatched heads, attention DP and TP to share the cache, keep NVFP4 (whose block-scale views are not resharded) as the remaining head-count refusal, and give the mocks is_dspark (including the derived-window costing test, which reaches the same gate). Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…orkers Without TRTLLM_DISAGG_SHARED_DRAFT_KV=1, a DFlash or standalone DSpark drafter on the TRTLLM backend keeps a separate draft KV cache on a disaggregated worker. A transferred request's drafter fills it from column 0, while the cache commits history at the request's real positions and retires the blocks behind a sliding window, which the drafter's table then maps to block 0. Once a request passes the window, acceptance length silently collapses to about 1 (long prompts on the Muse Spark Open DFlash2 drafter: 1.005, VANILLA 3.604). The startup check that refuses an unsatisfiable TRTLLM_DISAGG_SHARED_DRAFT_KV=1 now also refuses this fallback and points at the shared draft KV or attention_backend VANILLA. A worker cannot tell which disaggregated role it serves, so a context worker is refused as well. The llm_args warning that the drafter's context is not transferred no longer fires when the draft KV is shared. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…hared draft KV feat/musespark-open-agentx turned the refusal of DFlash on a disaggregated engine into a warning (264e601) but kept the test that expects the refusal, so it failed. A draft KV shared with the target now silences the warning as well. Test the current behaviour: the warning fires for VANILLA, and for TRTLLM without TRTLLM_DISAGG_SHARED_DRAFT_KV=1; it stays silent for TRTLLM with it. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Without SWA scratch slots, a per-layer attention-op table is a plain conversion of the request's host page-index row. Every step still copied the whole host table to the device, rebuilt all layers x IndexMapper rows x max_blocks_per_seq into a full-capacity staging buffer, and then copied the batch's rows out. That work scales with max_seq_len and capacity, not with the batch: for a unified Muse Spark cache (67 layers, 129 rows) at max_seq_len 997696 it costs a 2 GiB staging buffer and 3.5-4.2 ms/step. Snapshot only the batch's host rows, up to the widest row's num_blocks (the rest of a row is BAD_PAGE_INDEX), into a fresh pinned buffer, copy it asynchronously and convert it straight into the destination with one Triton kernel. The snapshot also fixes the ordering: the async H2D of the live host table read it when the copy ran, so host-table updates for the next step could leak into this step's tables (nvbug 6293536 on V1). The batch's rows are bit-identical to the full rebuild; other rows are not touched. The copy refuses to run under CUDA-graph capture, which would replay a stale snapshot. The SWA-scratch path (unused by Muse Spark) is left as is. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
A life cycle that mixes page sizes (a unified FP8 target / BF16 draft cache, heterogeneous head dims) fell back to one attention-op table per layer: 67 tables for Muse Spark, each max_num_sequences x max_blocks_per_seq wide (1 GiB per metadata copy at max_seq_len 997696), rebuilt every step. A per-layer table is base * scale + layer_offset. Layers of one life cycle and physical pool with the same page stride, scale, K/V distance and kv factor differ only by layer_offset, which, as on the SHARED path, moves to the kv_cache_pool_mapping offset the attention op already applies. Such layers now share one table: Muse Spark needs 3 (window FP8, window BF16, full FP8), the same layout the separate target and draft managers had. Page addresses and BAD_PAGE_INDEX entries are unchanged, and DFlash's decoded draft tables and extents are identical. SWA scratch reuse, NVFP4 and subclasses keep per-layer tables. Two readers that assume life-cycle tables now decline shared ones: the modeling_v2 Muse Spark block-offset stager, and the opt-in prims_ts FMHA, whose K/V page offset is the life cycle's. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…aft KV With the unified KV cache, a disaggregated worker that sets TRTLLM_DISAGG_SHARED_DRAFT_KV=1 transfers the draft KV with the target's, so AUTO no longer needs VANILLA there. AUTO now keeps VANILLA on a disaggregated worker only without the variable, where a windowed TRTLLM draft pool is refused at startup; with it, AUTO follows the checkpoint rule as on an aggregated worker, and the warning that the drafter's context is not transferred stays off as it does for TRTLLM. Ranks pin AUTO before the KV cache creator reads the drafter backend, so a pinned AUTO unifies the target and draft KV like TRTLLM; a test covers it. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
On this branch kv_cache_config.use_kv_cache_manager_v2 defaults to
"auto", which resolves to V1 for Qwen3ForCausalLM. Both workers then
refuse TRTLLM_DISAGG_SHARED_DRAFT_KV=1 at startup ("the target does not
use KVCacheManagerV2"). Request V2 explicitly; main defaults to it.
Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ree on draft KV layers A context and a generation worker that keep a different number of draft layers in their unified KV caches cannot move the drafter's prompt KV. Only the sender checked this, in its listener thread, where the error is logged and the request then waits out the KV transfer timeout. The receiver now runs the same check with its other layout gates, so every request to such a context peer fails immediately with the reason, and the peer is cached as incompatible. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…with the default AUTO backend The launcher turned on draft KV sharing only for a DFlash attention_backend set to TRTLLM. AUTO is the default, so a config with DFlash in both sections and no backend kept the drafter's prompt KV on the context server, and the generation server drafted without it: acceptance dropped with only a warning. AUTO and an unset backend now count like TRTLLM, as they already did for DSpark. An unwindowed drafter, whose AUTO resolves to VANILLA, is refused at startup instead of losing acceptance. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…fault A disaggregated worker kept its drafter's KV with the target's, so that the transceiver moves it, only with TRTLLM_DISAGG_SHARED_DRAFT_KV=1. Workers launched on their own (one trtllm-serve per server) had to set it by hand on every server; without it, AUTO fell back to VANILLA and the generation server drafted without the prompt context, with only a warning. Unset, the variable now leaves the choice to the drafter: DFlash with attention_backend AUTO or TRTLLM shares, so a context and a generation server given the same speculative_config agree on their own. 1 still requires sharing and 0 keeps the draft KV private. A pair that disagrees, such as a drafter on the generation server only, fails at the first KV transfer; to draft there alone, set attention_backend VANILLA or TRTLLM_DISAGG_SHARED_DRAFT_KV=0 on it. DSpark keeps needing the variable. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…atched copy The receiver's scatter worker issued one batched copy per fan-in writer and synchronized the stream once at the end. The batched copy reuses one pinned metadata buffer per stream and relies on a sync between calls. With several writers the host refilled that buffer for writer k+1 while writer k's async upload was still queued behind writer k-1's kernel, so writer k's kernel ran with another writer's fragment table and writer k's head slice was never scattered: silent KV corruption. Seen on ctx TEP4 -> gen DEP8 (4 writers, ~1 GB each) once those receives were allowed to coalesce: every coalesced request decoded from partly stale KV (acceptance length 7.7-7.9 at K=7, degenerate output). Concatenate every writer's fragment table, each still read from its own source address, and issue one batched copy and one sync per region. A sync after every writer also removes the race, but measured ~117 ms per 4 GiB receive against ~50 ms for the racy single-sync loop (on different nodes), so keep a single copy and a single sync. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…ock sizes reserve() fell back to the per-fragment path for every multi-writer receive whose layer groups differ in per-block bytes. The equal split can only overrun when the writers hold different layers (PP fan-in); a TP fan-in writer sends a head slice of every block of every group, so its share is exactly total // num_writers for any group sizes. Receiver.dispatch_task now passes writers_split_layers (a fan-in whose writers span several PP stages) and reserve() applies the uniform-slot check only then. The default stays True, so other callers keep the old behavior. Depends on the previous commit: before it, coalesced fan-in receives could be scattered with the wrong fragment table. Example: ctx TEP4 -> gen DEP8 (4 writers) on a model with a full-attention and a sliding-window layer group (1 MiB vs 2.9 MiB per block, 3.5 MiB with DFlash draft layers) never coalesced; 4 GiB receives took ~450 ms on the per-fragment path. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…e MuseSpark pre-merge stage Pre-merge CI on feat/musespark-open-agentx runs only the MuseSpark stages, so the DFlash TRTLLM-backend, shared draft KV and disaggregated KV transfer tests were not run on this branch. List them in l0_gb300_musespark. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
a595535 to
e0d534b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #80558 [ run ] triggered by Bot. Commit: |
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Reviewed the 17 branch-specific commits in full and the #20709 port (6b40260a1) in full; I did not line-review the #20486+#20178 port (1a1b710a7), which has three upstream approvals and a merged foundation. Worth noting in the description that #20709 has no upstream reviews, unlike #20486 -- the two are presented together, and both blocking findings below sit in that half.
|
PR_Github #80558 [ run ] completed with state
|
…eSpark CPU stage GPU stages deselect cpu_only tests, so the seven files whose tests are all cpu_only collected nothing on l0_gb300_musespark and failed with pytest's exit code 5, and test_kv_transfer.py's cpu_only tests ran nowhere. List those on l0_cpu_musespark, which runs cpu_only tests, and keep test_kv_transfer.py on both. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…unce A multi-writer bounce splits the region equally among the writers, which holds only if every writer sends a share. The gate checked the receiver's rank-wide head duplication factor but not the senders', so when the senders replicate KV heads and elect a subset to send, the elected ones overran their neighbours' slots: 4 KV heads with ctx TEP8 -> gen TEP2, or MLA with ctx TP8 -> gen TP4. Check both factors, as the per-view clause already does; those layouts take the per-fragment path. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…saggregated_mpi_worker The launcher sets one value for every worker it spawns, and wrote 1 whenever both sections could share, overriding an explicit 0. Now that sharing is the default, 0 is how an operator keeps the draft KV private on a pair that would otherwise share; keep it. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…mup offsets once The batched op-table conversion kernel now takes BAD_PAGE_INDEX as a constexpr instead of a hardcoded -1, and its wrapper asserts that params is contiguous, as the kernel's pointer arithmetic needs. The warmup slot reseed carries a running token offset instead of re-summing seq_lens for every row. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #80840 [ run ] triggered by Bot. Commit: |
Thanks @zhaoyangwang-nvidia for the review! This PR is a port to muse spark side branch for AgentX submission. I have fixed most of the things you commented, but some comments maybe it's better to be raised in original PRs in main. Thanks! |
|
PR_Github #80840 [ run ] completed with state |
@coderabbitai summary
Description
In disaggregated serving, the DFlash drafter on the generation server starts without the prompt's draft KV, because the transceiver moves only the target's KV. Acceptance drops (Muse coding AL 4.13 aggregated, 3.89 1P1D), and with a TRTLLM drafter long prompts collapse to ~1.
This PR keeps the draft layers in the target's KV cache manager, so their KV is transferred with the target's. It ports Allison's #20486 (with #20178) and #20709 as her commits, plus what Muse needs on this branch:
Usage: give both servers the same DFlash
speculative_config. To draft only on the generation server, setattention_backend: VANILLAthere.Results on GB300 (Muse Spark Open, RadixArk DFlash2, K=7): disaggregated AL equals aggregated within noise for every layout tried (TEP4, TEP8, DEP4, DEP8), up to ~1M tokens and on AgentX traffic. GSM8K, step time and memory are unchanged. Under
modeling_v2the drafter stays on VANILLA until #20812. Not validated: the C++ changes, which only NVFP4 KV needs.Test Coverage
New tests for the shared KV gate, page tables, DFlash context and dummy slot, KV-bounce fan-in and the context/generation mismatch, plus
TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv. These and the existing DFlash and KV transfer tests now run inl0_gb300_musespark, this branch's only pre-merge GPU stage.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.🤖 Generated with Claude Code