Integrate TE MoE block support for DSv4-related changes

Author: dinodeep

src/maxtext/configs/types.py

@@ -5697,10 +5697,11 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           raise ValueError("te_moe_quantize_before_fsdp_all_gather=True does not support 2D FSDP weight sharding.")
       if self.te_moe_block and not self.prefuse_moe_weights:
         raise ValueError("te_moe_block=True requires prefuse_moe_weights=True.")
-      if self.te_moe_block and self.routed_bias_update_rate > 0.0:
-        raise ValueError("te_moe_block=True does not currently support routed_bias_update_rate > 0.")
-      if self.te_moe_block and self.norm_topk_prob:
-        raise ValueError("te_moe_block=True does not currently support norm_topk_prob=True.")
+      if self.te_moe_block and self.norm_topk_prob and self.routed_score_func not in ["sqrtsoftplus", "sigmoid"]:
+        raise ValueError(
+            "te_moe_block=True does not currently support norm_topk_prob=True "
+            "when routed_score_func is not 'sqrtsoftplus' or 'sigmoid'."
+        )
       if self.te_moe_block and self.use_random_routing:
         raise ValueError("te_moe_block=True does not support use_random_routing=True.")
       if self.te_moe_block and self.decoder_block == DecoderBlockType.LLAMA4:

src/maxtext/layers/moe.py

@@ -401,8 +401,7 @@ def load_balance_updates_from_counts(expert_counts, num_experts, rate):
   total_tokens = jnp.sum(expert_counts, axis=-1, keepdims=True)
   average_load = total_tokens / num_experts
   direction = jnp.sign(average_load - expert_counts)
-  output = direction * rate
-  return output
+  return direction * rate
 
 
 def finalize_deferred_bias_signal(partial_counts, config):
@@ -1022,6 +1021,7 @@ def __init__(
       self.quant_einsums = None
 
     gate_weight_dtype = ctypes.get_weight_dtype(self.config, "gate")
+    self.use_gate_bias = self.config.routed_bias and not self.is_hash_routing
     self.gate = GateLogit(
         in_features_shape=self.moe_expert_input_dim,
         out_features_shape=self.num_experts,
@@ -1032,7 +1032,7 @@ def __init__(
         quant=ctypes.get_quant(self.config, "gate", self.quant),
         kernel_init=self.kernel_init,
         kernel_axes=self.kernel_axes,
-        use_bias=self.config.routed_bias and not self.is_hash_routing,
+        use_bias=self.use_gate_bias,
         bias_dtype=jnp.float32 if self.config.float32_gate_logits else gate_weight_dtype,
         # tpu-inference applies the score function in the fused_moe_gmm kernel,
         # so we don't apply it here to avoid redundant computation.
@@ -4867,6 +4867,8 @@ def _te_moe_block(
       w0_bias: jax.Array | None,
       w1_bias: jax.Array | None,
       wo_bias: jax.Array | None,
+      routing_indices: jax.Array | None = None,
+      routing_weights: jax.Array | None = None,
       out_sharding: NamedSharding | None = None,
   ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]:
     """Run TransformerEngine's fused EP MoEBlock using MaxText-owned params."""
@@ -4881,10 +4883,10 @@ def _te_moe_block(
 
     if self.quant is None or not hasattr(self.quant, "get_moe_block_quantizer_sets"):
       raise ValueError("te_moe_block=True requires TransformerEngine quantization or te_gmm_quantization=te_no_quant.")
+    if self.is_hash_routing and (routing_indices is None or routing_weights is None):
+      raise ValueError("TE MoE hash routing requires precomputed routing indices and weights.")
 
-    expert_bias = None
-    if self.config.routed_bias:
-      expert_bias = jnp.asarray(self.gate.bias[...], jnp.float32)
+    expert_bias = jnp.asarray(self.gate.bias[...], jnp.float32) if self.use_gate_bias else None
 
     fsdp_size = self.mesh.shape.get("fsdp", 1)
     ep_size = self.mesh.shape.get(self._expert_parallelism_name, 1)
@@ -4908,44 +4910,104 @@ def _te_moe_block(
     checkpoint_kwargs = _get_te_moe_checkpoint_kwargs(te_moe.moe)
     cudnn_gmm_fusion_kwargs = _get_te_moe_cudnn_gmm_fusion_kwargs(self.config.te_moe_disable_cudnn_gmm_fusion)
 
-    output, lb_loss, total_recv_tokens = te_moe.moe(
+    # Older TE versions take the router arguments directly and do not return expert counts.
+    supports_router_info = hasattr(te_moe, "RouterComputationInfo")
+    if not supports_router_info:
+      if self.is_hash_routing:
+        raise ImportError(
+            "te_moe_block=True with hash routing requires TransformerEngine with "
+            "moe.RoutingMapInfo support. Please upgrade TransformerEngine."
+        )
+      if self.should_update_load_balance():
+        raise ImportError(
+            "te_moe_block=True with routed_bias_update_rate > 0 requires TransformerEngine with "
+            "MoE expert-count support. Please upgrade TransformerEngine."
+        )
+      router_kwargs = {
+          "gate_kernel": gate_kernel,
+          "num_experts_per_tok": self.num_experts_per_tok,
+          "expert_bias": expert_bias,
+          "score_function": self.config.routed_score_func or "softmax",
+          "use_pre_softmax": False,
+          "num_groups": None if self.config.n_routing_groups <= 0 else self.config.n_routing_groups,
+          "group_topk": None if self.config.topk_routing_group <= 0 else self.config.topk_routing_group,
+          "scaling_factor": self.config.routed_scaling_factor,
+          "aux_loss_coeff": self.config.load_balance_loss_weight,
+          "gate_kernel_axes": self.kernel_axes,
+      }
+    elif self.is_hash_routing:
+      router_kwargs = {
+          "router_info": te_moe.RoutingMapInfo(
+              routing_indices=routing_indices,
+              routing_weights=routing_weights,
+          ),
+          "collect_expert_counts": False,
+      }
+    else:
+      router_kwargs = {
+          "router_info": te_moe.RouterComputationInfo(
+              gate_kernel=gate_kernel,
+              num_experts_per_tok=self.num_experts_per_tok,
+              expert_bias=expert_bias,
+              score_function=self.config.routed_score_func or "softmax",
+              use_pre_softmax=False,
+              num_groups=None if self.config.n_routing_groups <= 0 else self.config.n_routing_groups,
+              group_topk=None if self.config.topk_routing_group <= 0 else self.config.topk_routing_group,
+              scaling_factor=self.config.routed_scaling_factor,
+              aux_loss_coeff=self.config.load_balance_loss_weight,
+              gate_kernel_axes=self.kernel_axes,
+          ),
+          "collect_expert_counts": self.should_update_load_balance(),
+      }
+
+    # Pass the weights by keyword: the older interface has gate_kernel in the second positional slot.
+    te_moe_outputs = te_moe.moe(
         inputs,
-        gate_kernel,
-        wi_kernel,
-        wo_kernel,
-        w0_bias,
-        w1_bias,
-        wo_bias,
-        expert_bias,
+        wi=wi_kernel,
+        wo=wo_kernel,
+        wi_0_bias=w0_bias,
+        wi_1_bias=w1_bias,
+        wo_bias=wo_bias,
         num_experts=self.num_experts,
-        num_experts_per_tok=self.num_experts_per_tok,
         activation_type=self.config.mlp_activations[0],
-        score_function=self.config.routed_score_func or "softmax",
-        use_pre_softmax=False,
-        num_groups=None if self.config.n_routing_groups <= 0 else self.config.n_routing_groups,
-        group_topk=None if self.config.topk_routing_group <= 0 else self.config.topk_routing_group,
-        scaling_factor=self.config.routed_scaling_factor,
-        aux_loss_coeff=self.config.load_balance_loss_weight,
         apply_topk_weights_early=True,
         quantizer_sets=(fc1_quantizer_set, fc2_quantizer_set),
         mesh_resource=MeshResource(dp_resource="data", fsdp_resource="fsdp", ep_resource=self._expert_parallelism_name),
         input_axes=("activation_batch", "activation_norm_length", None),
-        gate_kernel_axes=self.kernel_axes,
         wi_kernel_axes=self.wi_kernel_axes,
         wo_kernel_axes=self.wo_kernel_axes,
         dtype=self.dtype,
         recv_capacity_per_rank=max_utils.get_te_moe_recv_capacity_per_rank(),
         quant_before_fsdp_ag=self.config.te_moe_quantize_before_fsdp_all_gather,
+        **router_kwargs,
         **cudnn_gmm_fusion_kwargs,
         **checkpoint_kwargs,
     )
+    if supports_router_info:
+      output, lb_loss, total_recv_tokens, expert_counts = te_moe_outputs
+    else:
+      output, lb_loss, total_recv_tokens = te_moe_outputs
+      expert_counts = None
     recv_capacity_per_rank = max_utils.get_te_moe_recv_capacity_per_rank()
     output = output.astype(self.dtype)
     if lb_loss is not None:
       lb_loss = lb_loss.astype(self.dtype)
     if out_sharding is not None:
       output = jax.lax.with_sharding_constraint(output, out_sharding)
-    return output, lb_loss, (total_recv_tokens, jnp.asarray(recv_capacity_per_rank, dtype=jnp.int32))
+    bias_updates = None
+    if self.should_update_load_balance():
+      bias_updates = load_balance_updates_from_counts(
+          expert_counts, self.num_experts, self.config.routed_bias_update_rate
+      )
+    return (
+        output,
+        lb_loss,
+        (
+            total_recv_tokens,
+            jnp.asarray(recv_capacity_per_rank, dtype=jnp.int32),
+            bias_updates,
+        ),
+    )
 
   def __call__(
       self,
@@ -4980,6 +5042,16 @@ def __call__(
     routing_inputs = inputs if gate_inputs is None else gate_inputs.astype(gate_dtype)
 
     if cfg.te_moe_block:
+      routing_indices = None
+      routing_weights = None
+      if self.is_hash_routing:
+        gate_logits, pre_bias_logits = self.gate(routing_inputs)
+        routing_weights, routing_indices = self.get_topk(
+            gate_logits,
+            pre_bias_logits,
+            self.rngs,
+            input_ids=input_ids,
+        )
       gate_kernel = jnp.asarray(self.gate.kernel[...], self.dtype)
       wi_kernel = jnp.asarray(self.wi[...], self.dtype)
       wo_kernel = jnp.asarray(self.wo[...], self.dtype)
@@ -4999,7 +5071,9 @@ def __call__(
           w0_bias,
           w1_bias,
           wo_bias,
-          out_sharding,
+          routing_indices=routing_indices,
+          routing_weights=routing_weights,
+          out_sharding=out_sharding,
       )
 
     gate_logits, pre_bias_logits = self.gate(routing_inputs)

src/maxtext/models/deepseek.py

@@ -249,15 +249,16 @@ def post_process(self, layer_output, load_balance_loss, moe_bias_updates, kv_cac
       self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss)
 
     if self.config.te_moe_block and moe_bias_updates is not None:
-      total_recv_tokens, recv_capacity_per_rank = moe_bias_updates
+      total_recv_tokens, recv_capacity_per_rank, routed_bias_updates = moe_bias_updates
       self.sow(nnx.Intermediate, "te_moe_total_recv_tokens", total_recv_tokens)
       self.sow(nnx.Intermediate, "te_moe_recv_capacity_per_rank", recv_capacity_per_rank)
       self.sow(
           nnx.Intermediate,
           "te_moe_capacity_overflow",
           jnp.any(total_recv_tokens > recv_capacity_per_rank),
       )
-    elif self.config.routed_bias and self.config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None:
+      moe_bias_updates = routed_bias_updates
+    if self.config.routed_bias and self.config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None:
       self.sow(nnx.Intermediate, "moe_bias_updates", moe_bias_updates)
 
     if getattr(self.config, "record_internal_nn_metrics", False):
@@ -600,18 +601,26 @@ def extract_fn(x):
           self.DeepSeekMoeBlock_0,
           x=intermediate_inputs,
           mhc_type=HyperConnectionType.MLP_MOE,
+          input_ids=decoder_input_tokens,
       )
       load_balance_loss = metadata["load_balance_loss"]
       moe_bias_updates = metadata["moe_bias_updates"]
     else:
-      mlp_lnx, load_balance_loss, moe_bias_updates = self.mlp_op(hidden_states, deterministic)
+      mlp_lnx, load_balance_loss, moe_bias_updates = self.mlp_op(
+          hidden_states,
+          deterministic,
+          input_ids=decoder_input_tokens,
+      )
       layer_output = mlp_lnx + intermediate_inputs
     layer_output = self.dropout_op(layer_output, deterministic=deterministic)
 
     return self.post_process(layer_output, load_balance_loss, moe_bias_updates, kv_cache)
 
   def mlp_op(self, x, deterministic, *args, **kwargs):
     mlp_lnx, load_balance_loss, moe_bias_updates = self.DeepSeekMoeBlock_0(
-        x, intermediate_sharding=self.mlp_intermediate_sharding, out_sharding=self.out_sharding
+        x,
+        intermediate_sharding=self.mlp_intermediate_sharding,
+        out_sharding=self.out_sharding,
+        input_ids=kwargs.get("input_ids"),
     )
     return self.with_logical_constraint(mlp_lnx), load_balance_loss, moe_bias_updates

tests/integration/train_tests.py

@@ -737,6 +737,50 @@ def test_gpu_te_moe_block(self):
         ]
     )
 
+  @pytest.mark.integration_test
+  @pytest.mark.gpu_only
+  def test_gpu_te_moe_block_deepseek4_hash_and_routed_bias(self):
+    """Runs real TE MoE kernels for DSv4 hash routing and top-k routed-bias updates."""
+    gpu_device = jax.devices("gpu")[0]
+    compute_capability = getattr(gpu_device, "compute_capability", None)
+    try:
+      if float(compute_capability) < 10.0:
+        pytest.skip("TransformerEngine MoEBlock is only supported on sm100+!")
+    except Exception:  # pylint: disable=broad-exception-caught
+      pytest.skip("TransformerEngine MoEBlock is only supported on sm100+!")
+
+    train_main(
+        TrainTests.CONFIGS["synthetic"]
+        + [
+            "model_name=deepseek4-tiny",
+            "override_model_config=True",
+            "base_num_decoder_layers=3",
+            "first_num_hash_layers=1",
+            "compress_ratios=[0,4,128]",
+            "attention=dot_product",
+            "quantization=te_no_quant",
+            "base_emb_dim=32",
+            "base_num_kv_heads=1",
+            "head_dim=32",
+            "num_experts=4",
+            "num_experts_per_tok=2",
+            "shared_experts=1",
+            "base_moe_mlp_dim=32",
+            "routed_bias=True",
+            "routed_bias_update_rate=0.125",
+            "load_balance_loss_weight=0.0",
+            "sparse_matmul=True",
+            "megablox=False",
+            "prefuse_moe_weights=True",
+            "te_moe_block=True",
+            "te_gmm_quantization=te_no_quant",
+            "hardware=gpu_multiprocess",
+            "ici_fsdp_parallelism=1",
+            "ici_expert_parallelism=-1",
+            "enable_tensorboard=False",
+        ]
+    )
+
   @pytest.mark.integration_test
   @pytest.mark.tpu_only
   def test_tpu_dropout(self):

tests/unit/configs_value_test.py

@@ -181,7 +181,11 @@ def test_te_moe_block_rejects_unsupported_options_during_config_validation(self)
         "override_logical_axis_rules": True,
     }
     invalid_configs = (
-        ({"norm_topk_prob": True}, "te_moe_block=True does not currently support norm_topk_prob=True."),
+        (
+            {"norm_topk_prob": True},
+            "te_moe_block=True does not currently support norm_topk_prob=True "
+            "when routed_score_func is not 'sqrtsoftplus' or 'sigmoid'.",
+        ),
         ({"use_random_routing": True}, "te_moe_block=True does not support use_random_routing=True."),
         (
             {"decoder_block": types.DecoderBlockType.LLAMA4},