vsa_kernel/index.py

## pytorch sdpa version of block sparse ##
from typing import Tuple

import triton
import triton.language as tl
import torch


@triton.jit
def topk_index_to_map_kernel(
    map_ptr,
    index_ptr,
    map_bs_stride,
    map_h_stride,
    map_q_stride,
    map_kv_stride,
    index_bs_stride,
    index_h_stride,
    index_q_stride,
    index_kv_stride,
    topk,
):
    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride

    for i in tl.static_range(topk):
        index = tl.load(index_ptr_base + i * index_kv_stride)
        tl.store(map_ptr_base + index * map_kv_stride, 1.0)


@triton.jit
def map_to_index_kernel(
    map_ptr,
    index_ptr,
    index_num_ptr,
    map_bs_stride,
    map_h_stride,
    map_q_stride,
    map_kv_stride,
    index_bs_stride,
    index_h_stride,
    index_q_stride,
    index_kv_stride,
    index_num_bs_stride,
    index_num_h_stride,
    index_num_q_stride,
    num_kv_blocks,
):
    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride

    num = 0
    for i in tl.range(num_kv_blocks):
        map_entry = tl.load(map_ptr_base + i * map_kv_stride)
        if map_entry:
            tl.store(index_ptr_base + num * index_kv_stride, i)
            num += 1

    tl.store(index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + q * index_num_q_stride, num)


def topk_index_to_map(index: torch.Tensor, num_kv_blocks: int, transpose_map: bool = False):
    """
    Convert topk indices to a map.
    
    Args:
        index: [bs, h, num_q_blocks, topk]
            The topk indices tensor.
        num_kv_blocks: int
            The number of key-value blocks in the block_map returned
        transpose_map: bool
            If True, the block_map will be transposed on the final two dimensions.
    
    Returns:
        block_map: [bs, h, num_q_blocks, num_kv_blocks]
            A binary map where 1 indicates that the q block attends to the kv block.
    """
    bs, h, num_q_blocks, topk = index.shape

    if transpose_map is False:
        block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks), dtype=torch.bool, device=index.device)
    else:
        block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks), dtype=torch.bool, device=index.device)
        block_map = block_map.transpose(2, 3)

    grid = (bs, h, num_q_blocks)
    topk_index_to_map_kernel[grid](
        block_map,
        index,
        block_map.stride(0),
        block_map.stride(1),
        block_map.stride(2),
        block_map.stride(3),
        index.stride(0),
        index.stride(1),
        index.stride(2),
        index.stride(3),
        topk=topk,
    )

    return block_map


def map_to_index(block_map: torch.Tensor):
    """
    Convert a block map to indices and counts.
    
    Args:
        block_map: [bs, h, num_q_blocks, num_kv_blocks]
            The block map tensor.
    
    Returns:
        index: [bs, h, num_q_blocks, num_kv_blocks]
            The indices of the blocks.
        index_num: [bs, h, num_q_blocks]
            The number of blocks for each q block.
    """
    bs, h, num_q_blocks, num_kv_blocks = block_map.shape

    index = torch.full((block_map.shape), -1, dtype=torch.int32, device=block_map.device)
    index_num = torch.empty((bs, h, num_q_blocks), dtype=torch.int32, device=block_map.device)

    grid = (bs, h, num_q_blocks)
    map_to_index_kernel[grid](
        block_map,
        index,
        index_num,
        block_map.stride(0),
        block_map.stride(1),
        block_map.stride(2),
        block_map.stride(3),
        index.stride(0),
        index.stride(1),
        index.stride(2),
        index.stride(3),
        index_num.stride(0),
        index_num.stride(1),
        index_num.stride(2),
        num_kv_blocks=num_kv_blocks,
    )

    return index, index_num


@triton.jit
def _invert_indices_kernel(
    q2k_idx_ptr,
    q2k_num_ptr,
    k2q_idx_ptr,
    k2q_num_ptr,
    q2k_idx_b,
    q2k_idx_h,
    q2k_idx_q,
    q2k_idx_k,
    q2k_num_b,
    q2k_num_h,
    q2k_num_q,
    k2q_idx_b,
    k2q_idx_h,
    k2q_idx_k,
    k2q_idx_q,
    k2q_num_b,
    k2q_num_h,
    k2q_num_k,
    MAX_KV_PER_Q: tl.constexpr,
):
    # One program per (b, h, q): reserve a slot in k2q via atomicAdd, write q.
    pid_b = tl.program_id(0)
    pid_h = tl.program_id(1)
    pid_q = tl.program_id(2)

    n = tl.load(q2k_num_ptr + pid_b * q2k_num_b + pid_h * q2k_num_h + pid_q * q2k_num_q)

    q2k_row = (q2k_idx_ptr + pid_b * q2k_idx_b + pid_h * q2k_idx_h + pid_q * q2k_idx_q)

    for i in tl.range(0, MAX_KV_PER_Q):
        if i < n:
            kv = tl.load(q2k_row + i * q2k_idx_k)
            count_ptr = (k2q_num_ptr + pid_b * k2q_num_b + pid_h * k2q_num_h + kv * k2q_num_k)
            pos = tl.atomic_add(count_ptr, 1)
            tl.store(
                k2q_idx_ptr + pid_b * k2q_idx_b + pid_h * k2q_idx_h + kv * k2q_idx_k + pos * k2q_idx_q,
                pid_q,
            )


def invert_indices(
    q2k_idx: torch.Tensor,
    q2k_num: torch.Tensor,
    num_kv_blocks: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Transpose a Q->KV index list into a K->Q one via atomic compaction (GPU)."""
    if q2k_idx.dim() != 4:
        raise ValueError(f"q2k_idx must be [B, H, Nq, Mk], got shape={tuple(q2k_idx.shape)}")
    if q2k_num.dim() != 3:
        raise ValueError(f"q2k_num must be [B, H, Nq], got shape={tuple(q2k_num.shape)}")
    if not q2k_idx.is_cuda or not q2k_num.is_cuda:
        raise RuntimeError("invert_indices requires CUDA tensors.")

    B, H, Nq, Mk = q2k_idx.shape
    if q2k_num.shape != (B, H, Nq):
        raise ValueError(f"q2k_num shape {tuple(q2k_num.shape)} does not match q2k_idx "
                         f"[B, H, Nq] = {(B, H, Nq)}")

    q2k_idx = q2k_idx.contiguous()
    q2k_num = q2k_num.contiguous()
    if q2k_idx.dtype != torch.int32:
        q2k_idx = q2k_idx.to(torch.int32)
    if q2k_num.dtype != torch.int32:
        q2k_num = q2k_num.to(torch.int32)

    # Any KV block is attended by at most Nq Q blocks (one per Q row), so
    # `Nq` is a tight upper bound on the compacted K->Q slots.
    k2q_idx = torch.empty(
        (B, H, num_kv_blocks, Nq),
        dtype=torch.int32,
        device=q2k_idx.device,
    )
    k2q_num = torch.zeros(
        (B, H, num_kv_blocks),
        dtype=torch.int32,
        device=q2k_idx.device,
    )

    grid = (B, H, Nq)
    _invert_indices_kernel[grid](
        q2k_idx,
        q2k_num,
        k2q_idx,
        k2q_num,
        q2k_idx.stride(0),
        q2k_idx.stride(1),
        q2k_idx.stride(2),
        q2k_idx.stride(3),
        q2k_num.stride(0),
        q2k_num.stride(1),
        q2k_num.stride(2),
        k2q_idx.stride(0),
        k2q_idx.stride(1),
        k2q_idx.stride(2),
        k2q_idx.stride(3),
        k2q_num.stride(0),
        k2q_num.stride(1),
        k2q_num.stride(2),
        MAX_KV_PER_Q=Mk,
    )

    return k2q_idx, k2q_num