[Model][Quant] Fused WNA16 GEMM for tied quantized lm_head logits

[KKothuri]

Purpose

Follow-up to #45535 (compressed-tensors WNA16 input embeddings + tied lm_head). This PR speeds up the logits matmul when a quantized embedding is reused as a tied lm_head.

In #45535, CompressedTensorsEmbeddingWNA16Int.apply() dequantizes the entire packed [vocab, hidden] table into a dense weight on every call and runs F.linear. This PR routes that matmul through vLLM's existing WNA16 Linear kernel (Marlin/Machete) instead — a fused dequant-GEMM that never materializes the dense weight.

Key insight: the embedding's packed weight is already in the compressed-tensors layout those kernels expect (packed along the input/hidden dim, identical to a WNA16 Linear), so process_weights_after_loading sets up the fused kernel directly. It's set up only for tied embeddings (flagged in tie_weights), keeping a gather-format copy for the input-lookup path. When no fused kernel is available it falls back to a full-table dequant (no index tensor) + F.linear.

Note

Stacked on #45535 — the tied-embedding apply() this optimizes only exists there, so the net-new change here is compressed_tensors_embedding.py; the rest of the diff belongs to #45535. Please review/merge after #45535. Kept as a separate PR to avoid growing the (already large) plumbing PR.

Not a duplicate: quantized-lm_head-via-Marlin already exists for the untied / separately quantized case — a ParallelLMHead independently quantized as a GPTQ/AWQ/compressed-tensors Linear already dispatches to Marlin (cf. #40999). This PR covers the distinct tied case, where the embedding is the lm_head and its quant method is the embedding method (CompressedTensorsEmbeddingWNA16Int), not a Linear method. The open ModelOpt lm_head/embedding PRs (#35660, #42791, #44671, #41000) target a different backend (NVFP4/FP8 in modelopt.py) and address loading/dispatch, not this fused tied-logits path. No open PR touches compressed_tensors_embedding.py.

Correctness

The fused path matches the dequant reference numerically (max relative diff 0.025%, fp16). The existing tied fixture test (tests/quantization/test_quantized_embedding.py::test_tied_quantized_embedding) exercises the fused Marlin path end-to-end and generates coherently.

Perf evals

RTX 5080 (sm120, CUDA 13), apply() microbenchmark, µs/call (lower is better):

dims	M	current (dequant + F.linear)	no-gather dequant	fused (Marlin)	fused vs current
vocab 151936 × 1024, W4-g64	1	1329	1317	240	5.5×
	32	1703	2395	619	2.8×
	256	3247	2460	2266	1.4×
vocab 151936 × 1024, W8-g128	1	2859	1374	329	8.7×
	32	1442	575	575	2.5×
	256	2971	2088	2088	1.4×
vocab 50304 × 512, W4-g64	1	141	142	19	7.6×
	32	192	190	123	1.6×

Biggest win at decode (M=1), the common serving case for logits. The no-gather dequant (kept as the fallback path) is marginal versus the current dequant, as expected — the fused GEMM is the real lever.

Test Plan

pytest tests/quantization/test_quantized_embedding.py

Test Result

• Both fixture tests pass on an RTX 5080; the tied test confirms dispatch to MarlinLinearKernel for the logits path (verified logits_kernel is set and gather copies are preserved).
• pre-commit (ruff + mypy) clean on the changed file.

---

This change was developed with AI assistance (Claude Code). All changed lines were reviewed by the submitter.

[mergify]

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

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

[sgurwinderr]

Drive-by data from an unrelated Turing (sm_75) experiment on the same tied-lm_head path — two things that may be useful here.

1. M=1 is purely bandwidth-bound. On an RTX 2060 Mobile (250 GB/s achievable, 336 GB/s spec) the fp16 logits GEMV for Qwen2.5-Coder-3B (151936 × 2048, tied) runs 622 MB / 2.51 ms = 248 GB/s, 99% of roofline. Quantizing the weight to int8 gave 2.88 → 1.42 ms = 2.03x, i.e. 98% of the half-bandwidth floor. So at M=1 the only lever is bytes moved, and 2x is the ceiling for a W8 path.

That makes the 8.7x in your W8-g128 M=1 row look like it's mostly removing a redundant pass (dequantize-to-dense then read) rather than kernel speed. Worth saying explicitly — it sets the expectation that the fused path is now at the memory floor.

2. Max-relative-diff understates argmax movement, and there's a nearly-free fix. Greedy decoding only cares about argmax, and one flip changes every later token. On 2632 captured decode states, int8 per-row symmetric gave 99.658% top-1 agreement (~1 flip per 292 tokens). Caveat: at 646 states I measured 646/646 and nearly called it lossless — the rate only shows up at a few thousand states.

Cheap fix, applicable to any weight-only lm_head scheme: take top-K from the quantized logits, recompute those K rows exactly, scatter back. 32 KB at K=8 next to a 311 MB weight — no measurable throughput change.

• int8 alone: 99.658%
• int8 + exact top-8 rescoring: 99.848%

Residual disagreements are all cases where the fp16 reference's own top1–top2 margin is ~0. K > 4 buys nothing.

3. Methodology note: a fp16-vs-fp16 control (same build, seed, max_num_seqs=1, temp 0) diverged from itself on 2/6 prompts, 2 flips / 2625 tokens. Bit-exact token comparison isn't a usable gate for long generations.

Context: this came from adding runtime int8 quantization of an unquantized tied lm_head on 6 GB Turing (Marlin needs ≥ 80), +10.2% end-to-end decode. Not proposing it upstream — happy to share the microbenchmark or argmax-agreement harness if useful.

[mergify]

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

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