Integrate TE MoE block support for DSv4

Author: dinodeep

src/maxtext/configs/types.py

@@ -4965,10 +4965,11 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("te_moe_block=True requires sparse_matmul=True.")
       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

@@ -373,8 +373,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):
@@ -924,6 +923,7 @@ def __init__(
     else:
       self.quant_einsums = None
 
+    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,
@@ -934,7 +934,7 @@ def __init__(
         quant=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,
         # tpu-inference applies the score function in the fused_moe_gmm kernel,
         # so we don't apply it here to avoid redundant computation.
         # See https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/layers/common/fused_moe_gmm.py#L58.
@@ -4398,6 +4398,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."""
@@ -4411,10 +4413,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)
@@ -4430,42 +4432,68 @@ def _te_moe_block(
         n_expert_groups=self.num_experts,
     )
 
-    output, lb_loss, total_recv_tokens = te_moe.moe(
+    if self.is_hash_routing:
+      router_info = te_moe.RoutingMapInfo(
+          routing_indices=routing_indices,
+          routing_weights=routing_weights,
+      )
+      collect_expert_counts = False
+    else:
+      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()
+
+    output, lb_loss, total_recv_tokens, expert_counts = te_moe.moe(
         inputs,
-        gate_kernel,
+        router_info,
         wi_kernel,
         wo_kernel,
         w0_bias,
         w1_bias,
         wo_bias,
-        expert_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),
         ep_axis=self._expert_parallelism_name,
         data_parallelism_axes=("fsdp",),
         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(),
+        collect_expert_counts=collect_expert_counts,
     )
     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,
@@ -4500,6 +4528,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)
@@ -4519,7 +4557,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

@@ -250,15 +250,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):
@@ -607,19 +608,27 @@ 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

@@ -637,6 +637,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

@@ -126,7 +126,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},