From 62466f3f58594e6a0cda97c06745933abe4b9028 Mon Sep 17 00:00:00 2001 From: arai713 Date: Tue, 14 Jul 2026 03:46:39 -0400 Subject: [PATCH 1/4] Guard gfx942 2D fp8-loader assert behind KV_FP8 Guard the fp8 chunk-count assert and fp8_chunks_per_thread with KV_FP8, to match the 3D builder. Add a codegen regression test covering the bf16 SWA-decode config (t=16, hd=64, threads=256). --- .../kernels/gfx942/attention_tiled_2d.py | 11 ++++--- .../library/tests/test_attention_builds.py | 31 +++++++++++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py index 6450bccf4fb8..d6e5137126fc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py @@ -3240,11 +3240,12 @@ def _dequant_fp8_lds_to_bf16( # store is one ``smem_store_vN(..., n=8)``. fp8_elems_per_chunk = 8 fp8_total_chunks = (T * HD) // fp8_elems_per_chunk - assert fp8_total_chunks % THREADS == 0, ( - f"fp8 loader: total chunks {fp8_total_chunks} must be divisible by " - f"THREADS={THREADS} (T={T}, HD={HD})" - ) - fp8_chunks_per_thread = fp8_total_chunks // THREADS + if KV_FP8: + assert fp8_total_chunks % THREADS == 0, ( + f"fp8 loader: total chunks {fp8_total_chunks} must be divisible by " + f"THREADS={THREADS} (T={T}, HD={HD})" + ) + fp8_chunks_per_thread = fp8_total_chunks // THREADS if KV_FP8 else 0 def _issue_fp8_dequant_loads( kv_tile_idx: Value, buf_idx: Value, lds_token: str diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py index 11457100b20e..bd9065105269 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py @@ -630,6 +630,37 @@ def test_unified_attention_2d_tiled_kernel_compiles(self): # `qq_bias_stride_0` is the very last kernel param. self.assertIn("i32 %qq_bias_stride_0", ll) + def test_gfx942_bf16_swa_decode_no_fp8_loader_assert(self): + """Regression: bf16 windowed-decode must not trip the fp8-loader assert. + + gfx942 2D built the fp8 chunk-count assert unconditionally. A bf16 + SWA/decode small-tile config (T=16, HD=64, THREADS=256) gives + fp8_total_chunks=128; 128 % 256 != 0 raised AssertionError even though + the fp8 loader is never used for bf16 (gfx942 has no fp8 K/V support). + Guarding the assert with KV_FP8 (matching the 3D builder) fixes it. + """ + from kernels.gfx942.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec, + build_unified_attention_2d_tiled, + ) + + spec = UnifiedAttention2DTiledSpec( + head_size=64, + block_size=16, + num_query_heads=64, + num_kv_heads=8, # gpt-oss 64/8 GQA + dtype="bf16", + use_sinks=True, + sliding_window=128, + has_softcap=False, + num_warps=4, + tile_size=16, + block_m_per_warp=16, + ) + k = build_unified_attention_2d_tiled(spec, arch="gfx942") + ll = lower_kernel_to_llvm(k, arch="gfx942") + self.assertIn("define amdgpu_kernel void", ll) + def test_unified_attention_2d_tiled_half_local_pv_compiles(self): """The R4 half-local PV variant emits 32x32 MFMA with its suffixes.""" from kernels import ( From 17c4db56e2ba24c89b3affd1f2b86fea9aaf154a Mon Sep 17 00:00:00 2001 From: arai713 Date: Tue, 14 Jul 2026 03:46:40 -0400 Subject: [PATCH 2/4] Tune gfx942 D64 decode to num_warps=1 and bench sinks Decode is memory-bound; num_warps=1 beats the shipped nw4 on gfx942 D64 with identical output. Apply only to decode (all_decode) so prefill keeps nw4. Add sink and sliding-window support to the decode benchmark plus a gpt-oss sink shapes file --- .../attention/decode/benchmark_decode_live.py | 28 ++++++++++++++++++- .../attention/decode/gpt_oss_sink_shapes.json | 23 +++++++++++++++ .../kernels/common/attention_unified.py | 2 ++ 3 files changed, 52 insertions(+), 1 deletion(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/gpt_oss_sink_shapes.json diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py index 5c8dbbe7f5a8..8175564a4ac7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py @@ -58,6 +58,8 @@ class DecodeShape: block_size: int dtype: str label: str + use_sinks: bool = False + sliding_window: int = 0 @property def signature(self) -> str: @@ -100,6 +102,8 @@ def load_decode_shapes(paths: List[Path]) -> List[DecodeShape]: block_size=int(merged["block_size"]), dtype=str(merged.get("dtype", "bf16")), label=str(merged.get("label", f"kv{merged['seqlen_k']}")), + use_sinks=bool(merged.get("use_sinks", False)), + sliding_window=int(merged.get("sliding_window", 0)), ) shapes.append(shape) return shapes @@ -180,6 +184,13 @@ def _make_inputs( if use_qq_bias else None ) + # Attention sinks: one learned logit per query head, loaded in the kernel + # dtype (the kernel casts to f32 internally). gpt-oss decode only. + sinks = ( + torch.randn(shape.num_query_heads, dtype=dtype, device="cuda") * 0.1 + if shape.use_sinks + else None + ) return dict( q=q, @@ -192,16 +203,26 @@ def _make_inputs( softcap=softcap, alibi_slopes=alibi_slopes, qq_bias=qq_bias, + sinks=sinks, ) def _run_triton( shape: DecodeShape, data: dict, *, warmup: int, iters: int ) -> Optional[float]: - """Time AITER Triton unified_attention. Returns ms or None on failure.""" + """Time AITER Triton unified_attention. Returns ms or None on failure. + + Sink shapes are skipped: AOTriton has no sink attention kernel, so the DSL + is the only path that computes them correctly. Running Triton with + ``sinks=None`` would time a different (wrong-math) kernel and report a + meaningless speedup, so we return None and leave the comparison blank. + """ from rocke.runtime import synchronize_and_release, time_launches import torch + if shape.use_sinks: + return None + try: from aiter.ops.triton.attention.unified_attention import unified_attention as tri # type: ignore except ImportError: @@ -277,6 +298,8 @@ def _run_dsl(shape: DecodeShape, data: dict, num_sms: int, *, warmup: int, iters dtype=shape.dtype, kv_block_size=shape.block_size, num_sms=num_sms, + use_sinks=shape.use_sinks, + sliding_window=shape.sliding_window, ) result = dispatch_attention(req) path = result.spec.path # "2d" or "3d" @@ -296,6 +319,8 @@ def _run_dsl(shape: DecodeShape, data: dict, num_sms: int, *, warmup: int, iters softcap=data["softcap"], use_alibi=data["alibi_slopes"] is not None, use_qq_bias=data["qq_bias"] is not None, + use_sinks=data["sinks"] is not None, + sliding_window=shape.sliding_window, num_sms=num_sms, ) @@ -311,6 +336,7 @@ def call_once(): softmax_scale=data["scale"], block_table=data["block_table"], softcap=data["softcap"], + sinks=data["sinks"], alibi_slopes=data["alibi_slopes"], qq_bias=data["qq_bias"], backend=run_backend, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/gpt_oss_sink_shapes.json b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/gpt_oss_sink_shapes.json new file mode 100644 index 000000000000..1fbbf3baa38e --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/gpt_oss_sink_shapes.json @@ -0,0 +1,23 @@ +{ + "meta": { + "model": "gpt-oss-120b/20b", + "arch": "gfx942", + "dtype": "bf16", + "num_query_heads": 64, + "num_kv_heads": 8, + "head_size": 64, + "block_size": 16, + "use_sinks": true, + "note": "gpt-oss sink-bias decode (seqlen_q=1). Two mask families: full-causal (sliding_window=0) and sliding-window SWA-128. bf16/D64/64-8 GQA. AOTriton has no sink kernel, so the Triton column is blank by design (DSL-only). kv_len sweep short->long; batch sweep to fill the GPU (decode is bandwidth-bound only once batched)." + }, + "shapes": [ + {"batch": 1, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b1_kv2048"}, + {"batch": 64, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b64_kv2048"}, + {"batch": 128, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b128_kv2048"}, + {"batch": 64, "seqlen_q": 1, "seqlen_k": 8192, "sliding_window": 0, "label": "full_sink_b64_kv8192"}, + {"batch": 1, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b1_kv2048"}, + {"batch": 64, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b64_kv2048"}, + {"batch": 128, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b128_kv2048"}, + {"batch": 64, "seqlen_q": 1, "seqlen_k": 8192, "sliding_window": 128, "label": "swa_sink_b64_kv8192"} + ] +} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 77dda48f3bd8..f849f5a33c42 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -993,6 +993,8 @@ def _select_2d_num_warps(problem: UnifiedAttentionProblem) -> int: # gfx942 D64 oracle: paired with ``T=block_size`` and mw=32, four waves fit # in the MI300X 64 KB LDS budget and match the direct gfx942 harness. if _resolve_attention_arch() == "gfx942" and problem.head_size == 64: + if problem.all_decode: + return 1 return 4 if _enable_combo_2d(problem): # bf16 sliding-window is prelude-bound -> nw2 (lighter prelude). But From c8baa477b2acf0f73c6eda3df8e1db17b8de9874 Mon Sep 17 00:00:00 2001 From: arai713 Date: Wed, 15 Jul 2026 20:43:56 -0400 Subject: [PATCH 3/4] Add gfx950 sink decode benchmark support Add attention-sink and sliding-window support to the gfx950 decode benchmark plus a gpt-oss sink shapes file, mirroring the setup for gfx942. --- .../attention/decode/benchmark_decode_live.py | 7 +++++- .../attention/decode/benchmark_decode_live.py | 25 ++++++++++++++++++- .../attention/decode/gpt_oss_sink_shapes.json | 23 +++++++++++++++++ .../kernels/gfx942/attention_tiled_2d.py | 6 ++--- .../library/tests/test_attention_builds.py | 6 ++--- .../common/attention_unified_selectors.cpp | 4 +++ 6 files changed, 63 insertions(+), 8 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/gpt_oss_sink_shapes.json diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py index 8175564a4ac7..2a0ebef2bd91 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py @@ -63,11 +63,16 @@ class DecodeShape: @property def signature(self) -> str: - return ( + sig = ( f"b{self.batch}_sq{self.seqlen_q}_sk{self.seqlen_k}" f"_nhq{self.num_query_heads}_nhk{self.num_kv_heads}" f"_hd{self.head_size}_bs{self.block_size}_{self.dtype}" ) + if self.use_sinks: + sig += "_sinks" + if self.sliding_window: + sig += f"_sw{self.sliding_window}" + return sig def load_decode_shapes(paths: List[Path]) -> List[DecodeShape]: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py index 4ce789280823..64cea2b02216 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py @@ -52,14 +52,21 @@ class DecodeShape: block_size: int dtype: str label: str + use_sinks: bool = False + sliding_window: int = 0 @property def signature(self) -> str: - return ( + sig = ( f"b{self.batch}_sq{self.seqlen_q}_sk{self.seqlen_k}" f"_nhq{self.num_query_heads}_nhk{self.num_kv_heads}" f"_hd{self.head_size}_bs{self.block_size}_{self.dtype}" ) + if self.use_sinks: + sig += "_sinks" + if self.sliding_window: + sig += f"_sw{self.sliding_window}" + return sig def load_decode_shapes(paths: List[Path]) -> List[DecodeShape]: @@ -94,6 +101,8 @@ def load_decode_shapes(paths: List[Path]) -> List[DecodeShape]: block_size=int(merged["block_size"]), dtype=str(merged.get("dtype", "bf16")), label=str(merged.get("label", f"kv{merged['seqlen_k']}")), + use_sinks=bool(merged.get("use_sinks", False)), + sliding_window=int(merged.get("sliding_window", 0)), ) shapes.append(shape) return shapes @@ -174,6 +183,11 @@ def _make_inputs( if use_qq_bias else None ) + sinks = ( + torch.randn(shape.num_query_heads, dtype=dtype, device="cuda") * 0.1 + if shape.use_sinks + else None + ) return dict( q=q, @@ -186,6 +200,7 @@ def _make_inputs( softcap=softcap, alibi_slopes=alibi_slopes, qq_bias=qq_bias, + sinks=sinks, ) @@ -196,6 +211,9 @@ def _run_triton( from rocke.runtime import synchronize_and_release, time_launches import torch + if shape.use_sinks: + return None + try: from aiter.ops.triton.attention.unified_attention import unified_attention as tri # type: ignore except ImportError: @@ -271,6 +289,8 @@ def _run_dsl(shape: DecodeShape, data: dict, num_sms: int, *, warmup: int, iters dtype=shape.dtype, kv_block_size=shape.block_size, num_sms=num_sms, + use_sinks=shape.use_sinks, + sliding_window=shape.sliding_window, ) result = dispatch_attention(req) path = result.spec.path # "2d" or "3d" @@ -290,6 +310,8 @@ def _run_dsl(shape: DecodeShape, data: dict, num_sms: int, *, warmup: int, iters softcap=data["softcap"], use_alibi=data["alibi_slopes"] is not None, use_qq_bias=data["qq_bias"] is not None, + use_sinks=data["sinks"] is not None, + sliding_window=shape.sliding_window, num_sms=num_sms, ) @@ -305,6 +327,7 @@ def call_once(): softmax_scale=data["scale"], block_table=data["block_table"], softcap=data["softcap"], + sinks=data["sinks"], alibi_slopes=data["alibi_slopes"], qq_bias=data["qq_bias"], backend=run_backend, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/gpt_oss_sink_shapes.json b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/gpt_oss_sink_shapes.json new file mode 100644 index 000000000000..3ff61f7b35dc --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/gpt_oss_sink_shapes.json @@ -0,0 +1,23 @@ +{ + "meta": { + "model": "gpt-oss-120b/20b", + "arch": "gfx950", + "dtype": "bf16", + "num_query_heads": 64, + "num_kv_heads": 8, + "head_size": 64, + "block_size": 16, + "use_sinks": true, + "note": "gpt-oss sink-bias decode (seqlen_q=1). Two mask families: full-causal (sliding_window=0) and sliding-window SWA-128. bf16/D64/64-8 GQA. AOTriton has no sink kernel, so the Triton column is blank by design (DSL-only). kv_len sweep short->long; batch sweep to fill the GPU (decode is bandwidth-bound only once batched)." + }, + "shapes": [ + {"batch": 1, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b1_kv2048"}, + {"batch": 64, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b64_kv2048"}, + {"batch": 128, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b128_kv2048"}, + {"batch": 64, "seqlen_q": 1, "seqlen_k": 8192, "sliding_window": 0, "label": "full_sink_b64_kv8192"}, + {"batch": 1, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b1_kv2048"}, + {"batch": 64, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b64_kv2048"}, + {"batch": 128, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b128_kv2048"}, + {"batch": 64, "seqlen_q": 1, "seqlen_k": 8192, "sliding_window": 128, "label": "swa_sink_b64_kv8192"} + ] +} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py index d6e5137126fc..e77291aee247 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py @@ -3242,9 +3242,9 @@ def _dequant_fp8_lds_to_bf16( fp8_total_chunks = (T * HD) // fp8_elems_per_chunk if KV_FP8: assert fp8_total_chunks % THREADS == 0, ( - f"fp8 loader: total chunks {fp8_total_chunks} must be divisible by " - f"THREADS={THREADS} (T={T}, HD={HD})" - ) + f"fp8 loader: total chunks {fp8_total_chunks} must be divisible by " + f"THREADS={THREADS} (T={T}, HD={HD})" + ) fp8_chunks_per_thread = fp8_total_chunks // THREADS if KV_FP8 else 0 def _issue_fp8_dequant_loads( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py index bd9065105269..de98a0f8209a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py @@ -648,13 +648,13 @@ def test_gfx942_bf16_swa_decode_no_fp8_loader_assert(self): head_size=64, block_size=16, num_query_heads=64, - num_kv_heads=8, # gpt-oss 64/8 GQA + num_kv_heads=8, # gpt-oss 64/8 GQA dtype="bf16", use_sinks=True, sliding_window=128, has_softcap=False, - num_warps=4, - tile_size=16, + num_warps=4, + tile_size=16, block_m_per_warp=16, ) k = build_unified_attention_2d_tiled(spec, arch="gfx942") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp index 57743b8ced42..821135ff3dbe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp @@ -495,6 +495,10 @@ int rocke_unified_attn_select_2d_num_warps(const rocke_unified_attn_problem_t* p /* gfx942 D64 oracle. */ if(arch_is("gfx942") && p->head_size == 64) { + if(p->max_seqlen_q == 1) + { + return 1; + } return 4; } if(enable_combo_2d(p)) From baab3ab763fe8bcc532bda8f5bf0e7e51165a767 Mon Sep 17 00:00:00 2001 From: arai713 Date: Wed, 29 Jul 2026 15:17:10 -0400 Subject: [PATCH 4/4] fix(rocke): fixes for PR review - gate the gfx942 D64 decode num_warps=1 selector on `not use_fp8` in Python + C++ selector. - guard the gfx950 2D fp8-loader chunk-count assert behind KV_FP8, matching the gfx942 fix. - decode benchmark backend gating (gfx942 + gfx950): * pass sinks + sliding-window through to AITER Triton instead of skipping sink shapes. * skip aoTriton and FlyDSL for sink / sliding-window shapes. * derive window_size from sliding_window instead of hardcoding (-1,-1). - use dataclasses.replace for --batch/--dtype overrides so use_sinks and sliding_window are preserved. - add selector unit test and harden the fp8-loader-assert regression test. - results JSON records use_sinks/sliding_window as keys - DecodeShape signature distinguishes sink and windowed shapes. --- .../attention/decode/benchmark_decode_live.py | 44 +++++++------ .../attention/decode/benchmark_decode_live.py | 66 +++++++------------ .../kernels/common/attention_unified.py | 11 +++- .../kernels/gfx950/attention_tiled_2d.py | 11 ++-- .../library/tests/test_attention_builds.py | 61 ++++++++++++++++- .../common/attention_unified_selectors.cpp | 2 +- 6 files changed, 121 insertions(+), 74 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py index 2a0ebef2bd91..a95549409690 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py @@ -189,8 +189,6 @@ def _make_inputs( if use_qq_bias else None ) - # Attention sinks: one learned logit per query head, loaded in the kernel - # dtype (the kernel casts to f32 internally). gpt-oss decode only. sinks = ( torch.randn(shape.num_query_heads, dtype=dtype, device="cuda") * 0.1 if shape.use_sinks @@ -215,19 +213,10 @@ def _make_inputs( def _run_triton( shape: DecodeShape, data: dict, *, warmup: int, iters: int ) -> Optional[float]: - """Time AITER Triton unified_attention. Returns ms or None on failure. - - Sink shapes are skipped: AOTriton has no sink attention kernel, so the DSL - is the only path that computes them correctly. Running Triton with - ``sinks=None`` would time a different (wrong-math) kernel and report a - meaningless speedup, so we return None and leave the comparison blank. - """ + """Time AITER Triton unified_attention. Returns ms or None on failure.""" from rocke.runtime import synchronize_and_release, time_launches import torch - if shape.use_sinks: - return None - try: from aiter.ops.triton.attention.unified_attention import unified_attention as tri # type: ignore except ImportError: @@ -235,6 +224,7 @@ def _run_triton( hip_stream = _bench_stream_handle() out = torch.empty_like(data["q"]) + window_size = (shape.sliding_window - 1, 0) if shape.sliding_window else (-1, -1) def call_once(): tri( @@ -248,7 +238,7 @@ def call_once(): max_seqlen_k=shape.seqlen_k, softmax_scale=data["scale"], causal=True, - window_size=(-1, -1), + window_size=window_size, block_table=data["block_table"], softcap=data["softcap"], q_descale=None, @@ -256,7 +246,7 @@ def call_once(): v_descale=None, alibi_slopes=data["alibi_slopes"], qq_bias=data["qq_bias"], - sinks=None, + sinks=data["sinks"], ) try: @@ -361,9 +351,9 @@ def _run_aoTriton( """Time AOTriton flash SDPA for decode. Returns ms or None on failure. Reconstructs dense [B, H, S_k, D] KV from the paged cache outside the - timed region. Returns None (skipped) when any bias operand is active - (softcap, alibi_slopes, or qq_bias) because - scaled_dot_product_attention does not support them. + timed region. Returns None (skipped) when any operand it cannot express + is active (softcap, alibi_slopes, qq_bias, attention sinks, or a sliding + window) because scaled_dot_product_attention does not support them. """ import torch from torch.nn.attention import SDPBackend, sdpa_kernel @@ -398,10 +388,16 @@ def _run_aoTriton( alibi_slopes = data["alibi_slopes"] softcap = data["softcap"] - # scaled_dot_product_attention does not support alibi, qq_bias or - # softcap - use_bias = alibi_slopes is not None or data["qq_bias"] is not None or softcap - if use_bias: + # scaled_dot_product_attention does not support alibi, qq_bias, + # softcap, attention sinks, or a sliding window. + unsupported = ( + alibi_slopes is not None + or data["qq_bias"] is not None + or softcap + or shape.use_sinks + or shape.sliding_window > 0 + ) + if unsupported: return None # Probe eligibility. @@ -500,6 +496,10 @@ def _flydsl_supported(shape: DecodeShape) -> tuple[bool, str]: return False, f"head_size={shape.head_size} (need 64 or 128)" if shape.dtype not in ("bf16", "fp16"): return False, f"dtype={shape.dtype} (need bf16 or fp16)" + if shape.use_sinks: + return False, "sinks unsupported" + if shape.sliding_window: + return False, "sliding_window unsupported" return True, "" @@ -794,6 +794,8 @@ def main() -> int: "head_size": shape.head_size, "block_size": shape.block_size, "dtype": shape.dtype, + "use_sinks": shape.use_sinks, + "sliding_window": shape.sliding_window, "triton_ms": tri_ms, "aoTriton_ms": aot_ms, "dsl": {str(sms): dsl_results[sms] for sms in args.num_sms_sweep}, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py index 64cea2b02216..3000a1af2b5d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py @@ -31,7 +31,7 @@ import math import sys import traceback -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Dict, List, Optional @@ -211,9 +211,6 @@ def _run_triton( from rocke.runtime import synchronize_and_release, time_launches import torch - if shape.use_sinks: - return None - try: from aiter.ops.triton.attention.unified_attention import unified_attention as tri # type: ignore except ImportError: @@ -221,6 +218,7 @@ def _run_triton( hip_stream = _bench_stream_handle() out = torch.empty_like(data["q"]) + window_size = (shape.sliding_window - 1, 0) if shape.sliding_window else (-1, -1) def call_once(): tri( @@ -234,7 +232,7 @@ def call_once(): max_seqlen_k=shape.seqlen_k, softmax_scale=data["scale"], causal=True, - window_size=(-1, -1), + window_size=window_size, block_table=data["block_table"], softcap=data["softcap"], q_descale=None, @@ -242,7 +240,7 @@ def call_once(): v_descale=None, alibi_slopes=data["alibi_slopes"], qq_bias=data["qq_bias"], - sinks=None, + sinks=data["sinks"], ) try: @@ -347,9 +345,9 @@ def _run_aoTriton( """Time AOTriton flash SDPA for decode. Returns ms or None on failure. Reconstructs dense [B, H, S_k, D] KV from the paged cache outside the - timed region. Returns None (skipped) when any bias operand is active - (softcap, alibi_slopes, or qq_bias) because - scaled_dot_product_attention does not support them. + timed region. Returns None (skipped) when any operand it cannot express + is active (softcap, alibi_slopes, qq_bias, attention sinks, or a sliding + window) because scaled_dot_product_attention does not support them. """ import torch from torch.nn.attention import SDPBackend, sdpa_kernel @@ -384,10 +382,16 @@ def _run_aoTriton( alibi_slopes = data["alibi_slopes"] softcap = data["softcap"] - # scaled_dot_product_attention does not support alibi, qq_bias or - # softcap - use_bias = alibi_slopes is not None or data["qq_bias"] is not None or softcap - if use_bias: + # scaled_dot_product_attention does not support alibi, qq_bias, + # softcap, attention sinks, or a sliding window. + unsupported = ( + alibi_slopes is not None + or data["qq_bias"] is not None + or softcap + or shape.use_sinks + or shape.sliding_window > 0 + ) + if unsupported: return None # Probe eligibility. @@ -482,6 +486,10 @@ def _flydsl_supported(shape: DecodeShape) -> tuple[bool, str]: return False, f"head_size={shape.head_size} (need 64 or 128)" if shape.dtype not in ("bf16", "fp16"): return False, f"dtype={shape.dtype} (need bf16 or fp16)" + if shape.use_sinks: + return False, "sinks unsupported" + if shape.sliding_window: + return False, "sliding_window unsupported" return True, "" @@ -656,35 +664,9 @@ def main() -> int: shapes = load_decode_shapes(args.shapes) if args.dtype is not None: - shapes = [ - DecodeShape( - batch=s.batch, - seqlen_q=s.seqlen_q, - seqlen_k=s.seqlen_k, - num_query_heads=s.num_query_heads, - num_kv_heads=s.num_kv_heads, - head_size=s.head_size, - block_size=s.block_size, - dtype=args.dtype, - label=s.label, - ) - for s in shapes - ] + shapes = [replace(s, dtype=args.dtype) for s in shapes] if args.batch is not None: - shapes = [ - DecodeShape( - batch=args.batch, - seqlen_q=s.seqlen_q, - seqlen_k=s.seqlen_k, - num_query_heads=s.num_query_heads, - num_kv_heads=s.num_kv_heads, - head_size=s.head_size, - block_size=s.block_size, - dtype=s.dtype, - label=s.label, - ) - for s in shapes - ] + shapes = [replace(s, batch=args.batch) for s in shapes] if args.limit is not None: shapes = shapes[: args.limit] @@ -817,6 +799,8 @@ def main() -> int: "head_size": shape.head_size, "block_size": shape.block_size, "dtype": shape.dtype, + "use_sinks": shape.use_sinks, + "sliding_window": shape.sliding_window, "triton_ms": tri_ms, "aoTriton_ms": aot_ms, "dsl": {str(sms): dsl_results[sms] for sms in args.num_sms_sweep}, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index f849f5a33c42..32e02406b954 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -990,10 +990,15 @@ def _select_2d_num_warps(problem: UnifiedAttentionProblem) -> int: # key, so a drift would silently launch the wrong CTA count, not rebuild). if _enable_gfx942_l4(problem): return _select_gfx942_flash_num_warps(problem) - # gfx942 D64 oracle: paired with ``T=block_size`` and mw=32, four waves fit - # in the MI300X 64 KB LDS budget and match the direct gfx942 harness. + # gfx942 D64 oracle. + # * prefill: num_warps=4 (BLOCK_M=128) + mw=32, four waves fit the MI300X + # 64 KB LDS budget and match the direct gfx942 harness. + # * decode (max_seqlen_q == 1, bf16): num_warps=1. Decode is memory-bound; + # nw=1 (BLOCK_M=32, T=2*block_size) measured markedly faster than nw=4 + # on gpt-oss sink decode. fp8 decode is excluded (dequant-bound, wants + # more warps to spread the fp8->bf16 dequant). if _resolve_attention_arch() == "gfx942" and problem.head_size == 64: - if problem.all_decode: + if problem.all_decode and not problem.use_fp8: return 1 return 4 if _enable_combo_2d(problem): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index e441a0b6dca4..634b25e88aa9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -2375,11 +2375,12 @@ def _dequant_fp8_lds_to_bf16( # store is one ``smem_store_vN(..., n=8)``. fp8_elems_per_chunk = 8 fp8_total_chunks = (T * HD) // fp8_elems_per_chunk - assert fp8_total_chunks % THREADS == 0, ( - f"fp8 loader: total chunks {fp8_total_chunks} must be divisible by " - f"THREADS={THREADS} (T={T}, HD={HD})" - ) - fp8_chunks_per_thread = fp8_total_chunks // THREADS + if KV_FP8: + assert fp8_total_chunks % THREADS == 0, ( + f"fp8 loader: total chunks {fp8_total_chunks} must be divisible by " + f"THREADS={THREADS} (T={T}, HD={HD})" + ) + fp8_chunks_per_thread = fp8_total_chunks // THREADS if KV_FP8 else 0 def _issue_fp8_dequant_loads( kv_tile_idx: Value, buf_idx: Value, lds_token: str diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py index de98a0f8209a..7a70e4fddd9b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py @@ -635,15 +635,24 @@ def test_gfx942_bf16_swa_decode_no_fp8_loader_assert(self): gfx942 2D built the fp8 chunk-count assert unconditionally. A bf16 SWA/decode small-tile config (T=16, HD=64, THREADS=256) gives - fp8_total_chunks=128; 128 % 256 != 0 raised AssertionError even though - the fp8 loader is never used for bf16 (gfx942 has no fp8 K/V support). - Guarding the assert with KV_FP8 (matching the 3D builder) fixes it. + fp8_total_chunks=128; 128 % 256 != 0 raised AssertionError. The fp8 + loader is never used for bf16, so the guard (``if KV_FP8:``, matching + the 3D builder) is load-bearing: gfx942 *does* support fp8 K/V via + ``kv_storage_dtype='fp8e4m3'`` (the 2D gate admits it), so the assert + must fire for fp8 and stay dormant for bf16. """ from kernels.gfx942.attention_tiled_2d import ( UnifiedAttention2DTiledSpec, build_unified_attention_2d_tiled, ) + # Precondition that makes this a regression test, not a smoke test: + # the config must be one the old unconditional assert rejected, i.e. + # (T*HD)//8 not divisible by THREADS=num_warps*64. T=16,HD=64,nw=4 -> + # 128 chunks, 256 threads -> 128 % 256 != 0. + T, HD, THREADS = 16, 64, 4 * 64 + self.assertNotEqual(((T * HD) // 8) % THREADS, 0) + spec = UnifiedAttention2DTiledSpec( head_size=64, block_size=16, @@ -660,6 +669,16 @@ def test_gfx942_bf16_swa_decode_no_fp8_loader_assert(self): k = build_unified_attention_2d_tiled(spec, arch="gfx942") ll = lower_kernel_to_llvm(k, arch="gfx942") self.assertIn("define amdgpu_kernel void", ll) + # The bf16 kernel must emit no fp8 dequant path. + self.assertNotIn("@llvm.amdgcn.cvt.f32.fp8", ll) + + # The guard must not neuter fp8 validation: the same geometry with + # fp8 K/V still trips the chunk-count assert (128 % 256 != 0). + import dataclasses + + fp8_spec = dataclasses.replace(spec, kv_storage_dtype="fp8e4m3") + with self.assertRaisesRegex(AssertionError, "fp8 loader"): + build_unified_attention_2d_tiled(fp8_spec, arch="gfx942") def test_unified_attention_2d_tiled_half_local_pv_compiles(self): """The R4 half-local PV variant emits 32x32 MFMA with its suffixes.""" @@ -1042,6 +1061,42 @@ def test_gfx942_l4_num_warps_matches_flash_selector(self): au._select_gfx942_flash_num_warps(p), ) + def test_gfx942_d64_decode_num_warps(self): + """gfx942 D64 decode picks num_warps=1; prefill keeps num_warps=4. + + Decode (max_seqlen_q == 1) is memory-bound and wins at nw=1 (BLOCK_M=32, + 4x the CTAs of nw=4). Prefill is compute-bound and stays at nw=4. fp8 + decode is excluded from the nw=1 lever (dequant-bound, wants more warps), + so it also stays at nw=4. This branch is the production geometry change; + pin it so a refactor can't silently revert it. The C++ selector mirrors + this exactly (see attention_unified_selectors.cpp); the run_all.py + byte-identity gate enforces the two agree. + """ + import kernels.common.attention_unified as au + + def _p(max_seqlen_q, use_fp8=False): + return UnifiedAttentionProblem( + total_q=max_seqlen_q, + num_seqs=1, + num_query_heads=64, + num_kv_heads=8, + head_size=64, + block_size=16, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=2048, + dtype="bf16", + use_sinks=True, + use_fp8=use_fp8, + ) + + with _patch_resolved_arch("gfx942"): + self.assertEqual(au._select_2d_num_warps(_p(1)), 1) # decode + self.assertEqual(au._select_2d_num_warps(_p(512)), 4) # prefill + self.assertEqual(au._select_2d_num_warps(_p(2048)), 4) # prefill + self.assertEqual( + au._select_2d_num_warps(_p(1, use_fp8=True)), 4 + ) # fp8 decode excluded + def test_tiled_3d_dispatch_gate_accepts_kwargs_per_arch(self): """Regression: the shared dispatch entry ``supports_native_unified_attention_3d_tiled`` forwards its kwargs to the diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp index 821135ff3dbe..41ac32de4473 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp @@ -495,7 +495,7 @@ int rocke_unified_attn_select_2d_num_warps(const rocke_unified_attn_problem_t* p /* gfx942 D64 oracle. */ if(arch_is("gfx942") && p->head_size == 64) { - if(p->max_seqlen_q == 1) + if(p->max_seqlen_q == 1 && !p->use_fp8) { return 1; }