diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index c1d1058a92..262cedf203 100644 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -546,6 +546,14 @@ def _resolve_pkv_names(layer_idx, layer_state): f"k_pe.{i}", ] ) + elif param == "index_keys": + if hasattr(self.model, "get_onnx_index_key_names"): + input_names.extend(self.model.get_onnx_index_key_names()) + elif isinstance(example_inputs.get("index_keys"), (list, tuple)): + for i in range(len(example_inputs["index_keys"])): + input_names.append(f"index_key.{i}") + else: + input_names.append(param) else: input_names.append(param) @@ -897,6 +905,14 @@ def _resolve_pkv_names(layer_idx, layer_state): for layer_offset in range(len(example_inputs["compressed_kvs"])): layer_idx = idx + layer_offset input_names.extend([f"compressed_kv.{layer_idx}", f"k_pe.{layer_idx}"]) + elif param == "index_keys": + if hasattr(self.model, "get_onnx_index_key_names"): + input_names.extend(self.model.get_onnx_index_key_names()) + elif isinstance(example_inputs.get("index_keys"), (list, tuple)): + for i in range(len(example_inputs["index_keys"])): + input_names.append(f"index_key.{i}") + else: + input_names.append(param) else: input_names.append(param) dynamic_axes = {k: v for k, v in dynamic_axes.items() if k in input_names} diff --git a/QEfficient/blocking/blocked_attention_forwards.py b/QEfficient/blocking/blocked_attention_forwards.py index 5874b548fc..75003c8141 100644 --- a/QEfficient/blocking/blocked_attention_forwards.py +++ b/QEfficient/blocking/blocked_attention_forwards.py @@ -408,7 +408,13 @@ def blocked_kv_attention_forward_headpar_offline( skip_future = None if skip_kv: - skip_future = (torch.tensor(start_index, device=query.device) > current_position).all() + # start_index is a plain Python int; comparing it directly (rather than via + # torch.tensor(start_index, device=query.device)) traces to a scalar-comparison + # op with the constant baked in as an attribute, not a separate lifted tensor + # placeholder -- avoids creating a device-bound constant that ends up on the + # meta device (and fails ONNX serialization) when query is a meta tensor, as + # it is for weight-free export. + skip_future = (start_index > current_position).all() # Eager mode Only if not torch.onnx.is_in_onnx_export() and not torch.jit.is_tracing(): if skip_future.item(): diff --git a/QEfficient/customop/__init__.py b/QEfficient/customop/__init__.py index 168e4c3d49..d97244da86 100644 --- a/QEfficient/customop/__init__.py +++ b/QEfficient/customop/__init__.py @@ -16,6 +16,7 @@ CtxScatterFunc3D, CtxScatterFunc3DGeneralized, CtxScatterFunc3DInt, + M3CtxScatterFunc, ) from QEfficient.customop.ctx_scatter_gather_cb import ( CtxGatherFuncBlockedKVCB, @@ -44,6 +45,7 @@ ctx_scatter_3d_int, ctx_scatter_cb, ctx_scatter_cb_3d, + m3_ctx_scatter, ) __all__ = [ @@ -56,6 +58,7 @@ "CtxScatterFunc3D", "CtxScatterFunc3DGeneralized", "CtxScatterFunc3DInt", + "M3CtxScatterFunc", "CtxGatherFunc", "CtxGatherFunc3D", "CtxGatherFunc3DGeneralized", @@ -70,6 +73,7 @@ "ctx_scatter_3d", "ctx_scatter_3d_generalized", "ctx_scatter_3d_int", + "m3_ctx_scatter", "ctx_gather", "ctx_gather_3d", "ctx_gather_3d_generalized", diff --git a/QEfficient/customop/ctx_scatter_gather.py b/QEfficient/customop/ctx_scatter_gather.py index e3b50c8ab1..075fcaecaa 100644 --- a/QEfficient/customop/ctx_scatter_gather.py +++ b/QEfficient/customop/ctx_scatter_gather.py @@ -393,3 +393,47 @@ def setup_context(ctx, inputs, outputs): @staticmethod def symbolic(g: torch.Graph, data: torch.Value, ctx_indices: torch.Value) -> torch.Value: return g.onnxscript_op(CtxGatherBlockedKVBatch, data, ctx_indices).setTypeAs(data) + + +# --------------------------------------------------------------------------- +# MiniMax-M3 indexer scatter +# Identical semantics to CtxScatter but casts position_ids to INT64 before +# building ScatterND indices so that all three index tensors (batch_idx, +# head_idx, ctx_idx) share the same INT64 element type. Without this cast +# ctx_idx stays INT32 while batch_idx/head_idx are INT64 (derived from +# Shape → Gather), producing a mixed-type Concat that triggers the QAIC +# compiler assertion "sameSameShapeExceptDim: Different types". +# --------------------------------------------------------------------------- +@qeff_custom_op("com.qualcomm.cloud", 1) +def M3CtxScatter(data: onnxscript.FLOAT, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT) -> onnxscript.FLOAT: + batch_size = ops.Gather(ops.Shape(data), [0]) + num_heads = ops.Gather(ops.Shape(data), [1]) + seq_len = ops.Gather(ops.Shape(position_ids), [1]) + zero = ops.Constant(value_ints=[0]) + one = ops.Constant(value_ints=[1]) + exp_shape = ops.Concat(batch_size, num_heads, seq_len, one, axis=0) + batch_idx = ops.Expand(ops.Unsqueeze(ops.Range(zero, batch_size, one), [1, 2, 3]), exp_shape) + head_idx = ops.Expand(ops.Unsqueeze(ops.Range(zero, num_heads, one), [0, 2, 3]), exp_shape) + ctx_idx = ops.Expand(ops.Unsqueeze(ops.Cast(position_ids, to=onnxscript.INT64.dtype), [1, 3]), exp_shape) + indices = ops.Concat(batch_idx, head_idx, ctx_idx, axis=3) + return ops.ScatterND(data, indices, updates) + + +class M3CtxScatterFunc(torch.autograd.Function): + """Scatter idx_k into the MiniMax-M3 index-key cache at position_ids.""" + + @staticmethod + def forward(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> torch.Tensor: + batch_idx = torch.arange(data.shape[0]).view(-1, 1, 1) + head_idx = torch.arange(data.shape[1]).view(1, -1, 1) + ctx_idx = position_ids.unsqueeze(1) + data[batch_idx, head_idx, ctx_idx] = updates + return data + + @staticmethod + def setup_context(ctx, inputs, output): + pass + + @staticmethod + def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updates: torch.Value) -> torch.Value: + return g.onnxscript_op(M3CtxScatter, data, position_ids, updates).setTypeAs(data) diff --git a/QEfficient/customop/dynamo_ops.py b/QEfficient/customop/dynamo_ops.py index 81318efbbf..0655c98b65 100644 --- a/QEfficient/customop/dynamo_ops.py +++ b/QEfficient/customop/dynamo_ops.py @@ -14,6 +14,7 @@ CtxScatter, CtxScatter3D, CtxScatter3DInt, + M3CtxScatter, ) from QEfficient.customop.ctx_scatter_gather_cb import ( # noqa: E402 CtxGatherBlockedKVCB, @@ -367,6 +368,25 @@ def _(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> return torch.empty_like(data) +# M3 INDEX-KEY SCATTER (CtxScatter with INT64 index cast — MiniMax-M3 indexer) +@torch.library.custom_op("qefficient::m3_ctx_scatter", mutates_args=()) +def m3_ctx_scatter_op(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> torch.Tensor: + """MiniMax-M3 index-key cache scatter. Same eager semantics as ctx_scatter but the + ONNX export uses M3CtxScatter which casts position_ids to INT64 before building + ScatterND indices, avoiding the INT32/INT64 Concat type mismatch.""" + result = data.clone() + batch_idx = torch.arange(result.shape[0]).view(-1, 1, 1) + head_idx = torch.arange(result.shape[1]).view(1, -1, 1) + ctx_idx = position_ids.unsqueeze(1) + result[batch_idx, head_idx, ctx_idx] = updates + return result + + +@m3_ctx_scatter_op.register_fake +def _(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> torch.Tensor: + return torch.empty_like(data) + + # --------------------------------------------------------------------------- # Translation table: torch.ops.qefficient.* → ONNX export classes. # Used by _export_via_dynamo via custom_translation_table. @@ -380,6 +400,7 @@ def _(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> torch.ops.qefficient.ctx_scatter_cb_3d.default: get_dynamo_onnxscript_func(CtxScatterCB3D), torch.ops.qefficient.ctx_scatter_3d_int.default: get_dynamo_onnxscript_func(CtxScatter3DInt), torch.ops.qefficient.ctx_scatter_3d_generalized.default: get_dynamo_onnxscript_func(CtxScatter3D), + torch.ops.qefficient.m3_ctx_scatter.default: get_dynamo_onnxscript_func(M3CtxScatter), torch.ops.qefficient.ctx_gather.default: get_dynamo_onnxscript_func(CtxGather), torch.ops.qefficient.ctx_gather_3d.default: get_dynamo_onnxscript_func(CtxGather3D), torch.ops.qefficient.ctx_gather_cb.default: get_dynamo_onnxscript_func(CtxGatherCB), diff --git a/QEfficient/customop/utils.py b/QEfficient/customop/utils.py index 7cadaa456d..bbc71592a8 100644 --- a/QEfficient/customop/utils.py +++ b/QEfficient/customop/utils.py @@ -78,6 +78,12 @@ def ctx_gather_blocked_kv(data: torch.Tensor, ctx_indices: torch.Tensor) -> torc return select_interface(CtxGatherFuncBlockedKV.apply, torch.ops.qefficient.ctx_gather_blocked_kv)(data, ctx_indices) +def m3_ctx_scatter(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> torch.Tensor: + from QEfficient.customop.ctx_scatter_gather import M3CtxScatterFunc + + return select_interface(M3CtxScatterFunc.apply, torch.ops.qefficient.m3_ctx_scatter)(data, position_ids, updates) + + # --------------------------------------------------------------------------- # Interface functions for ctx_scatter_gather_cb ops # --------------------------------------------------------------------------- diff --git a/QEfficient/exporter/weight_free/checkpoint_key_resolver.py b/QEfficient/exporter/weight_free/checkpoint_key_resolver.py index 62d255a8f8..a05af46c41 100644 --- a/QEfficient/exporter/weight_free/checkpoint_key_resolver.py +++ b/QEfficient/exporter/weight_free/checkpoint_key_resolver.py @@ -5,6 +5,7 @@ # # ---------------------------------------------------------------------------- +import re from pathlib import Path from typing import Dict, List, Optional @@ -40,6 +41,10 @@ "embed_scale", } +# MiniMax-M3's sparse-attention indexer: nested "indexer." submodule in this HF +# port vs. flat "index_" checkpoint leaves — see find_checkpoint_key. +_INDEXER_RE = re.compile(r"\.indexer\.(q_proj|k_proj|q_norm|k_norm)\.weight$") + def _collect_tied_weights(model: nn.Module) -> list[TiedWeightAlias]: """Return aliases for tied weights, keyed by the model's own tied-weights contract. @@ -127,14 +132,57 @@ def find_checkpoint_key( if prefix and stripped.startswith(f"{prefix}."): candidates.append(stripped[len(f"{prefix}.") :]) - if ".mlp." in stripped: - candidates.append(stripped.replace(".mlp.", ".block_sparse_moe.")) - - if stripped.endswith(".mlp.gate.weight"): - candidates.append(stripped[: -len(".gate.weight")] + ".router.weight") + # MiniMax-M3's published checkpoint wraps the whole text stack as a top-level + # "language_model.*" container (language_model.model.layers.N..., including MoE + # expert/gate/norm weights), the reverse of this HF port's module tree (text + # stack nested under model.language_model.*). This swap applies uniformly across + # every text-side weight — not just per-layer ones — so match it generically on + # any ONNX name containing "model.language_model." rather than per-key-suffix. + lang_model_swap = re.search(r"^(.*?)model\.language_model\.(.*)$", stripped) + swapped = "language_model.model." + lang_model_swap.group(2) if lang_model_swap else None + if swapped: + candidates.append(swapped) + + # The rules below (.mlp./.block_sparse_moe. renaming, gate/router aliasing, the + # sparse-attention indexer flattening, e_score_correction_bias flattening) are all + # independent of the language_model/model prefix swap above, so apply each to both + # the un-swapped and swapped forms — the checkpoint key may need any combination of + # these rewrites simultaneously. + for base in (stripped, swapped): + if base is None: + continue - if stripped.endswith(".mlp.router.weight"): - candidates.append(stripped[: -len(".router.weight")] + ".gate.weight") + renamed_variants = [base] + if ".mlp." in base: + renamed = base.replace(".mlp.", ".block_sparse_moe.") + candidates.append(renamed) + renamed_variants.append(renamed) + + if base.endswith(".mlp.gate.weight"): + candidates.append(base[: -len(".gate.weight")] + ".router.weight") + + if base.endswith(".mlp.router.weight"): + candidates.append(base[: -len(".router.weight")] + ".gate.weight") + + # MiniMax-M3's sparse-attention indexer is a nested "self_attn.indexer." + # submodule in this HF port, but the checkpoint stores those four tensors as + # flat "self_attn.index_" leaves with no separate indexer container. + flattened = _INDEXER_RE.sub(r".index_\1.weight", base) + if flattened != base: + candidates.append(flattened) + + # MiniMax-M3's router correction bias is a buffer nested under the + # "mlp.gate"/"block_sparse_moe.gate" router submodule in this HF port, but the + # checkpoint stores it one level up, directly on the MoE block (no "gate" + # segment). Apply to both the pre- and post-block_sparse_moe-rename forms. + for variant in renamed_variants: + if variant.endswith(".gate.e_score_correction_bias"): + candidates.append(variant[: -len(".gate.e_score_correction_bias")] + ".e_score_correction_bias") + + # lm_head sits directly on the outer model (no "language_model" segment in its + # own ONNX name), so it needs a separate rule from the swap above. + if stripped == "lm_head.weight" or stripped.endswith(".lm_head.weight"): + candidates.append("language_model.lm_head.weight") return _find_checkpoint_key(candidates, checkpoint_index, onnx_name) diff --git a/QEfficient/exporter/weight_free/checkpoint_transforms.py b/QEfficient/exporter/weight_free/checkpoint_transforms.py index 7c8516c332..fd2e18ee3b 100644 --- a/QEfficient/exporter/weight_free/checkpoint_transforms.py +++ b/QEfficient/exporter/weight_free/checkpoint_transforms.py @@ -335,13 +335,21 @@ class MoEExpertStackingCheckpointTransform(BaseCheckpointTransform): matching the derived parameter layout that OptimizedMoETransform creates, so promote_initializers_and_build_spec finds an exact key match. - Non-expert keys receive dtype conversion in the same pass. + + Non-expert keys receive dtype conversion in the same pass. Some of those + non-expert keys are themselves split gate_proj/up_proj pairs (dense MLP + layers, and MoE shared_experts sub-modules) whose QEff wrapper module + (e.g. MiniMax's MiniMaxM3VLDenseMLP) declares a single fused gate_up_proj + nn.Linear instead — those pairs are fused into one gate_up_proj tensor + (torch.cat([gate, up], dim=0), matching gate_up_proj(x).chunk(2, dim=-1) + splitting the output into gate-then-up) rather than copied through as-is. Parallelism: - Phase 1 (scan): one thread per shard, reads keys only (I/O bound, cheap). - Phase 2 (stack): one thread per layer, loads and stacks its experts. - - Phase 3 (base): one thread per shard, converts non-expert keys. + - Phase 3 (base): one thread per shard, converts non-expert keys and + fuses any split gate_proj/up_proj pairs found among them. Phases 2 and 3 run concurrently once phase 1 completes. """ @@ -350,6 +358,11 @@ class MoEExpertStackingCheckpointTransform(BaseCheckpointTransform): r"^(.+\.layers\.(\d+)\..+?\.experts)\.(\d+)\." r"(gate_proj|up_proj|down_proj|linear|linear_v|linear_1|w1|w2|w3)\.weight$" ) + # Split gate_proj/up_proj pair *not* under a numbered .experts.{idx}. path (those + # are handled by EXPERT_RE/stacking above) — dense MLP layers and MoE + # shared_experts sub-modules, whose QEff wrapper module expects one fused + # gate_up_proj tensor instead of two separate ones. + FUSABLE_GATE_UP_RE = re.compile(r"^(.+?)\.(gate_proj|up_proj)\.weight$") @classmethod def is_applicable(cls, weight_map: Dict[str, str], **kwargs) -> bool: @@ -496,10 +509,30 @@ def _stack_layer(layer_idx: int) -> Tuple[str, List[str]]: def _convert_base(shard_name: str, keys: List[str]) -> None: tensors: Dict[str, torch.Tensor] = {} + fusable: Dict[str, Dict[str, torch.Tensor]] = {} with safe_open(str(src / shard_name), framework="pt") as f: for key in keys: + fuse_m = cls.FUSABLE_GATE_UP_RE.match(key) t = f.get_tensor(key) - tensors[key] = t.to(target_dtype) if t.is_floating_point() else t + t = t.to(target_dtype) if t.is_floating_point() else t + if fuse_m: + prefix, kind = fuse_m.group(1), fuse_m.group(2) + fusable.setdefault(prefix, {})[kind] = t + else: + tensors[key] = t + + for prefix, parts in fusable.items(): + gate, up = parts.get("gate_proj"), parts.get("up_proj") + if gate is not None and up is not None: + # gate_up_proj(x).chunk(2, dim=-1) splits the *output* into gate-then-up, + # so the fused weight's output rows (dim 0) must be gate rows then up rows. + tensors[f"{prefix}.gate_up_proj.weight"] = torch.cat([gate, up], dim=0) + else: + # Only one half present under this prefix (not a fusable pair after all, + # e.g. a coincidental name match) — keep it under its original key. + for kind, t in parts.items(): + tensors[f"{prefix}.{kind}.weight"] = t + atomic_save(tensors, out / new_base_name_for[shard_name]) # Phase 3: mixed I/O + memory — one thread per shard, capped at CPU count. @@ -513,8 +546,21 @@ def _convert_base(shard_name: str, keys: List[str]) -> None: for fut in as_completed(futures_base): fut.result() + fused_prefixes: set = set() + for key in base_entries: + fuse_m = cls.FUSABLE_GATE_UP_RE.match(key) + if fuse_m: + prefix = fuse_m.group(1) + sibling = f"{prefix}.{'up_proj' if fuse_m.group(2) == 'gate_proj' else 'gate_proj'}.weight" + if sibling in base_entries: + fused_prefixes.add(prefix) + for key, shard_name in base_entries.items(): - new_weight_map[key] = new_base_name_for[shard_name] + fuse_m = cls.FUSABLE_GATE_UP_RE.match(key) + if fuse_m and fuse_m.group(1) in fused_prefixes: + new_weight_map[f"{fuse_m.group(1)}.gate_up_proj.weight"] = new_base_name_for[shard_name] + else: + new_weight_map[key] = new_base_name_for[shard_name] write_index(out, new_weight_map) sentinel.touch() diff --git a/QEfficient/transformers/cache_utils.py b/QEfficient/transformers/cache_utils.py index 550c314170..e60906a72f 100755 --- a/QEfficient/transformers/cache_utils.py +++ b/QEfficient/transformers/cache_utils.py @@ -25,6 +25,7 @@ ctx_scatter_3d, ctx_scatter_cb, ctx_scatter_cb_3d, + m3_ctx_scatter, ) @@ -80,6 +81,51 @@ def _remainder_with_symbolic_divisor(value: torch.Tensor, divisor) -> torch.Tens return torch.remainder(value, divisor_tensor) +def read_kv_cache_with_indices( + key_cache: torch.Tensor, + value_cache: torch.Tensor, + token_indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather selected key/value slices from a flat [batch, heads, ctx_len, head_dim] cache. + + token_indices: int32 tensor of shape [batch, heads, n_tokens] where invalid entries are 0. + Returns selected_k, selected_v both of shape [batch, heads, n_tokens, head_dim]. + """ + return ctx_gather_blocked_kv(key_cache, token_indices), ctx_gather_blocked_kv(value_cache, token_indices) + + +def update_and_read_index_key_cache( + index_key_cache: torch.Tensor, + position_ids: torch.Tensor, + idx_k: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Scatter idx_k into index_key_cache at position_ids, then read all ctx_len index keys. + + Future positions (beyond max position_id) are masked with INT32_MAX (ONNX export) or 0 (eager). + Returns (gathered_index_keys [batch, 1, ctx_len, head_dim], updated_index_key_cache). + """ + index_key_cache = m3_ctx_scatter(index_key_cache, position_ids.to(torch.int32), idx_k) + batch, _, ctx_len, _ = index_key_cache.shape + ctx_indices = torch.arange(ctx_len, device=index_key_cache.device)[None, None, :] + gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1) + invalid_idx = torch.iinfo(torch.int32).max if torch.onnx.is_in_onnx_export() else 0 + ctx_indices = torch.where(ctx_indices > gather_limit, invalid_idx, ctx_indices).to(torch.int32) + ctx_indices = ctx_indices.expand(batch, 1, ctx_len) + return ctx_gather_blocked_kv(index_key_cache, ctx_indices), index_key_cache + + +def scatter_kv_into_cache( + key_cache: torch.Tensor, + value_cache: torch.Tensor, + position_ids: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Write key_states and value_states into the flat KV cache at position_ids.""" + pid = position_ids.to(torch.int32) + return ctx_scatter(key_cache, pid, key_states), ctx_scatter(value_cache, pid, value_states) + + class QEffDynamicLayer(CacheLayerMixin): is_compileable = False @@ -234,7 +280,7 @@ def read_only_blocked_K(self, start_index, end_index, cache_kwargs): position_ids = cache_kwargs.get("position_ids") batch_index = cache_kwargs.get("batch_index", None) batch, num_kv_heads, _, _ = k_out.shape - ctx_indices = torch.arange(start=start_index, end=end_index)[None, None, ...] + ctx_indices = torch.arange(start=start_index, end=end_index, device=position_ids.device)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1) invalid_mask = ctx_indices > gather_limit @@ -303,7 +349,7 @@ def read_only_blocked_K_batch(self, start_index, end_index, cache_kwargs, folded k_out = folded_cache T_block = end_index - start_index - ctx_indices = torch.arange(start=start_index, end=end_index)[None, None, ...] + ctx_indices = torch.arange(start=start_index, end=end_index, device=position_ids.device)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values gather_limit = torch.cat([gather_limit] * Hkv, dim=1).reshape(1, -1, 1) invalid_mask = ctx_indices > gather_limit @@ -339,7 +385,7 @@ def read_only_blocked_V(self, start_index, end_index, cache_kwargs): position_ids = cache_kwargs.get("position_ids") batch_index = cache_kwargs.get("batch_index", None) batch, num_kv_heads, _, _ = v_out.shape - ctx_indices = torch.arange(start=start_index, end=end_index)[None, None, ...] + ctx_indices = torch.arange(start=start_index, end=end_index, device=position_ids.device)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1) invalid_mask = ctx_indices > gather_limit @@ -385,7 +431,7 @@ def read_only_blocked_V_batch(self, start_index, end_index, cache_kwargs, folded else: v_out = folded_cache T_block = end_index - start_index - ctx_indices = torch.arange(start=start_index, end=end_index)[None, None, ...] + ctx_indices = torch.arange(start=start_index, end=end_index, device=position_ids.device)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values gather_limit = torch.cat([gather_limit] * Hkv, dim=1).reshape(1, -1, 1) invalid_mask = ctx_indices > gather_limit @@ -1073,6 +1119,89 @@ def from_legacy_cache( return cache +class QEffMiniMaxSparseCache(QEffDynamicCache): + """ + QEffDynamicCache extended with per-layer index key caches for MiniMax-M3 sparse attention. + + Dense layers use the inherited write_only() / update() interface unchanged. + Sparse layers additionally use update_index_key_cache() and read_kv_with_block_indices(). + """ + + def __init__(self, ddp_cache_data=None, *args, **kwargs): + super().__init__(ddp_cache_data, *args, **kwargs) + self.index_keys: dict[int, Optional[torch.Tensor]] = {} + + def update_index_key_cache( + self, + idx_k: torch.Tensor, + layer_idx: int, + position_ids: torch.Tensor, + ) -> torch.Tensor: + """Scatter idx_k into the index key cache and return all ctx_len gathered index keys.""" + gathered, updated = update_and_read_index_key_cache( + self.index_keys.get(layer_idx), position_ids, idx_k + ) + self.index_keys[layer_idx] = updated + return gathered + + def read_kv_with_block_indices( + self, + layer_idx: int, + token_indices: torch.Tensor, + token_valid: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Gather K/V by block-sparse token indices. Returns (selected_k, selected_v, flat_valid).""" + flat_indices = token_indices.flatten(-2) + flat_valid = token_valid.flatten(-2) + layer = self.layers[layer_idx] + selected_k, selected_v = read_kv_cache_with_indices(layer.keys, layer.values, flat_indices) + return selected_k, selected_v, flat_valid + + @classmethod + def from_legacy_cache( + cls, + past_key_values: Optional[Tuple] = None, + index_keys: Optional[dict] = None, + ) -> "QEffMiniMaxSparseCache": + """Build from a per-layer tuple cache. + + Each layer entry is either a 2-tuple ``(key, value)`` for dense layers or a + 3-tuple ``(key, value, index_key)`` for sparse layers. The legacy + ``index_keys`` dict is accepted for backward-compat but the inline + 3-tuple form takes precedence. + """ + cache = cls() + if past_key_values is not None: + for i, layer_tuple in enumerate(past_key_values): + if len(layer_tuple) == 3: + key_states, value_states, index_key = layer_tuple + cache.index_keys[i] = index_key + else: + key_states, value_states = layer_tuple + cache.layers.append(QEffDynamicLayer.from_tensors(key_states, value_states)) + if index_keys is not None: + cache.index_keys.update(index_keys) + return cache + + def to_legacy_cache(self) -> Tuple: + """Return per-layer tuples; sparse layers carry a 3-tuple ``(key, value, index_key)``.""" + layers = [] + for i, layer in enumerate(self.layers): + if i in self.index_keys and self.index_keys[i] is not None: + layers.append((layer.keys, layer.values, self.index_keys[i])) + else: + layers.append((layer.keys, layer.values)) + return tuple(layers) + + def to_kv_only_cache(self) -> Tuple: + """Return per-layer 2-tuples ``(key, value)`` only, without index keys.""" + return tuple((layer.keys, layer.values) for layer in self.layers) + + def get_index_keys_tuple(self) -> Tuple: + """Return index keys as a flat tuple ordered by layer index.""" + return tuple(self.index_keys[i] for i in sorted(self.index_keys.keys())) + + class QEffEncoderDecoderCache(EncoderDecoderCache): """ Updated the `EncoderDecoderCache` to use the `QEffDynamicCache` for both self-attention and cross-attention caches. @@ -1613,7 +1742,7 @@ def read_only_blocked_K( batch, num_kv_heads, _, _ = k_out.shape - ctx_indices = torch.arange(start=start_idx, end=end_idx)[None, None, ...] + ctx_indices = torch.arange(start=start_idx, end=end_idx, device=position_ids.device)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1) invalid_mask = ctx_indices > gather_limit if torch.onnx.is_in_onnx_export(): @@ -1644,7 +1773,7 @@ def read_only_blocked_V( batch, num_kv_heads, _, _ = v_out.shape - ctx_indices = torch.arange(start=start_idx, end=end_idx)[None, None, ...] + ctx_indices = torch.arange(start=start_idx, end=end_idx, device=position_ids.device)[None, None, ...] gather_limit = position_ids.max(1, keepdim=True).values.unsqueeze(1) invalid_mask = ctx_indices > gather_limit if torch.onnx.is_in_onnx_export(): diff --git a/QEfficient/transformers/models/gemma3/modeling_gemma3.py b/QEfficient/transformers/models/gemma3/modeling_gemma3.py index ad86869fab..48275293ea 100644 --- a/QEfficient/transformers/models/gemma3/modeling_gemma3.py +++ b/QEfficient/transformers/models/gemma3/modeling_gemma3.py @@ -21,11 +21,13 @@ Gemma3ForConditionalGeneration, Gemma3TextConfig, Gemma3TextModel, - logger, repeat_kv, rotate_half, ) +import logging +logger = logging.getLogger(__name__) + from QEfficient.customop.rms_norm import CustomRMSNorm from QEfficient.transformers.cache_utils import QEffSlidingWindowCache from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask diff --git a/QEfficient/transformers/models/minimax_m3_vl/__init__.py b/QEfficient/transformers/models/minimax_m3_vl/__init__.py new file mode 100644 index 0000000000..1d70974f3d --- /dev/null +++ b/QEfficient/transformers/models/minimax_m3_vl/__init__.py @@ -0,0 +1,65 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +from transformers.models.minimax_m3_vl.configuration_minimax_m3_vl import ( + MiniMaxM3VLConfig, + MiniMaxM3VLTextConfig, + MiniMaxM3VLVisionConfig, +) +from transformers.models.minimax_m3_vl.modeling_minimax_m3_vl import ( + MiniMaxM3SparseForConditionalGeneration, + MiniMaxM3VLAttention, + MiniMaxM3VLDecoderLayer, + MiniMaxM3VLDenseMLP, + MiniMaxM3VLForCausalLM, + MiniMaxM3VLIndexer, + MiniMaxM3VLModel, + MiniMaxM3VLRMSNorm, + MiniMaxM3VLSparseMoeBlock, + MiniMaxM3VLTextModel, + MiniMaxM3VLTopKRouter, +) + +from QEfficient.transformers.models.minimax_m3_vl.modeling_minimax_m3_vl import ( + QEffMiniMaxM3SparseForConditionalGeneration, + QEffMiniMaxM3VLAttention, + QEffMiniMaxM3VLDecoderLayer, + QEffMiniMaxM3VLDenseMLP, + QEffMiniMaxM3VLForCausalLM, + QEffMiniMaxM3VLIndexer, + QEffMiniMaxM3VLRotaryEmbedding, + QEffMiniMaxM3VLSparseMoeBlock, + QEffMiniMaxM3VLTextModel, + QEffMiniMaxM3VLTopKRouter, +) + +__all__ = [ + "MiniMaxM3SparseForConditionalGeneration", + "MiniMaxM3VLAttention", + "MiniMaxM3VLConfig", + "MiniMaxM3VLDecoderLayer", + "MiniMaxM3VLDenseMLP", + "MiniMaxM3VLForCausalLM", + "MiniMaxM3VLIndexer", + "MiniMaxM3VLModel", + "MiniMaxM3VLRMSNorm", + "MiniMaxM3VLSparseMoeBlock", + "MiniMaxM3VLTextConfig", + "MiniMaxM3VLTextModel", + "MiniMaxM3VLTopKRouter", + "MiniMaxM3VLVisionConfig", + "QEffMiniMaxM3SparseForConditionalGeneration", + "QEffMiniMaxM3VLAttention", + "QEffMiniMaxM3VLDecoderLayer", + "QEffMiniMaxM3VLDenseMLP", + "QEffMiniMaxM3VLForCausalLM", + "QEffMiniMaxM3VLIndexer", + "QEffMiniMaxM3VLRotaryEmbedding", + "QEffMiniMaxM3VLSparseMoeBlock", + "QEffMiniMaxM3VLTextModel", + "QEffMiniMaxM3VLTopKRouter", +] \ No newline at end of file diff --git a/QEfficient/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py b/QEfficient/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py new file mode 100644 index 0000000000..2ae1cc4f4d --- /dev/null +++ b/QEfficient/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py @@ -0,0 +1,983 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +from typing import List, Optional, Tuple, Type, Union +from functools import partial +import logging +logging.getLogger("QEfficient").setLevel(logging.INFO) +import torch +import torch.nn.functional as F +from torch import nn +from transformers.cache_utils import Cache +from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast +from transformers.models.minimax_m3_vl.modeling_minimax_m3_vl import ( + MiniMaxM3SparseForConditionalGeneration, + MiniMaxM3VLAttention, + MiniMaxM3VLDecoderLayer, + MiniMaxM3VLDenseMLP, + MiniMaxM3VLExperts, + MiniMaxM3VLForCausalLM, + MiniMaxM3VLIndexer, + MiniMaxM3VLRotaryEmbedding, + MiniMaxM3VLSparseMoeBlock, + MiniMaxM3VLTextModel, + MiniMaxM3VLTopKRouter, + dynamic_rope_update, + maybe_autocast, + repeat_kv, +) + +from QEfficient.transformers.cache_utils import ( + QEffDynamicCache, + QEffMiniMaxSparseCache, + read_kv_cache_with_indices, + scatter_kv_into_cache, + update_and_read_index_key_cache, +) +from QEfficient.blocking.attention_blocking import ( + AttentionBlockingConfig, + BlockingMode, + generic_blocked_attention_interface, +) +from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask +from QEfficient.transformers.moe import ( + MoEFlavour, + MoEProfile, + QEffMoEBlockMixin, + build_canonical_expert_weights, + delete_module_attrs, + minimax_clamped_glu_mlp, +) +from QEfficient.utils import constants +from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config +from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + first, second = x.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +def qeff_apply_rotary_pos_emb( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + rotary_dim = cos.shape[-1] + rotated_q = q[..., :rotary_dim] + passthrough_q = q[..., rotary_dim:] + rotated_k = k[..., :rotary_dim] + passthrough_k = k[..., rotary_dim:] + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + rotated_q = rotated_q * cos + rotate_half(rotated_q) * sin + rotated_k = rotated_k * cos + rotate_half(rotated_k) * sin + return torch.cat((rotated_q, passthrough_q), dim=-1), torch.cat((rotated_k, passthrough_k), dim=-1) + +class QEffMiniMaxM3VLRotaryEmbedding(MiniMaxM3VLRotaryEmbedding): + @torch.no_grad() + @dynamic_rope_update + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = self.inv_freq.to(device=x.device, dtype=torch.float32) + position_ids_expanded = position_ids[..., None].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): + freqs = position_ids_expanded.float() * inv_freq.view(1, 1, -1) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def qeff_eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + if attention_mask.dtype == torch.bool: + attn_weights = torch.where( + attention_mask, + torch.tensor(MIN_MASKED_ATTENTION_VALUE, dtype=torch.float32, device=attn_weights.device), + attn_weights, + ) + else: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights + + +class QEffMiniMaxM3VLIndexer(MiniMaxM3VLIndexer): + def _select_blocks( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + past_key_values: "QEffMiniMaxSparseCache", + layer_idx: int, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + cfg = self.config + batch, seq_len = hidden_states.shape[0], hidden_states.shape[1] + ctx_len = past_key_values.layers[layer_idx].keys.shape[2] + num_blocks = (ctx_len + cfg.index_block_size - 1) // cfg.index_block_size + + idx_q = self.q_proj(hidden_states).view(batch, seq_len, cfg.index_n_heads, cfg.index_head_dim).transpose(1, 2) + idx_k = self.k_proj(hidden_states).view(batch, seq_len, 1, cfg.index_head_dim).transpose(1, 2) + idx_q = self.q_norm(idx_q) + idx_k = self.k_norm(idx_k) + idx_q, idx_k = qeff_apply_rotary_pos_emb( + idx_q, + idx_k, + cos[..., : cfg.index_head_dim], + sin[..., : cfg.index_head_dim], + cfg.index_head_dim // 2, + ) + idx_k = past_key_values.update_index_key_cache(idx_k, layer_idx, position_ids) + + scores = torch.matmul(idx_q.float(), idx_k.float().transpose(-1, -2)) + causal_mask = _create_causal_mask(position_ids=position_ids, target_length=ctx_len) + scores = scores.masked_fill(causal_mask, -1.0e30) + + padded_len = num_blocks * cfg.index_block_size + if padded_len != ctx_len: + scores = F.pad(scores, (0, padded_len - ctx_len), value=-1.0e30) + block_scores = scores.view(batch, cfg.index_n_heads, seq_len, num_blocks, cfg.index_block_size).amax(dim=-1) + + block_ids = torch.arange(num_blocks, device=hidden_states.device).view(1, 1, 1, -1) + # position_ids may be -1 (padding sentinel); floor_divide decomposes to a Sign node + # that AIC's Decode backend rejects (COMPILE_UNSUPPORTED_NODE_AFTER_OPTIMIZE). Trunc + # division avoids that op, and the -1 case is clamped to 0 immediately below either way. + q_block = torch.div(position_ids, cfg.index_block_size, rounding_mode="trunc") + for local_offset in range(cfg.index_local_blocks): + local_block = (q_block - local_offset).clamp(min=0) + block_scores = torch.where( + block_ids == local_block[:, None, :, None], + torch.full_like(block_scores, 1.0e30), + block_scores, + ) + # torch.export requires topk's k to be a fixed constant, but the exact value of + # min(index_topk_blocks, num_blocks) depends on where the symbolic ctx_len falls + # relative to index_topk_blocks * index_block_size — torch.export can't keep that + # branch symbolic (min()/abs() force it to specialize to the trace-time shape). + # Instead, pad the block dimension up to at least index_topk_blocks (a one-directional, + # branch-free relation via sym_max) and always select a constant k=index_topk_blocks; + # the existing token_valid masking below already discards padding/out-of-range blocks. + target_blocks = torch.sym_max(cfg.index_topk_blocks, num_blocks) + block_scores = F.pad(block_scores, (0, target_blocks - num_blocks), value=-1.0e30) + k = cfg.index_topk_blocks + topk_scores, block_indices = torch.topk(block_scores, k=k, dim=-1) + block_valid = topk_scores > -1.0e29 + offsets = torch.arange(cfg.index_block_size, device=hidden_states.device).view(1, 1, 1, 1, -1) + token_indices = block_indices.unsqueeze(-1) * cfg.index_block_size + offsets + token_valid = block_valid.unsqueeze(-1) & (token_indices < ctx_len) + token_valid = token_valid & (token_indices <= position_ids[:, None, :, None, None]) + token_indices = token_indices[:, :, 0] + token_valid = token_valid[:, :, 0] + safe_indices = torch.where(token_valid, token_indices, torch.zeros_like(token_indices)).to(torch.int32) + return safe_indices, token_valid + + +class QEffMiniMaxM3VLAttention(MiniMaxM3VLAttention): + def _baseline_attention( + self, + query_states: torch.Tensor, + selected_k: torch.Tensor, + selected_v: torch.Tensor, + flat_valid: torch.Tensor, + ) -> torch.Tensor: + batch, seq_len = query_states.shape[0], query_states.shape[2] + num_heads = self.config.index_n_heads + num_kv_groups = self.num_key_value_groups + q = query_states.view(batch, num_heads, num_kv_groups, seq_len, self.head_dim) + selected_k = selected_k.unsqueeze(2) + selected_v = selected_v.unsqueeze(2) + scores = torch.matmul(q.float(), selected_k.transpose(-1, -2).float()) * (self.head_dim**-0.5) + scores = scores.masked_fill(~flat_valid[:, :, None, None, :], -1.0e30) + probs = torch.softmax(scores, dim=-1, dtype=torch.float32).to(selected_v.dtype) + return torch.matmul(probs, selected_v) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Cache] = None, + position_ids: Optional[torch.Tensor] = None, + **kwargs, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + input_shape = hidden_states.shape[:-1] + query_shape = (*input_shape, self.config.num_attention_heads, self.head_dim) + key_value_shape = (*input_shape, self.config.num_key_value_heads, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states).view(query_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(key_value_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(key_value_shape).transpose(1, 2) + + cos, sin = position_embeddings + # to do - don't use constant, should get from config + query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos, sin, self.config.index_head_dim // 2) + + cache_kwargs = { + "position_ids": position_ids, + "batch_index": kwargs.get("batch_index"), + } + if attention_mask is not None: + cache_kwargs["CCL"] = attention_mask.shape[-1] + + if self.indexer is not None and isinstance(past_key_values, QEffMiniMaxSparseCache): + past_key_values.write_only(key_states, value_states, self.layer_idx, cache_kwargs) + token_indices, token_valid = self.indexer._select_blocks( + hidden_states, position_ids, past_key_values, self.layer_idx, cos, sin + ) + selected_k, selected_v, flat_valid = past_key_values.read_kv_with_block_indices( + self.layer_idx, token_indices, token_valid + ) + attn_output = self._baseline_attention(query_states, selected_k, selected_v, flat_valid) + attn_output = attn_output.reshape(*input_shape, self.config.num_attention_heads * self.head_dim) + else: + blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig()) + use_blocking = blocking_config is not None and blocking_config.mode != BlockingMode.NONE + if use_blocking: + past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0 + attn_output, _ = generic_blocked_attention_interface( + module=self, + query=query_states, + key=key_states, + value=value_states, + attention_mask=attention_mask, + scaling=self.scaling, + layer_idx=self.layer_idx, + past_key_value=past_key_values, + blocking_config=blocking_config, + comp_ctx_lengths=kwargs.get("comp_ctx_lengths"), + batch_index=kwargs.get("batch_index"), + position_ids=position_ids, + past_seen_tokens=past_seen_tokens, + prefill_only=blocking_config.mode.is_prefill, + ) + else: + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) + attn_output, _ = qeff_eager_attention_forward( + self, query_states, key_states, value_states, attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + ) + attn_output = attn_output.reshape(*input_shape, self.config.num_attention_heads * self.head_dim) + + return self.o_proj(attn_output.contiguous()), None + + +def _qeff_minimax_clamp(hidden_states: torch.Tensor, min_value=None, max_value=None) -> torch.Tensor: + # torch.clamp accepts plain Python scalars directly, unlike torch.maximum/minimum + # against a torch.tensor(..., device=hidden_states.device) constant -- the latter + # becomes a dataless meta-device lifted tensor when hidden_states is a meta tensor + # (weight-free export), which fails ONNX serialization. + return hidden_states.clamp(min=min_value, max=max_value) + + +class QEffMiniMaxM3VLDenseMLP(MiniMaxM3VLDenseMLP): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate_up = self.gate_up_proj(hidden_states) + gate, up = gate_up.chunk(2, dim=-1) + gate = _qeff_minimax_clamp(gate, max_value=self.swiglu_limit) + up = _qeff_minimax_clamp(up, min_value=-self.swiglu_limit, max_value=self.swiglu_limit) + glu = gate * torch.sigmoid(gate * self.swiglu_alpha) + return self.down_proj((up + 1.0) * glu) + + +class QEffMiniMaxM3VLExperts(MiniMaxM3VLExperts): + def __qeff_init__(self): + self.weights_transformed = False + + def transform_weights(self): + if getattr(self, "weights_transformed", False): + return self.moe_weights + self.moe_weights = build_canonical_expert_weights( + gate_up=self.gate_up_proj, + down=self.down_proj, + fused=True, + fused_split_dim=1, + transpose_gate_up=True, + transpose_down=True, + clone=True, + ) + delete_module_attrs(self, "gate_up_proj", "down_proj") + self.weights_transformed = True + return self.moe_weights + + +class QEffMiniMaxM3VLTopKRouter(MiniMaxM3VLTopKRouter): + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = nn.functional.linear(hidden_states.to(self.weight.dtype), self.weight) + routing_weights = nn.functional.sigmoid(router_logits.float()) + # e_score_correction_bias is a registered buffer, not a parameter; accelerate's + # init_empty_weights() (used for weight-free/meta export) doesn't move it to meta, + # so it can stay on cpu while routing_weights is on meta — move it explicitly. + scores_for_choice = routing_weights + self.e_score_correction_bias.to(device=routing_weights.device) + _, top_k_index = torch.topk(scores_for_choice, self.top_k, dim=1, sorted=False) + top_k_weights = routing_weights.gather(1, top_k_index) + denom = torch.einsum("tk->t", top_k_weights) + top_k_weights = top_k_weights / denom[:, None] + # routing_weights (and thus top_k_weights) is float32 from the sigmoid upcast above; + # cast back to the router's native dtype before returning, mirroring upstream HF's + # `current.to(final.dtype)` cast in MiniMaxM3VLExperts.forward. Without this, the + # float32 weight silently promotes moe_decode_bmm's output to float32 (via + # `experts_out = down * topk_weights`), which is invisible when the MoE layer is + # the last layer in a truncated model but fails ("expected float32 but found + # float16") the moment a later float16-weighted layer consumes it. + top_k_weights = top_k_weights.to(router_logits.dtype) + return router_logits, top_k_weights, top_k_index + + +class QEffMiniMaxM3VLSparseMoeBlock(QEffMoEBlockMixin, MiniMaxM3VLSparseMoeBlock): + _moe_return_router_logits = False + supported_moe_flavours = ( + MoEFlavour.SIMPLE_LOOP, + MoEFlavour.DECODE_BMM, + MoEFlavour.EXPERT_PARALLEL, + ) + + def __qeff_init__(self): + super().__qeff_init__() + self.top_k = getattr(self.gate, "top_k", None) + + def transform_weights(self): + if getattr(self, "weights_transformed", False): + return self.moe_weights + weights = self.experts.transform_weights() + # Plain attribute (bypassing nn.Module.__setattr__'s submodule registration) — + # `weights` is already registered as a submodule under self.experts.moe_weights. + # Re-registering the same MoEWeights instance under a second attribute path here + # would make torch.export/torch.onnx.export emit a duplicate set of gate/up/down + # initializers (one per FQN), and promote_initializers_and_build_spec (which dedups + # by named_parameters() identity) would only promote one of the two duplicate + # names, leaving the other as an un-promoted meta tensor that fails at ONNX save + # time ("Cannot copy out of meta tensor") — see QEffMiniMaxM3VLDecoderWrapper for + # the same failure mode with self.language_model. + object.__setattr__(self, "moe_weights", weights) + self.weights_transformed = True + return self.moe_weights + + @property + def moe_profile(self) -> MoEProfile: + return MoEProfile( + expert_mlp=partial( + minimax_clamped_glu_mlp, + limit=self.experts.swiglu_limit, + alpha=self.experts.swiglu_alpha, + ) + ) + + def route(self, x: torch.Tensor): + router_logits, top_w, top_i = self.gate(x) + return (top_i, top_w), router_logits + + def apply_shared_experts(self, out: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + return out * self.routed_scaling_factor + self.shared_experts(residual) + + +class QEffMiniMaxM3VLDecoderLayer(MiniMaxM3VLDecoderLayer): + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + **kwargs, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class QEffMiniMaxM3VLTextModel(MiniMaxM3VLTextModel): + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + **kwargs, + ) -> MoeModelOutputWithPast: + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + use_legacy_cache = False + if use_cache and not isinstance(past_key_values, Cache): + use_legacy_cache = True + past_key_values = QEffMiniMaxSparseCache.from_legacy_cache(past_key_values) + elif use_cache and past_key_values is None: + past_key_values = QEffMiniMaxSparseCache() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + if position_ids is None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + target_length = past_seen_tokens + inputs_embeds.shape[1] + if isinstance(attention_mask, torch.Tensor): + target_length = attention_mask.shape[-1] + elif past_key_values is not None and getattr(past_key_values, "layers", None): + first_layer = past_key_values.layers[0] + cached_keys = getattr(first_layer, "keys", None) + if cached_keys is not None: + target_length = cached_keys.shape[-2] + causal_mask = _create_causal_mask(position_ids=position_ids, target_length=target_length) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + + if use_legacy_cache: + past_key_values = past_key_values.to_legacy_cache() + + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +class QEffMiniMaxM3VLForCausalLM(MiniMaxM3VLForCausalLM): + def get_submodules_for_export(self) -> Type[nn.Module]: + return {QEffMiniMaxM3VLDecoderLayer} + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + **kwargs, + ) -> Union[Tuple, MoeCausalLMOutputWithPast]: + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + if position_ids is None: + hidden_states = outputs.last_hidden_state[:, -1:, :] + logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float() + else: + logit_idx = position_ids.to(torch.int32).argmax(1, keepdim=True) + hidden_states = outputs.last_hidden_state[torch.arange(position_ids.shape[0]).view(-1, 1), logit_idx] + logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float() + return MoeCausalLMOutputWithPast( + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + router_logits=getattr(outputs, "router_logits", None), + ) + + +class QEffMiniMaxM3VLEncoderWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.config = model.config + + def get_submodules_for_export(self) -> Type[nn.Module]: + layers = getattr(getattr(self.model.model, "vision_tower", None), "layers", None) + if layers: + return {layers[0].__class__} + return set() + + def forward(self, pixel_values, image_grid_thw): + image_outputs = self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw) + image_embeds = image_outputs.pooler_output if hasattr(image_outputs, "pooler_output") else image_outputs + if isinstance(image_embeds, (list, tuple)): + image_embeds = torch.cat(image_embeds, dim=0) + image_embeds = image_embeds.to(pixel_values.device, pixel_values.dtype) + bs = image_grid_thw.shape[0] + split_size = torch.floor_divide(torch.tensor(image_embeds.size(0), device=image_embeds.device), bs) + image_embeds = image_embeds.reshape(bs, split_size, image_embeds.size(-1)) + return image_embeds + + +class QEffMiniMaxM3VLDecoderWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + # Plain attribute (bypassing nn.Module.__setattr__'s submodule registration) — + # this is the same MiniMaxM3VLTextModel already reachable via self.model.model.language_model. + # Registering it as a second nn.Module attribute would make torch.export/torch.onnx.export + # emit a duplicate set of initializers under a second FQN for every text-model parameter, + # and promote_initializers_and_build_spec (which dedups by named_parameters() identity) + # would only promote one of the two duplicate initializer names, leaving the other as an + # un-promoted meta tensor that fails at ONNX save time ("Cannot copy out of meta tensor"). + object.__setattr__(self, "language_model", self.model.model.language_model) + self.config = model.config + + def get_submodules_for_export(self) -> Type[nn.Module]: + return {QEffMiniMaxM3VLDecoderLayer} + + def get_onnx_past_key_value_names(self, layer_idx: int, layer_state=None) -> List[str]: + return [f"past_key.{layer_idx}", f"past_value.{layer_idx}"] + + def get_onnx_index_key_names(self) -> List[str]: + layer_types = getattr(self.config.text_config, "layer_types", None) or getattr( + self.language_model.config, "layer_types", None + ) + if layer_types is None: + return [] + index_key_names = [] + for i in range(self.config.text_config.num_hidden_layers): + if layer_types[i] == "minimax_m3_sparse": + index_key_names.append(f"index_key.{i}" ) + return index_key_names + + def forward( + self, + input_ids=None, + vision_embeds=None, + position_ids=None, + image_idx=None, + past_key_values=None, + inputs_embeds: Optional[torch.FloatTensor] = None, + batch_index: Optional[torch.LongTensor] = None, + comp_ctx_lengths: Optional[List[int]] = None, + index_keys=None, + ): + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Exactly one of input_ids or inputs_embeds must be provided.") + + if inputs_embeds is None: + inputs_embeds = self.model.model.language_model.embed_tokens(input_ids) + _, _, hidden_dim = inputs_embeds.shape + selected = input_ids == self.config.image_token_index + indices1 = selected.to(torch.int64).cumsum(1) - 1 + indices1 = torch.where(indices1 != -1, indices1 + image_idx, indices1) + indices0 = torch.arange(selected.shape[0], device=selected.device).view(-1, 1) + image_features_expanded = vision_embeds.reshape(-1, hidden_dim).unsqueeze(0)[indices0, indices1] + image_input_embeds = torch.where(selected.unsqueeze(-1), image_features_expanded, inputs_embeds) + inputs_embeds = torch.where( + # Plain scalar constant for the shape comparison — must not be created on + # input_ids.device: for weight-free export the traced example inputs are on + # the meta device, so torch.tensor(1, device=input_ids.device) would create + # a dataless meta constant that torch.export lifts into the graph as-is, + # and ONNX serialization then fails ("Cannot copy out of meta tensor"). + input_ids.shape[1] == torch.tensor(1), inputs_embeds, image_input_embeds + ) + image_idx_output = (indices1.max() + 1).unsqueeze(0).unsqueeze(0) + else: + if image_idx is None: + image_idx = torch.zeros((1, 1), dtype=torch.int64, device=inputs_embeds.device) + image_idx_output = image_idx + + # Build a QEffMiniMaxSparseCache combining the KV cache (2-tuples) and the separate index keys. + sparse_layer_indices = [int(name.split(".")[-1]) for name in self.get_onnx_index_key_names()] + index_keys_dict = None + if index_keys is not None: + index_keys_dict = {sparse_layer_indices[j]: index_keys[j] for j in range(len(sparse_layer_indices))} + cache = QEffMiniMaxSparseCache.from_legacy_cache(past_key_values, index_keys=index_keys_dict) + + outputs = self.language_model( + inputs_embeds=inputs_embeds, + position_ids=position_ids, + past_key_values=cache, + comp_ctx_lengths=comp_ctx_lengths, + batch_index=batch_index, + use_cache=True, + ) + logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) + hidden_states = outputs.last_hidden_state[ + torch.arange(position_ids.shape[0], device=position_ids.device).view(-1, 1), logit_index + ] + logits = self.model.lm_head(hidden_states.to(self.model.lm_head.weight.dtype)).float() + + result_cache = outputs.past_key_values + if isinstance(result_cache, QEffMiniMaxSparseCache): + past_kv_out = result_cache.to_kv_only_cache() + index_keys_out = result_cache.get_index_keys_tuple() + else: + past_kv_out = tuple((t[0], t[1]) for t in result_cache) + index_keys_out = tuple(t[2] for t in result_cache if len(t) == 3) + + return logits, vision_embeds.clone(), image_idx_output, past_kv_out, index_keys_out + + +class QEffMiniMaxM3SparseForConditionalGeneration(MiniMaxM3SparseForConditionalGeneration): + def __qeff_init__(self): + # Plain attribute — see QEffMiniMaxM3VLDecoderWrapper.__init__ for why this must not + # be a second nn.Module registration of the same self.model.language_model submodule. + object.__setattr__(self, "language_model", self.model.language_model) + self.config._attn_implementation = "eager" + self.model.language_model.config._attn_implementation = "eager" + self.model.vision_tower.config._attn_implementation = "eager" + + def get_qeff_vision_encoder(self): + return QEffMiniMaxM3VLEncoderWrapper(self) + + def get_qeff_language_decoder(self): + return QEffMiniMaxM3VLDecoderWrapper(self) + + def forward( + self, + input_ids=None, + position_ids=None, + pixel_values=None, + image_grid_thw=None, + image_idx=None, + past_key_values=None, + comp_ctx_lengths: Optional[List[int]] = None, + batch_index: Optional[torch.LongTensor] = None, + **kwargs, + ): + if input_ids is None or position_ids is None or pixel_values is None or image_grid_thw is None: + raise ValueError("input_ids, position_ids, pixel_values, and image_grid_thw must be provided.") + if image_idx is None: + image_idx = torch.zeros((1, 1), dtype=torch.int64, device=input_ids.device) + + image_features = self.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw) + if hasattr(image_features, "pooler_output"): + image_features = image_features.pooler_output + image_features = image_features.to(device=input_ids.device, dtype=self.lm_head.weight.dtype) + + inputs_embeds = self.model.language_model.embed_tokens(input_ids) + _, _, hidden_dim = inputs_embeds.shape + selected = input_ids == self.config.image_token_index + indices1 = selected.to(torch.int64).cumsum(1) - 1 + indices1 = torch.where(indices1 != -1, indices1 + image_idx, indices1) + indices0 = torch.arange(selected.shape[0], device=selected.device).view(-1, 1) + image_features_expanded = image_features.reshape(-1, hidden_dim).unsqueeze(0)[indices0, indices1] + image_input_embeds = torch.where(selected.unsqueeze(-1), image_features_expanded, inputs_embeds) + inputs_embeds = torch.where( + # See QEffMiniMaxM3VLDecoderWrapper.forward for why this constant must not be + # created on input_ids.device. + input_ids.shape[1] == torch.tensor(1), inputs_embeds, image_input_embeds + ) + + if past_key_values is not None and not isinstance(past_key_values, Cache): + past_key_values = QEffDynamicCache.from_legacy_cache(past_key_values) + + outputs = self.model.language_model( + inputs_embeds=inputs_embeds, + position_ids=position_ids, + past_key_values=past_key_values, + comp_ctx_lengths=comp_ctx_lengths, + batch_index=batch_index, + use_cache=True, + ) + logit_index = position_ids.to(torch.int32).argmax(1, keepdim=True) + hidden_states = outputs.last_hidden_state[torch.arange(position_ids.shape[0]).view(-1, 1), logit_index] + logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float() + image_idx = (indices1.max() + 1).unsqueeze(0).unsqueeze(0) + + present = outputs.past_key_values + if isinstance(present, Cache): + if hasattr(present, "to_legacy_cache"): + present = present.to_legacy_cache() + elif hasattr(present, "layers"): + legacy_cache = () + for layer in present.layers: + legacy_cache += ((getattr(layer, "keys", None), getattr(layer, "values", None)),) + present = legacy_cache + return logits, pixel_values, image_idx, present + + def get_specializations( + self, + batch_size: int, + prefill_seq_len: int, + ctx_len: int, + comp_ctx_lengths_prefill: Optional[List[int]] = None, + comp_ctx_lengths_decode: Optional[List[int]] = None, + kv_offload: bool = False, + continuous_batching: bool = False, + kv_cache_batch_size: Optional[int] = None, + full_batch_size: Optional[int] = None, + **compiler_options, + ): + prefill_seq_len = prefill_seq_len if prefill_seq_len else constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + ctx_len = ctx_len if ctx_len else constants.ONNX_EXPORT_CTX_LEN + # img_size is accepted by generic VLM compile APIs, but MiniMax-M3 VLM + # specialization derives language/vision shapes from patch settings. + # Drop it to avoid leaking `-img-size=None` into qaic-compile flags. + compiler_options.pop("img_size", None) + num_image_patches = int(compiler_options.pop("num_image_patches", 4)) + num_images = int(compiler_options.pop("num_images", 1)) + vision_size = int(compiler_options.pop("vision_size", num_image_patches)) + + def _build_spec(seq_len, comp_ctx_lengths=None): + spec = { + "batch_size": full_batch_size if (continuous_batching and seq_len == 1) else batch_size, + "seq_len": seq_len, + "ctx_len": ctx_len, + "num_image_patches": num_image_patches, + "num_images": num_images, + } + if continuous_batching: + spec["full_batch_size"] = kv_cache_batch_size + if full_batch_size: + spec["full_batch_exec_size"] = full_batch_size + if comp_ctx_lengths is not None: + spec["comp_ctx_lengths"] = comp_ctx_lengths + return spec + + def _build_lang_spec(seq_len, comp_ctx_lengths=None): + spec = { + "batch_size": full_batch_size if (continuous_batching and seq_len == 1) else batch_size, + "seq_len": seq_len, + "ctx_len": ctx_len, + "vision_size": vision_size, + "vision_batch_size": batch_size, + } + if continuous_batching: + spec["full_batch_size"] = kv_cache_batch_size + if full_batch_size and seq_len != 1: + spec["full_batch_exec_size"] = full_batch_size + if comp_ctx_lengths is not None: + spec["comp_ctx_lengths"] = comp_ctx_lengths + return spec + + if comp_ctx_lengths_prefill and comp_ctx_lengths_decode: + specs = [_build_spec(prefill_seq_len, c) for c in comp_ctx_lengths_prefill] + specs.extend(_build_spec(1, c) for c in comp_ctx_lengths_decode) + else: + specs = [_build_spec(prefill_seq_len), _build_spec(1)] + + if kv_offload: + vision = [{"batch_size": batch_size, "num_image_patches": num_image_patches, "num_images": num_images}] + if comp_ctx_lengths_prefill and comp_ctx_lengths_decode: + lang = [_build_lang_spec(prefill_seq_len, c) for c in comp_ctx_lengths_prefill] + lang.extend(_build_lang_spec(1, c) for c in comp_ctx_lengths_decode) + else: + lang = [_build_lang_spec(prefill_seq_len), _build_lang_spec(1)] + return {"vision": vision, "lang": lang}, compiler_options + + return specs, compiler_options + + def get_onnx_dynamic_axes( + self, + comp_ctx_lengths: Optional[List[int]] = None, + kv_offload: bool = False, + continuous_batching: bool = False, + ): + vision_dynamic_axes = { + "pixel_values": {0: "num_image_patches"}, + "image_grid_thw": {0: "num_images"}, + } + + lang_dynamic_axes = { + "input_ids": {0: "batch_size", 1: "seq_len"}, + "position_ids": {0: "batch_size", 1: "seq_len"}, + "vision_embeds": {0: "vision_batch_size", 1: "vision_size"}, + "image_idx": {}, + } + + lm_config = self.model.language_model.config + layer_types = getattr(lm_config, "layer_types", None) or ["full_attention"] * lm_config.num_hidden_layers + for i in range(lm_config.num_hidden_layers): + lang_dynamic_axes[f"past_key.{i}"] = { + 0: "full_batch_size" if continuous_batching else "batch_size", + 2: "ctx_len", + } + lang_dynamic_axes[f"past_value.{i}"] = { + 0: "full_batch_size" if continuous_batching else "batch_size", + 2: "ctx_len", + } + if i < len(layer_types) and layer_types[i] == "minimax_m3_sparse": + lang_dynamic_axes[f"index_key.{i}"] = { + 0: "full_batch_size" if continuous_batching else "batch_size", + 2: "ctx_len", + } + if continuous_batching: + lang_dynamic_axes["batch_index"] = {0: "batch_size"} + if comp_ctx_lengths is not None: + lang_dynamic_axes["comp_ctx_lengths"] = {0: "comp_ctx_lengths"} + + if kv_offload: + return {"vision": vision_dynamic_axes, "lang": lang_dynamic_axes} + + dynamic_axes = {**vision_dynamic_axes, **lang_dynamic_axes} + dynamic_axes.pop("vision_embeds") + return dynamic_axes + + def get_output_names(self, kv_offload: bool = False): + lm_config = self.model.language_model.config + layer_types = getattr(lm_config, "layer_types", None) or ["full_attention"] * lm_config.num_hidden_layers + vision_output_names = ["vision_embeds"] + output_names = ["logits", "pixel_values_RetainedState", "image_idx_output"] + for i in range(lm_config.num_hidden_layers): + output_names.append(f"past_key.{i}_RetainedState") + output_names.append(f"past_value.{i}_RetainedState") + for i in range(lm_config.num_hidden_layers): + if i < len(layer_types) and layer_types[i] == "minimax_m3_sparse": + output_names.append(f"index_key.{i}_RetainedState") + if kv_offload: + lang_output_names = ["logits", "vision_embeds_RetainedState", "image_idx_output"] + for i in range(lm_config.num_hidden_layers): + lang_output_names.append(f"past_key.{i}_RetainedState") + lang_output_names.append(f"past_value.{i}_RetainedState") + for i in range(lm_config.num_hidden_layers): + if i < len(layer_types) and layer_types[i] == "minimax_m3_sparse": + lang_output_names.append(f"index_key.{i}_RetainedState") + return {"vision": vision_output_names, "lang": lang_output_names} + return output_names + + def get_dummy_pkv_cache(self, config, batch_size, seq_len, dtype=None): + dtype = dtype or getattr(config, "torch_dtype", torch.float32) or torch.float32 + kv_cache_shape = get_padding_shape_from_config(config=config, batch_size=batch_size, seq_len=seq_len) + past_key_values = [] + for _ in range(config.num_hidden_layers): + k = torch.zeros(kv_cache_shape, dtype=dtype) + v = torch.zeros(kv_cache_shape, dtype=dtype) + past_key_values.append([k, v]) + return past_key_values + + def get_dummy_index_keys(self, config, batch_size, seq_len, dtype=None): + dtype = dtype or getattr(config, "torch_dtype", torch.float32) or torch.float32 + layer_types = getattr(config, "layer_types", None) or ["full_attention"] * config.num_hidden_layers + index_key_shape = (batch_size, 1, seq_len, config.index_head_dim) + index_keys = [] + for i in range(config.num_hidden_layers): + if layer_types[i] == "minimax_m3_sparse": + index_keys.append(torch.zeros(index_key_shape, dtype=dtype)) + return index_keys + + def get_dummy_inputs( + self, + comp_ctx_lengths: Optional[List[int]] = None, + kv_offload: bool = False, + continuous_batching: bool = False, + **kwargs, + ): + prefill_seq_len = kwargs.get("prefill_seq_len") + if prefill_seq_len is None: + prefill_seq_len = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + prefill_seq_len = int(prefill_seq_len) + + batch_size = kwargs.get("batch_size", constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE) + fbs = constants.ONNX_EXPORT_EXAMPLE_FBS + dtype = getattr(self.config, "torch_dtype", torch.float32) or torch.float32 + + patch_dim = ( + self.config.vision_config.num_channels + * self.config.vision_config.temporal_patch_size + * self.config.vision_config.patch_size + * self.config.vision_config.patch_size + ) + image_grid_thw = torch.tensor([[1, 2, 2]], dtype=torch.int64) + num_image_patches = int(torch.prod(image_grid_thw).item()) + + inputs = { + "input_ids": torch.zeros((batch_size, prefill_seq_len), dtype=torch.int64), + "pixel_values": torch.zeros((num_image_patches, patch_dim), dtype=dtype), + "image_grid_thw": image_grid_thw, + "position_ids": torch.arange(prefill_seq_len, dtype=torch.int64) + .view(1, prefill_seq_len) + .repeat(batch_size, 1), + "image_idx": torch.zeros((1, 1), dtype=torch.int64), + } + inputs["input_ids"][:, 0] = self.config.image_token_index + past_key_values = self.get_dummy_pkv_cache( + config=self.model.language_model.config, + batch_size=fbs if continuous_batching else batch_size, + seq_len=prefill_seq_len, + dtype=dtype, + ) + index_keys = self.get_dummy_index_keys( + config=self.model.language_model.config, + batch_size=fbs if continuous_batching else batch_size, + seq_len=prefill_seq_len, + dtype=dtype, + ) + inputs["past_key_values"] = past_key_values + inputs["index_keys"] = index_keys + if continuous_batching: + inputs["batch_index"] = torch.arange(batch_size).view(batch_size, 1) + if comp_ctx_lengths is not None: + inputs["comp_ctx_lengths"] = torch.randint(0, 100, (40,), dtype=torch.int64) + if kv_offload: + vision_inputs = { + "pixel_values": inputs["pixel_values"], + "image_grid_thw": inputs["image_grid_thw"], + } + lang_inputs = { + "input_ids": inputs["input_ids"], + "vision_embeds": torch.zeros( + (batch_size, num_image_patches, self.model.language_model.config.hidden_size), dtype=dtype + ), + "position_ids": inputs["position_ids"], + "image_idx": inputs["image_idx"], + "past_key_values": past_key_values, + "index_keys": index_keys, + } + if continuous_batching: + lang_inputs["batch_index"] = inputs["batch_index"] + if comp_ctx_lengths is not None: + lang_inputs["comp_ctx_lengths"] = inputs["comp_ctx_lengths"] + return {"vision": vision_inputs, "lang": lang_inputs} + return inputs + + def get_inputs_info(self): + patch_dim = ( + self.config.vision_config.num_channels + * self.config.vision_config.temporal_patch_size + * self.config.vision_config.patch_size + * self.config.vision_config.patch_size + ) + return [ + IOInfo(name="input_ids", datatype=torch.int64, shape=("batch_size", "seq_len")), + IOInfo(name="pixel_values", datatype=self.config.torch_dtype, shape=("num_image_patches", patch_dim)), + IOInfo(name="image_grid_thw", datatype=torch.int64, shape=("num_images", 3)), + ] \ No newline at end of file diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 213986d41e..a2626b1c43 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -1156,12 +1156,17 @@ def export(self, inputs, output_names, dynamic_axes, export_dir=None, offload_pt If True, PyTorch weights will be offloaded after export. Default is True. use_onnx_subfunctions: bool, optional whether to enable ONNX subfunctions during export. Exporting PyTorch model to ONNX with modules as subfunctions helps to reduce export/compile time. Defaults to False + dynamo: bool, optional + whether to enable dynamo during export. Returns ------- str Path to the generated ONNX graph file for the vision encoder. """ + # Weight-free export always uses the dynamo (torch.export) path. + # Must be set here — @export_wrapper reads dynamo from kwargs before _export() body runs. + dynamo = kwargs.get("dynamo", False) or self._weight_free return self._export( inputs, output_names=output_names, @@ -1169,6 +1174,7 @@ def export(self, inputs, output_names, dynamic_axes, export_dir=None, offload_pt export_dir=export_dir, offload_pt_weights=offload_pt_weights, use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False), + dynamo=dynamo, ) def compile( @@ -1260,6 +1266,14 @@ class QEffCausalLMForTextImageToTextModel(QEFFBaseModel): ] _onnx_transforms = [] + _checkpoint_transforms = [ + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, + GraniteMoeFusedExpertSplitCheckpointTransform, + DtypeConversionCheckpointTransform, + ] + def __init__(self, model, qaic_config: Optional[dict] = None, **kwargs): """ Initializes the language decoder component for multimodal models. @@ -1333,6 +1347,8 @@ def export( If True, PyTorch weights will be offloaded after export. Default is True. use_onnx_subfunctions: bool, optional whether to enable ONNX subfunctions during export. Exporting PyTorch model to ONNX with modules as subfunctions helps to reduce export/compile time. Defaults to False + dynamo: bool, optional + whether to enable dynamo during export. Returns ------- @@ -1353,6 +1369,9 @@ def export( self.__update_prefill_transform(False, retain_full_kv=kwargs.get("retain_full_kv", False)) qaic_config = kwargs.pop("qaic_config", getattr(self.model, "qaic_config", None)) + # Weight-free export always uses the dynamo (torch.export) path. + # Must be set here — @export_wrapper reads dynamo from kwargs before _export() body runs. + dynamo = kwargs.get("dynamo", False) or self._weight_free if QEfficient.base.modeling_qeff.QEFFBaseModel._layerwise_active: return self._export_layerwise( @@ -1378,6 +1397,7 @@ def export( export_dir=export_dir, offload_pt_weights=offload_pt_weights, use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False), + dynamo=dynamo, ) def compile( @@ -1596,6 +1616,7 @@ def export( layerwise_window_size: int = 1, kv_cache_prefix: Optional[str] = None, offload_pt_weights: Optional[bool] = None, + dynamo: bool = False, **kwargs, ) -> str: """ @@ -1610,6 +1631,8 @@ def export( Directory path where the exported ONNX graphs will be saved. Default is None. use_onnx_subfunctions: bool, optional whether to enable ONNX subfunctions during export. Exporting PyTorch model to ONNX with modules as subfunctions helps to reduce export/compile time. Defaults to False + dynamo: bool, optional + whether to enable dynamo during export. **kwargs : Additional keyword arguments. @@ -1635,6 +1658,18 @@ def export( ) bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE seq_len: int = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + # TODO: Remove this hack ################## + if dynamo: + # torch.export requires example inputs to satisfy dynamic_shapes min=2 (batch_min). + bs = max(2, bs) + if getattr(self.model.config, "model_type", None) == "minimax_m3_vl": + # MiniMax's sparse-attention indexer requires ctx_len > index_block_size to + # export (see convert_dynamic_axes_to_dynamic_shapes); the KV-cache dummy here + # reuses seq_len as its traced ctx_len, so bump it past that boundary too. + index_block_size = getattr(self.model.config.text_config, "index_block_size", None) + if index_block_size is not None: + seq_len = max(seq_len, index_block_size + 1) + ########################################### qaic_config = kwargs.get("qaic_config", getattr(self.lang_model.model, "qaic_config", None)) # TODO: move this to a DA Serving utility class if self.model.config.model_type in SPECIALIZED_DISAGG_SERVING_MODEL_ARCH: @@ -1701,6 +1736,7 @@ def export( export_dir=export_dir, offload_pt_weights=False, use_onnx_subfunctions=use_onnx_subfunctions, + dynamo=dynamo, ) # TODO: remove the current pt weight offload capability once CustomLoader is in place @@ -1725,6 +1761,7 @@ def export( qaic_config=qaic_config, _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, + dynamo=dynamo, ) return self.onnx_path @@ -1909,6 +1946,7 @@ def compile( layerwise: bool = False, layerwise_window_size: int = 1, kv_cache_prefix: Optional[str] = None, + dynamo: bool = False, **compiler_options, ) -> str: """ @@ -2128,6 +2166,7 @@ def compile( _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, offload_pt_weights=offload_pt_weights, + dynamo=dynamo, ) if layerwise_cache_probe: return self.lang_model.onnx_path @@ -3373,6 +3412,7 @@ def from_pretrained( continuous_batching: bool = False, qaic_config: Optional[dict] = None, layerwise: bool = False, + weight_free: bool = False, **kwargs, ): """ @@ -3388,6 +3428,14 @@ def from_pretrained( If None, the default behavior of the internal classes is used (typically dual QPC). qaic_config : dict, optional A dictionary for QAIC-specific configurations. + weight_free : bool, optional + If True, builds the model on the meta device instead of loading real + checkpoint weights — no weights are materialized into RAM. This is + the single place to enable weight-free export; ``export()``/``compile()`` + automatically route through the weight-free path afterward, with no + further flag needed. The real checkpoint weights are supplied at + export time via ``pretrained_model_name_or_path``. Mutually exclusive + with ``layerwise=True``. Default is False. **kwargs : Additional arguments passed to HuggingFace's ``from_pretrained``. @@ -3404,6 +3452,13 @@ def from_pretrained( NotImplementedError If `continuous_batching` is provided as True. """ + if layerwise and weight_free: + raise ValueError( + "`layerwise=True` and `weight_free=True` are mutually exclusive; weight_free replaces layerwise mode." + ) + if weight_free: + validate_dynamo_export_requirements("weight_free=True") + enable_proxy = kwargs.pop("enable_proxy", False) # TODO: add a check to see if kv_offload is allowed for given model by loading the config and checking architecture or type of config here. @@ -3432,6 +3487,11 @@ def from_pretrained( # internally via the layer-wise driver, so the outer instance is # only used as a config holder. model = _build_meta_model(cls._hf_auto_class, pretrained_model_name_or_path, kwargs) + elif weight_free: + # Weight-free mode: build the model on the meta device so no + # checkpoint weights are ever materialized here. The real weights + # are supplied later at export time via pretrained_model_name_or_path. + model = _build_meta_model(cls._hf_auto_class, pretrained_model_name_or_path, kwargs) else: model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, **kwargs) @@ -3443,6 +3503,7 @@ def from_pretrained( continuous_batching=continuous_batching, pretrained_model_name_or_path=pretrained_model_name_or_path, qaic_config=qaic_config, + weight_free=weight_free, **kwargs, ) # Mark the wrapper so its compile() can default ``layerwise=True`` if diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index 7e2b3f80e6..9d4695b456 100644 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -151,6 +151,20 @@ MixtralRMSNorm, MixtralSparseMoeBlock, ) +from transformers.models.minimax_m3_vl.modeling_minimax_m3_vl import ( + MiniMaxM3SparseForConditionalGeneration, + MiniMaxM3VLAttention, + MiniMaxM3VLDecoderLayer, + MiniMaxM3VLDenseMLP, + MiniMaxM3VLExperts, + MiniMaxM3VLForCausalLM, + MiniMaxM3VLIndexer, + MiniMaxM3VLRMSNorm, + MiniMaxM3VLRotaryEmbedding, + MiniMaxM3VLSparseMoeBlock, + MiniMaxM3VLTextModel, + MiniMaxM3VLTopKRouter, +) from transformers.models.mllama.modeling_mllama import ( MllamaCrossAttentionDecoderLayer, MllamaForCausalLM, @@ -488,6 +502,19 @@ QEffMixtralModel, QEffMixtralSparseMoeBlock, ) +from QEfficient.transformers.models.minimax_m3_vl.modeling_minimax_m3_vl import ( + QEffMiniMaxM3SparseForConditionalGeneration, + QEffMiniMaxM3VLAttention, + QEffMiniMaxM3VLDecoderLayer, + QEffMiniMaxM3VLDenseMLP, + QEffMiniMaxM3VLExperts, + QEffMiniMaxM3VLForCausalLM, + QEffMiniMaxM3VLIndexer, + QEffMiniMaxM3VLRotaryEmbedding, + QEffMiniMaxM3VLSparseMoeBlock, + QEffMiniMaxM3VLTextModel, + QEffMiniMaxM3VLTopKRouter, +) from QEfficient.transformers.models.mllama.modeling_mllama import ( QEffMllamaCrossAttentionDecoderLayer, QEffMllamaForCausalLM, @@ -694,6 +721,7 @@ class CustomOpsTransform(ModuleMappingTransform): Qwen3VLMoeTextRMSNorm: CustomRMSNormAIC, Qwen3VLTextRMSNorm: CustomRMSNormAIC, Glm4MoeRMSNorm: CustomRMSNormAIC, + MiniMaxM3VLRMSNorm: GemmaCustomRMSNormAIC, Wav2Vec2Encoder: QEffWav2Vec2Encoder, Wav2Vec2EncoderStableLayerNorm: QEffWav2Vec2EncoderStableLayerNorm, # BERT-family: replace _create_attention_masks (uses create_bidirectional_mask, @@ -783,6 +811,16 @@ class KVCacheTransform(ModuleMappingTransform): Qwen3VLVisionModel: QEffQwen3VLVisionModel, Qwen3VLTextModel: QEffQwen3VLTextModel, Qwen3VLTextRotaryEmbedding: QEffQwen3VLTextRotaryEmbedding, + # MiniMaxM3VL + MiniMaxM3SparseForConditionalGeneration: QEffMiniMaxM3SparseForConditionalGeneration, + MiniMaxM3VLAttention: QEffMiniMaxM3VLAttention, + MiniMaxM3VLDecoderLayer: QEffMiniMaxM3VLDecoderLayer, + MiniMaxM3VLDenseMLP: QEffMiniMaxM3VLDenseMLP, + MiniMaxM3VLForCausalLM: QEffMiniMaxM3VLForCausalLM, + MiniMaxM3VLIndexer: QEffMiniMaxM3VLIndexer, + MiniMaxM3VLRotaryEmbedding: QEffMiniMaxM3VLRotaryEmbedding, + MiniMaxM3VLTextModel: QEffMiniMaxM3VLTextModel, + MiniMaxM3VLTopKRouter: QEffMiniMaxM3VLTopKRouter, # Gemma2 Gemma2Attention: QEffGemma2Attention, Gemma2DecoderLayer: QEffGemma2DecoderLayer, @@ -1311,6 +1349,18 @@ class KVCacheExternalModuleMapperTransform(ExternalModuleMapperTransform): "DeepseekV3RMSNorm": { "forward": QEffDeepseekV3CustomRMSNormAIC.forward, }, + "MiniMaxM3SparseForConditionalGeneration": {"forward": QEffMiniMaxM3SparseForConditionalGeneration.forward}, + "MiniMaxM3VLForCausalLM": { + "forward": QEffMiniMaxM3VLForCausalLM.forward, + "get_submodules_for_export": QEffMiniMaxM3VLForCausalLM.get_submodules_for_export, + }, + "MiniMaxM3VLTextModel": {"forward": QEffMiniMaxM3VLTextModel.forward}, + "MiniMaxM3VLDecoderLayer": {"forward": QEffMiniMaxM3VLDecoderLayer.forward}, + "MiniMaxM3VLDenseMLP": {"forward": QEffMiniMaxM3VLDenseMLP.forward}, + "MiniMaxM3VLAttention": {"forward": QEffMiniMaxM3VLAttention.forward}, + "MiniMaxM3VLRotaryEmbedding": {"forward": QEffMiniMaxM3VLRotaryEmbedding.forward}, + "MiniMaxM3VLTopKRouter": {"forward": QEffMiniMaxM3VLTopKRouter.forward}, + "MiniMaxM3VLSparseMoeBlock": {"forward": QEffMiniMaxM3VLSparseMoeBlock.forward}, } @@ -1517,6 +1567,9 @@ class OptimizedMoEMapperTransform(ModuleMappingTransform): GraniteMoeTopKGating: QEffGraniteMoeTopKGating, # Mixtral MixtralSparseMoeBlock: QEffMixtralSparseMoeBlock, + # MiniMaxM3VL + MiniMaxM3VLExperts: QEffMiniMaxM3VLExperts, + MiniMaxM3VLSparseMoeBlock: QEffMiniMaxM3VLSparseMoeBlock, } @classmethod diff --git a/QEfficient/transformers/moe/__init__.py b/QEfficient/transformers/moe/__init__.py index d6fe44271e..2cd1a734db 100644 --- a/QEfficient/transformers/moe/__init__.py +++ b/QEfficient/transformers/moe/__init__.py @@ -32,6 +32,7 @@ SILU_GLU_PROFILE, MoEProfile, gptoss_clamped_glu_mlp, + minimax_clamped_glu_mlp, silu_glu_mlp, ) from QEfficient.transformers.moe.weights import ( @@ -63,6 +64,7 @@ "SILU_GLU_PROFILE", "MoEProfile", "gptoss_clamped_glu_mlp", + "minimax_clamped_glu_mlp", "silu_glu_mlp", "MoEWeights", "as_parameters", diff --git a/QEfficient/transformers/moe/profiles.py b/QEfficient/transformers/moe/profiles.py index 641b90b79c..1846a23ca7 100644 --- a/QEfficient/transformers/moe/profiles.py +++ b/QEfficient/transformers/moe/profiles.py @@ -73,4 +73,28 @@ def gptoss_clamped_glu_mlp( return (intermediate @ W_d) + b_d.unsqueeze(-2) +def minimax_clamped_glu_mlp( + x: torch.Tensor, + W_g: torch.Tensor, + W_u: torch.Tensor, + W_d: torch.Tensor, + b_g: Optional[torch.Tensor] = None, + b_u: Optional[torch.Tensor] = None, + b_d: Optional[torch.Tensor] = None, + *, + limit: float, + alpha: float, +) -> torch.Tensor: + """MiniMax-M3 clamped GLU (no biases): ``(up + 1) * gate * sigmoid(gate * alpha)``. + + Same gating formula as GPT-OSS but MiniMax's experts have no bias terms and + the gate clamp has no explicit lower bound (only an upper bound at ``limit``). + """ + gate = (x @ W_g).clamp(max=limit) + up = (x @ W_u).clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + intermediate = (up + 1.0) * glu + return intermediate @ W_d + + SILU_GLU_PROFILE = MoEProfile(expert_mlp=silu_glu_mlp, has_bias=False) diff --git a/QEfficient/utils/_utils.py b/QEfficient/utils/_utils.py index 5e08a68b9b..857cc65e94 100755 --- a/QEfficient/utils/_utils.py +++ b/QEfficient/utils/_utils.py @@ -42,6 +42,7 @@ "k_pe.", "conv_state.", "recurrent_state.", + "index_key.", ) _RETAINED_STATE_SUFFIX = "_RetainedState" _INTERNAL_RETAINED_STATE_SUFFIX = "_InternalRetainedState" diff --git a/QEfficient/utils/export_utils.py b/QEfficient/utils/export_utils.py index 78bf2675be..da2c44cb2e 100644 --- a/QEfficient/utils/export_utils.py +++ b/QEfficient/utils/export_utils.py @@ -86,7 +86,7 @@ def build_dynamo_export_kwargs(export_kwargs): from QEfficient.utils import constants kwargs = dict(export_kwargs) - kwargs.setdefault("report", False) + kwargs.setdefault("report", True) kwargs.setdefault("optimize", False) kwargs["dynamo"] = True kwargs["opset_version"] = constants.ONNX_DYNAMO_EXPORT_OPSET @@ -125,10 +125,25 @@ def convert_dynamic_axes_to_dynamic_shapes( torch.export dynamic_shapes dict with Dim objects, suitable for torch.onnx.export(dynamic_shapes=...). """ - max_seq_len = getattr(model_config, "max_position_embeddings", 1024) + max_seq_len = getattr( + model_config, + "max_position_embeddings", + getattr(getattr(model_config, "text_config", None), "max_position_embeddings", 1024), + ) model_type = getattr(model_config, "model_type", None) batch_min = 1 if model_type == "gpt_oss" else 2 + # MiniMax's sparse-attention indexer reshapes ctx_len into (num_blocks, index_block_size) + # blocks and requires num_blocks >= 2 to export: torch.export inserts a broadcast-safety + # guard at the num_blocks==1 boundary that a single dynamic_shapes Dim can't satisfy across + # a range straddling it. ctx_len values <= index_block_size always give num_blocks == 1, so + # raise ctx_len's min past that boundary for this model only. + ctx_len_min = 2 + if model_type == "minimax_m3_vl": + index_block_size = getattr(getattr(model_config, "text_config", None), "index_block_size", None) + if index_block_size is not None: + ctx_len_min = index_block_size + 1 + dim_registry: Dict[str, Any] = {} def resolve_dim(dim_name: str): @@ -143,7 +158,7 @@ def resolve_dim(dim_name: str): elif "comp_ctx_lengths" in dim_name: dim_registry[dim_name] = Dim("comp_ctx_lengths", min=DYNAMO_DIM_MIN_COMP_CTX_LENGTHS, max=max_seq_len) elif "ctx_len" in dim_name: - dim_registry[dim_name] = Dim("ctx_len", min=2, max=max_seq_len) + dim_registry[dim_name] = Dim("ctx_len", min=ctx_len_min, max=max_seq_len) elif "sliding_window" in dim_name: dim_registry[dim_name] = Dim( "sliding_window", @@ -151,7 +166,7 @@ def resolve_dim(dim_name: str): max=getattr(model_config, "sliding_window", max_seq_len), ) else: - dim_registry[dim_name] = Dim.DYNAMIC + dim_registry[dim_name] = dim_name return dim_registry[dim_name] dynamic_shapes: Dict[str, Any] = {} @@ -159,6 +174,7 @@ def resolve_dim(dim_name: str): past_values: Dict[int, Any] = {} compressed_kv_layers: Dict[int, Any] = {} k_pe_layers: Dict[int, Any] = {} + index_key_layers: Dict[int, Any] = {} for input_name, axes_map in dynamic_axes.items(): resolved = {axis_idx: resolve_dim(dim_name) for axis_idx, dim_name in axes_map.items()} @@ -170,6 +186,8 @@ def resolve_dim(dim_name: str): compressed_kv_layers[int(input_name.split(".")[1])] = resolved elif input_name.startswith("k_pe."): k_pe_layers[int(input_name.split(".")[1])] = resolved + elif input_name.startswith("index_key."): + index_key_layers[int(input_name.split(".")[1])] = resolved else: dynamic_shapes[input_name] = resolved @@ -185,6 +203,13 @@ def resolve_dim(dim_name: str): (compressed_kv_layers.get(i, {}), k_pe_layers.get(i, {})) for i in range(max_layer + 1) ] + if index_key_layers: + # index_key.N only exists for sparse-attention layers (a subset of all decoder + # layers), so unlike past_key_values there is no gap-filling by full layer range — + # the aggregated list must match the compact "index_keys" list order from + # get_dummy_index_keys (ascending original layer index, sparse layers only). + dynamic_shapes["index_keys"] = [index_key_layers[i] for i in sorted(index_key_layers)] + return dynamic_shapes diff --git a/examples/text_generation/minimax_m3_decode_only.py b/examples/text_generation/minimax_m3_decode_only.py new file mode 100644 index 0000000000..74d083af5e --- /dev/null +++ b/examples/text_generation/minimax_m3_decode_only.py @@ -0,0 +1,208 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import os +import tempfile + +import torch +from transformers import AutoConfig, AutoProcessor, AutoTokenizer, AutoModelForImageTextToText + +from QEfficient import QEFFAutoModelForImageTextToText + +MODEL_ID = "MiniMaxAI/MiniMax-M3" + + +def _run_pytorch_parity_test( + model_id: str, + prompt: str, + export_dir: str, + ctx_len: int = 128, + num_cores: int = 16, + num_devices: int = 1, + expert_parallel_chunk_size: int = 256, + cores_per_expert: int = 2, + tree_reduce: bool = True, +) -> None: + """Compare HF PyTorch vs AIC on the last decode token of the prompt (prefill_seq_len=1).""" + # Load real model architecture with 4 layers for a fast test (random weights). + full_config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) + full_config.text_config.num_hidden_layers = 4 + + torch.manual_seed(42) + model_hf = AutoModelForImageTextToText.from_config(full_config).eval() + model_dir = os.path.join(export_dir, "minimax-m3-parity") + model_hf.save_pretrained(model_dir) + + # Tokenize the real prompt and take the last token as the single decode input. + processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) + messages = [[{"role": "user", "content": [{"type": "text", "text": prompt}]}]] + inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + last_token_ids = inputs["input_ids"][:, -1:] + + with torch.no_grad(): + hf_logits = model_hf.language_model(input_ids=last_token_ids, use_cache=False).logits[:, -1:, :] + expected_token = int(hf_logits.argmax(-1)[0, 0]) + + qeff_model = QEFFAutoModelForImageTextToText.from_pretrained(model_dir, torch_dtype=torch.float32) + qeff_model.compile( + batch_size=1, + prefill_seq_len=1, + ctx_len=ctx_len, + num_cores=num_cores, + num_devices=num_devices, + use_onnx_subfunctions=False, + skip_vision=True, + offload_pt_weights=False, + qaic_config={ + "moe_config": { + "flavour": "expert_parallel", + "expert_parallel_chunk_size": expert_parallel_chunk_size, + "cores_per_expert": cores_per_expert, + "tree_reduce": tree_reduce, + } + }, + ) + + aic_inputs = qeff_model.model.prepare_inputs_for_generation( + inputs={"input_ids": last_token_ids}, prefill_seq_len=1, batch_size=1 + ) + output = qeff_model.generate(inputs=aic_inputs, generation_len=1) + aic_token = int(output.generated_ids[0, 0]) + assert aic_token == expected_token, f"Parity check FAILED: expected {expected_token}, got {aic_token}" + print(f"[PASS] PyTorch vs AIC parity check passed (token={aic_token})") + + +def main(): + parser = argparse.ArgumentParser(description="MiniMax-M3 text-only decode (PL=1).") + parser.add_argument("--model-id", default=MODEL_ID) + parser.add_argument("--ctx-len", type=int, default=1000000) + parser.add_argument("--num-devices", type=int, default=16) + parser.add_argument("--num-cores", type=int, default=16) + parser.add_argument( + "--device-ids", + type=int, + nargs="+", + default=None, + help="Explicit QAIC device IDs to run generate() on (e.g. --device-ids 0 1 2 3). " + "Defaults to [0..num_devices-1]; set this if any device in that range is unhealthy.", + ) + parser.add_argument("--generation-len", type=int, default=32) + parser.add_argument("--prompt", default="Tell me about yourself.") + parser.add_argument("--num-layers", type=int, default=None) + parser.add_argument("--skip-generate", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument( + "--expert-parallel-chunk-size", + type=int, + default=256, + help="MoE expert-parallel chunk size (expert_parallel_chunk_size in moe_config).", + ) + parser.add_argument( + "--cores-per-expert", + type=int, + default=2, + help="Number of NSP cores assigned to each expert during decode.", + ) + parser.add_argument( + "--no-tree-reduce", + dest="tree_reduce", + action="store_false", + default=True, + help="Disable tree-reduce for MoE expert-parallel dispatch.", + ) + parser.add_argument( + "--test", + action="store_true", + help="Run PyTorch vs ONNX parity check using a tiny random model.", + ) + args = parser.parse_args() + + if args.test: + with tempfile.TemporaryDirectory() as tmp_dir: + _run_pytorch_parity_test( + model_id=args.model_id, + prompt=args.prompt, + export_dir=tmp_dir, + ctx_len=args.ctx_len, + num_cores=args.num_cores, + num_devices=args.num_devices, + expert_parallel_chunk_size=args.expert_parallel_chunk_size, + cores_per_expert=args.cores_per_expert, + tree_reduce=args.tree_reduce, + ) + return + + factory_kwargs = dict(kv_offload=True, dtype=torch.float16) + config = AutoConfig.from_pretrained(args.model_id) + if args.num_layers is not None: + config.text_config.num_hidden_layers = args.num_layers + factory_kwargs["config"] = config + + qeff_model = QEFFAutoModelForImageTextToText.from_pretrained(args.model_id, weight_free=True,**factory_kwargs) + + qpc_paths = qeff_model.compile( + batch_size=1, + prefill_seq_len=1, + ctx_len=args.ctx_len, + num_cores=args.num_cores, + num_devices=args.num_devices, + mxfp6_matmul=True, + mxint8_kv_cache=True, + use_onnx_subfunctions=False, + skip_vision=True, + offload_pt_weights=False, + dynamo=True, + qaic_config={ + "blocking_mode": "kv_headpar", + "num_kv_blocks": 2, + "moe_config": { + "flavour": "decode_bmm", + "expert_parallel_chunk_size": args.expert_parallel_chunk_size, + "cores_per_expert": args.cores_per_expert, + "tree_reduce": args.tree_reduce, + }, + }, + ) + print(f"QPC paths: {qpc_paths}") + + if args.skip_generate: + return + + processor = AutoProcessor.from_pretrained(args.model_id, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(args.model_id, trust_remote_code=True) + + messages = [ + [ + { + "role": "user", + "content": [{"type": "text", "text": args.prompt}], + } + ] + ] + inputs = tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + device_ids = args.device_ids if args.device_ids is not None else list(range(args.num_devices)) + output = qeff_model.generate(inputs=inputs, generation_len=args.generation_len, device_ids=device_ids) + + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(f"Generated: {output.generated_texts[0]}") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 5b6a258146..3607cb434b 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ classifiers = [ ] requires-python = ">=3.8,<3.13" dependencies = [ - "transformers==5.5.4", + "transformers==5.12.0", "diffusers==0.38.0", "huggingface-hub==1.7.1", "hf_transfer==0.1.9", diff --git a/tests/unit_test/models/test_model_quickcheck.py b/tests/unit_test/models/test_model_quickcheck.py index 5fe9217794..34aba1124a 100644 --- a/tests/unit_test/models/test_model_quickcheck.py +++ b/tests/unit_test/models/test_model_quickcheck.py @@ -63,6 +63,7 @@ Qwen3VLMoeVisionConfig, ) +from QEfficient.transformers.models.minimax_m3_vl import MiniMaxM3VLForCausalLM, MiniMaxM3VLTextConfig from QEfficient.transformers.models.modeling_auto import ( QEFFAutoModel, QEFFAutoModelForCausalLM, @@ -192,6 +193,33 @@ PREFIX_CACHING_MODEL_ID = "hf-internal-testing/tiny-random-GPT2LMHeadModel" +def _tiny_minimax_m3_text_config(dtype=torch.float32) -> MiniMaxM3VLTextConfig: + config = MiniMaxM3VLTextConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=16, + dense_intermediate_size=64, + shared_intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=32, + num_local_experts=4, + num_experts_per_tok=2, + routed_scaling_factor=1.0, + layer_types=["sparse", "full_attention"], + mlp_layer_types=["sparse", "dense"], + index_n_heads=2, + index_head_dim=8, + index_block_size=4, + index_topk_blocks=2, + index_local_blocks=1, + ) + config.torch_dtype = dtype + return config + + def _per_test_thread_budget() -> int: override = os.environ.get("QEFF_NUM_THREADS") if override: @@ -1092,6 +1120,113 @@ def test_causal_lm_cpu_runtime_parity_with_api_runner(model_type, model_id, tmp_ assert np.array_equal(kv_tokens, ort_tokens) +@pytest.mark.llm_model +def test_minimax_m3_text_config_derived_pt_and_onnx_runtime_parity(tmp_path): + torch.manual_seed(7) + config = _tiny_minimax_m3_text_config() + model_hf = MiniMaxM3VLForCausalLM(config).eval() + input_ids = torch.arange(4, dtype=torch.int64).view(1, 4) % config.vocab_size + position_ids = torch.arange(4, dtype=torch.int64).view(1, 4) + + with torch.no_grad(): + hf_logits = model_hf(input_ids=input_ids, position_ids=position_ids, use_cache=False).logits[:, -1:, :] + + qeff_model = QEFFAutoModelForCausalLM(model_hf) + with torch.no_grad(): + qeff_logits = qeff_model.model(input_ids=input_ids, position_ids=position_ids, use_cache=False).logits + + assert torch.allclose(hf_logits, qeff_logits, atol=1e-5, rtol=1e-5) + + past_key_values = tuple( + ( + torch.zeros((1, config.num_key_value_heads, input_ids.shape[1], config.head_dim)), + torch.zeros((1, config.num_key_value_heads, input_ids.shape[1], config.head_dim)), + ) + for _ in range(config.num_hidden_layers) + ) + past_key_values = tuple( + ( + torch.zeros((1, config.num_key_value_heads, input_ids.shape[1], config.head_dim)), + torch.zeros((1, config.num_key_value_heads, input_ids.shape[1], config.head_dim)), + ) + for _ in range(1, config.num_hidden_layers) + ) + index_past_key_values = [ + torch.zeros((1, 1, input_ids.shape[1], config.index_head_dim)) + for lt in config.layer_types + if lt == "sparse" + ] + with torch.no_grad(): + qeff_cached_logits = qeff_model.model( + input_ids=input_ids, + position_ids=position_ids, + past_key_values=past_key_values, + index_keys=index_past_key_values, + use_cache=True, + ).logits + + onnx_path = _exported_onnx_path(qeff_model.export(tmp_path / "minimax-m3-text", prefill_seq_len=4)) + session = _ort_session(onnx_path) + ort_inputs = {} + for input_meta in session.get_inputs(): + shape = [1 if not isinstance(dim, int) else dim for dim in input_meta.shape] + if input_meta.name == "input_ids": + ort_inputs[input_meta.name] = input_ids.numpy() + elif input_meta.name == "position_ids": + ort_inputs[input_meta.name] = position_ids.numpy() + elif input_meta.name.startswith(("past_key.", "past_value.")): + ort_inputs[input_meta.name] = np.zeros( + (1, config.num_key_value_heads, input_ids.shape[1], config.head_dim), dtype=np.float32 + ) + elif inputs_met.name.startswith("index_keys"): + ort_inputs[input_meta.name] = np.zeros( + (1, 1, input_ids.shape[1], config.head_dim), dtype=np.float32 + ) + else: + dtype = np.int64 if input_meta.type == "tensor(int64)" else np.float32 + ort_inputs[input_meta.name] = np.zeros(shape, dtype=dtype) + + ort_logits = session.run(None, ort_inputs)[0] + assert ort_logits.shape == (1, 1, config.vocab_size) + assert np.allclose(ort_logits, qeff_cached_logits.detach().numpy(), atol=1e-4, rtol=1e-4) + + onnx_model = onnx.load(onnx_path, load_external_data=False) + output_names = {output.name for output in onnx_model.graph.output} + assert any(name.startswith("past_key.") and name.endswith("_RetainedState") for name in output_names) + + +@pytest.mark.llm_model +def test_minimax_m3_text_hf_qeff_pytorch_parity(): + torch.manual_seed(7) + config = _tiny_minimax_m3_text_config() + model_hf = MiniMaxM3VLForCausalLM(config).eval() + # Preserve the original HF model before transforms mutate it in-place. + model_hf_orig = deepcopy(model_hf) + + input_ids = torch.arange(4, dtype=torch.int64).view(1, 4) % config.vocab_size + position_ids = torch.arange(4, dtype=torch.int64).view(1, 4) + seq_len = input_ids.shape[1] + + with torch.no_grad(): + hf_logits_no_cache = model_hf( + input_ids=input_ids, + position_ids=position_ids, + use_cache=False, + ).logits[:, -1:] + + qeff_model = QEFFAutoModelForCausalLM(model_hf) + with torch.no_grad(): + qeff_logits_no_cache = qeff_model.model( + input_ids=input_ids, + position_ids=position_ids, + use_cache=False, + ).logits + + assert torch.allclose(hf_logits_no_cache, qeff_logits_no_cache, atol=1e-5, rtol=1e-5), ( + f"HF vs QEff logit mismatch on no-cache prefill: " + f"max_diff={(hf_logits_no_cache - qeff_logits_no_cache).abs().max().item():.6f}" + ) + @pytest.mark.llm_model def test_vlm_text_side_runtime_parity_and_full_export(tmp_path): tokenizer = AutoTokenizer.from_pretrained(VLM_TEXT_RUNTIME_MODEL_ID, trust_remote_code=True)