feat(rocke): D256 gfx950 bf16 prefill 32x32 + softmax-MFMA interleave fast path

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

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]}

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(

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}")

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]