[SymmMem] op to get remote tensors by kwen2501 · Pull Request #167779 · pytorch/pytorch

kwen2501 · 2025-11-14T00:12:41Z

Stack from ghstack (oldest at bottom):

-> [SymmMem] op to get remote tensors #167779

To support use case in pytorch/helion#1122, i.e.

@helion.kernel
def foo(
    x: Tensor,
    group_name: str
):
    x_remotes = torch.ops.symm_mem.get_remote_tensors(x, group_name)
    for t in x_remotes:
        ...

Helion uses fake tensor to trace a program, thus we cannot use the following code in a Helion function:

hdl = rendezvous(tensor)
remote_tensors = tuple(
    hdl.get_remote_tensor(peer, ...) for peer in range(world_size)
)

The reason is that when tensor is fake, the returned hdl is None, thus any subsequent call on it will fail.

This PR wraps the above functionality as an op:

lib.define("get_remote_tensors(Tensor x, str group_name) -> Tensor[]")

so that things like hdl is not exposed to Helion. The op also provides a meta implementation so that Helion can trace it without actually running the rendezvous.

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci

[ghstack-poisoned]

pytorch-bot · 2025-11-14T00:12:45Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/167779

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (7 Unrelated Failures)

As of commit 28333cb with merge base 7d7ae4d ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

pull / linux-jammy-py3.10-clang12 / test (default, 4, 5, lf.linux.4xlarge) (gh) (disabled by #156921 but the issue was closed recently and a rebase is needed to make it pass)
dynamo/test_structured_trace.py::StructuredTraceTest::test_compiled_autograd_attribution
pull / linux-jammy-py3.10-clang18-asan / test (default, 5, 7, lf.linux.4xlarge) (gh) (disabled by #156921 but the issue was closed recently and a rebase is needed to make it pass)
dynamo/test_structured_trace.py::StructuredTraceTest::test_compiled_autograd_attribution
pull / linux-jammy-py3.10-gcc11 / test (default, 3, 5, lf.linux.2xlarge) (gh) (disabled by #156921 but the issue was closed recently and a rebase is needed to make it pass)
dynamo/test_structured_trace.py::StructuredTraceTest::test_compiled_autograd_attribution
pull / linux-jammy-py3.10-gcc11 / test (distributed, 1, 2, lf.linux.2xlarge) (gh) (disabled by #123294 but the issue was closed recently and a rebase is needed to make it pass)
distributed/tensor/test_dtensor_ops.py::TestDTensorOpsCPU::test_dtensor_op_db_nn_functional_multi_head_attention_forward_cpu_float32
pull / linux-jammy-py3.13-clang12 / test (default, 2, 5, lf.linux.4xlarge) (gh) (disabled by #156921 but the issue was closed recently and a rebase is needed to make it pass)
dynamo/test_structured_trace.py::StructuredTraceTest::test_compiled_autograd_attribution
trunk / linux-jammy-cuda12.8-py3.10-gcc11 / test (default, 4, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (disabled by #156921 but the issue was closed recently and a rebase is needed to make it pass)
dynamo/test_structured_trace.py::StructuredTraceTest::test_compiled_autograd_attribution
trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh) (disabled by #156921 but the issue was closed recently and a rebase is needed to make it pass)
dynamo/test_structured_trace.py::StructuredTraceTest::test_compiled_autograd_attribution

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ghstack-source-id: 2f964a0 Pull-Request: #167779

pytorch-bot · 2025-11-14T00:12:51Z