diff --git a/example/ck_tile/42_unified_attention/script/edge_test_sink.sh b/example/ck_tile/42_unified_attention/script/edge_test_sink.sh new file mode 100755 index 0000000000..c7c196f971 --- /dev/null +++ b/example/ck_tile/42_unified_attention/script/edge_test_sink.sh @@ -0,0 +1,130 @@ +#!/bin/bash +# edge_test_sink.sh - edge-case tests for the kernel-side per-Q-head +# attention sink path on the prefill tier (prefill_d64 / prefill_d128). +# +# These cases probe corners that the smoke test does not: +# 1. sink ≈ -infinity (collapses to no-sink output; sanity guard +# against a runaway sign / overflow in the m-init). +# 2. sink ≈ +infinity (sink absorbs all softmax mass; output → 0). +# 3. sink = 0 (sink contributes weight 1 per virtual key). +# 4. per-head random (exercises the per-row head indexing in both +# kernel and reference; default GQA-broadcast +# check baked in because baseB has 8 Q-heads). +# 5. per-head CSV (the host-side CSV parser path). +# 6. random + d128 (covers prefill_d128 with non-trivial per-head +# sinks). +# +# The "all-window-masked + sink" fixture from the plan (combine the SWA +# `b:0,0` window with `sink=const:0.0`) needs sink_local instances — +# those land in a later phase. The dispatcher currently fast-fails the +# (sink && is_local) combo to avoid silently routing to a non-sink +# instance; that fixture is documented in the script body but skipped +# below. +# +# Run with HIP_VISIBLE_DEVICES set; defaults to 6 on the shared dev +# node. +# +# Exit code = number of FAIL'd tests (0 = all PASS). + +set -uo pipefail + +export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-6}" + +EXE_NAME=tile_example_unified_attention +EXE="${EXE:-$(find . -name "$EXE_NAME" -type f -executable 2>/dev/null | head -n 1)}" +if [ -z "${EXE:-}" ] || [ ! -x "$EXE" ]; then + echo "ERROR: $EXE_NAME not found. Set EXE=/path/to/$EXE_NAME or run from build dir." >&2 + exit 2 +fi +echo "Using EXE=$EXE" +echo "Using HIP_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES" + +# Deterministic, verification-only fixtures. +COMMON="-prec=bf16 -seed=17 -verify=1 -warmup=0 -repeat=1 -varlen=0 -nb=1024 -page_blk_size=128" + +# baselineB: prefill_d64 (h_k=1, nqpkv=8 → GQA-8 fan-out). Same fixture +# as smoke_test_sink.sh so the two scripts cross-validate. +BASELINE_B="-d=64 -h_k=1 -nqpkv=8 -b=4 -s=512 -s_k=512 -query_lens=400,256,512,128 -kv_lens=400,256,512,128" + +# baselineD: prefill_d128 with non-trivial query lengths so the variant +# selector lands on prefill_d128 (not decode_d128_m128 which has no +# sink instance yet). +BASELINE_D="-d=128 -h_k=8 -nqpkv=1 -b=4 -s=512 -s_k=512 -query_lens=300,300,300,300 -kv_lens=300,300,300,300" + +TESTS=( + # 1. Sink ≈ -infinity: kernel must collapse the sink contribution to + # ~0 (m_init = -1e4/sm_scale dominated by max_j S_raw_j after the + # first iteration), reproducing the no-sink output. Both kernel + # and reference apply the same near-zero sink, so they match. + "baseB sink≈-inf |$BASELINE_B -mask=b -sink=const:-1e4" + + # 2. Sink ≈ +infinity: kernel must absorb almost all the softmax + # mass onto the sink (m_init = 1e4/sm_scale dominates max_j S_raw_j), + # driving the post-normalization V-weighted output to ≈ 0. The + # reference does the same, so they match within bf16 noise (the + # output is dominated by zero ± rounding). + "baseB sink≈+inf |$BASELINE_B -mask=b -sink=const:+1e4" + + # 3. Sink = 0: one virtual key with weight exp(0 - m_max); ~0.1-1% + # of the existing softmax mass. Output is materially different + # from no-sink but well-defined; both kernel and reference must + # agree. + "baseB sink=0 |$BASELINE_B -mask=b -sink=const:0.0" + + # 4. Per-head random sinks (GQA-8 fan-out): baseB has 8 Q-heads, so + # `random:17` draws 8 distinct sink values. The kernel's per-row + # indexing `sink_ptr_pre_offset[r % num_queries_per_kv]` must + # match the reference's `sink[q_head_idx]` lookup — the same head + # must receive the same sink across every q-token in its group. + "baseB sink=random:17 |$BASELINE_B -mask=b -sink=random:17" + + # 5. Explicit per-head CSV (8 distinct sinks). Same indexing check + # as (4), but with the CSV parser path instead of the seeded + # 'random:N' draw. + "baseB sink=CSV[8] |$BASELINE_B -mask=b -sink=0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8" + + # 6. prefill_d128 with per-head random sinks (h_k=8, nqpkv=1 → 8 + # Q-heads, one Q per KV group). Validates the d=128 sink + # instance and the GQA-1 indexing edge. + "baseD d128 sink=rand |$BASELINE_D -mask=b -sink=random:17" +) + +# (Skipped) All-window-masked + sink — needs sink_local instances which +# land in a later phase. Dispatcher currently fast-fails (sink && +# is_local) → {false, -1.f}, so the test would FAIL. Re-enable when +# Phase 9 ships: +# $BASELINE_B -mask=b:0,0 -sink=const:0.0 +# Expected outcome at that point: kernel writes 0 for sink-only Q-tiles +# (not NaN), reference does the same. + +n_pass=0 +n_fail=0 + +for entry in "${TESTS[@]}"; do + name="${entry%%|*}" + name="${name// /}" + args="${entry#*|}" + + printf '== %-22s :: %s\n' "$name" "$args" + set +e + "$EXE" $COMMON $args > /tmp/sink_edge.$$ 2>&1 + ret=$? + set -e + + if [ $ret -eq 0 ]; then + echo " PASS" + n_pass=$((n_pass + 1)) + else + echo " FAIL (rc=$ret). Tail of output:" + tail -3 /tmp/sink_edge.$$ | sed 's/^/ /' + n_fail=$((n_fail + 1)) + fi + rm -f /tmp/sink_edge.$$ +done + +echo +echo "Summary:" +printf ' PASS : %d\n' $n_pass +printf ' FAIL : %d\n' $n_fail + +exit $n_fail diff --git a/example/ck_tile/42_unified_attention/script/perf_test_sink.sh b/example/ck_tile/42_unified_attention/script/perf_test_sink.sh new file mode 100755 index 0000000000..aa2d1a63cb --- /dev/null +++ b/example/ck_tile/42_unified_attention/script/perf_test_sink.sh @@ -0,0 +1,188 @@ +#!/bin/bash +# perf_test_sink.sh — perf gate that locks in the "zero-overhead sinks" +# contract from Phase 4 of the sink rollout. +# +# Two regressions this script catches: +# +# 1. The no-sink instance's runtime growing past `BASELINE_TOLERANCE` +# (default 5%) of the captured Phase-4 baseline. This fires on +# catastrophic regressions of the kHasSink=false code path: +# barriers added in the K/V loop, extra HBM round-trips per tile, +# tile-storage allocated inside the iteration, register-spill +# cliffs from an over-grown pipeline operator(), etc. +# → time(no-sink) / BASELINE_NO_SINK must be ≤ BASELINE_TOLERANCE. +# +# 2. The sink-aware path adding more than `SINK_OVERHEAD_MAX` (default +# 10%) over the no-sink leg on the same shape. Same kinds of +# catastrophic regressions, but on the kHasSink=true path +# specifically — e.g. someone moves the per-row sink-init sweep +# inside the K/V loop, or makes `sink_ptr_pre_offset[r % +# num_qpkv]` a non-coalesced indirection. +# → time(sink) / time(no-sink) must be ≤ SINK_OVERHEAD_MAX. +# +# Both gates run on the two prefill shapes that mirror `perf_test_swa.sh`'s +# coverage: +# +# * d=128 MHA prefill (-d=128 -h_k=8 -nqpkv=1) q = kv = 8192 +# * d=64 GQA-8 prefill (-d=64 -h_k=1 -nqpkv=8) q = kv = 8192 +# +# Baselines were captured on MI355 (HIP_VISIBLE_DEVICES=6) right after +# Phase 4 landed. The procedure was: +# +# for shape in {d128_MHA, d64_GQA8}: +# for i in 1..3: +# run -mask=b -verify=0 -warmup=10 -repeat=30 -varlen=0 +# extract "X.XX ms" from the bench line +# baseline = mean of the 3 runs +# +# Observed numbers (ms) on MI355: +# +# d128 MHA no-sink : 0.4573, 0.4606, 0.4610 mean ≈ 0.460 +# d64 GQA-8 no-sink: 0.3370, 0.3388, 0.3381 mean ≈ 0.338 +# +# Both shapes show sink-on overhead of <2.5% on the same MI355, so the +# 10% SINK_OVERHEAD gate has ~4× headroom over the observed noise floor. +# +# What this gate does NOT catch — and why we are OK with that: +# +# The natural "branch-leak" regression mode (someone replaces +# `if constexpr (kHasSink)` with something that lets the per-row +# sweep run on no-sink instances) is empirically below the gate's +# noise floor on prefill. The sink init's cost is ≈ 1 fp32 division +# per thread (kBlockM=256, ~256 threads/CTA → ~1 row/thread); the +# K/V loop runs ~10^6 cycles. Forced sanity provocations on MI355: +# +# * `if constexpr(kHasSink)` → `if constexpr(true)` plus a null-guard +# on `sink_ptr_pre_offset`: no-sink stayed at 0.454 ms +# (1.3% under baseline) ❌ +# * 100× redundant sweep in the lambda: no-sink crept to 0.350 ms +# (compiler dead-coded most writes) (3.5% over baseline) ❌ +# * 50× `block_sync_lds()` before the +# set_tile: no-sink at 0.461 ms +# (0.2% over baseline) ❌ +# +# In all three cases the gate stays GREEN. That's because the gate +# is calibrated against the *kernel wall time*, not the init phase +# in isolation, and the init phase is fundamentally < 1 µs out of +# ~460 µs. The gate's actual value is catching regressions in the +# one place that dominates wall time: the K/V loop. The "did the +# if-constexpr collapse correctly" check belongs to a static / SASS +# diff, not a runtime gate. We accept that scope. +# +# Run with HIP_VISIBLE_DEVICES set to your assigned GPU (defaults to 6 +# on the shared MI355 dev node). Build with +# `ninja -j 50 tile_example_unified_attention` from the CK build dir +# (do not use `cmake --build`; see Sink-implementation-steps.md hard +# constraints). Exit code is the number of failed assertions across +# both shapes (0 = all PASS). + +set -uo pipefail + +export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-6}" + +EXE_NAME=tile_example_unified_attention +EXE="${EXE:-$(find . -name "$EXE_NAME" -type f -executable 2>/dev/null | head -n 1)}" +if [ -z "${EXE:-}" ] || [ ! -x "$EXE" ]; then + echo "ERROR: $EXE_NAME not found. Set EXE=/path/to/$EXE_NAME or run from build dir." >&2 + exit 2 +fi +echo "Using EXE=$EXE" +echo "Using HIP_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES" + +# Allow override via env vars; defaults match the Phase 5 plan thresholds. +BASELINE_TOLERANCE="${BASELINE_TOLERANCE:-1.05}" # no-sink within 5% of baseline +SINK_OVERHEAD_MAX="${SINK_OVERHEAD_MAX:-1.10}" # sink-on within 10% of no-sink + +# Bench-only config — same shape as perf_test_swa.sh's so the two +# scripts can share GPU-warmup state when run back-to-back. +COMMON="-prec=bf16 -seed=17 -verify=0 -warmup=10 -repeat=30 -varlen=0 -nb=512" + +# Each row is "NAME|BASELINE_MS|EXTRA_ARGS". BASELINE_MS is the +# captured no-sink runtime (see procedure above). Adding new shapes: +# capture the baseline with `EXE COMMON SHAPE_ARGS -mask=b` averaged +# over 3 runs and paste the mean here. +SHAPES=( + "prefill_d128 MHA q=kv=8192 |0.460 |-page_blk_size=128 -d=128 -h_k=8 -nqpkv=1 -b=1 -s=8192 -s_k=8192 -query_lens=8192 -kv_lens=8192" + "prefill_d64 GQA-8 q=kv=8192 |0.338 |-page_blk_size=128 -d=64 -h_k=1 -nqpkv=8 -b=1 -s=8192 -s_k=8192 -query_lens=8192 -kv_lens=8192" +) + +# Extract "X.XXX ms" from a benchmark output line. Same helper as +# perf_test_swa.sh so the two scripts stay in lockstep. +extract_ms() { + grep -oE '[0-9]+\.[0-9]+ ms' "$1" | head -1 | awk '{print $1}' +} + +n_pass=0 +n_fail=0 + +for row in "${SHAPES[@]}"; do + name="${row%%|*}" + rest="${row#*|}" + baseline="${rest%%|*}" + baseline="${baseline// /}" + args="${rest#*|}" + + printf '== %s\n' "$name" + + nosink_log=$(mktemp) + sink_log=$(mktemp) + "$EXE" $COMMON $args -mask=b > "$nosink_log" 2>&1 || true + "$EXE" $COMMON $args -mask=b -sink=random:17 > "$sink_log" 2>&1 || true + + t_nosink=$(extract_ms "$nosink_log") + t_sink=$(extract_ms "$sink_log") + + if [ -z "$t_nosink" ] || [ -z "$t_sink" ]; then + echo " FAIL: could not parse timing" + echo " no-sink tail:" + tail -3 "$nosink_log" | sed 's/^/ /' + echo " sink tail:" + tail -3 "$sink_log" | sed 's/^/ /' + n_fail=$((n_fail + 1)) + rm -f "$nosink_log" "$sink_log" + continue + fi + rm -f "$nosink_log" "$sink_log" + + # awk handles fp arithmetic; bash itself is integer-only. + baseline_ratio=$(awk -v t="$t_nosink" -v b="$baseline" 'BEGIN{printf "%.3f", t/b}') + overhead_ratio=$(awk -v s="$t_sink" -v n="$t_nosink" 'BEGIN{printf "%.3f", s/n}') + + baseline_pass=$(awk -v r="$baseline_ratio" -v g="$BASELINE_TOLERANCE" \ + 'BEGIN{print (r+0 <= g+0) ? 1 : 0}') + overhead_pass=$(awk -v r="$overhead_ratio" -v g="$SINK_OVERHEAD_MAX" \ + 'BEGIN{print (r+0 <= g+0) ? 1 : 0}') + + printf ' no-sink : %s ms (baseline %s ms, ratio %sx, gate ≤ %s)\n' \ + "$t_nosink" "$baseline" "$baseline_ratio" "$BASELINE_TOLERANCE" + printf ' sink : %s ms (sink/nosink %sx, gate ≤ %s)\n' \ + "$t_sink" "$overhead_ratio" "$SINK_OVERHEAD_MAX" + + shape_fail=0 + if [ "$baseline_pass" = "1" ]; then + echo " PASS baseline (no-sink instance unchanged)" + else + echo " FAIL baseline — kHasSink=false instance got slower" + shape_fail=$((shape_fail + 1)) + fi + + if [ "$overhead_pass" = "1" ]; then + echo " PASS overhead (sink init is near-zero cost)" + else + echo " FAIL overhead — sink-aware init got expensive" + shape_fail=$((shape_fail + 1)) + fi + + if [ "$shape_fail" = "0" ]; then + n_pass=$((n_pass + 1)) + else + n_fail=$((n_fail + shape_fail)) + fi +done + +echo +echo "Summary:" +printf ' PASS (shapes) : %d\n' "$n_pass" +printf ' FAIL (assertions across shapes) : %d\n' "$n_fail" + +exit "$n_fail"