[None][perf] fp8 block scale quant fusion in SM90 Cutlass MoE#16849
Open
amukkara wants to merge 3 commits into
Open
[None][perf] fp8 block scale quant fusion in SM90 Cutlass MoE#16849amukkara wants to merge 3 commits into
amukkara wants to merge 3 commits into
Conversation
Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
Signed-off-by: Anurag Mukkara <134339030+amukkara@users.noreply.github.com>
amukkara
changed the title
amukkara
marked this pull request as ready for review
July 25, 2026 01:23
amukkara
changed the title
Contributor
📝 WalkthroughWalkthroughAdds DeepSeek FP8 block-scale activation fusion for SM90 MoE execution, including prequantized GEMM runner queries, fused FP8 and scale-writing kernels, workspace and output routing changes, and an SM90 comparison test against unfused and BF16 reference paths. ChangesFP8 block-scale MoE fusion
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: 🚥 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}
Contributor
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
Comment on lines
+2725
to
+2734
| if constexpr (!IsNVFP4 && !IsMXFP8) | ||
| { | ||
| if (fp8_block_scale_out.fp8_out != nullptr) | ||
| { | ||
| return static_cast<KernelFnPtr>(&doActivationKernel<T, GemmOutputType, ScaleBiasType, | ||
| GLUAdaptor<cutlass::epilogue::thread::SiLu>, | ||
| TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, | ||
| num_rows_per_cta_v, false, true>); | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Fused pre-FC2 dispatch ignores activation_params.activation_type.
When fp8_block_scale_out.fp8_out != nullptr the kernel is unconditionally instantiated with GLUAdaptor<SiLu>. A block-scale MoE configured with Geglu, SwigluBias, or a non-gated activation would silently compute SwiGLU instead of erroring or falling back. The comment states only Swiglu reaches here, so make that an enforced precondition.
🐛 Proposed guard
if (fp8_block_scale_out.fp8_out != nullptr)
{
+ TLLM_CHECK_WITH_INFO(activation_type.activation_type == ActivationType::Swiglu,
+ "Fused FP8 block-scale activation quant only supports Swiglu");
return static_cast<KernelFnPtr>(&doActivationKernel<T, GemmOutputType, ScaleBiasType,📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if constexpr (!IsNVFP4 && !IsMXFP8) | |
| { | |
| if (fp8_block_scale_out.fp8_out != nullptr) | |
| { | |
| return static_cast<KernelFnPtr>(&doActivationKernel<T, GemmOutputType, ScaleBiasType, | |
| GLUAdaptor<cutlass::epilogue::thread::SiLu>, | |
| TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, | |
| num_rows_per_cta_v, false, true>); | |
| } | |
| } | |
| if constexpr (!IsNVFP4 && !IsMXFP8) | |
| { | |
| if (fp8_block_scale_out.fp8_out != nullptr) | |
| { | |
| TLLM_CHECK_WITH_INFO(activation_type.activation_type == ActivationType::Swiglu, | |
| "Fused FP8 block-scale activation quant only supports Swiglu"); | |
| return static_cast<KernelFnPtr>(&doActivationKernel<T, GemmOutputType, ScaleBiasType, | |
| GLUAdaptor<cutlass::epilogue::thread::SiLu>, | |
| TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE, | |
| num_rows_per_cta_v, false, true>); | |
| } | |
| } |
🤖 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 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.
Comment on lines
+4581
to
+4607
| // Fuse the pre-FC1 quant into the row expansion when the runner is prequantized: write fp8 A + 1x128 | ||
| // scales into permuted_data_ in the layout BlockScaleFC1 reads. | ||
| bool fused_fc1_expand_done = false; | ||
| if constexpr (std::is_same_v<InputType, __nv_bfloat16>) | ||
| { | ||
| if (use_fused_block_scale_quant) | ||
| { | ||
| int64_t const fc1_scale_leading_dim = blockscale_gemm_runner->getActScaleLeadingDim(); | ||
| auto* fp8_a = reinterpret_cast<__nv_fp8_e4m3*>(permuted_data_); | ||
| auto* fp8_scales = reinterpret_cast<float*>( | ||
| reinterpret_cast<char*>(permuted_data_) + fp8BlockScaleByteOffset(expanded_num_rows, hidden_size)); | ||
| expandInputRowsFp8BlockScaleKernelLauncher(input_activations, | ||
| Fp8BlockScaleActOutput{fp8_a, fp8_scales, fc1_scale_leading_dim}, token_topk_unpermuted_scales, | ||
| permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, | ||
| experts_per_token, num_experts_per_node, expert_first_token_offset_, stream); | ||
| fused_fc1_expand_done = true; | ||
| } | ||
| } | ||
| if (!fused_fc1_expand_done) | ||
| { | ||
| // Expand input and maybe apply prequant scale for AWQ | ||
| expandInputRowsKernelLauncher(input_activations, gemm1_input_expand, token_topk_unpermuted_scales, | ||
| permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, experts_per_token, | ||
| num_experts_per_node, quant_params, use_per_expert_act_scale, expert_first_token_offset_, | ||
| fc1_fp4_act_scale_, input_sf, swizzled_input_sf, | ||
| (use_w4afp8 && !use_fp8_input) ? quant_params.groupwise.fc1.act_scales : nullptr, stream); | ||
| } |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Fused FC1 quant can be silently skipped while FC2/FC1 still assume prequantized input.
use_fused_block_scale_quant is derived from the runner (Line 4364) but the fused expand is gated on std::is_same_v<InputType, __nv_bfloat16>. For any other InputType (e.g. the fp8-input path referenced at Line 4606) fused_fc1_expand_done stays false, expandInputRowsKernelLauncher writes plain activations into permuted_data_, yet BlockScaleFC1 → runBlockScaleMoeGemm still reinterprets that buffer as fp8 A plus 1x128 scales at fp8BlockScaleByteOffset(...). Result: garbage FC1 output and reads past the activation region, with no diagnostic. Fail loudly (or force the unfused runner) instead.
🐛 Proposed fix
bool fused_fc1_expand_done = false;
if constexpr (std::is_same_v<InputType, __nv_bfloat16>)
{
if (use_fused_block_scale_quant)
{
@@
fused_fc1_expand_done = true;
}
}
+ TLLM_CHECK_WITH_INFO(!use_fused_block_scale_quant || fused_fc1_expand_done,
+ "Fused FP8 block-scale activation quant requires a BF16 input activation type");
if (!fused_fc1_expand_done)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| // Fuse the pre-FC1 quant into the row expansion when the runner is prequantized: write fp8 A + 1x128 | |
| // scales into permuted_data_ in the layout BlockScaleFC1 reads. | |
| bool fused_fc1_expand_done = false; | |
| if constexpr (std::is_same_v<InputType, __nv_bfloat16>) | |
| { | |
| if (use_fused_block_scale_quant) | |
| { | |
| int64_t const fc1_scale_leading_dim = blockscale_gemm_runner->getActScaleLeadingDim(); | |
| auto* fp8_a = reinterpret_cast<__nv_fp8_e4m3*>(permuted_data_); | |
| auto* fp8_scales = reinterpret_cast<float*>( | |
| reinterpret_cast<char*>(permuted_data_) + fp8BlockScaleByteOffset(expanded_num_rows, hidden_size)); | |
| expandInputRowsFp8BlockScaleKernelLauncher(input_activations, | |
| Fp8BlockScaleActOutput{fp8_a, fp8_scales, fc1_scale_leading_dim}, token_topk_unpermuted_scales, | |
| permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, | |
| experts_per_token, num_experts_per_node, expert_first_token_offset_, stream); | |
| fused_fc1_expand_done = true; | |
| } | |
| } | |
| if (!fused_fc1_expand_done) | |
| { | |
| // Expand input and maybe apply prequant scale for AWQ | |
| expandInputRowsKernelLauncher(input_activations, gemm1_input_expand, token_topk_unpermuted_scales, | |
| permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, experts_per_token, | |
| num_experts_per_node, quant_params, use_per_expert_act_scale, expert_first_token_offset_, | |
| fc1_fp4_act_scale_, input_sf, swizzled_input_sf, | |
| (use_w4afp8 && !use_fp8_input) ? quant_params.groupwise.fc1.act_scales : nullptr, stream); | |
| } | |
| // Fuse the pre-FC1 quant into the row expansion when the runner is prequantized: write fp8 A + 1x128 | |
| // scales into permuted_data_ in the layout BlockScaleFC1 reads. | |
| bool fused_fc1_expand_done = false; | |
| if constexpr (std::is_same_v<InputType, __nv_bfloat16>) | |
| { | |
| if (use_fused_block_scale_quant) | |
| { | |
| int64_t const fc1_scale_leading_dim = blockscale_gemm_runner->getActScaleLeadingDim(); | |
| auto* fp8_a = reinterpret_cast<__nv_fp8_e4m3*>(permuted_data_); | |
| auto* fp8_scales = reinterpret_cast<float*>( | |
| reinterpret_cast<char*>(permuted_data_) + fp8BlockScaleByteOffset(expanded_num_rows, hidden_size)); | |
| expandInputRowsFp8BlockScaleKernelLauncher(input_activations, | |
| Fp8BlockScaleActOutput{fp8_a, fp8_scales, fc1_scale_leading_dim}, token_topk_unpermuted_scales, | |
| permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, | |
| experts_per_token, num_experts_per_node, expert_first_token_offset_, stream); | |
| fused_fc1_expand_done = true; | |
| } | |
| } | |
| TLLM_CHECK_WITH_INFO(!use_fused_block_scale_quant || fused_fc1_expand_done, | |
| "Fused FP8 block-scale activation quant requires a BF16 input activation type"); | |
| if (!fused_fc1_expand_done) | |
| { | |
| // Expand input and maybe apply prequant scale for AWQ | |
| expandInputRowsKernelLauncher(input_activations, gemm1_input_expand, token_topk_unpermuted_scales, | |
| permuted_token_final_scales_, permuted_row_to_unpermuted_row_, num_rows, hidden_size, experts_per_token, | |
| num_experts_per_node, quant_params, use_per_expert_act_scale, expert_first_token_offset_, | |
| fc1_fp4_act_scale_, input_sf, swizzled_input_sf, | |
| (use_w4afp8 && !use_fp8_input) ? quant_params.groupwise.fc1.act_scales : nullptr, stream); | |
| } |
🤖 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 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.
Comment on lines
+146
to
+154
| # Primary check: the fused quant reproduces the standalone quant bit-for-bit (same bf16 activation, | ||
| # same amax, same fp8 rounding, same GEMM), so outputs should be ~identical. | ||
| torch.testing.assert_close( | ||
| out_fused, | ||
| out_unfused, | ||
| rtol=5e-3, | ||
| atol=5e-3, | ||
| msg="Fused activation quant diverges from the unfused standalone path", | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Tolerances contradict the "bit-for-bit" claim.
The comment (and the module docstring at Lines 24-25) states the fused path is numerically identical, but the assertion allows 5e-3. If equality really is exact, use rtol=0, atol=0 so a numerical regression cannot hide; otherwise reword the claim to "closely matches".
🤖 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 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Bug Fixes
Tests
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_fp8_block_scale_fusion.pycompares MoE module outputs with and w/o fusion.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.