[https://nvbugs/6525059][fix] 128KiB-align TMA-OOB MoE workspace buffers on Blackwell

[sunnyqgg]

Description

On Blackwell (SM100/SM103) the FP8 block-scale MoE GEMM kernels drive their A/B/C loads/stores through the rgTma TMA-OOB descriptors, which run in OOB_ADDR_GEN_MODE_BASE_128kB: the reshaped descriptor's addressable space is >= 128 KiB, so the TMA unit may probe up to a 128 KiB-aligned boundary past the descriptor base even though software never reads there. The backing workspace must therefore start 128 KiB-aligned and own >= 128 KiB of mapped memory after that base.

maybeGetMinTokenCount already sizes gemm1_output / activation_output / gemm2_output to >= 128 KiB, but at::detail::empty_cuda returns 512 B-aligned slices of a shared caching-pool block. A workspace that lands near a pool-block end lets the 128 KiB-mode OOB probe cross into an unmapped neighbor, giving a sporadic cudaErrorIllegalAddress during warmup / small-decode batches (tile_tokens_dim 8 and 16). Same Blackwell TMA-OOB family as NVBug 5435459; the trtllm-gen test harness dodges it by allocating dedicated 128 KiB-padded cudaMalloc.

Fix: allocate the three rgTma workspace buffers 128 KiB-aligned with 128 KiB tail headroom. Scaling-factor buffers are unaffected (DeepSeek-FP8 loads them via LDGSTS).

Resolves NVBug 6525059 (tileN=8) and NVBug 6602176 (tileN=16, same root cause; its failing case test_moe_backend ... quant=FP8_BLOCK_SCALES runs the same fp8BlockScaleMoe.cpp path).

Root-cause evidence

• coredump: fault PC at cp.async.bulk.tensor (GEMM2 activation load), coords valid.
• dmesg: Xid 31 ... GPCCLIENT_T1 ... FAULT_PDE VIRT_READ at 128 KiB-aligned addresses.
• instrumentation: the faulting activation buffer was handed to TMA at ptr%128KiB=91136 (misaligned) in oob128kMode=1; after the fix every rgTma (dim=4) descriptor is 128 KiB-aligned and the run passes with no new Xid 31.

Test Coverage

• accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_dummy_load_format — unwaived on B200/B300/GB200/GB300 (was nvbugs/6525059).
• unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" — unwaived (was nvbugs/6602176).

Follow-ups (not in this PR)

• Same 128 KiB-aligned-workspace treatment should be applied to fp4BlockScaleMoe.cpp, mxFp4BlockScaleMoe.cpp, fp8PerTensorScaleMoe.cpp (same allocation pattern, same latent TMA-OOB exposure).

Dev Engineer Review

• Updated the three rgTma workspace buffers for Blackwell FP8 block-scale MoE GEMM.
• Added 128 KiB alignment and mapped tail space for TMA out-of-bounds descriptor access.
• Retained backing tensors until execution completes.
• Kept scaling-factor buffers unchanged.
• No public API changes are present.
• The allocation changes address NVBugs 6525059 and 6602176, including warmup and small-decode cases with tile sizes 8 and 16.

QA Engineer Review

• Removed five platform-specific waivers for TestQwen3_30B_A3B::test_dummy_load_format.
• Affected platforms are B200, B300, GB200, and GB300.
• Removed the generic TRTLLM MoE backend waiver.
• No test code changed.
• The test-list changes unwaive coverage for the fixed kernel behavior.
• Verdict: needs follow-up because CBTS coverage data is unavailable.

[sunnyqgg]

/bot run --disable-fail-fast

[coderabbitai]

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: d7d49d98-a3f6-4164-9bb5-1f50386a037d

📥 Commits

Reviewing files that changed from the base of the PR and between 6be4f23 and 115194e.

📒 Files selected for processing (2)

• cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

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

• cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp

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

---

Walkthrough

The FP8 block-scale MoE path now uses aligned intermediate buffers with tail headroom and passes their pointers through the workspace. Five platform-specific integration test waivers are removed.

Changes

FP8 MoE buffer alignment

Layer / File(s)	Summary
Aligned intermediate buffer allocation cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp	GEMM1, activation, and GEMM2 buffers now use 128 KiB-aligned CUDA byte allocations with tail headroom. Retained backing tensors keep the allocations valid during execution. The workspace uses the aligned pointers.
Integration test waiver removal tests/integration/test_lists/waives.txt	Five platform-specific TestQwen3_30B_A3B::test_dummy_load_format and generic TRTLLM MoE backend skip entries are removed.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 11519

This localized change aligns and pads Blackwell MoE workspace buffers to prevent illegal-memory accesses, with targeted regression coverage; no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: xxi-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the NVBug, fix type, 128 KiB alignment, TMA-OOB issue, MoE workspace buffers, and Blackwell scope.
Description check	✅ Passed	The description explains the root cause, solution, affected NVBugs, test coverage, and follow-up scope. It is mostly complete, although it does not include the template's PR Checklist section.
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.}

[sunnyqgg]

/bot run --disable-fail-fast

[sunnyqgg]

/bot run --disable-fail-fast

[xxi-nv]

Could you help to add a corresponding test case in the test_moe_backend.py？

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[sunnyqgg]

Could you help to add a corresponding test case in the test_moe_backend.py？

The exact failing case for nvbugs/6602176 already exists in test_moe_backend.py
(e8_k1_h512_i512-seq=8-...-quant=FP8_BLOCK_SCALES-routing=Renormalize) — it was skipped by the -k
"TRTLLM" waive, and this PR re-enables it (along with test_dummy_load_format for 6525059), so the
corresponding coverage is restored by the unwaive itself.

[sunnyqgg]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #68955 [ run ] triggered by Bot. Commit: 115194e Link to invocation

[tensorrt-cicd]

PR_Github #68955 [ run ] completed with state FAILURE. Commit: 115194e
/LLM/main/L0_MergeRequest_PR pipeline #56335 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

[sunnyqgg]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #69016 [ run ] triggered by Bot. Commit: 115194e Link to invocation

[tensorrt-cicd]

PR_Github #69016 [ run ] completed with state SUCCESS. Commit: 115194e
/LLM/main/L0_MergeRequest_PR pipeline #56389 completed with status: 'SUCCESS'

CI Report

Link to invocation