feat(rocke): optional LSE output on unified tiled 2D attention

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

@@ -238,9 +238,9 @@ pass the buffer as `tensors["lse"]`; `dispatch_attention(req).bind_torch(tensors
 then launches the LSE kernel. A candidate whose kernel cannot write LSE does not
 declare the `lse` feature, so an LSE request is refused rather than routed to a
 kernel without the output. The prefill CLI, hipDNN, and the standalone split-KV
-decode kernels do not write LSE; unified attention writes it on its split-KV path
-only (see `dispatch/attention/README.md`). The gfx942 dense kernel has the same
-opt-in.
+decode kernels do not write LSE; unified attention writes it on its split-KV and
+tiled 2D paths on gfx942/gfx950 (see `dispatch/attention/README.md`). The gfx942
+dense kernel has the same opt-in.
 
 Through the dispatcher (opt-in: pin the candidate by `algorithm` and `spec_id`,
 and optionally a swept point by `tuning_id`):

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

@@ -305,21 +305,37 @@ Selection pins a spec the same way for dense and unified:
   another variant is rejected before the search.
 - `tuning_id="auto"` selects the candidate's base spec.
 - `emit_lse=True` asks for the optional FP32 LSE output (natural-log softmax
-  normalizer). It is the `lse` request feature. The dense candidates, the RDNA
-  WMMA candidates, `attention_unified_3d` and the split-KV tuning candidates
-  (`*_u3d_*`) declare it; the unified 2D, d256 decode and strided decode
-  candidates do not. The split-KV support check then refuses LSE on
-  arches other than gfx942/gfx950 and with sinks. So an LSE request that cannot
-  be served is refused instead of reaching a kernel without the output. To
-  launch unified LSE through `bind_torch`, pin a split-KV tuning spec
-  (`algorithm="unified_tuning"`,
-  `spec_id="<arch>_u3d_..."`): like any unpinned unified request, the
-  `attention_unified_3d` result routes but cannot be bound. It is a dense output
-  field (`DENSE_OUTPUT_FIELDS`) and a unified semantic field (`_SEMANTIC_FIELDS`),
-  so sweeps never vary it, and with it off every spec hash and tuning id is the
-  one stored before the field existed. The caller passes the buffer to
-  `bind_torch` as `tensors["lse"]`; the runner validates its shape, dtype,
-  device, and overlap. The unified layout is `[total_q, Hq, 1]`.
+  normalizer). It is the `lse` request feature. These declare it:
+  - the dense candidates;
+  - the RDNA WMMA candidates;
+  - `attention_unified_2d` and `attention_unified_3d`;
+  - every unified tuning candidate (`*_u2d_*`, `*_u3d_*`) except the gfx942
+    4-warp GQA one (`gfx942_u2d_gfx942_4warp_nw1_mw32_t64_llvm`).
+
+  The d256 decode, strided decode and route-only `attention_gfx942_dense_pipe`
+  / `attention_gfx950_d256` candidates do not. Each path's own support check
+  (split-KV or tiled 2D) refuses LSE on arches other than gfx942/gfx950 and with
+  sinks. The scalar fallback writes no LSE, so it refuses it. When the tiled 2D
+  check declines an LSE request that `select_path()` sends to tiled 2D (gfx942
+  D256 prefill at block size 32 exceeds its LDS budget),
+  `attention_unified_3d` serves it with split-KV, as the runner's `auto` backend
+  does (`select_auto_path`). An LSE request that cannot be served is refused
+  instead of reaching a kernel without the output.
+
+  To launch unified LSE through `bind_torch`, pin a tuning spec
+  (`algorithm="unified_tuning"`, `spec_id="<arch>_u2d_..."` or `"<arch>_u3d_..."`).
+  Like any unpinned unified request, the `attention_unified_2d` / `_3d` result
+  routes but cannot be bound.
+
+  `emit_lse` is a dense output field (`DENSE_OUTPUT_FIELDS`) and a unified
+  semantic field (`_SEMANTIC_FIELDS`), so sweeps never vary it. It is kept out
+  of the config defaults fingerprint, so a pin's `tuning_id` is the same with
+  LSE on or off. With it off, every spec hash and tuning id is the one stored
+  before the field existed.
+
+  The caller passes the buffer to `bind_torch` as `tensors["lse"]`; the runner
+  validates its shape, dtype, device and overlap. The unified layout is
+  `[total_q, Hq, 1]`.
 
 A candidate resolves a request once: `support` and `select_spec` share a
 per-candidate cache, so a pin is not resolved twice.

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

@@ -108,6 +108,13 @@ def validate_tuning_attention_contract(request, problem, spec) -> None:
         raise ValueError("kernel spec K/V storage dtype disagrees with problem")
     if bool(spec.fp8_fnuz) != bool(problem.fp8_fnuz):
         raise ValueError("tuning wrapper FP8 encoding disagrees with problem")
+    if spec.path == "2d":
+        spec_lse = bool(getattr(kernel_spec, "emit_lse", False))
+        if spec_lse != problem.emit_lse:
+            raise ValueError(
+                f"kernel_spec.emit_lse={spec_lse} disagrees with "
+                f"problem.emit_lse={problem.emit_lse}"
+            )
 
     reduce_spec = spec.reduce_spec
     if reduce_spec is not None:

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

@@ -10,6 +10,8 @@
 
 from __future__ import annotations
 
+import dataclasses
+
 from typing import Callable, Mapping, Optional
 
 from kernels.common.attention_unified import (
@@ -44,11 +46,10 @@
 
 
 def _tuning_precheck(variant: AttentionGeometryVariant):
-    """Coverage for one geometry variant. The scalar coverage check knows nothing
-    about the LSE output, so an LSE request also has to pass the gate of the
-    variant's own path: the split-KV gate owns which arches and features can
-    write LSE, and the 2D gate refuses it. That holds even if a 2D variant ever
-    declared the ``lse`` feature by mistake."""
+    """Coverage for one geometry variant. The scalar coverage check refuses LSE
+    (the scalar kernel writes none), so it is asked about the LSE-off problem
+    and an LSE request must also pass the gate of the variant's own path: the
+    split-KV or tiled 2D gate owns which arches and features can write LSE."""
     path_gate = (
         supports_native_unified_attention_3d_tiled
         if variant.path == "3d"
@@ -57,8 +58,12 @@ def _tuning_precheck(variant: AttentionGeometryVariant):
 
     def precheck(req):
         problem = _tuning_problem(req)
-        ok, why = supports_native_unified_attention(problem, arch=variant.arch)
-        if ok and problem.emit_lse:
+        if not problem.emit_lse:
+            return supports_native_unified_attention(problem, arch=variant.arch)
+        ok, why = supports_native_unified_attention(
+            dataclasses.replace(problem, emit_lse=False), arch=variant.arch
+        )
+        if ok:
             return path_gate(problem, arch=variant.arch)
         return ok, why
 
@@ -88,6 +93,7 @@ def signature(spec):
             include_bt_stride=True,
             include_qq_bias_stride=True,
             kv_dtype=spec.kernel_spec.kv_storage_dtype,
+            emit_lse=bool(getattr(spec.kernel_spec, "emit_lse", False)),
         )
 
     def bind_torch(request, spec, tensors, **kwargs):
@@ -112,11 +118,12 @@ def bind_torch(request, spec, tensors, **kwargs):
                 ShapeRange("hdim_q", allowed=UNIFIED_HEAD_SIZES),
                 ShapeRange("kv_block_size", allowed=UNIFIED_BLOCK_SIZES),
             ),
-            # Only split-KV variants write LSE; 2D variants refuse the feature.
+            # Every unified variant writes LSE except the gfx942 4-warp GQA
+            # builder (see _GFX942_4WARP_LSE_REFUSAL).
             supports_features=(
-                ATTENTION_FEATURES | {"lse"}
-                if variant.path == "3d"
-                else ATTENTION_FEATURES
+                ATTENTION_FEATURES
+                if variant.builder_kind == "gfx942_4warp_gqa"
+                else ATTENTION_FEATURES | {"lse"}
             ),
         ),
         space=unified_space(variant),

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

@@ -201,7 +201,7 @@ class AttentionRequest(OperatorRequest):
     # Also write the natural-log softmax normalizer (LSE) into a caller-owned
     # FP32 ``tensors["lse"]``. Reported as the ``lse`` feature, so only
     # candidates whose kernels implement it (dense, RDNA WMMA, and the
-    # gfx942/gfx950 unified split-KV path) admit it.
+    # gfx942/gfx950 unified split-KV and tiled 2D paths) admit it.
     emit_lse: bool = False
     # Serve runtime head widths from one compiled bucket: ``hdim_q``/``hdim_v``
     # are the bucket maxima, and any widths the kernel takes up to them (RDNA

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

@@ -15,8 +15,10 @@
 from typing import Tuple
 
 from kernels.common.attention_unified import (
+    select_auto_path,
     supports_native_unified_attention,
     supports_native_unified_attention_3d_tiled,
+    supports_native_unified_attention_tiled,
 )
 from rocke.core.arch import known_arches
 from rocke.dispatch.core import (
@@ -63,9 +65,9 @@
     # the scalar kernel elsewhere (tiled gates reject it off gfx950).
     supports_features=ATTENTION_FEATURES,
 )
-# Only the split-KV (3D) path writes the optional LSE output; its support check
-# (the 3D gate below) still refuses it on arches/features that cannot.
-_UNIFIED_3D_CAPABILITY = dataclasses.replace(
+# Both unified paths write the optional LSE output; each path's own gate (in
+# ``support`` below) still refuses it on arches/features that cannot.
+_UNIFIED_LSE_CAPABILITY = dataclasses.replace(
     _UNIFIED_CAPABILITY, supports_features=ATTENTION_FEATURES | {"lse"}
 )
 
@@ -83,17 +85,34 @@ def support(req: OperatorRequest) -> Tuple[bool, str]:
         if not ok:
             return False, why
         problem = _problem(req)
-        ok, why = supports_native_unified_attention(problem, arch=req.arch)
+        # 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.
+        coverage = (
+            dataclasses.replace(problem, emit_lse=False)
+            if problem.emit_lse
+            else problem
+        )
+        ok, why = supports_native_unified_attention(coverage, arch=req.arch)
         if not ok:
             return False, why
-        if problem.select_path(req.arch) != path:
+        # Split-KV also takes the LSE requests that ``select_path(arch)`` sends
+        # to tiled 2D but whose tiled 2D gate declines, as auto does
+        # (``select_auto_path``); the tiled 2D candidate keeps its gate's reason.
+        if problem.select_path(req.arch) != path and not (
+            path == "3d" and select_auto_path(problem, req.arch) == "3d"
+        ):
             return False, (
                 f"problem routes to {problem.select_path(req.arch)!r} path, not {path!r}"
             )
         if problem.emit_lse:
-            # The scalar coverage check above knows nothing about LSE; the
-            # split-KV gate owns which arches and features can write it.
-            ok, why = supports_native_unified_attention_3d_tiled(problem, arch=req.arch)
+            # The path's own gate owns which arches and features can write LSE.
+            gate = (
+                supports_native_unified_attention_3d_tiled
+                if path == "3d"
+                else supports_native_unified_attention_tiled
+            )
+            ok, why = gate(problem, arch=req.arch)
             if not ok:
                 return False, why
         return True, "ok"
@@ -126,7 +145,7 @@ def select(req: OperatorRequest) -> AttentionSpec:
         spec_id=spec_id,
         abi_version=ATTENTION_ABI_VERSION,
         priority=priority,
-        capability=_UNIFIED_3D_CAPABILITY if path == "3d" else _UNIFIED_CAPABILITY,
+        capability=_UNIFIED_LSE_CAPABILITY,
         _supports=support,
         select_spec=select,
         signature=lambda _spec: (),

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

@@ -13,6 +13,7 @@
 from typing import Mapping, Optional, Tuple
 
 from kernels.common.attention_unified import (
+    _GFX942_4WARP_LSE_REFUSAL,
     UnifiedAttentionProblem,
     _kv_block_bytes,
     _kv_cache_policy_3d,
@@ -208,6 +209,8 @@ def make_explicit_attention_2d_spec(
     validate_explicit_fp8_encoding(
         arch=arch, use_fp8=problem.use_fp8, fp8_fnuz=problem.fp8_fnuz
     )
+    if problem.emit_lse and config.builder_kind == "gfx942_4warp_gqa":
+        raise ValueError(_GFX942_4WARP_LSE_REFUSAL)
     if config.num_warps not in NUM_WARPS:
         raise ValueError(f"num_warps must be one of {NUM_WARPS}")
     if config.block_m_per_warp not in BLOCK_M_PER_WARP:
@@ -231,6 +234,9 @@ def make_explicit_attention_2d_spec(
         block_m_per_warp=int(config.block_m_per_warp),
         tile_size=int(tile_size),
         waves_per_eu=config.waves_per_eu,
+        # The output the caller asked for, never a tuning knob (see
+        # _SEMANTIC_FIELDS); passed only when on, so LSE-off specs are as before.
+        **({"emit_lse": True} if problem.emit_lse else {}),
         **knobs,
     )
 

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

@@ -389,14 +389,15 @@ def defaults(self, base, kernel):
         recorded = self.recorded(base)
         return {
             "codepath": dict(self.fixed(base)),
+            # ``emit_lse`` is an output the problem selects, not a default that
+            # knobs are a delta against: keeping it out of both the 2D kernel
+            # and the split-KV reduce defaults leaves every config_key, and so
+            # every stored pin, as before the field existed.
             "kernel": {
                 name: value
                 for name, value in _declared_defaults(type(kernel.kernel_spec))
-                if name not in recorded
+                if name not in recorded and name != "emit_lse"
             },
-            # ``emit_lse`` is an output the problem selects, not a default that
-            # knobs are a delta against: keeping it out leaves every split-KV
-            # config_key, and so every stored pin, as before the field existed.
             "reduce": (
                 None
                 if reduce is None

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

@@ -9,9 +9,10 @@
   LSE store, ``dense_row_has_key`` / ``zero_keyless_row`` the empty-row output
   masking, and ``validate_dense_lse`` is the host-side contract for the
   caller-owned LSE buffer;
-* the unified split-KV reduce: the same IR helpers emit its LSE store, and
-  ``validate_unified_lse`` is its host-side buffer contract.  Both validators
-  sit on the common ``_check_lse_buffer``;
+* the unified split-KV reduce and the unified tiled 2D epilogues (gfx942 /
+  gfx950): the same IR helpers emit their LSE store, ``lse_row_has_mass`` the
+  ``!(l == 0)`` row test, and ``validate_unified_lse`` is their host-side buffer
+  contract.  Both validators sit on the common ``_check_lse_buffer``;
 * the MFMA / WMMA attention forward inner bodies: ``make_*_lse_epilogue``
   produce *epilogue callables* for
   :class:`rocke.helpers.attention_fwd_ext.AttnFwdExt.epilogue_hook`.
@@ -61,6 +62,7 @@
     "LSE_ROWS_PER_LANE",
     "dense_row_has_key",
     "lse_row_value",
+    "lse_row_has_mass",
     "lse_store",
     "make_lse_epilogue",
     "make_mfma_lse_epilogue",
@@ -90,6 +92,17 @@ def lse_store(
     b.global_store(dst, element_offset, value, align=4)
 
 
+def lse_row_has_mass(b: IRBuilder, l: Value) -> Value:
+    """``!(l == 0)``, NaN included: the exact inverse of the tiled 2D O guard
+    ``fcmp("oeq", l, 0)``. A keyless row (l == 0) gets -inf; a NaN row keeps NaN.
+
+    Built as ``one OR uno`` instead of ``lnot(oeq)``: the HIP lowerer maps i1 to
+    ``bool`` and lowers ``arith.not`` to ``~x``, which is true for both inputs.
+    """
+    zero = b.const_f32(0.0)
+    return b.lor(b.fcmp("one", l, zero), b.fcmp("uno", l, zero))
+
+
 def dense_row_has_key(
     b, query_row, kv_length, *, causal, diagonal, window, right_bound=0
 ):