[aryatsc]

Motivation

K3's vocab-parallel logits gather and its routed-expert up-projection gather each run a dim-0 all_gather, a movedim copy, and one or two adds.

Modifications

• out = bf16(bf16(all_gather(inp, dim=-1) + add_b) [+ add_c]) runs in one launch.
• Each rank reads its peers' shards directly from the registered IPC buffers.
• Rounding matches the unfused sequence, so the output is bit-identical.
• Under graph capture the input is read in place. In eager mode it is first staged into the registered pool.

Tests

op_tests/multigpu_tests/test_custom_allreduce.py checks bit-exactness against the unfused reference in eager and graph mode on MI355X TP8. Used by the Kimi-K3 attention-residual PR in SGLang.

PR Self-Review

Reviewed revision

b33ea4b

Behavior and intent

Adds all_gather_lastdim_add: out = all_gather(inp, dim=-1) + add_b [+ add_c] in one launch, reading peer shards from the registered IPC buffers. Today the same result needs a dim-0 all-gather, a movedim copy and one or two adds (3-4 launches and an extra round trip through HBM). Kimi-K3 runs it for the vocab-parallel logits gather and the routed-expert up-projection gather. Existing collectives are unchanged.

Validation and unverified scope

• torchrun --nproc_per_node=8 op_tests/multigpu_tests/test_custom_allreduce.py on 8x MI355X (gfx950), bf16 and fp16, eager and CUDA graph: all cases bit-exact (atol=rtol=0) against torch.cat(shards, -1) + add_b [+ add_c].
• Fused vs unfused (custom_all_gather(dim=0) + reshape + adds), TP8 MI355X, graph mode, 200 replays, max over ranks: 1.6x at 8-256 rows, 1.1-1.4x at 1 row (table in the review thread).
• TennyWang1223 reproduced bit-exactness on 8x MI300X (gfx942).
• Unverified: TP2/TP4 (allowed by the predicate, not run), gfx1250 (op not built there; the predicate returns False).

PR-specific question

Selected only through should_custom_ag_lastdim_add(inp) and the wrapper, which require: communicator enabled and op built (not gfx1250), world size 2/4/8 and fully connected above 2, bf16 or fp16, contiguous 2D inp with last dim % 8 == 0, 16-byte aligned, and inp no larger than the registered pool. out/add_b/add_c must be contiguous [rows, n * world_size], same dtype and device, 16-byte aligned, and out must not alias add_b/add_c. The C++ entry repeats these checks and raises, so direct callers can't reach the kernel with unsupported inputs.

Changed op tests

op_tests/multigpu_tests/test_custom_allreduce.py: adds the only test of all_gather_lastdim_add, at Kimi-K3 shapes, bf16 and fp16, eager and graph, with and without add_c, plus an input rewrite before graph replay and rejection of unsupported inputs (predicate False cases, aliased, misaligned and wrong-shape tensors). No existing test covers this op.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6960 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[TennyWang1223]

● Review Summary — PR #6960: all_gather_lastdim_add

Verdict: Request changes. The kernel itself is correct. What blocks merge is a lint failure in CI and two input-validation gaps in the wrapper that produce wrong results without raising.

Validation (local, PR head merged onto current main, 8× MI300X gfx942, TP8)

• The PR's own test (test_custom_allreduce.py, 5 shapes) passes bit-exact (atol=rtol=0) in both eager and graph mode.
• Extra cases I ran also come out bit-exact:
  • the add_c=None path, eager and graph;
  • changing the input in place, then replaying the captured graph.
• The bit-exact claim also holds on gfx942: module_custom_all_reduce builds with OPUS_FP32_to_BF16_DEFAULT=5 (round-to-nearest-even), so its rounding matches torch's.

CI

• ❌ Check Code Style with Ruff fails because of this PR, in aiter/ops/custom_all_reduce.py: No changes this session
  • I001: the import block is unsorted (from typing import Optional must come before import torch).
  • UP045: write torch.Tensor | None instead of Optional[torch.Tensor].
• That failure makes check-signal fail, so every GPU job is skipped, the ci:multi-gpu tests included, and the new test has not run in CI yet. After fixing the lint, re-run CI and confirm the Multi-GPU Tests (8 GPU) job passes.

Issues

1. ⚠️ [verified] A non-contiguous inp gives a silently wrong result.
   • custom_all_reduce.py:1909-1916 checks add_b, add_c and out for contiguity, but not inp. The C++ side then copies numel * element_size bytes from inp.data_ptr() as if it were
     contiguous.
   • Repro: inp = wide[:, :n] raises no error, and the max abs error is 7.16.
   • Fix: add assert is_weak_contiguous(inp), the same check should_custom_ag applies.
2. ⚠️ [verified] add_b, add_c and out have no dtype check.
   • Only inp is checked to be bf16. An fp32 add_b is read as bf16 with no error; the max abs error is 7.39.
   • Fix: assert that add_b, add_c and out have the same dtype and device as inp.
3. 📝 Test coverage gaps.
   • The test always passes add_c, so the kHasC=false kernel instantiation never runs in CI.
   • The graph test never changes the input between replays, although reading inp in place under graph capture is exactly what this PR adds.
   • Fix: add a case without add_c, and a case that modifies the input before replay.
4. 📝 [inferred] The registered capture path skips _capture_registered(inp).
   • Every other collective calls it, and it raises a clear error when inp cannot be IPC-exported (the expandable_segments case, issue #4174). Without it, that failure would show
     up later, at flush or replay, with an unclear message.
   • Fix: reuse _capture_registered(inp) here. I have not reproduced this one.

[aryatsc]

● Review Summary — PR #6960: all_gather_lastdim_add

Verdict: Request changes. The kernel itself is correct. What blocks merge is a lint failure in CI and two input-validation gaps in the wrapper that produce wrong results without raising.

Validation (local, PR head merged onto current main, 8× MI300X gfx942, TP8)

• The PR's own test (test_custom_allreduce.py, 5 shapes) passes bit-exact (atol=rtol=0) in both eager and graph mode.

• Extra cases I ran also come out bit-exact:

  • the add_c=None path, eager and graph;
  • changing the input in place, then replaying the captured graph.

• The bit-exact claim also holds on gfx942: module_custom_all_reduce builds with OPUS_FP32_to_BF16_DEFAULT=5 (round-to-nearest-even), so its rounding matches torch's.

CI

• ❌ Check Code Style with Ruff fails because of this PR, in aiter/ops/custom_all_reduce.py: No changes this session

  • I001: the import block is unsorted (from typing import Optional must come before import torch).
  • UP045: write torch.Tensor | None instead of Optional[torch.Tensor].

• That failure makes check-signal fail, so every GPU job is skipped, the ci:multi-gpu tests included, and the new test has not run in CI yet. After fixing the lint, re-run CI and confirm the Multi-GPU Tests (8 GPU) job passes.

Issues

1. ⚠️ [verified] A non-contiguous inp gives a silently wrong result.

   • custom_all_reduce.py:1909-1916 checks add_b, add_c and out for contiguity, but not inp. The C++ side then copies numel * element_size bytes from inp.data_ptr() as if it were
     contiguous.
   • Repro: inp = wide[:, :n] raises no error, and the max abs error is 7.16.
   • Fix: add assert is_weak_contiguous(inp), the same check should_custom_ag applies.

2. ⚠️ [verified] add_b, add_c and out have no dtype check.

   • Only inp is checked to be bf16. An fp32 add_b is read as bf16 with no error; the max abs error is 7.39.
   • Fix: assert that add_b, add_c and out have the same dtype and device as inp.

3. 📝 Test coverage gaps.

   • The test always passes add_c, so the kHasC=false kernel instantiation never runs in CI.
   • The graph test never changes the input between replays, although reading inp in place under graph capture is exactly what this PR adds.
   • Fix: add a case without add_c, and a case that modifies the input before replay.

4. 📝 [inferred] The registered capture path skips _capture_registered(inp).

   • Every other collective calls it, and it raises a clear error when inp cannot be IPC-exported (the expandable_segments case, issue [Bug] custom_all_reduce.cu:397 hipIpcGetMemHandle fails when caller uses PYTORCH_HIP_ALLOC_CONF=expandable_segments:True #4174). Without it, that failure would show
     up later, at flush or replay, with an unclear message.
   • Fix: reuse _capture_registered(inp) here. I have not reproduced this one.

Thanks for the thorough review. Fixed in fc6dbbd:

• ruff I001/UP045
• inp must be 2D and contiguous; add_b/add_c/out are checked for shape, contiguity, dtype and device against inp
• the registered capture path goes through _capture_registered(inp)
• the test also runs every case with add_c=None, and graph mode rewrites the input shard before a final replay
  Passes bit-exact on 8x MI355X (gfx950), eager and graph.

[TennyWang1223]

LGTM

[amd-songpiao]

1. No way for callers to know if the fast path is available. The new method never checks self.disabled. On a disabled communicator (cross-node, missing library), _pool and _ptr are never set, so it fails with an AttributeError. Other failures surface as exceptions too:

   • NotImplementedError on gfx1250.
   • C++ RuntimeErrors for an unsupported world size, rows that aren't a multiple of 8 elements, or an input larger than the
     staging pool.

   Unlike custom_all_gather and should_custom_ag, nothing tells the caller in advance. I'd ask for a
   should_custom_ag_lastdim_add(inp) check, or a custom_ variant that returns None on any of these. That way SGLang can fall
   back to the unfused path cleanly instead of catching exceptions.

2. Should out be allowed to be the same tensor as add_b or add_c? That's the natural in-place residual-add use. It works in practice, because each thread reads add_b[i] before writing out[i]. But the kernel marks these pointers restrict, which tells the compiler they don't overlap. Either document that sharing is supported and drop restrict on those two, or assert that they're different tensors.

3. No check that tensors are 16-byte aligned. The kernel loads 16 bytes at a time. The checks accept any contiguous tensor, including a view that starts at an odd element offset. A data_ptr() % 16 == 0 assert on inp, add_b, add_c and out would turn a possible fault into a clear error.

4. No performance numbers. The PR's motivation is performance, but the test only times the fused kernel. A comparison against the old dim-0 all-gather, copy and add sequence would justify the change, ideally at the Kimi-K3 shapes (n=896 and the vocab shard) on MI355X.

5. Eager mode always copies inp into the pool, even when it's already registered.

6. The C++ entry point doesn't check the shapes of out, add_b and add_c; only the Python wrapper does, so a direct aiter.all_gather_lastdim_add call skips them.

7. The op is bf16-only, though the template would work for fp16 as-is.

[aiter-bot]

NOT READY FOR HUMAN REVIEW

PR Self-Review

Revision b33ea4b1536a · mode shadow

Author self-review accepted. Readiness report:

Author answers

• PASS — current-SHA answers are present and structurally specific

Show the author's self-review answers

Behavior and intent

Adds `all_gather_lastdim_add`: `out = all_gather(inp, dim=-1) + add_b [+ add_c]` in one launch, reading peer shards from the registered IPC buffers. Today the same result needs a dim-0 all-gather, a movedim copy and one or two adds (3-4 launches and an extra round trip through HBM). Kimi-K3 runs it for the vocab-parallel logits gather and the routed-expert up-projection gather. Existing collectives are unchanged.

Validation and unverified scope

- `torchrun --nproc_per_node=8 op_tests/multigpu_tests/test_custom_allreduce.py` on 8x MI355X (gfx950), bf16 and fp16, eager and CUDA graph: all cases bit-exact (atol=rtol=0) against `torch.cat(shards, -1) + add_b [+ add_c]`.
- Fused vs unfused (`custom_all_gather(dim=0)` + reshape + adds), TP8 MI355X, graph mode, 200 replays, max over ranks: 1.6x at 8-256 rows, 1.1-1.4x at 1 row (table in the review thread).
- TennyWang1223 reproduced bit-exactness on 8x MI300X (gfx942).
- Unverified: TP2/TP4 (allowed by the predicate, not run), gfx1250 (op not built there; the predicate returns False).

PR-specific question

Selected only through `should_custom_ag_lastdim_add(inp)` and the wrapper, which require: communicator enabled and op built (not gfx1250), world size 2/4/8 and fully connected above 2, bf16 or fp16, contiguous 2D `inp` with last dim % 8 == 0, 16-byte aligned, and `inp` no larger than the registered pool. `out`/`add_b`/`add_c` must be contiguous `[rows, n * world_size]`, same dtype and device, 16-byte aligned, and `out` must not alias `add_b`/`add_c`. The C++ entry repeats these checks and raises, so direct callers can't reach the kernel with unsupported inputs.

Changed op tests

`op_tests/multigpu_tests/test_custom_allreduce.py`: adds the only test of `all_gather_lastdim_add`, at Kimi-K3 shapes, bf16 and fp16, eager and graph, with and without `add_c`, plus an input rewrite before graph replay and rejection of unsupported inputs (predicate False cases, aliased, misaligned and wrong-shape tensors). No existing test covers this op.

PR state

• PASS — pull request is marked ready for review

Copilot findings

• PASS — no active Copilot finding lacks a substantive author reply
• ADVISORY — Copilot did not review the current head SHA

Code and op-test hygiene

• PASS — no committed artifact or conflict marker detected

Required CI

• PASS — Black: passed
• PASS — Ruff: passed
• PASS — Repository dependency: passed
• FAIL — Aiter Test Gate: failed
• Required CI failed; fix it before requesting human review.

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[aryatsc]

Edit the PR Self-Review section in the PR description. Include the current revision SHA; this comment refreshes automatically.

Thanks, all fair. Changes in a95a988 and b33ea4b:

1. Added should_custom_ag_lastdim_add(inp). It returns False (no raise) when the communicator is disabled or the op isn't built (gfx1250), world size isn't 2/4/8, the group isn't fully connected (ws>2), dtype isn't bf16/fp16, inp isn't contiguous 2D, n isn't a multiple of 8 or inp isn't 16-byte aligned, or inp exceeds the pool. Calling the op when it's unavailable now raises a RuntimeError that points to the predicate.
2. out aliasing add_b/add_c is rejected in Python and C++. __restrict__ stays.
3. All four tensors are checked for 16-byte alignment, plus n % 8 == 0, in Python and C++.
4. TP8 MI355X, CUDA graph, one op per graph, mean of 200 replays, max over ranks. Unfused = custom_all_gather(dim=0) + reshape + 2 adds. Output is bit-identical to the fused op.

shape (rows, n)	fused us	unfused us	speedup
1 x 896	10.5-12.9	14.4-14.6	1.1-1.4x
8 x 896	11.0	17.6	1.59x
64 x 896	12.5	20.1	1.60x
128 x 896	14.2	23.3	1.64x
256 x 896	18.6	30.5	1.64x
1 x 20480	11.8	14.3	1.21x
8 x 20480	16.1	26.3	1.63x

bf16 is shown; fp16 is within noise of it. In Kimi-K3 this is the attention-residual / logits gather. Turning it off costs 1-6% end-to-end throughput (sgl-project/sglang#46656).

5. Kept as is on purpose. In eager mode inp is a regular allocation the peers can't read, so it has to be staged into the registered buffer, the same as all_gather_unreg. The graph path reads the registered input in place with no copy. Skipping the eager copy would need a caller that allocates inp in the IPC buffer, and none does today, so I'd rather not add that path without a user.
6. The C++ entry now checks dtype, 2D, contiguity, shape [rows, n*world_size], device, alignment and aliasing for inp/out/add_b/add_c, so direct callers get the same checks as the Python wrapper.
7. Added fp16 dispatch. The test runs bf16 and fp16, eager and graph, with and without add_c, plus input rewrite before replay, predicate false cases, and rejection of aliased/misaligned/wrong-shape tensors. All bit-exact on 8x MI355X.

[zufayu]

@TennyWang1223 The requested changes have been addressed. There are no CI failures so far, with a few checks still running. Could you please take another look and confirm whether this is ready to merge once CI is green?

[TennyWang1223]

Thanks for the follow-ups. All of songpiao's items look addressed, and I re-ran fp16/bf16 eager+graph at TP2 on MI300X (merged with latest main), all bit-exact. CI's only failure (test_custom_allreduce_capture_registration.py, HIPCachingAllocator assert in teardown) also happens on main and is unrelated.

One remaining issue: with an empty shard (rows == 0, e.g. an idle DP-attention rank, same case as #7302), should_custom_ag_lastdim_add returns True but all_gather_lastdim_add then raises AssertionError: ... out aliasing add_b or add_c, because all empty tensors have data_ptr() == 0. Please early-return out when inp.numel() == 0 (Python), and skip the alias check, the memcpy and the launch in C++ for that case. Optional: consider honoring AITER_CUSTOM_AG_MAX_SIZE (_cag_max_size) in the predicate the way should_custom_ag now does on main.