[None][fix] Skip no-op MXFP4 weight padding

[jiaganc]

Description

Avoid materializing a duplicate MXFP4 weight tensor when its dimensions already satisfy the requested row and column alignments.

maybe_pad_for_mxfp4 previously called torch.nn.functional.pad even when every padding width was zero. PyTorch still allocated a new tensor for that no-op, which duplicated already-aligned expert weights during Kimi K3 TP8 model loading and could stall server startup under memory pressure. Return the original tensor for the zero-padding cases while preserving the existing padding behavior for unaligned weights.

Test Coverage

• pre-commit run --files tensorrt_llm/_torch/modules/fused_moe/quantization.py
• Kimi K3 TP8 end-to-end job 359822 on GB300: server healthy, 10/10 benchmark requests successful, GSM8K exact match 0.9644

PR Checklist

• PR description clearly explains what and why.

• PR follows the TRT-LLM coding guidelines to the best of my knowledge.

• No API changes, new dependencies, ownership changes, or architecture changes.

• 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, comment /bot help.

Dev Engineer Review

maybe_pad_for_mxfp4 skips F.pad when both dimensions meet MXFP4 alignment requirements. This avoids unnecessary allocation during Kimi K3 TP8 model loading. Padding remains enabled for unaligned weights.

The MXFP4 TRTLLM-Gen w2 bias scaling now uses out-of-place division. This prevents mutation of shared bias tensors during reloads or retries.

The change is limited to tensorrt_llm/_torch/modules/fused_moe/quantization.py. No public declarations, configuration files, or test-list files changed.

The no-op path should be checked for contiguous-tensor requirements because F.pad previously returned a contiguous tensor.

Verdict: sufficient, with a follow-up check for downstream contiguous-tensor requirements.

QA Engineer Review

No test changes.

[jiaganc]

/bot run

[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: 4bed9726-02c1-41f9-aad0-d1cf0405b2d7

📥 Commits

Reviewing files that changed from the base of the PR and between 8e2e342 and c80388f.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/modules/fused_moe/quantization.py

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

• tensorrt_llm/_torch/modules/fused_moe/quantization.py

---

Walkthrough

Changes

MXFP4 quantization

Layer / File(s)	Summary
Quantization adjustments tensorrt_llm/_torch/modules/fused_moe/quantization.py	maybe_pad_for_mxfp4 skips F.pad when no padding is required. MXFP4 TRTLLM-Gen w2 bias scaling uses an out-of-place assignment.

Estimated code review effort: 1 (Trivial) | ~2 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the MXFP4 padding fix and follows the required ticket and type format.
Description check	✅ Passed	The description explains the issue, solution, test coverage, and checklist status with the required sections.
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.}

[tensorrt-cicd]

PR_Github #63918 [ run ] triggered by Bot. Commit: 8e2e342 Link to invocation

[BowenFu]

The optimization is right — F.pad with all-zero widths really does allocate and copy, and skipping it is worth doing. But the change also silently rewrites the function's contract, and one existing caller depends on the old one.

Before this PR maybe_pad_for_mxfp4 always returned a freshly allocated tensor. After it, when every pad width is zero, it returns the caller's own tensor object. NVFP4TRTLLMGenFusedMoEMethod.load_expert_w2_weight relies on the old guarantee (quantization.py:5006-5010):

w2_weight = maybe_pad_for_mxfp4(w2_weight, self.weight_alignment)

# Divide bias by tp_size as we shard along the hidden dimension.
w2_weight /= module.tp_size          # in-place

On main that /= lands on the private F.pad copy. After this PR, on the no-op path, it writes straight into the caller's tensor — and the caller is the shared driver at quantization.py:~392/:420, which passes w2_bias = weights["down_proj.bias"][expert_id], i.e. a view into the shared weights dict, and then hands the very same tensor to module._add_raw_shared_weights_for_unmap(...) and maybe_pageout_mmapped_cpu_weights(...).

This isn't a rare branch. That class has weight_alignment = 32 (quantization.py:4821) and the 1-D arm is the bias, whose length is a hidden dimension — a multiple of 32 for every real config — so col_pad_size == 0 essentially always, and the aliasing path is the normal one.

Concrete failure: NVFP4 TRTLLM-gen MoE with module.bias, tp_size > 1, and any second pass over the same weights (EPLB reload, the allow_partial_loading=True RLHF reload path, a retried load) — the w2 bias gets divided by tp_size twice. Silent accuracy loss, no error. Writing into a paged/mmapped checkpoint tensor is also a change in its own right.

The sibling site at :6021 is fine because it uses out-of-place .float() / module.tp_size.

Two ways to close it, either is fine:

• make :5010 out-of-place — w2_weight = w2_weight / module.tp_size; or
• keep the contract and return weight.clone() (or document loudly that the result may alias the input, and audit the callers).

Same question for contiguity, if you want to check it: F.pad always returned a contiguous tensor, and the fused gate/up loaders pass transpose/chunk views (w1_w3_weight.chunk(2, dim=0) off a .transpose(0, 1)). Anything downstream that used to get contiguity for free from the pad — dtype-changing .view(...) in particular — now gets a non-contiguous tensor on the no-op path.

[tensorrt-cicd]

PR_Github #63918 [ run ] completed with state SUCCESS. Commit: 8e2e342
/LLM/main/L0_MergeRequest_PR pipeline #51855 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

[longlee0622]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #63942 [ run ] triggered by Bot. Commit: 8e2e342 Link to invocation

[jiaganc]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #63985 [ run ] triggered by Bot. Commit: c80388f Link to invocation

[tensorrt-cicd]

PR_Github #63942 [ run ] completed with state ABORTED. Commit: 8e2e342

Link to invocation

[BowenFu]

The aliasing issue I raised is closed. w2_weight = w2_weight / module.tp_size (quantization.py:5011) is out-of-place, so the no-op path can no longer divide weights["down_proj.bias"][expert_id] in the shared dict. I audited all 28 maybe_pad_for_mxfp4 call sites at this head — no other one mutates the returned tensor in place.

Also checked and cleared the contiguity worry raised earlier: the eight .view(dst.dtype) sites all reinterpret one-byte storage as another one-byte dtype (uint8 / float8_e4m3fn, and float4_sf_dtype is uint8), and view(dtype) only imposes stride requirements when item sizes differ — so dropping the F.pad copy can't break them. The new guard also preserves every case that used to pad.

Not blocking on it, but the full L0 run on c80388f (PR_Github #63985) hadn't reported when I looked — worth a glance before merge.

[tensorrt-cicd]

PR_Github #63985 [ run ] completed with state SUCCESS. Commit: c80388f
/LLM/main/L0_MergeRequest_PR pipeline #51919 completed with status: 'FAILURE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

[longlee0622]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #64039 [ run ] triggered by Bot. Commit: c80388f Link to invocation

[tensorrt-cicd]

PR_Github #64039 [ run ] completed with state FAILURE. Commit: c80388f
/LLM/main/L0_MergeRequest_PR pipeline #51969 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 --disable-fail-fast

[tensorrt-cicd]

PR_Github #64085 [ run ] triggered by Bot. Commit: c80388f Link to invocation

[tensorrt-cicd]

PR_Github #64085 [ run ] completed with state FAILURE. Commit: c80388f
/LLM/main/L0_MergeRequest_PR pipeline #52012 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

[longlee0622]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #64143 [ run ] triggered by Bot. Commit: c80388f Link to invocation

[tensorrt-cicd]

PR_Github #64143 [ run ] completed with state FAILURE. Commit: c80388f
/LLM/main/L0_MergeRequest_PR pipeline #52061 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 --disable-fail-fast

[tensorrt-cicd]

PR_Github #64258 [ run ] triggered by Bot. Commit: c80388f Link to invocation