apple envy moe implementation

src/maxtext/common/common_types.py

@@ -114,6 +114,8 @@ class DecoderBlockType(enum.Enum):
   LLAMA4 = "llama4"
   OLMO3 = "olmo3"
   DEEPSEEK4 = "deepseek4"
+  ENVY = "envy"
+
 
 
 class VisionEncoderBlockType(enum.Enum):

src/maxtext/configs/models/envy-switch-base.yml

@@ -0,0 +1,38 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Model configuration for Envy Switch-Base (30B)
+
+decoder_block: "envy"
+mlp_activations: ["silu", "linear"]
+enable_dropout: false
+logits_via_embedding: false
+tokenizer_type: "huggingface"
+
+base_emb_dim: 1536
+base_num_decoder_layers: 12
+base_num_query_heads: 12
+base_num_kv_heads: 12
+base_mlp_dim: 6144
+base_moe_mlp_dim: 6144
+vocab_size: 32768
+normalization_layer_epsilon: 1e-05
+rope_max_timescale: 500000
+num_experts: 128
+num_experts_per_tok: 1
+capacity_factor: 2.0
+use_qk_norm: true
+interleave_moe_layer_step: 2
+inhomogeneous_layer_cycle_interval: 2
+max_target_length: 8192

src/maxtext/configs/models/envy-switch-large.yml

@@ -0,0 +1,38 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Model configuration for Envy Switch-Large (104B)
+
+decoder_block: "envy"
+mlp_activations: ["silu", "linear"]
+enable_dropout: false
+logits_via_embedding: false
+tokenizer_type: "huggingface"
+
+base_emb_dim: 2048
+base_num_decoder_layers: 24
+base_num_query_heads: 16
+base_num_kv_heads: 16
+base_mlp_dim: 8192
+base_moe_mlp_dim: 8192
+vocab_size: 32768
+normalization_layer_epsilon: 1e-05
+rope_max_timescale: 500000
+num_experts: 128
+num_experts_per_tok: 1
+capacity_factor: 2.0
+use_qk_norm: true
+interleave_moe_layer_step: 2
+inhomogeneous_layer_cycle_interval: 2
+max_target_length: 8192

src/maxtext/configs/models/envy-switch-xxl.yml

@@ -0,0 +1,38 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Model configuration for Envy Switch-XXL (520B)
+
+decoder_block: "envy"
+mlp_activations: ["silu", "linear"]
+enable_dropout: false
+logits_via_embedding: false
+tokenizer_type: "huggingface"
+
+base_emb_dim: 8192
+base_num_decoder_layers: 24
+base_num_query_heads: 64
+base_num_kv_heads: 8
+base_mlp_dim: 20480
+base_moe_mlp_dim: 20480
+vocab_size: 32768
+normalization_layer_epsilon: 1e-05
+rope_max_timescale: 500000
+num_experts: 64
+num_experts_per_tok: 1
+capacity_factor: 2.0
+use_qk_norm: true
+interleave_moe_layer_step: 2
+inhomogeneous_layer_cycle_interval: 2
+max_target_length: 8192

src/maxtext/configs/models/envy-test.yml

@@ -0,0 +1,38 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Model configuration for Envy-test
+
+decoder_block: "envy"
+mlp_activations: ["silu", "linear"]
+enable_dropout: false
+logits_via_embedding: false
+tokenizer_type: "huggingface"
+
+base_emb_dim: 8
+base_num_decoder_layers: 4
+base_num_query_heads: 4
+base_num_kv_heads: 2
+base_mlp_dim: 16
+base_moe_mlp_dim: 16
+vocab_size: 32
+normalization_layer_epsilon: 1e-05
+rope_max_timescale: 500000
+num_experts: 8
+num_experts_per_tok: 1
+capacity_factor: 2.0
+use_qk_norm: true
+interleave_moe_layer_step: 2
+inhomogeneous_layer_cycle_interval: 2
+max_target_length: 64

src/maxtext/configs/types.py

@@ -280,6 +280,10 @@ class ProfilerType(str, Enum):
     "olmo3-7b",
     "olmo3-7b-pt",
     "olmo3-32b",
+    "envy-test",
+    "envy-switch-base",
+    "envy-switch-large",
+    "envy-switch-xxl",
 ]
 
 
@@ -3470,6 +3474,12 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
             f"The number of decoder layers ({self.base_num_decoder_layers}) must be divisible by interleave moe layer step "
             f"({self.interleave_moe_layer_step})"
         )
+    if self.decoder_block == DecoderBlockType.ENVY:
+      if self.base_num_decoder_layers % self.interleave_moe_layer_step != 0:
+        raise ValueError(
+            f"The number of decoder layers ({self.base_num_decoder_layers}) must be divisible by interleave moe layer step "
+            f"({self.interleave_moe_layer_step})"
+        )
     if self.decoder_block in (
         DecoderBlockType.QWEN3_NEXT,
         DecoderBlockType.QWEN3_5,

src/maxtext/layers/decoders.py

@@ -44,6 +44,7 @@
     deepseek4,
     deepseek_batchsplit,
     deepseek_batchsplit_fp8,
+    envy,
     gemma,
     gemma2,
     gemma3,
@@ -493,6 +494,9 @@ def get_decoder_layers(self):
         return [llama4.Llama4ScannableBlockToLinen] if self.config.scan_layers else [llama4.Llama4DecoderLayerToLinen]
       case DecoderBlockType.OLMO3:
         return [olmo3.Olmo3ScannableBlockToLinen] if self.config.scan_layers else [olmo3.Olmo3DecoderLayerToLinen]
+      case DecoderBlockType.ENVY:
+        return [envy.EnvyScannableBlockToLinen] if self.config.scan_layers else [envy.EnvyDecoderLayerToLinen]
+
 
       case _:
         # Default case to handle any unknown decoder block types.
@@ -528,6 +532,7 @@ def get_scannable(normal_cls, scannable_cls):
         DecoderBlockType.DEEPSEEK: [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer],
         DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock),
         DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock),
+        DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock),
     }
 
     if cfg.decoder_block not in layer_map:
@@ -1078,6 +1083,11 @@ def __call__(
                 "nope_layer_interval": self.config.nope_layer_interval,
                 "interleave_moe_layer_step": self.config.interleave_moe_layer_step,
             }
+          if cfg.decoder_block == DecoderBlockType.ENVY:
+            layer_kwargs = {
+                "interleave_moe_layer_step": self.config.interleave_moe_layer_step,
+            }
+
           # Update broadcast_args and in_axes_tuple for vLLM RPA
           in_axes_tuple = (nn.broadcast,) * len(broadcast_args)
           current_broadcast_args = list(broadcast_args)
@@ -1203,6 +1213,11 @@ def __call__(
                   "is_nope_layer": llama4.determine_is_nope_layer(lyr, self.config.nope_layer_interval),
                   "is_moe_layer": llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step),
               }
+            if cfg.decoder_block == DecoderBlockType.ENVY:
+              layer_kwargs = {
+                  "is_moe_layer": (lyr + 1) % self.config.interleave_moe_layer_step == 0,
+              }
+
             if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, DecoderBlockType.DEEPSEEK4):
               layer_kwargs = {"layer_idx": lyr}
             if cfg.decoder_block == DecoderBlockType.DEEPSEEK4:

src/maxtext/layers/nnx_decoders.py

@@ -48,6 +48,7 @@
     deepseek4,
     deepseek_batchsplit,
     deepseek_batchsplit_fp8,
+    envy,
     gemma,
     gemma2,
     gemma3,
@@ -680,6 +681,11 @@ def _init_scanned_generic(self, decoder_block_classes, rngs):
           "nope_layer_interval": self.config.nope_layer_interval,
           "interleave_moe_layer_step": self.config.interleave_moe_layer_step,
       }
+    if config.decoder_block == DecoderBlockType.ENVY:
+      layer_kwargs = {
+          "interleave_moe_layer_step": self.config.interleave_moe_layer_step,
+      }
+
 
     if num_layers > 0:
       self.layers = self._create_scanned_layers(
@@ -726,6 +732,11 @@ def _init_sequential_generic(self, decoder_block_classes, rngs):
             "is_nope_layer": llama4.determine_is_nope_layer(lyr, self.config.nope_layer_interval),
             "is_moe_layer": llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step),
         }
+      elif config.decoder_block == DecoderBlockType.ENVY:
+        layer_kwargs = {
+            "is_moe_layer": (lyr + 1) % self.config.interleave_moe_layer_step == 0,
+        }
+
       elif config.decoder_block in {
           DecoderBlockType.QWEN3_NEXT,
           DecoderBlockType.QWEN3_5,
@@ -1053,6 +1064,7 @@ def get_deepseek():
         DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock),
         DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock),
         DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock),
+        DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock),
     }
 
     if cfg.decoder_block not in layer_map:
@@ -1212,6 +1224,7 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs):
         DecoderBlockType.SIMPLE_MLP,
         DecoderBlockType.LLAMA4,
         DecoderBlockType.OLMO3,
+        DecoderBlockType.ENVY,
     }:
       return functools.partial(
           RMSNorm,