[aoshen02]

Summary

vLLM keeps a storage convention for FP8 linear weights: non-block strategies store layer.weight as (K, N), block stores it as (N, K). The three process_fp8_weight_*_strategy helpers do everything except the final step that establishes it, so each call site ends with a hand-written weight.t() — or, for block, deliberately omits one.

That makes a three-way rule (tensor → transpose, channel → transpose, block → don't) invisible from the signatures and replicated at all six call sites. This PR moves the step into the helpers so the convention has one definition, and separately de-duplicates the MLA kv_b_proj absorption block, which exists in two identical copies.

No behavior change. 6 files, +36/−49.

FP8 weight layout (fp8_utils.py)

Helper	Returns	Call sites that previously transposed by hand
process_fp8_weight_tensor_strategy	(K, N)	2 of 2
process_fp8_weight_channel_strategy	(K, N)	2 of 2
process_fp8_weight_block_strategy	(N, K), unchanged	0 of 2

Each helper's docstring now states the layout it returns. .t() is a view, so this adds no copy and no memory.

To be clear about scope: this is vLLM's storage convention, not a kernel requirement. Kernels adapt from it on their own — aiter and the CPU prepack path transpose back to (N, K), XPU additionally needs .t().contiguous(). This PR does not change or simplify any of that. It only fixes where the convention is enforced.

input_scale is untouched. It appears in these signatures only because FNUZ normalization must scale it in lockstep with weight_scale; collapsing it for fused shards is an activation-side concern and stays at the call sites exactly as before.

Callers updated

Every call site of the two changed helpers, enumerated against origin/main:

$ git grep -n "= process_fp8_weight_tensor_strategy(\|= process_fp8_weight_channel_strategy(" origin/main
compressed_tensors_w8a8_fp8.py:150   tensor
compressed_tensors_w8a8_fp8.py:158   channel
fp8.py:410                           tensor
modelopt.py:2129                     channel

All four are in this PR, and each change is the removal of the now-redundant transpose:

• fp8.py (Fp8LinearMethod) — drops weight = weight.t(). One line.
• compressed_tensors_w8a8_fp8.py — drops two weight = weight.t(), and rewords the comment above the input_dim/output_dim retag, which previously described the transpose it was compensating for.
• modelopt.py (KFp8StaticChannel) — Parameter(weight.t(), ...) → Parameter(weight, ...).

process_fp8_weight_tensor_strategy_moe is a different function; it and its three call sites are untouched. The two process_fp8_weight_block_strategy call sites are untouched, since that helper's contract did not change.

MLA absorption (mla_attention.py)

New module-level split_kv_b_proj() — dequantize kv_b_proj, assert geometry, reshape, split into (W_UK, W_UV). MLAAttention.process_weights_after_loading and models/kimi_k3/nvidia/mla.py now both call it instead of carrying identical copies.

The extraction moves the dequantize to after the dcp_q_replicate validation block rather than before it. That block reads only scalar config attributes and only asserts or raises — its write set is empty, and get_and_maybe_dequant_weights has an empty write set too, so neither can observe the other. The only difference is which error surfaces when both would fail, and the new order reports the actionable configuration error instead of a downstream dequantize failure.

Not touched

Quark (quark_w8a8_fp8.py) deliberately skips _maybe_pad_fp8_weight on its channel path. Routing it through the shared helper would silently change weight strides on ROCm, so it is left to the Quark maintainers.

Why this is not duplicating an existing PR

gh pr list --repo vllm-project/vllm --state open --search "split_kv_b_proj" returns nothing, and --search "FP8 MLA weight transforms reusable" returns nothing.

Test commands run and results

Weight fingerprint A/B

Both arms run the same vllm/vllm-openai:nightly image on one GB200 node; only the changed .py files are bind-mounted over the image for the PR arm. For every tensor these transforms touch — weight, weight_scale, weight_scale_inv, input_scale, W_UK_T, W_UV — the dump records dtype, shape, stride and sha256 per TP rank.

Config	Model	TP	Tensors	Diffs
CT per-tensor, static input scale	nm-testing/Llama-3.2-1B-Instruct-FP8-KV	1	192	0
CT per-tensor, static input scale	nm-testing/Llama-3.2-1B-Instruct-FP8-KV	2	384	0
CT per-channel + MLA	inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC	1	14	0
CT per-channel + MLA	inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC	2	28	0
Block FP8 — control, must also not change	Qwen/Qwen3-0.6B-FP8	1	224	0

842 tensors, bit-identical, strides included — which is the property this PR is actually about. The comparator refuses to report a result unless the two arms really imported different source files, so an overlay that silently failed to apply cannot pass as a success.

MLA coverage is real rather than incidental: the per-channel dumps contain W_UK_T [128,128,512] and W_UV [128,512,128], and kv_b_proj.weight carries stride=[1,512], the transposed layout the channel path produces.

Fixed-token oracle

Same two arms, greedy decode, identical prompt:

Model	in-process repeat identical	token_ids identical across arms
nm-testing/Llama-3.2-1B-Instruct-FP8-KV	yes, both arms	yes
inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC	yes, both arms	yes, and logprobs too

For the first model the logprobs differ in their last digits across arms while token_ids match. That is cross-process nondeterminism, not a PR effect: a base-vs-base control — byte-identical builds, temperature=0, concurrency=1, same node — diverged on 200/200 GSM8K records with a 0.065 accuracy spread, wider than anything measured here. The fingerprints are the decisive evidence.

Equivalence check for the Kimi-K3 copy

moonshotai/Kimi-K3 is public but 1.5 TB, which does not fit the node used above, so that call site is covered mechanically instead: the pre-PR inline block was copied verbatim into a reference function and run against split_kv_b_proj on the same input. W_UK and W_UV, and the transpose(0,1) / permute(1,2,0) K3 applies downstream, are torch.equal both on an asymmetric shape — which would catch a swapped qk_nope_head_dim / v_head_dim argument, invisible at K3's real 128/128 dims — and on those real dims.

Coverage gap, stated explicitly

The line removed from fp8.py sits in Fp8LinearMethod's per-tensor, non-block branch, and none of the three models above reaches it: the two compressed-tensors models use CompressedTensorsW8A8Fp8, and Qwen3-0.6B-FP8 is block-quantized and returns early. Covering it needs a native per-tensor FP8 checkpoint such as neuralmagic/Meta-Llama-3-8B-Instruct-FP8. I can add that run if reviewers want it before merge.

Lint

pre-commit run --all-files   # passed (ruff-check, ruff-format, mypy-3.10, SPDX, forbidden-imports)

AI assistance

AI assistance was used in developing this PR. I have reviewed every changed line and run the tests and experiments reported above.

🤖 Generated with Claude Code

[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.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #90013 for commit 2f65b2176bb8.

[aoshen02]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #90013.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #90548 for commit a75991589bb6.

[aoshen02]

/ci run

[github-actions]

✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/90548

[aoshen02]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #90548.