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

def _sides_are_packed(k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool:
    """Whether a block holds both its K and V pages instead of one side."""
    return v_cache.shape[0] != k_cache.shape[0]

@triton.jit
def _write_sparse_block_table_row_from_values(
    blk,  # [BLOCK_SIZE_T] int32 selected logical block ids for this query
    bt_row,  # [max_blocks] int32, this request's logical page table
    sbt_row,  # [topk * PAGES_PER_BLOCK] int32, physical 16-page table
    ctx_ptr,  # int32, this query's attended context length
    abs_pos,  # absolute position of this query token, may be negative
    max_topk,
    SPARSE_BLOCK_SIZE_C: tl.constexpr,
    PAGES_PER_BLOCK: tl.constexpr,
    BLOCK_PAGE_STRIDE: tl.constexpr,
    BLOCK_SIZE_T: tl.constexpr,
):
    """Compact one query's selected logical blocks into physical page-16 ids.

    ``BLOCK_PAGE_STRIDE`` is how many page ids a block spans, which is
    ``PAGES_PER_BLOCK`` when each K/V side is its own dense plane and twice
    that when both sides share a block. Padded rows carry a negative
    ``abs_pos``, which clamps the causal range to empty.
    """
    causal_len = tl.maximum(abs_pos + 1, 0)
    self_blk = abs_pos // SPARSE_BLOCK_SIZE_C

    off_t = tl.arange(0, BLOCK_SIZE_T)
    valid = (off_t < max_topk) & (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
    is_tail = valid & (blk == self_blk)
    is_full = valid & (blk < self_blk)

    n_full = tl.sum(is_full.to(tl.int32), axis=0)
    n_valid = tl.sum(valid.to(tl.int32), axis=0)
    earlier_full = tl.cumsum(is_full.to(tl.int32), axis=0) - is_full.to(tl.int32)
    slot = tl.where(is_full, earlier_full, n_full)

    logical_page = tl.load(bt_row + blk, mask=valid, other=0).to(tl.int32)
    base_phys = logical_page * BLOCK_PAGE_STRIDE
    dst_base = slot * PAGES_PER_BLOCK

    for j in tl.static_range(PAGES_PER_BLOCK):
        tl.store(sbt_row + dst_base + j, base_phys + j, mask=valid)

    n_used = n_valid * PAGES_PER_BLOCK
    off_w = tl.arange(0, BLOCK_SIZE_T * PAGES_PER_BLOCK)
    row_width = max_topk * PAGES_PER_BLOCK
    tl.store(
        sbt_row + off_w,
        tl.zeros_like(off_w),
        mask=(off_w >= n_used) & (off_w < row_width),
    )

    tail_tokens = causal_len - self_blk * SPARSE_BLOCK_SIZE_C
    has_tail = tl.sum(is_tail.to(tl.int32), axis=0) > 0
    ctx = n_full * SPARSE_BLOCK_SIZE_C + tl.where(has_tail, tail_tokens, 0)
    ctx = tl.where(has_tail, ctx, tl.minimum(n_valid * SPARSE_BLOCK_SIZE_C, causal_len))
    tl.store(ctx_ptr, ctx)

@torch.no_grad()
def minimax_m3_build_sparse_block_table_decode(
    topk_idx: torch.Tensor,  # [1, batch * decode_query_len, topk]
    block_table: torch.Tensor,  # [batch, max_blocks]
    seq_lens: torch.Tensor,  # [batch]
    decode_query_len: int = 1,
    block_page_stride: int = PAGES_PER_SPARSE_BLOCK,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Build one page-16 sparse block table row per decode query token."""
    total_q = topk_idx.shape[1]
    expected_q = block_table.shape[0] * decode_query_len
    assert total_q == expected_q, (
        "MiniMax-M3 decode top-k rows must equal batch * decode_query_len: "
        f"{total_q} != {expected_q}"
    )
    sparse_bt, sparse_ctx = minimax_m3_alloc_sparse_block_table(topk_idx)
    topk = topk_idx.shape[-1]
    _build_sparse_block_table_kernel[(total_q,)](
        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_PAGE_STRIDE=block_page_stride,
        BLOCK_SIZE_T=triton.next_power_of_2(topk),
        DECODE_QUERY_LEN=decode_query_len,
    )
    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]
    block_page_stride: int = PAGES_PER_SPARSE_BLOCK,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Build one page-16 sparse block table row per prefill query token."""
    sparse_bt, sparse_ctx = minimax_m3_alloc_sparse_block_table(topk_idx)
    topk = topk_idx.shape[-1]
    _build_sparse_block_table_prefill_kernel[(topk_idx.shape[1],)](
        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_PAGE_STRIDE=block_page_stride,
        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,
    )

def minimax_m3_rebase_block_table_to_page16(
    block_table: torch.Tensor,
    out: torch.Tensor | None = None,
) -> torch.Tensor:
    """Rebase a logical page table onto AITER's page-16 page numbering.

    The page ids AITER's top-k emits for the attend are
    ``block_table[blk] * pages_per_block + j``, and ``pages_per_block`` is fixed
    at AITER build time to one side's pages. An interleaved block spans both
    sides, so the only way to reach its pages through that kernel is to hand it
    a table already scaled to the wider stride. The caller does this once per
    step, since every sparse layer resolves its selection through the same
    table.
    """
    if out is None:
        out = torch.empty_like(block_table)
    return torch.mul(block_table, PAGE16_SIDES_PER_BLOCK, out=out)

def minimax_m3_sparse_block_page_stride(
    k_cache: torch.Tensor, v_cache: torch.Tensor
) -> int:
    """How many page ids one sparse block spans.

    Each side owns ``PAGES_PER_SPARSE_BLOCK`` pages. When the sides are dense
    planes a block is exactly that wide; when they interleave, a block covers
    both sides' pages and V is reached from the same page id at a fixed offset.
    """
    return PAGES_PER_SPARSE_BLOCK * (2 if _sides_are_packed(k_cache, v_cache) else 1)