Enable PCP for GDN models

tests/layers/common/test_sharding.py

@@ -329,14 +329,14 @@ def test_initial_state_is_uninitialized(self):
     @patch("tpu_inference.layers.common.sharding.envs.NEW_MODEL_DESIGN", False)
     def test_use_2d_tp_selects_base(self):
         lazy = LazyShardingAxisName()
-        _ = lazy.SEQUENCE
+        _ = lazy.DENSE_DATA
         self.assertIs(lazy._cls, ShardingAxisNameBase)
 
     @patch("tpu_inference.layers.common.sharding.envs.USE_2D_TP", False)
     @patch("tpu_inference.layers.common.sharding.envs.NEW_MODEL_DESIGN", False)
     def test_both_false_selects_2d(self):
         lazy = LazyShardingAxisName()
-        _ = lazy.SEQUENCE
+        _ = lazy.DENSE_DATA
         self.assertIs(lazy._cls, ShardingAxisName2D)
 
     def test_cls_cached_after_first_access(self):
@@ -345,15 +345,15 @@ def test_cls_cached_after_first_access(self):
                    False), \
              patch("tpu_inference.layers.common.sharding.envs.NEW_MODEL_DESIGN",
                    False):
-            _ = lazy.SEQUENCE
+            _ = lazy.DENSE_DATA
             first_cls = lazy._cls
 
         # Even with different env vars, _cls should not change
         with patch("tpu_inference.layers.common.sharding.envs.USE_2D_TP",
                    True), \
              patch("tpu_inference.layers.common.sharding.envs.NEW_MODEL_DESIGN",
                    True):
-            _ = lazy.SEQUENCE
+            _ = lazy.DENSE_DATA
             self.assertIs(lazy._cls, first_cls)
 
     def test_initialized_after_env_change_uses_new_value(self):
@@ -372,7 +372,7 @@ def test_initialized_after_env_change_uses_new_value(self):
                    True), \
              patch("tpu_inference.layers.common.sharding.envs.NEW_MODEL_DESIGN",
                    False):
-            _ = lazy.SEQUENCE  # initialized here, after env changed
+            _ = lazy.DENSE_DATA  # initialized here, after env changed
             self.assertIs(lazy._cls, ShardingAxisNameBase)
 
 

tests/layers/vllm/test_fused_moe.py

@@ -119,7 +119,7 @@ class TestMaybeReduceSharedExpertOutput:
 
     def test_passthrough_when_shared_output_none(self):
         runner = _make_runner()
-        with patch.object(fm, "_all_reduce_over_tp") as reduce:
+        with patch.object(fm, "_sum_partials") as reduce:
             assert runner._maybe_reduce_shared_expert_output(None) is None
         reduce.assert_not_called()
 
@@ -129,7 +129,7 @@ def test_passthrough_under_sequence_parallel(self):
         shared = torch.ones(2, 3)
         with patch.object(fm.VllmMoERunner, "_fused_output_is_reduced",
                           new_callable=PropertyMock, return_value=True), \
-             patch.object(fm, "_all_reduce_over_tp") as reduce:
+             patch.object(fm, "_sum_partials") as reduce:
             out = runner._maybe_reduce_shared_expert_output(shared)
         reduce.assert_not_called()
         assert out is shared
@@ -141,7 +141,7 @@ def test_passthrough_when_fused_output_not_reduced(self):
         shared = torch.ones(2, 3)
         with patch.object(fm.VllmMoERunner, "_fused_output_is_reduced",
                           new_callable=PropertyMock, return_value=False), \
-             patch.object(fm, "_all_reduce_over_tp") as reduce:
+             patch.object(fm, "_sum_partials") as reduce:
             out = runner._maybe_reduce_shared_expert_output(shared)
         reduce.assert_not_called()
         assert out is shared
@@ -155,27 +155,48 @@ def test_reduces_shared_output_on_early_path(self):
                           new_callable=PropertyMock, return_value=True), \
              patch.object(fm, "_get_mesh", return_value=mesh), \
              patch.object(fm, "is_attn_dp", return_value=False), \
-             patch.object(fm, "_all_reduce_over_tp",
+             patch.object(fm, "_sum_partials",
                           return_value=reduced) as reduce:
             out = runner._maybe_reduce_shared_expert_output(shared)
-        reduce.assert_called_once_with(shared, mesh,
-                                       fm.ShardingAxisName.MLP_TENSOR)
+        reduce.assert_called_once_with(shared)
         assert out is reduced
 
     def test_reduces_shared_output_under_attention_dp(self):
+        # Under attention DP the fused kernel reduces its own output, so the
+        # shared-expert partial stack is summed on the early path (its
+        # leading axis is sharded over the in-group TP axis). Drives the real
+        # ``_fused_output_is_reduced`` property via ``is_attn_dp``.
         runner = _make_runner(shared_experts=object())
         shared = torch.ones(2, 3)
         reduced = torch.full((2, 3), 7.0)
         mesh = object()
         with patch.object(fm, "_get_mesh", return_value=mesh), \
              patch.object(fm, "is_attn_dp", return_value=True), \
-             patch.object(fm, "_all_reduce_over_tp",
+             patch.object(fm, "_sum_partials",
                           return_value=reduced) as reduce:
             out = runner._maybe_reduce_shared_expert_output(shared)
-        reduce.assert_called_once_with(shared, mesh,
-                                       fm.ShardingAxisName.ATTN_HEAD)
+        reduce.assert_called_once_with(shared)
         assert out is reduced
 
+    def test_pcp_is_plain_tp(self):
+        # PCP is not attention DP (tokens are replicated over pcp outside
+        # attention, weights shard over pcp as extra TP), so a GMM backend
+        # whose reduction matches the shared expert's takes the deferred path:
+        # nothing is reduced early, shared + fused are reduced together late.
+        runner = _make_runner(shared_experts=object())
+        shared = torch.ones(2, 3)
+        with patch.object(fm, "_get_mesh", return_value=object()), \
+             patch.object(fm, "is_attn_dp", return_value=False), \
+             patch.object(fm, "select_moe_backend_from_fused_moe_config",
+                          return_value=MoEBackend.GMM_EP), \
+             patch.object(fm, "_gmm_and_shared_reduce_over_same_axes",
+                          return_value=True), \
+             patch.object(fm, "_sum_partials") as reduce:
+            assert runner._fused_output_is_reduced is False
+            out = runner._maybe_reduce_shared_expert_output(shared)
+        reduce.assert_not_called()
+        assert out is shared
+
 
 # ---------------------------------------------------------------------------
 # VllmMoERunner._maybe_reduce_final_output  (late path)
@@ -189,7 +210,7 @@ def test_attention_dp_only_truncates(self):
         states = torch.arange(8, dtype=torch.float32).reshape(2, 4)
         with patch.object(fm, "_get_mesh", return_value=object()), \
              patch.object(fm, "is_attn_dp", return_value=True), \
-             patch.object(fm, "_all_reduce_over_tp") as reduce:
+             patch.object(fm, "_sum_partials") as reduce:
             out = runner._maybe_reduce_final_output(states, trunc_size=3)
         reduce.assert_not_called()
         torch.testing.assert_close(out, states[..., :3])
@@ -201,7 +222,7 @@ def test_sequence_parallel_only_truncates(self):
              patch.object(fm, "is_attn_dp", return_value=False), \
              patch.object(fm.VllmMoERunner, "_fused_output_is_reduced",
                           new_callable=PropertyMock, return_value=False), \
-             patch.object(fm, "_all_reduce_over_tp") as reduce:
+             patch.object(fm, "_sum_partials") as reduce:
             out = runner._maybe_reduce_final_output(states, trunc_size=2)
         reduce.assert_not_called()
         torch.testing.assert_close(out, states[..., :2])
@@ -215,7 +236,7 @@ def test_only_truncates_when_kernel_already_reduced(self):
              patch.object(fm, "is_attn_dp", return_value=False), \
              patch.object(fm.VllmMoERunner, "_fused_output_is_reduced",
                           new_callable=PropertyMock, return_value=True), \
-             patch.object(fm, "_all_reduce_over_tp") as reduce:
+             patch.object(fm, "_sum_partials") as reduce:
             out = runner._maybe_reduce_final_output(states, trunc_size=2)
         reduce.assert_not_called()
         torch.testing.assert_close(out, states[..., :2])
@@ -229,9 +250,9 @@ def test_reduces_then_truncates_on_late_path(self):
              patch.object(fm, "is_attn_dp", return_value=False), \
              patch.object(fm.VllmMoERunner, "_fused_output_is_reduced",
                           new_callable=PropertyMock, return_value=False), \
-             patch.object(fm, "_all_reduce_over_tp",
+             patch.object(fm, "_sum_partials",
                           return_value=reduced) as reduce:
             out = runner._maybe_reduce_final_output(states, trunc_size=3)
-        reduce.assert_called_once_with(states, mesh)
+        reduce.assert_called_once_with(states)
         # Reduction happens first, then the padding is stripped.
         torch.testing.assert_close(out, reduced[..., :3])

tests/layers/vllm/test_unquantized.py

@@ -40,7 +40,7 @@
 from tests.layers.common import utils as test_utils
 from tpu_inference.layers.common.moe import MoEBackend
 from tpu_inference.layers.common.quantization.configs import QuantLinearConfig
-from tpu_inference.layers.vllm.custom_ops.fused_moe import _all_reduce_over_tp
+from tpu_inference.layers.vllm.custom_ops.fused_moe import _sum_partials
 from tpu_inference.layers.vllm.interface.moe import FusedMoEFactory
 from tpu_inference.layers.vllm.quantization import get_tpu_quantization_config
 from tpu_inference.layers.vllm.quantization.unquantized import (
@@ -214,9 +214,10 @@ def test_row_parallel_linear(model, bias, num_devices, enable_sp,
 @pytest.mark.parametrize("num_devices", [1, jax.local_device_count()])
 def test_row_parallel_linear_defer_all_reduce(model, num_devices):
     """reduce_results=False routes through sharded_matmul: the layer returns
-    per-shard partial sums (the psum is actually skipped, which a plain einsum
+    the per-shard partial sums stacked under a leading axis sharded over the
+    contraction axis (the psum is actually skipped, which a plain einsum
     under GSPMD cannot do) and the caller's single deferred all-reduce
-    (VllmMoERunner._all_reduce_over_tp) reconstructs the full result.
+    (VllmMoERunner._sum_partials) reconstructs the full result.
 
     No bias: vLLM's RowParallelLinear rejects reduce_results=False with an
     in-layer bias."""
@@ -280,10 +281,10 @@ def test_row_parallel_linear_defer_all_reduce(model, num_devices):
         if num_devices > 1:
             # The psum was actually skipped: pre-reduction output is partial,
             # not the full matmul.
-            partial = j2t(deferred.to(torch.float32)).to(dtype)
+            partial = j2t(deferred[0].to(torch.float32)).to(dtype)
             assert not torch.allclose(partial, expected)
 
-        reduced = _all_reduce_over_tp(deferred, mesh)
+        reduced = _sum_partials(deferred)
         jax_output = j2t(reduced.to(torch.float32)).to(dtype)
 
     torch.testing.assert_close(expected, jax_output)

tests/runner/test_kv_cache.py

@@ -58,7 +58,8 @@ def test_create_kv_caches(mesh: Mesh):
                    ]), patch("tpu_inference.envs.NEW_MODEL_DESIGN", True):
         expected_sharding = NamedSharding(
             mesh,
-            PartitionSpec(ShardingAxisName.BATCH, ShardingAxisName.KV_CONTEXT,
+            PartitionSpec(ShardingAxisName.DENSE_DATA,
+                          ShardingAxisName.KV_CONTEXT,
                           ShardingAxisName.KV_HEAD))
         expected_shape = get_kv_cache_shape_with_mesh(mesh, num_blocks,
                                                       block_size, num_kv_heads,
@@ -99,8 +100,9 @@ def test_create_kv_caches_mla(mesh: Mesh):
 
     # For MLA, sharding is by the 'model' axis on the token dimension.
     expected_sharding = NamedSharding(
-        mesh, PartitionSpec(ShardingAxisName.BATCH,
-                            ShardingAxisName.KV_CONTEXT))
+        mesh,
+        PartitionSpec(ShardingAxisName.DENSE_DATA,
+                      ShardingAxisName.KV_CONTEXT))
     expected_dtype = jnp.bfloat16
     expected_shape = get_kv_cache_shape_with_mesh(
         mesh,

tpu_inference/kernels/experimental/rpa_v3_cp/pcp_vs_tp_benchmark.py

@@ -287,8 +287,9 @@ def make_pcp(pcp, tp):
         per_rank = rpa_v3_cp.get_kv_cache_shape(npages, page, NKV // tp, HD,
                                                 kv_dtype)
         cache_shape = (npages, gpage, per_rank[2] * tp) + tuple(per_rank[3:])
-        cache_spec = P(ShardingAxisName.BATCH, ShardingAxisName.KV_CONTEXT,
-                       ShardingAxisName.KV_HEAD, None, None)
+        cache_spec = P(ShardingAxisName.DENSE_DATA,
+                       ShardingAxisName.KV_CONTEXT, ShardingAxisName.KV_HEAD,
+                       None, None)
 
         def put(x, s):
             return jax.device_put(x, NamedSharding(mesh, s))
@@ -315,7 +316,7 @@ def put(x, s):
             pcp_cu[r, 2:] = C + treal
             pcp_qp[r, 0] = r * C
             pcp_qp[r, 1] = toff
-        pcp_spec = P(ShardingAxisName.PREFILL_CONTEXT, None)
+        pcp_spec = P(ShardingAxisName.PCP, None)
         pcp_cu = put(jnp.asarray(pcp_cu), pcp_spec)
         pcp_qp = put(jnp.asarray(pcp_qp), pcp_spec)
         fns = {}
@@ -350,7 +351,7 @@ def fn(cache, q, k, v, kvl, kvcl, _cp=cache_pages):
                     out = jax.shard_map(functools.partial(
                         layer_collectives,
                         axis=ShardingAxisName.ATTN_HEAD,
-                        gather_axis=ShardingAxisName.PREFILL_CONTEXT),
+                        gather_axis=ShardingAxisName.PCP),
                                         mesh=mesh,
                                         in_specs=q_spec,
                                         out_specs=q_spec,

tpu_inference/layers/common/attention_interface.py

@@ -38,6 +38,7 @@
     AttentionMetadata, SharedAttentionMetadata)
 from tpu_inference.layers.common.cp_attention import dcp_forward, pcp_forward
 from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.common.utils import replicate_kv_heads_for_tp
 from tpu_inference.logger import init_logger
 from tpu_inference.utils import get_megacore, get_mesh_shape_product
 
@@ -425,20 +426,10 @@ def sharded_ragged_paged_attention(
     decode_query_size: int = 1,
 ):
     """Shards along KV heads."""
-    # Handle GQA/MQA where num_kv_heads < tp_size
-    # We replicate KV heads to match tp_size so that we can shard them evenly.
+    # Handle GQA/MQA where num_kv_heads < tp_size: replicate KV heads so the
+    # head dim shards evenly.
     # TODO (ranlihao): This is not performant and introduces extra overhead during inference. We need to handle this during weight loading
-    tp_size = get_mesh_shape_product(mesh, ShardingAxisName.ATTN_HEAD)
-    if tp_size > 1:
-        num_kv_heads = k.shape[1]
-        if num_kv_heads < tp_size:
-            if tp_size % num_kv_heads != 0:
-                raise ValueError(
-                    f"For GQA/MQA, tp_size {tp_size} must be divisible by num_kv_heads {num_kv_heads}"
-                )
-            factor = tp_size // num_kv_heads
-            k = jnp.repeat(k, factor, axis=1)
-            v = jnp.repeat(v, factor, axis=1)
+    k, v = replicate_kv_heads_for_tp(mesh, k, v)
 
     qkv_spec = P(ShardingAxisName.ATTN_DATA, ShardingAxisName.ATTN_HEAD, None)
     kv_cache_spec = P(ShardingAxisName.ATTN_DATA, None,
@@ -651,14 +642,14 @@ def mla_attention(
         keyvalue_skh_sharding or P(ShardingAxisName.MLP_TENSOR, None),  # k
         keyvalue_skh_sharding
         or P(ShardingAxisName.MLP_TENSOR, None),  # k_rope
-        P(ShardingAxisName.BATCH),  # kv_cache
+        P(ShardingAxisName.DENSE_DATA),  # kv_cache
         P(ShardingAxisName.ATTN_DATA),  # md.seq_lens
         P(ShardingAxisName.ATTN_DATA),  # md.page_indices_flat
         P(ShardingAxisName.ATTN_DATA),  # md.query_start_loc
         P(ShardingAxisName.ATTN_DATA),  # md.distribution
     )
     out_specs = (
-        P(ShardingAxisName.BATCH),  # kv cache
+        P(ShardingAxisName.DENSE_DATA),  # kv cache
         attn_o_nth_sharding
         or P(None, ShardingAxisName.MLP_TENSOR, None)  # attn output
     )

tpu_inference/layers/common/cp_attention.py

@@ -22,6 +22,7 @@
 import tpu_inference.kernels.experimental.rpa_v3_cp.kernel as rpa_v3_cp
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.common.utils import replicate_kv_heads_for_tp
 from tpu_inference.logger import init_logger
 from tpu_inference.utils import get_mesh_shape_product
 
@@ -174,23 +175,13 @@ def dcp_forward(
     dcp_size = mesh.shape[dcp_axis]
 
     # GQA/MQA: replicate KV heads to match ATTN_HEAD sharding before shard_map.
-    tp_size = get_mesh_shape_product(mesh, ShardingAxisName.ATTN_HEAD)
-    if tp_size > 1:
-        num_kv_heads = k.shape[1]
-        if num_kv_heads < tp_size:
-            if tp_size % num_kv_heads != 0:
-                raise ValueError(
-                    f"tp_size {tp_size} must be divisible by num_kv_heads {num_kv_heads}"
-                )
-            factor = tp_size // num_kv_heads
-            k = jnp.repeat(k, factor, axis=1)
-            v = jnp.repeat(v, factor, axis=1)
+    k, v = replicate_kv_heads_for_tp(mesh, k, v)
 
     cp_rank_global = jnp.arange(dcp_size, dtype=jnp.int32)
 
     q_spec = P(ShardingAxisName.ATTN_DATA, ShardingAxisName.ATTN_HEAD, None)
     kv_spec = P(ShardingAxisName.ATTN_DATA, ShardingAxisName.ATTN_HEAD, None)
-    kv_cache_spec = P(ShardingAxisName.BATCH, ShardingAxisName.KV_CONTEXT,
+    kv_cache_spec = P(ShardingAxisName.DENSE_DATA, ShardingAxisName.KV_CONTEXT,
                       ShardingAxisName.KV_HEAD, None, None)
     print(f"page_size={kv_cache.shape[1]}")
 
@@ -291,8 +282,12 @@ def pcp_forward(
       2. current phase         local Q (head+tail) attends all-gathered current KV
       3. merge_attn_states     lse-weighted combine
     """
-    pcp_axis = ShardingAxisName.PREFILL_CONTEXT
+    pcp_axis = ShardingAxisName.PCP
     pcp_size = get_mesh_shape_product(mesh, pcp_axis)
+
+    # GQA/MQA: replicate KV heads to match ATTN_HEAD sharding, mirroring
+    # dcp_forward — tp may exceed the head count (e.g. NKV=2 at tp=4).
+    k, v = replicate_kv_heads_for_tp(mesh, k, v)
     two_p = 2 * pcp_size
     padded_q_len = q.shape[0]
     C = padded_q_len // two_p
@@ -306,7 +301,7 @@ def pcp_forward(
 
     q_spec = P(ShardingAxisName.ATTN_DATA, ShardingAxisName.ATTN_HEAD, None)
     kv_spec = P(ShardingAxisName.ATTN_DATA, ShardingAxisName.KV_HEAD, None)
-    kv_cache_spec = P(ShardingAxisName.BATCH, ShardingAxisName.KV_CONTEXT,
+    kv_cache_spec = P(ShardingAxisName.DENSE_DATA, ShardingAxisName.KV_CONTEXT,
                       ShardingAxisName.KV_HEAD, None, None)
 
     common = dict(sm_scale=sm_scale,