[None][perf] MXFP8 carrier and quantize folds - #17453
Draft
brb-nv wants to merge 7 commits into
Draft
Conversation
…ed activation The MXFP8 CUTLASS path quantizes its activation in a kernel of its own, immediately after whichever kernel produced that activation and still had it in registers. Give the quant method a way to accept the quantized form directly, so a producer that can emit it does not force a second pass over the data. MXFP8QuantizedTensor is the carrier, alongside the existing Fp4QuantizedTensor. It is rejected on the dequant reference path, which needs the high-precision input, and when the block scales arrive in the linear rather than swizzled layout, which is not what the GEMM reads. Backend selection for the GEMM is unchanged; only the quantize is skipped. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
A tl.func a Triton kernel can call on a tile it already holds in registers, so it can write the MXFP8 form alongside its normal output rather than leaving a standalone quantize to re-read the result from HBM. Matches the reference kernel bit for bit, including its flush-to-zero on the reciprocal of the block scale, which is what the accompanying test pins. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… MXFP8 The merge kernel finishes flash-decoding with the attention output in registers, and the o_proj that consumes it immediately quantizes that tensor to MXFP8. Have the merge write both forms while it still holds the data. It also writes the scale-factor rows that swizzling invents past the last live row, spreading them over the live programs for a few bytes each. That is what lets the caller pass a plain per-step allocation rather than a persistently zeroed buffer with a CUDA-graph-safe allocation point. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…se decode The decode o_proj runs a standalone mxfp8_quantize over a tensor the sparse merge kernel has just held in registers, so let that kernel emit the MXFP8 form alongside the bf16 one and hand o_proj both. o_proj sits outside the attention's compile boundary, so the choice of what to feed it cannot depend on the step: a traced graph would carry one step's answer into every other. So the buffers are allocated on the static configuration alone and are always passed to o_proj, and the steps the merge kernel does not own -- prefill, a mixed batch, a padded activation -- fill them with the same standalone quantize as before, from inside the boundary. That costs a copy on exactly the steps that were never the target. Only a pure-decode step qualifies. A mixed batch leaves the context prefix to fmha_sm100, whose rows would reach the GEMM unquantized, so the fold stands down there rather than producing a partially quantized tensor. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…h range Sweeps 1 to 4096 tokens at M3's TP4 shapes with x.to(e4m3) as a same-run bandwidth reference, which is what establishes that the standalone quantize is flat to 64 tokens and therefore worth folding into a producer rather than retuning. Two invariance tests come with it, pinning that the quantize does not depend on its grid shape and that the swizzled and linear layouts agree. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
A tensor-parallel layer whose consumer wants MXFP8 pays for a separate quantize pass over the normalized activation, which at decode sizes is pure launch and bandwidth on top of a collective that already held the data in registers. RESIDUAL_RMS_NORM_OUT_QUANT_MXFP8 reuses cvt_warp_fp16_to_mxfp8 on the same value the epilogue stores to norm_out, so the result is bitwise identical to quantizing that tensor afterwards. The swizzled scale-factor layout pads rows to a multiple of 128 and the block-scaled GEMM reads the whole tile, so both fusion kernels zero the invented rows themselves rather than requiring a memset of a buffer that is mostly overwritten. A hidden size that would also pad scale columns is rejected on the host, since the epilogue writes one scale per group of four threads and cannot reach them. Also fixes an unchecked map lookup in selectStrategyLookUpTable that made any fusion op missing from mapFusionOpToIndex undefined behavior under AllReduceStrategy::AUTO. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…ndary AllReduce Each layer ends in an AllReduce+RMSNorm whose result the next layer's qkv projection immediately quantizes to MXFP8. That quantize is a kernel of its own on the serial chain between the two layers, reading back what the epilogue had in registers. The epilogue emits it instead, and the carrier is threaded down to qkv_proj through the decoder layer. Everything else in the attention forward keeps reading the bf16 norm, which is still returned: the index projection on the compatibility path, the spec-decoding capture and the final model norm all need it. Whether to fold is decided from load-time facts only -- the quant method, the hidden size, LoRA, the fused GEMM+AllReduce path -- so a traced graph cannot bake one step's answer into every other. The last layer is excluded because its boundary norm feeds the LM head, not a qkv projection. The hidden-size check wants four scale columns per swizzled tile, not one block: the epilogue writes one scale per group of four threads, so a size that pads scale columns leaves some unreachable and the op rejects it outright. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.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.
@coderabbitai summary
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.