Skip to content

vllm.models.minimax_m3.common.varlen_decode

Ragged decode metadata shared by the MiniMax M3 indexer and sparse builders.

Under adaptive verification, each decode request verifies 1 to 1 + K tokens, so the builders pass query_start_loc and a token-to-request map to their kernels instead of one uniform per-request query length.

Classes:

Functions:

VarlenDecodeBuilderMixin

Builder hooks for ragged decode; subclasses opt in via supports_varlen_decode.

Source code in vllm/models/minimax_m3/common/varlen_decode.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class VarlenDecodeBuilderMixin:
    """Builder hooks for ragged decode; subclasses opt in via
    `supports_varlen_decode`."""

    if TYPE_CHECKING:
        # Provided by AttentionMetadataBuilder, which every user mixes in.
        @classmethod
        def get_cudagraph_support(
            cls, vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
        ) -> AttentionCGSupport: ...

    @classmethod
    def supports_varlen_decode(cls) -> bool:
        return False

    @classmethod
    def _use_varlen_decode(cls, vllm_config: VllmConfig) -> bool:
        spec_config = vllm_config.speculative_config
        return (
            cls.supports_varlen_decode()
            and spec_config is not None
            and spec_config.enable_adaptive_verification
        )

    @classmethod
    def get_varlen_cudagraph_max_query_len(
        cls,
        vllm_config: VllmConfig,
        kv_cache_spec: KVCacheSpec,
    ) -> int | None:
        # Varlen decode replays decode-only batches of 1 to 1 + K tokens per
        # request; mixed prefill+decode batches never replay FULL graphs.
        if not cls._use_varlen_decode(vllm_config) or (
            cls.get_cudagraph_support(vllm_config, kv_cache_spec)
            != AttentionCGSupport.UNIFORM_BATCH
        ):
            return None
        return max_decode_query_len(vllm_config)

build_decode_query_layout(common_attn_metadata, query_start_loc, num_decodes, num_decode_tokens, token_to_req_buffer, varlen)

Return (decode_query_len, query_start_loc, token_to_req_indices).

A ragged batch returns query length 0 plus the device layout; a uniform batch returns its query length and no layout tensors.

Source code in vllm/models/minimax_m3/common/varlen_decode.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def build_decode_query_layout(
    common_attn_metadata: CommonAttentionMetadata,
    query_start_loc: torch.Tensor,
    num_decodes: int,
    num_decode_tokens: int,
    token_to_req_buffer: torch.Tensor | None,
    varlen: bool,
) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:
    """Return `(decode_query_len, query_start_loc, token_to_req_indices)`.

    A ragged batch returns query length 0 plus the device layout; a uniform
    batch returns its query length and no layout tensors.
    """
    if varlen:
        assert token_to_req_buffer is not None
        token_to_req = common_attn_metadata.token_to_req_indices(token_to_req_buffer)
        return (
            0,
            query_start_loc[: num_decodes + 1],
            token_to_req[:num_decode_tokens],
        )
    qsl_cpu = common_attn_metadata.query_start_loc_cpu
    query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
    decode_query_len = int(query_lens_cpu[0].item())
    assert decode_query_len > 0
    assert torch.all((query_lens_cpu == decode_query_len) | (query_lens_cpu == 0))
    assert num_decode_tokens == num_decodes * decode_query_len
    return decode_query_len, None, None

make_token_to_req_buffer(vllm_config, device, varlen)

Token-to-request scratch for ragged decode; uniform builders skip it.

Source code in vllm/models/minimax_m3/common/varlen_decode.py 
63
64
65
66
67
68
69
70
71
72
73
def make_token_to_req_buffer(
    vllm_config: VllmConfig, device: torch.device, varlen: bool
) -> torch.Tensor | None:
    """Token-to-request scratch for ragged decode; uniform builders skip it."""
    if not varlen:
        return None
    return torch.empty(
        vllm_config.scheduler_config.max_num_batched_tokens,
        dtype=torch.int32,
        device=device,
    )