[ROCm][Kernel] Add HybridW4A16LinearKernel: Triton prefill + HIP skinny decode

[mgehre-amd]

Add a hybrid W4A16 linear kernel for ROCm gfx11/gfx12 that routes between two GEMM implementations based on batch size:

• M <= 5: HIP wvSplitK_int4_g skinny GEMM (optimized for single-token decode)
• M > 5: Triton fused dequant GEMM (optimized for prefill/batched inference)

Both paths share a single weight tensor in the AWQ-style packed shuffle layout [N, K//8]. Supports both symmetric (uint4b8, bias=8) and asymmetric (uint4, per-group zero-points) quantization with group sizes 32, 64 or 128.

The Triton kernel includes tuned tile-size heuristics for gfx1151 (Strix Halo, RDNA 3.5) and gfx1201 (RDNA 4) based on real weight shapes from Qwen3-4B and Llama-3.1-8B-AWQ, plus shape-specific overrides.

The HIP skinny decode kernel uses a marlin-style magic-number dequantization trick for bf16 models, avoiding an expensive fp32 round-trip on gfx1151.

Benchmarked on AMD Strix Halo (gfx1151, LPDDR5X-8000 128 GB) with --max-num-seqs 1:

Model: RedHatAI/Qwen3-4B-quantized.w4a16 (compressed-tensors, symmetric)

Config	Metric	Baseline	Kernel	Speedup
in=128, out=128	TTFT	82.33 ms	68 ms	1.21x
in=128, out=128	TPOT	54.81 ms	13.69 ms	4.00x
in=1920, out=128	TTFT	1222.24 ms	760 ms	1.61x
in=1920, out=128	TPOT	56.16 ms	15.29 ms	3.67x

Model: Qwen/Qwen2.5-7B-Instruct-AWQ (AWQ, asymmetric)

Config	Metric	Baseline	Kernel	Speedup
in=128, out=128	TTFT	118.64 ms	100.09 ms	1.19x
in=128, out=128	TPOT	82.04 ms	21.03 ms	3.90x
in=1920, out=128	TTFT	2004.68 ms	1183.48 ms	1.69x
in=1920, out=128	TPOT	84.69 ms	21.67 ms	3.91x

Model: RedHatAI/gemma-3-4b-it-quantized.w4a16 (compressed-tensors, symmetric)

Config	Metric	Baseline	Kernel	Speedup
in=128, out=128	TTFT	85.18 ms	66 ms	1.29x
in=128, out=128	TPOT	53.39 ms	15.59 ms	3.43x
in=1920, out=128	TTFT	1392.00 ms	727 ms	1.91x
in=1920, out=128	TPOT	61.36 ms	24.24 ms	2.53x

Model: RedHatAI/Qwen3-8B-quantized.w4a16 (compressed-tensors, asymmetric)

Config	Metric	Baseline	Kernel	Speedup
in=128, out=128	TTFT	150.51 ms	140 ms	1.08x
in=128, out=128	TPOT	116.99 ms	24.55 ms	4.77x
in=1920, out=128	TTFT	3181.00 ms	1529 ms	2.08x
in=1920, out=128	TPOT	119.72 ms	26.12 ms	4.58x

Model: trymirai/SmolLM2-1.7B-Instruct-AWQ (AWQ, asymmetric)

Config	Metric	Baseline	Kernel	Speedup
in=128, out=128	TTFT	39.02 ms	46 ms	0.85x
in=128, out=128	TPOT	22.95 ms	7.25 ms	3.17x
in=1920, out=128	TTFT	769.37 ms	413 ms	1.86x
in=1920, out=128	TPOT	26.07 ms	10.53 ms	2.48x

Accuracy (lm_eval gsm8k, 200 samples):
RedHatAI/Qwen3-4B-quantized.w4a16: 0.835 -> 0.855 (no regression)

[gemini-code-assist]

Code Review

This pull request introduces a Hybrid W4A16 kernel for ROCm, optimizing performance by dispatching to a HIP skinny GEMM kernel for decode and a Triton kernel for prefill using a shared ExLlama shuffle weight layout. A critical issue was identified in the ROCm C++ implementation where a hardcoded LDS size macro conflicts with dynamic host-side checks, potentially causing memory corruption on architectures with larger LDS (like MI300X/MI350) when the input size exceeds the static allocation.

[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 @mgehre-amd, 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.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[BadrBasowid]

@mgehre-amd some tests pass on the second try, if only a few tests are failing, tag @tjtanaa so he can help retry those specific tests.

[mgehre-amd]

@AndreasKaratzas @tjtanaa, only three test failing. Do you want to rerun those or force merge?

[tjtanaa]

@mgehre-amd I have rerun them. On of the failure is critical where CUDA CI didn't even run.

[mgehre-amd]

Failed this time:

tests/kernels/mamba/test_cpu_short_conv.py::test_short_conv_forward_native_prefill
FAILED lora/test_qwenvl.py::test_qwen2vl_multiple_lora_types - AssertionError: Generated text A view of the Tokyo doesn't match expected pattern A closeup shot of the Tokyo Skytree with pink flowers in the foreground.
FAILED lora/test_transformers_model.py::test_ilama_lora - RuntimeError: Engine core initialization failed. See root cause above. Failed core proc(s): {}

test_cpu_short_conv is a known problem on main and was fixed by #47848
The lora test failures are known problems on main and there is a PR to address them: #47798
None of the failed tests are due to this PR

[mgehre-amd]

Failed on this rerun:

model_executor/model_loader/test_reload.py::test_reload_weights[inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC-inference-optimization/DeepSeek-V3-debug-multiply-FP8_DYNAMIC-inference-optimization/DeepSeek-V3-debug-add-FP8_DYNAMIC-1]: AcceleratorError: CUDA error: operation not permitted when stream is capturing on MI300X. This PR only triggers for w4a16 models.
docker-build-arm64-image: denied: The repository with name 'vllm-ci-test-repo' in registry with id 'q9t5s3a7' already has the maximum allowed number of images which is '100000'
docker-build-cpu-image: denied: The repository with name 'vllm-ci-test-repo' in registry with id 'q9t5s3a7' already has the maximum allowed number of images which is '100000'
docker-build-image

[mgehre-amd]

After rebase, failed unrelated tests:

FAILED v1/kv_connector/extract_hidden_states_integration/test_extraction.py::test_extract_hidden_states_tp2 - Failed: Timeout >60.0s
FAILED model_executor/model_loader/test_reload.py::test_reload_weights[inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC-inference-optimization/DeepSeek-V3-debug-multiply-FP8_DYNAMIC-inference-optimization/DeepSeek-V3-debug-add-FP8_DYNAMIC-1]

[mgehre-amd]

List of tests failed after this rebase:

FAILED v1/cudagraph/test_cudagraph_mode.py::test_cudagraph_compilation_combo[FA2-FULL_DECODE_ONLY-0-True] - ValueError: Memory of devices devices=[0] not free after dur_s=120.01 (threshold='device=0: 0.100')
FAILED renderers/test_hf.py::test_resolve_content_format_hf_defined[Qwen/Qwen3.5-4B-openai] - ValueError: No example model defined for Qwen/Qwen3.5-4B; please update this file.

First one is already known issue on main with a revert tracked in #47232
Second one is also unrelated to this PR.