[None][perf] DSV4 o_proj: fuse fp8/UE8M0 quantize into o_a proj and introduce splitk cutedsl o_b proj - #16346
Conversation
|
/bot run |
WalkthroughAdds split-x support to fused mHC kernels and launchers, including split-major input validation and accumulation. It also adds Blackwell FP8 CuTe-DSL GEMM/BMM paths, packed-scale execution, dynamic MXF8 MMA helpers, and a fused DeepSeek-V4 projection route 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 labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #59071 [ ] completed with state |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
| float* y_acc_workspace, float* r_acc_workspace, int M, int hidden_size, int hc_mult, int num_k_splits, | ||
| int bigfuse_block_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, | ||
| int sinkhorn_repeat, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream) | ||
| int sinkhorn_repeat, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream, int x_num_splits) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Update the disabled-DeepGEMM stub to the new signature.
The #else definition at Lines 790-795 still ends at cudaStream_t stream. Non-DeepGEMM builds therefore lack the mhcFusedHcLaunch(..., int x_num_splits) symbol used by the header and Torch binding.
Proposed fix
- int sinkhorn_repeat, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream)
+ int sinkhorn_repeat, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream,
+ int x_num_splits)📝 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.
| int sinkhorn_repeat, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream, int x_num_splits) | |
| int sinkhorn_repeat, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream, | |
| int x_num_splits) |
🤖 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/mhcKernels/mhcFusedHcKernel.cu` at 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.
| def forward_fp8out( | ||
| self, | ||
| inputs: List[torch.Tensor], | ||
| tactic=None, | ||
| ) -> None: | ||
| """Run the fixed DSV4 o_a tactic with FP8 and packed-UE8M0 output.""" | ||
| fixed_tactic = (False, (128, 128), (1, 1)) | ||
| use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = ( | ||
| tactic if isinstance(tactic, tuple) else fixed_tactic) | ||
| if (use_2cta_instrs, mma_tiler_mn, | ||
| cluster_shape_mn) != fixed_tactic: | ||
| raise ValueError("FP8-output BMM supports only the fixed " | ||
| "(1-CTA, 128x128, 1x1 cluster) tactic") | ||
| if not self.use_tvm_ffi: | ||
| raise ValueError("FP8-output BMM requires TVM FFI") | ||
|
|
||
| (a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, c_tensor, | ||
| sf_out_tensor) = inputs | ||
| batch_size, m, k = a_tensor.shape | ||
| n = b_tensor.shape[1] | ||
| aligned_m = pad_up(m, 4) | ||
| sf_m = aligned_m | ||
| sf_k = ceil_div(k, 128) | ||
| sf_n = ceil_div(n, 128) | ||
| # SMEM cooperation wins through M=32. | ||
| use_fp8_smem_epilogue = m <= 32 | ||
| fp8_smem_row_iters = ceil_div(m, 16) if use_fp8_smem_epilogue else 1 | ||
|
|
||
| if c_tensor.dtype != torch.float8_e4m3fn or c_tensor.shape != ( | ||
| batch_size, m, n): | ||
| raise ValueError( | ||
| "FP8-output BMM output must have shape [batch, M, N] " | ||
| "and dtype float8_e4m3fn") | ||
| if sf_out_tensor.dtype != torch.int32 or sf_out_tensor.stride() != ( | ||
| 1, aligned_m): | ||
| raise ValueError( | ||
| "FP8-output BMM scales must be int32 with MN-major " | ||
| f"stride (1, {aligned_m})") | ||
|
|
||
| c_tmp = c_tensor.permute(1, 2, 0) | ||
| cache_key = ( | ||
| "fp8out", | ||
| use_2cta_instrs, | ||
| mma_tiler_mn, | ||
| cluster_shape_mn, | ||
| self.use_tvm_ffi, | ||
| fp8_smem_row_iters, | ||
| use_fp8_smem_epilogue, | ||
| ) | ||
| if cache_key not in self.__class__.kernel_cache: | ||
| a_ptr = make_ptr(cutlass.Float8E4M3FN, | ||
| a_tensor.data_ptr(), | ||
| cute.AddressSpace.gmem, | ||
| assumed_align=16) | ||
| b_ptr = make_ptr(cutlass.Float8E4M3FN, | ||
| b_tensor.data_ptr(), | ||
| cute.AddressSpace.gmem, | ||
| assumed_align=16) | ||
| a_sf_ptr = make_ptr(cutlass.Float32, | ||
| a_sf_tensor.data_ptr(), | ||
| cute.AddressSpace.gmem, | ||
| assumed_align=16) | ||
| b_sf_ptr = make_ptr(cutlass.Float32, | ||
| b_sf_tensor.data_ptr(), | ||
| cute.AddressSpace.gmem, | ||
| assumed_align=16) | ||
| c_cute_tensor = cute.runtime.from_dlpack( | ||
| c_tmp).mark_layout_dynamic(leading_dim=1) | ||
| sf_out_ptr = make_ptr(cutlass.Int32, | ||
| sf_out_tensor.data_ptr(), | ||
| cute.AddressSpace.gmem, | ||
| assumed_align=4) | ||
|
|
||
| gemm = self.__class__.kernel_class( | ||
| cutlass.Float32, | ||
| use_2cta_instrs=use_2cta_instrs, | ||
| mma_tiler_mn=mma_tiler_mn, | ||
| cluster_shape_mn=cluster_shape_mn, | ||
| ) | ||
| cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] | ||
| max_active_clusters = cutlass.utils.HardwareInfo( | ||
| ).get_max_active_clusters(cluster_size) | ||
| stream = cute.runtime.make_fake_stream( | ||
| use_tvm_ffi_env_stream=True) | ||
| compiled_gemm = cute.compile( | ||
| gemm.wrapper_fp8sf, | ||
| m, | ||
| n, | ||
| k, | ||
| sf_m, | ||
| sf_n, | ||
| sf_k, | ||
| batch_size, | ||
| a_ptr, | ||
| b_ptr, | ||
| a_sf_ptr, | ||
| b_sf_ptr, | ||
| c_cute_tensor, | ||
| max_active_clusters, | ||
| stream, | ||
| sf_out_ptr, | ||
| aligned_m, | ||
| sf_n, | ||
| fp8_smem_row_iters, | ||
| use_fp8_smem_epilogue, | ||
| options="--opt-level 2 --enable-tvm-ffi", | ||
| ) | ||
| self.__class__.kernel_cache[cache_key] = compiled_gemm | ||
| else: | ||
| compiled_gemm = self.__class__.kernel_cache[cache_key] | ||
|
|
||
| compiled_gemm( | ||
| m, | ||
| n, | ||
| k, | ||
| sf_m, | ||
| sf_n, | ||
| sf_k, | ||
| batch_size, | ||
| a_tensor.data_ptr(), | ||
| b_tensor.data_ptr(), | ||
| a_sf_tensor.data_ptr(), | ||
| b_sf_tensor.data_ptr(), | ||
| c_tmp, | ||
| sf_out_tensor.data_ptr(), | ||
| aligned_m, | ||
| sf_n, | ||
| ) | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Validate the complete tensor contract before launching raw pointers.
sf_out_tensor.shape is unchecked. For example, (M, 1) with stride (1, aligned_m) passes, but the kernel writes ceil_div(batch_size * ceil_div(N, 128), 4) columns, causing out-of-bounds CUDA stores.
Also validate all tensors’ CUDA device, common device, dtype, rank, compatible shapes, and required layouts before launch. Require:
sf_out.shape == (M, ceil_div(batch_size * ceil_div(N, 128), 4))
Add negative-contract coverage in tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py.
Also applies to: 4461-4496
🤖 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 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.
| if x_num_splits in (2, 4): | ||
| if B <= 32: | ||
| # CUDA-core FMA wins at tiny M. | ||
| best_tactic = _FUSED_HC_SPLIT_X_FMA_TACTIC | ||
| else: | ||
| # C++ selects split-K and BigFuse geometry from M. | ||
| best_tactic = _FUSED_HC_FALLBACK_TACTIC_MMA |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Honor MMA availability when selecting the split-x fallback.
For B > 32, this always selects MMA, even when SM100 MMA is unavailable or hidden_size is unsupported. Reuse the capability-aware fallback so split-x can select FMA instead.
Proposed fix
else:
# C++ selects split-K and BigFuse geometry from M.
- best_tactic = _FUSED_HC_FALLBACK_TACTIC_MMA
+ best_tactic = _get_fused_hc_fallback_tactic(hidden_size)📝 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.
| if x_num_splits in (2, 4): | |
| if B <= 32: | |
| # CUDA-core FMA wins at tiny M. | |
| best_tactic = _FUSED_HC_SPLIT_X_FMA_TACTIC | |
| else: | |
| # C++ selects split-K and BigFuse geometry from M. | |
| best_tactic = _FUSED_HC_FALLBACK_TACTIC_MMA | |
| if x_num_splits in (2, 4): | |
| if B <= 32: | |
| # CUDA-core FMA wins at tiny M. | |
| best_tactic = _FUSED_HC_SPLIT_X_FMA_TACTIC | |
| else: | |
| # C++ selects split-K and BigFuse geometry from M. | |
| best_tactic = _get_fused_hc_fallback_tactic(hidden_size) |
🤖 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/modules/mhc/mhc_cuda.py` around lines 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.
| def _should_use_fused_oproj(self) -> bool: | ||
| # Eligible Blackwell paths are on by default; keep an emergency fallback. | ||
| return ( | ||
| not _is_env_truthy("TRTLLM_DSV4_DISABLE_FUSED_OPROJ") | ||
| and is_sm_100f() | ||
| and IS_CUTLASS_DSL_AVAILABLE | ||
| and self.n_local_groups == self.num_groups | ||
| and getattr(self.o_b_proj, "tp_size", 1) == 1 | ||
| and self.o_a_proj.dtype == torch.float8_e4m3fn | ||
| and self.o_b_proj.has_fp8_block_scales | ||
| and not getattr(self.o_b_proj, "use_cute_dsl_blockscaling_mm", False) | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Gate the CuTe split-K path to BF16 output.
The custom op requires BF16 partials, but Line 1243 allocates them with self.dtype while the eligibility check accepts FP16 or unset dtypes. Such configurations enter fusion and fail with TypeError.
Proposed fix
and self.o_a_proj.dtype == torch.float8_e4m3fn
and self.o_b_proj.has_fp8_block_scales
+ and self.dtype == torch.bfloat16
and not getattr(self.o_b_proj, "use_cute_dsl_blockscaling_mm", False)Also applies to: 1242-1245
🤖 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/modules/mla.py` around lines 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.
| device = torch.device("cuda") | ||
| mla, attn_out_latent, position_ids, refs = _build_dsv4_o_proj_case(num_tokens, "fp8", device) | ||
|
|
||
| # Ensure the fused path is eligible. | ||
| assert mla.o_a_proj.dtype == torch.float8_e4m3fn | ||
| assert mla.n_local_groups == mla.num_groups | ||
| assert getattr(mla.o_b_proj, "tp_size", 1) == 1 | ||
| assert mla.o_b_proj.has_fp8_block_scales | ||
| assert not getattr(mla.o_b_proj, "use_cute_dsl_blockscaling_mm", False) | ||
|
|
||
| # Explicit kill switch retains the unfused fallback. | ||
| monkeypatch.setenv("TRTLLM_DSV4_DISABLE_FUSED_OPROJ", "1") | ||
| out_unfused = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids) | ||
|
|
||
| # Unset is the production default and must take the fused path. | ||
| monkeypatch.delenv("TRTLLM_DSV4_DISABLE_FUSED_OPROJ", raising=False) | ||
| mla.ob_split_k = 1 | ||
| out_fused = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Use a fixture that is eligible for fused O-projection.
_build_dsv4_o_proj_case(..., "fp8", ...) sets use_cute_dsl_blockscaling_mm=True, which propagates to o_b_proj. Consequently, Line 659 fails and _should_use_fused_oproj() returns false. Parameterize the helper and pass False for this test.
As per path instructions, coverage is currently insufficient because this test cannot reach the fused route.
🤖 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 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.
Source: Path instructions
| 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) |
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.
| 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.
|
/bot run |
|
PR_Github #59110 [ run ] triggered by Bot. Commit: |
|
PR_Github #59110 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59184 [ run ] triggered by Bot. Commit: |
|
PR_Github #59184 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59317 [ run ] triggered by Bot. Commit: |
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>
2a3f208 to
faad69e
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh (1)
236-265: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate the barrier-inventory comment for the new x-split barriers.
The file header (lines 42-50) documents the barrier set as
full_B,empty_B,full_input/empty_input("residual+x arrived"),full_cast/empty_cast,tmem_full. WithkXSplit > 1,full_input/empty_inputnow only cover residual bytes (SMEM_RES_PER_ISTG), and x is streamed through the newly-addedfull_x/empty_xpair — but this top-level doc block wasn't updated to reflect that split.📝 Suggested doc update
// Barriers: // full_B[N_B_STAGES] TMA -> MMA (B arrived in SMEM) // empty_B[N_B_STAGES] MMA -> TMA (B slot empty) -// full_input[N_INPUT] TMA -> pmap (residual+x arrived) +// full_input[N_INPUT] TMA -> pmap (residual arrived; also x when kXSplit==1) // empty_input[N_INPUT] pmap -> TMA (input slot empty) +// full_x[N_INPUT] TMA -> pmap (x split tile arrived, kXSplit>1 only) +// empty_x[N_INPUT] pmap -> TMA (x split slot empty, kXSplit>1 only) // full_cast[2] pmap -> MMA (A ready in TMEM) // empty_cast[2] MMA -> pmap (TMEM slot empty) // tmem_full[1] MMA -> epilogueSeparately (non-blocking):
full_x/empty_xare unconditionally defined viaPatternVisitoreven whenkXSplit == 1, at which point they alias the same memory asfull_cast/empty_cast/tmem_full_barrier(sincekXBarrierscollapses to 0). They're never dereferenced in that case, so this is harmless, but a short comment noting the aliasing would help future readers.🤖 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/mhcKernels/fused_tf32_pmap_gemm.cuh` around lines 236 - 265, Update the file-header barrier inventory documentation to state that full_input/empty_input cover only residual bytes and that, when kXSplit > 1, full_x/empty_x synchronize streamed x tiles. Leave the barrier allocation and PatternVisitor logic unchanged; the aliasing note for kXSplit == 1 is optional and non-blocking.
🤖 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.
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh`:
- Around line 236-265: Update the file-header barrier inventory documentation to
state that full_input/empty_input cover only residual bytes and that, when
kXSplit > 1, full_x/empty_x synchronize streamed x tiles. Leave the barrier
allocation and PatternVisitor logic unchanged; the aliasing note for kXSplit ==
1 is optional and non-blocking.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f481854a-2601-4a90-960c-5122f3b192b1
📒 Files selected for processing (4)
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.cuh
🚧 Files skipped from review as they are similar to previous changes (3)
- cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h
- cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh
- cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
|
/bot run --disable-fail-fast |
|
PR_Github #61827 [ run ] triggered by Bot. Commit: |
|
PR_Github #61827 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62153 [ run ] triggered by Bot. Commit: |
chienchunhung
left a comment
chienchunhung
left a comment
There was a problem hiding this comment.
Thanks for the PR!
The warmup contracts cover only the largest SK4/SK2 buckets, while get_tuning_config() maps runtime requests to every bucket (1/2/4/8/16 for SK4 and 32/64/128 for SK2). Consequently, requests such as SK4/M=1 or SK2/M=37 miss the cache and use the fallback tactic during inference.
Please consider warming every generated bucket (or map all requests to the warmed bucket) and add a cache-miss regression test for representative small and irregular M values.
|
PR_Github #62153 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62377 [ run ] triggered by Bot. Commit: |
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
|
PR_Github #62377 [ run ] completed with state |
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/custom_ops/cute_dsl_custom_ops.py (1)
4774-4789: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard the split-4/2 warmup contracts with the same K constraint.
k % (packed_scale_k * num_splits)only checks the active contract, but warmup always probes split-4 first and split-2 next. For inputs that are only valid for split-1/2, that can hand an incompatiblektochoose_one()and fail autotune startup.🤖 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 4774 - 4789, The warmup contract construction in the tuner.is_tuning_mode block must only add split-4 and split-2 contracts when k satisfies each contract’s packed_scale_k × split divisibility requirement. Apply the same K-validity guard used for the active (num_splits, m) contract before appending each lower-split warmup entry, while preserving split-1 and valid active-contract tuning.
🤖 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/custom_ops/cute_dsl_custom_ops.py`:
- Around line 4774-4789: The warmup contract construction in the
tuner.is_tuning_mode block must only add split-4 and split-2 contracts when k
satisfies each contract’s packed_scale_k × split divisibility requirement. Apply
the same K-validity guard used for the active (num_splits, m) contract before
appending each lower-split warmup entry, while preserving split-1 and valid
active-contract tuning.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6c2a41e5-f3ee-49dd-ba47-1723058c1b42
📒 Files selected for processing (2)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
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
8 participants
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 & consistency
o_afusion epilogue: Adds a fused BF16→FP8 + packed UE8M0 quantization for the DeepSeek-V4o_aprojection BMM epilogue. The fused path is disabled by default and gated byDSV4_FUSE_OPROJ=1withtp_size==1. Reported correctness includes matching fused FP8 outputs and packed UE8M0 scales vs the standalone baseline.o_bsplit-K CuTe DSL + partial reduction via mHC:o_b(dsv4_fp8_splitk_gemm) with partial sums reduced through the mHC integration.M=2048path that uses no DeepGEMM fallback.x_num_splits ∈ {1,2,4}):x_prev([SK*B, hidden]) by propagatingx_num_splitsinto:kXSplit), including TMA/smem barrier/descriptor sequencing changes forkXSplit>1.desc_xTMA outer-dimension scaling, shared-memory sizing adjustments, and dispatch/tactic selection updates.fused_pmap_gemm_fma_ksplit) withX_NUM_SPLITSand accumulates per-split contributions.sf_out) via new CuTe wrappers/helpers and a cooperative/alternate epilogue path.x_num_splitsto exported fused-HC launchers:mhcFusedHcLaunch(..., cudaStream_t stream, int x_num_splits=1)mhcFusedHcFmaLaunch(..., cudaStream_t stream, int x_num_splits=1)mhc_fused_hcschema to exposex_num_splits=1.fused_hccall path to validate shapes/dtypes/contiguity and reshape split-major partials into the expected split rows.o_bpath (including reuse of FP8 swap-ab warmup behavior, as reflected in commit messaging).Performance (as reported)
Mvalues:o_alatency reduction: 13.02%–24.15%o_bGEMM latency reduction: 0.29%–11.50%o_b+ mHC latency reduction: 0.09%–6.96%CI/operational notes
/bot runattempts had multiple failures (PR_Github runs#59071,#59110,#60325,#60577,#60626,#61827,#62153; also#61148failed;#61306aborted).#50544(PR_Github run#62377).QA Engineer Review
Test changes
Touched test code
tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py_build_dsv4_o_proj_case(...)test_dsv4_fused_oproj_requires_split_output_consumertest_dsv4_deep_gemm_ob_warmup_skips_cute_backendtest_select_dsv4_ob_split_k_auto_policytest_select_dsv4_ob_split_k_overridestest_dsv4_fmha_epilogue_output_uses_fused_oprojtest_dsv4_fmha_epilogue_output_uses_unsplit_fallback_without_fused_hctest_dsv4_ob_cute_tactics_are_runtime_tunedtest_dsv4_ob_tuning_bucket_mappingtest_dsv4_ob_tuning_bucket_generationtest_dsv4_ob_cute_compile_key_excludes_runtime_shapestest_dsv4_ob_tuning_input_hook_restores_packed_scale_layouttest_dsv4_deep_gemm_ob_warmup_uses_startup_autotunertest_dsv4_model_warmup_deduplicates_deep_gemm_signaturestest_dsv4_oa_fp8out_clears_packed_scale_paddingtest_dsv4_ob_split_k_one_uses_cute_dsl_and_caches_weight_scaletest_dsv4_ob_auto_split_uses_cute_dsltest_dsv4_ob_uses_deep_gemm_fallbacktest_dsv4_pro_fp8_splitk_gemm_partialstest_dsv4_pro_fp8_splitk_gemm_reuses_kernels_and_captures_cuda_graphtest_dsv4_pro_fp8_splitk_gemm_packed_scalestest_deepseek_v4_o_proj_fused_fp8_equivalencetest_deepseek_v4_o_proj(...) -> Nonetests/unittest/_torch/modules/test_mhc.pytest_mhc_fused_hc_reduces_split_x_for_production_backendsCoverage in tests/integration/test_lists/
tests/integration/test_lists/test-db/l0_b200.yml(includes the two modified test files)tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml(includestests/unittest/_torch/modules/test_mhc.py)