fix(moe): pad BF16 TRTLLM-Gen intermediates to 128 KiB by nvpohanh · Pull Request #4319 · flashinfer-ai/flashinfer

nvpohanh · 2026-08-03T05:10:35Z

[by Codex]

Summary

pad the BF16 TRTLLM-Gen MoE GEMM1 and GEMM2 intermediate allocations to at least 128 KiB
reuse Routing::maybeGetMinTokenCount, matching the existing FP8/MXFP launchers
calculate the GEMM1 and GEMM2 row counts independently because their hidden dimensions can differ

Why

On Blackwell, the TRTLLM-Gen BMM2 kernel can use TMA BASE_128KB address generation based on the logical tensor-map shape. With a small eager decode batch, the BF16 launcher currently allocates only the logical 16/32-KiB intermediate. That leaves less than 128 KiB mapped after the activation base and can produce an illegal memory access in UTMALDG.4D.

The quantized launchers already call maybeGetMinTokenCount for both intermediates. Applying the same allocation policy to the BF16 launcher gives the TMA transaction the required backing without changing the logical token count passed to routing or GEMM.

Validation

pre-commit run --all-files
B200x8 patched-JIT compile/load preflight
six consecutive Miles/SGLang rollout runs with the original reproducer setting:
- pinned image and model revision
- FlashInfer TRTLLM backend, TP8
- four train/rollout steps per run
- 256 one-token post-update health probes per step boundary
- no descriptor-bit override, TMA-OOB disablement, extra allocation-padding knob, or synchronization workaround

Result: six consecutive valid passes on B200x8. Every run exited normally, completed all four steps, passed the diagnostic evidence gate, and had zero IMA/sanitizer/segfault matches.

Summary by CodeRabbit

Bug Fixes
- Improved BF16 Mixture-of-Experts processing by ensuring intermediate buffers have sufficient capacity for required token mappings.
- Increased reliability for GEMM-based operations across varying tensor dimensions while preserving expected output shapes and device placement.
- Helps prevent capacity-related processing failures in workloads with larger or uneven token distributions.

Signed-off-by: Po-Han Huang <pohanh@nvidia.com>

gemini-code-assist · 2026-08-03T05:10:38Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

coderabbitai · 2026-08-03T05:10:43Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: c8d55d7c-436b-4d0a-999b-5e2842e76c15

📥 Commits

Reviewing files that changed from the base of the PR and between c8a3734 and d428307.

📒 Files selected for processing (1)

csrc/trtllm_fused_moe_kernel_launcher.cu

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

csrc/trtllm_fused_moe_kernel_launcher.cu

📝 Walkthrough

Walkthrough

BF16 MoE GEMM1 and GEMM2 output buffers now use separate minimum token capacities derived from output dimensions and BF16 element size.

Changes

BF16 MoE capacity allocation

Layer / File(s)	Summary
Calculate minimum buffer capacities `csrc/trtllm_fused_moe_kernel_launcher.cu`	`prepare_moe` calls `Routing::maybeGetMinTokenCount` for GEMM1 and GEMM2, then allocates both BF16 buffers with the adjusted capacities.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

flashinfer-ai/flashinfer#4281: Both changes modify the fused MoE kernel launcher and its GEMM output-buffer allocation logic.

Suggested labels: run-ci

Suggested reviewers: aleozlx, sricketts, yzh119

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description clearly explains the change, motivation, implementation, and validation results, although it does not use every template heading.
Title check	✅ Passed	The title concisely identifies the BF16 TRTLLM-Gen MoE intermediate allocation padding change.
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 💡 1

🛠️ Fix failing CI checks 💡

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

gemini-code-assist · 2026-08-03T05:25:29Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

nvpohanh · 2026-08-03T05:25:37Z

/bot run

flashinfer-bot · 2026-08-03T05:26:32Z

GitLab MR !1105 has been created, and the CI pipeline #60808548 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-08-04T00:08:28Z

[FAILED] Pipeline #60808548 — 11/18 executed test jobs passed

Compared with nightly #60712014.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Notes
5090	⚠️ Infra	⚠️ Infra	Infrastructure: `CI infrastructure failure` (2 jobs; CUDA 12.9, CUDA 13.0)
B300	✅ Pass	✅ Pass	—
GB200	✅ Pass	❌ New	New: `tests.comm.test_ulysses_communicator` (1 failure; CUDA 13.0)
GB300	✅ Pass	✅ Pass	—
H100	❌ New	🟡 Old	New: `tests.attention.test_hopper` (1818 failures; CUDA 12.9) Old: `tests.gemm.test_sm_constraint_gemm` (55296 failures; CUDA 12.9, CUDA 13.0) Old: `tests.gemm.test_mm_fp4` (47916 failures; CUDA 12.9, CUDA 13.0) … and 159 more
RTX Pro 6000 Blackwell	✅ Pass	✅ Pass	—

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 4/6 passed

GPU	CUDA 12.9	CUDA 13.0	Notes
B300 (multi-GPU)	❔ Failed	❔ Failed	—
GB200 (multi-node)	✅ Pass	✅ Pass	—
GB300 (multi-node)	✅ Pass	✅ Pass	—

Failure details

New relative to nightly (attribution uncertain)

tests.attention.test_hopper — 1818 failures on H100 / CUDA 12.9
- torch.AcceleratorError: CUDA error: unknown error Search for 'cudaErrorUnknown' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.…
tests.comm.test_ulysses_communicator — 1 failure on GB200 / CUDA 13.0
- failed on setup with "torch.distributed.DistNetworkError: The server socket has failed to listen on any local network address. port: 37005, useIpv6: false, code: -98, name: EADD…

Pre-existing failures

tests.gemm.test_sm_constraint_gemm — 55296 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.gemm.test_mm_fp4 — 47916 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.attention.test_rope — 30254 failures on H100 / CUDA 12.9
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.gemm.test_mm_bf16 — 15126 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.attention.test_trtllm_gen_attention_prefill — 14804 failures on H100 / CUDA 12.9
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.attention.test_trtllm_gen_mla — 14075 failures on H100 / CUDA 12.9
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.attention.test_sliding_window — 11792 failures on H100 / CUDA 12.9
- failed on setup with "RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program star…
tests.moe.test_trtllm_gen_routed_fused_moe — 7418 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.gdn.test_prefill_delta_rule — 7088 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.gemm.test_bmm_fp8 — 6912 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.gemm.test_groupwise_scaled_gemm_mxfp4 — 6912 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
tests.utils.test_norm — 5774 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- RuntimeError: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the availab…
… and 148 more failing test groups

Timeouts, infrastructure, or incomplete jobs

Infrastructure: CI infrastructure failure — 5090 / CUDA 12.9, 5090 / CUDA 13.0
- Jobs: unit_test_5090: [cu130], unit_test_5090: [cu129]
Infrastructure: GPU environment became unavailable — H100 / CUDA 12.9, H100 / CUDA 13.0
- Jobs: unit_test_h100: [cu130], unit_test_h100: [cu129]

nvpohanh · 2026-08-04T00:16:21Z

This will fix #3427

flashinfer-bot · 2026-08-04T00:17:03Z

GitLab MR !1105 has been updated with latest changes, and the CI pipeline #60945480 is currently running. I'll report back once the pipeline job completes.

nvpohanh · 2026-08-05T02:08:24Z

/bot run

flashinfer-bot · 2026-08-05T02:08:43Z

GitLab MR !1105 has been updated with latest changes, and the CI pipeline #61135895 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-08-05T20:31:13Z

[FAILED] Pipeline #61135895 — 8/18 executed test jobs passed

Compared with nightly #60831563.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Notes
5090	❔ Unknown	❔ Unknown	Unknown: `script failed before producing a JUnit report` (2 jobs; CUDA 12.9, CUDA 13.0)
B300	❌ New	❌ New	New: `tests.trace.test_template_registry` (2 failures; CUDA 12.9, CUDA 13.0) Infrastructure: `test infrastructure interrupted the job` (2 jobs; CUDA 12.9, CUDA 13.0)
GB200	⚠️ Infra	⚠️ Infra	Infrastructure: `test infrastructure interrupted the job` (2 jobs; CUDA 12.9, CUDA 13.0)
GB300	⚠️ Infra	❌ New	New: `tests.gemm.test_groupwise_scaled_gemm_fp8` (1 failure; CUDA 13.0) Infrastructure: `test infrastructure interrupted the job` (2 jobs; CUDA 12.9, CUDA 13.0)
H100	❌ New	❌ New	New: `tests.utils.test_nvfp4_per_token_quantize_large_m` (8 failures; CUDA 12.9, CUDA 13.0) New: `tests.trace.test_template_registry` (2 failures; CUDA 12.9, CUDA 13.0)
RTX Pro 6000 Blackwell	❌ New	❌ New	New: `tests.trace.test_template_registry` (2 failures; CUDA 12.9, CUDA 13.0)

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 4/6 passed

GPU	CUDA 12.9	CUDA 13.0	Notes
B300 (multi-GPU)	❔ Failed	❔ Failed	—
GB200 (multi-node)	✅ Pass	✅ Pass	—
GB300 (multi-node)	✅ Pass	✅ Pass	—

Failure details

New relative to nightly (attribution uncertain)

tests.utils.test_nvfp4_per_token_quantize_large_m — 8 failures on H100 / CUDA 12.9, H100 / CUDA 13.0
- cutlass.base_dsl.compiler.CompilerDiagnosticError: #x1B[91merror:#x1B[0m NVVM backend compilation failed #x1B[91merror:#x1B[0m libNVVM failed while compiling generated device IR…
tests.trace.test_template_registry — 6 failures on B300 / CUDA 12.9, B300 / CUDA 13.0, H100 / CUDA 12.9, H100 / CUDA 13.0, RTX Pro 6000 Blackwell / CUDA 12.9, RTX Pro 6000 Blackwell / CUDA 13.0
- AssertionError: assert {'flashinfer....cat_ops', ...} == {'flashinfer....cat_ops', ...} Extra items in the right set: 'flashinfer.kda' Full diff: { 'flashinfer.activation', 'fla…
tests.gemm.test_groupwise_scaled_gemm_fp8 — 1 failure on GB300 / CUDA 13.0
- AssertionError: Tensor-likes are not close! Mismatched elements: 124 / 8192 (1.5%) Greatest absolute difference: 1.71875 at index (8, 42) (up to 0.01 allowed) Greatest relative…

Timeouts, infrastructure, or incomplete jobs

Infrastructure: test infrastructure interrupted the job — B300 / CUDA 12.9, B300 / CUDA 13.0, GB200 / CUDA 12.9, GB200 / CUDA 13.0, GB300 / CUDA 12.9, GB300 / CUDA 13.0
- Jobs: unit_test_b300: [cu130], unit_test_b300: [cu129], unit_test_gb200: [cu130], unit_test_gb200: [cu129], unit_test_gb300: [cu130], unit_test_gb300: [cu129]
Unknown: script failed before producing a JUnit report — 5090 / CUDA 12.9, 5090 / CUDA 13.0
- Jobs: unit_test_5090: [cu130], unit_test_5090: [cu129]

nvpohanh · 2026-08-06T11:47:32Z

/bot run

flashinfer-bot · 2026-08-06T11:48:29Z

GitLab MR !1105 has been updated with latest changes, and the CI pipeline #61387766 is currently running. I'll report back once the pipeline job completes.

aleozlx · 2026-08-06T20:47:43Z

/bot run tests/moe

flashinfer-bot · 2026-08-06T20:48:17Z

GitLab MR !1105 has been created, and the CI pipeline #61461786 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-08-07T04:39:50Z

[FAILED] Pipeline #61461786 — 14/18 executed test jobs passed

Compared with nightly #61182354.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Notes
5090	✅ Pass	✅ Pass	—
B300	✅ Pass	✅ Pass	—
GB200	❔ Failed	❔ Failed	—
GB300	✅ Pass	✅ Pass	—
H100	✅ Pass	✅ Pass	—
RTX Pro 6000 Blackwell	✅ Pass	✅ Pass	—

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 4/6 passed

GPU	CUDA 12.9	CUDA 13.0	Notes
B300 (multi-GPU)	❔ Failed	❔ Failed	—
GB200 (multi-node)	✅ Pass	✅ Pass	—
GB300 (multi-node)	✅ Pass	✅ Pass	—

No individual test or infrastructure failures could be extracted.

b8zhong · 2026-08-09T01:34:20Z

Hi, is there anything blocking the merge of this PR? It occurs during GLM 5.2 draft MoE, so ideally I'm wondering if it can be merged soon

nvpohanh · 2026-08-10T02:06:04Z

/bot run tests/moe

flashinfer-bot · 2026-08-10T02:06:58Z

GitLab MR !1105 has been updated with latest changes, and the CI pipeline #61889462 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-08-10T04:32:10Z

[SUCCESS] Pipeline #61889462: 18/18 executed test jobs passed

henrylhtsang · 2026-08-19T23:23:28Z

@nvpohanh dumb question why not consider the cutlass fix for tma descriptor? since it seems like it is the same issue? or am I wrong?

  if (driver_version <= 13010) {
    if (cute::bits_to_bytes(
          cute::cosize(tensor_cwhdn.layout()) *
          cute::sizeof_bits<typename EngineA::value_type>::value) < 131072) {
      reinterpret_cast<uint64_t*>(&tma_desc)[1] &= ~(1llu << 21);
    }
  }

https://github.com/NVIDIA/cutlass/blob/6c68991985ca8b09594ac6fd43abbfd5830c4140/include/cute/atom/copy_traits_sm90_im2col.hpp#L474-L483

nvpohanh · 2026-08-20T06:58:37Z

TBH I think that fix is a little ugly 😹 And padding workspace to 128KiB should be almost harmless

nlima-prog · 2026-08-25T16:58:07Z

Hatch Session

This PR was worked on in a Hatch coding session.


Session	`6493ed5f-d798-40ac-88c2-392844ffd43c`
User	nlima-prog
Environment	Hatch (Local)
Agent	Abacus (Local)
Sandbox	local
Time	2026-08-25T16:58:07.185Z

fix(moe): pad BF16 TRTLLM-Gen intermediates to 128 KiB

bbfaff4

Signed-off-by: Po-Han Huang <pohanh@nvidia.com>

flashinfer-bot added the op: moe label Aug 3, 2026

nvpohanh marked this pull request as ready for review August 3, 2026 05:25

nvpohanh requested review from aleozlx, sricketts and yzh119 as code owners August 3, 2026 05:25

nvpohanh requested review from Aneureka, bkryu, cyx-6, dhiraj113, feih-nv, jiahanc, jimmyzho, kahyunnam, nv-yunzheq, qiching, qsang-nv, saltyminty, samuellees, yongwww and yyihuang as code owners August 3, 2026 05:25

nvpohanh mentioned this pull request Aug 3, 2026

FlashInfer TRTLLM BF16 MoE IMA on B200 with SGLang RL rollout, correlated with t128x8x128 + schPd2x1x2x3 BMM kernel at effbs=1024 #3427

Closed

Merge branch 'main' into fix/bf16-moe-tma-128k

1812142

Merge branch 'main' into fix/bf16-moe-tma-128k

c8a3734

leejnau added the priority: should have (P1) label Aug 5, 2026

Merge branch 'main' into fix/bf16-moe-tma-128k

d428307

aleozlx approved these changes Aug 6, 2026

View reviewed changes

aleozlx added the run-ci label Aug 6, 2026

leejnau added run-ci and removed run-ci labels Aug 7, 2026

Merge branch 'main' into fix/bf16-moe-tma-128k

ce984d8

leejnau merged commit 2fb785c into flashinfer-ai:main Aug 10, 2026
28 of 37 checks passed

coderabbitai Bot mentioned this pull request Aug 11, 2026

perf(moe): optimize CUTLASS fused MoE mem-bound kernels #3761

Merged

5 tasks

janbernloehr mentioned this pull request Aug 14, 2026

fix(trtllm-gen moe): persistent, CUDA-graph-capture-safe routing scratch #4238

Closed

1 task

b8zhong mentioned this pull request Sep 3, 2026

[Bug] GlmMoeDsa (GLM-5.2) FP4 + EAGLE: illegal memory access in flashinfer_trtllm bf16 batched-GEMM (nextn draft MoE); triton nextn default is HIP-gated after #30137 sgl-project/sglang#30209

Open

edwingao28 mentioned this pull request Sep 10, 2026

[AgentX] update B300 GLM image and extend concurrency coverage / 更新 B300 GLM 镜像并扩展并发覆盖 SemiAnalysisAI/InferenceX#2829

Merged

9 tasks

janbernloehr mentioned this pull request Sep 15, 2026

[Bug] IMA in trtllm_bf16_moe autotune sweep at the 1-token tuning profile (B200/sm100f, BF16 dynB GEMM2) #4157

Open

Uh oh!

Conversation

nvpohanh commented Aug 3, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Why

Validation

Summary by CodeRabbit

Uh oh!

gemini-code-assist Bot commented Aug 3, 2026

Uh oh!

coderabbitai Bot commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

gemini-code-assist Bot commented Aug 3, 2026

Uh oh!

nvpohanh commented Aug 3, 2026

Uh oh!

flashinfer-bot commented Aug 3, 2026

Uh oh!

flashinfer-bot commented Aug 4, 2026

Unit Tests

Multi-GPU and Multi-Node Tests — 4/6 passed

New relative to nightly (attribution uncertain)

Pre-existing failures

Timeouts, infrastructure, or incomplete jobs

Uh oh!

nvpohanh commented Aug 4, 2026

Uh oh!

flashinfer-bot commented Aug 4, 2026

Uh oh!

nvpohanh commented Aug 5, 2026

Uh oh!

flashinfer-bot commented Aug 5, 2026

Uh oh!

flashinfer-bot commented Aug 5, 2026

Unit Tests

Multi-GPU and Multi-Node Tests — 4/6 passed

New relative to nightly (attribution uncertain)

Timeouts, infrastructure, or incomplete jobs

Uh oh!

nvpohanh commented Aug 6, 2026

Uh oh!

flashinfer-bot commented Aug 6, 2026

Uh oh!

aleozlx commented Aug 6, 2026

Uh oh!

flashinfer-bot commented Aug 6, 2026

Uh oh!

flashinfer-bot commented Aug 7, 2026

Unit Tests

Multi-GPU and Multi-Node Tests — 4/6 passed

Uh oh!

b8zhong commented Aug 9, 2026

Uh oh!

nvpohanh commented Aug 10, 2026

Uh oh!

flashinfer-bot commented Aug 10, 2026

Uh oh!

flashinfer-bot commented Aug 10, 2026

Uh oh!

Uh oh!

henrylhtsang commented Aug 19, 2026

Uh oh!

nvpohanh commented Aug 20, 2026

Uh oh!

nlima-prog commented Aug 25, 2026

Hatch Session

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

nvpohanh commented Aug 3, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 3, 2026 •

edited

Loading