[None][perf] DFlash performance optimizations - #13995
Closed
ziyixiong-nv wants to merge 7 commits into
Closed
Conversation
Cherry-pick of NVIDIA#13782: - Fix mrope position_ids in DFlash - Fix GDN verify buffer size for DFlash - Fix GDN/Mamba2 state updates in DFlash - Add tests Original-author: Anurag Mukkara <134339030+amukkara@users.noreply.github.com> Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Replace per-layer context K/V recomputation on the padded target_hidden
with a persistent per-layer cache, and drop the legacy recompute path
entirely.
- Build _ctx_k_buf / _ctx_v_buf of shape
[max_batch, L, max_ctx+block_size, nkv, hd] at worker init. Each
decode iter runs one fused GEMM (stacked across drafter layers) on
the newly accepted K+1 tokens and scatters into the per-slot buffer;
subsequent iters read the cached tensors directly via
flash_attn_with_kvcache's cache_batch_idx — no gather, no re-project.
- Use fused q_norm + k_norm + RoPE per layer in dflash_forward via
torch.ops.trtllm.fused_qk_norm_rope.
- Use fused cross-layer RoPE in precompute_context_kv via
flashinfer_apply_rope_with_cos_sin_cache_inplace.
Drop the old fallback that recomputed context K/V per-layer from
target_hidden each step: every DFlash checkpoint has
_build_fused_kv_buffers, so the fast path is always taken. Removed
dflash_forward's target_hidden / context_positions params, the
use_ctx_cache branching, the per-layer recomputation, and the
DFlashForCausalLM._kv_buf_k / _kv_buf_v / _cache_seqlens /
_block_offsets plus DFlashWorker._ctx_buf / _ctx_pos_buf / _proj_dim
attributes that only served the fallback.
Qwen3-8B, 8xB200, TP=8, MTBench, mdl=3, tok/s per user:
DFlash(new) DFlash(old) Eagle3(best)
bs=1 977.2 701.1 909.4
bs=2 881.1 661.3 799.5
bs=4 801.9 608.8 748.3
bs=8 729.7 529.3 691.4
bs=16 592.3 317.5 564.4
bs=32 495.2 224.0 482.3
bs=64 452.5 156.0 438.6
Memory: +160 MB per TP rank for the ctx K/V buffers on Qwen3-8B.
Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
fused_qk_norm_rope in dflash_forward rebuilt YaRN frequencies from the HuggingFace formulation (honors truncate=false), while precompute_context_kv rotated context K through the DeepSeek-V2 style cache (always truncated). On gpt-oss-120b (truncate=false) the two paths produced different cos/sin for the same position, so context K and query Q / noise K were rotated inconsistently. Attention scores decayed to noise and AR collapsed from 0.35 to 0.012. Replace fused_qk_norm_rope with explicit q_norm + k_norm followed by flashinfer_apply_rope_with_cos_sin_cache_inplace reading the same cache as precompute_context_kv. Both paths now see identical cos/sin for any rope scaling. Qwen3-8B (no rope_scaling) unaffected. gpt-oss-120b recovers and beats Eagle3: mtbench bs=1 mdl=3: tps 357 -> 975, AR 0.012 -> 0.695 (Eagle3: 822, 0.437) mtbench bs=1 mdl=5: tps -> 1174, AR -> 0.582 codeEdit bs=1 mdl=3: tps -> 915, AR -> 0.636 (Eagle3: 888, 0.470) codeEdit bs=16 mdl=3: tps -> 422, AR -> 0.646 (Eagle3: 441, 0.390) Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…edding lookup Replaces the ~10 scattered small PyTorch ops (expand + clone + gather + scatter + arange broadcasts) that ran sequentially on the host per draft step with: 1. A few broadcast-add expressions for query_positions / ctx_positions. 2. A single embed_tokens call over [bonus..., mask_token_id], cutting the per-step vocab-sharded embedding gather+all-reduce from O(num_gens * block_size) to O(num_gens + 1) tokens. Routes through embed_tokens.forward (not .weight[...]) so TP-sharded vocabularies are handled correctly (each rank only owns a slice of the embedding table and must all-reduce). Verified on Qwen3-8B TP=8 B200 / MTBench, bs=64 mdl=8: AL and output length bit-identical to the old pure-Torch baseline; +1.4% tps/gpu, +2.4% tps/user. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
The YaRN RoPE fix disabled the fused q_norm + k_norm + RoPE kernel. For drafters with no rope_scaling and partial_rotary_factor=1.0 the fused kernel and the flashinfer fallback produce identical frequencies, so the fused kernel is safe and saves two launches per layer. Guard picks the fused kernel only when rope_params.scale_type is None (or RotaryScalingType.none) and partial_rotary_factor == 1.0. YaRN / long-rope / partial-rotary drafters keep the safe fallback path. AR unchanged: Qwen3-8B (no rope_scaling): bs=1 mdl=3 AR=0.702, tps ~960-975 gpt-oss-120b (YaRN): bs=1 mdl=3 AR=0.695, tps ~981 Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Two AR-preserving cleanups that target the large-batch draft path.
- precompute_context_kv: replace the Python for-loop of L per-layer
k_norm calls with a single F.rms_norm over [N, L, nkv, hd] followed
by a broadcast multiply with the stacked per-layer weights. Stacks
k_norm weights once in _build_fused_kv_buffers; asserts all layers
share variance_epsilon (reports the observed values on mismatch so
the error is self-diagnosable).
- prepare_1st_drafter_inputs: turn 'k_new = k_new * mask_bc' into
'k_new.mul_(mask_bc)' (and same for v_new). Avoids allocating a
new [num_gens*(K+1), L, nkv, hd] tensor each step.
Qwen3-8B, 8xB200, TP=8, MTBench:
bs=1 mdl=10: tps 1571 -> 1582, AR 0.4414 -> 0.4416
bs=64 mdl=3 : tps ~452 -> ~442 (2 runs: 434, 450; within run-to-run
noise), AR 0.709 -> 0.711
Neither change touches the math; AR diffs are at fp rounding level.
Perf impact is small at this scale (kernel-launch savings are a tiny
fraction of bs=64 step time), but the code is cleaner.
Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
DFlash was forwarding 2K tokens per gen request through the full target model (K+1 accepted tokens plus K-1 mask fillers), but only the K+1 prefix is consumed by acceptance and the context-store projection. The K-1 fillers hit every target layer and allreduce for no benefit, and at large batch that wasted target work is the dominant step-time gap against Eagle3. Reduce tokens_per_gen_step to K+1: target now runs exactly the tokens the algorithm needs, and the DFlash worker drops the paired slicing, padding, and reshape that only existed to bridge the 2K/K+1 mismatch. The draft model's block_size and mask-query generation are untouched, so acceptance length is preserved byte-for-byte. Qwen3-8B on B200, TP=4, MTBench: bs=64 mdl=3: 4422 -> 5064 TPS/gpu (+14.5%), +16.1% vs Eagle3 4363 bs=64 mdl=7: 4756 TPS/gpu, +4.6% vs Eagle3 4548 bs=32 mdl=7: 3377 -> 4162 TPS/gpu (+23.3%) bs=32 mdl=3: 3503 -> 3806 TPS/gpu (+8.6%) AL unchanged across all configs. Tests: test_dflash_qwen3_8b, test_dflash_qwen3_5_4b (both overlap- scheduler variants), TestQwen3_8B::test_dflash (GSM8K) all pass. Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Contributor
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (11)
📝 WalkthroughWalkthroughThis PR refactors DFlash speculative decoding from a 2K-token-pair model to a K+1 model (K draft tokens plus 1 bonus token). It adds precomputed context K/V caching, flashinfer-compatible RoPE support, and fixes position ID handling for 3D MRoPE tensors while extending validation tests for Qwen models. ChangesDFlash K+1 Refactor
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
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.
Summary
Performance optimizations for the DFlash speculative-decoding worker on the PyTorch backend. The headline change is dropping the K-1 mask-filler tokens that DFlash was forwarding through every target layer; the rest are supporting fusions that close the remaining gap on the draft side.
Qwen3-8B on B200, TP=4, MTBench:
Acceptance length is byte-for-byte preserved vs. the pre-branch baseline.
Commits
d517eedDFlash: drop K-1 filler tokens from target gen input — reducetokens_per_gen_stepfrom 2K to K+1, remove the paired slicing/padding. +14.5% TPS/gpu at bs=64 mdl=3 (main large-batch win).59ea073DFlash: fuse cross-layer k_norm and avoid mask-mul alloc.f609129DFlash: use fused qk-norm+rope on non-YaRN draft.785cdeeDFlash: fuse query-input build in pure Torch + fused embedding lookup.5c45f60DFlash: match RoPE between context and query for YaRN (bug fix).cbdb9f4DFlash: migrate draft path to per-layer cached K/V pool.4d73942DFlash: Qwen3.5 fixes (cherry-pick).Test plan
tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_8b(both overlap-scheduler variants) — passtests/unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b(both variants) — passtests/integration/defs/accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash(GSM8K) — pass/bot runSummary by CodeRabbit
New Features
Improvements