[None][perf] Optimize Blackwell fused MHC half-MMA kernel

[MengmSun]

Dev Engineer Review

• Updated cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh to optimize the Blackwell fused MHC pmap/cast pipeline used by the two-kernel fused_half_mma path via:

  • Adding packed float2 arithmetic helpers in namespace fused_mhc (with an arch-gated PTX path when available and scalar fallbacks otherwise), including:
    • bitcast_f32x2_to_u64, bitcast_u64_to_f32x2
    • mul_f32x2, fma_f32x2, fma_f32x2_vv
  • Refactoring the pmap prologue so the warp-group no longer stages post_mix/comb_mix rows into SMEM; instead it loads the needed rows directly from global memory into per-thread registers using validity-masked float4-style loads (deterministic zeroing for out-of-range).
  • Switching squared-sum accumulation from scalar to packed float2 (sqr2_u/sqr2_l), with later reduction derived from the packed components.
  • Refactoring the cast-stage residual/post-mapping flow into a tile-batched float2 sequence:
    • batching cast-stage waits across all cast stages at the tile boundary,
    • guarding overwrite of single-buffer residual scratch (smem_rc) by draining in-flight residual TMA stores with tma_store_wait<0>() when ht > 0,
    • applying bf16 rounding on intermediate packed results, storing routing outputs via STSM, and accumulating squared sums with packed float2 FMA.
  • Loop-end ordering tweaks: moving non-early-release empty_input[i_stage]->arrive() outside the inner restructuring and draining residual TMA once via tma_store_wait<0>(); computing warp-reduction inputs for sqr_sum from sqr2_u.x + sqr2_u.y / sqr2_l.x + sqr2_l.y.

• Scope preserved: kernel signatures, tactics, launch configuration, shared-memory allocation, APIs/dependencies, and build configuration remain unchanged.

• Correctness/performance review focus:

  • Equivalence of the packed float2 bitcasts and PTX mul.f32x2/fma.f32x2 behavior vs scalar fallbacks.
  • Validity-masked global loads for out-of-range coefficients returning deterministic zeros.
  • TMA/store synchronization: ensuring tma_store_wait<0>() occurs before reusing smem_rc to avoid residual race regressions.

• Follow-up change: renamed packed FP32x2 PTX helper wrapper names to generic forms (mul_f32x2, fma_f32x2, and related bitcast helpers), described as a functional no-op.

• Expected impact (per objective): on NVIDIA B300, half-MMA workloads for M=32–16384 improve by ~4.24%–19.86%, while half-FMA for M=1–16 stays within about -0.25% to +0.77% of baseline.

QA Engineer Review

No test changes.

Description

This PR optimizes the two-kernel fused_half_mma path used by fused MHC on Blackwell GPUs.

This change improves that kernel by:

1. Packed FP32x2 math. The per-element post-mapping chains use inline PTX mul.rn.ftz.f32x2 and fma.rn.ftz.f32x2 on SM100-family Blackwell GPUs, including SM103/B300. The scalar implementation remains available as the fallback for other architectures.
2. Loop reorder. The column-pair loop is moved outside the residual-stream loop so that BF16-to-FP32 conversions are reused.
3. Tile-batched cast-stage protocol. Four cast stages are claimed at tile start, TMEM stores are interleaved with pmap compute, and the stages are released together after the TMEM-store fence.
4. Register coefficient loads and residual TMA stores. The coefficient shared-memory prologue and its synchronization are replaced by direct float4 register loads. Residual output stores are issued through TMA after
   the tile compute.

The existing kernel signature, tactic definitions, launch configuration, and dynamic shared-memory allocation are preserved. There are no API, dependency, or build-system changes.

Test Coverage

• GPU: NVIDIA B300 SXM6 AC, CC 10.3, 148 SMs
• SM clock: locked at 1800 MHz for both baseline and MMA-optimized builds
• Operation: complete mhc_fused_hc full-op through TensorRT-LLM's MhcFusedHcRunner, including BigFuse
• Shape: H=7168, n=4, M=1,2,4,...,16384
• Measurement: 100 warmups per tactic, 5 timing blocks, adaptive iterations covering at least 1 second per block, CUDA-event median
• Baseline: e96c3309f3f3bae77e6b6a4ec95fa046fd40eced
• MMA-optimized: 6e214bd94f06a2a23108a3e4bbfd0c7769746f1b
• Tactics: all 476 tactics returned by MhcFusedHcRunner.get_valid_tactics(inputs, profile=None) for each build;
  all production FMA and MMA backends are measured
• Improvement: (baseline_us - optimized_us) / baseline_us x 100

M	Baseline best tactic	Baseline (us)	Optimized best tactic	Optimized (us)	Reduction (us)	Improvement
1	half-FMA TN=1 SK=8 BS=512	9.306	half-FMA TN=1 SK=8 BS=512	9.329	-0.023	-0.25%
2	half-FMA TN=3 SK=2 BS=512	9.476	half-FMA TN=2 SK=4 BS=256	9.403	0.073	0.77%
4	half-FMA TN=3 SK=2 BS=512	9.445	half-FMA TN=3 SK=2 BS=512	9.400	0.044	0.47%
8	half-FMA TN=3 SK=2 BS=512	9.406	half-FMA TN=3 SK=2 BS=512	9.402	0.003	0.04%
16	half-FMA TN=3 SK=2 BS=512	11.941	half-FMA TN=3 SK=2 BS=512	11.927	0.013	0.11%
32	half-MMA SK=112 BS=512	12.955	half-MMA SK=112 BS=512	12.405	0.550	4.24%
64	half-MMA SK=112 BS=512	14.292	half-MMA SK=112 BS=512	12.383	1.909	13.36%
128	half-MMA SK=56 BS=512	16.657	half-MMA SK=56 BS=512	13.484	3.173	19.05%
256	half-MMA SK=28 BS=512	21.252	half-MMA SK=28 BS=512	17.032	4.220	19.86%
512	all-MMA SK=16	29.451	half-MMA SK=16 BS=256	24.066	5.386	18.29%
1024	all-MMA SK=8	57.101	half-MMA SK=8 BS=512	50.688	6.413	11.23%
2048	half-MMA SK=4 BS=256	99.367	half-MMA SK=4 BS=256	91.102	8.264	8.32%
4096	half-MMA SK=2 BS=256	177.645	half-MMA SK=2 BS=256	167.224	10.421	5.87%
8192	half-MMA SK=1 BS=128	329.608	half-MMA SK=1 BS=128	314.911	14.696	4.46%
16384	half-MMA SK=1 BS=128	645.130	half-MMA SK=1 BS=128	611.906	33.224	5.15%

In this paired reference, the unchanged FMA-selected region (M=1..16) stays within -0.25% to +0.77%. The affected production-best MMA region(M=32..16384) improves every shape by 4.24% to 19.86%, with a 9.77% median and 10.98% mean improvement.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[MengmSun]

/bot run

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d6bd3658-1f3d-4a0f-b79b-f2a9768767c9

📥 Commits

Reviewing files that changed from the base of the PR and between 4a52d3d and fca43ef.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh

🚧 Files skipped from review as they are similar to previous changes (1)

• cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh

---

Walkthrough

Changes

The fused TF32 pmap GEMM kernel adds packed float2 arithmetic and masked coefficient loading, then refactors cast-stage synchronization, residual mixing, TMA draining, and squared-sum reduction.

Fused TF32 pmap GEMM

Layer / File(s)	Summary
Packed math and coefficient loading cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh	Adds float2 bit-casting and multiply/FMA helpers, loads valid coefficient rows directly, removes coefficient preloading, and initializes packed square-sum state.
Cast-stage pipeline cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh	Batches cast-stage waits, performs packed residual mixing with BF16-rounded values, reorganizes residual stores, signals stage completion, drains residual output TMAs, and reduces packed square sums.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant fused_tf32_pmap_gemm
  participant global_coefficient_memory
  participant cast_stage_pipeline
  participant residual_output_TMA
  fused_tf32_pmap_gemm->>global_coefficient_memory: Load valid post_mix and comb_mix rows
  global_coefficient_memory-->>fused_tf32_pmap_gemm: Return masked float4 coefficients
  fused_tf32_pmap_gemm->>cast_stage_pipeline: Synchronize cast stages and mix residuals
  cast_stage_pipeline->>residual_output_TMA: Drain outstanding residual stores
  residual_output_TMA-->>fused_tf32_pmap_gemm: Complete residual output

Loading

Suggested reviewers: zongfeijing

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and clearly describes the main change: optimizing the Blackwell fused MHC half-MMA kernel.
Description check	✅ Passed	The description follows the template with Description, Test Coverage, and PR Checklist sections and provides concrete implementation and validation details.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (2)

cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh (2)

441-458: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value

Vectorized float4 reinterpret onto plain float[4] locals relies on implicit compiler alignment.

*reinterpret_cast<float4*>(pm_u)/cm_u[j] requires 16-byte alignment, but pm_u/cm_u are declared as plain float[HC_MULT]/float[HC_MULT][HC_MULT] with no explicit alignment. This is a common working CUDA idiom for register-promoted locals, but is not guaranteed by the language. Consider declaring these as float4 (or alignas(16) float[...]) to make the alignment explicit and avoid relying on compiler behavior if the surrounding code changes and these locals stop being register-promoted.

🤖 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
441 - 458, Make the local buffers used by the float4 stores in the fused kernel
explicitly 16-byte aligned, preferably by declaring pm_u, pm_l, cm_u, and cm_l
with alignas(16) while preserving their existing float-array layout and
indexing. Keep the loadCoefficientRow calls and subsequent consumers unchanged.

---

240-240: 🚀 Performance & Scalability | 🔵 Trivial

Reclaim now-dead SMEM_POST_SIZE/SMEM_COMB_SIZE reservation.

Since post_mix/comb_mix are now loaded directly into registers via loadCoefficientRow (lines 436-458), smem_post/smem_comb are never referenced anywhere else in this kernel, yet the cursor still reserves SMEM_POST_SIZE + SMEM_COMB_SIZE (~5 KB per CTA for BLOCK_M=64/HC_MULT=4). Removing this reservation would free SMEM that could improve occupancy or enable other buffer growth — but check whether a companion host-side launcher computes the dynamic SMEM size from a shared formula that assumes this layout, since removing it here without that would need a matching update there.

♻️ Possible reclaim (verify host-side smem-size formula first)

-    cursor += SMEM_POST_SIZE + SMEM_COMB_SIZE;
     auto smem_rc = reinterpret_cast<nv_bfloat16*>(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16, single-buffered

🤖 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` at line 240,
Remove the unused SMEM_POST_SIZE and SMEM_COMB_SIZE reservation from the cursor
advancement in the fused TF32 pmap GEMM kernel, since post_mix and comb_mix are
loaded through loadCoefficientRow. Before finalizing, locate the companion
host-side launcher’s dynamic shared-memory calculation and update its shared
layout/size formula consistently so allocation matches the kernel.

🤖 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 441-458: Make the local buffers used by the float4 stores in the
fused kernel explicitly 16-byte aligned, preferably by declaring pm_u, pm_l,
cm_u, and cm_l with alignas(16) while preserving their existing float-array
layout and indexing. Keep the loadCoefficientRow calls and subsequent consumers
unchanged.
- Line 240: Remove the unused SMEM_POST_SIZE and SMEM_COMB_SIZE reservation from
the cursor advancement in the fused TF32 pmap GEMM kernel, since post_mix and
comb_mix are loaded through loadCoefficientRow. Before finalizing, locate the
companion host-side launcher’s dynamic shared-memory calculation and update its
shared layout/size formula consistently so allocation matches the kernel.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ec358972-5b52-435b-aac3-8f144996f774

📥 Commits

Reviewing files that changed from the base of the PR and between f4e692d and c878fac.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh

[coderabbitai]

🧹 Nitpick comments (1)

cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh (1)

240-240: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Release the now-unused coefficient SMEM region.

SMEM_POST_SIZE + SMEM_COMB_SIZE still reserves 5,120 bytes per CTA although this kernel no longer allocates or reads smem_post/smem_comb. Remove the cursor gap and reduce the corresponding dynamic-SMEM launch allocation.

Proposed change

-    constexpr uint32_t SMEM_POST_SIZE = BLOCK_M * HC_MULT * sizeof(float);
-    constexpr uint32_t SMEM_COMB_SIZE = BLOCK_M * HC_MULT * HC_MULT * sizeof(float);
...
-    // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)]
+    // SMEM layout: [cd, B stages, res stages, x stages, rc (HC_MULT slices)]
...
-    cursor += SMEM_POST_SIZE + SMEM_COMB_SIZE;
     auto smem_rc = reinterpret_cast<nv_bfloat16*>(cursor);

🤖 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` at line 240,
Remove the SMEM cursor advance for the obsolete smem_post/smem_comb region near
the affected kernel setup, and reduce the corresponding dynamic shared-memory
launch allocation by SMEM_POST_SIZE + SMEM_COMB_SIZE. Ensure cursor layout and
allocation remain consistent without reserving or accessing that unused region.

🤖 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`:
- Line 240: Remove the SMEM cursor advance for the obsolete smem_post/smem_comb
region near the affected kernel setup, and reduce the corresponding dynamic
shared-memory launch allocation by SMEM_POST_SIZE + SMEM_COMB_SIZE. Ensure
cursor layout and allocation remain consistent without reserving or accessing
that unused region.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: dbf530ed-eea6-4e0e-82a8-f909b89d45e8

📥 Commits

Reviewing files that changed from the base of the PR and between c878fac and bd9cf29.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh (1)

202-233: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win

smem_post/smem_comb are now dead shared-memory allocations in this kernel.

SMEM_POST_SIZE/SMEM_COMB_SIZE are still reserved and smem_post/smem_comb still declared here, but since the coefficient-loading prologue was removed (line 289) in favor of direct register loads via load_row4 (lines 439-459), these two pointers are never read anywhere else in fused_tf32_pmap_gemm_rout_atomic_impl. This wastes 5,120 B of dynamic shared memory per CTA (1,024 B + 4,096 B) that could otherwise improve occupancy or be reallocated — directly relevant to this PR's stated optimization goal. It will also likely trigger unused-variable warnings.

♻️ Proposed fix: remove the dead allocation

     cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG;
-    auto smem_post = reinterpret_cast<float*>(cursor);
-    cursor += SMEM_POST_SIZE;
-    auto smem_comb = reinterpret_cast<float*>(cursor);
-    cursor += SMEM_COMB_SIZE;
     auto smem_rc = reinterpret_cast<nv_bfloat16*>(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16, single-buffered

🤖 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
202 - 233, Remove the unused SMEM_POST_SIZE and SMEM_COMB_SIZE constants and the
smem_post/smem_comb declarations and cursor increments in
fused_tf32_pmap_gemm_rout_atomic_impl. Keep the shared-memory layout and
subsequent smem_rc allocation correctly contiguous after smem_x_stg, eliminating
the 5,120-byte dead allocation.

🤖 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 `@cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh`:
- Around line 202-233: Remove the unused SMEM_POST_SIZE and SMEM_COMB_SIZE
constants and the smem_post/smem_comb declarations and cursor increments in
fused_tf32_pmap_gemm_rout_atomic_impl. Keep the shared-memory layout and
subsequent smem_rc allocation correctly contiguous after smem_x_stg, eliminating
the 5,120-byte dead allocation.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ece6c2e2-2fac-4f1b-9047-54e684028a0b

📥 Commits

Reviewing files that changed from the base of the PR and between bd9cf29 and 6e214bd.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh

[MengmSun]

/bot run --disable-fail-fast

[mingyangHao]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61505 [ run ] triggered by Bot. Commit: 6e214bd Link to invocation

[tensorrt-cicd]

PR_Github #61505 [ run ] completed with state SUCCESS. Commit: 6e214bd
/LLM/main/L0_MergeRequest_PR pipeline #49724 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[MengmSun]

/bot run

[jiaganc]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61719 [ run ] triggered by Bot. Commit: 4a52d3d Link to invocation

[tensorrt-cicd]

PR_Github #61719 [ run ] completed with state SUCCESS. Commit: 4a52d3d
/LLM/main/L0_MergeRequest_PR pipeline #49924 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[jiaganc]

/bot run

[tensorrt-cicd]

PR_Github #61730 [ run ] triggered by Bot. Commit: 4a52d3d Link to invocation

[tensorrt-cicd]

PR_Github #61730 [ run ] completed with state SUCCESS. Commit: 4a52d3d
/LLM/main/L0_MergeRequest_PR pipeline #49936 completed with status: 'SUCCESS'

CI Report

Link to invocation

[mingyangHao]

LGTM

[mingyangHao]

/bot skip --comment 'Func name changing, has tested locally'

[tensorrt-cicd]

PR_Github #61775 [ skip ] triggered by Bot. Commit: fca43ef Link to invocation

[tensorrt-cicd]

PR_Github #61775 [ skip ] completed with state SUCCESS. Commit: fca43ef
Skipping testing for commit fca43ef

Link to invocation