vllm.models.minimax_m3.amd.ops.sparse_pa ¶

@torch.no_grad()
def minimax_m3_build_sparse_block_table(
    topk_idx: torch.Tensor,  # [1, batch, topk]
    block_table: torch.Tensor,  # [batch, max_blocks]
    seq_lens: torch.Tensor,  # [batch]
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compact selected logical sparse blocks into physical page-16 tables."""
    assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
    batch = topk_idx.shape[1]
    topk = topk_idx.shape[-1]
    width = topk * PAGES_PER_SPARSE_BLOCK
    sparse_bt = torch.empty((batch, width), dtype=torch.int32, device=topk_idx.device)
    sparse_ctx = torch.empty((batch,), dtype=torch.int32, device=topk_idx.device)
    _build_sparse_block_table_kernel[(batch,)](
        topk_idx,
        block_table,
        seq_lens,
        sparse_bt,
        sparse_ctx,
        topk,
        topk_idx.stride(1),
        topk_idx.stride(2),
        block_table.stride(0),
        sparse_bt.stride(0),
        SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
        PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
        BLOCK_SIZE_T=triton.next_power_of_2(topk),
    )
    return sparse_bt, sparse_ctx

@torch.no_grad()
def minimax_m3_build_sparse_block_table_prefill(
    topk_idx: torch.Tensor,  # [1, total_q, topk]
    block_table: torch.Tensor,  # [batch, max_blocks]
    query_req_id: torch.Tensor,  # [total_q]
    query_abs_pos: torch.Tensor,  # [total_q]
) -> tuple[torch.Tensor, torch.Tensor]:
    """Build one page-16 sparse block table row per prefill query token."""
    assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
    total_q = topk_idx.shape[1]
    topk = topk_idx.shape[-1]
    width = topk * PAGES_PER_SPARSE_BLOCK
    sparse_bt = torch.empty((total_q, width), dtype=torch.int32, device=topk_idx.device)
    sparse_ctx = torch.empty((total_q,), dtype=torch.int32, device=topk_idx.device)
    _build_sparse_block_table_prefill_kernel[(total_q,)](
        topk_idx,
        block_table,
        query_req_id,
        query_abs_pos,
        sparse_bt,
        sparse_ctx,
        topk,
        topk_idx.stride(1),
        topk_idx.stride(2),
        block_table.stride(0),
        sparse_bt.stride(0),
        SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
        PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
        BLOCK_SIZE_T=triton.next_power_of_2(topk),
    )
    return sparse_bt, sparse_ctx

@torch.no_grad()
def minimax_m3_insert_index_cache(
    index_k: torch.Tensor,
    index_cache: torch.Tensor,
    index_slot_mapping: torch.Tensor,
) -> None:
    """Scatter index keys into MiniMax-M3's key-only side cache."""
    if index_k.numel() == 0 or index_cache.numel() == 0:
        return
    if index_k.dim() != 2 or index_cache.dim() != 3:
        raise ValueError("MiniMax-M3 index cache insert expects [N,D] and [B,T,D]")
    if index_k.shape[1] != index_cache.shape[2]:
        raise ValueError("MiniMax-M3 index key dim must match index cache head dim")
    if index_slot_mapping.dim() != 1 or index_slot_mapping.shape[0] != index_k.shape[0]:
        raise ValueError("MiniMax-M3 index slot mapping must be a length-N vector")
    if index_cache.stride(2) != 1:
        raise ValueError("MiniMax-M3 index cache requires contiguous head dimension")

    head_dim = index_k.shape[1]
    _insert_index_cache_kernel[(index_k.shape[0],)](
        index_k,
        index_cache,
        index_slot_mapping,
        index_k.stride(0),
        index_k.stride(1),
        index_cache.stride(0),
        index_cache.stride(1),
        index_cache.stride(2),
        index_slot_mapping.stride(0),
        CACHE_BLOCK_SIZE=index_cache.shape[1],
        HEAD_DIM=head_dim,
        BLOCK_D=triton.next_power_of_2(head_dim),
        num_warps=4,
    )