From 0a42169ac3aa9fb375f593c6a3e7ac026794dbf5 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 3 Sep 2026 08:30:15 -0600 Subject: [PATCH] [ROCm][Perf] W4A16: packed fp16 prefill dequant and gfx1151 tile selection The Triton prefill path of RDNAHybridW4A16LinearKernel is VALU-issue-bound on RDNA3, and its tile ladder was tuned before the dequant changed shape. Packed fp16 dequant. OR-ing a 4-bit code n into the low mantissa of fp16 1024.0 (0x6400) bitcasts to exactly 1024+n. Applying that to a full 32-bit lane with one v_and_or_b32 dequants a nibble PAIR per instruction, against the scalar v_and_b16 + v_or_b16 pair Triton emits from the elementwise form. The ExLlama shuffle already stores val[2p] at bits [4p:4p+4] and val[2p+1] at [16+4p:20+4p], so a single pre-shift by 4p yields both halves of a half2 in K order and the downstream affine packs into v_pk_fma_f16. Dtype-aware gfx1151 tile selection. The packed fp16 path and the scalar bf16 path want different tiles -- most visibly BLOCK_N at deep M, 256 against 64 -- so a single M-ladder cannot serve both. Measured on gfx1151, 3 interleaved reps per arm, 2048-token prefill: Qwen3-8B-quantized.w4a16 fp16 asymmetric +7.8% prefill throughput / -7.3% TTFT; Qwen3-4B-quantized.w4a16 fp16 symmetric +13.0% / -11.5%. Both fp16 rep ranges are disjoint. The bf16 control lands at +0.02% with overlapping rep ranges, i.e. the null it should be. Changes: - The kernel resolves the unpack itself, from its own compile target and the activation dtype -- there is no flag to pass and no way for the host's view of the GPU to disagree with what is being compiled. That needs a @triton.constexpr_function helper, which in turn needs constexpr_function added to vLLM's Triton placeholder shim: this module is imported on every platform, and the shim exposed jit/autotune/heuristics/Config but not constexpr_function, so a module-scope use raised AttributeError on builds without Triton. A dummy decorator is sufficient there, since such bodies only run while a kernel is being compiled. - The packed dequant is enabled for the whole gfx11 family on fp16: fp16 is a hard requirement of the magic constant, and given that, the packed form is a pure instruction-count reduction producing bit-identical values, so there is nothing to tune per part. Only the tile table, which IS tuned, stays gfx1151-gated. Verified in the generated ISA: the fp16 kernel contains 65 v_and_or_b32, the bf16 one none from the dequant. - num_stages=1 is applied to the fp16 arm ONLY. Applying it to bf16 as well measured -15.8% end-to-end prefill on an asymmetric bf16 model: the packed fp16 path issues one per-group load and does not miss the pipelining, while the scalar asymmetric path issues two (scale and zero point) and needs the software pipeline to hide the second gather. - The bf16 arm is therefore byte-for-byte the pre-existing scalar-tuned ladder, its per-shape override table, and its pipeline depth. Verified identical to the previous selection across all 576 shape x group-size combinations swept. - Scope is gfx1151 only. Widening the tuned table to the whole gfx11 family was considered and rejected: it was swept on gfx1151, and gfx1100 is a 96-CU discrete part that cannot be measured here. - ROCm#923's K>=4096 and N>=4096 bf16 tile branch was evaluated and deliberately NOT carried over. At group_size 128 it is shadowed by the later, more specific per-shape override table for every shape that table covers; at group_size 32/64 its distinguishing BLOCK_K=128 collapses to the group size anyway. It would only ever fire on shapes nobody measured, in a region where the later sweep disagreed with it. - The kernel benchmark gains --dtype: it was fp16-hardcoded and so could not reach the scalar bf16 path at all. Numerics. The dequant arithmetic is unchanged: 1024+n is integer-exact in fp16, whose 11-bit mantissa covers every integer below 2048, so folding the 1024 into the subtrahend -- (b_raw - (1024 + zp)) == (nibble - zp) -- is exact and the multiply that follows rounds once, as before. End-to-end output is nonetheless not GUARANTEED identical on fp16, because the new tiles move BLOCK_K (64 -> 32 on most shapes) and so change the fp32 accumulation order. Measured: greedy decode over 6 fixed prompts is byte-identical on all three configurations, and GSM8K 5-shot over 500 questions is unchanged on the asymmetric model (0.884 -> 0.884) and moves by one question on the symmetric one (0.848 -> 0.850), against a ~1.5pp single-arm sampling error. Testing: pytest tests/kernels/quantization/test_rdna_hybrid_w4a16.py tests/kernels/quantization/test_w4a16_kernel_selection.py -- 103 passed on gfx1151 (Radeon 8060S, torch 2.11.0+rocm7.15, Triton 3.8.0). --- .../benchmark_rdna_hybrid_w4a16_gemm.py | 43 +-- .../quantization/test_rdna_hybrid_w4a16.py | 74 ++++++ .../mixed_precision/rdna_hybrid_w4a16.py | 248 +++++++++++++++--- vllm/triton_utils/importing.py | 4 + 4 files changed, 310 insertions(+), 59 deletions(-) diff --git a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py index 4ecd917f31ac..f393ca152bb3 100644 --- a/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py +++ b/benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py @@ -52,11 +52,11 @@ # --------------------------------------------------------------------------- # Weight packing # --------------------------------------------------------------------------- -def prepare_hybrid_weights(K, N, group_size, device="cuda"): +def prepare_hybrid_weights(K, N, group_size, dtype=torch.float16, device="cuda"): """Create random weights for benchmarking. - Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives - its int32 view from w_q_skinny, so no separate int32 buffer is returned. + Returns (w_q_skinny, w_s_skinny, w_dense, w_zp). The triton path derives its + int32 view from w_q_skinny, so no separate int32 buffer is returned. """ num_groups = K // group_size @@ -65,23 +65,23 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"): 0, 2**31, (N, K // 8), dtype=torch.int32, device=device ) w_q_skinny = w_q_skinny_i32.view(torch.int8).contiguous() - w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01 + w_s_skinny = torch.randn(N, num_groups, dtype=dtype, device=device) * 0.01 # Raw per-group zero-points for asymmetric benchmarks w_zp = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device).to( - torch.float16 + dtype ) - # FP16 baseline for F.linear - w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01 + # Unquantized baseline for F.linear + w_dense = torch.randn(N, K, dtype=dtype, device=device) * 0.01 - return w_q_skinny, w_s_skinny, w_fp16, w_zp + return w_q_skinny, w_s_skinny, w_dense, w_zp # --------------------------------------------------------------------------- # Benchmark # --------------------------------------------------------------------------- -PROVIDERS = ["torch-fp16", "hybrid-w4a16", "hybrid-w4a16-zp"] +PROVIDERS = ["torch-dense", "hybrid-w4a16", "hybrid-w4a16-zp"] @triton.testing.perf_report( @@ -93,22 +93,21 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"): line_vals=PROVIDERS, line_names=PROVIDERS, ylabel="TFLOP/s (larger is better)", - plot_name="FP16 vs Hybrid W4A16", + plot_name="Dense vs Hybrid W4A16", args={}, ) ) -def benchmark(batch_size, provider, N, K, group_size, weights): +def benchmark(batch_size, provider, N, K, group_size, dtype, weights): M = batch_size device = "cuda" - dtype = torch.float16 a = torch.randn((M, K), device=device, dtype=dtype) quantiles = [0.5, 0.2, 0.8] - if provider == "torch-fp16": - w_fp16 = weights["w_fp16"] + if provider == "torch-dense": + w_dense = weights["w_dense"] ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( - lambda: torch.nn.functional.linear(a, w_fp16), + lambda: torch.nn.functional.linear(a, w_dense), quantiles=quantiles, ) elif provider in ("hybrid-w4a16", "hybrid-w4a16-zp"): @@ -168,21 +167,28 @@ def prepare_shapes(args): ) parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1]) parser.add_argument("--group-size", type=int, default=128) + parser.add_argument( + "--dtype", type=str, default="float16", choices=["float16", "bfloat16"] + ) parser.add_argument("--save-path", type=str, default=None) args = parser.parse_args() + dtype = getattr(torch, args.dtype) + for K, N, model in prepare_shapes(args): group_size = args.group_size print(f"\n{'=' * 70}") - print(f"{model}, N={N} K={K}, group_size={group_size}") + print(f"{model}, N={N} K={K}, group_size={group_size}, dtype={args.dtype}") print(f"{'=' * 70}") - w_q_skinny, w_s_skinny, w_fp16, w_zp = prepare_hybrid_weights(K, N, group_size) + w_q_skinny, w_s_skinny, w_dense, w_zp = prepare_hybrid_weights( + K, N, group_size, dtype + ) weights = { "w_q_skinny": w_q_skinny, "w_s_skinny": w_s_skinny, - "w_fp16": w_fp16, + "w_dense": w_dense, "w_zp": w_zp, } @@ -195,6 +201,7 @@ def prepare_shapes(args): N=N, K=K, group_size=group_size, + dtype=dtype, weights=weights, ) diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index 0d007aa1170f..cc5365a4b25d 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -30,6 +30,8 @@ pack_int4_exllama_shuffle = hybrid_module.pack_int4_exllama_shuffle SUPPORTED_GROUP_SIZES = hybrid_module.SUPPORTED_GROUP_SIZES MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE +triton_w4a16_skinny_fmt_gemm = hybrid_module.triton_w4a16_skinny_fmt_gemm +select_skinny_gfx1151_config = hybrid_module._select_skinny_gfx1151_config # --------------------------------------------------------------------------- @@ -173,6 +175,78 @@ def test_rdna_hybrid_w4a16_apply_with_bias(dtype, M): torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2) +# --------------------------------------------------------------------------- +# Triton prefill path +# --------------------------------------------------------------------------- + + +def _make_prefill_case(M, K, N, G, dtype, has_zp): + """Random [M,K] activations + skinny [N,K//8] weights and their metadata.""" + x = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + w_int4 = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32) + b_q = pack_int4_exllama_shuffle(w_int4) + scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to( + dtype + ) + zp = ( + torch.randint(0, 16, (N, K // G), device=device, dtype=torch.int32).to(dtype) + if has_zp + else None + ) + return x, w_int4, b_q, scales, zp + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_zp", [False, True]) +@pytest.mark.parametrize( + "M,K,N,G", + [(17, 256, 512, 32), (32, 512, 256, 64), (33, 512, 512, 128), (64, 1024, 256, 128)], +) +def test_triton_prefill_gemm_matches_reference(dtype, has_zp, M, K, N, G): + """Prefill GEMM against a float32 oracle, over both unpacks and both the + asymmetric and symmetric dequants.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA/HIP device not available") + set_random_seed(0) + + x, w_int4, b_q, scales, zp = _make_prefill_case(M, K, N, G, dtype, has_zp) + out = triton_w4a16_skinny_fmt_gemm(a=x, b_q=b_q, scales=scales, group_size=G, zp=zp) + ref = _rdna_hybrid_w4a16_reference(x, w_int4, scales, zp, G, bias=None) + torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_gfx1151_tile_table_never_straddles_a_quant_group(dtype): + """BLOCK_K > group_size would give a tile's tail the wrong scale. + + The kernel loads one scale per BLOCK_K tile, so this is a correctness + invariant of the table, not a tuning preference. Checked in Python so it + holds for shapes no test has hardware for. + """ + for group_size in SUPPORTED_GROUP_SIZES: + for M in (1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096): + for N, K in [ + (512, 2048), + (4096, 4096), + (24576, 4096), + (4096, 12288), + (32768, 2048), + (1024, 8192), + ]: + _, _, block_k, _, _ = select_skinny_gfx1151_config( + M, N, K, group_size, dtype + ) + assert block_k <= group_size, ( + f"BLOCK_K={block_k} > group_size={group_size} " + f"at M={M} N={N} K={K} dtype={dtype}" + ) + assert block_k % 8 == 0, ( + f"BLOCK_K={block_k} must be a multiple of 8 " + f"(8 nibbles per packed int32)" + ) + + # --------------------------------------------------------------------------- # pack_int4_exllama_shuffle round-trips correctly # --------------------------------------------------------------------------- diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index 01c353c55474..ed65737af88e 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -71,6 +71,37 @@ def _on_gfx1151() -> bool: # --------------------------------------------------------------------------- +@triton.constexpr_function +def _target_is_gfx11() -> bool: + """True when the kernel is being compiled for RDNA3 (gfx11).""" + target = tl.target_info.current_target() + if target is None or target.backend != "hip": + return False + return str(target.arch).startswith("gfx11") + + +@triton.jit +def _int4_pair_to_fp16x2(x): + """Unpack two packed int4 nibbles into a uint32 holding two fp16 lanes, + each equal to 1024 + nibble, with one ``v_and_or_b32`` + (``(x & 0x000F000F) | 0x64006400``). + + OR-ing a 4-bit nibble into the low mantissa of fp16 1024.0 (0x6400) + bitcasts to exactly 1024+n. Doing it on a full 32-bit lane dequants two + nibbles per instruction, vs the scalar v_and_b16 + v_or_b16 pair Triton + emits from the elementwise form. + """ + mask = tl.full(x.shape, 0x000F000F, tl.int32) + return tl.inline_asm_elementwise( + asm="v_and_or_b32 $0, $1, $2, 0x64006400", + constraints="=v,v,v", + args=[x, mask], + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + + @triton.jit def _triton_w4a16_skinny_fmt_kernel( # Pointers @@ -106,6 +137,11 @@ def _triton_w4a16_skinny_fmt_kernel( When HAS_ZP=True, raw zero-points zp_raw are loaded from zp_ptr [N, K//G] and subtracted directly: (nibble - zp_raw) * scale. When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric). + + On the fp16 path the nibble arrives as ``b_raw`` = 1024 + nibble (the + magic-constant unpack), so the subtrahend absorbs the 1024: the arithmetic + is unchanged and every intermediate stays exact, since fp16 represents every + integer below 2048. """ pid_m = tl.program_id(0) pid_n = tl.program_id(1) @@ -137,10 +173,30 @@ def _triton_w4a16_skinny_fmt_kernel( mask_b = (offs_n[:, None] < N) & (offs_k8[None, :] < K8) b_packed = tl.load(b_ptrs, mask=mask_b, other=0) - b = tl.interleave(b_packed, b_packed) - b = tl.interleave(b, b) - b = tl.interleave(b, b) - b = (b >> shifts_full) & 0xF + if a.dtype == tl.float16 and _target_is_gfx11(): + # The ExLlama int32 holds the paired nibbles val[2p] @ bits[4p:4p+4] + # and val[2p+1] @ bits[16+4p:20+4p], so for pre-shift 4p (p=0..3), + # (x >> 4p) & 0x000F000F | 0x64006400 + # is one v_and_or_b32 producing a half2 = (1024+val[2p], + # 1024+val[2p+1]) in K order (signed shift is fine: the sign fill + # lands above bit 20 and is masked out). The interleave(lo, hi) lays + # b_raw out as half2 so the downstream affine also packs into + # v_pk_fma_f16. The dequant inner loop is VALU-issue-bound on gfx11, + # so this ~halves the dequant instruction count per WMMA. + shifts4 = (tl.arange(0, 4) * 4)[None, None, :] + bp_shift = tl.reshape( + b_packed[:, :, None] >> shifts4, (BLOCK_N, BLOCK_K // 2) + ) + packed_hl = _int4_pair_to_fp16x2(bp_shift) # u32 half2: 1024+nibble + lo = (packed_hl & 0xFFFF).to(tl.uint16).to(tl.float16, bitcast=True) + hi = (packed_hl >> 16).to(tl.uint16).to(tl.float16, bitcast=True) + b_raw = tl.interleave(lo, hi) # [BLOCK_N, BLOCK_K] fp16 = 1024+nibble + else: + # ExLlama unshuffle: replicate each int32 8x then per-lane shift+mask. + b = tl.interleave(b_packed, b_packed) + b = tl.interleave(b, b) + b = tl.interleave(b, b) + b = (b >> shifts_full) & 0xF # [BLOCK_N, BLOCK_K] g_idx = (k_start * BLOCK_K) // group_size scale_ptrs = scales_ptr + offs_n * num_groups + g_idx @@ -150,9 +206,24 @@ def _triton_w4a16_skinny_fmt_kernel( if HAS_ZP: zp_ptrs = zp_ptr + offs_n * num_groups + g_idx zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0) - b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None] + + if a.dtype == tl.float16: + # The magic unpack yields b_raw = 1024 + nibble, so fold the 1024 + # into the subtrahend: (b_raw - (1024 + zp)) == (nibble - zp), + # exactly, and the multiply that follows rounds once as before. + if not _target_is_gfx11(): + b_raw = (b | 0x6400).to(tl.uint16).to(tl.float16, bitcast=True) + c1024 = tl.full((), 1024.0, tl.float16) + if HAS_ZP: + b_fp = (b_raw - (c1024 + zp_raw)[:, None]) * scales[:, None] + else: + b_fp = (b_raw - (c1024 + ZP_BIAS)) * scales[:, None] else: - b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None] + # bf16 keeps the scalar int-domain subtract before the cast. + if HAS_ZP: + b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None] + else: + b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None] b_fp_t = tl.trans(b_fp) accumulator += tl.dot(a, b_fp_t, out_dtype=tl.float32) @@ -164,11 +235,15 @@ def _triton_w4a16_skinny_fmt_kernel( # Per-shape (group_size, K, N) -> (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, -# num_stages) tile-config overrides for prefill (M <= 128) on gfx1x. +# num_stages) tile-config overrides for prefill (M <= 128) on gfx1151. Applies +# to the SCALAR (bf16) dequant path only; the packed fp16 path is tuned by the +# ladder in _select_skinny_gfx1151_config and needs no per-shape entries. # Picked by sweeping benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py + a # per-config sweep script; only added when better than the generic heuristic # by > 20% at M=128. Re-run benchmarks after edits. -_GFX1X_PREFILL_OVERRIDES: dict[tuple[int, int, int], tuple[int, int, int, int, int]] = { +_GFX1151_BF16_PREFILL_OVERRIDES: dict[ + tuple[int, int, int], tuple[int, int, int, int, int] +] = { # SmolLM2-1.7B-Instruct-AWQ (gs=32, K=2048; gs forces BLOCK_K to 32 so # widen BLOCK_M and let Triton pipeline 4 stages to amortize the small # K-tile). @@ -186,6 +261,122 @@ def _triton_w4a16_skinny_fmt_kernel( } +# Explicit gfx1151 prefill tile selection -- DTYPE-AWARE. The kernel takes the +# packed v_and_or/v_pk_fma dequant for fp16 and the scalar dequant for bf16, and +# the two paths want different tiles (most visibly BLOCK_N at deep M: 256 for +# packed fp16 vs 64 for scalar bf16). +# +# fp16 (packed) -- tuned under do_bench_cudagraph with rotating cold weights +# over a broad shape catalog: +# * M <= 16: BLOCK_M=16 (more M-tiles fill the CUs at tiny M). +# * 17..64: BLOCK_M=32; small BLOCK_N keeps the grid large (a wide BLOCK_N +# leaves only ceil(N/BN) workgroups -- an M-blind BLOCK_N=256 was a 1.6-3x +# regression here). Square mid shapes take BLOCK_N=128/BLOCK_K=64. +# * 65..256: square -> BLOCK_N=128 (BLOCK_K=32 nw=8 at M>=128); tall -> 128. +# * 257..2047: the wide distilled BLOCK_N=256/BLOCK_M=128 tile. +# * M >= 2048: distilled BLOCK_N=256; BLOCK_M=64 for narrow+deep K (N<=2048 +# and K>=4096), else 128. +# +# bf16 (scalar) -- byte-for-byte the pre-existing scalar-tuned ladder, its +# per-shape overrides, and its pipeline depth, so bf16 is bit-for-bit unchanged: +# the packed fp16 table regresses bf16 by up to ~40% at deep M, where scalar bf16 +# wants BLOCK_N=64, not 256. num_stages=1 is deliberately NOT applied here -- +# measured at -15.8% end-to-end prefill on an asymmetric bf16 model. The packed +# fp16 path issues one per-group load and does not miss the pipelining; the +# scalar asymmetric path issues two (scale and zero point) and needs the software +# pipeline to hide the second gather. +# +# BLOCK_K is capped to group_size so a K-block never straddles a quant group +# (scale aliasing); gs=128 -- the bulk -- passes the table BLOCK_K through. +def _select_skinny_gfx1151_config( + M: int, N: int, K: int, group_size: int, dtype: torch.dtype +) -> tuple[int, int, int, int, int | None]: + """Return (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) for gfx1151. + + num_stages None means "leave Triton's default pipeline depth alone". + """ + num_stages: int | None = None + if dtype == torch.float16: + # The packed path issues a single per-group load, so the pipeline buys + # nothing and only costs registers. Not applied to bf16 -- see above. + num_stages = 1 + tall = K >= 2 * N # tall-K (down_proj-like) + # Very wide N with small K (e.g. gemma gate_up 32768x2048): memory-bound, + # wants the small square tile at tiny M, not BLOCK_M=16. + vwide_smallk = N >= 8192 and K <= 2048 + if M <= 16: # BLOCK_M=16: more M-tiles fill the CUs at tiny M + if N <= 1024 or vwide_smallk: + block_m, block_n, block_k, num_warps = 32, 32, 128, 4 + else: + block_m, block_n, block_k, num_warps = 16, 64, 128, 4 + elif M <= 32: + if vwide_smallk: + block_m, block_n, block_k, num_warps = 32, 32, 128, 4 + else: + block_m, block_n, block_k, num_warps = 32, 64, 128, 4 + elif M <= 64: + if tall or N >= 4 * K: # tall or very wide + block_m, block_n, block_k, num_warps = 32, 64, 128, 4 + else: # square mid + block_m, block_n, block_k, num_warps = 32, 128, 64, 4 + elif M <= 128: + if tall: + block_m, block_n, block_k, num_warps = 32, 128, 64, 4 + elif N >= 32768 and K <= 2048: + # Extremely wide + tiny K (e.g. gemma gate_up 32768x2048): + # BLOCK_N=128 collapses to 0.6x, needs 64. + block_m, block_n, block_k, num_warps = 128, 64, 64, 8 + elif N >= 16384: # very wide N (K>2048): BLOCK_N=128 wins + block_m, block_n, block_k, num_warps = 128, 128, 32, 8 + elif K <= 2048: # small-K square needs BLOCK_K=128 + block_m, block_n, block_k, num_warps = 32, 64, 128, 4 + else: # larger square + block_m, block_n, block_k, num_warps = 64, 128, 32, 4 + elif M <= 256: + block_m, block_n, block_k, num_warps = 128, 128, 32, 8 + elif M < 2048: # 257..2047 (mostly 512, 1024): wide distilled tile + block_m, block_n, block_k, num_warps = 128, 256, 32, 8 + else: # M >= 2048 (deep prefill) + if N <= 2048 and K >= 4096: # narrow + deep: halved BM saturates + block_m, block_n, block_k, num_warps = 64, 256, 32, 8 + else: + block_m, block_n, block_k, num_warps = 128, 256, 32, 8 + # Very narrow N at small/mid M: a wide BLOCK_N leaves too few N-tiles to + # fill the CUs, so clamp it. At M>=1024 the M-tiles already saturate. + if N <= 1024 and M <= 512: + block_n = min(block_n, 32) + else: + # Scalar-dequant path (bf16): the pre-existing scalar-tuned ladder. + key = (group_size, K, N) + override = _GFX1151_BF16_PREFILL_OVERRIDES.get(key) if M <= 128 else None + if override is not None: + block_m, block_n, block_k, num_warps, num_stages = override + elif M <= 32: + block_m, block_n, block_k, num_warps = 32, 32, 128, 4 + elif M <= 64: + block_m, block_n, block_k, num_warps = 64, 64, 32, 4 + elif M <= 128: + if K >= 2 * N: # tall K (down_proj) + block_m, block_n, block_k, num_warps = 64, 16, 64, 1 + elif N > K: # wide N (qkv / gate_up) + block_m, block_n, block_k, num_warps = 64, 64, 64, 4 + else: # N ~= K (o_proj) + block_m, block_n, block_k, num_warps = 64, 32, 64, 4 + elif M <= 1024: + if K >= 2 * N: # tall K (down_proj) + block_m, block_n, block_k, num_warps = 64, 64, 64, 4 + elif N >= 4 * K: # very wide N (gate_up) + block_m, block_n, block_k, num_warps = 128, 64, 64, 8 + else: + block_m, block_n, block_k, num_warps = 64, 128, 32, 4 + else: # M > 1024 + if K >= 2 * N: # tall K (down_proj) + block_m, block_n, block_k, num_warps = 128, 512, 32, 16 + else: + block_m, block_n, block_k, num_warps = 128, 64, 64, 8 + return block_m, block_n, min(block_k, group_size), num_warps, num_stages + + def triton_w4a16_skinny_fmt_gemm( a: torch.Tensor, # [M, K] fp16/bf16 b_q: torch.Tensor, # [N, K//8] int32 (ExLlama shuffle packed) @@ -232,8 +423,8 @@ def triton_w4a16_skinny_fmt_gemm( c = torch.empty((M, N), dtype=a.dtype, device=a.device) - # num_stages stays None unless a per-shape override sets it, so the - # generic heuristics fall back to Triton's default pipeline depth. + # num_stages stays None unless the tile table sets it, so the generic + # heuristics fall back to Triton's default pipeline depth. num_stages: int | None = None if _on_gfx12x(): # Tuned on gfx1201 (Radeon AI PRO R9700, 32 CUs, 32-wide wavefronts) @@ -269,39 +460,14 @@ def triton_w4a16_skinny_fmt_gemm( else: BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 128, 32, 8 elif _on_gfx1151(): - # Tuned on gfx1151 (Strix Halo, 40 CUs, 32-wide wavefronts) - # using Qwen3-4B weight shapes with group_size=128. - # Per-shape overrides for known prefill regressions live in a small - # lookup table — see _GFX1X_PREFILL_OVERRIDES below. Re-run + # gfx1151 (Strix Halo, 40 CUs, 32-wide wavefronts): per-(M, N, K) tile + # config from the dtype-aware table, since the packed fp16 dequant and + # the scalar bf16 dequant want different tiles. See + # _select_skinny_gfx1151_config; re-run # benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py after edits. - override = ( - _GFX1X_PREFILL_OVERRIDES.get((group_size, K, N)) if M <= 128 else None + BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = ( + _select_skinny_gfx1151_config(M, N, K, group_size, a.dtype) ) - if override is not None: - BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = override - elif M <= 32: - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 32, 32, 128, 4 - elif M <= 64: - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 64, 32, 4 - elif M <= 128: - if K >= 2 * N: # tall K (e.g. down_proj) - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 16, 64, 1 - elif N > K: # wide N (e.g. qkv_proj, gate_up_proj) - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 64, 64, 4 - else: # N ~= K (e.g. o_proj) - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 32, 64, 4 - elif M <= 1024: - if K >= 2 * N: # tall K (e.g. down_proj) - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 64, 64, 4 - elif N >= 4 * K: # very wide N (e.g. gate_up_proj) - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 64, 64, 8 - else: - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 64, 128, 32, 4 - else: - if K >= 2 * N: # tall K (e.g. down_proj) - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 512, 32, 16 - else: - BLOCK_M, BLOCK_N, BLOCK_K, num_warps = 128, 64, 64, 8 else: num_warps = 4 if M <= 32: diff --git a/vllm/triton_utils/importing.py b/vllm/triton_utils/importing.py index 966ab55f326b..ead5512ffacb 100644 --- a/vllm/triton_utils/importing.py +++ b/vllm/triton_utils/importing.py @@ -108,6 +108,10 @@ def __init__(self): self.jit = self._dummy_decorator("jit") self.autotune = self._dummy_decorator("autotune") self.heuristics = self._dummy_decorator("heuristics") + # Bodies of constexpr_function helpers only run while a kernel is being + # compiled, so the dummy decorator is enough: without Triton nothing + # compiles, and the module must still import. + self.constexpr_function = self._dummy_decorator("constexpr_function") self.Config = self._dummy_decorator("Config") self.cdiv = cdiv self.language = TritonLanguagePlaceholder()