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

@@ -33,7 +33,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
 
@@ -201,7 +201,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] = {}
@@ -236,7 +236,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"
@@ -297,7 +300,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,
@@ -433,6 +442,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(
@@ -463,7 +482,8 @@ def main() -> int:
     print(f"device: {torch.cuda.get_device_name(0)}")
     print(f"shapes: {len(shapes)}  variants: {args.variants}")
 
-    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 = []
     for i, shape in enumerate(shapes, 1):
         sw = shape.window_size[0] + 1 if shape.window_size[0] >= 0 else 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,
@@ -60,6 +60,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,
@@ -213,7 +218,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,
@@ -276,6 +281,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(