feat(rocke): top-left causal and attention band coverage on gfx950

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

@@ -87,7 +87,10 @@ def _requests(args):
                     hdim_v=d,
                     arch=args.arch,
                     dtype=args.dtype,
-                    mask_type=1 if args.causal else 0,
+                    # ``_reference`` anchors the mask on the bottom-right diagonal
+                    # (``sk - sq``) for both the causal and the windowed case.
+                    mask_type=2 if args.causal else 0,
+                    diagonal_alignment=1,
                     kv_block_size=args.kv_block_size,
                     sliding_window=args.sliding_window,
                     num_cus=args.num_cus,

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

@@ -57,7 +57,11 @@ def run_sweep(
     timed entry, which records the engine ``spec.name``s that mapped to it.
     """
     import torch
-    from dispatch.attention import AttentionRequest, attention_sweep_space
+    from dispatch.attention import (
+        AttentionMaskType,
+        AttentionRequest,
+        attention_sweep_space,
+    )
     from dispatch.attention.bindings import (
         validate_tuning_attention_contract,
         validate_tuning_attention_tensors,
@@ -82,6 +86,10 @@ def run_sweep(
         num_cus=bench.num_sms,
         use_fp8=bool(problem.use_fp8),
         fp8_fnuz=bool(problem.fp8_fnuz),
+        # The paged unified kernels (and this harness' reference) are bottom-right
+        # causal; AttentionRequest defaults to NO_MASK, which now means full
+        # attention on gfx950 and rejects a sliding window.
+        mask_type=AttentionMaskType.BOTTOM_RIGHT_CAUSAL,
     )
 
     specs = attention_sweep_space(

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

@@ -285,7 +285,11 @@ def _run_dsl(shape: DecodeShape, data: dict, num_cus: int, *, warmup: int, iters
     import torch
 
     try:
-        from dispatch.attention import AttentionRequest, dispatch_attention
+        from dispatch.attention import (
+            AttentionMaskType,
+            AttentionRequest,
+            dispatch_attention,
+        )
         from kernels import UnifiedAttentionProblem, run_unified_attention_torch, prepare_unified_attention_torch  # type: ignore
         from kernels.common.attention_unified import _resolve_attention_arch
     except ImportError:
@@ -312,6 +316,10 @@ def _run_dsl(shape: DecodeShape, data: dict, num_cus: int, *, warmup: int, iters
             sliding_window=shape.sliding_window,
             use_fp8=shape.use_fp8,
             fp8_fnuz=shape.fp8_fnuz,
+            # The paged unified kernels (and this harness' reference) are bottom-right
+            # causal; AttentionRequest defaults to NO_MASK, which now means full
+            # attention on gfx950 and rejects a sliding window.
+            mask_type=AttentionMaskType.BOTTOM_RIGHT_CAUSAL,
         )
         result = dispatch_attention(req)
         path = result.spec.path  # "2d" or "3d"

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py

@@ -1142,7 +1142,11 @@ def _run_prod(shape, data, sw, is_fp8, bench, *, warmup, iters, backend="auto"):
     code path as the production provider.
     """
     import torch
-    from dispatch.attention import AttentionRequest, dispatch_attention
+    from dispatch.attention import (
+        AttentionMaskType,
+        AttentionRequest,
+        dispatch_attention,
+    )
     from kernels import run_unified_attention_torch
     from kernels.common.attention_unified import _resolve_attention_arch
     from rocke.runtime import synchronize_and_release, time_launches
@@ -1162,6 +1166,10 @@ def _run_prod(shape, data, sw, is_fp8, bench, *, warmup, iters, backend="auto"):
         sliding_window=sw,
         kv_block_size=shape.block_size,
         num_cus=bench.num_cus,
+        # The paged unified kernels (and this harness' reference) are bottom-right
+        # causal; AttentionRequest defaults to NO_MASK, which now means full
+        # attention on gfx950 and rejects a sliding window.
+        mask_type=AttentionMaskType.BOTTOM_RIGHT_CAUSAL,
     )
 
     # Force path when caller requests a specific one (e.g. backend="3d").

dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py

@@ -296,6 +296,8 @@ def _spec_gfx950_generic(problem: UnifiedAttentionProblem, arch: str):
     _spec = UnifiedAttention2DTiledSpec(
         **_base_2d_generic_fields(problem, arch),
         **_schedule_fields,
+        causal_top_left=problem.causal_top_left,
+        right_bound=problem.right_bound,
     )
     if _kau._d256_gfx950_fast(problem, arch):
         # D256 gfx950 bf16 prefill fast route -- pins the 32x32 transposed + FA3
@@ -385,6 +387,10 @@ def _spec_generic_3d(problem: UnifiedAttentionProblem, arch: str):
         use_invariant_hoist=_enable_gfx942_3d_invariant_hoist(problem, arch),
         use_wide_kv_load=_enable_gfx942_3d_wide_kv_load(problem, arch),
         use_i64_kv_addr=_enable_i64_kv_addr(problem),
+        # Only the gfx950 spec class declares these; the 3D tiled gate rejects any
+        # non-default mask on every other arch before a spec is built.
+        **({"causal_top_left": True} if problem.causal_top_left else {}),
+        **({"right_bound": problem.right_bound} if problem.right_bound != 0 else {}),
     )
 
 

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md

@@ -278,10 +278,91 @@ For bottom-right masking, set `mask_type=AttentionMaskType.BOTTOM_RIGHT_CAUSAL`
 (exported by `dispatch.attention`). With unequal Q/K lengths, only the grid
 algorithm of this standalone gfx950 dense path admits it; the persistent
 candidates reject it. Equal lengths preserve the equivalent top-left path.
-`algorithm="auto"` still uses the existing unified 2D/3D paths or their eligible
-dense-pipe/D256 candidates: those kernels already shift the causal diagonal by
-each sequence's runtime KV/query length difference. The standalone gfx942 dense
-and gfx1250 WMMA candidates still reject a moving bottom-right diagonal.
+`TOP_LEFT_CAUSAL` (unshifted `key_pos <= query_pos`) is served on gfx950 for
+unequal Q/K lengths too, including non-tile-aligned (ragged) lengths in both
+directions (`seqlen_q < seqlen_k` and `seqlen_q > seqlen_k`).
+
+**The attention band.** `algorithm="auto"` uses the unified 2D/3D paths or their
+eligible dense-pipe/D256 candidates. On gfx950 those kernels compute the cuDNN /
+hipDNN band `(left, right, diagonal_alignment)`: key `k` is visible to the query at
+diagonal position `d` iff `d - k < left` (when set) and `k <= d + right`. An
+`AttentionRequest` spells it with `sliding_window` (left, 0 = unbounded),
+`right_bound` (-1 unbounded = default, 0 causal, R > 0 lookahead) and
+`diagonal_alignment` (unset = TOP_LEFT, or the causal kind's own alignment; 0
+TOP_LEFT, 1 BOTTOM_RIGHT; only matters when `seqlen_q != seqlen_k`). Setting it
+against a causal kind's alignment, mixing it or `right_bound` with the explicit
+`window_left` / `window_right` / `window_bottom_right` fields, or naming
+`SLIDING_WINDOW` without a window or a lookahead is rejected, not ignored:
+
+| mask_type | band |
+|---|---|
+| `NO_MASK` | the trio as given: no bounds = full attention (the SDPA default in cuDNN, FlashAttention, PyTorch); `sliding_window` alone = a left-only (non-causal) window; with `right_bound` = a two-sided window or a lookahead |
+| `SLIDING_WINDOW` | the same trio, but it must name a window or a lookahead (`sliding_window > 0` and/or `right_bound > 0`); a bare causal mask is `TOP_LEFT_CAUSAL` |
+| `TOP_LEFT_CAUSAL` / `BOTTOM_RIGHT_CAUSAL` | right bound 0, alignment from the kind (a causal window = these + `sliding_window`) |
+
+Full attention at `seqlen_q == 1` (single-token decode) is the same computation as
+bottom-right causal, so it is dispatched to the causal kernels and the decode-only
+candidates (strided-KV, d256 decode) keep serving the default `NO_MASK` request. A
+decode *window* is a different band: `NO_MASK` + `sliding_window` is the top-left
+non-causal window, so a decode caller that wants the last-`W`-keys window passes
+`BOTTOM_RIGHT_CAUSAL` (or `diagonal_alignment=1`). Candidates without band support
+reject any real band with a `cannot serve features ['band']` capability reason.
+
+Kernel variants are distinct bodies and names: `_tl` (top-left), `_rbu` / `_rb<R>`
+(unbounded / R-key right bound, plus `_sw<W>` for a non-causal window); at equal
+length the alignments coincide, so the default body is kept. Window, right bound
+and the sinks / softcap options compose; ALiBi and QQ-bias are only defined for
+causal masks. At `seqlen_q == 1` top-left attends only key 0, which is the
+definition of top-left, not a bug. A row whose band holds no key outputs zeros (the
+same in the unified and dense kernels, and what PyTorch SDPA returns for a fully
+masked row).
+
+The standalone dense kernel (`algorithm="attention_dense_grid"` / `"attention_dense_persist"`) implements the same band
+on the top-left diagonal -- full attention (including ragged cross-attention
+`seqlen_q != seqlen_k`), left-only / two-sided windows and lookahead, in both the
+grid and persistent algorithms (kernel names `_full`, `_causal`, `_swa<W>`, `_rb<R>`).
+It does not implement a window or lookahead at bottom-right alignment with
+`seqlen_q != seqlen_k`; that request is rejected with a reason, and `algorithm="auto"`
+routes it to the unified kernels.
+
+Per-sequence alignment: for the paged unified kernels `seqlen_q` / `seqlen_k` are
+the batch maxima and the diagonal is derived per sequence at runtime, so on gfx950
+`TOP_LEFT_CAUSAL` (or the top-left alignment of the band fields) selects the `_tl`
+body even when the two maxima are equal; it is never folded into bottom-right. Only
+dense, which takes uniform lengths, folds the alignments at equal length. On arches
+without a top-left kernel (gfx942 / gfx1250) the legacy equal-length shortcut remains
+and serves top-left with the bottom-right body, which is correct only when every
+sequence has `q_len == kv_len`.
+
+Known gap: the C++ engine (`platform/cpp/instances/`, the tiled 2D / 3D and unified
+bodies, with the `library/tests/parity/*_emit` pairs) does not yet mirror the
+`causal_top_left` / `right_bound` bodies. The default mask is unchanged, so the
+Python/C++ byte-identity gate still holds, and the new variants' lowered IR is pinned
+by representative cases in `platform/tests/golden/rocke_representative_ir_sha256.json`
+(Python only). Mirroring them in C++ and adding emit pairs is tracked as a follow-up.
+
+Every alias above (causal kinds fix the right bound at 0, full attention ignores the
+alignment, single-token full attention is causal, equal-length top-left folds for the
+dense kernel only) is resolved in one place, `normalize_band` in
+[common.py](dispatch/attention/common.py), into a frozen `AttentionBand(left_bound,
+right_bound, alignment)`. The unified problem fields, the feature set and the dense
+band all read that value. The feature vocabulary is `causal`, `full_attention`
+(no bound on either side), and `band` (a lookahead or non-causal window); a candidate
+that does not declare one is excluded for a request that needs it.
+
+Known gap: the gfx942 and gfx1250 tiled kernels are still causal-only. `NO_MASK` is
+the request default, so on those arches it keeps being served by the existing
+causal body (a window on it is a causal window, as before) instead of re-routing
+every default request to the slow scalar kernel; any other band is rejected with a
+reason, and a directly built non-default `UnifiedAttentionProblem` is rejected by
+their tiled gates. This is the only place the logical mask depends on the arch; the
+stacked follow-up (#13211) implements the band on the gfx942 tiled and dense kernels
+and removes it. A caller that maps hipDNN `(left, right)` bounds onto this
+request is responsible for passing the right bound through (the ingestor pack
+declines any bound it cannot serve).
+
+The standalone gfx942 dense and gfx1250 WMMA candidates still
+reject a moving bottom-right diagonal.
 That rejection does not disable gfx942 dense sliding-window attention: top-left
 and equal-length bottom-right requests retain the windowed path on both grids.