[None][perf] fp8 block scale quant fusion in SM90 Cutlass MoE - #16849
[None][perf] fp8 block scale quant fusion in SM90 Cutlass MoE#16849amukkara wants to merge 1 commit into
Conversation
amukkara
changed the title
amukkara
changed the title
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughThe change adds prequantized activation metadata to the FP8 block-scale GEMM interface. It adds fused FP8 block-scale quantization for DeepSeek MoE activation paths and updates runner selection, workspace handling, GEMM routing, padding, and output buffers. ChangesFP8 block-scale MoE fusion
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant MoEExecution
participant ActivationKernels
participant FC1GEMM
participant SwiGLU
participant FC2GEMM
MoEExecution->>ActivationKernels: expand rows and quantize BF16 inputs
ActivationKernels->>FC1GEMM: provide packed FP8 values and block scales
FC1GEMM->>SwiGLU: produce FC1 output
SwiGLU->>FC2GEMM: provide fused FP8 activation and scales
FC2GEMM->>MoEExecution: return FC2 output
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (5)
cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h (1)
125-128: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winDocument/guard the ordering contract of
getActScaleLeadingDim().The returned value is only meaningful after
getWorkspaceSize()has populatedmax_shape_m_32_align_padded_; a premature call returns a stale/zero leading dim, which silently corrupts 1x128 scale addressing in the fused MoE path. A Doxygen note plus aTLLM_CHECK/assert on non-zero would make the misuse loud instead of silent.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h` around lines 125 - 128, Update getActScaleLeadingDim() to document that getWorkspaceSize() must run first to initialize max_shape_m_32_align_padded_, then guard the return with the established TLLM_CHECK/assert mechanism to reject a zero or uninitialized leading dimension before returning it.cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu (3)
1748-1767: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStyle: brace the loop bodies and name the quantization constants.
448.f, the128block size, and the16/15half-warp masks are magic literals, and the three#pragma unrollloops have unbraced bodies. Same pattern at Lines 1828-1829 and Lines 2473-2474.As per coding guidelines: "use Allman braces, braced control-flow bodies" and "Avoid magic literals except
0,nullptr,true, andfalse; initialize named constants instead, usingk-prefixed camelCase names".♻️ Sketch
+static constexpr float kFp8E4M3Max = 448.f; +static constexpr int kBlockScaleVecSize = 128; +static constexpr unsigned kHalfWarpMask = 0xFFFFu; @@ `#pragma` unroll for (int e = 0; e < N; ++e) + { amax = fmaxf(amax, fabsf(vals[e])); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 1748 - 1767, Update fp8BlockScaleQuantize and the corresponding quantization loops around the referenced later locations to use Allman-style braces for every loop body. Replace the magic block size, half-warp width/mask values, and FP8 maximum constant with appropriately named k-prefixed camelCase constants, while preserving the existing quantization and shuffle behavior.Source: Coding guidelines
3286-3302: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winRelying on
getWorkspaceSize()as the setter for the scale leading dim is fragile.Line 3297 calls
getWorkspaceSize(...)purely for its side effect and discards the result; if that call is ever removed/reordered,getActScaleLeadingDim()(read later at Lines 3548 and 4588) silently returns a stale value and every scale lands on the wrong row. An explicit initializer on the runner would make the dependency visible.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 3286 - 3302, Initialize the activation-scale leading dimension explicitly on the blockscale GEMM runner instead of relying on the discarded getWorkspaceSize call in the isActivationPrequantized path. Update the runner API or setup flow around getDeepSeekBlockScaleGemmRunner and getActScaleLeadingDim so the dimension is assigned from num_rows, experts_per_token, and num_experts_per_node before it is read, while preserving the existing workspace sizing behavior.
3144-3152: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTighten this env check to the shared bool helper.
env[0] == '1'also treats values like10or1abcas enabled;tensorrt_llm::common::getBoolEnv()matches the repo’s exact-1convention and keeps this flag consistent with the other env-backed booleans.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 3144 - 3152, Update useFp8BlockScaleActFusion() to use tensorrt_llm::common::getBoolEnv() for TLLM_MOE_DISABLE_FP8_BLOCK_SCALE_ACT_FUSION instead of manually checking env[0], preserving the existing SM-version gate and inversion semantics while enforcing exact-1 boolean parsing.tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py (1)
62-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage gaps and test cost worth a follow-up.
Assessment: coverage is adequate for the primary fused-vs-unfused equality claim but thin around the fusion's preconditions.
- The test cannot detect that fusion was actually exercised — an env-name typo or a runner-selection regression makes both configs identical and the assertion still passes. Consider asserting the selected runner variant (or that the two paths differ in some observable way, e.g. a debug counter).
- Only
ActivationType.Swiglu+ bf16 + single rank is covered; the fused epilogue hardcodes SwiGLU, so aGeglu/SwigluBiascase would be the useful negative test.num_experts=256plus a bf16 reference module of the same shapes allocates several GB per parametrization; a smaller expert count exercises the same kernels far more cheaply.- If this file needs to run in CI beyond default unittest collection, add it to the relevant list (e.g.
tests/integration/test_lists/test-db/*.yml).As per path instructions: "suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR".
Also applies to: 96-141
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py` around lines 62 - 63, The fusion test coverage needs follow-up: verify the selected fused and unfused runner variants so configuration regressions cannot produce a false equality pass, and add a negative Geglu or SwigluBias case for the SwiGLU-only epilogue. Reduce num_experts from 256 to a smaller representative value to lower parametrization memory cost. If this test is intended for CI outside unittest discovery, add it to the applicable tests/integration/test_lists/test-db/*.yml list; coverage is otherwise adequate for the primary equality claim.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu`:
- Around line 4581-4607: Update the FC1 expansion logic around
fused_fc1_expand_done so use_fused_block_scale_quant cannot proceed unless the
fused FP8 block-scale expansion is supported for the current InputType. For
unsupported input types, fail loudly or select the unfused FC1 runner before
expandInputRowsKernelLauncher writes plain activations, ensuring BlockScaleFC1
never reinterprets an unfused buffer as FP8 data and scales.
- Around line 2725-2734: In the fused pre-FC2 dispatch branch of the relevant
kernel-selection function, validate that activation_params.activation_type is
the supported SwiGLU/Silu-gated type before returning the
GLUAdaptor<cutlass::epilogue::thread::SiLu> kernel. For Geglu, SwigluBias,
non-gated, or any other activation, reject the path or fall back using the
surrounding dispatch behavior instead of instantiating the SwiGLU kernel
unconditionally.
In `@tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py`:
- Around line 146-154: The fused-versus-unfused assertion in the test should
enforce the documented bit-for-bit behavior by changing
torch.testing.assert_close to use zero relative and absolute tolerances. Keep
the existing comparison and message unchanged, and ensure the module
docstring/comment remains consistent with exact equality.
---
Nitpick comments:
In
`@cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h`:
- Around line 125-128: Update getActScaleLeadingDim() to document that
getWorkspaceSize() must run first to initialize max_shape_m_32_align_padded_,
then guard the return with the established TLLM_CHECK/assert mechanism to reject
a zero or uninitialized leading dimension before returning it.
In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu`:
- Around line 1748-1767: Update fp8BlockScaleQuantize and the corresponding
quantization loops around the referenced later locations to use Allman-style
braces for every loop body. Replace the magic block size, half-warp width/mask
values, and FP8 maximum constant with appropriately named k-prefixed camelCase
constants, while preserving the existing quantization and shuffle behavior.
- Around line 3286-3302: Initialize the activation-scale leading dimension
explicitly on the blockscale GEMM runner instead of relying on the discarded
getWorkspaceSize call in the isActivationPrequantized path. Update the runner
API or setup flow around getDeepSeekBlockScaleGemmRunner and
getActScaleLeadingDim so the dimension is assigned from num_rows,
experts_per_token, and num_experts_per_node before it is read, while preserving
the existing workspace sizing behavior.
- Around line 3144-3152: Update useFp8BlockScaleActFusion() to use
tensorrt_llm::common::getBoolEnv() for
TLLM_MOE_DISABLE_FP8_BLOCK_SCALE_ACT_FUSION instead of manually checking env[0],
preserving the existing SM-version gate and inversion semantics while enforcing
exact-1 boolean parsing.
In `@tests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py`:
- Around line 62-63: The fusion test coverage needs follow-up: verify the
selected fused and unfused runner variants so configuration regressions cannot
produce a false equality pass, and add a negative Geglu or SwigluBias case for
the SwiGLU-only epilogue. Reduce num_experts from 256 to a smaller
representative value to lower parametrization memory cost. If this test is
intended for CI outside unittest discovery, add it to the applicable
tests/integration/test_lists/test-db/*.yml list; coverage is otherwise adequate
for the primary equality claim.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2ea6f944-67cd-4c13-87c0-2b645e1d8197
📒 Files selected for processing (3)
cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.hcpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cutests/unittest/_torch/modules/moe/test_fp8_block_scale_fusion.py
1b93f4b to
4502f2b
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (7)
cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu (7)
1723-1735: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winName the 1x128 block size and the scale alignment.
The literals
15,16, and128appear here, and128also appears at Line 1852, Line 3538, and implicitly as>> 4/& 15insidefp8BlockScaleQuantize. Define named constants once so the layout and the quantizer cannot diverge.♻️ Proposed refactor
+// Channels per 1x128 activation scale block. +static constexpr int64_t kFp8BlockScaleBlockSize = 128; +// Alignment of the float scale region that follows the fp8 activations. +static constexpr size_t kFp8BlockScaleAlignment = 16; + static inline size_t fp8BlockScaleByteOffset(int64_t rows, int64_t cols) { - return ((static_cast<size_t>(rows) * cols * sizeof(__nv_fp8_e4m3)) + 15) / 16 * 16; + return tensorrt_llm::common::alignSize( + static_cast<size_t>(rows) * cols * sizeof(__nv_fp8_e4m3), kFp8BlockScaleAlignment); } @@ return fp8BlockScaleByteOffset(rows, cols) - + static_cast<size_t>(scale_leading_dim) * tensorrt_llm::common::ceilDiv(cols, static_cast<int64_t>(128)) + + static_cast<size_t>(scale_leading_dim) * tensorrt_llm::common::ceilDiv(cols, kFp8BlockScaleBlockSize) * sizeof(float);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 1723 - 1735, Define shared named constants for the 1x128 scale block size and 16-byte scale alignment, then replace the literals in fp8BlockScaleByteOffset, fp8BlockScaleRegionBytes, the references near the later 128 usages, and the implicit >> 4/& 15 calculations in fp8BlockScaleQuantize. Ensure all layout and quantization paths derive their block and alignment behavior from these single constants.Source: Coding guidelines
1754-1766: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBrace the new loop and
ifbodies.Lines 1755-1756, 1759-1760, 1763-1764, and 1765-1766 use unbraced bodies. The same pattern appears at Lines 1828-1829 and Lines 2473-2474. The repository style requires braces on every
if/elseand loop body.As per coding guidelines: "Use Allman brace style; always brace if/else, loop, and switch bodies".♻️ Proposed change for this block
`#pragma` unroll for (int e = 0; e < N; ++e) + { amax = fmaxf(amax, fabsf(vals[e])); + } unsigned const group_mask = 0xFFFFu << (lane & 16u); `#pragma` unroll for (int m = 1; m < 16; m <<= 1) + { amax = fmaxf(amax, __shfl_xor_sync(group_mask, amax, m, 32)); + } float const scale = 448.f / fmaxf(static_cast<float>(static_cast<__nv_bfloat16>(amax)), 1e-10f); `#pragma` unroll for (int e = 0; e < N; ++e) + { out[e] = static_cast<__nv_fp8_e4m3>(vals[e] * scale); + } if (write_scale) + { block_scales[kb * scale_leading_dim + scale_row] = 1.f / scale; + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 1754 - 1766, Apply Allman-style braces to every unbraced loop and if body in the shown block, including the loops around amax, scale reduction, and output conversion and the write_scale conditional. Make the same brace-only formatting change at the corresponding loop bodies near the other identified occurrences, without altering behavior.Source: Coding guidelines
1852-1857: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd a numeric equivalence test for the fused path.
The fused path changes the activation-quantization layout for both FC1 and FC2. A defect in the scale row index, the scale leading dim, or the half-warp amax reduction produces wrong numbers rather than a crash. The fused path also runs only on SM90, so a CI runner on another architecture never exercises it.
Add a test that forces both runner variants for the same inputs and compares the FC1 and FC2 outputs. Cover
hidden_sizeandinter_sizeat the 128 boundary and a token count that leaves a partial expert tile.I can draft that test if you want.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 1852 - 1857, Add a numeric equivalence test for the fused activation-quantization path that forces both runner variants on identical inputs, compares FC1 and FC2 outputs, and executes on SM90. Cover hidden_size and inter_size values at the 128 boundary and use a token count producing a partial expert tile; exercise the relevant MoE runner/test symbols rather than only validating kernel completion.
3535-3541: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUpdate the comment about why the scales fit.
The comment states that
outputis "bf16-sized so the scales fit above the fp8". The buffer is no longer sized that way. Lines 3294-3296 size the overlapped inputs buffer explicitly withfp8BlockScaleRegionBytes. Point the comment at that sizing, so a later change does not remove it as redundant.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 3535 - 3541, The comment above the fused pre-FC2 quantization incorrectly describes output as bf16-sized; update it to reference the explicit fp8 scale-region sizing via fp8BlockScaleRegionBytes used for the overlapped inputs buffer. Preserve the explanation that scales are stored above the fp8 data and that FC2’s standalone scale_1x128 is removed.
1749-1751: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the
laneprecondition onfp8BlockScaleQuantize.
lanemust satisfylane % 32 == laneidso thatlane & 16uselects the correct half-warp shuffle mask. The FC1 caller passesthreadIdx.xwithblockDim = (256,1,1), and the epilogue caller passesthreadIdx.ywithblockDim = (1,256,1). Both hold today only because the block is effectively 1D. Add a Doxygen note stating the requirement, so a later block-shape change does not silently corrupt the amax reduction.As per coding guidelines: "document new interfaces with Doxygen".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 1749 - 1751, Document the lane precondition in the Doxygen comments for fp8BlockScaleQuantize: lane must satisfy lane % 32 == laneid so lane & 16u selects the correct half-warp shuffle mask. Mention that callers must preserve this relationship when changing block shapes.Source: Coding guidelines
2464-2480: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winMake the silent fall-through detectable.
The branch condition includes
ACTIVATION_ELEM_PER_THREAD == 8. If an instantiation ever reaches this code withWriteFp8BlockScale == trueand a different element count, control falls to theelseat Line 2481. That store writesTvalues into a buffer the FC2 GEMM reads as packed fp8 plus 1x128 scales, and the result is silently wrong. Nothing fails at compile time or at run time.Add a host-side guard where
WriteFp8BlockScaleis selected, so the unsupported combination aborts instead of producing wrong numbers.🛡️ Proposed guard at the dispatch site (Line 2792)
if (fp8_block_scale_out.fp8_out != nullptr) { + static_assert(ACTIVATION_ELEM_PER_THREAD == 8 || !std::is_same_v<T, GemmOutputType>, + "Fused FP8 block-scale quant requires 8 channels per thread"); + TLLM_CHECK_WITH_INFO(ACTIVATION_ELEM_PER_THREAD == 8, + "Fused FP8 block-scale activation quant requires 8 channels per thread"); // Fuse pre-FC2 1x128 activation quantization.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 2464 - 2480, Add a host-side validation at the dispatch site where WriteFp8BlockScale is selected: when it is enabled, require ACTIVATION_ELEM_PER_THREAD to equal 8 and abort otherwise. Ensure the unsupported instantiation cannot reach the fallback store path in the MoE GEMM dispatch.
3284-3297: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winSeparate scale-leading-dimension initialization from workspace sizing
getActScaleLeadingDim()depends only on the padded M dimension, so the FC1/FC2 shape difference is safe. However, this call relies ongetWorkspaceSize()to initialize runner state while discarding its return value. Use an explicit initialization API or document this contract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 3284 - 3297, Separate activation-scale leading-dimension initialization from workspace sizing in the isActivationPrequantized branch. Replace the discarded getWorkspaceSize call with the runner’s explicit initialization API if available; otherwise document that getWorkspaceSize initializes the required state before getActScaleLeadingDim, while preserving the existing sizing behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu`:
- Around line 3143-3159: Store the device ID selected by
useFp8BlockScaleActFusion() in CutlassMoeFCRunner and validate it against
input.get_device() at the start of FusedMoeRunner::runMoe(). Reject execution
when they differ, before selecting the DeepSeek FP8 activation format or
workspace layout.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu`:
- Around line 1723-1735: Define shared named constants for the 1x128 scale block
size and 16-byte scale alignment, then replace the literals in
fp8BlockScaleByteOffset, fp8BlockScaleRegionBytes, the references near the later
128 usages, and the implicit >> 4/& 15 calculations in fp8BlockScaleQuantize.
Ensure all layout and quantization paths derive their block and alignment
behavior from these single constants.
- Around line 1754-1766: Apply Allman-style braces to every unbraced loop and if
body in the shown block, including the loops around amax, scale reduction, and
output conversion and the write_scale conditional. Make the same brace-only
formatting change at the corresponding loop bodies near the other identified
occurrences, without altering behavior.
- Around line 1852-1857: Add a numeric equivalence test for the fused
activation-quantization path that forces both runner variants on identical
inputs, compares FC1 and FC2 outputs, and executes on SM90. Cover hidden_size
and inter_size values at the 128 boundary and use a token count producing a
partial expert tile; exercise the relevant MoE runner/test symbols rather than
only validating kernel completion.
- Around line 3535-3541: The comment above the fused pre-FC2 quantization
incorrectly describes output as bf16-sized; update it to reference the explicit
fp8 scale-region sizing via fp8BlockScaleRegionBytes used for the overlapped
inputs buffer. Preserve the explanation that scales are stored above the fp8
data and that FC2’s standalone scale_1x128 is removed.
- Around line 1749-1751: Document the lane precondition in the Doxygen comments
for fp8BlockScaleQuantize: lane must satisfy lane % 32 == laneid so lane & 16u
selects the correct half-warp shuffle mask. Mention that callers must preserve
this relationship when changing block shapes.
- Around line 2464-2480: Add a host-side validation at the dispatch site where
WriteFp8BlockScale is selected: when it is enabled, require
ACTIVATION_ELEM_PER_THREAD to equal 8 and abort otherwise. Ensure the
unsupported instantiation cannot reach the fallback store path in the MoE GEMM
dispatch.
- Around line 3284-3297: Separate activation-scale leading-dimension
initialization from workspace sizing in the isActivationPrequantized branch.
Replace the discarded getWorkspaceSize call with the runner’s explicit
initialization API if available; otherwise document that getWorkspaceSize
initializes the required state before getActScaleLeadingDim, while preserving
the existing sizing behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f199f697-9aef-4d4e-bd36-7dd655b22fb4
📒 Files selected for processing (2)
cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.hcpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu
🚧 Files skipped from review as they are similar to previous changes (1)
- cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h
|
/bot run --disable-fail-fast |
|
PR_Github #63943 [ run ] triggered by Bot. Commit: |
|
PR_Github #63943 [ run ] completed with state
|
|
/bot run |
4502f2b to
0bcbd68
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #64408 [ run ] triggered by Bot. Commit: |
|
PR_Github #64408 [ run ] completed with state
|
4 participants
Summary
doActivationepilogue.<fp8, fp8, bf16>GEMM selection and non-SM90 fallback.Dev Engineer Review
CutlassFp8BlockScaleGemmRunnerInterfaceimplementation provides the new pure virtual methods.QA Engineer Review
No test changes.
Description
Qwen3.5-35B-A3B-FP8, H100 NVL 400W, TP=1, ISL=10K, OSL=32, BS=1.
Test Coverage
tests/unittest/_torch/modules/moe/test_moe_backend.pyExisting model accuracy tests cover e2e accuracy.
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.