model.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,892 lines 78.1 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""Inference-only MiniMax M3 (text backbone) model — AMD ROCm implementation.
4
5	Self-contained per-platform impl (mirrors ``deepseek_v4/amd``). It is identical
6	to ``../nvidia/model.py`` except for RMS normalization: FlashInfer's Gemma
7	RMSNorm kernels are CUDA-only, so ``MiniMAXGemmaRMSNorm`` here uses a native
8	(FlashInfer-free) implementation.
9
10	The MiniMax-M3-preview config selects a single set of branches:
11	* qk_norm_type == "per_head"
12	* hidden_act == "swigluoai"
13	* use_gemma_norm == True -> Gemma-style RMSNorm everywhere
14	* attention_output_gate == False
15	* scoring_func == "sigmoid" with a routing-bias correction term
16	* sparse_attention_config present -> a subset of layers run the extra
17	"index" attention branch.
18	"""
19
20	from collections.abc import Iterable
21
22	import torch
23	from torch import nn
24	from transformers import PreTrainedConfig
25
26	from vllm import _custom_ops as ops
27	from vllm import envs
28	from vllm._aiter_ops import rocm_aiter_ops
29	from vllm.compilation.breakable_cudagraph import eager_break_during_capture
30	from vllm.config import (
31	CacheConfig,
32	VllmConfig,
33	get_current_vllm_config,
34	)
35	from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36	from vllm.forward_context import get_forward_context
37	from vllm.logger import init_logger
38	from vllm.model_executor.layers.attention import Attention
39	from vllm.model_executor.layers.attention.attention import set_default_quant_scales
40	from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
41	from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
42	fused_allreduce_gemma_rms_norm,
43	)
44	from vllm.model_executor.layers.fused_moe import (
45	FusedMoEFactory,
46	GateLinear,
47	fused_moe_make_expert_params_mapping,
48	)
49	from vllm.model_executor.layers.fused_moe.utils import (
50	is_model_fused_shared_expert_compatible,
51	)
52	from vllm.model_executor.layers.linear import (
53	MergedColumnParallelLinear,
54	MinimaxM3QKVParallelLinearWithIndexer,
55	QKVParallelLinear,
56	RowParallelLinear,
57	)
58	from vllm.model_executor.layers.logits_processor import LogitsProcessor
59	from vllm.model_executor.layers.quantization import QuantizationConfig
60	from vllm.model_executor.layers.quantization.utils.config_utils import (
61	is_shared_expert_quant_fse_compatible,
62	)
63	from vllm.model_executor.layers.rotary_embedding import get_rope
64	from vllm.model_executor.layers.vocab_parallel_embedding import (
65	ParallelLMHead,
66	VocabParallelEmbedding,
67	)
68	from vllm.model_executor.model_loader.weight_utils import (
69	default_weight_loader,
70	maybe_remap_kv_scale_name,
71	)
72	from vllm.model_executor.models.interfaces import (
73	EagleModelMixin,
74	MultiModalEmbeddings,
75	SupportsEagle3,
76	SupportsMultiModal,
77	SupportsPP,
78	)
79	from vllm.model_executor.models.utils import (
80	AutoWeightsLoader,
81	PPMissingLayer,
82	WeightsMapper,
83	init_vllm_registered_model,
84	is_pp_missing_parameter,
85	make_empty_intermediate_tensors_factory,
86	make_layers,
87	maybe_prefix,
88	spec_decode_needs_target_embed,
89	)
90	from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
91	from vllm.models.minimax_m3.amd.indexer_aiter import (
92	MiniMaxM3AiterIndexer,
93	select_aiter_indexer_impl_cls,
94	)
95	from vllm.models.minimax_m3.amd.ops import (
96	gemma_fused_add_rmsnorm,
97	gemma_rmsnorm,
98	swiglu_oai_split,
99	)
100	from vllm.models.minimax_m3.amd.ops.sparse_pa import (
101	ASM_PAGE_SIZE,
102	minimax_m3_alloc_sparse_block_table,
103	minimax_m3_sparse_block_page_stride,
104	)
105	from vllm.models.minimax_m3.amd.sparse_attention_msa import (
106	MiniMaxM3SparseAiterPADecodeMetadata,
107	MiniMaxM3SparseAiterPAImpl,
108	MiniMaxM3SparseAiterPAPrefillMetadata,
109	)
110	from vllm.models.minimax_m3.common.encoder_cudagraph import (
111	MiniMaxM3EncoderCudaGraphMixin,
112	)
113	from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer
114	from vllm.models.minimax_m3.common.mm_preprocess import (
115	MiniMaxM3VLDummyInputsBuilder,
116	MiniMaxM3VLMultiModalProcessor,
117	MiniMaxM3VLProcessingInfo,
118	)
119	from vllm.models.minimax_m3.common.sparse_attention import (
120	MiniMaxM3SparseBackend,
121	MiniMaxM3SparseImpl,
122	MiniMaxM3SparseMetadata,
123	minimax_m3_rebase_slots_to_page16,
124	minimax_m3_use_aiter_sparse_pa,
125	select_main_backend_and_impl_cls,
126	)
127	from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
128	from vllm.multimodal import MULTIMODAL_REGISTRY
129	from vllm.sequence import IntermediateTensors
130	from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
131	from vllm.v1.kv_cache_interface import (
132	FullAttentionSpec,
133	KVCacheSpec,
134	get_kv_quant_mode,
135	is_quantized_kv_cache,
136	)
137
138	logger = init_logger(__name__)
139
140
141	def _sparse_attention_layer_ids(config: PreTrainedConfig) -> set[int]:
142	"""Layer ids whose attention runs the extra sparse "index" branch."""
143	cfg = getattr(config, "sparse_attention_config", None)
144	if not cfg:
145	return set()
146	freq = cfg.get("sparse_attention_freq")
147	if freq is None:
148	return set()
149	return {i for i, f in enumerate(freq) if f != 0}
150
151
152	def _sparse_attention_layer_ordinals(config: PreTrainedConfig) -> dict[int, int]:
153	"""Map each sparse-attention layer id to its ordinal among sparse layers."""
154	return {
155	lid: ordinal
156	for ordinal, lid in enumerate(sorted(_sparse_attention_layer_ids(config)))
157	}
158
159
160	def _should_skip_index_topk(config: PreTrainedConfig, layer_id: int) -> bool:
161	"""ATOM ``index_topk_freq`` (cross-layer index sharing).
162
163	Only 1 of every ``index_topk_freq`` sparse-attention layers recomputes the
164	lightning-indexer top-k block selection; the rest reuse the selection the
165	preceding compute layer wrote into the shared ``topk_indices_buffer`` this
166	same forward pass. This cuts the indexer score + top-k cost ~``freq``x with
167	negligible accuracy impact (adjacent sparse layers pick nearly the same
168	blocks; ATOM validated GSM8K with freq=4). Gated by ``use_index_cache``;
169	enable via ``--hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}'``.
170	"""
171	if not getattr(config, "use_index_cache", False):
172	return False
173	freq = int(getattr(config, "index_topk_freq", 1) or 1)
174	if freq <= 1:
175	return False
176	ordinal = _sparse_attention_layer_ordinals(config).get(layer_id)
177	if ordinal is None:
178	return False
179	offset = int(getattr(config, "index_skip_topk_offset", 0) or 0)
180	return max(ordinal - offset, 0) % freq != 0
181
182
183	def _is_moe_layer(config: PreTrainedConfig, layer_id: int) -> bool:
184	"""Whether this layer's MLP is a sparse MoE block (vs a dense MLP)."""
185	moe_layer_freq = getattr(config, "moe_layer_freq", None)
186	if moe_layer_freq is None:
187	return True
188	return moe_layer_freq[layer_id] != 0
189
190
191	def _build_rotary_emb(config: PreTrainedConfig, head_dim: int):
192	"""Build the (partial NeoX) RoPE, honoring an optional ``rope_scaling`` config.
193
194	Without scaling the cos/sin cache is sized to ``max_position_embeddings``
195	(524288 native); a request whose positions exceed that reads the cache out of
196	bounds and the worker hard-crashes (no Python traceback). When ``rope_scaling``
197	is set (e.g. YaRN ``factor: 2`` to reach 1M), thread it into ``get_rope`` so the
198	proper scaled embedding is built and its cache covers
199	``original_max_position_embeddings * factor`` positions. Default behavior
200	(no scaling) is unchanged. Shared by the dense and sparse attention layers, and
201	the index branch reuses the returned module.
202
203	Note: for the VL checkpoint, set ``rope_scaling`` on the text config
204	(``--hf-overrides '{"text_config":{"rope_scaling":{...}}}'``) -- that is the
205	config the decoder reads here; a top-level override does not reach it.
206	"""
207	rope_parameters = {
208	"rope_theta": config.rope_theta,
209	"partial_rotary_factor": config.partial_rotary_factor,
210	}
211	max_position = config.max_position_embeddings
212	rope_scaling = getattr(config, "rope_scaling", None)
213	if rope_scaling:
214	rope_parameters.update(rope_scaling)
215	# HF uses "rope_type" (older configs: "type"); get_rope reads "rope_type".
216	if "rope_type" not in rope_parameters and "type" in rope_scaling:
217	rope_parameters["rope_type"] = rope_scaling["type"]
218	rope_parameters.setdefault(
219	"original_max_position_embeddings", config.max_position_embeddings
220	)
221	factor = float(rope_scaling.get("factor", 1.0))
222	# Cover the extended range (informational for get_rope's default branch;
223	# the YaRN embedding sizes its own cache from original * factor).
224	max_position = int(rope_parameters["original_max_position_embeddings"] * factor)
225	return get_rope(
226	head_dim,
227	max_position=max_position,
228	rope_parameters=rope_parameters,
229	)
230
231
232	class MiniMAXGemmaRMSNorm(nn.Module):
233	"""Gemma-style RMS normalization (native ROCm implementation).
234
235	Normalizes in fp32 and scales by ``(1 + weight)`` — numerically equivalent
236	to the FlashInfer ``gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm`` kernels
237	used in the NVIDIA path, which are unavailable on ROCm. When ``residual`` is
238	given, the fused add + norm returns the updated ``(normed, residual)`` pair.
239
240	The fp32 normalize + scale + (optional) residual-add run in a single fused
241	Triton pass (``amd.ops.gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm``) instead
242	of a chain of elementwise PyTorch kernels.
243	"""
244
245	def __init__(
246	self,
247	hidden_size: int,
248	eps: float = 1e-6,
249	) -> None:
250	super().__init__()
251	self.weight = nn.Parameter(torch.zeros(hidden_size))
252	self.variance_epsilon = eps
253
254	def forward(
255	self,
256	x: torch.Tensor,
257	residual: torch.Tensor \| None = None,
258	) -> torch.Tensor \| tuple[torch.Tensor, torch.Tensor]:
259	if residual is None:
260	return gemma_rmsnorm(x, self.weight, self.variance_epsilon)
261	return gemma_fused_add_rmsnorm(x, residual, self.weight, self.variance_epsilon)
262
263
264	class MiniMaxM3MLP(nn.Module):
265	"""Dense SwiGLU-OAI MLP (used by the leading dense layers)."""
266
267	def __init__(
268	self,
269	config: PreTrainedConfig,
270	intermediate_size: int,
271	quant_config: QuantizationConfig \| None = None,
272	reduce_results: bool = True,
273	prefix: str = "",
274	) -> None:
275	super().__init__()
276	self.gate_up_proj = MergedColumnParallelLinear(
277	config.hidden_size,
278	[intermediate_size] * 2,
279	bias=False,
280	quant_config=quant_config,
281	prefix=f"{prefix}.gate_up_proj",
282	)
283	self.down_proj = RowParallelLinear(
284	intermediate_size,
285	config.hidden_size,
286	bias=False,
287	quant_config=quant_config,
288	reduce_results=reduce_results,
289	prefix=f"{prefix}.down_proj",
290	)
291	if config.hidden_act != "swigluoai":
292	raise ValueError(
293	f"Unsupported activation: {config.hidden_act}. "
294	"Only swigluoai is supported."
295	)
296	# gate * sigmoid(alpha * gate) * (up + beta), with both halves clamped.
297	# Kept as our fp32 Triton kernel (not the #22 SWIGLUOAI_UNINTERLEAVE op
298	# ``silu_and_mul_with_clamp``): that op IS built on ROCm but rounds
299	# intermediates to bf16 (rel ~3e-3 vs our fp32 ~1e-6), which costs gsm8k
300	# accuracy since this activation feeds the MXFP8 quant + MoE.
301	self.swiglu_alpha = config.swiglu_alpha
302	self.swiglu_beta = config.swiglu_beta
303	self.swiglu_limit = config.swiglu_limit
304
305	def forward(self, x: torch.Tensor) -> torch.Tensor:
306	gate_up, _ = self.gate_up_proj(x)
307	x = swiglu_oai_split(
308	gate_up,
309	alpha=self.swiglu_alpha,
310	beta=self.swiglu_beta,
311	limit=self.swiglu_limit,
312	)
313	x, _ = self.down_proj(x)
314	return x
315
316
317	def _aiter_moe_fused_shared_experts_enabled(
318	is_fused_shared_expert_enabled: bool,
319	) -> bool:
320	"""Whether the fused shared expert routes through aiter's grouped top-k MoE.
321
322	A strict sub-case of `is_fused_shared_expert_enabled`: shared-expert
323	fusion must already be opted in (``VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS``)
324	and allowed (not under expert parallelism). When additionally on gfx950 with
325	an active aiter MoE backend, the shared expert is appended inside aiter's
326	biased grouped top-k kernel (``num_fused_shared_experts``) instead of the
327	vLLM router's torch concat. Otherwise FSE still runs via the vLLM top-k bias
328	router.
329	"""
330	from vllm.platforms.rocm import on_gfx950
331
332	return (
333	on_gfx950()
334	and is_fused_shared_expert_enabled
335	and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
336	)
337
338
339	class MiniMaxM3MoE(nn.Module):
340	"""Sigmoid-routed MoE block with a routing-bias correction and a shared
341	expert."""
342
343	def __init__(
344	self,
345	config: PreTrainedConfig,
346	layer_id: int,
347	quant_config: QuantizationConfig \| None = None,
348	prefix: str = "",
349	) -> None:
350	super().__init__()
351	self.tp_size = get_tensor_model_parallel_world_size()
352	if self.tp_size > config.num_local_experts:
353	raise ValueError(
354	f"Tensor parallel size {self.tp_size} is greater than "
355	f"the number of experts {config.num_local_experts}."
356	)
357
358	self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
359	self.n_shared_experts = getattr(config, "n_shared_experts", None)
360
361	# Sigmoid routing uses a per-expert score-correction bias for selection.
362	self.use_routing_bias = getattr(config, "use_routing_bias", False)
363	if self.use_routing_bias:
364	self.e_score_correction_bias = nn.Parameter(
365	torch.empty(config.num_local_experts, dtype=torch.float32)
366	)
367	self.e_score_correction_bias.weight_loader = (
368	MiniMaxM3MoE.ebias_weight_loader
369	)
370	else:
371	self.e_score_correction_bias = None
372
373	# Router weights are stored in fp32; GateLinear upcasts the bf16
374	# activations and computes the gate in fp32 (fp32 router logits).
375	self.gate = GateLinear(
376	config.hidden_size,
377	config.num_local_experts,
378	bias=False,
379	params_dtype=torch.float32,
380	out_dtype=torch.float32,
381	prefix=f"{prefix}.gate",
382	)
383
384	# Shared-expert fusion (opt-in via VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS,
385	# off under expert parallelism) folds the shared expert into the routed
386	# MoE call as the last expert slot, so we don't build a separate module.
387	# On gfx950 with aiter MoE the append is fused inside aiter's grouped
388	# top-k kernel; otherwise it goes through the vLLM top-k bias router.
389	# It is disabled under expert parallelism (the shared slot is appended to
390	# the routed top-k, which the EP expert-mapping path does not handle).
391	# TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1`
392	# is checked to enable FSE for MiniMax-M3, despite AITER not being used.
393	# This should be cleaned up and use `resolve_layer_fused_shared_expert`.
394	fse_requested = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
395
396	self.is_fused_shared_expert_enabled = False
397	if (
398	fse_requested
399	and bool(getattr(config, "n_shared_experts", None))
400	and not get_current_vllm_config().parallel_config.enable_expert_parallel
401	):
402	fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
403	quant_config,
404	f"{prefix}.experts",
405	f"{prefix}.shared_experts",
406	)
407	if not fse_compatible:
408	logger.warning(
409	"VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
410	"cannot be enabled: %s.",
411	fse_reason,
412	)
413	else:
414	self.is_fused_shared_expert_enabled = True
415
416	# When additionally on gfx950 with an active aiter MoE backend, the shared
417	# expert is appended inside aiter's an active aiter MoE backend, the shared
418	# expert is appended inside aiter's biased grouped top-k kernel
419	# (``num_fused_shared_experts``) instead of the vLLM router's torch concat.
420	# Otherwise FSE still runs via the vLLM top-k bias router.
421	# TODO: `on_gfx950()` check here should not be MiniMax-M3 specific, and
422	# the check should be done on resolved MOE backend directly.
423	self.use_aiter_moe_fse = _aiter_moe_fused_shared_experts_enabled(
424	self.is_fused_shared_expert_enabled
425	)
426
427	self.shared_experts: MiniMaxM3MLP \| None = None
428	if self.n_shared_experts and not self.is_fused_shared_expert_enabled:
429	self.shared_experts = MiniMaxM3MLP(
430	config=config,
431	intermediate_size=config.intermediate_size * self.n_shared_experts,
432	quant_config=quant_config,
433	reduce_results=False,
434	prefix=f"{prefix}.shared_experts",
435	)
436
437	# The aiter MoE fused path goes through aiter's biased grouped top-k
438	# (GroupedTopKRouter, as in DeepSeek-V4): M3 is not group-routed, so a
439	# trivial single group (num_expert_group=topk_group=1) reduces to plain
440	# top-k while applying the sigmoid + bias correction and appending the
441	# always-on shared expert; aiter applies the routed scaling internally.
442	# Every other path (vLLM top-k bias router, or no fusion) applies the
443	# routed scaling to the MoE output here.
444	self.experts = FusedMoEFactory(
445	num_experts=config.num_local_experts,
446	top_k=config.num_experts_per_tok,
447	hidden_size=config.hidden_size,
448	intermediate_size=config.intermediate_size,
449	intermediate_pad=0,
450	scoring_func=config.scoring_func,
451	e_score_correction_bias=self.e_score_correction_bias,
452	renormalize=True,
453	use_grouped_topk=self.use_aiter_moe_fse,
454	num_expert_group=1 if self.use_aiter_moe_fse else None,
455	topk_group=1 if self.use_aiter_moe_fse else None,
456	activation="swigluoai_uninterleave",
457	swiglu_limit=config.swiglu_limit,
458	swiglu_alpha=config.swiglu_alpha,
459	swiglu_beta=config.swiglu_beta,
460	routed_scaling_factor=self.routed_scaling_factor,
461	apply_routed_scale_to_output=not self.use_aiter_moe_fse,
462	router_logits_dtype=self.gate.out_dtype,
463	shared_experts=self.shared_experts,
464	n_shared_experts=(
465	self.n_shared_experts if self.is_fused_shared_expert_enabled else None
466	),
467	fuse_shared_experts=self.is_fused_shared_expert_enabled,
468	quant_config=quant_config,
469	prefix=f"{prefix}.experts",
470	)
471
472	@staticmethod
473	def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
474	assert param.size() == loaded_weight.size()
475	param.data.copy_(loaded_weight.to(torch.float32))
476
477	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
478	num_tokens, hidden_dim = hidden_states.shape
479	hidden_states = hidden_states.view(-1, hidden_dim)
480
481	# router_logits: (num_tokens, n_experts); GateLinear casts to fp32.
482	router_logits, _ = self.gate(hidden_states)
483	final_hidden_states = self.experts(
484	hidden_states=hidden_states, router_logits=router_logits
485	)
486
487	return final_hidden_states.view(num_tokens, hidden_dim)
488
489
490	class MiniMaxM3Attention(nn.Module):
491	"""Dense attention with per-head QK norm and partial RoPE."""
492
493	def __init__(
494	self,
495	config: PreTrainedConfig,
496	layer_id: int,
497	quant_config: QuantizationConfig \| None = None,
498	prefix: str = "",
499	cache_config: CacheConfig \| None = None,
500	) -> None:
501	super().__init__()
502	self.hidden_size = config.hidden_size
503	tp_size = get_tensor_model_parallel_world_size()
504
505	self.total_num_heads = config.num_attention_heads
506	assert self.total_num_heads % tp_size == 0
507	self.num_heads = self.total_num_heads // tp_size
508	self.total_num_kv_heads = config.num_key_value_heads
509	if self.total_num_kv_heads >= tp_size:
510	assert self.total_num_kv_heads % tp_size == 0
511	else:
512	assert tp_size % self.total_num_kv_heads == 0
513	self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
514	self.head_dim = config.head_dim
515	self.q_size = self.num_heads * self.head_dim
516	self.kv_size = self.num_kv_heads * self.head_dim
517	self.scaling = self.head_dim**-0.5
518
519	self.qkv_proj = QKVParallelLinear(
520	self.hidden_size,
521	self.head_dim,
522	self.total_num_heads,
523	self.total_num_kv_heads,
524	bias=False,
525	quant_config=quant_config,
526	prefix=f"{prefix}.qkv_proj",
527	)
528	# reduce_results=False: the attention all-reduce is fused with the
529	# following post_attention_layernorm (GemmaRMSNorm) in the decoder layer
530	# via fused_allreduce_gemma_rms_norm.
531	self.o_proj = RowParallelLinear(
532	self.total_num_heads * self.head_dim,
533	self.hidden_size,
534	bias=False,
535	reduce_results=False,
536	quant_config=quant_config,
537	prefix=f"{prefix}.o_proj",
538	)
539
540	# Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True).
541	self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
542	self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
543
544	# Partial RoPE: rotary_dim == head_dim * partial_rotary_factor. Honors
545	# config.rope_scaling (e.g. YaRN) so long-context positions are covered.
546	self.rotary_emb = _build_rotary_emb(config, self.head_dim)
547
548	self.attn = Attention(
549	self.num_heads,
550	self.head_dim,
551	self.scaling,
552	num_kv_heads=self.num_kv_heads,
553	cache_config=cache_config,
554	quant_config=quant_config,
555	prefix=f"{prefix}.attn",
556	)
557
558	def forward(
559	self,
560	positions: torch.Tensor,
561	hidden_states: torch.Tensor,
562	) -> torch.Tensor:
563	qkv, _ = self.qkv_proj(hidden_states)
564	# Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place (dense
565	# mode: no index branch, no KV-cache insert). Matches nvidia/model.py and
566	# replaces the unfused split -> q_norm/k_norm -> rotary_emb chain; verified
567	# bit-equivalent on ROCm (q/k rel ~2e-3 bf16 noise, v untouched).
568	ops.fused_minimax_m3_qknorm_rope_kv_insert(
569	qkv,
570	self.q_norm.weight,
571	self.k_norm.weight,
572	self.rotary_emb.cos_sin_cache,
573	positions,
574	self.num_heads,
575	self.num_kv_heads,
576	self.rotary_emb.rotary_dim,
577	self.q_norm.variance_epsilon,
578	kv_cache_dtype="auto",
579	)
580	q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
581	attn_output = self.attn(q, k, v)
582	output, _ = self.o_proj(attn_output)
583	return output
584
585
586	# Widest row pitch aiter's reshape_and_cache can carry in its `int` stride args.
587	_MAX_KV_INSERT_ROW_STRIDE = 2**31 - 1
588
589
590	def _kv_insert_operand(t: torch.Tensor) -> torch.Tensor:
591	"""K or V operand for ``aiter.reshape_and_cache``, without a needless copy.
592
593	The AITER sparse-PA insert takes K and V as column slices of the fused
594	``[q \| k \| v \| index_q \| index_k]`` projection, so they are row-strided
595	views and never contiguous. ``reshape_and_cache`` does not need them to be:
596	it passes ``key.stride(0)`` to the kernel as the row pitch and reads
597	``key[token_idx * key_stride + i]`` for ``i`` over ``num_heads *
598	head_size``, so the only requirement is that the trailing
599	``(num_heads, head_size)`` block be contiguous within a row, which such a
600	slice always satisfies. Any other layout falls back to a copy.
601	"""
602	if (
603	t.dim() == 3
604	and t.stride(2) == 1
605	and t.stride(1) == t.shape[2]
606	and t.stride(0) <= _MAX_KV_INSERT_ROW_STRIDE
607	):
608	return t
609	return t.contiguous()
610
611
612	class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
613	"""Block-sparse attention layer with the lightning-indexer branch.
614
615	This is a merged attention layer: it owns the projections (qkv + index
616	q/k), per-head QK norms and RoPE, and the attention-backend wiring that a
617	generic ``Attention`` layer would normally provide — it binds the
618	``MiniMaxM3SparseBackend`` + main impl, registers the main paged K/V cache,
619	and owns the lightning indexer (``MiniMaxM3Indexer``), which holds the
620	index-key side cache.
621
622	The index branch (index_{q,k}_proj + index_{q,k}_norm) feeds the sparse
623	top-k block selection. M3 always disables the index value/output
624	projections (``sparse_disable_index_value`` set for every sparse layer), so
625	``index_{v,o}_proj`` are never created.
626	"""
627
628	def __init__(
629	self,
630	config: PreTrainedConfig,
631	layer_id: int,
632	quant_config: QuantizationConfig \| None = None,
633	prefix: str = "",
634	cache_config: CacheConfig \| None = None,
635	topk_indices_buffer: torch.Tensor \| None = None,
636	sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] \| None = None,
637	) -> None:
638	super().__init__()
639	self.hidden_size = config.hidden_size
640	tp_size = get_tensor_model_parallel_world_size()
641
642	self.total_num_heads = config.num_attention_heads
643	assert self.total_num_heads % tp_size == 0
644	self.num_heads = self.total_num_heads // tp_size
645	self.total_num_kv_heads = config.num_key_value_heads
646	if self.total_num_kv_heads >= tp_size:
647	assert self.total_num_kv_heads % tp_size == 0
648	else:
649	assert tp_size % self.total_num_kv_heads == 0
650	self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
651	self.head_dim = config.head_dim
652	self.q_size = self.num_heads * self.head_dim
653	self.kv_size = self.num_kv_heads * self.head_dim
654	self.scaling = self.head_dim**-0.5
655
656	# Cross-layer index sharing (ATOM index_topk_freq): when True this sparse
657	# layer reuses the previous compute layer's top-k block selection from the
658	# shared topk_indices_buffer instead of recomputing it. Static per layer
659	# -> cudagraph-capture-safe.
660	self.skip_index_topk = _should_skip_index_topk(config, layer_id)
661
662	# Sparse "index" branch dims. index_q has the same head count as the KV
663	# heads (sparse_num_index_heads == num_key_value_heads), so it shards
664	# identically -- including replication when tp_size > num_key_value_heads.
665	sparse_cfg = config.sparse_attention_config
666	self.total_idx_heads = sparse_cfg["sparse_num_index_heads"]
667	self.num_idx_heads = self.num_kv_heads
668	self.idx_head_dim = sparse_cfg["sparse_index_dim"]
669	self.index_q_size = self.num_idx_heads * self.idx_head_dim
670
671	# Single fused projection: q, k, v, index_q, index_k in one GEMM.
672	self.qkv_proj = MinimaxM3QKVParallelLinearWithIndexer(
673	self.hidden_size,
674	self.head_dim,
675	self.total_num_heads,
676	self.total_num_kv_heads,
677	self.total_idx_heads,
678	self.idx_head_dim,
679	bias=False,
680	quant_config=quant_config,
681	prefix=f"{prefix}.qkv_proj",
682	)
683	# reduce_results=False: the attention all-reduce is fused with the
684	# following post_attention_layernorm (GemmaRMSNorm) in the decoder layer
685	# via fused_allreduce_gemma_rms_norm.
686	self.o_proj = RowParallelLinear(
687	self.total_num_heads * self.head_dim,
688	self.hidden_size,
689	bias=False,
690	reduce_results=False,
691	quant_config=quant_config,
692	prefix=f"{prefix}.o_proj",
693	)
694
695	# Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True).
696	self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
697	self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
698
699	# Partial RoPE: rotary_dim == head_dim * partial_rotary_factor. Honors
700	# config.rope_scaling (e.g. YaRN) so long-context positions are covered.
701	self.rotary_emb = _build_rotary_emb(config, self.head_dim)
702
703	self.index_q_norm = MiniMAXGemmaRMSNorm(
704	self.idx_head_dim, eps=config.rms_norm_eps
705	)
706	self.index_k_norm = MiniMAXGemmaRMSNorm(
707	self.idx_head_dim, eps=config.rms_norm_eps
708	)
709	self.index_rotary_emb = self.rotary_emb
710
711	# Attention-backend wiring.
712	vllm_config = get_current_vllm_config()
713	self.layer_name = f"{prefix}.attn"
714	self.kv_cache_dtype = (
715	cache_config.cache_dtype if cache_config is not None else "auto"
716	)
717	self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
718	self.kv_cache_dtype, vllm_config.model_config
719	)
720	# MiniMax-M3 sparse attention owns its KV-cache insert/read path instead
721	# of wrapping the generic Attention module. Keep the same runtime scale
722	# attributes so FP8 KV reads can honor vLLM's per-layer descale contract.
723	set_default_quant_scales(self, register_buffer=True)
724
725	# Indexer side-cache dtype, mirroring --kv-cache-dtype for the main cache
726	# (--attention-config '{"indexer_kv_dtype": ...}'). fp8 e4m3 is what the
727	# AITER indexer needs; bf16 keeps the Triton indexer.
728	self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
729	"bf16"
730	)
731
732	# Shared top-k buffer: the indexer writes the selected blocks into it and
733	# the attend impl reads them back (no Python value crosses the break).
734	self.topk_indices_buffer = topk_indices_buffer
735	# Indexer and main attention are separate impls. The AITER indexer is
736	# ROCm-only, so it is selected here rather than inside the neutral
737	# MiniMaxM3Indexer, which knows nothing about it and would pick Triton.
738	indexer_impl_cls = select_aiter_indexer_impl_cls(
739	topk_blocks=sparse_cfg["sparse_topk_blocks"],
740	sparse_block_size=sparse_cfg["sparse_block_size"],
741	num_index_heads=self.num_idx_heads,
742	index_head_dim=self.idx_head_dim,
743	indexer_kv_dtype=self.indexer_kv_dtype,
744	score_type=sparse_cfg.get("sparse_score_type", "max"),
745	)
746	# The attend gate below needs this: an emitted table is what lets it
747	# serve more than one KV head per rank. Buffers to write the table into
748	# only arrive when the model already resolved the attend to the AITER
749	# path the table addresses, so that is not rechecked here.
750	self.indexer_emits_table = (
751	indexer_impl_cls is not None and sparse_table_buffers is not None
752	)
753	# The backend names the metadata builder, which for the AITER path is
754	# the one that rebases the block table the indexer's top-k resolves
755	# through, so both have to come from the same selection.
756	self.attn_backend, main_impl_cls = select_main_backend_and_impl_cls(
757	topk_blocks=sparse_cfg["sparse_topk_blocks"],
758	kv_cache_dtype=self.kv_cache_dtype,
759	num_kv_heads=self.num_kv_heads,
760	emits_sparse_block_table=self.indexer_emits_table,
761	)
762	# impl is AttentionImplBase (broader than AttentionLayerBase's annotation).
763	self.impl: MiniMaxM3SparseImpl = main_impl_cls( # type: ignore[assignment]
764	self.num_heads,
765	self.head_dim,
766	self.scaling,
767	self.num_kv_heads,
768	kv_cache_dtype=self.kv_cache_dtype,
769	topk_blocks=sparse_cfg["sparse_topk_blocks"],
770	sparse_block_size=sparse_cfg["sparse_block_size"],
771	)
772	self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
773	self.num_kv_heads, emits_sparse_block_table=self.indexer_emits_table
774	)
775	self.kv_cache_k = torch.tensor([])
776	self.kv_cache_v = torch.tensor([])
777	self._aiter_sparse_pa_cache_data_ptr = 0
778	self._aiter_sparse_pa_block_page_stride = 0
779	indexer_kwargs = dict(
780	num_kv_heads=self.num_kv_heads,
781	scale=self.scaling,
782	topk_blocks=sparse_cfg["sparse_topk_blocks"],
783	sparse_block_size=sparse_cfg["sparse_block_size"],
784	num_index_heads=self.num_idx_heads,
785	index_head_dim=self.idx_head_dim,
786	prefix=self.layer_name,
787	init_blocks=sparse_cfg.get("sparse_init_block", 0),
788	local_blocks=sparse_cfg.get("sparse_local_block", 0),
789	score_type=sparse_cfg.get("sparse_score_type", "max"),
790	cache_config=cache_config,
791	indexer_kv_dtype=self.indexer_kv_dtype,
792	topk_indices_buffer=topk_indices_buffer,
793	)
794	self.sparse_bt_buffer: torch.Tensor \| None = None
795	self.sparse_ctx_buffer: torch.Tensor \| None = None
796	if indexer_impl_cls is None:
797	# Self-contained nn.Module: owns its side cache, selects its impl.
798	self.indexer = MiniMaxM3Indexer(**indexer_kwargs)
799	else:
800	if self.indexer_emits_table:
801	assert sparse_table_buffers is not None
802	self.sparse_bt_buffer, self.sparse_ctx_buffer = sparse_table_buffers
803	self.indexer = MiniMaxM3AiterIndexer(
804	impl_cls=indexer_impl_cls,
805	sparse_bt_buffer=self.sparse_bt_buffer,
806	sparse_ctx_buffer=self.sparse_ctx_buffer,
807	**indexer_kwargs,
808	)
809
810	# Register the main K/V cache so the KV-cache manager allocates it.
811	compilation_config = vllm_config.compilation_config
812	if self.layer_name in compilation_config.static_forward_context:
813	raise ValueError(f"Duplicate layer name: {self.layer_name}")
814	compilation_config.static_forward_context[self.layer_name] = self
815	self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache
816
817	def get_attn_backend(self) -> type[MiniMaxM3SparseBackend]:
818	return self.attn_backend
819
820	def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec \| None:
821	# Main GQA K/V cache. Block size may change after load, refresh it.
822	# AITER sparse PA wants K and V as two separate head groups
823	# (H folded into the content dim).
824	sparse_pa = self.use_aiter_sparse_pa
825	kv_bytes = (
826	self.num_kv_heads * self.head_dim * self.kv_cache_torch_dtype.itemsize
827	)
828	return FullAttentionSpec(
829	block_size=vllm_config.cache_config.block_size,
830	num_kv_heads=self.num_kv_heads,
831	head_size=self.head_dim,
832	head_size_v=self.head_dim,
833	dtype=self.kv_cache_torch_dtype,
834	kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
835	num_head_slots=2 if sparse_pa else None,
836	state_content_bytes=kv_bytes if sparse_pa else None,
837	)
838
839	def _ensure_aiter_sparse_pa_kv_cache(self) -> None:
840	if self.kv_cache.numel() == 0:
841	return
842	if self._aiter_sparse_pa_cache_data_ptr == self.kv_cache.data_ptr():
843	return
844
845	kv_cache = self.kv_cache
846	if is_quantized_kv_cache(self.kv_cache_dtype):
847	kv_cache = kv_cache.view(self.impl.kv_cache_fp8_dtype)
848	key_cache, value_cache = kv_cache.unbind(1)
849
850	x = 16 // kv_cache.element_size()
851	if self.head_dim % x != 0:
852	raise RuntimeError(
853	"MiniMax-M3 AITER sparse PA requires head_dim divisible by "
854	f"16 / dtype_size, got head_dim={self.head_dim}, x={x}"
855	)
856
857	num_blocks, _, block_size, _ = kv_cache.shape
858	pages_in_side = block_size // 16
859	if key_cache.is_contiguous() and value_cache.is_contiguous():
860	k_src, v_src = key_cache, value_cache
861	block_pages, v_page_offset = pages_in_side, 0
862	elif kv_cache.is_contiguous():
863	k_src = v_src = kv_cache
864	block_pages, v_page_offset = 2 * pages_in_side, pages_in_side
865	else:
866	raise RuntimeError(
867	"MiniMax-M3 AITER sparse PA needs each K/V head slot stored as "
868	"whole 16-token pages, but the resolved KV cache layout gives "
869	"neither dense planes nor block-contiguous pages."
870	)
871
872	num_phys16 = num_blocks * block_pages
873	self.kv_cache_k = k_src.view(
874	num_phys16,
875	self.num_kv_heads,
876	self.head_dim // x,
877	16,
878	x,
879	)
880	# Offsetting V by half a block lets both sides share one page id.
881	self.kv_cache_v = v_src.view(
882	num_phys16,
883	self.num_kv_heads,
884	16 // x,
885	self.head_dim,
886	x,
887	)[v_page_offset:]
888	self._aiter_sparse_pa_block_page_stride = minimax_m3_sparse_block_page_stride(
889	self.kv_cache_k, self.kv_cache_v
890	)
891	self._aiter_sparse_pa_cache_data_ptr = self.kv_cache.data_ptr()
892
893	def get_aiter_sparse_pa_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]:
894	self._ensure_aiter_sparse_pa_kv_cache()
895	return self.kv_cache_k, self.kv_cache_v
896
897	def _get_aiter_sparse_pa_slot_mapping(
898	self,
899	slot_mapping: torch.Tensor,
900	key_cache: torch.Tensor,
901	value_cache: torch.Tensor,
902	) -> torch.Tensor:
903	if value_cache.shape[0] == key_cache.shape[0]:
904	return slot_mapping
905
906	attn_metadata = get_forward_context().attn_metadata
907	page16_slot_mapping = None
908	if isinstance(attn_metadata, dict):
909	main_md = attn_metadata[self.layer_name]
910	assert isinstance(main_md, MiniMaxM3SparseMetadata)
911	page16_slot_mapping = main_md.page16_slot_mapping
912	if (
913	page16_slot_mapping is None
914	or page16_slot_mapping.shape != slot_mapping.shape
915	):
916	page16_slot_mapping = minimax_m3_rebase_slots_to_page16(
917	slot_mapping, self.kv_cache.shape[2]
918	)
919	return page16_slot_mapping
920
921	def _insert_aiter_sparse_pa_kv(
922	self,
923	k: torch.Tensor,
924	v: torch.Tensor,
925	index_k: torch.Tensor \| None,
926	slot_mapping: torch.Tensor,
927	index_slot_mapping: torch.Tensor \| None,
928	key_cache: torch.Tensor,
929	value_cache: torch.Tensor,
930	) -> None:
931	if self.kv_cache.numel() == 0:
932	return
933	from aiter import reshape_and_cache
934
935	from vllm.models.minimax_m3.amd.ops.sparse_pa import (
936	minimax_m3_insert_index_cache,
937	)
938
939	kv_cache_dtype = (
940	self.kv_cache_dtype
941	if is_quantized_kv_cache(self.kv_cache_dtype)
942	else "auto"
943	)
944	reshape_and_cache(
945	_kv_insert_operand(k),
946	_kv_insert_operand(v),
947	key_cache,
948	value_cache,
949	slot_mapping,
950	kv_cache_dtype=kv_cache_dtype,
951	k_scale=getattr(self, "_k_scale", None),
952	v_scale=getattr(self, "_v_scale", None),
953	asm_layout=True,
954	)
955	if index_k is None or index_slot_mapping is None:
956	return
957	index_cache = self.indexer.index_cache.kv_cache
958	if index_cache.numel() == 0:
959	return
960	minimax_m3_insert_index_cache(index_k, index_cache, index_slot_mapping)
961
962	def forward(
963	self,
964	positions: torch.Tensor,
965	hidden_states: torch.Tensor,
966	) -> torch.Tensor:
967	# Single fused projection emitting [q \| k \| v \| index_q \| index_k].
968	qkv, _ = self.qkv_proj(hidden_states)
969
970	# Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on the
971	# main (q/k) and index (index_q/index_k) branches, all read straight out
972	# of the single fused ``qkv`` tensor. Once the paged caches are bound the
973	# kernel also inserts k/v and the index key into them (each with its own
974	# slot_mapping); the memory-profiling run (caches unbound, no slot_mapping)
975	# short-circuits to zeros below. The main and index slot mappings are read
976	# from the forward context's slot_mapping dict, matching the
977	# breakable-cudagraph path -- see nvidia/model.py.
978	cos_sin_cache = self.rotary_emb.cos_sin_cache
979	rotary_dim = self.rotary_emb.rotary_dim
980	eps = self.q_norm.variance_epsilon
981	num_tokens = qkv.shape[0]
982
983	fwd_slot_mapping = get_forward_context().slot_mapping
984	if (
985	not isinstance(fwd_slot_mapping, dict)
986	or self.layer_name not in fwd_slot_mapping
987	):
988	# Memory-profiling run: caches not yet bound, slot_mapping is empty.
989	return qkv.new_zeros((num_tokens, self.hidden_size))
990
991	main_slot_mapping = fwd_slot_mapping[self.layer_name]
992	q = qkv.new_empty((num_tokens, self.q_size))
993	key_cache = value_cache = page16_slot_mapping = None
994	k_scale = v_scale = None
995	use_fused_qknorm = False
996	if self.use_aiter_sparse_pa:
997	key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache()
998	page16_slot_mapping = self._get_aiter_sparse_pa_slot_mapping(
999	main_slot_mapping, key_cache, value_cache
1000	)
1001	k_scale = getattr(self, "_k_scale", None)
1002	v_scale = getattr(self, "_v_scale", None)
1003	use_fused_qknorm = rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled(
1004	self.kv_cache_dtype, k_scale, v_scale
1005	)
1006	if self.skip_index_topk:
1007	index_q = None
1008	if self.use_aiter_sparse_pa:
1009	assert key_cache is not None
1010	assert value_cache is not None
1011	assert page16_slot_mapping is not None
1012	if use_fused_qknorm:
1013	rocm_aiter_ops.fused_qknorm_idxrqknorm(
1014	qkv,
1015	self.q_norm.weight,
1016	self.k_norm.weight,
1017	cos_sin_cache,
1018	positions,
1019	self.num_heads,
1020	self.num_kv_heads,
1021	rotary_dim,
1022	eps,
1023	page16_slot_mapping,
1024	key_cache,
1025	value_cache,
1026	q,
1027	self.kv_cache_dtype,
1028	k_scale,
1029	v_scale,
1030	num_index_heads=self.num_idx_heads,
1031	skip_index_branch=True,
1032	)
1033	else:
1034	ops.fused_minimax_m3_qknorm_rope_kv_insert(
1035	qkv,
1036	self.q_norm.weight,
1037	self.k_norm.weight,
1038	cos_sin_cache,
1039	positions,
1040	self.num_heads,
1041	self.num_kv_heads,
1042	rotary_dim,
1043	eps,
1044	num_index_heads=self.num_idx_heads,
1045	q_out=q,
1046	skip_index_branch=True,
1047	)
1048	k_start = self.q_size
1049	v_start = k_start + self.kv_size
1050	k = qkv[:, k_start:v_start].view(
1051	num_tokens, self.num_kv_heads, self.head_dim
1052	)
1053	v = qkv[:, v_start : v_start + self.kv_size].view(
1054	num_tokens, self.num_kv_heads, self.head_dim
1055	)
1056	self._insert_aiter_sparse_pa_kv(
1057	k,
1058	v,
1059	None,
1060	page16_slot_mapping,
1061	None,
1062	key_cache,
1063	value_cache,
1064	)
1065	else:
1066	ops.fused_minimax_m3_qknorm_rope_kv_insert(
1067	qkv,
1068	self.q_norm.weight,
1069	self.k_norm.weight,
1070	cos_sin_cache,
1071	positions,
1072	self.num_heads,
1073	self.num_kv_heads,
1074	rotary_dim,
1075	eps,
1076	num_index_heads=self.num_idx_heads,
1077	slot_mapping=main_slot_mapping,
1078	kv_cache=self.kv_cache,
1079	block_size=self.kv_cache.size(2), # paged-cache block size
1080	q_out=q,
1081	kv_cache_dtype=self.kv_cache_dtype,
1082	skip_index_branch=True,
1083	)
1084	else:
1085	index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
1086	# index_q matches the index-K cache dtype (e4m3 for the AITER fp8
1087	# score path); the fused kernel emits fp8 directly when this buffer
1088	# is e4m3.
1089	index_q = qkv.new_empty(
1090	(num_tokens, self.index_q_size),
1091	dtype=self.indexer.index_cache.dtype,
1092	)
1093	if self.use_aiter_sparse_pa:
1094	assert key_cache is not None
1095	assert value_cache is not None
1096	assert page16_slot_mapping is not None
1097	index_cache = self.indexer.index_cache.kv_cache
1098	if use_fused_qknorm:
1099	rocm_aiter_ops.fused_qknorm_idxrqknorm(
1100	qkv,
1101	self.q_norm.weight,
1102	self.k_norm.weight,
1103	cos_sin_cache,
1104	positions,
1105	self.num_heads,
1106	self.num_kv_heads,
1107	rotary_dim,
1108	eps,
1109	page16_slot_mapping,
1110	key_cache,
1111	value_cache,
1112	q,
1113	self.kv_cache_dtype,
1114	k_scale,
1115	v_scale,
1116	self.index_q_norm.weight,
1117	self.index_k_norm.weight,
1118	self.num_idx_heads,
1119	index_cache,
1120	index_q,
1121	index_slot_mapping,
1122	)
1123	else:
1124	ops.fused_minimax_m3_qknorm_rope_kv_insert(
1125	qkv,
1126	self.q_norm.weight,
1127	self.k_norm.weight,
1128	cos_sin_cache,
1129	positions,
1130	self.num_heads,
1131	self.num_kv_heads,
1132	rotary_dim,
1133	eps,
1134	self.index_q_norm.weight,
1135	self.index_k_norm.weight,
1136	self.num_idx_heads,
1137	q_out=q,
1138	index_q_out=index_q,
1139	kv_cache_dtype=self.kv_cache_dtype,
1140	)
1141	k_start = self.q_size
1142	v_start = k_start + self.kv_size
1143	index_k_start = v_start + self.kv_size + self.index_q_size
1144	k = qkv[:, k_start:v_start].view(
1145	num_tokens, self.num_kv_heads, self.head_dim
1146	)
1147	v = qkv[:, v_start : v_start + self.kv_size].view(
1148	num_tokens, self.num_kv_heads, self.head_dim
1149	)
1150	index_k = qkv[
1151	:, index_k_start : index_k_start + self.idx_head_dim
1152	].view(num_tokens, self.idx_head_dim)
1153	self._insert_aiter_sparse_pa_kv(
1154	k,
1155	v,
1156	index_k,
1157	page16_slot_mapping,
1158	index_slot_mapping,
1159	key_cache,
1160	value_cache,
1161	)
1162	else:
1163	ops.fused_minimax_m3_qknorm_rope_kv_insert(
1164	qkv,
1165	self.q_norm.weight,
1166	self.k_norm.weight,
1167	cos_sin_cache,
1168	positions,
1169	self.num_heads,
1170	self.num_kv_heads,
1171	rotary_dim,
1172	eps,
1173	self.index_q_norm.weight,
1174	self.index_k_norm.weight,
1175	self.num_idx_heads,
1176	main_slot_mapping,
1177	index_slot_mapping,
1178	self.kv_cache,
1179	self.indexer.index_cache.kv_cache,
1180	self.kv_cache.size(2), # paged-cache block size
1181	q,
1182	index_q,
1183	self.kv_cache_dtype,
1184	)
1185
1186	output = torch.empty_like(q)
1187	attn_output = self._run_attention(q, index_q, output)
1188	output, _ = self.o_proj(attn_output)
1189	return output
1190
1191	@eager_break_during_capture
1192	def _run_attention(
1193	self,
1194	query: torch.Tensor,
1195	index_query: torch.Tensor \| None,
1196	output: torch.Tensor,
1197	) -> torch.Tensor:
1198	# Single eager break around both: their split-K kernels read per-request
1199	# metadata and can't be captured into a cudagraph. The indexer writes its
1200	# top-k into the shared ``topk_indices_buffer``; the attend reads it back.
1201	# When skip_index_topk is set (ATOM index_topk_freq), reuse the selection
1202	# the preceding compute layer wrote into the shared buffer this forward.
1203	decode_sparse_table: tuple[torch.Tensor, torch.Tensor] \| None = None
1204	if not self.skip_index_topk:
1205	assert index_query is not None
1206	# The AITER indexer emits the attend table into its persistent
1207	# buffers, resolving its selection through the page-16 rebase of the
1208	# attend's block table. The Triton indexer can instead fuse decode
1209	# top-k with main's per-forward sparse-table allocation.
1210	if self.indexer_emits_table:
1211	attn_metadata = get_forward_context().attn_metadata
1212	decode_page16 = prefill_page16 = None
1213	if isinstance(attn_metadata, dict):
1214	main_md = attn_metadata[self.layer_name]
1215	assert isinstance(main_md, MiniMaxM3SparseMetadata)
1216	# Rebased once per step by the AITER attend's builder, which
1217	# this path selected along with the impl.
1218	d, p = main_md.decode, main_md.prefill
1219	if d is not None:
1220	assert isinstance(d, MiniMaxM3SparseAiterPADecodeMetadata)
1221	decode_page16 = d.page16_block_table
1222	if p is not None:
1223	assert isinstance(p, MiniMaxM3SparseAiterPAPrefillMetadata)
1224	prefill_page16 = p.page16_block_table
1225	self.indexer(
1226	index_query,
1227	decode_page16_block_table=decode_page16,
1228	prefill_page16_block_table=prefill_page16,
1229	)
1230	else:
1231	attn_metadata = get_forward_context().attn_metadata
1232	if self.use_aiter_sparse_pa and isinstance(attn_metadata, dict):
1233	main_md = attn_metadata[self.layer_name]
1234	assert isinstance(main_md, MiniMaxM3SparseMetadata)
1235	if main_md.num_decodes > 0:
1236	d = main_md.decode
1237	assert d is not None
1238	topk = self.topk_indices_buffer
1239	assert topk is not None
1240	decode_sparse_table = minimax_m3_alloc_sparse_block_table(
1241	topk[:, : main_md.num_decode_tokens, :]
1242	)
1243	block_page_stride = self._aiter_sparse_pa_block_page_stride
1244	assert block_page_stride > 0
1245	self.indexer(
1246	index_query,
1247	attention_block_table=d.block_table,
1248	sparse_block_table_out=decode_sparse_table[0],
1249	sparse_context_lens_out=decode_sparse_table[1],
1250	block_page_stride=block_page_stride,
1251	)
1252	else:
1253	self.indexer(index_query)
1254	else:
1255	self.indexer(index_query)
1256	if self.use_aiter_sparse_pa:
1257	assert isinstance(self.impl, MiniMaxM3SparseAiterPAImpl)
1258	return self.impl.forward(
1259	self,
1260	query,
1261	self.kv_cache,
1262	output,
1263	decode_sparse_table=decode_sparse_table,
1264	)
1265	return self.impl.forward(self, query, self.kv_cache, output)
1266
1267
1268	class MiniMaxM3DecoderLayer(nn.Module):
1269	def __init__(
1270	self,
1271	config: PreTrainedConfig,
1272	prefix: str,
1273	cache_config: CacheConfig \| None = None,
1274	quant_config: QuantizationConfig \| None = None,
1275	force_sparse_attn: bool = False,
1276	force_moe: bool = False,
1277	topk_indices_buffer: torch.Tensor \| None = None,
1278	sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] \| None = None,
1279	) -> None:
1280	super().__init__()
1281	self.hidden_size = config.hidden_size
1282	# DecoderLayers are created with `make_layers` which passes the prefix
1283	# with the layer's index.
1284	layer_id = int(prefix.split(sep=".")[-1])
1285	self.layer_id = layer_id
1286
1287	is_sparse_attention_layer = (
1288	force_sparse_attn or layer_id in _sparse_attention_layer_ids(config)
1289	)
1290
1291	if is_sparse_attention_layer:
1292	self.self_attn = MiniMaxM3SparseAttention(
1293	config=config,
1294	layer_id=layer_id,
1295	quant_config=quant_config,
1296	prefix=f"{prefix}.self_attn",
1297	cache_config=cache_config,
1298	topk_indices_buffer=topk_indices_buffer,
1299	sparse_table_buffers=sparse_table_buffers,
1300	)
1301	else:
1302	self.self_attn = MiniMaxM3Attention(
1303	config=config,
1304	layer_id=layer_id,
1305	quant_config=quant_config,
1306	prefix=f"{prefix}.self_attn",
1307	cache_config=cache_config,
1308	)
1309
1310	# Dense layers store the FFN under `mlp`; MoE layers under
1311	# `block_sparse_moe` -- matching the checkpoint's naming.
1312	self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id)
1313	if self.is_moe_layer:
1314	self.block_sparse_moe = MiniMaxM3MoE(
1315	config=config,
1316	layer_id=layer_id,
1317	quant_config=quant_config,
1318	prefix=f"{prefix}.block_sparse_moe",
1319	)
1320	else:
1321	self.mlp = MiniMaxM3MLP(
1322	config=config,
1323	intermediate_size=config.dense_intermediate_size,
1324	quant_config=quant_config,
1325	prefix=f"{prefix}.mlp",
1326	)
1327
1328	# config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm.
1329	self.input_layernorm = MiniMAXGemmaRMSNorm(
1330	config.hidden_size, eps=config.rms_norm_eps
1331	)
1332	self.post_attention_layernorm = MiniMAXGemmaRMSNorm(
1333	config.hidden_size, eps=config.rms_norm_eps
1334	)
1335
1336	def forward(
1337	self,
1338	positions: torch.Tensor,
1339	hidden_states: torch.Tensor,
1340	residual: torch.Tensor \| None,
1341	) -> tuple[torch.Tensor, torch.Tensor]:
1342	# Self Attention
1343	if residual is None:
1344	residual = hidden_states
1345	hidden_states = self.input_layernorm(hidden_states)
1346	else:
1347	hidden_states, residual = self.input_layernorm(hidden_states, residual)
1348	hidden_states = self.self_attn(
1349	positions=positions,
1350	hidden_states=hidden_states,
1351	)
1352
1353	hidden_states, residual = fused_allreduce_gemma_rms_norm(
1354	hidden_states, residual, self.post_attention_layernorm
1355	)
1356	ffn = self.block_sparse_moe if self.is_moe_layer else self.mlp
1357	hidden_states = ffn(hidden_states)
1358	return hidden_states, residual
1359
1360
1361	class MiniMaxM3Model(nn.Module, EagleModelMixin):
1362	fall_back_to_pt_during_load = False
1363
1364	def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1365	super().__init__()
1366
1367	config = vllm_config.model_config.hf_text_config
1368	cache_config = vllm_config.cache_config
1369	quant_config = vllm_config.quant_config
1370	self.config = config
1371
1372	self.vocab_size = config.vocab_size
1373
1374	if get_pp_group().is_first_rank or spec_decode_needs_target_embed(
1375	vllm_config, include_mtp=vllm_config.use_v2_model_runner
1376	):
1377	self.embed_tokens = VocabParallelEmbedding(
1378	config.vocab_size,
1379	config.hidden_size,
1380	quant_config=quant_config,
1381	prefix=f"{prefix}.embed_tokens",
1382	)
1383	else:
1384	self.embed_tokens = PPMissingLayer()
1385
1386	# Reserved top-k indices buffer shared by all sparse-attention indexer
1387	# layers (mirrors DeepseekV4); the indexer writes its per-head decode/
1388	# prefill block selection into it, the attend reads it back.
1389	sparse_cfg = getattr(config, "sparse_attention_config", None)
1390	self.sparse_table_buffers: tuple[torch.Tensor, torch.Tensor] \| None = None
1391	if sparse_cfg is not None:
1392	tp_size = get_tensor_model_parallel_world_size()
1393	num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
1394	max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
1395	self.topk_indices_buffer = torch.empty(
1396	num_index_heads,
1397	max_tokens,
1398	sparse_cfg["sparse_topk_blocks"],
1399	dtype=torch.int32,
1400	)
1401
1402	num_kv_heads = max(1, config.num_key_value_heads // tp_size)
1403	indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
1404	"bf16"
1405	)
1406	# Probed with the head count the layer will pass (it derives
1407	# num_idx_heads from num_kv_heads, not from sparse_num_index_heads),
1408	# so this cannot land on the other side of the MFMA column limit
1409	# from the selection the layer makes.
1410	emits_table = (
1411	select_aiter_indexer_impl_cls(
1412	topk_blocks=sparse_cfg["sparse_topk_blocks"],
1413	sparse_block_size=sparse_cfg["sparse_block_size"],
1414	num_index_heads=num_kv_heads,
1415	index_head_dim=sparse_cfg["sparse_index_dim"],
1416	indexer_kv_dtype=indexer_kv_dtype,
1417	score_type=sparse_cfg.get("sparse_score_type", "max"),
1418	)
1419	is not None
1420	)
1421	if emits_table and minimax_m3_use_aiter_sparse_pa(
1422	num_kv_heads, emits_sparse_block_table=True
1423	):
1424	# One row per (token, kv head): the table addresses the attend's
1425	# cache, not the index cache.
1426	rows = max_tokens * num_kv_heads
1427	self.sparse_table_buffers = (
1428	torch.empty(
1429	rows,
1430	sparse_cfg["sparse_topk_blocks"]
1431	* (sparse_cfg["sparse_block_size"] // ASM_PAGE_SIZE),
1432	dtype=torch.int32,
1433	),
1434	torch.empty(rows, dtype=torch.int32),
1435	)
1436	else:
1437	self.topk_indices_buffer = None
1438
1439	self.start_layer, self.end_layer, self.layers = make_layers(
1440	config.num_hidden_layers,
1441	lambda prefix: MiniMaxM3DecoderLayer(
1442	config,
1443	prefix,
1444	cache_config=cache_config,
1445	quant_config=quant_config,
1446	topk_indices_buffer=self.topk_indices_buffer,
1447	sparse_table_buffers=self.sparse_table_buffers,
1448	),
1449	prefix=f"{prefix}.layers",
1450	)
1451	self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
1452	self.layers,
1453	MiniMaxM3MoE,
1454	"block_sparse_moe",
1455	)
1456
1457	if get_pp_group().is_last_rank:
1458	self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1459	else:
1460	self.norm = PPMissingLayer()
1461	self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
1462	["hidden_states", "residual"], config.hidden_size
1463	)
1464
1465	def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1466	return self.embed_tokens(input_ids)
1467
1468	def forward(
1469	self,
1470	input_ids: torch.Tensor \| None,
1471	positions: torch.Tensor,
1472	intermediate_tensors: IntermediateTensors \| None,
1473	inputs_embeds: torch.Tensor \| None = None,
1474	) -> torch.Tensor \| IntermediateTensors \| tuple[torch.Tensor, list[torch.Tensor]]:
1475	if get_pp_group().is_first_rank:
1476	if inputs_embeds is not None:
1477	hidden_states = inputs_embeds
1478	else:
1479	hidden_states = self.embed_input_ids(input_ids)
1480	residual = None
1481	else:
1482	assert intermediate_tensors is not None
1483	hidden_states = intermediate_tensors["hidden_states"]
1484	residual = intermediate_tensors["residual"]
1485
1486	# EAGLE3 is not yet compatible with pipeline parallel
1487	aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
1488	for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
1489	hidden_states, residual = layer(positions, hidden_states, residual)
1490	self._maybe_add_hidden_state(
1491	aux_hidden_states, idx + 1, hidden_states, residual
1492	)
1493
1494	if not get_pp_group().is_last_rank:
1495	return IntermediateTensors(
1496	{"hidden_states": hidden_states, "residual": residual}
1497	)
1498
1499	hidden_states, _ = self.norm(hidden_states, residual)
1500
1501	if len(aux_hidden_states) > 0:
1502	return hidden_states, aux_hidden_states
1503	return hidden_states
1504
1505	def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1506	# Checkpoint experts use w1=gate, w2=down, w3=up. When fusing the shared
1507	# expert, include the appended slot (id == num_local_experts).
1508	n_shared = getattr(self.config, "n_shared_experts", 0) or 0
1509	num_experts = self.config.num_local_experts + (
1510	n_shared if self.is_fused_shared_expert_enabled else 0
1511	)
1512	return fused_moe_make_expert_params_mapping(
1513	self,
1514	ckpt_gate_proj_name="w1",
1515	ckpt_down_proj_name="w2",
1516	ckpt_up_proj_name="w3",
1517	num_experts=num_experts,
1518	)
1519
1520	def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1521	# q/k/v_proj -> fused qkv_proj; gate_proj/up_proj -> fused gate_up_proj
1522	# (dense MLP and shared expert). On sparse layers the indexer
1523	# index_q/index_k_proj fold into the same fused qkv_proj
1524	# (MinimaxM3QKVParallelLinearWithIndexer); these entries simply never match on
1525	# dense layers, whose checkpoints have no index_*_proj weights. Leading
1526	# dots keep `q_proj`/`k_proj` from matching `index_q_proj`/`index_k_proj`
1527	# (preceded by `_`, not `.`).
1528	stacked_params_mapping: list[tuple[str, str, int \| str]] = [
1529	# (param_name, shard_name, shard_id)
1530	(".qkv_proj", ".q_proj", "q"),
1531	(".qkv_proj", ".k_proj", "k"),
1532	(".qkv_proj", ".v_proj", "v"),
1533	(".qkv_proj", ".index_q_proj", "index_q"),
1534	(".qkv_proj", ".index_k_proj", "index_k"),
1535	(".gate_up_proj", ".gate_proj", 0),
1536	(".gate_up_proj", ".up_proj", 1),
1537	]
1538
1539	# (param_name, weight_name, expert_id, shard_id)
1540	expert_params_mapping = self.get_expert_mapping()
1541
1542	params_dict = dict(self.named_parameters())
1543	loaded_params: set[str] = set()
1544	for name, loaded_weight in weights:
1545	# The MTP module is not modeled yet.
1546	if "mtp." in name:
1547	continue
1548
1549	# The checkpoint stores block scales as ``weight_scale_inv``; the
1550	# ModelOpt MXFP8 layers expose them as ``weight_scale``.
1551	if "weight_scale_inv" in name:
1552	name = name.replace("weight_scale_inv", "weight_scale")
1553
1554	# Shared-expert fusion: redirect the checkpoint shared expert into
1555	# routed-expert slot ``num_local_experts`` (gate->w1, up->w3,
1556	# down->w2) so it loads via the routed expert loader. Runs before the
1557	# stacked/dense mappings so shared_experts.gate_proj/up_proj are not
1558	# captured by the dense gate_up_proj mapping.
1559	if self.is_fused_shared_expert_enabled and ".shared_experts." in name:
1560	sid = self.config.num_local_experts
1561	name = name.replace(".shared_experts.gate_proj.", f".experts.{sid}.w1.")
1562	name = name.replace(".shared_experts.up_proj.", f".experts.{sid}.w3.")
1563	name = name.replace(".shared_experts.down_proj.", f".experts.{sid}.w2.")
1564
1565	for param_name, weight_name, shard_id in stacked_params_mapping:
1566	if weight_name not in name:
1567	continue
1568	# Routed experts (w1/w2/w3) are handled below; don't let the
1569	# stacked mapping rewrite them.
1570	if ("block_sparse_moe.experts." in name) and name not in params_dict:
1571	continue
1572	name = name.replace(weight_name, param_name)
1573	if name.endswith(".bias") and name not in params_dict:
1574	continue
1575	if is_pp_missing_parameter(name, self):
1576	continue
1577	if name not in params_dict:
1578	continue
1579	param = params_dict[name]
1580	weight_loader = param.weight_loader
1581	weight_loader(param, loaded_weight, shard_id)
1582	break
1583	else:
1584	for (
1585	param_name,
1586	weight_name,
1587	expert_id,
1588	expert_shard_id,
1589	) in expert_params_mapping:
1590	if weight_name not in name:
1591	continue
1592	name = name.replace(weight_name, param_name)
1593	if is_pp_missing_parameter(name, self):
1594	continue
1595	if name not in params_dict:
1596	continue
1597	param = params_dict[name]
1598	weight_loader = param.weight_loader
1599	weight_loader(
1600	param,
1601	loaded_weight,
1602	name,
1603	shard_id=expert_shard_id,
1604	expert_id=expert_id,
1605	)
1606	break
1607	else:
1608	if name.endswith(".bias") and name not in params_dict:
1609	continue
1610	remapped = maybe_remap_kv_scale_name(name, params_dict)
1611	if remapped is None:
1612	continue
1613	name = remapped
1614	if is_pp_missing_parameter(name, self):
1615	continue
1616	# Modules not modeled yet (e.g. attention) are skipped until
1617	# they are ported.
1618	if name not in params_dict:
1619	continue
1620	param = params_dict[name]
1621	weight_loader = getattr(
1622	param, "weight_loader", default_weight_loader
1623	)
1624	weight_loader(param, loaded_weight)
1625	loaded_params.add(name)
1626	return loaded_params
1627
1628
1629	class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
1630	"""MiniMax M3 (sparse/dense backbone) for causal language modeling."""
1631
1632	packed_modules_mapping = {
1633	"qkv_proj": ["q_proj", "k_proj", "v_proj"],
1634	"gate_up_proj": ["gate_proj", "up_proj"],
1635	}
1636
1637	def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1638	super().__init__()
1639	config = vllm_config.model_config.hf_text_config
1640	quant_config = vllm_config.quant_config
1641	self.config = config
1642	self.quant_config = quant_config
1643	self.model = MiniMaxM3Model(
1644	vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
1645	)
1646	if get_pp_group().is_last_rank:
1647	self.lm_head = ParallelLMHead(
1648	config.vocab_size,
1649	config.hidden_size,
1650	quant_config=quant_config,
1651	prefix=maybe_prefix(prefix, "lm_head"),
1652	)
1653	else:
1654	self.lm_head = PPMissingLayer()
1655	self.logits_processor = LogitsProcessor(config.vocab_size)
1656	self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1657	self.model.make_empty_intermediate_tensors
1658	)
1659
1660	def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1661	return self.model.embed_input_ids(input_ids)
1662
1663	def forward(
1664	self,
1665	input_ids: torch.Tensor \| None,
1666	positions: torch.Tensor,
1667	intermediate_tensors: IntermediateTensors \| None = None,
1668	inputs_embeds: torch.Tensor \| None = None,
1669	**kwargs,
1670	) -> torch.Tensor \| IntermediateTensors:
1671	return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
1672
1673	def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor \| None:
1674	return self.logits_processor(self.lm_head, hidden_states)
1675
1676	def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1677	return self.model.get_expert_mapping()
1678
1679	def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1680	loader = AutoWeightsLoader(self)
1681	return loader.load_weights(weights)
1682
1683
1684	# TODO(refactor): this VL wrapper is platform-agnostic and byte-identical to the
1685	# NVIDIA copy — it only orchestrates the shared vision tower + the per-platform
1686	# language model (resolved via ``init_vllm_registered_model``). Hoist it into
1687	# ``common/`` to drop the amd/nvidia duplication once the split stabilizes.
1688	@MULTIMODAL_REGISTRY.register_processor(
1689	MiniMaxM3VLMultiModalProcessor,
1690	info=MiniMaxM3VLProcessingInfo,
1691	dummy_inputs=MiniMaxM3VLDummyInputsBuilder,
1692	)
1693	class MiniMaxM3SparseForConditionalGeneration(
1694	nn.Module,
1695	SupportsMultiModal,
1696	MiniMaxM3EncoderCudaGraphMixin,
1697	SupportsPP,
1698	SupportsEagle3,
1699	):
1700	"""Top-level (VL) entry point for MiniMax M3.
1701
1702	Owns the shared MiniMax-M3 vision tower on ROCm and delegates text
1703	generation to the AMD language-model path.
1704	"""
1705
1706	# The vision tower runs replicated per rank under ``--mm-encoder-tp-mode
1707	# data``; ``run_dp_sharded_mrope_vision_model`` shards the work across
1708	# ranks (see ``_process_image_input`` / ``_process_video_input``).
1709	supports_encoder_tp_data = True
1710
1711	packed_modules_mapping = {
1712	"qkv_proj": ["q_proj", "k_proj", "v_proj"],
1713	"gate_up_proj": ["gate_proj", "up_proj"],
1714	}
1715
1716	hf_to_vllm_mapper = WeightsMapper(
1717	orig_to_new_prefix={
1718	"multi_modal_projector.": "vision_tower.multi_modal_projector.",
1719	"patch_merge_mlp.": "vision_tower.patch_merge_mlp.",
1720	},
1721	orig_to_new_substr={
1722	".mlp.fc1": ".fc1",
1723	".mlp.fc2": ".fc2",
1724	},
1725	)
1726
1727	@classmethod
1728	def get_placeholder_str(cls, modality: str, i: int) -> str \| None:
1729	if modality == "image":
1730	return MiniMaxM3VLProcessingInfo.IMAGE_TOKEN
1731	if modality == "video":
1732	return MiniMaxM3VLProcessingInfo.VIDEO_TOKEN
1733	raise ValueError(f"Unsupported modality: {modality!r}")
1734
1735	def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1736	super().__init__()
1737	config = vllm_config.model_config.hf_config
1738	self.config = config
1739	self.quant_config = vllm_config.quant_config
1740	self.multimodal_config = vllm_config.model_config.multimodal_config
1741	assert self.multimodal_config is not None
1742	self.use_data_parallel = self.multimodal_config.mm_encoder_tp_mode == "data"
1743
1744	text_hidden_size = getattr(config.text_config, "hidden_size", None)
1745	assert text_hidden_size is not None, "text_config.hidden_size is required"
1746	projector_hidden_size = getattr(config, "projector_hidden_size", None)
1747
1748	with self._mark_tower_model(vllm_config, {"image", "video"}):
1749	vision_config = config.vision_config
1750	self.vision_tower = MiniMaxVLVisionModel(
1751	config=PreTrainedConfig.from_dict(vision_config),
1752	text_hidden_size=text_hidden_size,
1753	projector_hidden_size=projector_hidden_size,
1754	quant_config=self.quant_config,
1755	prefix=maybe_prefix(prefix, "vision_tower"),
1756	)
1757
1758	self.language_model = init_vllm_registered_model(
1759	vllm_config=vllm_config,
1760	hf_config=config.text_config,
1761	prefix=maybe_prefix(prefix, "language_model"),
1762	architectures=["MiniMaxM3SparseForCausalLM"],
1763	)
1764	self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1765	self.language_model.make_empty_intermediate_tensors
1766	)
1767
1768	def _parse_and_validate_image_input(self, **kwargs: object) -> dict \| None:
1769	pixel_values = kwargs.pop("pixel_values", None)
1770	image_grid_thw = kwargs.pop("image_grid_thw", None)
1771	if pixel_values is None:
1772	return None
1773	return {"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}
1774
1775	def _parse_and_validate_video_input(self, **kwargs: object) -> dict \| None:
1776	pixel_values_videos = kwargs.pop("pixel_values_videos", None)
1777	video_grid_thw = kwargs.pop("video_grid_thw", None)
1778	if pixel_values_videos is None:
1779	return None
1780	return {
1781	"pixel_values_videos": pixel_values_videos,
1782	"video_grid_thw": video_grid_thw,
1783	}
1784
1785	def _process_image_input(self, image_input: dict) -> tuple[torch.Tensor, ...]:
1786	pixel_values: torch.Tensor = image_input["pixel_values"].type(
1787	self.vision_tower.dtype
1788	)
1789	grid_thw: torch.Tensor = image_input["image_grid_thw"]
1790	assert grid_thw.ndim == 2
1791
1792	if self.use_data_parallel:
1793	# Already returns a per-item tuple of embeddings.
1794	return run_dp_sharded_mrope_vision_model(
1795	self.vision_tower,
1796	pixel_values,
1797	grid_thw.tolist(),
1798	rope_type="rope_3d",
1799	)
1800
1801	image_embeds = self.vision_tower(
1802	pixel_values=pixel_values,
1803	grid_thw=grid_thw.tolist(),
1804	)
1805
1806	# Split the concatenated output into one tensor per image item.
1807	merge_size = self.vision_tower.spatial_merge_size
1808	sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist()
1809	return image_embeds.split(sizes)
1810
1811	def _process_video_input(self, video_input: dict) -> tuple[torch.Tensor, ...]:
1812	pixel_values: torch.Tensor = video_input["pixel_values_videos"].type(
1813	self.vision_tower.dtype
1814	)
1815	grid_thw: torch.Tensor = video_input["video_grid_thw"]
1816	assert grid_thw.ndim == 2
1817
1818	if self.use_data_parallel:
1819	# Already returns a per-item tuple of embeddings.
1820	return run_dp_sharded_mrope_vision_model(
1821	self.vision_tower,
1822	pixel_values,
1823	grid_thw.tolist(),
1824	rope_type="rope_3d",
1825	)
1826
1827	video_embeds = self.vision_tower(
1828	pixel_values=pixel_values,
1829	grid_thw=grid_thw.tolist(),
1830	)
1831
1832	# Split the concatenated output into one tensor per video item.
1833	merge_size = self.vision_tower.spatial_merge_size
1834	sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist()
1835	return video_embeds.split(sizes)
1836
1837	def _parse_and_validate_multimodal_inputs(
1838	self, **kwargs: object
1839	) -> dict[str, dict]:
1840	mm_input_by_modality: dict[str, dict] = {}
1841	for input_key in kwargs:
1842	if input_key == "pixel_values" and "image" not in mm_input_by_modality:
1843	image_input = self._parse_and_validate_image_input(**kwargs)
1844	if image_input is not None:
1845	mm_input_by_modality["image"] = image_input
1846	if (
1847	input_key == "pixel_values_videos"
1848	and "video" not in mm_input_by_modality
1849	):
1850	video_input = self._parse_and_validate_video_input(**kwargs)
1851	if video_input is not None:
1852	mm_input_by_modality["video"] = video_input
1853	return mm_input_by_modality
1854
1855	def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
1856	mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
1857	if not mm_input_by_modality:
1858	return []
1859
1860	multimodal_embeddings: list[torch.Tensor] = []
1861	for modality in mm_input_by_modality:
1862	multimodal_input = mm_input_by_modality[modality]
1863	if modality == "image":
1864	image_embeddings = self._process_image_input(multimodal_input)
1865	multimodal_embeddings.extend(image_embeddings)
1866	if modality == "video":
1867	video_embeddings = self._process_video_input(multimodal_input)
1868	multimodal_embeddings.extend(video_embeddings)
1869
1870	return tuple(multimodal_embeddings)
1871
1872	def forward(
1873	self,
1874	input_ids: torch.Tensor \| None,
1875	positions: torch.Tensor,
1876	intermediate_tensors: IntermediateTensors \| None = None,
1877	inputs_embeds: torch.Tensor \| None = None,
1878	**kwargs,
1879	) -> torch.Tensor:
1880	return self.language_model(
1881	input_ids, positions, intermediate_tensors, inputs_embeds
1882	)
1883
1884	def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor \| None:
1885	return self.language_model.compute_logits(hidden_states)
1886
1887	def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1888	return self.language_model.get_expert_mapping()
1889
1890	def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1891	loader = AutoWeightsLoader(self)
1892	return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)