136 lines
6.0 KiB
Python
136 lines
6.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Routing tests for the opt-in sm_100a dispatch in ``block_sparse_attn_from_indices``.
|
|
|
|
``FASTVIDEO_VSA_SM100A=1`` routes the forward to the sm_100a extension when
|
|
``block_sparse_attn_sm100a.is_supported`` passes, pairing it with the Triton
|
|
backward (the sm_100a lse is already in Triton's M format). Everything else --
|
|
env unset, unsupported input, ``FASTVIDEO_VSA_TRITON`` override -- must keep
|
|
the pre-existing selection, bit-for-bit.
|
|
|
|
Run with: python -m pytest tests/test_block_sparse_sm100a_dispatch.py -v
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo_kernel import block_sparse_attn_sm100a as vsa
|
|
from fastvideo_kernel.block_sparse_attn import block_sparse_attn_from_indices
|
|
|
|
HEAD_DIM = 128
|
|
ENV = "FASTVIDEO_VSA_SM100A"
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0)
|
|
or not vsa._HAS_VSA_SM100A,
|
|
reason="requires Blackwell (sm_100a) and a built fastvideo_kernel extension",
|
|
)
|
|
|
|
|
|
def make_case(block, num_blocks=8, topk=4, heads=4, batch=1, seed=0, requires_grad=False):
|
|
torch.manual_seed(seed)
|
|
S = num_blocks * block
|
|
shape = (batch, heads, S, HEAD_DIM) if vsa.BHSD else (batch, S, heads, HEAD_DIM)
|
|
q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16,
|
|
requires_grad=requires_grad) for _ in range(3))
|
|
|
|
# from_indices takes Triton-shaped metadata: [B, H, Nq, Mk] / [B, H, Nq].
|
|
# The sm_100a extension reads both flat, so the same tensors serve both paths.
|
|
idx = torch.empty((batch * heads * num_blocks, topk), dtype=torch.int32, device="cuda")
|
|
for r in range(idx.shape[0]):
|
|
idx[r] = torch.randperm(num_blocks, device="cuda")[:topk].to(torch.int32).sort().values
|
|
idx = idx.view(batch, heads, num_blocks, topk)
|
|
num = torch.full((batch, heads, num_blocks), topk, dtype=torch.int32, device="cuda")
|
|
vbs = torch.full((num_blocks, ), block, dtype=torch.int32, device="cuda")
|
|
return q, k, v, idx, num, vbs
|
|
|
|
|
|
def triton_reference(q, k, v, idx, num, vbs, monkeypatch):
|
|
monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
|
|
out = block_sparse_attn_from_indices(q, k, v, idx, num, vbs)
|
|
monkeypatch.delenv("FASTVIDEO_VSA_TRITON")
|
|
return out
|
|
|
|
|
|
@pytest.mark.parametrize("block", [64, 128])
|
|
def test_env_routes_to_sm100a(block, monkeypatch):
|
|
"""With the env set on supported input, from_indices runs the sm_100a op:
|
|
bitwise-equal to the wrapper called directly, allclose to Triton."""
|
|
q, k, v, idx, num, vbs = make_case(block)
|
|
want_o, want_m = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs)
|
|
|
|
monkeypatch.setenv(ENV, "1")
|
|
got_o, got_m = block_sparse_attn_from_indices(q, k, v, idx, num, vbs)
|
|
assert torch.equal(got_o, want_o) and torch.equal(got_m, want_m)
|
|
|
|
if block == 64: # Triton's forward is hardcoded to 64-token blocks
|
|
ref_o, ref_m = triton_reference(q, k, v, idx, num, vbs, monkeypatch)
|
|
assert torch.allclose(got_o.float(), ref_o.float(), atol=2e-2, rtol=2e-2)
|
|
assert torch.allclose(got_m, ref_m, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_env_unset_keeps_default_backend(monkeypatch):
|
|
monkeypatch.delenv(ENV, raising=False)
|
|
q, k, v, idx, num, vbs = make_case(64)
|
|
got = block_sparse_attn_from_indices(q, k, v, idx, num, vbs)
|
|
ref = triton_reference(q, k, v, idx, num, vbs, monkeypatch)
|
|
assert torch.equal(got[0], ref[0]) and torch.equal(got[1], ref[1])
|
|
|
|
|
|
def test_unsupported_input_falls_back_to_triton(monkeypatch):
|
|
"""Env set but is_supported False (odd block count): silent Triton fallback."""
|
|
monkeypatch.setenv(ENV, "1")
|
|
q, k, v, idx, num, vbs = make_case(64, num_blocks=7, topk=3)
|
|
assert not vsa.is_supported(q, vbs)
|
|
got = block_sparse_attn_from_indices(q, k, v, idx, num, vbs)
|
|
ref = triton_reference(q, k, v, idx, num, vbs, monkeypatch)
|
|
assert torch.equal(got[0], ref[0]) and torch.equal(got[1], ref[1])
|
|
|
|
|
|
def test_unsupported_blk128_raises_not_implemented(monkeypatch):
|
|
"""128-token metadata cannot fall back to Triton's 64-token kernels."""
|
|
monkeypatch.setenv(ENV, "1")
|
|
q, k, v, idx, num, vbs = make_case(128, num_blocks=7, topk=3)
|
|
assert not vsa.is_supported(q, vbs)
|
|
with pytest.raises(NotImplementedError, match="128-token"):
|
|
block_sparse_attn_from_indices(q, k, v, idx, num, vbs)
|
|
|
|
|
|
def test_force_triton_overrides_sm100a(monkeypatch):
|
|
monkeypatch.setenv(ENV, "1")
|
|
monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
|
|
q, k, v, idx, num, vbs = make_case(64)
|
|
got = block_sparse_attn_from_indices(q, k, v, idx, num, vbs)
|
|
monkeypatch.delenv("FASTVIDEO_VSA_TRITON")
|
|
sm100a_o, _ = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs)
|
|
ref = triton_reference(q, k, v, idx, num, vbs, monkeypatch)
|
|
assert torch.equal(got[0], ref[0])
|
|
assert not torch.equal(got[0], sm100a_o)
|
|
|
|
|
|
def test_backward_runs_triton_and_matches(monkeypatch):
|
|
"""sm_100a forward + Triton backward: grads match the all-Triton path."""
|
|
monkeypatch.setenv(ENV, "1")
|
|
q, k, v, idx, num, vbs = make_case(64, requires_grad=True)
|
|
out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs)
|
|
out.float().square().sum().backward()
|
|
got = [t.grad.float().clone() for t in (q, k, v)]
|
|
|
|
q2, k2, v2 = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
|
|
monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
|
|
out2, _ = block_sparse_attn_from_indices(q2, k2, v2, idx, num, vbs)
|
|
out2.float().square().sum().backward()
|
|
ref = [t.grad.float() for t in (q2, k2, v2)]
|
|
|
|
for g, r, name in zip(got, ref, "qkv"):
|
|
assert torch.allclose(g, r, atol=5e-2, rtol=5e-2), \
|
|
f"d{name} max|diff|={(g - r).abs().max().item()}"
|
|
|
|
|
|
def test_blk128_backward_raises(monkeypatch):
|
|
"""128-token blocks: forward runs, backward refuses (Triton bwd is 64-block only)."""
|
|
monkeypatch.setenv(ENV, "1")
|
|
q, k, v, idx, num, vbs = make_case(128, requires_grad=True)
|
|
out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs)
|
|
with pytest.raises(RuntimeError, match="64-token"):
|
|
out.float().sum().backward()
|