"""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 sizes = block_sizes.cpu().tolist() for actual in sizes: 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 sizes = block_sizes.cpu().tolist() for actual in sizes: 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)) 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