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

[kernel] GB10 (DGX Spark) CUDA VSA forward for 128-token tiles

Browse Source 
FastH3's tile-128 route runs the ptx_blackwell CUDA forward on sm_100/sm_103 and, on every other GPU,
splits each 128-token tile into two 64-token halves on the Triton kernels. On the DGX Spark that split
route is what OmniRef's BF16 attention runs today. This adds a 128-token forward built for sm_121a
(mma.sync + TMA; the GB10 has no tcgen05 or tensor memory) and routes the tile-128 call to it on a
GB10, ahead of the split.

- csrc/attention/block_sparse_kernel_blk128_sm121.cuh and block_sparse_blk128_sm121.cu: the kernel and
  its binding block_sparse_sm121_blk128_fwd(q, k, v, q2k_idx, q2k_num, variable_block_sizes) -> {out},
  inference only, no LSE.
- CMakeLists.txt: built when TORCH_CUDA_ARCH_LIST has 12.1a or 12.1 and CUDA >= 13, so both the
  aarch64 wheel and the GB10 Docker image carry it.
- fastvideo_kernel.block_sparse_attn_sm121: is_supported and an opaque custom op with a fake. It is a
  separate module because block_sparse_attn_ptx_blackwell.is_supported also gates an LSE-requiring
  autograd path.
- video_sparse_attn_h3: the tile-128 route tries it after ptx_blackwell and before split_tile128.
- Tests: kernel correctness against an FP32 reference (ragged blocks, empty selections, odd block
  counts, many waves, a NaN-filled allocator, determinism), and the H3 route (fakes, plus a real-GB10
  oracle test).
This commit is contained in:
SolitaryThinker
2026-10-20 05:39:54 +00:00
parent 376e445f0e
commit 32d8a8765b
11 changed files with 1116 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/inference/fasth3_omniref.md
+7 -2
View File
@@ -97,8 +97,13 @@ The Blackwell YAML examples also select FP4 attention and the prefix precision
through `pipeline.experimental.h3_vsa_fp4`. Explicit YAML values take precedence
over the corresponding environment flags. A missing native extension raises an
error. To choose BF16 VSA for OmniRef's 128-token tiles, set
`h3_vsa_fp4.enabled: false` and
`FASTVIDEO_H3_VSA_FINE_ATTENTION=emulator-bf16`.
`h3_vsa_fp4.enabled: false`. On the DGX Spark this runs the 128-token GB10
CUDA kernel (`fastvideo_kernel.block_sparse_attn_sm121`, built when
`TORCH_CUDA_ARCH_LIST` contains `12.1a` or `12.1` and CUDA is 13 or newer);
leave `FASTVIDEO_H3_VSA_FINE_ATTENTION` unset there, since it replaces that
kernel. On the other consumer cards, add
`FASTVIDEO_H3_VSA_FINE_ATTENTION=emulator-bf16`, or leave it unset to run each
128-token tile as two 64-token halves on the Triton kernels.
The Blackwell profiles select the FP8 prefix split to protect the conditioning
queries from FP4 rounding. The split runs prefix rows densely in bf16 after
fastvideo-kernel/CMakeLists.txt
+22
View File
@@ -486,6 +486,24 @@ if(BUILD_CXX_KERNELS)
csrc/attention/block_sparse_blk128_blackwell.cu PROPERTIES
COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true")
endif()
# VSA block-sparse attention FORWARD for the GB10 (DGX Spark), 128-token blocks: an mma.sync + TMA kernel
# for sm_121a, the device the tcgen05 kernel above cannot run on. Inference only (no LSE). "12.1" without the
# suffix counts too (the GB10 Docker image builds that), since this one source carries its own sm_121a image.
set(ENABLE_VSA_SM121A OFF)
if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(12\\.1a?|121a?|sm_121a?)([; ,]|$)")
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
set(ENABLE_VSA_SM121A ON)
else()
message(STATUS "fastvideo-kernel: skipping block_sparse_sm121 (sm_121a needs CUDA >= 13.0, have "
"${CMAKE_CUDA_COMPILER_VERSION})")
endif()
endif()
if(ENABLE_VSA_SM121A)
message(STATUS "fastvideo-kernel: building block_sparse_sm121 (sm_121a; 128-token blocks, forward)")
list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_blk128_sm121.cu)
set_source_files_properties(csrc/attention/block_sparse_blk128_sm121.cu PROPERTIES
COMPILE_OPTIONS "-gencode;arch=compute_121a,code=sm_121a")
endif()
# VSA block-sparse attention BACKWARD, 64- and 128-token blocks: the same images as the forward.
if(_VSA_FWD_GENCODE)
list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_blackwell.cu
@@ -513,6 +531,9 @@ if(BUILD_CXX_KERNELS)
if(ENABLE_VSA_SM100A OR ENABLE_VSA_SM103A)
list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_PTX_BLACKWELL)
endif()
if(ENABLE_VSA_SM121A)
list(APPEND COMPILE_DEFS FASTVIDEO_KERNEL_COMPILE_BLOCK_SPARSE_VSA_SM121)
endif()
if(ENABLE_TK_KERNELS)
list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE)
endif()
@@ -640,6 +661,7 @@ message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}")
message(STATUS "fastvideo_kernel_ops: ON (turbodiffusion int8-gemm/quant/rmsnorm/layernorm, all listed archs)")
message(STATUS " + TK sta/block_sparse (sm_90a only): ${ENABLE_TK_KERNELS}")
message(STATUS " + Ulysses NCCL-device all-to-all: ${ENABLE_ULYSSES_A2A}")
message(STATUS " + VSA tile-128 forward, GB10 (sm_121a): ${ENABLE_VSA_SM121A}")
message(STATUS "fp4attn/fp4quant (sm_120a CUDA >= 12.8; sm_121a CUDA >= 13): ${ENABLE_ATTN_QAT_INFER}")
message(STATUS "Triton fallbacks ship in python/fastvideo_kernel/triton_kernels regardless.")
message(STATUS "============================================================")
fastvideo-kernel/README.md
+2 -1
View File
@@ -10,7 +10,8 @@ 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*.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 | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*_blackwell.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 | MiniMax-H3 block-sparse VSA forward, 128-token blocks, inference (`block_sparse_sm121_blk128_fwd`, no LSE) | `csrc/attention/block_sparse*_sm121.cu*` | GB10 / DGX Spark `sm_121a` | ON iff `12.1a` or `12.1` is in `TORCH_CUDA_ARCH_LIST` and CUDA ≥ 13 |
| 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/workstation Blackwell `sm_120a` (CUDA ≥ 12.8) + `sm_121a` (DGX Spark GB10; CUDA ≥ 13) | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` or `12.1a` is in the arch list) |
fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu
+82
View File
@@ -0,0 +1,82 @@
// SPDX-License-Identifier: Apache-2.0
// block_sparse_blk128_sm121.cu -- torch binding for the GB10 (DGX Spark, sm_121a) VSA block-sparse forward,
// 128-token blocks (block_sparse_kernel_blk128_sm121.cuh).
//
// The inference counterpart of block_sparse_ptx_blackwell_blk128_fwd with need_lse=False: it returns {out} only.
// out is allocated here and written whole by the kernel (padding rows and tiles that select nothing are written
// zero), so it never needs to be zero-initialized.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <limits>
#include <vector>
#include "block_sparse_kernel_blk128_sm121.cuh"
namespace {
constexpr int64_t kBlock = 128;
constexpr int64_t kHeadDim = 128;
void check_tensor(const torch::Tensor& t, const char* name, at::ScalarType dtype, const at::Device& device) {
TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
TORCH_CHECK(t.device() == device, name, " is on ", t.device(), ", but q is on ", device);
TORCH_CHECK(t.scalar_type() == dtype, name, " must be ", dtype, ", not ", t.scalar_type());
TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
}
} // namespace
// q, k, v: [B, H, S, 128] bf16, contiguous (BHSD), S = num_blocks * 128. q2k_idx [B, H, num_blocks, max_kv],
// q2k_num [B, H, num_blocks] and variable_block_sizes [num_blocks]: int32, contiguous. Returns {out}.
std::vector<torch::Tensor> block_sparse_sm121_blk128_fwd(torch::Tensor q, torch::Tensor k, torch::Tensor v,
torch::Tensor q2k_idx, torch::Tensor q2k_num,
torch::Tensor variable_block_sizes) {
const at::Device device = q.device();
check_tensor(q, "q", at::kBFloat16, device);
check_tensor(k, "k", at::kBFloat16, device);
check_tensor(v, "v", at::kBFloat16, device);
TORCH_CHECK(q.dim() == 4, "q must be [batch, heads, seq, 128], got ", q.sizes());
TORCH_CHECK(k.sizes() == q.sizes() && v.sizes() == q.sizes(), "q, k and v must have one shape, got ", q.sizes(),
", ", k.sizes(), " and ", v.sizes());
const int64_t batch = q.size(0), heads = q.size(1), seq = q.size(2);
TORCH_CHECK(q.size(3) == kHeadDim, "the head dimension must be ", kHeadDim, ", got ", q.size(3));
check_tensor(q2k_idx, "q2k_idx", at::kInt, device);
check_tensor(q2k_num, "q2k_num", at::kInt, device);
check_tensor(variable_block_sizes, "variable_block_sizes", at::kInt, device);
TORCH_CHECK(variable_block_sizes.dim() == 1, "variable_block_sizes must be [num_blocks], got ",
variable_block_sizes.sizes());
const int64_t blocks = variable_block_sizes.size(0);
TORCH_CHECK(seq == blocks * kBlock, "seqlen ", seq, " must equal num_blocks (", blocks, ") * ", kBlock,
"; FastVideo pads the sequence up to whole blocks");
TORCH_CHECK(q2k_idx.dim() == 4 && q2k_idx.size(0) == batch && q2k_idx.size(1) == heads && q2k_idx.size(2) == blocks,
"q2k_idx must be [", batch, ", ", heads, ", ", blocks, ", max_kv], got ", q2k_idx.sizes());
TORCH_CHECK(q2k_num.dim() == 3 && q2k_num.size(0) == batch && q2k_num.size(1) == heads && q2k_num.size(2) == blocks,
"q2k_num must be [", batch, ", ", heads, ", ", blocks, "], got ", q2k_num.sizes());
// The grid is (num_blocks, batch * heads).
TORCH_CHECK(batch * heads <= 65535, "batch * heads must be at most 65535, got ", batch * heads);
TORCH_CHECK(blocks <= std::numeric_limits<int>::max() && q2k_idx.size(3) <= std::numeric_limits<int>::max(),
"too many blocks: ", blocks, " (max_kv ", q2k_idx.size(3), ")");
auto out = torch::empty_like(q);
if (q.numel() == 0) return {out};
const at::cuda::CUDAGuard guard(device);
vsa_sm121::Vsa128Args args;
args.q = q.data_ptr();
args.k = k.data_ptr();
args.v = v.data_ptr();
args.out = out.data_ptr();
args.q2k_idx = q2k_idx.data_ptr<int32_t>();
args.q2k_num = q2k_num.data_ptr<int32_t>();
args.block_sizes = variable_block_sizes.data_ptr<int32_t>();
args.batch_heads = static_cast<int>(batch * heads);
args.tiles = static_cast<int>(blocks);
args.max_kv = static_cast<int>(q2k_idx.size(3));
const cudaError_t err = vsa_sm121::launch(args, at::cuda::getCurrentCUDAStream());
TORCH_CHECK(err == cudaSuccess, "block_sparse_sm121_blk128_fwd launch failed: ", cudaGetErrorString(err));
return {out};
}
fastvideo-kernel/csrc/attention/block_sparse_kernel_blk128_sm121.cuh
+617
View File
@@ -0,0 +1,617 @@
// SPDX-License-Identifier: Apache-2.0
// block_sparse_kernel_blk128_sm121.cuh -- the VSA block-sparse attention forward for 128-token blocks on the
// NVIDIA GB10 (DGX Spark, sm_121a), where the data-center Blackwell kernel (tcgen05, tensor memory) cannot run.
//
// What it computes is block_sparse_ptx_blackwell_blk128_fwd's output with need_lse=False: BHSD bf16, head_dim
// 128; for each query tile i that selects q2k_num[i] key tiles q2k_idx[i, 0..count), every valid row r <
// variable_block_sizes[i] of the tile is softmax(q_r K^T / sqrt(128)) V over the first variable_block_sizes[j]
// keys of each selected tile j, and every other row (padding rows, tiles that select nothing) is written zero.
// Entries of q2k_idx past a tile's count are never read. No LSE: inference forwards only.
//
// How, on mma.sync (the GB10 has no wgmma, tcgen05 or tensor memory): one CTA of 512 threads per (query tile,
// batch * head), one CTA per SM (99 KB of shared memory).
// - Two groups of eight warps, 16 query rows each: the QK group computes S = Q K^T and the online softmax and
// publishes P (bf16) and the rescale factors through shared memory; the PV group accumulates O += P V. Each
// PV warp waits for its QK partner half-way through the partner's next QK burst, so the two groups' mma
// bursts overlap instead of finishing together.
// - K and V stream in 64-key chunks through TMA (two 64x64 boxes, 128-byte swizzle) into separate rings;
// the last warp to finish reading a slot refills it.
// - Each CTA walks its selection starting (launch index mod SMs) / SMs of the way through, so the CTAs that
// run at once read the key tiles they share at the same time while those are in the L2; a query tile that
// selects more than half of the tiles walks from the front.
// The order of the walk changes no result beyond the rounding of the softmax's sums, and depends only on the
// launch index, the counts and the SM count, so every call is deterministic.
#pragma once
#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <climits>
#include <cmath>
#include <cstdint>
namespace vsa_sm121 {
// What one call reads and writes. All pointers are device pointers to contiguous tensors; BH = batch * heads,
// T = tiles, S = T * 128 tokens.
struct Vsa128Args {
const void* q; // [BH, S, 128] bf16; a tile's padding rows may hold anything, they are not read
const void* k; // [BH, S, 128] bf16
const void* v; // [BH, S, 128] bf16
void* out; // [BH, S, 128] bf16; every element is written
const int32_t* q2k_idx; // [BH, T, max_kv]: each query tile's selected key tiles; entries past its count are not read
const int32_t* q2k_num; // [BH, T]: how many key tiles each query tile selects
const int32_t* block_sizes; // [T]: valid tokens of each tile, which are its first ones
int batch_heads; // BH
int tiles; // T
int max_kv; // the last dimension of q2k_idx
};
namespace ptx {
// Copies 16 bytes from global to shared memory without going through L1 (.cg). Only the first `src_bytes`
// are read; the rest of the 16 are zero-filled, so src_bytes == 0 writes zeros and reads nothing.
__device__ __forceinline__ void cp_async_16(uint32_t dst, const void* src, uint32_t src_bytes) {
asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"r"(dst), "l"(src), "r"(src_bytes) : "memory");
}
// Closes the group of the cp.async copies this thread issued since the last commit (an empty group if none).
// UNUSED by the current kernel, deliberately. The main loop publishes a stage's copies through that stage's
// mbarrier (`cp.async.mbarrier.arrive.noinc`), which waits for every thread's copies into the stage and makes
// them visible -- something a commit group cannot do, because a group is private to the thread that issued it
// and `cp.async.wait_group` below retires only this thread's own. Both wrappers stay because they are the right
// tool for a thread-private wait; re-introducing cp_async_wait() beside the loop's per-stage waits would not
// restore anything, it would only add a wait that cannot see the other warps' copies.
__device__ __forceinline__ void cp_async_commit() {
asm volatile("cp.async.commit_group;\n" ::: "memory");
}
// Waits until at most `kPending` of this thread's most recent groups are still in flight. The copies are
// visible to the other threads of the CTA only after a barrier that follows the wait.
template <int kPending>
__device__ __forceinline__ void cp_async_wait() {
asm volatile("cp.async.wait_group %0;\n" ::"n"(kPending) : "memory");
}
// Initialises the shared-memory mbarrier at `addr` to flip its phase every `count` arrivals.
__device__ __forceinline__ void mbar_init(uint32_t addr, uint32_t count) {
asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(addr), "r"(count) : "memory");
}
// Makes the initialisations of this thread visible to the whole CTA before any thread arrives.
__device__ __forceinline__ void mbar_init_fence() {
asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory");
}
// One arrival at the mbarrier `addr`.
__device__ __forceinline__ void mbar_arrive(uint32_t addr) {
uint64_t state;
asm volatile("mbarrier.arrive.shared::cta.b64 %0, [%1];\n" : "=l"(state) : "r"(addr) : "memory");
}
// The mbarrier `addr` is signalled when this thread's prior cp.async copies have completed and are visible to
// the threads that observe the barrier's phase flip. One arrival per thread per pass.
__device__ __forceinline__ void cp_async_mbar_arrive(uint32_t addr) {
asm volatile("cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n" ::"r"(addr) : "memory");
}
// True once phase `phase` (0 or 1, alternating) of the mbarrier `addr` has completed; does not block.
__device__ __forceinline__ bool mbar_try_wait(uint32_t addr, uint32_t phase) {
uint32_t done;
asm volatile("{\n"
"\t.reg .pred p;\n"
"\tmbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n"
"\tselp.b32 %0, 1, 0, p;\n"
"}\n"
: "=r"(done)
: "r"(addr), "r"(phase)
: "memory");
return done != 0;
}
// Waits until phase `phase` of the mbarrier `addr` has completed.
__device__ __forceinline__ void mbar_wait(uint32_t addr, uint32_t phase) {
while (!mbar_try_wait(addr, phase)) {
}
}
// Brings the 128-byte line holding `p` into L2.
__device__ __forceinline__ void prefetch_l2(const void* p) {
asm volatile("prefetch.global.L2 [%0];\n" ::"l"(p));
}
// Loads four 8x8 matrices of 16-bit elements: lanes 8i..8i+7 give the addresses of the 8 rows (16 bytes each)
// of matrix i, and r[i] receives, in every lane, elements (lane / 4, 2 * (lane % 4) + {0, 1}) of matrix i.
__device__ __forceinline__ void ldmatrix_x4(uint32_t (&r)[4], uint32_t addr) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
: "r"(addr));
}
// As ldmatrix_x4, with each matrix transposed: r[i] receives elements (2 * (lane % 4) + {0, 1}, lane / 4).
__device__ __forceinline__ void ldmatrix_x4_trans(uint32_t (&r)[4], uint32_t addr) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
: "r"(addr));
}
// d += a * b for a 16x16 bf16 A (row-major), a 16x8 bf16 B (column-major) and a 16x8 fp32 D, held as the
// warp-wide fragments of mma.m16n8k16. With g = lane / 4 and c = 2 * (lane % 4):
// a: {A[g][c..c+1], A[g+8][c..c+1], A[g][c+8..c+9], A[g+8][c+8..c+9]}, two bf16 per register, lower k first
// b: {B[c..c+1][g], B[c+8..c+9][g]}
// d: {D[g][c], D[g][c+1], D[g+8][c], D[g+8][c+1]}
__device__ __forceinline__ void mma_bf16_16816(float (&d)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) {
asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, "
"{%0, %1, %2, %3};\n"
: "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1));
}
// Two floats rounded to nearest-even bf16 and packed, `lo` in the lower half: the element with the lower index.
__device__ __forceinline__ uint32_t pack_bf16x2(float lo, float hi) {
uint32_t r;
asm("cvt.rn.bf16x2.f32 %0, %1, %2;\n" : "=r"(r) : "f"(hi), "f"(lo));
return r;
}
// 2^x by the special function unit (relative error about 2^-22; exp2(-inf) = +0).
__device__ __forceinline__ float exp2_approx(float x) {
float y;
asm("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(y) : "f"(x));
return y;
}
} // namespace ptx
namespace split128 {
using bf16 = __nv_bfloat16;
constexpr int Tile = 128, Dim = 128, Chunk = 64;
constexpr int PhaseLead = 8;
constexpr int GroupThreads = 256, Threads = 512;
constexpr int OperandBytes = 16384, BoxBytes = 8192;
constexpr int KOffset = 0, VOffset = 32768, POffset = 81920;
constexpr int AlphaOffset = 98304;
constexpr int SumOffset = AlphaOffset + 512;
constexpr int BarOffset = SumOffset + 512;
constexpr int CounterOffset = BarOffset + 104; // 5 operand-full + 8 P-full
constexpr int SharedBytes = CounterOffset + 24; // 5 counters, 8-byte rounding
static_assert(SharedBytes <= 99 * 1024, "GB10 shared-memory budget");
constexpr float ScaleLog2 = 1.4426950408889634f * 0.088388347648318441f;
__device__ __forceinline__ int clamp_int(int a, int lo, int hi) { return min(max(a,lo),hi); }
// Q and output retain the original 256-byte-row swizzle. K/V alone use TMA boxes.
__device__ __forceinline__ uint32_t swizzle(int row, int unit) {
return row * 256 + ((unit ^ (row & 7)) << 4);
}
__device__ __forceinline__ uint32_t p_swizzle(int row, int unit) {
return row * 128 + ((unit ^ (row & 7)) << 4);
}
__device__ __forceinline__ uint32_t lane_offset(int key, int u0, int j) {
return 128 * key + 16 * ((2 * j + u0) ^ (key & 7));
}
// The walk of one CTA's selection. It reads entries start..count-1, then
// 0..start-1: every entry once, in a rotated order (see `walk_start`).
struct Chunks {
const int* selected;
const int* sizes;
// Low eight bits: tile valid count; bit eight: second half; bits nine
// and up: the entry the walk starts at. Keep the replicated walker one
// register smaller alongside the PV accumulator.
int count, tiles, pos, tile, state;
__device__ Chunks(const int* s, const int* b, int n, int t, int start)
: selected(s), sizes(b), count(n), tiles(t), pos(0), tile(0), state(256 | (start << 9)) {}
__device__ __forceinline__ bool next() {
const int now = state & 511;
if (now > Chunk && now < 256) { state |= 256; return true; }
while (pos < count) {
int at = pos++ + (state >> 9);
if (at >= count) at -= count;
tile = __ldg(selected + at);
const int valid = static_cast<unsigned>(tile) < static_cast<unsigned>(tiles)
? clamp_int(__ldg(sizes + tile), 0, Tile) : 0;
state = (state & ~511) | valid;
if (valid) return true;
}
return false;
}
__device__ __forceinline__ int keys() const { return min((state & 255) - (((state >> 8) & 1) * Chunk), Chunk); }
__device__ __forceinline__ int first_row() const { return tile * Tile + (((state >> 8) & 1) * Chunk); }
};
// Where a CTA's walk starts. CTAs start in launch order, `slots` at a time
// (one per SM), and each walks its selection over its whole life. Started
// at entry 0, the CTAs that run together are spread over every phase of
// their walks; the n-th of each `slots` consecutive CTAs instead starts
// n/slots of the way through its own selection, so that the CTAs running
// together are at about the same phase of their walks and read the key
// tiles they share at about the same time, while those are in the L2.
// A CTA that selects more than half of the tiles walks from entry 0: those
// start together and share one stream in order (phasing them measured
// -13 % at the long rows, PR41). The order changes no result beyond the
// rounding of the online softmax's sums, and depends only on the launch
// index, the count and the SM count, so a call is deterministic.
__device__ __forceinline__ int walk_start(int64_t launch, int count, int tiles, int slots) {
return 2 * count <= tiles ? static_cast<int>((launch % slots) * count / slots) : 0;
}
__device__ __forceinline__ void tma_load_2d(uint32_t dst, const void* map,
int col, int row, uint32_t full) {
asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes "
"[%0], [%1, {%2, %3}], [%4];"
:: "r"(dst), "l"(map), "r"(col), "r"(row), "r"(full) : "memory");
}
// Called by one lane only. The full barrier has arrival count one and 16 KB
// expected transactions (two 64x64 boxes); masks still use the valid key count.
__device__ __forceinline__ void load_operand(uint32_t dst, const void* map,
const Chunks& c, int head_row, uint32_t full) {
// Host checks BH*S <= INT32_MAX; tensor coordinates are rows, never elements.
const int row = head_row + c.first_row();
asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
asm volatile("{ .reg .b64 st; mbarrier.arrive.expect_tx.shared::cta.b64 st, [%0], %1; }"
:: "r"(full), "r"(OperandBytes) : "memory");
tma_load_2d(dst, map, 0, row, full);
tma_load_2d(dst + BoxBytes, map, 64, row, full);
}
// Every warp finishes all of its reads before its one counter arrival. The
// CTA-scoped acq_rel RMW chains those reads to the elected warp's TMA issue.
// All lanes receive the result so the last-warp branch is uniform.
__device__ __forceinline__ bool last_reader(uint32_t counter, int lane) {
__syncwarp();
uint32_t old = 0;
if (lane == 0) {
asm volatile("atom.acq_rel.cta.shared::cta.add.u32 %0, [%1], 1;"
: "=r"(old) : "r"(counter) : "memory");
}
old = __shfl_sync(0xffffffffu, old, 0);
return old == 7;
}
__device__ __forceinline__ void reset_counter(uint32_t counter) {
asm volatile("st.shared.u32 [%0], 0;" :: "r"(counter) : "memory");
}
__device__ __forceinline__ void load_q(uint32_t dst, const bf16* q, int rows, int tid) {
#pragma unroll
for (int i=0; i<4; ++i) {
const int u = tid + i * Threads;
const int row = u/16, unit = u%16;
const bool valid = row < rows;
ptx::cp_async_16(dst+swizzle(row,unit), q+(valid?row:0)*Dim+unit*8, valid?16:0);
}
}
__device__ __forceinline__ void q_fragments(uint32_t (&qf)[8][4], uint32_t base, int warp, int lane) {
const int row = warp*16+(lane&7)+((lane>>3)&1)*8;
const int unit = lane>>4;
#pragma unroll
for (int kk=0; kk<8; ++kk) ptx::ldmatrix_x4(qf[kk],base+swizzle(row,2*kk+unit));
}
template<int Begin,int End>
__device__ __forceinline__ void scores(float (&s)[8][4], const uint32_t (&qf)[8][4], const uint32_t (&a)[4]) {
if constexpr (Begin==0) {
#pragma unroll
for (int j=0; j<8; ++j) s[j][0]=s[j][1]=s[j][2]=s[j][3]=0.f;
}
#pragma unroll
for (int kk=Begin; kk<End; ++kk) {
#pragma unroll
for (int nn=0; nn<4; ++nn) {
uint32_t b[4];
ptx::ldmatrix_x4(b,a[kk&3]+2048*nn+BoxBytes*(kk>>2));
ptx::mma_bf16_16816(s[2*nn],qf[kk],b[0],b[1]);
ptx::mma_bf16_16816(s[2*nn+1],qf[kk],b[2],b[3]);
}
}
}
__device__ __forceinline__ void mask(float (&s)[8][4], int keys, int lane) {
const int col = 2*(lane&3);
#pragma unroll
for (int j=0; j<8; ++j) {
#pragma unroll
for (int e=0; e<2; ++e) {
if (8*j+col+e>=keys) s[j][e]=s[j][2+e]=-INFINITY;
}
}
}
__device__ __forceinline__ void softmax(float (&s)[8][4], float (&m)[2], float (&l)[2], float (&alpha)[2]) {
float mx[2]={m[0],m[1]};
#pragma unroll
for (int j=0; j<8; ++j) {
mx[0]=fmaxf(mx[0],fmaxf(s[j][0],s[j][1]));
mx[1]=fmaxf(mx[1],fmaxf(s[j][2],s[j][3]));
}
#pragma unroll
for (int r=0; r<2; ++r) {
mx[r]=fmaxf(mx[r],__shfl_xor_sync(0xffffffffu,mx[r],1));
mx[r]=fmaxf(mx[r],__shfl_xor_sync(0xffffffffu,mx[r],2));
}
float bias[2];
#pragma unroll
for (int r=0; r<2; ++r) {
alpha[r]=ptx::exp2_approx((m[r]-mx[r])*ScaleLog2);
bias[r]=-mx[r]*ScaleLog2;
m[r]=mx[r];
}
float sum[2]={0.f,0.f};
#pragma unroll
for (int j=0; j<8; ++j) {
s[j][0]=ptx::exp2_approx(fmaf(s[j][0],ScaleLog2,bias[0]));
s[j][1]=ptx::exp2_approx(fmaf(s[j][1],ScaleLog2,bias[0]));
s[j][2]=ptx::exp2_approx(fmaf(s[j][2],ScaleLog2,bias[1]));
s[j][3]=ptx::exp2_approx(fmaf(s[j][3],ScaleLog2,bias[1]));
sum[0]+=s[j][0]+s[j][1];
sum[1]+=s[j][2]+s[j][3];
}
#pragma unroll
for (int r=0; r<2; ++r) l[r]=l[r]*alpha[r]+sum[r];
}
__device__ __forceinline__ void store_p(uint32_t dst, const float (&p)[8][4], int warp, int lane) {
const int row=warp*16+(lane&7)+((lane>>3)&1)*8;
const int unit=lane>>4;
#pragma unroll
for (int kk=0; kk<4; ++kk) {
const uint32_t a0=ptx::pack_bf16x2(p[2*kk][0],p[2*kk][1]);
const uint32_t a1=ptx::pack_bf16x2(p[2*kk][2],p[2*kk][3]);
const uint32_t a2=ptx::pack_bf16x2(p[2*kk+1][0],p[2*kk+1][1]);
const uint32_t a3=ptx::pack_bf16x2(p[2*kk+1][2],p[2*kk+1][3]);
const uint32_t addr=dst+p_swizzle(row,2*kk+unit);
asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};"
:: "r"(addr), "r"(a0), "r"(a1), "r"(a2), "r"(a3) : "memory");
}
}
__device__ __forceinline__ void pv(float (&o)[16][4], const uint32_t (&p)[4][4], const uint32_t (&v_addr)[4]) {
#pragma unroll
for (int kk=0; kk<4; ++kk) {
#pragma unroll
for (int dd=0; dd<8; ++dd) {
uint32_t b[4];
ptx::ldmatrix_x4_trans(b,v_addr[dd&3]+2048*kk+BoxBytes*(dd>>2));
ptx::mma_bf16_16816(o[2*dd],p[kk],b[0],b[1]);
ptx::mma_bf16_16816(o[2*dd+1],p[kk],b[2],b[3]);
}
}
}
__device__ __forceinline__ void output(bf16* out, uint8_t* staging, const float (&o)[16][4],
const float* sums, int rows, int warp, int lane) {
const int row0=warp*16+(lane>>2);
bool keep[2]; float inv[2];
#pragma unroll
for (int r=0; r<2; ++r) {
const float l=sums[row0+8*r];
keep[r]=row0+8*r<rows && l>0.f;
inv[r]=keep[r]?1.f/l:0.f;
}
#pragma unroll
for (int d=0; d<16; ++d) {
#pragma unroll
for (int r=0; r<2; ++r) {
const uint32_t val=keep[r]?ptx::pack_bf16x2(o[d][2*r]*inv[r],o[d][2*r+1]*inv[r]):0u;
*reinterpret_cast<uint32_t*>(staging+swizzle(row0+8*r,d)+4*(lane&3))=val;
}
}
__syncwarp();
#pragma unroll
for (int i=0; i<8; ++i) {
const int row=warp*16+2*i+(lane>>4), unit=lane&15;
const uint4 val=*reinterpret_cast<const uint4*>(staging+swizzle(row,unit));
*reinterpret_cast<uint4*>(out+row*Dim+unit*8)=val;
}
}
__device__ __forceinline__ void pair_sync(int warp) {
asm volatile("bar.sync %0, 64;" :: "r"(warp+1) : "memory");
}
// slots: the CTAs that run at once, one per SM; it sets the walks' phases.
__global__ __launch_bounds__(Threads,1) void forward(const Vsa128Args args,
const __grid_constant__ CUtensorMap kmap, const __grid_constant__ CUtensorMap vmap, const int slots) {
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ == 1210
// 128B TMA swizzle repeats over 1024B; all slot starts are multiples of 16KB.
extern __shared__ __align__(1024) uint8_t smem[];
const uint32_t base=static_cast<uint32_t>(__cvta_generic_to_shared(smem));
const int tid=threadIdx.x, lane=tid&31, local_tid=tid&255, warp=local_tid>>5;
const int tile=blockIdx.x, bh=blockIdx.y, tiles=args.tiles;
const int head_row=bh*tiles*Tile; // host-checked TMA row coordinate range
const int64_t qo=(int64_t(head_row)+int64_t(tile)*Tile)*Dim;
const int64_t query=int64_t(bh)*tiles+tile;
const int rows=clamp_int(__ldg(args.block_sizes+tile),0,Tile);
const int count=clamp_int(__ldg(args.q2k_num+query),0,args.max_kv);
bf16* out=static_cast<bf16*>(args.out)+qo;
if (rows==0 || count==0) {
#pragma unroll
for (int i=0;i<4;++i) reinterpret_cast<uint4*>(out)[tid+i*Threads]=make_uint4(0,0,0,0);
return;
}
const int* selected=args.q2k_idx+query*args.max_kv;
const uint32_t bar=base+BarOffset, counters=base+CounterOffset;
const auto operand_full=[bar](int op,int stage) { return bar+(op*2+stage)*8; };
const auto operand_counter=[counters](int op,int stage) { return counters+(op*2+stage)*4; };
const uint32_t p_full=bar+40+warp*8;
if (tid==0) {
#pragma unroll
for (int i=0;i<5;++i) {
ptx::mbar_init(bar+i*8,1);
reset_counter(counters+i*4);
}
#pragma unroll
for (int i=0;i<8;++i) ptx::mbar_init(bar+40+i*8,32);
ptx::mbar_init_fence();
}
// Q aliases both K slots only in the prologue. Its original cp.async and
// swizzle are unchanged; all QK warps rendezvous after capturing Q.
load_q(base+KOffset,static_cast<const bf16*>(args.q)+qo,rows,tid);
ptx::cp_async_commit();
ptx::cp_async_wait<0>();
__syncthreads();
const bool active=warp*16<rows;
// The grid is (tiles, heads) with tiles fastest, so `query` is the launch index.
Chunks fill(selected,args.block_sizes,count,tiles,walk_start(query,count,tiles,slots));
if (tid<GroupThreads) {
uint32_t qf[8][4];
if (active) q_fragments(qf,base+KOffset,warp,lane);
asm volatile("bar.sync 9, 256;" ::: "memory");
int pending=0, keys0=0, keys1=0;
if (fill.next()) {
keys0=fill.keys(); ++pending;
if (local_tid==0) load_operand(base+KOffset,&kmap,fill,head_row,operand_full(0,0));
}
if (fill.next()) {
keys1=fill.keys(); ++pending;
if (local_tid==0) load_operand(base+KOffset+OperandBytes,&kmap,fill,head_row,operand_full(0,1));
}
float m[2]={-INFINITY,-INFINITY}, l[2]={0.f,0.f};
const int k_key=(lane&7)+((lane>>4)&1)*8, k_u0=(lane>>3)&1;
uint32_t klane[4];
#pragma unroll
for (int j=0;j<4;++j) klane[j]=lane_offset(k_key,k_u0,j);
int iter=0;
while (pending) {
const int stage=iter&1, phase=(iter>>1)&1;
ptx::mbar_wait(operand_full(0,stage),phase);
float s[8][4];
uint32_t kaddr[4];
#pragma unroll
for (int j=0;j<4;++j) kaddr[j]=base+KOffset+stage*OperandBytes+klane[j];
if (active) scores<0,PhaseLead>(s,qf,kaddr);
if (iter) pair_sync(warp);
if (active) scores<PhaseLead,8>(s,qf,kaddr);
float alpha[2];
if (active) {
const int keys=stage?keys1:keys0;
if (keys<Chunk) mask(s,keys,lane);
softmax(s,m,l,alpha);
store_p(base+POffset,s,warp,lane);
if ((lane&3)==0) {
float* a=reinterpret_cast<float*>(smem+AlphaOffset);
a[warp*16+(lane>>2)]=alpha[0];
a[warp*16+(lane>>2)+8]=alpha[1];
}
}
ptx::mbar_arrive(p_full);
// The fill walkers advance identically in every lane and warp.
// On iteration i the just-consumed K slot takes chunk i+2.
if (fill.next()) {
const int keys=fill.keys();
if (stage) keys1=keys; else keys0=keys;
const uint32_t counter=operand_counter(0,stage);
if (last_reader(counter,lane)) {
if (lane==0) {
reset_counter(counter);
load_operand(base+KOffset+stage*OperandBytes,&kmap,fill,head_row,operand_full(0,stage));
}
}
} else --pending;
++iter;
}
if (iter) pair_sync(warp);
float* sums=reinterpret_cast<float*>(smem+SumOffset);
#pragma unroll
for (int r=0;r<2;++r) {
float sum=l[r];
sum+=__shfl_xor_sync(0xffffffffu,sum,1);
sum+=__shfl_xor_sync(0xffffffffu,sum,2);
if ((lane&3)==0) sums[warp*16+(lane>>2)+8*r]=sum;
}
// Every issued operand has been consumed; this publishes sums and
// prevents PV output staging from overwriting QK's K alias too early.
asm volatile("bar.sync 0, 512;" ::: "memory");
} else {
int pending=0;
#pragma unroll
for (int s=0;s<3;++s) {
if (fill.next()) {
++pending;
if (local_tid==0) load_operand(base+VOffset+s*OperandBytes,&vmap,fill,head_row,operand_full(1,s));
}
}
float o[16][4];
#pragma unroll
for (int d=0;d<16;++d) o[d][0]=o[d][1]=o[d][2]=o[d][3]=0.f;
const int v_key=(lane&7)+((lane>>3)&1)*8, v_u0=lane>>4;
uint32_t vlane[4];
#pragma unroll
for (int j=0;j<4;++j) vlane[j]=lane_offset(v_key,v_u0,j);
const uint32_t plane=p_swizzle(warp*16+v_key,v_u0);
int iter=0, stage=0;
uint32_t phase=0;
while (pending) {
ptx::mbar_wait(operand_full(1,stage),phase);
ptx::mbar_wait(p_full,iter&1);
uint32_t pf[4][4];
float a0=0.f,a1=0.f;
if (active) {
const float* a=reinterpret_cast<const float*>(smem+AlphaOffset);
a0=a[warp*16+(lane>>2)]; a1=a[warp*16+(lane>>2)+8];
#pragma unroll
for (int kk=0;kk<4;++kk) ptx::ldmatrix_x4(pf[kk],base+POffset+(plane^(32*kk)));
}
pair_sync(warp);
if (active) {
#pragma unroll
for (int d=0;d<16;++d) {
o[d][0]*=a0; o[d][1]*=a0;
o[d][2]*=a1; o[d][3]*=a1;
}
uint32_t vaddr[4];
#pragma unroll
for (int j=0;j<4;++j) vaddr[j]=base+VOffset+stage*OperandBytes+vlane[j];
pv(o,pf,vaddr);
}
// All three V slots were prefetched. Refill exactly the slot just
// read with i+3, not A3's i+2 slot (stage+2)%3.
if (fill.next()) {
const uint32_t counter=operand_counter(1,stage);
if (last_reader(counter,lane)) {
if (lane==0) {
reset_counter(counter);
load_operand(base+VOffset+stage*OperandBytes,&vmap,fill,head_row,operand_full(1,stage));
}
}
} else --pending;
if (++stage==3) { stage=0; phase^=1; }
++iter;
}
asm volatile("bar.sync 0, 512;" ::: "memory");
output(out,smem,o,reinterpret_cast<const float*>(smem+SumOffset),rows,warp,lane);
}
#endif // sm_121: every other device pass compiles an empty kernel
}
CUresult encode_map(CUtensorMap* map, const void* ptr, int64_t rows) {
const cuuint64_t dims[2]={Dim,static_cast<cuuint64_t>(rows)};
const cuuint64_t strides[1]={Dim*2};
const cuuint32_t box[2]={64,64}, elem_strides[2]={1,1};
return cuTensorMapEncodeTiled(map,CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,2,const_cast<void*>(ptr),
dims,strides,box,elem_strides,CU_TENSOR_MAP_INTERLEAVE_NONE,CU_TENSOR_MAP_SWIZZLE_128B,
CU_TENSOR_MAP_L2_PROMOTION_L2_256B,CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
}
} // namespace split128
inline cudaError_t launch(const Vsa128Args& args,cudaStream_t stream) {
if (args.batch_heads<=0 || args.tiles<=0) return cudaSuccess;
const int64_t rows=int64_t(args.batch_heads)*args.tiles*split128::Tile;
if (rows>INT32_MAX) return cudaErrorInvalidValue;
// The walker keeps its start entry (< count <= tiles) in bits 9..30 of an int.
if (args.tiles>=(1<<22) || args.max_kv>=(1<<22)) return cudaErrorInvalidValue;
CUtensorMap kmap,vmap;
if (split128::encode_map(&kmap,args.k,rows)!=CUDA_SUCCESS ||
split128::encode_map(&vmap,args.v,rows)!=CUDA_SUCCESS) return cudaErrorInvalidValue;
cudaError_t err=cudaFuncSetAttribute(split128::forward,cudaFuncAttributeMaxDynamicSharedMemorySize,split128::SharedBytes);
if (err!=cudaSuccess) return err;
// One CTA per SM fits (99 KB of shared memory each): the SMs are the
// CTAs that run at once. A device property, read once per device.
int device=0;
err=cudaGetDevice(&device);
if (err!=cudaSuccess) return err;
static int slots_device=-1, slots=1;
if (device!=slots_device) {
int sms=0;
err=cudaDeviceGetAttribute(&sms,cudaDevAttrMultiProcessorCount,device);
if (err!=cudaSuccess) return err;
slots=sms>0 ? sms : 1;
slots_device=device;
}
split128::forward<<<dim3(args.tiles,args.batch_heads),split128::Threads,split128::SharedBytes,stream>>>(args,kmap,vmap,slots);
return cudaGetLastError();
}
} // namespace vsa_sm121
fastvideo-kernel/csrc/common_extension.cpp
+12
View File
@@ -52,6 +52,12 @@ extern std::vector<torch::Tensor> block_sparse_ptx_blackwell_blk128_bwd(
torch::Tensor variable_block_sizes, double sm_scale);
#endif
#ifdef FASTVIDEO_KERNEL_COMPILE_BLOCK_SPARSE_VSA_SM121
extern std::vector<torch::Tensor> block_sparse_sm121_blk128_fwd(
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor q2k_idx, torch::Tensor q2k_num,
torch::Tensor variable_block_sizes);
#endif
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.doc() = "FastVideo CUDA Kernels";
@@ -70,6 +76,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
"VSA block-sparse attention backward, 128-token blocks (Blackwell sm100a/sm103a)");
#endif
#ifdef FASTVIDEO_KERNEL_COMPILE_BLOCK_SPARSE_VSA_SM121
m.def("block_sparse_sm121_blk128_fwd",
torch::wrap_pybind_function(block_sparse_sm121_blk128_fwd),
"VSA block-sparse attention forward, 128-token blocks, inference only (GB10 / DGX Spark, sm121a)");
#endif
#ifdef TK_COMPILE_ST_ATTN
m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention (Hopper)");
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm121.py
+113
View File
@@ -0,0 +1,113 @@
# SPDX-License-Identifier: Apache-2.0
"""GB10 (DGX Spark, sm_121a) CUDA block-sparse VSA forward, 128-token blocks.
The GB10 counterpart of ``block_sparse_attn_ptx_blackwell``'s 128-token forward, which is a
tcgen05/tensor-memory kernel the GB10 cannot run. This one is an mma.sync + TMA kernel built for
sm_121a (``csrc/attention/block_sparse_kernel_blk128_sm121.cuh``). It computes the same output
for the same inputs, and nothing else: it is inference-only and returns no LSE, so it has no
backward and is never reached by grad-tracking calls.
It is a separate module, not another capability in ``block_sparse_attn_ptx_blackwell``, because
that module's ``is_supported`` also gates ``block_sparse_attn_from_indices``'s autograd op, which
always asks for the LSE.
"""
from typing import Tuple
import torch
try:
# The pybind symbols live on fastvideo_kernel_ops, not on the _C namespace package.
from fastvideo_kernel._C import fastvideo_kernel_ops as _C
_FWD = getattr(_C, "block_sparse_sm121_blk128_fwd", None)
except ImportError: # pragma: no cover - extension not built
_C = None
_FWD = None
_HAS_VSA_SM121 = _FWD is not None
_SUPPORTED_COMPUTE_CAPABILITIES = {(12, 1)}
BLOCK = 128
HEAD_DIM = 128
def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
"""True iff this build can run these tensors on this device; otherwise the caller falls back.
Static facts only (shapes, dtypes, device, layout), so it never synchronizes with the GPU:
* the extension carries the sm_121a kernel and q is on a GB10 (compute capability 12.1);
* q is a contiguous 4-D bf16 BHSD tensor with head_dim 128 and seqlen == num_blocks * 128,
num_blocks = variable_block_sizes.numel(); any num_blocks, odd included;
* variable_block_sizes is an integer CUDA tensor.
What the kernel accepts beyond that, and is tested to handle: per-row counts q2k_num in
[0, max_kv], non-uniform across rows (a row with count 0 gets exactly-zero output); q2k_idx
entries are only read below each row's count (map_to_index's -1 padding is never read); any
variable_block_sizes in [0, 128] (keys at or past a block's count are masked, and query rows
past it are written zero).
"""
if not _HAS_VSA_SM121 or not q.is_cuda:
return False
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 or not q.is_contiguous():
return False
num_blocks = variable_block_sizes.numel()
if num_blocks == 0 or q.shape[2] != num_blocks * BLOCK:
return False
if not variable_block_sizes.is_cuda or variable_block_sizes.dtype not in (torch.int32, torch.int64):
return False
return True
@torch.library.custom_op(
"fastvideo_kernel::block_sparse_attn_sm121_inference",
mutates_args=(),
device_types="cuda",
)
def _block_sparse_attn_sm121_inference(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
q2k_idx: torch.Tensor,
q2k_num: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> torch.Tensor:
"""Opaque launch, so that Dynamo does not descend into the raw pybind function."""
return _FWD(q, k, v, q2k_idx, q2k_num, variable_block_sizes)[0]
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm121_inference")
def _block_sparse_attn_sm121_inference_fake(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
q2k_idx: torch.Tensor,
q2k_num: torch.Tensor,
variable_block_sizes: torch.Tensor,
) -> torch.Tensor:
# The C++ binding allocates its output with torch::empty_like(q).
return torch.empty_like(q)
def block_sparse_attn_sm121(
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, None]:
"""Forward pass, inference only. Returns ``(out, None)``; ``out`` has q's [B, H, S, 128] layout.
The signature matches ``block_sparse_attn_ptx_blackwell(..., need_lse=False)``: q2k_idx is
[B, H, num_blocks, max_kv] and q2k_num [B, H, num_blocks], as ``map_to_index`` returns them.
"""
if _FWD is None:
raise RuntimeError("fastvideo_kernel was built without the sm_121a block-sparse forward "
"(build with TORCH_CUDA_ARCH_LIST containing 12.1a and CUDA >= 13)")
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
idx = q2k_idx.to(torch.int32).contiguous()
num = q2k_num.to(torch.int32).contiguous()
vbs = variable_block_sizes.to(torch.int32).contiguous()
return _block_sparse_attn_sm121_inference(q, k, v, idx, num, vbs), None
fastvideo-kernel/tests/test_block_sparse_sm121.py
+146
View File
@@ -0,0 +1,146 @@
# SPDX-License-Identifier: Apache-2.0
"""Correctness tests for the GB10 (DGX Spark, sm_121a) block-sparse VSA forward, 128-token blocks.
Compared against an explicit FP32 PyTorch reference of the semantics the kernel implements: each
valid query row attends to the valid keys of its block's selected key blocks; query rows at or past
their block's variable_block_sizes, and every row of a block that selects nothing, are zero (H3
discards them; the kernel writes them anyway, since its output comes from torch::empty_like). So
that a missed store cannot hide behind a recycled zero buffer, the tests that check the zero rows
first fill the allocator with NaN.
Run with: python -m pytest fastvideo-kernel/tests/test_block_sparse_sm121.py -v
"""
import pytest
import torch
from fastvideo_kernel import block_sparse_attn_sm121 as vsa
BLOCK = 128
HEAD_DIM = 128
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 1) or not vsa._HAS_VSA_SM121,
reason="requires a GB10 (DGX Spark, sm_121) and a fastvideo_kernel build with the sm_121a forward",
)
def make_case(num_blocks=8, topk=4, heads=4, batch=1, ragged=False, counts=None, empty_blocks=(), seed=0):
"""q/k/v [B, H, S, 128] with zero padding rows, map_to_index-shaped q2k_idx/q2k_num, and vbs.
``counts`` cycles over the rows instead of ``topk``; ``empty_blocks`` get variable_block_sizes 0.
"""
gen = torch.Generator(device="cuda").manual_seed(seed)
seq = num_blocks * BLOCK
q, k, v = (torch.randn((batch, heads, seq, HEAD_DIM), device="cuda", dtype=torch.bfloat16, generator=gen)
for _ in range(3))
rows = batch * heads * num_blocks
idx = torch.full((rows, num_blocks), -1, dtype=torch.int32, device="cuda")
num = torch.empty(rows, dtype=torch.int32, device="cuda")
for r in range(rows):
n = min(num_blocks, topk if counts is None else int(counts[r % len(counts)]))
idx[r, :n] = torch.randperm(num_blocks, device="cuda", generator=gen)[:n].sort().values.to(torch.int32)
num[r] = n
if ragged:
vbs = torch.randint(1, BLOCK + 1, (num_blocks, ), dtype=torch.int32, device="cuda", generator=gen)
else:
vbs = torch.full((num_blocks, ), BLOCK, dtype=torch.int32, device="cuda")
for b in empty_blocks:
vbs[b] = 0
valid = (torch.arange(BLOCK, device="cuda")[None] < vbs[:, None]).reshape(-1)
for t in (q, k, v):
t[:, :, ~valid] = 0
return q, k, v, idx.view(batch, heads, num_blocks, num_blocks), num.view(batch, heads, num_blocks), vbs
def reference(q, k, v, idx, num, vbs):
"""FP32 attention restricted to the selected blocks' valid keys; rows without any key are zero."""
B, H, S, D = q.shape
keep = torch.zeros((B, H, S, S), dtype=torch.bool, device=q.device)
sizes = vbs.tolist()
for b in range(B):
for h in range(H):
for i, rows in enumerate(sizes):
for j in idx[b, h, i, :int(num[b, h, i])].tolist():
keep[b, h, i * BLOCK:i * BLOCK + rows, j * BLOCK:j * BLOCK + sizes[j]] = True
scores = (q.float() @ k.float().transpose(-1, -2)) / D**0.5
has = keep.any(-1, keepdim=True)
p = torch.softmax(scores.masked_fill(~keep, float("-inf")).masked_fill(~has, 0.0), dim=-1) * has
return p @ v.float(), has.squeeze(-1)
def run_with_nan_pool(q, k, v, idx, num, vbs):
"""Call the kernel after leaving a NaN block of the output's size in the caching allocator."""
sentinel = torch.full_like(q, float("nan"))
del sentinel
out, lse = vsa.block_sparse_attn_sm121(q, k, v, idx, num, vbs)
torch.cuda.synchronize()
assert lse is None
return out
def check(case, atol=0.02):
q, k, v, idx, num, vbs = case
assert vsa.is_supported(q, vbs)
out = run_with_nan_pool(q, k, v, idx, num, vbs)
ref, has = reference(q, k, v, idx, num, vbs)
assert out.shape == q.shape and out.dtype == torch.bfloat16
assert torch.isfinite(out).all(), "the kernel left elements unwritten"
assert (out[~has] == 0).all(), "rows without a key (padding rows, empty selections) must be zero"
torch.testing.assert_close(out.float()[has], ref[has], atol=atol, rtol=0.02)
def test_forward_matches_reference():
check(make_case())
def test_forward_matches_reference_ragged():
"""variable_block_sizes is what FastVideo always passes: padded keys masked, padded rows zero."""
check(make_case(ragged=True, seed=1))
@pytest.mark.parametrize("topk", [1, 2, 3, 5, 7, 8])
def test_any_selection_size(topk):
check(make_case(ragged=True, topk=topk, seed=topk))
@pytest.mark.parametrize("num_blocks", [1, 3, 7, 9])
def test_odd_block_counts(num_blocks):
"""Unlike the ptx_blackwell kernel, this one needs no even block count."""
check(make_case(num_blocks=num_blocks, topk=max(1, num_blocks // 2), ragged=True, seed=num_blocks))
def test_zero_count_rows_and_non_uniform_counts():
check(make_case(num_blocks=8, counts=(0, 3, 0, 8, 1, 5), ragged=True, seed=7))
def test_selected_blocks_without_valid_keys():
"""A selected block whose variable_block_sizes is 0 contributes nothing; a row left without keys is zero."""
check(make_case(num_blocks=8, topk=3, ragged=True, empty_blocks=(2, 5), seed=11))
def test_batch_and_heads():
check(make_case(num_blocks=6, topk=3, heads=3, batch=2, ragged=True, seed=13))
def test_many_waves():
"""More CTAs than SMs, long selections: the phased walk starts mid-selection and wraps."""
check(make_case(num_blocks=64, topk=24, heads=2, ragged=True, seed=17))
def test_deterministic():
case = make_case(num_blocks=16, topk=6, ragged=True, seed=19)
first = vsa.block_sparse_attn_sm121(*case)[0]
second = vsa.block_sparse_attn_sm121(*case)[0]
assert torch.equal(first, second)
def test_is_supported_rejects_what_it_cannot_run():
q, k, v, idx, num, vbs = make_case(num_blocks=2, topk=1)
assert vsa.is_supported(q, vbs)
assert not vsa.is_supported(q.half(), vbs)
assert not vsa.is_supported(q[..., :64].contiguous(), vbs)
assert not vsa.is_supported(q.transpose(1, 2), vbs)
assert not vsa.is_supported(q[:, :, :BLOCK], vbs)
assert not vsa.is_supported(q, vbs.float())
assert not vsa.is_supported(q.cpu(), vbs)
fastvideo/attention/backends/video_sparse_attn_h3.py
+44 -8
View File
@@ -54,10 +54,12 @@ logs one warning and falls back.
Tile 128 prefers the same ptx_blackwell forward, which carries a 128-token
block instantiation. It needs no opt-in, pads an odd tile count with the same
zero-valid partner tile, and runs no-grad forwards only. When the extension or
device cannot run it, each 128-token tile splits into two contiguous 64-token
halves on the Triton kernels; if those are not installed either, the call fails
closed.
zero-valid partner tile, and runs no-grad forwards only. On a GB10 (DGX Spark,
sm_121), where that kernel cannot run, the 128-token sm_121a forward
(``fastvideo_kernel.block_sparse_attn_sm121``) takes the same padded call. When
neither the extension nor the device can run one of them, each 128-token tile
splits into two contiguous 64-token halves on the Triton kernels; if those are
not installed either, the call fails closed.
"""
import functools
@@ -85,6 +87,12 @@ try:
except ImportError:
_ptx_blackwell = None
try:
# Optional, the same way: the GB10 (sm_121a) 128-token forward.
from fastvideo_kernel import block_sparse_attn_sm121 as _sm121
except ImportError:
_sm121 = None
import fastvideo.envs as envs
from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
AttentionMetadataBuilder, layer_idx_from_prefix)
@@ -561,6 +569,11 @@ def _is_datacenter_blackwell(device: torch.device) -> bool:
return device.type == "cuda" and torch.cuda.get_device_capability(device)[0] == 10
def _sm121_can_run(query_bhsd: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool:
"""Whether the GB10 (sm_121a) 128-token forward can take a tile-128 call (needs map_to_index too)."""
return _sm121 is not None and map_to_index is not None and _sm121.is_supported(query_bhsd, variable_block_sizes)
@functools.cache
def _fine_attention_from_name(name: str) -> Any:
kind, _, mode = name.partition("-")
@@ -616,6 +629,9 @@ class MiniMaxH3VSAImpl(AttentionImpl):
# 128-token blocks, an even tile count with the partner tile, integer
# tile sizes), which every call meets, so it is evaluated once.
self._tile128_route: dict[tuple[torch.device, torch.dtype, int], str | None] = {}
# Per the same key, when ptx_blackwell cannot run it: whether the GB10
# (sm_121a) 128-token forward can, under the same contract.
self._tile128_sm121: dict[tuple[torch.device, torch.dtype, int], bool] = {}
# Optional replacement of the fine block-sparse kernel only (selection,
# mask and the gated compression branch are unchanged). Called as
# ``fn(query, key, value, mask, metadata) -> [B, S, H, D]`` on the
@@ -917,10 +933,12 @@ class MiniMaxH3VSAImpl(AttentionImpl):
# fwd/bwd pairing is possible, but it is not built here.
grad_mode = torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad)
use_ptx_blackwell = False
use_sm121 = False
split_tile128 = False
if tile_elems == 128:
# The CUDA kernel's 128-token instantiation needs no opt-in; without it,
# the Triton tile-64 kernels run each tile as two halves.
# The CUDA kernel's 128-token instantiation needs no opt-in; on a GB10
# its sm_121a counterpart runs the same call; without either, the
# Triton tile-64 kernels run each tile as two halves.
if grad_mode:
raise NotImplementedError("VSA-H3 tile 128 supports no-grad inference forwards only.")
route_key = (q_bhsd.device, q_bhsd.dtype, int(q_bhsd.shape[-1]))
@@ -929,6 +947,8 @@ class MiniMaxH3VSAImpl(AttentionImpl):
ptx_blackwell_variable_block_sizes, grad_mode)
if reason is None and map_to_index is None:
reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable"
if reason is not None:
self._tile128_sm121[route_key] = _sm121_can_run(q_bhsd, ptx_blackwell_variable_block_sizes)
if reason is not None and not compiling and _is_datacenter_blackwell(q_bhsd.device):
# These GPUs should have the native kernel; a broken build
# would otherwise only show up as a slow run.
@@ -938,6 +958,8 @@ class MiniMaxH3VSAImpl(AttentionImpl):
reason = self._tile128_route[route_key]
if reason is None:
use_ptx_blackwell = True
elif self._tile128_sm121.get(route_key, False):
use_sm121 = True
elif block_sparse_attn_64_bhsd is None:
raise RuntimeError(f"VSA-H3 tile 128 requires the ptx_blackwell block-sparse kernel or "
f"the Triton tile-64 kernels: {reason}.")
@@ -1003,6 +1025,20 @@ class MiniMaxH3VSAImpl(AttentionImpl):
ptx_blackwell_variable_block_sizes.to(torch.int32),
need_lse=False,
)
elif use_sm121:
# The GB10 counterpart of the call above: the same padded tiles,
# mask and per-row counts, output only.
if not compiling:
logger.info_once("MiniMax-H3 VSA tile-128 forward: using the GB10 (sm_121a) block-sparse kernel")
q2k_idx, q2k_num = map_to_index(ptx_blackwell_mask)
out_bhsd, _ = _sm121.block_sparse_attn_sm121(
q_bhsd,
k_bhsd,
v_bhsd,
q2k_idx,
q2k_num,
ptx_blackwell_variable_block_sizes.to(torch.int32),
)
else:
if has_ptx_blackwell_pair:
q_bhsd = q_bhsd[:, :, :logical_seq_len]
@@ -1011,7 +1047,7 @@ class MiniMaxH3VSAImpl(AttentionImpl):
if not sm89_strided:
q_bhsd, k_bhsd, v_bhsd = (t.contiguous() for t in (q_bhsd, k_bhsd, v_bhsd))
if split_tile128:
# No sm_100a kernel on this device: run each 128-token tile as two
# No sm_100a or sm_121a kernel on this device: run each 128-token tile as two
# contiguous 64-token halves on the Triton kernels. Every half inherits
# its tile's mask entries, and a partial tile's valid tokens fill the
# first half before the second, so the selection is unchanged.
@@ -1037,7 +1073,7 @@ class MiniMaxH3VSAImpl(AttentionImpl):
mask,
attn_metadata.variable_block_sizes,
)
if has_ptx_blackwell_pair and use_ptx_blackwell:
if has_ptx_blackwell_pair and (use_ptx_blackwell or use_sm121):
out_bhsd = out_bhsd[:, :, :logical_seq_len]
out = out_bhsd.transpose(1, 2).contiguous()
# Fine attention is complete. Release its layout copies before the
fastvideo/fastvideo_args.py
+3 -2
View File
@@ -299,7 +299,7 @@ class FastVideoArgs:
# VSA parameters. None means unset: __post_init__ fills the checkpoint's
# trained value when the checkpoint fixes one, else 0.0 and 256.
VSA_sparsity: float | None = None # inference/validation sparsity
VSA_tile_size: int | None = None # VSA-H3 tokens per tile (256, 128, 64); 128 = ptx_blackwell only
VSA_tile_size: int | None = None # VSA-H3 tokens per tile (256, 128, 64); 128 = CUDA kernel or split Triton
# V-MoBA parameters
moba_config_path: str | None = None
@@ -908,7 +908,8 @@ class FastVideoArgs:
type=int,
default=FastVideoArgs.VSA_tile_size,
help="VSA-H3 tile size in tokens (256, 128 or 64; default: the checkpoint's trained value, else 256); "
"64 runs the native Triton block-sparse path, 128 requires the ptx_blackwell kernel",
"64 runs the native Triton block-sparse path, 128 the ptx_blackwell (sm_100/sm_103) or GB10 "
"(sm_121) CUDA kernel, else each tile as two halves on the Triton path",
)
# Master port for distributed training/inference
fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py
+68
View File
@@ -424,6 +424,21 @@ class _FakePtxBlackwell:
return q.clone(), None
class _FakeSm121:
"""Stands in for fastvideo_kernel.block_sparse_attn_sm121 (the GB10 128-token forward)."""
def __init__(self, supported=True):
self.supported = supported
self.calls = []
def is_supported(self, q, variable_block_sizes):
return self.supported
def block_sparse_attn_sm121(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes):
self.calls.append(dict(q=q, q2k_num=q2k_num, vbs=variable_block_sizes))
return q.clone(), None
def _fake_map_to_index(block_map):
num = block_map.sum(dim=-1, dtype=torch.int32)
idx = torch.full(block_map.shape, -1, dtype=torch.int32)
@@ -441,6 +456,8 @@ def tile128(monkeypatch, env_overrides):
raise AssertionError("tile 128 must never fall back to the Triton-64 kernel")
monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake)
# The GB10 kernel is absent unless a test installs its fake.
monkeypatch.setattr(vsa_h3, "_sm121", None)
monkeypatch.setattr(vsa_h3, "map_to_index", _fake_map_to_index)
monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", forbid_triton)
# Tile 128 has no opt-in: it routes to the CUDA kernel with the switch off.
@@ -501,6 +518,43 @@ def test_forward_tile128_falls_back_to_triton_when_cuda_kernel_unavailable(tile1
assert out.shape == (1, n_tiles * 128, _HEADS, _DIM)
def test_forward_tile128_uses_sm121_kernel_when_ptx_blackwell_unavailable(tile128, monkeypatch):
"""On a GB10 the sm_121a forward takes the ptx_blackwell call: the padded partner tile and per-row counts."""
meta = _build_r2v(sparsity=0.5, tile_size=128, ref_keep_rate=0.25)
n_tiles = meta.variable_block_sizes.numel()
assert n_tiles % 2 == 1
tile128.supported = False
sm121 = _FakeSm121()
monkeypatch.setattr(vsa_h3, "_sm121", sm121)
_, out = _run_tile128(meta)
assert tile128.calls == [] and len(sm121.calls) == 1
call = sm121.calls[0]
assert call["q"].shape[2] == (n_tiles + 1) * 128
assert call["vbs"].dtype == torch.int32 and call["vbs"].numel() == n_tiles + 1 and call["vbs"][-1] == 0
assert (call["q2k_num"][..., n_tiles] == 0).all()
assert out.shape == (1, n_tiles * 128, _HEADS, _DIM)
def test_forward_tile128_prefers_ptx_blackwell_over_sm121(tile128, monkeypatch):
meta = _build_r2v(sparsity=0.5, tile_size=128)
sm121 = _FakeSm121()
monkeypatch.setattr(vsa_h3, "_sm121", sm121)
_run_tile128(meta)
assert len(tile128.calls) == 1 and sm121.calls == []
def test_forward_tile128_falls_back_to_triton_when_sm121_unsupported(tile128, monkeypatch):
meta = _build_r2v(sparsity=0.5, tile_size=128)
tile128.supported = False
sm121 = _FakeSm121(supported=False)
monkeypatch.setattr(vsa_h3, "_sm121", sm121)
calls = []
monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd",
lambda q, k, v, mask, vbs: (calls.append(q.shape), (torch.zeros_like(q), None))[1])
_run_tile128(meta)
assert len(calls) == 1 and sm121.calls == []
@pytest.mark.parametrize("datacenter", [False, True])
def test_forward_tile128_fallback_warns_only_on_datacenter_blackwell(tile128, monkeypatch, datacenter):
meta = _build_r2v(sparsity=0.5, tile_size=128)
@@ -585,6 +639,20 @@ def test_forward_tile128_real_ptx_blackwell_kernel_matches_oracle(env_overrides,
forwards = getattr(vsa_h3._ptx_blackwell, "_FWD_BY_BLOCK", {}) if vsa_h3._ptx_blackwell is not None else {}
if forwards.get(128) is None or vsa_h3.map_to_index is None:
pytest.skip("requires a fastvideo_kernel build with the 128-token ptx_blackwell forward")
_check_real_tile128_forward_against_oracle(env_overrides, sparsity, ref_keep_rate)
@pytest.mark.parametrize("sparsity,ref_keep_rate", [(0.5, 0.25), (0.9, 0.1)])
def test_forward_tile128_real_sm121_kernel_matches_oracle(env_overrides, sparsity, ref_keep_rate):
"""The real GB10 (sm_121a) 128-token forward on a multi-region layout against an FP32 masked-SDPA oracle."""
if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 1):
pytest.skip("requires a GB10 GPU (DGX Spark, sm_121)")
if vsa_h3._sm121 is None or not vsa_h3._sm121._HAS_VSA_SM121 or vsa_h3.map_to_index is None:
pytest.skip("requires a fastvideo_kernel build with the sm_121a block-sparse forward")
_check_real_tile128_forward_against_oracle(env_overrides, sparsity, ref_keep_rate)
def _check_real_tile128_forward_against_oracle(env_overrides, sparsity, ref_keep_rate):
env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None))
device = torch.device("cuda")
meta = _build_r2v(sparsity=sparsity, tile_size=128, ref_keep_rate=ref_keep_rate, device=device)