Minimax m3 mask rope cache

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}

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():

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",

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)

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),

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
 # ---------------------------------------------------------------------------

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.<name>" submodule in this HF
+# port vs. flat "index_<name>" 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.<name>"
+        # submodule in this HF port, but the checkpoint stores those four tensors as
+        # flat "self_attn.index_<name>" 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)
 

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()