feat(ffn): add deterministic distributed Triton FFN for ROCm

[frank-2077]

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:

Question	Baseline	Metric
Does distributed execution preserve determinism?	This PR's deterministic Triton FFN at TP=1	Element mismatch count for forward output, training output, dHidden, and dWeights; acceptance is 0
What is the deterministic performance cost?	Official Hugging Face Transformers Qwen3MLP at TP=1	Median FFN latency only; no accuracy comparison is mixed into the speed result
What is the simple FP16 precision observation?	Official Qwen3MLP at TP=1 in FP32	FP16 output relative-L2, max-absolute, and mean-absolute error

Native ROCm/RCCL is not used as a numerical-accuracy reference in this report.

Design

• Implement the bias-free gated Qwen3 FFN directly with ROCm-native Triton
  kernels; there is no CUDA-generated HIP source in this PR.
• Use a canonical FP32-leaf/BF16-node midpoint reduction tree in deterministic
  GEMM and preserve BF16 stage boundaries in forward and backward.
• Make each contiguous TP K shard the same subtree used by TP=1.
• Use RCCL for fixed rank-order tensor transport, followed by a fixed balanced
  BF16 rank reduction tree.
• Gather complete CP token sequences before weight-gradient GEMMs so their K
  tree matches CP=1.
• Support TP all-reduce, SP all-gather/reduce-scatter, CP all-gather, and all
  corresponding backward paths.

Operator test matrix

Weights use Hugging Face [out, in] layout. No model-level benchmark and no
separate gate/up/down projection benchmark is included.

Experiment	Shape / dtype	Parallel configurations
Single-GPU FFN speed	(M,H,I)=(1/8/32,4096,12288), BF16	Triton TP1 vs official Qwen3MLP TP1; forward and forward+backward
Distributed FFN speed	(M,H,I)=(32,4096,12288), BF16	TP2, TP2+SP, TP4, TP2+CP2, TP2+CP2+SP, TP8, TP4+CP2, TP4+CP2+SP; every row vs official TP1
Distributed exactness	Same full logical input and weights, BF16	Every TP/CP/SP layout vs deterministic Triton TP1 exact slices
FP16/FP32 observation	(M,H,I)=(8,4096,12288)	Official Qwen3MLP TP1 FP16 vs the same operator in FP32

ROCm environment

Item	Value
GPU	8 × AMD Instinct MI300X
Architecture	gfx942
PyTorch	2.12.0+rocm7.14.0a20260608
ROCm runtime	7.14.60850
Transformers	5.10.4
Benchmark implementation commit	08f47d97d0443c5998b8da6b41a22fdf3848da8f
Result and figure commit	e64abab

Correctness results

Validation	Result
Single-GPU FFN plus real RCCL TP/CP/SP topology suite	17 passed
Formal TP/CP/SP forward output vs Triton TP1	0 mismatch
Formal TP/CP/SP training output vs Triton TP1	0 mismatch
Formal TP/CP/SP dHidden vs Triton TP1	0 mismatch
Formal TP/CP/SP sharded dWeights vs Triton TP1	0 mismatch
Repeated execution and training/inference forward	0 mismatch

Commands used:

NCCL_IB_DISABLE=1 pytest -q \
  tests/test_qwen_ffn.py \
  tests/distributed/test_qwen_ffn_topology.py

NCCL_IB_DISABLE=1 python benchmarks/benchmark_rocm_ffn.py \
  --warmup 3 \
  --samples 10 \
  --training-samples 5 \
  --output-dir benchmarks/results/pr325_rocm_mi300x

Performance results

All speed numbers compare complete FFN calls. The official baseline is upstream
Transformers Qwen3MLP with unsharded weights and input at TP=1. Distributed
timing uses synchronized wall time and the slowest rank per sample. These rows
compare speed only.

Scope	Deterministic Triton / official Qwen3MLP TP1 median latency
Single GPU, forward, M=1/8/32	9.03-22.86x
Single GPU, forward+backward, M=1/8/32	7.38-11.56x
Distributed, forward, eight TP/CP/SP layouts	8.76-15.79x
Distributed, forward+backward, eight TP/CP/SP layouts	7.45-14.06x

The separate dtype observation runs only official Qwen3MLP TP1:

Candidate	Reference	Relative L2	Max abs	Mean abs
FP16	FP32	6.544e-4 (0.06544%)	2.046e-6	3.742e-7

Full combined report ·
Raw JSON

Communication overlap assessment

The reported implementation serializes dependencies; the measurements do not
claim communication/computation overlap.

• Forward SP all-gather must complete before gate/up computation, and the final
  TP reduction consumes the down-projection output. These are hard dependencies.
• In backward, the gate and up contributions to dHidden are independent until
  their final ordered addition. A future implementation can reduce one on a
  second stream while computing the other.
• That optimization must preserve rank order, the reduction tree, wait points,
  BF16 stage boundaries, and gate-then-up addition order. It is accepted only if
  every TP1 mismatch column remains zero.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 0fbb7c99-630f-4eb5-95aa-e19fa31dfecd

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}