nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[refactor] VSA: rename the sm100a block-sparse kernel to ptx_blackwell (#1956)

Browse Source 
Co-authored-by: root <root@ptyche0124.ptyche.clusters.nvidia.com>
This commit is contained in:
lpc0220
2026-10-17 21:10:19 -07:00
committed by GitHub
co-authored by root
parent 377c79fa7e
commit d4ed61cc7c
75 changed files with 613 additions and 586 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py
+48 -46
View File
@@ -49,18 +49,21 @@ def _force_tk() -> bool:
return os.environ.get("FASTVIDEO_VSA_TK", "0") == "1"
def _force_sm100a() -> bool:
def _force_ptx_blackwell() -> bool:
"""True iff the data-center Blackwell forward is explicitly opted into.
Opt-in only (same legacy-named env the H3 backend honors). Honored only when
``block_sparse_attn_sm100a.is_supported`` passes; the backward then runs the
sm_100a/sm_103a CUDA backward when ``block_sparse_attn_bwd_sm100a.is_supported``
Opt-in only (same env the H3 backend honors). Honored only when
``block_sparse_attn_ptx_blackwell.is_supported`` passes; the backward then runs the
ptx_blackwell backward when ``block_sparse_attn_bwd_ptx_blackwell.is_supported``
passes and Triton otherwise (the forward's lse is already in Triton's M
format). 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.
compatible fallback. ``FASTVIDEO_VSA_TRITON`` still wins. The deprecated name
``FASTVIDEO_VSA_SM100A`` is read when ``FASTVIDEO_VSA_PTX_BLACKWELL`` is unset.
"""
return _env_bool("FASTVIDEO_VSA_SM100A")
if "FASTVIDEO_VSA_PTX_BLACKWELL" not in os.environ and "FASTVIDEO_VSA_SM100A" in os.environ:
return _env_bool("FASTVIDEO_VSA_SM100A")
return _env_bool("FASTVIDEO_VSA_PTX_BLACKWELL")
def _env_bool(name: str) -> bool:
@@ -79,12 +82,12 @@ def _env_bool(name: str) -> bool:
"or an empty string")
def _sm100a_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
def _ptx_blackwell_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
try:
from fastvideo_kernel import block_sparse_attn_sm100a as vsa_sm100a
from fastvideo_kernel import block_sparse_attn_ptx_blackwell as vsa_ptx_blackwell
except Exception:
return False
return vsa_sm100a.is_supported(q, variable_block_sizes)
return vsa_ptx_blackwell.is_supported(q, variable_block_sizes)
def _infer_block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int:
@@ -386,10 +389,10 @@ def _backward_sm90(ctx, grad_o, grad_lse):
block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90)
# ---------------------------------------------------------------------------
# Data-center Blackwell backend custom op (index-native; legacy sm100a API name)
# Data-center Blackwell backend custom op (index-native; the ptx_blackwell 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- or 128-token
# Forward runs the ptx_blackwell extension. Backward runs the ptx_blackwell
# backward when block_sparse_attn_bwd_ptx_blackwell.is_supported passes (64- or 128-token
# blocks, data-center Blackwell device, extension built with that block's op) and
# the Triton kernels otherwise. The native forward emits lse in exactly Triton's M
# format (max*log2e + log2(l)), so either pairing needs no conversion. The Triton
@@ -399,11 +402,11 @@ block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_co
@torch.library.custom_op(
"fastvideo_kernel::block_sparse_attn_backward_sm100a",
"fastvideo_kernel::block_sparse_attn_backward_ptx_blackwell",
mutates_args=(),
device_types="cuda",
)
def block_sparse_attn_backward_sm100a(
def block_sparse_attn_backward_ptx_blackwell(
grad_o: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
@@ -414,19 +417,19 @@ def block_sparse_attn_backward_sm100a(
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, )
from fastvideo_kernel.block_sparse_attn_bwd_ptx_blackwell import (
block_sparse_attn_backward_ptx_blackwell_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(
dq, dk, dv = block_sparse_attn_backward_ptx_blackwell_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(
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_ptx_blackwell")
def _block_sparse_attn_backward_ptx_blackwell_fake(
grad_o: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
@@ -440,20 +443,20 @@ def _block_sparse_attn_backward_sm100a_fake(
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:
def _ptx_blackwell_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
from fastvideo_kernel import block_sparse_attn_bwd_ptx_blackwell as vsa_bwd_ptx_blackwell
except ImportError: # pragma: no cover - extension not built
return False
return vsa_bwd_sm100a.is_supported(q, variable_block_sizes)
return vsa_bwd_ptx_blackwell.is_supported(q, variable_block_sizes)
@torch.library.custom_op(
"fastvideo_kernel::block_sparse_attn_sm100a",
"fastvideo_kernel::block_sparse_attn_ptx_blackwell",
mutates_args=(),
device_types="cuda",
)
def block_sparse_attn_sm100a_op(
def block_sparse_attn_ptx_blackwell_op(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -461,15 +464,15 @@ def block_sparse_attn_sm100a_op(
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
from fastvideo_kernel.block_sparse_attn_ptx_blackwell import block_sparse_attn_ptx_blackwell
o, M = block_sparse_attn_sm100a(q, k, v, q2k_idx, q2k_num, variable_block_sizes,
need_lse=True)
o, M = block_sparse_attn_ptx_blackwell(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(
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_ptx_blackwell")
def _block_sparse_attn_ptx_blackwell_fake(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -486,32 +489,32 @@ def _block_sparse_attn_sm100a_fake(
return o, M
def _setup_context_sm100a(ctx, inputs, output):
def _setup_context_ptx_blackwell(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):
def _backward_ptx_blackwell(ctx, grad_o, grad_M):
q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors
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)
if _ptx_blackwell_backward_is_supported(q, variable_block_sizes):
dq, dk, dv = block_sparse_attn_backward_ptx_blackwell(grad_o, q, k, v, o, M, q2k_idx,
q2k_num, variable_block_sizes)
else:
block = q.shape[2] // variable_block_sizes.numel()
if block != 64:
raise RuntimeError(
"block_sparse_attn_sm100a backward: no sm_100a/sm_103a backward for "
"block_sparse_attn_ptx_blackwell backward: no ptx_blackwell backward for "
f"{block}-token blocks on this build/device, and the Triton backward is "
"hardcoded to 64-token blocks. Run this metadata without grad, or build "
"the extension with block_sparse_sm100a_blk128_bwd.")
"the extension with block_sparse_ptx_blackwell_blk128_bwd.")
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)
block_sparse_attn_ptx_blackwell_op.register_autograd(_backward_ptx_blackwell,
setup_context=_setup_context_ptx_blackwell)
# ---------------------------------------------------------------------------
# Public API
@@ -538,9 +541,8 @@ def block_sparse_attn_from_indices(
# Backend resolution:
# - FASTVIDEO_VSA_TRITON forces Triton everywhere.
# - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward (CUDA
# backward when supported, else Triton). The environment name is retained
# for compatibility.
# - FASTVIDEO_VSA_PTX_BLACKWELL opts into the data-center Blackwell forward (CUDA
# backward when supported, else Triton).
# 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
@@ -548,13 +550,13 @@ def block_sparse_attn_from_indices(
# - 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)
elif _force_ptx_blackwell():
if _ptx_blackwell_is_supported(q, variable_block_sizes):
return block_sparse_attn_ptx_blackwell_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/sm_103a forward; "
"128-token block-sparse attention requires the ptx_blackwell forward; "
"the Triton fallback only supports 64-token blocks, and the "
"native data-center Blackwell route is unavailable for this input.")
use_sm90 = sm90_available
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py → fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_ptx_blackwell.py
+20 -20
View File
@@ -1,12 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
"""sm_100a/sm_103a (data-center Blackwell) CUDA block-sparse VSA backward.
"""ptx_blackwell block-sparse VSA backward for data-center Blackwell (sm_100a / sm_103a).
Companion of ``block_sparse_attn_sm100a`` (the forward): consumes the forward's ``lse`` in the
Companion of ``block_sparse_attn_ptx_blackwell`` (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). Two kernels, picked by the
metadata's block size: 64-token blocks (``block_sparse_sm100a_bwd``) and 128-token blocks
(``block_sparse_sm100a_blk128_bwd``); any other configuration falls back to Triton via
metadata's block size: 64-token blocks (``block_sparse_ptx_blackwell_bwd``) and 128-token blocks
(``block_sparse_ptx_blackwell_blk128_bwd``); any other configuration falls back to Triton via
``is_supported``.
"""
@@ -19,14 +19,14 @@ try:
# (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_BY_BLOCK = {
64: getattr(_C, "block_sparse_sm100a_bwd", None),
128: getattr(_C, "block_sparse_sm100a_blk128_bwd", None),
64: getattr(_C, "block_sparse_ptx_blackwell_bwd", None),
128: getattr(_C, "block_sparse_ptx_blackwell_blk128_bwd", None),
}
_HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values())
_HAS_VSA_BWD_PTX_BLACKWELL = any(_BWD_BY_BLOCK.values())
except ImportError: # pragma: no cover - extension not built
_C = None
_BWD_BY_BLOCK = {}
_HAS_VSA_BWD_SM100A = False
_HAS_VSA_BWD_PTX_BLACKWELL = False
_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)}
HEAD_DIM = 128
@@ -36,20 +36,20 @@ BHSD = True
def set_extension(module) -> None:
"""Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd`` and / or
``block_sparse_sm100a_blk128_bwd``.
"""Use an already-loaded extension module exposing ``block_sparse_ptx_blackwell_bwd`` and / or
``block_sparse_ptx_blackwell_blk128_bwd``.
A standalone build of the binding .cu files (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_BY_BLOCK, _HAS_VSA_BWD_SM100A
global _C, _BWD_BY_BLOCK, _HAS_VSA_BWD_PTX_BLACKWELL
_C = module
_BWD_BY_BLOCK = {
64: getattr(module, "block_sparse_sm100a_bwd", None),
128: getattr(module, "block_sparse_sm100a_blk128_bwd", None),
64: getattr(module, "block_sparse_ptx_blackwell_bwd", None),
128: getattr(module, "block_sparse_ptx_blackwell_blk128_bwd", None),
}
_HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values())
_HAS_VSA_BWD_PTX_BLACKWELL = any(_BWD_BY_BLOCK.values())
def _seqlen(q: torch.Tensor) -> int:
@@ -77,7 +77,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
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:
if not _HAS_VSA_BWD_PTX_BLACKWELL or not q.is_cuda:
return False
if torch.cuda.get_device_capability(q.device) not in _SUPPORTED_COMPUTE_CAPABILITIES:
return False
@@ -99,7 +99,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
return True
def block_sparse_attn_backward_sm100a_from_k2q(
def block_sparse_attn_backward_ptx_blackwell_from_k2q(
grad_o: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
@@ -119,7 +119,7 @@ def block_sparse_attn_backward_sm100a_from_k2q(
block = _block_size(q, variable_block_sizes)
backward = _BWD_BY_BLOCK.get(block)
if backward is None:
raise RuntimeError(f"block_sparse_attn_backward_sm100a: no kernel for {block}-token "
raise RuntimeError(f"block_sparse_attn_backward_ptx_blackwell: no kernel for {block}-token "
f"blocks (built: {sorted(b for b, f in _BWD_BY_BLOCK.items() if f)})")
sm_scale = 1.0 / (q.shape[-1]**0.5)
idx = k2q_idx.to(torch.int32).contiguous()
@@ -130,7 +130,7 @@ def block_sparse_attn_backward_sm100a_from_k2q(
return res[0], res[1], res[2]
def block_sparse_attn_backward_sm100a(
def block_sparse_attn_backward_ptx_blackwell(
grad_o: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
@@ -158,5 +158,5 @@ def block_sparse_attn_backward_sm100a(
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)
return block_sparse_attn_backward_ptx_blackwell_from_k2q(grad_o, q, k, v, o, lse, k2q_idx, k2q_num,
variable_block_sizes)
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py → fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_ptx_blackwell.py
+25 -25
View File
@@ -1,14 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
"""Data-center Blackwell CUDA block-sparse VSA forward.
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.
``ptx_blackwell`` is the kernel's name, not an architecture: 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. 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- and 128-token blocks, sm_100a / sm_103a) run against it
fp32 -- so both ``block_sparse_attn_backward_triton`` and the ptx_blackwell backward
(``block_sparse_attn_bwd_ptx_blackwell``, 64- and 128-token blocks, sm_100a / sm_103a) run against it
unchanged.
The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks
@@ -27,14 +27,14 @@ try:
# kernel built and present.
from fastvideo_kernel._C import fastvideo_kernel_ops as _C
_FWD_BY_BLOCK = {
64: getattr(_C, "block_sparse_sm100a_fwd", None),
128: getattr(_C, "block_sparse_sm100a_blk128_fwd", None),
64: getattr(_C, "block_sparse_ptx_blackwell_fwd", None),
128: getattr(_C, "block_sparse_ptx_blackwell_blk128_fwd", None),
}
_HAS_VSA_SM100A = any(_FWD_BY_BLOCK.values())
_HAS_VSA_PTX_BLACKWELL = any(_FWD_BY_BLOCK.values())
except ImportError: # pragma: no cover - extension not built
_C = None
_FWD_BY_BLOCK = {}
_HAS_VSA_SM100A = False
_HAS_VSA_PTX_BLACKWELL = False
_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)}
HEAD_DIM = 128
@@ -69,9 +69,9 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
(= q2k_idx.shape[-1]) must be >= 1, which the host launcher re-checks.
* variable_block_sizes: per-KV-block valid-token counts in [0, block]; keys at or past
a block's count are masked. Integer metadata is converted to int32/contiguous by
``block_sparse_attn_sm100a`` itself, so int64 inputs merely cost a cast.
``block_sparse_attn_ptx_blackwell`` itself, so int64 inputs merely cost a cast.
"""
if not _HAS_VSA_SM100A or not q.is_cuda:
if not _HAS_VSA_PTX_BLACKWELL or not q.is_cuda:
return False
if torch.cuda.get_device_capability(q.device) not in _SUPPORTED_COMPUTE_CAPABILITIES:
return False
@@ -91,11 +91,11 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
@torch.library.custom_op(
"fastvideo_kernel::block_sparse_attn_sm100a_inference",
"fastvideo_kernel::block_sparse_attn_ptx_blackwell_inference",
mutates_args=(),
device_types="cuda",
)
def _block_sparse_attn_sm100a_inference(
def _block_sparse_attn_ptx_blackwell_inference(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -103,7 +103,7 @@ def _block_sparse_attn_sm100a_inference(
q2k_num: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> torch.Tensor:
"""Opaque no-LSE launch used by the inference-only sm_100a/sm_103a route.
"""Opaque no-LSE launch used by the inference-only ptx_blackwell route.
The extension is exposed as a raw pybind function rather than a dispatcher
op. Calling it directly makes Dynamo descend through a Python/C++ boundary
@@ -117,8 +117,8 @@ def _block_sparse_attn_sm100a_inference(
return res[0]
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_inference")
def _block_sparse_attn_sm100a_inference_fake(
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_ptx_blackwell_inference")
def _block_sparse_attn_ptx_blackwell_inference_fake(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -131,18 +131,18 @@ def _block_sparse_attn_sm100a_inference_fake(
@torch.library.custom_op(
"fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference",
"fastvideo_kernel::block_sparse_attn_ptx_blackwell_from_mask_inference",
mutates_args=(),
device_types="cuda",
)
def _block_sparse_attn_sm100a_from_mask_inference(
def _block_sparse_attn_ptx_blackwell_from_mask_inference(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> torch.Tensor:
"""Opaque mask compaction plus no-LSE sm_100a/sm_103a launch for inference.
"""Opaque mask compaction plus no-LSE ptx_blackwell launch for inference.
H3 naturally produces a bool block map. Its Triton ``map_to_index`` call
must live behind the same opaque boundary as the raw pybind launch;
@@ -152,7 +152,7 @@ def _block_sparse_attn_sm100a_from_mask_inference(
from fastvideo_kernel.triton_kernels.index import map_to_index
q2k_idx, q2k_num = map_to_index(block_map)
return _block_sparse_attn_sm100a_inference(
return _block_sparse_attn_ptx_blackwell_inference(
q,
k,
v,
@@ -162,8 +162,8 @@ def _block_sparse_attn_sm100a_from_mask_inference(
)
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference")
def _block_sparse_attn_sm100a_from_mask_inference_fake(
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_ptx_blackwell_from_mask_inference")
def _block_sparse_attn_ptx_blackwell_from_mask_inference_fake(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -173,7 +173,7 @@ def _block_sparse_attn_sm100a_from_mask_inference_fake(
return torch.empty_like(q)
def block_sparse_attn_sm100a_from_mask(
def block_sparse_attn_ptx_blackwell_from_mask(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -181,7 +181,7 @@ def block_sparse_attn_sm100a_from_mask(
variable_block_sizes: torch.Tensor,
) -> Tuple[torch.Tensor, None]:
"""Inference forward from a bool block map, with compaction kept opaque."""
out = _block_sparse_attn_sm100a_from_mask_inference(
out = _block_sparse_attn_ptx_blackwell_from_mask_inference(
q.contiguous(),
k.contiguous(),
v.contiguous(),
@@ -191,7 +191,7 @@ def block_sparse_attn_sm100a_from_mask(
return out, None
def block_sparse_attn_sm100a(
def block_sparse_attn_ptx_blackwell(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -212,7 +212,7 @@ def block_sparse_attn_sm100a(
# This is the production inference path. The custom op keeps the raw
# pybind launch opaque to Dynamo while its fake kernel carries output
# metadata through full-graph capture.
return _block_sparse_attn_sm100a_inference(q, k, v, idx, num, vbs), None
return _block_sparse_attn_ptx_blackwell_inference(q, k, v, idx, num, vbs), None
# Preserve the established LSE-producing path for correctness tests and
# any future forward/backward pairing; only inference needs the opaque op.
fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py
+2 -2
View File
@@ -14,8 +14,8 @@ import math
import torch
VSA_TILE_SIZE = (4, 4, 4)
# 128 is served by the sm_100a/sm_103a CUDA backend (legacy API name
# block_sparse_attn_sm100a); 64 and 256 by
# 128 is served by the ptx_blackwell backend
# (block_sparse_attn_ptx_blackwell); 64 and 256 by
# Triton and the CuTe-DSL path. A volume here only needs a backend that accepts it.
_SUPPORTED_VSA_BLOCK_VOLUMES = (64, 128, 256)