[None][perf] DSV4 o_proj: fuse fp8/UE8M0 quantize into o_a proj and introduce splitk cutedsl o_b proj - #16346
Open
mingyangHao wants to merge 20 commits into
Open
Conversation
Collaborator
Author
|
/bot run |
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds split-x support to fused mHC kernels and launchers. It validates split-major inputs and accumulates x splits. It also adds Blackwell FP8 CuTe-DSL GEMM/BMM paths, packed-scale execution, dynamic MXF8 MMA helpers, and fused DeepSeek-V4 projection routing with tests. ChangesmHC split-x execution
Blackwell FP8 and DeepSeek-V4 paths
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant MLA
participant FP8BMM
participant SplitKGEMM
participant DeepGEMM
MLA->>FP8BMM: produce FP8 O_a intermediate and scales
MLA->>SplitKGEMM: execute eligible split-K O_b projection
MLA->>DeepGEMM: execute out-of-bucket O_b projection
SplitKGEMM->>MLA: return projection partials
DeepGEMM->>MLA: return projection output
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
|
PR_Github #59071 [ ] completed with state |
![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: 6
🧹 Nitpick comments (2)
tests/unittest/_torch/modules/test_mhc.py (1)
817-824: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover both supported split-x specializations.
Parameterize
x_splitsover(2, 4); XS=2 and XS=4 instantiate distinct CUDA kernels and barrier schedules.As per path instructions, assess TensorRT-LLM test coverage and provide a concrete file-level follow-up; coverage in
tests/unittest/_torch/modules/test_mhc.pyis insufficient for XS=2.🤖 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/test_mhc.py` around lines 817 - 824, Parameterize the test case around x_partials and x_reduced over x_splits values 2 and 4 so both CUDA kernel and barrier-schedule specializations execute, preserving the existing setup for each parameter. Also assess TensorRT-LLM coverage and add a concrete follow-up for the relevant file to cover XS=2 there.Source: Path instructions
tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py (1)
49-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new test and helper functions.
Add parameter and return annotations, including
pytest.MonkeyPatch,-> None, and a precisetuple[...]return for_build_dsv4_o_proj_case.As per coding guidelines, “Annotate every function.”
Also applies to: 54-54, 62-62, 100-100, 105-105, 127-135, 160-160, 182-186, 200-200, 219-226, 247-247, 272-272, 383-383, 580-580, 639-639
🤖 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/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py` at line 49, Add complete type annotations to every new test and helper function in this file, including parameter types, pytest.MonkeyPatch for monkeypatch parameters, -> None for tests without return values, and the precise tuple[...] return annotation for _build_dsv4_o_proj_case. Ensure all listed functions follow the project’s annotation conventions.Source: Coding guidelines
🤖 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/mhcKernels/mhcFusedHcKernel.cu`:
- Line 493: Update the disabled-DeepGEMM `mhcFusedHcLaunch` stub to accept the
trailing `int x_num_splits` parameter, matching the declaration and enabled
implementation used by the header and Torch binding. Keep the stub’s existing
behavior unchanged.
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 4278-4406: Update forward_fp8out and the corresponding FP8-output
launch path to validate every input tensor’s CUDA/common device, dtype, rank,
compatible shape, and required layout before creating raw pointers or launching
the kernel. Require sf_out_tensor to have shape (M, ceil_div(batch_size *
ceil_div(N, 128), 4)) in addition to its existing dtype and stride checks. Add
negative-contract coverage for invalid shapes, devices, dtypes, ranks, and
layouts in the DeepSeek V4 o-projection tests.
In `@tensorrt_llm/_torch/modules/mhc/mhc_cuda.py`:
- Around line 1064-1070: Update the x_num_splits branch in the tactic-selection
logic so the B > 32 path reuses the capability-aware fallback rather than
unconditionally selecting _FUSED_HC_FALLBACK_TACTIC_MMA. Preserve the existing B
<= 32 FMA choice while allowing unsupported SM100 MMA or hidden_size
configurations to select the FMA fallback.
In `@tensorrt_llm/_torch/modules/mla.py`:
- Around line 1133-1144: Update _should_use_fused_oproj to require BF16 output
eligibility, matching the BF16 partial allocation near the fused CuTe split-K
path; do not allow FP16 or unset self.dtype configurations to enter fusion.
In
`@tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py`:
- Around line 651-668: Update _build_dsv4_o_proj_case to accept a
use_cute_dsl_blockscaling_mm parameter, and pass False from this test’s fp8
fixture so mla.o_b_proj.use_cute_dsl_blockscaling_mm remains disabled. Preserve
the existing fixture behavior for other callers while ensuring
_should_use_fused_oproj() reaches the fused path.
- Around line 701-721: Update the test around
CuteDSLFp8BlackwellBmmRunner.kernel_cache to isolate the class-global cache per
case using monkeypatch, so fp8out_keys contains only the current compilation
before asserting its length. In the num_tokens == 16 validation, compare
out_cached against out_fused instead of out_unfused while preserving the cache
identity check.
---
Nitpick comments:
In
`@tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py`:
- Line 49: Add complete type annotations to every new test and helper function
in this file, including parameter types, pytest.MonkeyPatch for monkeypatch
parameters, -> None for tests without return values, and the precise tuple[...]
return annotation for _build_dsv4_o_proj_case. Ensure all listed functions
follow the project’s annotation conventions.
In `@tests/unittest/_torch/modules/test_mhc.py`:
- Around line 817-824: Parameterize the test case around x_partials and
x_reduced over x_splits values 2 and 4 so both CUDA kernel and barrier-schedule
specializations execute, preserving the existing setup for each parameter. Also
assess TensorRT-LLM coverage and add a concrete follow-up for the relevant file
to cover XS=2 there.
🪄 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: 5c94ea48-177a-413b-9fcf-ff9d5a16faba
📒 Files selected for processing (14)
cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuhcpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cucpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.hcpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuhcpp/tensorrt_llm/thop/mhcOp.cpptensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/dynamic_mxf8_mainloop.pytensorrt_llm/_torch/modules/mhc/hyper_connection.pytensorrt_llm/_torch/modules/mhc/mhc_cuda.pytensorrt_llm/_torch/modules/mla.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pytests/unittest/_torch/modules/test_mhc.py
Comment on lines
+701
to
+721
| expected_smem_epilogue = num_tokens <= 32 | ||
| expected_smem_row_iters = (num_tokens + 15) // 16 if expected_smem_epilogue else 1 | ||
| fp8out_keys = [ | ||
| key | ||
| for key in cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache | ||
| if key[0] == "fp8out" | ||
| and key[-2] == expected_smem_row_iters | ||
| and key[-1] == expected_smem_epilogue | ||
| ] | ||
| assert len(fp8out_keys) == 1 | ||
|
|
||
| if num_tokens == 16: | ||
| compiled_gemm = cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache[ | ||
| fp8out_keys[0] | ||
| ] | ||
| out_cached = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids) | ||
| assert ( | ||
| cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache[fp8out_keys[0]] | ||
| is compiled_gemm | ||
| ) | ||
| torch.testing.assert_close(out_cached, out_unfused, rtol=0, atol=0) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Isolate the kernel cache and compare against the fused run.
kernel_cache is class-global, making len(fp8out_keys) == 1 order-dependent. Reset it with monkeypatch for each case. The cached result should also be compared with out_fused, not the numerically distinct unfused baseline.
Proposed fix
+ monkeypatch.setattr(
+ cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner,
+ "kernel_cache",
+ {},
+ )
...
- torch.testing.assert_close(out_cached, out_unfused, rtol=0, atol=0)
+ torch.testing.assert_close(out_cached, out_fused, rtol=0, atol=0)📝 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
| expected_smem_epilogue = num_tokens <= 32 | |
| expected_smem_row_iters = (num_tokens + 15) // 16 if expected_smem_epilogue else 1 | |
| fp8out_keys = [ | |
| key | |
| for key in cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache | |
| if key[0] == "fp8out" | |
| and key[-2] == expected_smem_row_iters | |
| and key[-1] == expected_smem_epilogue | |
| ] | |
| assert len(fp8out_keys) == 1 | |
| if num_tokens == 16: | |
| compiled_gemm = cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache[ | |
| fp8out_keys[0] | |
| ] | |
| out_cached = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids) | |
| assert ( | |
| cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache[fp8out_keys[0]] | |
| is compiled_gemm | |
| ) | |
| torch.testing.assert_close(out_cached, out_unfused, rtol=0, atol=0) | |
| monkeypatch.setattr( | |
| cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner, | |
| "kernel_cache", | |
| {}, | |
| ) | |
| expected_smem_epilogue = num_tokens <= 32 | |
| expected_smem_row_iters = (num_tokens + 15) // 16 if expected_smem_epilogue else 1 | |
| fp8out_keys = [ | |
| key | |
| for key in cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache | |
| if key[0] == "fp8out" | |
| and key[-2] == expected_smem_row_iters | |
| and key[-1] == expected_smem_epilogue | |
| ] | |
| assert len(fp8out_keys) == 1 | |
| if num_tokens == 16: | |
| compiled_gemm = cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache[ | |
| fp8out_keys[0] | |
| ] | |
| out_cached = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids) | |
| assert ( | |
| cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner.kernel_cache[fp8out_keys[0]] | |
| is compiled_gemm | |
| ) | |
| torch.testing.assert_close(out_cached, out_fused, rtol=0, atol=0) |
🤖 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/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py`
around lines 701 - 721, Update the test around
CuteDSLFp8BlackwellBmmRunner.kernel_cache to isolate the class-global cache per
case using monkeypatch, so fp8out_keys contains only the current compilation
before asserting its length. In the num_tokens == 16 validation, compare
out_cached against out_fused instead of out_unfused while preserving the cache
identity check.
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #59110 [ run ] triggered by Bot. Commit: |
mingyangHao
commented
Jul 14, 2026
yuxianq
reviewed
Jul 14, 2026
yuxianq
reviewed
Jul 14, 2026
Collaborator
|
PR_Github #59110 [ run ] completed with state
|
yuxianq
reviewed
Jul 14, 2026
yuxianq
reviewed
Jul 14, 2026
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #59184 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59184 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #59317 [ run ] triggered by Bot. Commit: |
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/models/modeling_deepseekv4.py (1)
1763-1784: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPreserve the required default-off O-projection gate.
_resolve_enable_fused_hc()can resolve toTrueby default, and this value is now passed asallow_dsv4_split_output. The suppliedMLA._should_use_fused_oproj()predicate only checks theTRTLLM_DSV4_DISABLE_FUSED_OPROJkill switch, so eligible layers can enter the fused path without the PR’s requiredDSV4_FUSE_OPROJ=1gate. Enforce the positive rollout gate here or in the downstream predicate.Based on the supplied
mla.py:1249-1262contract and PR objective.🤖 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 `@tensorrt_llm/_torch/models/modeling_deepseekv4.py` around lines 1763 - 1784, Ensure the fused O-projection path remains disabled unless the explicit DSV4_FUSE_OPROJ=1 rollout gate is enabled. Update the enablement flow involving _resolve_enable_fused_hc(), DeepseekV4Attention’s allow_split_output, or the downstream MLA._should_use_fused_oproj() predicate so the positive gate is required in addition to the existing kill switch.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_deepseekv4.py`:
- Around line 1763-1784: Ensure the fused O-projection path remains disabled
unless the explicit DSV4_FUSE_OPROJ=1 rollout gate is enabled. Update the
enablement flow involving _resolve_enable_fused_hc(), DeepseekV4Attention’s
allow_split_output, or the downstream MLA._should_use_fused_oproj() predicate so
the positive gate is required in addition to the existing kill switch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 46fd0f4d-c4d4-4c42-af98-29fae0fbe3da
📒 Files selected for processing (3)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
mingyangHao
force-pushed
the
user/mingyangh/dsv4-oproj-fp8-fuse
branch
from
August 3, 2026 01:46
7554020 to
79a23ad
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (6)
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.py (4)
2591-2609: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winParameterize the existing helper instead of duplicating it.
epilog_bf16_smem_copy_and_partitionrepeatsepilog_smem_copy_and_partition(Lines 2553-2589) verbatim. The only difference is the dtype passed tosm100_utils.get_smem_store_op:cutlass.BFloat16here,self.c_dtypethere.Add an optional dtype parameter to the existing method and delete the copy.
♻️ Proposed consolidation
def epilog_smem_copy_and_partition( self, tiled_copy_t2r: cute.TiledCopy, tTR_rC: cute.Tensor, tidx: cutlass.Int32, sC: cute.Tensor, + store_dtype: Optional[Type[cutlass.Numeric]] = None, ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:copy_atom_r2s = sm100_utils.get_smem_store_op( - self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r + self.c_layout, + self.c_dtype if store_dtype is None else store_dtype, + self.acc_dtype, + tiled_copy_t2r, )Then delete
epilog_bf16_smem_copy_and_partitionand update the call site at Lines 1935-1940:) = self.epilog_bf16_smem_copy_and_partition( tiled_copy_t2r, tTR_rBf16, epi_tidx, sEpiBf16, + store_dtype=cutlass.BFloat16, )🤖 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 `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.py` around lines 2591 - 2609, Parameterize epilog_smem_copy_and_partition with an optional dtype argument, defaulting to self.c_dtype, and pass that argument to sm100_utils.get_smem_store_op. Remove epilog_bf16_smem_copy_and_partition and update its call site to invoke the existing helper with cutlass.BFloat16.
134-175: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd docstrings to the two new public DSL helpers.
ceil_to_ue8m0_deviceandstg_u8_raweach carry a docstring.abs_f32_deviceandmax_nan_f32_devicedo not. Both are module-level public names.Ruff's
Drule set is enabled for this repository, so a missing docstring on a public function can fail lint.♻️ Proposed docstrings
def abs_f32_device( value: cutlass.Float32, *, loc=None, ip=None, ) -> cutlass.Float32: + """Return the absolute value of an FP32 register.""" return cutlass.Float32(def max_nan_f32_device( lhs: cutlass.Float32, rhs: cutlass.Float32, *, loc=None, ip=None, ) -> cutlass.Float32: + """Return the NaN-propagating maximum of two FP32 registers.""" return cutlass.Float32(🤖 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 `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.py` around lines 134 - 175, Add docstrings to the public module-level helpers abs_f32_device and max_nan_f32_device, matching the concise documentation style used by ceil_to_ue8m0_device and stg_u8_raw. Leave their implementations and behavior unchanged.Source: Learnings
3103-3126: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the arguments of
wrapper_fp8sf.The sibling
wrapperat Lines 3026-3061 documents every argument.wrapper_fp8sfhas a one-line docstring only. It is a public entry point with a wider ABI, including the non-obvioussf_out_ptraddress-carrying tensor,sf_aligned_mn,n_tiles_per_group, and the twoConstexprepilogue controls.The coding guidelines require Google-style docstrings that document public function arguments.
🤖 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 `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.py` around lines 3103 - 3126, Expand the docstring for wrapper_fp8sf to Google-style documentation covering every parameter, matching the sibling wrapper’s argument descriptions where applicable. Explicitly document sf_out_ptr as the address-carrying output-scale pointer, sf_aligned_mn and n_tiles_per_group, and the fp8_smem_epilogue Constexpr controls, while preserving the existing function behavior.Source: Coding guidelines
478-507: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the one-stage trade for the BF16 staging tile explicit.
Line 480 drops one A/B stage to pay for
sEpiBf16. The trade is exact only for the enforced configuration: with a 128x128 FP8 tile, one A/B stage is128*128 + 128*128 = 32 KiB, and the BF16 staging tile is128 * 128 * 2 = 32 KiB. The validation at Lines 772-785 pins that configuration, so the arithmetic holds today.The invariant is implicit. A later change to
mma_tiler_mnor tofp8_bf16_stage_countwould silently overflow shared memory or leavenum_ab_stageat zero. Add an assertion that ties the two sizes together.♻️ Proposed guard
if self.fp8_smem_epi_mode: # Trade one A/B stage for the BF16 epilogue tile. + if self.num_ab_stage < 2: + raise ValueError( + "The cooperative FP8 epilogue needs at least two A/B " + f"stages, got {self.num_ab_stage}" + ) self.num_ab_stage -= 1🤖 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 `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.py` around lines 478 - 507, In the fp8_smem_epi_mode path near fp8_bf16_stage_count and the staged shared-memory layouts, add an assertion verifying that the shared-memory size released by exactly one A/B stage equals the BF16 epilogue staging tile size. Keep the existing num_ab_stage decrement, and ensure the assertion also protects against an invalid zero or negative remaining A/B stage under the enforced configuration.tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (1)
4515-4550: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBuild the CuTe pointers and DLPack tensors only when they are needed.
In the TVM-FFI path the launch at Lines 4591-4603 passes raw
data_ptr()values and torch tensors. It never usesa_ptr,b_ptr,a_sf_ptr,b_sf_ptr,c_cute_tensor, oralpha_cute_tensor. These objects are still built on every call, including twocute.runtime.from_dlpackconversions. This adds per-invocation host overhead on a decode-latency path.
CuteDSLFp8BlackwellBmmRunner.forward(Lines 4876-4905) andforward_fp8out(Lines 5137-5159) already guard the equivalent construction. Apply the same pattern here.Note that
alphaitself must stay outside the branch, because the TVM-FFI launch passes the torch tensor.♻️ Proposed restructuring
- a_ptr = make_ptr( - cutlass.Float8E4M3FN, - kernel_a.data_ptr(), - cute.AddressSpace.gmem, - assumed_align=16, - ) - b_ptr = make_ptr( - cutlass.Float8E4M3FN, - kernel_b.data_ptr(), - cute.AddressSpace.gmem, - assumed_align=16, - ) - a_sf_ptr = make_ptr( - cutlass.Uint32, - kernel_sfa.data_ptr(), - cute.AddressSpace.gmem, - assumed_align=16, - ) - b_sf_ptr = make_ptr( - cutlass.Uint32, - kernel_sfb.data_ptr(), - cute.AddressSpace.gmem, - assumed_align=16, - ) - c_cute_tensor = cute.runtime.from_dlpack( - output).mark_layout_dynamic(leading_dim=1) alpha = self.__class__.alpha_cache.get(device_key) if alpha is None: alpha = torch.ones((1, ), device=device_key, dtype=torch.float32) self.__class__.alpha_cache[device_key] = alpha - alpha_cute_tensor = cute.runtime.from_dlpack(alpha) - stream = (cute.runtime.make_fake_stream( - use_tvm_ffi_env_stream=True) if self.use_tvm_ffi else - cuda.CUstream(torch.cuda.current_stream().cuda_stream)) + + need_cute_args = (cache_key not in self.__class__.kernel_cache + or not self.use_tvm_ffi) + if need_cute_args: + a_ptr = make_ptr(cutlass.Float8E4M3FN, kernel_a.data_ptr(), + cute.AddressSpace.gmem, assumed_align=16) + b_ptr = make_ptr(cutlass.Float8E4M3FN, kernel_b.data_ptr(), + cute.AddressSpace.gmem, assumed_align=16) + a_sf_ptr = make_ptr(cutlass.Uint32, kernel_sfa.data_ptr(), + cute.AddressSpace.gmem, assumed_align=16) + b_sf_ptr = make_ptr(cutlass.Uint32, kernel_sfb.data_ptr(), + cute.AddressSpace.gmem, assumed_align=16) + c_cute_tensor = cute.runtime.from_dlpack( + output).mark_layout_dynamic(leading_dim=1) + alpha_cute_tensor = cute.runtime.from_dlpack(alpha) + stream = (cute.runtime.make_fake_stream( + use_tvm_ffi_env_stream=True) if self.use_tvm_ffi else + cuda.CUstream( + torch.cuda.current_stream().cuda_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 `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py` around lines 4515 - 4550, In the forward path around the pointer and DLPack setup, keep alpha tensor creation outside the backend branch, but construct a_ptr, b_ptr, a_sf_ptr, b_sf_ptr, c_cute_tensor, and alpha_cute_tensor only when use_tvm_ffi is false. Mirror the conditional construction pattern used by CuteDSLFp8BlackwellBmmRunner.forward and forward_fp8out, while preserving the TVM-FFI launch’s use of raw data_ptr values and the torch alpha tensor.tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py (1)
65-68: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winShare the packed-scale constants with the custom-op layer.
_SCALE_BLOCK_K,_SCALES_PER_WORD, and_PACKED_SCALE_Kare redefined intensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pyat Lines 4339-4341 asCuteDSLFp8SplitKGemmRunner._SCALE_BLOCK_K,_SCALES_PER_WORD, and_PACKED_SCALE_K.These values form a producer-consumer contract. The custom op validates
k % (packed_scale_k * num_splits) == 0with its copy (Line 4658) and derivespacked_k = k // packed_scale_k(Line 4663). The kernel derivespacked_k_per_split = split_k // _PACKED_SCALE_Kwith this copy (Line 494). If the two copies ever diverge, the scale tensor layout is wrong and the GEMM produces incorrect results with no error.Export the constants from this module and import them in the custom-op module.
♻️ Proposed consolidation
In
cute_dsl_custom_ops.py, replace the local copies:- _SCALE_BLOCK_K = 128 - _SCALES_PER_WORD = 4 - _PACKED_SCALE_K = _SCALE_BLOCK_K * _SCALES_PER_WORD + _SCALE_BLOCK_K = _SCALE_BLOCK_K + _SCALES_PER_WORD = _SCALES_PER_WORD + _PACKED_SCALE_K = _PACKED_SCALE_Kwith an import of the kernel-module constants at the top of the file, for example:
from ..cute_dsl_kernels.blackwell.dense_blockscaled_gemm_persistent import ( _PACKED_SCALE_K, _SCALE_BLOCK_K, _SCALES_PER_WORD)Consider renaming them without the leading underscore and adding them to
__all__, since they are now a cross-module 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 `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py` around lines 65 - 68, Export the shared constants _SCALE_BLOCK_K, _SCALES_PER_WORD, and _PACKED_SCALE_K from the dense blockscaled GEMM kernel module, then remove their duplicate definitions in CuteDSLFp8SplitKGemmRunner and import the kernel-module values in the custom-op module. Update references as needed so validation and packed-scale calculations use the same constants; consider public names and __all__ only if required by the module’s export conventions.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 5065-5099: Add a capacity validation beside the existing sf_out
stride check to require storage for all padded rows: verify sf_out’s backing
storage can hold at least sf_m * packed_sf_n int32 elements, not merely its
logical shape. Keep the existing shape and MN-major stride checks, and raise
ValueError when the available storage is insufficient.
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.py`:
- Around line 2246-2293: Update the FP8 validation block around the existing
validation logic to enforce that the computed epilogue tile N dimension,
epi_tile_n, equals 32, matching the four cooperative epilogue subtiles assumed
by fp8_subtile_idx and the TMA drain loop. Reject unsupported configurations
before the staging and transfer paths execute, using the existing validation
mechanism.
- Around line 2027-2109: Reject non-identity epilogue_op when FP8 scaling output
is enabled, since both FP8 register and cooperative paths bypass the callback.
Define a module-level _identity_epilogue_op sentinel, use it as the __call__
default instead of the inline lambda, and add validation in the existing
__call__ validation block to accept only that sentinel for the sf_out_tensor/FP8
path.
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py`:
- Around line 2239-2242: Update the epilogue sequence around acc_vec and
epilogue_op so apply_alpha multiplies the accumulator in its higher-precision
type before converting to self.c_dtype. Preserve the existing conversion and
epilogue flow for non-alpha paths, ensuring alpha_value is not applied after
output-type rounding.
- Around line 793-801: Handle packed MXF4 SFA staging consistently at the
packed-scale sizing logic in dense_blockscaled_gemm_persistent.py: either reject
packed Float4 inputs or use a packed SFA layout whose stage stride is
cta_tile_shape_mnk[0] multiplied by _PACKED_SCALE_BYTES, matching
make_smem_layout_sfa and sSFA_raw allocation. Apply the same correction at the
anchor scale-storage calculation around lines 793-801 and the sibling packed-SFA
handling around lines 1061-1078; the existing Float8 packed caller must remain
unchanged.
---
Nitpick comments:
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 4515-4550: In the forward path around the pointer and DLPack
setup, keep alpha tensor creation outside the backend branch, but construct
a_ptr, b_ptr, a_sf_ptr, b_sf_ptr, c_cute_tensor, and alpha_cute_tensor only when
use_tvm_ffi is false. Mirror the conditional construction pattern used by
CuteDSLFp8BlackwellBmmRunner.forward and forward_fp8out, while preserving the
TVM-FFI launch’s use of raw data_ptr values and the torch alpha tensor.
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.py`:
- Around line 2591-2609: Parameterize epilog_smem_copy_and_partition with an
optional dtype argument, defaulting to self.c_dtype, and pass that argument to
sm100_utils.get_smem_store_op. Remove epilog_bf16_smem_copy_and_partition and
update its call site to invoke the existing helper with cutlass.BFloat16.
- Around line 134-175: Add docstrings to the public module-level helpers
abs_f32_device and max_nan_f32_device, matching the concise documentation style
used by ceil_to_ue8m0_device and stg_u8_raw. Leave their implementations and
behavior unchanged.
- Around line 3103-3126: Expand the docstring for wrapper_fp8sf to Google-style
documentation covering every parameter, matching the sibling wrapper’s argument
descriptions where applicable. Explicitly document sf_out_ptr as the
address-carrying output-scale pointer, sf_aligned_mn and n_tiles_per_group, and
the fp8_smem_epilogue Constexpr controls, while preserving the existing function
behavior.
- Around line 478-507: In the fp8_smem_epi_mode path near fp8_bf16_stage_count
and the staged shared-memory layouts, add an assertion verifying that the
shared-memory size released by exactly one A/B stage equals the BF16 epilogue
staging tile size. Keep the existing num_ab_stage decrement, and ensure the
assertion also protects against an invalid zero or negative remaining A/B stage
under the enforced configuration.
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py`:
- Around line 65-68: Export the shared constants _SCALE_BLOCK_K,
_SCALES_PER_WORD, and _PACKED_SCALE_K from the dense blockscaled GEMM kernel
module, then remove their duplicate definitions in CuteDSLFp8SplitKGemmRunner
and import the kernel-module values in the custom-op module. Update references
as needed so validation and packed-scale calculations use the same constants;
consider public names and __all__ only if required by the module’s export
conventions.
🪄 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: 98063236-5277-4a27-9948-8a22ae782857
📒 Files selected for processing (10)
cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuhcpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cucpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.hcpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuhcpp/tensorrt_llm/thop/mhcOp.cpptensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockwise_gemm/blockwise_gemm.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/dynamic_mxf8_mainloop.pytensorrt_llm/_torch/models/modeling_deepseekv4.py
🚧 Files skipped from review as they are similar to previous changes (7)
- cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh
- cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h
- cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh
- cpp/tensorrt_llm/thop/mhcOp.cpp
- tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dynamic_mxf8_mainloop.py
- cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
- tensorrt_llm/_torch/models/modeling_deepseekv4.py
Comment on lines
+5065
to
+5099
| sf_m = pad_up(m, 4) | ||
| sf_k = ceil_div(k, 128) | ||
| sf_n = ceil_div(n, 128) | ||
| packed_sf_n = ceil_div(batch_size * sf_n, 4) | ||
| expected_shapes = ( | ||
| ("weight", b, (batch_size, n, k)), | ||
| ("input_scale", a_sf, (batch_size, sf_k, sf_m)), | ||
| ("weight_scale", b_sf, (batch_size, sf_n, sf_k)), | ||
| ("output_fp8", output, (m, batch_size * n)), | ||
| ("sf_out", sf_out, (m, packed_sf_n)), | ||
| ) | ||
| for name, tensor, shape in expected_shapes: | ||
| if tensor.shape != shape: | ||
| raise ValueError( | ||
| f"{name} must have shape {shape}, got {tuple(tensor.shape)}" | ||
| ) | ||
|
|
||
| for name, tensor in ( | ||
| ("input", a), | ||
| ("weight", b), | ||
| ("input_scale", a_sf), | ||
| ("weight_scale", b_sf), | ||
| ): | ||
| if not tensor.is_contiguous(): | ||
| raise ValueError(f"{name} must be contiguous") | ||
| if tensor.data_ptr() % 16 != 0: | ||
| raise ValueError(f"{name} must be 16-byte aligned") | ||
|
|
||
| if not output.is_contiguous(): | ||
| raise ValueError("output_fp8 must be contiguous") | ||
| if output.data_ptr() % 16 != 0: | ||
| raise ValueError("output_fp8 must be 16-byte aligned") | ||
| if sf_out.stride() != (1, sf_m): | ||
| raise ValueError( | ||
| f"sf_out must have MN-major stride (1, {sf_m})") |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Also validate that sf_out has storage for the padded rows.
The kernel writes padded scale rows beyond M. In blockwise_gemm.py the register path (Lines 2091-2097) and the cooperative path (Lines 2271-2278) store to runtime_m + 0..2 while padding_m < sf_aligned_mn. sf_aligned_mn is the aligned_m value passed at Line 5213, which is pad_up(m, 4).
The current checks accept sf_out with shape (m, packed_sf_n) and stride (1, sf_m). That describes a strided view. The checks do not confirm that the backing storage holds sf_m * packed_sf_n int32 elements. If a caller allocates exactly m * packed_sf_n elements and applies as_strided, the padded-row stores in the last column write past the allocation. That is an out-of-bounds device store.
Add a capacity check next to the stride check.
🛡️ Proposed capacity check
if sf_out.stride() != (1, sf_m):
raise ValueError(
f"sf_out must have MN-major stride (1, {sf_m})")
+ required_sf_out_elems = sf_m * packed_sf_n
+ available_sf_out_elems = (sf_out.untyped_storage().nbytes() //
+ sf_out.element_size() -
+ sf_out.storage_offset())
+ if available_sf_out_elems < required_sf_out_elems:
+ raise ValueError(
+ f"sf_out must back {required_sf_out_elems} int32 elements "
+ f"for M padding, got {available_sf_out_elems}")📝 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
| sf_m = pad_up(m, 4) | |
| sf_k = ceil_div(k, 128) | |
| sf_n = ceil_div(n, 128) | |
| packed_sf_n = ceil_div(batch_size * sf_n, 4) | |
| expected_shapes = ( | |
| ("weight", b, (batch_size, n, k)), | |
| ("input_scale", a_sf, (batch_size, sf_k, sf_m)), | |
| ("weight_scale", b_sf, (batch_size, sf_n, sf_k)), | |
| ("output_fp8", output, (m, batch_size * n)), | |
| ("sf_out", sf_out, (m, packed_sf_n)), | |
| ) | |
| for name, tensor, shape in expected_shapes: | |
| if tensor.shape != shape: | |
| raise ValueError( | |
| f"{name} must have shape {shape}, got {tuple(tensor.shape)}" | |
| ) | |
| for name, tensor in ( | |
| ("input", a), | |
| ("weight", b), | |
| ("input_scale", a_sf), | |
| ("weight_scale", b_sf), | |
| ): | |
| if not tensor.is_contiguous(): | |
| raise ValueError(f"{name} must be contiguous") | |
| if tensor.data_ptr() % 16 != 0: | |
| raise ValueError(f"{name} must be 16-byte aligned") | |
| if not output.is_contiguous(): | |
| raise ValueError("output_fp8 must be contiguous") | |
| if output.data_ptr() % 16 != 0: | |
| raise ValueError("output_fp8 must be 16-byte aligned") | |
| if sf_out.stride() != (1, sf_m): | |
| raise ValueError( | |
| f"sf_out must have MN-major stride (1, {sf_m})") | |
| sf_m = pad_up(m, 4) | |
| sf_k = ceil_div(k, 128) | |
| sf_n = ceil_div(n, 128) | |
| packed_sf_n = ceil_div(batch_size * sf_n, 4) | |
| expected_shapes = ( | |
| ("weight", b, (batch_size, n, k)), | |
| ("input_scale", a_sf, (batch_size, sf_k, sf_m)), | |
| ("weight_scale", b_sf, (batch_size, sf_n, sf_k)), | |
| ("output_fp8", output, (m, batch_size * n)), | |
| ("sf_out", sf_out, (m, packed_sf_n)), | |
| ) | |
| for name, tensor, shape in expected_shapes: | |
| if tensor.shape != shape: | |
| raise ValueError( | |
| f"{name} must have shape {shape}, got {tuple(tensor.shape)}" | |
| ) | |
| for name, tensor in ( | |
| ("input", a), | |
| ("weight", b), | |
| ("input_scale", a_sf), | |
| ("weight_scale", b_sf), | |
| ): | |
| if not tensor.is_contiguous(): | |
| raise ValueError(f"{name} must be contiguous") | |
| if tensor.data_ptr() % 16 != 0: | |
| raise ValueError(f"{name} must be 16-byte aligned") | |
| if not output.is_contiguous(): | |
| raise ValueError("output_fp8 must be contiguous") | |
| if output.data_ptr() % 16 != 0: | |
| raise ValueError("output_fp8 must be 16-byte aligned") | |
| if sf_out.stride() != (1, sf_m): | |
| raise ValueError( | |
| f"sf_out must have MN-major stride (1, {sf_m})") | |
| required_sf_out_elems = sf_m * packed_sf_n | |
| available_sf_out_elems = ( | |
| sf_out.untyped_storage().nbytes() // sf_out.element_size() | |
| - sf_out.storage_offset() | |
| ) | |
| if available_sf_out_elems < required_sf_out_elems: | |
| raise ValueError( | |
| f"sf_out must back {required_sf_out_elems} int32 elements " | |
| f"for M padding, got {available_sf_out_elems}" | |
| ) |
🤖 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 `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py` around lines 5065 -
5099, Add a capacity validation beside the existing sf_out stride check to
require storage for all padded rows: verify sf_out’s backing storage can hold at
least sf_m * packed_sf_n int32 elements, not merely its logical shape. Keep the
existing shape and MN-major stride checks, and raise ValueError when the
available storage is insufficient.
mingyangHao
force-pushed
the
user/mingyangh/dsv4-oproj-fp8-fuse
branch
from
August 3, 2026 04:43
5e00659 to
04b99e3
Compare
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
…actic pool The runner always queried get_max_active_clusters(2), so (1,1)-cluster tactics launched with a 74-cluster persistent grid on 148 SMs (half the machine). Query per tactic cluster size instead, cached per (device, cluster size). Expand _SPLIT_K1_TACTICS from 10 to 17: add noswap (256,112/160/224/240) and swap (256,208/224/240) tiles. Fine tile-N noswap variants raise SM fill at mid M (tile-M count 28 under swap never fills the grid), and the 224/240 tiles recover the M>=4096 range where the previous pool lost to DeepGEMM under realistic data. Fallback entries are unchanged and the runtime AutoTuner selects winners per token bucket. Paired-capture sweep vs the previous pool (ratio vs DeepGEMM within shared captures, B200): decode M=1-32 +4.7-6.8%, M=256 +3.1%, M=4096/8192/16384 +1.9/+3.0/+1.9%; no regressions at 15 sweep points. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
The packed-scale pipeline's empty-buffer release used ab_pipeline.consumer_mask, which spans the whole A/B TMA multicast group. The scale pipeline is pair-local (producers are the MMA pair's scale warps, consumer is the pair leader), so for any cluster larger than the pair the release lands spurious arrivals on the other pair's scale barriers, corrupting their phases and deadlocking the kernel. (2,1) clusters worked only because the two masks coincide there. Release with the pair-local V-mode image mask instead. For (2,1) the mask value is identical by construction; all shipped tactics validated byte-exact with unchanged latency, and previously-hanging (2,2) and (4,1) configurations now run and validate byte-exact against DeepGEMM. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
mingyangHao
force-pushed
the
user/mingyangh/dsv4-oproj-fp8-fuse
branch
from
August 5, 2026 08:19
f6560f2 to
e565933
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #64000 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64000 [ run ] completed with state
|
mikeiovine
approved these changes
Aug 26, 2026
mikeiovine
left a comment
mikeiovine
left a comment
Collaborator
There was a problem hiding this comment.
Stamp on behalf of runtime devs, delegating review to @NVIDIA/trt-llm-torch-attention-devs
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
9 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fold the o_lora bf16->fp8 + 1x128 packed-UE8M0 quantization into the o_a blockwise GEMM epilogue (single-pass, tree-reduce amax), so o_a emits fp8+sf directly and o_b (DeepGEMM 1d1d) consumes it without the separate fp8_quantize_1x128_packed kernel and the bf16 o_lora HBM round-trip.
blockwise_gemm.py: optional fp8+SF epilogue path, gated by sf_out tensor; the bf16 path is byte-identical when off. Single-pass register cache, tree-reduce amax (byte-exact), packed-UE8M0 MN-major sf store, predicated for tail-M, config-guarded to (128,128)/1-CTA.
cute_dsl_custom_ops.py: new trtllm::cute_dsl_fp8_bmm_blackwell_fp8out op.
attention.py: gated DSV4_FUSE_OPROJ branch in _deepseek_v4_o_proj (tp_size==1), falls back to the existing path when the flag is off.
Validated byte-exact (G=1/G=16/tail-M/negative guards) and net-positive (+1.6us/o_a in decode, nsys hardware-timestamp measured). Default OFF -> zero functional change unless DSV4_FUSE_OPROJ=1.
The O_a baseline is the BF16-output CuTe BMM followed by the standalone
fp8_quantize_1x128_packed_ue8m0kernel. The fused path emits FP8 and packedUE8M0 scales directly from the BMM epilogue. The benchmark uses the production
[M, G, N]backing storage and[G, M, N]BMM transpose view.The O_b baseline is upstream DeepGEMM followed by mHC. The candidate is the
CuTe DSL GEMM followed by mHC, with split-partial reduction folded into mHC.
Therefore the O_b table includes the split reduction cost.
O_a Projection + Quantization
O_a is positive at all 15 points, with a latency reduction of 13.02% to
24.15%. FP8 output bytes and packed UE8M0 scale words matched the standalone
quantization baseline exactly at every M.
O_b Projection + mHC (
M=1to16384)The O_b sweep uses the same power-of-two M points as O_a:
1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384.O_b improves at all 15 points. The production selector uses SK4 for
M <= 16, SK2 for16 < M <= 128, and SK1 above 128. Isolated GEMM gainsrange from 0.29% to 11.50%, and the O_b+mHC chain gains range from 0.09% to
6.96%. The tuned M=2048 path is 1.42% faster before mHC and 0.55% faster end
to end, with no DeepGEMM fallback. CuTe partial sums matched DeepGEMM, and all
four mHC outputs matched the reduced-input reference across the sweep. Useful
SOL is
2*M*N*K / latency; it reaches 90.6% to 93.4% of the nominal4.5 PFLOP/s dense FP8 peak at
M >= 4096.Combined O-Projection
This sum counts O_a quantization once and includes O_b split reduction through
mHC.
Combined O-projection latency improves at all points by 5.75% to 10.14%.
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.Dev Engineer Review
Correctness and consistency
o_aBF16-to-FP8 quantization with packed UE8M0 scales.o_bprojection and integrates partial reduction with mHC.x_prevsupport for mHC withx_num_splitsvalues of1,2, or4.DSV4_FUSE_OPROJ=1withtp_size==1.M=2048path avoids DeepGEMM fallback.x_num_splits=1defaults.Performance
o_alatency reduction: 13.02%–24.15%.o_bGEMM latency reduction: 0.29%–11.50%.o_bplus mHC latency reduction: 0.09%–6.96%.CI
/bot runjobs failed, and one job was aborted.#62377completed successfully with merge-request pipeline#50544.QA Engineer Review
Test changes
test_mhc_fused_hc_reduces_split_x_for_production_backends.test_deepseek_v4_o_projthrough_build_dsv4_o_proj_case.test_deepseek_v4_o_proj.Test-list coverage
tests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_gb200_multi_gpus.ymlVerdict
Sufficient