[charon14]

Dev Engineer Review

The kernel adds a compile-time epilogue for GroupedWithOffset when BLOCK_N contains whole groups of eight 16-byte chunks. That path reads B scales from global memory and uses XOR-swizzled shared-memory staging, warp-group barriers, and coalesced global stores. It checks M boundaries and N tails. Other configurations retain the existing epilogue.

The change affects scale loading and output synchronization for the selected path. Validate scale indexing, partial M/N tiles, and warp-group synchronization. The supplied summary reports bitwise-identical outputs and H20 performance gains, but the test results are not independently available here.

QA Engineer Review

No test code changed. tests/integration/test_lists/waives.txt adds two SKIP entries: test_flashinfer_gdn_decode.py and one test_mm_encoder_standalone.py parameterization. The diff does not establish changes to CI or manual-QA test membership. The stated kernel validation and compute-sanitizer checks are author-reported; TensorRT-LLM CI was not run locally. Coverage verdict: needs follow-up, because this PR adds no test coverage for the changed kernel path.

Per-File QA Perspective

• cpp/include/tensorrt_llm/deep_gemm/fp8_gemm_impl.cuh: Verify selected grouped-with-offset cases, scale loading, boundary rows, N tails, and output correctness. Confirm unaffected paths continue to use the original epilogue.
• tests/integration/test_lists/waives.txt: Two unrelated tests are newly waived with bug references. These entries change waiver behavior; they do not add coverage for the kernel changes.

Description

On Hopper, the CUTLASS fused MoE with FP8 block scales runs both expert GEMMs through fp8_grouped_gemm_run → grouped_gemm_dispatch, which uses the JIT-compiled DeepGEMM fp8_gemm_kernel (GemmType::GroupedWithOffset) once expected_m (the token count) reaches 32, or 64 on 78-SM GPUs. For MoE shapes it picks 128×128 tiles, so two math warp groups share each output tile. Much of this kernel's time goes to a fixed cost per output tile, which weighs most when K is short. On H20, the FC2 GEMM of a Qwen3-Next-style layer (N=2560, K=640, 128 local experts, ~1175 rows per expert) reaches 66% of FP8 peak (valid rows only). Two things cost the most:

1. Bank conflicts in the STSM write-back. smem_d is a plain row-major tile (256-byte rows for BLOCK_N = 128) because the TMA store uses SWIZZLE_NONE. Each 8×8 STSM matrix covers the same 16-byte column of 8 consecutive rows, so every STSM is an 8-way bank conflict (about 9% of that GEMM's time).
2. Two 256-thread barriers per tile. One follows the staging of the tile's B scales in shared memory, and one precedes the TMA store, after which thread 0 waits until the store has read smem_d. They keep the two math warp groups in lockstep (removing the tile-start barrier with the B-scale staging saves about 7%, removing the store wait about 5%).

This PR changes only the non-swap-AB kernel, and only on the GroupedWithOffset path with BLOCK_N a multiple of 64 (64 or 128 of the 16..128 candidates):

• B scales are read with __ldg for each K block before waiting for that block's TMA data and are used after its WGMMA, so the load hides behind the MMA. They are no longer staged in shared memory, and the tile-start barrier is gone.
• Each math warp group STSMs its own 64 rows into its half of smem_d, with 16-byte chunks XOR-swizzled by row % 8 (conflict-free), syncs only its own 128 threads, and copies the rows out with coalesced 16-byte global stores, masking rows past the expert (m_boundary) and columns past N. A second sync of the same 128 threads protects the rows from the next tile's STSM. There is no TMA store, no store wait and no synchronization between the two math warp groups.

Normal, StridedBatched, other BLOCK_N values (e.g. 120) and the swap-AB kernel are unchanged. Outputs are bitwise identical.

The same file is vendored in FlashInfer, where the change is proposed in flashinfer-ai/flashinfer#5664 with the full data. Kernel time on H20 (78 SMs), CUDA 13.3, 1175 rows per expert, tile config chosen by get_best_gemm_config (128×128, 5 stages here, also on 132 SMs), L2 flushed before each launch, median of 5 interleaved rounds:

GEMM (N × K, experts)	main (ms)	this PR (ms)	change
Qwen3-Next-style, EP4 FC1: 1280 × 2560, 128	4.130	3.820	−7.5%
Qwen3-Next-style, EP4 FC2: 2560 × 640, 128	2.536	1.918	−24.4%
Qwen3-30B-A3B FC1: 1536 × 2048, 128	4.072	3.744	−8.1%
Qwen3-30B-A3B FC2: 2048 × 768, 128	2.392	1.978	−17.3%
Qwen3-235B-A22B FC1: 3072 × 4096, 128	15.267	14.421	−5.5%
Qwen3-235B-A22B FC2: 4096 × 1536, 128	8.347	7.305	−12.5%
DeepSeek-V3, EP8 FC1: 4096 × 7168, 32	8.731	8.511	−2.5%
DeepSeek-V3, EP8 FC2: 7168 × 2048, 32	4.712	4.312	−8.5%

The kernel compiled from this PR's header produces the same SASS as the FlashInfer version these numbers were measured with. End to end, the change raised serving throughput of Qwen3.8-Flash-Next-FP8 on 4×H20 (vLLM with FlashInfer's copy of this kernel) by 2.8%.

JIT cache note: with the default NVRTC path compiled kernels are not written to disk, so they pick up the change. With TRTLLM_DG_JIT_USE_NVCC=1 or TRTLLM_DG_JIT_DUMP_CUBIN=1, cubins under $TRTLLM_DG_CACHE_DIR/cache are named only by shape and tile config and would keep being loaded until that directory is cleared.

Test Coverage

• tests/unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu with FP8_BLOCK_SCALES on the CUTLASS backend (Hopper) runs this kernel. I could not run the TensorRT-LLM test suite locally (no TensorRT-LLM build on this machine); please trigger CI.
• Kernel-level validation of this header, with the kernel built exactly as deep_gemm::jit::Compiler builds it and host code mirroring runGemm: 392 cases over BLOCK_M 64/128 × BLOCK_N 64/128 (plus 120/112/96/16, which keep the original code), N with and without a partial last block (2560, 1280, 2112, 2080, 1000 and long-K shapes up to K = 7168), 1 to 128 experts, and per-expert row counts including 0, 1, tile-boundary values ±1, runs of empty experts and partial last tiles, with NaN in the stale rows the tiles read past each expert. Every valid output element is bitwise equal to an exact reference and to the original kernel, and no row past the valid ones is written.
• compute-sanitizer memcheck, racecheck, synccheck and initcheck: no errors or hazards; 1000 back-to-back launches give identical outputs.
• Through FlashInfer's cutlass_fused_moe (same kernel), outputs are bitwise identical before and after the change on 16 (config, batch) cases, and a batch large enough to use this kernel matches the same tokens run 16 at a time on the swap-AB kernel bit for bit.

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.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 26279694-1edb-4000-a5ff-0f71c0b116bf

📥 Commits

Reviewing files that changed from the base of the PR and between f08592a and ba06804.

📒 Files selected for processing (1)

• cpp/include/tensorrt_llm/deep_gemm/fp8_gemm_impl.cuh

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The FP8 GEMM kernel adds a warp-group-local epilogue for qualifying grouped-with-offset tiles. It loads B scales directly from global memory on that path and retains the existing shared-memory output epilogues for other configurations.

Changes

FP8 GEMM epilogue

Layer / File(s)	Summary
Path selection and B-scale loading cpp/include/tensorrt_llm/deep_gemm/fp8_gemm_impl.cuh	A compile-time predicate enables the local epilogue for qualifying GroupedWithOffset tiles. That path reads B scales from global memory for each K stage. Other configurations retain shared-memory scale staging.
Local and retained output epilogues cpp/include/tensorrt_llm/deep_gemm/fp8_gemm_impl.cuh	The local epilogue stages bfloat16 output in XOR-swizzled shared memory, then writes bounded 16-byte chunks to global memory. Other configurations use the existing output path, including its TMA and direct-store cases.

Priority: ➖ Normal

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

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant MathWarpGroup
  participant SharedMemory
  participant GlobalMemory
  MathWarpGroup->>SharedMemory: Stage bfloat16 output in XOR-swizzled layout
  MathWarpGroup->>GlobalMemory: Write bounded 16-byte output chunks

Loading

Suggested reviewers: bowenfu

Merge Risk: ⚪ Minimal · up to ba068

No actionable issue remains from this review. The PR is mergeable subject to normal checks.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the performance optimization, its target kernel, and the applicable warp-group-local write-back path. It follows the repository format with [None][perf].
Description check	✅ Passed	The description is complete and relevant. It explains the motivation, implementation scope, performance results, correctness validation, test coverage, limitations, and checklist status. It also clear…
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0…
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 a new PR

---

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