Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/publish-kernel.yml
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ jobs:
torch-cuda-short: 'cu130'
platform:
# x86_64 builds the full cu126 + cu130 set. cu130 ships the
# data-center Blackwell sm_100a VSA and consumer sm_120a FP4
# data-center Blackwell sm_100a/sm_103a VSA and consumer sm_120a FP4
# kernels.
- os: ubuntu-22.04
arch: x86_64
Expand Down Expand Up @@ -169,7 +169,7 @@ jobs:
# covers sm_120a; turbodiffusion covers sm_100a+sm_120a. The sm_100 FP4
# forward is the FA4 CuTe DSL path in the fastvideo package (PR #1221),
# JIT-compiled at runtime — not built into this wheel.
# * x86_64 cu130 = Hopper TK + data-center Blackwell sm_100a VSA
# * x86_64 cu130 = Hopper TK + data-center Blackwell sm_100a/sm_103a VSA
# + consumer Blackwell sm_120a FP4.
# * x86_64 cu126 = Hopper TK only (older drivers; CUDA < 12.8 has no FP4).
# The per-arch split in CMakeLists pins the FP4 targets to sm_120a and builds
Expand Down
24 changes: 12 additions & 12 deletions fastvideo-kernel/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ if(NOT GPU_BACKEND STREQUAL "ROCM")
else()
# Best-effort probe of the visible GPU (mirrors build.sh detect_with_torch).
execute_process(
COMMAND "${Python_EXECUTABLE}" -c "import torch; assert torch.cuda.is_available(); mj, mn = torch.cuda.get_device_capability(0); print(f'{mj}.{mn}a' if (mj, mn) in ((9, 0), (12, 0)) else f'{mj}.{mn}')"
COMMAND "${Python_EXECUTABLE}" -c "import torch; assert torch.cuda.is_available(); mj, mn = torch.cuda.get_device_capability(0); print(f'{mj}.{mn}a' if (mj, mn) in ((9, 0), (10, 0), (10, 3), (12, 0)) else f'{mj}.{mn}')"
OUTPUT_VARIABLE _FV_ARCH_LIST
OUTPUT_STRIP_TRAILING_WHITESPACE
RESULT_VARIABLE _fv_detect_rc
Expand Down Expand Up @@ -382,11 +382,11 @@ if(BUILD_CXX_KERNELS)
# Combined FastVideo Extension
# Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace
# ---------------------------------------------------------------------------
# VSA block-sparse attention forward, Blackwell (sm_100a) only.
# VSA block-sparse attention forward, data-center Blackwell (sm_100a/sm_103a).
#
# NOTE the "a" suffix: -arch=sm_100a is NOT enough -- it emits a plain sm_100 target and
# ptxas rejects every tcgen05 / setmaxnreg instruction. The explicit gencode spelling
# below is required, and matches the 10.0a entry in TORCH_CUDA_ARCH_LIST.
# NOTE the "a" suffix: these kernels use architecture-conditional tcgen05 / setmaxnreg
# instructions. Explicit gencode entries are required for both the GB200 (sm_100a) and
# B300/GB300 (sm_103a) images. The source body is shared; ptxas specializes each image.
#
# Built for 64-token sparse blocks -- FastVideo's default (4,4,4) tiling, so no
# tile-size change and no top-k granularity change is needed. VSA_BHSD
Expand All @@ -396,21 +396,21 @@ if(BUILD_CXX_KERNELS)
# exported (build.sh, and `pip install` with it set) the branch above only prints it --
# the cmake variable stays empty, so testing that alone silently skips the kernel and
# leaves a build that succeeds with the op missing.
set(ENABLE_VSA_SM100A OFF)
set(ENABLE_VSA_SM100_FAMILY OFF)
set(_VSA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}")
if(NOT _VSA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST})
set(_VSA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}")
endif()
if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)")
set(ENABLE_VSA_SM100A ON)
if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.(0|3)a|10(0|3)a|sm_10(0|3)a)([; ,]|$)")
set(ENABLE_VSA_SM100_FAMILY ON)
endif()
if(ENABLE_VSA_SM100A)
message(STATUS "fastvideo-kernel: building block_sparse_sm100a (Blackwell, 64- and 128-token blocks)")
if(ENABLE_VSA_SM100_FAMILY)
message(STATUS "fastvideo-kernel: building block_sparse_sm100a (sm_100a/sm_103a, 64- and 128-token blocks)")
list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_sm100a.cu
csrc/attention/block_sparse_blk128_sm100a.cu)
set_source_files_properties(csrc/attention/block_sparse_sm100a.cu
csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES
COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-DVSA_BHSD=true")
COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-gencode;arch=compute_103a,code=sm_103a;-DVSA_BHSD=true")
endif()

Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI
Expand All @@ -428,7 +428,7 @@ if(BUILD_CXX_KERNELS)

# Build compile definitions list
set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops)
if(ENABLE_VSA_SM100A)
if(ENABLE_VSA_SM100_FAMILY)
list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A)
endif()
if(ENABLE_TK_KERNELS)
Expand Down
4 changes: 2 additions & 2 deletions fastvideo-kernel/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ Compiled CUDA extensions (CMake, see the build summary printed at the end of eve
|---|---|---|---|---|
| `fastvideo_kernel._C.fastvideo_kernel_ops` | TurboDiffusion INT8 GEMM, quant, RMSNorm, LayerNorm | `csrc/turbodiffusion/` | every arch in `TORCH_CUDA_ARCH_LIST` | always built |
| same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO = ON iff `9.0a` is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there) |
| same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*_sm100a.cu` | Blackwell `sm_100a` only | ON iff `10.0a` is in `TORCH_CUDA_ARCH_LIST` |
| same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*_sm100a.cu` | Data-center Blackwell `sm_100a`/`sm_103a` | ON iff `10.0a` or `10.3a` is in `TORCH_CUDA_ARCH_LIST` |
| same extension, optional part | fused NVLink Ulysses all-to-all | `csrc/comm/ulysses_all_to_all.cu` | CUDA | `FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A` (AUTO = ON with NCCL 2.29+ device headers and library; always OFF on ROCm) |
| `fp4attn_cuda`, `fp4quant_cuda` | FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) | `attn_qat_infer/` | consumer Blackwell `sm_120a` only, CUDA ≥ 12.8 | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` is in the arch list) |

Expand All @@ -37,7 +37,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
Notes:

- No Docker image ships the FP4 kernels; only the x86_64/aarch64 cu130 wheels do.
- Both cu130 PyPI wheels ship the MiniMax-H3 sm_100a VSA forward.
- Both cu130 PyPI wheels ship native sm_100a and sm_103a images for the MiniMax-H3 VSA forward.
- On arm64 images (GH200 included) STA/VSA run on the Triton fallbacks, since TK never builds on aarch64.
- Ulysses AUTO builds only when CMake finds a NCCL library and device-API headers with the 2.29 initializers. Use
`-DFASTVIDEO_KERNEL_BUILD_ULYSSES_A2A=ON` to require it or `OFF` to test the portable build.
Expand Down
4 changes: 4 additions & 0 deletions fastvideo-kernel/build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,10 @@ if [ "${GPU_BACKEND}" = "CUDA" ]; then
if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then
if [ "${cc_major}" = "9" ] && [ "${cc_minor}" = "0" ]; then
export TORCH_CUDA_ARCH_LIST="9.0a"
elif [ "${cc_major}" = "10" ] && { [ "${cc_minor}" = "0" ] || [ "${cc_minor}" = "3" ]; }; then
# Data-center Blackwell VSA uses architecture-conditional tcgen05
# instructions and therefore requires the 'a' target.
export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}a"
elif [ "${cc_major}" = "12" ] && [ "${cc_minor}" = "0" ]; then
# Blackwell sm_120 needs the arch-conditional 'a' suffix so CMake's
# AUTO gate (matches 12.0a/120a/sm_120a) builds the attn_qat_infer
Expand Down
22 changes: 11 additions & 11 deletions fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
// block_sparse_kernel_sm100a.cuh -- VSA block-sparse FMHA forward (per-q-block top-k), sm_100a.
// block_sparse_kernel_sm100a.cuh -- VSA block-sparse FMHA forward (per-q-block top-k),
// data-center Blackwell sm_100a/sm_103a. The filename is retained for API compatibility.
// Warp-specialized: load / MMA (tcgen05) / softmax / correction / epilogue / scheduler.
// Writes O and, when asked, the log-sum-exp the backward consumes.
//
Expand Down Expand Up @@ -202,15 +203,14 @@ fmha_context_bf16_gen_kernel(const __grid_constant__ CUtensorMap tmap_q,
const int* __restrict__ variable_block_sizes,
float* __restrict__ lse_out) {
// Multi-arch builds: torch's cmake appends -gencode for EVERY entry of
// TORCH_CUDA_ARCH_LIST to this TU on top of the pinned compute_100a pass, and
// tcgen05/setmaxnreg do not exist outside sm_100a -- ptxas rejects the sm_120a
// (or plain sm_100) pass outright. Keep the body only where it can compile:
// the host pass (no __CUDA_ARCH__, needed for launch plumbing) and the
// sm_100a device pass (arch 1000 WITH the family-specific feature set that the
// "a" suffix defines). Every other device pass gets an empty stub; the Python
// is_supported() / host launcher never dispatch here off sm_100, so the stub
// is unreachable at runtime.
#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL))
// TORCH_CUDA_ARCH_LIST to this TU on top of the pinned data-center Blackwell
// passes. tcgen05/setmaxnreg require an architecture-conditional target, so
// plain sm_100/sm_103 and consumer sm_120 passes cannot compile this body.
// Keep it for the host pass (needed for launch plumbing) plus sm_100a/sm_103a;
// every other device pass gets an unreachable empty stub.
#if !defined(__CUDA_ARCH__) || \
((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \
(__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL)))

const int total_workitems = num_samples * num_heads * packed_mtiles_per_seq;

Expand Down Expand Up @@ -1070,7 +1070,7 @@ fmha_context_bf16_gen_kernel(const __grid_constant__ CUtensorMap tmap_q,
}
__syncthreads();
if (warp_id == 0) tcgen05_dealloc<1>(tmem_base, TMEM_TOTAL);
#endif // host pass or sm_100a device pass (multi-arch guard; see note at the top of the body)
#endif // host pass or sm_100a/sm_103a device pass (see multi-arch note above)
}

} // namespace VSA_NAMESPACE
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
#ifndef BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH
#define BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH

// Launch surface for the sm_100a VSA block-sparse FMHA forward.
// Launch surface for the sm_100a/sm_103a VSA block-sparse FMHA forward.
//
// Everything a caller needs: a POD argument struct, a predicate saying whether this build can
// run those arguments, and one launch entry point. The benchmark in
Expand Down
2 changes: 1 addition & 1 deletion fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// block_sparse_sm100a.cu -- torch binding for the sm_100a VSA block-sparse FMHA forward.
// block_sparse_sm100a.cu -- torch binding for the sm_100a/sm_103a VSA block-sparse FMHA forward.
//
// Forward only: returns (out, lse) so FastVideo's existing Triton backward keeps working
// unchanged. lse is exactly the M tensor triton_block_sparse_attn_forward writes --
Expand Down
3 changes: 1 addition & 2 deletions fastvideo-kernel/csrc/attention/primitives.cuh
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// primitives.cuh -- device primitives for the sm_100a VSA block-sparse attention
// primitives.cuh -- device primitives for the sm_100a/sm_103a VSA block-sparse attention
// forward: tcgen05 (alloc / mma / ld / st / commit / wait / fence), TMA load / store /
// tensormap, mbarrier, cluster launch control, setmaxnreg, fast math, and the FMHA helpers.
//
Expand Down Expand Up @@ -874,4 +874,3 @@ __device__ __forceinline__
void sts_f32(uint32_t smem_addr, float val) {
asm volatile("st.shared.f32 [%0], %1;" :: "r"(smem_addr), "f"(val) : "memory");
}

4 changes: 2 additions & 2 deletions fastvideo-kernel/csrc/common_extension.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A
m.def("block_sparse_sm100a_fwd",
torch::wrap_pybind_function(block_sparse_sm100a_fwd),
"VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a)");
"VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a/sm103a)");
m.def("block_sparse_sm100a_blk128_fwd",
torch::wrap_pybind_function(block_sparse_sm100a_blk128_fwd),
"VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a)");
"VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a/sm103a)");
#endif

#ifdef TK_COMPILE_ST_ATTN
Expand Down
19 changes: 10 additions & 9 deletions fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,9 +50,9 @@ def _force_tk() -> bool:


def _force_sm100a() -> bool:
"""True iff the sm_100a (Blackwell) forward is explicitly opted into.
"""True iff the data-center Blackwell forward is explicitly opted into.

Opt-in only (same env the H3 backend honors): the sm_100a extension is
Opt-in only (same legacy-named env the H3 backend honors): the 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
Expand Down Expand Up @@ -370,10 +370,10 @@ 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)
# Data-center Blackwell backend custom op (index-native; legacy sm100a API name)
#
# 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 +
# Forward runs the sm_100a/sm_103a CUDA extension; backward reuses the Triton kernels.
# 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
# hardcoded to 64-token blocks, hence the block-size assert below.
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -428,7 +428,7 @@ 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 forward with the "
"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}. "
"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,
Expand Down Expand Up @@ -464,7 +464,8 @@ 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).
# - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward
# (Triton backward). The environment name is retained for compatibility.
# 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
Expand All @@ -478,9 +479,9 @@ def block_sparse_attn_from_indices(
variable_block_sizes)
if _infer_block_size(q, variable_block_sizes) == 128:
raise NotImplementedError(
"128-token block-sparse attention requires the sm_100a forward; "
"128-token block-sparse attention requires the sm_100a/sm_103a forward; "
"the Triton fallback only supports 64-token blocks, and the "
"sm_100a route is unavailable for this input.")
"native data-center Blackwell route is unavailable for this input.")
use_sm90 = sm90_available
elif _force_tk():
use_sm90 = sm90_available
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
"""sm_100a (Blackwell) CUDA block-sparse VSA forward.
"""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.

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``
Expand Down Expand Up @@ -31,7 +34,7 @@
_FWD_BY_BLOCK = {}
_HAS_VSA_SM100A = False

_SM100 = (10, 0)
_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)}
HEAD_DIM = 128
# Must match the -DVSA_BHSD the extension was compiled with (see CMakeLists).
BHSD = True
Expand All @@ -52,7 +55,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
enough to sit on a per-layer dispatch path.

What the kernel accepts (and is tested to handle):
* q/k/v: contiguous 4-D bf16 CUDA tensors on an sm_100 device, head_dim 128, laid out
* q/k/v: contiguous 4-D bf16 CUDA tensors on an sm_100/sm_103 device, head_dim 128, laid out
as compiled (BHSD here); seqlen == num_blocks * block with an EVEN num_blocks (a CTA
owns an adjacent pair of query blocks) and a 64- or 128-token build present.
* q2k_num: any per-row counts in [0, max_kv], NON-uniform across rows included. Rows
Expand All @@ -67,7 +70,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
"""
if not _HAS_VSA_SM100A or not q.is_cuda:
return False
if torch.cuda.get_device_capability(q.device) != _SM100:
if torch.cuda.get_device_capability(q.device) not in _SUPPORTED_COMPUTE_CAPABILITIES:
return False
if q.dtype != torch.bfloat16 or q.dim() != 4 or q.shape[-1] != HEAD_DIM:
return False
Expand Down Expand Up @@ -97,7 +100,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 route.
"""Opaque no-LSE launch used by the inference-only sm_100a/sm_103a 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
Expand Down Expand Up @@ -136,7 +139,7 @@ def _block_sparse_attn_sm100a_from_mask_inference(
block_map: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> torch.Tensor:
"""Opaque mask compaction plus no-LSE sm_100a launch for inference.
"""Opaque mask compaction plus no-LSE sm_100a/sm_103a 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;
Expand Down
3 changes: 2 additions & 1 deletion fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@
import torch

VSA_TILE_SIZE = (4, 4, 4)
# 128 is served by the sm_100a CUDA backend (block_sparse_attn_sm100a); 64 and 256 by
# 128 is served by the sm_100a/sm_103a CUDA backend (legacy API name
# block_sparse_attn_sm100a); 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)

Expand Down
Loading
Loading