From 9286b6df96ca3a9dfe1c3d5c5bf0d5d150849ea5 Mon Sep 17 00:00:00 2001 From: Damien Lejeune Date: Wed, 27 May 2026 15:06:50 +0000 Subject: [PATCH] Add GPT-OSS shapes instances + checks for non-page-alignment --- ...d_attention_d64_bf16_mask_decode_local.cpp | 15 +++++++ ...attention_d64_bf16_mask_decode_t_local.cpp | 14 ++++++ ...d_attention_d64_fp16_mask_decode_local.cpp | 11 +++++ ...attention_d64_fp16_mask_decode_t_local.cpp | 11 +++++ .../script/edge_test_swa.sh | 43 +++++++++++++++++++ .../script/perf_test_swa.sh | 23 +++++++--- .../unified_attention.cpp | 15 +++++-- 7 files changed, 122 insertions(+), 10 deletions(-) create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_local.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_t_local.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_local.cpp create mode 100644 example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_t_local.cpp diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_local.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_local.cpp new file mode 100644 index 0000000000..51070491d6 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_local.cpp @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +// Sliding-Window-Attention decode_d64_m128 bf16 (IsLocal=true, runtime +// page_size). Pairs with the existing INST_UNIFIED_ATTENTION_DISPATCH +// `_mask_decode` instance — same variant, same dtype, just the IsLocal= +// true cousin that honours both window bounds. +INST_UNIFIED_ATTENTION_DISPATCH_PS_LOCAL(decode_d64_m128, bf16, true, 0, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_t_local.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_t_local.cpp new file mode 100644 index 0000000000..20d1eaf42b --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_t_local.cpp @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +// Sliding-Window-Attention decode_d64_m16 bf16 — the q=1 tiny tier. This +// is the primary GPT-OSS decode shape (single-token generation step with +// a 4k-128k context window). +INST_UNIFIED_ATTENTION_DISPATCH_PS_LOCAL(decode_d64_m16, bf16, true, 0, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_local.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_local.cpp new file mode 100644 index 0000000000..d024699635 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_local.cpp @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +INST_UNIFIED_ATTENTION_DISPATCH_PS_LOCAL(decode_d64_m128, fp16, true, 0, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_t_local.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_t_local.cpp new file mode 100644 index 0000000000..709f956126 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_t_local.cpp @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +INST_UNIFIED_ATTENTION_DISPATCH_PS_LOCAL(decode_d64_m16, fp16, true, 0, true) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/script/edge_test_swa.sh b/example/ck_tile/42_unified_attention/script/edge_test_swa.sh index c6d61b71e7..378f81b41b 100755 --- a/example/ck_tile/42_unified_attention/script/edge_test_swa.sh +++ b/example/ck_tile/42_unified_attention/script/edge_test_swa.sh @@ -83,6 +83,49 @@ TESTS=( # firing on the trailing partial page. "d64 odd s_k=480 (xb:64) |-d=64 -h_k=1 -nqpkv=8 -b=2 -s=480 -s_k=480 -query_lens=480,480 -kv_lens=480,480 -mask=xb:64" "d128 odd s_k=480 (xb:64) |-d=128 -h_k=8 -nqpkv=1 -b=2 -s=480 -s_k=480 -query_lens=257,480 -kv_lens=257,480 -mask=xb:64" + + # --- Phase 5: GPT-OSS shapes (page_blk_size=32) ----------------------- + # The primary motivator for SWA in unified attention. GPT-OSS uses + # d=64 with GQA-8, page_blk_size=32, and three operating regimes: + # * q=1 decode (single-token generation, routes to decode_d64_m16) + # * q≈128 medium (short prefill / continuation, decode_d64_m128) + # * q∈[256,1024] (full prefill, prefill_d64) + # Each is tested with both window styles GPT-OSS configures. + "DECODE_BS32_Q1 xb:128 |-d=64 -h_k=1 -nqpkv=8 -b=4 -s=512 -s_k=512 -query_lens=1,1,1,1 -kv_lens=512,512,512,512 -page_blk_size=32 -mask=xb:128" + "DECODE_BS32_Q1 b:127,0 |-d=64 -h_k=1 -nqpkv=8 -b=4 -s=512 -s_k=512 -query_lens=1,1,1,1 -kv_lens=512,512,512,512 -page_blk_size=32 -mask=b:127,0" + "DECODE_BS32_Q128 xb:128 |-d=64 -h_k=1 -nqpkv=8 -b=4 -s=1024 -s_k=1024 -query_lens=128,128,128,128 -kv_lens=1024,1024,1024,1024 -page_blk_size=32 -mask=xb:128" + "DECODE_BS32_Q128 b:127,0 |-d=64 -h_k=1 -nqpkv=8 -b=4 -s=1024 -s_k=1024 -query_lens=128,128,128,128 -kv_lens=1024,1024,1024,1024 -page_blk_size=32 -mask=b:127,0" + "DECODE_BS32_QM xb:128 |-d=64 -h_k=1 -nqpkv=8 -b=4 -s=1024 -s_k=1024 -query_lens=512,1024,512,1024 -kv_lens=1024,1024,1024,1024 -page_blk_size=32 -mask=xb:128" + "DECODE_BS32_QM b:127,0 |-d=64 -h_k=1 -nqpkv=8 -b=4 -s=1024 -s_k=1024 -query_lens=512,1024,512,1024 -kv_lens=1024,1024,1024,1024 -page_blk_size=32 -mask=b:127,0" + + # --- Phase 5.4: non-page-aligned stress ------------------------------- + # For prefill_d64 bf16, kPageBlockSize=32 (kernel tile in tokens). The + # *runtime* page_size is set by -page_blk_size. When page_size > + # kPageBlockSize, each cache page holds multiple kernel tiles, and Step + # D's `num_blocks_start` (in kernel-tile units) can land *mid-page*: + # (num_blocks_start * kPageBlockSize) % page_size != 0 + # That triggers the `logical_token / page_size` math path inside + # refresh_*_offsets to resolve both the right page AND the right + # within-page row. The shapes below pick window + Q-tile combos that + # force `num_blocks_start` to an odd multiple of kPageBlockSize: + # + # * page_size=64 (2 tiles/page), -mask=xb:64 + # Q-tile 1 (rows 256..511) → num_blocks_start = 7 → 7*32 = 224 = + # 3.5 pages → mid-page start. + # * page_size=128 (4 tiles/page), -mask=xb:64 + # Q-tile 1 → num_blocks_start = 7 = 1.75 pages → mid-page start. + # (This is the same case the Phase 3 smoke test already hits, + # kept here for explicitness.) + # * page_size=64, -mask=b:48,0 (window 48, asymmetric) + # Q-tile 1 → num_blocks_start = 6 (page-aligned) but the per-Q-tile + # start *boundary* varies across the batch — keeps both alignment + # paths covered in a single run. + "non-align ps=64 xb:64 |-d=64 -h_k=1 -nqpkv=8 -b=2 -s=512 -s_k=512 -query_lens=512,512 -kv_lens=512,512 -page_blk_size=64 -mask=xb:64" + "non-align ps=128 xb:64 |-d=64 -h_k=1 -nqpkv=8 -b=2 -s=512 -s_k=512 -query_lens=512,512 -kv_lens=512,512 -page_blk_size=128 -mask=xb:64" + "non-align ps=64 b:48,0 |-d=64 -h_k=1 -nqpkv=8 -b=2 -s=512 -s_k=512 -query_lens=512,512 -kv_lens=512,512 -page_blk_size=64 -mask=b:48,0" + # And on prefill_d128 (kPageBlockSize=16): page_size=128 = 8 tiles/page, + # so virtually every Step D clip will be mid-page. + "non-align d128 ps=128 xb:64 |-d=128 -h_k=8 -nqpkv=1 -b=2 -s=512 -s_k=512 -query_lens=257,512 -kv_lens=257,512 -page_blk_size=128 -mask=xb:64" ) n_pass=0 diff --git a/example/ck_tile/42_unified_attention/script/perf_test_swa.sh b/example/ck_tile/42_unified_attention/script/perf_test_swa.sh index 7d97a5d1db..a80b78adf8 100755 --- a/example/ck_tile/42_unified_attention/script/perf_test_swa.sh +++ b/example/ck_tile/42_unified_attention/script/perf_test_swa.sh @@ -41,21 +41,30 @@ fi echo "Using EXE=$EXE" echo "Using HIP_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES" -# Speedup threshold: 5× headroom on the observed ~10–21× ratios. Below +# Speedup threshold: 5× headroom on the observed ~10–40× ratios. Below # this means Step D was disabled or the SWA envelope blew open. RATIO_MIN="${RATIO_MIN:-5.0}" # Bench-only config — verify off, ample warmup, 30 timed repeats so the # faster runs are dominated by the kernel and not the launch overhead. -COMMON="-prec=bf16 -seed=17 -verify=0 -warmup=10 -repeat=30 -varlen=0 \ - -nb=128 -page_blk_size=128 -b=1 -s=8192 -s_k=8192 \ - -query_lens=8192 -kv_lens=8192" +COMMON="-prec=bf16 -seed=17 -verify=0 -warmup=10 -repeat=30 -varlen=0 -nb=512" # Each row is "NAME|EXTRA_ARGS". Both -mask=b and -mask=xb:128 are appended -# per run by the loop; we never time the kernel with verify on. +# per run by the loop; we never time the kernel with verify on. The shapes +# cover all four SWA-capable instances landed so far: +# +# prefill_d128 / prefill_d64 (Phase 3 — long prefill) +# decode_d64_m128 (Phase 5 — short prefill, q≈128) +# decode_d64_m16 (Phase 5 — q=1 generation step, GPT-OSS) +# +# Phase 5 rows use -page_blk_size=32 to mirror GPT-OSS's KV cache page +# layout (vs the 128 used for the prefill rows where ps is irrelevant to +# the Step D math). SHAPES=( - "d128 MHA prefill q=kv=8192 win=128 |-d=128 -h_k=8 -nqpkv=1" - "d64 GQA-8 prefill q=kv=8192 win=128|-d=64 -h_k=1 -nqpkv=8" + "prefill_d128 MHA q=kv=8192 win=128 |-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 win=128 |-page_blk_size=128 -d=64 -h_k=1 -nqpkv=8 -b=1 -s=8192 -s_k=8192 -query_lens=8192 -kv_lens=8192" + "decode_m128 GQA-8 q=128 kv=8192 |-page_blk_size=32 -d=64 -h_k=1 -nqpkv=8 -b=1 -s=8192 -s_k=8192 -query_lens=128 -kv_lens=8192" + "decode_m16 GQA-8 q=1 kv=8192 |-page_blk_size=32 -d=64 -h_k=1 -nqpkv=8 -b=1 -s=8192 -s_k=8192 -query_lens=1 -kv_lens=8192" ) # Extract the "X.XXX ms" from a single benchmark line such as diff --git a/example/ck_tile/42_unified_attention/unified_attention.cpp b/example/ck_tile/42_unified_attention/unified_attention.cpp index 144ba9b59a..9017430461 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/unified_attention.cpp @@ -137,9 +137,18 @@ std::pair dispatch_local(const unified_attention_args& args, const stream_config& config) { using DT = unified_attention_args::data_type_enum; - if constexpr((V == KernelVariant::prefill_d128 || - V == KernelVariant::prefill_d64) && - (DType == DT::bf16 || DType == DT::fp16)) + // SWA availability matrix as of Phase 5: + // prefill_d128 / prefill_d64 — large-tier prefill (Phase 3) + // decode_d64_m128 / decode_d64_m16 — GPT-OSS shapes (Phase 5) + // For all of them we only compile {bf16, fp16}; fp8 SWA is a future + // phase. Anything else (decode_d128_*, decode_d64_m64, …) still falls + // back to {false, 0.f} so the aiter wrapper picks Triton. + constexpr bool is_supported_variant = + (V == KernelVariant::prefill_d128 || + V == KernelVariant::prefill_d64 || + V == KernelVariant::decode_d64_m128 || + V == KernelVariant::decode_d64_m16); + if constexpr(is_supported_variant && (DType == DT::bf16 || DType == DT::fp16)) { return unified_attention_kernel_dispatch< unified_attention_kernel_traits