feat(rocke/gfx950): add causal_bottom_right to the dense attention kernel

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

@@ -27,13 +27,15 @@ rocke IR DSL. Forward-only, bf16/fp16, head_dim 64/128, MHA or GQA.
   slab-padded K/V layouts. IGLP-1 owns this loop schedule and K-major PV traversal
   keeps the 256-VGPR kernel spill-free.
 
-Heads / head_dim / causal / dtype are baked at build time. On the aligned path
+Heads / head_dim / causal / dtype are baked at build time. On the aligned, unshifted path
 `batch`, `seqlen_q`, and `seqlen_kv` are **runtime kernel params** — the spec
 declares them in `runtime_param_fields`, the launcher cache key excludes them, and
 one compiled binary therefore serves every shape (see
 [runtime param field](../../../../../platform/dsl_docs/instances/attention.md#runtime-param-fields)).
 Sub-modes that still bake seqlen into the body — persistent, ragged, varlen,
-paged, sliding-window — keep per-shape identity and a statically-sized ABI.
+paged, sliding-window, and moving bottom-right causal — keep per-shape identity.
+Non-persistent kernels still take the shape parameters; that ABI does not make
+a baked diagonal safe to reuse for a different sequence-length difference.
 
 Tile/resource knobs are `block_n`, `waves_per_eu`, and `lds_k_group_pad`;
 persistent scheduling knobs are `num_persistent`, `persist_decode`, `interleave`,
@@ -203,6 +205,17 @@ spec = dense_spec_for_request(req)             # launch-ready best-config Attent
 run_attention_dense_torch(spec=spec, q=q, k=k, v=v, out=out, scale=1/128**0.5)
 ```
 
+For bottom-right masking, set `mask_type=AttentionMaskType.BOTTOM_RIGHT_CAUSAL`
+(exported by `dispatch.attention`). With unequal Q/K lengths, this standalone
+gfx950 dense path uses a non-persistent grid under `dense_persistent="auto"` and
+rejects an explicit `"on"`. 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.
+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.
+
 `dense_persistent="auto"` turns on the persistent grid-stride variant once there is
 enough work to fill the grid (`⌈Sq/256⌉·Hq·B >= num_persistent`) — i.e. the large-Sq
 prefill regime — so the dispatcher reaches the persistent path, not the default

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

@@ -40,6 +40,7 @@
     ATTENTION_FEATURES,
     UNIFIED_BLOCK_SIZES,
     UNIFIED_HEAD_SIZES,
+    AttentionMaskType,
     AttentionRequest,
     AttentionSpec,
     FAMILY,
@@ -164,6 +165,7 @@ def dispatch_attention(
     "ATTENTION_REGISTRY",
     "UNIFIED_BLOCK_SIZES",
     "UNIFIED_HEAD_SIZES",
+    "AttentionMaskType",
     "AttentionRequest",
     "AttentionSpec",
     "attention_candidates",

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

@@ -45,6 +45,8 @@
 from __future__ import annotations
 
 from dataclasses import asdict, dataclass
+from enum import IntEnum
+from operator import index
 from typing import Tuple
 
 from kernels.common.attention_unified import (
@@ -65,6 +67,33 @@
 ATTENTION_ABI_VERSION = "hipkg-attention-unified/v1"
 
 
+class AttentionMaskType(IntEnum):
+    """Attention-mask ordinals matching ``plan_utils::MaskType``.
+
+    These are mask kinds, not hipDNN ``DiagonalAlignment`` ordinals.
+    """
+
+    NO_MASK = 0
+    TOP_LEFT_CAUSAL = 1
+    BOTTOM_RIGHT_CAUSAL = 2
+    SLIDING_WINDOW = 3
+
+
+_ATTENTION_MASK_ORDINALS = tuple(mask_type.value for mask_type in AttentionMaskType)
+
+
+def _parse_attention_mask_type(value: object) -> AttentionMaskType:
+    """Return a validated mask enum without truncating or parsing strings."""
+    try:
+        ordinal = index(value)
+        return AttentionMaskType(ordinal)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(
+            "mask_type must be an exact integer ordinal in "
+            f"{_ATTENTION_MASK_ORDINALS}, got {value!r}"
+        ) from exc
+
+
 @dataclass(frozen=True)
 class AttentionRequest(OperatorRequest):
     """Normalized scaled-dot-product-attention request."""
@@ -77,7 +106,7 @@ class AttentionRequest(OperatorRequest):
     hdim_q: int
     hdim_v: int
     arch: str
-    mask_type: int = 0  # 0=none, 1=causal/top-left, ...
+    mask_type: AttentionMaskType | int = AttentionMaskType.NO_MASK
     use_sinks: bool = False
     sliding_window: int = 0
     kv_block_size: int = 16  # paged KV block_size (modulus); {16,32,64}
@@ -109,6 +138,8 @@ class AttentionRequest(OperatorRequest):
     def normalized(self) -> dict:
         d = asdict(self)
         d["dtype"] = self.dtype.lower()
+        # IntEnum and raw-int callers describe the same request/cache identity.
+        d["mask_type"] = _parse_attention_mask_type(self.mask_type).value
         return d
 
     def dims(self) -> dict[str, int]:
@@ -125,8 +156,18 @@ def dims(self) -> dict[str, int]:
 
     def features(self) -> frozenset[str]:
         active = set()
-        if int(self.mask_type) != 0:
+        try:
+            mask_type = _parse_attention_mask_type(self.mask_type)
+        except ValueError:
+            # Let _request_errors report the invalid ordinal. In particular, do
+            # not silently classify an arbitrary nonzero value as causal.
+            mask_type = None
+        if mask_type is not None and mask_type != AttentionMaskType.NO_MASK:
             active.add("causal")
+        if mask_type == AttentionMaskType.BOTTOM_RIGHT_CAUSAL and int(
+            self.seqlen_q
+        ) != int(self.seqlen_k):
+            active.add("causal_bottom_right")
         if int(self.sliding_window) > 0:
             active.add("sliding_window")
         if bool(self.use_sinks):
@@ -147,7 +188,9 @@ def features(self) -> frozenset[str]:
     "kv_block_size",
 )
 
-ATTENTION_FEATURES = frozenset({"causal", "sliding_window", "sinks", "fp8"})
+ATTENTION_FEATURES = frozenset(
+    {"causal", "causal_bottom_right", "sliding_window", "sinks", "fp8"}
+)
 
 
 def _request_errors(req: OperatorRequest) -> list[str]:
@@ -163,6 +206,10 @@ def _request_errors(req: OperatorRequest) -> list[str]:
         errors.append("only hdim_q == hdim_v is supported")
     if int(req.nhead_q) % int(req.nhead_k):
         errors.append("nhead_q must be divisible by nhead_k (GQA grouping)")
+    try:
+        _parse_attention_mask_type(req.mask_type)
+    except ValueError as exc:
+        errors.append(str(exc))
     try:
         ArchTarget.from_gfx(req.arch)
     except KeyError as e:

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

@@ -53,6 +53,8 @@
         ShapeRange("hdim_q", allowed=UNIFIED_HEAD_SIZES),
         ShapeRange("kv_block_size", allowed=UNIFIED_BLOCK_SIZES),
     ),
+    # Unified kernels already shift the causal diagonal by the runtime
+    # difference between each sequence's KV and query lengths.
     supports_features=ATTENTION_FEATURES,
 )
 
@@ -179,7 +181,7 @@ def select(req: OperatorRequest) -> AttentionSpec:
                 ShapeRange("hdim_q", allowed=(256,)),
                 ShapeRange("kv_block_size", allowed=UNIFIED_BLOCK_SIZES),
             ),
-            supports_features=frozenset({"causal"}),
+            supports_features=frozenset({"causal", "causal_bottom_right"}),
         ),
         _supports=support,
         select_spec=select,

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

@@ -29,9 +29,11 @@
     ATTENTION_FEATURES,
     UNIFIED_BLOCK_SIZES,
     UNIFIED_HEAD_SIZES,
+    AttentionMaskType,
     AttentionRequest,
     AttentionSpec,
     FAMILY,
+    _parse_attention_mask_type,
     _problem,
     _request_errors,
     _selector_matches,
@@ -118,10 +120,8 @@ def select(req: OperatorRequest) -> AttentionSpec:
                 ShapeRange("hdim_q", allowed=UNIFIED_HEAD_SIZES),
                 ShapeRange("kv_block_size", allowed=UNIFIED_BLOCK_SIZES),
             ),
-            # ``_enable_gfx942_fp16_flash`` is the real narrowing; nothing here
-            # claims a feature it turns down. fp8 is the one exception -- this is
-            # an fp16-only flash path with no fp8 dequant kernel -- so it is
-            # dropped rather than left for the dtype gate to catch by accident.
+            # ``_enable_gfx942_fp16_flash`` is the real narrowing; fp8 is
+            # unsupported, but the unified body already shifts causal masking.
             supports_features=ATTENTION_FEATURES - {"fp8"},
         ),
         _supports=support,
@@ -175,6 +175,11 @@ def _dense_spec(req: OperatorRequest):
             f"gfx942 dense spec factory requires arch='gfx942', got {req.arch!r}"
         )
     sq, sk = int(req.seqlen_q), int(req.seqlen_k)
+    mask_type = _parse_attention_mask_type(req.mask_type)
+    causal = mask_type != AttentionMaskType.NO_MASK
+    causal_bottom_right = (
+        mask_type == AttentionMaskType.BOTTOM_RIGHT_CAUSAL and sq != sk
+    )
     bm = int(DENSE_TILE_GEOMETRIES["default"]["block_m"])
     bn = _DENSE_BLOCK_N
     head_size = int(req.hdim_q)
@@ -205,7 +210,8 @@ def _dense_spec(req: OperatorRequest):
         num_query_heads=int(req.nhead_q),
         num_kv_heads=int(req.nhead_k),
         head_size=head_size,
-        causal=(int(req.mask_type) != 0),
+        causal=causal,
+        causal_bottom_right=causal_bottom_right,
         dtype=dtype,
         block_m=bm,
         block_n=bn,
@@ -306,9 +312,10 @@ def select(req: OperatorRequest) -> AttentionSpec:
         capability=Capability(
             arches=("gfx942",),
             dtypes=("bf16", "fp16"),
-            # Dense: causal + sliding-window; no sinks yet. Causal is a mask, not a
-            # feature this path turns down. Head size stays out -- D64/D128 coverage
-            # is ``supports_attention_dense``'s call, and it reads the built spec
+            # Dense: causal + sliding-window; no sinks or moving bottom-right
+            # diagonal. The latter is a distinct request feature, absent here.
+            # Head size stays out -- D64/D128 coverage is
+            # ``supports_attention_dense``'s call, and it reads the built spec
             # (LDS budget, block_n divisibility), which a ShapeRange cannot.
             supports_features=frozenset({"causal", "sliding_window"}),
         ),

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

@@ -19,9 +19,11 @@
 from .common import (
     ATTENTION_ABI_VERSION,
     UNIFIED_BLOCK_SIZES,
+    AttentionMaskType,
     AttentionRequest,
     AttentionSpec,
     FAMILY,
+    _parse_attention_mask_type,
     _problem,
     _request_errors,
     _selector_matches,
@@ -61,19 +63,30 @@ def _dense_spec(req: OperatorRequest):
     bm = int(geometry["block_m"])
     bn = int(geometry["block_n"])
     decode = req.dense_persist_decode.strip().lower()
-    # on-chip ragged padding for ragged self-attention lengths (seqlen_q==seqlen_kv,
-    # not a 256/block_n multiple). Cross-attention ragged is left to the validator.
-    ragged = (sq == sk) and ((sq % bm != 0) or (sk % bn != 0))
+    mask_type = _parse_attention_mask_type(req.mask_type)
+    causal = mask_type != AttentionMaskType.NO_MASK
+    moving_bottom_right = (
+        mask_type == AttentionMaskType.BOTTOM_RIGHT_CAUSAL and sq != sk
+    )
+    # Cross-length ragged attention is valid only when bottom-right supplies the
+    # shifted diagonal. Equal-length bottom-right is ordinary causal attention.
+    ragged = (sq == sk or moving_bottom_right) and ((sq % bm != 0) or (sk % bn != 0))
     nqb = (sq + bm - 1) // bm
     work = nqb * int(req.nhead_q) * int(req.batch)
     np = int(req.dense_num_persistent)
     mode = req.dense_persistent.strip().lower()
     if mode == "on":
+        if moving_bottom_right:
+            raise ValueError(
+                "dense_persistent='on' is not supported with moving "
+                "bottom-right causal attention"
+            )
         persistent = True
     elif mode == "off":
         persistent = False
     elif mode == "auto":
-        persistent = work >= np  # enough work to fill the persistent grid
+        # The persistent builder has no shifted-diagonal implementation.
+        persistent = False if moving_bottom_right else work >= np
     else:
         raise ValueError(
             f"dense_persistent must be 'auto'/'on'/'off', got {req.dense_persistent!r}"
@@ -87,10 +100,11 @@ def _dense_spec(req: OperatorRequest):
         and int(req.hdim_q) == 128
         and int(req.hdim_v) == 128
         and req.dtype.lower() in ("fp16", "bf16")
-        and int(req.mask_type) != 0
+        and causal
         and sw == 0
         and not use_sinks
         and not ragged
+        and not moving_bottom_right
     )
     return Gfx950AttentionDenseSpec(
         batch=int(req.batch),
@@ -99,7 +113,7 @@ def _dense_spec(req: OperatorRequest):
         num_query_heads=int(req.nhead_q),
         num_kv_heads=int(req.nhead_k),
         head_size=int(req.hdim_q),
-        causal=(int(req.mask_type) != 0),
+        causal=causal,
         dtype=req.dtype.lower(),
         block_m=bm,
         block_n=bn,
@@ -111,6 +125,7 @@ def _dense_spec(req: OperatorRequest):
         sliding_window=sw,
         use_sinks=use_sinks,
         wide_lds_dma=wide_lds_dma,
+        causal_bottom_right=moving_bottom_right,
     )
 
 
@@ -191,7 +206,9 @@ def select(req: OperatorRequest) -> AttentionSpec:
             arches=("gfx950",),
             dtypes=("bf16", "fp16"),
             # Dense: Causal is a mask, not a feature this path turns down.
-            supports_features=frozenset({"causal", "sliding_window", "sinks"}),
+            supports_features=frozenset(
+                {"causal", "causal_bottom_right", "sliding_window", "sinks"}
+            ),
         ),
         _supports=support,
         select_spec=select,
@@ -275,7 +292,7 @@ def select(req: OperatorRequest) -> AttentionSpec:
                 ShapeRange("hdim_q", allowed=(256,)),
                 ShapeRange("kv_block_size", allowed=UNIFIED_BLOCK_SIZES),
             ),
-            supports_features=frozenset({"causal"}),
+            supports_features=frozenset({"causal", "causal_bottom_right"}),
         ),
         _supports=support,
         select_spec=select,

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, fields as _dataclass_fields
+from dataclasses import dataclass, field, fields as _dataclass_fields
 from types import MappingProxyType
 
 from rocke.core.ir import BF16, F16
@@ -73,6 +73,8 @@ class AttentionDenseSpec:
     block_size: int = 0
     num_kv_blocks: int = 0
     use_sinks: bool = False
+    # Appended for positional compatibility with existing concrete specs.
+    causal_bottom_right: bool = field(default=False, kw_only=True)
 
     def supported_persist_decodes(self) -> frozenset[str]:
         """Decode values the concrete kernel type can actually emit."""
@@ -97,13 +99,22 @@ def __post_init__(self) -> None:
                 "elements (16 bytes) so the K group pitch stays "
                 f"ds_read_b128-aligned, got {self.lds_k_group_pad}"
             )
+        if self.causal_bottom_right:
+            if not self.causal:
+                raise ValueError("causal_bottom_right requires causal=True")
+            if self.seqlen_q > self.seqlen_kv:
+                raise ValueError(
+                    "causal_bottom_right requires seqlen_q <= seqlen_kv, got "
+                    f"{self.seqlen_q} > {self.seqlen_kv}"
+                )
 
         if self.ragged:
             if self.seqlen_q <= 0 or self.seqlen_kv <= 0:
                 raise ValueError("ragged requires positive seqlen_q/seqlen_kv")
-            if self.seqlen_q != self.seqlen_kv:
+            if self.seqlen_q != self.seqlen_kv and not self.causal_bottom_right:
                 raise ValueError(
-                    "ragged is self-attention only (seqlen_q == seqlen_kv), got "
+                    "ragged is self-attention only (seqlen_q == seqlen_kv) unless "
+                    "causal_bottom_right is set, got "
                     f"{self.seqlen_q} != {self.seqlen_kv}"
                 )
             if self.varlen:
@@ -294,6 +305,8 @@ def kernel_name(self) -> str:
         parts.extend(self._layout_name_parts())
         parts.extend(self._shape_name_parts())
         parts.append("causal" if self.causal else "full")
+        if self.causal_bottom_right:
+            parts.append("br")
         if self.ragged:
             parts.append("ragged")
         if self.sliding_window > 0: