vllm.models.minimax_m3.common.dspark

MiniMax-M3.1 DSpark draft model.

The Qwen3 DSpark block drafter (context-KV precompute, non-causal query block, Markov and confidence heads) built from MiniMax-M3 layers: Gemma RMSNorm (x * (1 + w)) everywhere, partial RoPE, and a dense SwiGLU-OAI MLP. The checkpoint ships no token embedding or LM head; both are shared with the target.

Classes:

• MiniMaxM3DSparkMLP –

  Dense SwiGLU-OAI MLP: gate * sigmoid(alpha * gate) * (up + beta).

• MiniMaxM3DSparkModel –

Functions:

• dspark_qknorm_rope_kv_insert –

  The fused QK-norm/RoPE on qkv, also storing the draft tokens' k/v.

MiniMaxM3DSparkMLP

Bases: Module

Dense SwiGLU-OAI MLP: gate * sigmoid(alpha * gate) * (up + beta).

Source code in vllm/models/minimax_m3/common/dspark.py

class MiniMaxM3DSparkMLP(nn.Module):
    """Dense SwiGLU-OAI MLP: gate * sigmoid(alpha * gate) * (up + beta)."""

    def __init__(
        self,
        config: MiniMaxM3TextConfig,
        quant_config: QuantizationConfig | None,
        prefix: str,
    ) -> None:
        super().__init__()
        if config.hidden_act != "swigluoai":
            raise ValueError(
                f"Unsupported MiniMax DSpark activation {config.hidden_act!r}; "
                "only swigluoai is supported."
            )
        intermediate_size = config.dspark_ffn_hidden_size
        self.gate_up_proj = MergedColumnParallelLinear(
            config.hidden_size,
            [intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            config.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
        self.act_fn = SiluAndMulWithClamp(
            swiglu_limit=config.swiglu_limit,
            alpha=config.swiglu_alpha,
            beta=config.swiglu_beta,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up, _ = self.gate_up_proj(x)
        x, _ = self.down_proj(self.act_fn(gate_up))
        return x

MiniMaxM3DSparkModel

Bases: Qwen3DSparkModel

Methods:

• precompute_and_store_context_kv –

  One kernel norms, ropes and stores the context K/V of every layer,

Source code in vllm/models/minimax_m3/common/dspark.py

class MiniMaxM3DSparkModel(Qwen3DSparkModel):
    decoder_layer_cls = MiniMaxM3DSparkDecoderLayer
    quant_config_mapper = WeightsMapper(
        orig_to_new_prefix={_CHECKPOINT_PREFIX: "model."}
    )

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        start_layer_id: int = 0,
        prefix: str = "",
    ) -> None:
        super().__init__(
            vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix
        )
        hidden_size, eps = self.config.hidden_size, self.config.rms_norm_eps
        self.hidden_norm = _gemma_rms_norm(hidden_size, eps=eps)
        self.norm = _gemma_rms_norm(hidden_size, eps=eps)

    def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor:
        return self.hidden_norm(context_states)

    def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor:
        return torch.stack(
            [layer.self_attn.k_norm(k) for layer, k in zip(self.layers, all_k)]
        )

    def _context_kv_writer(self) -> DSparkContextKVWriter | None:
        """The fused context K/V writer, or None where the base path runs."""
        if not hasattr(self, "_ctx_kv_writer"):
            if any(a.kv_cache.numel() == 0 for a in self._attn_layers):
                return None  # KV caches not bound yet: decide once they are
            rotary_dim = self._rope_cos_sin_cache.shape[-1]
            supported = (
                self._rope_is_neox
                and self.layers[0].self_attn.v_scale is None
                and rotary_dims_supported(rotary_dim, self._head_dim)
                and self._rope_cos_sin_cache.dtype == torch.bfloat16
                and DSparkContextKVWriter.supported(self._attn_layers)
            )
            self._ctx_kv_writer = (
                DSparkContextKVWriter(self._attn_layers) if supported else None
            )
        return self._ctx_kv_writer

    def precompute_and_store_context_kv(
        self,
        context_states: torch.Tensor,
        context_positions: torch.Tensor,
        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
    ) -> None:
        """One kernel norms, ropes and stores the context K/V of every layer,
        straight from the projection output."""
        slot_mapping = _shared_slot_mapping(context_slot_mapping)
        writer = (
            self._context_kv_writer()
            if slot_mapping is not None
            and hasattr(self, "_num_attn_layers")
            and context_states.dtype == torch.bfloat16
            else None
        )
        if writer is None:
            super().precompute_and_store_context_kv(
                context_states, context_positions, context_slot_mapping
            )
            return
        all_kv = self._project_context_kv_token_major(
            context_states,
            context_states.shape[0],
            self._num_attn_layers,
            self._num_kv_heads,
            self._head_dim,
        )
        writer(
            all_kv,
            self._k_norm_weights,
            self._rms_norm_eps,
            context_positions,
            self._rope_cos_sin_cache,
            self._rope_cos_sin_cache.shape[-1],
            slot_mapping,
        )

_context_kv_writer()

The fused context K/V writer, or None where the base path runs.

Source code in vllm/models/minimax_m3/common/dspark.py

def _context_kv_writer(self) -> DSparkContextKVWriter | None:
    """The fused context K/V writer, or None where the base path runs."""
    if not hasattr(self, "_ctx_kv_writer"):
        if any(a.kv_cache.numel() == 0 for a in self._attn_layers):
            return None  # KV caches not bound yet: decide once they are
        rotary_dim = self._rope_cos_sin_cache.shape[-1]
        supported = (
            self._rope_is_neox
            and self.layers[0].self_attn.v_scale is None
            and rotary_dims_supported(rotary_dim, self._head_dim)
            and self._rope_cos_sin_cache.dtype == torch.bfloat16
            and DSparkContextKVWriter.supported(self._attn_layers)
        )
        self._ctx_kv_writer = (
            DSparkContextKVWriter(self._attn_layers) if supported else None
        )
    return self._ctx_kv_writer

precompute_and_store_context_kv(context_states, context_positions, context_slot_mapping=None)

One kernel norms, ropes and stores the context K/V of every layer, straight from the projection output.

Source code in vllm/models/minimax_m3/common/dspark.py

def precompute_and_store_context_kv(
    self,
    context_states: torch.Tensor,
    context_positions: torch.Tensor,
    context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
) -> None:
    """One kernel norms, ropes and stores the context K/V of every layer,
    straight from the projection output."""
    slot_mapping = _shared_slot_mapping(context_slot_mapping)
    writer = (
        self._context_kv_writer()
        if slot_mapping is not None
        and hasattr(self, "_num_attn_layers")
        and context_states.dtype == torch.bfloat16
        else None
    )
    if writer is None:
        super().precompute_and_store_context_kv(
            context_states, context_positions, context_slot_mapping
        )
        return
    all_kv = self._project_context_kv_token_major(
        context_states,
        context_states.shape[0],
        self._num_attn_layers,
        self._num_kv_heads,
        self._head_dim,
    )
    writer(
        all_kv,
        self._k_norm_weights,
        self._rms_norm_eps,
        context_positions,
        self._rope_cos_sin_cache,
        self._rope_cos_sin_cache.shape[-1],
        slot_mapping,
    )

_gemma_rms_norm(hidden_size, eps)

The target's Gemma RMSNorm (fused FlashInfer kernels on CUDA).

vLLM's GemmaRMSNorm hands the IR op an fp32 1 + w weight, which the vllm_c kernel rejects for bf16 inputs, so every call ran the unfused native path (about ten kernels).

Source code in vllm/models/minimax_m3/common/dspark.py

def _gemma_rms_norm(hidden_size: int, eps: float) -> nn.Module:
    """The target's Gemma RMSNorm (fused FlashInfer kernels on CUDA).

    vLLM's GemmaRMSNorm hands the IR op an fp32 ``1 + w`` weight, which the
    vllm_c kernel rejects for bf16 inputs, so every call ran the unfused native
    path (about ten kernels).
    """
    if current_platform.is_rocm():
        from vllm.models.minimax_m3.amd.model import MiniMAXGemmaRMSNorm
    else:
        from vllm.models.minimax_m3.nvidia.model import MiniMAXGemmaRMSNorm
    return MiniMAXGemmaRMSNorm(hidden_size, eps=eps)

_shared_slot_mapping(slot_mapping)

The context slot mapping when every draft layer shares it, else None (dummy runs, and drafts whose layers span several KV cache groups).

Source code in vllm/models/minimax_m3/common/dspark.py

def _shared_slot_mapping(
    slot_mapping: torch.Tensor | list[torch.Tensor | None] | None,
) -> torch.Tensor | None:
    """The context slot mapping when every draft layer shares it, else None
    (dummy runs, and drafts whose layers span several KV cache groups)."""
    if slot_mapping is None or isinstance(slot_mapping, torch.Tensor):
        return slot_mapping
    first = slot_mapping[0]
    if first is None or any(
        s is None or s.data_ptr() != first.data_ptr() for s in slot_mapping
    ):
        return None
    return first

dspark_qknorm_rope_kv_insert(qkv, q_norm_weight, k_norm_weight, cos_sin_cache, positions, num_heads, num_kv_heads, rotary_dim, eps, q_fp8_out, layer_name)

The fused QK-norm/RoPE on qkv, also storing the draft tokens' k/v.

A custom op so the slot mapping is read when it runs: read in the compiled forward, it is fixed at trace time, when the draft layers have none.

Source code in vllm/models/minimax_m3/common/dspark.py

def dspark_qknorm_rope_kv_insert(
    qkv: torch.Tensor,
    q_norm_weight: torch.Tensor,
    k_norm_weight: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    positions: torch.Tensor,
    num_heads: int,
    num_kv_heads: int,
    rotary_dim: int,
    eps: float,
    q_fp8_out: torch.Tensor | None,
    layer_name: LayerNameType,
) -> None:
    """The fused QK-norm/RoPE on ``qkv``, also storing the draft tokens' k/v.

    A custom op so the slot mapping is read when it runs: read in the compiled
    forward, it is fixed at trace time, when the draft layers have none.
    """
    _, layer, kv_cache, slot_mapping = get_attention_context(
        _resolve_layer_name(layer_name)
    )
    insert = slot_mapping is not None
    ops.fused_minimax_m3_qknorm_rope_kv_insert(
        qkv,
        q_norm_weight,
        k_norm_weight,
        cos_sin_cache,
        positions,
        num_heads,
        num_kv_heads,
        rotary_dim,
        eps,
        slot_mapping=slot_mapping,
        kv_cache=kv_cache if insert else None,
        block_size=kv_cache.size(2) if insert else 0,
        kv_cache_dtype="fp8" if insert and kv_cache.dtype == torch.uint8 else "auto",
        q_fp8_out=q_fp8_out,
        q_fp8_scale=layer._q_scale_float,
    )