[Bugfix] Canonicalize FP8 weight layout to (K, N) at the source

[mgoin]

Summary

Fixes #44110 at the root cause, as an alternative to #44113.

MarlinFP8ScaledMMLinearKernel.process_weights_after_loading previously
tried to detect whether its incoming weight was (N, K) or (K, N) and
transpose accordingly. The shape-based heuristic was a no-op when
N == K, silently corrupting square layers. #44113 swaps the heuristic
from shape to is_contiguous(), but that just trades one fragile
implicit contract in the kernel for another.

The real issue is that the kernel boundary has no agreed-on layout:

• CutlassFP8ScaledMMLinearKernel expects (K, N) and does not transpose.
• ModelOptFp8{,PcPt}LinearMethod already pre-transposes to (K, N).
• Fp8LinearMethod (non-marlin), Fp8OnlineLinearMethod (non-marlin) pre-transpose to (K, N).
• Fp8LinearMethod (use_marlin) and CompressedTensorsW8A16Fp8 skip the transpose and let Marlin guess.

This PR makes every FP8 linear caller canonicalize to (K, N) before
delegating, and removes the detection heuristic from Marlin entirely.
This is the canonicalization step that the TODO referenced in #33314
was asking for, scoped to the FP8 paths that share this kernel.

Changes

• Fp8LinearMethod (use_marlin, non-block): transpose before delegating.
• Fp8OnlineLinearMethod: collapse the marlin/non-marlin branches into a single transpose + delegate.
• CompressedTensorsW8A16Fp8 (non-block): transpose before delegating.
• MarlinFP8ScaledMMLinearKernel: drop the conditional transpose from the non-block branch.

Net diff: 3 files, +11 / -29.

Why this is preferable to #44113

• Removes the implicit "kernel detects layout from weight metadata" contract entirely.
• Makes the FP8 callers consistent with each other and with the Cutlass kernel's existing expectation.
• The is_contiguous() switch in [Bugfix] Fix MarlinFP8 weight transpose silently skipped for square matrices (N==K) #44113 still breaks if any future caller pre-transposes and then calls .contiguous(), or loads a checkpoint that happens to land non-contiguous.

Test plan

• Square (N==K, 4096×4096) regression case verified end-to-end on A40 (sm_86) with VLLM_TEST_FORCE_FP8_MARLIN=1, both the checkpoint-layout path (CompressedTensors-style) and the pre-transposed path (ModelOpt-style). Relative error < 0.005 in all cases.
• Non-square shapes (1024×4096, 4096×12800) verified for both paths.
• pre-commit run --files on the changed files — all hooks pass (ruff, mypy, etc.).
• Existing tests/quantization/test_fp8.py and tests/evals/gsm8k/test_gsm8k_correctness.py runs in CI.

AI assistance (Claude) was used to draft the change; I (mgoin) reviewed and tested every line.

Co-authored-by: Claude noreply@anthropic.com

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Hi @mgoin, 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.