remove no exp

NuojCheng · NuojCheng · commit 87525554b94a · 2026-04-06T16:57:29.000Z
diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py
@@ -32,7 +32,6 @@
 AxisIdxes = tuple[int, ...]
 
 BATCH = "activation_batch"
-BATCH_NO_EXP = "activation_batch_no_exp"
 
 ATTN_LENGTH = "activation_attn_length"
 ATTN_LENGTH_NO_EXP = "activation_attn_length_no_exp"
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
@@ -450,7 +450,6 @@ mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'co
 logical_axis_rules: [
                       ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
-                      ['activation_batch_no_exp', ['data', 'fsdp', 'fsdp_transpose']],
                       ['activation_batch_no_exp_moe', ['data', 'fsdp', 'fsdp_transpose']],
                       ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'expert']],
diff --git a/src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml b/src/maxtext/configs/custom_mesh_and_rule/pipeline-large-moe.yml
@@ -31,7 +31,6 @@ data_sharding: [['data', 'stage', 'fsdp', 'context', 'tensor', 'expert']]
 logical_axis_rules: [
                       ['activation_batch', ['data', 'fsdp', 'expert']],
                       ['activation_batch_moe', ['data', 'fsdp', 'expert']],
-                      ['activation_batch_no_exp', ['data', 'fsdp']],
                       ['activation_batch_no_exp_moe', ['data', 'fsdp']],
                       ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'expert']],
                       ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'context', 'expert']],
diff --git a/src/maxtext/configs/custom_mesh_and_rule/pure-fsdp.yml b/src/maxtext/configs/custom_mesh_and_rule/pure-fsdp.yml
@@ -18,7 +18,6 @@ mesh_axes: ['fsdp']
 data_sharding: [['fsdp']]
 logical_axis_rules: [
                       ['activation_batch', ['fsdp']],
-                      ['activation_batch_no_exp', ['fsdp']],
                       ['activation_batch_moe', ['fsdp']],
                       ['activation_batch_no_exp_moe', ['fsdp']],
                       ['activation_embed_and_logits_batch', ['fsdp']],
diff --git a/src/maxtext/configs/inference/inference.yml b/src/maxtext/configs/inference/inference.yml
@@ -2,7 +2,6 @@ base_config: "base.yml"
 
 logical_axis_rules: [
                       ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
-                      ['activation_batch_no_exp', ['data', 'fsdp', 'fsdp_transpose']],
                       ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_heads', ['tensor', 'tensor_transpose', 'sequence','tensor_sequence']],
                       ['activation_kv_heads', ['tensor', 'tensor_transpose', 'sequence','tensor_sequence']],
diff --git a/src/maxtext/configs/inference/vllm.yml b/src/maxtext/configs/inference/vllm.yml
@@ -31,7 +31,6 @@ mesh_axes: ['data', 'attn_dp', 'model', 'expert', 'attn_dp_expert']
 logical_axis_rules: [
                       ['activation_batch', ['data']],
                       ['activation_batch_moe', []],
-                      ['activation_batch_no_exp', []],
                       ['activation_batch_no_exp_moe', []],
                       ['activation_embed_and_logits_batch', ['data', 'expert']],
                       ['activation_embed_and_logits_batch_sequence', ['data', 'expert']],
diff --git a/src/maxtext/configs/post_train/rl_mt_jt.yml b/src/maxtext/configs/post_train/rl_mt_jt.yml
@@ -18,7 +18,6 @@ logical_axis_rules: [
                       ['prefill_activation_length', ['data']],
                       ['prefill_activation_norm_length', ['data']],
                       ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
-                      ['activation_batch_no_exp', ['data', 'fsdp', 'fsdp_transpose']],
                       ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_heads', ['tensor', 'tensor_transpose', 'sequence','tensor_sequence']],
                       ['activation_kv_heads', ['tensor', 'tensor_transpose', 'sequence','tensor_sequence']],
diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py
@@ -37,7 +37,6 @@
     AxisIdxes,
     AxisNames,
     BATCH,
-    BATCH_NO_EXP,
     CACHE_BATCH,
     CACHE_BATCH_PREFILL,
     CACHE_SEQUENCE,
@@ -432,9 +431,7 @@ def mla_as_linen(
     ep_key_axis_names: AxisNames = (KV_BATCH_NO_EXP, LENGTH, KV_HEAD, KV_HEAD_DIM),
     ep_value_axis_names: AxisNames = (KV_BATCH_NO_EXP, LENGTH, KV_HEAD, KV_HEAD_DIM),
     input_axis_names: AxisNames = (BATCH, LENGTH_NO_EXP, EMBED),
-    ep_input_axis_names: AxisNames = (BATCH_NO_EXP, LENGTH, EMBED),
     out_axis_names: AxisNames = (BATCH, LENGTH_NO_EXP, HEAD, D_KV),
-    ep_out_axis_names: AxisNames = (BATCH_NO_EXP, LENGTH, HEAD, D_KV),
     prefill_input_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, EMBED),
     decode_input_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, EMBED),
     prefill_out_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV),
@@ -503,9 +500,7 @@ def mla_as_linen(
       ep_key_axis_names=ep_key_axis_names,
       ep_value_axis_names=ep_value_axis_names,
       input_axis_names=input_axis_names,
-      ep_input_axis_names=ep_input_axis_names,
       out_axis_names=out_axis_names,
-      ep_out_axis_names=ep_out_axis_names,
       prefill_input_axis_names=prefill_input_axis_names,
       decode_input_axis_names=decode_input_axis_names,
       prefill_out_axis_names=prefill_out_axis_names,
@@ -580,9 +575,7 @@ def __init__(
       ep_key_axis_names: AxisNames = (KV_BATCH_NO_EXP, LENGTH, KV_HEAD, KV_HEAD_DIM),
       ep_value_axis_names: AxisNames = (KV_BATCH_NO_EXP, LENGTH, KV_HEAD, KV_HEAD_DIM),
       input_axis_names: AxisNames = (BATCH, LENGTH_NO_EXP, EMBED),
-      ep_input_axis_names: AxisNames = (BATCH_NO_EXP, LENGTH, EMBED),
       out_axis_names: AxisNames = (BATCH, LENGTH_NO_EXP, HEAD, D_KV),
-      ep_out_axis_names: AxisNames = (BATCH_NO_EXP, LENGTH, HEAD, D_KV),
       prefill_input_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, EMBED),
       decode_input_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, EMBED),
       prefill_out_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV),
@@ -668,9 +661,7 @@ def __init__(
         ep_key_axis_names=ep_key_axis_names,
         ep_value_axis_names=ep_value_axis_names,
         input_axis_names=input_axis_names,
-        ep_input_axis_names=ep_input_axis_names,
         out_axis_names=out_axis_names,
-        ep_out_axis_names=ep_out_axis_names,
         prefill_input_axis_names=prefill_input_axis_names,
         decode_input_axis_names=decode_input_axis_names,
         prefill_out_axis_names=prefill_out_axis_names,
@@ -1178,10 +1169,6 @@ def __call__(
       inputs_q = self._maybe_shard_with_logical(inputs_q, self.prefill_input_axis_names)
       inputs_kv = self._maybe_shard_with_logical(inputs_kv, self.prefill_input_axis_names)
       out_logical_name = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV)
-    elif model_mode == MODEL_MODE_TRAIN and self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      inputs_q = self._maybe_shard_with_logical(inputs_q, self.ep_input_axis_names)
-      inputs_kv = self._maybe_shard_with_logical(inputs_kv, self.ep_input_axis_names)
-      out_logical_name = (BATCH_NO_EXP, LENGTH, HEAD, D_KV)
     else:
       inputs_q = self._maybe_shard_with_logical(inputs_q, self.input_axis_names)
       inputs_kv = self._maybe_shard_with_logical(inputs_kv, self.input_axis_names)
diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py
@@ -39,7 +39,6 @@
     AxisIdxes,
     AxisNames,
     BATCH,
-    BATCH_NO_EXP,
     CACHE_BATCH,
     CACHE_BATCH_PREFILL,
     CACHE_HEADS,
@@ -61,7 +60,6 @@
     HEAD,
     KV_LENGTH,
     LENGTH,
-    LENGTH_NO_EXP,
     MODEL_MODE_AUTOREGRESSIVE,
     MODEL_MODE_PREFILL,
     MODEL_MODE_TRAIN,
@@ -302,12 +300,9 @@ def attention_op_as_linen(
     float32_qk_product: bool = False,
     max_prefill_predict_length: int = -1,
     float32_logits: bool = False,
-    flash_axis_names_q: AxisNames = (BATCH, HEAD, LENGTH_NO_EXP, D_KV),
-    flash_axis_names_q_ep: AxisNames = (BATCH_NO_EXP, HEAD, LENGTH, D_KV),
+    flash_axis_names_q: AxisNames = (BATCH, HEAD, LENGTH, D_KV),
     flash_axis_names_kv: AxisNames = (BATCH, HEAD, KV_LENGTH, D_KV),
-    flash_axis_names_kv_ep: AxisNames = (BATCH_NO_EXP, HEAD, KV_LENGTH, D_KV),
-    flash_axis_names_splash_kernel: AxisNames = (HEAD, LENGTH_NO_EXP),
-    flash_axis_names_splash_kernel_ep: AxisNames = (HEAD, LENGTH),
+    flash_axis_names_splash_kernel: AxisNames = (HEAD, LENGTH),
     prefill_cache_logical_axis_names: AxisNames = (
         CACHE_BATCH_PREFILL,
         CACHE_SEQUENCE,
@@ -364,11 +359,8 @@ def attention_op_as_linen(
       max_prefill_predict_length=max_prefill_predict_length,
       float32_logits=float32_logits,
       flash_axis_names_q=flash_axis_names_q,
-      flash_axis_names_q_ep=flash_axis_names_q_ep,
       flash_axis_names_kv=flash_axis_names_kv,
-      flash_axis_names_kv_ep=flash_axis_names_kv_ep,
       flash_axis_names_splash_kernel=flash_axis_names_splash_kernel,
-      flash_axis_names_splash_kernel_ep=flash_axis_names_splash_kernel_ep,
       prefill_cache_logical_axis_names=prefill_cache_logical_axis_names,
       cache_logical_axis_names=cache_logical_axis_names,
       cache_scale_logical_axis_names=cache_scale_logical_axis_names,
@@ -405,12 +397,9 @@ def __init__(
       float32_qk_product: bool = False,
       max_prefill_predict_length: int = -1,
       float32_logits: bool = False,
-      flash_axis_names_q: AxisNames = (BATCH, HEAD, LENGTH_NO_EXP, D_KV),
-      flash_axis_names_q_ep: AxisNames = (BATCH_NO_EXP, HEAD, LENGTH, D_KV),
+      flash_axis_names_q: AxisNames = (BATCH, HEAD, LENGTH, D_KV),
       flash_axis_names_kv: AxisNames = (BATCH, HEAD, KV_LENGTH, D_KV),
-      flash_axis_names_kv_ep: AxisNames = (BATCH_NO_EXP, HEAD, KV_LENGTH, D_KV),
-      flash_axis_names_splash_kernel: AxisNames = (HEAD, LENGTH_NO_EXP),
-      flash_axis_names_splash_kernel_ep: AxisNames = (HEAD, LENGTH),
+      flash_axis_names_splash_kernel: AxisNames = (HEAD, LENGTH),
       prefill_cache_logical_axis_names: AxisNames = (
           CACHE_BATCH_PREFILL,
           CACHE_SEQUENCE,
@@ -492,11 +481,8 @@ def __init__(
     self.max_prefill_predict_length = max_prefill_predict_length
     self.float32_logits = float32_logits
     self.flash_axis_names_q = flash_axis_names_q
-    self.flash_axis_names_q_ep = flash_axis_names_q_ep
     self.flash_axis_names_kv = flash_axis_names_kv
-    self.flash_axis_names_kv_ep = flash_axis_names_kv_ep
     self.flash_axis_names_splash_kernel = flash_axis_names_splash_kernel
-    self.flash_axis_names_splash_kernel_ep = flash_axis_names_splash_kernel_ep
     self.prefill_cache_logical_axis_names = prefill_cache_logical_axis_names
     self.cache_logical_axis_names = cache_logical_axis_names
     self.cache_scale_logical_axis_names = cache_scale_logical_axis_names
@@ -1150,23 +1136,13 @@ def tpu_flash_attention(
     segment_axis_names_kv = None
     sink_axis_names = self._logical_to_mesh_axes((HEAD,))
     if decoder_segment_ids is not None:
-      if self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-        segment_axis_names_q = self._logical_to_mesh_axes((BATCH_NO_EXP, Q_LENGTH))
-        segment_axis_names_kv = self._logical_to_mesh_axes((BATCH_NO_EXP, KV_LENGTH))
-      else:
-        segment_axis_names_q = self._logical_to_mesh_axes((BATCH, Q_LENGTH_NO_EXP))
-        segment_axis_names_kv = self._logical_to_mesh_axes((BATCH, KV_LENGTH))
-
-    if self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      axis_names_splash_kernel = self._logical_to_mesh_axes(self.flash_axis_names_splash_kernel_ep)
-      axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q_ep)
-      axis_names_kv = self._logical_to_mesh_axes(self.flash_axis_names_kv_ep)
-      indexer_mask_axis_names = self._logical_to_mesh_axes((BATCH_NO_EXP, Q_LENGTH, KV_LENGTH))
-    else:
-      axis_names_splash_kernel = self._logical_to_mesh_axes(self.flash_axis_names_splash_kernel)
-      axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q)
-      axis_names_kv = self._logical_to_mesh_axes(self.flash_axis_names_kv)
-      indexer_mask_axis_names = self._logical_to_mesh_axes((BATCH, Q_LENGTH, KV_LENGTH))
+      segment_axis_names_q = self._logical_to_mesh_axes((BATCH, Q_LENGTH_NO_EXP))
+      segment_axis_names_kv = self._logical_to_mesh_axes((BATCH, KV_LENGTH))
+
+    axis_names_splash_kernel = self._logical_to_mesh_axes(self.flash_axis_names_splash_kernel)
+    axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q)
+    axis_names_kv = self._logical_to_mesh_axes(self.flash_axis_names_kv)
+    indexer_mask_axis_names = self._logical_to_mesh_axes((BATCH, Q_LENGTH, KV_LENGTH))
 
     global global_block_q, global_block_kv, global_block_kv_compute, global_block_q_dkv, global_block_kv_dkv
     global global_block_kv_dkv_compute, global_block_q_dq, global_block_kv_dq, global_use_fused_bwd_kernel
diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py
@@ -28,7 +28,6 @@
 from maxtext.common.common_types import (
     DecoderBlockType,
     BATCH,
-    BATCH_NO_EXP,
     HEAD,
     PREFILL_LENGTH,
     D_KV,
@@ -149,9 +148,7 @@ def attention_as_linen(
     ep_key_axis_names: AxisNames = (KV_BATCH_NO_EXP, ATTN_LENGTH, KV_HEAD, KV_HEAD_DIM),
     ep_value_axis_names: AxisNames = (KV_BATCH_NO_EXP, ATTN_LENGTH, KV_HEAD, KV_HEAD_DIM),
     input_axis_names: AxisNames = (BATCH, ATTN_LENGTH_NO_EXP, ATTN_EMBED),
-    ep_input_axis_names: AxisNames = (BATCH_NO_EXP, ATTN_LENGTH, ATTN_EMBED),
     out_axis_names: AxisNames = (BATCH, ATTN_LENGTH_NO_EXP, HEAD, D_KV),
-    ep_out_axis_names: AxisNames = (BATCH_NO_EXP, ATTN_LENGTH, HEAD, D_KV),
     prefill_input_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, ATTN_EMBED),
     decode_input_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, ATTN_EMBED),
     prefill_out_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV),
@@ -215,9 +212,7 @@ def attention_as_linen(
       ep_key_axis_names=ep_key_axis_names,
       ep_value_axis_names=ep_value_axis_names,
       input_axis_names=input_axis_names,
-      ep_input_axis_names=ep_input_axis_names,
       out_axis_names=out_axis_names,
-      ep_out_axis_names=ep_out_axis_names,
       prefill_input_axis_names=prefill_input_axis_names,
       decode_input_axis_names=decode_input_axis_names,
       prefill_out_axis_names=prefill_out_axis_names,
@@ -316,9 +311,7 @@ def __init__(
       ep_key_axis_names: AxisNames = (KV_BATCH_NO_EXP, ATTN_LENGTH, KV_HEAD, KV_HEAD_DIM),
       ep_value_axis_names: AxisNames = (KV_BATCH_NO_EXP, ATTN_LENGTH, KV_HEAD, KV_HEAD_DIM),
       input_axis_names: AxisNames = (BATCH, ATTN_LENGTH_NO_EXP, ATTN_EMBED),
-      ep_input_axis_names: AxisNames = (BATCH_NO_EXP, ATTN_LENGTH, ATTN_EMBED),
       out_axis_names: AxisNames = (BATCH, ATTN_LENGTH_NO_EXP, HEAD, D_KV),
-      ep_out_axis_names: AxisNames = (BATCH_NO_EXP, ATTN_LENGTH, HEAD, D_KV),
       prefill_input_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, ATTN_EMBED),
       decode_input_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, ATTN_EMBED),
       prefill_out_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV),
@@ -424,9 +417,7 @@ def __init__(
     self.ep_key_axis_names = ep_key_axis_names
     self.ep_value_axis_names = ep_value_axis_names
     self.input_axis_names = input_axis_names
-    self.ep_input_axis_names = ep_input_axis_names
     self.out_axis_names = out_axis_names
-    self.ep_out_axis_names = ep_out_axis_names
     self.prefill_input_axis_names = prefill_input_axis_names
     self.decode_input_axis_names = decode_input_axis_names
     self.prefill_out_axis_names = prefill_out_axis_names
@@ -1100,8 +1091,6 @@ def __call__(
     """
     if model_mode == MODEL_MODE_PREFILL:
       input_axis_names = self.prefill_input_axis_names
-    elif model_mode == MODEL_MODE_TRAIN and self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      input_axis_names = self.ep_input_axis_names
     elif model_mode == MODEL_MODE_TRAIN:
       input_axis_names = self.input_axis_names
     else:
@@ -1219,8 +1208,6 @@ def __call__(
     out = jax.ad_checkpoint.checkpoint_name(out, "attention_out")
     if model_mode == MODEL_MODE_PREFILL:
       out = self._maybe_shard_with_logical(out, self.prefill_out_axis_names)
-    elif model_mode == MODEL_MODE_TRAIN and self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      out = self._maybe_shard_with_logical(out, self.ep_out_axis_names)
     elif model_mode == MODEL_MODE_TRAIN:
       out = self._maybe_shard_with_logical(out, self.out_axis_names)
     else:
diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py
@@ -27,7 +27,7 @@
 from jax.ad_checkpoint import checkpoint_name
 import jax.numpy as jnp
 from jax.sharding import Mesh
-from maxtext.common.common_types import Config, DecoderBlockType, EP_AS_CONTEXT, ShardMode
+from maxtext.common.common_types import Config, DecoderBlockType, ShardMode
 from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN
 from maxtext.inference import page_manager
 from maxtext.layers import linears
@@ -106,8 +106,6 @@ def __call__(
 
     if self.model_mode == MODEL_MODE_PREFILL:
       logical_axis_names = ("activation_batch", "prefill_activation_length", "activation_embed")
-    elif self.config.expert_shard_attention_option == EP_AS_CONTEXT and self.model_mode == MODEL_MODE_TRAIN:
-      logical_axis_names = ("activation_batch_no_exp", "activation_length", "activation_embed")
     else:
       logical_axis_names = ("activation_batch", "activation_length_no_exp", "activation_embed")
 
diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py
@@ -30,7 +30,7 @@
 import flax.linen as nn
 
 from maxtext.common.common_types import DecoderBlockType, ShardMode, DType, Array, Config
-from maxtext.common.common_types import MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, EP_AS_CONTEXT
+from maxtext.common.common_types import MODEL_MODE_PREFILL
 from maxtext.layers import nnx_wrappers, quantizations
 from maxtext.layers import normalizations
 from maxtext.layers.initializers import NdInitializer, nd_dense_init, default_bias_init, variable_to_logically_partitioned
@@ -404,8 +404,6 @@ def __init__(
 
     if self.model_mode == MODEL_MODE_PREFILL:
       self.intermediate_logical = ("activation_batch", "prefill_activation_length", "activation_mlp")
-    elif config.expert_shard_attention_option == EP_AS_CONTEXT and self.model_mode == MODEL_MODE_TRAIN:
-      self.intermediate_logical = ("activation_batch_no_exp", "activation_length", "activation_mlp")
     else:
       self.intermediate_logical = ("activation_batch", "activation_length_no_exp", "activation_mlp")
 
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
@@ -30,7 +30,6 @@
 from jax.sharding import Mesh
 
 from maxtext.common.common_types import (
-    EP_AS_CONTEXT,
     MODEL_MODE_AUTOREGRESSIVE,
     MODEL_MODE_PREFILL,
     MODEL_MODE_TRAIN,
@@ -169,8 +168,6 @@ def __call__(
 
     if self.model_mode == MODEL_MODE_PREFILL:
       logical_axis_names = ("activation_batch", "prefill_activation_length", "activation_embed")
-    elif self.config.expert_shard_attention_option == EP_AS_CONTEXT and self.model_mode == MODEL_MODE_TRAIN:
-      logical_axis_names = ("activation_batch_no_exp", "activation_length", "activation_embed")
     else:
       logical_axis_names = ("activation_batch", "activation_length_no_exp", "activation_embed")
 
diff --git a/src/maxtext/layers/pipeline.py b/src/maxtext/layers/pipeline.py
@@ -28,7 +28,7 @@
 from flax import linen as nn
 from flax.linen.spmd import LogicallyPartitioned
 
-from maxtext.common.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT, ShardMode
+from maxtext.common.common_types import Config, MODEL_MODE_TRAIN, ShardMode
 from maxtext.utils.sharding import (
     maybe_shard_with_logical,
     maybe_shard_with_name,
@@ -56,12 +56,8 @@ def setup(self):
     self.microbatches_per_stage = microbatches_per_stage
     self.use_circ_storage = self.need_circ_storage()
 
-    if self.config.expert_shard_attention_option == EP_AS_CONTEXT:
-      self.batch_axis_name = "activation_batch_no_exp"
-      self.seq_len_axis_name = "activation_length"
-    else:
-      self.batch_axis_name = "activation_batch"
-      self.seq_len_axis_name = "activation_length_no_exp"
+    self.batch_axis_name = "activation_batch"
+    self.seq_len_axis_name = "activation_length_no_exp"
 
     self.spmd_axis_name = "stage" if self.config.shard_mode == ShardMode.AUTO else None
 
diff --git a/src/maxtext/layers/pipeline_deprecated.py b/src/maxtext/layers/pipeline_deprecated.py
diff --git a/tests/utils/attention_test_util.py b/tests/utils/attention_test_util.py
