diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 74e9e7739f64..c900c291245a 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -443,6 +443,7 @@ th {
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
+| `MiniMaxM3SparseForCausalLM` | MiniMax-M3 | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ |
| `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
| `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
| `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ |
@@ -593,7 +594,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + IE+ + VE+ + A+ | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
| `MiniCPMO` | MiniCPM-O | T + IE+ + VE+ + AE+ | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
| `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
-| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I+ + V+ | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | |
+| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I+ + V+ | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | ✅︎ |
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + IE+ | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
| `MolmoForCausalLM` | Molmo | T + I+ | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py
index c171b4bfe8c5..4f3528806c9d 100644
--- a/vllm/models/minimax_m3/amd/model.py
+++ b/vllm/models/minimax_m3/amd/model.py
@@ -30,7 +30,7 @@
VllmConfig,
get_current_vllm_config,
)
-from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
from vllm.forward_context import get_forward_context
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -64,12 +64,15 @@
MultiModalEmbeddings,
SupportsEagle3,
SupportsMultiModal,
+ SupportsPP,
)
from vllm.model_executor.models.utils import (
AutoWeightsLoader,
+ PPMissingLayer,
WeightsMapper,
init_vllm_registered_model,
is_pp_missing_parameter,
+ make_empty_intermediate_tensors_factory,
make_layers,
maybe_prefix,
)
@@ -92,6 +95,7 @@
)
from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
@@ -772,12 +776,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
self.vocab_size = config.vocab_size
- self.embed_tokens = VocabParallelEmbedding(
- config.vocab_size,
- config.hidden_size,
- quant_config=quant_config,
- prefix=f"{prefix}.embed_tokens",
- )
+ if get_pp_group().is_first_rank:
+ self.embed_tokens = VocabParallelEmbedding(
+ config.vocab_size,
+ config.hidden_size,
+ quant_config=quant_config,
+ prefix=f"{prefix}.embed_tokens",
+ )
+ else:
+ self.embed_tokens = PPMissingLayer()
# Reserved top-k indices buffer shared by all sparse-attention indexer
# layers (mirrors DeepseekV4); the indexer writes its per-head decode/
@@ -807,7 +814,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
prefix=f"{prefix}.layers",
)
- self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ if get_pp_group().is_last_rank:
+ self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ else:
+ self.norm = PPMissingLayer()
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+ ["hidden_states", "residual"], config.hidden_size
+ )
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
@@ -816,13 +829,19 @@ def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None,
inputs_embeds: torch.Tensor | None = None,
- ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
- if inputs_embeds is not None:
- hidden_states = inputs_embeds
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
+ if get_pp_group().is_first_rank:
+ if inputs_embeds is not None:
+ hidden_states = inputs_embeds
+ else:
+ hidden_states = self.embed_input_ids(input_ids)
+ residual = None
else:
- hidden_states = self.embed_input_ids(input_ids)
- residual = None
+ assert intermediate_tensors is not None
+ hidden_states = intermediate_tensors["hidden_states"]
+ residual = intermediate_tensors["residual"]
# EAGLE3 is not yet compatible with pipeline parallel
aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
@@ -832,6 +851,11 @@ def forward(
aux_hidden_states, idx + 1, hidden_states, residual
)
+ if not get_pp_group().is_last_rank:
+ return IntermediateTensors(
+ {"hidden_states": hidden_states, "residual": residual}
+ )
+
hidden_states, _ = self.norm(hidden_states, residual)
if len(aux_hidden_states) > 0:
@@ -946,7 +970,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
return loaded_params
-class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
+class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
"""MiniMax M3 (sparse/dense backbone) for causal language modeling."""
packed_modules_mapping = {
@@ -963,13 +987,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
self.model = MiniMaxM3Model(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
- self.lm_head = ParallelLMHead(
- config.vocab_size,
- config.hidden_size,
- quant_config=quant_config,
- prefix=maybe_prefix(prefix, "lm_head"),
- )
+ if get_pp_group().is_last_rank:
+ self.lm_head = ParallelLMHead(
+ config.vocab_size,
+ config.hidden_size,
+ quant_config=quant_config,
+ prefix=maybe_prefix(prefix, "lm_head"),
+ )
+ else:
+ self.lm_head = PPMissingLayer()
self.logits_processor = LogitsProcessor(config.vocab_size)
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
+ self.model.make_empty_intermediate_tensors
+ )
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
@@ -978,10 +1008,11 @@ def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs,
- ) -> torch.Tensor:
- return self.model(input_ids, positions, inputs_embeds)
+ ) -> torch.Tensor | IntermediateTensors:
+ return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
return self.logits_processor(self.lm_head, hidden_states)
@@ -1004,7 +1035,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
dummy_inputs=MiniMaxM3VLDummyInputsBuilder,
)
class MiniMaxM3SparseForConditionalGeneration(
- nn.Module, SupportsMultiModal, SupportsEagle3
+ nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3
):
"""Top-level (VL) entry point for MiniMax M3.
@@ -1070,6 +1101,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
prefix=maybe_prefix(prefix, "language_model"),
architectures=["MiniMaxM3SparseForCausalLM"],
)
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
+ self.language_model.make_empty_intermediate_tensors
+ )
def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None:
pixel_values = kwargs.pop("pixel_values", None)
@@ -1179,10 +1213,13 @@ def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
- return self.language_model(input_ids, positions, inputs_embeds)
+ return self.language_model(
+ input_ids, positions, intermediate_tensors, inputs_embeds
+ )
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
return self.language_model.compute_logits(hidden_states)
diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py
index c27ded3b83ef..a30bc335fcfb 100644
--- a/vllm/models/minimax_m3/nvidia/model.py
+++ b/vllm/models/minimax_m3/nvidia/model.py
@@ -21,7 +21,7 @@
from vllm import _custom_ops as ops
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
-from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
from vllm.forward_context import get_forward_context
from vllm.model_executor.layers.activation import SiluAndMulWithClamp
from vllm.model_executor.layers.attention import Attention
@@ -56,12 +56,15 @@
MultiModalEmbeddings,
SupportsEagle3,
SupportsMultiModal,
+ SupportsPP,
)
from vllm.model_executor.models.utils import (
AutoWeightsLoader,
+ PPMissingLayer,
WeightsMapper,
init_vllm_registered_model,
is_pp_missing_parameter,
+ make_empty_intermediate_tensors_factory,
make_layers,
maybe_prefix,
)
@@ -79,6 +82,7 @@
)
from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
@@ -754,12 +758,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
self.vocab_size = config.vocab_size
- self.embed_tokens = VocabParallelEmbedding(
- config.vocab_size,
- config.hidden_size,
- quant_config=quant_config,
- prefix=f"{prefix}.embed_tokens",
- )
+ if get_pp_group().is_first_rank:
+ self.embed_tokens = VocabParallelEmbedding(
+ config.vocab_size,
+ config.hidden_size,
+ quant_config=quant_config,
+ prefix=f"{prefix}.embed_tokens",
+ )
+ else:
+ self.embed_tokens = PPMissingLayer()
# Reserved top-k indices buffer shared by all sparse-attention indexer
# layers (mirrors DeepseekV4); kept at a stable address so the indexer's
@@ -792,7 +799,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
prefix=f"{prefix}.layers",
)
- self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ if get_pp_group().is_last_rank:
+ self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ else:
+ self.norm = PPMissingLayer()
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+ ["hidden_states", "residual"], config.hidden_size
+ )
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
@@ -801,13 +814,19 @@ def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None,
inputs_embeds: torch.Tensor | None = None,
- ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
- if inputs_embeds is not None:
- hidden_states = inputs_embeds
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
+ if get_pp_group().is_first_rank:
+ if inputs_embeds is not None:
+ hidden_states = inputs_embeds
+ else:
+ hidden_states = self.embed_input_ids(input_ids)
+ residual = None
else:
- hidden_states = self.embed_input_ids(input_ids)
- residual = None
+ assert intermediate_tensors is not None
+ hidden_states = intermediate_tensors["hidden_states"]
+ residual = intermediate_tensors["residual"]
# EAGLE3 is not yet compatible with pipeline parallel
aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
@@ -817,6 +836,11 @@ def forward(
aux_hidden_states, idx + 1, hidden_states, residual
)
+ if not get_pp_group().is_last_rank:
+ return IntermediateTensors(
+ {"hidden_states": hidden_states, "residual": residual}
+ )
+
hidden_states, _ = self.norm(hidden_states, residual)
if len(aux_hidden_states) > 0:
@@ -931,7 +955,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
return loaded_params
-class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
+class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
"""MiniMax M3 (sparse/dense backbone) for causal language modeling."""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -943,13 +967,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
self.model = MiniMaxM3Model(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
- self.lm_head = ParallelLMHead(
- config.vocab_size,
- config.hidden_size,
- quant_config=quant_config,
- prefix=maybe_prefix(prefix, "lm_head"),
- )
+ if get_pp_group().is_last_rank:
+ self.lm_head = ParallelLMHead(
+ config.vocab_size,
+ config.hidden_size,
+ quant_config=quant_config,
+ prefix=maybe_prefix(prefix, "lm_head"),
+ )
+ else:
+ self.lm_head = PPMissingLayer()
self.logits_processor = LogitsProcessor(config.vocab_size)
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
+ self.model.make_empty_intermediate_tensors
+ )
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
@@ -958,10 +988,11 @@ def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs,
- ) -> torch.Tensor:
- return self.model(input_ids, positions, inputs_embeds)
+ ) -> torch.Tensor | IntermediateTensors:
+ return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
return self.logits_processor(self.lm_head, hidden_states)
@@ -980,7 +1011,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
dummy_inputs=MiniMaxM3VLDummyInputsBuilder,
)
class MiniMaxM3SparseForConditionalGeneration(
- nn.Module, SupportsMultiModal, SupportsEagle3
+ nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3
):
"""Top-level (VL) entry point for MiniMax M3.
@@ -1042,6 +1073,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
prefix=maybe_prefix(prefix, "language_model"),
architectures=["MiniMaxM3SparseForCausalLM"],
)
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
+ self.language_model.make_empty_intermediate_tensors
+ )
# Expose language model / lm_head for EAGLE3 spec decode.
@property
@@ -1160,10 +1194,13 @@ def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
- return self.language_model(input_ids, positions, inputs_embeds)
+ return self.language_model(
+ input_ids, positions, intermediate_tensors, inputs_embeds
+ )
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
return self.language_model.compute_logits(hidden_states)