From a43dae842bd355792febaf667a58d1d54e88a2ae Mon Sep 17 00:00:00 2001 From: Canonik Date: Wed, 1 Jul 2026 22:44:54 +0200 Subject: [PATCH 1/3] Added minimal TransformerBridge support skeleton for Qwen2MoeForCausalLM --- .../model_bridge/test_qwen2_moe_bridge.py | 101 ++++++++++++ .../test_qwen2_moe_adapter.py | 151 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 1 + .../model_bridge/sources/transformers.py | 6 +- .../supported_architectures/__init__.py | 4 + .../supported_architectures/qwen2_moe.py | 67 ++++++++ .../tools/model_registry/__init__.py | 2 + .../model_registry/data/supported_models.json | 18 ++- .../tools/model_registry/generate_report.py | 1 + 10 files changed, 350 insertions(+), 3 deletions(-) create mode 100644 tests/integration/model_bridge/test_qwen2_moe_bridge.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/qwen2_moe.py diff --git a/tests/integration/model_bridge/test_qwen2_moe_bridge.py b/tests/integration/model_bridge/test_qwen2_moe_bridge.py new file mode 100644 index 000000000..318c357f6 --- /dev/null +++ b/tests/integration/model_bridge/test_qwen2_moe_bridge.py @@ -0,0 +1,101 @@ +"""Integration tests for the Qwen2-MoE TransformerBridge.""" + +import copy + +import torch +from transformers import Qwen2MoeConfig, Qwen2MoeForCausalLM + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + GatedMLPBridge, + MoEBridge, + PositionEmbeddingsAttentionBridge, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_from_module, +) + + +def tiny_qwen2_moe_config() -> Qwen2MoeConfig: + return Qwen2MoeConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=96, + moe_intermediate_size=32, + shared_expert_intermediate_size=96, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + decoder_sparse_step=1, + mlp_only_layers=[], + ) + + +def tiny_qwen2_moe_bridge() -> tuple[TransformerBridge, Qwen2MoeForCausalLM]: + torch.manual_seed(0) + cfg = tiny_qwen2_moe_config() + cfg._attn_implementation = "eager" + hf_reference = Qwen2MoeForCausalLM(cfg).eval() + hf_reference.config._attn_implementation = "eager" + for layer in hf_reference.model.layers: + layer.self_attn.config._attn_implementation = "eager" + + hf_model = Qwen2MoeForCausalLM(copy.deepcopy(cfg)).eval() + hf_model.load_state_dict(hf_reference.state_dict()) + bridge = build_bridge_from_module( + hf_model, + "Qwen2MoeForCausalLM", + hf_config=copy.deepcopy(cfg), + tokenizer=None, + device="cpu", + ).eval() + bridge.adapter.setup_component_testing(hf_model, bridge_model=bridge) + return bridge, hf_reference + + +def _tokens() -> torch.Tensor: + return torch.tensor([[1, 2, 3, 4, 5]]) + + +class TestQwen2MoeBridge: + def test_bridge_structure(self) -> None: + bridge, _ = tiny_qwen2_moe_bridge() + + assert len(bridge.blocks) == 2 + assert isinstance(bridge.blocks[0].attn, PositionEmbeddingsAttentionBridge) + assert isinstance(bridge.blocks[0].mlp, MoEBridge) + assert isinstance(bridge.blocks[0].mlp.shared_expert, GatedMLPBridge) + assert hasattr(bridge.blocks[0].mlp, "shared_expert_gate") + + def test_forward_matches_hf(self) -> None: + bridge, hf_reference = tiny_qwen2_moe_bridge() + tokens = _tokens() + + with torch.no_grad(): + bridge_out = bridge(tokens) + hf_out = hf_reference(tokens).logits + + assert bridge_out.shape == (1, 5, 128) + assert not torch.isnan(bridge_out).any() + assert not torch.isinf(bridge_out).any() + max_diff = (bridge_out - hf_out).abs().max().item() + assert max_diff < 1e-5, f"Bridge vs HF max diff = {max_diff}" + + def test_run_with_cache_captures_moe_hooks(self) -> None: + bridge, _ = tiny_qwen2_moe_bridge() + + _, cache = bridge.run_with_cache(_tokens()) + + for layer_idx in range(len(bridge.blocks)): + for hook_name in ( + "gate.hook_out", + "experts.hook_out", + "shared_expert.hook_out", + "shared_expert_gate.hook_out", + "hook_out", + ): + key = f"blocks.{layer_idx}.mlp.{hook_name}" + assert key in cache, f"Missing cache key: {key}" diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py new file mode 100644 index 000000000..f405b2fd1 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py @@ -0,0 +1,151 @@ +"""Unit tests for the Qwen2MoeArchitectureAdapter.""" + +import pytest +import torch +from transformers import Qwen2MoeConfig + +from transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion import ( + RearrangeTensorConversion, +) +from transformer_lens.conversion_utils.param_processing_conversion import ( + ParamProcessingConversion, +) +from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, +) +from transformer_lens.model_bridge.generalized_components import ( + GatedMLPBridge, + LinearBridge, + MoEBridge, + PositionEmbeddingsAttentionBridge, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) +from transformer_lens.model_bridge.sources.transformers import ( + determine_architecture_from_hf_config, +) +from transformer_lens.model_bridge.supported_architectures.qwen2_moe import ( + Qwen2MoeArchitectureAdapter, +) + + +def _tiny_config() -> Qwen2MoeConfig: + return Qwen2MoeConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=96, + moe_intermediate_size=32, + shared_expert_intermediate_size=96, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + decoder_sparse_step=1, + mlp_only_layers=[], + ) + + +@pytest.fixture +def bridge_cfg(): + return build_bridge_config_from_hf( + _tiny_config(), + "Qwen2MoeForCausalLM", + "qwen2-moe-test", + torch.float32, + ) + + +@pytest.fixture +def adapter(bridge_cfg) -> Qwen2MoeArchitectureAdapter: + return Qwen2MoeArchitectureAdapter(bridge_cfg) + + +class TestQwen2MoeDetection: + def test_model_type_detects_qwen2_moe(self) -> None: + assert determine_architecture_from_hf_config(_tiny_config()) == "Qwen2MoeForCausalLM" + + def test_factory_selects_adapter(self, bridge_cfg) -> None: + selected = ArchitectureAdapterFactory.select_architecture_adapter(bridge_cfg) + assert isinstance(selected, Qwen2MoeArchitectureAdapter) + + +class TestQwen2MoeConfigMapping: + def test_core_fields_map_from_hf_config(self, bridge_cfg) -> None: + hf_cfg = _tiny_config() + assert bridge_cfg.d_vocab == hf_cfg.vocab_size + assert bridge_cfg.d_model == hf_cfg.hidden_size + assert bridge_cfg.n_layers == hf_cfg.num_hidden_layers + assert bridge_cfg.n_heads == hf_cfg.num_attention_heads + assert bridge_cfg.n_key_value_heads == hf_cfg.num_key_value_heads + assert bridge_cfg.d_mlp == hf_cfg.intermediate_size + + def test_moe_fields_map_from_hf_config(self, bridge_cfg) -> None: + hf_cfg = _tiny_config() + assert bridge_cfg.num_experts == hf_cfg.num_experts + assert bridge_cfg.experts_per_token == hf_cfg.num_experts_per_tok + assert getattr(bridge_cfg, "moe_intermediate_size") == hf_cfg.moe_intermediate_size + assert ( + getattr(bridge_cfg, "shared_expert_intermediate_size") + == hf_cfg.shared_expert_intermediate_size + ) + + +class TestQwen2MoeComponentMapping: + def test_reuses_qwen2_attention_mapping(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + attn = blocks.submodules["attn"] + assert isinstance(attn, PositionEmbeddingsAttentionBridge) + assert set(attn.submodules.keys()) == {"q", "k", "v", "o"} + assert attn.submodules["q"].name == "q_proj" + assert attn.submodules["k"].name == "k_proj" + assert attn.submodules["v"].name == "v_proj" + assert attn.submodules["o"].name == "o_proj" + + def test_mlp_is_moe_bridge(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert isinstance(mlp, MoEBridge) + assert mlp.name == "mlp" + + def test_moe_submodules(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert set(mlp.submodules.keys()) == { + "gate", + "experts", + "shared_expert", + "shared_expert_gate", + } + assert isinstance(mlp.submodules["gate"], LinearBridge) + assert isinstance(mlp.submodules["experts"], MoEBridge) + assert isinstance(mlp.submodules["shared_expert"], GatedMLPBridge) + assert isinstance(mlp.submodules["shared_expert_gate"], LinearBridge) + + def test_moe_hf_paths(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert mlp.submodules["gate"].name == "gate" + assert mlp.submodules["experts"].name == "experts" + assert mlp.submodules["shared_expert"].name == "shared_expert" + assert mlp.submodules["shared_expert_gate"].name == "shared_expert_gate" + + def test_shared_expert_submodules(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + shared_expert = mlp.submodules["shared_expert"] + assert set(shared_expert.submodules.keys()) == {"gate", "in", "out"} + assert shared_expert.submodules["gate"].name == "gate_proj" + assert shared_expert.submodules["in"].name == "up_proj" + assert shared_expert.submodules["out"].name == "down_proj" + + +class TestQwen2MoeWeightConversions: + def test_kv_rearrange_uses_n_key_value_heads( + self, adapter: Qwen2MoeArchitectureAdapter + ) -> None: + conversions = adapter.weight_processing_conversions + assert conversions is not None + for slot in ("k", "v"): + conv = conversions[f"blocks.{{i}}.attn.{slot}.weight"] + assert isinstance(conv, ParamProcessingConversion) + assert isinstance(conv.tensor_conversion, RearrangeTensorConversion) + assert conv.tensor_conversion.axes_lengths["n"] == 2 diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index c7aa2acf2..cf18ffcdd 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -67,6 +67,7 @@ PhiArchitectureAdapter, PhiMoEArchitectureAdapter, Qwen2ArchitectureAdapter, + Qwen2MoeArchitectureAdapter, Qwen3_5ArchitectureAdapter, Qwen3_5MultimodalArchitectureAdapter, Qwen3ArchitectureAdapter, @@ -144,6 +145,7 @@ "PhiMoEForCausalLM": PhiMoEArchitectureAdapter, "QwenForCausalLM": QwenArchitectureAdapter, "Qwen2ForCausalLM": Qwen2ArchitectureAdapter, + "Qwen2MoeForCausalLM": Qwen2MoeArchitectureAdapter, "Qwen3ForCausalLM": Qwen3ArchitectureAdapter, "Qwen3MoeForCausalLM": Qwen3MoeArchitectureAdapter, "Qwen3NextForCausalLM": Qwen3NextArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 55f53ada8..05e9fd5bf 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -52,6 +52,7 @@ # Hybrid/MoE architectures "layer_types", "moe_intermediate_size", + "shared_expert_intermediate_size", "norm_eps", "attention_bias", "lm_head_bias", diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index bce60d96f..1e6e6c90f 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -180,7 +180,9 @@ def map_default_transformer_lens_config(hf_config): tl_config.eps = source_config.layer_norm_epsilon elif hasattr(source_config, "norm_eps"): tl_config.eps = source_config.norm_eps - if hasattr(source_config, "num_local_experts"): + if hasattr(source_config, "num_experts"): + tl_config.num_experts = source_config.num_experts + elif hasattr(source_config, "num_local_experts"): tl_config.num_experts = source_config.num_local_experts if hasattr(source_config, "num_experts_per_tok"): tl_config.experts_per_token = source_config.num_experts_per_tok @@ -248,6 +250,7 @@ def determine_architecture_from_hf_config(hf_config): "phi3": "Phi3ForCausalLM", "qwen": "QwenForCausalLM", "qwen2": "Qwen2ForCausalLM", + "qwen2_moe": "Qwen2MoeForCausalLM", "qwen3": "Qwen3ForCausalLM", # qwen3_5 is the top-level multimodal config type; qwen3_5_text is # the text-only sub-config. Both map to the text-only adapter so @@ -544,6 +547,7 @@ def boot( # Hybrid/MoE architectures "layer_types", "moe_intermediate_size", + "shared_expert_intermediate_size", "norm_eps", "attention_bias", "lm_head_bias", diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index efa25f5d6..ba5b56430 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -177,6 +177,9 @@ from transformer_lens.model_bridge.supported_architectures.qwen2 import ( Qwen2ArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.qwen2_moe import ( + Qwen2MoeArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.qwen3 import ( Qwen3ArchitectureAdapter, ) @@ -268,6 +271,7 @@ "PhiMoEArchitectureAdapter", "QwenArchitectureAdapter", "Qwen2ArchitectureAdapter", + "Qwen2MoeArchitectureAdapter", "Qwen3ArchitectureAdapter", "Qwen3MoeArchitectureAdapter", "Qwen3NextArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py b/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py new file mode 100644 index 000000000..060bae2ad --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py @@ -0,0 +1,67 @@ +"""Qwen2-MoE architecture adapter.""" + +from typing import Any + +import torch + +from transformer_lens.model_bridge.generalized_components import ( + GatedMLPBridge, + LinearBridge, + MoEBridge, +) +from transformer_lens.model_bridge.supported_architectures.qwen2 import ( + Qwen2ArchitectureAdapter, +) + + +class Qwen2MoeRouterBridge(LinearBridge): + """Bridge Qwen2-MoE router logits while preserving HF's tuple return.""" + + def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> tuple[Any, ...]: + if self.original_component is None: + raise RuntimeError( + f"Original component not set for {self.name}. Call set_original_component() first." + ) + input = self.hook_in(input) + output = self.original_component(input, *args, **kwargs) + if not isinstance(output, tuple) or len(output) == 0: + return self.hook_out(output) + router_logits = self.hook_out(output[0]) + return (router_logits,) + output[1:] + + +class Qwen2MoeArchitectureAdapter(Qwen2ArchitectureAdapter): + """Architecture adapter for Qwen2-MoE models. + + Qwen2-MoE uses the Qwen2 attention stack plus a sparse MoE MLP with an + always-on shared expert path. + """ + + def __init__(self, cfg: Any) -> None: + """Initialize the Qwen2-MoE architecture adapter.""" + super().__init__(cfg) + + self.cfg.attn_implementation = "eager" + + if self.component_mapping is None: + raise ValueError("Qwen2 component mapping was not initialized") + + blocks = self.component_mapping["blocks"] + blocks.submodules["mlp"] = MoEBridge( + name="mlp", + config=self.cfg, + submodules={ + "gate": Qwen2MoeRouterBridge(name="gate"), + "experts": MoEBridge(name="experts", config=self.cfg), + "shared_expert": GatedMLPBridge( + name="shared_expert", + config=self.cfg, + submodules={ + "gate": LinearBridge(name="gate_proj"), + "in": LinearBridge(name="up_proj"), + "out": LinearBridge(name="down_proj"), + }, + ), + "shared_expert_gate": LinearBridge(name="shared_expert_gate"), + }, + ) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index e4b2e93d6..b0be24999 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -101,6 +101,7 @@ "PhiMoEForCausalLM", "QwenForCausalLM", "Qwen2ForCausalLM", + "Qwen2MoeForCausalLM", "Qwen3ForCausalLM", "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM", @@ -173,6 +174,7 @@ "PhiMoEForCausalLM": ["microsoft"], "PhiForCausalLM": ["microsoft"], "Qwen2ForCausalLM": ["Qwen", "nvidia"], + "Qwen2MoeForCausalLM": ["Qwen"], "Qwen3ForCausalLM": ["Qwen", "nvidia"], "Qwen3MoeForCausalLM": ["Qwen"], "Qwen3NextForCausalLM": ["Qwen"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index f8b282aa1..a01926111 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -6,10 +6,24 @@ "min_downloads": 500, "scan_duration_seconds": 8.1 }, - "total_architectures": 65, - "total_models": 12899, + "total_architectures": 66, + "total_models": 12900, "total_verified": 1017, "models": [ + { + "architecture_id": "Qwen2MoeForCausalLM", + "model_id": "hyper-accel/ci-random-qwen2-moe-a3b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, { "architecture_id": "HunYuanDenseV1ForCausalLM", "model_id": "tencent/Hunyuan-0.5B-Instruct", diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index 5bfc96825..5e7f4f065 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -44,6 +44,7 @@ "GlmMoeDsaForCausalLM": "Z.ai's GLM-5 MoE model with DeepSeek Sparse Attention", "Glm4MoeForCausalLM": "Z.ai's GLM-4.5/4.6/4.7 sparse Mixture-of-Experts causal LM", "Qwen2ForCausalLM": "Alibaba's Qwen2 multilingual model", + "Qwen2MoeForCausalLM": "Alibaba's Qwen2 sparse Mixture-of-Experts model with shared experts", "Qwen3ForCausalLM": "Alibaba's Qwen3 latest generation", "Qwen3_5ForConditionalGeneration": "Alibaba's Qwen3.5 vision-language model", "BloomForCausalLM": "BigScience's BLOOM multilingual model", From 0408f0411cc162a9f8cf5c81496cb2d16b42cfe9 Mon Sep 17 00:00:00 2001 From: Canonik Date: Thu, 2 Jul 2026 08:16:02 +0200 Subject: [PATCH 2/3] Fix return type annotation in Qwen2MoeRouterBridge.forward method --- .../model_bridge/supported_architectures/qwen2_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py b/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py index 060bae2ad..f3e6191a3 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py @@ -17,7 +17,7 @@ class Qwen2MoeRouterBridge(LinearBridge): """Bridge Qwen2-MoE router logits while preserving HF's tuple return.""" - def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> tuple[Any, ...]: + def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> Any: if self.original_component is None: raise RuntimeError( f"Original component not set for {self.name}. Call set_original_component() first." From 909c4680a10bec2efe5578db1a2c03b6c821639c Mon Sep 17 00:00:00 2001 From: Canonik Date: Thu, 2 Jul 2026 09:08:26 +0200 Subject: [PATCH 3/3] Enhance tests for Qwen2-MoE bridge to validate cache capturing of MoE hooks --- .../model_bridge/test_qwen2_moe_bridge.py | 38 ++++++++++++++----- 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/tests/integration/model_bridge/test_qwen2_moe_bridge.py b/tests/integration/model_bridge/test_qwen2_moe_bridge.py index 318c357f6..05ee6a351 100644 --- a/tests/integration/model_bridge/test_qwen2_moe_bridge.py +++ b/tests/integration/model_bridge/test_qwen2_moe_bridge.py @@ -86,16 +86,34 @@ def test_forward_matches_hf(self) -> None: def test_run_with_cache_captures_moe_hooks(self) -> None: bridge, _ = tiny_qwen2_moe_bridge() + tokens = _tokens() + batch, seq_len = tokens.shape + flat_tokens = batch * seq_len + d_model = bridge.cfg.d_model + num_experts = bridge.cfg.num_experts - _, cache = bridge.run_with_cache(_tokens()) + _, cache = bridge.run_with_cache(tokens) for layer_idx in range(len(bridge.blocks)): - for hook_name in ( - "gate.hook_out", - "experts.hook_out", - "shared_expert.hook_out", - "shared_expert_gate.hook_out", - "hook_out", - ): - key = f"blocks.{layer_idx}.mlp.{hook_name}" - assert key in cache, f"Missing cache key: {key}" + gate_key = f"blocks.{layer_idx}.mlp.gate.hook_out" + assert gate_key in cache, f"Missing cache key: {gate_key}" + assert cache[gate_key].shape == (flat_tokens, num_experts) + + experts_key = f"blocks.{layer_idx}.mlp.experts.hook_out" + assert experts_key in cache, f"Missing cache key: {experts_key}" + assert cache[experts_key].shape == (flat_tokens, d_model) + + shared_expert_key = f"blocks.{layer_idx}.mlp.shared_expert.hook_out" + assert shared_expert_key in cache, f"Missing cache key: {shared_expert_key}" + assert cache[shared_expert_key].shape == (flat_tokens, d_model) + + shared_expert_gate_key = f"blocks.{layer_idx}.mlp.shared_expert_gate.hook_out" + assert shared_expert_gate_key in cache, f"Missing cache key: {shared_expert_gate_key}" + assert cache[shared_expert_gate_key].shape == (flat_tokens, 1) + + mlp_out_key = f"blocks.{layer_idx}.mlp.hook_out" + assert mlp_out_key in cache, f"Missing cache key: {mlp_out_key}" + assert cache[mlp_out_key].shape == (batch, seq_len, d_model) + + router_scores_key = f"blocks.{layer_idx}.mlp.hook_router_scores" + assert router_scores_key not in cache