Marcusr/aiesw 32176 w4a16 ck wmma

[marcusr-amd]

Purpose

Closes AIESW-32176.

Adds a CK WMMA W4A16 b_scale GEMM kernel and dispatches it from
HybridW4A16LinearKernel for the Qwen3-4B gate_up_proj prefill shape on
Strix Halo (gfx1151). This is the largest GEMM by FLOPs in Qwen3-4B prefill
(M=3968, N=19456, K=2560, group_size=128) — roughly 2x the next largest — so
optimizing it gives the biggest TTFT win for AWQ / W4A16 Qwen3-4B.

Two variants land in this PR:

• Symmetric (torch.ops._rocm_C.ck_w4a16_b_scale_gemm) for
  compressed_tensors uint4b8 checkpoints (e.g.
  RedHatAI/Qwen3-4B-quantized.w4a16).
• Asymmetric (torch.ops._rocm_C.ck_w4a16_b_scale_zp_gemm) for AWQ
  checkpoints with per-group zero points (e.g. Qwen/Qwen3-4B-AWQ — the
  AC#3 model). Reuses the symmetric kernel via the identity
  (nibble - zp) * scale = (nibble - 8) * scale - (zp - 8) * scale, with
  scaled_zp = (zp - 8) * scale precomputed once at weight load. Adds one
  fp16 fma per dequant pack — both variants share all tile sizing,
  scheduler, and threadmap config.

Build is gated on -DVLLM_CK_INCLUDE_DIR + -DVLLM_CK_BUILD_INCLUDE_DIR
(CK's ck/config.h is generated by CK's own configure step). Without those
flags csrc/rocm/ck_w4a16.cu is skipped and the dispatcher falls through to
Triton. VLLM_DISABLE_CK_W4A16=1 forces fall-through at runtime for A/B
benchmarking.

Dispatch decision lives inside the existing hybrid_w4a16_apply custom op
(extended with optional w_q_ck, ck_target_m, w_scaled_zp_ck args),
keeping it opaque to dynamo. CK-format weights are precomputed once in
process_weights_after_loading; the runtime M check is a plain int compare
against a per-layer cached value.

Caveats

• Depends on six matching CK header changes that are NOT in this PR.
  vLLM consumes CK via include paths at build time; the new
  csrc/rocm/ck_w4a16.cu references a DequantPack8WithZp element-op
  and an optional p_b_zero_point arg on
  DeviceGemm_BScale_Wmma_CShuffleV3::MakeArgument that don't exist in
  upstream CK yet. The CK changes are additive and if constexpr-gated
  (zero impact on existing symmetric CK callers, verified bit-identical),
  but they need to land somewhere this build can find them before this PR
  will compile.

  Options being discussed: (a) open a companion PR against
  ROCm/composable_kernel; (b) vendor the modified CK headers under
  csrc/rocm/external/. Either way, this PR is draft/WIP awaiting the
  CK side.

  Build behavior in the meantime:

  • VLLM_CK_INCLUDE_DIR unset (default): csrc/rocm/ck_w4a16.cu is
    skipped, ops are unregistered, dispatcher falls through to Triton.
    Build succeeds; no behavior change vs gfx11 baseline.
  • VLLM_CK_INCLUDE_DIR set, pointing at a CK with the matching
    changes applied: kernel compiles, ops register, dispatch fires for
    the target shape.
  • VLLM_CK_INCLUDE_DIR set, pointing at upstream CK without the
    changes: compile error in csrc/rocm/ck_w4a16.cu.

• Single-shape (column) specialization: only the Qwen3-4B gate_up_proj
  column (N=19456, K=2560) on gfx1151 routes to CK; the M dimension is
  validated at M=2048 (default chunked-prefill chunk size) and M=3968 (full
  prompt with --max-num-batched-tokens 4096), both ~30 TFLOPS standalone.
  Other columns (qkv_proj, o_proj, down_proj, other Qwen variants) stay on
  Triton — left for follow-up tuning. Per-layer dispatch carries a small
  list of validated M values (ck_target_ms) so adding shapes is a
  one-line change to _CK_W4A16_TARGET_SHAPES.

Test Plan

1. Standalone CK kernel benchmark at the target shape, comparing against
   Triton W4A16 (via benchmarks/kernels/benchmark_hybrid_w4a16_gemm.py,
   edited to add batch_size=3968 + bf16 providers) and the hipBLASLt fp16
   TN roofline (hipblaslt-bench).
2. Numerical correctness smoke tests: random fp16 weights → ExLlama
   pack → repack to CK layout → call new op → compare to torch fp16
   reference using the dequantized weights. Both symmetric and asymmetric
   variants.
3. E2E A/B with VLLM_DISABLE_CK_W4A16 env-var toggle on the same
   vllm bench serve invocation:
   • Qwen/Qwen3-4B-AWQ (asymmetric, AC#3 model)
   • RedHatAI/Qwen3-4B-quantized.w4a16 (symmetric reference)
4. lm_eval MMLU on Qwen3-4B-AWQ comparing CK kernel vs Triton baseline
   per AC#1 — TBD.
5. Numerical equivalence at full vLLM model output (AC#3 verbatim:
   "produces the same output as with the triton kernel") — currently
   confirmed via the per-layer fp16 tolerance smoke tests; full-model
   output diff TBD.

Test Result

Measured on Strix Halo (gfx1151, Radeon 8060S) against ROCm 7.13.

Standalone GEMM (N=19456, K=2560, group=128, fp16, gfx1151)

At the target shape M=3968:

Path	TFLOPS	vs roofline	vs Triton
hipBLASLt fp16 TN (roofline)	38.6	100%	—
Triton W4A16 (baseline)	15.6	41%	1.00×
CK b_scale (symmetric)	30.0	78%	1.92×
CK b_scale_zp (asymmetric)	~26.4	~68%	~1.69×

The same kernel config holds across the M dimension on this column —
both M=2048 (chunked-prefill default chunk size) and M=3968 dispatch
to CK and hit ~30 TFLOPS:

M	CK TFLOPS	% of roofline
2048 (default chunk)	30.86	80%
3968 (full prompt @ chunk=4096)	30.00	78%
4096	30.10	78%

Asymmetric kernel is ~3.7% slower than symmetric per call — exactly the
design promise of one extra fp16 fma per dequant pack.

E2E (vllm bench serve, max_num_seqs=1, fp16, num_prompts=10)

Qwen/Qwen3-4B-AWQ (asymmetric, AC#3 model). The CK dispatch is
keyed per-layer by a min-M threshold (ck_min_m=256 for the gate_up
column, where standalone sweep shows the same kernel config holds
28-31 TFLOPS uniformly across M=256-16384). The default
chunked-prefill chunk (max_num_batched_tokens=2048) now hits CK
on both prefill chunks (M=2048 and M=1920) — the prior
--max-num-batched-tokens 4096 workaround is not required:

Config	TTFT mean	TPOT
chunk=2048 (default), CK ON (min_m=256: M=2048+M=1920 both CK)	1922 ms	16.3
chunk=2048 (default), CK OFF (Triton)	2139 ms	17.1
chunk=4096 (override), CK ON (M=3968 single dispatch)	2166 ms	16.4
chunk=4096 (override), CK OFF (Triton)	2682 ms	17.5

The default chunk + CK ON config (1922 ms) wins both A/Bs:

• vs VLLM_DISABLE_CK_W4A16=1 at the same chunk: −217 ms (−10%)
• vs the prior chunk=4096 + CK ON config: −244 ms (−11%) —
  splitting the prefill into two chunks (both CK) is faster than one
  big CK call. Per-dispatch profile shows Triton W4A16's qkv/o_proj/
  down_proj scale super-linearly with M, so chunked prefill saves
  ~270 ms of downstream Triton GEMM time per prompt by feeding those
  ops smaller M dimensions.

RedHatAI/Qwen3-4B-quantized.w4a16 (symmetric reference, prior
chunk=4096 measurement):

Config	TTFT mean	TPOT
chunk=4096, CK ON	2226 ms	16.0
chunk=4096, CK OFF	2630 ms	15.9
Delta	−404 ms (−15%)	~0

Decode (TPOT) is unchanged in all A/Bs — CK only fires on prefill
chunks at the registered M values, leaving the decode path on the
existing skinny / Triton kernels.

Numerical correctness smoke tests

Variant	Max abs / max ref
Symmetric (vs torch fp16 ref using dequantized weights)	5.4e−04
Asymmetric (vs torch fp16 ref using dequantized weights)	5.4e−04

Both well within the fp16 GEMM tolerance class used by
tests/kernels/quantization/.

MMLU regression check (AC#1)

TBD — will post lm_eval MMLU before/after on Qwen/Qwen3-4B-AWQ.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.