feat(hip-kernel-provider): route SDPA to unified and dense attention kernels

dnn-providers/hip-kernel-provider/rocke/KERNEL_AUTHORING.md

@@ -318,6 +318,7 @@ down to [platform/](platform/). It distills the conventions already encoded in
 | Kernel emitters | [library/kernels/common/](library/kernels/common/), [library/kernels/gfx950/](library/kernels/gfx950/) | Build the typed SSA `KernelDef` (e.g. [attention_unified.py](library/kernels/common/attention_unified.py), [attention_tiled_2d.py](library/kernels/gfx950/attention_tiled_2d.py), [attention_tiled_3d.py](library/kernels/gfx950/attention_tiled_3d.py)) |
 | Spec seam | [library/builders/common/attention_spec_builder.py](library/builders/common/attention_spec_builder.py) | Maps `UnifiedAttentionProblem` → arch tiled-spec; owns every knob decision |
 | Dispatch | [library/dispatch/attention/](library/dispatch/attention/) | `ATTENTION_REGISTRY` + `dispatch_attention`; selects `(path, head_size, block_size)` |
+| Routed families | [library/dispatch/sdpa/](library/dispatch/sdpa/) | hipDNN SDPA request contract; routes to `dispatch.attention`'s unified/dense candidates via a gate per family ([H](#h-route-a-family-over-dispatchattention)) — registers no kernel of its own |
 | Builders / harnesses | [library/builders/gfx950/attention/](library/builders/gfx950/attention/) | Parity + benchmark drivers, captured shape JSONs, case studies |
 | Benchmarks | [library/benchmarks/](library/benchmarks/) | Standalone perf sweeps (e.g. [benchmark_rocke_unified_attention.py](library/benchmarks/common/benchmark_rocke_unified_attention.py)) |
 | Parity emitters | [library/tests/parity/](library/tests/parity/) | `.py`/`.c` emitter pairs, byte-compared |
@@ -350,7 +351,11 @@ scoped to library.
    ([sweep coverage](#a-new-kernel-family)). If it introduces a new
    selection axis, extend `select_path` / `supports_native_unified_attention`
    (keep the decision a **pure** function of the problem so it mirrors on the C
-   side).
+   side). Registering it here is not enough by itself for a family routed over
+   `dispatch.attention` (SDPA) to reach it: the routed family's gate names
+   which algorithms it routes among, so the new candidate stays unreachable
+   through that family until it is added there too
+   ([H](#h-route-a-family-over-dispatchattention)).
 5. **Add the parity pair.** Create `tests/parity/<family>_emit.py` (Python
    reference, driven by `run_emit` from
    [platform/tests/instances/parity/_emit_common.py](platform/tests/instances/parity/_emit_common.py))
@@ -583,6 +588,68 @@ paged addressing primitive (`global_ptr_add` / `offset_i64*`) — several of the
 already live in platform; new ones that recur across families follow the same
 path.
 
+### H. Route a family over `dispatch.attention`
+
+A family that has no kernel of its own but serves its requests by translating
+them into `AttentionRequest` and asking `ATTENTION_ROUTE_REGISTRY` — SDPA
+(`library/dispatch/sdpa/`) is the worked example — follows [E](#e-new-operator-with-a-hipdnn-graph-node)
+and [F](#f-new-operand-or-output-on-an-existing-operator) for its own request
+contract, plus this binding-layer checklist for the gate that sits in front of
+`dispatch.attention`:
+
+- [ ] **One `bind_<family>.py` per attention launch family** (e.g.
+      `bind_unified.py`, `bind_dense.py`), each exposing `<family>_admits`,
+      `<family>_candidates` and `<family>_bind` behind one shared reject
+      function, so admission and candidate enumeration cannot drift apart.
+      Translation only: the module adds no shape/dtype/head-size rule of its
+      own — "can it run" stays the attention candidates' own `admits()`. Dense
+      (opt-in) candidates are each probed with their own `algorithm` / `spec_id`
+      pinned (`opt_in_probe`), since they do not serve by default.
+- [ ] **Every feature classified once per family**, either carried
+      (`SDPA_<FAMILY>_FEATURES`) or declined with a reason; a shortcut decline
+      (a feature no candidate implements yet, as opposed to an ABI limit) is
+      tagged `NOT_YET_IMPLEMENTED:`. A test asserts carried ∪ declined covers
+      every feature, so a newly added request field cannot silently fall
+      through both sets.
+- [ ] **Check what the kernel computes, not just that it admits.** A candidate
+      that accepts the translated request may still build a different band
+      (e.g. bottom-right causal when top-left was asked) or the wrong scale
+      range; the gate must compare what gets computed against what was asked,
+      not only call `admits()`.
+- [ ] **Decline a pinned dispatch choice** (`algorithm`, `spec_id`, `tuning_id`,
+      `tuning_knobs`) rather than forwarding it — the family's launch takes
+      only the automatic route, so a pin would be silently ignored.
+- [ ] **The never-overclaim test**: nothing is admitted and then dropped. For
+      every admitted request, a kernel is actually reachable with it.
+- [ ] **Four test layers**: wiring (one decline case per declined feature, the
+      value checks, the band comparison), a case-table test (every entry in
+      the family's shared case table gets a verdict), a CPU binding test (the
+      kernel replaced with the reference oracle, every layout/length mode), and
+      a GPU numeric test (binds every admitted case, compares against the
+      reference), enrolled in the GPU test lane.
+- [F](#f-new-operand-or-output-on-an-existing-operator) **for a routed
+  family**: move the feature from the gate's decline table into
+  `SDPA_<FAMILY>_FEATURES`, and bind the request field in both the
+  `SdpaRequest` → `AttentionRequest` translation (`to_attention_request`) and
+  the launch plan that copies tensors through.
+- **[A](#a-new-kernel-family) step 4, for a routed family**: registering a new
+  candidate in `ATTENTION_ROUTE_REGISTRY` is not enough by itself — a routed
+  family's gate names which algorithms it routes among (e.g.
+  `bind_unified.py`'s `_UNIFIED_ALGORITHMS`, built from the exported
+  `dispatch.attention.UNIFIED_ALGORITHMS`; `bind_dense.py`'s
+  `_DENSE_ALGORITHMS`, built from `dispatch.attention.DENSE_ALGORITHMS`), so a
+  new route candidate stays unreachable through the family until it is added
+  there too.
+
+Known gaps in the current SDPA routing (tracked as "Not in this PR" on the
+introducing change, not silently): choosing between the unified and dense
+families for one request is not modeled — a request both admit goes to
+whichever the caller picks; `sdpa_dense_bind` always launches the first
+serving dense candidate in registry order, not the fastest; and the unified
+binding always copies K and V into the paged cache rather than viewing them in
+place when the layout and block size would allow it (Q already has the
+in-place path).
+
 ### Attention specifics — quick reference
 
 - **Parity families:** the 34 matched `library/tests/parity/*_emit.{py,c}` pairs

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/attention_combo_sweep.py

@@ -41,7 +41,7 @@
     attention_dispatch_result,
     iter_dispatch_attention_all,
 )
-from dispatch.attention.common import _problem
+from dispatch.attention.common import attention_problem
 from rocke.dispatch.core import spec_identity
 from rocke.dispatch.tuning.walk import sample_count, sweep_level
 from benchmarks.common.attention_flops import attention_flops
@@ -206,7 +206,7 @@ def _unified_tensors(req, seed: int):
         "cu_seqlens_q": cu,
         "seqused_k": used,
         "block_table": table,
-        "problem": _problem(req),
+        "problem": attention_problem(req),
         "_dense_q": q_dense,
         "_dense_k": k_dense,
         "_dense_v": v_dense,

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/decode_dispatch_diagnose.py

@@ -10,7 +10,7 @@
 # exercise it; leave the variable unset to see the production default.
 
 from dispatch.attention import AttentionRequest, registered_attention_combos
-from dispatch.attention.common import _problem
+from dispatch.attention.common import attention_problem
 from kernels.common.attention_unified import (
     _enable_3d_graph_replay,
     _num_segments,
@@ -50,7 +50,7 @@ def dump_shape(label, req):
         print(
             f"    spec.kernel_name={spec.kernel_name() if hasattr(spec, 'kernel_name') else spec}"
         )
-    p = _problem(req)
+    p = attention_problem(req)
     n2d = p.total_num_q_blocks_upper_bound * p.num_kv_heads
     want_2d = use_2d_kernel(
         head_size=p.head_size,
@@ -90,7 +90,7 @@ def probe_graph():
     from rocke.runtime import no_fence, time_launches
 
     req = _req()
-    p = _problem(req)
+    p = attention_problem(req)
     dt = torch.bfloat16
     B, Hq, Hkv, D, Sk, page = 1, 32, 8, 128, 1024, 16
     q = torch.randn(B, Hq, D, dtype=dt, device="cuda")

dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/__init__.py

@@ -48,9 +48,13 @@
 )
 from .candidate import (
     DENSE_ALGORITHM,
+    DENSE_ALGORITHMS,
     DENSE_GRID_ALGORITHM,
     DENSE_PERSIST_ALGORITHM,
 )
+from .generic import D256_DECODE_ALGORITHM, UNIFIED_2D_ALGORITHM, UNIFIED_3D_ALGORITHM
+from .gfx942_unified import GFX942_DENSE_PIPE_ALGORITHM
+from .gfx950_unified import D256_GFX950_ALGORITHM
 from .common import (
     ATTENTION_ABI_VERSION,
     ATTENTION_DIM_VOCABULARY,
@@ -63,11 +67,29 @@
     AttentionTuningSpec,
     FAMILY,
     _device_num_cus,
-    _problem,
     _request_errors,
     _resolve_num_cus,
     _selector_matches,
+    attention_band,
+    attention_problem,
     canonical_arch,
+    dense_scale_range,
+)
+
+# Every algorithm that routes through the unified launch path
+# (``run_unified_attention_torch``, no tuning spec): the two generic tiled
+# paths, the two per-arch fast cohorts, and D256 decode. A routed family's
+# gate (e.g. SDPA's ``bind_unified.py``) builds its own "which algorithms do
+# I serve" set from this one, so a new unified route candidate cannot go
+# unnoticed by only renaming a literal nobody re-checks.
+UNIFIED_ALGORITHMS = frozenset(
+    {
+        UNIFIED_2D_ALGORITHM,
+        UNIFIED_3D_ALGORITHM,
+        D256_DECODE_ALGORITHM,
+        D256_GFX950_ALGORITHM,
+        GFX942_DENSE_PIPE_ALGORITHM,
+    }
 )
 
 _FAMILY = FAMILY
@@ -539,9 +561,16 @@ def dispatch_for_arches(
 
 
 __all__ = [
+    "D256_DECODE_ALGORITHM",
+    "D256_GFX950_ALGORITHM",
     "DENSE_ALGORITHM",
+    "DENSE_ALGORITHMS",
     "DENSE_GRID_ALGORITHM",
     "DENSE_PERSIST_ALGORITHM",
+    "GFX942_DENSE_PIPE_ALGORITHM",
+    "UNIFIED_2D_ALGORITHM",
+    "UNIFIED_3D_ALGORITHM",
+    "UNIFIED_ALGORITHMS",
     "ATTENTION_ABI_VERSION",
     "ATTENTION_DIM_VOCABULARY",
     "ATTENTION_EXECUTION_REGISTRY",
@@ -554,11 +583,14 @@ def dispatch_for_arches(
     "AttentionRequest",
     "AttentionSpec",
     "AttentionTuningSpec",
+    "attention_band",
     "attention_candidates",
     "attention_execution_candidates",
     "attention_dispatch_result",
+    "attention_problem",
     "attention_sweep_space",
     "attention_tuning_spec",
+    "dense_scale_range",
     "dispatch_attention",
     "dispatch_for_arches",
     "dispatch_attention_all",

dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py

@@ -43,6 +43,14 @@
 DENSE_ALGORITHM = "attention_dense"
 DENSE_GRID_ALGORITHM = "attention_dense_grid"
 DENSE_PERSIST_ALGORITHM = "attention_dense_persist"
+# Every algorithm that routes through the dense launch path (SDPA's dense
+# gate, ``dispatch.sdpa.bind_dense``, builds its own "which algorithms do I
+# serve" set from this one; so does the attention route-coverage test), so a
+# new dense candidate cannot go unnoticed by only adding a literal nobody
+# re-checks.
+DENSE_ALGORITHMS = frozenset(
+    {DENSE_ALGORITHM, DENSE_GRID_ALGORITHM, DENSE_PERSIST_ALGORITHM}
+)
 
 
 def _tuning_precheck(variant: AttentionGeometryVariant):

dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py

@@ -755,26 +755,27 @@ def normalize_band(
     return AttentionBand(left_bound, right_bound, align)
 
 
-def _attention_band(req: AttentionRequest) -> Tuple[bool, int]:
-    """``(causal_top_left, right_bound)`` the unified kernels compute for ``req``.
+def attention_band(req: AttentionRequest) -> AttentionBand:
+    """The band the unified kernels compute for ``req``.
 
     Requests the unified family does not serve (other arches, and the explicit
     ``window_*`` fields, which carry their own bounds) get the default mask: the
     problem is then only a geometry carrier, e.g. for the gfx1151 launch grid.
+    The single home for this decision -- callers outside this package (e.g. the
+    SDPA gates) use the same function internal callers do.
     """
     explicit = req.window_left is not None or (
         req.window_right is not None and int(req.window_right) >= 0
     )
     if not _uses_band_dispatch(req, explicit):
-        return False, 0
-    band = normalize_band(req)
-    return band.top_left, band.kernel_right_bound
+        return AttentionBand(None, 0, Alignment.BOTTOM_RIGHT)
+    return normalize_band(req)
 
 
 def _problem_with_num_cus(
     req: AttentionRequest, num_cus: int
 ) -> UnifiedAttentionProblem:
-    top_left, right_bound = _attention_band(req)
+    band = attention_band(req)
     # total_q = batch * seqlen_q (the flattened query rows). num_seqs = batch.
     return UnifiedAttentionProblem(
         total_q=int(req.batch) * int(req.seqlen_q),
@@ -792,14 +793,18 @@ def _problem_with_num_cus(
         fp8_fnuz=bool(req.fp8_fnuz),
         num_cus=int(num_cus),
         target_ctas=int(req.target_ctas),
-        causal_top_left=top_left,
-        right_bound=right_bound,
+        causal_top_left=band.top_left,
+        right_bound=band.kernel_right_bound,
         emit_lse=bool(req.emit_lse),
     )
 
 
-def _problem(req: AttentionRequest) -> UnifiedAttentionProblem:
-    """The production-routing problem, including the live-device CU policy."""
+def attention_problem(req: AttentionRequest) -> UnifiedAttentionProblem:
+    """The production-routing problem, including the live-device CU policy.
+
+    The single home for this decision -- callers outside this package (e.g.
+    the SDPA gates) use the same function internal callers do.
+    """
     return _problem_with_num_cus(req, _resolve_num_cus(req))
 
 
@@ -919,6 +924,21 @@ def _dense_kernel_module(arch: str):
     return attention_dense
 
 
+def dense_scale_range(arch: str) -> Tuple[float, float] | None:
+    """``(low, high)`` the ``arch`` dense kernels accept for ``scale``, or
+    ``None`` when the kernel module declares no range.
+
+    Public, so a caller outside this package (e.g. the SDPA gate) names the
+    range instead of reaching for a kernel module's private ``_MIN_SCALE`` /
+    ``_MAX_SCALE`` attributes, which would silently stop being checked if
+    those names moved.
+    """
+    module = _dense_kernel_module(arch)
+    low = getattr(module, "_MIN_SCALE", None)
+    high = getattr(module, "_MAX_SCALE", None)
+    return None if low is None or high is None else (low, high)
+
+
 @dataclass(frozen=True)
 class AttentionTuningSpec:
     """Concrete dispatcher-owned spec for every executable tuned candidate.

dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/generic.py

@@ -42,7 +42,7 @@
     AttentionRequest,
     AttentionSpec,
     FAMILY,
-    _problem,
+    attention_problem,
     _request_errors,
     _selector_matches,
 )
@@ -78,6 +78,16 @@
     _UNIFIED_CAPABILITY, supports_features=ATTENTION_FEATURES | {"lse"}
 )
 
+# The two generic unified-path algorithms this module registers, named so a
+# caller outside this package (the SDPA gate) does not have to reconstruct
+# ``f"unified_{path}"`` itself.
+UNIFIED_2D_ALGORITHM = "unified_2d"
+UNIFIED_3D_ALGORITHM = "unified_3d"
+_UNIFIED_PATH_ALGORITHMS = {"2d": UNIFIED_2D_ALGORITHM, "3d": UNIFIED_3D_ALGORITHM}
+
+# The D256 bf16 decode candidate's algorithm (see _make_d256_decode_candidate).
+D256_DECODE_ALGORITHM = "d256_decode"
+
 
 def _make_candidate(*, path: str, priority: int) -> KernelCandidate:
     spec_id = f"unified_{path}"
@@ -91,7 +101,7 @@ def support(req: OperatorRequest) -> Tuple[bool, str]:
         ok, why = _selector_matches(req, candidate)
         if not ok:
             return False, why
-        problem = _problem(req)
+        problem = attention_problem(req)
         # The scalar gate is the coverage check; it refuses LSE (the scalar
         # kernel writes none), so ask it about the LSE-off problem and let the
         # path's own gate rule on LSE below.
@@ -129,7 +139,7 @@ def select(req: OperatorRequest) -> AttentionSpec:
         if not ok:
             raise ValueError(f"{name} does not support request: {why}")
         assert isinstance(req, AttentionRequest)
-        problem = _problem(req)
+        problem = attention_problem(req)
         return AttentionSpec(
             path=path,
             head_size=problem.head_size,
@@ -148,7 +158,7 @@ def select(req: OperatorRequest) -> AttentionSpec:
     candidate = KernelCandidate(
         name=name,
         family=FAMILY,
-        algorithm=f"unified_{path}",
+        algorithm=_UNIFIED_PATH_ALGORITHMS[path],
         spec_id=spec_id,
         abi_version=ATTENTION_ABI_VERSION,
         priority=priority,
@@ -186,7 +196,7 @@ def support(req: OperatorRequest) -> Tuple[bool, str]:
         ok, why = _selector_matches(req, candidate)
         if not ok:
             return False, why
-        problem = _problem(req)
+        problem = attention_problem(req)
         ok, why = supports_native_unified_attention(problem, arch=req.arch)
         if not ok:
             return False, why
@@ -203,7 +213,7 @@ def select(req: OperatorRequest) -> AttentionSpec:
         if not ok:
             raise ValueError(f"{name} does not support request: {why}")
         assert isinstance(req, AttentionRequest)
-        problem = _problem(req)
+        problem = attention_problem(req)
         return AttentionSpec(
             path="3d",
             head_size=problem.head_size,
@@ -220,7 +230,7 @@ def select(req: OperatorRequest) -> AttentionSpec:
     candidate = KernelCandidate(
         name=name,
         family=FAMILY,
-        algorithm="d256_decode",
+        algorithm=D256_DECODE_ALGORITHM,
         spec_id=spec_id,
         abi_version=ATTENTION_ABI_VERSION,
         priority=5,