diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index 40617e993..5aab846d8 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -49,6 +49,36 @@ def _force_tk() -> bool: return os.environ.get("FASTVIDEO_VSA_TK", "0") == "1" +def _force_sm100a() -> bool: + """True iff the sm_100a (Blackwell) forward is explicitly opted into. + + Opt-in only (same env the H3 backend honors): the sm_100a extension is + forward-only, so this routing pairs it with the Triton backward -- its lse + is already in Triton's M format. Honored only when + ``block_sparse_attn_sm100a.is_supported`` passes. Unsupported 64-token + metadata falls through to the default selection; unsupported 128-token + metadata raises because Triton has no compatible fallback. + ``FASTVIDEO_VSA_TRITON`` still wins. + """ + return os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1" + + +def _sm100a_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: + try: + from fastvideo_kernel import block_sparse_attn_sm100a as vsa_sm100a + except Exception: + return False + return vsa_sm100a.is_supported(q, variable_block_sizes) + + +def _infer_block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int: + num_blocks = variable_block_sizes.numel() + seq_len = q.shape[2] + if num_blocks == 0 or seq_len % num_blocks != 0: + return 0 + return seq_len // num_blocks + + # --------------------------------------------------------------------------- # Index helpers # --------------------------------------------------------------------------- @@ -339,6 +369,76 @@ def _backward_sm90(ctx, grad_o, grad_lse): block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90) +# --------------------------------------------------------------------------- +# SM100A backend custom op (index-native) +# +# Forward runs the sm_100a CUDA extension; backward reuses the Triton kernels. +# The sm_100a forward emits lse in exactly Triton's M format (max*log2e + +# log2(l)), so the pairing needs no conversion. The Triton backward is +# hardcoded to 64-token blocks, hence the block-size assert below. +# --------------------------------------------------------------------------- + + +@torch.library.custom_op( + "fastvideo_kernel::block_sparse_attn_sm100a", + mutates_args=(), + device_types="cuda", +) +def block_sparse_attn_sm100a_op( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + from fastvideo_kernel.block_sparse_attn_sm100a import block_sparse_attn_sm100a + + o, M = block_sparse_attn_sm100a(q, k, v, q2k_idx, q2k_num, variable_block_sizes, + need_lse=True) + return o, M + + +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a") +def _block_sparse_attn_sm100a_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + o = torch.empty_like(q) + M = torch.empty( + (q.shape[0], q.shape[1], q.shape[2]), + device=q.device, + dtype=torch.float32, + ) + return o, M + + +def _setup_context_sm100a(ctx, inputs, output): + q, k, v, q2k_idx, q2k_num, variable_block_sizes = inputs + o, M = output + ctx.save_for_backward(q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) + + +def _backward_sm100a(ctx, grad_o, grad_M): + q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors + block = q.shape[2] // variable_block_sizes.numel() + if block != 64: + raise RuntimeError( + "block_sparse_attn_sm100a backward pairs the sm_100a forward with the " + f"Triton backward, which is hardcoded to 64-token blocks; got {block}. " + "Run 128-token-block metadata without grad, or use the Triton forward.") + dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx, + q2k_num, variable_block_sizes) + return dq, dk, dv, None, None, None + + +block_sparse_attn_sm100a_op.register_autograd(_backward_sm100a, + setup_context=_setup_context_sm100a) + # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- @@ -364,11 +464,24 @@ def block_sparse_attn_from_indices( # Backend resolution: # - FASTVIDEO_VSA_TRITON forces Triton everywhere. + # - FASTVIDEO_VSA_SM100A opts into the sm_100a forward (Triton backward). + # Unsupported 64-token metadata falls through; unsupported 128-token + # metadata raises because Triton cannot consume it. # - FASTVIDEO_VSA_TK requests sm_90 TK; honored only when it's actually # available (else falls through to the default below). # - Otherwise: TK on sm_90 if available, else Triton. if _force_triton(): use_sm90 = False + elif _force_sm100a(): + if _sm100a_is_supported(q, variable_block_sizes): + return block_sparse_attn_sm100a_op(q, k, v, q2k_idx, q2k_num, + variable_block_sizes) + if _infer_block_size(q, variable_block_sizes) == 128: + raise NotImplementedError( + "128-token block-sparse attention requires the sm_100a forward; " + "the Triton fallback only supports 64-token blocks, and the " + "sm_100a route is unavailable for this input.") + use_sm90 = sm90_available elif _force_tk(): use_sm90 = sm90_available else: diff --git a/tests/test_block_sparse_sm100a_dispatch.py b/tests/test_block_sparse_sm100a_dispatch.py new file mode 100644 index 000000000..75b7723ee --- /dev/null +++ b/tests/test_block_sparse_sm100a_dispatch.py @@ -0,0 +1,135 @@ +# 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()