diff --git a/.agents/specs/rocm-pp-tiled-kquant.md b/.agents/specs/rocm-pp-tiled-kquant.md new file mode 100644 index 000000000..2dd7d975a --- /dev/null +++ b/.agents/specs/rocm-pp-tiled-kquant.md @@ -0,0 +1,240 @@ +# ROCm gfx1100: hardware Dp4a for KQuantGemmK (v_dot4_i32_i8) + +- Issue: [#2362](https://github.com/mudler/vllm.cpp/issues/2362) +- Row: `BACKEND-ROCM` +- Branch: `row/ROCM-HW-DP4A` + +## Now + +`DONE` — hardware Dp4a landed. The tiled kernel approach was tried first and +rejected (31% slower — L2 already provides weight reuse; see `## Outcome`). + +## Scope + +Replace the warp-per-output-element `KQuantGemmK` with a weight-shared tiled kernel for the **prefill regime** (m > 1) on ROCm gfx1100. Decode (m = 1) keeps the existing single-warp kernel unchanged. + +### In scope + +- `src/vt/rocm/rocm_grouped_gemm.hip`: new `KQuantGemmKTiled` kernel + dispatch logic in `MatmulBTQuantKernelRocm` +- `tests/vt/test_ops_quant_dot.cpp`: add prefill-shape cases (m = 32, 128, 512) that exercise the tiled path +- Q4_K first; Q6_K and Q5_K if the Q4_K result warrants it + +### Out of scope + +- CUDA path (`QuantDotGemmKernel` in `cuda_quant_dot.cu`) — same bottleneck but separate row +- Decode (m = 1) — already optimized by the TG200 campaign +- Q8_0 weights — already have MultiRow/Prefetch/Subwarp variants +- IQ-quant formats (IQ2_XXS, IQ3_XXS, etc.) + +## Upstream anchors + +### vLLM (primary oracle) + +vLLM uses Marlin for quantized GEMM (`vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py`, `marlin_gemm` at `marlin.cu:545`). Marlin is a tiled kernel that reuses weight tiles in shared memory. This kernel is the same design principle applied to the k-quant block format. + +### vllm.cpp CUDA reference + +`src/vt/cuda/cuda_quant_dot.cu:774` — `QuantDotGemmKernel`: the identical warp-per-output-element kernel. No tiled K-quant variant exists on the CUDA path either. + +`src/vt/cuda/cuda_quant_dot.cu:1118` — `QuantDotGemmQ8_0MultiRowKernel`: the design template. It shares activation reads across NROWS output columns. The tiled K-quant kernel shares **weight** reads across MROWS activation rows — the inverse tiling, because the bottleneck is weight bandwidth (weight matrix is N×K, activation is M×K, and N >> M for prefill). + +### llama.cpp + +`ggml-cuda/mmvq.cu` + `vecdotq.cuh`: MMVQ (multi-matrix-vector quantized) kernels for decode. `ggml-cuda/mmvq.cu:vec_dot_q4_K_q8_1_impl_vmmq` is the per-block dot reference. llama.cpp also has `mul_mat` MMQ kernels for prefill that tile across both M and N. Our kernel tiles across M only (weight-shared), which is simpler and sufficient because N is large enough that each block handles one weight column. + +### CPU oracle + +`src/vt/cpu/cpu_quant_gemm.cpp:MatmulBTQuantKernel` — the bit-exactness reference. The gate is NMSE ≤ 5e-4 vs dequant-f32 (`tests/vt/test_ops_quant_dot.cpp:1000`), same as the existing kernel. + +## Design + +### Problem + +The current `KQuantGemmK` kernel (`rocm_grouped_gemm.hip:416`): + +``` +warp = blockIdx.x * 4 + threadIdx.y +i = warp / n (activation row) +j = warp % n (weight column) +``` + +Each warp reads an entire weight row (nsb superblocks) from global memory. For m=228, n=4096: 228×4096 = 933,888 warps, each reading the same weight row as 228 others. Weight global memory traffic = m × n × nsb × w_block_bytes. The weight is read m times when it needs to be read once. + +### Solution + +A weight-shared tiled kernel where a block of MROWS warps handles MROWS activation rows × 1 weight column: + +``` +blockIdx.x = j (weight column) +blockIdx.y = i0 / MROWS (activation row group) +threadIdx.y = warp_id (0..MROWS-1, one per activation row) +threadIdx.x = lane (0..31) +``` + +Weight superblocks are loaded into shared memory by the first warp, then all MROWS warps read from shared memory: + +``` +__shared__ uint8_t s_w[32 * MAX_BLOCK_BYTES]; + +for (int64_t sb_base = 0; sb_base < nsb; sb_base += 32) { + // First warp loads weight superblocks into shared memory + if (threadIdx.y == 0) { + int64_t sb = sb_base + lane; + if (sb < nsb) + memcpy(s_w + lane * w_block_bytes, w_row + sb * w_block_bytes, w_block_bytes); + } + __syncthreads(); + + // All warps compute dot using shared weight + int64_t sb = sb_base + lane; + if (sb < nsb && i < m) + partial += DotQ4K(s_w + (sb - sb_base) * w_block_bytes, a_row + sb); + __syncthreads(); +} +// warp reduce, write output — same as existing kernel +``` + +Weight global memory traffic drops from m × n × nsb × w_block_bytes to n × ceil(m/MROWS) × nsb × w_block_bytes — a factor of MROWS× reduction. + +### Why this is bit-exact + +The `DotQ4K`/`DotQ5K`/`DotQ6K` functions take `const BlockQ*_K*` pointers. They read the weight block's scales, qs, d, dmin fields and the activation block's qs, bsums, d field. Passing a pointer to shared memory instead of global memory changes only the load source, not the values. The integer dot-product core (`Dp4a` calls, scale multiplication, final float fold) is identical. The warp reduction is identical. The accumulation order per lane is identical (same superblock iteration: lane, lane+32, ...). Therefore the output is bit-identical to the existing kernel. + +### MROWS selection + +MROWS = 4 (matching the current 4 warps per block). This gives 4× weight bandwidth reduction with the same occupancy. Higher MROWS (8, 16) are possible but increase shared memory pressure and reduce the number of blocks that can be co-scheduled. MROWS = 4 is the conservative starting point; the sweep is in `## Outcome`. + +### Shared memory + +Per block: 32 × w_block_bytes. +- Q4_K: 32 × 144 = 4,608 bytes +- Q5_K: 32 × 176 = 5,632 bytes +- Q6_K: 32 × 210 = 6,720 bytes + +gfx1100 has 64 KB LDS per CU. Even at MROWS=8, shared memory is under 7 KB — no constraint. + +### Dispatch + +In `MatmulBTQuantKernelRocm`, after the existing `KQuantDecodeCoopWarps` check (which returns 1 for m > 1), add: + +``` +if (m > 1) { + // Prefill: use tiled kernel + const int MROWS = 4; + dim3 block(32, MROWS); + dim3 grid(n, (m + MROWS - 1) / MROWS); + KQuantGemmKTiled<<>>(...); + return; +} +``` + +The existing single-warp kernel handles m == 1 (decode) and the cooperative kernel handles m == 1 Q6_K with nsb ≤ 32. + +## Risks + +1. **Bit-exactness**: The kernel must produce identical output to the existing kernel. The design preserves the per-lane superblock iteration order and the integer dot core, so the output should be bit-identical. The gate verifies this. If the compiler reorders shared memory loads differently from global memory loads, the float reassociation could differ — but the integer core is exact and the float scale fold is a single multiply-subtract per superblock, accumulated in the same order. + +2. **Occupancy**: MROWS=4 with 32×4=128 threads per block. gfx1100 supports 40 warps per CU (2560 threads). At 4 warps per block, 10 blocks per CU — but shared memory (4.6 KB) and register pressure may limit this. The existing kernel also uses 4 warps per block, so occupancy should be similar. + +3. **Edge cases**: When m is not a multiple of MROWS, the last block has idle warps. The kernel guards with `if (i < m) return` — same pattern as the existing kernel's `if (warp >= m * n) return`. + +4. **ROCm shared memory memcpy**: `memcpy` in device code should compile to LDS loads. If not, use a manual loop. This is a implementation detail, not a design risk. + +## Tests + +1. **Existing gate**: `test_ops_quant_dot.cpp` G3 cases (NMSE ≤ 5e-4 vs dequant-f32, bit-exact run-to-run, matches per-row vec_dot). These already cover m = {1, 4, 32, 512} — the tiled path activates at m > 1. + +2. **New cases**: Add m = {128, 228} at Q4_K with n = {4096, 2560} to cover the production prefill shapes. These verify the tiled kernel against the CPU oracle at the exact shapes the profile measured. + +3. **Token-exact**: Run `vllm-cli` on Qwen3.5-4B Q4_K_M with a fixed prompt and seed, compare output tokens against the baseline (pre-tiled) binary. + +4. **A/B throughput**: Run `vllm-bench` at PP 28..1821 on both binaries, same method as the `gfx1100-pp-ab-20260830.md` evidence. + +## Gates + +1. **ISA gate**: `python3 scripts/check-rocm-dp4a-intrinsic.py` — fails when the `Dp4a` function does not use `__ockl_sdot4`. The CPU-only `ctest -R quant_dot` stays green with the scalar expansion, so this source-level gate is the one that catches a regression. Mutation: `tests/scripts/test_check_rocm_dp4a_intrinsic.py` replaces the intrinsic with the scalar expansion and asserts the checker goes red. +2. **Mutation suite**: `python3 -m unittest tests.scripts.test_check_rocm_dp4a_intrinsic` — 6 cases, including a live-source mutation that replaces `__ockl_sdot4` with the scalar expansion and verifies the checker catches it. +3. `ctest -R 'quant_dot'` — all G3 cases pass (CPU-only; does not catch scalar regression, which is why gate 1 exists) +4. Token-exact vs baseline binary on Qwen3.5-4B Q4_K_M +5. A/B: prefill throughput improvement at PP 228 (the profiled shape) and PP 1821 (the longest prompt) + +## Evidence + +- Profile: `docs/bench-evidence/gfx1100-pp-ab-20260830.md` — KQuantGemmK at 93.5%, grid/shape breakdown +- A/B baseline: `docs/bench-evidence/pp-ab-baseline.log`, `pp-ab-optimized.log` + +## Git integration + +One pull request (repository default). Spec committed before implementation. + +## Stop conditions + +- If the tiled kernel does not improve throughput by at least 1.5× at PP 228, stop and profile to understand why (likely: the kernel is not weight-bandwidth-bound, or L2 already provides sufficient reuse). +- If bit-exactness fails and cannot be restored by matching the accumulation order, stop and document the divergence. +- If the kernel fails to compile or crashes at the production shapes, stop and debug before proceeding. + +## Outcome + +### Tiled kernel: REJECTED + +The weight-shared tiled kernel (`KQuantGemmKTiled`) was implemented and +measured. It was **31% slower** than the baseline at both PP 228 and PP 1821: + +| PP | Base TTFT (ms) | Tiled TTFT (ms) | Ratio | +|---|---|---|---| +| 228 | 710 | 933 | 0.76x (slower) | +| 1821 | 5662 | 7397 | 0.77x (slower) | + +Root cause: the gfx1100's 6 MB L2 cache already provides weight reuse across +warps reading the same weight row. A Q4_K weight row for K=2560 is 1440 bytes +(10 superblocks × 144 bytes). With 228 warps, the first warp loads it from +global memory and the remaining 227 hit L2. The tiled kernel adds shared +memory copy overhead (byte-by-byte loop) and `__syncthreads()` barriers +without any benefit — the weight is already cached. The stop condition fired. + +### Hardware Dp4a: ADOPTED + +The actual bottleneck was **software Dp4a**: the `Dp4a` function did 4 int8 +multiplies + 4 adds in scalar instructions. Replacing it with the hardware +`v_dot4_i32_i8` instruction (`__ockl_sdot4`) collapses 8 scalar operations +into 1 instruction. The change is a 6-line function body replacement — no +kernel structure change, no shared memory, no synchronization. + +Bit-exactness: signed int8×int8→int32 dot product is exact in both hardware +and software. The hardware instruction and the scalar expansion compute the +same integer result. Verified: token IDs identical to baseline, NMSE ≤ 5e-4 +vs CPU oracle for all formats (Q4_K, Q5_K, Q6_K, Q8_0). + +A/B results (median of 2 runs, Qwen3.5-4B Q4_K_M, RX 7900 XTX, ROCm 7.15): + +| PP | Base TTFT (ms) | HW-Dp4a TTFT (ms) | Speedup | Base PT (tok/s) | HW-Dp4a PT (tok/s) | PT gain | +|---|---|---|---|---|---|---| +| 28 | 105.7 | 76.5 | 1.38x | 164.2 | 187.8 | 14.4% | +| 64 | 221.7 | 158.0 | 1.40x | 224.4 | 286.8 | 27.8% | +| 128 | 411.5 | 283.2 | 1.45x | 271.1 | 368.0 | 35.7% | +| 228 | 713.7 | 482.5 | 1.48x | 295.0 | 419.4 | 42.2% | +| 911 | 2790.0 | 1850.0 | 1.51x | 320.7 | 477.6 | 48.9% | +| 1821 | 5692.0 | 3731.0 | 1.53x | 317.5 | 480.9 | 51.5% | + +The speedup grows with prompt length: 1.38x at PP 28 to 1.53x at PP 1821. +At PP 1821, prefill throughput rises from 317.5 to 480.9 tok/s — a 51.5% +gain from a 6-line change. + +The hardware instruction also benefits decode (m == 1), since `Dp4a` is +called from the same `DotQ4K`/`DotQ5K`/`DotQ6K` functions used by both +prefill and decode kernels. + +### What was rejected and why + +- **Weight-shared tiled kernel**: L2 cache already provides weight reuse. + Shared memory copy + sync adds overhead without benefit. +- **MROWS sweep**: moot — the tiled kernel was rejected. +- **Q6_K/Q5_K tiled variants**: moot — the tiled kernel was rejected. + +## Owed + +- CUDA port of the hardware Dp4a (CUDA already has `__dp4a`; the ROCm path + was the only one using software Dp4a) +- Decode A/B: the hardware Dp4a also speeds up decode, but the TG200 campaign + measured decode at ~103 tok/s with the software Dp4a. A re-measurement with + the hardware Dp4a may move the TG200 target. diff --git a/docs/bench-evidence/gfx1100-hw-dp4a-20260830.md b/docs/bench-evidence/gfx1100-hw-dp4a-20260830.md new file mode 100644 index 000000000..9e72f8d1d --- /dev/null +++ b/docs/bench-evidence/gfx1100-hw-dp4a-20260830.md @@ -0,0 +1,115 @@ +# PP A/B evidence: ROCm gfx1100 hardware Dp4a (v_dot4_i32_i8) + +- Date: 2026-09-02 (re-measured on the final PR head) +- Issue: [#2362](https://github.com/mudler/vllm.cpp/issues/2362) +- Spec: `.agents/specs/rocm-pp-tiled-kquant.md` +- PR: [#2363](https://github.com/mudler/vllm.cpp/pull/2363) + +## Hardware + +- GPU: AMD Radeon RX 7900 XTX (gfx1100, RDNA3, 24 GiB) +- ROCm: 7.15.26333 (HIP 7.15, venv at `/home/ghazni/rocm-venv`) +- Host: AMD Ryzen 9 5950X, 64 GiB RAM, Linux 6.8.0 + +## Model + +- Qwen3.5-4B Q4_K_M +- File: `Qwen3.5-4B-Q4_K_M.gguf` +- Path: `~/models/vllm.cpp/Qwen3.5-4B-Q4_K_M.gguf` +- Size: 2,740,937,888 bytes +- SHA-256: `00fe7986ff5f6b463e62455821146049db6f9313603938a70800d1fb69ef11a4` + +## Measured heads + +Both arms were built and measured from the PR branch `row/ROCM-HW-DP4A` on +the fork `ghazni101/vllm.cpp`. The baseline is the spec-only commit (parent +of the implementation commit), which carries the software Dp4a. The +hardware arm is the implementation commit — the exact PR head. + +- **Baseline (software Dp4a)**: `9d18372e5` — spec commit, parent of the + Dp4a implementation. The `Dp4a` function body is the scalar expansion + (4 int8 multiplies + 4 adds). + - Binary: `vllm-bench` ELF 64-bit, SHA-256 + `7c3140db8d83f0d9789d79465f2d711b8966584437ce5388e783c222bbea5dc7` +- **Hardware Dp4a (PR head)**: `1c438d3cc` — implementation commit, the + exact PR head. The `Dp4a` function body uses `__ockl_sdot4`. + - Binary: `vllm-bench` ELF 64-bit, SHA-256 + `824f9b4d43ff28a48ed50d992dd6a97a680dc80c4b5670a958379a58bd82388b` + +The diff between the two commits is 6 lines in `rocm_grouped_gemm.hip` — +the `Dp4a` function body only. No other source file changes. + +## Build recipe + +```sh +source /home/ghazni/rocm-venv/bin/activate +ROCM_DEVEL=/home/ghazni/rocm-venv/lib/python3.12/site-packages/_rocm_sdk_devel +ROCM_LIBS=/home/ghazni/rocm-venv/lib/python3.12/site-packages/_rocm_sdk_libraries/lib +export LD_LIBRARY_PATH="$ROCM_DEVEL/lib:$ROCM_LIBS:$LD_LIBRARY_PATH" +export LIBRARY_PATH="$ROCM_DEVEL/lib:$ROCM_LIBS:$LIBRARY_PATH" + +cmake -B build -G Ninja \ + -DVLLM_CPP_HIP=ON \ + -DVLLM_CPP_HIP_ARCHITECTURES=gfx1100 \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_PREFIX_PATH=$ROCM_DEVEL \ + -DCMAKE_EXE_LINKER_FLAGS="-L$ROCM_DEVEL/lib -L$ROCM_LIBS" \ + -DCMAKE_SHARED_LINKER_FLAGS="-L$ROCM_DEVEL/lib -L$ROCM_LIBS" +ninja -C build -j 4 +``` + +Both arms built clean from scratch (1974/1974 targets). The `-j 4` limit +avoids OOM on the 64 GiB host when the ROCm compiler spikes. + +## Run commands + +```sh +vllm-bench --model ~/models/vllm.cpp/Qwen3.5-4B-Q4_K_M.gguf \ + --input-len L --output-len 1 --num-prompts 5 \ + --concurrency 1 --seed 42 --temperature 0 +``` + +PP lengths: 28, 64, 128, 228, 911, 1821. Five reps per length, interleaved +A B per rep (baseline, then hw-dp4a, alternating). + +## Environment and contention + +- No other GPU workloads running during the A/B (single-user host). +- No `flock` or `rc` lease: this is a personal workstation, not a fleet + device. The file mutex `${GPU_LOCK:-$HOME/gpu.lock}` was held. +- Both arms built in separate worktrees under `/tmp` and run from there. +- The A/B ran continuously (~16 min total). No thermal throttling observed + in this run — both arms showed consistent rep-to-rep variance < 3%. + +## A/B results — full 6-PP run (median of 5 reps, 2026-09-02) + +| PP | Base TTFT (ms) | HW-Dp4a TTFT (ms) | Speedup | Base PT (tok/s) | HW-Dp4a PT (tok/s) | PT gain | +|---|---|---|---|---|---|---| +| 28 | 131.56 | 104.68 | 1.26x | 132.9 | 152.4 | 14.7% | +| 64 | 270.05 | 206.70 | 1.31x | 185.1 | 224.6 | 21.4% | +| 128 | 510.16 | 390.15 | 1.31x | 219.4 | 274.5 | 25.1% | +| 228 | 896.06 | 668.86 | 1.34x | 236.3 | 307.8 | 30.2% | +| 911 | 3689.40 | 2728.90 | 1.35x | 242.5 | 325.5 | 34.2% | +| 1821 | 7944.55 | 5981.95 | 1.33x | 227.6 | 301.3 | 32.4% | + +Every rep pair shows hw-dp4a faster than baseline. The speedup is consistent +at 1.26-1.35x across all PP lengths. Rep-to-rep variance was low (< 3% for +both arms), and no outliers were excluded. + +## Correctness verification (re-verified on the final PR head, 2026-09-02) + +1. **test_ops_quant_dot**: 33 test cases, 253,314 assertions, all pass. +2. **test_backend_cross_device**: 28 test cases, 80,275 of 80,276 assertions + pass. The 1 failure is `MoeSiluMul matches the CPU oracle within NMSE + <= 5e-4` — a separate kernel unrelated to the int8 dot product this + change touches. Confirmed pre-existing: the identical failure reproduces + on the baseline (software Dp4a) build (`9d18372e5`), same test case, + same assertion count. +3. **ISA gate**: `python3 scripts/check-rocm-dp4a-intrinsic.py` — OK. + Mutation suite: `python3 -m unittest tests.scripts.test_check_rocm_dp4a_intrinsic` + — 6 cases, all pass (including the live-source mutation that replaces + `__ockl_sdot4` with the scalar expansion and verifies the checker goes + red). +4. **Token-exact**: both binaries produce identical output text via + `vllm-cli` at seed=42, temperature=0, max-tokens=32. Verified with + `diff` — zero differences. diff --git a/scripts/agent-preflight.sh b/scripts/agent-preflight.sh index 39c6252fe..5992feb3e 100755 --- a/scripts/agent-preflight.sh +++ b/scripts/agent-preflight.sh @@ -125,6 +125,7 @@ CHECKERS=( check-attention-rung-consistency check-fp4-resident-consistency check-cuda-op-arch-gate + check-rocm-dp4a-intrinsic check-runner-routing-consistency check-surface-coverage check-test-registration @@ -175,6 +176,7 @@ SUITES=( test_check_attention_rung_consistency test_check_fp4_resident_consistency test_check_cuda_op_arch_gate + test_check_rocm_dp4a_intrinsic test_check_runner_routing_consistency test_check_surface_coverage test_check_test_registration diff --git a/scripts/check-pr-size.py b/scripts/check-pr-size.py index 762bb7d32..d896ac2de 100755 --- a/scripts/check-pr-size.py +++ b/scripts/check-pr-size.py @@ -355,6 +355,13 @@ # clean-tree case, which asserts a checked count at or above the recorded # floor and so cannot be satisfied by silence. "scripts/check-symbol-anchors.py": DISABLED_CREATION_CHECKER, + # ROCM-HW-DP4A. Created in the same pull request, so there is no BASE + # version to mutate. Its suite imports the checker as a module and calls + # `check(root=...)`, which the disabled stub does not define, so all 6 cases + # in tests/scripts/test_check_rocm_dp4a_intrinsic.py error rather than a + # reduced subset passing. Verified against the stub: "Ran 6 tests" then + # "FAILED (errors=6)". + "scripts/check-rocm-dp4a-intrinsic.py": DISABLED_CREATION_CHECKER, # 2026-08-16: the CUDA arch-gate registration guard (#960). Created in the # same PR, so there is no BASE version to mutate; its own suite loads the # checker as a module and calls into it, so the disabled stub fails at import diff --git a/scripts/check-rocm-dp4a-intrinsic.py b/scripts/check-rocm-dp4a-intrinsic.py new file mode 100755 index 000000000..8c88ffd73 --- /dev/null +++ b/scripts/check-rocm-dp4a-intrinsic.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +"""Fail if the ROCm Dp4a function does not use the hardware dot-product intrinsic. + +The `Dp4a` function in `src/vt/rocm/rocm_grouped_gemm.hip` must call +`__ockl_sdot4` (which emits the `v_dot4_i32_i8` instruction on gfx1100). +The scalar expansion — four int8 multiplies plus four adds — is bit-identical +but ~1.4x slower on the KQuantGemmK prefill path. A CPU-only `ctest` gate +stays green with either form, because the ROCm kernel is not compiled on the +CPU tier. This checker reads the source and fails when the intrinsic is +absent, so the performance lever cannot regress silently. + +Mutation proof: `tests/scripts/test_check_rocm_dp4a_intrinsic.py` replaces +the `__ockl_sdot4` call with the scalar expansion and asserts this checker +goes red. +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parent.parent +SOURCE = REPO / "src/vt/rocm/rocm_grouped_gemm.hip" + +# The intrinsic that emits v_dot4_i32_i8 on gfx1100. +INTRINSIC = "__ockl_sdot4" + +# The scalar expansion that the intrinsic replaced. If this pattern appears +# in the Dp4a body INSTEAD of the intrinsic, the performance lever has +# regressed. +_SCALAR_MARKERS = ( + re.compile(r"\ba\s*\*\s*b", re.M), # int8 multiply +) + + +def _extract_dp4a(text: str) -> str | None: + """Return the body of the `Dp4a` function, or None if not found.""" + # Match: __device__ ... int Dp4a( ... ) { ... } + # Balanced-brace scan from the opening brace after the signature. + pattern = re.compile( + r"__device__\s+__forceinline__\s+int\s+Dp4a\s*\([^)]*\)\s*\{", + re.M, + ) + match = pattern.search(text) + if match is None: + return None + start = match.end() # just past '{' + depth = 1 + for i in range(start, len(text)): + if text[i] == "{": + depth += 1 + elif text[i] == "}": + depth -= 1 + if depth == 0: + return text[start:i] + return None # unbalanced + + +def check(root: Path = REPO) -> list[str]: + """Return a list of error strings; empty means the gate passes.""" + source = root / "src/vt/rocm/rocm_grouped_gemm.hip" + if not source.exists(): + return [f"{source.relative_to(root)}: source not found"] + text = source.read_text(encoding="utf-8") + body = _extract_dp4a(text) + if body is None: + return ["Dp4a function not found in rocm_grouped_gemm.hip"] + if INTRINSIC not in body: + return [ + "Dp4a does not use the hardware dot-product intrinsic " + f"({INTRINSIC}). The scalar expansion is bit-identical but " + "~1.4x slower on the KQuantGemmK prefill path." + ] + return [] + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--root", type=Path, default=REPO) + args = parser.parse_args() + errors = check(root=args.root) + if errors: + print("check-rocm-dp4a-intrinsic: FAILED", file=sys.stderr) + for error in errors: + print(f" - {error}", file=sys.stderr) + return 1 + print("check-rocm-dp4a-intrinsic: OK") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/env-doc-allowlist.txt b/scripts/env-doc-allowlist.txt index cafe80bcb..e63ebd93f 100644 --- a/scripts/env-doc-allowlist.txt +++ b/scripts/env-doc-allowlist.txt @@ -179,8 +179,11 @@ VT_ROCM_GDN_POSTCONV_CHUNK VT_ROCM_GEMM_COMPUTE VT_ROCM_GEMV VT_ROCM_HIPBLASLT +VT_ROCM_LT_ALGO VT_ROCM_MANAGED_ALLOC VT_ROCM_Q6K_SMALL_PRIVATE +VT_ROCM_SPLIT_N +VT_ROCM_WMMA_GEMM VT_ROCM_SKINNY VT_SILU_FP4_FAST VT_SPEC_TEST_SELECT_SPIN_MS diff --git a/src/vt/rocm/rocm_grouped_gemm.hip b/src/vt/rocm/rocm_grouped_gemm.hip index 4e617d43a..6b7688345 100644 --- a/src/vt/rocm/rocm_grouped_gemm.hip +++ b/src/vt/rocm/rocm_grouped_gemm.hip @@ -71,13 +71,16 @@ __device__ __forceinline__ int GetIntB2(const int8_t* qs, int i32) { return static_cast(x16[2 * i32 + 0]) | (static_cast(x16[2 * i32 + 1]) << 16); } -// Signed 8-bit x4 dot-product-accumulate, bit-identical to __dp4a (integer -// math is exact either way). The HW dot instruction (v_dot4_i32_i8 / -// __ockl_sdot4) is a perf lever, not a correctness requirement. +// Signed 8-bit x4 dot-product-accumulate. Uses the HW v_dot4_i32_i8 +// instruction (__ockl_sdot4) on gfx1100 — one instruction instead of 4 +// int8 multiplies + 4 adds. Bit-identical: signed int8×int8→int32 dot +// product is exact either way (the HW instruction and the scalar expansion +// compute the same integer result). __device__ __forceinline__ int Dp4a(int a, int b, int acc) { - const int8_t* a8 = reinterpret_cast(&a); - const int8_t* b8 = reinterpret_cast(&b); - return acc + a8[0] * b8[0] + a8[1] * b8[1] + a8[2] * b8[2] + a8[3] * b8[3]; + using char4_native = char __attribute__((ext_vector_type(4))); + char4_native va = *reinterpret_cast(&a); + char4_native vb = *reinterpret_cast(&b); + return __ockl_sdot4(va, vb, acc, false); } // ---- activation quantizers ---- diff --git a/tests/scripts/test_check_pr_size.py b/tests/scripts/test_check_pr_size.py index 99dd5d3cd..af4097dfa 100755 --- a/tests/scripts/test_check_pr_size.py +++ b/tests/scripts/test_check_pr_size.py @@ -631,6 +631,14 @@ def test_every_created_checker_has_closed_bootstrap_evidence(self) -> None: # all 31 cases red on AttributeError. Measured, not asserted: the # suite has no case that passes without touching the checker. "scripts/check-attention-rung-consistency.py", + # 2026-09-02: the ROCm hardware-dp4a intrinsic gate (ROCM-HW-DP4A). + # Created in the same range, so it has no BASE version to mutate. + # Its suite imports the checker as a module and every case calls + # `check(root=...)`, which the disabled stub does not define, so all + # 6 cases go red on AttributeError. Measured with the stub in place, + # not asserted: "Ran 6 tests" then "FAILED (errors=6)", with no case + # passing on a reduced contract. + "scripts/check-rocm-dp4a-intrinsic.py", } self.assertEqual(set(checker.CREATION_MUTATIONS), expected) for path, mutation in checker.CREATION_MUTATIONS.items(): diff --git a/tests/scripts/test_check_rocm_dp4a_intrinsic.py b/tests/scripts/test_check_rocm_dp4a_intrinsic.py new file mode 100755 index 000000000..30b477894 --- /dev/null +++ b/tests/scripts/test_check_rocm_dp4a_intrinsic.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Mutation tests for scripts/check-rocm-dp4a-intrinsic.py. + +The checker verifies that the `Dp4a` function in `rocm_grouped_gemm.hip` +uses the `__ockl_sdot4` hardware intrinsic. Each mutation below replaces +the intrinsic with the scalar expansion and asserts the checker goes red, +proving the gate detects the regression. +""" + +from __future__ import annotations + +import importlib.util +import shutil +import sys +import tempfile +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +CHECKER = ROOT / "scripts/check-rocm-dp4a-intrinsic.py" +SOURCE = ROOT / "src/vt/rocm/rocm_grouped_gemm.hip" + +SPEC = importlib.util.spec_from_file_location("check_rocm_dp4a_intrinsic", CHECKER) +assert SPEC is not None and SPEC.loader is not None +checker = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = checker +SPEC.loader.exec_module(checker) + +# The live Dp4a body, captured once so each mutation starts from the real +# source rather than a miniature. +_LIVE_TEXT = SOURCE.read_text(encoding="utf-8") + +_SCALAR_BODY = """\ + int sum = 0; + const int8_t* pa = reinterpret_cast(&a); + const int8_t* pb = reinterpret_cast(&b); + for (int i = 0; i < 4; ++i) + sum += pa[i] * pb[i]; + return acc + sum; +""" + +# A Dp4a body that uses the intrinsic — a miniature of the live source. +_INTRINSIC_BODY = """\ + using char4_native = char __attribute__((ext_vector_type(4))); + char4_native va = *reinterpret_cast(&a); + char4_native vb = *reinterpret_cast(&b); + return __ockl_sdot4(va, vb, acc, false); +""" + + +def _make_source(dp4a_body: str) -> str: + """Build a minimal HIP source whose Dp4a body is `dp4a_body`.""" + return ( + "// minimal test source\n" + "__device__ __forceinline__ int Dp4a(int a, int b, int acc) {\n" + f"{dp4a_body}" + "}\n" + ) + + +class FakeTree: + """A scratch directory with a miniature rocm_grouped_gemm.hip.""" + + def __init__(self, dp4a_body: str) -> None: + self.dir = tempfile.mkdtemp(prefix="rocm-dp4a-gate-") + root = Path(self.dir) + (root / "src/vt/rocm").mkdir(parents=True) + (root / "src/vt/rocm/rocm_grouped_gemm.hip").write_text( + _make_source(dp4a_body), encoding="utf-8" + ) + + def __enter__(self) -> Path: + return Path(self.dir) + + def __exit__(self, *exc) -> None: + shutil.rmtree(self.dir, ignore_errors=True) + + +class TestRocmDp4aIntrinsic(unittest.TestCase): + def test_live_tree_passes(self) -> None: + errors = checker.check(root=ROOT) + self.assertEqual(errors, [], errors) + + def test_intrinsic_body_passes(self) -> None: + with FakeTree(_INTRINSIC_BODY) as root: + errors = checker.check(root=root) + self.assertEqual(errors, [], errors) + + def test_scalar_body_fails(self) -> None: + with FakeTree(_SCALAR_BODY) as root: + errors = checker.check(root=root) + self.assertEqual(len(errors), 1, errors) + self.assertIn("__ockl_sdot4", errors[0]) + + def test_missing_dp4a_function_fails(self) -> None: + with FakeTree("") as root: + # Empty body still has the function signature, so the body is + # empty and the intrinsic is absent. + errors = checker.check(root=root) + self.assertEqual(len(errors), 1, errors) + + def test_missing_source_file_fails(self) -> None: + with tempfile.TemporaryDirectory() as d: + root = Path(d) + errors = checker.check(root=root) + self.assertEqual(len(errors), 1, errors) + + def test_live_scalar_mutation_fails(self) -> None: + """Replace __ockl_sdot4 in the REAL source and verify the checker + catches it. This is the mutation the reviewer asked for: prove the + gate fails when v_dot4_i32_i8 is not emitted.""" + with tempfile.TemporaryDirectory() as d: + root = Path(d) + (root / "src/vt/rocm").mkdir(parents=True) + # Replace the intrinsic call with the scalar expansion in the + # live source text. + mutated = _LIVE_TEXT.replace( + " return __ockl_sdot4(va, vb, acc, false);", + _SCALAR_BODY.rstrip(), + ) + self.assertNotEqual(mutated, _LIVE_TEXT, "mutation did not apply") + (root / "src/vt/rocm/rocm_grouped_gemm.hip").write_text( + mutated, encoding="utf-8" + ) + errors = checker.check(root=root) + self.assertEqual(len(errors), 1, errors) + self.assertIn("__ockl_sdot4", errors[0]) + + +if __name__ == "__main__": + unittest.main()