[test]: VSA-256 forward / cross-backend / vbs / triton parity tests

Four GPU correctness tests covering the new VSA-256 path:

- test_vsa256_forward.py:        CuTe forward vs dense torch reference
                                 (qk_equal and qk_diff shapes).
- test_vsa256_forward_vbs.py:    CuTe forward with variable KV block sizes
                                 (<256) vs token-masked torch reference.
- test_vsa256_triton.py:         Route-A Triton fwd + bwd vs token-masked
                                 torch reference (forces FASTVIDEO_VSA_TRITON).
- test_vsa256_forward_cross.py:  Three-way parity: torch ref vs CuTe vs
                                 Triton on the same inputs.

All tests skip cleanly when CUDA is unavailable. They exercise the public
fastvideo_kernel.video_sparse_attn entrypoint with block_size=(4, 8, 8)
and rely on its built-in block_elements=256 dispatch.

Author: alexzms

fastvideo-kernel/tests/test_vsa256_forward.py

@@ -0,0 +1,116 @@
+"""VSA-256 CuTe forward correctness vs. a dense torch reference."""
+
+from __future__ import annotations
+
+import math
+from typing import Tuple
+
+import pytest
+import torch
+
+from fastvideo_kernel import video_sparse_attn
+
+from .utils import create_full_mask_from_block_mask
+
+
+def _dense_reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    full_mask: torch.Tensor,
+) -> torch.Tensor:
+    # q,k,v: [B,H,S,D]; full_mask: [H,Sq,Skv].
+    qf = q.float()
+    kf = k.float()
+    vf = v.float()
+    attn = torch.matmul(qf, kf.transpose(-2, -1)) / math.sqrt(q.shape[-1])
+    attn = attn.masked_fill(~full_mask.unsqueeze(0), float("-inf"))
+    prob = torch.softmax(attn, dim=-1)
+    return torch.matmul(prob, vf).to(q.dtype)
+
+
+def _run_case(
+    heads: int,
+    head_dim: int,
+    q_blocks_256: int,
+    kv_blocks_256: int,
+    topk_logical: int,
+) -> Tuple[float, float]:
+    assert torch.cuda.is_available()
+    # Default dispatch routes (4,8,8) -> CuTe; no env var needed.
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    batch = 1
+    q_block = 256
+    kv_block_logical = 256
+
+    sq = q_blocks_256 * q_block
+    skv = kv_blocks_256 * kv_block_logical
+    q = torch.randn(batch, heads, sq, head_dim, device=device, dtype=dtype)
+    k = torch.randn(batch, heads, skv, head_dim, device=device, dtype=dtype)
+    v = torch.randn(batch, heads, skv, head_dim, device=device, dtype=dtype)
+
+    q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device)
+    kv_var = torch.full((kv_blocks_256,), kv_block_logical, dtype=torch.int32, device=device)
+
+    # Reproduce the compression branch's per-block average so we can pre-compute
+    # the top-k mask that the kernel will see and feed it to the dense reference.
+    q_c = q.view(batch, heads, q_blocks_256, q_block, head_dim)
+    k_c = k.view(batch, heads, kv_blocks_256, kv_block_logical, head_dim)
+    v_c = v.view(batch, heads, kv_blocks_256, kv_block_logical, head_dim)
+    q_c = (q_c.float().sum(dim=3) / q_var.view(1, 1, -1, 1)).to(q.dtype)
+    k_c = (k_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(k.dtype)
+    v_c = (v_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(v.dtype)
+
+    scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (head_dim ** 0.5)
+    attn = torch.softmax(scores, dim=-1)
+    out_c = torch.matmul(attn, v_c)
+    out_c = (
+        out_c.view(batch, heads, q_blocks_256, 1, head_dim)
+        .repeat(1, 1, 1, q_block, 1)
+        .view(batch, heads, sq, head_dim)
+    )
+    topk_idx = torch.topk(scores, topk_logical, dim=-1).indices
+    mask_256 = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)[0]
+    full_mask = create_full_mask_from_block_mask(mask_256, q_var, kv_var, device=device)
+    out_ref = out_c + _dense_reference(q, k, v, full_mask)
+
+    out = video_sparse_attn(
+        q, k, v,
+        kv_var,
+        q_var,
+        topk_logical,
+        block_size=(4, 8, 8),
+        compress_attn_weight=None,
+    )
+
+    assert torch.isfinite(out).all().item(), "NaN/Inf in kernel output"
+    diff = (out_ref - out).abs()
+    avg_abs = diff.mean().item()
+    max_rel = (diff.max() / (out_ref.abs().mean() + 1e-6)).item()
+    return avg_abs, max_rel
+
+
+@pytest.mark.cuda
+def test_vsa256_forward_qk_equal() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    avg_abs, max_rel = _run_case(
+        heads=8, head_dim=128, q_blocks_256=8, kv_blocks_256=8, topk_logical=2,
+    )
+    print(f"[vsa256 qk_equal] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
+    assert avg_abs < 5e-2
+    assert max_rel < 2.0
+
+
+@pytest.mark.cuda
+def test_vsa256_forward_qk_diff() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    avg_abs, max_rel = _run_case(
+        heads=8, head_dim=128, q_blocks_256=8, kv_blocks_256=12, topk_logical=2,
+    )
+    print(f"[vsa256 qk_diff] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
+    assert avg_abs < 5e-2
+    assert max_rel < 2.0

fastvideo-kernel/tests/test_vsa256_forward_cross.py

@@ -0,0 +1,124 @@
+"""Three-way parity for the VSA-256 forward: torch reference vs CuTe vs Triton."""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+import torch
+
+from fastvideo_kernel import video_sparse_attn
+
+
+def _torch_vsa256_reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_var: torch.Tensor,
+    kv_var: torch.Tensor,
+    topk_logical: int,
+) -> torch.Tensor:
+    bsz, heads, _sq, dim = q.shape
+    q_blocks = q_var.numel()
+    kv_blocks = kv_var.numel()
+    q_block = int(q_var[0].item())
+    kv_block = int(kv_var[0].item())
+
+    q_c = q.view(bsz, heads, q_blocks, q_block, dim)
+    k_c = k.view(bsz, heads, kv_blocks, kv_block, dim)
+    v_c = v.view(bsz, heads, kv_blocks, kv_block, dim)
+    q_c = (q_c.float().sum(dim=3) / q_var.view(1, 1, -1, 1)).to(q.dtype)
+    k_c = (k_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(k.dtype)
+    v_c = (v_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(v.dtype)
+
+    scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / math.sqrt(dim)
+    attn = torch.softmax(scores, dim=-1)
+    out_c = torch.matmul(attn, v_c)
+    out_c = (
+        out_c.view(bsz, heads, q_blocks, 1, dim)
+        .repeat(1, 1, 1, q_block, 1)
+        .view_as(q)
+    )
+
+    topk_idx = torch.topk(scores, topk_logical, dim=-1).indices
+    block_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
+    token_mask = (
+        block_mask
+        .repeat_interleave(q_block, dim=2)
+        .repeat_interleave(kv_block, dim=3)
+        .to(torch.bool)
+    )
+
+    qf, kf, vf = q.float(), k.float(), v.float()
+    logits = torch.matmul(qf, kf.transpose(-2, -1)) / math.sqrt(dim)
+    logits = logits.masked_fill(~token_mask, float("-inf"))
+    prob = torch.softmax(logits, dim=-1)
+    out_s = torch.matmul(prob, vf).to(q.dtype)
+    return out_c + out_s
+
+
+def _metrics(a: torch.Tensor, b: torch.Tensor) -> tuple[float, float]:
+    diff = (a - b).abs()
+    avg_abs = diff.mean().item()
+    max_rel = (diff.max() / (a.abs().mean() + 1e-6)).item()
+    return avg_abs, max_rel
+
+
+@pytest.mark.cuda
+def test_vsa256_forward_cross_torch_cute_triton(monkeypatch) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    bsz, heads, dim = 1, 8, 128
+    q_blocks_256, kv_blocks_256 = 8, 12
+    topk_logical = 2
+    q_block = 256
+    kv_block = 256
+
+    sq = q_blocks_256 * q_block
+    skv = kv_blocks_256 * kv_block
+    q = torch.randn(bsz, heads, sq, dim, device=device, dtype=dtype)
+    k = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype)
+    v = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype)
+    q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device)
+    kv_var = torch.full((kv_blocks_256,), kv_block, dtype=torch.int32, device=device)
+
+    out_torch = _torch_vsa256_reference(q, k, v, q_var, kv_var, topk_logical)
+
+    # CuTe (default).
+    out_cute = video_sparse_attn(
+        q, k, v,
+        kv_var, q_var, topk_logical,
+        block_size=(4, 8, 8),
+        compress_attn_weight=None,
+    )
+
+    # Triton via route-A 256->64 expansion.
+    monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
+    out_triton = video_sparse_attn(
+        q, k, v,
+        kv_var, q_var, topk_logical,
+        block_size=(4, 8, 8),
+        compress_attn_weight=None,
+    )
+
+    for t in (out_torch, out_cute, out_triton):
+        assert torch.isfinite(t).all().item()
+
+    torch_vs_cute = _metrics(out_torch, out_cute)
+    torch_vs_triton = _metrics(out_torch, out_triton)
+    cute_vs_triton = _metrics(out_cute, out_triton)
+    print(
+        "[cross-forward] "
+        f"torch_vs_cute(avg_abs={torch_vs_cute[0]:.6e}, max_rel={torch_vs_cute[1]:.6e}), "
+        f"torch_vs_triton(avg_abs={torch_vs_triton[0]:.6e}, max_rel={torch_vs_triton[1]:.6e}), "
+        f"cute_vs_triton(avg_abs={cute_vs_triton[0]:.6e}, max_rel={cute_vs_triton[1]:.6e})"
+    )
+
+    assert torch_vs_cute[0] < 1e-3 and torch_vs_cute[1] < 0.2
+    assert torch_vs_triton[0] < 1e-3 and torch_vs_triton[1] < 0.2
+    assert cute_vs_triton[0] < 1e-3 and cute_vs_triton[1] < 0.2

fastvideo-kernel/tests/test_vsa256_forward_vbs.py

@@ -0,0 +1,115 @@
+"""VSA-256 CuTe correctness with variable KV block sizes (<256)."""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+import torch
+
+from fastvideo_kernel import video_sparse_attn
+
+
+def _torch_vsa256_reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_var: torch.Tensor,
+    kv_var: torch.Tensor,
+    topk_logical: int,
+) -> torch.Tensor:
+    bsz, heads, _sq, dim = q.shape
+    q_blocks = q_var.numel()
+    kv_blocks = kv_var.numel()
+    q_block = q.shape[2] // q_blocks
+    kv_block = k.shape[2] // kv_blocks
+
+    q_c = q.view(bsz, heads, q_blocks, q_block, dim)
+    k_c = k.view(bsz, heads, kv_blocks, kv_block, dim)
+    v_c = v.view(bsz, heads, kv_blocks, kv_block, dim)
+    q_c = (q_c.float().sum(dim=3) / q_var.view(1, 1, -1, 1)).to(q.dtype)
+    k_c = (k_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(k.dtype)
+    v_c = (v_c.float().sum(dim=3) / kv_var.view(1, 1, -1, 1)).to(v.dtype)
+
+    scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / math.sqrt(dim)
+    attn = torch.softmax(scores, dim=-1)
+    out_c = torch.matmul(attn, v_c)
+    out_c = (
+        out_c.view(bsz, heads, q_blocks, 1, dim)
+        .repeat(1, 1, 1, q_block, 1)
+        .view_as(q)
+    )
+
+    with torch.no_grad():
+        topk_idx = torch.topk(scores.detach(), topk_logical, dim=-1).indices
+        block_mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
+        block_token_idx = torch.arange(kv_block, device=kv_var.device, dtype=torch.int32)
+        kv_token_valid_by_block = (
+            block_token_idx.view(1, -1) < kv_var.to(torch.int32).view(-1, 1)
+        ).to(torch.bool)
+        kv_token_valid = kv_token_valid_by_block.reshape(1, 1, 1, kv_blocks * kv_block)
+        token_mask = (
+            block_mask
+            .repeat_interleave(q_block, dim=2)
+            .repeat_interleave(kv_block, dim=3)
+        )
+        token_mask = token_mask & kv_token_valid
+
+    qf, kf, vf = q.float(), k.float(), v.float()
+    logits = torch.matmul(qf, kf.transpose(-2, -1)) / math.sqrt(dim)
+    logits = logits.masked_fill(~token_mask, float("-inf"))
+    prob = torch.softmax(logits, dim=-1)
+    out_s = torch.matmul(prob, vf).to(q.dtype)
+    return out_c + out_s
+
+
+@pytest.mark.cuda
+def test_vsa256_cute_variable_block_size_vs_torch_ref() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    bsz, heads, dim = 1, 8, 128
+    q_blocks_256, kv_blocks_256 = 8, 12
+    q_block = 256
+    kv_block = 256
+    topk_logical = 2
+    sq = q_blocks_256 * q_block
+    skv = kv_blocks_256 * kv_block
+
+    q = torch.randn(bsz, heads, sq, dim, device=device, dtype=dtype)
+    k = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype)
+    v = torch.randn(bsz, heads, skv, dim, device=device, dtype=dtype)
+
+    q_var = torch.full((q_blocks_256,), q_block, dtype=torch.int32, device=device)
+    kv_var = torch.randint(16, kv_block + 1, (kv_blocks_256,), dtype=torch.int32, device=device)
+    token_idx = torch.arange(kv_block, device=device, dtype=torch.int32)
+    kv_valid = token_idx.view(1, -1) < kv_var.view(-1, 1)
+    kv_valid = kv_valid.view(1, 1, kv_blocks_256, kv_block, 1)
+    kv_valid = kv_valid.expand(bsz, heads, kv_blocks_256, kv_block, dim).reshape(bsz, heads, skv, dim)
+    k = k * kv_valid.to(k.dtype)
+    v = v * kv_valid.to(v.dtype)
+
+    out = video_sparse_attn(
+        q, k, v,
+        kv_var,
+        q_var,
+        topk_logical,
+        block_size=(4, 8, 8),
+        compress_attn_weight=None,
+    )
+    out_ref = _torch_vsa256_reference(q, k, v, q_var, kv_var, topk_logical)
+
+    diff = (out - out_ref).abs()
+    avg_abs = diff.mean().item()
+    max_rel = (diff.max() / (out_ref.abs().mean() + 1e-6)).item()
+    print(
+        f"[vsa256-cute-vbs] kv_var[min={int(kv_var.min().item())}, "
+        f"max={int(kv_var.max().item())}], "
+        f"avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}"
+    )
+
+    assert avg_abs < 1e-3 and max_rel < 0.2