User/brb/m3 decode perf clean - #17441
Draft
brb-nv wants to merge 38 commits into
Draft
Conversation
…ogram top-k
Rows wider than 128 blocks fell to a kernel that gave each row one warp,
which walked the row 32 blocks at a time maintaining a 16-deep sorted
register array. Its cost scales with the row width times the insertion
depth, so it degrades as context grows.
Replace it with a histogram select: one CTA per row, an fp16 histogram
as a first cut, then up to three exact fp32 refinement steps over
whatever landed in the boundary bin, finished by a sort over the staged
boundary candidates. The algorithm is TRT-LLM's own indexerTopK
(topKPerRowJob), which the MSA port in 3rdparty vendors as well; taking
it from the in-tree original avoids carrying a fork of our own code.
Four changes were needed to fit this op's contract:
- The row is read through blockStride, so there is no transpose kernel
and no workspace; blockStride == 1 still gets float4 loads.
- rowEnd is per row (nValidBlocks[query]), so a mixed-length decode
batch needs no -inf padding.
- The boundary sort breaks ties towards the smaller block index. The
vendored version ranks ties by staging slot, which is atomic-order
dependent; the selector's contract is torch.topk's tie-break.
- Selected blocks whose forced score is -inf are emitted as -1.
More than 2048 bit-identical scores in one row still resolves
arbitrarily among the tied indices, which the init/local sentinels
cannot reach.
Measured on a B200 under CUDA graph replay, full rows: the histogram is
flat at ~6 us from 32 blocks up to 4096 for batches to 128 rows, where
the old kernel grew from ~10 us at 160 blocks to ~140 us at 8192. It is
ahead from 160 blocks up at every row count from 1 to 2048, so the old
kernel is dominated over the whole range it served and is removed.
The register-resident bitonic paths keep rows up to 128 blocks, which is
both the widest row they can hold and, measured, the point where they
stop winning.
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The ported Triton sparse decode kernel dequantizes the paged FP8 cache into q's dtype, so it needs bf16 q, while the fused QK-norm+RoPE producer emits e4m3. The producer succeeding is therefore what forces a standalone widening kernel onto the serial per-layer decode chain, not what avoids it. Which branch runs cannot be read off the producer, since it sits behind a long chain of runtime capability checks. Report the outcome once per dtype pair so a build can be confirmed from a serve log. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The fused QK-norm+RoPE producer emits E4M3 q, but the ported Triton sparse decode kernel derives its compute dtype from q and dequantizes the paged cache into it, so the caller materialized a widened BF16 copy first. That copy is a standalone kernel between the block selector and the decode kernel, where nothing overlaps it: ~3.5 us on every sparse layer, ~0.2 ms per decode step at TP4 on B300. Widen inside the kernel instead, immediately after the load and before anything reads q.dtype. Doing it later would leave K, V and the softmax probabilities cast to E4M3 and run the whole attention in FP8. E4M3 -> BF16 is exact, so the result is bitwise identical to the copy it replaces, which the new test asserts as an equality. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… batches The static block kernel claims every batch up to BlockKernelMaxNumTokens, so the dynamic-block kernel added after it has never run in that range. That is backwards relative to intent: the dynamic kernel gives each token its own warp instead of looping over a batch of them, and replaces the CUB block scan with a fused warp scan, both of which target exactly the small batches the static kernel is holding. At TP4 decode with a 3-token draft, the static kernel is a single-CTA launch costing ~6.1 us per sparse layer, ~0.35 ms per step, and it gates the routed expert GEMMs. Gate the choice on TRTLLM_MOE_ROUTING_PREFER_DYN_BLOCK so both kernels can be timed from one build; the flag goes away once the threshold is settled by measurement. Only the raw-scores entry point is switched. RoutingFromTopKIds.cu makes the same choice for pre-computed topK and is left alone until this is measured. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…oice The trtllm-gen routing dispatch hands every batch of <= 4 tokens to routingIndicesBlockKernel, so routingIndicesDynBlockKernel never runs in the decode regime it was written for. TRTLLM_MOE_ROUTING_PREFER_DYN_BLOCK already lets a single build take either path; this times them against each other so the threshold can be set by measurement rather than assumption. Drives the routing kernel in isolation through moe_topk_sort under CUDA graph replay, and re-runs itself once per arm because the override is read into a function-local static. Batch sizes above 4 are included as a control, where both arms must agree. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… the kernel With one call per graph, every batch size and both arms reported 6.15 us: graph.replay() costs about that much in host time, and the routing kernel finishes before the next launch arrives, so the loop measured Python rather than the GPU. Capture many calls per graph so the replay cost amortizes, and time a trivial kernel through the same path to publish the floor, flagging any row that sits near it. Also check that the override string is present in the loaded libtensorrt_llm.so, since a rebuilt but uninstalled tree puts both arms on the same branch and reads as a flat 1.00x. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…t 2x The probe was added to settle by measurement whether routingIndicesDynBlockKernel beats routingIndicesBlockKernel below 5 tokens, on the strength of the dynamic kernel's claim to a warp per token and half the barriers. It does win, but the margin decays with batch size: 1.09x at 1 token, 1.03x by 4, which is the production point of one request plus three Eagle3 draft tokens. That is 0.11 us over 57 sparse layers, ~6 us/step, against a projected 114 us. Measured through moe_topk_sort with many calls per CUDA graph, against a 1.22 us harness floor; batches above 4 held as a control and agreed to 0.01 us. Routing also turns out to cost 3.84 us intrinsically at 4 tokens while showing 6.14 us in the profile. The gap is PDL spin waiting on the gate and quantize chain ahead of it, so its apparent cost belongs to the router window rather than to this kernel. Removes the benchmark with the flag it exercised; both are recoverable from history if the routing kernels are revisited. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…plit-K The router keeps its weight in fp32 to match SGLang, and cuBLAS has no good answer for that at decode shapes. A 4x6144x128 gate becomes a 32-way split-K TF32 GEMM plus a splitKreduce, preceded by a standalone cast of the bf16 hidden states to fp32 because that is the activation cuBLAS wants: three kernels and ~11.6 us per sparse layer, all of it serial ahead of the routed expert GEMMs. Neither existing small-router kernel helps. dsv3_router_gemm_op hard-codes 256 experts, so M3's 128 would silently reach its cublas fallback, and both it and tinygemm2 require a bf16 weight, which is the one thing an fp32 router cannot give up. Add a GEMV shaped for the decode window: one CTA per expert reducing over the hidden dimension with plain fp32 FMA, widening the bf16 activation in-register. That folds the cast in, leaves no partials to reduce, and keeps the operands off the tensor cores, which is where the TF32 rounding was coming from. Batches above 32 tokens keep F.linear, so prefill is untouched. Numerics move, toward the reference rather than away: the tests hold the result against an fp64 reference and require it to be at least as close as the TF32 path it replaces. This wants an eval, not a bitwise A/B. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Times the GEMV against the full cuBLAS sequence it replaces, cast included, since the cast only exists to give the fp32 weight an fp32 activation. Reports both the TF32 and true-fp32 cuBLAS variants, because the profile shows cuBLAS picking TF32 in production, and converts the per-call delta into us/step at 57 sparse layers. Many calls per graph, with a trivial kernel timed through the same path as the floor: a single-call graph measures the host cost of graph.replay() rather than the kernel. --tune sweeps BLOCK_K and num_warps for choosing the launch config. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The 32-token cap was a guess and it was wrong in the direction that matters: on B200 at 128 experts / 6144 hidden the GEMV runs 2.8x at 1 token, 2.0x at 4 and 1.2x at 8, but loses at 16 and is 2.9x slower at 32. Left as it was, batches of 16 or 32 would have regressed. Cap the window at 8. The tuning sweep is monotonic in BLOCK_K up to the largest size tried, which is the signature of a latency-bound kernel at one CTA per expert: what helps is loads in flight per thread. Raising the register tile to 8192 elements puts a 4-token batch at BLOCK_K=2048, worth 4.26 -> 3.12 us, and 2048 divides 6144 exactly. Extend the sweep past 2048 since it had not turned over yet. Narrow the accuracy tests onto the window the kernel now serves, so they cover the GEMV rather than the fallback. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The BLOCK_K sweep never turned over: 2048 beat 1024, 4096 beat 2048, and 8192 beat 4096 at 2.53 us for a 4-token batch. That is the shape of a latency-bound kernel. With one CTA per expert there are fewer CTAs than SMs, so occupancy is already capped and the only lever left is loads in flight per thread; the largest block wins because it issues the whole reduction at once. Raise the register tile to 32768 elements so a 4-token batch takes all 6144 in a single iteration, and clamp BLOCK_K at the hidden dimension, past which the tail only masks off. Production 4-token gate: 8.50 us of cuBLAS down to 2.53, ~341 us/step across 57 sparse layers. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…ing sweep The sweep bypasses the token-count cap to reach the kernel directly, so it is the only way to see what the GEMV can do above the window it currently serves. Printing the baseline for the same token count makes it answer the question that sets the cap, rather than just ranking configs against each other. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…a bad block size The 8-token cap came from measurements taken with a register tile that starved BLOCK_K exactly as the batch grew: 16 tokens ran at BLOCK_K=256 and 32 at 128, the latter below the 256-thread block, so most threads idled. Those configs said the GEMV lost at 16. It does not. Swept against the cuBLAS baseline at the shipped tile, the GEMV runs 3.4x at 4 tokens, 1.6x at 12 and 1.5x at 16, turning over only by 32 at 0.85x. The tile heuristic already picks the sweep's best BLOCK_K at every one of those points; only the warp count was off, and a wide block wants 8 warps where a narrow one measured faster with 4. The turn at 32 is the register tile, not the approach: the activation block is BLOCK_M x BLOCK_K, so BLOCK_K eventually shrinks faster than the extra rows pay for. Streaming K per token, as dsv3RouterGemm does, would push the crossover out further and is the move if wider batches ever matter. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…ar SF layout The fused AllReduce+residual+RMSNorm+NVFP4-quant epilogue only ever emitted block scale factors in the swizzled 128x4 layout, which the CUTLASS FP4 GEMMs want. The trtllm-gen MoE kernels want the linear row-major layout instead and assert on a swizzled Fp4QuantizedTensor, so a MoE input could not be quantized in an upstream epilogue and had to pay for its own quantize kernel. The AR fusion kernel already honours a layout field and its offset helper already has a linear branch; only the host side hardcoded swizzled. Expose that through a new fusion op rather than a layout flag on four op schemas, so no existing caller changes and the UB and MNNVL paths, which allow-list the ops they support, reject it rather than silently returning swizzled scales. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…ce epilogue Between the attention o_proj AllReduce and the first routed-expert GEMM sits a serial window on the main stream, every sparse layer, and one of the kernels in it is the standalone NVFP4 quantize of the MoE input. The AllReduce that opens the window already computes the bf16 RMSNorm result that quantize reads, so it can emit the quantized activation from the same pass. Uses the OUT_ variant because the router gate and the shared expert both still need the bf16 norm. The fold is bitwise-identical: the epilogue quantizes from the same bf16 values it writes to norm_out, through the same conversion the standalone kernel uses. It stands down when the experts are not NVFP4, when NVFP4_AWQ needs a per-channel pre-quant scale the epilogue cannot apply, or when the MoE would pad the activation past the hidden size. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The point of the linear-SF epilogue is to stand in for a standalone quantize of the norm output, so the test asserts that substitution directly: the fp4 data and scale factors must equal fp4_quantize on the epilogue's own bf16 norm output exactly, and everything that does not depend on scale factor layout must match the swizzled epilogue exactly. TRTLLMGenFusedMoE.forward_impl also opened with an unconditional x.dtype == bfloat16 assert while typing x as Union[Tensor, Fp4QuantizedTensor], which has no dtype, so a pre-quantized activation could never reach this backend's quantize_input despite the branch there for it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
An exception raised in an MPI worker has to pickle to get back to the parent, and several common ones cannot. AttributeError has retained the object it was raised on since 3.11, so anything raised off torch.ops carries an unpicklable _Ops with it and the executor reports a pickling error in place of the actual failure. Send the formatted traceback instead. Also compare the FP4 and scale-factor outputs byte-wise rather than through torch.testing.assert_close, which is built around closeness on ordinary floating point types; what this test wants from these two is identity. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
These worker functions are shipped to the ranks by value, and cloudpickle pulls in any sys.modules entry under a referenced module whose name components all appear in the function's co_names. torch.ops is such an entry, and its type subclasses ModuleType rather than being it, so it misses cloudpickle's module reducer and fails as a plain unpicklable object. Reading it from a nested helper, as the other workers in this file already do, keeps the name off the outer code object. Drop the reporting wrapper added while diagnosing this: the failure was in the parent, not the rank, so it never applied. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…d it Median consumer-start minus producer-end on the main stream was -64 ns for the QKV GEMM into the fused producer and for the index-score kernel into the block selector, against roughly -400 ns on the edges that already overlap. Both consumers now launch with programmatic stream serialization and wait immediately before their first dependent load, so their index arithmetic runs in the producer's tail. Worth about 0.35 us per edge across 57 sparse layers. Behaviour is unchanged: cudaGridDependencySynchronize is a no-op when the kernel was not launched with PDL, and the launches follow getEnvEnablePDL as the rest of the codebase does. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The scorer read seq_lens before griddepcontrol.wait. num_blocks derived from it bounds an unmasked block_table load and an unmasked score store, so a stale length is an out-of-bounds write, and seq_lens is patched every step by on_update_kv_lens. That was safe only because nothing between that patch and this kernel writes it. Move the load past the wait and leave the barrier init and descriptor prefetch ahead of it as the work PDL overlaps. Also drop the launch_dependents at kernel entry. It is not a visibility bug, since a consumer's wait blocks until this grid completes, but the only consumer is the block selector, which reduces over every block this grid writes and so cannot start early. Now that the selector is PDL-launched, triggering at entry only left it resident and spinning in its own wait for this kernel's whole lifetime. The implicit trigger as the CTAs exit still gives it the launch overlap. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…-bound Items 5 and 6 want the qkv and o_proj MXFP8 quantizes folded into their producers' epilogues, together worth ~330 us/step. For qkv that means writing an MXFP8 epilogue into the AllReduce fusion kernel, which has no MXFP8 path today, so the premise is worth checking first: at 4 tokens the qkv quantize reads 48 KB yet the profile charges it ~5 us per call, which would make it per-launch rather than per-byte. Time it against the harness floor and against a plain e4m3 cast over the same bytes, so a launch-bound result argues for the fusion and a slow-kernel result argues for tuning the kernel instead. Also drop a getattr fallback for has_nvfp4 in the routed-expert input scale resolution. Every MoE exposes the property, from the base class or from ConfigurableMoE's delegating override, so the default could only have masked a non-MoE experts object by silently disabling the fold. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The benchmark asked whether the quantize is launch-bound at one operating point, which is the wrong question: it invites tuning the kernel when the real case for folding it into the producer epilogue is that fusion removes an HBM round-trip of the activation at every batch size, not just a launch at small ones. Sweep 1 to 4096 tokens and report achieved bandwidth beside the launch floor, so the small-batch and large-batch ends can be read separately. Rename accordingly, since the question is no longer about small batches. Also record in the tracking doc which items generalize past concurrency 1. Items 1, 3, 5 and 6 remove a round-trip and hold up as batch grows; item 2 falls back above 16 tokens by construction and item 7's PDL edges are worth a fixed amount per layer, so neither should be defended on throughput configurations. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
torch.ops.trtllm only exists once the TensorRT-LLM extension module is imported, so the availability guard reported the op as missing on builds that do contain it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…enchmark The swizzled layout pads rows up to 128 and sizes the grid from the padded count, so a one-token quantize still launches 128 blocks and writes a full 128 rows of scale-factor padding. That is indistinguishable from launch-bound behavior in a token sweep alone, yet it points at a different fix. Add a linear-layout column, which skips row padding, and sample 128 and 129 tokens, which straddle the padding boundary. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
invokeMxFP8Quantization laid its grid out over rows alone, so a one-token quantize put the whole row on a single block: 147 of 148 SMs sat idle and the kernel ran at the latency of that block's own loads. Measured on B200 at one token, MiniMax-M3's qkv shape costs 1.30 us above the launch floor against 0.11 us for a plain cast over the same bytes, and o_proj costs 0.62 us. The two differ by 2x under identical launch conditions, which is per-block work rather than launch overhead. A padding-row probe (linear layout, and 128 vs 129 tokens) ruled out the scale-factor padding as the cause. Split the column dimension across gridDim.y and narrow the block when the row count cannot fill the device, keeping the total block count at the budget the row grid used on its own. Callers that leave gridDim.y at 1, which is both FP4 launchers, get the walk they had before. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…he device The first cut split columns unconditionally and let padding rows follow the split, which regressed both ends of the batch range on B200. Splitting once the rows already fill the device cost a 4096-row qkv quantize 20.9 -> 27.3 us, and fanning 128 padding rows across 12 column blocks launched 1536 blocks where 1524 only zero scale factors, widening the swizzled-vs-linear gap at one token from 0.18 to 1.56 us. o_proj was unaffected in both cases because its narrower rows leave it with a single column block, which isolates the cause. Split only when the row count falls short of the SM count, and leave padding rows to the first column block, where they walk the row as before. The real data path keeps the gain that motivated the change: one-token qkv in the linear layout went 1.97 -> 1.48 us. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Narrowing blocks to win more column blocks trades the wrong way. On B200 it took the one-token qkv data path from 1.97 to 1.49 us, measured in the linear layout, but padding rows walk the same width and got slower by more: the swizzled cost rose from 2.15 to 2.32 us. o_proj isolates it, since its linear time never moved across the change while its swizzled time lost 0.26 us. Split at the width the launcher already chose. Padding then behaves exactly as it did before the split existed, and a row that underfills the device still gets more than one block. Shapes whose rows already cover the SMs, and shapes narrow enough to need a single column block, are left untouched. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Three variants of splitting the quantize grid over columns never converted the data-path gain into an end-to-end win. The row itself did get faster, from 1.97 to 1.48 us at one token measured in the linear layout, but padding rows walk the same block width, so buying column parallelism by narrowing blocks costs more on the padding than it gains on the row. The width-preserving variant is worth ~0.20 us at one token for qkv and nothing elsewhere, which is inside the run-to-run noise: rows whose launch config is provably unchanged moved by 0.4 us in the same comparison. Restore the kernel exactly. Keep the two invariance tests, which hold on the original grid, and record the measurements so the question is not reopened. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…ed activation MXFP8LinearMethod calls mxfp8_quantize itself before every GEMM. At decode that is a standalone launch reading an activation the producing kernel held in registers a moment earlier, and the microbenchmark puts it at 129 us/step for qkv and 88 us/step for o_proj at small batch. Grid tuning did not move it, so the cost has to come out through the producer epilogue. MXFP8QuantizedTensor is the carrier for that, mirroring Fp4QuantizedTensor. It only lets the linear layer skip its own quantize; the producer folds come next and are what actually removes the launch. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
A Triton kernel that already holds an activation can emit the E4M3 values and UE8M0 block scales itself instead of leaving them for a standalone mxfp8_quantize launch. mxfp8_quantize_tile is that epilogue, and the M3 sparse decode merge kernel is its first caller. The numerics follow cvt_warp_fp16_to_mxfp8 exactly, including emitting rcp.approx.ftz.f32 rather than dividing by 448: the two differ by an ulp, and the round-up to a power of two turns that into a whole exponent whenever the block amax is 1.75 * 2^k. The scale round-up is integer bit arithmetic for the same reason. The tests pin all of this against the op bitwise. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
… MXFP8 The merge kernel writes the attention output that o_proj then re-reads through a standalone mxfp8_quantize, worth 88 us/step across the sparse layers at small batch. One program owns one head, so with head_dim a multiple of 32 its columns are whole scale-factor blocks and the quantize folds into the epilogue with no cross-program reduction. It quantizes the narrowed value rather than the fp32 accumulator, because the standalone op would have read what the store leaves in memory and the two forms are used interchangeably layer by layer. The buffers are optional and cover the whole o_proj input, so a caller running only the generation suffix of a mixed batch must leave them out; the rows this kernel skips would otherwise reach the GEMM unquantized. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
…iton MXFP8 tile Scale byte 0 stands for 2^-127, which is denormal, so the reference takes rcp.approx.ftz of it and gets infinity rather than 2^127: the block saturates. Only a vanishing block amax reaches this, but the fused and standalone forms are used interchangeably across layers, so they have to agree there too. Also reach the two module constants through tl.constexpr mirrors, which is the only way a jit'ed function can read a global. 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. The merge kernel now 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 buffer be a plain per-step allocation rather than a persistently zeroed one with a CUDA-graph-safe allocation point. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The guard moves ahead of the K/V write so a caller that got it wrong is told before the step has moved, and the cases it turns on -- no span, a mixed batch, a dense layer -- get a test each, alongside the predicate the model reads. 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.
…ndary AllReduce Each layer boundary already runs an AllReduce with a fused residual add and RMSNorm whose output feeds the next layer's qkv projection. When that projection is MXFP8 on the CUTLASS path, the boundary can hand it an activation that is already quantized. The carrier crosses a layer boundary, so the decoder layer returns it alongside the bf16 norm and the model loop threads it to the next layer. The bf16 result is still needed there: the index projection, the spec-decoding capture and the final norm all read it. Layer 0 has no boundary AllReduce and the last layer's norm feeds the LM head, so the fold covers the layers in between. The decision reads only load-time state, so it is stable across steps and needs no in-graph fallback.
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.