Repository navigation
[release/2.15] [ROCm] Enable native AsyncTP (#177961, #200054) - #3868
Merged
Merged
Conversation
This PR enables the native AsyncTP path of `fused_all_gather_matmul` on ROCm (gfx942 / MI300X and gfx950 / MI355X). The native path overlaps the all-gather of the activation shard `A` with `A @ B`, using a persistent GEMM that waits on a per-chunk signal before reading each gathered chunk of `A`. On CUDA this GEMM is a CUTLASS kernel; this PR adds the ROCm equivalent using the Composable Kernel (CK) tile `PersistentAsyncInputScheduler`, which implements the same signal protocol.
Before this PR, `_async_input_mm` had no ROCm implementation and the native test was skipped on ROCm. The CUDA code path is unchanged.
A hand-written HIP file (not hipified) that provides `c10d::cuda::detail::async_input_mm{,_out}` on ROCm. The existing `_async_input_mm` registration in `CUDASymmetricMemoryOps.cu` binds to it.
- CK-tile persistent universal GEMM, bf16 in/out with fp32 accumulation. Block tile 256x256x32, 2x2x1 warps, 32x32x16 warp tile, compute-V3 pipeline. A `static_assert` checks the kernel is persistent.
- `DoubleSmemBuffer = true` does not double-buffer LDS: CompV3's main loop uses one LDS buffer either way. The flag only makes the CShuffle epilogue use smaller tiles. Turning it off made the GEMM alone 14-32% slower on MI300X and about 1.5% faster on MI355X, so it stays on.
- `kPadN` / `kPadK` handle N and K that are not tile multiples. M is not padded: a tile straddling two chunks could read a chunk before its signal is set, so the chunk size must be a multiple of 256 rows (checked).
- B is made row-major (`b.contiguous()`) before launch. With the current CK pin the column-major-B config is 1.6x (gfx942) to 6x (gfx950) slower, far more than one copy of B.
- The epilogue rounds fp32 to bf16 to nearest even (`CK_TILE_FLOAT_TO_BFLOAT16_STANDARD_ASM`), as hipBLASLt and the CUDA kernel do. ck_tile truncates by default on gfx942, and PyTorch does not override it: in the native test's GEMM on MI300X about half of the 4.19M outputs were one bf16 step from the correctly rounded value, against at most 588 for hipBLASLt. gfx950 converts with `v_cvt_pk_bf16_f32`, which rounds to nearest even whatever the setting, so its code is unchanged. The C++ `STANDARD` mode compiles to branches and made the GEMM 3-32% slower on MI300X; the asm mode costs 2-6% at K=512 and up to 3% at K=1024, and is within noise from K=2048 (M=4096, N=1024-8192, the kernel alone on one GPU).
- Persistent grid: blocks spin until their chunk is signaled, and HIP runs the signal writes (`stream_write_value32`, `memset32`) and peer copies of at most 1 MiB as kernels, so the grid leaves one CU free on every shader engine (4 per XCC) or the GEMM can starve them. The free CU is taken from the per-engine CU count rounded down: 256 blocks on MI300X (38 CUs per XCC) and 224 on MI355X. Reserving one CU per engine from the total instead gave 272 blocks on MI300X, and any GEMM with at least 272 output tiles then deadlocked (e.g. M=4096 with N>=4352 at 2 and 4 ranks).
- Tiles are visited in M order (`TilePartitionerGroupNum = 1`), the order the chunks arrive in; the CUDA scheduler also works through the chunks in order. With CK's default of 8 groups, only the blocks of the last chunk's groups worked on it once it arrived. When the GEMM has more tiles than the grid, the GEMM takes 4-33% less time; with fewer tiles than the grid it is unchanged.
- Only gfx942 / gfx950 are supported (checked at runtime via `gcnArchName`), and tensors must have at most `INT32_MAX` elements because CK indexes with 32-bit integers. The other input checks are the CUDA kernel's.
On ROCm, `hip/AsyncMM.hip` replaces `cuda/AsyncMM.cu` in the HIP sources and always provides the `async_input_mm` symbol the op registration needs. The kernel itself (`USE_ROCM_ASYNC_MM`) is compiled only when `PYTORCH_ROCM_ARCH` contains gfx942 or gfx950 and the CK headers are on the include path, i.e. `(USE_FLASH_ATTENTION AND USE_ROCM_CK_SDPA) OR USE_ROCM_CK_GEMM`. Without it, `_async_input_mm` raises, as the CUDA op does in builds without its kernel (e.g. on Windows).
`_fused_all_gather_matmul_native` dispatches to `_fused_all_gather_matmul_native_rocm` on ROCm. It has the CUDA path's barriers and stream waits, with three differences:
- **Peer copies are issued before the GEMM.** The GEMM spins until the peer copies signal their chunks. HIP can put `current_stream` and the backend stream on one in-order hardware queue, where a GEMM issued first waits on copies queued behind it: with `GPU_MAX_HW_QUEUES=1`, CUDA's order hung in 3 of 3 runs at 2 ranks on MI300X and MI355X. HIP also replays the branches of a small captured graph one after another, in capture order, so the peer copies are captured before this rank's own copy too.
- **`A_signals` is allocated before the backend stream is ordered after the current stream.** It is zeroed on the current stream and set on the backend stream; in a captured graph, zeroing it after the stream waits puts the zeroing on the GEMM's branch, which HIP can replay after the peer signals.
- **Peer shards are copied in pieces of at most 1 MiB**, which HIP runs as kernels instead of on a DMA engine; that makes a call 2-3x faster. Local shards are copied whole, since HIP runs local copies as kernels at any size.
At 2 and 4 ranks on both GPUs, a captured graph hangs on its first replay with CUDA's order, with the own copy first, or with `A_signals` zeroed after the stream waits. With this order it replays correctly under every HIP graph scheduling mode and with one hardware queue.
On ROCm, `_should_use_fused_all_gather_matmul_native` adds one condition to CUDA's: `local_M % 256 == 0`. Each 256-row tile of the kernel waits on one chunk's signal, so a chunk must be a whole number of tiles (see above). CUDA's kernel has the same requirement for its 128-row tiles, but CUDA's selector does not check it. Other unsupported inputs raise in `_async_input_mm`, as on CUDA, and the native path stays opt-in through `TORCH_SYMM_MEM_ENABLE_NATIVE_ASYNC_TP`.
There is no K condition. Compared with the decomposition fallback (the implementation `fused_all_gather_matmul` would otherwise use), whether native is faster depends on the world size as much as on K: at M=4096 and 2 ranks it is slower for many shapes with K >= 1024, while at 8 ranks on MI300X it was faster for every shape measured, K from 512 to 8192. A K threshold taken from 2 and 4 ranks would have turned the native path off where it helps most.
- `test_fused_all_gather_matmul_native` runs on ROCm with the same shapes as on CUDA. To check that the native path ran, it looks for the CPU event of `symm_mem::_async_input_mm` on ROCm instead of the GEMM's kernel record, which the ROCm profiler can drop: in 20 fresh 2-rank runs on MI300X the record was missing in 3, while both outputs matched the fallback. The test now profiles CPU activity on all platforms; the CUDA profiler check is unchanged.
- The native GEMM's output is checked per element against a float64 reference instead of with `assert_close` against the fallback: `|out - ref| <= |fallback - ref| + sqrt(K) * eps_fp32 * (|A| @ |B|)`. The two GEMMs sum K in different orders, so near-zero outputs of the K=1024 reduction legitimately differ by more than bf16's atol of 1e-5. That failed one element at 8 ranks on MI300X and at 2 ranks on the MI350 CI runners, where hipBLASLt's result was 7 ULP from float64 and the CK kernel's 2 ULP. The slack is each element's fp32 accumulation error bound, so it stays tight on small outputs. The check fails a truncating bf16 epilogue, which `assert_close` passed, and one row computed with truncation, which a check of the whole tensor's mean and max error passed. It applies on CUDA too and passes on the H100 symm-mem CI lane.
- `test_fused_all_gather_matmul_native_graph_capture` (ROCm only) captures the native op in a CUDA graph and checks three replays with the same check.
- Both set `TORCH_SYMM_MEM_ENABLE_NATIVE_ASYNC_TP` with `mock.patch.dict`. The native test used to set it to `"0"` when done, but the selector checks only that it is set, so later tests in the process kept taking the native path.
- The native test runs on every ROCm version, including the ROCm 10.0 CI image. ROCM-30947 crashes copies of more than 1 MiB out of a peer's mapping, which HIP runs on a DMA engine, while a memory-copy tracing tool is registered. The native path copies peer shards in pieces of at most 1 MiB, and HIP runs local copies as kernels. Under memory-copy tracing on a ROCm 10.2 nightly that has the bug (10.2.0a20260916), a torch-free reproducer crashes on a 4 MiB copy out of a peer's mapping and passes 1 MiB peer copies and 4 and 16 MiB copies to and from a local mapping.
- `AsyncTPTest` is skipped on MI350 only before ROCm 10.1, instead of on all ROCm versions.
- `test_fused_scaled_matmul_reduce_scatter` runs on ROCm. Its skip said the fallback's API had changed, but the test already calls the fallback with its current signature, and its four variants pass at 2, 4 and 8 ranks on MI300X.
- Upstream's other ROCm skips in this file are unchanged.
On MI300X (gfx942), ROCm 10.2 nightly (20260927), with 2, 4 and 8 GPUs visible:
```
python -m pytest test/distributed/test_symmetric_memory.py -k AsyncTPTest
python -m pytest test/distributed/test_symmetric_memory.py -k "AsyncTPTest and test_fused_all_gather_matmul_native"
python -m pytest test/distributed/test_symmetric_memory.py -k SymmMemSingleProcTest
```
On MI355X (gfx950), same ROCm, the five native tests pass both natively and with the device properties of the CI runners' DPX partition (128 CUs, 4 XCCs). A sweep emulating both ranks on one GPU, 100 seeds over the four native variants plus 400 more seeds for the graph-capture configuration, had no failures under SPX, DPX, QPX and CPX device properties; the worst case used 33% of the bound.
In this PR's CI, `AsyncTPTest` gives 25 passed, 1 skipped on MI300X at 4 ranks (ROCm 10.0), on MI350 DPX at 2 ranks (rocm-preview, ROCm 10.2.0a20260916) and on H100. It is skipped on MI200 (gfx90a has no symmetric memory support) and on the trunk MI350 lane (ROCm 10.0). The remaining failures also fail on main.
PYFMT and RUFF report nothing on the two changed Python files.
Also checked locally, not part of this PR's tests: a GEMM with more output tiles than the grid at 2 and 4 ranks, and the eager and captured native op with `GPU_MAX_HW_QUEUES=1`.
Fixes pytorch#168524
Fixes pytorch#168522
Co-authored-by: @albmalamd
PR authored with AI assistance(Cursor + Claude)
Pull Request resolved: pytorch#177961
Approved by: https://github.com/jeffdaily
Co-authored-by: Albert Malewski <albert.malewski@amd.com>
(cherry picked from commit a476c63)
(cherry picked from commit 0603b4a)
## Summary pytorch#177961 sets `USE_ROCM_ASYNC_MM` for `hip/AsyncMM.hip` when `PYTORCH_ROCM_ARCH` contains gfx942 or gfx950. The file is then compiled for every arch in `PYTORCH_ROCM_ARCH`, and the `ck_tile` headers it includes do not compile for archs CK tile does not know. Multi-arch builds that also include gfx900, gfx906, gfx90c, gfx1010, gfx1011 or gfx1012 fail: ``` ck_tile/core/arch/arch.hpp:506:40: error: member reference base type 'void' is not a structure or union 20 errors generated when compiling for gfx1010. ``` `get_compiler_target()` has no branch for gfx101x, so its return type is deduced as `void`. gfx900/906/90c fail a different `static_assert` in `ck_tile/core/config.hpp`. Builds whose arch list contains only supported archs, like upstream CI's, are not affected. ## Fix Pass `--no-offload-arch=<arch>` to `AsyncMM.hip` for every arch in `PYTORCH_ROCM_ARCH` other than gfx942/gfx950. Only this file is affected; the rest of `torch_hip` still builds for all archs. The kernel already rejects other archs at runtime through its `gcnArchName` check. ## Test plan - Compiling `AsyncMM.hip` against the pinned CK (`5a74dec0`) with the 28-arch `PYTORCH_ROCM_ARCH` of the TheRock nightly reproduces the 20 errors for gfx1010. With this change the file builds, and its fatbin contains only gfx942 and gfx950 code objects. - In a minimal CMake project, the per-source `COMPILE_OPTIONS` drop the archs for the one file only; other sources in the same target keep all archs. Fixes ROCm/TheRock#8817 Pull Request resolved: pytorch#200054 Approved by: https://github.com/pragupta, https://github.com/jeffdaily (cherry picked from commit 9418553) (cherry picked from commit c31f1e7)
Resolving the cherry-pick conflict of pytorch#177961 onto release/2.15 kept both sides of the common_utils import block, importing get_cycles_per_ms and instantiate_parametrized_tests twice (RUFF F811), and took upstream's `skip` import from unittest, which this branch does not use (F401). Drop the duplicates and the unused import; no behavior change. Test Plan: ``` lintrunner --skip CLANGTIDY,CLANGTIDY_EXECUTORCH_COMPATIBILITY,PYREFLY,CODEOWNERS_TAXONOMY \ test/distributed/test_symmetric_memory.py torch/distributed/_symmetric_memory/__init__.py \ caffe2/CMakeLists.txt torch/csrc/distributed/c10d/hip/AsyncMM.hip ``` This commit was authored with an AI assistant (Claude Code). Co-Authored-By: Claude <noreply@anthropic.com>
jithunnair-amd
pushed a commit
that referenced
this pull request
Nov 4, 2026
…200054) (#3868) Cherry-picks two upstream commits onto release/2.15 to bring ROCm native AsyncTP support (the CK-tile fused_all_gather_matmul persistent-GEMM path) to this branch: - pytorch#177961 "[ROCm] Enable native AsyncTP" -- adds the CK-tile GEMM kernel (hip/AsyncMM.hip) and ROCm host dispatch for the native AsyncTP path of fused_all_gather_matmul, gated to gfx942/gfx950. - pytorch#200054 "[ROCm] Build AsyncMM.hip only for gfx942/gfx950" -- follow-up fix for multi-arch builds (e.g. TheRock's nightly arch list): the CK-tile headers AsyncMM.hip includes don't compile for every arch, so this passes --no-offload-arch=<arch> for every arch outside gfx942/gfx950. A third commit on top fixes two lint-only issues introduced while cherry-picking pytorch#177961's import-list diff against release/2.15's already-diverged test_symmetric_memory.py: a duplicate instantiate_parametrized_tests import (RUFF F811) and an unused get_cycles_per_ms import pulled in as unrelated diff context (RUFF F401). No functional changes. Both cherry-picks were verified hunk-for-hunk against their upstream commits and are otherwise byte-identical. Same pattern as #3840 (the release/2.14 backport of the same two commits). ## Test plan Built with `pip install -e . -v --no-build-isolation`, `PYTORCH_ROCM_ARCH=gfx950`, on an 8x MI355X (gfx950) box, ROCm 7.17. Confirmed `USE_ROCM_ASYNC_MM` compiled into `build/compile_commands.json`. Ran `AsyncTPTest` in test/distributed/test_symmetric_memory.py at world size 2, 4, and 8 (`HIP_VISIBLE_DEVICES=0,1` / `0,1,2,3` / `0,1,2,3,4,5,6,7`): ``` python test/distributed/test_symmetric_memory.py -k AsyncTPTest ``` 25 passed / 1 skipped (multimem, no multicast support on this box) at each world size, repeated 3x at ws=8. `lintrunner` on the touched files: clean. ## Agent details This PR was prepared with an AI assistant (Claude Code): rebasing the cherry-picks onto the current release/2.15 tip, running the build and the AsyncTPTest suite at ws=2/4/8, and drafting this description. All test runs and lint were verified directly, not taken on faith from the assistant. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Albert Malewski <albert.malewski@amd.com> Co-authored-by: Marko Radosavljevic <maradosa@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cherry-picks two upstream commits onto release/2.15 to bring ROCm native AsyncTP
support (the CK-tile fused_all_gather_matmul persistent-GEMM path) to this branch:
(hip/AsyncMM.hip) and ROCm host dispatch for the native AsyncTP path of
fused_all_gather_matmul, gated to gfx942/gfx950.
for multi-arch builds (e.g. TheRock's nightly arch list): the CK-tile
headers AsyncMM.hip includes don't compile for every arch, so this passes
--no-offload-arch= for every arch outside gfx942/gfx950.
A third commit on top fixes two lint-only issues introduced while cherry-picking
pytorch#177961's import-list diff against release/2.15's already-diverged
test_symmetric_memory.py: a duplicate instantiate_parametrized_tests import
(RUFF F811) and an unused get_cycles_per_ms import pulled in as unrelated diff
context (RUFF F401). No functional changes.
Both cherry-picks were verified hunk-for-hunk against their upstream commits and
are otherwise byte-identical. Same pattern as #3840 (the release/2.14 backport
of the same two commits).