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..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 @@ -58,14 +58,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]: @@ -100,6 +107,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 +189,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, @@ -192,6 +206,7 @@ def _make_inputs( softcap=softcap, alibi_slopes=alibi_slopes, qq_bias=qq_bias, + sinks=sinks, ) @@ -209,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( @@ -222,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, @@ -230,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: @@ -277,6 +293,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 +314,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 +331,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, @@ -330,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 @@ -367,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. @@ -469,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, "" @@ -763,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/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/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..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 @@ -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, ) @@ -203,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( @@ -216,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, @@ -224,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: @@ -271,6 +287,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 +308,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 +325,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, @@ -324,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 @@ -361,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. @@ -459,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, "" @@ -633,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] @@ -794,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/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/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 77dda48f3bd8..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,9 +990,16 @@ 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 and not problem.use_fp8: + return 1 return 4 if _enable_combo_2d(problem): # bf16 sliding-window is prelude-bound -> nw2 (lighter prelude). But 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..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 @@ -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/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 11457100b20e..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 @@ -630,6 +630,56 @@ 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. 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, + 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) + # 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.""" from kernels import ( @@ -1011,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 57743b8ced42..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,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 && !p->use_fp8) + { + return 1; + } return 4; } if(enable_combo_2d(p))