[ROCm][W4A16] Cache dequantized bf16 weights for the linear prefill GEMM

[roberteg16]

Summary

This adds an opt-in path (VLLM_W4A16_PREFILL_BF16=1, default off) that trades VRAM for compute on the ROCm W4A16 linear prefill path (gfx11/gfx1151). At load time each W4A16 weight is dequantized once to a dense bf16/fp16 copy; prefill (batch M > MAX_SKINNY_BATCH_SIZE) then runs a dense hipBLASLt GEMM (F.linear) on that copy instead of the fused in-kernel int4 unpack, while decode (M <= 5) is untouched and still uses the int4 wvSplitK skinny path. The motivation is that the fused prefill kernel is compute-bound (~26 TFLOPS, ~61% of peak) because of the in-kernel ExLlama unshuffle + per-group dequant, whereas a dense bf16 GEMM reaches ~33-39 TFLOPS at the same shapes. The feature is fully A/B-able: with the env unset nothing changes, and even when set, layers that do not fit the memory budget transparently keep the int4 prefill path.

What this adds

• vllm/envs.py: registers VLLM_W4A16_PREFILL_BF16 (bool, default off). The bf16 copy is only kept while there is room left in the gpu_memory_utilization budget; there is no separate reserve knob.
• vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py: in process_weights_after_loading, when the env is set, dequantize each W4A16 weight to a dense copy and store it as a non-persistent buffer (_hybrid_w_bf16). The copy is gated per-weight on the memory budget (see Memory behavior). The decode-vs-prefill-vs-bf16 dispatch is performed inside the existing torch.ops.vllm.hybrid_w4a16_apply custom op (which now also takes the optional w_bf16 tensor), so the M-based selection is evaluated at runtime rather than baked at trace time.

Design notes

The memory budget gate is anchored to the configured gpu_memory_utilization rather than a raw free-VRAM number. With budget = total * gpu_memory_utilization, a weight is cached only if free_mem - total * (1 - gpu_memory_utilization) >= bf16_bytes at load time, otherwise it keeps the int4 prefill path. Free memory is read via MemorySnapshot, which falls back to psutil on integrated/UMA GPUs (e.g. gfx1151 Strix Halo, where cudaMemGetInfo underreports free memory), so the gate stays consistent with how vLLM sizes the KV cache. The check uses live free memory and runs before the memory profiler, so it is self-limiting (each copy shrinks the spendable budget until it is exhausted) and the KV-cache sizing automatically accounts for the copies. A warn-once message is emitted if the budget is exhausted.

The decode/prefill selection lives inside the opaque custom op on purpose, and this is load-bearing. An earlier version branched in Python in apply_weights (if w_bf16 is not None and M > MAX_SKINNY_BATCH_SIZE: F.linear(...)). Under VLLM_COMPILE that branch is traced with a prefill example (M large), baked as taken, and then reused at decode (M=1) with dynamic_shapes_config.evaluate_guards=False, which forced the dense bf16 GEMM onto the memory-bound decode path and caused a large TPOT/E2E regression. Moving the M-check inside the custom op (which torch.compile treats as opaque) makes the selection run for real on the actual batch size every call, so decode reliably stays on int4.

Performance

End-to-end on gfx1151 (Strix Halo), Qwen3.6-35B-A3B-W4A16, single prompt, 100 input tokens + one 1280x720 image, --output-len 128, --no-cudagraph, VLLM_MOE_HIP=1, ROCM_AITER_FA. Single run each (treat small deltas as indicative).

metric	baseline (off)	bf16 cache on
TTFT (ms)	598.1	583.5
prefill (tok/s)	1662	1703
TPOT (ms)	30.86	31.45
E2E latency (ms)	4517	4577
KV cache	7.23 GiB / 210,066 tok	4.85 GiB / 140,434 tok

Prefill is ~2.4% faster (TTFT) / ~2.5% higher throughput; decode (TPOT) and E2E are unchanged within run-to-run noise, confirming the int4 decode path is preserved. Profile inspection (tools/roofline_trace.py --report-by-op --kind ALL --gpu gfx1151) confirms the swap: with the env on, the W4A16 linear prefill no longer shows _triton_w4a16_skinny_fmt_kernel (replaced by dense bf16 aten::mm/addmm at the prefill shapes), while at decode the W4A16 layers (e.g. the GDN in_proj, M=1, N=12288) run _rocm_C::wvSplitK_int4_g exactly as in the baseline (227.6 ms total, ~59.7 us/call in both runs).

Memory behavior

The bf16 copy is roughly 4x the int4 weight, so it is opt-in and budget-gated. On the run above it consumed ~2.38 GiB (KV cache 7.23 -> 4.85 GiB) at gpu_memory_utilization=0.9. This is the expected cost; the feature is worthwhile only when prefill latency matters more than KV-cache capacity for the deployment.

Testing

• Lint/type: pre-commit run ruff-check, ruff-format, and mypy-local pass on the changed files.
• End-to-end smoke + A/B on gfx1151 against Qwen3.6-35B-A3B-W4A16 using the command above, run three ways: env off (baseline), env on, and env on after the compile-safe dispatch fix. All three complete with Failed requests: 0; the numbers and profile evidence are in the Performance section.
• Profile verification via tools/roofline_trace.py on the exported traces confirmed (a) decode stays on wvSplitK_int4_g, (b) prefill uses the dense bf16 GEMM, and (c) the budget gate emits no warning at gpu_memory_utilization=0.9 (all copies fit).

[mgehre-amd]

Looks good!

[mgehre-amd]

Please merge origin/gfx11 to fix build failure

[roberteg16]

Found a bug, converting to draft