Repository navigation
Conversation
…lits The VANILLA DFlash drafter attends through flash_attn_with_kvcache over a dense per-slot context arena whose seqlen dimension is min(max_seq_len, max_position_embeddings). FA2's split-KV heuristic sizes num_splits from k_cache.size(1), i.e. the arena capacity, not the live context length. With a 262k arena and ~38k live tokens it picks ~4-6 splits of ~44-65k tokens, so only split 0 has work: ~70-100 CTAs on a 212-SM GPU, each streaming the whole context serially, in every drafter layer on every step. Pass num_splits=128 explicitly. It does not depend on the live length, so it is CUDA-graph safe; splits past the live length exit early and the combine kernel is a few microseconds. Measured on Kimi-K3 + DSpark drafter (draft 7), Rubin VR200 (sm107, 212 SMs), disaggregated 1P1D TP8/TP8 C16 with full-prompt drafter context (NVIDIA#19733): per-call drafter attention 612 us -> ~150 us, gen step -1.2 ms at batch 8 and -2.3 ms at batch 10-12, output tok/s/user +6.3%, acceptance length unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Viraat Chandra <viraatc@nvidia.com>
… index H2D snapshot_managed_draft_history and the managed-history / block-count refreshes build small index tensors with torch.tensor(list, device=cuda). That is a pageable H2D copy, which synchronizes the host with the stream on every step and defeats the overlap scheduler: the next step's prepare can no longer run while the current forward is on the GPU. Build them in pinned host memory and copy non_blocking, as the rest of dflash.py already does (prefer_pinned()). Measured on Kimi-K3 + DSpark drafter (draft 7), Rubin VR200, disaggregated 1P1D TP8/TP8 C16 with full-prompt drafter context, on top of the FA2 num_splits fix: nsys showed a 26 ms cudaStreamSynchronize per step (the whole forward) right after the graph launch, and GPU idle of 6.0 vs 1.5 ms/step without drafter context. With this change the gen step drops 2.2-2.5 ms at equal batch, output tok/s/user +15% and total tok/s/GPU +6%. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Viraat Chandra <viraatc@nvidia.com>
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
1 participant
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.
Stacked on NVIDIA#19733 (unified DSpark/DFlash draft KV, by @allisonlim-nv). The base branch
viraatc/pr19733-baseis NVIDIA#19733's head (ca1646af61), so this diff shows only the two fixes. It is for sharing; it does not touch NVIDIA#19733.On Rubin VR200 (Kimi-K3, DSpark draft 7, disaggregated 1P1D TP8/TP8 C16), NVIDIA#19733 alone raises acceptance length by 16% (3.35 → 3.89). As built, though, the gen step is 26% slower, so output speed drops 10%. These two commits make it a net win.
num_splitsThe win holds on 1P3D (1× TP8 ctx + 3× TP8 gen, 32 GPUs): +10.7% X at C24 and +10.2% at C38, with +2.3–2.4% total.
num_splits=128inmodeling_dflash.py, on the VANILLA path. Without it, FA2 sizes the split count from the max_seq_len-dense arena (262k), not the live ~38k context, so attention ran effectively unsplit at 612 µs per call. This is the same change as [None][perf] DFlash VANILLA draft attention: pass explicit FA2 num_splits NVIDIA/TensorRT-LLM#20120, which targetsmain.dflash.py:snapshot_managed_draft_history, the managed-history slot map, and the block counts. The pageabletorch.tensor(list, device=cuda)copies synchronized the host on every step: a 26 mscudaStreamSynchronizeper step in nsys. That defeated the overlap scheduler.Not run: unit tests. The measurements above are on silicon, with the two patched files bind-mounted into the K3 Rubin image.
🤖 Generated with Claude Code