@@ -45,6 +45,7 @@ fi
 
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/distributed/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py"
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_symmetric_memory.xml $TE_PATH/tests/pytorch/distributed/test_symmetric_memory.py || test_fail "test_symmetric_memory.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py"
 
@@ -0,0 +1,68 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import os
+import tempfile
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from transformer_engine.pytorch import distributed as te_dist
+
+
+def _check_cache(rank, init_file):
+    torch.cuda.set_device(rank)
+    dist.init_process_group("nccl", init_method=f"file://{init_file}", rank=rank, world_size=2)
+    try:
+        group = dist.group.WORLD
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            # Ordinary all-reduces with varying row counts must not retain one buffer per size.
+            for rows in range(1, 33):
+                tensor = torch.full((rows, 256), rank + 1, device="cuda", dtype=torch.bfloat16)
+                te_dist.symmetric_all_reduce(tensor, group, all_reduce_type="two_shot")
+                assert torch.equal(tensor, torch.full_like(tensor, 3))
+            buffers = list(te_dist.symmetric_mem_cache.values())
+            assert sum(t.numel() for t in buffers) < 8 * tensor.numel()
+
+            # Capture two sizes sharing a capacity, then grow before replaying the old graphs.
+            graphs = []
+            for rows in (23, 24):
+                inp = torch.empty(rows, 256, device="cuda", dtype=torch.bfloat16)
+                out = torch.empty_like(inp)
+                for _ in range(3):
+                    inp.fill_(rank + 1)
+                    te_dist.symmetric_all_reduce(inp, group, all_reduce_type="two_shot")
+                torch.cuda.synchronize()
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph, stream=stream):
+                    inp.fill_(rank + 1)
+                    te_dist.symmetric_all_reduce(inp, group, all_reduce_type="two_shot")
+                    out.copy_(inp)
+                graphs.append((graph, inp, out))
+            large = torch.full((128, 256), rank + 1, device="cuda", dtype=torch.bfloat16)
+            te_dist.symmetric_all_reduce(large, group, all_reduce_type="two_shot")
+            for _ in range(3):
+                for graph, _, out in reversed(graphs):
+                    graph.replay()
+                    torch.cuda.synchronize()
+                    assert torch.equal(out, torch.full_like(out, 3))
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.cuda.synchronize()
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two peer-accessible GPUs")
+def test_symmetric_cache_capacity_and_graph_replay():
+    """Exercise cache growth and graph replay in fresh distributed processes."""
+    if not te_dist.HAS_TORCH_SYMMETRIC:
+        pytest.skip("Requires PyTorch symmetric memory support")
+    if not (torch.cuda.can_device_access_peer(0, 1) and torch.cuda.can_device_access_peer(1, 0)):
+        pytest.skip("Requires peer access between GPUs 0 and 1")
+    with tempfile.TemporaryDirectory() as directory:
+        mp.spawn(_check_cache, args=(os.path.join(directory, "rendezvous"),), nprocs=2, join=True)
@@ -2018,6 +2018,9 @@ def symmetric_all_reduce(
     created symmetric memory tensors not in the cache currently will not be able to
     avoid the extra copies.
 
+    Note: Scratch buffers for uncached inputs are rounded up to a bucketed capacity, so
+    each one may use up to 25% more symmetric memory than the input tensor.
+
     Parameters
     ----------
     inp : torch.Tensor
@@ -2089,8 +2092,14 @@ def symmetric_all_reduce(
             group_name,
         )
     else:
-        # Get symmetric memory tensor. Build or retrieve from cache.
-        msg = get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_group)
+        # Round up to one of four capacities per power of two. This bounds retained
+        # scratch space by the largest request while padding each buffer by at most 25%.
+        # Keep older capacities alive because captured graphs still reference them.
+        numel = max(1, tensor_numel)
+        step = 1 << max(0, numel.bit_length() - 3)
+        capacity = -(-numel // step) * step
+        msg = get_symmetric_memory_tensor(capacity, tensor_dtype, tensor_device, tp_group)
+        msg = msg[:tensor_numel]
 
         msg.copy_(inp.reshape(-1))