feat(ffn): add deterministic distributed Triton FFN for ROCm - #325
Merged
Conversation
frank-2077
requested review from
Flink-ddd,
KJLdefeated,
bitborne and
inaniloquentee
as code owners
August 20, 2026 15:23
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
Flink-ddd
changed the base branch from
codex/ws2-rocm-strict-attention
to
main
August 21, 2026 15:37
frank-2077
changed the base branch from
main
to
codex/ws2-rocm-strict-attention
August 21, 2026 15:40
CPU host numbers for the two paths that exist without a device: torch SDPA and NativeAttentionOp. S=512/1024/2048, bf16 and fp16. Timing is wall clock and peak memory an RSS high-water delta from /proc, so the host figures approximate and are not directly comparable to the device columns. The budget is now enforced twice: a pre-flight projection from the one untimed call that already runs to capture outputs, and a wall clock inside the sampling loops, because the projection under-estimates once the materialized score matrix leaves cache. A cell that runs out of budget is truncated and flagged rather than dropped, and one that cannot start is listed under "Skipped cells" with its observed single-call cost. Worth knowing before reading the host column: fp16 pytorch-native is 21.1 s per forward at S=2048 against 228 ms for the same shape in bf16. That is PyTorch having no optimized fp16 CPU matmul, not a property of this operator. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01824WkgMDxtBKD4Ex2NkfYA
Signed-off-by: maxiaosong1124 <maxiaosong7890@outlook.com>
Port the HIP IPC fixed-tree transport and packed reduce-scatter path from PR RL-Align#357, with RCCL fallback and focused ROCm coverage.
Record the PR RL-Align#357 collective measurements, preserve the four-way same-topology comparison, and publish the refreshed MI300X artifacts.
feat(cuda): promote deterministic cross-config runtime and kernel validation to main
# Conflicts: # csrc/ops.cpp # rl_engine/distributed/collectives.py # rl_engine/kernels/ops/pytorch/ffn/ffn.py
…tch path CP orchestration lived only in StrictCUDAAttentionRuntime, so ROCm had nowhere to put it: the AITER/CK core is single-rank arithmetic, the Vime provider failed closed at CP>1, and the only working AG/core/RS sequence was in the benchmark script. StrictRocmAttentionRuntime mirrors the CUDA runtime over the RCCL AG/RS transport. Two things differ from the CUDA runtime and both are load-bearing: - The core is launched once per (batch row, KV group) rather than once per sequence. AITER/CK's reduction order depends on how many heads shared the launch, so a head shard computed under TP=N is otherwise not bit-identical to the same shard under a different TP degree. FA4 has no such dependence. - RCCL moves tensors but never reduces them. The cross-rank (out, lse) combine is the transport's fixed balanced rank tree, not RCCL's own algorithm selection, which varies with message size and topology. The sequence reorder and position validation are bound from the CUDA runtime rather than reimplemented, so the two runtimes cannot drift into two different global orderings. The per-KV-group launch loop moves out of the Vime provider into the runtime, so CP=1 and CP>1 now share one schedule instead of keeping a second copy in the provider. Opening CP also required the registry to stop rejecting it: cp_world_sizes was (1,) and deterministic_cp_merge was False, so AttentionBackendCapability rejected CP>1 twice over. cp_world_sizes now matches the world sizes the RCCL transport accepts and deterministic_cp_merge is True because the merge order is ours. A test pins the two together so the declaration cannot drift from the transport. zigzag fails closed: the strict CP plan describes one contiguous block per rank, and a zigzag rank owns two discontiguous runs. Measured on 8xMI300X through attention_provider, not the transport directly, so the test also pins that CP is reachable from the production dispatch path: CP=2/4/8 are bitwise against a CP=1 run of the same core on the same logical sequence, 0 mismatched elements on out and lse, repeat-bitwise on every rank. Also corrects the stale TP comment the PR description flagged in section 7: the shipped policy removes the degree dependence with the per-KV-group launch rather than binding the degree to avoid a ~3x cost. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NPU7GosGZSdj6pKBWX7n2Y
… collective The ROCm attention CP path owned a second transport implementation. RCCLAGRSAttentionCPCommunication overrode _get_collective to construct _RCCLRankOrderedTransport, which bypassed the collective_for_group factory that the CUDA adapter goes through and that already dispatches to RCCLDeterministicCollective on HIP. Both copies evaluated the same balanced rank tree, so the two platforms were bit-identical -- but only by coincidence. Nothing pinned them together, so a later change to the shared collective's reduction order would have left the attention path on the old tree with no test failing. Delete _RCCLRankOrderedTransport and the override. ROCm now inherits the CUDA adapter's resolution, so one implementation serves both platforms. The reduction expression is unchanged, so this is not expected to move any bit. _RootReduceScatterSequence falls back from scatter() to reduce_scatter() because the shared collective exposes no scatter entrypoint. That is the same branch the CUDA path has always taken and it is semantically equivalent -- non-root ranks zero their input, so the tree sum returns the root's value and adding zero is exact. It costs one extra all-gather plus tree per call. The error messages the adapters raise are now keyed off a collective_label class attribute so the inherited path still reports RCCL on ROCm instead of mislabeling itself as CUDA. Three tests pin the arrangement: the two adapters must share one _get_collective, the ROCm adapter must resolve through collective_for_group, and the registry's cp_world_sizes must equal the shared collective's _SUPPORTED_WORLD_SIZES so the capability declaration cannot drift from what the transport accepts. Not verified on MI300X. The reduction expression is unchanged, but the reduce_scatter fallback and the shared collective's capacity and signature validation are new to this path, so a CP=2/4/8 bitwise run is still owed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: zhangj1an <jianmusings@gmail.com>
Signed-off-by: vensen <vensenmu@gmail.com>
Signed-off-by: vensen <vensenmu@gmail.com>
Signed-off-by: maxiaosong1124 <maxiaosong7890@outlook.com>
…codex/ws2-rocm-strict-attention Brings in the ROCm HIP IPC deterministic collective so ROCm gets a native transport instead of the RCCL-only Python path. Because this branch already routes the attention CP adapter through collective_for_group, the ROCm attention CP path now resolves to that HIP IPC collective with no further change. Conflict resolutions: * cp_comm.py -- PR RL-Align#357 optimizes _RCCLRankOrderedTransport.scatter; this branch deleted that class in favour of the shared collective. Kept the deletion: the optimization targets code that no longer exists, and the shared collective supersedes it. * collectives.py -- kept both sides' module constants (they are additive: CUDA staging-frame sizes and ROCm IPC tuning thresholds). reduce_scatter_many had diverged signatures, so the merged one takes the union: PR RL-Align#357's inputs/outs plus this branch's validate_signature, forwarding both. * ffn.py -- PR RL-Align#357 restructured the backward to compute both gate and up input gradients up front and pack them into one reduce_scatter_many, while this branch renamed _gemm_fwd to _linear_da/_linear_dw. Took PR RL-Align#357's structure with this branch's helper names; the old second reduce-scatter for the up lane is gone. * setup.py -- kept this branch's Ascend build imports (sysconfig, CompileError, find_executable, Extension are used further down the file) and added PR RL-Align#357's ROCm .hip source to cuda_sources. * ops.cpp, _C.pyi -- both additive; kept both sides. tests/distributed and tests/test_build_platform_collectives: 46 passed, 5 skipped. tests/test_qwen_ffn.py fails 24 here, but 25 of the same tests fail on the pre-merge tree in this environment: the extension was built without KERNEL_ALIGN_DET_GEMM_SM90=1, so strict GEMM refuses to run. The merge removes one of those failures and adds none. Not verified on MI300X. The HIP IPC path has no coverage in this environment, so the ROCm CP bitwise run is still owed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: zhangj1an <jianmusings@gmail.com>
… into The CUDA core checks the FA4 CuTe API by parameter name before it runs, so a renamed or dropped strict control fails at load. ROCm had no equivalent: inspect.signature reports (*args, **kwargs) for AITER's JIT wrapper, so the only guard was a SHA-256 of the module source. That catches "something changed" but cannot say what, and it fires on unrelated edits. Read the registered Torch schema instead (torch.ops.aiter.<op>.default). The check is an ordered prefix, not a name set, because the two call sites pass positionally. An argument inserted upstream would shift the meaning of every later argument while the call still type-checks -- dropout_p, the two window sizes and sink_size are all int/bool, so nothing would raise and the kernel would run with silently reinterpreted controls. Name presence alone does not catch that; the FA4 path is exempt only because it calls by keyword. This also pins something that was previously unprovable: the True at backward position 11 is the schema's `deterministic`. ROCm's backward was already deterministic, but nothing tied that literal to its parameter. The source fingerprint stays as a second line: the schema check describes what changed, the fingerprint still catches a same-schema implementation change. Verified against the installed AITER; the forward prefix matches exactly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: zhangj1an <jianmusings@gmail.com>
StrictCUDAAttentionRuntime has forward_paged_with_lse; the ROCm runtime had no decode entry point at all, so decode-stage KV-cache replay existed only in the device-neutral comparison harness with no ROCm path behind it. No AITER paged kernel can serve this contract. paged_attention_rocm, _ragged and _v1 are all two-pass partition reducers -- partition_size, with exp_sums/max_logits/tmp_out partials -- so the partition count tracks the cached length and Split-KV cannot be turned off. AITER's flash_attn_varlen_func takes a block_table but exposes no num_splits to pin, unlike CUDA's FA4. Either way the contract could not prove Split-KV disabled, which attention_binding checks from both runtime evidence and the contract. So the pages are gathered into logical order and handed to the same dense core the prefill path uses, at the same one-launch-per (batch row, KV group) granularity. The arithmetic is then identical to a CP=1 prefill over the same logical sequence, which is what makes decode replay comparable against it. The cost is materializing the cached KV; a native paged kernel avoids that and can replace this once AITER can pin its split count. The registry deliberately does NOT gain AttentionMode.DECODE. Nothing routes to the new entry point: the Vime provider always calls forward_with_lse and builds its contract with kv_cache=None, and its request carries no page table. Declaring the mode now would let the binding layer accept a decode path that never executes. A test pins the omission so it flips together with the dispatch wiring rather than drifting ahead of it. Tests inject the core, so they run without ROCm. They pin the part that is ours rather than AITER's: a shuffled page table still yields logical KV order, the gather truncates to seqused_k instead of exposing the page tail, each launch still sees exactly one KV group, and the provenance says paged_kernel=none so no reader mistakes this for a native paged path. Not verified on MI300X. The core arithmetic is unchanged, but the gather's index_select/reshape/permute and the claimed bitwise equality with a CP=1 prefill over the same tokens both need a real run. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: zhangj1an <jianmusings@gmail.com>
Added a new function to precompile strict attention training for better performance when using RL_KERNEL. Modified the integration entry point to include this precompilation step based on the plan's implementation.
…l-decode-cudagraph
# Conflicts: # csrc/ops.cpp # rl_engine/distributed/collectives.py # rl_engine/kernels/attention_contract.py # rl_engine/kernels/ops/cuda/attention/__init__.py # rl_engine/kernels/ops/cuda/attention/deterministic_attn.py # rl_engine/kernels/ops/cuda/attention/flash_attn.py # rl_engine/kernels/ops/pytorch/attention/cp_attention.py # scripts/ws2_p2p_nccl_attention_reference_check.py # tests/test_det_gemm.py # tests/test_flashinfer_pr7_attention.py # tests/test_qwen_ffn.py
frank-2077
changed the base branch from
codex/ws2-rocm-strict-attention
to
test
September 1, 2026 03:36
maxiaosong1124
self-requested a review
September 3, 2026 05:20
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
5 participants
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.
Summary
This PR adds a ROCm-native deterministic distributed Qwen3 FFN implemented in
Triton. It supports FFN forward/backward across tensor parallelism (TP), context
parallelism (CP), and sequence parallelism (SP), with fixed-order RCCL tensor
transport.
Note
Validation is operator-only. It uses seeded tensors and does not load or
benchmark a model, checkpoint, tokenizer, dataset, or serving engine.
Comparison contract
The three experiments are intentionally independent:
dHidden, anddWeights; acceptance is 0Qwen3MLPat TP=1Qwen3MLPat TP=1 in FP32Native ROCm/RCCL is not used as a numerical-accuracy reference in this report.
Design
kernels; there is no CUDA-generated HIP source in this PR.
GEMM and preserve BF16 stage boundaries in forward and backward.
BF16 rank reduction tree.
tree matches CP=1.
corresponding backward paths.
Operator test matrix
Weights use Hugging Face
[out, in]layout. No model-level benchmark and noseparate gate/up/down projection benchmark is included.
(M,H,I)=(1/8/32,4096,12288), BF16Qwen3MLPTP1; forward and forward+backward(M,H,I)=(32,4096,12288), BF16(M,H,I)=(8,4096,12288)Qwen3MLPTP1 FP16 vs the same operator in FP32ROCm environment
08f47d97d0443c5998b8da6b41a22fdf3848da8fe64ababCorrectness results
dHiddenvs Triton TP1dWeightsvs Triton TP1Commands used:
Performance results
All speed numbers compare complete FFN calls. The official baseline is upstream
Transformers
Qwen3MLPwith unsharded weights and input at TP=1. Distributedtiming uses synchronized wall time and the slowest rank per sample. These rows
compare speed only.
9.03-22.86x7.38-11.56x8.76-15.79x7.45-14.06xThe separate dtype observation runs only official
Qwen3MLPTP1:6.544e-4(0.06544%)2.046e-63.742e-7Full combined report ·
Raw JSON
Communication overlap assessment
The reported implementation serializes dependencies; the measurements do not
claim communication/computation overlap.
TP reduction consumes the down-projection output. These are hard dependencies.
dHiddenare independent untiltheir final ordered addition. A future implementation can reduce one on a
second stream while computing the other.
BF16 stage boundaries, and gate-then-up addition order. It is accepted only if
every TP1 mismatch column remains zero.
Communication implementation provenance
The ROCm deterministic communication operator used by this PR is adopted from PR #357. The current path uses the optimized HIP IPC fixed-tree implementation with RCCL fallback, including the packed
reduce_scatter_manypath for the independent sequence-parallel FFN backward lanes.The checked benchmark report records this implementation provenance and keeps the main performance figure as a four-way same-topology comparison without adding a PR-specific series.