ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add GPT-OSS shapes instances + checks for non-page-alignment

Browse Source
This commit is contained in:
Damien Lejeune
2026-05-27 15:06:50 +00:00
parent 480b2c5d79
commit 9286b6df96
7 changed files with 122 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_local.cpp Normal file
View File

@@ -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

14
example/ck_tile/42_unified_attention/instances/unified_attention_d64_bf16_mask_decode_t_local.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_local.cpp Normal file
View File

@@ -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

11
example/ck_tile/42_unified_attention/instances/unified_attention_d64_fp16_mask_decode_t_local.cpp Normal file
View File

@@ -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

43
example/ck_tile/42_unified_attention/script/edge_test_swa.sh
View File

@@ -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

23
example/ck_tile/42_unified_attention/script/perf_test_swa.sh
View File

@@ -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

15
example/ck_tile/42_unified_attention/unified_attention.cpp
View File

@@ -137,9 +137,18 @@ std::pair<bool, float> 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<V, DType, /*IsMasking=*/true,