feat(rocke): add SDPA padding host adapter, mask spec, and oracle sup…

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

@@ -0,0 +1,54 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""Specification + host reference for the SDPA padding mask (ticket 10).
+
+Given hipDNN's per-sequence valid lengths (SEQ_LEN_Q / SEQ_LEN_KV, normalized
+by :mod:`kernels.common._sdpa_seqlen`), this says which key columns contribute
+to each query's softmax and which query rows carry a real output. It is the
+mask that sits ON TOP of the arbitrary-length buffer clipping (ticket 8):
+clipping makes a partial tile safe to *read*; this makes the padded positions
+not *count*. A key read as 0 by a clipped load still scores 0 and would take
+softmax weight, so the mask below is a separate, still-required step.
+
+Two pieces, matching where each lands in the shell (``fmha_sdpa_fwd.py``, ST4):
+
+1. KV score mask -- the shell's ``extra_mask_predicate(b, k_idx)`` hook. Keep
+   key column ``k`` iff ``k < len_kv``; masked columns are driven to ``neg_inf``
+   by the hook's existing ``b.select(keep, score, neg_inf)``. The IR emitter is
+   a one-liner mirroring :func:`rocke.helpers.attention.causal_mask` --
+   ``b.cmp_lt(key_pos, len_kv)`` -- added to the shell with its byte-identity
+   C++ mirror when the shell lands. Until then this module is the host-side
+   reference that emitter must match.
+
+2. Query-row output validity -- the epilogue. A row with ``q >= len_q`` is
+   padding and is written as ``O = 0`` (LSE ``-inf``). A row left with no valid
+   key (e.g. ``len_kv == 0``) is fully masked and reaches ``O = 0`` through the
+   epilogue's existing ``l == 0`` guard; no extra kernel logic is needed for it.
+
+Causal / sliding-window masking is NOT here -- the shell already applies it
+through its own ``mask_mode``; padding composes with it. These predicates are
+pure Python so the padding semantics are locked and cross-checked against
+``dense_attention_reference`` on the host, with no GPU.
+"""
+
+from __future__ import annotations
+
+
+def keep_key(key_pos, len_kv):
+    """KV padding score mask: keep key column ``key_pos`` iff it is valid.
+
+    ``key_pos`` is the key column index and ``len_kv`` the sequence's valid key
+    count; accepts scalars or numpy arrays. Emitter form when wired into the
+    shell's ``extra_mask_predicate`` hook: ``b.cmp_lt(key_pos, len_kv)``.
+    """
+    return key_pos < len_kv
+
+
+def row_is_valid(query_pos, len_q):
+    """Query-row output validity: a row at or past ``len_q`` is padding.
+
+    Returns ``False`` for padded rows, which the epilogue writes as ``O = 0``
+    (LSE ``-inf``). Accepts scalars or numpy arrays.
+    """
+    return query_pos < len_q

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

@@ -0,0 +1,52 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""Host-side adapter for hipDNN per-sequence padding lengths (SEQ_LEN_Q/KV).
+
+hipDNN supplies padding lengths as int32 tensors of shape ``[B, 1, 1, 1]`` --
+one valid (unpadded) count per batch item. The SDPA stride shell
+(``fmha_sdpa_fwd.py``, SH-01) reads flat per-batch length arrays. This module is
+the CPU-side validation + translation between the two, run before the kernel
+launches; it emits no kernel IR.
+
+The shell's exact kernarg contract is not yet frozen (ST4). The returned
+per-batch ``int32`` arrays are the stable part of the interface; the call
+signature may be adjusted when the shell lands.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def _validate_lengths(lengths, *, batch, s_max, name):
+    arr = np.asarray(lengths)
+    if not np.issubdtype(arr.dtype, np.integer):
+        raise ValueError(f"{name}: lengths must be an integer dtype, got {arr.dtype}")
+    if arr.size != batch:
+        raise ValueError(f"{name}: expected {batch} per-batch lengths, got {arr.size}")
+    flat = arr.reshape(batch).astype(np.int32, copy=False)
+    if flat.min(initial=0) < 0 or flat.max(initial=0) > s_max:
+        raise ValueError(
+            f"{name}: every length must satisfy 0 <= len <= {s_max}; "
+            f"got min={int(flat.min())}, max={int(flat.max())}"
+        )
+    return np.ascontiguousarray(flat, dtype=np.int32)
+
+
+def sdpa_seqlens(seqlen_q, seqlen_kv, *, batch, s_q_max, s_kv_max):
+    """Validate hipDNN SEQ_LEN_Q / SEQ_LEN_KV and return flat per-batch arrays.
+
+    ``seqlen_q`` / ``seqlen_kv`` are array-likes of ``batch`` integer counts
+    (hipDNN supplies shape ``[B, 1, 1, 1]``); a count is the valid (unpadded)
+    length for that batch item. Returns ``(len_q, len_kv)`` as contiguous
+    ``int32`` arrays of shape ``(batch,)`` for the shell to read.
+
+    Raises ``ValueError`` if a tensor is not integer-typed, does not hold
+    exactly ``batch`` counts, or carries a length outside ``[0, S_max]``.
+    """
+    len_q = _validate_lengths(seqlen_q, batch=batch, s_max=s_q_max, name="SEQ_LEN_Q")
+    len_kv = _validate_lengths(
+        seqlen_kv, batch=batch, s_max=s_kv_max, name="SEQ_LEN_KV"
+    )
+    return len_q, len_kv

dnn-providers/hip-kernel-provider/rocke/library/tests/test_sdpa_padding_mask.py

@@ -0,0 +1,86 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""Host-only tests for the SDPA padding-mask spec.
+
+``kernels.common._sdpa_padding_mask`` is the host-side specification of the
+padding half of the SDPA mask -- the ``keep_key`` / ``row_is_valid`` predicates
+the shell's mask hook and epilogue will emit. These tests lock the boundaries
+(``<`` not ``<=``) and prove the predicates compose to the same result as
+``dense_attention_reference`` with padding, so the kernel emitter has an exact
+host reference to match. No GPU.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+from kernels.common._sdpa_padding_mask import keep_key, row_is_valid
+from rocke.numeric.references import dense_attention_reference
+
+
+def _qkv(Sq=5, Sk=6, H=3, D=8, seed=0xBEEF):
+    rng = np.random.default_rng(seed)
+    Q = (rng.standard_normal((Sq, H, D)) * 0.3).astype(np.float16)
+    K = (rng.standard_normal((Sk, H, D)) * 0.3).astype(np.float16)
+    V = (rng.standard_normal((Sk, H, D)) * 0.3).astype(np.float16)
+    return Q, K, V
+
+
+def _emulate_kernel_padding(Q, K, V, *, len_q, len_kv):
+    """Apply the padding mask the way the shell will -- per-column ``keep_key``
+    into the score, per-row ``row_is_valid`` at the epilogue -- and run softmax,
+    so the predicates can be checked against the vectorized oracle."""
+    Sq, _, D = Q.shape
+    Sk = K.shape[0]
+    scores = np.einsum("ihd,jhd->ihj", Q.astype(np.float32), K.astype(np.float32))
+    scores /= math.sqrt(D)
+    keep = keep_key(np.arange(Sk), len_kv)  # (Sk,) bool
+    scores = np.where(keep[None, None, :], scores, -1e30)
+    scores -= scores.max(axis=-1, keepdims=True)
+    probs = np.exp(scores)
+    probs /= probs.sum(axis=-1, keepdims=True)
+    out = np.einsum("ihj,jhd->ihd", probs, V.astype(np.float32))
+    if not keep.any():  # fully masked (len_kv == 0): O = 0
+        out[:] = 0.0
+    valid = row_is_valid(np.arange(Sq), len_q)  # padded rows -> 0
+    return np.where(valid[:, None, None], out, 0.0)
+
+
+def test_keep_key_boundary():
+    # keep iff key < len_kv (strict), so len_kv itself is masked.
+    np.testing.assert_array_equal(
+        keep_key(np.arange(5), 3), [True, True, True, False, False]
+    )
+    assert keep_key(0, 1) and not keep_key(0, 0)
+
+
+def test_row_is_valid_boundary():
+    np.testing.assert_array_equal(
+        row_is_valid(np.arange(4), 2), [True, True, False, False]
+    )
+    assert row_is_valid(0, 1) and not row_is_valid(1, 1)
+
+
+def test_predicates_compose_to_oracle():
+    for len_q, len_kv in ((5, 4), (3, 6), (2, 2), (5, 6)):
+        Q, K, V = _qkv()
+        emulated = _emulate_kernel_padding(Q, K, V, len_q=len_q, len_kv=len_kv)
+        oracle = dense_attention_reference(
+            Q, K, V, causal=False, seqlen_q=len_q, seqlen_kv=len_kv
+        )
+        np.testing.assert_allclose(emulated, oracle, atol=1e-6, rtol=0)
+
+
+def test_fully_masked_and_padded_rows_are_zero():
+    Q, K, V = _qkv()
+    # len_kv = 0 -> every row fully masked -> all zero.
+    np.testing.assert_array_equal(
+        _emulate_kernel_padding(Q, K, V, len_q=5, len_kv=0),
+        np.zeros_like(Q, dtype=np.float32),
+    )
+    # len_q = 2 -> rows 2.. are padded -> zero; earlier rows nonzero.
+    out = _emulate_kernel_padding(Q, K, V, len_q=2, len_kv=6)
+    assert np.all(out[2:] == 0.0) and not np.all(out[:2] == 0.0)

dnn-providers/hip-kernel-provider/rocke/library/tests/test_sdpa_seqlen.py

@@ -0,0 +1,81 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""Host-only tests for the SDPA padding-length adapter.
+
+``kernels.common._sdpa_seqlen.sdpa_seqlens`` validates the hipDNN SEQ_LEN_Q /
+SEQ_LEN_KV padding tensors at the host boundary and normalizes them into the
+flat per-batch int32 arrays the SDPA shell reads. No GPU and no torch.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from kernels.common._sdpa_seqlen import sdpa_seqlens
+
+
+def test_valid_lengths_pass_through():
+    q = np.array([3, 5, 2, 8], dtype=np.int32)
+    kv = np.array([7, 7, 1, 8], dtype=np.int32)
+    len_q, len_kv = sdpa_seqlens(q, kv, batch=4, s_q_max=8, s_kv_max=8)
+    assert len_q.dtype == np.int32 and len_kv.dtype == np.int32
+    assert len_q.shape == (4,) and len_kv.shape == (4,)
+    assert len_q.flags["C_CONTIGUOUS"] and len_kv.flags["C_CONTIGUOUS"]
+    np.testing.assert_array_equal(len_q, q)
+    np.testing.assert_array_equal(len_kv, kv)
+
+
+def test_accepts_hipdnn_bshd_shape():
+    # hipDNN supplies SEQ_LEN as [B, 1, 1, 1].
+    q = np.array([2, 4], dtype=np.int32).reshape(2, 1, 1, 1)
+    kv = np.array([4, 4], dtype=np.int32).reshape(2, 1, 1, 1)
+    len_q, len_kv = sdpa_seqlens(q, kv, batch=2, s_q_max=4, s_kv_max=4)
+    assert len_q.shape == (2,) and len_kv.shape == (2,)
+    np.testing.assert_array_equal(len_q, [2, 4])
+
+
+def test_boundary_lengths_zero_and_max_allowed():
+    q = np.array([0, 6], dtype=np.int32)
+    kv = np.array([6, 0], dtype=np.int32)
+    len_q, len_kv = sdpa_seqlens(q, kv, batch=2, s_q_max=6, s_kv_max=6)
+    np.testing.assert_array_equal(len_q, [0, 6])
+    np.testing.assert_array_equal(len_kv, [6, 0])
+
+
+def test_length_above_max_rejected():
+    q = np.array([3, 9], dtype=np.int32)  # 9 > s_q_max
+    kv = np.array([4, 4], dtype=np.int32)
+    with pytest.raises(ValueError, match="SEQ_LEN_Q"):
+        sdpa_seqlens(q, kv, batch=2, s_q_max=8, s_kv_max=8)
+
+
+def test_negative_length_rejected():
+    q = np.array([3, 2], dtype=np.int32)
+    kv = np.array([-1, 4], dtype=np.int32)
+    with pytest.raises(ValueError, match="SEQ_LEN_KV"):
+        sdpa_seqlens(q, kv, batch=2, s_q_max=8, s_kv_max=8)
+
+
+def test_wrong_batch_count_rejected():
+    q = np.array([3, 2, 1], dtype=np.int32)  # 3 != batch
+    kv = np.array([4, 4], dtype=np.int32)
+    with pytest.raises(ValueError, match="SEQ_LEN_Q"):
+        sdpa_seqlens(q, kv, batch=2, s_q_max=8, s_kv_max=8)
+
+
+def test_non_integer_dtype_rejected():
+    q = np.array([3.0, 2.0], dtype=np.float32)
+    kv = np.array([4, 4], dtype=np.int32)
+    with pytest.raises(ValueError, match="integer dtype"):
+        sdpa_seqlens(q, kv, batch=2, s_q_max=8, s_kv_max=8)
+
+
+def test_q_and_kv_bounded_independently():
+    # Q valid up to s_q_max, KV up to a different s_kv_max.
+    q = np.array([4, 4], dtype=np.int32)
+    kv = np.array([16, 10], dtype=np.int32)
+    len_q, len_kv = sdpa_seqlens(q, kv, batch=2, s_q_max=4, s_kv_max=16)
+    np.testing.assert_array_equal(len_q, [4, 4])
+    np.testing.assert_array_equal(len_kv, [16, 10])

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/numeric/references.py

@@ -19,14 +19,25 @@
 __all__ = ["dense_attention_reference"]
 
 
-def dense_attention_reference(Q, K, V, *, causal: bool, out_dtype=None):
+def dense_attention_reference(
+    Q, K, V, *, causal: bool, seqlen_q=None, seqlen_kv=None, out_dtype=None
+):
     """Dense softmax-attention reference (fp32 math).
 
     ``Q``/``K``/``V`` are a single batch of shape
     ``(seqlen, heads, head_size)`` with KV heads already expanded to the query
     head count -- GQA/MQA expansion (``np.repeat`` over the head axis) is the
     caller's responsibility, as is looping the batch axis.
 
+    ``seqlen_q`` / ``seqlen_kv`` are this sequence's valid (unpadded) query and
+    key counts; positions at or beyond them are padding. ``None`` (the default
+    for both) keeps the full-length, unpadded behavior so existing callers are
+    byte-unchanged. When set, keys ``k >= seqlen_kv`` are masked out of every
+    softmax, a query row left with no valid key is fully masked and yields
+    ``O = 0`` (hipDNN's ``LSE = -inf`` convention), and padded query rows
+    ``q >= seqlen_q`` are zeroed. The padded-query-row zeroing is the one piece
+    that should be confirmed against the hipDNN padding reference.
+
     Returns fp32 unless ``out_dtype`` is given, in which case the result is
     cast to it (e.g. ``np.float16`` to model a dtype-truncated reference).
     """
@@ -37,8 +48,21 @@ def dense_attention_reference(Q, K, V, *, causal: bool, out_dtype=None):
         q_pos = np.arange(Q.shape[0])[:, None, None]
         k_pos = np.arange(K.shape[0])[None, None, :]
         scores = np.where(k_pos <= q_pos, scores, -1e30)
+    if seqlen_kv is not None:
+        k_pos = np.arange(K.shape[0])[None, None, :]
+        scores = np.where(k_pos < seqlen_kv, scores, -1e30)
     scores -= scores.max(axis=-1, keepdims=True)
     probs = np.exp(scores)
     probs /= probs.sum(axis=-1, keepdims=True)
     out = np.einsum("ihj,jhd->ihd", probs, V.astype(np.float32))
+    if seqlen_kv is not None:
+        # A row with no valid key is fully masked -> O = 0; the -1e30 softmax
+        # above would otherwise average all keys uniformly instead of zero.
+        k_idx = np.arange(K.shape[0])[None, None, :]
+        keep = k_idx < seqlen_kv
+        if causal:
+            keep = keep & (k_idx <= np.arange(Q.shape[0])[:, None, None])
+        out = np.where(keep.any(axis=-1, keepdims=True), out, 0.0)
+    if seqlen_q is not None:
+        out = np.where(np.arange(Q.shape[0])[:, None, None] < seqlen_q, out, 0.0)
     return out if out_dtype is None else out.astype(out_dtype)