vllm - 💡(How to fix) Fix [Feature]: triton_scaled_mm uses an AMD-tuned fixed tile heuristic — leaves up to 1.82x on NVIDIA H800 (Hopper sm_90)

🚀 The feature, motivation and pitch

Motivation

triton_scaled_mm (vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py) selects tile sizes with a fixed hand-written heuristic introduced in #11698 ("avg 2.8x speedup for int8 models"), which was tuned on AMD. It keys only on M (four buckets), ignores N, and never tunes num_warps / num_stages.

On an NVIDIA H800 (Hopper, sm_90), benchmarking the exact current kernel against a variant that is a faithful superset of it — identical int32 accumulation, the same two-step scale epilogue, and the same optional bias add, with the only differences being @triton.autotune over tile/num_warps/num_stages (instead of the fixed heuristic) plus L2-friendly grouped program ordering (PID swizzling, as in the abandoned #20163) — shows the fixed heuristic leaves up to 1.82x on the table in the compute-bound regime (M ≥ 256) for representative W8A8 INT8 GEMM shapes. Outputs are bit-identical (max rel err = 0.0, verified for both the no-bias and with-bias paths), so the gap is purely tile/scheduling, not precision or a simpler kernel. The effect is stable across two triton versions (3.6.0 and 3.5.0) and reproduces on L20 (Ada, sm_89) too (peak 1.65x there) — so this is the heuristic, not a version or a single device.

Results (H800)

Speedup = current_ms / autotuned_ms (faithful-superset kernel, bit-identical output). Numbers below are from the triton 3.6.0 run.

Absolute TFLOPS for the headline cells (current → autotuned):

• down M=256: 358 → 653 TFLOPS (1.82x)
• o_proj M=256: 278 → 439 (1.58x)
• gate_up M=1024: 732 → 1121 (1.53x), M=4096: 769 → 1081 (1.41x)

Observations:

• Large-N layers plateau under the fixed tile. On gate_up (N=28672) at large M the current (128,128,128)-for-all-M>128 choice plateaus around 732–769 TFLOPS while a tuned config reaches 1081–1121 — a persistent 1.41–1.53x. The heuristic ignores N entirely.
• The compute-bound regime (M ≥ 256, prefill / chunked-prefill) concentrates the largest gaps (1.31–1.82x).
• Unlike on L20, even small-M (M ≤ 64) shows 1.1–1.25x on several shapes (qkv/o_proj/down) on H800 — the AMD-tuned (64,64,256) / lack of warp/stage tuning is suboptimal for Hopper there too.
• Large M=4096 mostly converges (≈1.0–1.06x): both are near roofline, so the fixed tile is fine in that corner.

Corroboration on L20 (Ada): same kernel/sweep — up to 1.65x (gate_up M=256), 1.2–1.5x typical in the compute-bound regime, ≈1.0x for small-M decode.

Honest caveats

• These are microbenchmark TFLOPS, not end-to-end serving throughput. The prior attempt #20163 saw microbench wins not translate to e2e (Python/dispatch overhead); e2e impact here is not yet measured.
• On NVIDIA the INT8 W8A8 path defaults to CUTLASS; this Triton kernel is a fallback (ROCm, non-cutlass-compatible shapes, non-cutlass fp8 in _custom_ops.cutlass_scaled_mm). Practical NVIDIA benefit is bounded by how often the fallback is hit — hence filing an issue to gauge value rather than a PR.
• The autotune used here is online (not torch.compile-safe), purely to measure headroom. A real fix must be compile / cudagraph-safe (see #26993).

Proposed direction (for discussion, not a commitment)

Replace the fixed heuristic with a compile/cudagraph-safe autotuned tile selection (e.g. an ahead-of-time tuned config table keyed on shape, resolved at init — not a per-call decorator), keeping the AMD path intact. Happy to send a PR with H800 numbers if maintainers agree this fallback is worth improving and on the mechanism to stay compile-safe.

Standalone reproduction script (verbatim copy of the current kernel vs the faithful-superset autotuned variant): https://gist.github.com/DuGuLuOwO/b4a8c09938b211e6ea8bdef5610025a8

Alternatives

Prior art / related (to avoid duplication)

• #11698 — origin of the current AMD-tuned heuristic.
• #20163 — earlier Draft adding NVIDIA autotuning + PID swizzling to this kernel; closed unmerged when the author's fork was deleted (never reviewed). This issue re-surfaces the problem with concrete H800 + L20 data and the compile-safety constraint.
• #34275 (open) — adds an AMD gfx1100 tile heuristic to the same function (would textually conflict; the _get_tile_config() seam there is a natural plug point).
• #26993 — dynamo/torch.compile incompatibility with naive runtime autotune (motivates the compile-safe constraint below).
• #38697 (open) — a real-world report that llmcompressor W8A8 decode is slower than FP16; corroborates that this W8A8 path's efficiency matters in practice.
• Helion scaled_mm (#33651 / #44020 / #44459) — the longer-term direction; an autotuned Triton heuristic is a near-term, compile-safe win independent of that migration.

Additional context

Questions for maintainers

1. Is improving this Triton fallback still worthwhile given the Helion direction?
2. Preferred mechanism to stay compile/cudagraph-safe — AOT config table vs decorator?

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.