bug fix after rebase

Author: bzgoogle

tpu_inference/layers/jax/moe/deepseek_v3_moe.py

@@ -19,7 +19,7 @@
     manually_quantize_qwix_activation, manually_quantize_qwix_weight)
 
 modeling_flax_utils = FlaxUtils()
-
+jax.config.update("jax_ragged_dot_use_ragged_dot_instruction", True),
 
 @dataclass
 class DeepSeekV3Router(nnx.Module):
@@ -329,8 +329,9 @@ def _unpermute(self, processed_tokens: jax.Array, sort_indices: jax.Array,
         with jax.named_scope("unpermute"):
             unsorted_tokens_tD = self._sort_activations(
                 processed_tokens, jnp.argsort(sort_indices))
+            D = unsorted_tokens_tD.shape[-1]
             reshaped_tokens_TXD = unsorted_tokens_tD.reshape(
-                -1, self.num_experts_per_tok, self.hidden_size)
+                -1, self.num_experts_per_tok, D)
         with jax.named_scope("combine_weights"):
             output_TD = jnp.einsum(
                 "TXD,TX -> TD",
@@ -394,10 +395,10 @@ def _distributed_sparse_moe_fwd(
 
         # TODO: update to 'expert' after we enable expert parallelism, currently experts are sharded along model axis
         # or we sould derive it from the model init
-        expert_shard_id = jax.lax.axis_index(self.expert_axis_name)
-        local_expert_size = self.num_local_experts // self.num_expert_parallelism
 
         if self.num_expert_parallelism > 1:
+            expert_shard_id = jax.lax.axis_index(self.expert_axis_name)
+            local_expert_size = self.num_local_experts // self.num_expert_parallelism
             if self.is_batch_sharded_by_expert:
                 # When token sharded in devices
                 # In this path, we assume the data(tokens) are fully sharded on expert, namely data_axis_name == expert_axis_name

tpu_inference/models/jax/deepseek_v3.py

@@ -363,10 +363,6 @@ def __init__(self, vllm_config: VllmConfig, num_layers, hidden_size,
             "is_verbose", None) is not None
         self.num_routed_experts = num_local_experts
         self.model_dtype = model_dtype
-<<<<<<< HEAD:tpu_inference/models/jax/deepseek_v3.py
-
-=======
->>>>>>> 641cb6d4 (Fix dtype bug of weight_loading. It not only occurs for sparsematmul, but densematmul):tpu_commons/models/jax/deepseek_v3.py
         self._transpose_map = {
             # dense mlp
             r"mlp\.down_proj": (1, 0),