From fb57e1ae9f50391272112a0d2b97b4388ce71529 Mon Sep 17 00:00:00 2001 From: ParkiratS Date: Mon, 15 Jun 2026 23:04:05 -0400 Subject: [PATCH 1/3] MOE --- .../convert/cactus_adapters/config_utils.py | 9 + .../cactus/convert/model_adapters/adapters.py | 349 +++++++ .../convert/tests/test_gemma4_adapter.py | 73 ++ .../cactus/convert/tests/test_lfm2_adapter.py | 45 + .../cactus/convert/tests/test_qwen_adapter.py | 64 ++ python/cactus/transpile/aten_ops.py | 6 + .../cactus/transpile/canonicalize/cleanup.py | 3 + python/cactus/transpile/component_plan.py | 13 +- python/cactus/transpile/importers.py | 143 +++ python/cactus/transpile/lower.py | 125 +++ python/cactus/transpile/model_adapters.py | 929 ++++++++++++++++-- python/cactus/transpile/model_profiles.py | 8 +- python/cactus/transpile/ops.py | 37 + python/tests/test_component_plan.py | 36 + python/tests/test_transpile_aten_ops.py | 324 ++++++ 15 files changed, 2092 insertions(+), 72 deletions(-) create mode 100644 python/cactus/convert/tests/test_gemma4_adapter.py create mode 100644 python/cactus/convert/tests/test_qwen_adapter.py create mode 100644 python/tests/test_component_plan.py diff --git a/python/cactus/convert/cactus_adapters/config_utils.py b/python/cactus/convert/cactus_adapters/config_utils.py index b17b7c229..5ec4ccfc1 100644 --- a/python/cactus/convert/cactus_adapters/config_utils.py +++ b/python/cactus/convert/cactus_adapters/config_utils.py @@ -492,6 +492,12 @@ def extract_complex_gemma_config(cfg, root_config): cfg_get(root_config, 'num_global_key_value_heads', None)) expert_intermediate_size = cfg_get(cfg, 'expert_intermediate_size', cfg_get(root_config, 'expert_intermediate_size', None)) + if expert_intermediate_size is None: + expert_intermediate_size = cfg_get(cfg, 'moe_intermediate_size', + cfg_get(root_config, 'moe_intermediate_size', None)) + top_k_experts = cfg_get(cfg, 'top_k_experts', cfg_get(root_config, 'top_k_experts', None)) + if top_k_experts is None: + top_k_experts = cfg_get(cfg, 'num_experts_per_tok', cfg_get(root_config, 'num_experts_per_tok', None)) attention_k_eq_v = bool(cfg_get(cfg, 'attention_k_eq_v', cfg_get(root_config, 'attention_k_eq_v', False))) vocab_size_per_layer_input = cfg_get(cfg, 'vocab_size_per_layer_input', @@ -515,6 +521,9 @@ def extract_complex_gemma_config(cfg, root_config): 'attention_k_eq_v': attention_k_eq_v, 'enable_moe_block': enable_moe_block, } + if top_k_experts is not None: + result['num_top_experts'] = int(top_k_experts) + result['num_experts_per_tok'] = int(top_k_experts) if global_head_dim is not None: result['global_head_dim'] = int(global_head_dim) if num_global_kv_heads is not None: diff --git a/python/cactus/convert/model_adapters/adapters.py b/python/cactus/convert/model_adapters/adapters.py index a29592d6a..eda98d722 100644 --- a/python/cactus/convert/model_adapters/adapters.py +++ b/python/cactus/convert/model_adapters/adapters.py @@ -3,6 +3,7 @@ from dataclasses import dataclass, field, replace import json from pathlib import Path +import re from typing import Any import numpy as np @@ -16,6 +17,7 @@ extract_audio_config, extract_base_config, extract_complex_gemma_config, + extract_lfm2_config, extract_parakeet_config, extract_parakeet_tdt_config, extract_vision_config, @@ -136,6 +138,36 @@ def __init__(self) -> None: super().__init__() self.kv_shared_from_layer = 15 + def normalize_state_dict(self, state_dict: dict[str, Any]) -> NormalizedState: + out = dict(state_dict) + provenance: dict[str, TensorProvenance] = {} + down_pattern = re.compile( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.moe\.down_proj" + ) + for source_name, tensor in list(state_dict.items()): + if down_pattern.fullmatch(source_name) is None: + continue + scale_name = source_name[: -len(".down_proj")] + ".per_expert_scale" + scale = state_dict.get(scale_name) + if scale is None: + continue + shape = _tensor_shape(tensor) + scale_shape = _tensor_shape(scale) + if len(shape) != 3 or len(scale_shape) != 1 or int(shape[0]) != int(scale_shape[0]): + continue + if torch is not None and isinstance(tensor, torch.Tensor): + scale_tensor = scale if isinstance(scale, torch.Tensor) else torch.as_tensor(scale) + scale_value = scale_tensor.to(device=tensor.device, dtype=tensor.dtype).reshape(-1, 1, 1) + out[source_name] = (tensor * scale_value).contiguous() + else: + out[source_name] = (np.asarray(tensor) * np.asarray(scale).reshape(-1, 1, 1)).copy() + provenance[source_name] = TensorProvenance( + [source_name, scale_name], + "gemma4_moe_down_proj_per_expert_scale", + "hf_key", + ) + return NormalizedState(out, provenance) + def runtime_config(self, cfg: Any) -> dict[str, Any]: text_cfg = _cfg_get(cfg, "text_config", None) base_cfg = text_cfg if text_cfg is not None else cfg @@ -167,10 +199,20 @@ def scale_factor(self, output_name: str) -> float: return gemma4_scale_factor(output_name) def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) -> TensorPolicy: + if match.output_name and ( + "moe_router" in match.output_name + or "moe_expert_bias" in match.output_name + or "moe_per_expert_scale" in match.output_name + or match.source_name.endswith(".router.proj.weight") + or match.source_name.endswith(".router.scale") + ): + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "moe routing tensor") policy = super().policy(match, shape, requested_bits) name = match.source_name if name == "model.embed_vision.embedding_projection.weight": return policy + if match.output_name and "moe_expert_" in match.output_name: + return replace(policy, use_gptq=False) if ".self_attn." in name and (name.endswith(".k_proj.weight") or name.endswith(".v_proj.weight")): parts = name.split(".") try: @@ -181,6 +223,91 @@ def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) return replace(policy, use_gptq=False, fallback_reason=policy.fallback_reason or "shared KV tensor has no per-layer hook module") return policy + def module_target_name(self, source_name: str, _model: Any) -> str | None: + if ".moe.gate_up_proj" in source_name or ".moe.down_proj" in source_name: + return None + return super().module_target_name(source_name, _model) + + def name_tensor(self, source_name: str, tensor: Any, num_layers: int | None) -> NameMatch: + match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.moe\.(gate_up_proj|down_proj)", + source_name, + ) + if match is not None: + layer_index = int(match.group(1)) + kind = match.group(2) + if kind == "gate_up_proj": + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_gate_up.weights" + else: + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_w2.weights" + return NameMatch(source_name, output_name, "language", True, hf_name=source_name, adapter_name=source_name) + return super().name_tensor(source_name, tensor, num_layers) + + def expand_tensor(self, match: NameMatch, tensor: Any) -> list[TensorEmission]: + if match.output_name is None: + return [] + source_match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.moe\.(gate_up_proj|down_proj)", + match.source_name, + ) + if source_match is None: + return super().expand_tensor(match, tensor) + + layer_index = int(source_match.group(1)) + kind = source_match.group(2) + shape = _tensor_shape(tensor) + if len(shape) != 3: + raise ValueError(f"Gemma4-MoE expert tensor expects rank 3, got {shape}") + num_experts = int(shape[0]) + emissions: list[TensorEmission] = [] + + def _piece(value, expert_idx: int, start: int | None = None, end: int | None = None): + if torch is not None and isinstance(value, torch.Tensor): + selected = value[expert_idx] if start is None else value[expert_idx, start:end, :] + return selected.contiguous() + array = np.asarray(value) + selected = array[expert_idx] if start is None else array[expert_idx, start:end, :] + return selected.copy() + + def _source_names(weight_name: str, expert_idx: int) -> list[str]: + return [ + f"model.language_model.layers.{layer_index}.moe.{weight_name}.{expert_idx}", + f"model.layers.{layer_index}.moe.{weight_name}.{expert_idx}", + f"layers.{layer_index}.moe.{weight_name}.{expert_idx}", + f"backbone.layers.{layer_index}.moe.{weight_name}.{expert_idx}", + ] + + if kind == "gate_up_proj": + if int(shape[1]) % 2 != 0: + raise ValueError(f"Gemma4-MoE gate_up_proj second dim must be even, got {shape}") + intermediate = int(shape[1]) // 2 + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w1.weights", + _piece(tensor, expert_idx, 0, intermediate), + "gemma4_moe_gate_up_split_gate", + source_names=_source_names("w1_weights", expert_idx), + qdq_restore="runtime_key", + )) + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w3.weights", + _piece(tensor, expert_idx, intermediate, int(shape[1])), + "gemma4_moe_gate_up_split_up", + source_names=_source_names("w3_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w2.weights", + _piece(tensor, expert_idx), + "gemma4_moe_down_proj_split_scaled", + source_names=_source_names("w2_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + def module_target_name(self, source_name: str, _model: Any) -> str | None: target = super().module_target_name(source_name, _model) if target and source_name.startswith("model.vision_tower."): @@ -196,6 +323,20 @@ def module_target_name(self, source_name: str, _model: Any) -> str | None: class QwenAdapter(FamilyAdapter): family = "qwen" + def runtime_config(self, cfg: Any) -> dict[str, Any]: + config = super().runtime_config(cfg) + model_type = str(_cfg_get(cfg, "model_type", "") or "").lower() + if model_type == "qwen2_moe" or _cfg_get(cfg, "moe_intermediate_size", None) is not None: + config.update({ + "moe_intermediate_size": int(_cfg_get(cfg, "moe_intermediate_size", 0) or 0), + "shared_expert_intermediate_size": int(_cfg_get(cfg, "shared_expert_intermediate_size", 0) or 0), + "decoder_sparse_step": int(_cfg_get(cfg, "decoder_sparse_step", 0) or 0), + "num_experts": int(_cfg_get(cfg, "num_experts", config.get("num_experts", 0)) or 0), + "num_experts_per_tok": int(_cfg_get(cfg, "num_experts_per_tok", config.get("num_experts_per_tok", 0)) or 0), + "norm_topk_prob": bool(_cfg_get(cfg, "norm_topk_prob", False)), + }) + return config + def model_class(self, cfg: Any): from transformers import AutoModel, AutoModelForCausalLM @@ -210,6 +351,115 @@ def model_class(self, cfg: Any): return AutoModel return AutoModelForCausalLM + def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) -> TensorPolicy: + if match.output_name and ( + "moe_router" in match.output_name + or "moe_expert_" in match.output_name + or "shared_expert_gate" in match.output_name + ): + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "moe routing/expert tensor") + return super().policy(match, shape, requested_bits) + + def module_target_name(self, source_name: str, _model: Any) -> str | None: + if ".mlp.experts.gate_up_proj" in source_name or ".mlp.experts.down_proj" in source_name: + return None + return super().module_target_name(source_name, _model) + + def name_tensor(self, source_name: str, tensor: Any, num_layers: int | None) -> NameMatch: + match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.mlp\.experts\.(gate_up_proj|down_proj)", + source_name, + ) + if match is not None: + layer_index = int(match.group(1)) + kind = match.group(2) + if kind == "gate_up_proj": + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_gate_up.weights" + else: + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_w2.weights" + return NameMatch(source_name, output_name, "language", True, hf_name=source_name, adapter_name=source_name) + + router_match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.mlp\.gate\.weight", + source_name, + ) + if router_match is not None: + layer_index = int(router_match.group(1)) + return NameMatch( + source_name, + f"layer_{layer_index}_moe_router.weights", + "language", + True, + hf_name=source_name, + adapter_name=source_name, + ) + + return super().name_tensor(source_name, tensor, num_layers) + + def expand_tensor(self, match: NameMatch, tensor: Any) -> list[TensorEmission]: + if match.output_name is None: + return [] + source_match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.mlp\.experts\.(gate_up_proj|down_proj)", + match.source_name, + ) + if source_match is None: + return super().expand_tensor(match, tensor) + + layer_index = int(source_match.group(1)) + kind = source_match.group(2) + shape = _tensor_shape(tensor) + if len(shape) != 3: + raise ValueError(f"Qwen2-MoE expert tensor expects rank 3, got {shape}") + num_experts = int(shape[0]) + emissions: list[TensorEmission] = [] + + def _piece(value, expert_idx: int, start: int | None = None, end: int | None = None): + if torch is not None and isinstance(value, torch.Tensor): + selected = value[expert_idx] if start is None else value[expert_idx, start:end, :] + return selected.contiguous() + array = np.asarray(value) + selected = array[expert_idx] if start is None else array[expert_idx, start:end, :] + return selected.copy() + + def _source_names(weight_name: str, expert_idx: int) -> list[str]: + return [ + f"model.layers.{layer_index}.mlp.{weight_name}.{expert_idx}", + f"layers.{layer_index}.mlp.{weight_name}.{expert_idx}", + f"backbone.layers.{layer_index}.mlp.{weight_name}.{expert_idx}", + ] + + if kind == "gate_up_proj": + if int(shape[1]) % 2 != 0: + raise ValueError(f"Qwen2-MoE gate_up_proj second dim must be even, got {shape}") + intermediate = int(shape[1]) // 2 + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w1.weights", + _piece(tensor, expert_idx, 0, intermediate), + "qwen2_moe_gate_up_split_gate", + source_names=_source_names("w1_weights", expert_idx), + qdq_restore="runtime_key", + )) + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w3.weights", + _piece(tensor, expert_idx, intermediate, int(shape[1])), + "qwen2_moe_gate_up_split_up", + source_names=_source_names("w3_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w2.weights", + _piece(tensor, expert_idx), + "qwen2_moe_down_proj_split", + source_names=_source_names("w2_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + class WhisperAdapter(FamilyAdapter): family = "whisper" @@ -435,6 +685,10 @@ def normalize_state_dict(self, state_dict: dict[str, Any]) -> NormalizedState: class Lfm2Adapter(FamilyAdapter): family = "lfm2" + def __init__(self) -> None: + super().__init__() + self.num_experts = 0 + def model_class(self, cfg: Any): from transformers import AutoModelForCausalLM @@ -449,6 +703,12 @@ def model_class(self, cfg: Any): pass return AutoModelForCausalLM + def runtime_config(self, cfg: Any) -> dict[str, Any]: + config = super().runtime_config(cfg) + config.update(extract_lfm2_config(cfg)) + self.num_experts = int(_cfg_get(cfg, "num_experts", config.get("num_experts", 0)) or 0) + return config + def load_processor(self, model_id_or_path: str): processor = super().load_processor(model_id_or_path) if processor is not None and hasattr(processor, "tokenizer"): @@ -484,16 +744,105 @@ def load_processor(self, model_id_or_path: str): return Lfm2VlProcessor(image_processor=image_processor, tokenizer=tokenizer) def policy(self, match: NameMatch, shape: tuple[int, ...], requested_bits: int) -> TensorPolicy: + if match.output_name and ( + "moe_router" in match.output_name + or "moe_expert_bias" in match.output_name + ): + return TensorPolicy("fallback", "FP16", None, match.component, False, "none", "moe routing tensor") policy = super().policy(match, shape, requested_bits) if match.source_name.endswith("vision_model.embeddings.patch_embedding.weight"): return replace(policy, use_gptq=False, fallback_reason=policy.fallback_reason or "vision patch embedding has no linear Hessian target") + if match.output_name and "moe_expert_" in match.output_name: + return replace(policy, use_gptq=False) return policy def module_target_name(self, source_name: str, _model: Any) -> str | None: + if ".feed_forward.experts.gate_up_proj" in source_name or ".feed_forward.experts.down_proj" in source_name: + return None if source_name.startswith("model.vision_tower.vision_model."): source_name = source_name.replace("model.vision_tower.vision_model.", "model.vision_tower.", 1) return super().module_target_name(source_name, _model) + def name_tensor(self, source_name: str, tensor: Any, num_layers: int | None) -> NameMatch: + match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.feed_forward\.experts\.(gate_up_proj|down_proj)", + source_name, + ) + if match is not None: + layer_index = int(match.group(1)) + kind = match.group(2) + if kind == "gate_up_proj": + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_gate_up.weights" + else: + output_name = f"layer_{layer_index}_moe_expert_{{channel}}_w2.weights" + return NameMatch(source_name, output_name, "language", True, hf_name=source_name, adapter_name=source_name) + return super().name_tensor(source_name, tensor, num_layers) + + def expand_tensor(self, match: NameMatch, tensor: Any) -> list[TensorEmission]: + if match.output_name is None: + return [] + source_match = re.fullmatch( + r"(?:(?:model|language_model|model\.language_model)\.)?layers\.(\d+)\.feed_forward\.experts\.(gate_up_proj|down_proj)", + match.source_name, + ) + if source_match is None: + return super().expand_tensor(match, tensor) + + layer_index = int(source_match.group(1)) + kind = source_match.group(2) + shape = _tensor_shape(tensor) + if len(shape) != 3: + raise ValueError(f"LFM2-MoE expert tensor expects rank 3, got {shape}") + num_experts = int(shape[0]) + self.num_experts = self.num_experts or num_experts + emissions: list[TensorEmission] = [] + + def _piece(value, expert_idx: int, start: int | None = None, end: int | None = None): + if torch is not None and isinstance(value, torch.Tensor): + selected = value[expert_idx] if start is None else value[expert_idx, start:end, :] + return selected.contiguous() + array = np.asarray(value) + selected = array[expert_idx] if start is None else array[expert_idx, start:end, :] + return selected.copy() + + def _source_names(weight_name: str, expert_idx: int) -> list[str]: + return [ + f"model.layers.{layer_index}.feed_forward.{weight_name}.{expert_idx}", + f"layers.{layer_index}.feed_forward.{weight_name}.{expert_idx}", + f"backbone.layers.{layer_index}.feed_forward.{weight_name}.{expert_idx}", + ] + + if kind == "gate_up_proj": + if int(shape[1]) % 2 != 0: + raise ValueError(f"LFM2-MoE gate_up_proj second dim must be even, got {shape}") + intermediate = int(shape[1]) // 2 + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w1.weights", + _piece(tensor, expert_idx, 0, intermediate), + "lfm2_moe_gate_up_split_gate", + source_names=_source_names("w1_weights", expert_idx), + qdq_restore="runtime_key", + )) + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w3.weights", + _piece(tensor, expert_idx, intermediate, int(shape[1])), + "lfm2_moe_gate_up_split_up", + source_names=_source_names("w3_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + + for expert_idx in range(num_experts): + emissions.append(TensorEmission( + f"layer_{layer_index}_moe_expert_{expert_idx}_w2.weights", + _piece(tensor, expert_idx), + "lfm2_moe_down_proj_split", + source_names=_source_names("w2_weights", expert_idx), + qdq_restore="runtime_key", + )) + return emissions + class Gemma3Adapter(FamilyAdapter): family = "gemma3" diff --git a/python/cactus/convert/tests/test_gemma4_adapter.py b/python/cactus/convert/tests/test_gemma4_adapter.py new file mode 100644 index 000000000..a93a5cc86 --- /dev/null +++ b/python/cactus/convert/tests/test_gemma4_adapter.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import torch + +from cactus.convert.model_adapters.adapters import adapter_for_family + + +def test_gemma4_moe_down_proj_folds_per_expert_scale_before_split(): + adapter = adapter_for_family("gemma4") + down = torch.arange(2 * 4 * 3, dtype=torch.float16).reshape(2, 4, 3) + scale = torch.tensor([2.0, 0.5], dtype=torch.float16) + + normalized = adapter.normalize_state_dict({ + "model.language_model.layers.5.moe.down_proj": down, + "model.language_model.layers.5.moe.per_expert_scale": scale, + }) + scaled = normalized.state_dict["model.language_model.layers.5.moe.down_proj"] + match = adapter.name_tensor( + "model.language_model.layers.5.moe.down_proj", + scaled, + num_layers=30, + ) + emissions = adapter.expand_tensor(match, scaled) + + assert torch.equal(emissions[0].tensor, down[0] * scale[0]) + assert torch.equal(emissions[1].tensor, down[1] * scale[1]) + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w2.weights", + "layer_5_moe_expert_1_w2.weights", + ] + assert "model.language_model.layers.5.moe.w2_weights.0" in emissions[0].source_names + assert normalized.provenance["model.language_model.layers.5.moe.down_proj"].source_names == [ + "model.language_model.layers.5.moe.down_proj", + "model.language_model.layers.5.moe.per_expert_scale", + ] + + +def test_gemma4_moe_gate_up_splits_into_runtime_experts(): + adapter = adapter_for_family("gemma4") + tensor = torch.arange(2 * 6 * 4, dtype=torch.float16).reshape(2, 6, 4) + + match = adapter.name_tensor( + "model.language_model.layers.5.moe.gate_up_proj", + tensor, + num_layers=30, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w1.weights", + "layer_5_moe_expert_0_w3.weights", + "layer_5_moe_expert_1_w1.weights", + "layer_5_moe_expert_1_w3.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0, :3, :]) + assert torch.equal(emissions[1].tensor, tensor[0, 3:, :]) + assert "model.language_model.layers.5.moe.w1_weights.0" in emissions[0].source_names + assert "model.language_model.layers.5.moe.w3_weights.1" in emissions[3].source_names + + +def test_gemma4_moe_router_policy_stays_fp16(): + adapter = adapter_for_family("gemma4") + tensor = torch.randn(128, 2816, dtype=torch.float16) + match = adapter.name_tensor( + "model.language_model.layers.5.router.proj.weight", + tensor, + num_layers=30, + ) + policy = adapter.policy(match, tuple(tensor.shape), requested_bits=4) + + assert match.output_name == "layer_5_router_proj.weights" + assert policy.action == "fallback" + assert policy.precision == "FP16" diff --git a/python/cactus/convert/tests/test_lfm2_adapter.py b/python/cactus/convert/tests/test_lfm2_adapter.py index 49d5798c9..1a11f66ff 100644 --- a/python/cactus/convert/tests/test_lfm2_adapter.py +++ b/python/cactus/convert/tests/test_lfm2_adapter.py @@ -1,5 +1,7 @@ from __future__ import annotations +import torch + from cactus.convert.model_adapters.adapters import adapter_for_family @@ -59,3 +61,46 @@ def test_lfm2_processor_fallback_handles_tokenizers_backend(tmp_path): assert isinstance(processor, Lfm2VlProcessor) assert processor.image_token == "" assert processor.image_token_id == 3 + + +def test_lfm2_moe_packed_gate_up_splits_into_runtime_experts(): + adapter = adapter_for_family("lfm2") + tensor = torch.arange(2 * 6 * 4, dtype=torch.float16).reshape(2, 6, 4) + + match = adapter.name_tensor( + "model.layers.5.feed_forward.experts.gate_up_proj", + tensor, + num_layers=24, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w1.weights", + "layer_5_moe_expert_0_w3.weights", + "layer_5_moe_expert_1_w1.weights", + "layer_5_moe_expert_1_w3.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0, :3, :]) + assert torch.equal(emissions[1].tensor, tensor[0, 3:, :]) + assert "model.layers.5.feed_forward.w1_weights.0" in emissions[0].source_names + assert "model.layers.5.feed_forward.w3_weights.1" in emissions[3].source_names + + +def test_lfm2_moe_packed_down_splits_into_runtime_experts(): + adapter = adapter_for_family("lfm2") + tensor = torch.arange(2 * 4 * 3, dtype=torch.float16).reshape(2, 4, 3) + + match = adapter.name_tensor( + "model.layers.5.feed_forward.experts.down_proj", + tensor, + num_layers=24, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w2.weights", + "layer_5_moe_expert_1_w2.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0]) + assert torch.equal(emissions[1].tensor, tensor[1]) + assert "model.layers.5.feed_forward.w2_weights.0" in emissions[0].source_names diff --git a/python/cactus/convert/tests/test_qwen_adapter.py b/python/cactus/convert/tests/test_qwen_adapter.py new file mode 100644 index 000000000..e3e3e48f9 --- /dev/null +++ b/python/cactus/convert/tests/test_qwen_adapter.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import torch + +from cactus.convert.model_adapters.adapters import adapter_for_family + + +def test_qwen2_moe_packed_gate_up_splits_into_runtime_experts(): + adapter = adapter_for_family("qwen") + tensor = torch.arange(2 * 6 * 4, dtype=torch.float16).reshape(2, 6, 4) + + match = adapter.name_tensor( + "model.layers.5.mlp.experts.gate_up_proj", + tensor, + num_layers=24, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w1.weights", + "layer_5_moe_expert_0_w3.weights", + "layer_5_moe_expert_1_w1.weights", + "layer_5_moe_expert_1_w3.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0, :3, :]) + assert torch.equal(emissions[1].tensor, tensor[0, 3:, :]) + assert "model.layers.5.mlp.w1_weights.0" in emissions[0].source_names + assert "model.layers.5.mlp.w3_weights.1" in emissions[3].source_names + + +def test_qwen2_moe_packed_down_splits_into_runtime_experts(): + adapter = adapter_for_family("qwen") + tensor = torch.arange(2 * 4 * 3, dtype=torch.float16).reshape(2, 4, 3) + + match = adapter.name_tensor( + "model.layers.5.mlp.experts.down_proj", + tensor, + num_layers=24, + ) + emissions = adapter.expand_tensor(match, tensor) + + assert [emission.output_name for emission in emissions] == [ + "layer_5_moe_expert_0_w2.weights", + "layer_5_moe_expert_1_w2.weights", + ] + assert torch.equal(emissions[0].tensor, tensor[0]) + assert torch.equal(emissions[1].tensor, tensor[1]) + assert "model.layers.5.mlp.w2_weights.0" in emissions[0].source_names + + +def test_qwen2_moe_router_uses_runtime_router_name_and_fp16_policy(): + adapter = adapter_for_family("qwen") + tensor = torch.randn(60, 2048, dtype=torch.float16) + + match = adapter.name_tensor( + "model.layers.5.mlp.gate.weight", + tensor, + num_layers=24, + ) + policy = adapter.policy(match, tuple(tensor.shape), requested_bits=4) + + assert match.output_name == "layer_5_moe_router.weights" + assert policy.action == "fallback" + assert policy.precision == "FP16" diff --git a/python/cactus/transpile/aten_ops.py b/python/cactus/transpile/aten_ops.py index 90cb587ec..e4878886f 100644 --- a/python/cactus/transpile/aten_ops.py +++ b/python/cactus/transpile/aten_ops.py @@ -47,6 +47,12 @@ def canonical_torch_op(target: Any) -> str: "aten.flatten.using_ints": "flatten", "aten.t.default": "transpose", "aten.gelu.erf": "gelu_erf", + "cactus_transpile.lfm2_moe_layer_gated.default": "lfm2_moe_layer_gated", + "lfm2_moe_layer_gated.default": "lfm2_moe_layer_gated", + "cactus_transpile.qwen2_moe_layer_gated.default": "qwen2_moe_layer_gated", + "qwen2_moe_layer_gated.default": "qwen2_moe_layer_gated", + "cactus_transpile.gemma4_moe_layer_gated.default": "gemma4_moe_layer_gated", + "gemma4_moe_layer_gated.default": "gemma4_moe_layer_gated", # Keep the legacy importer key until the importer itself is renamed. "aten.diff.default": "aten.diff.default", } diff --git a/python/cactus/transpile/canonicalize/cleanup.py b/python/cactus/transpile/canonicalize/cleanup.py index b7f4de2aa..c68e00447 100644 --- a/python/cactus/transpile/canonicalize/cleanup.py +++ b/python/cactus/transpile/canonicalize/cleanup.py @@ -129,6 +129,9 @@ "attention_block", "conv_module", "dense_mlp_tq_fused", + "lfm2_moe_layer_gated", + "qwen2_moe_layer_gated", + "gemma4_moe_layer_gated", "rope", } diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py index c9e63cabd..5c8fdd751 100644 --- a/python/cactus/transpile/component_plan.py +++ b/python/cactus/transpile/component_plan.py @@ -26,6 +26,8 @@ def _plan_from_profile( ) -> ComponentPlan | None: if profile.default_task is None or not profile.default_components: return None + has_vision: bool | None = None + has_audio: bool | None = None if config is not None and profile.needs_image: has_vision = _has_dict_config(config, "vision_config", "visual_config", "image_config") if not has_vision: @@ -42,12 +44,19 @@ def _plan_from_profile( components=("decoder", "decoder_step"), force_component_pipeline=True, ) + if config is not None and profile.needs_audio: + has_audio = _has_dict_config(config, "audio_config", "speech_config", "acoustic_config") + components = profile.default_components + needs_audio = profile.needs_audio + if profile.needs_image and profile.needs_audio and has_vision and has_audio is False: + components = tuple(component for component in components if component != "audio_encoder") + needs_audio = False return ComponentPlan( task=profile.default_task, - components=profile.default_components, + components=components, default_max_new_tokens=profile.default_max_new_tokens, needs_image=profile.needs_image, - needs_audio=profile.needs_audio, + needs_audio=needs_audio, force_component_pipeline=profile.force_component_pipeline, ) diff --git a/python/cactus/transpile/importers.py b/python/cactus/transpile/importers.py index d606d32f5..62f2cbb55 100644 --- a/python/cactus/transpile/importers.py +++ b/python/cactus/transpile/importers.py @@ -1774,6 +1774,146 @@ def import_getitem(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[i register_node(ir, ir_node, shape=shape, dtype=dtype) +def import_lfm2_moe_layer_gated( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + if len(node.args) != 12: + ctx.fail(f"lfm2_moe_layer_gated expected 12 args, got {len(node.args)}") + w1_inputs = extract_input_ids(node.args[3], ctx) + w3_inputs = extract_input_ids(node.args[4], ctx) + w2_inputs = extract_input_ids(node.args[5], ctx) + num_experts = int(extract_literals(node.args[6])) + num_experts_per_tok = int(extract_literals(node.args[7])) + if len(w1_inputs) != num_experts or len(w3_inputs) != num_experts or len(w2_inputs) != num_experts: + ctx.fail( + "lfm2_moe_layer_gated expert input count mismatch: " + f"num_experts={num_experts} w1={len(w1_inputs)} w3={len(w3_inputs)} w2={len(w2_inputs)}" + ) + attrs = { + "num_experts": num_experts, + "num_experts_per_tok": num_experts_per_tok, + "use_expert_bias": bool(extract_literals(node.args[8])), + "normalize_routing": bool(extract_literals(node.args[9])), + "epsilon": float(extract_literals(node.args[10])), + "routed_scaling_factor": float(extract_literals(node.args[11])), + "activation": "silu", + } + ir_node = IRNode( + id=node_id(node), + op="lfm2_moe_layer_gated", + inputs=[ + value_id(node.args[0], ctx), + value_id(node.args[1], ctx), + value_id(node.args[2], ctx), + *w1_inputs, + *w3_inputs, + *w2_inputs, + ], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_qwen2_moe_layer_gated( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + if len(node.args) != 10: + ctx.fail(f"qwen2_moe_layer_gated expected 10 args, got {len(node.args)}") + w1_inputs = extract_input_ids(node.args[2], ctx) + w3_inputs = extract_input_ids(node.args[3], ctx) + w2_inputs = extract_input_ids(node.args[4], ctx) + num_experts = int(extract_literals(node.args[5])) + num_experts_per_tok = int(extract_literals(node.args[6])) + if len(w1_inputs) != num_experts or len(w3_inputs) != num_experts or len(w2_inputs) != num_experts: + ctx.fail( + "qwen2_moe_layer_gated expert input count mismatch: " + f"num_experts={num_experts} w1={len(w1_inputs)} w3={len(w3_inputs)} w2={len(w2_inputs)}" + ) + attrs = { + "num_experts": num_experts, + "num_experts_per_tok": num_experts_per_tok, + "normalize_routing": bool(extract_literals(node.args[7])), + "epsilon": float(extract_literals(node.args[8])), + "routed_scaling_factor": float(extract_literals(node.args[9])), + "activation": "silu", + } + ir_node = IRNode( + id=node_id(node), + op="qwen2_moe_layer_gated", + inputs=[ + value_id(node.args[0], ctx), + value_id(node.args[1], ctx), + *w1_inputs, + *w3_inputs, + *w2_inputs, + ], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + +def import_gemma4_moe_layer_gated( + ir: IRGraph, + node: Any, + ctx: ImportContext, + *, + shape: tuple[int, ...] | None, + dtype: str | None, + torch_op: str, +) -> None: + if len(node.args) != 10: + ctx.fail(f"gemma4_moe_layer_gated expected 10 args, got {len(node.args)}") + w1_inputs = extract_input_ids(node.args[2], ctx) + w3_inputs = extract_input_ids(node.args[3], ctx) + w2_inputs = extract_input_ids(node.args[4], ctx) + num_experts = int(extract_literals(node.args[5])) + num_experts_per_tok = int(extract_literals(node.args[6])) + if len(w1_inputs) != num_experts or len(w3_inputs) != num_experts or len(w2_inputs) != num_experts: + ctx.fail( + "gemma4_moe_layer_gated expert input count mismatch: " + f"num_experts={num_experts} w1={len(w1_inputs)} w3={len(w3_inputs)} w2={len(w2_inputs)}" + ) + attrs = { + "num_experts": num_experts, + "num_experts_per_tok": num_experts_per_tok, + "normalize_routing": bool(extract_literals(node.args[7])), + "epsilon": float(extract_literals(node.args[8])), + "routed_scaling_factor": float(extract_literals(node.args[9])), + "activation": "silu", + } + ir_node = IRNode( + id=node_id(node), + op="gemma4_moe_layer_gated", + inputs=[ + value_id(node.args[0], ctx), + value_id(node.args[1], ctx), + *w1_inputs, + *w3_inputs, + *w2_inputs, + ], + outputs=[value_id(node, ctx)], + attrs=attrs, + meta=_base_meta(shape, dtype, torch_op, node), + ) + register_node(ir, ir_node, shape=shape, dtype=dtype) + + OP_IMPORTERS = { "arange": import_arange, "identity": import_identity, @@ -1867,4 +2007,7 @@ def import_getitem(ir: IRGraph, node: Any, ctx: ImportContext, *, shape: tuple[i "batch_norm": import_batch_norm, "contiguous": import_contiguous, "getitem": import_getitem, + "lfm2_moe_layer_gated": import_lfm2_moe_layer_gated, + "qwen2_moe_layer_gated": import_qwen2_moe_layer_gated, + "gemma4_moe_layer_gated": import_gemma4_moe_layer_gated, } diff --git a/python/cactus/transpile/lower.py b/python/cactus/transpile/lower.py index ceeee03b1..04be578dc 100644 --- a/python/cactus/transpile/lower.py +++ b/python/cactus/transpile/lower.py @@ -1106,6 +1106,131 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) -> out = g.precision_cast(out, output_dtype) return [out] + if op == "lfm2_moe_layer_gated": + num_experts = int(node.attrs["num_experts"]) + num_experts_per_tok = int(node.attrs["num_experts_per_tok"]) + expected_inputs = 3 + 3 * num_experts + if len(node.inputs) != expected_inputs: + raise NotImplementedError( + f"lfm2_moe_layer_gated expects {expected_inputs} inputs for {num_experts} experts, " + f"got {len(node.inputs)}" + ) + hidden = _ensure_tensor_dtype(g, _tensor(env, node.inputs[0]), Graph.FP16) + router_logits = _ensure_tensor_dtype(g, _tensor(env, node.inputs[1]), Graph.FP16) + expert_bias = _tensor(env, node.inputs[2]) + w1_start = 3 + w3_start = w1_start + num_experts + w2_start = w3_start + num_experts + w1_weights = [_tensor(env, value_id) for value_id in node.inputs[w1_start:w3_start]] + w3_weights = [_tensor(env, value_id) for value_id in node.inputs[w3_start:w2_start]] + w2_weights = [_tensor(env, value_id) for value_id in node.inputs[w2_start:]] + + routing_probs = g.sigmoid(router_logits) + topk_source = routing_probs + if bool(node.attrs.get("use_expert_bias", False)): + expert_bias = _ensure_tensor_dtype(g, expert_bias, routing_probs.dtype) + topk_source = _lower_binary_op(g, routing_probs, expert_bias, "add") + topk_output = g.topk(topk_source, num_experts_per_tok) + topk_indices = _ensure_tensor_dtype(g, g.index(topk_output, 0, axis=0), Graph.FP32) + out = g.moe_layer_gated( + hidden, + routing_probs, + topk_indices, + w1_weights, + w3_weights, + w2_weights, + num_experts, + num_experts_per_tok, + normalize_routing=bool(node.attrs.get("normalize_routing", False)), + epsilon=float(node.attrs.get("epsilon", 1.0e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + output_value = ir.values.get(node.outputs[0]) + output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else out.dtype + if out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + + if op == "qwen2_moe_layer_gated": + num_experts = int(node.attrs["num_experts"]) + num_experts_per_tok = int(node.attrs["num_experts_per_tok"]) + expected_inputs = 2 + 3 * num_experts + if len(node.inputs) != expected_inputs: + raise NotImplementedError( + f"qwen2_moe_layer_gated expects {expected_inputs} inputs for {num_experts} experts, " + f"got {len(node.inputs)}" + ) + hidden = _ensure_tensor_dtype(g, _tensor(env, node.inputs[0]), Graph.FP16) + router_logits = _ensure_tensor_dtype(g, _tensor(env, node.inputs[1]), Graph.FP16) + w1_start = 2 + w3_start = w1_start + num_experts + w2_start = w3_start + num_experts + w1_weights = [_tensor(env, value_id) for value_id in node.inputs[w1_start:w3_start]] + w3_weights = [_tensor(env, value_id) for value_id in node.inputs[w3_start:w2_start]] + w2_weights = [_tensor(env, value_id) for value_id in node.inputs[w2_start:]] + + routing_probs = g.softmax(router_logits, axis=-1) + topk_output = g.topk(routing_probs, num_experts_per_tok) + topk_indices = _ensure_tensor_dtype(g, g.index(topk_output, 0, axis=0), Graph.FP32) + out = g.moe_layer_gated( + hidden, + routing_probs, + topk_indices, + w1_weights, + w3_weights, + w2_weights, + num_experts, + num_experts_per_tok, + normalize_routing=bool(node.attrs.get("normalize_routing", False)), + epsilon=float(node.attrs.get("epsilon", 1.0e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + output_value = ir.values.get(node.outputs[0]) + output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else out.dtype + if out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + + if op == "gemma4_moe_layer_gated": + num_experts = int(node.attrs["num_experts"]) + num_experts_per_tok = int(node.attrs["num_experts_per_tok"]) + expected_inputs = 2 + 3 * num_experts + if len(node.inputs) != expected_inputs: + raise NotImplementedError( + f"gemma4_moe_layer_gated expects {expected_inputs} inputs for {num_experts} experts, " + f"got {len(node.inputs)}" + ) + hidden = _ensure_tensor_dtype(g, _tensor(env, node.inputs[0]), Graph.FP16) + router_logits = _ensure_tensor_dtype(g, _tensor(env, node.inputs[1]), Graph.FP16) + w1_start = 2 + w3_start = w1_start + num_experts + w2_start = w3_start + num_experts + w1_weights = [_tensor(env, value_id) for value_id in node.inputs[w1_start:w3_start]] + w3_weights = [_tensor(env, value_id) for value_id in node.inputs[w3_start:w2_start]] + w2_weights = [_tensor(env, value_id) for value_id in node.inputs[w2_start:]] + + routing_probs = g.softmax(router_logits, axis=-1) + topk_output = g.topk(routing_probs, num_experts_per_tok) + topk_indices = _ensure_tensor_dtype(g, g.index(topk_output, 0, axis=0), Graph.FP32) + out = g.moe_layer_gated( + hidden, + routing_probs, + topk_indices, + w1_weights, + w3_weights, + w2_weights, + num_experts, + num_experts_per_tok, + normalize_routing=bool(node.attrs.get("normalize_routing", True)), + epsilon=float(node.attrs.get("epsilon", 1.0e-6)), + routed_scaling_factor=float(node.attrs.get("routed_scaling_factor", 1.0)), + ) + output_value = ir.values.get(node.outputs[0]) + output_dtype = _map_ir_dtype(output_value.dtype) if output_value is not None and output_value.dtype is not None else out.dtype + if out.dtype != output_dtype: + out = g.precision_cast(out, output_dtype) + return [out] + if op == "addmm": bias = _tensor(env, node.inputs[0]) lhs = _tensor(env, node.inputs[1]) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index b6e2cee0d..7dbaa957c 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -437,6 +437,700 @@ def _torch_is_compiling() -> bool: return False +_CACTUS_TORCH_LIBRARIES: list[object] = [] +_LFM2_MOE_CUSTOM_OP_REGISTERED = False +_LFM2_MOE_CUSTOM_OP = "cactus_transpile::lfm2_moe_layer_gated" +_QWEN2_MOE_CUSTOM_OP_REGISTERED = False +_QWEN2_MOE_CUSTOM_OP = "cactus_transpile::qwen2_moe_layer_gated" +_GEMMA4_MOE_CUSTOM_OP_REGISTERED = False +_GEMMA4_MOE_CUSTOM_OP = "cactus_transpile::gemma4_moe_layer_gated" + + +def _ignore_duplicate_torch_registration(exc: RuntimeError) -> bool: + text = str(exc).lower() + return "already" in text or "duplicate" in text or "previously" in text + + +def _ensure_lfm2_moe_custom_op_registered() -> None: + global _LFM2_MOE_CUSTOM_OP_REGISTERED + if _LFM2_MOE_CUSTOM_OP_REGISTERED: + return + + schema = ( + "lfm2_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, Tensor expert_bias, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool use_expert_bias, " + "bool normalize_routing, float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + try: + library = torch.library.Library("cactus_transpile", "FRAGMENT") + library.define(schema) + _CACTUS_TORCH_LIBRARIES.append(library) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + + def _impl( + hidden: torch.Tensor, + router_logits: torch.Tensor, + expert_bias: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + use_expert_bias: bool, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + routing_probs = router_logits.sigmoid() + if use_expert_bias: + scores_for_routing = routing_probs + expert_bias.to(device=routing_probs.device, dtype=routing_probs.dtype) + _, selected_experts = torch.topk(scores_for_routing, k=int(num_experts_per_tok), dim=-1) + routing_weights = torch.gather(routing_probs, dim=1, index=selected_experts).type_as(router_logits) + else: + routing_weights, selected_experts = torch.topk(routing_probs, k=int(num_experts_per_tok), dim=-1) + + if normalize_routing: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + float(epsilon)) + routing_weights = routing_weights * float(routed_scaling_factor) + + output = torch.zeros_like(hidden) + for expert_idx in range(int(num_experts)): + selected = selected_experts == expert_idx + if not bool(selected.any()): + continue + token_idx, topk_pos = torch.where(selected) + current = hidden[token_idx] + gate = F.linear(current, w1_weights[expert_idx]) + up = F.linear(current, w3_weights[expert_idx]) + expert_out = F.linear(F.silu(gate) * up, w2_weights[expert_idx]) + expert_out = expert_out * routing_weights[token_idx, topk_pos, None] + output.index_add_(0, token_idx, expert_out.to(dtype=output.dtype)) + return output + + def _fake( + hidden: torch.Tensor, + router_logits: torch.Tensor, + expert_bias: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + use_expert_bias: bool, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden.new_empty(hidden.shape) + + try: + torch.library.impl(_LFM2_MOE_CUSTOM_OP, "CompositeExplicitAutograd")(_impl) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + try: + torch.library.register_fake(_LFM2_MOE_CUSTOM_OP)(_fake) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + _LFM2_MOE_CUSTOM_OP_REGISTERED = True + + +def _ensure_qwen2_moe_custom_op_registered() -> None: + global _QWEN2_MOE_CUSTOM_OP_REGISTERED + if _QWEN2_MOE_CUSTOM_OP_REGISTERED: + return + + schema = ( + "qwen2_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool normalize_routing, " + "float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + try: + library = torch.library.Library("cactus_transpile", "FRAGMENT") + library.define(schema) + _CACTUS_TORCH_LIBRARIES.append(library) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + + def _impl( + hidden: torch.Tensor, + router_logits: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32).to(dtype=router_logits.dtype) + routing_weights, selected_experts = torch.topk(routing_probs, k=int(num_experts_per_tok), dim=-1) + if normalize_routing: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + float(epsilon)) + routing_weights = routing_weights * float(routed_scaling_factor) + + output = torch.zeros_like(hidden) + for expert_idx in range(int(num_experts)): + selected = selected_experts == expert_idx + if not bool(selected.any()): + continue + token_idx, topk_pos = torch.where(selected) + current = hidden[token_idx] + gate = F.linear(current, w1_weights[expert_idx]) + up = F.linear(current, w3_weights[expert_idx]) + expert_out = F.linear(F.silu(gate) * up, w2_weights[expert_idx]) + expert_out = expert_out * routing_weights[token_idx, topk_pos, None] + output.index_add_(0, token_idx, expert_out.to(dtype=output.dtype)) + return output + + def _fake( + hidden: torch.Tensor, + router_logits: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden.new_empty(hidden.shape) + + try: + torch.library.impl(_QWEN2_MOE_CUSTOM_OP, "CompositeExplicitAutograd")(_impl) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + try: + torch.library.register_fake(_QWEN2_MOE_CUSTOM_OP)(_fake) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + _QWEN2_MOE_CUSTOM_OP_REGISTERED = True + + +def _ensure_gemma4_moe_custom_op_registered() -> None: + global _GEMMA4_MOE_CUSTOM_OP_REGISTERED + if _GEMMA4_MOE_CUSTOM_OP_REGISTERED: + return + + schema = ( + "gemma4_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool normalize_routing, " + "float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + try: + library = torch.library.Library("cactus_transpile", "FRAGMENT") + library.define(schema) + _CACTUS_TORCH_LIBRARIES.append(library) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + + def _impl( + hidden: torch.Tensor, + router_logits: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32).to(dtype=router_logits.dtype) + routing_weights, selected_experts = torch.topk(routing_probs, k=int(num_experts_per_tok), dim=-1) + if normalize_routing: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + float(epsilon)) + routing_weights = routing_weights * float(routed_scaling_factor) + + output = torch.zeros_like(hidden) + for expert_idx in range(int(num_experts)): + selected = selected_experts == expert_idx + if not bool(selected.any()): + continue + token_idx, topk_pos = torch.where(selected) + current = hidden[token_idx] + gate = F.linear(current, w1_weights[expert_idx]) + up = F.linear(current, w3_weights[expert_idx]) + expert_out = F.linear(F.silu(gate) * up, w2_weights[expert_idx]) + expert_out = expert_out * routing_weights[token_idx, topk_pos, None] + output.index_add_(0, token_idx, expert_out.to(dtype=output.dtype)) + return output + + def _fake( + hidden: torch.Tensor, + router_logits: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden.new_empty(hidden.shape) + + try: + torch.library.impl(_GEMMA4_MOE_CUSTOM_OP, "CompositeExplicitAutograd")(_impl) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + try: + torch.library.register_fake(_GEMMA4_MOE_CUSTOM_OP)(_fake) + except RuntimeError as exc: + if not _ignore_duplicate_torch_registration(exc): + raise + _GEMMA4_MOE_CUSTOM_OP_REGISTERED = True + + +class _Lfm2MoeRuntimeMoeBlock(torch.nn.Module): + """Export-friendly LFM2-MoE block backed by the Cactus sparse MoE runtime op.""" + + def __init__(self, block: torch.nn.Module, *, layer_index: int): + super().__init__() + _ensure_lfm2_moe_custom_op_registered() + experts = getattr(block, "experts", None) + gate_up = getattr(experts, "gate_up_proj", None) + down = getattr(experts, "down_proj", None) + gate = getattr(block, "gate", None) + if not isinstance(gate, torch.nn.Module) or not isinstance(gate_up, torch.Tensor) or not isinstance(down, torch.Tensor): + raise TypeError("LFM2-MoE runtime block requires gate, gate_up_proj, and down_proj") + + self.layer_index = int(layer_index) + self.gate = gate + self.num_experts = int(getattr(block, "num_experts", int(gate_up.shape[0])) or int(gate_up.shape[0])) + self.top_k = int(getattr(block, "top_k", 0) or 0) + self.use_expert_bias = bool(getattr(block, "use_expert_bias", False)) + self.normalize_routing = bool(getattr(block, "norm_topk_prob", False)) + self.routed_scaling_factor = float(getattr(block, "routed_scaling_factor", 1.0) or 1.0) + self.epsilon = 1.0e-6 + self.hidden_dim = int(gate_up.shape[2]) + self.intermediate_dim = int(gate_up.shape[1]) // 2 + + if self.num_experts <= 0 or self.top_k <= 0: + raise ValueError("LFM2-MoE runtime block requires positive num_experts and top_k") + if int(gate_up.shape[0]) != self.num_experts or int(down.shape[0]) != self.num_experts: + raise ValueError("LFM2-MoE expert tensor count mismatch") + if int(gate_up.shape[1]) != 2 * self.intermediate_dim: + raise ValueError("LFM2-MoE gate_up_proj must have an even expert dimension") + + self.w1_weights = torch.nn.ParameterList() + self.w3_weights = torch.nn.ParameterList() + self.w2_weights = torch.nn.ParameterList() + for expert_idx in range(self.num_experts): + self.w1_weights.append(torch.nn.Parameter( + gate_up[expert_idx, : self.intermediate_dim, :].detach().contiguous(), + requires_grad=False, + )) + self.w3_weights.append(torch.nn.Parameter( + gate_up[expert_idx, self.intermediate_dim :, :].detach().contiguous(), + requires_grad=False, + )) + self.w2_weights.append(torch.nn.Parameter(down[expert_idx].detach().contiguous(), requires_grad=False)) + + expert_bias = getattr(block, "expert_bias", None) + if self.use_expert_bias and isinstance(expert_bias, torch.Tensor): + bias_value = expert_bias.detach().to(dtype=torch.float32) + else: + bias_value = torch.zeros(self.num_experts, dtype=torch.float32, device=gate_up.device) + self.register_buffer("expert_bias", bias_value, persistent=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_flat = hidden_states.reshape(-1, hidden_dim) + router_logits = self.gate(hidden_flat) + output = torch.ops.cactus_transpile.lfm2_moe_layer_gated( + hidden_flat, + router_logits, + self.expert_bias, + list(self.w1_weights), + list(self.w3_weights), + list(self.w2_weights), + self.num_experts, + self.top_k, + self.use_expert_bias, + self.normalize_routing, + self.epsilon, + self.routed_scaling_factor, + ) + return output.reshape(batch_size, sequence_length, hidden_dim) + + +def _is_lfm2_moe_sparse_block(block: object) -> bool: + experts = getattr(block, "experts", None) + return ( + isinstance(block, torch.nn.Module) + and isinstance(getattr(block, "gate", None), torch.nn.Module) + and isinstance(getattr(experts, "gate_up_proj", None), torch.Tensor) + and isinstance(getattr(experts, "down_proj", None), torch.Tensor) + ) + + +def _is_lfm2_moe_model(model: torch.nn.Module) -> bool: + config = getattr(model, "config", None) + model_type = str(getattr(config, "model_type", "") or "").lower() + if model_type == "lfm2_moe": + return True + return type(model).__module__.startswith("transformers.models.lfm2_moe.") + + +def _lfm2_backbone_from_model(model: torch.nn.Module) -> torch.nn.Module | None: + model_root = getattr(model, "model", None) + language_model = getattr(model_root, "language_model", None) + backbone = language_model if isinstance(language_model, torch.nn.Module) else model_root + return backbone if isinstance(backbone, torch.nn.Module) else None + + +def _patch_lfm2_moe_feed_forwards(model: torch.nn.Module) -> None: + if not _is_lfm2_moe_model(model): + return + backbone = _lfm2_backbone_from_model(model) + layers = getattr(backbone, "layers", None) + if not isinstance(layers, torch.nn.ModuleList): + return + for layer_index, layer in enumerate(layers): + feed_forward = getattr(layer, "feed_forward", None) + if isinstance(feed_forward, _Lfm2MoeRuntimeMoeBlock): + continue + if _is_lfm2_moe_sparse_block(feed_forward): + setattr(layer, "feed_forward", _Lfm2MoeRuntimeMoeBlock(feed_forward, layer_index=layer_index)) + + +class _Qwen2MoeRuntimeMoeBlock(torch.nn.Module): + """Export-friendly Qwen2-MoE block backed by the Cactus sparse MoE runtime op.""" + + def __init__(self, block: torch.nn.Module, *, layer_index: int): + super().__init__() + _ensure_qwen2_moe_custom_op_registered() + experts = getattr(block, "experts", None) + gate_up = getattr(experts, "gate_up_proj", None) + down = getattr(experts, "down_proj", None) + gate = getattr(block, "gate", None) + gate_weight = getattr(gate, "weight", None) + shared_expert = getattr(block, "shared_expert", None) + shared_expert_gate = getattr(block, "shared_expert_gate", None) + if ( + not isinstance(gate, torch.nn.Module) + or not isinstance(gate_weight, torch.Tensor) + or not isinstance(gate_up, torch.Tensor) + or not isinstance(down, torch.Tensor) + or not isinstance(shared_expert, torch.nn.Module) + or not isinstance(shared_expert_gate, torch.nn.Module) + ): + raise TypeError( + "Qwen2-MoE runtime block requires gate, packed experts, shared_expert, and shared_expert_gate" + ) + + self.layer_index = int(layer_index) + self.gate = gate + self.shared_expert = shared_expert + self.shared_expert_gate = shared_expert_gate + self.num_experts = int(getattr(gate, "num_experts", int(gate_up.shape[0])) or int(gate_up.shape[0])) + self.top_k = int(getattr(gate, "top_k", 0) or 0) + self.normalize_routing = bool(getattr(gate, "norm_topk_prob", False)) + self.routed_scaling_factor = float(getattr(block, "routed_scaling_factor", 1.0) or 1.0) + self.epsilon = 1.0e-6 + self.hidden_dim = int(gate_up.shape[2]) + self.intermediate_dim = int(gate_up.shape[1]) // 2 + + if self.num_experts <= 0 or self.top_k <= 0: + raise ValueError("Qwen2-MoE runtime block requires positive num_experts and top_k") + if int(gate_up.shape[0]) != self.num_experts or int(down.shape[0]) != self.num_experts: + raise ValueError("Qwen2-MoE expert tensor count mismatch") + if int(gate_up.shape[1]) != 2 * self.intermediate_dim: + raise ValueError("Qwen2-MoE gate_up_proj must have an even expert dimension") + + self.w1_weights = torch.nn.ParameterList() + self.w3_weights = torch.nn.ParameterList() + self.w2_weights = torch.nn.ParameterList() + for expert_idx in range(self.num_experts): + self.w1_weights.append(torch.nn.Parameter( + gate_up[expert_idx, : self.intermediate_dim, :].detach().contiguous(), + requires_grad=False, + )) + self.w3_weights.append(torch.nn.Parameter( + gate_up[expert_idx, self.intermediate_dim :, :].detach().contiguous(), + requires_grad=False, + )) + self.w2_weights.append(torch.nn.Parameter(down[expert_idx].detach().contiguous(), requires_grad=False)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_flat = hidden_states.reshape(-1, hidden_dim) + router_logits = F.linear(hidden_flat, self.gate.weight) + routed_output = torch.ops.cactus_transpile.qwen2_moe_layer_gated( + hidden_flat, + router_logits, + list(self.w1_weights), + list(self.w3_weights), + list(self.w2_weights), + self.num_experts, + self.top_k, + self.normalize_routing, + self.epsilon, + self.routed_scaling_factor, + ) + shared_output = self.shared_expert(hidden_flat) + shared_output = torch.sigmoid(self.shared_expert_gate(hidden_flat)) * shared_output + output = routed_output + shared_output + return output.reshape(batch_size, sequence_length, hidden_dim) + + +def _is_qwen2_moe_sparse_block(block: object) -> bool: + experts = getattr(block, "experts", None) + gate = getattr(block, "gate", None) + return ( + isinstance(block, torch.nn.Module) + and isinstance(gate, torch.nn.Module) + and isinstance(getattr(gate, "weight", None), torch.Tensor) + and isinstance(getattr(experts, "gate_up_proj", None), torch.Tensor) + and isinstance(getattr(experts, "down_proj", None), torch.Tensor) + and isinstance(getattr(block, "shared_expert", None), torch.nn.Module) + and isinstance(getattr(block, "shared_expert_gate", None), torch.nn.Module) + ) + + +def _is_qwen2_moe_model(model: torch.nn.Module) -> bool: + config = getattr(model, "config", None) + model_type = str(getattr(config, "model_type", "") or "").lower() + if model_type == "qwen2_moe": + return True + return type(model).__module__.startswith("transformers.models.qwen2_moe.") + + +def _qwen2_moe_backbone_from_model(model: torch.nn.Module) -> torch.nn.Module | None: + backbone = getattr(model, "model", None) + return backbone if isinstance(backbone, torch.nn.Module) else None + + +def _patch_qwen2_moe_mlps(model: torch.nn.Module) -> None: + if not _is_qwen2_moe_model(model): + return + backbone = _qwen2_moe_backbone_from_model(model) + layers = getattr(backbone, "layers", None) + if not isinstance(layers, torch.nn.ModuleList): + return + for layer_index, layer in enumerate(layers): + mlp = getattr(layer, "mlp", None) + if isinstance(mlp, _Qwen2MoeRuntimeMoeBlock): + continue + if _is_qwen2_moe_sparse_block(mlp): + setattr(layer, "mlp", _Qwen2MoeRuntimeMoeBlock(mlp, layer_index=layer_index)) + + +class _Gemma4RuntimeMoeBlock(torch.nn.Module): + """Export-friendly Gemma4 MoE block with per-expert scale folded into w2.""" + + def __init__(self, block: torch.nn.Module, *, layer_index: int): + super().__init__() + _ensure_gemma4_moe_custom_op_registered() + gate_up = getattr(block, "gate_up_proj", None) + down = getattr(block, "down_proj", None) + per_expert_scale = getattr(block, "per_expert_scale", None) + if ( + not isinstance(gate_up, torch.Tensor) + or not isinstance(down, torch.Tensor) + or not isinstance(per_expert_scale, torch.Tensor) + ): + raise TypeError("Gemma4 MoE runtime block requires gate_up_proj, down_proj, and per_expert_scale") + + self.layer_index = int(layer_index) + self.num_experts = int(getattr(block, "num_experts", int(gate_up.shape[0])) or int(gate_up.shape[0])) + self.hidden_dim = int(gate_up.shape[2]) + self.intermediate_dim = int(gate_up.shape[1]) // 2 + if self.num_experts <= 0: + raise ValueError("Gemma4 MoE runtime block requires positive num_experts") + if int(gate_up.shape[0]) != self.num_experts or int(down.shape[0]) != self.num_experts: + raise ValueError("Gemma4 MoE expert tensor count mismatch") + if int(gate_up.shape[1]) != 2 * self.intermediate_dim: + raise ValueError("Gemma4 MoE gate_up_proj must have an even expert dimension") + + self.w1_weights = torch.nn.ParameterList() + self.w3_weights = torch.nn.ParameterList() + self.w2_weights = torch.nn.ParameterList() + for expert_idx in range(self.num_experts): + self.w1_weights.append(torch.nn.Parameter( + gate_up[expert_idx, : self.intermediate_dim, :].detach().contiguous(), + requires_grad=False, + )) + self.w3_weights.append(torch.nn.Parameter( + gate_up[expert_idx, self.intermediate_dim :, :].detach().contiguous(), + requires_grad=False, + )) + scale = per_expert_scale[expert_idx].detach().to(device=down.device, dtype=down.dtype) + self.w2_weights.append(torch.nn.Parameter( + (down[expert_idx].detach() * scale).contiguous(), + requires_grad=False, + )) + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + *, + top_k: int, + normalize_routing: bool = True, + epsilon: float = 1.0e-6, + ) -> torch.Tensor: + return torch.ops.cactus_transpile.gemma4_moe_layer_gated( + hidden_states, + router_logits, + list(self.w1_weights), + list(self.w3_weights), + list(self.w2_weights), + self.num_experts, + int(top_k), + bool(normalize_routing), + float(epsilon), + 1.0, + ) + + +class _Gemma4MoeRuntimeDecoderLayer(torch.nn.Module): + """Gemma4 decoder layer that preserves HF math while exporting MoE as one semantic op.""" + + def __init__(self, layer: torch.nn.Module, *, layer_index: int): + super().__init__() + self.config = layer.config + self.hidden_size = layer.hidden_size + self.layer_idx = layer.layer_idx + self.attention_type = layer.attention_type + self.self_attn = layer.self_attn + self.mlp = layer.mlp + self.input_layernorm = layer.input_layernorm + self.post_attention_layernorm = layer.post_attention_layernorm + self.pre_feedforward_layernorm = layer.pre_feedforward_layernorm + self.post_feedforward_layernorm = layer.post_feedforward_layernorm + self.register_buffer("layer_scalar", layer.layer_scalar.detach().clone(), persistent=True) + self.hidden_size_per_layer_input = getattr(layer, "hidden_size_per_layer_input", 0) + if self.hidden_size_per_layer_input: + self.act_fn = layer.act_fn + self.per_layer_input_gate = layer.per_layer_input_gate + self.per_layer_projection = layer.per_layer_projection + self.post_per_layer_input_norm = layer.post_per_layer_input_norm + self.enable_moe_block = True + self.router = layer.router + self.moe = _Gemma4RuntimeMoeBlock(layer.moe, layer_index=layer_index) + self.post_feedforward_layernorm_1 = layer.post_feedforward_layernorm_1 + self.post_feedforward_layernorm_2 = layer.post_feedforward_layernorm_2 + self.pre_feedforward_layernorm_2 = layer.pre_feedforward_layernorm_2 + self.top_k_experts = int(getattr(self.config, "top_k_experts", 0) or 0) + self.router_epsilon = float(getattr(self.config, "rms_norm_eps", 1.0e-6) or 1.0e-6) + + def _router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + router_states = self.router.norm(hidden_states) + router_states = router_states * self.router.root_size.to(router_states.dtype) + router_states = router_states * self.router.scale.to(router_states.dtype) + return self.router.proj(router_states) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: torch.Tensor = None, + per_layer_input: torch.Tensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: object | None = None, + use_cache: bool | None = False, + **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, + use_cache=use_cache, + **kwargs, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states_1 = self.pre_feedforward_layernorm(hidden_states) + hidden_states_1 = self.mlp(hidden_states_1) + hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states_1) + + hidden_states_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + router_logits = self._router_logits(hidden_states_flat) + hidden_states_2 = self.pre_feedforward_layernorm_2(hidden_states_flat) + hidden_states_2 = self.moe( + hidden_states_2, + router_logits, + top_k=self.top_k_experts, + normalize_routing=True, + epsilon=self.router_epsilon, + ) + hidden_states_2 = hidden_states_2.reshape(hidden_states.shape) + hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2) + + hidden_states = hidden_states_1 + hidden_states_2 + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + if self.hidden_size_per_layer_input: + residual = hidden_states + hidden_states = self.per_layer_input_gate(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = torch.multiply(hidden_states, per_layer_input) + hidden_states = self.per_layer_projection(hidden_states) + hidden_states = self.post_per_layer_input_norm(hidden_states) + hidden_states = residual + hidden_states + + hidden_states = hidden_states * self.layer_scalar + return hidden_states + + +def _is_gemma4_moe_decoder_layer(layer: object) -> bool: + return ( + isinstance(layer, torch.nn.Module) + and bool(getattr(layer, "enable_moe_block", False)) + and hasattr(layer, "router") + and hasattr(layer, "moe") + and not isinstance(layer, _Gemma4MoeRuntimeDecoderLayer) + ) + + +def _gemma4_backbone_from_model(model: torch.nn.Module) -> torch.nn.Module | None: + model_root = getattr(model, "model", None) + backbone = getattr(model_root, "language_model", model_root) + return backbone if isinstance(backbone, torch.nn.Module) else None + + +def _patch_gemma4_moe_layers(model: torch.nn.Module) -> None: + if _family_key(model) != "gemma4": + return + backbone = _gemma4_backbone_from_model(model) + layers = getattr(backbone, "layers", None) + if not isinstance(layers, torch.nn.ModuleList): + return + for layer_index, layer in enumerate(layers): + if isinstance(layer, _Gemma4MoeRuntimeDecoderLayer): + continue + if _is_gemma4_moe_decoder_layer(layer): + layers[layer_index] = _Gemma4MoeRuntimeDecoderLayer(layer, layer_index=layer_index) + + def _gemma4_cpu_safe_text_mlp_enabled(layer: torch.nn.Module, hidden_states: torch.Tensor) -> bool: if hidden_states.device.type != "cpu" or hidden_states.dtype != torch.float16: return False @@ -1519,12 +2213,14 @@ def get_transpile_metadata(self): class Lfm2CausalLMLogitsAdapter(torch.nn.Module): def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): super().__init__() + _patch_lfm2_moe_feed_forwards(model) self.model = model model_root = getattr(model, "model", None) language_model = getattr(model_root, "language_model", None) self.backbone = language_model if isinstance(language_model, torch.nn.Module) else model_root self.lm_head = getattr(model, "lm_head", None) self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + self.adapter_family = _family_key(model) from transformers.models.lfm2.modeling_lfm2 import create_causal_mask # type: ignore self._create_causal_mask = create_causal_mask @@ -1580,7 +2276,7 @@ def get_transpile_metadata(self): "graph": { **_transpile_graph_meta( self.model, - adapter_family="lfm2", + adapter_family=self.adapter_family, adapter_type=type(self).__name__, input_names=("input_ids",), ), @@ -1592,12 +2288,14 @@ def get_transpile_metadata(self): class Lfm2CausalLMStepAdapter(torch.nn.Module): def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): super().__init__() + _patch_lfm2_moe_feed_forwards(model) self.model = model model_root = getattr(model, "model", None) language_model = getattr(model_root, "language_model", None) self.backbone = language_model if isinstance(language_model, torch.nn.Module) else model_root self.lm_head = getattr(model, "lm_head", None) self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + self.adapter_family = _family_key(model) def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: backbone = self.backbone @@ -1635,7 +2333,7 @@ def get_transpile_metadata(self): "graph": { **_transpile_graph_meta( self.model, - adapter_family="lfm2", + adapter_family=self.adapter_family, adapter_type=type(self).__name__, input_names=("input_ids", "position_ids"), ), @@ -1653,6 +2351,7 @@ def _lfm2_vl_model_root(model: torch.nn.Module) -> torch.nn.Module: def _lfm2_language_backbone(model: torch.nn.Module) -> torch.nn.Module: + _patch_lfm2_moe_feed_forwards(model) root = getattr(model, "model", None) language_model = getattr(root, "language_model", None) backbone = language_model if isinstance(language_model, torch.nn.Module) else root @@ -2162,6 +2861,7 @@ def _additive_mask_mapping(self, inputs_embeds: torch.Tensor) -> dict[str, torch class Gemma4CausalLMLogitsAdapter(torch.nn.Module): def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): super().__init__() + _patch_gemma4_moe_layers(model) self.model = model model_backbone = model.model self.backbone = getattr(model_backbone, "language_model", model_backbone) @@ -2342,6 +3042,7 @@ def _gemma4_apply_final_logit_softcapping(model: torch.nn.Module, logits: torch. class Gemma4MultimodalCausalLMLogitsAdapter(BoundInputAdapter): def __init__(self, model: torch.nn.Module, *, input_names: tuple[str, ...], weights_dir: str | None = None): + _patch_gemma4_moe_layers(model) super().__init__( model, input_names=input_names, @@ -2820,6 +3521,7 @@ def __init__( native_image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None, ): super().__init__() + _patch_gemma4_moe_layers(model) self.model = model self.input_names = input_names self.multimodal_backbone = model.model @@ -2983,14 +3685,19 @@ def _prepare_decoder_inputs( attention_mask: torch.Tensor | None, token_type_ids: torch.Tensor, image_features: torch.Tensor, - audio_features: torch.Tensor, + audio_features: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: if self._native_merge_plan is not None: + audio_features_for_merge = audio_features + if audio_features_for_merge is None: + audio_features_for_merge = image_features.new_empty( + (int(input_ids.shape[0]), 0, int(image_features.shape[-1])) + ) inputs_embeds, per_layer_inputs_tokens = _gemma4_apply_native_merge_plan( self.model, input_ids=input_ids, image_features=image_features, - audio_features=audio_features, + audio_features=audio_features_for_merge, merge_plan=self._native_merge_plan, pli_token_ids=self._native_merge_pli_token_ids, ) @@ -3020,10 +3727,11 @@ def _prepare_decoder_inputs( image_mask.unsqueeze(-1).expand_as(inputs_embeds), image_features.to(inputs_embeds.device, inputs_embeds.dtype), ) - inputs_embeds = inputs_embeds.masked_scatter( - audio_mask.unsqueeze(-1).expand_as(inputs_embeds), - audio_features.to(inputs_embeds.device, inputs_embeds.dtype), - ) + if audio_features is not None: + inputs_embeds = inputs_embeds.masked_scatter( + audio_mask.unsqueeze(-1).expand_as(inputs_embeds), + audio_features.to(inputs_embeds.device, inputs_embeds.dtype), + ) per_layer_inputs_tokens = input_ids * text_mask.to(dtype=input_ids.dtype) per_layer_inputs: torch.Tensor @@ -3130,10 +3838,15 @@ def __init__( *, weights_dir: str | None = None, native_merge_plan: _Gemma4NativeMergePlan | None = None, + include_audio: bool = True, ): + self.include_audio = bool(include_audio) + input_names = ("input_ids", "attention_mask", "token_type_ids", "image_features") + if self.include_audio: + input_names = (*input_names, "audio_features") super().__init__( model, - input_names=("input_ids", "attention_mask", "token_type_ids", "image_features", "audio_features"), + input_names=input_names, weights_dir=weights_dir, native_merge_plan=native_merge_plan, ) @@ -3144,7 +3857,7 @@ def forward( attention_mask: torch.Tensor | None, token_type_ids: torch.Tensor, image_features: torch.Tensor, - audio_features: torch.Tensor, + audio_features: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.LongTensor]: return self._prepare_decoder_inputs( input_ids=input_ids, @@ -3158,7 +3871,7 @@ def get_transpile_metadata(self): return { "graph": self._base_graph_meta( adapter_type=type(self).__name__, - input_names=("input_ids", "attention_mask", "token_type_ids", "image_features", "audio_features"), + input_names=self.input_names, ), } @@ -3541,18 +4254,27 @@ def _build_gemma4_multimodal_component_specs( components: tuple[str, ...] | None = None, cache_context_length: str | int | None = None, ) -> list[ComponentModuleSpec]: - pixel_values = named_tensors["pixel_values"] + pixel_values = named_tensors.get("pixel_values") pixel_position_ids = named_tensors.get("pixel_position_ids") - input_features = named_tensors["input_features"] - input_features_mask = named_tensors["input_features_mask"] + input_features = named_tensors.get("input_features") + input_features_mask = named_tensors.get("input_features_mask") input_ids = named_tensors["input_ids"] attention_mask = named_tensors.get("attention_mask") token_type_ids = named_tensors["token_type_ids"] - requested_components = tuple(components or ("vision_encoder", "audio_encoder", "lm_encoder", "decoder")) + has_vision_inputs = pixel_values is not None + has_audio_inputs = input_features is not None and input_features_mask is not None + default_components = ["lm_encoder", "decoder"] + if has_vision_inputs: + default_components.insert(0, "vision_encoder") + if has_audio_inputs: + insert_at = 1 if has_vision_inputs else 0 + default_components.insert(insert_at, "audio_encoder") + requested_components = tuple(components or tuple(default_components)) requested_set = set(requested_components) if not requested_set: return [] + include_audio = "audio_encoder" in requested_set or (components is None and has_audio_inputs) expanded_components: list[str] = [] @@ -3561,8 +4283,10 @@ def _require(component: str) -> None: expanded_components.append(component) if "decoder" in requested_set: - _require("vision_encoder") - _require("audio_encoder") + if has_vision_inputs or "vision_encoder" in requested_set: + _require("vision_encoder") + if include_audio: + _require("audio_encoder") _require("lm_encoder") _require("decoder_prefill_chunk") _require("decoder_embed_chunk") @@ -3572,8 +4296,10 @@ def _require(component: str) -> None: _require("lm_encoder_media_step") _require("decoder_step") if "lm_encoder" in requested_set: - _require("vision_encoder") - _require("audio_encoder") + if has_vision_inputs or "vision_encoder" in requested_set: + _require("vision_encoder") + if include_audio: + _require("audio_encoder") _require("lm_encoder") if "decoder_prefill_chunk" in requested_set: _require("decoder_prefill_chunk") @@ -3593,30 +4319,39 @@ def _require(component: str) -> None: _require("lm_encoder_step") _require("decoder_step") - vision_tower = getattr(getattr(model, "model", model), "vision_tower", None) - pooling_kernel_size = int(_module_or_config_attr(vision_tower, "pooling_kernel_size", 3) or 3) - native_image_soft_token_counts = _gemma4_static_image_soft_token_counts( - pixel_position_ids, - pooling_kernel_size=pooling_kernel_size, - ) - native_image_pool_shapes = _gemma4_static_image_pool_shapes( - pixel_position_ids, - pooling_kernel_size=pooling_kernel_size, - ) + if "vision_encoder" in expanded_components and pixel_values is None: + raise RuntimeError("Gemma4 vision component requested but pixel_values were not prepared") + if "audio_encoder" in expanded_components and (input_features is None or input_features_mask is None): + raise RuntimeError("Gemma4 audio component requested but input_features/input_features_mask were not prepared") - vision_encoder = Gemma4VisionEncoderAdapter( - model, - weights_dir=weights_dir, - native_image_soft_token_counts=native_image_soft_token_counts, - native_image_pool_shapes=native_image_pool_shapes, - ).eval() - vision_encoder_npu = Gemma4VisionEncoderAdapter( - model, - weights_dir=weights_dir, - native_image_soft_token_counts=native_image_soft_token_counts, - native_image_pool_shapes=None, - ).eval() - audio_encoder = Gemma4AudioEncoderAdapter(model, weights_dir=weights_dir).eval() + vision_encoder: Gemma4VisionEncoderAdapter | None = None + vision_encoder_npu: Gemma4VisionEncoderAdapter | None = None + audio_encoder: Gemma4AudioEncoderAdapter | None = None + if "vision_encoder" in expanded_components: + vision_tower = getattr(getattr(model, "model", model), "vision_tower", None) + pooling_kernel_size = int(_module_or_config_attr(vision_tower, "pooling_kernel_size", 3) or 3) + native_image_soft_token_counts = _gemma4_static_image_soft_token_counts( + pixel_position_ids, + pooling_kernel_size=pooling_kernel_size, + ) + native_image_pool_shapes = _gemma4_static_image_pool_shapes( + pixel_position_ids, + pooling_kernel_size=pooling_kernel_size, + ) + vision_encoder = Gemma4VisionEncoderAdapter( + model, + weights_dir=weights_dir, + native_image_soft_token_counts=native_image_soft_token_counts, + native_image_pool_shapes=native_image_pool_shapes, + ).eval() + vision_encoder_npu = Gemma4VisionEncoderAdapter( + model, + weights_dir=weights_dir, + native_image_soft_token_counts=native_image_soft_token_counts, + native_image_pool_shapes=None, + ).eval() + if "audio_encoder" in expanded_components: + audio_encoder = Gemma4AudioEncoderAdapter(model, weights_dir=weights_dir).eval() image_features: torch.Tensor | None = None audio_features: torch.Tensor | None = None @@ -3633,27 +4368,33 @@ def _require(component: str) -> None: "lm_encoder" in expanded_components or "decoder" in expanded_components ): + if vision_encoder is None or pixel_values is None: + raise RuntimeError("Gemma4 vision features requested but vision encoder inputs are unavailable") image_features = vision_encoder(pixel_values, pixel_position_ids) if "audio_encoder" in expanded_components and ( "lm_encoder" in expanded_components or "decoder" in expanded_components ): + if audio_encoder is None or input_features is None or input_features_mask is None: + raise RuntimeError("Gemma4 audio features requested but audio encoder inputs are unavailable") audio_features = audio_encoder(input_features, input_features_mask) if "lm_encoder" in expanded_components or "decoder" in expanded_components: - if image_features is None: + if image_features is None and vision_encoder is not None and pixel_values is not None: image_features = vision_encoder(pixel_values, pixel_position_ids) - if audio_features is None: + if include_audio and audio_features is None and audio_encoder is not None and input_features is not None and input_features_mask is not None: audio_features = audio_encoder(input_features, input_features_mask) - native_merge_plan = _gemma4_build_native_merge_plan( - getattr(model, "model", model), - input_ids, - image_feature_count=_gemma4_feature_token_count(image_features), - audio_feature_count=_gemma4_feature_token_count(audio_features), - ) + if image_features is not None: + native_merge_plan = _gemma4_build_native_merge_plan( + getattr(model, "model", model), + input_ids, + image_feature_count=_gemma4_feature_token_count(image_features), + audio_feature_count=_gemma4_feature_token_count(audio_features), + ) lm_encoder = Gemma4LMEncoderAdapter( model, weights_dir=weights_dir, native_merge_plan=native_merge_plan, + include_audio=include_audio, ).eval() encoder_chunk_size = max(1, int(os.environ.get("CACTUS_GEMMA4_ENCODER_CHUNK", "128") or "128")) encoder_chunk_size = min(encoder_chunk_size, int(input_ids.shape[1])) @@ -3668,13 +4409,27 @@ def _require(component: str) -> None: with torch.no_grad(): if "decoder" in expanded_components: - decoder_inputs = lm_encoder( - input_ids, - attention_mask, - token_type_ids, - image_features, - audio_features, - ) + if image_features is None: + raise RuntimeError("Gemma4 decoder spec requires precomputed image features") + if include_audio and audio_features is None: + raise RuntimeError("Gemma4 decoder spec requires precomputed audio features") + lm_encoder_inputs: tuple[torch.Tensor | None, ...] + if include_audio: + lm_encoder_inputs = ( + input_ids, + attention_mask, + token_type_ids, + image_features, + audio_features, + ) + else: + lm_encoder_inputs = ( + input_ids, + attention_mask, + token_type_ids, + image_features, + ) + decoder_inputs = lm_encoder(*lm_encoder_inputs) if "lm_encoder_step" in expanded_components or "decoder_step" in expanded_components: step_input_ids = input_ids[:, -1:].contiguous() step_position_ids = torch.full( @@ -3735,6 +4490,8 @@ def _require(component: str) -> None: } specs: list[ComponentModuleSpec] = [] if "vision_encoder" in expanded_components: + if vision_encoder is None or vision_encoder_npu is None or pixel_values is None: + raise RuntimeError("Gemma4 vision_encoder spec requires prepared pixel_values") specs.append(ComponentModuleSpec( component="vision_encoder", module=vision_encoder, @@ -3747,6 +4504,8 @@ def _require(component: str) -> None: npu_runtime_input_count=2, )) if "audio_encoder" in expanded_components: + if audio_encoder is None or input_features is None or input_features_mask is None: + raise RuntimeError("Gemma4 audio_encoder spec requires prepared audio features") specs.append(ComponentModuleSpec( component="audio_encoder", module=audio_encoder, @@ -3757,13 +4516,22 @@ def _require(component: str) -> None: metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, )) if "lm_encoder" in expanded_components: - if image_features is None or audio_features is None: - raise RuntimeError("Gemma4 lm_encoder spec requires precomputed image/audio features") + if image_features is None: + raise RuntimeError("Gemma4 lm_encoder spec requires precomputed image features") + if include_audio and audio_features is None: + raise RuntimeError("Gemma4 lm_encoder spec requires precomputed audio features") + lm_encoder_example_inputs: tuple[torch.Tensor | None, ...] + lm_encoder_input_keys = ("input_ids", "attention_mask", "token_type_ids", "image_features") + if include_audio: + lm_encoder_example_inputs = (input_ids, attention_mask, token_type_ids, image_features, audio_features) + lm_encoder_input_keys = (*lm_encoder_input_keys, "audio_features") + else: + lm_encoder_example_inputs = (input_ids, attention_mask, token_type_ids, image_features) specs.append(ComponentModuleSpec( component="lm_encoder", module=lm_encoder, - example_inputs=(input_ids, attention_mask, token_type_ids, image_features, audio_features), - input_keys=("input_ids", "attention_mask", "token_type_ids", "image_features", "audio_features"), + example_inputs=lm_encoder_example_inputs, + input_keys=lm_encoder_input_keys, output_keys=_GEMMA4_DECODER_PIPELINE_IO_KEYS, graph_meta={**common_graph_meta, "component": "lm_encoder"}, metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, @@ -4373,6 +5141,7 @@ def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): self.model = model self.backbone = model.model self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + self.adapter_family = _family_key(model) def forward(self, input_ids: torch.Tensor): return self.debug_forward(input_ids)[0] @@ -4434,7 +5203,7 @@ def get_transpile_metadata(self): "graph": { **_transpile_graph_meta( self.model, - adapter_family="qwen3", + adapter_family=self.adapter_family, adapter_type=type(self).__name__, input_names=("input_ids",), ), @@ -4451,6 +5220,7 @@ def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): self.model = model self.backbone = model.model self.pad_token_id = pad_token_id if pad_token_id is not None else _resolve_model_pad_token_id(model) + self.adapter_family = _family_key(model) def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: inputs_embeds = self.backbone.embed_tokens(input_ids) @@ -4495,7 +5265,7 @@ def get_transpile_metadata(self): "graph": { **_transpile_graph_meta( self.model, - adapter_family="qwen3", + adapter_family=self.adapter_family, adapter_type=type(self).__name__, input_names=("input_ids", "position_ids"), ), @@ -4506,6 +5276,18 @@ def get_transpile_metadata(self): } +class Qwen2MoeCausalLMLogitsAdapter(Qwen3CausalLMLogitsAdapter): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + _patch_qwen2_moe_mlps(model) + super().__init__(model, pad_token_id=pad_token_id) + + +class Qwen2MoeCausalLMStepAdapter(Qwen3CausalLMStepAdapter): + def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): + _patch_qwen2_moe_mlps(model) + super().__init__(model, pad_token_id=pad_token_id) + + class Qwen3LMEncoderStepAdapter(torch.nn.Module): def __init__(self, model: torch.nn.Module): super().__init__() @@ -4619,6 +5401,9 @@ def _build_qwen_causal_lm_component_specs( if family == "qwen3_5": decoder = Qwen35CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() decoder_step = Qwen35CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() + elif family == "qwen2_moe": + decoder = Qwen2MoeCausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + decoder_step = Qwen2MoeCausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() elif family == "qwen3": decoder = Qwen3CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() decoder_step = Qwen3CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() @@ -5352,6 +6137,8 @@ def _family_key(model: torch.nn.Module) -> str: return "qwen3_5" if module_name.startswith("transformers.models.qwen3_vl."): return "qwen3_5" + if module_name.startswith("transformers.models.qwen2_moe."): + return "qwen2_moe" if module_name.startswith("transformers.models.qwen3."): return "qwen3" if module_name.startswith("transformers.models.lfm2_vl."): @@ -5363,6 +6150,8 @@ def _family_key(model: torch.nn.Module) -> str: model_type = str(getattr(getattr(model, "config", None), "model_type", "") or "").lower() if model_type == "needle": return "needle" + if model_type == "qwen2_moe": + return "qwen2_moe" if "nomic" in module_name.lower() or "nomic" in model_type: return "nomic" return "generic" @@ -5711,12 +6500,13 @@ def _build_lfm2_causal_lm_component_specs( return None requested = tuple(components or ("decoder", "decoder_step")) requested_set = set(requested) + family = _family_key(model) common_graph_meta = { "weights_dir": weights_dir, "task": "causal_lm_logits", - "adapter_family": "lfm2", + "adapter_family": family, } - metadata = {"family": "lfm2", "task": "causal_lm_logits"} + metadata = {"family": family, "task": "causal_lm_logits"} chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk"} if chunk_components & requested_set: return _lfm2_chunked_pipeline_specs( @@ -6083,7 +6873,7 @@ def build_component_module_specs( components=components, cache_context_length=cache_context_length, ) - if family in {"qwen3", "qwen3_5"} and task == "causal_lm_logits": + if family in {"qwen2_moe", "qwen3", "qwen3_5"} and task == "causal_lm_logits": return _build_qwen_causal_lm_component_specs( model, named_tensors=named_tensors, @@ -6115,7 +6905,7 @@ def build_component_module_specs( components=components, cache_context_length=cache_context_length, ) - if family == "lfm2" and task == "causal_lm_logits": + if family in {"lfm2", "lfm2_moe"} and task == "causal_lm_logits": return _build_lfm2_causal_lm_component_specs( model, named_tensors=named_tensors, @@ -6191,6 +6981,11 @@ def canonicalize_model_interface( inner_model, pad_token_id=padding_token_id_value, ) + elif family == "qwen2_moe": + adapter_factory = lambda inner_model: Qwen2MoeCausalLMLogitsAdapter( # type: ignore[assignment] + inner_model, + pad_token_id=padding_token_id_value, + ) elif family == "qwen3": adapter_factory = lambda inner_model: Qwen3CausalLMLogitsAdapter( # type: ignore[assignment] inner_model, diff --git a/python/cactus/transpile/model_profiles.py b/python/cactus/transpile/model_profiles.py index 75580eb47..7d132182b 100644 --- a/python/cactus/transpile/model_profiles.py +++ b/python/cactus/transpile/model_profiles.py @@ -74,7 +74,7 @@ class ModelProfile: LFM2_PROFILE = ModelProfile( family="lfm2", - model_types=("lfm2_vl", "lfm2"), + model_types=("lfm2_vl", "lfm2", "lfm2_moe"), multimodal_context_tokens=256, input_combinations=((), ("image",)), model_id_aliases=(("lfm", "LiquidAI/LFM2-VL-450M"),), @@ -94,6 +94,8 @@ class ModelProfile: text_component_plans=( ("lfm2forcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk")), + ("lfm2moeforcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", + "decoder_prefill_chunk")), ), needs_image=True, force_component_pipeline=True, @@ -172,7 +174,7 @@ class ModelProfile: QWEN_PROFILE = ModelProfile( family="qwen", - model_types=("qwen", "qwen2", "qwen3", "qwen3_5", "qwen3.5", "qwen3_vl"), + model_types=("qwen", "qwen2", "qwen2_moe", "qwen3", "qwen3_5", "qwen3.5", "qwen3_vl"), multimodal_context_tokens=512, input_combinations=((), ("image",)), model_id_aliases=( @@ -183,7 +185,7 @@ class ModelProfile: ("qwen3", "Qwen/Qwen3-1.7B"), ), model_id_markers=("qwen",), - family_aliases=("qwen2", "qwen3", "qwen3_5", "qwen3.5", "qwen3_vl"), + family_aliases=("qwen2", "qwen2_moe", "qwen3", "qwen3_5", "qwen3.5", "qwen3_vl"), stop_tokens=("<|im_end|>",), avoid_native_loader=True, cached_step_components=("decoder_media_step",), diff --git a/python/cactus/transpile/ops.py b/python/cactus/transpile/ops.py index e1a870f8b..5e3ee5436 100644 --- a/python/cactus/transpile/ops.py +++ b/python/cactus/transpile/ops.py @@ -107,6 +107,43 @@ class OpSchema: # CUSTOM / FUSED OPS OpSchema("dense_mlp_tq_fused", num_inputs=4, attrs=("product_scale",), backend_op="dense_mlp_tq_fused"), + OpSchema( + "lfm2_moe_layer_gated", + num_inputs=None, + attrs=( + "num_experts", + "num_experts_per_tok", + "use_expert_bias", + "normalize_routing", + "epsilon", + "routed_scaling_factor", + "activation", + ), + ), + OpSchema( + "qwen2_moe_layer_gated", + num_inputs=None, + attrs=( + "num_experts", + "num_experts_per_tok", + "normalize_routing", + "epsilon", + "routed_scaling_factor", + "activation", + ), + ), + OpSchema( + "gemma4_moe_layer_gated", + num_inputs=None, + attrs=( + "num_experts", + "num_experts_per_tok", + "normalize_routing", + "epsilon", + "routed_scaling_factor", + "activation", + ), + ), OpSchema("rms_norm", num_inputs=2, attrs=("eps",), backend_op="rms_norm"), OpSchema("rope", num_inputs=1, attrs=("theta", "position_offset"), backend_op="rope"), OpSchema("rel_pos_bias", num_inputs=2, attrs=("scale",), backend_op="rel_pos_bias"), diff --git a/python/tests/test_component_plan.py b/python/tests/test_component_plan.py new file mode 100644 index 000000000..4bd7bbb73 --- /dev/null +++ b/python/tests/test_component_plan.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from cactus.transpile.component_plan import infer_component_plan_from_config + + +def test_gemma4_vision_only_config_does_not_request_audio_component() -> None: + config = { + "model_type": "gemma4", + "architectures": ["Gemma4ForConditionalGeneration"], + "vision_config": {"hidden_size": 1152}, + "text_config": {"num_hidden_layers": 30}, + } + + plan = infer_component_plan_from_config(config, model_id="google/gemma-4-26B-A4B") + + assert plan is not None + assert plan.task == "multimodal_causal_lm_logits" + assert plan.components == ("vision_encoder", "lm_encoder", "decoder") + assert plan.needs_image + assert not plan.needs_audio + + +def test_gemma4_image_audio_config_keeps_audio_component() -> None: + config = { + "model_type": "gemma4", + "architectures": ["Gemma4ForConditionalGeneration"], + "vision_config": {"hidden_size": 1152}, + "audio_config": {"hidden_size": 1024}, + "text_config": {"num_hidden_layers": 30}, + } + + plan = infer_component_plan_from_config(config, model_id="google/gemma-4-E2B-it") + + assert plan is not None + assert plan.components == ("vision_encoder", "audio_encoder", "lm_encoder", "decoder") + assert plan.needs_audio diff --git a/python/tests/test_transpile_aten_ops.py b/python/tests/test_transpile_aten_ops.py index b0f0ad7f4..935e6cbcf 100644 --- a/python/tests/test_transpile_aten_ops.py +++ b/python/tests/test_transpile_aten_ops.py @@ -122,3 +122,327 @@ def test_dense_mlp_fusion_uses_topology_not_layer_names() -> None: assert fuse_dense_mlp_tq(graph) is True assert graph.nodes["down_linear"].op == "dense_mlp_tq_fused" assert graph.nodes["down_linear"].inputs == ["x", "pear_weight", "banana_weight", "plum_weight"] + + +_TEST_LFM2_MOE_OP_REGISTERED = False +_TEST_QWEN2_MOE_OP_REGISTERED = False +_TEST_GEMMA4_MOE_OP_REGISTERED = False +_TEST_TORCH_LIBRARIES: list[object] = [] + + +def _ignore_duplicate_registration(exc: RuntimeError) -> bool: + text = str(exc).lower() + return "already" in text or "duplicate" in text or "previously" in text + + +def _ensure_test_lfm2_moe_op_registered() -> None: + global _TEST_LFM2_MOE_OP_REGISTERED + if _TEST_LFM2_MOE_OP_REGISTERED: + return + try: + library = torch.library.Library("cactus_transpile", "FRAGMENT") + library.define( + "lfm2_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, Tensor expert_bias, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool use_expert_bias, " + "bool normalize_routing, float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + _TEST_TORCH_LIBRARIES.append(library) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + + def impl( + hidden: torch.Tensor, + router_logits: torch.Tensor, + expert_bias: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + use_expert_bias: bool, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden + router_logits.sum() * 0.0 + + def fake( + hidden: torch.Tensor, + router_logits: torch.Tensor, + expert_bias: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + use_expert_bias: bool, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden.new_empty(hidden.shape) + + try: + torch.library.impl("cactus_transpile::lfm2_moe_layer_gated", "CompositeExplicitAutograd")(impl) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + try: + torch.library.register_fake("cactus_transpile::lfm2_moe_layer_gated")(fake) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + _TEST_LFM2_MOE_OP_REGISTERED = True + + +def _ensure_test_qwen2_moe_op_registered() -> None: + global _TEST_QWEN2_MOE_OP_REGISTERED + if _TEST_QWEN2_MOE_OP_REGISTERED: + return + try: + library = torch.library.Library("cactus_transpile", "FRAGMENT") + library.define( + "qwen2_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool normalize_routing, " + "float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + _TEST_TORCH_LIBRARIES.append(library) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + + def impl( + hidden: torch.Tensor, + router_logits: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden + router_logits.sum() * 0.0 + + def fake( + hidden: torch.Tensor, + router_logits: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden.new_empty(hidden.shape) + + try: + torch.library.impl("cactus_transpile::qwen2_moe_layer_gated", "CompositeExplicitAutograd")(impl) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + try: + torch.library.register_fake("cactus_transpile::qwen2_moe_layer_gated")(fake) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + _TEST_QWEN2_MOE_OP_REGISTERED = True + + +def _ensure_test_gemma4_moe_op_registered() -> None: + global _TEST_GEMMA4_MOE_OP_REGISTERED + if _TEST_GEMMA4_MOE_OP_REGISTERED: + return + try: + library = torch.library.Library("cactus_transpile", "FRAGMENT") + library.define( + "gemma4_moe_layer_gated(" + "Tensor hidden, Tensor router_logits, " + "Tensor[] w1_weights, Tensor[] w3_weights, Tensor[] w2_weights, " + "int num_experts, int num_experts_per_tok, bool normalize_routing, " + "float epsilon, float routed_scaling_factor" + ") -> Tensor" + ) + _TEST_TORCH_LIBRARIES.append(library) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + + def impl( + hidden: torch.Tensor, + router_logits: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden + router_logits.sum() * 0.0 + + def fake( + hidden: torch.Tensor, + router_logits: torch.Tensor, + w1_weights: list[torch.Tensor], + w3_weights: list[torch.Tensor], + w2_weights: list[torch.Tensor], + num_experts: int, + num_experts_per_tok: int, + normalize_routing: bool, + epsilon: float, + routed_scaling_factor: float, + ) -> torch.Tensor: + return hidden.new_empty(hidden.shape) + + try: + torch.library.impl("cactus_transpile::gemma4_moe_layer_gated", "CompositeExplicitAutograd")(impl) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + try: + torch.library.register_fake("cactus_transpile::gemma4_moe_layer_gated")(fake) + except RuntimeError as exc: + if not _ignore_duplicate_registration(exc): + raise + _TEST_GEMMA4_MOE_OP_REGISTERED = True + + +class _FakeLfm2MoeBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.gate = torch.nn.Linear(4, 2, bias=False) + self.w1 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w3 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w2 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(4, 3)) for _ in range(2)]) + self.register_buffer("expert_bias", torch.zeros(2, dtype=torch.float32)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + return torch.ops.cactus_transpile.lfm2_moe_layer_gated( + hidden_flat, + self.gate(hidden_flat), + self.expert_bias, + list(self.w1), + list(self.w3), + list(self.w2), + 2, + 1, + True, + True, + 1.0e-6, + 1.0, + ).reshape(hidden_states.shape) + + +class _FakeQwen2MoeBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.gate = torch.nn.Linear(4, 2, bias=False) + self.w1 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w3 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w2 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(4, 3)) for _ in range(2)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + return torch.ops.cactus_transpile.qwen2_moe_layer_gated( + hidden_flat, + self.gate(hidden_flat), + list(self.w1), + list(self.w3), + list(self.w2), + 2, + 1, + False, + 1.0e-6, + 1.0, + ).reshape(hidden_states.shape) + + +class _FakeGemma4MoeBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.router = torch.nn.Linear(4, 2, bias=False) + self.w1 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w3 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(3, 4)) for _ in range(2)]) + self.w2 = torch.nn.ParameterList([torch.nn.Parameter(torch.randn(4, 3)) for _ in range(2)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + return torch.ops.cactus_transpile.gemma4_moe_layer_gated( + hidden_flat, + self.router(hidden_flat), + list(self.w1), + list(self.w3), + list(self.w2), + 2, + 1, + True, + 1.0e-6, + 1.0, + ).reshape(hidden_states.shape) + + +def test_lfm2_moe_runtime_block_exports_single_semantic_op() -> None: + _ensure_test_lfm2_moe_op_registered() + block = _FakeLfm2MoeBlock().eval() + captured = capture_model(block, (torch.randn(1, 2, 4),), strict=False) + + moe_nodes = [ + captured.ir_graph.nodes[node_id] + for node_id in captured.ir_graph.order + if captured.ir_graph.nodes[node_id].op == "lfm2_moe_layer_gated" + ] + + assert len(moe_nodes) == 1 + assert moe_nodes[0].attrs["num_experts"] == 2 + assert moe_nodes[0].attrs["num_experts_per_tok"] == 1 + assert moe_nodes[0].attrs["use_expert_bias"] is True + assert len(moe_nodes[0].inputs) == 3 + 3 * 2 + + +def test_qwen2_moe_runtime_block_exports_single_semantic_op() -> None: + _ensure_test_qwen2_moe_op_registered() + block = _FakeQwen2MoeBlock().eval() + captured = capture_model(block, (torch.randn(1, 2, 4),), strict=False) + + moe_nodes = [ + captured.ir_graph.nodes[node_id] + for node_id in captured.ir_graph.order + if captured.ir_graph.nodes[node_id].op == "qwen2_moe_layer_gated" + ] + + assert len(moe_nodes) == 1 + assert moe_nodes[0].attrs["num_experts"] == 2 + assert moe_nodes[0].attrs["num_experts_per_tok"] == 1 + assert moe_nodes[0].attrs["normalize_routing"] is False + assert len(moe_nodes[0].inputs) == 2 + 3 * 2 + + +def test_gemma4_moe_runtime_block_exports_single_semantic_op() -> None: + _ensure_test_gemma4_moe_op_registered() + block = _FakeGemma4MoeBlock().eval() + captured = capture_model(block, (torch.randn(1, 2, 4),), strict=False) + + moe_nodes = [ + captured.ir_graph.nodes[node_id] + for node_id in captured.ir_graph.order + if captured.ir_graph.nodes[node_id].op == "gemma4_moe_layer_gated" + ] + + assert len(moe_nodes) == 1 + assert moe_nodes[0].attrs["num_experts"] == 2 + assert moe_nodes[0].attrs["num_experts_per_tok"] == 1 + assert moe_nodes[0].attrs["normalize_routing"] is True + assert len(moe_nodes[0].inputs) == 2 + 3 * 2 From e28cc81d85aabb05e3263185a35ff40775b7a0d9 Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Tue, 23 Jun 2026 23:36:59 -0700 Subject: [PATCH 2/3] made lfm moe and gemma4 moe work --- cactus-engine/src/model.cpp | 2 - cactus-engine/tests/run.cpp | 54 ++- cactus-graph/cactus_graph.h | 2 +- cactus-graph/src/builder.cpp | 11 +- cactus-graph/src/graph_ffi.cpp | 4 +- cactus-graph/src/ops_nn.cpp | 49 +- python/cactus/bindings/cactus.py | 6 +- python/cactus/cli/__init__.py | 5 + python/cactus/cli/convert.py | 20 +- python/cactus/cli/model.py | 18 +- python/cactus/convert/cli.py | 7 +- .../cactus/convert/model_adapters/adapters.py | 40 +- .../convert/tests/test_gemma4_adapter.py | 47 ++ python/cactus/transpile/aten_ops.py | 1 + python/cactus/transpile/hf_model.py | 297 +++++++++--- python/cactus/transpile/importers.py | 196 +++----- python/cactus/transpile/lower.py | 186 ++++--- python/cactus/transpile/model_adapters.py | 459 +++++++++--------- python/cactus/transpile/optimize_graph.py | 5 +- python/cactus/transpile/weight_binding.py | 50 +- python/cactus/transpile/weight_compat.py | 21 + python/tests/test_cli_transpile_defaults.py | 13 + python/tests/test_graph.py | 27 ++ python/tests/test_transpile_aten_ops.py | 188 +------ python/tests/test_transpile_importers.py | 32 ++ python/tests/test_transpile_weight_compat.py | 25 + 26 files changed, 1033 insertions(+), 732 deletions(-) create mode 100644 python/tests/test_transpile_importers.py diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 0a8f8b490..3d551f766 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -999,7 +999,6 @@ void Model::move_cache_states(Component& source, Component& target, size_t logic logical_current != std::numeric_limits::max()) { auto* meta = static_cast(target.graph->get_output(static_cast(dst_node))); if (meta && logical_current < meta[0]) { - meta[4] = std::min(meta[4], logical_current); meta[0] = logical_current; } } @@ -1014,7 +1013,6 @@ void Model::set_cache_current_len(Component& comp, size_t len) { if (comp.graph->get_node_op_type(static_cast(node_id)) != OpType::KV_CACHE_STATE) continue; auto* meta = static_cast(comp.graph->get_output(static_cast(node_id))); if (!meta || len >= meta[0]) continue; - meta[4] = std::min(meta[4], len); meta[0] = len; } } diff --git a/cactus-engine/tests/run.cpp b/cactus-engine/tests/run.cpp index d28aa57f6..a36718c3d 100644 --- a/cactus-engine/tests/run.cpp +++ b/cactus-engine/tests/run.cpp @@ -247,6 +247,9 @@ struct TokenPrinter { std::string pending; bool in_code_fence = false; bool color = false; + bool suppress_thinking_stream = false; // inside a <|channel>… reasoning block + bool show_thinking = false; // --thinking: render reasoning dimmed instead of hiding it + std::string stream_carry; // holds a possibly-partial control marker across chunks bool label_printed = false; std::thread spinner_thread; std::atomic spinner_running{false}; @@ -261,10 +264,12 @@ struct TokenPrinter { in_code_fence = false; label_printed = false; color = stdout_is_terminal(); + suppress_thinking_stream = false; + stream_carry.clear(); } void start_thinking() { - if (!color) return; + if (!color || spinner_running.load()) return; spinner_running = true; spinner_thread = std::thread([this]() { const char* frames[] = {"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}; @@ -434,25 +439,63 @@ struct TokenPrinter { } void on_token(const char* text) { - print_label(); if (!saw_first) { first = std::chrono::steady_clock::now(); saw_first = true; } ++count; if (!text) return; - pending += text; + feed(text); + std::cout << std::flush; + } + + // Strip Gemma4 turn/channel control markers from the live stream and hide the + // <|channel>… reasoning block (rendered inline only under --thinking). + void feed(const std::string& chunk) { + static const std::vector markers = { + "<|channel>", "", "<|turn>", "", "<|think|>", "", + }; + stream_carry += chunk; + size_t i = 0; + while (i < stream_carry.size()) { + const std::string* matched = nullptr; + bool maybe_partial = false; + for (const auto& m : markers) { + if (stream_carry.compare(i, m.size(), m) == 0) { matched = &m; break; } + size_t avail = stream_carry.size() - i; + if (avail < m.size() && stream_carry.compare(i, avail, m, 0, avail) == 0) maybe_partial = true; + } + if (matched) { + if (*matched == "<|channel>") { suppress_thinking_stream = true; if (!show_thinking) start_thinking(); } + else if (*matched == "") { suppress_thinking_stream = false; stop_thinking(); } + i += matched->size(); + continue; + } + if (maybe_partial) break; // a marker may complete on the next chunk + if (!suppress_thinking_stream || show_thinking) emit_visible(stream_carry[i]); + ++i; + } + stream_carry.erase(0, i); + } + + void emit_visible(char c) { + print_label(); + pending += c; size_t nl; while ((nl = pending.find('\n')) != std::string::npos) { render_line(pending.substr(0, nl)); std::cout << "\n"; pending.erase(0, nl + 1); } - std::cout << std::flush; } // Flush the trailing partial line once generation completes (no trailing newline). void finish() { + stop_thinking(); + if (!stream_carry.empty() && (!suppress_thinking_stream || show_thinking)) { + for (char c : stream_carry) emit_visible(c); + } + stream_carry.clear(); print_label(); if (!pending.empty()) { render_line(pending); @@ -891,6 +934,7 @@ int main(int argc, char** argv) { std::vector current_pcm; TokenPrinter printer; g_printer = &printer; + printer.show_thinking = thinking; bool auto_send = !initial_prompt.empty() || !current_audio.empty() || !current_image.empty(); while (true) { @@ -988,7 +1032,7 @@ int main(int argc, char** argv) { + ",\"auto_handoff\":" + (auto_handoff ? "true" : "false") + ",\"confidence_threshold\":" + std::to_string(confidence_threshold) + ",\"cloud_timeout_ms\":" + std::to_string(cloud_timeout_ms) - + ",\"stop_sequences\":[\"<|im_end|>\",\"\"]}"; + + ",\"stop_sequences\":[\"<|im_end|>\",\"\",\"\"]}"; if (!current_image.empty()) std::cout << "[image: " << current_image << "]\n"; if (!current_audio.empty()) std::cout << "[audio: " << current_audio << "]\n"; diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 68a3ae46d..d8006e6e8 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -1050,7 +1050,7 @@ CACTUS_FFI_EXPORT int cactus_graph_gaussian_topk( CACTUS_FFI_EXPORT int cactus_graph_moe_layer_gated( cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices, const cactus_node_t* w1_weights, const cactus_node_t* w3_weights, const cactus_node_t* w2_weights, - size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, cactus_node_t* out); + size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_moe_layer_ungated( cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices, const cactus_node_t* w1_weights, const cactus_node_t* w2_weights, diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index 35b9b1d3f..b132db45f 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -1220,7 +1220,12 @@ size_t CactusGraph::add_node(OpType op_type, const std::vector& inputs, } Precision result_precision = params.output_precision; - if (op_type == OpType::PRECISION_CAST || op_type == OpType::EMBEDDING) { + if (op_type == OpType::PRECISION_CAST || + op_type == OpType::EMBEDDING || + op_type == OpType::TOPK || + op_type == OpType::SCATTER_TOPK || + op_type == OpType::SAMPLE || + op_type == OpType::GAUSSIAN_TOPK) { result_precision = params.output_precision; } else if (!inputs.empty()) { result_precision = nodes_[node_index_map_[inputs[0]]]->output_buffer.precision; @@ -1301,8 +1306,12 @@ size_t CactusGraph::altup_correct(size_t coefs, size_t innovation, const size_t* size_t CactusGraph::gaussian_topk(size_t input, float ppf) { const auto& in_buf = get_output_buffer(input); + if (in_buf.precision != Precision::FP16) { + throw std::runtime_error("gaussian_topk only supports FP16 input"); + } OpParams params; params.scalar = ppf; + params.output_precision = in_buf.precision; return add_node(OpType::GAUSSIAN_TOPK, {input}, in_buf.shape, params); } diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index 02861ab63..3509d77ca 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/cactus-graph/src/graph_ffi.cpp @@ -1315,7 +1315,7 @@ int cactus_graph_gaussian_topk(cactus_graph_t graph, cactus_node_t input, float int cactus_graph_moe_layer_gated(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices, const cactus_node_t* w1_weights, const cactus_node_t* w3_weights, const cactus_node_t* w2_weights, - size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, cactus_node_t* out) { + size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out) { if (!graph || !w1_weights || !w3_weights || !w2_weights || !out) return fail_invalid("Invalid args to cactus_graph_moe_layer_gated"); try { std::vector w1(num_experts), w3(num_experts), w2(num_experts); @@ -1324,7 +1324,7 @@ int cactus_graph_moe_layer_gated(cactus_graph_t graph, cactus_node_t hidden, cac w3[i] = static_cast(w3_weights[i]); w2[i] = static_cast(w2_weights[i]); } - *out = static_cast(as_graph(graph)->graph.moe_layer(static_cast(hidden), static_cast(routing_probs), static_cast(topk_indices), w1, w3, w2, num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor)); + *out = static_cast(as_graph(graph)->graph.moe_layer(static_cast(hidden), static_cast(routing_probs), static_cast(topk_indices), w1, w3, w2, num_experts, num_experts_per_tok, normalize_routing, epsilon, routed_scaling_factor, static_cast(activation))); return 0; } catch (const std::exception& e) { last_error_message = e.what(); diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index 23afc2f20..d2cd5e2cc 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -83,6 +83,7 @@ void compute_matmul_node(GraphNode& node, const std::vector moe_compact_hidden_buf; thread_local std::vector<__fp16> moe_gate_buf; + thread_local std::vector<__fp16> moe_gate_pad_buf; thread_local std::vector<__fp16> moe_up_buf; thread_local std::vector<__fp16> moe_expert_out_buf; thread_local std::vector moe_expert_offsets_buf; @@ -185,18 +186,21 @@ void compute_moe_layer_node(GraphNode& node, const std::vector size_t { + if (!std::isfinite(raw_idx) || raw_idx < 0.0f) { + throw std::runtime_error("moe_layer got invalid expert index"); + } + size_t idx = static_cast(raw_idx + 0.5f); + if (idx >= num_experts) { + throw std::runtime_error("moe_layer got expert index out of range"); + } + return idx; + }; + std::memset(expert_offsets, 0, (num_experts + 1) * sizeof(size_t)); for (size_t tok = 0; tok < token_count; ++tok) { for (size_t k = 0; k < top_k; ++k) { - float raw_idx = topk_idx[tok * top_k + k]; - if (!std::isfinite(raw_idx)) { - throw std::runtime_error("moe_layer got non-finite expert index"); - } - size_t idx = static_cast(raw_idx + 0.5f); - if (idx >= num_experts) { - throw std::runtime_error("moe_layer got expert index out of range"); - } - expert_offsets[idx + 1]++; + expert_offsets[expert_index(topk_idx[tok * top_k + k]) + 1]++; } } @@ -210,7 +214,7 @@ void compute_moe_layer_node(GraphNode& node, const std::vector(topk_idx[tok * top_k + k] + 0.5f); + size_t idx = expert_index(topk_idx[tok * top_k + k]); expert_tokens_flat[moe_write_cursors[idx]++] = tok; } } @@ -220,7 +224,7 @@ void compute_moe_layer_node(GraphNode& node, const std::vector(topk_idx[tok * top_k + k] + 0.5f); + size_t idx = expert_index(topk_idx[tok * top_k + k]); sum_probs += routing_prob(tok, idx); } routing_denom[tok] = sum_probs + eps; @@ -277,7 +281,25 @@ void compute_moe_layer_node(GraphNode& node, const std::vector 768): the quant matmul strides + // activations by the weight's K, so unpadded rows misalign when selected_count > 1. + const size_t w2_k = w2_buffer.shape.size() == 2 ? w2_buffer.shape[1] : 0; + if (w2_k < expert_intermediate_dim) { + throw std::runtime_error("moe_layer down-proj weight K smaller than expert intermediate dim"); + } + const __fp16* w2_input = gate; + if (w2_k != expert_intermediate_dim) { + if (moe_gate_pad_buf.size() < selected_count * w2_k) moe_gate_pad_buf.resize(selected_count * w2_k); + std::memset(moe_gate_pad_buf.data(), 0, selected_count * w2_k * sizeof(__fp16)); + for (size_t i = 0; i < selected_count; ++i) { + std::memcpy(moe_gate_pad_buf.data() + i * w2_k, + gate + i * expert_intermediate_dim, + expert_intermediate_dim * sizeof(__fp16)); + } + w2_input = moe_gate_pad_buf.data(); + } + + moe_matmul(w2_input, selected_count, w2_k, w2_buffer, expert_out, hidden_dim); for (size_t i = 0; i < selected_count; ++i) { const size_t tok = selected_tokens[i]; @@ -928,6 +950,3 @@ void compute_groupnorm_node(GraphNode& node, const std::vector