[kernel] Add varlen support for block-sparse attention

fastvideo-kernel/python/fastvideo_kernel/__init__.py

@@ -24,11 +24,16 @@
     int8_quant,
 )
 
+from fastvideo_kernel.block_sparse_attn_varlen import (
+    block_sparse_attn_varlen,
+)
+
 __all__ = [
     "sliding_tile_attention",
     "video_sparse_attn",
     "block_sparse_attn",
     "block_sparse_attn_from_indices",
+    "block_sparse_attn_varlen",
     "moba_attn_varlen",
     "process_moba_input",
     "process_moba_output",

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_varlen.py

@@ -0,0 +1,207 @@
+"""Variable-length block-sparse attention via sequence packing.
+
+Packs multiple variable-length sequences into a single [1, H, T_total, D]
+tensor and delegates to the existing block_sparse_attn_from_indices kernel
+in a single launch. No kernel modifications required.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import torch
+
+from .block_sparse_attn import block_sparse_attn_from_indices
+
+BLOCK_SIZE = 64
+
+
+def _scatter_to_padded(
+    src: torch.Tensor,
+    block_sizes: torch.Tensor,
+    block_size: int,
+    dst: torch.Tensor,
+    dst_offset: int,
+    src_start: int,
+    src_end: int,
+) -> None:
+    """Copy tokens from a flat source into block-aligned positions in dst.
+
+    Each block occupies exactly `block_size` slots in dst.  The first
+    `block_sizes[b]` slots of block *b* receive real tokens; the remainder
+    stays zero (padding the kernel expects).
+
+    src:         [total_tokens, H, D]
+    dst:         [1, H, total_padded, D]
+    block_sizes: [num_blocks] int32, actual token count per block.
+    """
+    src_pos = src_start
+    dst_pos = dst_offset
+    for b in range(block_sizes.numel()):
+        actual = int(block_sizes[b].item())
+        actual = min(actual, src_end - src_pos)
+        if actual > 0:
+            dst[:, :, dst_pos:dst_pos + actual, :] = (
+                src[src_pos:src_pos + actual].transpose(0, 1).unsqueeze(0)
+            )
+        src_pos += actual
+        dst_pos += block_size
+
+
+def _gather_from_padded(
+    src: torch.Tensor,
+    block_sizes: torch.Tensor,
+    block_size: int,
+    dst: torch.Tensor,
+    src_offset: int,
+    dst_start: int,
+    dst_end: int,
+) -> None:
+    """Inverse of _scatter_to_padded: extract real tokens from padded blocks.
+
+    src: [1, H, total_padded, D]
+    dst: [total_tokens, H, D]
+    """
+    src_pos = src_offset
+    dst_pos = dst_start
+    for b in range(block_sizes.numel()):
+        actual = int(block_sizes[b].item())
+        actual = min(actual, dst_end - dst_pos)
+        if actual > 0:
+            dst[dst_pos:dst_pos + actual] = (
+                src[0, :, src_pos:src_pos + actual, :].transpose(0, 1)
+            )
+        dst_pos += actual
+        src_pos += block_size
+
+
+def block_sparse_attn_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_kv: torch.Tensor,
+    q2k_idx_list: Sequence[torch.Tensor],
+    q2k_num_list: Sequence[torch.Tensor],
+    variable_block_sizes_list: Sequence[torch.Tensor],
+    q_variable_block_sizes_list: Sequence[torch.Tensor] | None = None,
+    block_size: int = BLOCK_SIZE,
+) -> torch.Tensor:
+    """Block-sparse attention over packed variable-length sequences.
+
+    Args:
+        q: [total_q_tokens, H, D] packed query tensor.
+        k: [total_kv_tokens, H, D] packed key tensor.
+        v: [total_kv_tokens, H, D] packed value tensor.
+        cu_seqlens_q: [N+1] int32, cumulative Q token offsets.
+        cu_seqlens_kv: [N+1] int32, cumulative KV token offsets.
+        q2k_idx_list: Per-sequence q2k_idx tensors, each [1, H, Nq_i, Mk].
+        q2k_num_list: Per-sequence q2k_num tensors, each [1, H, Nq_i].
+        variable_block_sizes_list: Per-sequence KV block sizes, each [Nkv_i].
+        q_variable_block_sizes_list: Per-sequence Q block sizes, each [Nq_i].
+            If None, each Q block is assumed to be exactly `block_size` tokens.
+        block_size: Attention block size (default 64).
+
+    Returns:
+        out: [total_q_tokens, H, D] packed output tensor.
+    """
+    device = q.device
+    dtype = q.dtype
+    num_heads = q.shape[1]
+    head_dim = q.shape[2]
+    num_seqs = cu_seqlens_q.shape[0] - 1
+
+    cu_q = cu_seqlens_q.cpu().tolist()
+    cu_kv = cu_seqlens_kv.cpu().tolist()
+
+    padded_q_lens = []
+    padded_kv_lens = []
+    q_block_offsets = [0]
+    kv_block_offsets = [0]
+    q_vbs_resolved = []
+
+    for i in range(num_seqs):
+        n_q_blocks = q2k_num_list[i].shape[-1]
+        n_kv_blocks = variable_block_sizes_list[i].numel()
+        padded_q_lens.append(n_q_blocks * block_size)
+        padded_kv_lens.append(n_kv_blocks * block_size)
+        q_block_offsets.append(q_block_offsets[-1] + n_q_blocks)
+        kv_block_offsets.append(kv_block_offsets[-1] + n_kv_blocks)
+
+        if q_variable_block_sizes_list is not None:
+            q_vbs_resolved.append(q_variable_block_sizes_list[i])
+        else:
+            q_vbs_resolved.append(
+                torch.full((n_q_blocks,), block_size, dtype=torch.int32, device=device)
+            )
+
+    total_padded_q = sum(padded_q_lens)
+    total_padded_kv = sum(padded_kv_lens)
+
+    q_packed = torch.zeros(1, num_heads, total_padded_q, head_dim, device=device, dtype=dtype)
+    k_packed = torch.zeros(1, num_heads, total_padded_kv, head_dim, device=device, dtype=dtype)
+    v_packed = torch.zeros(1, num_heads, total_padded_kv, head_dim, device=device, dtype=dtype)
+
+    q_offset = 0
+    kv_offset = 0
+    for i in range(num_seqs):
+        _scatter_to_padded(
+            q, q_vbs_resolved[i], block_size,
+            q_packed, q_offset, cu_q[i], cu_q[i + 1],
+        )
+        _scatter_to_padded(
+            k, variable_block_sizes_list[i], block_size,
+            k_packed, kv_offset, cu_kv[i], cu_kv[i + 1],
+        )
+        _scatter_to_padded(
+            v, variable_block_sizes_list[i], block_size,
+            v_packed, kv_offset, cu_kv[i], cu_kv[i + 1],
+        )
+        q_offset += padded_q_lens[i]
+        kv_offset += padded_kv_lens[i]
+
+    total_q_blocks = q_block_offsets[-1]
+    max_kv_per_q = max(t.shape[-1] for t in q2k_idx_list)
+
+    global_q2k_idx = torch.zeros(
+        1, num_heads, total_q_blocks, max_kv_per_q,
+        dtype=torch.int32, device=device,
+    )
+    global_q2k_num = torch.zeros(
+        1, num_heads, total_q_blocks,
+        dtype=torch.int32, device=device,
+    )
+    global_vbs_parts = []
+
+    for i in range(num_seqs):
+        qb_start = q_block_offsets[i]
+        qb_end = q_block_offsets[i + 1]
+        n_q_blocks = qb_end - qb_start
+        kv_offset_blocks = kv_block_offsets[i]
+
+        idx = q2k_idx_list[i]
+        num = q2k_num_list[i]
+        vbs = variable_block_sizes_list[i]
+
+        mk = idx.shape[-1]
+        global_q2k_idx[:, :, qb_start:qb_end, :mk] = idx[:, :, :n_q_blocks, :] + kv_offset_blocks
+        global_q2k_num[:, :, qb_start:qb_end] = num[:, :, :n_q_blocks]
+        global_vbs_parts.append(vbs)
+
+    global_vbs = torch.cat(global_vbs_parts, dim=0).to(torch.int32).contiguous()
+
+    out_packed, _ = block_sparse_attn_from_indices(
+        q_packed, k_packed, v_packed,
+        global_q2k_idx, global_q2k_num, global_vbs,
+    )
+
+    out = torch.zeros(cu_q[-1], num_heads, head_dim, device=device, dtype=dtype)
+    q_offset = 0
+    for i in range(num_seqs):
+        _gather_from_padded(
+            out_packed, q_vbs_resolved[i], block_size,
+            out, q_offset, cu_q[i], cu_q[i + 1],
+        )
+        q_offset += padded_q_lens[i]
+
+    return out