intel / intel/torch-xpu-ops

[xccl][symm_mem] Gap checklist for enabling upstream symmetric-memory tests on XPU

Open
#5,404 0 comments 0 reactions 0 assignees View on GitHub
module: distributed
Dominant language
Python
Stars
113
Forks
128
Avg merge
5d 13h
Merged PRs (30d)
107

Description

> Drafted with assistance from an AI coding agent (Claude Code). Static analysis is against
> torch-xpu-ops `main` (1a18a331) and pytorch `d232af7b2f8`. The "Measured on hardware"
> sections were verified by running the ported tests on 4x Intel Data Center GPU Max 1100
> with the `xccl` backend, 4 ranks.

### What already works

`XPUSymmetricMemory` implements the `SymmetricMemory` base interface and
`XPUSymmetricMemoryAllocator` is registered via
`register_allocator(c10::DeviceType::XPU, allocator)`: `alloc`, `free`,
`get_alloc_size`, `rendezvous`, `get_buffer_ptrs`, `get_signal_pad_ptrs`,
`get_buffer_ptrs_dev`, `get_signal_pad_ptrs_dev`, `get_buffer_size`,
`get_offset`, `barrier`, `put_signal`, `wait_signal`, `get_rank`,
`get_world_size`, `get_device`.

Five Python ops are registered for XPU: `get_remote_tensors`,
`fused_all_gather_matmul`, `fused_all_gather_scaled_matmul`,
`fused_matmul_reduce_scatter`, `fused_scaled_matmul_reduce_scatter`.

**Measured on hardware.** Basic alloc / rendezvous / barrier / signalling do work, and
`fused_all_gather_matmul` works. Two corrections to the expectation that all the fused
async-TP ops work:

- `fused_matmul_reduce_scatter` returns **NaN** on XPU -- see gap 7.5 below.
- `fused_all_gather_scaled_matmul` and `fused_scaled_matmul_reduce_scatter` are
**unverified**: every test covering them skips behind a CUDA compute-capability gate
(`Requires compute capability >= 8.9`, `only supports sm>=90`), so they were never
exercised on XPU.

`has_allocation` is also missing from the allocator (gap 7.1).

### 1. Multicast

- [ ] `XPUSymmetricMemory::has_multicast_support()` returns `false` (`src/xccl/XPUSymmetricMemory.cpp:128`)
- [ ] `XPUSymmetricMemory::get_multicast_ptr()` returns `nullptr` (`:132`)
- [ ] `XPUSymmetricMemoryAllocator::has_multicast_support()` returns `false` (`:453`)

**Measured on hardware.** `_SymmetricMemory.has_multicast_support(DeviceType.XPU, 0)`
returns `False`. This single item accounts for **47** of the 118 skips in
`test_symmetric_memory.py` -- by far the largest blocker by test count.

### 2. Symmetric-memory collective kernels (no XPU registration)

Registered only under `TORCH_LIBRARY_IMPL(symm_mem, CUDA)` in
`CUDASymmetricMemoryOps.cu`; torch-xpu-ops has no `TORCH_LIBRARY_IMPL(symm_mem, XPU)` at all.

- [ ] `one_shot_all_reduce` / `_out` / `_copy` / `_copy_out`
- [ ] `two_shot_all_reduce_` / `_out`
- [ ] `multimem_one_shot_all_reduce`, `multimem_one_shot_reduce_out`
- [ ] `multimem_all_reduce_`, `multimem_all_gather_out`
- [ ] `reduce_scatter_out`
- [ ] `memset32_`, `stream_write_value32_`, `memcpy_to_multicast_`
- [ ] `_async_input_mm`

**Measured on hardware.** Each op below now has a failing test attached, all raising
`NotImplementedError: The operator '' is not currently implemented for the XPU device`:

| Missing op | Failing test |
|---|---|
| `one_shot_all_reduce_copy` | `SymmMemCollectiveTest.test_one_shot_all_reduce` |
| `one_shot_all_reduce_out` | 6 x `LoweringTest`: `test_comm_buffer_reuse`, `test_comm_buffer_inplace_prevention`, `test_lowering_one_shot_all_reduce`, `test_output_buffer_reuse`, `test_external_allocation_fallback`, `test_symm_mem_alloc_id_no_collision_after_reseed` |
| `one_shot_all_reduce_copy_out` | `LoweringTest.test_output_buffer_reuse_copy_variant` |
| `two_shot_all_reduce_` | `SymmMemCollectiveTest.test_two_shot_all_reduce` |
| `reduce_scatter_out` | `SymmMemCollectiveTest.test_reduce_scatter`, `test_reduce_scatter_corner_cases` |
| `memset32_` | `SymmMemSingleProcTest.test_memset32` |
| `stream_write_value32_` | `SymmMemSingleProcTest.test_stream_write_value32` |

### 3. Python-level ops registered for CUDA only

- [ ] `_low_contention_all_gather`
- [ ] `_low_contention_all_gather_ce_multicast` / `_out`
- [ ] `_low_contention_reduce_scatter`

**Measured on hardware.** `_low_contention_all_gather` blocks
`SymmetricMemoryTest.test_low_contention_all_gather` (2 params);
`_low_contention_reduce_scatter` blocks
`SymmetricMemoryTest.test_low_contention_reduce_scatter` (4 params).

### 4. No XPU SHMEM backend

- [ ] `symm_mem.is_nvshmem_available()` has no XPU equivalent; the binding is
`_is_nvshmem_available` -> `c10d::nvshmem_extension::is_nvshmem_available`
(`init.cpp:1188`). torch-xpu-ops contains no `nvshmem`/`rocshmem`-style
SHMEM layer (strict token search over `src/` returns nothing).
- [ ] `_shmem_triton.get_shmem_backend_module()` dispatches to
`_rocshmem_triton` on ROCm else `_nvshmem_triton` - there is no XPU branch,
so `@requires_shmem` Triton kernels cannot target XPU.
- [ ] NVSHMEM op family, all `TORCH_LIBRARY_IMPL(symm_mem, CUDA)` in
`nvshmem_extension.cu`: `nvshmem_broadcast`, `nvshmem_all_to_all`,
`nvshmem_get`, `nvshmem_put`, `nvshmem_put_with_signal`,
`nvshmem_wait_for_signal`, `all_to_all_vdev`, `all_to_all_vdev_2d`,
`all_to_all_vdev_2d_offset`, `tile_reduce`, `multi_root_tile_reduce`

### 5. Backend naming

- [ ] `symm_mem.set_backend()` accepts only `Literal["NVSHMEM", "CUDA", "NCCL"]`
(`_symmetric_memory/__init__.py:2296`). No `"XCCL"` name, so tests pinning a
backend explicitly have no XPU equivalent.

### 6. Token switch

- [ ] `torch.distributed._token_switch` provides only `TokenSwitchNCCL`, and
`_import_nccl_ep()` requires the NCCL expert-parallel package. No XCCL
equivalent class or EP backend.

### 7. Gaps found only by running the tests

These are not visible from op-registration analysis and were each reproduced on hardware.

- [ ] **7.1 `is_symm_mem_tensor()` is always `False` on XPU.**
`XPUSymmetricMemoryAllocator` does not override `has_allocation`, whose base
implementation returns `false`
(`torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp:129`). A tensor
allocated by `symm_mem.empty(128, device="xpu")` reports
`is_symm_mem_tensor(t) == False`. Blocks
`SymmetricMemoryTest.test_is_symm_mem_tensor` and
`LoweringTest.test_custom_op_symm_mem_realization`. Same shape as the
`ProcessGroupXCCL` Backend-virtuals gap: the binding is present, the virtual is
not overridden, and a permissive default hides it.
- [ ] **7.2 No MemPool backend for XPU.** `symm_mem.get_mem_pool(device)` raises
`RuntimeError: SymmetricMemory MemPool did not find backend for device type xpu`
from `get_mempool_allocator` (`_symmetric_memory/__init__.py:2329`). Blocks 5 of
6 `SymmMemPoolTest` tests.
- [ ] **7.3 Invalid arguments abort the process instead of raising.** Five
`SymmMemNegativeTest` tests expect a catchable `RuntimeError` (e.g.
`assertRaisesRegex(RuntimeError, r"must be in \[0")` for an out-of-bounds channel
or rank); on XPU the child dies with SIGABRT (exit `-6`), so the negative paths
cannot be tested at all.
- [ ] **7.4 `wait_signal` honours no timeout.**
`SymmMemNegativeTest.test_wait_signal_timeout` and `SymmetricMemoryTest.test_get`
hang until the harness kills them (300s / 420s) instead of timing out.
- [ ] **7.5 `fused_matmul_reduce_scatter` produces NaN.** 4 failing tests
(`AsyncTPTest.test_fused_matmul_reduce_scatter` x3 params and
`test_fused_matmul_reduce_scatter_bfloat16_custom_reduce`): 512/512 elements
mismatched, `Greatest absolute difference: nan`. Note
`fused_all_gather_matmul` passes, so this is specific to the reduce-scatter
direction.
- [ ] **7.6 `ProcessGroupXCCL.Options` has no `use_pg_for_symm_mem_rendezvous`.** The
options object exposes only `backend`, `enable_reconfigure`,
`global_ranks_in_group`, `group_name`, `is_high_priority_stream`, so PG-based
symmetric-memory rendezvous cannot be requested. Blocks
`SymmetricMemoryTest.test_pg_rendezvous_abort_after`.

`SymmetricMemoryTest.test_rendezvous_via_pg_allgather` is blocked by the XCCL flight
recorder gap in #5381 rather than by anything here: rendezvous records zero
`xccl:_all_gather_base` entries.

### Test impact (measured)

The earlier estimate in this issue counted test *methods*; the numbers below count
parametrized test *cases*, which is what the runner reports, so the totals are larger.

`test_symmetric_memory.py`, after a device-agnostic port: **140 cases -- 22 pass,
118 skip, 0 fail**, with 39 cases gated `@skipIfXpu` referencing this issue.

| Class | Cases | Pass | Skip |
|---|---|---|---|
| `SymmetricMemoryTest` | 30 | 9 | 21 |
| `AsyncTPTest` | 25 | 6 | 19 |
| `SymmMemCollectiveTest` | 49 | 0 | 49 |
| `LoweringTest` | 10 | 1 | 9 |
| `SymmMemPoolTest` | 6 | 1 | 5 |
| `SymmMemEmptySetDeviceTest` | 4 | 4 | 0 |
| `SymmMemNegativeTest` | 6 | 0 | 6 |
| `SymmMemSingleProcTest` | 3 | 1 | 2 |
| `ExternalNcclCommRegistrationTest` | 4 | 0 | 4 |
| `TorchCommsCudaSymmMemTest` | 2 | 0 | 2 |
| `SymmetricMemoryTestCudaGraph` | 1 | 0 | 1 |

Of the 118 skips, only 39 are XPU gaps from this issue. The rest are multicast (47,
gap 1), CUDA compute-capability gates (14), NCCL-only facilities (10) and other
CUDA-specific paths.

A caveat on any earlier XPU results from this file: before the port it reported `OK`
for all 140 cases on XPU **while executing nothing**. `device_type` was hardcoded to
`"cuda"` at module scope, so `world_size` came from `torch.cuda.device_count()` == 0
and no subprocess was ever spawned -- the harness logs `Note: no subprocesses were
spawned, test was likely skipped` and still prints `ok`.

`test_nccl.py`: **37 cases, all gated**. Every one is NCCL-symmetric-memory specific
(`symm_mem.set_backend("NCCL")`, the torchcomms `nccl2` backend, NCCL CFT
logical-endpoint handles), so they are skipped rather than ported. This supersedes the
earlier "25 tests, 15 portable" estimate.

`test_nvshmem.py`, `test_shmem_triton.py` and `test_token_switch.py` are blocked
at module level and cannot be enabled at all until items 4 and 6 land.

`TestNCCL` in `test_nccl.py` (7 tests) is **out of scope**: it exercises the raw
`torch.cuda.nccl` binding (`unique_id`, `init_rank`, `broadcast`, `reduce`,
`all_reduce`, `all_gather`, `reduce_scatter`, `version`) and there is no
`torch.xpu.xccl` module, so it should stay NCCL-gated. It is already
`instantiate_device_type_tests(..., only_for="cuda")`, so it generates no XPU cases.

### Two PyTorch-side CUDA hardcodes surfaced by this port

Not torch-xpu-ops work, recorded here so they are not mistaken for XPU gaps:

- `torch/_inductor/fx_passes/low_contention_collectives.py` falls back to
`torch.cuda.current_device()`, raising `AssertionError: Torch not compiled with CUDA
enabled` from inside an Inductor FX pass on any non-CUDA build.
- `torch/_inductor/codegen/wrapper.py` emits `torch.device("cuda:N")` into the
generated `empty_strided_p2p` call regardless of the buffer's actual device.

**Before closing this ticket, remove the corresponding `@skipIfXpu` lines from the upstream tests.**

Contributor guide

Open the contributing guide

Research direction

Start with test_symmetric_memory.py and the relevant gap in src/xccl/XPUSymmetricMemory.cpp; run the device-agnostic symmetric-memory cases on XPU to reproduce the listed behavior. Choose and scope one checklist item with maintainers, then implement it, add or update the corresponding test coverage, and remove its @skipIfXpu lines once the cases pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend, distributed-systems, testing-qa
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.