Skip to content

vllm.models.minimax_m3.amd.indexer_cp

Indexer-only context parallelism for MiniMax M3 (AMD MSA indexer).

Every rank projects all H index heads and scores 1/P of the context blocks, then one exchange gives it each shard's candidates for the heads it owns, cutting the replicated index-cache read to 1/P. Shards cut to the full k so selection stays exact; candidates carry a generation tag, so all ranks must call this in lockstep.

Classes:

Functions:

IndexerCpPeers dataclass

The buffers the CP selector exchanges through.

cand_ptrs is every rank's candidate buffer; cp_gen is this rank's own.

Source code in vllm/models/minimax_m3/amd/indexer_cp.py 
44
45
46
47
48
49
50
51
52
@dataclass(frozen=True)
class IndexerCpPeers:
    """The buffers the CP selector exchanges through.

    ``cand_ptrs`` is every rank's candidate buffer; ``cp_gen`` is this rank's own.
    """

    cand_ptrs: torch.Tensor
    cp_gen: torch.Tensor

_ipc_alloc(nbytes, group)

Allocate a zeroed, IPC-shared nbytes on every rank.

Returns a [world] int64 CPU tensor of their device addresses, never freed.

Source code in vllm/models/minimax_m3/amd/indexer_cp.py 
32
33
34
35
36
37
38
39
40
41
def _ipc_alloc(nbytes: int, group) -> torch.Tensor:
    """Allocate a zeroed, IPC-shared ``nbytes`` on every rank.

    Returns a ``[world]`` int64 CPU tensor of their device addresses, never freed.
    """
    from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce

    nbytes = (nbytes + _HIP_PAGE - 1) // _HIP_PAGE * _HIP_PAGE
    peers = CustomAllreduce.create_shared_buffer(nbytes, group=group.cpu_group)
    return torch.tensor(peers, dtype=torch.int64, device="cpu")

get_indexer_cp_group()

The group the candidate exchange runs over.

Source code in vllm/models/minimax_m3/amd/indexer_cp.py 
21
22
23
def get_indexer_cp_group():
    """The group the candidate exchange runs over."""
    return get_tp_group()

get_indexer_cp_peers(owned_heads, max_rows, topk)

The buffers the CP selector exchanges through, allocated once per key.

Cached and shared by every layer, sized to max_rows so captured graphs always address the same buffer. Every rank must reach this with the same key.

Source code in vllm/models/minimax_m3/amd/indexer_cp.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def get_indexer_cp_peers(owned_heads: int, max_rows: int, topk: int) -> IndexerCpPeers:
    """The buffers the CP selector exchanges through, allocated once per key.

    Cached and shared by every layer, sized to ``max_rows`` so captured graphs
    always address the same buffer. Every rank must reach this with the same key.
    """
    from aiter.ops.msa_block_select import CP_MAX_BLOCKS, topk_cp_candidate_numel

    group = get_indexer_cp_group()
    key = (group.world_size, owned_heads, max_rows, topk)
    peers = _PEERS.get(key)
    if peers is not None:
        return peers

    numel = topk_cp_candidate_numel(group.world_size, owned_heads, max_rows, topk)
    peers = IndexerCpPeers(
        cand_ptrs=_ipc_alloc(numel * 8, group),
        cp_gen=torch.zeros(CP_MAX_BLOCKS, dtype=torch.int32, device="cuda"),
    )
    _PEERS[key] = peers
    logger.info_once(
        "MiniMax M3 indexer CP: %d KiB candidate buffer per rank mapped over "
        "%d ranks [owned_heads=%d, rows=%d, topk=%d, generations=2]",
        numel * 8 // 1024,
        group.world_size,
        owned_heads,
        max_rows,
        topk,
    )
    return peers