Repository navigation
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
There was a problem hiding this comment.
🟡 Changes recommended
Tuned thresholds belong in configuration, and sliding-window routing currently bypasses the specialized 2D path without gaining KV splitting.
1 open finding
What changed in this PR
Adjusts gfx950 Gluon unified-attention routing to favor KV-split 3D execution for underfilled speculative-decode workloads.
Changes:
- Adds dtype-specific occupancy thresholds.
- Routes low-program non-decode workloads through the 3D kernel.
| File | Description |
|---|---|
aiter/ops/triton/attention/unified_attention.py |
Updates gfx950 Gluon kernel routing heuristics. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if params.kv_cache_dtype == e4m3_dtype: | ||
| split_limit = params.num_sms * 2 | ||
| else: | ||
| split_limit = params.num_sms * 3 // 4 |
@ChuanLi1101 Could we prioritize earlier PR (#6336) over this? I have some extra optimizations over this PR. |
Will check this #PR6336. Will update here with e2e results |
I think that PR (precedes this one), should fix this issue. If you also observe that, can we close this PR? |
|
Closing due to inactivity. If it’s still relevant, feel free to reopen it, and thanks! |
Summary
On gfx950,
use_2d_kernel()sent every batch withmax_seqlen_q > 1to the Gluon 2d (unsplit) kernel, so an MTP verify step -- long context, a handful of query blocks -- parallelized over query blocks x KV heads alone and left most of the GPU idle. The routing now also asks whether those query blocks can fill (num_2d_prgms > num_sms // 2); when they can't, the 3d kernel adds a KV-segment grid axis and recovers the parallelism from the context length.Validation
Run once on the base commit and once with the patch, then compare the two files:
GPU=4 bash pr_bench.sh > before.txt
GPU=4 bash pr_bench.sh > after.txt
`set -u
cd "$(dirname "$0")/aiter" 2>/dev/null || cd /app/aiter
BENCH=op_tests/op_benchmarks/triton/bench_unified_attention.py
COMMON="-hq 16 -hk 1 -d 128 -block_size 128 -equal_seqlens --metric time -test -backend gluon"
run () { # dtype_label b sq sk decode_p extra_flags...
$COMMON "$ @" 2>/dev/null)$dp $ {ms:-nan} ${chk:-NA}"
local dt=$1 b=$2 sq=$3 sk=$4 dp=$5; shift 5
local out ms chk
out=$(HIP_VISIBLE_DEVICES=${GPU:-0} python $BENCH -b "$b" -sq "$sq" -sk "$sk"
ms=$(echo "$out" | awk '/^0 /{print $NF}')
chk=$(echo "$out" | grep -oE '\b(PASS|FAIL)\b' | head -1)
echo "$dt $b $sq $sk
}
for fp8 in "" "-fp8"; do
dt=$([ -n "$fp8" ] && echo fp8 || echo bf16)
for sq in 1 4; do # pure decode, then spec-decode verify
for b in 8 32 64 128; do
for sk in 8192 16384 65536 131072; do
run "$dt" "$b" "$sq" "$sk" 0.0 $fp8
done
done
done
for sq in 16384 32768; do # prefill
run "$dt" 1 "$sq" "$sq" 0.0 $fp8
done
run "$dt" 8 2048 16384 0.75 $fp8 -decode 0.75 # mixed
run "$dt" 32 1024 16384 0.75 $fp8 -decode 0.75
done`