Skip to content

vllm.models.minimax_m3.amd.mono.drafter.dispatch

Which MiniMax-M3.1 DSpark draft steps run on the drafter kernels (VLLM_ROCM_MINIMAX_M3_MONO_DRAFTER=1).

A draft step of 1..MAX_REQS requests (7 query rows each) runs its 5 layers and final norm in launches of up to 8 requests; other steps run the stock layers. Every condition is decided from TP-replicated state: the kernels exchange partials in-kernel, so ranks that disagreed would wait on each other.

Classes:

  • DrafterDecode –

    Routes supported draft steps of one MiniMaxM3DSparkModel to the kernels.

DrafterDecode

Routes supported draft steps of one MiniMaxM3DSparkModel to the kernels.

Methods:

  • maybe_forward –

    The step's final hidden states from the kernels, or None for the stock

Source code in vllm/models/minimax_m3/amd/mono/drafter/dispatch.py 
 56
 57
 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
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
class DrafterDecode:
    """Routes supported draft steps of one ``MiniMaxM3DSparkModel`` to the kernels."""

    def __init__(self, model, vllm_config: VllmConfig) -> None:
        self._model = model
        self._runner: DrafterRunner | None = None
        why = _config_refusal(vllm_config)
        self._enabled = why is None
        if why is not None:
            logger.info("MiniMax-M3 DSpark drafter kernels off: %s", why)

    def _disable(self, why: str) -> None:
        logger.warning("MiniMax-M3 DSpark drafter kernels off: %s", why)
        self._enabled = False

    def _ensure_runner(self) -> bool:
        """Build (or rebuild after a KV cache re-bind) the runner outside capture."""
        if self._runner is not None and self._runner.stale():
            if torch.cuda.is_current_stream_capturing():
                return False
            self._runner.close()
            self._runner = None
            logger.info("MiniMax-M3 DSpark drafter kernels: KV cache re-bound")
        if self._runner is not None:
            return True
        if torch.cuda.is_current_stream_capturing():
            return False
        from vllm.models.minimax_m3.amd.mono.drafter.runner import DrafterRunner

        try:
            self._runner = DrafterRunner(self._model)
        except MonoUnsupported as why:
            self._disable(str(why))
            return False
        logger.info("MiniMax-M3 DSpark drafter kernels on")
        return True

    def maybe_forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | None:
        """The step's final hidden states from the kernels, or None for the stock
        layers."""
        if not self._enabled:
            return None
        n = positions.numel()
        if n % layout.QPR or not 0 < n // layout.QPR <= layout.MAX_REQS:
            return None
        if positions.dtype != torch.int64:
            return None
        if not is_forward_context_available():
            return None
        if get_forward_context().attn_metadata is None:
            return None
        if not self._ensure_runner():
            return None
        assert self._runner is not None
        if inputs_embeds is None:
            inputs_embeds = self._model.embed_input_ids(input_ids)
        return self._runner.forward(positions, inputs_embeds)

_ensure_runner()

Build (or rebuild after a KV cache re-bind) the runner outside capture.

Source code in vllm/models/minimax_m3/amd/mono/drafter/dispatch.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def _ensure_runner(self) -> bool:
    """Build (or rebuild after a KV cache re-bind) the runner outside capture."""
    if self._runner is not None and self._runner.stale():
        if torch.cuda.is_current_stream_capturing():
            return False
        self._runner.close()
        self._runner = None
        logger.info("MiniMax-M3 DSpark drafter kernels: KV cache re-bound")
    if self._runner is not None:
        return True
    if torch.cuda.is_current_stream_capturing():
        return False
    from vllm.models.minimax_m3.amd.mono.drafter.runner import DrafterRunner

    try:
        self._runner = DrafterRunner(self._model)
    except MonoUnsupported as why:
        self._disable(str(why))
        return False
    logger.info("MiniMax-M3 DSpark drafter kernels on")
    return True

maybe_forward(input_ids, positions, inputs_embeds=None)

The step's final hidden states from the kernels, or None for the stock layers.

Source code in vllm/models/minimax_m3/amd/mono/drafter/dispatch.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def maybe_forward(
    self,
    input_ids: torch.Tensor | None,
    positions: torch.Tensor,
    inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor | None:
    """The step's final hidden states from the kernels, or None for the stock
    layers."""
    if not self._enabled:
        return None
    n = positions.numel()
    if n % layout.QPR or not 0 < n // layout.QPR <= layout.MAX_REQS:
        return None
    if positions.dtype != torch.int64:
        return None
    if not is_forward_context_available():
        return None
    if get_forward_context().attn_metadata is None:
        return None
    if not self._ensure_runner():
        return None
    assert self._runner is not None
    if inputs_embeds is None:
        inputs_embeds = self._model.embed_input_ids(input_ids)
    return self._runner.forward(positions, inputs_embeds)

_config_refusal(vllm_config)

Why this deployment cannot use the drafter kernels at all, or None.

Source code in vllm/models/minimax_m3/amd/mono/drafter/dispatch.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def _config_refusal(vllm_config: VllmConfig) -> str | None:
    """Why this deployment cannot use the drafter kernels at all, or None."""
    pc = vllm_config.parallel_config
    spec = vllm_config.speculative_config
    checks = (
        (on_gfx950(), "not gfx950"),
        (_flydsl_available(), "FlyDSL not installed"),
        (pc.tensor_parallel_size == 4, f"TP {pc.tensor_parallel_size} != 4"),
        (pc.data_parallel_size == 1, "DP > 1"),
        (pc.pipeline_parallel_size == 1, "PP > 1"),
        (spec is not None and spec.method == "dspark", "no DSpark draft"),
        (
            spec is not None and spec.num_speculative_tokens == layout.QPR,
            f"{layout.QPR} speculative tokens required",
        ),
    )
    for ok, why in checks:
        if not ok:
            return why
    return None