diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py index a6c5005ca43e..c1e2f6fcb01f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py @@ -46,7 +46,7 @@ import sys import traceback from pathlib import Path -from typing import Any +from typing import Any, Optional from rocke.assets import shape_utils_dir @@ -381,7 +381,7 @@ def _variant_flags(name: str, *, sliding_window: int, dtype: str, is_fp8: bool) class CkVariantBench: - def __init__(self, *, compile_backend: str = "llvm", num_sms: int = 256): + def __init__(self, *, compile_backend: Optional[str] = None, num_sms: int = 256): self.compile_backend = compile_backend self.num_sms = num_sms self._launchers: dict[tuple, Any] = {} @@ -416,7 +416,10 @@ def build(self, shape, variant: str, sliding_window: int, is_fp8: bool): build_unified_attention_2d_tiled, supports_tiled_2d, ) - from kernels.common.attention_unified import _attn_signature + from kernels.common.attention_unified import ( + _attn_signature, + _select_2d_compile_backend, + ) from rocke.runtime import KernelLauncher dtype = "bf16" if shape.q_dtype == "torch.bfloat16" else "fp16" @@ -477,7 +480,13 @@ def build(self, shape, variant: str, sliding_window: int, is_fp8: bool): key = (shape.signature, variant, spec.kernel_name(), self.compile_backend) if key not in self._launchers: kernel = build_unified_attention_2d_tiled(spec) - artifact = compile_kernel(kernel, capture_ir_text=False) + backend = self.compile_backend or _select_2d_compile_backend(problem) + if backend == "hipcc": + from rocke.helpers.compile import compile_kernel_via_hipcc + + artifact = compile_kernel_via_hipcc(kernel) + else: + artifact = compile_kernel(kernel, capture_ir_text=False) self._launchers[key] = ( KernelLauncher( hsaco=artifact.hsaco, @@ -718,6 +727,16 @@ def main() -> int: # cap=8192 vs 1.11x at cap=65536. ap.add_argument("--cap-blocks", type=int, default=65536) ap.add_argument("--num-sms", type=int, default=256) + ap.add_argument( + "--compile-backend", + choices=("auto", "llvm", "hipcc"), + default="auto", + help="Compile backend for all CK DSL lanes (prod + build() variants). " + "'auto' (default) matches production dispatch " + "(_select_2d_compile_backend); 'llvm'/'hipcc' force a backend for A/B. " + "Forcing 'llvm' mismeasures cohorts whose production backend is hipcc " + "(e.g. D256 large prefill) by ~3x.", + ) ap.add_argument("--tol", type=float, default=5e-2) ap.add_argument("--shape-utils-path", type=Path, default=DEFAULT_SHAPE_UTILS) ap.add_argument( @@ -753,7 +772,8 @@ def main() -> int: print(f"device: {torch.cuda.get_device_name(0)}") print(f"shapes: {len(shapes)} variants: {args.variants} flydsl={args.flydsl}") - bench = CkVariantBench(num_sms=args.num_sms) + _cb = None if args.compile_backend == "auto" else args.compile_backend + bench = CkVariantBench(num_sms=args.num_sms, compile_backend=_cb) results = [] n_fly_supported = 0 n_fly_correct = 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json new file mode 100644 index 000000000000..5a30aaabe1a3 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json @@ -0,0 +1,3 @@ +{"call_idx": 0, "kind": "2d", "ALL_DECODE": false, "num_seqs": 1, "num_query_heads": 16, "num_kv_heads": 2, "num_queries_per_kv": 8, "head_size": 256, "block_size": 16, "max_seqlen_q": 2048, "max_seqlen_k": 2048, "softmax_scale": 0.0625, "softcap": 0.0, "has_sinks": false, "has_alibi": false, "has_output_scale": false, "window_size": [-1, -1], "q_dtype": "torch.bfloat16", "k_dtype": "torch.bfloat16", "v_dtype": "torch.bfloat16", "out_dtype": "torch.bfloat16", "q_shape": [2048, 16, 256], "k_shape": [128, 16, 2, 256], "v_shape": [128, 16, 2, 256], "out_shape": [2048, 16, 256], "block_table_shape": [1, 128], "cu_seqlens_q_shape": [2], "seqused_k_shape": [1]} +{"call_idx": 1, "kind": "2d", "ALL_DECODE": false, "num_seqs": 1, "num_query_heads": 16, "num_kv_heads": 2, "num_queries_per_kv": 8, "head_size": 256, "block_size": 16, "max_seqlen_q": 4096, "max_seqlen_k": 4096, "softmax_scale": 0.0625, "softcap": 0.0, "has_sinks": false, "has_alibi": false, "has_output_scale": false, "window_size": [-1, -1], "q_dtype": "torch.bfloat16", "k_dtype": "torch.bfloat16", "v_dtype": "torch.bfloat16", "out_dtype": "torch.bfloat16", "q_shape": [4096, 16, 256], "k_shape": [256, 16, 2, 256], "v_shape": [256, 16, 2, 256], "out_shape": [4096, 16, 256], "block_table_shape": [1, 256], "cu_seqlens_q_shape": [2], "seqused_k_shape": [1]} +{"call_idx": 2, "kind": "2d", "ALL_DECODE": false, "num_seqs": 1, "num_query_heads": 16, "num_kv_heads": 2, "num_queries_per_kv": 8, "head_size": 256, "block_size": 16, "max_seqlen_q": 8192, "max_seqlen_k": 8192, "softmax_scale": 0.0625, "softcap": 0.0, "has_sinks": false, "has_alibi": false, "has_output_scale": false, "window_size": [-1, -1], "q_dtype": "torch.bfloat16", "k_dtype": "torch.bfloat16", "v_dtype": "torch.bfloat16", "out_dtype": "torch.bfloat16", "q_shape": [8192, 16, 256], "k_shape": [512, 16, 2, 256], "v_shape": [512, 16, 2, 256], "out_shape": [8192, 16, 256], "block_table_shape": [1, 512], "cu_seqlens_q_shape": [2], "seqused_k_shape": [1]} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py index 2040286e3374..95620341a2da 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py @@ -14,7 +14,7 @@ """ from __future__ import annotations -from dataclasses import fields +from dataclasses import fields, replace from kernels.common.attention_unified import ( UnifiedAttentionProblem, @@ -61,6 +61,11 @@ _tiled_3d_impl, ) +# Imported as a module (not a bound symbol) so tests that +# ``mock.patch.object(attention_unified, "_d256_gfx950_fast", ...)`` still steer +# the builder's fast-route branch below (a bound import would freeze the ref). +from kernels.common import attention_unified as _kau + def _tiled_spec_from_problem( problem: UnifiedAttentionProblem, @@ -225,7 +230,7 @@ def _tiled_spec_from_problem( # no-fp8 preconditions so it can never fire on an incompatible spec. if "use_k_single_buffer" in _spec_field_names and _enable_k_single_buffer(problem): _gfx950_schedule_fields["use_k_single_buffer"] = True - return UnifiedAttention2DTiledSpec( + _spec = UnifiedAttention2DTiledSpec( head_size=problem.head_size, block_size=problem.block_size, num_query_heads=problem.num_query_heads, @@ -274,6 +279,11 @@ def _tiled_spec_from_problem( and not problem.use_fp8 and problem.num_query_heads == 64 and problem.num_kv_heads == 8 + # self-consistency: fast_paged_kv_desc requires T==64. Only enable it + # when the tile selector actually picks 64 for this shape, so the flag + # can never be set with an incompatible tile (which trips the spec + # validator). _select_2d_tile_size forces T=64 for this family. + and _select_2d_tile_size(problem) == 64 ), use_register_pv=_enable_register_pv(problem), use_fp8_mfma_qk=_enable_fp8_mfma_qk(problem), @@ -288,6 +298,17 @@ def _tiled_spec_from_problem( # below -- gfx942's spec class does not declare it.) **_gfx950_schedule_fields, ) + if _kau._d256_gfx950_fast(problem): + # D256 gfx950 bf16 prefill fast route. Authored here in the builder + # (was a post-build override in kernels.common ``_tiled_spec_from_problem``, + # so the winning spec is created in the builder, not + # baked into the dispatch layer). Geometry (num_warps / tile_size / + # block_m_per_warp) already comes from the gated selectors above; this + # pins the 32x32 transposed + FA3 softmax<->MFMA-interleave codegen + # constellation. The cohort is discriminated in ``_tiled_cache_key`` by + # ``_d256_gfx950_fast`` so the key stays faithful to the built kernel. + _spec = replace(_spec, **_kau._d256_gfx950_spec_overrides()) + return _spec def _tiled_3d_spec_from_problem( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/occupancy_probe.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/occupancy_probe.py new file mode 100644 index 000000000000..b48e8ce48617 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/occupancy_probe.py @@ -0,0 +1,46 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Shared static occupancy probe for the attention parity/bench harnesses. + +Wraps the DSL static occupancy probe (``dsl_probes/probe_occupancy.py``, +``llvm-readelf`` on the HSACO notes -- no kernel launch, safe during the build +step) so the per-arch harnesses consume one implementation instead of drifting +into near-identical copies. +""" +from __future__ import annotations + + +def print_occupancy(hsaco: bytes, num_warps: int, *, arch: str = "gfx950") -> None: + """Print VGPR/AGPR/SGPR/LDS + coarse waves-per-CU for a compiled kernel. + + Static (reads the HSACO notes); no launch, so it is safe to call during the + build step. ``arch`` selects the occupancy model (``gfx950`` / ``gfx942`` / + ``gfx90a``); unknown values fall back to gfx950. + """ + try: + import sys as _sys + + from rocke.assets import dsl_docs_dir + + probe_dir = ( + dsl_docs_dir() / "optimization" / "utilities" / "tools" / "dsl_probes" + ) + if str(probe_dir) not in _sys.path: + _sys.path.insert(0, str(probe_dir)) + import probe_occupancy as _po # type: ignore + + arch_caps = getattr(_po, f"ARCH_{arch.upper()}", _po.ARCH_GFX950) + notes = _po.parse_hsaco_notes(hsaco) + occ = _po.estimate_occupancy( + notes=notes, waves_per_wg=num_warps, arch=arch_caps + ) + print( + f" [occupancy] vgpr={notes.get('vgpr_count')} " + f"agpr={notes.get('agpr_count')} sgpr={notes.get('sgpr_count')} " + f"spill={notes.get('vgpr_spill_count')} " + f"lds={notes.get('lds_size')}B " + f"waves/CU={occ['waves_per_cu']} wg/CU={occ['wgs_per_cu']} " + f"limit={occ['limited_by']}" + ) + except Exception as e: # noqa: BLE001 + print(f" [occupancy] unavailable: {type(e).__name__}: {e}") diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py new file mode 100644 index 000000000000..e7362740cac8 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py @@ -0,0 +1,126 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Shared torch flash-attention reference timer for the parity/bench harnesses. + +Reconstructs the dense per-sequence causal Q/K/V from a paged KV cache (so the +numbers match the reference paged attention), then times torch's +``F.scaled_dot_product_attention`` (flash backend) on it. Kept arch-agnostic and +decoupled from any harness-local ``Scenario``/``data`` shape so the gfx942/gfx950 +harnesses can consume one implementation instead of near-identical copies. +""" +from __future__ import annotations + +from typing import Callable + + +def run_torch_flash( + *, + query, + key_cache, + value_cache, + block_tables, + query_lens, + kv_lens_list, + num_query_heads: int, + num_kv_heads: int, + head_size: int, + block_size: int, + scale: float, + warmup: int, + attempts: int, + bench_stream: int, + time_lane_ms: Callable, +): + """Time torch F.sdpa (flash) on the same shape as the paged reference. + + ``time_lane_ms(call_once, *, warmup, attempts, stream)`` and ``bench_stream`` + are injected by the caller so this stays free of harness-local timing infra. + Returns ``(out[total_q, H, D], ms, backend_name)``. + """ + import torch + import torch.nn.functional as F + from torch.nn.attention import sdpa_kernel, SDPBackend + + H = num_query_heads + Hkv = num_kv_heads + D = head_size + dev = query.device + nqpkv = H // Hkv + + # Gather dense per-seq tensors once (outside the timed region). + seqs = [] + start = 0 + for i, (ql, kvl) in enumerate(zip(query_lens, kv_lens_list)): + if ql <= 0: + continue + q_i = query[start : start + ql] # [ql, H, D] + start += ql + nblk = (kvl + block_size - 1) // block_size + blk = block_tables[i, :nblk] + k_i = key_cache[blk].reshape(-1, Hkv, D)[:kvl] # [kvl, Hkv, D] + v_i = value_cache[blk].reshape(-1, Hkv, D)[:kvl] + # [1, H, L, D] + qh = q_i.permute(1, 0, 2).unsqueeze(0) + kh = k_i.permute(1, 0, 2).unsqueeze(0) + vh = v_i.permute(1, 0, 2).unsqueeze(0) + if nqpkv > 1: + kh = kh.repeat_interleave(nqpkv, dim=1) + vh = vh.repeat_interleave(nqpkv, dim=1) + seqs.append((qh, kh, vh, ql, kvl)) + + # Every sequence was empty (all query_lens <= 0) -> nothing to time. Fail + # explicitly here rather than letting the seqs[0] backend probe below + # IndexError into a misleading "no working torch sdpa backend". + if not seqs: + raise ValueError( + "run_torch_flash: no sequence with query_len > 0 " + f"(query_lens={list(query_lens)}); nothing to time" + ) + + # Pick a working flash-first backend once. + backend = None + for be, name in ( + (SDPBackend.FLASH_ATTENTION, "flash"), + (SDPBackend.EFFICIENT_ATTENTION, "efficient"), + (SDPBackend.MATH, "math"), + ): + try: + with sdpa_kernel(be): + _ = F.scaled_dot_product_attention( + seqs[0][0][:, :, :64], + seqs[0][1][:, :, :64], + seqs[0][2][:, :, :64], + is_causal=True, + scale=scale, + ) + backend = (be, name) + break + except Exception: # noqa: BLE001 + continue + if backend is None: + raise NotImplementedError("no working torch sdpa backend") + + outs = [None] * len(seqs) + + def call_once(): + with sdpa_kernel(backend[0]): + for j, (qh, kh, vh, ql, kvl) in enumerate(seqs): + if ql == kvl: + o = F.scaled_dot_product_attention( + qh, kh, vh, is_causal=True, scale=scale + ) + else: + # query is a suffix of the kv stream (aiter convention) + mask = torch.ones(ql, kvl, device=dev, dtype=torch.bool) + mask = torch.triu(mask, diagonal=kvl - ql + 1).logical_not() + o = F.scaled_dot_product_attention( + qh, kh, vh, attn_mask=mask, scale=scale + ) + outs[j] = o + + ms = time_lane_ms(call_once, warmup=warmup, attempts=attempts, stream=bench_stream) + + # Assemble [total_q, H, D] in scenario order for the correctness compare. + flat = [o.squeeze(0).permute(1, 0, 2) for o in outs] # each [L, H, D] + out = torch.cat(flat, dim=0).to(query.dtype) + return out, ms, backend[1] diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index 006122c73c7c..169c88ccd4d5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -264,6 +264,38 @@ def default_scenarios() -> List[Scenario]: block_size=16, dtype=torch.float16, ), + # D256 gfx950 bf16 prefill fast path (Qwen3-Next-80B-A3B cohort, + # GQA 16/2). Exercises the _d256_gfx950_fast route (32x32 transposed + # + FA3-style softmax<->MFMA interleave) at the two measured points. + Scenario( + name="prefill_d256_gqa8_b16_sq4096", + seq_lens=[(4096, 4096)], + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + ), + Scenario( + name="prefill_d256_gqa8_b16_sq8192", + seq_lens=[(8192, 8192)], + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + ), + # bs=32 D256: takes the same fast route (T=64 forced -> 2 blocks/tile), + # a distinct multi-block ratio from the bs=16 (4/tile) validated points. + Scenario( + name="prefill_d256_gqa8_b32_sq4096", + seq_lens=[(4096, 4096)], + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=32, + dtype=torch.bfloat16, + ), Scenario( name="prefill_d128_b16", seq_lens=[(64, 64), (128, 256), (32, 256)], @@ -898,6 +930,104 @@ def creative_scenarios() -> List[Scenario]: block_size=16, dtype=torch.bfloat16, ), + # --- PR #9233: D256 gfx950 bf16 prefill fast-path cohort (Sq bs64) --- + # These are the shapes the PR perf table reports (+9.1% @ Sq4096, + # +10.1% @ Sq8192 over the 32x32 base). bf16, head_size=256, causal + # prefill (query_len == kv_len), batch 64. num_blocks covers 64 seqs * + # (Sq / block_size) unique blocks so the paged block-table is dense. + Scenario( + name="pr9233_d256_bf16_prefill_sq4096_bs64", + seq_lens=[(4096, 4096)] * 64, + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + num_blocks=64 * (4096 // 16), + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_bs64", + seq_lens=[(8192, 8192)] * 64, + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + num_blocks=64 * (8192 // 16), + ), + # bs64 == block_size=64 interpretation, single long sequence (batch 1): + # a low-occupancy compute-bound prefill where the exposed softmax v_exp + # is the bottleneck the interleave targets. + Scenario( + name="pr9233_d256_bf16_prefill_sq4096_b1_blk64", + seq_lens=[(4096, 4096)], + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=4096 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_blk64", + seq_lens=[(8192, 8192)], + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=8192 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_h8", + seq_lens=[(8192, 8192)], + num_query_heads=8, + num_kv_heads=8, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=8192 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_gqa16x2", + seq_lens=[(8192, 8192)], + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=8192 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_h1", + seq_lens=[(8192, 8192)], + num_query_heads=1, + num_kv_heads=1, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=8192 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq4096_b1_blk16", + seq_lens=[(4096, 4096)], + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + num_blocks=4096 // 16, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_blk16", + seq_lens=[(8192, 8192)], + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + num_blocks=8192 // 16, + ), Scenario( name="creative_d256_prefill", seq_lens=[(128, 1024), (256, 2048)], @@ -1166,7 +1296,40 @@ def call_once(): _force_triton_path("auto") -def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): +def _with_kq_xor_swizzle(spec): + """Return a harness-only ProbeSpec with the measurement-only KQ XOR swizzle ON. + + The XOR swizzle produces INTENTIONALLY WRONG numerics (a read-only + bank-conflict probe), so the flag lives ONLY on this harness-local subclass -- + never on the production ``UnifiedAttention2DTiledSpec`` -- and the kernel emit + reads it via ``getattr(spec, "use_kq_xor_swizzle", False)``. + """ + import dataclasses + + from kernels import UnifiedAttention2DTiledSpec + + @dataclasses.dataclass(frozen=True) + class _ProbeTiledSpec(UnifiedAttention2DTiledSpec): + use_kq_xor_swizzle: bool = False + + fields = { + f.name: getattr(spec, f.name) + for f in dataclasses.fields(UnifiedAttention2DTiledSpec) + } + return _ProbeTiledSpec(**fields, use_kq_xor_swizzle=True) + + +def _run_rocke( + s: Scenario, + data, + *, + path: str, + warmup: int, + attempts: int, + kq_swizzle: bool = False, + kq_pad: int = 0, + probe_occupancy: bool = False, +): """Run CK DSL `run_unified_attention_torch` with the requested path forced. Both backends share the default bench stream (torch's current @@ -1211,16 +1374,15 @@ def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): hip_stream = _bench_stream_handle() - # NOTE: the ``"2d"`` lane force-builds the hand-tuned MFMA-32x32 / - # half-local-PV 2D kernel directly so we can measure the *best* 2D - # variant for a shape (this is the aspirational ceiling for the 2D - # path, including the d64/b32/h64kv8 trace family). The ``"auto"`` - # lane must instead exercise the *production* dispatcher + # NOTE: the ``"2d"`` lane builds and launches the PRODUCTION 2D spec + # (``_tiled_spec_from_problem``) directly -- GPU-time only, no + # ``run_unified_attention_torch`` dispatcher overhead -- so it measures the + # exact kernel production selects for a shape. The ``"auto"`` lane instead + # exercises the *production* dispatcher # (``run_unified_attention_torch(backend="auto")`` -> ``select_path``), - # otherwise we'd be reporting forced-2D timings as if they were what - # production launches -- which mis-measures decode shapes by ~6x - # (production correctly routes them to the 3D split-KV path). Keep the - # two lanes strictly separate. + # otherwise we'd report forced-2D timings as if they were what production + # launches -- which mis-measures decode shapes by ~6x (production correctly + # routes them to the 3D split-KV path). Keep the two lanes strictly separate. if path == "2d": from rocke import compile_kernel from kernels import ( @@ -1232,9 +1394,87 @@ def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): _attn_signature, _attn_values, _select_2d_compile_backend, + _d256_gfx950_fast, + _tiled_spec_from_problem, ) from rocke.runtime import KernelLauncher, LaunchConfig + # PR #9233 direct-launch lane: for the D256 gfx950 bf16 prefill cohort, + # build and launch the PRODUCTION fast spec directly (GPU-time only, no + # run_unified_attention_torch dispatcher overhead). --no-interleave + # (applied in main via a scoped patch of _tiled_spec_from_problem) makes + # this lane measure both the interleave fast path and the "32x32 base". + if _d256_gfx950_fast(problem): + spec = _tiled_spec_from_problem(problem) + if kq_swizzle: + spec = _with_kq_xor_swizzle(spec) + print( + " [2d-direct] KQ XOR swizzle ON (read-only probe; " + "numerics intentionally wrong)" + ) + if kq_pad > 0: + from dataclasses import replace as _replace + + spec = _replace(spec, use_kq_lds_pad=True, kq_lds_pad_halves=kq_pad) + print( + f" [2d-direct] KQ slab-gran pad ON (pad={kq_pad} halves, " + "correct: numerics preserved)" + ) + kernel = build_unified_attention_2d_tiled(spec) + if _select_2d_compile_backend(problem) == "hipcc": + from rocke.helpers.compile import compile_kernel_via_hipcc + + artifact = compile_kernel_via_hipcc(kernel) + else: + artifact = compile_kernel(kernel, capture_ir_text=False) + launcher = KernelLauncher( + hsaco=artifact.hsaco, + kernel_name=artifact.kernel_name, + signature=_attn_signature( + dtype_str, include_bt_stride=True, include_qq_bias_stride=True + ), + cache_key=("d256_fast_direct", spec.kernel_name()), + ) + if probe_occupancy: + from builders.common.occupancy_probe import print_occupancy + + print_occupancy(artifact.hsaco, spec.num_warps, arch="gfx950") + vals = _attn_values( + problem=problem, + q=q, + k=data["key_cache"], + v=data["value_cache"], + out=output, + cu_seqlens_q=data["cu_q"], + seqused_k=data["kv_lens"], + softmax_scale=data["scale"], + block_table=data["block_tables"], + softcap=float(s.softcap), + sinks=data["sinks"], + bt_stride=int(data["block_tables"].stride(0)), + include_bt_stride=True, + alibi_slopes=data["alibi_slopes"], + qq_bias=qq_bias, + qq_bias_stride_0=qq_bias_stride_0, + include_qq_bias_stride=True, + ) + block_q = spec.block_q + total_num_q_blocks = q.shape[0] // block_q + len(s.seq_lens) + cfg = LaunchConfig( + grid=(int(s.num_kv_heads), int(total_num_q_blocks), 1), + block=(64 * spec.num_warps, 1, 1), + stream=hip_stream, + ) + print(f" [2d-direct] {spec.kernel_name()}") + + def call_once(): + launcher(vals, config=cfg) + + ms = _time_lane_ms( + call_once, warmup=warmup, attempts=attempts, stream=hip_stream + ) + return output, ms + ok, reason = supports_tiled_2d( head_size=s.head_size, block_size=s.block_size, @@ -1251,66 +1491,7 @@ def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): if not ok: raise NotImplementedError(reason) - def use_hlpv_variant() -> bool: - if dtype_str != "bf16": - return False - if s.head_size != 64 or s.block_size != 32: - return False - if s.num_query_heads != 64 or s.num_kv_heads != 8: - return False - if ( - s.softcap > 0 - or problem.use_fp8 - or problem.use_alibi - or problem.use_qq_bias - ): - return False - if data["max_query_len"] <= 256: - return False - # The half-local PV path regresses on the high-num-seq SW tail. - if (s.sliding_window or 0) > 0 and len(s.seq_lens) >= 450: - return False - return True - - use_hlpv = use_hlpv_variant() - # Measured best local policy for the d64/b32/h64kv8 bf16+sinks family: - # * no-SW: R4_s1mask_hlpv + mask-limit + fast paged-KV + skip legacy Q - # * SW: R4_s1mask_hlpv + fast paged-KV + skip legacy Q - # * SW high-num-seq tail falls back to plain R4 in use_hlpv_variant(). - use_transposed_mask_limit = use_hlpv and (s.sliding_window or 0) == 0 - use_mfma32_skip_legacy_qreg = use_hlpv - use_fast_paged_kv_desc = use_hlpv - # AGPR0 is still experimental and did not improve this path broadly; keep - # it as an explicit environment opt-in for microbench work only. - use_agpr_alloc_zero = ( - use_hlpv and os.environ.get("ROCKE_ATTENTION_AGPR_ALLOC_ZERO") == "1" - ) - spec = UnifiedAttention2DTiledSpec( - head_size=s.head_size, - block_size=s.block_size, - num_query_heads=s.num_query_heads, - num_kv_heads=s.num_kv_heads, - dtype=dtype_str, - use_sinks=data["sinks"] is not None, - sliding_window=s.sliding_window or 0, - has_softcap=s.softcap > 0, - use_alibi=data["alibi_slopes"] is not None, - use_qq_bias=qq_bias is not None, - num_seqs=len(s.seq_lens), - num_warps=4, - waves_per_eu=2, - tile_size=2 * s.block_size, - block_m_per_warp=32, - use_mfma_32x32=True, - use_transposed_qk_32x32=True, - use_transposed_scalar_state=use_hlpv, - use_transposed_mask_once=use_hlpv, - use_transposed_half_local_pv=use_hlpv, - use_mfma32_skip_legacy_qreg=use_mfma32_skip_legacy_qreg, - use_transposed_mask_limit=use_transposed_mask_limit, - use_fast_paged_kv_desc=use_fast_paged_kv_desc, - use_agpr_alloc_zero=use_agpr_alloc_zero, - ) + spec = _tiled_spec_from_problem(problem) kernel = build_unified_attention_2d_tiled(spec) if _select_2d_compile_backend(problem) == "hipcc": from rocke.helpers.compile import compile_kernel_via_hipcc @@ -1324,14 +1505,7 @@ def use_hlpv_variant() -> bool: signature=_attn_signature( dtype_str, include_bt_stride=True, include_qq_bias_stride=True ), - cache_key=( - "r4_hlpv_parity", - spec.kernel_name(), - use_agpr_alloc_zero, - use_mfma32_skip_legacy_qreg, - use_transposed_mask_limit, - use_fast_paged_kv_desc, - ), + cache_key=("prod_2d_direct", spec.kernel_name()), ) vals = _attn_values( problem=problem, @@ -1392,6 +1566,39 @@ def call_once(): return output, ms +def _run_torch_flash(s: Scenario, data, *, warmup: int, attempts: int): + """Time torch flash SDPA on the same shape (the "vs torch-flash" reference). + + Thin arch-local adapter over the shared + ``builders.common.torch_flash_reference.run_torch_flash`` (kept generic so + the gfx942/decode harnesses can reuse it). Returns (out[total_q, H, D], ms). + """ + from builders.common.torch_flash_reference import run_torch_flash + + out, ms, backend = run_torch_flash( + query=data["query"], + key_cache=data["key_cache"], + value_cache=data["value_cache"], + block_tables=data["block_tables"], + query_lens=data["query_lens"], + kv_lens_list=data["kv_lens_list"], + num_query_heads=s.num_query_heads, + num_kv_heads=s.num_kv_heads, + head_size=s.head_size, + block_size=s.block_size, + scale=data["scale"], + warmup=warmup, + attempts=attempts, + bench_stream=_bench_stream_handle(), + time_lane_ms=_time_lane_ms, + ) + _LAST_TORCH_FLASH["backend"] = backend + return out, ms + + +_LAST_TORCH_FLASH = {"backend": None} + + def run_unified(backend: str, s: Scenario, data, warmup: int = 3, attempts: int = 10): """Backwards-compatible wrapper: backend in {triton, rocke}; path auto.""" if backend == "triton": @@ -1500,12 +1707,57 @@ def _row_print(label: str, out_ms, ref_out, t_out): def main() -> int: + # Guarantee any --no-interleave patch of the production + # _tiled_spec_from_problem is restored, even on exceptions (no global leak). + import kernels.common.attention_unified as _au + + _orig_tiled_spec = _au._tiled_spec_from_problem + try: + return _main_impl() + finally: + _au._tiled_spec_from_problem = _orig_tiled_spec + + +def _main_impl() -> int: parser = argparse.ArgumentParser() parser.add_argument("--scenario", default=None, action="append") parser.add_argument("--attempts", type=int, default=10) parser.add_argument("--warmup", type=int, default=3) parser.add_argument("--report", type=Path, default=None) parser.add_argument("--skip-ck", action="store_true") + parser.add_argument( + "--no-interleave", + action="store_true", + help="build the D256 fast spec with softmax<->MFMA interleave OFF (the " + "'32x32 base'); run with/without to recover the +N%% interleave delta " + "(was PARITY_NO_INTERLEAVE)", + ) + parser.add_argument( + "--kq-swizzle", + action="store_true", + help="direct-launch 2D lane: KQ XOR swizzle ON (read-only probe; " + "numerics intentionally wrong; was PARITY_KQ_SWIZZLE)", + ) + parser.add_argument( + "--kq-pad", + type=int, + default=0, + metavar="HALVES", + help="direct-launch 2D lane: KQ slab-granularity LDS pad width in " + "halves (0=off, correct; was PARITY_KQ_PAD / PARITY_KQ_PAD_W)", + ) + parser.add_argument( + "--probe-occupancy", + action="store_true", + help="direct-launch 2D lane: print static VGPR/LDS/waves occupancy for " + "the built kernel (was PARITY_PROBE_OCC)", + ) + parser.add_argument( + "--torch-flash", + action="store_true", + help="also time torch F.sdpa (flash) and report the CK-auto vs " + "torch-flash ratio (PR #9233 'vs torch-flash' column)", + ) parser.add_argument( "--skip-triton", action="store_true", @@ -1540,6 +1792,37 @@ def main() -> int: return 1 print("device:", torch.cuda.get_device_name(0)) + # PR #9233 reproduction lever (--no-interleave): force the softmax<->MFMA + # interleave OFF on the production ("auto") dispatch AND the direct-launch + # lane so the D256 gfx950 fast path degrades to the "32x32 base" it is + # measured against. Run the harness twice (with and without) and take the + # latency ratio to recover the "+N% over 32x32 base" number. Only the + # interleave codegen flags are dropped; the gated geometry + # (num_warps/tile_size/block_m_per_warp) is unchanged. The patch is restored + # in main()'s finally (no global leak). + if args.no_interleave: + from dataclasses import replace as _replace + import kernels.common.attention_unified as _au + + _orig_tiled_spec = _au._tiled_spec_from_problem + + def _tiled_spec_no_interleave(problem): + spec = _orig_tiled_spec(problem) + if getattr(spec, "use_softmax_mfma_interleave", False): + spec = _replace( + spec, + use_softmax_mfma_interleave=False, + softmax_interleave_mode=0, + softmax_interleave_groups=1, + ) + return spec + + _au._tiled_spec_from_problem = _tiled_spec_no_interleave + print( + "[parity] --no-interleave: softmax<->MFMA interleave DISABLED " + "(auto + direct-launch -> 32x32 base)" + ) + if args.set == "default": scenarios = default_scenarios() elif args.set == "creative": @@ -1650,6 +1933,38 @@ def main() -> int: row["ck_auto_status"] = err_ck _row_print("ck-auto", ck_auto, ref_out, None) _isolate_benchmark_lane() + + if args.torch_flash: + tf, err_tf = _safe_run( + lambda: _run_torch_flash( + s, + data, + warmup=args.warmup, + attempts=args.attempts, + ) + ) + if tf: + tf_out, tf_ms = tf + row["torch_flash_ms"] = tf_ms + row["torch_flash_backend"] = _LAST_TORCH_FLASH["backend"] + tf_diff = compare(ref_out, tf_out) + row["torch_flash_vs_ref"] = tf_diff["max_abs"] + print( + f" {'torch-flash':14s}: {tf_ms * 1000:9.2f} us " + f"max_abs={tf_diff['max_abs']:.5f} " + f"backend={_LAST_TORCH_FLASH['backend']}" + ) + if row.get("ck_auto_ms"): + ratio = tf_ms / row["ck_auto_ms"] + row["ck_auto_vs_torch_flash"] = ratio + print( + f" {'ck-auto vs torch-flash':22s}: {ratio:.3f}x " + f"(ck {'faster' if ratio > 1 else 'slower'})" + ) + elif err_tf: + row["torch_flash_status"] = err_tf + print(f" torch-flash FAILED: {err_tf}") + _isolate_benchmark_lane() continue # Force-path: Triton on `path`, CK DSL on `path`. @@ -1680,6 +1995,9 @@ def main() -> int: path=p, warmup=args.warmup, attempts=args.attempts, + kq_swizzle=args.kq_swizzle, + kq_pad=args.kq_pad, + probe_occupancy=args.probe_occupancy, ) ) _row_print(f"ck-{path}", ck_p, ref_out, None) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py index de7717159937..8e600866f043 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py @@ -150,6 +150,11 @@ class AttentionSpec: num_query_heads: int num_kv_heads: int name: str = "rocke_attention_unified" + # Optional pinned codegen knobs a specialized candidate wants the builder to + # apply (sorted (key, value) pairs; empty for generic candidates). See + # ``attention_unified._d256_gfx950_spec_overrides``; the builder consumes them + # via ``_tiled_spec_from_problem(problem, overrides=...)``. + tiled_overrides: Tuple[Tuple[str, object], ...] = () def kernel_name(self) -> str: from rocke.helpers.spec import kernel_name_join @@ -301,6 +306,85 @@ def select(req: OperatorRequest) -> AttentionSpec: ATTENTION_REGISTRY.register(_make_gfx942_dense_pipe_candidate()) +def _make_gfx950_d256_candidate() -> KernelCandidate: + """Fast gfx950 bf16 head_size-256 prefill kernel — 32x32 transposed stack + with FA3-style softmax<->MFMA interleave (mode2/g4) + slab-padded K_lds. + + Registered at priority 5 so it outranks the generic unified_2d candidate + (priority 10) for the gfx950 bf16 D256 prefill cohort. The registry sorts + ascending (lower = higher precedence); gfx950-only, so it never competes + with the gfx942 dense_pipe candidate. Callers can also force this path + explicitly via algorithm="d256_gfx950". + + The cohort is the single source of truth + ``kernels.common.attention_unified._d256_gfx950_cohort`` — the same predicate + the orchestrator's ``_d256_gfx950_fast`` override uses — so dispatch selection + and the built spec cannot drift. Only the arch gate differs (request arch + here vs resolved device arch there). + """ + spec_id = "gfx950_d256" + name = "attention_gfx950_d256" + + def support(req: OperatorRequest) -> Tuple[bool, str]: + errors = _request_errors(req) + if errors: + return False, "; ".join(errors) + assert isinstance(req, AttentionRequest) + if req.arch != "gfx950": + return False, f"d256_gfx950 requires gfx950 (got {req.arch!r})" + if req.dtype != "bf16": + return False, f"d256_gfx950 is bf16-only (got {req.dtype!r})" + ok, why = _selector_matches(req, candidate) + if not ok: + return False, why + problem = _problem(req) + ok, why = supports_native_unified_attention(problem) + if not ok: + return False, why + if problem.select_path() != "2d": + return False, "problem routes to 3D, not 2D" + from kernels.common.attention_unified import _d256_gfx950_cohort + + if not _d256_gfx950_cohort(problem): + return False, "not the gfx950 bf16 D256 prefill fast-path cohort" + return True, "ok" + + def select(req: OperatorRequest) -> AttentionSpec: + ok, why = support(req) + if not ok: + raise ValueError(f"{name} does not support request: {why}") + assert isinstance(req, AttentionRequest) + problem = _problem(req) + from kernels.common.attention_unified import _d256_gfx950_spec_overrides + + return AttentionSpec( + path="2d", + head_size=problem.head_size, + block_size=problem.block_size, + dtype=problem.dtype, + num_query_heads=problem.num_query_heads, + num_kv_heads=problem.num_kv_heads, + name="rocke_attention_gfx950_d256", + tiled_overrides=tuple(sorted(_d256_gfx950_spec_overrides().items())), + ) + + candidate = KernelCandidate( + name=name, + family=_FAMILY, + algorithm="d256_gfx950", + spec_id=spec_id, + abi_version=ATTENTION_ABI_VERSION, + priority=5, + supports=support, + select_spec=select, + signature=lambda _spec: (), + grid=lambda spec, req: (0, 0, 0), + block=lambda spec: (0, 0, 0), + sweep_space=lambda req: (select(req),) if support(req)[0] else (), + ) + return candidate + + def _make_d256_decode_candidate() -> KernelCandidate: """D256 bf16 decode candidate for gfx950 and gfx942 — 3D split-KV path. @@ -370,6 +454,7 @@ def select(req: OperatorRequest) -> AttentionSpec: return candidate +ATTENTION_REGISTRY.register(_make_gfx950_d256_candidate()) ATTENTION_REGISTRY.register(_make_d256_decode_candidate()) 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 1f631e071fa7..b2d1896eb4a4 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 @@ -585,6 +585,78 @@ def _enable_k_single_buffer(problem: UnifiedAttentionProblem) -> bool: return _enable_d128_small_tile(problem) and problem.block_size >= 32 +def _d256_gfx950_cohort(problem: "UnifiedAttentionProblem") -> bool: + """Arch-agnostic cohort for the gfx950 bf16 D256 prefill fast path. + + The arch gate is applied by callers (``_d256_gfx950_fast`` uses the resolved + device arch; the ``dispatch.attention`` candidate uses the request arch), so + this stays CPU-pure and is the single source of truth for the cohort shape. + """ + return ( + problem.head_size == 256 + and problem.dtype == "bf16" + and not problem.use_fp8 + and problem.sliding_window == 0 + and problem.softcap == 0 + and not problem.use_sinks + and not problem.use_alibi + and not problem.use_qq_bias + and problem.max_seqlen_q > 1 + ) + + +def _d256_gfx950_fast(problem: "UnifiedAttentionProblem") -> bool: + """Route the D256 gfx950 bf16 prefill cohort to the 32x32 + transposed fast path + FA3-style softmax<->MFMA interleave. + + Gates the three launch-critical geometry selectors (num_warps / tile_size / + block_m_per_warp) so the builder spec, ``_tiled_cache_key`` and + ``_get_2d_launch_meta`` all agree; the remaining codegen flags (32x32 MFMA, + transposed stack, register_pv off, K single-buffer, interleave mode2/g4) are + applied by the spec override in ``_tiled_spec_from_problem``. Narrow + + arch/feature-guarded: every other shape is byte-identical. + + Perf (MI355X, direct-launch sweep, fp32-ref max_abs=1.5625e-02): +interleave + mode2 g4 gives +9.1% @ Sq4096 and +10.1% @ Sq8192 over the 32x32 base. + """ + return _resolve_attention_arch() == "gfx950" and _d256_gfx950_cohort(problem) + + +def _d256_gfx950_spec_overrides() -> dict: + """Single source of truth for the D256 gfx950 bf16 prefill knob constellation. + + The 32x32 transposed + FA3-style softmax<->MFMA-interleave (mode2/g4) + + slab-padded K_lds codegen pins. Consumed by BOTH the builder override in + ``builders.common.attention_spec_builder._tiled_spec_from_problem`` (the + production path) and the ``dispatch.attention`` gfx950 D256 candidate, so the + dispatched spec and the built kernel cannot drift. + """ + return dict( + use_mfma_32x32=True, + use_transposed_qk_32x32=True, + use_q_direct_reg=True, + use_transposed_half_local_pv=True, + use_transposed_scalar_state=True, + use_transposed_mask_once=True, + use_transposed_mask_limit=True, + use_mask_phase_split=True, + use_register_pv=False, + use_k_single_buffer=True, + use_v_double_buffer=False, + use_early_v_schedule=False, + use_sched_barrier=False, + use_softmax_mfma_interleave=True, + softmax_interleave_mode=2, + softmax_interleave_groups=4, + use_fast_paged_kv_desc=False, + use_mfma32_skip_legacy_qreg=False, + # Slab-granularity K_lds pad (16 halves): breaks the row-aliased bank + # conflict on the QK K read for a ~25% Sq8192 latency win. + use_kq_lds_pad=True, + kq_lds_pad_halves=16, + ) + + def _enable_softmax_mfma_interleave(problem: UnifiedAttentionProblem) -> bool: """gfx950 single-batch d128 prefill: interleave the softmax VALU into the MFMA window via ``iglp_opt(1)`` and widen to num_warps=4. @@ -687,6 +759,8 @@ def _select_2d_tile_size(problem: UnifiedAttentionProblem) -> int: ``T = block_size`` (single block per iter, 32 tokens) — measured 1.15-1.30× win on every FP8 SW long-prefill shape. """ + if _d256_gfx950_fast(problem): + return 64 if _resolve_attention_arch() == "gfx1250": # gfx1250 v1 consumes exactly one 32-token paged-KV block per WMMA # iteration; wider T needs separate multi-block block-table handling. @@ -738,6 +812,17 @@ def _select_2d_tile_size(problem: UnifiedAttentionProblem) -> int: # GQA-8 d64 S2048 (1.19-2.68x over prod). if _enable_single_batch_combo(problem): if problem.head_size == 64: + # fast_paged_kv_desc (needs T=64) is a ~1.3-1.4x win for the bf16 + # h64kv8 no-SW family -> beats the T=128 early-V path (measured on + # MI355X). Force T=64 there so the combo's fast_paged enablement is + # valid; other d64 single-batch shapes keep the wider T=128 tile. + if ( + problem.dtype == "bf16" + and problem.num_query_heads == 64 + and problem.num_kv_heads == 8 + and problem.sliding_window == 0 + ): + return 64 return 128 # d128 occupancy lever (supersedes the small-tile pick for the # LONG-context holdout): at num_warps=2 the d128 combo is LDS-bound @@ -829,6 +914,8 @@ def _select_2d_num_warps(problem: UnifiedAttentionProblem) -> int: - an LDS-budget check (``<= 96 KiB`` so we keep >= 1 CTA/CU on MI355X comfortably). """ + if _d256_gfx950_fast(problem): + return 2 if _resolve_attention_arch() == "gfx1250": # A gfx1250 workgroup is one wave32 in the v1 WMMA tiled path. return 1 @@ -1077,6 +1164,13 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: _select_2d_tile_size(problem), _select_2d_waves_per_eu(problem), _select_2d_block_m_per_warp(problem), + # Layer-2 override marker: flips the key exactly when the D256 gfx950 + # codegen override applies in `_tiled_spec_from_problem`, so the key stays + # faithful to the built kernel (key === built kernel) without duplicating + # the override's flags here. The override is a hand-pinned exception that + # deliberately lives ABOVE the autotuner-regenerable `_enable_*` helpers + # below; keying on the predicate (not the flags) keeps that layering. + _d256_gfx950_fast(problem), _enable_mfma_32x32(problem), _enable_transposed_qk_32x32(problem), _enable_transposed_half_local_pv(problem), @@ -2150,12 +2244,16 @@ def _consider(label, cand, why): def _tiled_spec_from_problem( problem: UnifiedAttentionProblem, + *, + overrides=None, ): from builders.common.attention_spec_builder import ( _tiled_spec_from_problem as _impl, ) _spec = _impl(problem) + if overrides is not None: + _spec = replace(_spec, **overrides) return _resolve_lds_budget(_spec) @@ -2200,6 +2298,8 @@ def _select_2d_block_m_per_warp(problem: UnifiedAttentionProblem) -> int: Gate: ``max_seqlen_q > 256 and num_seqs >= 2``. Below this, mw=16 is consistently within noise of mw=32 in the per-shape sweep. """ + if _d256_gfx950_fast(problem): + return 32 if _resolve_attention_arch() == "gfx1250": return 16 # mw=32 (BLOCK_M = 32 * num_warps) only pays off when a path actually 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 cb724841ce94..e441a0b6dca4 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 @@ -342,6 +342,18 @@ class UnifiedAttention2DTiledSpec: # K-axis mostly lives in registers. This flag tracks that orientation # independently while it is brought up. use_transposed_qk_32x32: bool = False + # Slab-granularity K_lds pad: pad the K_lds allocation once per 2-row + # async-DMA slab (pad width ``kq_lds_pad_halves``) to break the row-aliased + # bank conflict on the QK K read (drops ``SQ_LDS_BANK_CONFLICT``). DMA-safe + # and numerics-preserving -- the async-DMA K write and the ``smem_load_vN`` + # read use the same padded addressing. The padded LDS stays within the + # 2-WG/CU budget (<= 80 KB/WG). Enabled by the D256 gfx950 fast route + # (``kq_lds_pad_halves=16``) for a ~25% Sq8192 latency win. + use_kq_lds_pad: bool = False + # Pad width in halves for the slab-granularity K_lds pad (must be a multiple + # of 8 to keep ds_read/ds_write b128 16-byte alignment). Only used when + # ``use_kq_lds_pad`` is set. + kq_lds_pad_halves: int = 8 # Transposed 32x32 softmax has one online-softmax state per query lane, # not one per output-dimension accumulator register. Keep a single m/l # loop-carried state and broadcast alpha across the 16 output regs. @@ -1098,6 +1110,29 @@ def build_unified_attention_2d_tiled( FP8_NATIVE_QK = False REGISTER_PV = spec.use_register_pv TRANSPOSED_QK_32X32 = spec.use_transposed_qk_32x32 + KQ_XOR_SWIZZLE = getattr(spec, "use_kq_xor_swizzle", False) + # Slab-granularity K_lds pad. The async DMA writes one WAVE_BYTES (=1024 B = + # SLAB_ROWS*HD*2) contiguous slab per wave per call, so padding BETWEEN slabs + # keeps the write contiguous while breaking the row-aliased bank conflict on + # the QK read. Correct (write+read consistent) -> numerics preserved. + KQ_LDS_PAD = int(spec.kq_lds_pad_halves) if spec.use_kq_lds_pad else 0 + # SLAB_ROWS = rows filled by one wave's 1024-B (WAVE=64 lanes x 16 B) slab. + KQ_SLAB_ROWS = (512 // HD) if KQ_LDS_PAD else 0 + # The slab pad's remap is only implemented for the non-fp8 K read (see the K + # read path below). The async-DMA write pads whenever KQ_LDS_PAD is set, but + # the fp8 K read does NOT remap -- an fp8+pad combo would address different + # slots -> silent numeric corruption. So disable the pad under fp8 (and under + # unsupported geometry) so write and read stay consistent. Byte-identical for + # the supported bf16 D256 path (FP8_MFMA_QK is False there). + if KQ_LDS_PAD and ( + FP8_MFMA_QK + or 512 % HD != 0 + or KQ_SLAB_ROWS == 0 + or T % KQ_SLAB_ROWS != 0 + or KQ_LDS_PAD % 8 != 0 + ): + KQ_LDS_PAD = 0 + KQ_SLAB_ROWS = 0 KV_BYTES = 1 if KV_FP8 else 2 kv_io_dtype = FP8E4M3 if KV_FP8 else dtype @@ -1358,7 +1393,21 @@ def build_unified_attention_2d_tiled( # with the ring (>2) and V-double-buffer paths in __post_init__. K_SINGLE_BUFFER = bool(spec.use_k_single_buffer) K_BUFS = 1 if K_SINGLE_BUFFER else (KV_RING_DEPTH if KV_RING_DEPTH > 2 else 2) - K_lds = b.smem_alloc(K_LDS_DTYPE, [K_BUFS, T, HD], name_hint="Klds") + if KQ_LDS_PAD and K_BUFS != 1: + # Slab pad only wired for the single-buffer K path; disable otherwise. + KQ_LDS_PAD = 0 + KQ_SLAB_ROWS = 0 + if KQ_LDS_PAD: + # Slab-granularity layout: [K_BUFS, T/SLAB_ROWS, SLAB_ROWS*HD + PAD]. + # Each DMA wave-slab (SLAB_ROWS contiguous rows) lands in one padded + # slab; the read remaps (row,col) -> (slab, row%SLAB_ROWS*HD + col). + _KQ_NSLAB = T // KQ_SLAB_ROWS + _KQ_SLAB_COLS = KQ_SLAB_ROWS * HD + KQ_LDS_PAD + K_lds = b.smem_alloc( + K_LDS_DTYPE, [K_BUFS, _KQ_NSLAB, _KQ_SLAB_COLS], name_hint="Klds" + ) + else: + K_lds = b.smem_alloc(K_LDS_DTYPE, [K_BUFS, T, HD], name_hint="Klds") # V double-buffer: 2 slots let V[i+1] prefetch into the alternate slot # during iter i (reusing ``cur_buf``), giving V the same 1-deep prefetch # K has. Costs only LDS (async DMA bypasses VGPR), occupancy-neutral. @@ -1863,6 +1912,26 @@ def _acc_idx(n: int, atom: int) -> int: wave_lds_offset_i32 = b.to_sgpr_u32(b.mul(wave_id, b.const_i32(WAVE_BYTES))) wave_lds_offset_i64 = b.zext(wave_lds_offset_i32, I64) + # K-specific slab-padded strides. One wave-slab == WAVE_BYTES contiguous B; + # inserting KQ_LDS_PAD*2 bytes between slabs shifts each slab's banks while + # leaving the intra-slab write contiguous (DMA-safe). V is untouched. + if KQ_LDS_PAD: + _KQ_PAD_BYTES = KQ_LDS_PAD * 2 + WAVE_BYTES_K = WAVE_BYTES + _KQ_PAD_BYTES + bytes_per_call_K = NUM_WARPS * WAVE_BYTES_K + bytes_per_buf_K = (T // KQ_SLAB_ROWS) * WAVE_BYTES_K + if NUM_WARPS == 1: + wave_lds_offset_K_i64 = b.const_i64(0) + else: + wave_lds_offset_K_i64 = b.zext( + b.to_sgpr_u32(b.mul(wave_id, b.const_i32(WAVE_BYTES_K))), I64 + ) + else: + WAVE_BYTES_K = WAVE_BYTES + bytes_per_call_K = bytes_per_call + bytes_per_buf_K = bytes_per_buf + wave_lds_offset_K_i64 = wave_lds_offset_i64 + # ---- Paged KV byte descriptor (full transform DAG) ---- # The paged-KV cache is laid out ``[num_blocks, BS, NUM_KV, HD]`` with # *byte* strides. The kernel addresses it via a chain of coordinate @@ -2023,10 +2092,10 @@ def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: offset (lanes 64..127 have `tid*8 / HD` advanced by T/NUM_WARPS), the cooperative load fills the full `[T, HD]` slab correctly. """ - buf_off_i32 = b.mul(buf_idx, b.const_i32(bytes_per_buf)) + buf_off_i32 = b.mul(buf_idx, b.const_i32(bytes_per_buf_K)) buf_off_i64 = b.zext(buf_off_i32, I64) K_buf_base = b.smem_ptr_add(K_lds_addr, buf_off_i64) - K_wave_base = b.smem_ptr_add(K_buf_base, wave_lds_offset_i64) + K_wave_base = b.smem_ptr_add(K_buf_base, wave_lds_offset_K_i64) if FAST_PAGED_KV_DESC: fast_block0, fast_block1 = _fast_paged_kv_blocks(kv_tile_idx) for call in range(kv_calls_per_tile): @@ -2059,7 +2128,7 @@ def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: linear_half=linear_half, kv_head=kv_head_idx, ) - k_dst = b.smem_ptr_add(K_wave_base, b.const_i64(call * bytes_per_call)) + k_dst = b.smem_ptr_add(K_wave_base, b.const_i64(call * bytes_per_call_K)) # CACHE_STREAM (SLC): one-shot streaming load, never re-read # within this kernel. Documented in # ``dsl_docs/primitives/intrinsics_and_primitives.md`` as the @@ -2625,6 +2694,24 @@ def _read_k8_mfma_operand(buf_idx: Value, k_row: Value, k_off: Value) -> Value: Otherwise read bf16 directly. Shared by the 16x16 and 32x32 QK paths so both get the fp8 LDS-footprint win. """ + if KQ_XOR_SWIZZLE: + # Measurement-only vanilla XOR swizzle: permute the read column by + # the low 3 bits of the row so 8 consecutive rows land on distinct + # banks (col ^= (row & 7) << 3, i.e. (row%8)*8 halves; stays 8-half + # aligned + in [0, HD)). Read-only -> numerically wrong; used solely + # to measure the SQ_LDS_BANK_CONFLICT delta. + k_off = b.xor(k_off, b.mul(b.mod(k_row, b.const_i32(8)), b.const_i32(8))) + if KQ_LDS_PAD and not K_FP8_MFMA: + # Slab-granularity padded layout: K_lds is + # [K_BUFS, T/SLAB_ROWS, SLAB_ROWS*HD + PAD]. Map the logical + # (row, col) to (slab, row_in_slab*HD + col) so the read hits the + # same padded slot the DMA wrote (write+read consistent). + slab = b.div(k_row, b.const_i32(KQ_SLAB_ROWS)) + col_in_slab = b.add( + b.mul(b.mod(k_row, b.const_i32(KQ_SLAB_ROWS)), b.const_i32(HD)), + k_off, + ) + return b.smem_load_vN(K_lds, buf_idx, slab, col_in_slab, dtype=dtype, n=8) if not K_FP8_MFMA: return b.smem_load_vN(K_lds, buf_idx, k_row, k_off, dtype=dtype, n=8) if FP8_NATIVE_QK: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py new file mode 100644 index 000000000000..1997d4f4a5cc --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py @@ -0,0 +1,156 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Routing tests for the gfx950 bf16 D256 prefill dispatch candidate. + +CPU-only: exercises the ``attention_gfx950_d256`` candidate's selection through +``dispatch_attention`` (filter -> priority -> pick). No GPU required. +""" + +from __future__ import annotations + +import unittest + +from dispatch.attention import AttentionRequest, _problem, dispatch_attention +from kernels.common import attention_unified as au + + +def _d256(arch="gfx950", **kw): + base = dict( + batch=2, + nhead_q=32, + nhead_k=8, + seqlen_q=4096, + seqlen_k=4096, + hdim_q=256, + hdim_v=256, + arch=arch, + dtype="bf16", + mask_type=1, # causal prefill + ) + base.update(kw) + return AttentionRequest(**base) + + +def _routed_spec_id(req): + """spec_id of the winning candidate, or None if nothing supports the req.""" + try: + return dispatch_attention(req).candidate.spec_id + except ValueError: + return None + + +class TestD256Gfx950Routing(unittest.TestCase): + def test_selects_gfx950_d256_for_cohort(self): + r = dispatch_attention(_d256()) + self.assertEqual(r.candidate.spec_id, "gfx950_d256") + self.assertEqual(r.spec.path, "2d") + + def test_outranks_generic_2d(self): + # priority 5 must beat the generic unified_2d (priority 10). + self.assertEqual(dispatch_attention(_d256()).candidate.priority, 5) + + def test_force_by_algorithm(self): + r = dispatch_attention(_d256(algorithm="d256_gfx950")) + self.assertEqual(r.candidate.spec_id, "gfx950_d256") + + def test_rejects_non_gfx950(self): + self.assertNotEqual(_routed_spec_id(_d256(arch="gfx942")), "gfx950_d256") + + def test_rejects_non_bf16(self): + self.assertNotEqual(_routed_spec_id(_d256(dtype="fp16")), "gfx950_d256") + + def test_rejects_non_d256(self): + self.assertNotEqual( + _routed_spec_id(_d256(hdim_q=128, hdim_v=128)), "gfx950_d256" + ) + + def test_rejects_decode(self): + # q=1 decode over long kv routes to the 3D split-KV path, not our cohort. + self.assertNotEqual( + _routed_spec_id( + _d256(batch=1, nhead_q=16, nhead_k=16, seqlen_q=1, seqlen_k=8192) + ), + "gfx950_d256", + ) + + def test_rejects_sliding_window(self): + self.assertNotEqual(_routed_spec_id(_d256(sliding_window=256)), "gfx950_d256") + + +class TestD256Gfx950SpecOverrides(unittest.TestCase): + """MVP: the dispatcher declares the knobs; the orchestrator accepts them.""" + + EXPECTED_KEYS = { + "use_mfma_32x32", + "use_transposed_qk_32x32", + "use_q_direct_reg", + "use_transposed_half_local_pv", + "use_transposed_scalar_state", + "use_transposed_mask_once", + "use_transposed_mask_limit", + "use_mask_phase_split", + "use_register_pv", + "use_k_single_buffer", + "use_v_double_buffer", + "use_early_v_schedule", + "use_sched_barrier", + "use_softmax_mfma_interleave", + "softmax_interleave_mode", + "softmax_interleave_groups", + "use_fast_paged_kv_desc", + "use_mfma32_skip_legacy_qreg", + "use_kq_lds_pad", + "kq_lds_pad_halves", + } + + def setUp(self): + self._saved = getattr(au, "_RESOLVED_ATTENTION_ARCH", None) + au._RESOLVED_ATTENTION_ARCH = "gfx950" + + def tearDown(self): + au._RESOLVED_ATTENTION_ARCH = self._saved + + def test_ssot_constellation(self): + ov = au._d256_gfx950_spec_overrides() + self.assertEqual(set(ov), self.EXPECTED_KEYS) + self.assertIs(ov["use_softmax_mfma_interleave"], True) + self.assertEqual(ov["softmax_interleave_mode"], 2) + self.assertIs(ov["use_register_pv"], False) + + def test_dispatched_spec_carries_overrides(self): + # the candidate declares the knobs on the dispatched spec (via the SSOT). + spec = dispatch_attention(_d256()).spec + self.assertEqual( + spec.tiled_overrides, + tuple(sorted(au._d256_gfx950_spec_overrides().items())), + ) + + def test_generic_candidate_has_no_overrides(self): + # a non-cohort shape (hd128) routes to the generic candidate -> empty. + spec = dispatch_attention(_d256(hdim_q=128, hdim_v=128)).spec + self.assertEqual(spec.tiled_overrides, ()) + + def test_orchestrator_accepts_overrides(self): + problem = _problem(_d256()) + # no-arg path: _impl already applies the D256 override, and + # _resolve_lds_budget still runs (a valid spec is returned). + base = au._tiled_spec_from_problem(problem) + self.assertIs(base.use_softmax_mfma_interleave, True) + self.assertIs(base.use_register_pv, False) + # passing the SSOT overrides explicitly reproduces the same valid spec + # (proves overrides= is plumbed + applied, and the resolver still runs). + ov = au._tiled_spec_from_problem( + problem, overrides=au._d256_gfx950_spec_overrides() + ) + self.assertEqual(ov, base) + + def test_overrides_reach_the_spec(self): + # an override that violates a spec invariant is raised by the spec's + # __post_init__ -> proves overrides= is really applied, not ignored. + problem = _problem(_d256()) + with self.assertRaises(ValueError): + au._tiled_spec_from_problem(problem, overrides={"use_register_pv": True}) + + +if __name__ == "__main__": + unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/d256_attention_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/d256_attention_ir_sha256.json new file mode 100644 index 000000000000..c3beaa509869 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/d256_attention_ir_sha256.json @@ -0,0 +1,25 @@ +{ + "flavors": { + "llvm20": { + "cases": { + "d256_pad_interleave/gfx950": { + "arch": "gfx950", + "bytes": 609602, + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "4650f01f089d4cd21714515984b697f6dd307bf506c7993ede6a04286f50a14f" + } + } + }, + "llvm22": { + "cases": { + "d256_pad_interleave/gfx950": { + "arch": "gfx950", + "bytes": 609706, + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "50a532b844ea16f58bdc3bbc5b8bc53434642fa4bfeb4355f47b8c4f58e093c4" + } + } + } + }, + "schema": "rocke.d256_attention_ir_sha256/v1" +} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py new file mode 100755 index 000000000000..141d9ac9739e --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python3 +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Library-side IR golden for the gfx950 D256 bf16 prefill *fast* spec. + +Pins the emitted LLVM IR (sha256, both ``llvm20``/``llvm22`` flavors) of the +D256 fast spec: 32x32-transposed + ``use_kq_lds_pad`` (slab-granularity K_lds +pad) + ``use_softmax_mfma_interleave``. The slab pad's async-DMA write and its +padded read remap must agree byte-for-byte or numerics silently corrupt, but +only the spec dataclass is otherwise asserted -- an addressing drift would pass +every CPU test and fail only on GPU (#9233 review). + +Lives in the LIBRARY test layer (not ``platform/tests``) because it builds a +library kernel (``kernels`` / ``builders``); platform tests must not import +library (one-way ``platform <- library`` dependency). It imports only the +pure-Python lowerer from ``rocke.core.lower_llvm`` (``library -> platform`` is +the allowed direction) with an EXPLICIT llvm flavor, so the recorded sha is a +function of the committed IR alone -- reproducible in CI on any host, GPU or +not (no ``rocke_engine`` / comgr / real device needed). + +Bless (only from a verified-good tree, alongside a reviewed IR change):: + + python3 library/tests/test_d256_attention_ir_golden.py --write +""" +from __future__ import annotations + +import hashlib +import json +import unittest +from pathlib import Path + +GOLDEN = Path(__file__).resolve().parent / "golden" / "d256_attention_ir_sha256.json" +GOLDEN_FLAVORS = ("llvm20", "llvm22") +GOLDEN_SCHEMA = "rocke.d256_attention_ir_sha256/v1" + + +def _sha(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def _current_flavor() -> str: + """The llvm flavor this host autodetects (the golden stores both; the gate + compares only this one, so one committed golden is valid on every host).""" + from rocke.core.lower_llvm import _resolve_llvm_flavor + + return _resolve_llvm_flavor() + + +def _build_d256_fast(arch: str): + """Build the *production* D256 fast-spec kernel under a pinned arch. + + ``_d256_gfx950_fast`` gates on the *memoized runtime device arch* via + ``_resolve_attention_arch`` -- NOT any per-case argument -- and the spec + builder holds a *bound* import of it. Both bindings are patched wholesale + (per its docstring: "Tests that monkeypatch this function replace it + wholesale") to ``arch`` so the fast spec is selected deterministically on + any host; otherwise a non-gfx950 host silently pins the FALLBACK spec (no + pad/interleave) and the golden becomes host-dependent. Raises if the fast + spec / pad+interleave is not selected, so the golden can never pin the + fallback by accident. + """ + import builders.common.attention_spec_builder as asb + import kernels.common.attention_unified as au + from kernels import build_unified_attention_2d_tiled + from kernels.common.attention_unified import ( + UnifiedAttentionProblem, + _d256_gfx950_fast, + _tiled_spec_from_problem, + ) + + # Validated D256 cohort point (GQA 16/2, hd256, bs16, sq4096 bf16). + problem = UnifiedAttentionProblem( + total_q=4096, + num_seqs=1, + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + max_seqlen_q=4096, + max_seqlen_k=4096, + dtype="bf16", + num_sms=120, + ) + pin = lambda: arch # noqa: E731 + o_au, o_asb = au._resolve_attention_arch, asb._resolve_attention_arch + au._resolve_attention_arch = pin + asb._resolve_attention_arch = pin + try: + if not _d256_gfx950_fast(problem): + raise RuntimeError( + f"D256 fast spec not selected under pinned arch {arch!r}; " + "golden would pin the fallback (no pad/interleave)" + ) + spec = _tiled_spec_from_problem(problem) + if not (spec.use_kq_lds_pad and spec.use_softmax_mfma_interleave): + raise RuntimeError( + "expected pad+interleave in the D256 fast spec; got " + f"pad={spec.use_kq_lds_pad} " + f"interleave={spec.use_softmax_mfma_interleave}" + ) + return build_unified_attention_2d_tiled(spec, arch=arch) + finally: + au._resolve_attention_arch = o_au + asb._resolve_attention_arch = o_asb + + +def _cases(): + """case_id -> (arch, build thunk). Add future D256 IR pins here.""" + return { + "d256_pad_interleave/gfx950": ("gfx950", lambda: _build_d256_fast("gfx950")), + } + + +def _lower(arch, build, flavor): + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + kernel = build() + llvm = _lower_kernel_to_llvm_python(kernel, arch=arch, llvm_flavor=flavor) + return { + "arch": arch, + "kernel_name": getattr(kernel, "name", ""), + "sha256": _sha(llvm), + "bytes": len(llvm.encode("utf-8")), + } + + +def _run(flavor): + return {cid: _lower(arch, build, flavor) for cid, (arch, build) in _cases().items()} + + +def build_golden() -> dict: + return { + "schema": GOLDEN_SCHEMA, + "flavors": {fl: {"cases": _run(fl)} for fl in GOLDEN_FLAVORS}, + } + + +class TestD256AttentionIrGolden(unittest.TestCase): + def test_ir_matches_golden(self): + """CPU-only: recompute the D256 fast-spec IR and compare its sha256 to + the blessed golden for this host's llvm flavor. Any drift in the + emitted addressing (slab-pad write/read, interleave) fails here.""" + self.assertTrue( + GOLDEN.is_file(), + f"missing golden {GOLDEN}; bless with " + "`python3 library/tests/test_d256_attention_ir_golden.py --write`", + ) + doc = json.loads(GOLDEN.read_text()) + flavor = _current_flavor() + base = doc.get("flavors", {}).get(flavor) + self.assertIsNotNone( + base, + f"golden has no entry for flavor {flavor!r} (have {sorted(doc.get('flavors', {}))})", + ) + cur = _run(flavor) + self.assertEqual( + set(base["cases"]), + set(cur), + "D256 IR golden case set changed; re-bless with --write", + ) + for cid, brec in base["cases"].items(): + crec = cur[cid] + self.assertEqual( + brec["sha256"], + crec["sha256"], + f"IR DRIFT [{flavor}] {cid}: {brec['sha256']} -> {crec['sha256']} " + f"(kernel {crec['kernel_name']}); re-bless only with a reviewed, " + "expected IR change", + ) + + +if __name__ == "__main__": + import sys + + if "--write" in sys.argv: + GOLDEN.parent.mkdir(parents=True, exist_ok=True) + GOLDEN.write_text(json.dumps(build_golden(), indent=2, sort_keys=True) + "\n") + summary = ", ".join(f"{fl}={len(_run(fl))} case(s)" for fl in GOLDEN_FLAVORS) + print(f"wrote {GOLDEN}: {summary}") + else: + unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py index 123ac52bbc22..b38a429a6a05 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py @@ -39,6 +39,17 @@ def _d256_bf16_long_prefill(block_size: int = 64, max_seqlen: int = 4096): class TestLdsBudgetResolver(unittest.TestCase): + def setUp(self): + # The production D256 gfx950 fast route (``_d256_gfx950_fast``) targets + # this exact cohort (D256 / gfx950 / bf16 prefill) and overrides + # ``use_register_pv=False``, which would bypass the resolver under test. + # Disable it here so these unit tests exercise the register-PV LDS-shrink + # path directly, decoupled from the production routing decision -- which + # is covered separately by ``TestD256ProductionRouting`` below. + _p = mock.patch.object(au, "_d256_gfx950_fast", return_value=False) + _p.start() + self.addCleanup(_p.stop) + def test_d256_gfx950_bf16_prefill_shrinks_to_fit(self): """D256 bf16 register-PV overflows at the default geometry; the resolver must single-buffer K so the resolved spec fits the gfx950 160 KB cap.""" @@ -141,5 +152,32 @@ def test_infeasible_budget_raises_with_lever_diagnostics(self): self.assertIn("T=64", msg) +class TestD256ProductionRouting(unittest.TestCase): + """Production routing (``_d256_gfx950_fast`` live, not mocked): the + D256 / gfx950 / bf16 prefill cohort is served by the 32x32 transposed fast + path, NOT the register-PV LDS-shrink resolver. This is the routing the PR + introduces; the resolver becomes a no-op for it (register-PV disabled).""" + + def test_d256_gfx950_prefill_routes_to_fast_path(self): + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + problem = _d256_bf16_long_prefill() + # The cohort matches the fast-route predicate ... + self.assertTrue(au._d256_gfx950_fast(problem)) + spec = au._tiled_spec_from_problem(problem) + # ... so the built spec is the 32x32 interleave fast path, not + # register-PV: register-PV is off and the 32x32 stack is on. + self.assertFalse(spec.use_register_pv) + self.assertTrue(spec.use_mfma_32x32) + self.assertTrue(spec.use_transposed_qk_32x32) + self.assertTrue(spec.use_softmax_mfma_interleave) + self.assertEqual(spec.softmax_interleave_mode, 2) + self.assertEqual(spec.softmax_interleave_groups, 4) + # The fast route single-buffers K itself, so the spec fits the cap + # and the resolver is a strict no-op (register-PV off -> same object). + self.assertTrue(spec.use_k_single_buffer) + self.assertIs(au._resolve_lds_budget(spec), spec) + self.assertLessEqual(au._lds_bytes_regpv(spec), au._lds_capacity_bytes()) + + if __name__ == "__main__": unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index 4619dcd9609d..261083b8044d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -46,6 +46,36 @@ For the specific shape you are optimizing: has told you which config to redesign *from* and against which resource budget. +> **Caveat — the shipped preset bench is NOT this sweep.** The live benchmark +> most people reach for +> (`library/benchmarks/gfx{942,950}/attention/prefill/benchmark_prefill2d_live.py`) +> sweeps a *curated menu of named presets* (`_variant_flags`: +> `prod`/`combo`/`fallback`/`r4_t32`/…), **not** the full lever space above. That +> grammar exposes geometry (`num_warps`, `tile_size`, `block_m_per_warp`, MFMA +> atom, transpose flags) but **cannot express several default-off flags** — +> notably `use_softmax_mfma_interleave` (+ `softmax_interleave_mode` / +> `softmax_interleave_groups`), `use_k_single_buffer`, and `use_register_pv`. So +> running the preset bench does **not** satisfy Step 0: it systematically skips +> exactly the default-off levers this step flags as most likely mis-picked. The +> space it measures is `union(preset menu, production-dispatcher output)`; +> anything outside that — including winning configs — stays invisible until a +> human adds it as a candidate. +> +> To actually exhaust the space, sweep the **raw flags**, bypassing both the +> preset grammar and the dispatcher: build the spec from arbitrary flags and +> launch it directly +> ([`utilities/tools/dsl_probes/probe_config_sweep.py`](utilities/tools/dsl_probes/probe_config_sweep.py), +> or the `_select_2d_*` monkey-patch pattern in [§4.3 Attention](#43-attention)), +> or drive the autotuner ([§12 Autotune](#12-autotuning-strategy)) over the full +> cartesian product. +> +> *Real failure:* the D256 gfx950 softmax↔MFMA interleave win (+9–10% at +> Sq 4096/8192, correctness-clean) was invisible to the preset-bench sweep for a +> full optimization loop, because `use_softmax_mfma_interleave` is not in the +> `_variant_flags` grammar. A raw-lever direct-launch sweep found it in one pass +> once the profiler pointed at the exposed softmax on a low-occupancy +> (4 waves/CU) kernel. + This ordering is not optional. Real failures from skipping it: a selector routed a shape to a ~2× slower configuration that an exhaustive sweep beat immediately; a gap assumed "structural" turned out to be a single mis-set