vllm.models.minimax_m3.amd.ops.indexer_context_parallel

Opt-in M3 compute-only index context partitioning, independent of global DCP.

Functions:

• check_context_score_inputs –

  Validate the layout indexer_context_scores assumes; call once.

• indexer_context_scores –

  Return [heads,tokens,ceil(blocks/world)] with round-robin logical blocks.

_decode_score_chunks(batch, local_blocks)

Chunks one score row is split across.

A count, not a size: the grid is (request, chunk), so this is the second grid dim. The round trip through the size drops chunk counts that do not divide the blocks and would launch programs that only return.

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

def _decode_score_chunks(batch: int, local_blocks: int) -> int:
    """Chunks one score row is split across.

    A count, not a size: the grid is ``(request, chunk)``, so this is the
    second grid dim. The round trip through the size drops chunk counts that
    do not divide the blocks and would launch programs that only return.
    """
    if local_blocks <= 0:
        return 1
    target = max(1, DECODE_SCORE_TARGET_GRID // max(1, batch))
    chunks = min(1 << (target.bit_length() - 1), local_blocks)
    chunks = min(chunks, max(1, triton.cdiv(local_blocks, DECODE_SCORE_MIN_BLOCKS)))
    return triton.cdiv(local_blocks, triton.cdiv(local_blocks, chunks))

check_context_score_inputs(idx_q, index_cache, block_table, seq_lens)

Validate the layout indexer_context_scores assumes; call once.

These hold for every step of a served model, so they are not re-checked on the hot path.

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

def check_context_score_inputs(
    idx_q: torch.Tensor,
    index_cache: torch.Tensor,
    block_table: torch.Tensor,
    seq_lens: torch.Tensor,
) -> None:
    """Validate the layout ``indexer_context_scores`` assumes; call once.

    These hold for every step of a served model, so they are not re-checked
    on the hot path.
    """
    if (
        idx_q.ndim != 3
        or idx_q.shape[2] != 128
        or idx_q.dtype != torch.bfloat16
        or idx_q.stride(2) != 1
    ):
        raise ValueError(
            "index queries must be BF16 [tokens,heads,128] with contiguous D"
        )
    if (
        index_cache.ndim != 3
        or index_cache.shape[1:] != (128, 128)
        or not index_cache.is_contiguous()
        or index_cache.shape[0] < 1
    ):
        raise ValueError("index cache must be contiguous non-empty [pages,128,128]")
    if index_cache.dtype != torch.bfloat16:
        # The score loop applies no per-token scale.
        raise ValueError(
            f"index cache must be bf16 for CP scoring, got {index_cache.dtype}"
        )
    if (
        block_table.ndim != 2
        or block_table.stride(1) != 1
        or block_table.dtype != torch.int32
        or seq_lens.ndim != 1
        or seq_lens.dtype != torch.int32
        or not seq_lens.is_contiguous()
    ):
        raise ValueError("block table and lengths must be int32 with unit stride")
    if any(
        not x.is_cuda or x.device != idx_q.device
        for x in (index_cache, block_table, seq_lens)
    ):
        raise ValueError("all score inputs must be on the same GPU")

indexer_context_scores(idx_q, index_cache, block_table, seq_lens, max_seq_len, rank, world_size, max_query_len, sm_scale, out=None)

Return [heads,tokens,ceil(blocks/world)] with round-robin logical blocks.

max_seq_len sizes the grid and must be a capture-stable bound, not the batch's own longest sequence: the grid and CHUNK are baked into a CUDA graph at capture, where the dummy batch is one token long, and a replay sized off that would score only the first block of each shard. Pass max_model_len; the per-request causal bound clips the real work and chunks past the end return immediately.

out, when given, is a [heads, >=tokens, >=local] buffer written in place (stable address for CUDA graph capture).

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

def indexer_context_scores(
    idx_q,
    index_cache,
    block_table,
    seq_lens,
    max_seq_len,
    rank,
    world_size,
    max_query_len,
    sm_scale,
    out=None,
):
    """Return [heads,tokens,ceil(blocks/world)] with round-robin logical blocks.

    ``max_seq_len`` sizes the grid and must be a capture-stable bound, not the
    batch's own longest sequence: the grid and ``CHUNK`` are baked into a CUDA
    graph at capture, where the dummy batch is one token long, and a replay
    sized off that would score only the first block of each shard. Pass
    ``max_model_len``; the per-request causal bound clips the real work and
    chunks past the end return immediately.

    ``out``, when given, is a ``[heads, >=tokens, >=local]`` buffer written in
    place (stable address for CUDA graph capture).
    """
    tokens, heads, _ = idx_q.shape
    blocks = triton.cdiv(max_seq_len, 128)
    # Layout invariants are checked once by ``check_context_score_inputs``;
    # only what can change per step is checked here.
    assert tokens == seq_lens.numel() * max_query_len, "rows != batch * query_len"
    assert block_table.shape[1] >= blocks, "block table narrower than max_seq_len"
    num_pages = int(index_cache.shape[0])
    local = triton.cdiv(blocks, world_size)
    out_ok = (
        out is not None
        and out.ndim == 3
        and out.shape[0] == heads
        and out.shape[1] >= tokens
        and out.shape[2] >= local
        and out.dtype == torch.float32
        and out.device == idx_q.device
    )
    if out_ok:
        scores = out[:, :tokens, :local]
    else:
        # Profiling dummy batches can exceed the persistent decode buffer;
        # do not record a fresh allocation into a CUDA graph.
        if out is not None and torch.cuda.is_current_stream_capturing():
            raise ValueError(
                "out must be fp32 [heads, >=tokens, >=local] on device: "
                f"got {tuple(out.shape)} {out.dtype} {out.device}; "
                f"need heads={heads} tokens={tokens} local={local} "
                f"device={idx_q.device}"
            )
        scores = torch.empty(
            (heads, tokens, local), dtype=torch.float32, device=idx_q.device
        )
    if tokens:
        batch = seq_lens.numel()
        chunks = min(local, _decode_score_chunks(batch, local))
        # CHUNK bounds the score loop and the pipelining this kernel is built
        # around reads that bound, so it is rounded to a power of two while
        # the strides stay runtime. Rounding up never drops a block: a wider
        # chunk over correspondingly fewer programs still spans the shard.
        chunk = triton.next_power_of_2(triton.cdiv(local, chunks))
        _context_score[(batch, triton.cdiv(local, chunk))](
            idx_q,
            index_cache,
            block_table,
            seq_lens,
            scores,
            Q_TOKEN_STRIDE=idx_q.stride(0),
            Q_HEAD_STRIDE=idx_q.stride(1),
            TABLE_STRIDE=block_table.stride(0),
            SCORE_HEAD_STRIDE=scores.stride(0),
            SCORE_ROW_STRIDE=scores.stride(1),
            SCORE_BLK_STRIDE=scores.stride(2),
            LOCAL_BLOCKS=local,
            GLOBAL_BLOCKS=blocks,
            NUM_PAGES=num_pages,
            HEADS=heads,
            QUERY_LEN=max_query_len,
            RANK=rank,
            WORLD=world_size,
            CHUNK=chunk,
            N=max(16, triton.next_power_of_2(heads * max_query_len)),
            SCALE=sm_scale * 1.4426950409,
            num_stages=3,
        )
    return scores