[kernel] sm_100a CUDA backward for VSA block-sparse attention (blk64) (#1819)
This commit is contained in:
lpc0220
@@ -372,13 +372,65 @@ block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_co
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data-center Blackwell backend custom op (index-native; legacy sm100a API name)
|
||||
#
|
||||
# Forward runs the sm_100a/sm_103a CUDA extension; backward reuses the Triton kernels.
|
||||
# Forward runs the sm_100a/sm_103a CUDA extension. Backward runs the sm_100a CUDA
|
||||
# backward when block_sparse_attn_bwd_sm100a.is_supported passes (64-token blocks,
|
||||
# sm_100a device, extension built with the op) and the Triton kernels otherwise.
|
||||
# The native forward emits lse in exactly Triton's M format (max*log2e +
|
||||
# log2(l)), so the pairing needs no conversion. The Triton backward is
|
||||
# log2(l)), so either pairing needs no conversion. Both backwards are
|
||||
# hardcoded to 64-token blocks, hence the block-size assert below.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo_kernel::block_sparse_attn_backward_sm100a",
|
||||
mutates_args=(),
|
||||
device_types="cuda",
|
||||
)
|
||||
def block_sparse_attn_backward_sm100a(
|
||||
grad_o: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
o: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
from fastvideo_kernel.block_sparse_attn_bwd_sm100a import (
|
||||
block_sparse_attn_backward_sm100a_from_k2q, )
|
||||
|
||||
num_kv_blocks = variable_block_sizes.numel()
|
||||
k2q_idx, k2q_num = _invert_indices_for_backward(q2k_idx, q2k_num, num_kv_blocks)
|
||||
dq, dk, dv = block_sparse_attn_backward_sm100a_from_k2q(
|
||||
grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(), o.contiguous(),
|
||||
lse.contiguous(), k2q_idx, k2q_num, variable_block_sizes)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_sm100a")
|
||||
def _block_sparse_attn_backward_sm100a_fake(
|
||||
grad_o: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
o: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
|
||||
|
||||
|
||||
def _sm100a_backward_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
|
||||
try:
|
||||
from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd_sm100a
|
||||
except ImportError: # pragma: no cover - extension not built
|
||||
return False
|
||||
return vsa_bwd_sm100a.is_supported(q, variable_block_sizes)
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"fastvideo_kernel::block_sparse_attn_sm100a",
|
||||
mutates_args=(),
|
||||
@@ -428,11 +480,15 @@ def _backward_sm100a(ctx, grad_o, grad_M):
|
||||
block = q.shape[2] // variable_block_sizes.numel()
|
||||
if block != 64:
|
||||
raise RuntimeError(
|
||||
"block_sparse_attn_sm100a backward pairs the sm_100a/sm_103a forward with the "
|
||||
f"Triton backward, which is hardcoded to 64-token blocks; got {block}. "
|
||||
"block_sparse_attn_sm100a backward pairs the sm_100a/sm_103a forward with a "
|
||||
f"backward that 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)
|
||||
if _sm100a_backward_is_supported(q, variable_block_sizes):
|
||||
dq, dk, dv = block_sparse_attn_backward_sm100a(grad_o, q, k, v, o, M, q2k_idx,
|
||||
q2k_num, variable_block_sizes)
|
||||
else:
|
||||
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
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""sm_100a (Blackwell) CUDA block-sparse VSA backward.
|
||||
|
||||
Companion of ``block_sparse_attn_sm100a`` (the forward): consumes the forward's ``lse`` in the
|
||||
Triton "M format" (``max(qk * sm_scale * log2e) + log2(l)``, ``[B, H, S]`` fp32) unchanged and
|
||||
FastVideo's k2q index metadata, returns ``(dq, dk, dv)`` in bf16 with the inputs' layout and the
|
||||
Triton backward's scaling (dq and dk carry sm_scale, dv does not). 64-token blocks only; any
|
||||
other configuration falls back to Triton via ``is_supported``.
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
# The pybind symbols live on fastvideo_kernel_ops, NOT on the _C package that contains it
|
||||
# (its __init__ is empty, so hasattr on the package fails with the kernel built and present).
|
||||
from fastvideo_kernel._C import fastvideo_kernel_ops as _C
|
||||
_BWD = getattr(_C, "block_sparse_sm100a_bwd", None)
|
||||
_HAS_VSA_BWD_SM100A = _BWD is not None
|
||||
except ImportError: # pragma: no cover - extension not built
|
||||
_C = None
|
||||
_BWD = None
|
||||
_HAS_VSA_BWD_SM100A = False
|
||||
|
||||
_SM100 = (10, 0)
|
||||
HEAD_DIM = 128
|
||||
BLOCK = 64
|
||||
# Must match the -DVSA_BHSD the extension was compiled with (FastVideo builds with true).
|
||||
BHSD = True
|
||||
|
||||
|
||||
def set_extension(module) -> None:
|
||||
"""Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd``.
|
||||
|
||||
A standalone build of ``block_sparse_bwd_sm100a.cu`` (for example through
|
||||
``torch.utils.cpp_extension.load`` with a ten-line pybind wrapper) can be injected here, so
|
||||
the backend can be exercised without rebuilding the fastvideo_kernel wheel.
|
||||
"""
|
||||
global _C, _BWD, _HAS_VSA_BWD_SM100A
|
||||
_C = module
|
||||
_BWD = getattr(module, "block_sparse_sm100a_bwd", None)
|
||||
_HAS_VSA_BWD_SM100A = _BWD is not None
|
||||
|
||||
|
||||
def _seqlen(q: torch.Tensor) -> int:
|
||||
return q.shape[2] if BHSD else q.shape[1]
|
||||
|
||||
|
||||
def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
|
||||
"""True iff this build can run these tensors; otherwise the caller uses Triton.
|
||||
|
||||
Static facts only (shapes, dtypes, arch, layout), never tensor contents, so it is cheap
|
||||
enough for a per-layer dispatch path. The kernel is fixed at 64-token blocks with
|
||||
head_dim 128 and needs seqlen == 64 * num_blocks with an even num_blocks (its preprocess
|
||||
works in 128-token blocks). Per-row k2q counts may be anything in [0, num_q_blocks],
|
||||
including 0: unselected kv blocks get exactly-zero dk/dv rows.
|
||||
"""
|
||||
if not _HAS_VSA_BWD_SM100A or not q.is_cuda:
|
||||
return False
|
||||
if torch.cuda.get_device_capability(q.device) != _SM100:
|
||||
return False
|
||||
if q.dtype != torch.bfloat16 or q.dim() != 4 or q.shape[-1] != HEAD_DIM:
|
||||
return False
|
||||
if not q.is_contiguous():
|
||||
return False
|
||||
# Metadata must be integer-typed so the wrapper's int32 conversion is value-preserving.
|
||||
if not variable_block_sizes.is_cuda or variable_block_sizes.dtype not in (torch.int32,
|
||||
torch.int64):
|
||||
return False
|
||||
num_blocks = variable_block_sizes.numel()
|
||||
if num_blocks == 0 or num_blocks % 2 != 0:
|
||||
return False
|
||||
if _seqlen(q) != BLOCK * num_blocks:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def block_sparse_attn_backward_sm100a_from_k2q(
|
||||
grad_o: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
o: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
k2q_idx: torch.Tensor,
|
||||
k2q_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Backward from k2q metadata already in hand (``invert_indices`` layout).
|
||||
|
||||
``k2q_idx`` is ``[B, H, num_kv_blocks, max_q_blocks]`` (or the flat 2-D view) of LOCAL q64
|
||||
block ids, ``k2q_num`` ``[B, H, num_kv_blocks]``; entries past a row's count are never read.
|
||||
"""
|
||||
sm_scale = 1.0 / (q.shape[-1]**0.5)
|
||||
idx = k2q_idx.to(torch.int32).contiguous()
|
||||
num = k2q_num.to(torch.int32).contiguous()
|
||||
vbs = variable_block_sizes.to(torch.int32).contiguous()
|
||||
res = _BWD(grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(),
|
||||
o.contiguous(), lse.contiguous(), idx, num, vbs, sm_scale)
|
||||
return res[0], res[1], res[2]
|
||||
|
||||
|
||||
def block_sparse_attn_backward_sm100a(
|
||||
grad_o: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
o: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
q2k_idx: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Backward pass from the forward's q2k metadata. Returns ``(dq, dk, dv)``.
|
||||
|
||||
Mirrors ``block_sparse_attn_backward_triton``: the k2q inversion is recomputed here with
|
||||
FastVideo's Triton ``invert_indices`` rather than saved by the forward.
|
||||
"""
|
||||
from fastvideo_kernel.triton_kernels.index import invert_indices
|
||||
|
||||
num_kv_blocks = variable_block_sizes.numel()
|
||||
batch = q.shape[0]
|
||||
heads = q.shape[1] if BHSD else q.shape[2]
|
||||
idx = q2k_idx.to(torch.int32).contiguous()
|
||||
num = q2k_num.to(torch.int32).contiguous()
|
||||
if idx.dim() != 4:
|
||||
idx = idx.view(batch, heads, -1, idx.shape[-1])
|
||||
if num.dim() != 3:
|
||||
num = num.view(batch, heads, -1)
|
||||
k2q_idx, k2q_num = invert_indices(idx, num, num_kv_blocks)
|
||||
return block_sparse_attn_backward_sm100a_from_k2q(grad_o, q, k, v, o, lse, k2q_idx, k2q_num,
|
||||
variable_block_sizes)
|
||||
@@ -4,10 +4,12 @@
|
||||
The historical ``sm100a`` module and symbol names are retained for compatibility, but the
|
||||
extension carries native sm_100a and sm_103a images and supports both device generations.
|
||||
|
||||
A third backend behind the same VSA op as the Triton and CuTe-DSL paths. Forward only: it
|
||||
returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward``
|
||||
writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]`` fp32 -- so
|
||||
``block_sparse_attn_backward_triton`` runs against it unchanged.
|
||||
A third backend behind the same VSA op as the Triton and CuTe-DSL paths. This module is the
|
||||
forward: it returns ``(out, lse)`` with ``lse`` in exactly the form
|
||||
``triton_block_sparse_attn_forward`` writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]``
|
||||
fp32 -- so both ``block_sparse_attn_backward_triton`` and the sm_100a CUDA backward
|
||||
(``block_sparse_attn_bwd_sm100a``, 64-token blocks, sm_100a devices only) run against it
|
||||
unchanged.
|
||||
|
||||
The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks
|
||||
(tile volumes 64 and 128 in ``build_vsa_metadata``); the block size is inferred from the
|
||||
|
||||
Reference in New Issue
Block a user