tests/test_kernels.py

"""Triton kernel correctness against CLUSTEN-free reference implementations.

The pre-existing forward/backward tests diff Triton against the compiled
CLUSTEN CUDA extension, which most installs will not have built, so they skip.
These tests use the pure-PyTorch references that ship in the kernel modules
instead, giving real kernel coverage on any CUDA device.
"""

import pytest
import torch

from conftest import requires_cuda
from affmae.ops.deform_attn_triton import dense_top4_knn
from affmae.ops.knn_torch import reference_dense_top4_knn
from affmae.ops.nbhd_attn_triton import (
    FlashNeighborhoodAttentionFunction,
    KVEdgeIndexCache,
    flash_nbhd_attn_reference_forward,
)

pytestmark = requires_cuda


def _neighbor_distances(coords, nn4, H, W):
    """Squared distances from each grid cell to the 4 neighbours chosen for it.

    Indices are not directly comparable between implementations: many
    candidates are equidistant, so Triton and the reference make different but
    equally valid tie-breaks. The distances are the real invariant.

    Note the output is indexed by *grid cell* ``p = y * W + x``, not by KV
    index, so the query position is derived from ``p`` rather than looked up in
    ``coords``.

    Args:
        coords: [1, Nk, 2] KV coordinates as (x, y).
        nn4: [1, H*W, 4] chosen neighbour indices into ``coords``.
        H: int, grid height.
        W: int, grid width.
    Returns:
        [H*W, 4] squared distances, sorted ascending per row.
    """
    pos = coords[0].float()
    p = torch.arange(H * W, device=coords.device)
    cells = torch.stack([(p % W).float(), (p // W).float()], dim=-1)
    neighbours = pos[nn4[0].long()]
    d2 = ((neighbours - cells[:, None, :]) ** 2).sum(-1)
    return d2.sort(dim=-1).values


def _grid_coords(grid, device):
    """All integer (x, y) positions of a ``grid x grid`` lattice, as [N, 2]."""
    ys, xs = torch.meshgrid(
        torch.arange(grid, device=device),
        torch.arange(grid, device=device),
        indexing="ij",
    )
    return torch.stack([xs.reshape(-1), ys.reshape(-1)], dim=-1)


class TestDenseTop4KNN:
    """The dense KNN table that feeds deformable attention."""

    @pytest.mark.parametrize("grid", [8, 16, 64])
    def test_matches_reference_on_full_grid(self, grid):
        device = "cuda"
        coords = _grid_coords(grid, device).unsqueeze(0).to(torch.int32)

        got = dense_top4_knn(coords, H=grid, W=grid)
        want = reference_dense_top4_knn(coords, H=grid, W=grid)

        assert got.shape == want.shape, f"{got.shape} != {want.shape}"
        torch.testing.assert_close(_neighbor_distances(coords, got, grid, grid),
                                   _neighbor_distances(coords, want, grid, grid),
                                   rtol=0, atol=0)

    def test_matches_reference_on_sparse_positions(self):
        """A subsampled point set, i.e. Nk far smaller than H*W."""
        device, grid = "cuda", 32
        gen = torch.Generator(device=device).manual_seed(0)
        coords = torch.randint(0, grid, (2, 200, 2), device=device,
                               generator=gen, dtype=torch.int32)

        got = dense_top4_knn(coords, H=grid, W=grid)
        want = reference_dense_top4_knn(coords, H=grid, W=grid)
        for b in range(coords.shape[0]):
            torch.testing.assert_close(
                _neighbor_distances(coords[b:b + 1], got[b:b + 1], grid, grid),
                _neighbor_distances(coords[b:b + 1], want[b:b + 1], grid, grid),
                rtol=0, atol=0)

    def test_indices_are_in_range(self):
        """Every returned index must address a real KV token."""
        device, grid, n_kv = "cuda", 16, 50
        gen = torch.Generator(device=device).manual_seed(1)
        coords = torch.randint(0, grid, (1, n_kv, 2), device=device,
                               generator=gen, dtype=torch.int32)
        # uint16 does not support reductions on CUDA, so widen first.
        nn4 = dense_top4_knn(coords, H=grid, W=grid).to(torch.int64)
        assert nn4.min().item() >= 0
        assert nn4.max().item() < n_kv, (
            f"index {nn4.max().item()} out of range for Nk={n_kv}")


class TestFlashNeighborhoodAttention:
    """The fused neighborhood attention used by the AFF encoder."""

    @staticmethod
    def _inputs(b, h, n, d, m, device, seed=0):
        gen = torch.Generator(device=device).manual_seed(seed)
        q = torch.randn(b, h, n, d, device=device, generator=gen)
        k = torch.randn(b, h, n, d, device=device, generator=gen)
        v = torch.randn(b, h, n, d, device=device, generator=gen)
        member = torch.randint(0, n, (b, n, m), device=device, generator=gen,
                               dtype=torch.int32)
        bias = torch.randn(b, h, n, m, device=device, generator=gen)
        blank_k = torch.randn(h, d, device=device, generator=gen)
        blank_v = torch.randn(h, d, device=device, generator=gen)
        return q, k, v, member, bias, blank_k, blank_v

    @pytest.mark.parametrize("b,h,n,d,m", [(1, 2, 64, 32, 16), (2, 4, 128, 32, 32)])
    def test_forward_matches_reference(self, b, h, n, d, m):
        device = "cuda"
        q, k, v, member, bias, bk, bv = self._inputs(b, h, n, d, m, device)
        scale = d ** -0.5

        got = FlashNeighborhoodAttentionFunction.apply(
            q, k, v, member, bias, None, bk, bv, float(scale))
        want, _lse = flash_nbhd_attn_reference_forward(
            q, k, v, member, scale, bias=bias, mask=None,
            blank_k=bk, blank_v=bv)

        torch.testing.assert_close(got.to(torch.float64), want.to(torch.float64),
                                   rtol=2e-3, atol=2e-3)

    def test_broadcast_mask_matches_expanded_mask(self):
        """[B,1,N,M] and [B,H,N,M] masks must give identical results.

        ClusterAttention passes the size-1 form to avoid materializing the
        per-head copy; this pins that the kernels really do broadcast it.
        """
        device = "cuda"
        b, h, n, d, m = 2, 4, 64, 32, 16
        q, k, v, member, bias, bk, bv = self._inputs(b, h, n, d, m, device, seed=3)
        gen = torch.Generator(device=device).manual_seed(4)
        mask_1 = torch.rand(b, 1, n, m, device=device, generator=gen) > 0.3
        mask_h = mask_1.expand(b, h, n, m).contiguous()
        scale = float(d ** -0.5)

        out_1 = FlashNeighborhoodAttentionFunction.apply(
            q, k, v, member, bias, mask_1, bk, bv, scale)
        out_h = FlashNeighborhoodAttentionFunction.apply(
            q, k, v, member, bias, mask_h, bk, bv, scale)

        torch.testing.assert_close(out_1, out_h, rtol=0, atol=0)

    def test_backward_gradients_match_reference(self):
        device = "cuda"
        b, h, n, d, m = 1, 2, 64, 32, 16
        q, k, v, member, bias, bk, bv = self._inputs(b, h, n, d, m, device, seed=7)
        scale = float(d ** -0.5)

        def run(fn):
            qs, ks, vs = (t.clone().detach().requires_grad_(True) for t in (q, k, v))
            if fn is FlashNeighborhoodAttentionFunction:
                out = fn.apply(qs, ks, vs, member, bias, None, bk, bv, scale)
            else:
                out, _ = fn(qs, ks, vs, member, scale, bias=bias, mask=None,
                            blank_k=bk, blank_v=bv)
            out.sum().backward()
            return qs.grad, ks.grad, vs.grad

        got = run(FlashNeighborhoodAttentionFunction)
        want = run(flash_nbhd_attn_reference_forward)

        for name, g, w in zip(("dq", "dk", "dv"), got, want):
            torch.testing.assert_close(
                g.to(torch.float64), w.to(torch.float64),
                rtol=5e-3, atol=5e-3, msg=lambda s, n=name: f"{n}: {s}")


class TestDeformBackendParity:
    """Both deform backends must exist and agree.

    ``deform_point_attn`` mapped ``backend="atomic"`` onto
    ``dv_backend="backend3_kv_atomic"``, but the dispatch had been renamed to
    ``"kv_atomic"`` without updating the caller, so the atomic path raised
    ``ValueError: Unknown dv_backend`` in the backward pass. With both live they
    cross-validate each other.
    """

    @staticmethod
    def _inputs(device, grid=8, n_heads=2, n_points=4, channels=16, seed=7):
        from affmae.ops.deform_attn_triton import dense_top4_knn

        coords = _grid_coords(grid, device).float()
        n = coords.shape[0]
        qpos = coords.unsqueeze(0).contiguous()
        kvpos = coords.unsqueeze(0).contiguous()
        nn4 = dense_top4_knn(
            kvpos.round().clamp(0, grid - 1).to(torch.int32), H=grid, W=grid)
        gen = torch.Generator(device=device).manual_seed(seed)
        return dict(
            query_pos=qpos, kv_pos=kvpos, nn4_idx=nn4, grid=grid,
            offsets=torch.randn(1, n, n_heads, n_points, 2, device=device, generator=gen),
            logits=torch.randn(1, n, n_heads, n_points, device=device, generator=gen),
            values=torch.randn(1, n, n_heads, channels, device=device, generator=gen),
            tau=torch.tensor(3.0, device=device),
        )

    def _run(self, backend, data):
        from affmae.ops.deform_attn_triton import deform_point_attn

        offs, logits, vals = (t.clone().detach().requires_grad_(True)
                              for t in (data["offsets"], data["logits"], data["values"]))
        out = deform_point_attn(
            query_pos=data["query_pos"], kv_pos=data["kv_pos"],
            sampling_offsets=offs, attn_logits=logits, values=vals,
            tau=data["tau"], nn4_idx=data["nn4_idx"],
            grid_h=data["grid"], grid_w=data["grid"], backend=backend)
        out.sum().backward()
        return out.detach(), vals.grad.clone()

    @pytest.mark.parametrize("backend", ["csr_knn_cached", "atomic"])
    def test_backend_runs_forward_and_backward(self, backend):
        out, dv = self._run(backend, self._inputs("cuda"))
        assert torch.isfinite(out).all()
        assert torch.isfinite(dv).all()

    def test_backends_agree(self):
        data = self._inputs("cuda")
        out_csr, dv_csr = self._run("csr_knn_cached", data)
        out_atomic, dv_atomic = self._run("atomic", data)
        torch.testing.assert_close(out_csr, out_atomic, rtol=1e-3, atol=1e-3)
        torch.testing.assert_close(dv_csr, dv_atomic, rtol=1e-3, atol=1e-3)

    def test_unknown_backend_is_rejected(self):
        from affmae.ops.deform_attn_triton import deform_point_attn

        data = self._inputs("cuda")
        with pytest.raises(ValueError, match="Unknown deform backend"):
            deform_point_attn(
                query_pos=data["query_pos"], kv_pos=data["kv_pos"],
                sampling_offsets=data["offsets"], attn_logits=data["logits"],
                values=data["values"], tau=data["tau"], nn4_idx=data["nn4_idx"],
                grid_h=data["grid"], grid_w=data["grid"], backend="nope")


def test_autotune_disk_cache_is_off_by_default():
    """One shared cache file plus many ranks is a write race.

    Also keeps ranks from independently selecting different kernels.
    """
    from affmae.ops import nbhd_attn_triton

    assert nbhd_attn_triton._CACHE_TO_DISK is False, (
        "CACHE_TO_DISK must default off; opt in per-run via the env var")


class TestInferenceSkipsBackwardOnlyWork:
    """The KV-owner edge index is read only by the backward pass.

    ``affmae/layers/aff.py`` installs a cache unconditionally, so before this
    guard every block built the index under ``torch.no_grad()`` -- about 0.7 ms
    per block, which made the v2 kernel 1.12x slower than v1 at inference
    despite being the faster kernel when training.
    """

    @staticmethod
    def _inputs(requires_grad, device="cuda", b=1, h=2, n=64, d=32, m=16):
        torch.manual_seed(3)
        make = lambda *shape: torch.randn(
            *shape, device=device, dtype=torch.float16
        ).requires_grad_(requires_grad)
        return dict(
            q=make(b, h, n, d), k=make(b, h, n, d), v=make(b, h, n, d),
            member=torch.randint(0, n, (b, n, m), device=device,
                                 dtype=torch.int32),
            bias=torch.randn(b, h, n, m, device=device, dtype=torch.float16),
            bk=make(h, d), bv=make(h, d), scale=float(d ** -0.5))

    def _apply(self, args):
        return FlashNeighborhoodAttentionFunction.apply(
            args["q"], args["k"], args["v"], args["member"], args["bias"],
            None, args["bk"], args["bv"], args["scale"])

    def test_no_grad_builds_no_edge_index(self):
        cache = KVEdgeIndexCache(max_entries=4)
        FlashNeighborhoodAttentionFunction.set_kv_edge_cache(cache)
        try:
            args = self._inputs(requires_grad=False)
            with torch.no_grad():
                self._apply(args)
            assert len(cache._cache) == 0, (
                "an edge index was built under no_grad; nothing reads it")
        finally:
            FlashNeighborhoodAttentionFunction.set_kv_edge_cache(None)

    def test_training_builds_it_in_the_forward_pass(self):
        """Skipping it when grad *is* needed makes backward rebuild it.

        Asserting only after ``.backward()`` is not enough: the backward
        launcher populates the cache itself, so that check passes even when the
        forward pass wrongly skipped it. Inspect the cache *before* backward.
        """
        cache = KVEdgeIndexCache(max_entries=4)
        FlashNeighborhoodAttentionFunction.set_kv_edge_cache(cache)
        try:
            args = self._inputs(requires_grad=True)
            out = self._apply(args)
            assert len(cache._cache) == 1, (
                "the forward pass skipped the edge index while training, so "
                "backward has to rebuild it")
            out.sum().backward()
            assert args["q"].grad is not None
        finally:
            FlashNeighborhoodAttentionFunction.set_kv_edge_cache(None)

    def test_forward_is_bit_identical_either_way(self):
        cache = KVEdgeIndexCache(max_entries=4)
        FlashNeighborhoodAttentionFunction.set_kv_edge_cache(cache)
        try:
            args = self._inputs(requires_grad=False)
            with torch.no_grad():
                without = self._apply(args)
            grad_args = dict(args)
            for name in ("q", "k", "v", "bk", "bv"):
                grad_args[name] = args[name].detach().requires_grad_(True)
            with_grad = self._apply(grad_args)
            torch.testing.assert_close(without, with_grad.detach(),
                                       rtol=0, atol=0)
        finally:
            FlashNeighborhoodAttentionFunction.set_kv_edge_cache(None)

class TestGradModeIsReadOutsideTheKernel:
    """Whether backward state is needed cannot be detected inside forward.

    Two traps, both of which produced a guard that looked right and was not:

    * ``torch.is_grad_enabled()`` is always False inside an
      ``autograd.Function.forward``, so a guard using it skips the work in
      training too -- and the cost merely moves into backward, which rebuilds
      the index and warns.
    * ``ctx.needs_input_grad`` is True even under ``no_grad`` when an input
      arrives as an ``nn.Parameter``, because a Parameter's ``requires_grad``
      flag does not depend on grad mode. Neighbourhood attention passes
      ``blank_k``/``blank_v`` straight in, so this fired on every inference call.

    The caller therefore reads grad mode and passes it down.
    """

    @staticmethod
    def _count_builds(fn):
        """Run ``fn`` with a stage cache installed, counting index builds.

        The cache is what ``affmae/layers/aff.py`` installs per encoder stage;
        without it the forward pass has nowhere to put an index and skips it
        regardless of grad mode, which would make these assertions vacuous.
        """
        import affmae.ops.nbhd_attn_triton as kernels

        calls = {"n": 0}
        original = kernels._build_kv_edge_index

        def counting(*args, **kwargs):
            calls["n"] += 1
            return original(*args, **kwargs)

        kernels._build_kv_edge_index = counting
        FlashNeighborhoodAttentionFunction.set_kv_edge_cache(
            KVEdgeIndexCache(max_entries=4))
        try:
            fn()
        finally:
            kernels._build_kv_edge_index = original
            FlashNeighborhoodAttentionFunction.set_kv_edge_cache(None)
        return calls["n"]

    @staticmethod
    def _module():
        from affmae.layers.attention import ClusterAttention

        torch.manual_seed(0)
        return ClusterAttention(dim=64, num_heads=4,
                                backend="flash_nbhd_attn").cuda().half()

    @staticmethod
    def _inputs(n=64, m=8):
        return (torch.randn(1, n, 64, device="cuda", dtype=torch.float16),
                torch.randint(0, n, (1, n, m), device="cuda", dtype=torch.int32),
                torch.randint(0, 200, (1, n, m), device="cuda", dtype=torch.int32))

    def test_inference_builds_none(self):
        module = self._module().eval()
        feat, member, pe = self._inputs()

        def run():
            with torch.no_grad():
                module(feat, member, None, pe, False)

        assert self._count_builds(run) == 0

    def test_training_builds_them_in_forward(self):
        module = self._module().train()
        feat, member, pe = self._inputs()

        def run():
            module(feat, member, None, pe, False)

        assert self._count_builds(run) > 0

    def test_output_does_not_depend_on_the_guard(self):
        module = self._module().eval()
        feat, member, pe = self._inputs()
        with torch.no_grad():
            without = module(feat, member, None, pe, False)
        with_grad = module(feat, member, None, pe, False)
        torch.testing.assert_close(without, with_grad.detach(), rtol=0, atol=0)


class TestDenseTop4LaunchParameters:
    """Warps are matched to the vector width, which must not change results."""

    def test_warps_cover_the_block_width(self):
        from affmae.ops.deform_attn_triton import _warps_for_width

        # Every value in the kernel is a [BLOCK_P] vector, so a program wider
        # than the data wastes lanes: 32 elements across 4 warps uses 8 of each
        # warp's 32 lanes.
        assert _warps_for_width(32) == 1
        assert _warps_for_width(64) == 2
        assert _warps_for_width(128) == 4

    def test_never_returns_zero_warps(self):
        """A block narrower than one warp still needs a warp."""
        from affmae.ops.deform_attn_triton import _warps_for_width

        assert _warps_for_width(16) == 1
        assert _warps_for_width(1) == 1

    @pytest.mark.parametrize("grid,n_kv", [(16, 64), (32, 512), (64, 4096)])
    def test_launch_parameters_do_not_change_the_table(self, grid, n_kv):
        """Bit-identical across warp counts, so this is purely a speed choice."""
        import affmae.ops.deform_attn_triton as kernels
        from affmae.ops.knn import clamp_to_grid

        torch.manual_seed(0)
        coords = clamp_to_grid(
            torch.rand(1, n_kv, 2, device="cuda") * (grid - 1), grid, grid)

        original = kernels._warps_for_width
        try:
            tables = []
            for warps in (1, 2, 4):
                kernels._warps_for_width = lambda _width, w=warps: w
                tables.append(kernels.dense_top4_knn(coords, H=grid, W=grid))
        finally:
            kernels._warps_for_width = original

        for other in tables[1:]:
            torch.testing.assert_close(tables[0], other, rtol=0, atol=0)


class TestTileDepthNeverUnderrunsHeadDim:
    """BLOCK_D >= HEAD_DIM is correctness, not tuning.

    The neighbourhood-attention kernel does not tile its depth loop, so a
    BLOCK_D smaller than HEAD_DIM never computes the remaining channels. The
    config pruner said as much in a comment and then violated it in its own
    fallback: it minimized over (BLOCK_Q, BLOCK_MN, BLOCK_D) together, so
    whenever the tile filters emptied the candidate list it happily returned the
    32-wide tile for a 64-wide head.

    No config offers BLOCK_MN below 8, so any neighbourhood smaller than 8
    emptied that list. head_dim=64 with a neighbourhood of 4 returned all-NaN;
    head_dim=32 was correct only because 32 happens to match. Reproduced on a
    GH200, so this was never ROCm-specific -- it surfaced during ROCm testing
    only because that sweep varied head_dim, which no test had.
    """

    @staticmethod
    def _configs():
        import affmae.ops.nbhd_attn_triton as kernels

        return next(v for name, v in vars(kernels).items()
                    if name.endswith("CONFIGS") and isinstance(v, list) and v
                    and "BLOCK_Q" in getattr(v[0], "kwargs", {}))

    @pytest.mark.parametrize("head_dim", [8, 16, 32, 64])
    @pytest.mark.parametrize("n_nb", [1, 2, 4, 8, 16, 32])
    def test_pruner_never_returns_a_shallow_tile(self, head_dim, n_nb):
        from affmae.ops.nbhd_attn_triton import _prune_tile_configs

        chosen = _prune_tile_configs(
            self._configs(),
            {"N_CTX": 64, "NEIGHBOR_SIZE": n_nb, "HEAD_DIM": head_dim})
        assert chosen, "the pruner must always return at least one config"
        for cfg in chosen:
            block_d = int(cfg.kwargs["BLOCK_D"])
            assert block_d >= head_dim, (
                f"BLOCK_D={block_d} < HEAD_DIM={head_dim} at neighbourhood "
                f"{n_nb}: channels {block_d}..{head_dim - 1} would never be "
                f"computed")

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
    @pytest.mark.parametrize("head_dim,n_nb", [(64, 4), (64, 2), (64, 1),
                                               (32, 4), (16, 4), (8, 4)])
    def test_output_is_finite_and_matches_the_reference(self, head_dim, n_nb):
        """The end-to-end symptom, at the shapes that used to produce NaN."""
        from affmae.ops.nbhd_attn import nbhd_attn
        from affmae.ops.nbhd_attn_torch import neighborhood_attention

        b, h, n, d, m = 1, 2, 64, head_dim, n_nb
        gen = torch.Generator(device="cuda").manual_seed(0)
        make = lambda *shape: torch.randn(*shape, device="cuda", generator=gen)
        q, k, v = make(b, h, n, d), make(b, h, n, d), make(b, h, n, d)
        member = torch.randint(0, n, (b, n, m), device="cuda", generator=gen,
                               dtype=torch.int32)
        bias, blank_k, blank_v = make(b, h, n, m), make(h, d), make(h, d)
        scale = float(d ** -0.5)

        triton_out = nbhd_attn(q, k, v, member, scale, bias=bias,
                               blank_k=blank_k, blank_v=blank_v,
                               backend="triton")
        reference = neighborhood_attention(q, k, v, member, scale, bias=bias,
                                           blank_k=blank_k, blank_v=blank_v)
        assert torch.isfinite(triton_out).all(), "non-finite kernel output"
        torch.testing.assert_close(triton_out, reference, rtol=2e-2, atol=5e-3)


class TestWarpCountFitsTheWavefront:
    """`num_warps` must not exceed what the tile can fill on a 64-lane device.

    The configs were templated for 32-lane NVIDIA warps. On an MI300X a
    wavefront is 64 lanes, so a config asking for 2 warps over an 8x8 tile
    reserves 128 lanes for 64 elements and the cross-warp reductions in the
    online softmax return garbage.

    Measured on that hardware, fp16, BLOCK_Q=BLOCK_MN=8: num_warps 1 finite,
    2 and 4 NaN in 60-70% of elements, identical at BLOCK_D 16 and 32 and at
    num_stages 1 and 2 -- so it is the warp count, not the tile depth or the
    pipeline depth. Clamping took a 270-cell ROCm sweep from 56 non-finite to 0.
    """

    def test_it_is_a_no_op_on_cuda(self):
        """The shipped configs are tuned at 32 lanes; do not silently retune."""
        from affmae.ops.dispatch import clamp_num_warps, is_hip

        if is_hip():
            pytest.skip("this asserts the CUDA branch")
        for requested, elements in ((1, 64), (2, 64), (4, 64), (8, 256), (4, 16)):
            assert clamp_num_warps(requested, elements) == requested

    def test_it_clamps_to_the_lane_budget_on_rocm(self, monkeypatch):
        from affmae.ops import dispatch

        monkeypatch.setattr(dispatch, "is_hip", lambda: True)
        # 64 elements is one 64-lane wavefront, so 2 or 4 warps cannot be filled.
        assert dispatch.clamp_num_warps(2, 64) == 1
        assert dispatch.clamp_num_warps(4, 64) == 1
        # 256 elements is four wavefronts, so up to 4 warps is honest.
        assert dispatch.clamp_num_warps(4, 256) == 4
        assert dispatch.clamp_num_warps(8, 256) == 4
        # Smaller than a single wavefront still needs one warp, never zero.
        assert dispatch.clamp_num_warps(4, 16) == 1

    def test_never_returns_zero(self, monkeypatch):
        from affmae.ops import dispatch

        monkeypatch.setattr(dispatch, "is_hip", lambda: True)
        for elements in (0, 1, 8, 63):
            assert dispatch.clamp_num_warps(4, elements) >= 1

    def test_every_shipped_config_fits_its_tile_on_rocm(self, monkeypatch):
        """Applied where templates become configs, so no config can violate it."""
        from affmae.ops import dispatch

        monkeypatch.setattr(dispatch, "is_hip", lambda: True)
        import importlib

        import affmae.ops.nbhd_attn_triton as kernels
        importlib.reload(kernels)
        try:
            configs = next(
                v for name, v in vars(kernels).items()
                if name.endswith("CONFIGS") and isinstance(v, list) and v
                and "BLOCK_Q" in getattr(v[0], "kwargs", {}))
            for cfg in configs:
                tile = int(cfg.kwargs["BLOCK_Q"]) * int(cfg.kwargs["BLOCK_MN"])
                assert cfg.num_warps * 64 <= max(tile, 64), (
                    f"{cfg.kwargs} asks for {cfg.num_warps} warps "
                    f"({cfg.num_warps * 64} lanes) over {tile} elements")
        finally:
            monkeypatch.undo()
            importlib.reload(kernels)