vllm.models.minimax_m3.amd.ops.indexer_cp_exchange

Candidate exchange for MiniMax-M3 Triton indexer CP (graph-safe).

Each rank keeps, per index head, its own top-k over a 1/P block shard, all-gathers the packed keys (payload independent of context length), then merges per head. The merge is exact only when every shard scored the same head, so callers must score every global head on every rank. Forced init/local blocks are pinned on both the shard and the merge so they cannot be dropped before the exchange.

Functions:

• aiter_all_gather –

  AITER IPC all-gather over the TP group as [world, *x.shape], or None.

• all_gather –

  All-gather contiguous x over group to [world, *x.shape].

• local_topk_keys –

  Write [heads, tokens, topk] packed int64 keys for this shard.

• merge_topk_keys –

  Merge gathered [world, heads, tokens, topk] keys into global top-k.

aiter_all_gather(x)

AITER IPC all-gather over the TP group as [world, *x.shape], or None.

The payload is viewed as int32 so AITER's pybind layer (which has no float64 after it relabels integers as floats) accepts it -- int64 keys and bf16 index queries alike. Concatenates along dim 0, then views back to x's dtype.

Source code in vllm/models/minimax_m3/amd/ops/indexer_cp_exchange.py

def aiter_all_gather(x: torch.Tensor) -> torch.Tensor | None:
    """AITER IPC all-gather over the TP group as ``[world, *x.shape]``, or None.

    The payload is viewed as int32 so AITER's pybind layer (which has no
    float64 after it relabels integers as floats) accepts it -- int64 keys and
    bf16 index queries alike. Concatenates along dim 0,
    then views back to ``x``'s dtype.
    """
    from vllm._aiter_ops import rocm_aiter_ops

    comm = rocm_aiter_ops.get_aiter_allreduce()
    if comm is None or comm.disabled:
        return None
    if not x.is_contiguous() or (x.shape[-1] * x.element_size()) % 4:
        return None
    packed = x.view(torch.int32)
    if not comm.should_custom_ag(packed):
        return None
    gathered = comm.custom_all_gather(packed, dim=0)
    if gathered is None:
        return None
    world = get_tp_group().world_size
    return gathered.view(x.dtype).view(world, *x.shape)

all_gather(x, group)

All-gather contiguous x over group to [world, *x.shape].

AITER's IPC all-gather is tried first: on ROCm the TP communicator's all_gather stays on the RCCL collective, and its AITER path is only enabled for DP groups, so it would not pick AITER for these small per-layer payloads.

Source code in vllm/models/minimax_m3/amd/ops/indexer_cp_exchange.py

def all_gather(x: torch.Tensor, group: GroupCoordinator) -> torch.Tensor:
    """All-gather contiguous ``x`` over ``group`` to ``[world, *x.shape]``.

    AITER's IPC all-gather is tried first: on ROCm the TP communicator's
    ``all_gather`` stays on the RCCL collective, and its AITER path is only
    enabled for DP groups, so it would not pick AITER for these small
    per-layer payloads.
    """
    gathered = aiter_all_gather(x)
    if gathered is None:
        gathered = group.all_gather(x, dim=0)
    return gathered.view(group.world_size, *x.shape)

local_topk_keys(scores, seq_lens, *, topk, rank, world, query_len, global_blocks, init_blocks, local_blocks, out)

Write [heads, tokens, topk] packed int64 keys for this shard.

Source code in vllm/models/minimax_m3/amd/ops/indexer_cp_exchange.py

def local_topk_keys(
    scores: torch.Tensor,
    seq_lens: torch.Tensor,
    *,
    topk: int,
    rank: int,
    world: int,
    query_len: int,
    global_blocks: int,
    init_blocks: int,
    local_blocks: int,
    out: torch.Tensor,
) -> torch.Tensor:
    """Write ``[heads, tokens, topk]`` packed int64 keys for this shard."""
    heads, tokens, local = scores.shape
    if global_blocks >= 0xFFFF:
        raise ValueError("packed CP keys support fewer than 65535 sparse blocks")
    if tokens == 0:
        return out
    width = min(
        DECODE_TOPK_TILE,
        max(triton.next_power_of_2(max(local, 1)), triton.next_power_of_2(topk)),
    )
    _local_topk_keys[(tokens, heads)](
        scores,
        out,
        seq_lens,
        scores.stride(0),
        scores.stride(1),
        scores.stride(2),
        out.stride(0),
        out.stride(1),
        QUERY_LEN=query_len,
        LOCAL_BLOCKS=local,
        GLOBAL_BLOCKS=global_blocks,
        RANK=rank,
        WORLD=world,
        TOPK=topk,
        INIT_BLOCKS=init_blocks,
        LOCAL_KEEP=local_blocks,
        BLOCK_SIZE_K=width,
        BLOCK_SIZE_T=triton.next_power_of_2(topk),
        num_warps=DECODE_TOPK_NUM_WARPS,
    )
    return out

merge_topk_keys(keys, topk_idx, seq_lens, *, query_len, init_blocks, local_blocks, attention_block_table=None, sparse_block_table_out=None, sparse_context_lens_out=None, block_page_stride=None)

Merge gathered [world, heads, tokens, topk] keys into global top-k.

Source code in vllm/models/minimax_m3/amd/ops/indexer_cp_exchange.py

def merge_topk_keys(
    keys: torch.Tensor,
    topk_idx: torch.Tensor,
    seq_lens: torch.Tensor,
    *,
    query_len: int,
    init_blocks: int,
    local_blocks: int,
    attention_block_table: torch.Tensor | None = None,
    sparse_block_table_out: torch.Tensor | None = None,
    sparse_context_lens_out: torch.Tensor | None = None,
    block_page_stride: int | None = None,
) -> torch.Tensor:
    """Merge gathered ``[world, heads, tokens, topk]`` keys into global top-k."""
    shards, heads, tokens, per_shard = keys.shape
    topk = topk_idx.shape[2]
    candidates = shards * per_shard
    if tokens == 0:
        return topk_idx
    emit = attention_block_table is not None
    dummy = topk_idx
    _merge_topk_keys[(tokens, heads)](
        keys,
        topk_idx,
        seq_lens,
        attention_block_table if emit else dummy,
        sparse_block_table_out if emit else dummy,
        sparse_context_lens_out if emit else dummy,
        keys.stride(0),
        keys.stride(1),
        keys.stride(2),
        topk_idx.stride(0),
        topk_idx.stride(1),
        attention_block_table.stride(0) if attention_block_table is not None else 0,
        sparse_block_table_out.stride(0) if sparse_block_table_out is not None else 0,
        QUERY_LEN=query_len,
        TOPK=topk,
        INIT_BLOCKS=init_blocks,
        LOCAL_KEEP=local_blocks,
        KEYS_PER_SHARD=per_shard,
        REAL_CANDIDATES=candidates,
        CANDIDATES=triton.next_power_of_2(max(candidates, 1)),
        BLOCK_SIZE_T=triton.next_power_of_2(topk),
        EMIT_SPARSE_TABLE=emit,
        BLOCK_SIZE=128,
        PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
        BLOCK_PAGE_STRIDE=block_page_stride or PAGES_PER_SPARSE_BLOCK,
        num_warps=DECODE_TOPK_NUM_WARPS,
    )
    return topk_idx