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

[kernel] VSA kernel: one sm_100a / sm_103a image per listed arch; un-gate backward on sm_103a (#1833)

Browse Source 
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
lpc0220
2026-09-12 16:49:36 -07:00
committed by GitHub
co-authored by Claude Fable 5.1
parent a943220c11
commit aef4e9b3b1
7 changed files with 48 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/publish-kernel.yml
+2 -2
View File
@@ -176,11 +176,11 @@ jobs:
# the main extension for the full arch list. CMAKE_BUILD_PARALLEL_LEVEL caps
# Ninja so heavy CUTLASS/TK template TUs don't OOM the 16 GB runner (exit 143).
if [ "${{ matrix.platform.arch }}" = "aarch64" ]; then
export TORCH_CUDA_ARCH_LIST="10.0a;12.0a"
export TORCH_CUDA_ARCH_LIST="10.0a;10.3a;12.0a"
export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON"
export CMAKE_BUILD_PARALLEL_LEVEL=1
elif [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then
export TORCH_CUDA_ARCH_LIST="9.0a;10.0a;12.0a"
export TORCH_CUDA_ARCH_LIST="9.0a;10.0a;10.3a;12.0a"
export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
# A single FP4 TU (attn_qat_infer) can use ~8-12 GB on its own, so serialize.
export CMAKE_BUILD_PARALLEL_LEVEL=1
fastvideo-kernel/CMakeLists.txt
+25 -11
View File
@@ -396,27 +396,41 @@ 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_SM100_FAMILY OFF)
set(ENABLE_VSA_SM100A OFF)
set(ENABLE_VSA_SM103A 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\\.(0|3)a|10(0|3)a|sm_10(0|3)a)([; ,]|$)")
set(ENABLE_VSA_SM100_FAMILY ON)
if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)")
set(ENABLE_VSA_SM100A ON)
endif()
if(ENABLE_VSA_SM100_FAMILY)
message(STATUS "fastvideo-kernel: building block_sparse_sm100a (sm_100a/sm_103a, 64- and 128-token blocks)")
if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.3a|103a|sm_103a)([; ,]|$)")
set(ENABLE_VSA_SM103A ON)
endif()
set(_VSA_FWD_GENCODE "")
set(_VSA_FWD_ARCHS "")
if(ENABLE_VSA_SM100A)
list(APPEND _VSA_FWD_GENCODE "-gencode;arch=compute_100a,code=sm_100a")
list(APPEND _VSA_FWD_ARCHS "sm_100a")
endif()
if(ENABLE_VSA_SM103A)
list(APPEND _VSA_FWD_GENCODE "-gencode;arch=compute_103a,code=sm_103a")
list(APPEND _VSA_FWD_ARCHS "sm_103a")
endif()
if(_VSA_FWD_GENCODE)
message(STATUS "fastvideo-kernel: building block_sparse_sm100a (${_VSA_FWD_ARCHS}; 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;-gencode;arch=compute_103a,code=sm_103a;-DVSA_BHSD=true")
# VSA block-sparse attention BACKWARD, 64-token blocks. sm_100a only for now: validated
# on GB200, not yet on B300/GB300, so no sm_103a image is built and the Python side keeps
# the Triton backward for sm_103a devices.
COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true")
endif()
# VSA block-sparse attention BACKWARD, 64-token blocks: the same images as the forward.
if(_VSA_FWD_GENCODE)
list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu)
set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu PROPERTIES
COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-DVSA_BHSD=true")
COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true")
endif()
Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI
@@ -434,7 +448,7 @@ if(BUILD_CXX_KERNELS)
# Build compile definitions list
set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops)
if(ENABLE_VSA_SM100_FAMILY)
if(ENABLE_VSA_SM100A OR ENABLE_VSA_SM103A)
list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A)
endif()
if(ENABLE_TK_KERNELS)
fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh
+12 -4
View File
@@ -123,7 +123,9 @@ __global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) vsa
const int* __restrict__ k2q_num, const int* __restrict__ workitem_remap,
const int* __restrict__ variable_block_sizes, int max_q_blocks, int num_samples, int num_heads,
int seqlen, float scale_log2, float sm_scale) {
#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL))
#if !defined(__CUDA_ARCH__) || \
((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \
(__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL)))
using DQ = DQConfig<DQ_DTYPE>;
const int num_kv_blocks_per_seq = seqlen / BLOCK;
@@ -863,7 +865,9 @@ __global__ void __launch_bounds__(ORDER_THREADS, 1)
vsa_bwd_order_kernel(const int* __restrict__ k2q_idx, const int* __restrict__ k2q_num,
int max_q_blocks, int num_kv_blocks_per_seq, int order_bin, bool snake,
int* __restrict__ order_out) {
#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL))
#if !defined(__CUDA_ARCH__) || \
((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \
(__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL)))
extern __shared__ int order_smem[];
int* sbin = order_smem;
int* smid = order_smem + num_kv_blocks_per_seq;
@@ -908,7 +912,9 @@ __global__ void __launch_bounds__(256, 1)
__nv_bfloat16* __restrict__ dk, __nv_bfloat16* __restrict__ dv,
const int* __restrict__ k2q_num, int num_samples, int num_heads,
int seqlen) {
#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL))
#if !defined(__CUDA_ARCH__) || \
((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \
(__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL)))
__shared__ __align__(128) __nv_bfloat16 tile[PRE_TOKENS][SUB_COLS_BF16 + 4];
const int num_kv_blocks_per_seq = seqlen / BLOCK;
const int token_block_id = (int)blockIdx.x;
@@ -1038,7 +1044,9 @@ template <bool BHSD = false, typename DQ_DTYPE = float>
__global__ void __launch_bounds__(128, 1)
vsa_bwd_postprocess_kernel(const DQ_DTYPE* __restrict__ dqaccum, __nv_bfloat16* __restrict__ dq,
int num_heads, int seqlen, float sm_scale) {
#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL))
#if !defined(__CUDA_ARCH__) || \
((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \
(__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL)))
const int q_block_id = (int)blockIdx.x;
const int batch_head = (int)blockIdx.y;
const int batch = batch_head / num_heads, head = batch_head % num_heads;
fastvideo-kernel/csrc/common_extension.cpp
+1 -1
View File
@@ -60,7 +60,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
"VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a/sm103a)");
m.def("block_sparse_sm100a_bwd",
torch::wrap_pybind_function(block_sparse_sm100a_bwd),
"VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a)");
"VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a/sm103a)");
#endif
#ifdef TK_COMPILE_ST_ATTN
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py
+3 -3
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
"""sm_100a (Blackwell) CUDA block-sparse VSA backward.
"""sm_100a/sm_103a (data-center 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
@@ -23,7 +23,7 @@ except ImportError: # pragma: no cover - extension not built
_BWD = None
_HAS_VSA_BWD_SM100A = False
_SM100 = (10, 0)
_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)}
HEAD_DIM = 128
BLOCK = 64
# Must match the -DVSA_BHSD the extension was compiled with (FastVideo builds with true).
@@ -58,7 +58,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
"""
if not _HAS_VSA_BWD_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
tests/test_block_sparse_bwd_sm100a.py
+2 -2
View File
@@ -31,9 +31,9 @@ REL_MAX_TOL = 1e-2 # max|got - ref| / max|ref|
MEAN_ABS_TOL = 1e-3 # mean|got - ref|
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0)
not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}
or not bwd._HAS_VSA_BWD_SM100A,
reason="requires a compute-capability (10, 0) GPU (sm_100a) and a fastvideo_kernel "
reason="requires a data-center Blackwell GPU (sm_100a/sm_103a) and a fastvideo_kernel "
"extension built with block_sparse_sm100a_bwd",
)
tests/test_block_sparse_sm100a_dispatch.py
+3 -2
View File
@@ -175,8 +175,9 @@ def test_backward_uses_sm100a_kernel_when_built(monkeypatch):
"""Guards against a silent Triton fallback: with the op built on an sm_100a device the
CUDA backward must be the one that runs."""
from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd
if not vsa_bwd._HAS_VSA_BWD_SM100A or torch.cuda.get_device_capability() != (10, 0):
pytest.skip("sm_100a backward not built for this device")
if (not vsa_bwd._HAS_VSA_BWD_SM100A
or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}):
pytest.skip("sm_100a/sm_103a backward not built for this device")
_, _, sm100a_backward = _grads_sm100a_route_vs_triton(monkeypatch)
assert sm100a_backward