qwen3 muon fork

src/dependencies/dockerfiles/maxtext_runner.Dockerfile

@@ -14,6 +14,9 @@ ENV MAXTEXT_REPO_ROOT=/deps
 # Set the working directory in the container
 WORKDIR /deps
 
+# Install GDN v3 Tokamax commit
+RUN pip install --no-deps --no-cache-dir --force-reinstall git+https://github.com/openxla/tokamax.git@b626dd8b54d708047788cf2ec538cba63a4e3739
+
 # Copy assets separately
 COPY ${PACKAGE_DIR}/maxtext/assets/ "${MAXTEXT_ASSETS_ROOT}"
 

src/maxtext/configs/base.yml

@@ -1271,6 +1271,10 @@ gdn_num_value_heads: 32
 gdn_chunk_size: 64
 # Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.
 use_qk_norm_in_gdn: true
+# Whether to use GDN Pallas kernel
+use_gdn_kernel: false
+# Whether to use hybrid GDN v3 Tokamax forward + Custom VJP backward
+use_hybrid_gdn: false
 # The ratio of dimension to apply ROPE on
 partial_rotary_factor: 1.0
 

src/maxtext/configs/models/qwen3-next-80b-a3b.yml

@@ -18,35 +18,41 @@
 decoder_block: "qwen3_next"
 
 # Core Architectural Parameters
-base_emb_dim: 2048
-base_num_decoder_layers: 48
-base_num_query_heads: 16
-base_num_kv_heads: 2
-head_dim: 256
-vocab_size: 151936
+base_emb_dim: 3072
+base_num_decoder_layers: 40
+base_num_query_heads: 64
+base_num_kv_heads: 8
+head_dim: 128
+vocab_size: 128008
 normalization_layer_epsilon: 1.0e-6
 
 # MoE Specific Parameters
-# Set base_mlp_dim to match base_moe_mlp_dim to pass validation for fully MoE models.
-base_mlp_dim: 512
-base_moe_mlp_dim: 512
-num_experts: 512
+# base_mlp_dim sizes the dense-prefix layer's MLP
+# base_moe_mlp_dim sizes every other (MoE) layer's routed + shared experts.
+base_mlp_dim: 8192
+base_moe_mlp_dim: 1536
+num_experts: 128
 shared_experts: 1
-num_experts_per_tok: 10
+num_experts_per_tok: 8
 norm_topk_prob: true
 
+# The first layer is a dense MLP (no MoE) and always uses full attention.
+first_num_dense_layers: 1
+
 # Qwen3-Next Specific Parameters for Linear Attention (Gated Delta Net)
-inhomogeneous_layer_cycle_interval: 4
+inhomogeneous_layer_cycle_interval: 3
 gdn_conv_kernel_dim: 4
 gdn_key_head_dim: 128
 gdn_value_head_dim: 128
-gdn_num_key_heads: 16
-gdn_num_value_heads: 32
+gdn_num_key_heads: 32
+gdn_num_value_heads: 64
 gdn_chunk_size: 64
 
 # RoPE Settings
-rope_max_timescale: 10000000
-partial_rotary_factor: 0.25
+rope_max_timescale: 10000
+partial_rotary_factor: 1.0
+
+mhc_expansion_rate: 4
 
 # General Model Settings
 enable_dropout: false

src/maxtext/configs/types.py

@@ -1052,6 +1052,14 @@ class Qwen3Next(BaseModel):
       True,
       description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.",
   )
+  use_gdn_kernel: bool = Field(
+      False,
+      description="Whether to use GDN Pallas kernel.",
+  )
+  use_hybrid_gdn: bool = Field(
+      False,
+      description="Whether to use hybrid GDN v3 Tokamax forward + Custom VJP backward.",
+  )
   partial_rotary_factor: float = Field(1.0, description="The ratio of dimension to apply ROPE on")
 
 
@@ -3769,6 +3777,10 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         DecoderBlockType.DEEPSEEK,
         DecoderBlockType.DEEPSEEK4,
         DecoderBlockType.QWEN3,
+        DecoderBlockType.QWEN3_NEXT,
+        DecoderBlockType.QWEN3_MOE,
+        DecoderBlockType.QWEN3_5,
+        DecoderBlockType.QWEN3_CUSTOM_MOE,
         DecoderBlockType.GEMMA3,
         DecoderBlockType.LLAMA2,
     ]:

src/maxtext/layers/decoders.py

@@ -1070,6 +1070,18 @@ def __call__(
               kv_caches=kv_caches,
               attention_metadata=attention_metadata,
           )
+        elif cfg.decoder_block == DecoderBlockType.QWEN3_NEXT:
+          y = self._apply_qwen3_next_scanned_blocks(
+              y,
+              decoder_segment_ids,
+              decoder_positions,
+              deterministic,
+              model_mode,
+              previous_chunk,
+              slot,
+              kv_caches=kv_caches,
+              attention_metadata=attention_metadata,
+          )
         elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4:
           y = self._apply_deepseek4_scanned_blocks(
               y,
@@ -1422,6 +1434,123 @@ def _apply_gemma3_scanned_blocks(
 
     return y
 
+  def _apply_qwen3_next_scanned_blocks(
+      self,
+      y,
+      decoder_segment_ids,
+      decoder_positions,
+      deterministic,
+      model_mode,
+      previous_chunk,
+      slot,
+      kv_caches=None,
+      attention_metadata=None,
+  ):
+    """Applies Qwen3-Next scanned decoder blocks, handling main scan and remainders."""
+
+    cfg = self.config
+    mesh = self.mesh
+
+    # Define the repeating pattern length and calculate how many full blocks to scan
+    block_pattern_len = cfg.inhomogeneous_layer_cycle_interval
+    num_full_blocks = cfg.num_decoder_layers // block_pattern_len
+    remainder_layers = cfg.num_decoder_layers % block_pattern_len
+
+    if num_full_blocks > 0:
+      ScannableBlockToLinen = qwen3.Qwen3NextScannableBlockToLinen
+      policy = self.get_remat_policy()
+
+      kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan(
+          kv_caches, num_full_blocks, block_pattern_len, stack=True
+      )
+
+      broadcast_args_spec = [
+          (decoder_segment_ids, nn.broadcast),
+          (decoder_positions, nn.broadcast),
+          (deterministic, nn.broadcast),
+          (model_mode, nn.broadcast),
+          (slot, nn.broadcast),
+          (None, nn.broadcast),  # page_state
+          (previous_chunk, nn.broadcast),
+          (None, nn.broadcast),  # bidirectional_mask
+          (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast),
+          (attention_metadata, nn.broadcast),
+      ]
+      broadcast_args = tuple(arg for arg, _ in broadcast_args_spec)
+      in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec)
+
+      # For a fully scanned block, apply it inside an nn.scan over the calculated number of full blocks
+      y, returned_kv_cache = nn.scan(
+          ScannableBlockToLinen,
+          variable_axes={
+              "params": cfg.param_scan_axis,
+              "cache": 0,
+              "intermediates": 0,
+              "aqt": 0,
+              "_overwrite_with_gradient": 0,
+          },
+          split_rngs={"params": True, "dropout": cfg.enable_dropout},
+          in_axes=in_axes_tuple,
+          length=num_full_blocks,
+          unroll=num_full_blocks,
+          metadata_params={
+              nn.PARTITION_NAME: "layers",
+              "abstract_init": False,
+          },
+      )(
+          config=cfg,
+          mesh=mesh,
+          quant=self.quant,
+          model_mode=model_mode,
+          num_of_layers=block_pattern_len,
+          remat_policy_fn=policy,
+          apply_internal_remat=True,
+          name="scanned_blocks",
+      )(
+          y, *broadcast_args
+      )
+
+      maxtext_utils.update_kv_caches_after_scan(
+          kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True
+      )
+
+    # Process any remaining layers that don't fit into a full scanned block
+    for layer_id in range(cfg.num_decoder_layers - remainder_layers, cfg.num_decoder_layers):
+      layer = qwen3.Qwen3NextDecoderLayerToLinen(
+          config=cfg,
+          mesh=mesh,
+          model_mode=model_mode,
+          quant=self.quant,
+          layer_idx=layer_id,
+      )
+      kv_cache = kv_caches[layer_id] if kv_caches is not None else None
+
+      remainder_args = (
+          decoder_segment_ids,
+          decoder_positions,
+          deterministic,
+          model_mode,
+          previous_chunk,
+          None,  # page_state
+          slot,
+          None,  # bidirectional_mask
+          kv_cache,
+          attention_metadata,
+      )
+
+      y_and_kv = layer(y, *remainder_args)
+      if isinstance(y_and_kv, tuple):
+        y = y_and_kv[0]
+        new_kv = y_and_kv[1]
+      else:
+        y = y_and_kv
+        new_kv = None
+
+      if kv_caches is not None and new_kv is not None:
+        kv_caches[layer_id] = new_kv
+
+    return y
+
   def _apply_gemma4_scanned_blocks(
       self,
       y,

src/maxtext/layers/moe.py

@@ -1554,11 +1554,7 @@ def is_batch_sharded_by_ep(input_activation):
       return input_activation.shape[0] > 1
 
     def explicitly_weight_ag(shard_exp_on_fsdp):
-      if shard_exp_on_fsdp:
-        quantization_rule = qpl.get_current_rule("gmm")
-        if quantization_rule and quantization_rule.weight_calibration_method.startswith("fixed"):
-          return True
-      return False
+      return bool(shard_exp_on_fsdp)
 
     def maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec):
       if isinstance(w0_kernel, aqt.QTensor):
@@ -1596,17 +1592,9 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids):
       # w0, w1, wo needs to be un sharded on fsdp / fsdp_transpose axis, so use
       # mlp_no_fsdp axis
       if self.config.shard_exp_on_fsdp:
-        quantization_rule = qpl.get_current_rule("gmm")
-        if quantization_rule and quantization_rule.weight_calibration_method.startswith("fixed"):
-          # special sharding when using static scaling for weights in quantization with shard_exp_on_fsdp
-          w0_pspec = self._logical_to_mesh_axes(self.wi_kernel_axes)
-          w1_pspec = self._logical_to_mesh_axes(self.wi_kernel_axes)
-          wo_pspec = self._logical_to_mesh_axes(self.wo_kernel_axes)
-        else:
-          # special sharding for dsv3 to remove overhead between gmm/AG
-          w0_pspec = self._logical_to_mesh_axes((None, None, "mlp_no_fsdp"))
-          w1_pspec = self._logical_to_mesh_axes((None, None, "mlp_no_fsdp"))
-          wo_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None))
+        w0_pspec = self._logical_to_mesh_axes(self.wi_kernel_axes)
+        w1_pspec = self._logical_to_mesh_axes(self.wi_kernel_axes)
+        wo_pspec = self._logical_to_mesh_axes(self.wo_kernel_axes)
       elif self.config.use_2d_fsdp_sharding:
         w0_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None))
         w1_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None))