Pad FP8 Marlin weights to valid thread tiles

[mikekg]

Purpose

FP8 and MXFP8 Marlin repack and GEMM require rank-local N/K extents that
match a supported Marlin thread-tile family.

A fixed padding rule can be valid but overpad some partitions. This change
selects the lower-overhead padded rank-local extent from the two Marlin tile
families that are valid for both small and large batches:

• N multiple 128 with K multiple 64
• N multiple 64 with K multiple 128

The padding is applied after tensor-parallel partitioning, so global
checkpoint tensors are not expanded before TP slicing.

Changes

• Pad FP8 and MXFP8 packed weights before Marlin repack.
• Pad block/group scales consistently with the padded K extent.
• Pad bias before Marlin bias permutation when bias is present.
• Store the padded N/K extents on the layer.
• Use those padded extents for Marlin GEMM.
• Zero-pad activation K for the GEMM call.
• Slice padded output columns back to the logical output extent.

Test Plan

Run Nemotron 3 Super 120B NVFP4 with TP4 on A100. This exercises the
dense FP8 Marlin path that failed in gptq_marlin_repack().

Repro vllm serve command:

  vllm serve nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 \
    --quantization modelopt_mixed \
    --tensor-parallel-size 4 \
    --trust-remote-code \
    --max-model-len 4096 \
    --kv-cache-dtype fp8_e4m3

Test Result

Before this change, the TP4 A100 run failed during dense FP8 Marlin repack:

Selected MarlinFP8ScaledMMLinearKernel
RuntimeError: size_n = 4640 is not divisible by tile_n_size = 64

With this change, the dense FP8 Marlin repack failure no longer occurs. The
same run progressed past the FP8 dense path and then exposed a separate NVFP4
MoE Marlin padding issue:

RuntimeError: size_n = 672 is not divisible by tile_n_size = 64

That later failure is in the NVFP4 MoE path and is separate from the dense
FP8 Marlin padding fixed here. It is addressed by #43908.

After applying both the present PR and #43908, the TP4 NVFP4 serve sanity run completed successfully with Nemotron 3 Super:

TEST=nemo3super
MODEL_REPO=nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4
TP=4
QUANTIZATION=modelopt
server ready after 245s
...
Job Complete - Exit code: 0

This also requires #42610 to pass FP4 weights between Marlin operators when torch.compile() creates compiled execution paths.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[vadiklyutiy]

It's worth to add unit tests

[mikekg]

@vadiklyutiy

It's worth to add unit tests

Done. Tested in tests/model_executor/test_fp8_marlin_padding.py
BTW, because I don't have 4 merged PRs yet, I can't run CI/CD on here, and need PRs tagged as "ready" to run ci/cd.

[vadiklyutiy]

@claude review

[vadiklyutiy]

LGTM

[mergify]

Hi @mikekg, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mikekg]

Hi @mikekg, the pre-commit checks have failed. Please run:

Thank you for flagging this. After rebasing, the files changed by this PR pass pre-commit; the remaining pre-commit run --all-files failure is in .buildkite/scripts/ci-bake-rocm.sh, outside this PR’s diff.

The failing line appears to have been introduced on main by PR #36949 (87954eb50, [ROCm][CI] Optimize ROCm Docker build...). I opened a separate fix attempt here: #44937.

How would you like me to proceed? My preference is to keep this PR focused and handle that unrelated pre-commit failure separately. The CI failures I inspected also appear unrelated to this PR’s diff.

[mergify]

Hi @mikekg, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mikekg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mikekg]

This was merged into omnibus Marlin padding patch #45295 and the issue is now closed