hao-ai-lab-FastVideo/fastvideo-kernel/tests/test_vsa128_backward.py

"""VSA-128 CuTe/Triton forward and backward parity on Blackwell."""

from __future__ import annotations

import math

import pytest
import torch

from fastvideo_kernel import video_sparse_attn, video_sparse_attn_bshd
from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128

from .test_vsa256_triton import _metrics, _torch_vsa256_reference

_BLOCK = 128
_BLOCK_SIZE_3D = (2, 8, 8)


def _select_backend(monkeypatch, backend: str) -> None:
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")
    if backend == "cute":
        pytest.importorskip(
            "flash_attn.cute.block_sparsity",
            reason="optional FA4 CuTe build (flash_attn.cute) not installed",
        )
        monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
        monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False)
        monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False)
    else:
        monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
        monkeypatch.delenv("FASTVIDEO_VSA_CUTEDSL", raising=False)


def _dense_sparse_reference(q, k, v, block_map, variable_block_sizes):
    token_mask = block_map.repeat_interleave(_BLOCK, dim=2).repeat_interleave(_BLOCK, dim=3)
    kv_valid = torch.arange(_BLOCK, device=k.device) < variable_block_sizes[:, None]
    token_mask = token_mask & kv_valid.reshape(1, 1, 1, -1)
    logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) / math.sqrt(q.shape[-1])
    probabilities = torch.softmax(logits.masked_fill(~token_mask, float("-inf")), dim=-1)
    return torch.matmul(probabilities, v.float()).to(q.dtype)


def _check(tag: str, expected: torch.Tensor, actual: torch.Tensor, avg_tol: float, rel_tol: float) -> None:
    assert torch.isfinite(actual).all().item(), f"{tag}: non-finite values"
    avg_abs, max_rel = _metrics(expected, actual)
    print(f"  {tag}: avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
    assert avg_abs < avg_tol
    assert max_rel < rel_tol


@pytest.mark.cuda
@pytest.mark.parametrize("backend", ["cute", "triton"])
def test_vsa128_explicit_routes_forward_backward(backend: str, monkeypatch) -> None:
    """Adjacent Q128 blocks must keep independent routes instead of merging."""
    _select_backend(monkeypatch, backend)
    torch.manual_seed(53)
    shape = (1, 1, 3 * _BLOCK, 128)
    base = [torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)]
    grad_output = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
    variable_block_sizes = torch.tensor([128, 91, 37], device="cuda", dtype=torch.int32)
    block_map = torch.eye(3, device="cuda", dtype=torch.bool).view(1, 1, 3, 3)

    actual_inputs = [tensor.detach().clone().requires_grad_() for tensor in base]
    actual, _ = block_sparse_attn_128(*actual_inputs, block_map, variable_block_sizes)
    (actual * grad_output).sum().backward()

    reference_inputs = [tensor.detach().clone().requires_grad_() for tensor in base]
    expected = _dense_sparse_reference(*reference_inputs, block_map, variable_block_sizes)
    (expected * grad_output).sum().backward()

    print(f"[vsa128-explicit-{backend}]")
    _check("out", expected, actual, 1e-3, 0.2)
    for name, reference, candidate in zip(("dq", "dk", "dv"), reference_inputs, actual_inputs, strict=True):
        _check(name, reference.grad, candidate.grad, 2e-2, 0.5)


def _zero_kv_tail(x: torch.Tensor, variable_block_sizes: torch.Tensor) -> torch.Tensor:
    valid = torch.arange(_BLOCK, device=x.device) < variable_block_sizes[:, None]
    valid = valid.view(1, 1, -1, _BLOCK, 1).expand_as(x.view(1, x.shape[1], -1, _BLOCK, x.shape[-1]))
    return x * valid.reshape_as(x).to(x.dtype)


@pytest.mark.cuda
@pytest.mark.parametrize("backend", ["cute", "triton"])
@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
def test_vsa128_wrapper_forward_backward(backend: str, layout: str, monkeypatch) -> None:
    _select_backend(monkeypatch, backend)
    torch.manual_seed(59)
    batch, heads, dim = 1, 2, 128
    q_blocks, kv_blocks, topk = 3, 4, 2
    q_shape = (batch, heads, q_blocks * _BLOCK, dim)
    kv_shape = (batch, heads, kv_blocks * _BLOCK, dim)
    q_base = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16)
    kv_sizes = torch.tensor([128, 91, 37, 128], device="cuda", dtype=torch.int32)
    q_sizes = torch.full((q_blocks, ), _BLOCK, device="cuda", dtype=torch.int32)
    k_base = _zero_kv_tail(torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16), kv_sizes)
    v_base = _zero_kv_tail(torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16), kv_sizes)
    gate_base = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16) * 0.1
    grad_output = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16)

    if layout == "bhsd":
        actual_inputs = [tensor.detach().clone().requires_grad_() for tensor in (q_base, k_base, v_base)]
        actual_gate = gate_base.detach().clone().requires_grad_()
        actual = video_sparse_attn(
            *actual_inputs,
            kv_sizes,
            q_sizes,
            topk,
            block_size=_BLOCK_SIZE_3D,
            compress_attn_weight=actual_gate,
        )
        actual_grads = actual_inputs
    else:
        bshd_inputs = [tensor.transpose(1, 2).contiguous().detach().requires_grad_()
                       for tensor in (q_base, k_base, v_base)]
        bshd_gate = gate_base.transpose(1, 2).contiguous().detach().requires_grad_()
        actual = video_sparse_attn_bshd(
            *bshd_inputs,
            kv_sizes,
            q_sizes,
            topk,
            block_size=_BLOCK_SIZE_3D,
            compress_attn_weight=bshd_gate,
        ).transpose(1, 2)
        actual_grads = bshd_inputs
    (actual * grad_output).sum().backward()

    reference_inputs = [tensor.detach().clone().requires_grad_() for tensor in (q_base, k_base, v_base)]
    reference_gate = gate_base.detach().clone().requires_grad_()
    expected = _torch_vsa256_reference(
        *reference_inputs,
        q_sizes,
        kv_sizes,
        topk,
        compress_attn_weight=reference_gate,
    )
    (expected * grad_output).sum().backward()

    print(f"[vsa128-wrapper-{backend}-{layout}]")
    _check("out", expected, actual, 1e-3, 0.2)
    for name, reference, candidate in zip(("dq", "dk", "dv"), reference_inputs, actual_grads, strict=True):
        candidate_grad = candidate.grad if layout == "bhsd" else candidate.grad.transpose(1, 2)
        _check(name, reference.grad, candidate_grad, 2e-2, 0.5)
    actual_gate_grad = actual_gate.grad if layout == "bhsd" else bshd_gate.grad.transpose(1, 2)
    _check("dgate", reference_gate.grad, actual_gate_grad, 1e-3, 0.2)