tpu-inference 0.12.0.dev20251222__py3-none-any.whl → 0.12.0.dev20251224__py3-none-any.whl

tpu_inference/core/sched/dp_scheduler.py

@@ -18,6 +18,7 @@ from collections import defaultdict, deque
 from dataclasses import dataclass
 from enum import Enum
 from multiprocessing import Process, Queue
+from time import time
 from typing import Any, Dict, List, Optional, Tuple
 
 import cloudpickle
@@ -102,7 +103,7 @@ def _disable_cloudpickle():
 def _scheduler_worker_process(
     rank: int,
     input_queue: Queue,
-    output_queue: Queue,
+    output_queues: Dict[str, Queue],
     vllm_config: Any,
     kv_cache_config: Any,
     structured_output_manager: Any,
@@ -135,55 +136,55 @@ def _scheduler_worker_process(
                 case SchedulerCommand.ADD_REQUEST:
                     request = data
                     scheduler.add_request(request)
-                    output_queue.put(None)  # Signal completion
+                    output_queues[command.value].put(None)  # Signal completion
 
                 case SchedulerCommand.SCHEDULE:
                     output = scheduler.schedule()
-                    output_queue.put(output)
+                    output_queues[command.value].put(output)
 
                 case SchedulerCommand.FINISH_REQUESTS:
                     request_ids, finished_status = data
                     scheduler.finish_requests(request_ids, finished_status)
-                    output_queue.put(None)  # Signal completion
+                    output_queues[command.value].put(None)  # Signal completion
 
                 case SchedulerCommand.UPDATE_DRAFT_TOKEN_IDS:
                     draft_token_ids = data
                     scheduler.update_draft_token_ids(draft_token_ids)
-                    output_queue.put(None)  # Signal completion
+                    output_queues[command.value].put(None)  # Signal completion
 
                 case SchedulerCommand.UPDATE_FROM_OUTPUT:
                     scheduler_output, model_runner_output = data
                     result = scheduler.update_from_output(
                         scheduler_output, model_runner_output)
-                    output_queue.put(result)
+                    output_queues[command.value].put(result)
 
                 case SchedulerCommand.GET_GRAMMAR_BITMASK:
                     scheduler_output = data
                     result = scheduler.get_grammar_bitmask(scheduler_output)
-                    output_queue.put(result)
+                    output_queues[command.value].put(result)
 
                 case SchedulerCommand.MAKE_STATS:
                     spec_decoding_stats, kv_connector_stats = data
                     result = scheduler.make_stats(spec_decoding_stats,
                                                   kv_connector_stats)
-                    output_queue.put(result)
+                    output_queues[command.value].put(result)
 
                 case SchedulerCommand.RESET_PREFIX_CACHE:
                     result = scheduler.reset_prefix_cache()
-                    output_queue.put(result)
+                    output_queues[command.value].put(result)
 
                 case SchedulerCommand.GET_NUM_UNFINISHED_REQUESTS:
                     result = scheduler.get_num_unfinished_requests()
-                    output_queue.put(result)
+                    output_queues[command.value].put(result)
 
                 case SchedulerCommand.HAS_FINISHED_REQUESTS:
                     result = scheduler.has_finished_requests()
-                    output_queue.put(result)
+                    output_queues[command.value].put(result)
 
                 case SchedulerCommand.GET_REQUEST_COUNTS:
                     running = len(scheduler.running)
                     waiting = len(scheduler.waiting)
-                    output_queue.put((running, waiting))
+                    output_queues[command.value].put((running, waiting))
 
                 case SchedulerCommand.GET_TOKEN_COUNT:
                     # Calculate total tokens across running and waiting requests
@@ -192,30 +193,29 @@ def _scheduler_worker_process(
                         total_tokens += len(req.all_token_ids)
                     for req in scheduler.waiting:
                         total_tokens += len(req.all_token_ids)
-                    output_queue.put(total_tokens)
+                    output_queues[command.value].put(total_tokens)
 
                 case SchedulerCommand.GET_COMPUTED_BLOCKS:
                     request = data
                     blocks, cached_tokens = scheduler.kv_cache_manager.get_computed_blocks(
                         request)
-                    output_queue.put((blocks, cached_tokens))
+                    output_queues[command.value].put((blocks, cached_tokens))
 
                 case SchedulerCommand.SHUTDOWN:
                     scheduler.shutdown()
-                    output_queue.put(None)  # Signal completion
+                    output_queues[command.value].put(None)  # Signal completion
                     break
                 case _:
                     error = SchedulerWorkerError(
                         rank, f"Unknown command: {command}")
-                    output_queue.put(error)
+                    output_queues[command.value].put(error)
                     raise error
 
         except Exception as e:
             logger.error(f"Error in scheduler worker {rank}: {e}",
                          exc_info=True)
-            # Put error on output queue
             error = SchedulerWorkerError(rank, str(e))
-            output_queue.put(error)
+            output_queues[command.value].put(error)
 
 
 @dataclass
@@ -276,26 +276,29 @@ class DPScheduler(SchedulerInterface):
         # Enable cloudpickle for multiprocessing to handle local functions
         _enable_cloudpickle()
 
-        # Create worker processes with one input and one output queue each
+        # Create worker processes with separate output queues for each command type
         import multiprocessing
         ctx = multiprocessing.get_context('fork')
         self.input_queues: List[Queue] = []
-        self.output_queues: List[Queue] = []
+        self.output_queues: Dict[Tuple[int, str], Queue] = {}
         self.processes: List[Process] = []
 
         for rank in range(self.dp_size):
             input_queue = ctx.Queue()
-            output_queue = ctx.Queue()
-
             self.input_queues.append(input_queue)
-            self.output_queues.append(output_queue)
+
+            output_queues_for_rank: Dict[str, Queue] = {}
+            for cmd in SchedulerCommand:
+                output_queues_for_rank[cmd.value] = ctx.Queue()
+                self.output_queues[(
+                    rank, cmd.value)] = output_queues_for_rank[cmd.value]
 
             process = ctx.Process(
                 target=_scheduler_worker_process,
                 args=(
                     rank,
                     input_queue,
-                    output_queue,
+                    output_queues_for_rank,
                     self.vllm_config,
                     self.per_rank_kv_cache_configs[rank],
                     structured_output_manager,
@@ -323,8 +326,24 @@ class DPScheduler(SchedulerInterface):
             rank_config.num_blocks = kv_cache_config.num_blocks // self.dp_size
             self.per_rank_kv_cache_configs.append(rank_config)
 
-    def _get_result_from_queue(self, queue: Queue) -> Any:
-        result = queue.get()
+    def _get_result_from_queue(self, rank: int,
+                               command: SchedulerCommand) -> Any:
+        """Get result from the output queue for a specific rank and command type."""
+        queue_obj = self.output_queues[(rank, command.value)]
+        try:
+            start_time = time()
+            result = queue_obj.get()
+            end_time = time()
+            if end_time - start_time > 1.0:
+                logger.warning(
+                    f"Long wait time ({end_time - start_time:.2f}s) for rank {rank} "
+                    f"command {command.value} response.")
+        except EOFError as e:
+            raise RuntimeError(
+                f"Queue error for rank {rank} command {command.value}: "
+                "Worker process terminated unexpectedly. "
+                "This may indicate a crash in the scheduler worker process."
+            ) from e
         if isinstance(result, SchedulerWorkerError):
             raise result
         return result
@@ -337,7 +356,8 @@ class DPScheduler(SchedulerInterface):
 
         rank_tokens = {}
         for rank in range(self.dp_size):
-            token_count = self._get_result_from_queue(self.output_queues[rank])
+            token_count = self._get_result_from_queue(
+                rank, SchedulerCommand.GET_TOKEN_COUNT)
             rank_tokens[rank] = token_count
 
         return rank_tokens
@@ -355,7 +375,7 @@ class DPScheduler(SchedulerInterface):
         best_cache_tokens = 0
         for rank in range(self.dp_size):
             blocks, cached_tokens = self._get_result_from_queue(
-                self.output_queues[rank])
+                rank, SchedulerCommand.GET_COMPUTED_BLOCKS)
             if cached_tokens > best_cache_tokens:
                 best_cache_tokens = cached_tokens
                 best_cache_rank = rank
@@ -382,7 +402,7 @@ class DPScheduler(SchedulerInterface):
         self.assigned_dp_rank[request.request_id] = rank
 
         self.input_queues[rank].put((SchedulerCommand.ADD_REQUEST, request))
-        self._get_result_from_queue(self.output_queues[rank])
+        self._get_result_from_queue(rank, SchedulerCommand.ADD_REQUEST)
 
     @time_function
     def schedule(self) -> DPSchedulerOutput:
@@ -402,7 +422,8 @@ class DPScheduler(SchedulerInterface):
         # Collect outputs from all workers (blocking)
         rank_outputs = []
         for rank in range(self.dp_size):
-            output = self._get_result_from_queue(self.output_queues[rank])
+            output = self._get_result_from_queue(rank,
+                                                 SchedulerCommand.SCHEDULE)
             rank_outputs.append(output)
 
         # Cache scheduler outputs to use in `update_from_output`
@@ -531,7 +552,7 @@ class DPScheduler(SchedulerInterface):
                                          rank_scheduler_outputs[rank]))
         for rank in range(self.dp_size):
             grammar_output = self._get_result_from_queue(
-                self.output_queues[rank])
+                rank, SchedulerCommand.GET_GRAMMAR_BITMASK)
             if grammar_output is not None:
                 combined_structured_output_request_ids.extend(
                     grammar_output.structured_output_request_ids)
@@ -572,7 +593,7 @@ class DPScheduler(SchedulerInterface):
         combined_engine_outputs = defaultdict(list)
         for rank in range(self.dp_size):
             rank_engine_outputs = self._get_result_from_queue(
-                self.output_queues[rank])
+                rank, SchedulerCommand.UPDATE_FROM_OUTPUT)
             for client_idx, engine_output in rank_engine_outputs.items():
                 combined_engine_outputs[client_idx].append(engine_output)
 
@@ -640,7 +661,7 @@ class DPScheduler(SchedulerInterface):
         for rank, req_ids in rank_request_ids.items():
             self.input_queues[rank].put(
                 (SchedulerCommand.FINISH_REQUESTS, (req_ids, finished_status)))
-            self._get_result_from_queue(self.output_queues[rank])
+            self._get_result_from_queue(rank, SchedulerCommand.FINISH_REQUESTS)
 
     def get_num_unfinished_requests(self) -> int:
         """Get total number of unfinished requests across all DP ranks."""
@@ -650,7 +671,8 @@ class DPScheduler(SchedulerInterface):
 
         total = 0
         for rank in range(self.dp_size):
-            count = self._get_result_from_queue(self.output_queues[rank])
+            count = self._get_result_from_queue(
+                rank, SchedulerCommand.GET_NUM_UNFINISHED_REQUESTS)
             total += count
         return total
 
@@ -663,7 +685,7 @@ class DPScheduler(SchedulerInterface):
         has_finished_any = False
         for rank in range(self.dp_size):
             has_finished_any |= self._get_result_from_queue(
-                self.output_queues[rank])
+                rank, SchedulerCommand.HAS_FINISHED_REQUESTS)
         return has_finished_any
 
     def get_request_counts(self) -> Tuple[int, int]:
@@ -676,7 +698,7 @@ class DPScheduler(SchedulerInterface):
         total_waiting = 0
         for rank in range(self.dp_size):
             running, waiting = self._get_result_from_queue(
-                self.output_queues[rank])
+                rank, SchedulerCommand.GET_REQUEST_COUNTS)
             total_running += running
             total_waiting += waiting
         return total_running, total_waiting
@@ -689,7 +711,8 @@ class DPScheduler(SchedulerInterface):
 
         all_success = True
         for rank in range(self.dp_size):
-            success = self._get_result_from_queue(self.output_queues[rank])
+            success = self._get_result_from_queue(
+                rank, SchedulerCommand.RESET_PREFIX_CACHE)
             all_success &= success
         return all_success
 
@@ -715,7 +738,8 @@ class DPScheduler(SchedulerInterface):
                                                kv_connector_stats)))
 
         for rank in range(self.dp_size):
-            rank_stats = self._get_result_from_queue(self.output_queues[rank])
+            rank_stats = self._get_result_from_queue(
+                rank, SchedulerCommand.MAKE_STATS)
             if rank_stats is None:
                 continue
 
@@ -776,7 +800,8 @@ class DPScheduler(SchedulerInterface):
             self.input_queues[rank].put(
                 (SchedulerCommand.UPDATE_DRAFT_TOKEN_IDS,
                  rank_draft_token_ids))
-            self._get_result_from_queue(self.output_queues[rank])
+            self._get_result_from_queue(
+                rank, SchedulerCommand.UPDATE_DRAFT_TOKEN_IDS)
 
     def shutdown(self) -> None:
         """Shutdown all DP rank scheduler worker processes."""
@@ -786,7 +811,7 @@ class DPScheduler(SchedulerInterface):
 
         # Wait for acknowledgment (blocking)
         for rank in range(self.dp_size):
-            self._get_result_from_queue(self.output_queues[rank])
+            self._get_result_from_queue(rank, SchedulerCommand.SHUTDOWN)
 
         # Terminate and join all processes
         for process in self.processes:

tpu_inference/kernels/mla/v1/kernel.py

@@ -822,36 +822,17 @@ def _mla_ragged_paged_attention_kernel(
         return q_nope_vec, q_rope_vec
 
     def load_bkv(bkv_sem_idx, *, bkvc_mask, bkpe_mask):
-        bitwidth = 32 // kv_packing
-        repack_ty = jnp.dtype(f"uint{bitwidth}")
         bkvc_ref = (bkvc_x2_ref.bitcast(jnp.uint32).at[bkv_sem_idx].reshape(
             bkv_sz_per_kv_packing, lkv_dim))
-        bkvc_vec = bkvc_ref[...]
-        bkvc_vecs = []
-        for i in range(kv_packing):
-            masked_bkvc_vec = bkvc_vec >> (i * bitwidth)
-            bkvc_vecs.append(masked_bkvc_vec)
-        concated_bkvc_vec = jnp.concatenate(bkvc_vecs, axis=-1)
-        concated_bkvc_vec = concated_bkvc_vec.reshape(bkv_sz, lkv_dim)
-        concated_bkvc_vec = lax.select(bkvc_mask, concated_bkvc_vec,
-                                       jnp.zeros_like(concated_bkvc_vec))
-        concated_bkvc_vec = pltpu.bitcast(concated_bkvc_vec.astype(repack_ty),
-                                          kv_dtype)
+        bkvc_vec = pltpu.bitcast(bkvc_ref[...], kv_dtype)
+        bkvc_vec = lax.select(bkvc_mask, bkvc_vec, jnp.zeros_like(bkvc_vec))
+
         bkpe_ref = (bkpe_x2_ref.bitcast(jnp.uint32).at[bkv_sem_idx].reshape(
             bkv_sz_per_kv_packing, r_dim))
-        bkpe_vec = bkpe_ref[...]
-        bkpe_vecs = []
-        for i in range(kv_packing):
-            masked_bkpe_vec = bkpe_vec >> (i * bitwidth)
-            bkpe_vecs.append(masked_bkpe_vec)
-        concated_bkpe_vec = jnp.concatenate(bkpe_vecs, axis=-1)
-        concated_bkpe_vec = concated_bkpe_vec.reshape(bkv_sz, r_dim)
-        concated_bkpe_vec = lax.select(bkpe_mask, concated_bkpe_vec,
-                                       jnp.zeros_like(concated_bkpe_vec))
-        concated_bkpe_vec = pltpu.bitcast(concated_bkpe_vec.astype(repack_ty),
-                                          kv_dtype)
-
-        return concated_bkvc_vec, concated_bkpe_vec
+        bkpe_vec = pltpu.bitcast(bkpe_ref[...], kv_dtype)
+        bkpe_vec = lax.select(bkpe_mask, bkpe_vec, jnp.zeros_like(bkpe_vec))
+
+        return bkvc_vec, bkpe_vec
 
     def broadcast_minor(src, shape):
         if src.shape == shape:

tpu_inference/layers/common/sharding.py

@@ -40,7 +40,7 @@ class ShardingAxisNameBase:
     MLP_TENSOR = ('attn_dp', 'model', 'expert')
     MOE_TENSOR = ('attn_dp', 'model')
     EXPERT = ('attn_dp', 'expert', 'model')
-    VOCAB = ('expert', 'model')
+    VOCAB = ('expert', 'attn_dp', 'model')
 
 
 class ShardingAxisName2D:
@@ -141,6 +141,11 @@ class ShardingConfigManager:
             kv_dtype = utils.get_jax_dtype_from_str_dtype(
                 cache_dtype) or jnp.bfloat16
             packing = 4 // jnp.dtype(kv_dtype).itemsize
+
+            # The default head dim is 128 but 64 is also supported as a special case.
+            if vllm_config.model_config.get_head_size() == 64:
+                packing *= 2
+
             # When num_kv_heads * 2 / packing < TP, tensor parallelism would
             # duplicate KV heads across devices, wasting kv cache memory.
             # Use attention DP instead to reduce per-device num_kv_heads and
@@ -186,8 +191,8 @@ class ShardingConfigManager:
         if sharding_strategy.attention_data_parallelism > 1:
             if not envs.NEW_MODEL_DESIGN:
                 raise ValueError(
-                    "Must run Attention DP with NEW_MODEL_DESIGN enabled. Please set the "
-                    "NEW_MODEL_DESIGN=True.")
+                    "Must run Attention DP with NEW_MODEL_DESIGN enabled. Please set "
+                    "NEW_MODEL_DESIGN=True")
 
     @property
     def total_dp_size(self) -> int:

tpu_inference/layers/jax/attention/deepseek_v3_attention.py

@@ -30,6 +30,7 @@ from tpu_inference.kernels.ragged_paged_attention.v3.kernel import \
 from tpu_inference.kernels.ragged_paged_attention.v3.tuned_block_sizes import \
     get_tuned_block_sizes
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.quantization import quantize_kv
 from tpu_inference.layers.common.sharding import ShardingAxisName
 from tpu_inference.layers.jax.base import create_param
 from tpu_inference.layers.jax.layers import RMSNorm
@@ -310,9 +311,8 @@ class MLA(nnx.Module):
                     # TODO(kyuyeunk/jacobplatin): Enable w8a8 when VREG spill issue is resolved.
                     k_scale = self._k_scale
                     v_scale = self._v_scale
-                    k_SNH, v_SNH = utils.quantize_kv(
-                        k_SNH, v_SNH, self.kv_cache_quantized_dtype, k_scale,
-                        v_scale)
+                    k_SNH, v_SNH = quantize_kv(self.kv_cache_quantized_dtype,
+                                               k_SNH, v_SNH, k_scale, v_scale)
 
                 new_kv_cache, outputs_TNH = self.attention(
                     is_prefill,

tpu_inference/layers/jax/attention/gpt_oss_attention.py

@@ -26,6 +26,7 @@ from tpu_inference import utils
 from tpu_inference.kernels.ragged_paged_attention.v3.kernel_hd64 import \
     ragged_paged_attention_hd64
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.quantization import quantize_kv
 from tpu_inference.layers.jax.base import create_param
 from tpu_inference.layers.jax.rope import GptOssRotaryEmbedding
 
@@ -248,9 +249,8 @@ class GptOssAttention(nnx.Module):
             # q_scale = self._q_scale
             k_scale = self._k_scale
             v_scale = self._v_scale
-            k_TKH, v_TKH = utils.quantize_kv(k_TKH, v_TKH,
-                                             self.kv_cache_quantized_dtype,
-                                             k_scale, v_scale)
+            k_TKH, v_TKH = quantize_kv(self.kv_cache_quantized_dtype, k_TKH,
+                                       v_TKH, k_scale, v_scale)
 
         with jax.named_scope("attn_op"):
             new_kv_cache, attn_out_TNH = self.attention(

tpu_inference/layers/jax/attention/llama4_attention.py

@@ -19,8 +19,8 @@ import jax.numpy as jnp
 from flax import nnx
 from jax.sharding import Sharding
 
-from tpu_inference import utils
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.quantization import quantize_kv
 from tpu_inference.layers.jax.attention.attention import Attention, KVCache
 from tpu_inference.layers.jax.rope_interface import apply_rope
 from tpu_inference.logger import init_logger
@@ -128,9 +128,8 @@ class Llama4Attention(Attention):
             # q_scale = self._q_scale
             k_scale = self._k_scale
             v_scale = self._v_scale
-            k_SKH, v_SKH = utils.quantize_kv(k_SKH, v_SKH,
-                                             self.kv_cache_quantized_dtype,
-                                             k_scale, v_scale)
+            k_SKH, v_SKH = quantize_kv(self.kv_cache_quantized_dtype, k_SKH,
+                                       v_SKH, k_scale, v_scale)
 
         with jax.named_scope("attn_op"):
             new_kv_cache, outputs_TNH = self.attention(

tpu_inference/layers/jax/sample/sampling.py

@@ -42,7 +42,7 @@ def sample(
     if tpu_sampling_metadata.do_sampling:
         # Unshard the logits explicity to avoid latency increase.
         logits = jax.lax.with_sharding_constraint(
-            logits, NamedSharding(mesh, P(ShardingAxisName.ATTN_DATA, None)))
+            logits, NamedSharding(mesh, P(ShardingAxisName.MLP_DATA, None)))
     greedy_sampled = jnp.argmax(logits, axis=-1)
     if not tpu_sampling_metadata.do_sampling:
         return greedy_sampled

tpu_inference/layers/vllm/fused_moe.py

@@ -16,12 +16,14 @@ import functools
 
 import jax
 from jax import numpy as jnp
-from jax.sharding import Mesh
+from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
 
 from tpu_inference.kernels.megablox.gmm import gmm
+from tpu_inference.layers.common.sharding import ShardingAxisName
 from tpu_inference.layers.vllm.linear_common import \
     slice_sharded_tensor_for_concatenation
+from tpu_inference.utils import get_mesh_shape_product
 
 
 def activation_fn(activation: str, x1: jax.Array, x2: jax.Array) -> jax.Array:
@@ -137,25 +139,23 @@ def tensor_sharded_gmm_merged_column_parallel(
             group_offset=jnp.array(0),
         )
 
-    rhs_scale_spec = None if rhs_scale is None else P(None, None, None,
-                                                      "model")
-    rhs_bias_spec = None if rhs_bias is None else P(None, None, "model")
+    rhs_scale_spec = None if rhs_scale is None else P(
+        None, None, None, ShardingAxisName.MLP_TENSOR)
+    rhs_bias_spec = None if rhs_bias is None else P(
+        None, None, ShardingAxisName.MLP_TENSOR)
 
     gmm_result = jax.shard_map(
         _gmm,
         mesh=mesh,
-        in_specs=(
-            P("data", None),
-            P(None, "model", None),
-            rhs_scale_spec,
-            rhs_bias_spec,
-            P("data"),
-        ),
-        out_specs=(P("data", "model")),
+        in_specs=(P(ShardingAxisName.MLP_DATA,
+                    None), P(None, ShardingAxisName.MLP_TENSOR,
+                             None), rhs_scale_spec, rhs_bias_spec,
+                  P(ShardingAxisName.MLP_DATA)),
+        out_specs=(P(ShardingAxisName.MLP_DATA, ShardingAxisName.MLP_TENSOR)),
         check_vma=False,
     )(lhs, rhs, rhs_scale, rhs_bias, group_sizes)
 
-    tp_size = mesh.shape["model"]
+    tp_size = get_mesh_shape_product(mesh, ShardingAxisName.MLP_TENSOR)
     intermediate_size = gmm_result.shape[-1] // 2
     output_sizes = [intermediate_size, intermediate_size]
     return slice_sharded_tensor_for_concatenation(gmm_result, output_sizes,
@@ -175,7 +175,7 @@ def tensor_sharded_gmm_row_parallel(
         m, g, n, k = lhs.shape[0], *rhs.shape
         tm, tk, tn = _get_tiling_size_for_gmm_kernel(m, k, n, g)
         if rhs_bias is not None:
-            shard_id = jax.lax.axis_index("model")
+            shard_id = jax.lax.axis_index(ShardingAxisName.MLP_TENSOR).sum()
             rhs_bias = jnp.where(shard_id == 0, rhs_bias, 0)
         out = gmm(
             lhs,
@@ -188,22 +188,19 @@ def tensor_sharded_gmm_row_parallel(
             transpose_rhs=True,
             group_offset=jnp.array(0),
         )
-        return jax.lax.psum(out, axis_name="model")
+        return jax.lax.psum(out, axis_name=ShardingAxisName.MLP_TENSOR)
 
     num_blocks = 1 if rhs_scale is None else rhs_scale.shape[1]
-    rhs_scale_spec = None if num_blocks == 1 else P(None, "model", None, None)
+    rhs_scale_spec = None if num_blocks == 1 else P(
+        None, ShardingAxisName.MLP_TENSOR, None, None)
     rhs_bias_spec = None if rhs_bias is None else P(None, None, None)
     gmm_result = jax.shard_map(
         _gmm_all_reduce,
         mesh=mesh,
-        in_specs=(
-            P("data", "model"),
-            P(None, None, "model"),
-            rhs_scale_spec,
-            rhs_bias_spec,
-            P("data"),
-        ),
-        out_specs=(P("data")),
+        in_specs=(P(ShardingAxisName.MLP_DATA, ShardingAxisName.MLP_TENSOR),
+                  P(None, None, ShardingAxisName.MLP_TENSOR), rhs_scale_spec,
+                  rhs_bias_spec, P(ShardingAxisName.MLP_DATA)),
+        out_specs=(P(ShardingAxisName.MLP_DATA)),
         check_vma=False,
     )(lhs, rhs, rhs_scale, rhs_bias, group_sizes)
 
@@ -219,8 +216,8 @@ def expert_sharded_gmm(
     is_last_expert: bool,
     mesh: Mesh,
 ) -> jax.Array:
-    ep_size = mesh.shape["model"]
-
+    ep_size = get_mesh_shape_product(mesh, ShardingAxisName.MLP_TENSOR)
+    ep_p_spec = P(ShardingAxisName.EXPERT)
     num_experts = rhs.shape[0]
     num_experts_per_shard = num_experts // ep_size
     group_offset = jnp.arange(0, num_experts, num_experts_per_shard)
@@ -260,21 +257,22 @@ def expert_sharded_gmm(
     #       0, 0, 0, 0     0, 0, 0, 0     0, 0, 0, 0     D, D, D, D
     #        shard-0        shard-1        shard-2        shard-3
     # Each shards has 3 (row A), 2 (row B), 5 (row C) and 4 (row D).
-    lhs_spec = P("model") if is_last_expert else P()
-    rhs_scale_spec = None if rhs_scale is None else P("model")
-    rhs_bias_spec = None if rhs_bias is None else P("model")
+    lhs_spec = ep_p_spec if is_last_expert else P()
+    rhs_spec = ep_p_spec
+    rhs_scale_spec = None if rhs_scale is None else ep_p_spec
+    rhs_bias_spec = None if rhs_bias is None else ep_p_spec
     gmm_res = jax.shard_map(
         _gmm,
         mesh=mesh,
         in_specs=(
             lhs_spec,
-            P("model", None, None),
+            rhs_spec,
             rhs_scale_spec,
             rhs_bias_spec,
             P(),
-            P("model"),
+            ep_p_spec,
         ),
-        out_specs=(P("model", None)),
+        out_specs=ep_p_spec,
         check_vma=False,
     )(lhs, rhs, rhs_scale, rhs_bias, group_sizes, group_offset)
 
@@ -320,15 +318,13 @@ def expert_sharded_gmm(
         # send_sizes_of_shard          [3, 3, 3, 3]    [2, 2, 2, 2]    [5, 5, 5, 5]    [4, 4, 4, 4 ]
         # output_offsets_of_shard      [0, 0, 0, 0]    [0, 0, 0, 0]    [0, 0, 0, 0]    [10,10,10,10]
         # recv_sizes_of_shard          [3, 2, 5, 4]    [3, 2, 5, 4]    [3, 2, 5, 4]    [3, 2, 5, 4]
-        return jax.lax.ragged_all_to_all(
-            operand,
-            output,
-            input_offsets_of_shard,
-            send_sizes_of_shard,
-            output_offsets_of_shard,
-            recv_sizes_of_shard,
-            axis_name="model",
-        )
+        return jax.lax.ragged_all_to_all(operand,
+                                         output,
+                                         input_offsets_of_shard,
+                                         send_sizes_of_shard,
+                                         output_offsets_of_shard,
+                                         recv_sizes_of_shard,
+                                         axis_name=ShardingAxisName.EXPERT)
 
     # Use ragged_all_to_all to send the result from gmm for each expert to all
     # the shards.  In the working example, the result would be:
@@ -350,8 +346,8 @@ def expert_sharded_gmm(
     return jax.shard_map(
         _ragged_all_to_all,
         mesh=mesh,
-        in_specs=(P("model", None), P("model"), P("model"), P("model"), P()),
-        out_specs=(P()),
+        in_specs=(ep_p_spec, ep_p_spec, ep_p_spec, ep_p_spec, P()),
+        out_specs=(P(ShardingAxisName.MLP_DATA)),
         check_vma=False,
     )(gmm_res, input_offsets, send_sizes, output_offsets, recv_sizes)
 
@@ -412,6 +408,9 @@ def fused_moe_func(
     assert gating_output.shape == (num_tokens, global_num_experts)
 
     topk_weights = jax.nn.softmax(gating_output.astype(jnp.float32), axis=-1)
+    # All-gather topk weights for attention dp
+    topk_weights = jax.lax.with_sharding_constraint(
+        topk_weights, NamedSharding(mesh, P(ShardingAxisName.MLP_DATA, None)))
     topk_weights, topk_indices = jax.lax.top_k(topk_weights, k=topk)
     if renormalize:
         topk_weights = topk_weights / topk_weights.sum(axis=-1, keepdims=True)
@@ -434,8 +433,10 @@ def fused_moe_func(
     x, group_sizes, topk_argsort_revert_indices = jax.shard_map(
         _process_tokens_locally,
         mesh=mesh,
-        in_specs=(P("data", None), P("data", None)),
-        out_specs=(P("data", None), P("data"), P("data")),
+        in_specs=(P(ShardingAxisName.MLP_DATA,
+                    None), P(ShardingAxisName.MLP_DATA, None)),
+        out_specs=(P(ShardingAxisName.MLP_DATA, None),
+                   P(ShardingAxisName.MLP_DATA), P(ShardingAxisName.MLP_DATA)),
     )(hidden_states, topk_indices)
 
     x = jnp.pad(x, ((0, 0), (0, padded_hidden_size - hidden_size)))
@@ -495,8 +496,11 @@ def fused_moe_func(
     x = jax.shard_map(
         _finalize_output,
         mesh=mesh,
-        in_specs=(P("data", None), P("data"), P("data", None)),
-        out_specs=(P("data", None)),
+        in_specs=(P(ShardingAxisName.MLP_DATA,
+                    None), P(ShardingAxisName.MLP_DATA),
+                  P(ShardingAxisName.MLP_DATA, None)),
+        out_specs=(P(ShardingAxisName.ATTN_DATA, None)),
+        check_vma=False,
     )(x, topk_argsort_revert_indices, topk_weights)
 
     return x[:num_tokens, :hidden_size]