diff --git a/.gitignore b/.gitignore index 2c4e135e58dc..cf5ff4a763a7 100644 --- a/.gitignore +++ b/.gitignore @@ -249,3 +249,6 @@ vllm/grpc/vllm_engine_pb2.pyi # Ignore generated cpu headers csrc/cpu/cpu_attn_dispatch_generated.h + +# Hand-tuned AMDGCN build artifact (assembled from .amdgcn at build time) +*.hsaco diff --git a/CMakeLists.txt b/CMakeLists.txt index aae50f1a6dac..626fb5bad75b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1376,6 +1376,25 @@ if(VLLM_GPU_LANG STREQUAL "HIP") endif() endif() +# Hand-editable AMDGCN kernels (gfx1151): assemble .amdgcn -> .hsaco at build +# time so editable installs (and wheels) ship a ready binary next to the source +# .amdgcn. The runtime loader (asm_w4a16.py) re-assembles on demand if the +# .amdgcn is edited without a rebuild, so this is a convenience, not a hard dep. +if (VLLM_GPU_LANG STREQUAL "HIP" AND VLLM_GPU_ARCHES MATCHES "gfx1151") + set(_HANDASM_DIR + ${CMAKE_CURRENT_SOURCE_DIR}/vllm/model_executor/kernels/linear/mixed_precision/asm) + set(_HANDASM_SRC ${_HANDASM_DIR}/hybrid_w4a16_skinny_gfx1151.amdgcn) + set(_HANDASM_OUT ${_HANDASM_DIR}/hybrid_w4a16_skinny_gfx1151.hsaco) + add_custom_command( + OUTPUT ${_HANDASM_OUT} + COMMAND ${CMAKE_HIP_COMPILER} -target amdgcn-amd-amdhsa -mcpu=gfx1151 + -x assembler ${_HANDASM_SRC} -o ${_HANDASM_OUT} + DEPENDS ${_HANDASM_SRC} + COMMENT "Assembling hand-tuned W4A16 AMDGCN -> hsaco (gfx1151)" + VERBATIM) + add_custom_target(vllm_handasm_w4a16 ALL DEPENDS ${_HANDASM_OUT}) +endif() + # Must run after the last HIP `define_extension_target` so every extension # has registered its sources. if (VLLM_GPU_LANG STREQUAL "HIP") diff --git a/benchmarks/kernels/benchmark_hybrid_w4a16.py b/benchmarks/kernels/benchmark_hybrid_w4a16.py new file mode 100644 index 000000000000..24ec2f6cf95d --- /dev/null +++ b/benchmarks/kernels/benchmark_hybrid_w4a16.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CLI benchmark for the dense hybrid_w4a16 GEMM kernel. + +Drives the same measurement rig as the perf regression test +(``tests/kernels/quantization/test_hybrid_w4a16_perf.py``): packed int4 +weights via the production ``pack_skinny_int4`` helper, cold-weight MALL +rotation, and ``do_bench_cudagraph`` median timing -- so the reported +time/call matches what a real prefill forward pass sees (and what shows up +in a torch profiler trace). + +Defaults reproduce the Qwen3-VL-4B-Instruct-AWQ-4bit fused gate_up_proj +GEMM at the 1322-token prefill shape from the gfx1151 roofline analysis: +M=1322, N=19456, K=2560, group_size=32, symmetric (no zero-point), bf16. + +Usage:: + + # Default: the up_proj shape at M=1322 + python benchmarks/kernels/benchmark_hybrid_w4a16.py + + # All Qwen3-VL-4B AWQ decoder GEMMs across a batch sweep + python benchmarks/kernels/benchmark_hybrid_w4a16.py --model qwen3vl-4b-awq + + # A custom shape / quant config + python benchmarks/kernels/benchmark_hybrid_w4a16.py \\ + --k 2560 --n 19456 --group-size 32 --provider hybrid-w4a16-bf16 \\ + --batches 1024 1322 2048 + +The provider name follows the test convention +``hybrid-w4a16[-zp][-bf16]``: ``-zp`` selects the asymmetric per-group +zero-point dequant path; ``-bf16`` runs with bfloat16 activations/scales +instead of float16. Symmetric fp16 is the bare ``hybrid-w4a16``. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +import torch + +# Make the in-tree tests helper importable from a stock checkout, mirroring +# bench_hybrid_w4a16_moe.py. +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from tests.kernels.quantization.test_hybrid_w4a16_perf import ( # noqa: E402 + PROVIDERS, + _compute_packed_scale_zp, + _provider_dtype, + _provider_use_zp, + measure_tflops, + prepare_hybrid_weights, +) + +# Named (K, N) decoder GEMM presets. K=in_features, N=out_features. +# qwen3vl-4b-awq: group_size=32, symmetric, bf16 (see model config). +MODEL_SHAPES: dict[str, list[dict]] = { + "qwen3vl-4b-awq": [ + {"k": 2560, "n": 19456, "group_size": 32, "comment": "gate_up_proj"}, + {"k": 9728, "n": 2560, "group_size": 32, "comment": "down_proj"}, + {"k": 2560, "n": 6144, "group_size": 32, "comment": "qkv_proj"}, + {"k": 2560, "n": 4096, "group_size": 32, "comment": "o_proj"}, + ], +} + + +def _bench_one( + m: int, k: int, n: int, group_size: int, provider: str +) -> tuple[str, float, float]: + """Return (kernel, tflops, ms) for one shape/provider/batch.""" + weights = prepare_hybrid_weights(k, n, group_size, dtype=_provider_dtype(provider)) + measure_tflops(m, weights, k, n, group_size, provider) # warmup + kernel, tflops = measure_tflops(m, weights, k, n, group_size, provider) + ms = (2 * m * n * k) * 1e-12 / (tflops * 1e-3) + return kernel, tflops, ms + + +def _apply(a, weights, group_size, provider): + """Run the production apply path once (respects VLLM_W4A16_HANDASM).""" + from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import ( + _hybrid_w4a16_apply_impl, + ) + from vllm.utils.platform_utils import num_compute_units + + use_zp = _provider_use_zp(provider) + packed_scale_zp = _compute_packed_scale_zp( + weights["w_s_skinny"], weights["w_zp"] if use_zp else None, a.dtype + ) + return _hybrid_w4a16_apply_impl( + a, + weights["w_q_skinny"], + weights["w_s_skinny"], + weights["w_q_skinny_i32"], + weights["w_zp"] if use_zp else None, + None, # bias + num_compute_units(), + group_size, + packed_scale_zp, + ) + + +def _check_one( + m: int, k: int, n: int, group_size: int, provider: str +) -> tuple[bool, float]: + """Compare the active (hand-asm) kernel against the Triton reference on the + SAME inputs. Returns (handasm_actually_used, cosine_similarity). + + Uses Triton as the trusted reference: run once with VLLM_W4A16_HANDASM=0 + (force Triton) and once with =1 (hand-asm if its config matches this shape), + then compare. When the hand-asm config does not match, both runs are Triton + and the check is trivially exact (reported as handasm_used=False). + """ + from vllm.model_executor.kernels.linear.mixed_precision import asm_w4a16 + from vllm.platforms.rocm import on_gfx1151 + + dtype = _provider_dtype(provider) + weights = prepare_hybrid_weights(k, n, group_size, dtype=dtype) + a = torch.randn((m, k), device="cuda", dtype=dtype) + + # Did the dispatcher's pinned config match this launch? + block = _select_skinny_config(m, n, k, group_size, dtype) + handasm_used = ( + on_gfx1151() + and block is not None + and asm_w4a16.config_matches( + dtype, _provider_use_zp(provider), group_size, *block + ) + ) + + os.environ["VLLM_W4A16_HANDASM"] = "0" + c_ref = _apply(a, weights, group_size, provider).float() + os.environ["VLLM_W4A16_HANDASM"] = "1" + c = _apply(a, weights, group_size, provider).float() + + cos = torch.nn.functional.cosine_similarity( + c.flatten(), c_ref.flatten(), dim=0 + ).item() + return handasm_used, cos + + +def _select_skinny_config(m, n, k, group_size, dtype): + """The tile (block_m, block_n, block_k, num_warps) the dispatcher would pick, + or None off gfx11.""" + from vllm.model_executor.kernels.linear.mixed_precision.hybrid_w4a16 import ( + _select_skinny_gfx11_config, + ) + from vllm.platforms.rocm import on_gfx1x + + if not on_gfx1x(): + return None + return _select_skinny_gfx11_config(m, n, k, group_size, dtype) + + +def main() -> None: + p = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument( + "--model", + choices=sorted(MODEL_SHAPES), + help="Benchmark all decoder GEMMs for a named model preset.", + ) + p.add_argument("--k", type=int, default=2560, help="in_features (default 2560)") + p.add_argument("--n", type=int, default=19456, help="out_features (default 19456)") + p.add_argument("--group-size", type=int, default=32, help="quant group size") + p.add_argument( + "--provider", + default="hybrid-w4a16-bf16", + choices=PROVIDERS, + help="quant variant (default hybrid-w4a16-bf16: bf16, symmetric)", + ) + p.add_argument( + "--batches", + type=int, + nargs="+", + default=[1322], + help="M (token) batch sizes to sweep (default: 1322)", + ) + p.add_argument( + "--check", + action="store_true", + help="Verify the active kernel matches the Triton reference (cos > " + "0.999) on identical inputs before/instead of timing. Exits non-zero " + "on mismatch. Use this after editing the hand-asm .amdgcn.", + ) + p.add_argument( + "--check-tol", + type=float, + default=0.999, + help="Minimum cosine similarity for --check (default 0.999)", + ) + args = p.parse_args() + + if args.model: + shapes = MODEL_SHAPES[args.model] + else: + shapes = [ + {"k": args.k, "n": args.n, "group_size": args.group_size, "comment": ""} + ] + + if args.check: + failures = 0 + for s in shapes: + k, n, gs = s["k"], s["n"], s["group_size"] + for m in args.batches: + used, cos = _check_one(m, k, n, gs, args.provider) + ok = cos >= args.check_tol + failures += not ok + where = "hand-asm" if used else "triton-only (no hand-asm)" + print( + f" CHECK {m}x{n}x{k} g{gs} {args.provider}: " + f"cos={cos:.6f} [{where}] -> {'PASS' if ok else 'FAIL'}" + ) + torch.accelerator.synchronize() + if failures: + raise SystemExit(f"{failures} correctness check(s) FAILED") + return + + hdr = f"{'shape (MxNxK)':>22} {'g':>4} {'kernel':>22} {'TFLOP/s':>9} {'ms/call':>9}" + print(hdr) + print("-" * len(hdr)) + for s in shapes: + k, n, gs = s["k"], s["n"], s["group_size"] + for m in args.batches: + kernel, tflops, ms = _bench_one(m, k, n, gs, args.provider) + tag = f" # {s['comment']}" if s["comment"] else "" + shape_str = f"{m}x{n}x{k}" + print( + f"{shape_str:>22} {gs:>4} {kernel:>22} {tflops:>9.2f} {ms:>9.3f}{tag}" + ) + torch.accelerator.synchronize() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 117bc925e9dd..2399eb5b44a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,7 +131,8 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/tool-parser/src/gemma4.rs", "rust/src/text/src/output/decoded.rs", - "rust/src/tokenizer/src/incremental.rs", "rust/src/reasoning-parser/src/tests.rs"] + "rust/src/tokenizer/src/incremental.rs", "rust/src/reasoning-parser/src/tests.rs", + "*.amdgcn"] ignore-hidden = false [tool.typos.default] diff --git a/vllm/model_executor/kernels/linear/mixed_precision/asm/hybrid_w4a16_skinny_gfx1151.amdgcn b/vllm/model_executor/kernels/linear/mixed_precision/asm/hybrid_w4a16_skinny_gfx1151.amdgcn new file mode 100644 index 000000000000..22d2431461c1 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mixed_precision/asm/hybrid_w4a16_skinny_gfx1151.amdgcn @@ -0,0 +1,956 @@ + .amdgcn_target "amdgcn-amd-amdhsa--gfx1151" + .amdhsa_code_object_version 5 + +; --- W4A16 dequant macros (assembler text-expansion; zero runtime cost) --- +.macro DQ_EXT_LO dst, src ; even-K: (bf16(128+n) low, bf16(-136) high) + v_and_or_b32 \dst, \src, s35, 0xc3084300 +.endm +.macro DQ_EXT_HI dst, src ; odd-K: (bf16(-136) low, bf16(128+n) HIGH); nibble already @bit16 + v_and_or_b32 \dst, \src, s36, 0x4300c308 +.endm +.macro DQ_PACK dst, klo, khi ; two dot2: klo->low half, khi->high half of dst + v_dot2_bf16_bf16 \dst, \klo, v124, 0 + v_dot2_bf16_bf16 \dst, \khi, v124, 0 op_sel:[0,0,0,1] +.endm + + .text + .globl _triton_w4a16_skinny_fmt_kernel ; -- Begin function _triton_w4a16_skinny_fmt_kernel + .p2align 8 + .type _triton_w4a16_skinny_fmt_kernel,@function +_triton_w4a16_skinny_fmt_kernel: ; @_triton_w4a16_skinny_fmt_kernel +.Lfunc_begin0: + .file 1 "/scratch/mgehre/vllm/vllm/model_executor/kernels/linear/mixed_precision" "hybrid_w4a16.py" + .cfi_sections .debug_frame + .cfi_startproc +; %bb.0: + s_clause 0x1 + s_load_b32 s26, s[0:1], 0x30 + s_load_b128 s[20:23], s[0:1], 0x20 + v_dual_mov_b32 v8, 0 :: v_dual_and_b32 v33, 15, v0 +.Ltmp0: + v_and_b32_e32 v34, 0xc0, v0 + v_and_b32_e32 v35, 32, v0 + s_lshl_b32 s24, s2, 7 + s_delay_alu instid0(VALU_DEP_3) + v_mov_b32_e32 v7, v8 + v_mov_b32_e32 v6, v8 + v_mov_b32_e32 v5, v8 + v_mov_b32_e32 v4, v8 + v_mov_b32_e32 v3, v8 + v_mov_b32_e32 v2, v8 + v_mov_b32_e32 v1, v8 + v_mov_b32_e32 v16, v8 + v_mov_b32_e32 v15, v8 + v_mov_b32_e32 v14, v8 + v_mov_b32_e32 v13, v8 + v_mov_b32_e32 v12, v8 + v_mov_b32_e32 v11, v8 + v_mov_b32_e32 v10, v8 + v_mov_b32_e32 v9, v8 + v_mov_b32_e32 v24, v8 + v_mov_b32_e32 v23, v8 + v_mov_b32_e32 v22, v8 + v_mov_b32_e32 v21, v8 + v_mov_b32_e32 v20, v8 + v_mov_b32_e32 v19, v8 + v_mov_b32_e32 v18, v8 + v_mov_b32_e32 v17, v8 + v_mov_b32_e32 v32, v8 + v_mov_b32_e32 v31, v8 + v_mov_b32_e32 v30, v8 + v_mov_b32_e32 v29, v8 + v_mov_b32_e32 v28, v8 + v_mov_b32_e32 v27, v8 + v_mov_b32_e32 v26, v8 + v_mov_b32_e32 v25, v8 +.Ltmp1: + .file 2 "/scratch/mgehre/vllm/.venv/lib/python3.12/site-packages/triton/language" "standard.py" + s_waitcnt lgkmcnt(0) + s_add_i32 s2, s26, 31 +.Ltmp2: + s_lshl_b32 s25, s3, 6 + s_cmp_lt_i32 s2, 32 + s_cbranch_scc1 .LBB0_3 +; %bb.1: ; %.lr.ph + s_clause 0x2 + s_load_b128 s[4:7], s[0:1], 0x34 + s_load_b128 s[16:19], s[0:1], 0x0 + s_load_b64 s[8:9], s[0:1], 0x10 + v_lshrrev_b32_e32 v1, 2, v0 +.Ltmp3: + s_ashr_i32 s0, s2, 31 + v_dual_mov_b32 v25, 0 :: v_dual_and_b32 v36, 3, v0 + s_lshr_b32 s0, s0, 27 + s_delay_alu instid0(VALU_DEP_2) | instskip(SKIP_1) | instid1(VALU_DEP_2) +.Ltmp4: + v_or_b32_e32 v5, s24, v1 +.Ltmp5: + s_add_i32 s2, s2, s0 + v_dual_mov_b32 v29, v25 :: v_dual_lshlrev_b32 v4, 3, v0 + v_dual_mov_b32 v26, v25 :: v_dual_lshlrev_b32 v3, 1, v0 + s_delay_alu instid0(VALU_DEP_3) | instskip(SKIP_1) | instid1(VALU_DEP_4) +.Ltmp6: + v_or_b32_e32 v8, 64, v5 + v_cmp_gt_i32_e32 vcc_lo, s22, v5 + v_dual_mov_b32 v31, v25 :: v_dual_and_b32 v4, 48, v4 + v_dual_mov_b32 v27, v25 :: v_dual_lshlrev_b32 v2, 4, v0 + s_waitcnt lgkmcnt(0) + s_abs_i32 s27, s7 + v_or_b32_e32 v6, s25, v1 + s_cvt_f32_u32 s0, s27 + v_dual_mov_b32 v30, v25 :: v_dual_and_b32 v3, 48, v3 + v_dual_mov_b32 v32, v25 :: v_dual_add_nc_u32 v7, s25, v1 + s_delay_alu instid0(SALU_CYCLE_1) | instskip(SKIP_4) | instid1(VALU_DEP_3) + v_rcp_iflag_f32_e32 v5, s0 + v_dual_mov_b32 v18, v25 :: v_dual_add_nc_u32 v1, s24, v1 + v_lshl_or_b32 v4, v33, 6, v4 +.Ltmp7: + s_ashr_i32 s28, s2, 5 + v_xor_b32_e32 v2, v2, v3 + v_dual_mov_b32 v20, v25 :: v_dual_add_nc_u32 v3, 64, v1 +.Ltmp8: + v_cmp_gt_i32_e64 s1, s22, v8 + v_lshl_or_b32 v8, v34, 4, v4 + s_delay_alu instid0(TRANS32_DEP_1) + v_readfirstlane_b32 s2, v5 + v_lshl_or_b32 v4, v35, 5, v4 + v_mul_lo_u32 v7, s5, v7 + v_mul_lo_u32 v1, s26, v1 + s_mul_f32 s2, s2, 0x4f7ffffe + v_cmp_gt_i32_e64 s0, s23, v6 + v_mul_lo_u32 v38, s6, v6 + v_xor_b32_e32 v5, 16, v8 + v_xor_b32_e32 v6, 32, v8 + v_xor_b32_e32 v9, 48, v8 + v_xor_b32_e32 v10, 16, v4 + v_xor_b32_e32 v11, 32, v4 + v_xor_b32_e32 v12, 48, v4 + s_cvt_u32_f32 s2, s2 + v_dual_mov_b32 v22, v25 :: v_dual_lshlrev_b32 v13, 2, v36 + v_dual_mov_b32 v17, v25 :: v_dual_lshlrev_b32 v14, 4, v36 + s_sub_i32 s3, 0, s27 + v_dual_mov_b32 v28, v25 :: v_dual_lshlrev_b32 v37, 3, v36 + s_mul_i32 s3, s3, s2 + v_lshl_add_u32 v39, v7, 2, v13 + v_lshl_add_u32 v40, v1, 1, v14 + v_dual_mov_b32 v19, v25 :: v_dual_add_nc_u32 v42, 0, v2 + v_dual_mov_b32 v24, v25 :: v_dual_add_nc_u32 v43, 0, v8 + v_dual_mov_b32 v21, v25 :: v_dual_add_nc_u32 v44, 0, v5 + v_add_nc_u32_e32 v45, 0, v6 + v_dual_mov_b32 v23, v25 :: v_dual_add_nc_u32 v46, 0, v9 + v_add_nc_u32_e32 v47, 0, v4 + v_dual_mov_b32 v9, v25 :: v_dual_add_nc_u32 v48, 0, v10 + v_dual_mov_b32 v14, v25 :: v_dual_add_nc_u32 v49, 0, v11 + v_dual_mov_b32 v11, v25 :: v_dual_add_nc_u32 v82, 0, v12 + v_mov_b32_e32 v10, v25 + v_mov_b32_e32 v12, v25 + v_mov_b32_e32 v13, v25 + v_mov_b32_e32 v15, v25 + v_mov_b32_e32 v16, v25 + v_mov_b32_e32 v1, v25 + v_mov_b32_e32 v2, v25 + v_mov_b32_e32 v3, v25 + v_mov_b32_e32 v4, v25 + v_mov_b32_e32 v5, v25 + v_mov_b32_e32 v6, v25 + v_mov_b32_e32 v7, v25 + v_mov_b32_e32 v8, v25 + s_mov_b32 s15, 0x31027000 + ; --- per-buffer num_records (bytes): hardware clamps OOB loads to 0 (no fault). + ; --- dims still live here: M=s22 N=s23 K=s26 stride_bn=s5 num_groups=s6. + s_mul_i32 s14, s22, s26 ; activation: M*K + s_lshl_b32 s14, s14, 1 ; *2 (bf16) -> s[12:15] dword2 + s_lshl_b32 s39, s26, 7 ; A1 voffset delta = 128*K (64 rows); share base+soffset + s_mul_i32 s37, s23, s5 ; weight: N*stride_bn + s_lshl_b32 s37, s37, 2 ; *4 (int32) -> s[16:19] dword2 (temp) + s_mul_i32 s38, s23, s6 ; scale: N*num_groups + s_lshl_b32 s38, s38, 1 ; *2 (bf16) -> s[8:11] dword2 (temp) + s_and_b32 s17, s17, 0xffff + s_and_b32 s19, s19, 0xffff + s_mul_hi_u32 s3, s2, s3 + s_and_b32 s9, s9, 0xffff + s_ashr_i32 s6, s7, 31 + s_mov_b32 s5, 0 + s_add_i32 s7, s2, s3 + s_mov_b32 s12, s16 + s_mov_b32 s13, s17 + s_mov_b32 s16, s18 + s_mov_b32 s17, s19 + s_mov_b32 s18, s37 ; weight num_records = N*stride_bn*4 + s_mov_b32 s19, s15 + s_mov_b32 s10, s38 ; scale num_records = N*num_groups*2 + s_mov_b32 s11, s15 + s_mov_b32 s35, 15 ; mask_lo (SGPR; survives the loop, ds_load only writes VGPRs) + s_mov_b32 s36, 0xf0000 ; mask_hi + + + + ; =========================================================================== + ; main loop: K-reduction over in_features (K), BLOCK_K=32 per iteration + ; iterates ceil(K / BLOCK_K) tiles; each tile: prefetched global loads -> + ; int4 dequant -> bf16 WMMA accumulate into the output fragment + ; =========================================================================== + ; --- prologue: prefetch iter 0 --- + v_add_nc_u32_e32 v125, s5, v37 + v_cmp_gt_i32_e64 s2, s4, v36 + s_ashr_i32 s3, s5, 31 + s_abs_i32 s29, s5 + s_xor_b32 s31, s3, s6 + v_cmp_gt_i32_e64 s3, s26, v125 + s_mul_hi_u32 s30, s29, s7 + s_and_b32 s2, s0, s2 + s_mul_i32 s33, s30, s27 + v_cndmask_b32_e64 v125, 0x80000000, v39, s2 + s_sub_i32 s29, s29, s33 + s_and_b32 s2, vcc_lo, s3 + s_add_i32 s34, s30, 1 + s_sub_i32 s33, s29, s27 + v_cndmask_b32_e64 v126, 0x80000000, v40, s2 + s_and_b32 s2, s1, s3 + s_cmp_ge_u32 s29, s27 + v_cndmask_b32_e64 v127, 0x80000000, v40, s2 ; A1 mask on shared v40 base + buffer_load_b32 v123, v125, s[16:19], 0 offen + s_cselect_b32 s2, s34, s30 + s_cselect_b32 s3, s33, s29 + s_add_i32 s29, s2, 1 + s_cmp_ge_u32 s3, s27 + buffer_load_b128 v[115:118], v126, s[12:15], 0 offen + s_cselect_b32 s2, s29, s2 + s_xor_b32 s2, s2, s31 + s_sub_i32 s2, s2, s31 + v_add_lshl_u32 v125, s2, v38, 1 + s_add_i32 s5, s5, 32 + v_cndmask_b32_e64 v125, 0x80000000, v125, s0 + buffer_load_u16 v124, v125, s[8:11], 0 offen + buffer_load_b128 v[119:122], v127, s[12:15], s39 offen ; act A1 = v40+128*K +.LBB0_2: ; =>This Inner Loop Header: Depth=1 + s_waitcnt vmcnt(0) + s_barrier + ds_store_b128 v42, v[115:118] + ds_store_b128 v42, v[119:122] offset:4096 + v_lshrrev_b32_e32 v95, 4, v123 ; view>>4 + v_lshrrev_b32_e32 v96, 8, v123 ; view>>8 + v_lshrrev_b32_e32 v97, 12, v123 ; view>>12 + DQ_EXT_LO v84, v123 ; K0 + DQ_EXT_HI v85, v123 ; K1 + DQ_EXT_LO v86, v95 ; K2 + DQ_EXT_HI v87, v95 ; K3 + DQ_EXT_LO v88, v96 ; K4 + DQ_EXT_HI v89, v96 ; K5 + DQ_EXT_LO v90, v97 ; K6 + DQ_EXT_HI v91, v97 ; K7 + v_lshl_or_b32 v124, v124, 16, v124 ; (scale, scale) + s_waitcnt lgkmcnt(0) + s_barrier + ds_load_b128 v[50:53], v43 + ds_load_b128 v[58:61], v43 offset:4096 + ds_load_b128 v[54:57], v44 + ds_load_b128 v[62:65], v44 offset:4096 + ds_load_b128 v[66:69], v45 + ds_load_b128 v[74:77], v45 offset:4096 + ds_load_b128 v[70:73], v46 + ds_load_b128 v[78:81], v46 offset:4096 + DQ_PACK v94, v84, v85 + DQ_PACK v95, v86, v87 + DQ_PACK v96, v88, v89 + DQ_PACK v97, v90, v91 + ; --- prefetch iter k+1 (overlaps ds_store B + ds_load B + WMMA) --- + ; OOB guards removed: group_size(32)==BLOCK_K divides K => no K-tail ever; + ; M-tail + speculative last prefetch are clamped by the buffer descriptor. + s_lshr_b32 s37, s5, 1 ; weight K byte offset = s5/2 (4-bit packed) + s_lshl_b32 s38, s5, 1 ; act A0 K byte offset = s5*2 (bf16) + s_add_i32 s40, s38, s39 ; act A1 K byte offset = A0 + 128*K + s_ashr_i32 s3, s5, 31 + s_abs_i32 s29, s5 + s_xor_b32 s31, s3, s6 + s_mul_hi_u32 s30, s29, s7 + s_mul_i32 s33, s30, s27 + s_sub_i32 s29, s29, s33 + s_add_i32 s34, s30, 1 + s_sub_i32 s33, s29, s27 + s_cmp_ge_u32 s29, s27 + buffer_load_b32 v123, v39, s[16:19], s37 offen ; weight (base + s5/2) + s_cselect_b32 s2, s34, s30 + s_cselect_b32 s3, s33, s29 + s_add_i32 s29, s2, 1 + s_cmp_ge_u32 s3, s27 + buffer_load_b128 v[115:118], v40, s[12:15], s38 offen ; act A0 (base + s5*2) + s_cselect_b32 s2, s29, s2 + s_xor_b32 s2, s2, s31 + s_sub_i32 s2, s2, s31 + v_add_lshl_u32 v125, s2, v38, 1 ; scale group addr + s_add_i32 s28, s28, -1 + s_add_i32 s5, s5, 32 + s_cmp_lg_u32 s28, 0 + buffer_load_u16 v124, v125, s[8:11], 0 offen ; scale + buffer_load_b128 v[119:122], v40, s[12:15], s40 offen ; act A1 (base + s5*2 + 128*K) + s_waitcnt lgkmcnt(0) + s_barrier + ds_store_b128 v42, v[94:97] + s_waitcnt lgkmcnt(0) + s_barrier + ds_load_b128 v[83:86], v47 + ds_load_b128 v[91:94], v47 offset:2048 + ds_load_b128 v[87:90], v48 + ds_load_b128 v[95:98], v48 offset:2048 + ds_load_b128 v[99:102], v49 + ds_load_b128 v[107:110], v49 offset:2048 + ds_load_b128 v[103:106], v82 + ds_load_b128 v[111:114], v82 offset:2048 + s_waitcnt lgkmcnt(5) + v_wmma_f32_16x16x16_bf16 v[25:32], v[83:90], v[50:57], v[25:32] + s_waitcnt lgkmcnt(4) + v_wmma_f32_16x16x16_bf16 v[17:24], v[91:98], v[50:57], v[17:24] + v_wmma_f32_16x16x16_bf16 v[9:16], v[83:90], v[58:65], v[9:16] + v_wmma_f32_16x16x16_bf16 v[1:8], v[91:98], v[58:65], v[1:8] + s_waitcnt lgkmcnt(1) + v_wmma_f32_16x16x16_bf16 v[25:32], v[99:106], v[66:73], v[25:32] + s_waitcnt lgkmcnt(0) + v_wmma_f32_16x16x16_bf16 v[17:24], v[107:114], v[66:73], v[17:24] + v_wmma_f32_16x16x16_bf16 v[9:16], v[99:106], v[74:81], v[9:16] + v_wmma_f32_16x16x16_bf16 v[1:8], v[107:114], v[74:81], v[1:8] + s_cbranch_scc1 .LBB0_2 ; backedge: branch to next K-tile + ; --- end of main loop (fall through to epilogue) --- + + + +.LBB0_3: ; %._crit_edge + v_lshrrev_b32_e32 v34, 2, v34 + v_lshrrev_b32_e32 v0, 4, v0 + v_lshrrev_b32_e32 v35, 1, v35 + s_mul_i32 s0, s23, s24 + v_bfe_u32 v63, v26, 16, 1 + v_or_b32_e32 v33, v34, v33 + v_bfe_u32 v34, v25, 16, 1 + v_and_or_b32 v39, v0, 1, v35 + s_add_i32 s0, s0, s25 + s_lshl_b32 s1, s23, 6 + v_or_b32_e32 v66, s24, v33 + v_mul_lo_u32 v33, s23, v33 + v_or_b32_e32 v0, 46, v39 + v_add3_u32 v34, v25, v34, 0x7fff + v_or_b32_e32 v68, s25, v39 + v_or_b32_e32 v62, 2, v39 + v_cmp_gt_i32_e64 s16, s22, v66 + v_or_b32_e32 v43, s25, v0 + v_lshrrev_b32_e32 v34, 16, v34 + v_add_nc_u32_e32 v99, s0, v33 + v_cmp_gt_i32_e64 s15, s23, v68 + v_cmp_o_f32_e64 s17, v25, v25 + v_or_b32_e32 v65, s25, v62 + v_add3_u32 v63, v26, v63, 0x7fff + v_add3_u32 v33, s0, s1, v33 + v_cmp_gt_i32_e64 s0, s23, v43 + v_add_lshl_u32 v43, v99, v39, 1 + v_or_b32_e32 v61, 4, v39 + v_bfe_u32 v67, v27, 16, 1 + v_cndmask_b32_e64 v25, 0x7fff, v34, s17 + s_and_b32 s17, s16, s15 + v_or_b32_e32 v35, 44, v39 + v_lshrrev_b32_e32 v63, 16, v63 + v_cmp_gt_i32_e64 s14, s23, v65 + v_cndmask_b32_e64 v34, 0x80000000, v43, s17 + v_cmp_o_f32_e64 s17, v26, v26 + v_or_b32_e32 v64, s25, v61 + v_add3_u32 v67, v27, v67, 0x7fff + v_add_lshl_u32 v43, v99, v62, 1 + v_or_b32_e32 v44, s25, v35 + v_or_b32_e32 v56, 6, v39 + v_bfe_u32 v70, v28, 16, 1 + v_cndmask_b32_e64 v26, 0x7fff, v63, s17 + s_and_b32 s17, s16, s14 + v_or_b32_e32 v36, 42, v39 + v_lshrrev_b32_e32 v67, 16, v67 + v_cmp_gt_i32_e64 s13, s23, v64 + v_cndmask_b32_e64 v43, 0x80000000, v43, s17 + v_cmp_o_f32_e64 s17, v27, v27 + v_or_b32_e32 v60, s25, v56 + v_add3_u32 v70, v28, v70, 0x7fff + v_cmp_gt_i32_e64 s1, s23, v44 + v_add_lshl_u32 v44, v99, v61, 1 + v_or_b32_e32 v45, s25, v36 + v_or_b32_e32 v55, 8, v39 + v_bfe_u32 v71, v29, 16, 1 + v_cndmask_b32_e64 v27, 0x7fff, v67, s17 + s_and_b32 s17, s16, s13 + v_or_b32_e32 v37, 40, v39 + v_lshrrev_b32_e32 v70, 16, v70 + v_cmp_gt_i32_e64 s12, s23, v60 + v_cndmask_b32_e64 v44, 0x80000000, v44, s17 + v_cmp_o_f32_e64 s17, v28, v28 + v_or_b32_e32 v59, s25, v55 + v_add3_u32 v71, v29, v71, 0x7fff + v_cmp_gt_i32_e64 s2, s23, v45 + v_add_lshl_u32 v45, v99, v56, 1 + v_or_b32_e32 v38, 38, v39 + v_or_b32_e32 v40, 36, v39 + v_or_b32_e32 v41, 34, v39 + v_or_b32_e32 v42, 32, v39 + v_or_b32_e32 v46, s25, v37 + v_or_b32_e32 v49, 14, v39 + v_or_b32_e32 v52, 12, v39 + v_or_b32_e32 v53, 10, v39 + v_bfe_u32 v72, v30, 16, 1 + v_cndmask_b32_e64 v28, 0x7fff, v70, s17 + s_and_b32 s17, s16, s12 + v_lshrrev_b32_e32 v71, 16, v71 + v_cmp_gt_i32_e64 s11, s23, v59 + v_cndmask_b32_e64 v45, 0x80000000, v45, s17 + v_cmp_o_f32_e64 s17, v29, v29 + v_or_b32_e32 v47, s25, v38 + v_or_b32_e32 v48, s25, v40 + v_or_b32_e32 v50, s25, v41 + v_or_b32_e32 v51, s25, v42 + v_or_b32_e32 v54, s25, v49 + v_or_b32_e32 v57, s25, v52 + v_or_b32_e32 v58, s25, v53 + v_or_b32_e32 v69, 64, v66 + v_add3_u32 v72, v30, v72, 0x7fff + v_cmp_gt_i32_e64 s3, s23, v46 + v_add_lshl_u32 v46, v99, v55, 1 + v_bfe_u32 v73, v31, 16, 1 + v_cndmask_b32_e64 v29, 0x7fff, v71, s17 + s_and_b32 s17, s16, s11 + v_lshrrev_b32_e32 v72, 16, v72 + v_cmp_gt_i32_e32 vcc_lo, s22, v69 + v_cmp_gt_i32_e64 s10, s23, v58 + v_cmp_gt_i32_e64 s9, s23, v57 + v_cmp_gt_i32_e64 s8, s23, v54 + v_cmp_gt_i32_e64 s7, s23, v51 + v_cmp_gt_i32_e64 s6, s23, v50 + v_cmp_gt_i32_e64 s5, s23, v48 + v_cmp_gt_i32_e64 s4, s23, v47 + s_and_b32 s21, s21, 0xffff + s_mov_b32 s23, 0x31027000 + s_mov_b32 s22, 0x7ffffffe + v_cndmask_b32_e64 v46, 0x80000000, v46, s17 + v_cmp_o_f32_e64 s17, v30, v30 + v_add3_u32 v73, v31, v73, 0x7fff + s_clause 0x4 + buffer_store_b16 v25, v34, s[20:23], 0 offen + buffer_store_b16 v26, v43, s[20:23], 0 offen + buffer_store_b16 v27, v44, s[20:23], 0 offen + buffer_store_b16 v28, v45, s[20:23], 0 offen + buffer_store_b16 v29, v46, s[20:23], 0 offen + v_add_lshl_u32 v25, v99, v53, 1 + v_bfe_u32 v74, v32, 16, 1 + v_cndmask_b32_e64 v26, 0x7fff, v72, s17 + s_and_b32 s17, s16, s10 + v_lshrrev_b32_e32 v73, 16, v73 + v_cndmask_b32_e64 v25, 0x80000000, v25, s17 + v_cmp_o_f32_e64 s17, v31, v31 + v_add3_u32 v74, v32, v74, 0x7fff + v_add_lshl_u32 v27, v99, v52, 1 + v_bfe_u32 v75, v17, 16, 1 + v_add_lshl_u32 v29, v99, v49, 1 + v_cndmask_b32_e64 v28, 0x7fff, v73, s17 + s_and_b32 s17, s16, s9 + v_lshrrev_b32_e32 v74, 16, v74 + v_cndmask_b32_e64 v27, 0x80000000, v27, s17 + v_cmp_o_f32_e64 s17, v32, v32 + v_add3_u32 v75, v17, v75, 0x7fff + v_bfe_u32 v76, v18, 16, 1 + v_add_lshl_u32 v31, v99, v42, 1 + v_bfe_u32 v77, v19, 16, 1 + v_cndmask_b32_e64 v30, 0x7fff, v74, s17 + s_and_b32 s17, s16, s8 + v_lshrrev_b32_e32 v75, 16, v75 + v_cndmask_b32_e64 v29, 0x80000000, v29, s17 + v_cmp_o_f32_e64 s17, v17, v17 + v_add3_u32 v76, v18, v76, 0x7fff + v_add3_u32 v77, v19, v77, 0x7fff + s_clause 0x2 + buffer_store_b16 v26, v25, s[20:23], 0 offen + buffer_store_b16 v28, v27, s[20:23], 0 offen + buffer_store_b16 v30, v29, s[20:23], 0 offen + v_add_lshl_u32 v25, v99, v41, 1 + v_cndmask_b32_e64 v17, 0x7fff, v75, s17 + s_and_b32 s17, s16, s7 + v_lshrrev_b32_e32 v76, 16, v76 + v_cndmask_b32_e64 v31, 0x80000000, v31, s17 + v_cmp_o_f32_e64 s17, v18, v18 + v_bfe_u32 v78, v20, 16, 1 + v_lshrrev_b32_e32 v77, 16, v77 + v_bfe_u32 v79, v21, 16, 1 + buffer_store_b16 v17, v31, s[20:23], 0 offen + v_cndmask_b32_e64 v17, 0x7fff, v76, s17 + s_and_b32 s17, s16, s6 + v_add3_u32 v78, v20, v78, 0x7fff + v_cndmask_b32_e64 v18, 0x80000000, v25, s17 + v_cmp_o_f32_e64 s17, v19, v19 + v_add_lshl_u32 v25, v99, v40, 1 + v_add3_u32 v79, v21, v79, 0x7fff + v_lshrrev_b32_e32 v78, 16, v78 + v_add_lshl_u32 v26, v99, v38, 1 + v_cndmask_b32_e64 v19, 0x7fff, v77, s17 + s_and_b32 s17, s16, s5 + v_bfe_u32 v80, v22, 16, 1 + v_cndmask_b32_e64 v25, 0x80000000, v25, s17 + v_cmp_o_f32_e64 s17, v20, v20 + v_lshrrev_b32_e32 v79, 16, v79 + v_add_lshl_u32 v27, v99, v37, 1 + v_add3_u32 v80, v22, v80, 0x7fff + v_bfe_u32 v81, v23, 16, 1 + v_cndmask_b32_e64 v20, 0x7fff, v78, s17 + s_and_b32 s17, s16, s4 + v_add_lshl_u32 v28, v99, v36, 1 + v_cndmask_b32_e64 v26, 0x80000000, v26, s17 + v_cmp_o_f32_e64 s17, v21, v21 + v_lshrrev_b32_e32 v80, 16, v80 + v_add3_u32 v81, v23, v81, 0x7fff + v_bfe_u32 v82, v24, 16, 1 + v_bfe_u32 v83, v9, 16, 1 + v_cndmask_b32_e64 v21, 0x7fff, v79, s17 + s_and_b32 s17, s16, s3 + v_lshrrev_b32_e32 v81, 16, v81 + v_cndmask_b32_e64 v27, 0x80000000, v27, s17 + v_cmp_o_f32_e64 s17, v22, v22 + v_add3_u32 v82, v24, v82, 0x7fff + v_add3_u32 v83, v9, v83, 0x7fff + v_bfe_u32 v84, v10, 16, 1 + v_bfe_u32 v85, v11, 16, 1 + v_cndmask_b32_e64 v22, 0x7fff, v80, s17 + s_and_b32 s17, s16, s2 + v_lshrrev_b32_e32 v82, 16, v82 + v_cndmask_b32_e64 v28, 0x80000000, v28, s17 + v_cmp_o_f32_e64 s17, v23, v23 + s_clause 0x4 + buffer_store_b16 v17, v18, s[20:23], 0 offen + buffer_store_b16 v19, v25, s[20:23], 0 offen + buffer_store_b16 v20, v26, s[20:23], 0 offen + buffer_store_b16 v21, v27, s[20:23], 0 offen + buffer_store_b16 v22, v28, s[20:23], 0 offen + v_add_lshl_u32 v17, v99, v35, 1 + v_add_lshl_u32 v19, v99, v0, 1 + v_lshrrev_b32_e32 v83, 16, v83 + v_cndmask_b32_e64 v18, 0x7fff, v81, s17 + s_and_b32 s17, s16, s1 + s_and_b32 s16, s16, s0 + v_cndmask_b32_e64 v17, 0x80000000, v17, s17 + v_cmp_o_f32_e64 s17, v24, v24 + v_cndmask_b32_e64 v19, 0x80000000, v19, s16 + v_add_lshl_u32 v21, v33, v39, 1 + v_cmp_o_f32_e64 s16, v9, v9 + v_add3_u32 v84, v10, v84, 0x7fff + v_bfe_u32 v86, v12, 16, 1 + s_and_b32 s15, vcc_lo, s15 + v_bfe_u32 v87, v13, 16, 1 + v_cndmask_b32_e64 v20, 0x7fff, v82, s17 + v_add_lshl_u32 v22, v33, v62, 1 + v_bfe_u32 v88, v14, 16, 1 + v_cndmask_b32_e64 v9, 0x7fff, v83, s16 + v_cndmask_b32_e64 v21, 0x80000000, v21, s15 + v_bfe_u32 v89, v15, 16, 1 + v_lshrrev_b32_e32 v84, 16, v84 + v_add3_u32 v85, v11, v85, 0x7fff + v_cmp_o_f32_e64 s15, v10, v10 + s_and_b32 s14, vcc_lo, s14 + v_add3_u32 v86, v12, v86, 0x7fff + s_clause 0x2 + buffer_store_b16 v18, v17, s[20:23], 0 offen + buffer_store_b16 v20, v19, s[20:23], 0 offen + buffer_store_b16 v9, v21, s[20:23], 0 offen + v_add_lshl_u32 v9, v33, v61, 1 + v_add3_u32 v87, v13, v87, 0x7fff + v_cndmask_b32_e64 v22, 0x80000000, v22, s14 + v_cmp_o_f32_e64 s14, v11, v11 + v_add_lshl_u32 v11, v33, v56, 1 + v_add3_u32 v88, v14, v88, 0x7fff + v_add_lshl_u32 v17, v33, v55, 1 + v_add3_u32 v89, v15, v89, 0x7fff + s_and_b32 s13, vcc_lo, s13 + v_add_lshl_u32 v18, v33, v53, 1 + v_lshrrev_b32_e32 v85, 16, v85 + v_cndmask_b32_e64 v10, 0x7fff, v84, s15 + s_and_b32 s12, vcc_lo, s12 + v_lshrrev_b32_e32 v86, 16, v86 + v_cndmask_b32_e64 v9, 0x80000000, v9, s13 + v_cmp_o_f32_e64 s13, v12, v12 + s_and_b32 s11, vcc_lo, s11 + v_lshrrev_b32_e32 v87, 16, v87 + v_cndmask_b32_e64 v11, 0x80000000, v11, s12 + v_cmp_o_f32_e64 s12, v13, v13 + s_and_b32 s10, vcc_lo, s10 + v_lshrrev_b32_e32 v88, 16, v88 + v_cndmask_b32_e64 v17, 0x80000000, v17, s11 + v_cmp_o_f32_e64 s11, v14, v14 + v_bfe_u32 v90, v16, 16, 1 + v_lshrrev_b32_e32 v89, 16, v89 + v_add_lshl_u32 v19, v33, v52, 1 + v_cndmask_b32_e64 v18, 0x80000000, v18, s10 + v_cmp_o_f32_e64 s10, v15, v15 + v_bfe_u32 v91, v1, 16, 1 + buffer_store_b16 v10, v22, s[20:23], 0 offen + v_cndmask_b32_e64 v10, 0x7fff, v85, s14 + v_bfe_u32 v92, v2, 16, 1 + v_cndmask_b32_e64 v12, 0x7fff, v86, s13 + v_cndmask_b32_e64 v13, 0x7fff, v87, s12 + s_and_b32 s9, vcc_lo, s9 + v_bfe_u32 v93, v3, 16, 1 + v_cndmask_b32_e64 v14, 0x7fff, v88, s11 + v_add3_u32 v90, v16, v90, 0x7fff + v_cndmask_b32_e64 v15, 0x7fff, v89, s10 + v_cndmask_b32_e64 v19, 0x80000000, v19, s9 + v_add3_u32 v91, v1, v91, 0x7fff + s_clause 0x4 + buffer_store_b16 v10, v9, s[20:23], 0 offen + buffer_store_b16 v12, v11, s[20:23], 0 offen + buffer_store_b16 v13, v17, s[20:23], 0 offen + buffer_store_b16 v14, v18, s[20:23], 0 offen + buffer_store_b16 v15, v19, s[20:23], 0 offen + v_add_lshl_u32 v9, v33, v49, 1 + v_add3_u32 v92, v2, v92, 0x7fff + v_add_lshl_u32 v11, v33, v42, 1 + v_add3_u32 v93, v3, v93, 0x7fff + s_and_b32 s8, vcc_lo, s8 + v_add_lshl_u32 v12, v33, v41, 1 + v_lshrrev_b32_e32 v90, 16, v90 + v_cmp_o_f32_e64 s9, v16, v16 + s_and_b32 s7, vcc_lo, s7 + v_lshrrev_b32_e32 v91, 16, v91 + v_cndmask_b32_e64 v9, 0x80000000, v9, s8 + v_cmp_o_f32_e64 s8, v1, v1 + v_lshrrev_b32_e32 v92, 16, v92 + v_cndmask_b32_e64 v11, 0x80000000, v11, s7 + v_cmp_o_f32_e64 s7, v2, v2 + s_and_b32 s6, vcc_lo, s6 + v_lshrrev_b32_e32 v93, 16, v93 + v_add_lshl_u32 v13, v33, v40, 1 + v_cndmask_b32_e64 v12, 0x80000000, v12, s6 + v_cmp_o_f32_e64 s6, v3, v3 + v_cndmask_b32_e64 v10, 0x7fff, v90, s9 + v_cndmask_b32_e64 v1, 0x7fff, v91, s8 + v_bfe_u32 v94, v4, 16, 1 + v_cndmask_b32_e64 v2, 0x7fff, v92, s7 + s_and_b32 s5, vcc_lo, s5 + v_bfe_u32 v95, v5, 16, 1 + v_bfe_u32 v96, v6, 16, 1 + v_cndmask_b32_e64 v3, 0x7fff, v93, s6 + v_cndmask_b32_e64 v13, 0x80000000, v13, s5 + v_bfe_u32 v97, v7, 16, 1 + s_clause 0x2 + buffer_store_b16 v10, v9, s[20:23], 0 offen + buffer_store_b16 v1, v11, s[20:23], 0 offen + buffer_store_b16 v2, v12, s[20:23], 0 offen + v_add_lshl_u32 v1, v33, v38, 1 + v_bfe_u32 v98, v8, 16, 1 + v_add3_u32 v94, v4, v94, 0x7fff + v_add3_u32 v95, v5, v95, 0x7fff + s_and_b32 s4, vcc_lo, s4 + v_add3_u32 v96, v6, v96, 0x7fff + buffer_store_b16 v3, v13, s[20:23], 0 offen + v_add_lshl_u32 v3, v33, v37, 1 + v_add3_u32 v97, v7, v97, 0x7fff + v_cndmask_b32_e64 v1, 0x80000000, v1, s4 + v_cmp_o_f32_e64 s4, v5, v5 + v_add_lshl_u32 v5, v33, v36, 1 + v_add3_u32 v98, v8, v98, 0x7fff + v_add_lshl_u32 v9, v33, v35, 1 + v_lshrrev_b32_e32 v94, 16, v94 + v_cmp_o_f32_e64 s5, v4, v4 + s_and_b32 s3, vcc_lo, s3 + v_lshrrev_b32_e32 v95, 16, v95 + s_and_b32 s2, vcc_lo, s2 + v_lshrrev_b32_e32 v96, 16, v96 + v_cndmask_b32_e64 v3, 0x80000000, v3, s3 + v_cmp_o_f32_e64 s3, v6, v6 + s_and_b32 s1, vcc_lo, s1 + v_lshrrev_b32_e32 v97, 16, v97 + v_cndmask_b32_e64 v5, 0x80000000, v5, s2 + v_cmp_o_f32_e64 s2, v7, v7 + v_lshrrev_b32_e32 v98, 16, v98 + v_add_lshl_u32 v0, v33, v0, 1 + v_cndmask_b32_e64 v9, 0x80000000, v9, s1 + v_cmp_o_f32_e64 s1, v8, v8 + v_cndmask_b32_e64 v2, 0x7fff, v94, s5 + v_cndmask_b32_e64 v4, 0x7fff, v95, s4 + v_cndmask_b32_e64 v6, 0x7fff, v96, s3 + s_and_b32 vcc_lo, vcc_lo, s0 + v_cndmask_b32_e64 v7, 0x7fff, v97, s2 + v_cndmask_b32_e64 v8, 0x7fff, v98, s1 + v_cndmask_b32_e32 v0, 0x80000000, v0, vcc_lo + s_clause 0x4 + buffer_store_b16 v2, v1, s[20:23], 0 offen + buffer_store_b16 v4, v3, s[20:23], 0 offen + buffer_store_b16 v6, v5, s[20:23], 0 offen + buffer_store_b16 v7, v9, s[20:23], 0 offen + buffer_store_b16 v8, v0, s[20:23], 0 offen + s_nop 0 + s_sendmsg sendmsg(MSG_DEALLOC_VGPRS) + s_endpgm +.Ltmp9: + .section .rodata,"a",@progbits + .p2align 6, 0x0 + .amdhsa_kernel _triton_w4a16_skinny_fmt_kernel + .amdhsa_group_segment_fixed_size 0 + .amdhsa_private_segment_fixed_size 0 + .amdhsa_kernarg_size 88 + .amdhsa_user_sgpr_count 2 + .amdhsa_user_sgpr_dispatch_ptr 0 + .amdhsa_user_sgpr_queue_ptr 0 + .amdhsa_user_sgpr_kernarg_segment_ptr 1 + .amdhsa_user_sgpr_dispatch_id 0 + .amdhsa_user_sgpr_private_segment_size 0 + .amdhsa_wavefront_size32 1 + .amdhsa_uses_dynamic_stack 0 + .amdhsa_enable_private_segment 0 + .amdhsa_system_sgpr_workgroup_id_x 1 + .amdhsa_system_sgpr_workgroup_id_y 1 + .amdhsa_system_sgpr_workgroup_id_z 0 + .amdhsa_system_sgpr_workgroup_info 0 + .amdhsa_system_vgpr_workitem_id 0 + .amdhsa_next_free_vgpr 128 + .amdhsa_next_free_sgpr 41 + .amdhsa_reserve_vcc 1 + .amdhsa_float_round_mode_32 0 + .amdhsa_float_round_mode_16_64 0 + .amdhsa_float_denorm_mode_32 3 + .amdhsa_float_denorm_mode_16_64 3 + .amdhsa_dx10_clamp 1 + .amdhsa_ieee_mode 1 + .amdhsa_fp16_overflow 0 + .amdhsa_workgroup_processor_mode 0 + .amdhsa_memory_ordered 1 + .amdhsa_forward_progress 1 + .amdhsa_shared_vgpr_count 0 + .amdhsa_inst_pref_size 36 + .amdhsa_exception_fp_ieee_invalid_op 0 + .amdhsa_exception_fp_denorm_src 0 + .amdhsa_exception_fp_ieee_div_zero 0 + .amdhsa_exception_fp_ieee_overflow 0 + .amdhsa_exception_fp_ieee_underflow 0 + .amdhsa_exception_fp_ieee_inexact 0 + .amdhsa_exception_int_div_zero 0 + .end_amdhsa_kernel + .text +.Lfunc_end0: + .size _triton_w4a16_skinny_fmt_kernel, .Lfunc_end0-_triton_w4a16_skinny_fmt_kernel + .cfi_endproc + ; -- End function + .set _triton_w4a16_skinny_fmt_kernel.num_vgpr, 115 + .set _triton_w4a16_skinny_fmt_kernel.num_agpr, 0 + .set _triton_w4a16_skinny_fmt_kernel.numbered_sgpr, 35 + .set _triton_w4a16_skinny_fmt_kernel.num_named_barrier, 0 + .set _triton_w4a16_skinny_fmt_kernel.private_seg_size, 0 + .set _triton_w4a16_skinny_fmt_kernel.uses_vcc, 1 + .set _triton_w4a16_skinny_fmt_kernel.uses_flat_scratch, 0 + .set _triton_w4a16_skinny_fmt_kernel.has_dyn_sized_stack, 0 + .set _triton_w4a16_skinny_fmt_kernel.has_recursion, 0 + .set _triton_w4a16_skinny_fmt_kernel.has_indirect_call, 0 + .section .AMDGPU.csdata,"",@progbits +; Kernel info: +; codeLenInByte = 4548 +; TotalNumSgprs: 37 +; NumVgprs: 115 +; ScratchSize: 0 +; MemoryBound: 0 +; FloatMode: 240 +; IeeeMode: 1 +; LDSByteSize: 0 bytes/workgroup (compile time only) +; SGPRBlocks: 0 +; VGPRBlocks: 14 +; NumSGPRsForWavesPerEU: 37 +; NumVGPRsForWavesPerEU: 115 +; Occupancy: 12 +; WaveLimiterHint : 0 +; COMPUTE_PGM_RSRC2:SCRATCH_EN: 0 +; COMPUTE_PGM_RSRC2:USER_SGPR: 2 +; COMPUTE_PGM_RSRC2:TRAP_HANDLER: 0 +; COMPUTE_PGM_RSRC2:TGID_X_EN: 1 +; COMPUTE_PGM_RSRC2:TGID_Y_EN: 1 +; COMPUTE_PGM_RSRC2:TGID_Z_EN: 0 +; COMPUTE_PGM_RSRC2:TIDIG_COMP_CNT: 0 + .text + .p2alignl 7, 3214868480 + .fill 96, 4, 3214868480 + .section .AMDGPU.gpr_maximums,"",@progbits + .set amdgpu.max_num_vgpr, 0 + .set amdgpu.max_num_agpr, 0 + .set amdgpu.max_num_sgpr, 0 + .text + .section .debug_abbrev,"",@progbits + .byte 1 ; Abbreviation Code + .byte 17 ; DW_TAG_compile_unit + .byte 1 ; DW_CHILDREN_yes + .byte 37 ; DW_AT_producer + .byte 14 ; DW_FORM_strp + .byte 19 ; DW_AT_language + .byte 5 ; DW_FORM_data2 + .byte 3 ; DW_AT_name + .byte 14 ; DW_FORM_strp + .byte 16 ; DW_AT_stmt_list + .byte 23 ; DW_FORM_sec_offset + .byte 27 ; DW_AT_comp_dir + .byte 14 ; DW_FORM_strp + .byte 17 ; DW_AT_low_pc + .byte 1 ; DW_FORM_addr + .byte 18 ; DW_AT_high_pc + .byte 6 ; DW_FORM_data4 + .byte 0 ; EOM(1) + .byte 0 ; EOM(2) + .byte 2 ; Abbreviation Code + .byte 46 ; DW_TAG_subprogram + .byte 0 ; DW_CHILDREN_no + .byte 3 ; DW_AT_name + .byte 14 ; DW_FORM_strp + .byte 32 ; DW_AT_inline + .byte 11 ; DW_FORM_data1 + .byte 0 ; EOM(1) + .byte 0 ; EOM(2) + .byte 3 ; Abbreviation Code + .byte 46 ; DW_TAG_subprogram + .byte 1 ; DW_CHILDREN_yes + .byte 17 ; DW_AT_low_pc + .byte 1 ; DW_FORM_addr + .byte 18 ; DW_AT_high_pc + .byte 6 ; DW_FORM_data4 + .byte 49 ; DW_AT_abstract_origin + .byte 19 ; DW_FORM_ref4 + .byte 0 ; EOM(1) + .byte 0 ; EOM(2) + .byte 4 ; Abbreviation Code + .byte 29 ; DW_TAG_inlined_subroutine + .byte 0 ; DW_CHILDREN_no + .byte 49 ; DW_AT_abstract_origin + .byte 19 ; DW_FORM_ref4 + .byte 85 ; DW_AT_ranges + .byte 23 ; DW_FORM_sec_offset + .byte 88 ; DW_AT_call_file + .byte 11 ; DW_FORM_data1 + .byte 89 ; DW_AT_call_line + .byte 11 ; DW_FORM_data1 + .byte 87 ; DW_AT_call_column + .byte 11 ; DW_FORM_data1 + .byte 0 ; EOM(1) + .byte 0 ; EOM(2) + .byte 0 ; EOM(3) + .section .debug_info,"",@progbits +.Lcu_begin0: + .long .Ldebug_info_end0-.Ldebug_info_start0 ; Length of Unit +.Ldebug_info_start0: + .short 4 ; DWARF version number + .long .debug_abbrev ; Offset Into Abbrev. Section + .byte 8 ; Address Size (in bytes) + .byte 1 ; Abbrev [1] 0xb:0x44 DW_TAG_compile_unit + .long .Linfo_string0 ; DW_AT_producer + .short 2 ; DW_AT_language + .long .Linfo_string1 ; DW_AT_name + .long .Lline_table_start0 ; DW_AT_stmt_list + .long .Linfo_string2 ; DW_AT_comp_dir + .quad .Lfunc_begin0 ; DW_AT_low_pc + .long .Lfunc_end0-.Lfunc_begin0 ; DW_AT_high_pc + .byte 2 ; Abbrev [2] 0x2a:0x6 DW_TAG_subprogram + .long .Linfo_string3 ; DW_AT_name + .byte 1 ; DW_AT_inline + .byte 3 ; Abbrev [3] 0x30:0x1e DW_TAG_subprogram + .quad .Lfunc_begin0 ; DW_AT_low_pc + .long .Lfunc_end0-.Lfunc_begin0 ; DW_AT_high_pc + .long 42 ; DW_AT_abstract_origin + .byte 4 ; Abbrev [4] 0x41:0xc DW_TAG_inlined_subroutine + .long 42 ; DW_AT_abstract_origin + .long .Ldebug_ranges0 ; DW_AT_ranges + .byte 1 ; DW_AT_call_file + .byte 141 ; DW_AT_call_line + .byte 39 ; DW_AT_call_column + .byte 0 ; End Of Children Mark + .byte 0 ; End Of Children Mark +.Ldebug_info_end0: + .section .debug_ranges,"",@progbits +.Ldebug_ranges0: + .quad .Ltmp1-.Lfunc_begin0 + .quad .Ltmp2-.Lfunc_begin0 + .quad .Ltmp3-.Lfunc_begin0 + .quad .Ltmp4-.Lfunc_begin0 + .quad .Ltmp5-.Lfunc_begin0 + .quad .Ltmp6-.Lfunc_begin0 + .quad .Ltmp7-.Lfunc_begin0 + .quad .Ltmp8-.Lfunc_begin0 + .quad 0 + .quad 0 + .section .debug_str,"MS",@progbits,1 +.Linfo_string0: + .asciz "triton" ; string offset=0 +.Linfo_string1: + .asciz "hybrid_w4a16.py" ; string offset=7 +.Linfo_string2: + .asciz "/scratch/mgehre/vllm/vllm/model_executor/kernels/linear/mixed_precision" ; string offset=23 +.Linfo_string3: + .asciz "_triton_w4a16_skinny_fmt_kernel" ; string offset=95 + .section ".note.GNU-stack","",@progbits + .amdgpu_metadata +--- +amdhsa.kernels: + - .args: + - .address_space: global + .offset: 0 + .size: 8 + .value_kind: global_buffer + - .address_space: global + .offset: 8 + .size: 8 + .value_kind: global_buffer + - .address_space: global + .offset: 16 + .size: 8 + .value_kind: global_buffer + - .address_space: global + .offset: 24 + .size: 8 + .value_kind: global_buffer + - .address_space: global + .offset: 32 + .size: 8 + .value_kind: global_buffer + - .offset: 40 + .size: 4 + .value_kind: by_value + - .offset: 44 + .size: 4 + .value_kind: by_value + - .offset: 48 + .size: 4 + .value_kind: by_value + - .offset: 52 + .size: 4 + .value_kind: by_value + - .offset: 56 + .size: 4 + .value_kind: by_value + - .offset: 60 + .size: 4 + .value_kind: by_value + - .offset: 64 + .size: 4 + .value_kind: by_value + - .address_space: global + .offset: 72 + .size: 8 + .value_kind: global_buffer + - .address_space: global + .offset: 80 + .size: 8 + .value_kind: global_buffer + .group_segment_fixed_size: 0 + .kernarg_segment_align: 8 + .kernarg_segment_size: 88 + .max_flat_workgroup_size: 256 + .name: _triton_w4a16_skinny_fmt_kernel + .private_segment_fixed_size: 0 + .sgpr_count: 43 + .sgpr_spill_count: 0 + .symbol: _triton_w4a16_skinny_fmt_kernel.kd + .uniform_work_group_size: 1 + .uses_dynamic_stack: false + .vgpr_count: 128 + .vgpr_spill_count: 0 + .wavefront_size: 32 + .workgroup_processor_mode: 0 +amdhsa.target: amdgcn-amd-amdhsa--gfx1151 +amdhsa.version: + - 1 + - 2 +... + + .end_amdgpu_metadata + .section .debug_line,"",@progbits +.Lline_table_start0: diff --git a/vllm/model_executor/kernels/linear/mixed_precision/asm/hybrid_w4a16_skinny_gfx1151.meta.json b/vllm/model_executor/kernels/linear/mixed_precision/asm/hybrid_w4a16_skinny_gfx1151.meta.json new file mode 100644 index 000000000000..ef40574ef0bf --- /dev/null +++ b/vllm/model_executor/kernels/linear/mixed_precision/asm/hybrid_w4a16_skinny_gfx1151.meta.json @@ -0,0 +1,32 @@ +{ + "_comment": "Pinned ABI + constexpr config for the hand-editable AMDGCN of the W4A16 skinny GEMM. Dumped from Triton 3.6.0 _triton_w4a16_skinny_fmt_kernel for the Qwen3-VL-4B-AWQ up/gate_proj-class shapes (M>1024, N<2K, group_size=32) on gfx1151. Edit the .amdgcn next to this file; the loader asserts the launch matches every field here.", + "kernel_name": "_triton_w4a16_skinny_fmt_kernel", + "arch": "gfx1151", + "dtype": "bfloat16", + "has_zp": false, + "group_size": 32, + "block_m": 128, + "block_n": 64, + "block_k": 32, + "num_warps": 8, + "warp_size": 32, + "shared_bytes": 8192, + "kernarg_size": 88, + "kernarg_align": 8, + "kernarg_layout": [ + {"name": "a_ptr", "offset": 0, "kind": "ptr"}, + {"name": "b_ptr", "offset": 8, "kind": "ptr"}, + {"name": "scales_ptr", "offset": 16, "kind": "ptr"}, + {"name": "packed_scale_zp_ptr","offset": 24, "kind": "ptr"}, + {"name": "c_ptr", "offset": 32, "kind": "ptr"}, + {"name": "M", "offset": 40, "kind": "i32"}, + {"name": "N", "offset": 44, "kind": "i32"}, + {"name": "K", "offset": 48, "kind": "i32"}, + {"name": "K8", "offset": 52, "kind": "i32"}, + {"name": "stride_bn", "offset": 56, "kind": "i32"}, + {"name": "num_groups", "offset": 60, "kind": "i32"}, + {"name": "group_size", "offset": 64, "kind": "i32"}, + {"name": "_hidden_scratch0", "offset": 72, "kind": "ptr_null"}, + {"name": "_hidden_scratch1", "offset": 80, "kind": "ptr_null"} + ] +} diff --git a/vllm/model_executor/kernels/linear/mixed_precision/asm_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/asm_w4a16.py new file mode 100644 index 000000000000..99a2a6b0b14a --- /dev/null +++ b/vllm/model_executor/kernels/linear/mixed_precision/asm_w4a16.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Direct-HIP loader/launcher for the hand-editable AMDGCN W4A16 skinny GEMM. + +This bypasses the Triton JIT entirely: it loads a prebuilt ``.hsaco`` (assembled +from a hand-editable ``.amdgcn`` under ``asm/``) with ``libamdhip64`` and launches +it via ``hipModuleLaunchKernel`` using the kernarg ABI pinned in +``asm/hybrid_w4a16_skinny_gfx1151.meta.json``. + +Workflow: edit ``asm/hybrid_w4a16_skinny_gfx1151.amdgcn`` -> rebuild vLLM (CMake +assembles it) OR just rerun (this module re-assembles on demand when the +``.amdgcn`` is newer than the cached ``.hsaco``). The dispatcher in +``hybrid_w4a16.py`` routes the matching launch here and ``assert``s that every +constexpr/config field in the meta matches the request. + +Enabled only when ``VLLM_W4A16_HANDASM=1``. +""" + +from __future__ import annotations + +import ctypes +import functools +import glob +import json +import math +import os +import struct +import subprocess +from pathlib import Path + +import torch + +_ASM_DIR = Path(__file__).parent / "asm" +_STEM = "hybrid_w4a16_skinny_gfx1151" +_AMDGCN = _ASM_DIR / f"{_STEM}.amdgcn" +_HSACO = _ASM_DIR / f"{_STEM}.hsaco" +_META = _ASM_DIR / f"{_STEM}.meta.json" + +# hipModuleLaunchKernel "extra" sentinels. +_HIP_LAUNCH_PARAM_BUFFER_POINTER = ctypes.c_void_p(0x01) +_HIP_LAUNCH_PARAM_BUFFER_SIZE = ctypes.c_void_p(0x02) +_HIP_LAUNCH_PARAM_END = ctypes.c_void_p(0x03) + + +@functools.lru_cache(maxsize=1) +def load_meta() -> dict: + return json.loads(_META.read_text()) + + +def _find_clang() -> str: + """Locate an amdgcn-capable clang for on-demand (re)assembly.""" + cand = os.environ.get("VLLM_ROCM_CLANG") + if cand and Path(cand).exists(): + return cand + # ROCm SDK wheel shipped in the active venv (next to torch in + # site-packages), then system ROCm. + site = os.path.dirname(os.path.dirname(torch.__file__)) # site-packages + patterns = [ + os.path.join(site, "_rocm_sdk_devel", "lib", "llvm", "bin", "clang"), + os.path.join(site, "_rocm_sdk_core", "lib", "llvm", "bin", "clang"), + "/opt/rocm*/llvm/bin/clang", + ] + for pat in patterns: + for hit in glob.glob(pat): + if Path(hit).exists(): + return str(Path(hit).resolve()) + raise FileNotFoundError( + "No amdgcn clang found to assemble the W4A16 .amdgcn. Set " + "VLLM_ROCM_CLANG to a ROCm clang, or build vLLM (CMake assembles it)." + ) + + +def assemble(force: bool = False) -> Path: + """Assemble the .amdgcn -> .hsaco if missing/stale. Returns the .hsaco path. + + This is the dev-loop fast path; the CMake build performs the same step at + build time so installed wheels ship a ready .hsaco. + """ + if ( + not force + and _HSACO.exists() + and _HSACO.stat().st_mtime >= _AMDGCN.stat().st_mtime + ): + return _HSACO + clang = _find_clang() + arch = load_meta()["arch"] + cmd = [ + clang, + "-target", + "amdgcn-amd-amdhsa", + f"-mcpu={arch}", + "-x", + "assembler", + str(_AMDGCN), + "-o", + str(_HSACO), + ] + subprocess.run(cmd, check=True, capture_output=True, text=True) + return _HSACO + + +def _open_hip() -> ctypes.CDLL: + """Open libamdhip64 (already in-process via torch's ROCm SDK) by full path. + + The bare soname is not on the loader path in the SDK-wheel layout, so resolve + the versioned file directly; dlopen-ing the same path just shares torch's + handle and HIP context. + """ + site = os.path.dirname(os.path.dirname(torch.__file__)) + for pat in ( + os.path.join(site, "_rocm_sdk_core", "lib", "libamdhip64.so*"), + os.path.join(os.path.dirname(torch.__file__), "lib", "libamdhip64.so*"), + "libamdhip64.so", + ): + for hit in sorted(glob.glob(pat)) or ([pat] if "*" not in pat else []): + try: + return ctypes.CDLL(hit) + except OSError: + continue + raise OSError("could not locate libamdhip64 for the hand-asm W4A16 loader") + + +class _HandAsmKernel: + """Loads the hsaco once and launches it via hipModuleLaunchKernel.""" + + def __init__(self) -> None: + self.meta = load_meta() + self.hip = _open_hip() + self.hip.hipModuleLoadData.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_void_p, + ] + self.hip.hipModuleGetFunction.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_void_p, + ctypes.c_char_p, + ] + self.hip.hipModuleLaunchKernel.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_uint, # grid + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_uint, # block + ctypes.c_uint, # shared + ctypes.c_void_p, # stream + ctypes.POINTER(ctypes.c_void_p), # kernelParams + ctypes.POINTER(ctypes.c_void_p), # extra + ] + image = bytearray(assemble().read_bytes()) + buf = (ctypes.c_char * len(image)).from_buffer(image) + self._image_keepalive = image # hipModuleLoadData copies, but be safe + module = ctypes.c_void_p() + self._check(self.hip.hipModuleLoadData(ctypes.byref(module), buf)) + func = ctypes.c_void_p() + self._check( + self.hip.hipModuleGetFunction( + ctypes.byref(func), module, self.meta["kernel_name"].encode() + ) + ) + self.module = module + self.func = func + + def _check(self, err: int) -> None: + if err != 0: + raise RuntimeError(f"HIP error {err} in hand-asm W4A16 loader") + + def launch( + self, a, b_q, scales, c, M, N, K, K8, stride_bn, num_groups, group_size, stream + ) -> None: + m = self.meta + ksize = m["kernarg_size"] + kbuf = bytearray(ksize) + # Pack per the pinned kernarg_layout. ptr_null args stay zero. + vals = { + "a_ptr": a.data_ptr(), + "b_ptr": b_q.data_ptr(), + "scales_ptr": scales.data_ptr(), + # symmetric: carrier unused, mirror production's dummy (scales). + "packed_scale_zp_ptr": scales.data_ptr(), + "c_ptr": c.data_ptr(), + "M": M, + "N": N, + "K": K, + "K8": K8, + "stride_bn": stride_bn, + "num_groups": num_groups, + "group_size": group_size, + } + for arg in m["kernarg_layout"]: + off, kind = arg["offset"], arg["kind"] + if kind == "ptr": + struct.pack_into(" _HandAsmKernel: + return _HandAsmKernel() + + +def config_matches( + dtype, has_zp, group_size, block_m, block_n, block_k, num_warps +) -> bool: + """True iff this launch is exactly what the pinned .amdgcn was built for.""" + m = load_meta() + return ( + str(dtype) == f"torch.{m['dtype']}" + and bool(has_zp) == m["has_zp"] + and group_size == m["group_size"] + and block_m == m["block_m"] + and block_n == m["block_n"] + and block_k == m["block_k"] + and num_warps == m["num_warps"] + ) + + +def launch_skinny_w4a16( + a, b_q, scales, c, M, N, K, K8, stride_bn, num_groups, group_size +) -> None: + """Assert-checked entry point used by the hybrid_w4a16 dispatcher.""" + m = load_meta() + assert a.dtype == getattr(torch, m["dtype"]), ( + f"hand-asm W4A16 built for {m['dtype']}, got {a.dtype}" + ) + assert group_size == m["group_size"], ( + f"hand-asm W4A16 built for group_size={m['group_size']}, got {group_size}" + ) + assert num_groups == K // group_size + assert K8 == K // 8 + stream = torch.cuda.current_stream().cuda_stream + _kernel().launch( + a, b_q, scales, c, M, N, K, K8, stride_bn, num_groups, group_size, stream + ) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py index 108e2111c5b8..e6f74cdf1b79 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/hybrid_w4a16.py @@ -12,6 +12,7 @@ weight copy. The triton kernel transposes tiles in-register. """ +import os from contextlib import nullcontext import torch @@ -43,6 +44,14 @@ LDS_CAPACITY_ELEMENTS = 64 * 1024 // 2 # 32768 fp16 elements +def _handasm_w4a16_enabled() -> bool: + """Hand-tuned AMDGCN override for the skinny W4A16 GEMM (gfx1151). + + On by default; set VLLM_W4A16_HANDASM=0 to fall back to the Triton kernel. + """ + return os.environ.get("VLLM_W4A16_HANDASM", "1") == "1" and on_gfx1151() + + # --------------------------------------------------------------------------- # Triton kernel for the prefill path (reads skinny-format weights [N, K//8]) # --------------------------------------------------------------------------- @@ -447,6 +456,37 @@ def triton_w4a16_skinny_fmt_gemm( M, N, K, group_size, a.dtype ) grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n)) + # Opt-in hand-tuned AMDGCN override (VLLM_W4A16_HANDASM=1): dispatch to a + # prebuilt .hsaco assembled from a hand-editable .amdgcn when the selected + # tile/config exactly matches what that ISA was built for. The loader + # asserts every pinned constexpr/ABI field; otherwise we fall through to + # the Triton kernel below. See asm_w4a16.py / asm/. + if _handasm_w4a16_enabled(): + from .asm_w4a16 import config_matches, launch_skinny_w4a16 + + if config_matches( + a.dtype, + has_zp, + group_size, + block_m, + block_n, + block_k, + num_warps, + ): + launch_skinny_w4a16( + a, + b_q, + scales, + c, + M, + N, + K, + K8, + stride_bn, + num_groups, + group_size, + ) + return c # The kernel picks the packed (fp16/gfx1x) vs scalar unpack itself. _triton_w4a16_skinny_fmt_kernel[grid]( a,