fix(rocke): honour the block-table stride in the scalar attention fallback

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

@@ -797,18 +797,19 @@ def _cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple:
     ``arch`` sits at index 1 like every other key family in this module, so
     ``_assert_cache_key_arch`` can check key-vs-compile agreement uniformly and
     two arches can never alias to one cached kernel.
+
+    Like the tiled keys, it holds only fields the scalar 2D build reads. The
+    dynamic shape (``total_q``, ``num_seqs``, ``max_seqlen_q``, ``max_seqlen_k``)
+    reaches the kernel as kernargs or the launch grid, never as IR, so one binary
+    serves every batch. ``test_attention_scalar_cache_key.py`` enforces this.
     """
     return (
         "scalar",
         arch,
-        problem.total_q,
-        problem.num_seqs,
         problem.num_query_heads,
         problem.num_kv_heads,
         problem.head_size,
         problem.block_size,
-        problem.max_seqlen_q,
-        problem.max_seqlen_k,
         problem.dtype,
         problem.q_dtype,
         problem.sliding_window,
@@ -5745,7 +5746,7 @@ def _get_scalar_launcher(
         kernel_name=kname,
         signature=_attn_signature(
             problem.dtype,
-            include_bt_stride=False,
+            include_bt_stride=True,
             kv_dtype=_kv_storage_dtype(problem),
         ),
         cache_key=("scalar",) + cache_key,
@@ -6364,7 +6365,7 @@ def run_unified_attention_torch(
         softcap=softcap,
         sinks=sinks,
         bt_stride=bt_stride,
-        include_bt_stride=False,
+        include_bt_stride=True,
         k_scale=k_scale,
         v_scale=v_scale,
         out_scale=out_scale,
@@ -6406,7 +6407,6 @@ def kernel_name(self) -> str:
         p = self.problem
         return kernel_name_join(
             self.name,
-            f"q{p.total_q}",
             f"h{p.num_query_heads}",
             f"kv{p.num_kv_heads}",
             f"d{p.head_size}",
@@ -6511,6 +6511,7 @@ def build_unified_attention_2d(
     _out_scale = b.param("out_scale", F32)
     softcap = b.param("softcap", F32)
     num_seqs = b.param("num_seqs", I32)
+    bt_stride = b.param("block_table_stride", I32)
 
     q_tok = b.block_id_x()
     q_head = b.block_id_y()
@@ -6567,16 +6568,8 @@ def build_unified_attention_2d(
         iv_name="kpos",
     )
     with loop as (kpos, (m_val, l_val, acc_val)):
-        block_idx, token_in_block = _magic_div_mod(b, kpos, p.block_size)
-        physical = b.global_load_i32(
-            block_tables,
-            b.add(
-                b.mul(
-                    seq_idx,
-                    b.const_i32((p.max_seqlen_k + p.block_size - 1) // p.block_size),
-                ),
-                block_idx,
-            ),
+        physical, token_in_block = _physical_block_and_token(
+            b, p, block_tables, bt_stride, seq_idx, kpos
         )
 
         # Vectorised QK dot product. Both Q[q_tok, q_head, :] and
@@ -6690,7 +6683,6 @@ def kernel_name(self) -> str:
         p = self.problem
         return kernel_name_join(
             self.name,
-            f"q{p.total_q}",
             f"h{p.num_query_heads}",
             f"kv{p.num_kv_heads}",
             f"d{p.head_size}",
@@ -6748,6 +6740,7 @@ def build_unified_attention_3d(
     scale = abi["scale"]
     _softcap = b.param("softcap", F32)
     num_seqs = b.param("num_seqs", I32)
+    bt_stride = b.param("block_table_stride", I32)
 
     q_tok = b.block_id_x()
     q_head = b.block_id_y()
@@ -6798,6 +6791,7 @@ def build_unified_attention_3d(
             query,
             key,
             block_tables,
+            bt_stride,
             seq_idx,
             q_tok,
             q_head,
@@ -6823,7 +6817,9 @@ def build_unified_attention_3d(
         alpha = b.exp2(b.fsub(m_val, new_m))
         prob = b.exp2(b.fsub(score, new_m))
         new_l = b.fadd(b.fmul(l_val, alpha), prob)
-        vv = _emit_v_load(b, p, dtype, value, block_tables, seq_idx, kv_head, kpos, dim)
+        vv = _emit_v_load(
+            b, p, dtype, value, block_tables, bt_stride, seq_idx, kv_head, kpos, dim
+        )
         new_acc = b.fadd(b.fmul(acc_val, alpha), b.fmul(prob, vv))
         b.scf_yield(new_m, new_l, new_acc)
 
@@ -6860,7 +6856,6 @@ def kernel_name(self) -> str:
         p = self.problem
         return kernel_name_join(
             self.name,
-            f"q{p.total_q}",
             f"h{p.num_query_heads}",
             f"d{p.head_size}",
             f"seg{self.num_segments}",
@@ -7099,6 +7094,7 @@ def _emit_qk_score(
     query: Value,
     key: Value,
     block_tables: Value,
+    bt_stride: Value,
     seq_idx: Value,
     q_tok: Value,
     q_head: Value,
@@ -7127,7 +7123,7 @@ def _emit_qk_score(
     """
     score = b.const_f32(0.0)
     physical, token_in_block = _physical_block_and_token(
-        b, p, block_tables, seq_idx, kpos
+        b, p, block_tables, bt_stride, seq_idx, kpos
     )
     q_desc = _q_descriptor(p)
     kv_desc = _paged_kv_descriptor(p)
@@ -7178,13 +7174,14 @@ def _emit_v_load(
     dtype: Type,
     value: Value,
     block_tables: Value,
+    bt_stride: Value,
     seq_idx: Value,
     kv_head: Value,
     kpos: Value,
     dim: Value,
 ) -> Value:
     physical, token_in_block = _physical_block_and_token(
-        b, p, block_tables, seq_idx, kpos
+        b, p, block_tables, bt_stride, seq_idx, kpos
     )
     v_off = _paged_kv_descriptor(p).offset(
         b,
@@ -7200,12 +7197,17 @@ def _physical_block_and_token(
     b: IRBuilder,
     p: UnifiedAttentionProblem,
     block_tables: Value,
+    bt_stride: Value,
     seq_idx: Value,
     kpos: Value,
 ) -> Tuple[Value, Value]:
+    """Look up ``kpos``'s physical KV block in row ``seq_idx`` of the block table.
+
+    ``bt_stride`` is the table's row stride kernarg, not a value derived from
+    ``max_seqlen_k``: callers may allocate rows wider than the longest sequence.
+    """
     block_idx, token_in_block = _magic_div_mod(b, kpos, p.block_size)
-    max_blocks = (p.max_seqlen_k + p.block_size - 1) // p.block_size
     physical = b.global_load_i32(
-        block_tables, b.add(b.mul(seq_idx, b.const_i32(max_blocks)), block_idx)
+        block_tables, b.add(b.mul(seq_idx, bt_stride), block_idx)
     )
     return physical, token_in_block

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

@@ -0,0 +1,185 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""Scalar fallback launches with block-table rows wider than the batch needs.
+
+Callers usually size block-table rows for the longest sequence they admit, not
+the longest one in the current batch. The scalar kernels index rows by the
+table's real row stride (the ``block_table_stride`` kernarg), so every sequence
+after the first must still find its own blocks. Unused table slots point at a
+poison block, so a wrong row stride corrupts the output instead of passing by
+luck.
+
+The launch goes through ``run_unified_attention_torch(backend="scalar")`` with
+``rocke.runtime`` device buffers instead of torch tensors, so the test runs in
+GPU lanes without torch. Reference attention is computed independently in FP32
+with numpy from the logical K/V, after rounding inputs to fp16. Skips without a
+GPU, and on devices whose target id the attention builders do not know.
+"""
+
+import ctypes
+import itertools
+import math
+
+import numpy as np
+import pytest
+
+from kernels.common import attention_unified as au
+from rocke.core.arch.target import known_arches
+from rocke.runtime.hip_module import get_device_arch
+
+_HQ, _HKV, _D, _BS = 8, 2, 64, 16
+_ATOL = 1e-2
+
+# The launch compiles for the device's own target id, which the builders must
+# know (gfx11 devices are only known through their generic target).
+_ARCH = get_device_arch()
+pytestmark = [
+    pytest.mark.gpu,
+    pytest.mark.skipif(
+        _ARCH not in known_arches(),
+        reason=f"needs a GPU with a known attention target, found {_ARCH}",
+    ),
+]
+
+
+@pytest.fixture(scope="module")
+def rt():
+    from rocke.runtime.hip_module import Runtime
+
+    return Runtime()
+
+
+def _u8(array):
+    return (ctypes.c_uint8 * array.nbytes).from_buffer(array)
+
+
+class _DeviceArray:
+    """A device copy of a numpy array, exposing the tensor attributes
+    ``run_unified_attention_torch`` reads. ``shape`` may be narrower than the
+    host array along the last axis, which models a column view of it."""
+
+    device = None
+
+    def __init__(self, rt, host, *, shape=None):
+        from rocke.runtime.launcher import DeviceMem
+
+        self._rt = rt
+        self._host = np.ascontiguousarray(host)
+        self._mem = DeviceMem(self._host.nbytes)
+        rt.memcpy_h2d(self._mem.ptr(), _u8(self._host), self._host.nbytes)
+        self.shape = self._host.shape if shape is None else shape
+
+    def stride(self, dim):
+        return self._host.strides[dim] // self._host.itemsize
+
+    def data_ptr(self):
+        return self._mem.ptr()
+
+    def read(self):
+        out = np.empty_like(self._host)
+        self._rt.memcpy_d2h(_u8(out), self._mem.ptr(), out.nbytes)
+        return out
+
+
+def _case(q_lens, kv_lens, row_width, *, seed):
+    """A paged case whose block-table rows hold ``row_width`` slots."""
+    rng = np.random.default_rng(seed)
+    needed = [math.ceil(n / _BS) for n in kv_lens]
+    num_blocks = sum(needed) + 1
+    poison = num_blocks - 1
+    order = iter(rng.permutation(num_blocks - 1))
+    table = np.full((len(kv_lens), row_width), poison, dtype=np.int32)
+    for s, n in enumerate(needed):
+        table[s, :n] = [next(order) for _ in range(n)]
+
+    def randn(*shape):
+        return rng.standard_normal(shape).astype(np.float16)
+
+    q = randn(sum(q_lens), _HQ, _D)
+    k = randn(num_blocks, _BS, _HKV, _D)
+    v = randn(num_blocks, _BS, _HKV, _D)
+    k[poison] = 64.0
+    v[poison] = 64.0
+    cu = np.array([0, *itertools.accumulate(q_lens)], dtype=np.int32)
+    problem = au.UnifiedAttentionProblem(
+        total_q=sum(q_lens),
+        num_seqs=len(q_lens),
+        num_query_heads=_HQ,
+        num_kv_heads=_HKV,
+        head_size=_D,
+        block_size=_BS,
+        max_seqlen_q=max(q_lens),
+        max_seqlen_k=max(kv_lens),
+        dtype="fp16",
+    )
+    return problem, q, k, v, cu, table
+
+
+def _reference(q, k, v, cu, table, q_lens, kv_lens):
+    out = np.zeros(q.shape, dtype=np.float32)
+    ratio = _HQ // _HKV
+    for s, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)):
+        blocks = table[s, : math.ceil(kv_len / _BS)]
+        keys = k[blocks].reshape(-1, _HKV, _D)[:kv_len].astype(np.float32)
+        values = v[blocks].reshape(-1, _HKV, _D)[:kv_len].astype(np.float32)
+        context = kv_len - q_len
+        for i in range(q_len):
+            tok = int(cu[s]) + i
+            visible = context + i + 1
+            for h in range(_HQ):
+                scores = keys[:visible, h // ratio] @ q[tok, h].astype(np.float32)
+                scores /= math.sqrt(_D)
+                probs = np.exp(scores - scores.max())
+                out[tok, h] = (probs / probs.sum()) @ values[:visible, h // ratio]
+    return out
+
+
+def _launch(rt, problem, q, k, v, cu, kv_lens, block_table):
+    out = _DeviceArray(rt, np.zeros_like(q))
+    au.run_unified_attention_torch(
+        problem=problem,
+        q=_DeviceArray(rt, q),
+        k=_DeviceArray(rt, k),
+        v=_DeviceArray(rt, v),
+        out=out,
+        cu_seqlens_q=_DeviceArray(rt, cu),
+        seqused_k=_DeviceArray(rt, np.array(kv_lens, dtype=np.int32)),
+        softmax_scale=1.0 / math.sqrt(_D),
+        block_table=block_table,
+        softcap=0.0,
+        backend="scalar",
+    )
+    return out.read().astype(np.float32)
+
+
+@pytest.mark.parametrize("view_of_wider", (False, True), ids=("wide", "view"))
+def test_wide_block_table_rows(rt, view_of_wider):
+    q_lens, kv_lens = [5, 1, 9], [37, 20, 61]
+    row_width = math.ceil(max(kv_lens) / _BS) + 6
+    alloc_width = row_width + 5 if view_of_wider else row_width
+    problem, q, k, v, cu, table = _case(q_lens, kv_lens, alloc_width, seed=0)
+    dev = _DeviceArray(rt, table, shape=(len(q_lens), row_width))
+    assert dev.stride(0) > math.ceil(problem.max_seqlen_k / _BS)
+    got = _launch(rt, problem, q, k, v, cu, kv_lens, dev)
+    want = _reference(q, k, v, cu, table, q_lens, kv_lens)
+    max_abs = float(np.abs(got - want).max())
+    assert max_abs <= _ATOL, f"max_abs {max_abs:.3e} vs fp32 reference"
+
+
+def test_one_binary_serves_different_batches(rt):
+    """Shapes that differ only in batch geometry share a key; the second
+    launch must reuse the first binary and still match the reference."""
+    first = ([4, 2], [30, 17], 4)
+    second = ([1, 7, 3, 2], [90, 64, 15, 33], 12)
+    launchers = None
+    for i, (q_lens, kv_lens, width) in enumerate((first, second)):
+        problem, q, k, v, cu, table = _case(q_lens, kv_lens, width, seed=i)
+        got = _launch(rt, problem, q, k, v, cu, kv_lens, _DeviceArray(rt, table))
+        if launchers is None:
+            launchers = len(au._SCALAR_LAUNCHERS)
+        else:
+            assert len(au._SCALAR_LAUNCHERS) == launchers
+        want = _reference(q, k, v, cu, table, q_lens, kv_lens)
+        max_abs = float(np.abs(got - want).max())
+        assert max_abs <= _ATOL, f"case {i}: max_abs {max_abs:.3e}"