mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-17 00:58:44 +00:00
UA Sink: add perf scripts
This commit is contained in:
Damien Lejeune
130
example/ck_tile/42_unified_attention/script/edge_test_sink.sh
Executable file
130
example/ck_tile/42_unified_attention/script/edge_test_sink.sh
Executable file
@@ -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
|
||||
188
example/ck_tile/42_unified_attention/script/perf_test_sink.sh
Executable file
188
example/ck_tile/42_unified_attention/script/perf_test_sink.sh
Executable file
@@ -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"
|
||||
Reference in New Issue
Block a user