[kernel] GB10 VSA tile-128: second review pass (device check, NaN rows, routing, linkage)

- Binding: require compute capability 12.1. Other archs in the same build get
  an empty kernel that returned uninitialized output without an error.
- Kernel: a NaN softmax sum now gives a NaN output row (l != 0 instead of
  l > 0), as on the ptx_blackwell and Triton routes; rows without keys stay 0.
- Kernel header: the kernel, encode_map and launch get internal linkage, so a
  second translation unit can include the header (it failed to link before).
  The launch limits move into the header (MaxBatchHeads, MaxRows, MaxTiles),
  used by both launch() and the binding and mirrored in Python.
- Header comments: describe the QK/PV overlap that PhaseLead=8 gives; drop
  the stale cp.async comment and the unused cp_async_mbar_arrive/prefetch_l2.
- H3 backend: an odd tile count without tile()'s partner tile falls back to
  split instead of reaching ptx_blackwell; tile() no longer resolves the
  tile-128 route (or warns) under an eager fine-attention override.
- Python module docstring states the replacement contract for the kernel.

docs/inference/fasth3_omniref.md

@@ -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
@@ -154,10 +159,11 @@ missing native kernel raises rather than silently selecting another mode.
    query tiles, per-token FP8 Q/K for prefix tiles, bf16 P.V and K smoothing.
    Needs no consumer-Blackwell extension but is slower, and its numerics differ
    from the native kernel (which also quantizes P.V).
-3. **bf16 VSA**: set `FASTVIDEO_H3_VSA_FP4=0` and
+3. **bf16 VSA**: set `FASTVIDEO_H3_VSA_FP4=0`. On the DGX Spark, leave
+   `FASTVIDEO_H3_VSA_FINE_ATTENTION` unset so the GB10 CUDA kernel runs the
+   128-token tiles. On the other consumer CUDA GPUs, including sm_89, add
    `FASTVIDEO_H3_VSA_FINE_ATTENTION=emulator-bf16`. This runs unquantized
-   attention on the same tiles with Triton on consumer CUDA GPUs, including
-   sm_89.
+   attention on the same tiles with Triton.
 
 A nonempty fine-attention override takes priority over
 `FASTVIDEO_H3_VSA_FP4=1`. These overrides use the eager path; disable inference

fastvideo-kernel/CMakeLists.txt

@@ -486,6 +486,26 @@ 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). Every spelling of
+    # the GB10 counts: "12.1" without the suffix (the GB10 Docker image builds that), "12.1f" and "12.1+PTX",
+    # since this one source carries its own sm_121a image.
+    set(ENABLE_VSA_SM121A OFF)
+    if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(12\\.1[af]?|121[af]?|sm_121[af]?)(\\+PTX)?([; ,]|$)")
+        if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
+            set(ENABLE_VSA_SM121A ON)
+        else()
+            # A warning: the GB10 then runs tile-128 VSA on the slower split-Triton route.
+            message(WARNING "fastvideo-kernel: skipping block_sparse_sm121 (sm_121a needs CUDA >= 13.0, have "
+                            "${CMAKE_CUDA_COMPILER_VERSION}); tile-128 VSA on a GB10 falls back to Triton")
+        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 +533,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 +663,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

@@ -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 `TORCH_CUDA_ARCH_LIST` names the GB10 (`12.1a`, `12.1`, `12.1f`, with or without `+PTX`) 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

@@ -0,0 +1,88 @@
+// 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 <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);
+  // Other archs in the same build get an empty kernel (it is compiled for sm_121a only), which would return
+  // empty_like garbage without an error.
+  const cudaDeviceProp* prop = at::cuda::getDeviceProperties(device.index());
+  TORCH_CHECK(prop->major == 12 && prop->minor == 1, "block_sparse_sm121_blk128_fwd needs a GB10 (sm_121), got sm_",
+              prop->major, prop->minor);
+  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());
+  TORCH_CHECK(batch * heads <= vsa_sm121::MaxBatchHeads, "batch * heads must be at most ", vsa_sm121::MaxBatchHeads,
+              ", got ", batch * heads);
+  TORCH_CHECK(batch * heads * seq <= vsa_sm121::MaxRows, "batch * heads * seqlen must be at most ",
+              vsa_sm121::MaxRows, ", got ", batch * heads * seq);
+  TORCH_CHECK(blocks <= vsa_sm121::MaxTiles && q2k_idx.size(3) <= vsa_sm121::MaxTiles, "too many blocks: ", blocks,
+              " (max_kv ", q2k_idx.size(3), "); both must be at most ", vsa_sm121::MaxTiles);
+
+  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

@@ -0,0 +1,642 @@
+// 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. A
+//     QK warp computes the next chunk's scores while its PV partner multiplies the previous P by V; the pair
+//     meets (pair_sync) only before P is overwritten, so the two groups' mma bursts overlap.
+//   - 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 <atomic>
+#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; a tile's padding keys may hold anything, K and V both
+	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
+};
+
+// The largest call one launch covers. The binding checks them before launching, and
+// fastvideo_kernel/block_sparse_attn_sm121.py's _MAX_* mirror them for the callers that route around a launch.
+constexpr int64_t MaxBatchHeads = 65535;     // batch_heads: the grid's y dimension
+constexpr int64_t MaxRows = INT32_MAX;       // batch_heads * tiles * 128: TMA addresses K and V by int32 row
+constexpr int64_t MaxTiles = (1 << 22) - 1;  // tiles and max_kv: the walker packs its start entry into 22 bits
+
+
+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).
+// Only the Q load uses cp.async; K and V stream through TMA.
+__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");
+}
+
+// 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)) {
+	}
+}
+
+// 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;
+// QK k-steps (of 8) a warp issues before waiting for its PV partner to release P; 8 waits after the whole burst.
+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");
+    }
+}
+// A B fragment register's bf16 pair of keys (the lower key in the lower half), given how many keys from the
+// pair's first one on are valid: both kept, the lower one only, or neither.
+__device__ __forceinline__ uint32_t valid_keys(uint32_t b, int valid) {
+    return valid>=2 ? b : (valid==1 ? (b&0x0000ffffu) : 0u);
+}
+// Masked: the chunk holds fewer than 64 valid keys. P is zero past them, but the V rows there are whatever the
+// caller's padding holds, and 0 * NaN or 0 * Inf would still reach O: skip the 16-key steps that lie wholly past
+// the valid keys, and zero the invalid rows of the one that straddles them.
+template<bool Masked>
+__device__ __forceinline__ void pv(float (&o)[16][4], const uint32_t (&p)[4][4], const uint32_t (&v_addr)[4],
+                                   int keys, int lane) {
+#pragma unroll
+    for (int kk=0; kk<4; ++kk) {
+        if constexpr (Masked) {
+            if (16*kk>=keys) break;
+        }
+        // b[0], b[2] hold keys 16kk + c + {0, 1}; b[1], b[3] those 8 further on (c = 2 * (lane % 4)). `valid`
+        // counts the valid keys from 16kk + c on.
+        const int valid=keys-16*kk-2*(lane&3);
+#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));
+            if constexpr (Masked) {
+                b[0]=valid_keys(b[0],valid); b[1]=valid_keys(b[1],valid-8);
+                b[2]=valid_keys(b[2],valid); b[3]=valid_keys(b[3],valid-8);
+            }
+            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; // NaN stays NaN, as on the other routes
+        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.
+static __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 {
+        // Each V slot's valid key count, eight bits per slot.
+        int pending=0, vkeys=0;
+#pragma unroll
+        for (int s=0;s<3;++s) {
+            if (fill.next()) {
+                ++pending;
+                vkeys|=fill.keys()<<(8*s);
+                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];
+                const int keys=(vkeys>>(8*stage))&255;
+                if (keys<Chunk) pv<true>(o,pf,vaddr,keys,lane);
+                else pv<false>(o,pf,vaddr,keys,lane);
+            }
+            // 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()) {
+                vkeys=(vkeys&~(255<<(8*stage)))|(fill.keys()<<(8*stage));
+                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
+}
+
+static 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
+
+// The kernel's whole host interface. Internal linkage, like the kernel and its helpers: a translation unit that
+// includes this header gets its own copy, so a second includer (a benchmark, a newer kernel built beside this
+// one) links without duplicate symbols.
+static 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 (args.batch_heads>MaxBatchHeads || rows>MaxRows || args.tiles>MaxTiles || args.max_kv>MaxTiles)
+        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;
+    // One CTA per SM fits (99 KB of shared memory each): the SMs are the
+    // CTAs that run at once. Read once per device, after the per-device
+    // shared-memory attribute is set; 0 means not yet. Threads that race on a
+    // device's first launch each set the same attribute and store the same
+    // count, and a thread never sees another device's count.
+    constexpr int MaxDevices=64;
+    static std::atomic<int> sm_count[MaxDevices];
+    int device=0;
+    cudaError_t err=cudaGetDevice(&device);
+    if (err!=cudaSuccess) return err;
+    if (device<0 || device>=MaxDevices) return cudaErrorInvalidDevice;
+    int slots=sm_count[device].load(std::memory_order_acquire);
+    if (slots==0) {
+        err=cudaFuncSetAttribute(split128::forward,cudaFuncAttributeMaxDynamicSharedMemorySize,split128::SharedBytes);
+        if (err!=cudaSuccess) return err;
+        int sms=0;
+        err=cudaDeviceGetAttribute(&sms,cudaDevAttrMultiProcessorCount,device);
+        if (err!=cudaSuccess) return err;
+        slots=sms>0 ? sms : 1;
+        sm_count[device].store(slots,std::memory_order_release);
+    }
+    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

@@ -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

@@ -0,0 +1,136 @@
+# 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.
+
+Replacing the kernel: callers (the VSA-H3 backend) use only ``is_supported``,
+``within_launch_limits`` and ``block_sparse_attn_sm121``. A new kernel keeps them, together with the
+``block_sparse_sm121_blk128_fwd`` binding, by providing ``vsa_sm121::Vsa128Args``, the ``Max*``
+limits and ``vsa_sm121::launch`` in the header; ``fastvideo-kernel/tests/test_block_sparse_sm121.py``
+is its acceptance test.
+"""
+
+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
+# Launch limits, mirroring vsa_sm121::MaxBatchHeads, MaxRows and MaxTiles in the kernel header.
+_MAX_BATCH_HEADS = 65535
+_MAX_ROWS = 2**31 - 1
+_MAX_BLOCKS = 2**22 - 1
+
+
+def within_launch_limits(q: torch.Tensor, num_blocks: int) -> bool:
+    """True iff one launch can cover q's [B, H, S, 128] shape with num_blocks blocks.
+
+    The part of ``is_supported`` that depends on the size of the call, not on the device, dtype or
+    layout, so callers that cache ``is_supported`` per device must still check it per call.
+    ``map_to_index`` makes q2k_idx's last dimension num_blocks, which this bounds too.
+    """
+    batch_heads = q.shape[0] * q.shape[1]
+    return batch_heads <= _MAX_BATCH_HEADS and batch_heads * q.shape[2] <= _MAX_ROWS and num_blocks <= _MAX_BLOCKS
+
+
+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;
+      * the call is within the launch limits (``within_launch_limits``): batch * heads <= 65535,
+        batch * heads * seqlen < 2^31 and num_blocks < 2^22.
+    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); padding rows of q, k and v that hold anything, NaN and Inf included.
+    """
+    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 within_launch_limits(q, num_blocks)
+
+
+@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

@@ -0,0 +1,190 @@
+# 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)
+
+
+@pytest.mark.parametrize("fill", [float("nan"), float("inf")])
+def test_padding_keys_holding_nan_or_inf_do_not_leak(fill):
+    """Padding rows of K and V may hold anything: P is zero there, and 0 * NaN must not reach the output."""
+    q, k, v, idx, num, vbs = make_case(num_blocks=6, topk=3, ragged=True, seed=23)
+    assert (vbs % 64 != 0).any(), "the case must have partial 64-key chunks"
+    expected, has = reference(q, k, v, idx, num, vbs)
+    valid = (torch.arange(BLOCK, device="cuda")[None] < vbs[:, None]).reshape(-1)
+    for t in (q, k, v):
+        t[:, :, ~valid] = fill
+    out = run_with_nan_pool(q, k, v, idx, num, vbs)
+    assert torch.isfinite(out).all()
+    assert (out[~has] == 0).all()
+    torch.testing.assert_close(out.float()[has], expected[has], atol=0.02, rtol=0.02)
+
+
+def test_nan_in_a_valid_query_row_reaches_its_output_row():
+    """Like the other routes, a NaN input row gives a NaN output row instead of a silent zero."""
+    q, k, v, idx, num, vbs = make_case(num_blocks=4, topk=2, seed=29)
+    expected, has = reference(q, k, v, idx, num, vbs)
+    q[0, 0, 5] = float("nan")
+    out = run_with_nan_pool(q, k, v, idx, num, vbs)
+    assert torch.isnan(out[0, 0, 5]).all()
+    rest = has.clone()
+    rest[0, 0, 5] = False
+    assert torch.isfinite(out[rest]).all() and (out[~has] == 0).all()
+    torch.testing.assert_close(out.float()[rest], expected[rest], atol=0.02, rtol=0.02)
+
+
+def test_launch_limits(monkeypatch):
+    """Sizes one launch cannot cover: is_supported says so instead of the binding raising."""
+    def meta(*shape):
+        return torch.empty(shape, device="meta", dtype=torch.bfloat16)
+
+    assert vsa.within_launch_limits(meta(1, 65535, 2 * BLOCK, HEAD_DIM), 2)
+    assert not vsa.within_launch_limits(meta(2, 32768, 2 * BLOCK, HEAD_DIM), 2)
+    assert not vsa.within_launch_limits(meta(1, 2, 2**30, HEAD_DIM), 2**23)
+    assert not vsa.within_launch_limits(meta(1, 1, 2**22 * BLOCK, HEAD_DIM), 2**22)
+    assert vsa.within_launch_limits(meta(1, 1, (2**22 - 1) * BLOCK, HEAD_DIM), 2**22 - 1)
+    q, _, _, _, _, vbs = make_case(num_blocks=2, topk=1, heads=4)
+    assert vsa.is_supported(q, vbs)
+    monkeypatch.setattr(vsa, "_MAX_BATCH_HEADS", 3)
+    assert not vsa.is_supported(q, vbs)
+
+
+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

@@ -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 call, which needs no
+partner tile. 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,23 @@ def _is_datacenter_blackwell(device: torch.device) -> bool:
     return device.type == "cuda" and torch.cuda.get_device_capability(device)[0] == 10
 
 
+def _is_gb10(device: torch.device) -> bool:
+    """sm_121 (DGX Spark), the device the sm_121a tile-128 kernel targets."""
+    return device.type == "cuda" and torch.cuda.get_device_capability(device) == (12, 1)
+
+
+def _sm121_unavailable_reason(query_bhsd: torch.Tensor, variable_block_sizes: torch.Tensor) -> str | None:
+    """Why the GB10 (sm_121a) 128-token forward cannot run here, or None if it can."""
+    if _sm121 is None:
+        return "fastvideo_kernel.block_sparse_attn_sm121 is not installed"
+    if map_to_index is None:
+        return "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable"
+    if not _sm121.is_supported(query_bhsd, variable_block_sizes):
+        return ("block_sparse_attn_sm121.is_supported returned False (needs a GB10 and a fastvideo-kernel built "
+                "with TORCH_CUDA_ARCH_LIST containing 12.1a and CUDA >= 13, bf16, head_dim 128)")
+    return None
+
+
 @functools.cache
 def _fine_attention_from_name(name: str) -> Any:
     kind, _, mode = name.partition("-")
@@ -610,12 +635,15 @@ class MiniMaxH3VSAImpl(AttentionImpl):
         # request-time env/probe/fallback behavior; only Dynamo capture reads
         # the prepared, static route.
         self._regional_compile_ptx_blackwell_enabled: bool | None = None
-        # Tile-128 route per (device, dtype, head size): None when the CUDA
-        # kernel can run it, else why not. The kernel's predicate otherwise
-        # depends only on the tile-128 buffer contract (contiguous BHSD,
-        # 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] = {}
+        # Tile-128 route per (device, dtype, head size): the kernel that runs
+        # it ("ptx_blackwell", "sm121" or "split", see _tile128_route) and why
+        # the faster ones cannot (None for ptx_blackwell). The kernels'
+        # predicates otherwise depend only on the tile-128 buffer contract
+        # (contiguous BHSD, 128-token blocks, the partner tile where
+        # ptx_blackwell needs one, integer tile sizes), which every call meets,
+        # so they are evaluated once; only the sm121 kernel's launch-size
+        # limits are checked per call.
+        self._tile128_routes: dict[tuple[torch.device, torch.dtype, int], tuple[str, str | None]] = {}
         # 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
@@ -679,6 +707,45 @@ class MiniMaxH3VSAImpl(AttentionImpl):
             logger.warning_once(f"VSA-H3 regional compile is unavailable and will stay eager: {reason}")
         return reason
 
+    def _tile128_route(self, device: torch.device, dtype: torch.dtype, head_size: int,
+                       compiling: bool) -> tuple[str, str | None]:
+        """The no-grad tile-128 route on this device: ``(kernel, reason)``.
+
+        ``kernel`` is "ptx_blackwell" (sm_100/sm_103, an even tile count),
+        "sm121" (GB10, any tile count) or "split" (each tile as two halves on
+        the Triton tile-64 kernels); ``reason`` says why the faster kernels
+        cannot run, None for ptx_blackwell. Resolved once per key on a
+        two-tile probe, before ``tile()`` decides whether to add the partner
+        tile only ptx_blackwell needs.
+        """
+        route_key = (device, dtype, head_size)
+        route = self._tile128_routes.get(route_key)
+        if route is not None:
+            return route
+        probe_query = torch.empty((1, 1, 2 * 128, head_size), device=device, dtype=dtype)
+        probe_block_sizes = torch.full((2, ), 128, device=device, dtype=torch.int32)
+        reason = _ptx_blackwell_unavailable_reason(_ptx_blackwell, probe_query, probe_block_sizes, grad_mode=False)
+        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 None:
+            route = ("ptx_blackwell", None)
+        else:
+            sm121_reason = _sm121_unavailable_reason(probe_query, probe_block_sizes)
+            if sm121_reason is None:
+                route = ("sm121", reason)
+            else:
+                route = ("split", f"ptx_blackwell: {reason}; sm_121a: {sm121_reason}")
+                # These GPUs should have a native kernel; a broken build would
+                # otherwise only show up as a slow run.
+                if not compiling and _is_datacenter_blackwell(device):
+                    logger.warning_once(f"VSA-H3 tile 128 on a data-center Blackwell GPU is falling back to "
+                                        f"the slower Triton tile-64 kernels: {reason}.")
+                elif not compiling and _is_gb10(device):
+                    logger.warning_once(f"VSA-H3 tile 128 on a GB10 is falling back to the slower Triton tile-64 "
+                                        f"kernels: {sm121_reason}.")
+        self._tile128_routes[route_key] = route
+        return route
+
     def tile(self, x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor:
         """Scatter rows into the padded tile buffer (pad positions stay zero).
 
@@ -687,9 +754,9 @@ class MiniMaxH3VSAImpl(AttentionImpl):
         sites in ``forward()`` read it immediately). A grad-tracking forward
         instead receives a fresh buffer and leaves the holder untouched, so
         the builder never retains autograd state across steps. Odd no-grad
-        ptx_blackwell requests (tile 128 always, tile 64 when opted in) carry one
-        additional all-zero tile internally; metadata and all observable
-        outputs retain the logical geometry.
+        ptx_blackwell requests (tile 128 when it routes there, tile 64 when
+        opted in) carry one additional all-zero tile internally; metadata and
+        all observable outputs retain the logical geometry.
         """
         compiling = torch.compiler.is_compiling()
         regional_compiling = compiling and self._regional_compile_ptx_blackwell_enabled is True
@@ -706,9 +773,15 @@ class MiniMaxH3VSAImpl(AttentionImpl):
             ptx_blackwell_requested = False
         else:
             ptx_blackwell_requested = envs.FASTVIDEO_VSA_PTX_BLACKWELL.get()
-        # Tile 128 has no route other than the ptx_blackwell CUDA kernel.
-        ptx_blackwell_route = attn_metadata.tile_elems == 128 or (attn_metadata.tile_elems == 64
-                                                                  and ptx_blackwell_requested)
+        # Tile 128 picks its kernel by device; only ptx_blackwell needs the partner tile.
+        # An eager fine-attention override reads the logical rows and runs no kernel.
+        if attn_metadata.tile_elems == 128 and self.fine_attention_override is not None and not compiling:
+            ptx_blackwell_route = False
+        elif attn_metadata.tile_elems == 128:
+            ptx_blackwell_route = (not grad_mode and not regional_compiling and self._tile128_route(
+                x.device, x.dtype, int(x.shape[-1]), compiling)[0] == "ptx_blackwell")
+        else:
+            ptx_blackwell_route = attn_metadata.tile_elems == 64 and ptx_blackwell_requested
         if regional_compiling:
             # Pad arithmetically so odd and even tile counts share one graph.
             needs_ptx_blackwell_pair, kernel_tiles = False, n_tiles + n_tiles % 2
@@ -917,30 +990,32 @@ 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 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]))
-                if route_key not in self._tile128_route:
-                    reason = _ptx_blackwell_unavailable_reason(_ptx_blackwell, q_bhsd,
-                                                               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 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.
-                        logger.warning_once(f"VSA-H3 tile 128 on a data-center Blackwell GPU is falling back to "
-                                            f"the slower Triton tile-64 kernels: {reason}.")
-                    self._tile128_route[route_key] = reason
-                reason = self._tile128_route[route_key]
-                if reason is None:
+                kernel, reason = self._tile128_route(q_bhsd.device, q_bhsd.dtype, int(q_bhsd.shape[-1]), compiling)
+                if kernel == "sm121" and not _sm121.within_launch_limits(q_bhsd, n_tiles + pair_tiles):
+                    # Too large for one launch (batch * heads, rows or tiles); this call only.
+                    if not compiling:
+                        logger.warning_once("VSA-H3 tile 128: a call exceeds the GB10 kernel's launch limits; "
+                                            "it runs on the Triton tile-64 kernels")
+                    kernel, reason = "split", f"{reason}; sm_121a: the call exceeds its launch limits"
+                if kernel == "ptx_blackwell" and (n_tiles + pair_tiles) % 2 != 0:
+                    # Only tile() adds the partner tile; a call that skipped it cannot run there.
+                    kernel, reason = "split", "ptx_blackwell: an odd tile count without the partner tile"
+                if kernel == "ptx_blackwell":
                     use_ptx_blackwell = True
+                elif kernel == "sm121":
+                    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}.")
+                    raise RuntimeError(f"VSA-H3 tile 128 requires the ptx_blackwell block-sparse kernel "
+                                       f"(sm_100/sm_103), the GB10 block-sparse kernel (fastvideo-kernel built for "
+                                       f"sm_121a) or the Triton tile-64 kernels: {reason}.")
                 else:
                     split_tile128 = True
             elif regional_compiling:
@@ -1003,6 +1078,21 @@ 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 mask and
+                # per-row counts, output only. It takes any tile count, so
+                # tile() adds no partner tile on this route.
+                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 +1101,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 +1127,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

@@ -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_heads_first_tile.py

@@ -38,6 +38,9 @@ def _metadata(seq: int, n_tiles: int, tile_elems: int) -> SimpleNamespace:
 def _impl() -> vsa.MiniMaxH3VSAImpl:
     impl = object.__new__(vsa.MiniMaxH3VSAImpl)
     impl._regional_compile_ptx_blackwell_enabled = None
+    # tile() resolves the tile-128 route to decide on the ptx_blackwell partner tile.
+    impl._tile128_routes = {}
+    impl.fine_attention_override = None
     return impl
 
 

fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py

@@ -424,6 +424,25 @@ 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, within_limits=True):
+        self.supported = supported
+        self.within_limits = within_limits
+        self.calls = []
+
+    def is_supported(self, q, variable_block_sizes):
+        return self.supported
+
+    def within_launch_limits(self, q, num_blocks):
+        return self.within_limits
+
+    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 +460,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.
@@ -482,6 +503,7 @@ def test_forward_tile128_even_tile_count_no_partner_tile(tile128):
 def test_forward_tile128_falls_back_to_triton_when_cuda_kernel_unavailable(tile128, monkeypatch):
     meta = _build_r2v(sparsity=0.5, tile_size=128)
     n_tiles = meta.variable_block_sizes.numel()
+    assert n_tiles % 2 == 1
     tile128.supported = False
     calls = []
 
@@ -490,7 +512,9 @@ def test_forward_tile128_falls_back_to_triton_when_cuda_kernel_unavailable(tile1
         return torch.zeros_like(q), None
 
     monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", triton)
-    _, out = _run_tile128(meta)
+    tiled, out = _run_tile128(meta)
+    # The split route needs no partner tile either.
+    assert tiled.shape[1] == n_tiles * 128
     assert len(calls) == 1
     q, mask, vbs = calls[0]
     assert q.shape[2] == n_tiles * 128
@@ -501,26 +525,119 @@ def test_forward_tile128_falls_back_to_triton_when_cuda_kernel_unavailable(tile1
     assert out.shape == (1, n_tiles * 128, _HEADS, _DIM)
 
 
-@pytest.mark.parametrize("datacenter", [False, True])
-def test_forward_tile128_fallback_warns_only_on_datacenter_blackwell(tile128, monkeypatch, datacenter):
+def test_forward_tile128_uses_sm121_kernel_when_ptx_blackwell_unavailable(tile128, monkeypatch):
+    """On a GB10 the sm_121a forward takes the call with the logical tile count: no partner tile, even when odd."""
+    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)
+    tiled, out = _run_tile128(meta)
+    assert tiled.shape[1] == n_tiles * 128
+    assert tile128.calls == [] and len(sm121.calls) == 1
+    call = sm121.calls[0]
+    assert call["q"].shape[2] == n_tiles * 128
+    assert call["vbs"].dtype == torch.int32 and call["vbs"].numel() == n_tiles
+    assert call["q2k_num"].shape[-1] == n_tiles
+    assert out.shape == (1, n_tiles * 128, _HEADS, _DIM)
+
+
+def test_forward_tile128_sm121_call_past_launch_limits_runs_split(tile128, monkeypatch):
+    """A call too large for one GB10 launch takes the Triton split route instead of failing in the binding."""
+    meta = _build_r2v(sparsity=0.5, tile_size=128, ref_keep_rate=0.25)
+    n_tiles = meta.variable_block_sizes.numel()
+    tile128.supported = False
+    sm121 = _FakeSm121(within_limits=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])
+    _, out = _run_tile128(meta)
+    assert sm121.calls == [] and len(calls) == 1 and calls[0][2] == n_tiles * 128
+    assert out.shape == (1, n_tiles * 128, _HEADS, _DIM)
+
+
+def test_forward_tile128_odd_call_without_partner_tile_runs_split(tile128, monkeypatch):
+    """ptx_blackwell needs an even tile count; a call that skipped tile()'s partner tile falls back to Triton."""
+    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
+    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])
+    impl = _impl(_DIM)
+    raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16)
+    query, key, value = impl.preprocess_qkv(raw, meta)[:, :n_tiles * 128].chunk(3, dim=0)
+    out = impl.forward(query, key, value, None, meta)
+    assert tile128.calls == [] and len(calls) == 1 and calls[0][2] == n_tiles * 128
+    assert out.shape == (1, n_tiles * 128, _HEADS, _DIM)
+
+
+def test_tile128_fine_attention_override_skips_kernel_routing(tile128):
+    """An eager override runs no kernel: tile() neither probes the route nor adds the partner tile."""
+    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
+    probes, calls = [], []
+    tile128.is_supported = lambda q, variable_block_sizes: probes.append(q.device) or True
+    impl = _impl(_DIM)
+    impl.fine_attention_override = lambda q, k, v, mask, metadata: calls.append(q.shape) or torch.zeros_like(q)
+    raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16)
+    tiled = impl.preprocess_qkv(raw, meta)
+    assert tiled.shape[1] == n_tiles * 128 and probes == []
+    query, key, value = tiled.chunk(3, dim=0)
+    out = impl.forward(query, key, value, None, meta)
+    assert len(calls) == 1 and tile128.calls == [] and probes == []
+    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("device_kind", ["other", "datacenter", "gb10"])
+def test_forward_tile128_fallback_warns_only_where_a_native_kernel_belongs(tile128, monkeypatch, device_kind):
     meta = _build_r2v(sparsity=0.5, tile_size=128)
     tile128.supported = False
     warnings = []
     monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", lambda q, k, v, mask, vbs: (torch.zeros_like(q), None))
-    monkeypatch.setattr(vsa_h3, "_is_datacenter_blackwell", lambda device: datacenter)
+    monkeypatch.setattr(vsa_h3, "_is_datacenter_blackwell", lambda device: device_kind == "datacenter")
+    monkeypatch.setattr(vsa_h3, "_is_gb10", lambda device: device_kind == "gb10")
     monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append)
     _run_tile128(meta)
-    assert len(warnings) == int(datacenter)
-    if datacenter:
-        assert "falling back to the slower Triton tile-64 kernels" in warnings[0]
+    assert len(warnings) == int(device_kind != "other")
+    if device_kind == "datacenter":
+        assert "data-center Blackwell GPU is falling back to the slower Triton tile-64 kernels" in warnings[0]
+    elif device_kind == "gb10":
+        assert "on a GB10 is falling back to the slower Triton tile-64 kernels" in warnings[0]
+        assert "block_sparse_attn_sm121 is not installed" in warnings[0]
 
 
 def test_forward_tile128_unavailable_cuda_kernel_and_triton(tile128, monkeypatch):
     meta = _build_r2v(sparsity=0.5, tile_size=128)
     tile128.supported = False
     monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", None)
-    with pytest.raises(RuntimeError, match="tile 128 requires the ptx_blackwell block-sparse kernel"):
+    with pytest.raises(RuntimeError, match="tile 128 requires the ptx_blackwell block-sparse kernel") as error:
         _run_tile128(meta)
+    assert "GB10 block-sparse kernel (fastvideo-kernel built for sm_121a)" in str(error.value)
+    assert "sm_121a: fastvideo_kernel.block_sparse_attn_sm121 is not installed" in str(error.value)
     monkeypatch.setattr(vsa_h3, "_ptx_blackwell", None)
     with pytest.raises(RuntimeError, match="not installed"):
         _run_tile128(meta)
@@ -578,14 +695,44 @@ def test_build_tile128_uses_4x4x8_tiles():
 
 
 @pytest.mark.parametrize("sparsity,ref_keep_rate", [(0.5, 0.25), (0.9, 0.1)])
-def test_forward_tile128_real_ptx_blackwell_kernel_matches_oracle(env_overrides, sparsity, ref_keep_rate):
+def test_forward_tile128_real_ptx_blackwell_kernel_matches_oracle(env_overrides, monkeypatch, sparsity, ref_keep_rate):
     """The real 128-token CUDA forward on a multi-region layout against an FP32 masked-SDPA oracle."""
     if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}:
         pytest.skip("requires an sm_100a/sm_103a GPU (B200, B300, GB200, GB300)")
     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, monkeypatch, sparsity, ref_keep_rate,
+                                               kernel_module=vsa_h3._ptx_blackwell,
+                                               entry="block_sparse_attn_ptx_blackwell",
+                                               route="ptx_blackwell")
+
+
+@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, monkeypatch, 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, monkeypatch, sparsity, ref_keep_rate,
+                                               kernel_module=vsa_h3._sm121,
+                                               entry="block_sparse_attn_sm121",
+                                               route="sm121")
+
+
+def _check_real_tile128_forward_against_oracle(env_overrides, monkeypatch, sparsity, ref_keep_rate, *, kernel_module,
+                                               entry, route):
+    """Also proves the route: the Triton split route matches the oracle too, so a silent fallback must fail."""
     env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None))
+    launches = []
+    real_entry = getattr(kernel_module, entry)
+
+    def spy(*args, **kwargs):
+        launches.append(args[0].shape)
+        return real_entry(*args, **kwargs)
+
+    monkeypatch.setattr(kernel_module, entry, spy)
     device = torch.device("cuda")
     meta = _build_r2v(sparsity=sparsity, tile_size=128, ref_keep_rate=ref_keep_rate, device=device)
     n_tiles = meta.variable_block_sizes.numel()
@@ -595,9 +742,12 @@ def test_forward_tile128_real_ptx_blackwell_kernel_matches_oracle(env_overrides,
     raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, device=device, dtype=torch.bfloat16)
     with torch.inference_mode():
         tiled = impl.preprocess_qkv(raw, meta)
-        assert tiled.shape[1] == (n_tiles + n_tiles % 2) * 128
+        assert impl._tile128_routes[(raw.device, torch.bfloat16, _DIM)][0] == route
+        # Only ptx_blackwell needs an even tile count.
+        assert tiled.shape[1] == (n_tiles + (n_tiles % 2 if route == "ptx_blackwell" else 0)) * 128
         query, key, value = tiled.chunk(3, dim=0)
         actual = impl.postprocess_output(impl.forward(query, key, value, None, meta), meta)
+        assert len(launches) == 1, f"the {route} kernel must run exactly once, ran {len(launches)} times"
         # The mask the backend selected: pooled scores of the logical tiles.
         logical_query, logical_key, logical_value = (t[:, :logical] for t in (query, key, value))
         scores = torch.matmul(_pool_tiles(logical_query, meta.variable_block_sizes, 128),