feat(rocke): attention band coverage on gfx942

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

@@ -88,7 +88,7 @@ def run_sweep(
         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.
+        # attention on gfx950 and gfx942 and rejects a sliding window.
         mask_type=AttentionMaskType.BOTTOM_RIGHT_CAUSAL,
     )
 

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

@@ -270,7 +270,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:
@@ -295,6 +299,10 @@ def _run_dsl(shape: DecodeShape, data: dict, num_cus: int, *, warmup: int, iters
             num_cus=num_cus,
             use_sinks=shape.use_sinks,
             sliding_window=shape.sliding_window,
+            # 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 gfx942 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/gfx942/attention/prefill/benchmark_prefill2d_live.py

@@ -1212,7 +1212,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 rocke.runtime import synchronize_and_release, time_launches
 
@@ -1230,6 +1234,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 gfx942 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

@@ -74,6 +74,22 @@
 from kernels.common import attention_unified as _kau
 
 
+def _mask_spec_fields(problem: UnifiedAttentionProblem) -> dict:
+    """Band fields (``causal_top_left`` / ``right_bound``) for a non-gfx950 tiled spec.
+
+    Splatted only when set, so the default bottom-right causal mask builds exactly
+    the spec (and IR) it always did. gfx942's 2D/3D tiled specs declare both; every
+    other non-gfx950 arch's tiled gate (``_reject_mask_on_arch``) rejects a
+    non-default mask before a spec is built, so their classes never see them.
+    """
+    fields: dict = {}
+    if problem.causal_top_left:
+        fields["causal_top_left"] = True
+    if problem.right_bound != 0:
+        fields["right_bound"] = problem.right_bound
+    return fields
+
+
 def _spec_gfx942_fp16_flash(problem: UnifiedAttentionProblem, arch: str):
     """gfx942 fp16 transposed-x8 flash geometry (the ``gfx942_dense_pipe`` engine).
 
@@ -121,6 +137,7 @@ def _spec_gfx942_fp16_flash(problem: UnifiedAttentionProblem, arch: str):
         use_q_direct_global=_enable_gfx942_flash_q_direct(problem, arch),
         kv_cache_policy=_gfx942_flash_kv_cache_policy(problem),
         use_i64_kv_addr=_enable_i64_kv_addr(problem),
+        **_mask_spec_fields(problem),
     )
 
 
@@ -186,6 +203,7 @@ def _spec_gfx942_bf16_flash(problem: UnifiedAttentionProblem, arch: str):
         use_q_direct_global=_enable_gfx942_flash_q_direct(problem, arch),
         kv_cache_policy=_gfx942_flash_kv_cache_policy(problem),
         use_i64_kv_addr=_enable_i64_kv_addr(problem),
+        **_mask_spec_fields(problem),
     )
 
 
@@ -272,7 +290,9 @@ def _spec_generic_2d_non_gfx950(problem: UnifiedAttentionProblem, arch: str):
     unchanged.
     """
     UnifiedAttention2DTiledSpec, _, _ = _tiled_2d_impl(arch)
-    return UnifiedAttention2DTiledSpec(**_base_2d_generic_fields(problem, arch))
+    return UnifiedAttention2DTiledSpec(
+        **_base_2d_generic_fields(problem, arch), **_mask_spec_fields(problem)
+    )
 
 
 def _spec_gfx950_generic(problem: UnifiedAttentionProblem, arch: str):
@@ -418,10 +438,7 @@ def _spec_generic_3d(problem: UnifiedAttentionProblem, arch: str):
             if arch in ("gfx942", "gfx950")
             else {}
         ),
-        # 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 {}),
+        **_mask_spec_fields(problem),
     )
 
 

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

@@ -146,6 +146,16 @@ prefill shapes are not included.
   graphs decode (`max_seqlen_q == 1`) and short prefill (`<= 768`) — where host
   launch overhead is a large fraction — and ungraphs long prefill (kernel-bound).
   Toggle with `HIPDNN_GFX942_2D_GRAPH` / `HIPDNN_GFX942_3D_GRAPH` (both default-on).
+- **The attention band matches gfx950.** The tiled 2D/3D kernels compute the
+  cuDNN / hipDNN band: top-left alignment (`causal_top_left`, name token `tl`),
+  full attention (`NO_MASK`, `rbu`), lookahead (`rb<R>`) and non-causal windows.
+  With top-left, `context_len` is 0, so the ALiBi / QQ-bias offsets follow it;
+  `right_bound != 0` is rejected together with ALiBi / QQ-bias. The 4-warp GQA
+  cohorts (D256 causal, D128 sliding-window) keep top-left but decline a right
+  bound other than 0, so those problems build the generic tiled 2D spec. The
+  default mask keeps every kernel name, cache key and IR unchanged. See
+  [`../../../dispatch/attention/README.md`](../../../dispatch/attention/README.md)
+  ("Attention band"). gfx1250 remains the causal-only gap.
 - **The 3D decode kernel is transposed too.** The split-KV segment kernel
   (`attention_tiled_3d`) computes $S^\top = K Q^\top$ on the narrow `16x16x16` atom
   with $P^\top$ register-resident, as the flash path above does with `32x32x8`. It

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/attention_dense_prefill.py

@@ -207,6 +207,16 @@ def add_dense_tuning_args(ap: argparse.ArgumentParser) -> None:
         action="store_true",
         help="enable attention sinks (one [Hq] logit per query head)",
     )
+    # The band, also a request property.
+    ap.add_argument(
+        "--right-bound",
+        dest="right_bound",
+        type=int,
+        default=None,
+        help="attention-band right bound (-1 = unbounded, 0 = causal, R > 0 = R keys "
+        "of lookahead); when set the request is NO_MASK + this band (the cuDNN "
+        "vocabulary) and overrides the causal flag; default = the causal flag",
+    )
 
 
 def dense_request(
@@ -229,6 +239,11 @@ def dense_request(
     full causal) and ``use_sinks`` (``--use-sinks``, read from ``args``) are
     request properties, so dispatch ships the SWA-pruned, sink-seeded spec.
     """
+    mask_type = 1 if causal else 0
+    right_bound = getattr(args, "right_bound", None)
+    if right_bound is not None:
+        # The band comes from (sliding_window, right_bound) under NO_MASK.
+        mask_type = 0
     return AttentionRequest(
         batch=int(batch),
         nhead_q=int(num_query_heads),
@@ -238,13 +253,14 @@ def dense_request(
         hdim_q=int(head_size),
         hdim_v=int(head_size),
         arch=_ARCH,
-        mask_type=1 if causal else 0,
+        mask_type=mask_type,
         dtype=str(dtype).lower(),
         sliding_window=int(sliding_window),
         use_sinks=bool(getattr(args, "use_sinks", False)),
         # Opt-in selector: this is the candidate whose spec we are measuring.
         algorithm="attention_dense",
         spec_id=_SPEC_ID,
+        **({} if right_bound is None else {"right_bound": int(right_bound)}),
     )