[TRTLLM-14597][perf] Fuse the DSv4 MLA prologue: kv_a_layernorm, q_nope FP8 quant and Q RoPE - #17273
Draft
dc3671 wants to merge 13 commits into
Draft
[TRTLLM-14597][perf] Fuse the DSv4 MLA prologue: kv_a_layernorm, q_nope FP8 quant and Q RoPE#17273dc3671 wants to merge 13 commits into
dc3671 wants to merge 13 commits into
Conversation
Context phase only, mirroring the existing fused FP8-Q path. Removes two launches per layer: the standalone 512-wide kv_a_layernorm and the concat([compressed_kv, k_pe]) that follows it. Kernel. applyMLARopeAndAssignQKVKernelOptContext gains a kFuseKvNorm template arg. Its KV region gives one WARP a whole kv_lora_rank + qk_rope_head_dim row, so the sum-of-squares is a single warp shuffle -- the same shape that makes deepseekV4QNormFusedKernel work on the Q side. One pass loads the row to registers and reduces; a second pass scales by the norm weight, rotates the rope tail, quantizes and scatters to the paged cache, all without the values leaving registers. That is only possible because the Q region stops touching fuse_buf in this mode. Today every one of the head_num blocks loads and rotates k_pe while only head_idx == 0 writes it, so 127 of 128 blocks discard the result. Under kFuseKvNorm the k load/rotate/write is compiled out and the KV region owns k_pe outright. The generation kernel cannot do any of this: it splits one latent row across blockIdx.y regions (dims 448..511 in head_idx == head_num, dims 0..447 in head_num+1..+8), so no block there holds the data the RMS denominator needs. Strides. latent_cache on this path is a last-dim slice of the kv_a_proj output, so its row stride is q_lora_rank + kv_lora_rank + qk_rope_head_dim, not the packed width -- the concat being removed was also silently compacting. The kernel therefore takes the row stride as a parameter, read from latent_cache.stride(0); calling .contiguous() in Python would just reintroduce the copy this change exists to delete. Only the innermost dim must be unit-stride, which is what the 16B vector loads require. Gating mirrors _is_fused_q_fp8_quant_enabled: DSv4, kv_lora_rank == 448 and qk_rope_head_dim == 64 (the kernel describes the row with its K_DIM/ROPE_DIM template constants, and a static_assert rejects any layout whose row does not divide evenly across a warp), FP8 KV cache, absorption mode, and num_generations == 0. TRTLLM_DISABLE_FUSED_KV_NORM=1 is the kill switch. With the gate off the emitted code is byte-identical to before. Builds for sm_100f. Not yet validated at runtime: no perf number and no numerics check. layer_wise_benchmarks cannot run DSv4 (see previous commit), so the A/B needs a different vehicle. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com> (cherry picked from commit 08c56d5)
The context phase already folds kv_a_layernorm into the MLA RoPE kernel (`kFuseKvNorm`). Generation could not do the same in place: its RoPE kernel splits one 512-wide latent row across two disjoint blockIdx.y regions -- the rope tail at `y == head_num` and the nope segment at `y 129..136` -- so no block, let alone warp, sees the whole row and the RMS denominator cannot be formed. On top of that an FP8 KV cache multiplies that grid by 9. Move the KV work into a standalone kernel instead: * `mlaKvNormRopeQuantGenerationKernel` gives one WARP a whole latent row (64 x 16B vectors / 32 lanes = 2 per lane), so the sum-of-squares is a single `__shfl_xor_sync` reduction. Pass 1 loads the row and reduces; pass 2 scales by the norm weight, rotates the rope tail, quantizes and scatters to the paged cache, with values never leaving registers. Its grid is sized from rows alone, with no coupling to head_num or to the FP8 grid expansion. * `applyMLARopeAndAssignQKVKernelGeneration` gains `kSkipKv`, compiling out both KV regions. The `seqQOffset` stamp that lived in the nope region moves to block (0,0); the skipped blocks trigger programmatic launch completion before returning so PDL is unaffected. * Python then skips the standalone RMSNorm and the `concat([compressed_kv, k_pe])` and hands the kernels the RAW kv_a_proj slice, whose row stride is `q_lora_rank + 512` -- passed explicitly rather than materialized with `.contiguous()`, which would reintroduce the copy the fusion removes. Also hoist the cumulative Q/KV sequence lengths out of the per-layer kernel: they are layer-invariant, so the attention metadata computes them once per iteration into fixed-address buffers (`precomputed_cu_seqlens` then skips the in-kernel stamp and the `cub::BlockScan`). `cu_kv` is built in `prepare()` from the same pre-extra-tokens kv_lens the kernel's scan consumed. Gated by the existing `TRTLLM_DISABLE_FUSED_KV_NORM` kill switch plus the V4 layout checks; with the gate off the emitted code is unchanged. Measured on GB200, DSv4-Pro dep32/bs32/MTP3, rank 0 over 50 iterations: the standalone RMSNorm (3200 launches, 7.14 ms) and the latent concat (3200, 8.24 ms) disappear, replaced by 3200 launches / 14.72 ms of the fused kernel. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com> (cherry picked from commit b51f551)
`forward_absorption_context` cleared `_fused_kv_norm_active` at its end, but in a mixed batch it runs BEFORE `forward_absorption_generation`, which reads the flag to decide whether to pass `kv_norm_weight` down to `mlaKvNormRopeQuantGenerationKernel`. The generation half therefore saw False while `latent_cache_gen` was still the RAW kv_a_proj slice, and wrote un-normalized latent rows into the KV cache for every generation token in a mixed batch. Generation-only batches kept the flag and were correct, which is why the kernel matches its reference bit-for-bit in isolation. Clear the flag in `forward_impl_with_deepseek_v4` instead -- the only place that sets it -- once both halves have consumed it. GSM8K on DeepSeek-V4-Pro, 2x GB200 (tp8, fp8 KV, MTP1), threshold 92.797: un-fused 96.513 pass fused, before this fix 73.237 / 76.308 / 77.142 fail context fusion only 96.209 pass fused, after this fix 96.209 pass Also adds two env kill switches used to bisect this: TRTLLM_FUSED_KV_NORM_CTX_ONLY=1 un-fuses any batch carrying generation tokens, TRTLLM_DISABLE_PRECOMPUTED_CU_SEQLENS=1 restores the in-kernel cu_seqlens fill. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com> (cherry picked from commit 0d5b1df)
Two things in the DSv4 generation path.
1. PDL race in `mlaKvNormRopeQuantGenerationKernel`. The launch sets
`programmaticStreamSerializationAllowed`, and PDL defaults to ON for SM >= 90
(`getEnvEnablePDL`, only TRTLLM_ENABLE_PDL=0 turns it off), so the kernel was
free to start before the kv_a_proj GEMM that produces `fuse_buf` had finished
writing it. Add the missing `cudaGridDependencySynchronize()` before the
latent reads, plus a `cudaTriggerProgrammaticLaunchCompletion()` at the end so
the RoPE kernel that follows can still overlap this one's tail.
2. The FMHA scheduler prologue -- zeroing `fmha_tile_counter` and deriving the
bmm1/bmm2 scales -- moves out of block (0,0) of the MLA RoPE kernel and into
`_deepseek_v4_local_to_global_kernel`. That Triton kernel is the LAST one
launched before FMHA and already runs once per layer per forward for exactly
FMHA's benefit, so it is the natural owner; the RoPE kernel was doing this
work two launches early. Program 0 writes it ahead of the kernel's own
`gdc_wait()`, since it touches none of the index inputs.
`precomputed_fmha_scheduler` tells the C++ side to stand down (the fused KV
launcher passes null scheduler pointers, the RoPE kernel skips its prologue),
so there is exactly one writer. Together with the `precomputed_cu_seqlens`
hoist, block (0,0) of the generation RoPE kernel is now empty on the DSv4 FP8
path -- the kernel is pure Q work.
Also: the three scheduler scalars become persistent per-layer buffers instead of
`torch.empty` per layer per forward, and the fused KV kernel picks its block size
(32/64/128/256 threads) so the grid covers the SMs -- at batch 32 with MTP3 the
fixed 256-thread block gave 16 blocks on 148 SMs, leaving the kernel launch-bound.
`TRTLLM_MLA_KVNORM_GEN_ROWS_PER_BLOCK` pins it for tuning.
Numerics, single GPU, fused vs a torch-RMSNorm reference through the real
attention layer -- KV cache, rotated q, fp8 q buffer, cu_seqlens, tile counter
and both bmm scales all bit-identical (bf16 KV differs by one ULP, as expected):
generation: bf16; fp8; fp8 gen_len=2 (MTP1); fp8 8 seqs
context: bf16; fp8; fp8 + fused-q-fp8; 1k ctx; 8 seqs
A dedicated test feeds the RoPE op throwaway scale buffers and NaN-poisons the
ones handed to the attention forward, so only the Triton kernel can produce the
checked values; it reproduces bmm1/bmm2 exactly and clears the counter poison.
GSM8K on DeepSeek-V4-Pro, 2x GB200 (tp8, fp8 KV, MTP1), same build:
fused + Triton prologue 96.475 pass (repeat: 95.982, 0.018 under the
test's ref gate of 96.000)
un-fused control 96.285 pass
Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com>
(cherry picked from commit e9f6910)
…cupancy
The adaptive launch shape sized the grid to cover the SMs
(divUp(rows, sm_count)). ncu says that is the wrong knob: the kernel is
latency-bound, not occupancy-bound (achieved occupancy 1.6-13%), and the
heuristic lands on the slower option at both ends of the decode range.
Median of 6 launches on GB200, MTP3:
rows/block block 128 rows (b32) 896 rows (b224)
1 32 6688 ns 7168 ns <- heuristic picks at b32
4 128 6432 ns 6784 ns <- new default
8 256 7328 ns 6960 ns <- heuristic picks at b224
So drop the SM query and default to 4 warps, keeping
TRTLLM_MLA_KVNORM_GEN_ROWS_PER_BLOCK for re-tuning. Re-measured after the
change: 6192 ns at b32, 6640 ns at b224.
For the record, the fusion as a whole against the un-fused path
(RMSNorm + concat + RoPE kernel owning the KV half), same method:
batch 32, MTP3 18080 ns -> 15424 ns -14.7%
batch 224, MTP3 45664 ns -> 40848 ns -10.5%
per layer per decode step. The RoPE kernel itself is unchanged by kSkipKv --
the KV region was never its bottleneck -- so the win is three kernels becoming
two, with the two removed ones (9.2 us at b32) replaced by one 6.4 us kernel.
Numerics re-verified after the change: generation and context A/B both
bit-identical to the reference, Triton prologue attribution unchanged.
Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com>
(cherry picked from commit c5f2915)
… kernel `applyMLARopeAndAssignQKVKernelGeneration` launches `grid.y = head_num + 9 + head_num * 8` when the KV cache is FP8. At head_num 128 that is 1161 block rows, of which 1024 -- 88% -- are the trailing q_nope FP8 quantize-copy region, re-quantizing a buffer q_b_layernorm had just produced in bf16. Only 128 rows do the Q RoPE the kernel is named for. Context already avoids this: `deepseek_v4_q_norm_fused_fp8` writes the FP8 nope segment straight out of q_b_layernorm at full 512 row stride, leaving the rope slots for the RoPE kernel, so the context kernel never quantizes q_nope. Generation could not, because `_is_fused_q_fp8_quant_enabled` refused any batch carrying generation tokens: the fused path returns a placeholder bf16 q_buf and the generation half derived both `q_nope` and `q_pe` as views into it. That is now only true of MIXED batches. `q_nope` is dead on the DSv4 branch -- it feeds the non-DSv4 bmm alone -- and `q_pe` has a normalized replacement in `_fused_q_pe`, so a decode-only batch can take the same path context does. The gate therefore rejects mixed batches instead of all generation, and the launcher drops the q_nope rows when the caller pre-filled the buffer, making that region unreachable without a new template instance. The rope segment is quantized with `quant_scale_qkv` so both halves of a Q row share one scale; they are both 1.0 today, but that was a coincidence rather than a contract. Measured on GB200, dep32 / batch 32 / MTP3, decode steady state: grid.y 1161 -> 137 applyMLARopeAndAssignQKVKernelGeneration 10611 -> 3361 ns (-68.3%) deepseekV4QNorm[Fused]Kernel 5874 -> 5864 ns ( -0.2%) mlaKvNormRopeQuantGenerationKernel 4565 -> 4534 ns ( -0.7%) MLA prologue per layer-instance 21050 -> 13759 ns (-34.6%) Absorbing the quantize costs the q-norm kernel nothing: it already read those rows and now writes FP8 instead of bf16. 64 launches per iteration x 7.29 us predicts -1.59% device step time; measured p10 is -1.71% over two independent runs (28.696 -> 28.20 ms). GSM8K full accuracy 96.361. Also lands, off by default, the split that made this measurable: the generation rope op can now launch its KV and Q halves separately (`kv_only` / `kv_done_elsewhere`), so the fused kv-norm kernel can be hoisted onto `aux_stream` ahead of the Q branch. That hoist is gated behind `TRTLLM_MLA_KV_NORM_HOIST=1` because it loses: the 23.4 us span in front of the KV kernel is idle *stream* time, not idle SMs, so overlapping it with the Q branch slowed both kernels (KV +25.5%, q_b_layernorm +26.6%) and step time by 0.45 ms. The split itself is what lets the Q kernel run `kSkipKv`. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com> (cherry picked from commit 8cd26c8)
With the q_nope FP8 requantize region gone, `applyMLARopeAndAssignQKVKernelGeneration` is down to 137 of its former 1161 block rows and does nothing but rotate q_pe and write it into the FP8 Q buffer. `deepseekV4QNormFusedKernel` already holds that data in registers one step earlier, and its layout suits the rotation exactly -- `static_assert(kRopePairs == kWarpSize)` means lane `l` owns precisely rope pair `l`, so the rotation is register-local and costs one float2 of cos/sin. Doing it there leaves the whole Q row complete and the RoPE kernel with no work at all, so `mla_rope_generation` is called `kv_only` and that kernel stops being launched on the DSv4 decode path. Positions are derived exactly as the kernel being replaced derived them. Two forms: generation has uniform query length, so batch = token / seq_len; context is ragged, so the owning sequence is found by binary search over `cu_q_seqlens` and the cached offset (`kv_cache_len - current_seq_len`) is added, which is what makes chunked prefill correct. Every lane of a warp shares its row, so the search is warp-uniform. Note `cu_q_seqlens` counts TOKENS here, unlike the generation `seqQOffset`, which counts Q rows. Measured on GB200, dep32 / batch 32 / MTP3, decode steady state, one build, the two arms differing only by `TRTLLM_DISABLE_FUSED_Q_ROPE`: applyMLARopeAndAssignQKVKernelGeneration 3364 -> not launched deepseekV4QNormFusedKernel 5907 -> 7246 ns (+22.7%) mlaKvNormRopeQuantGenerationKernel 4576 -> 4106 ns MLA prologue per layer-instance 13848 -> 11352 ns (-18.0%) Device step time is NOT measurably improved: 64 launches x 2.50 us predicts -0.56%, which is under the ~0.49% run-to-run floor, and the measurement straddles zero (p10 -0.31%, p50 +0.45%). So this buys 18% of the MLA prologue's GPU time and one fewer launch per layer per step, not wall clock at this shape. GSM8K full accuracy 96.626. The context form of the position derivation is implemented and plumbed but stays behind `TRTLLM_ENABLE_FUSED_Q_ROPE_CTX=1`: in fused mode the kernel writes the rotated rope segment as FP8 and leaves bf16 `q_pe` stale, and the context RoPE kernel's Q region would rotate that stale buffer over the correct FP8 slots. Enabling it needs a matching skip in `applyMLARopeAndAssignQKVKernelOptContext`. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com> (cherry picked from commit cf7ef9c)
The generation half of this already landed. Context differs only in how a row's position is found: query lengths are ragged, so `batch = token / seq_len` does not hold. The kernel instead binary-searches `cu_q_seqlens` for the sequence owning a token and computes `position = (token - seq_begin) + (kv_cache_len[b] - current_seq_len)`; the second term is the chunked-prefill cached offset, matching what `applyMLARopeAndAssignQKVKernelOptContext` derives. Every lane of a warp shares its row, so the search is warp-uniform. Note `cu_q_seqlens` counts TOKENS here, unlike the generation `seqQOffset`, which counts Q rows. Once q_b_layernorm has rotated the rope segment into the FP8 Q buffer, the context RoPE kernel must NOT run its Q region: the bf16 `q_pe` it would rotate is left stale by the fused kernel, so rotating it would overwrite good FP8 rope slots with garbage. `MlaParams::q_rope_done` carries that, threaded from `mla.py` through `AttentionForwardArgs` -> nanobind -> `thop::attention` -> `invokeMLARopeContext`, and the Q region early-returns. The return sits after the `kOutputFp8Q` bmm-scale prologue, which runs on block (0,0,0) outside that branch, so the scheduler scales are still emitted. The token-wise cumsum comes from a new `mla_prepare_ctx_cu_seqlens`, built once per iteration into a fixed-address buffer. It deliberately does not reuse `ctx_uncached_token_indptr`, which holds the same values but only exists when `enable_context_mla_with_cached_kv` is set: with chunked prefill and block reuse both off, that buffer is absent and the fusion silently fell back to the old path -- GSM8K still passed, and only the kernel trace showed it had never engaged. Measured on GB200, CTX rank 0, ISL 8192, one build, arms differing only by `TRTLLM_DISABLE_FUSED_Q_ROPE`: applyMLARopeAndAssignQKVKernelOptContext 70443 -> 23565 ns (-66.5%) deepseekV4QNormFusedKernel 205310 -> 225194 ns (+9.7%) combined 275752 -> 248759 ns (-9.8%) share of CTX GPU time 2.7% -> 2.4% Prefill is MoE-communication bound, so this is ~0.3pp of context GPU time, not a throughput win; it is worth having for symmetry with the generation path and for the launch it removes. GSM8K full accuracy 96.399. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com> (cherry picked from commit 270c968)
Once the context Q RoPE moved into `deepseekV4QNormFusedKernel`, the Q region of `applyMLARopeAndAssignQKVKernelOptContext` became `if (q_rope_done) return;` -- but the launcher still sized the grid for it. At head_num 128 that is 136 z-rows launched to do the work of 8: 128 of them exist only to exit. Add `mlaKvNormRopeQuantContextKernel`, the context sibling of `mlaKvNormRopeQuantGenerationKernel`. Same warp-per-row shape -- one warp owns a whole kv_lora_rank + qk_rope_head_dim latent row, so the RMS reduction is a warp shuffle -- differing only in addressing: context walks `cu_q_seqlens` and a per-sequence `cached_offset` where decode uses a uniform `seq_len`. It also carries the bmm-scale prologue, which the Q-norm kernel does not emit. On the DSv4 fused path (`fuse_kv_norm_in_rope && q_rope_done`) the old kernel is no longer launched at all and the grid matches the work. `rope_append` and the non-fused DSv4 paths still use it, so the symbol stays for now. Not yet compiled or validated: the cluster's project inode quota is 4.2x over its hard limit and every build fails writing into the conda env. Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com> (cherry picked from commit a31d1a0)
Three kernels, one theme: cut per-lane memory instructions and integer
divides, and stop holding registers the decomposition does not need.
mlaKvNormRopeQuant{Generation,Context}Kernel
- `cudaGridDependencySynchronize` now precedes every global read. The
generation kernel read `quant_scale_kv[0]` and the FMHA-scheduler
scale tensors ahead of the barrier; those are static today, so it was
correct, but it would silently take a stale value the moment a
per-iteration KV scale is wired in. The PDL trigger stays after the
main loop.
- `kv_norm_weight` is row-invariant per lane, so it is hoisted out of
the row loop into registers as float, pre-multiplied by the quant
scale. It was indexed per element, which nvcc could not prove aligned,
so it emitted ELTS_PER_VEC 2B loads per vector per row -- each one
spanning four L1 wavefronts.
- The nope segment now goes bf16 -> float -> FP8 instead of
bf16 -> f32 -> bf16 -> f32 -> fp8. The round trip existed only so
`quantCopy` could take a T pointer; `quantStoreFromFloat` takes the
floats directly. Rope vectors keep the T path because
`rotary_embedding_transform` is typed on GPTJEltT.
- `/` and `%` by the runtime `seq_len` were two emulated 32-bit divides
per row. It is a power of two in every shipping config (1 + MTP
depth), so the host passes a flag and the kernel uses shift/mask,
falling back to the divide otherwise.
- The page-table load is issued before the reduction instead of after:
it depends on nothing the row loads produce, so it overlaps them.
- Both kernels are templated on the cache dtype, which removes the
per-vector branch, and the context kernel hoists its three per-batch
loads out of the token loop -- nothing could hoist them before,
because the FP8 stores alias through `unsigned char`.
- `TRTLLM_MLA_KVNORM_GEN_ROWS_PER_BLOCK` is read once, not per launch.
deepseekV4QNormFusedKernel
- Adds `kWideVec`, selecting the row decomposition per phase.
Generation takes the wide path (16B loads, 8B FP8 stores: 15 memory
instructions per lane per row become 4). Context keeps the narrow
pair-at-a-time path: wide needs 3 more registers, which drops
reg-limited occupancy from 16 to 14 blocks/SM, and the context
instance moves 1.6 GB per launch at ~81% of HBM SOL, where resident
waves are the bandwidth.
- Position derivation is hoisted out of the store loop (it is
warp-uniform) and the `row / num_heads`, `token / seq_len`,
`token % seq_len` divides take the same shift/mask treatment.
Measured on GB200 (poly), DSv4-Pro MXFP4, 4 layers, mean over 484
launches:
mlaKvNormRopeQuantGenerationKernel 3.91 -> 3.11 us (-20.5%)
mlaKvNormRopeQuantContextKernel 7.61 -> 7.02 us (-7.8%)
deepseekV4QNormFusedKernel, gen 3.10 -> 2.83 us (-8.7%)
deepseekV4QNormFusedKernel, ctx 248.41 -> 233.44 us (-6.0%)
The context q-norm figure is the narrow path with the phase-independent
fixes; forcing it onto the wide path instead measured 278.70 us, which
is what the occupancy argument above predicts. Its achieved bandwidth
goes 6.49 -> 6.90 TB/s, about 81% -> 86% of SOL.
16/16 DSv4 MLA unit tests pass.
(cherry picked from commit 2824df2)
…rnel by default
Two independent placement changes in the DSv4 generation prologue, both on by
default, plus five alternative KV placements behind one env knob for A/B work.
Compressor pre-launch (TRTLLM_MLA_COMPRESSOR_PRELAUNCH, default 1). The outer
compressor only reads hidden_states + attn_metadata, so it never needed to wait
for the kv_a_proj -> LN -> split chain. It was already pre-launched on CSA
layers as a side effect of the indexer overlap; the gate is now independent of
the indexer, so HCA layers get it too. On the trace the HCA compressor's first
kernel moves from +44.8 us to +14.4 us into the prologue window, which lets it
run underneath the kv_a_proj GEMM instead of losing the SM race to q_b_proj.
Fused KV kernel hoist (TRTLLM_MLA_KV_NORM_HOIST, now default 1). The July
measurement that made this off-by-default was taken WITHOUT the pre-launch and
read +1.55% on p50 device step time. With the pre-launch in place the ordering
reverses. Modes 2-6 explore the remaining placements: paired with q_a_layernorm
with an immediate join (2) or a deferred one (3), queued behind the compressor
on its stream (4), fired inside _q_branch alongside q_b_layernorm (5), and (6)
mode 5 with the join delegated to the attention backend so
_deepseek_v4_local_to_global_kernel is no longer ordered behind a KV kernel it
does not read.
Measured on GB200 / poly, batch 32, 4 layers, MXFP4 ckpt, layer-wise benchmarks.
Attention prologue per iteration, MoE excluded by windowing:
MTP0 upstream 328.8 us -> 306.2 us (-22.6, -6.9%)
MTP3 upstream 431.3 us -> 395.9 us (-35.3, -8.2%)
The KV placement itself contributes little: modes 1-6 land within ~1 us of each
other, because the prologue is a serial chain whose 68 us is 44% weight-
streaming GEMM (q_b_proj alone is 100 MB of FP8 weights at 66% of HBM roofline)
and the 3 us KV kernel already fits in its shadow. The win above is almost
entirely the compressor pre-launch.
CAVEAT: those are single-sample-per-arm numbers. Repeat runs on this cluster
drifted 5-20 us over a few hours on unchanged code, so differences below ~8 us
are not resolvable and arms must be compared within one submission batch.
(cherry picked from commit a8cb7c1)
The fusion work carried a kill switch per step plus a six-way placement matrix
for the hoisted KV kernel, all of it scaffolding for the A/B runs. Fix the
behaviour to what those runs settled on and delete the rest:
TRTLLM_DISABLE_FUSED_KV_NORM, TRTLLM_FUSED_KV_NORM_CTX_ONLY,
TRTLLM_DISABLE_FUSED_Q_ROPE, TRTLLM_DISABLE_FUSED_Q_FP8_QUANT,
TRTLLM_DISABLE_PRECOMPUTED_CU_SEQLENS, TRTLLM_MLA_COMPRESSOR_PRELAUNCH
-> always on
TRTLLM_MLA_KVNORM_GEN_ROWS_PER_BLOCK -> constexpr 4
TRTLLM_MLA_KV_NORM_HOIST -> mode 1 only (aux_stream, after q_a_layernorm)
Modes 0 and 2-6 go with their support code: _q_a_layernorm_with_fused_kv_norm,
_make_fused_kv_norm_gen_launch, _pending_kv_launch, _kv_join_delegated, the
one-shot hoist diagnostic, and the delegated join in the attention backend.
The three hoist helpers collapse into one _launch_fused_kv_norm_gen.
No functional change at the default settings.
Signed-off-by: Zhenhuan Chen <zhenhuanc@nvidia.com>
dc3671
force-pushed
the
user/zhenhuanc/dsv4-kv-norm-fusion-main
branch
from
August 5, 2026 03:04
a8cb7c1 to
a771f0f
Compare
dc3671
changed the title
The fused kv-norm work had reached into `applyMLARopeAndAssignQKVKernelOptContext`
and `applyMLARopeAndAssignQKVKernelGeneration` -- a `kFuseKvNorm` region, a
`kSkipKv` template arg and a `q_rope_done` early-out -- so DeepSeek-V3/V3.2, which
run those same kernels, carried the diff without using it.
DSv4 now uses its own kernels end to end, and both shared kernels go back to
their upstream form (`OptContext` byte for byte; `Generation` keeps only the
`precomputed_cu_seqlens` / `precomputed_fmha_scheduler` flags from the scheduler
prologue rehome). `q_rope_done` is gone from MlaParams, the thop signature, the
nanobind binding, `AttentionForwardArgs` and the fallback backend.
Reaching that meant covering the one case that still fell back: a mixed batch.
Context rows take positions from a ragged token cumsum and generation rows from a
uniform query length, so `_fused_q_rope_specs` now returns one spec per phase and
the Q-norm kernel is launched once per spec over disjoint row ranges of the same
output buffers. No kernel change was needed -- the op already derives its row
count from the input and reads output strides off the tensors.
Two consequences worth stating:
* the kv-norm fusion and the Q RoPE fold are now coupled. The KV kernels hand
the un-fused RoPE kernels the RAW latent, whose Q region would read it
un-normalized, so enabling one without the other is a correctness bug rather
than a slower path. Both are resolved together, ahead of the Q branch.
* the fused KV kernel is launched alongside q_a_layernorm instead of after it,
with the join left where it was, in front of FMHA.
The eight debug/experiment env gates are dropped with their support code; the
behaviour is fixed at what the A/B runs settled on.
Test Coverage: 16 sparse-MLA backend tests; a new MLA-level dispatch suite
(10 tests) that pins which path each batch shape takes -- the backend suites
build a bf16 KV cache and so never enter a fused path at all; GSM8K on 2x GB200
at 96.171 / 96.133 / 96.285. An aggregated nsys census confirms zero instances of
either shared kernel while the DSv4 kernels run.
Signed-off-by: Zhenhuan Chen <zhenhuanc@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.
Description
Fuses the DeepSeek-V4 MLA prologue so the per-layer kernel chain shrinks in both phases.
kv_a_layernormfolds into the KV path — a warp owns a whole latent row, so the RMS reduction is a warp shuffle. Context and generation each get a dedicated kernel (mlaKvNormRopeQuantContextKernel/...GenerationKernel), replacing the standalone RMSNorm plus the latent concat.deepseekV4QNormFusedKernel, which already had the data. That region was 1024 of the kernel's 1161 block rows.q_b_layernorm, soapplyMLARopeAndAssignQKVKernel{Generation,OptContext}is no longer launched on the DSv4 path.Every step is gated by an env switch, so the upstream path is still reachable on this branch.
Perf — DeepSeek-V4-Pro, GB200, MTP3, all four runs in one submission batch, nsys off:
Numbers come from
bench-trtllm-disaggprocess_data/get_gen_only_perf.pyandget_ctx_throughput.py.Test Coverage
tests/unittest/_torch/modules/test_mla_registry.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.pyaccuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracyon 2x GB200 — 96.17-96.63 across the fusion steps, against 96.25-96.51 for un-fused controls.No new unit test: these are kernel fusions behind existing call paths, and the tests above already exercise both phases plus mixed batches. Kernel selection was verified from nsys traces — the fused kernels appear and the kernels they replace are absent.
Not covered:
cached_offset > 0(chunked prefill / block reuse) for the context kernels.PR Checklist