diff --git a/test/prototype/pat/test_pat_coupled_sharded_parity.py b/test/prototype/pat/test_pat_coupled_sharded_parity.py
new file mode 100644
index 0000000000..d7dbed011e
--- /dev/null
+++ b/test/prototype/pat/test_pat_coupled_sharded_parity.py
@@ -0,0 +1,335 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from contextlib import ExitStack
+from unittest import mock
+
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, distribute_tensor
+from torch.distributed.tensor.placement_types import Replicate, Shard
+from torch.testing._internal.common_utils import run_tests
+
+from test.prototype.pat.test_pat_subset_mesh import _GlooMultiProcessTestCase
+from torchao.prototype.pat.group import (
+    AttentionHeadGrouperDim0,
+    AttentionHeadGrouperDim1,
+    Dim0Grouper,
+    Dim1Grouper,
+)
+from torchao.prototype.pat.optim import CoupledMinSparsityConstraint
+from torchao.prototype.pat.optim import prox_executor as executor
+
+
+def _channel_view(p, grouper_cls, kwargs):
+    if grouper_cls in (Dim1Grouper, AttentionHeadGrouperDim1):
+        channels = kwargs.get("num_heads", p.size(1))
+        return p.reshape(p.size(0), channels, -1).permute(1, 0, 2).reshape(channels, -1)
+    return p.reshape(kwargs.get("num_heads", p.size(0)), -1)
+
+
+def _fixture(groupers):
+    # Different tensors dominate different channels, with large score gaps.
+    amplitudes = 10.0 ** torch.tensor([3, 0, 6, 1, 4, 2, 7, 5])
+    specs = []
+    for i, cls in enumerate(groupers):
+        width = 5 + 2 * i
+        scale = torch.where(torch.arange(8) % len(groupers) == i, 1.0, 0.001)
+        view = (amplitudes * scale)[:, None] * torch.arange(1.0, width + 1)[None, :]
+        view[-1, -1] = 0.0  # An isolated zero outside the usual selected channels.
+        p = view.t().contiguous() if cls is Dim1Grouper else view.clone()
+        specs.append((p, cls, {}))
+    return specs
+
+
+def _dense_reference(specs, budget):
+    views = [_channel_view(p, cls, kwargs) for p, cls, kwargs in specs]
+    squared_scores = sum(view.float().square().sum(1) for view in views)
+    chosen = math.ceil(budget * views[0].size(0))
+    indices = squared_scores.argsort()[:chosen]
+    expected = [(p.clone(), cls, kwargs) for p, cls, kwargs in specs]
+    for p, cls, kwargs in expected:
+        axis = int(cls in (Dim1Grouper, AttentionHeadGrouperDim1))
+        width = p.size(axis) // views[0].size(0)
+        packed_indices = (indices[:, None] * width + torch.arange(width)).flatten()
+        p.index_fill_(axis, packed_indices, 0.0)
+    return expected, chosen
+
+
+class _CoupledParityMixin:
+    def _distribute(self, originals, mesh, placements):
+        return [
+            (
+                distribute_tensor(p.clone(), mesh, placements, src_data_rank=None),
+                cls,
+                kwargs,
+            )
+            for p, cls, kwargs in originals
+        ]
+
+    def _apply(self, specs, budget, score_type, route):
+        prox = CoupledMinSparsityConstraint(0.0, budget, score_type)
+        with ExitStack() as stack:
+            if route == "fast":
+                self.assertIsNotNone(executor._coupled_fast_path(specs))
+                sharded = stack.enter_context(
+                    mock.patch.object(
+                        executor,
+                        "_coupled_prox_sharded",
+                        wraps=executor._coupled_prox_sharded,
+                    )
+                )
+                stack.enter_context(
+                    mock.patch.object(
+                        executor,
+                        "grouped_view",
+                        side_effect=AssertionError("materialized grouped view"),
+                    )
+                )
+                stack.enter_context(
+                    mock.patch.object(
+                        DTensor,
+                        "full_tensor",
+                        side_effect=AssertionError("materialized DTensor"),
+                    )
+                )
+                collectives = []
+                for name in (
+                    "all_reduce",
+                    "all_gather_into_tensor",
+                    "all_gather",
+                    "broadcast",
+                    "reduce",
+                    "reduce_scatter_tensor",
+                    "all_to_all_single",
+                ):
+                    collectives.append(
+                        stack.enter_context(
+                            mock.patch.object(dist, name, wraps=getattr(dist, name))
+                        )
+                    )
+            else:
+                if route == "forced":
+                    stack.enter_context(
+                        mock.patch.object(
+                            executor, "_coupled_fast_path", return_value=None
+                        )
+                    )
+                else:
+                    self.assertIsNone(executor._coupled_fast_path(specs))
+                stack.enter_context(
+                    mock.patch.object(
+                        executor,
+                        "_coupled_prox_sharded",
+                        side_effect=AssertionError("unexpected fast path"),
+                    )
+                )
+                grouped = stack.enter_context(
+                    mock.patch.object(
+                        executor, "grouped_view", wraps=executor.grouped_view
+                    )
+                )
+                redistributed = stack.enter_context(
+                    mock.patch.object(
+                        executor, "distribute_tensor", wraps=executor.distribute_tensor
+                    )
+                )
+            result = executor.apply_coupled_prox(specs, prox, budget)
+            if route == "fast":
+                self.assertEqual(sharded.call_count, 1)
+                self.assertGreater(sum(call.call_count for call in collectives), 0)
+                self.assertLessEqual(sum(call.call_count for call in collectives), 3)
+                shard_group = executor._coupled_fast_path(specs)[0]
+                for collective in collectives:
+                    for call in collective.call_args_list:
+                        self.assertIs(call.kwargs["group"], shard_group)
+            else:
+                self.assertEqual(grouped.call_count, len(specs))
+                self.assertEqual(redistributed.call_count, len(specs))
+                for call in redistributed.call_args_list:
+                    self.assertIn("src_data_rank", call.kwargs)
+                    self.assertIsNone(call.kwargs["src_data_rank"])
+        return result
+
+    def _check_result(self, result, specs, expected, chosen):
+        self.assertEqual(result.zero_channels, chosen)
+        self.assertEqual(result.channels, 8)
+        self.assertEqual(len(result.parameters), len(specs))
+        zeros = []
+        sizes = []
+        for entry, (p, _, _), (dense, cls, kwargs) in zip(
+            result.parameters, specs, expected
+        ):
+            full = p.full_tensor()
+            self.assertEqual(full.eq(0), dense.eq(0))
+            self.assertEqual(full, dense, rtol=0, atol=0)
+            self.assertIs(entry.parameter, p)
+            zeros.append(int(dense.eq(0).sum()))
+            sizes.append(dense.numel())
+            self.assertEqual(entry.zero_elts, zeros[-1])
+            self.assertEqual(entry.numel, sizes[-1])
+            self.assertEqual(
+                int(_channel_view(full, cls, kwargs).eq(0).all(1).sum()), chosen
+            )
+        self.assertEqual(result.zero_elts, sum(zeros))
+        self.assertEqual(result.numel, sum(sizes))
+
+    def _parity(self, originals, mesh, placements, route="fast", budgets=(0.26,)):
+        for score_type in ("rms", "l2", "param_cost"):
+            for budget in budgets:
+                with self.subTest(
+                    score_type=score_type, budget=budget, placements=placements
+                ):
+                    expected, chosen = _dense_reference(originals, budget)
+                    actual = self._distribute(originals, mesh, placements)
+                    forced = self._distribute(originals, mesh, placements)
+                    a = self._apply(actual, budget, score_type, route)
+                    b = self._apply(forced, budget, score_type, "forced")
+                    self._check_result(a, actual, expected, chosen)
+                    self._check_result(b, forced, expected, chosen)
+                    self.assertEqual(
+                        (a.zero_elts, a.numel, a.zero_channels, a.channels),
+                        (b.zero_elts, b.numel, b.zero_channels, b.channels),
+                    )
+
+
+class TestPATCoupledShardedParity(_CoupledParityMixin, _GlooMultiProcessTestCase):
+    @property
+    def world_size(self):
+        return 4
+
+    def test_mixed_readers_writers_and_replicated_axes(self):
+        self._init_process_group()
+        try:
+            originals = _fixture(
+                (Dim1Grouper, Dim0Grouper, Dim1Grouper, Dim0Grouper, Dim0Grouper)
+            )
+            mesh = DeviceMesh("cpu", [0, 1, 2, 3])
+            self._parity(originals, mesh, (Shard(0),), budgets=(0.0, 0.26, 1.0))
+            double_specs = [(p.double(), cls, kwargs) for p, cls, kwargs in originals]
+            self._parity(double_specs, mesh, (Shard(0),))
+            mesh = DeviceMesh("cpu", [[0, 1], [2, 3]])
+            for placements in ((Shard(0), Replicate()), (Replicate(), Shard(0))):
+                self._parity(originals, mesh, placements)
+        finally:
+            dist.destroy_process_group()
+
+    def test_readers_only_uneven_rows_and_writers_only(self):
+        self._init_process_group()
+        try:
+            mesh = DeviceMesh("cpu", [0, 1, 2, 3])
+            for groupers in (
+                (Dim1Grouper,) * 3,
+                (Dim0Grouper,) * 3,
+                (Dim1Grouper, Dim0Grouper),
+            ):
+                self._parity(_fixture(groupers), mesh, (Shard(0),))
+        finally:
+            dist.destroy_process_group()
+
+    def test_unsupported_layouts_and_groupers_materialize(self):
+        self._init_process_group()
+        try:
+            originals = _fixture((Dim1Grouper, Dim0Grouper))
+            mesh = DeviceMesh("cpu", [0, 1, 2, 3])
+            self._parity(originals, mesh, (Shard(1),), "fallback")
+            kwargs_specs = [
+                (p, cls, {"start_dim": 0, "end_dim": 1}) for p, cls, _ in originals
+            ]
+            self._parity(kwargs_specs, mesh, (Shard(0),), "fallback")
+            head_specs = [
+                (
+                    p.repeat_interleave(2, dim=1 if cls is Dim1Grouper else 0),
+                    AttentionHeadGrouperDim1
+                    if cls is Dim1Grouper
+                    else AttentionHeadGrouperDim0,
+                    {"num_heads": 8},
+                )
+                for p, cls, _ in originals
+            ]
+            self._parity(head_specs, mesh, (Shard(0),), "fallback")
+            mesh = DeviceMesh("cpu", [[0, 1], [2, 3]])
+            self._parity(originals, mesh, (Shard(0), Shard(0)), "fallback")
+        finally:
+            dist.destroy_process_group()
+
+    def test_tied_and_near_tied_selection_agrees_on_all_ranks(self):
+        self._init_process_group()
+        try:
+            mesh = DeviceMesh("cpu", [[0, 1], [2, 3]])
+            for placements in ((Shard(0), Replicate()), (Replicate(), Shard(0))):
+                for near_tied in (False, True):
+                    for route in ("fast", "forced"):
+                        with self.subTest(
+                            placements=placements, near_tied=near_tied, route=route
+                        ):
+                            scales = torch.ones(8)
+                            if near_tied:
+                                scales += (
+                                    torch.arange(8) * torch.finfo(torch.float32).eps
+                                )
+                            originals = [
+                                (scales.repeat(8, 1), Dim1Grouper, {}),
+                                (scales[:, None].repeat(1, 4), Dim0Grouper, {}),
+                            ]
+                            specs = self._distribute(originals, mesh, placements)
+                            result = self._apply(specs, 0.5, "rms", route)
+                            mask = specs[0][0].to_local().eq(0).all(0)
+                            self.assertEqual(int(mask.sum()), 4)
+                            shard_dim = next(
+                                i
+                                for i, p in enumerate(placements)
+                                if isinstance(p, Shard)
+                            )
+                            shard_rank = mesh.get_coordinate()[shard_dim]
+                            self.assertEqual(
+                                specs[1][0].to_local().eq(0).all(1),
+                                mask.chunk(2)[shard_rank],
+                            )
+                            record = (
+                                mask.tolist(),
+                                result.zero_channels,
+                                result.zero_elts,
+                                result.numel,
+                            )
+                            records = [None] * self.world_size
+                            dist.all_gather_object(records, record)
+                            self.assertEqual(records, [record] * self.world_size)
+                            # Dense tie-breaking need not match; shared masks must.
+                            expected = [
+                                (p.clone(), cls, kwargs) for p, cls, kwargs in originals
+                            ]
+                            expected[0][0][:, mask] = 0.0
+                            expected[1][0][mask] = 0.0
+                            self._check_result(result, specs, expected, 4)
+        finally:
+            dist.destroy_process_group()
+
+
+class TestPATCoupledUnevenShards(_CoupledParityMixin, _GlooMultiProcessTestCase):
+    @property
+    def world_size(self):
+        return 3
+
+    def test_uneven_writers_fall_back_but_uneven_readers_are_fast(self):
+        self._init_process_group()
+        try:
+            mesh = DeviceMesh("cpu", [0, 1, 2])
+            self._parity(
+                _fixture((Dim1Grouper, Dim0Grouper, Dim0Grouper)),
+                mesh,
+                (Shard(0),),
+                "fallback",
+            )
+            self._parity(_fixture((Dim1Grouper,) * 3), mesh, (Shard(0),))
+        finally:
+            dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/prototype/pat/test_pat_subset_mesh.py b/test/prototype/pat/test_pat_subset_mesh.py
index 0f9ecf3b91..bef79fbb01 100644
--- a/test/prototype/pat/test_pat_subset_mesh.py
+++ b/test/prototype/pat/test_pat_subset_mesh.py
@@ -54,7 +54,7 @@ def _step(optimizer, params):
             param.grad = torch.zeros_like(param)
         optimizer.step()
 
-    def test_coupled_subset_mesh_materialized_state(self):
+    def test_coupled_subset_mesh_fast_and_materialized_state(self):
         self._init_process_group()
         try:
             mesh = DeviceMesh("cpu", [1, 2])
@@ -72,95 +72,130 @@ def test_coupled_subset_mesh_materialized_state(self):
             counts = [int(p.eq(0).sum()) for p in expected]
             total = sum(p.numel() for p in expected)
             groupers = (Dim1Grouper, Dim0Grouper)
-            direct = [
-                distribute_tensor(p.clone(), mesh, [Shard(0)], src_data_rank=None)
-                for p in originals
-            ]
-            params = [
-                torch.nn.Parameter(
-                    distribute_tensor(p.clone(), mesh, [Shard(0)], src_data_rank=None)
-                )
-                for p in originals
-            ]
-            groups = [
-                {
-                    "params": [p],
-                    "group_type": cls.__name__,
-                    "prox_type": "CoupledMinSparsityConstraint",
-                    "couple_key": "residual",
-                    "min_sparsity": 0.5,
-                }
-                for p, cls in zip(params, groupers)
-            ]
-            optimizer = PruneOptimizer(torch.optim.SGD(groups, lr=0.0))
-            optimizer.relative_sparsity = 0.375
-            optimizer.relative_factored_frac = 0.25
-            with mock.patch.object(
-                executor, "distribute_tensor", wraps=executor.distribute_tensor
-            ) as redistributed:
-                result = executor.apply_coupled_prox(
-                    [(p, cls, {}) for p, cls in zip(direct, groupers)],
-                    CoupledMinSparsityConstraint(0.0, 0.5),
-                    0.5,
-                )
-                for _ in range(2):
-                    self._step(optimizer, params)
-                    if participant:
+            for forced in (False, True):
+                with self.subTest(forced=forced):
+                    direct = [
+                        distribute_tensor(
+                            p.clone(), mesh, [Shard(0)], src_data_rank=None
+                        )
+                        for p in originals
+                    ]
+                    params = [
+                        torch.nn.Parameter(
+                            distribute_tensor(
+                                p.clone(), mesh, [Shard(0)], src_data_rank=None
+                            )
+                        )
+                        for p in originals
+                    ]
+                    groups = [
+                        {
+                            "params": [p],
+                            "group_type": cls.__name__,
+                            "prox_type": "CoupledMinSparsityConstraint",
+                            "couple_key": "residual",
+                            "min_sparsity": 0.5,
+                        }
+                        for p, cls in zip(params, groupers)
+                    ]
+                    optimizer = PruneOptimizer(torch.optim.SGD(groups, lr=0.0))
+                    optimizer.relative_sparsity = 0.375
+                    optimizer.relative_factored_frac = 0.25
+                    route_patch = (
+                        mock.patch.object(
+                            executor, "_coupled_fast_path", return_value=None
+                        )
+                        if forced
+                        else mock.patch.object(
+                            executor,
+                            "_coupled_prox_sharded",
+                            wraps=executor._coupled_prox_sharded,
+                        )
+                    )
+                    with (
+                        route_patch as route,
+                        mock.patch.object(
+                            executor,
+                            "distribute_tensor",
+                            wraps=executor.distribute_tensor,
+                        ) as redistributed,
+                    ):
+                        result = executor.apply_coupled_prox(
+                            [(p, cls, {}) for p, cls in zip(direct, groupers)],
+                            CoupledMinSparsityConstraint(0.0, 0.5),
+                            0.5,
+                        )
+                        for _ in range(2):
+                            self._step(optimizer, params)
+                            if participant:
+                                self.assertEqual(
+                                    optimizer.relative_sparsity, sum(counts) / total
+                                )
+                            else:
+                                self.assertEqual(optimizer.relative_sparsity, 0.375)
+                                self.assertEqual(optimizer.relative_factored_frac, 0.25)
+                                for p in params:
+                                    self.assertNotIn(
+                                        "sparsity_frac", optimizer.state[p]
+                                    )
+                    if forced:
                         self.assertEqual(
-                            optimizer.relative_sparsity, sum(counts) / total
+                            redistributed.call_count, 6 if participant else 0
                         )
+                        for call in redistributed.call_args_list:
+                            self.assertIn("src_data_rank", call.kwargs)
+                            self.assertIsNone(call.kwargs["src_data_rank"])
                     else:
-                        self.assertEqual(optimizer.relative_sparsity, 0.375)
-                        self.assertEqual(optimizer.relative_factored_frac, 0.25)
-                        for p in params:
-                            self.assertNotIn("sparsity_frac", optimizer.state[p])
-            self.assertEqual(redistributed.call_count, 6 if participant else 0)
-            for call in redistributed.call_args_list:
-                self.assertIn("src_data_rank", call.kwargs)
-                self.assertIsNone(call.kwargs["src_data_rank"])
-            if participant:
-                self.assertEqual(result.zero_channels, 4)
-                self.assertEqual(result.channels, 8)
-                self.assertEqual(result.zero_elts, sum(counts))
-                self.assertEqual(result.numel, total)
-                self.assertEqual(len(result.parameters), 2)
-                for i, (p, q, dense) in enumerate(zip(direct, params, expected)):
-                    self.assertIs(result.parameters[i].parameter, p)
-                    self.assertEqual(result.parameters[i].zero_elts, counts[i])
-                    self.assertEqual(result.parameters[i].numel, dense.numel())
-                    self.assertEqual(p.full_tensor(), dense, rtol=0, atol=0)
-                    self.assertEqual(q.full_tensor(), dense, rtol=0, atol=0)
+                        self.assertEqual(route.call_count, 3 if participant else 0)
+                        self.assertEqual(redistributed.call_count, 0)
+                    if participant:
+                        self.assertEqual(result.zero_channels, 4)
+                        self.assertEqual(result.channels, 8)
+                        self.assertEqual(result.zero_elts, sum(counts))
+                        self.assertEqual(result.numel, total)
+                        self.assertEqual(len(result.parameters), 2)
+                        for i, (p, q, dense) in enumerate(
+                            zip(direct, params, expected)
+                        ):
+                            self.assertIs(result.parameters[i].parameter, p)
+                            self.assertEqual(result.parameters[i].zero_elts, counts[i])
+                            self.assertEqual(result.parameters[i].numel, dense.numel())
+                            self.assertEqual(p.full_tensor(), dense, rtol=0, atol=0)
+                            self.assertEqual(q.full_tensor(), dense, rtol=0, atol=0)
+                            self.assertEqual(
+                                optimizer.state[q]["sparsity_frac"],
+                                counts[i] / dense.numel(),
+                            )
+                        self.assertEqual(optimizer.relative_factored_frac, 0.0)
+                        record = (
+                            True,
+                            result.zero_channels,
+                            result.zero_elts,
+                            optimizer.relative_sparsity,
+                        )
+                    else:
+                        self.assertEqual(result.parameters, ())
+                        self.assertEqual(
+                            (
+                                result.zero_elts,
+                                result.numel,
+                                result.zero_channels,
+                                result.channels,
+                            ),
+                            (0, 0, 0, 0),
+                        )
+                        record = (
+                            False,
+                            optimizer.relative_sparsity,
+                            optimizer.relative_factored_frac,
+                        )
+                    records = [None] * self.world_size
+                    dist.all_gather_object(records, record)
+                    self.assertEqual(records[0], (False, 0.375, 0.25))
+                    self.assertEqual(records[1], records[2])
                     self.assertEqual(
-                        optimizer.state[q]["sparsity_frac"], counts[i] / dense.numel()
+                        records[1], (True, 4, sum(counts), sum(counts) / total)
                     )
-                self.assertEqual(optimizer.relative_factored_frac, 0.0)
-                record = (
-                    True,
-                    result.zero_channels,
-                    result.zero_elts,
-                    optimizer.relative_sparsity,
-                )
-            else:
-                self.assertEqual(result.parameters, ())
-                self.assertEqual(
-                    (
-                        result.zero_elts,
-                        result.numel,
-                        result.zero_channels,
-                        result.channels,
-                    ),
-                    (0, 0, 0, 0),
-                )
-                record = (
-                    False,
-                    optimizer.relative_sparsity,
-                    optimizer.relative_factored_frac,
-                )
-            records = [None] * self.world_size
-            dist.all_gather_object(records, record)
-            self.assertEqual(records[0], (False, 0.375, 0.25))
-            self.assertEqual(records[1], records[2])
-            self.assertEqual(records[1], (True, 4, sum(counts), sum(counts) / total))
         finally:
             if dist.is_initialized():
                 dist.destroy_process_group()
diff --git a/torchao/prototype/pat/README.md b/torchao/prototype/pat/README.md
index 755e4e4beb..589331c34d 100644
--- a/torchao/prototype/pat/README.md
+++ b/torchao/prototype/pat/README.md
@@ -137,7 +137,7 @@ An existing non-SVD prox can add a nested `coupled` dictionary containing `coupl
 
 The caller must enumerate every relevant tensor and its channel axis. Coupling does not discover residual dependencies, handle normalization/bias dependencies automatically, resize modules, or produce a compact inference model.
 
-For DTensors, coupled selection materializes every clustered tensor on each mesh participant, selects one shared mask, and redistributes locally with `src_data_rank=None`. Peak dense memory is proportional to the sum of clustered tensors. Subset-mesh nonparticipants skip processing, and replica dimensions do not duplicate counts. CPU Gloo tests cover the materializing reference on a subset excluding rank zero, including uneven row sharding and cached nonparticipant metrics. These correctness tests do not establish NCCL/FSDP2 throughput or model-training performance.
+For 2-D DTensors with exact default `Dim0Grouper`/`Dim1Grouper` layouts, one `Shard(0)` mesh axis, replication elsewhere, and evenly split writers, coupling uses local masking and at most three mesh-scoped collectives per cluster, independent of tensor count. Other layouts fall back to full materialization; peak dense memory is then proportional to the sum of clustered tensors. Subset-mesh nonparticipants skip processing, and replica dimensions do not duplicate counts. CPU Gloo tests cover four-rank 1-D and 2-D meshes, both replicated/sharded axis orders, three-rank uneven fallback, and a subset excluding rank zero. Reduction order can perturb tied or near-tied selection across layouts; exact cross-layout mask parity is only asserted for well-separated scores. These correctness tests do not establish NCCL/FSDP2 throughput or model-training performance.
 
 SVD decompositions can dominate optimizer cost on wide tensors. Set `prox_freq: N` on a parameter group to run its grouper and proximal map every `N` optimizer steps. SVD groups using the hard `MinRankConstraint` reapply the proximal map during healing by default so the base optimizer cannot refill removed singular values. Soft SVD maps such as `ProxNuclearNorm` retain their existing behavior unless `prox_through_heal: true` is set explicitly, while `MinRankConstraint` can opt out with `prox_through_heal: false`. Setting `prox_through_heal: true` on a non-SVD grouper is invalid and rejected during optimizer construction.
 
diff --git a/torchao/prototype/pat/optim/prox_executor.py b/torchao/prototype/pat/optim/prox_executor.py
index fb9f63dc0d..76e9ae7a84 100644
--- a/torchao/prototype/pat/optim/prox_executor.py
+++ b/torchao/prototype/pat/optim/prox_executor.py
@@ -11,12 +11,14 @@
 from typing import Any, Sequence
 
 import torch
+import torch.distributed as dist
 from torch import Tensor
 from torch.distributed.tensor import distribute_tensor
 from torch.distributed.tensor.experimental import local_map
 from torch.distributed.tensor.placement_types import Partial, Shard
 
 from ..distributed_utils import _is_dtensor
+from ..group.dim import Dim0Grouper, Dim1Grouper
 from ..group.grouper import ElemGrouper, LayerGrouper
 from ..utils import get_index_linspace
 from .group_lasso import ProxGroupLasso, ProxGroupLassoVectorized
@@ -357,6 +359,147 @@ def _validate_coupled_specs(specs) -> bool:
     return True
 
 
+def _shard_mesh_dim(p: Tensor) -> int | None:
+    """Find a single plain Shard(0) mesh axis, with replication elsewhere."""
+    found = None
+    for mesh_dim, placement in enumerate(p.placements):
+        if type(placement) is Shard:
+            if placement.dim != 0 or found is not None:
+                return None
+            found = mesh_dim
+        elif not placement.is_replicate():
+            return None
+    return found
+
+
+def _coupled_fast_path(specs) -> tuple[Any, int, int] | None:
+    """Find eligible mesh-local reader reductions and writer gathers.
+
+    Readers (Dim1) are row-sharded and hold partial norms for every channel.
+    Writers (Dim0) own complete channels, requiring evenly sized row blocks.
+    Reshaped groupers and other layouts use the materializing reference path.
+    """
+    if not specs:
+        return None
+    mesh = None
+    mesh_dim = None
+    n_channels = None
+    defaults = {"start_dim": 1, "end_dim": -1}
+    for p, grouper_cls, grouper_kwargs in specs:
+        if grouper_cls not in (Dim0Grouper, Dim1Grouper):
+            return None
+        if any(
+            key not in defaults or value != defaults[key]
+            for key, value in grouper_kwargs.items()
+        ):
+            return None
+        if not _is_dtensor(p) or p.dim() != 2:
+            return None
+        p_mesh_dim = _shard_mesh_dim(p)
+        if p_mesh_dim is None:
+            return None
+        if mesh is None:
+            mesh, mesh_dim = p.device_mesh, p_mesh_dim
+        elif p.device_mesh != mesh or p_mesh_dim != mesh_dim:
+            return None
+        channels = p.size(1) if grouper_cls is Dim1Grouper else p.size(0)
+        if n_channels is None:
+            n_channels = channels
+        elif channels != n_channels:
+            return None
+    n_shard = mesh.size(mesh_dim)
+    coordinate = mesh.get_coordinate()
+    if coordinate is None:
+        return None
+    if any(cls is Dim0Grouper and p.size(0) % n_shard for p, cls, _ in specs):
+        return None
+    return mesh.get_group(mesh_dim), coordinate[mesh_dim], n_shard
+
+
+def _coupled_prox_sharded(
+    specs, prox_map, min_sparsity, shard_group, shard_rank, n_shard
+) -> CoupledProxResult:
+    """Apply a coupled mask with at most three mesh-scoped collectives.
+
+    Aggregate all readers' partials together, gather all writers' channel
+    blocks together, then reduce one vector of literal-zero counts. Replica
+    mesh dimensions never participate in these collectives.
+    """
+    reader_sq = None
+    writer_blocks = []
+    group_size_sum = 0
+    for p, cls, _ in specs:
+        local = p.to_local()
+        if cls is Dim1Grouper:
+            part = local.float().square().sum(dim=0)
+            reader_sq = part if reader_sq is None else reader_sq + part
+            group_size_sum += p.size(0)
+        else:
+            writer_blocks.append(local.float().square().sum(dim=1))
+            group_size_sum += p.size(1)
+
+    if reader_sq is not None:
+        dist.all_reduce(reader_sq, group=shard_group)
+    writer_sq = None
+    if writer_blocks:
+        stacked = torch.stack(writer_blocks).contiguous()
+        n_writers, local_channels = stacked.shape
+        gathered = torch.empty(
+            n_shard * n_writers,
+            local_channels,
+            dtype=stacked.dtype,
+            device=stacked.device,
+        )
+        dist.all_gather_into_tensor(gathered, stacked, group=shard_group)
+        writer_sq = (
+            gathered.view(n_shard, n_writers, local_channels)
+            .permute(1, 0, 2)
+            .reshape(n_writers, n_shard * local_channels)
+            .sum(dim=0)
+        )
+    sq_norm_sum = (
+        reader_sq
+        if writer_sq is None
+        else (writer_sq if reader_sq is None else reader_sq + writer_sq)
+    )
+    scores = prox_map.combine_scores(sq_norm_sum, group_size_sum)
+    n_channels = scores.numel()
+    n_zero = math.ceil(min_sparsity * n_channels)
+    if n_zero <= 0:
+        zero_idx = torch.empty(0, dtype=torch.long, device=scores.device)
+    elif n_zero >= n_channels:
+        zero_idx = torch.arange(n_channels, device=scores.device)
+    else:
+        _, zero_idx = torch.topk(scores, k=n_zero, largest=False, sorted=False)
+
+    zero_counts = []
+    for p, cls, _ in specs:
+        local = p.to_local()
+        if cls is Dim1Grouper:
+            local[:, zero_idx] = 0.0
+        else:
+            offset = shard_rank * local.size(0)
+            local_idx = (
+                zero_idx[(zero_idx >= offset) & (zero_idx < offset + local.size(0))]
+                - offset
+            )
+            local[local_idx] = 0.0
+        zero_counts.append((local == 0).sum())
+    counts = torch.stack(zero_counts).to(torch.long)
+    dist.all_reduce(counts, group=shard_group)
+    parameters = tuple(
+        ParameterSparsity(parameter=p, zero_elts=zeros, numel=p.numel())
+        for (p, _, _), zeros in zip(specs, counts.tolist())
+    )
+    return CoupledProxResult(
+        parameters=parameters,
+        zero_elts=sum(p.zero_elts for p in parameters),
+        numel=sum(p.numel for p in parameters),
+        zero_channels=zero_idx.numel(),
+        channels=n_channels,
+    )
+
+
 def apply_coupled_prox(
     specs: Sequence[tuple[Tensor, Any, dict[str, Any]]],
     prox_map,
@@ -373,9 +516,10 @@ def apply_coupled_prox(
     Reported element counts include all final literal zeros, including zeros
     from a prior prox; callers must account each participating parameter once.
 
-    DTensor execution caches every grouped full tensor, so peak dense memory
-    is the sum of clustered tensors. Every mesh participant selects from the
-    same materialized scores; nonparticipants skip processing.
+    Eligible Shard(0) layouts avoid full tensors. The fallback caches every
+    grouped full tensor, so peak dense memory is the sum of clustered tensors.
+    Floating-point reduction order can affect selection at tied or near-tied
+    score boundaries across layouts; selection is consistent within a mesh.
     """
     if (
         materialization_policy
@@ -393,6 +537,9 @@ def apply_coupled_prox(
         return empty
     if not _validate_coupled_specs(specs):
         return empty
+    fast_path = _coupled_fast_path(specs) if dist.is_initialized() else None
+    if fast_path is not None:
+        return _coupled_prox_sharded(specs, prox_map, min_sparsity, *fast_path)
 
     with contextlib.ExitStack() as stack:
         entries = []