[WS2][GEMM] Add Qwen3 FFN orchestration and consistency validation

[bitborne]

Part of #239.

What changed

• Add a rank-local Qwen3 FFN module:
  • Gate GEMM
  • Up GEMM
  • SwiGLU
  • Down GEMM
• Support two consistent backends:
  • CUDA det_gemm + CUDA SwiGLU
  • Triton det_gemm + Triton SwiGLU
• Add a PyTorch fast path for performance comparison.
• Add an uninterrupted FP32 reference for forward and backward validation.
• Add Qwen3-8B TP=2 local shape coverage:
  • Gate/Up weights: [4096, 6144]
  • Down weight: [6144, 4096]
• Add a benchmark for comparing the consistent and fast paths.

Boundary

This PR only covers rank-local FFN arithmetic.

The module consumes the token tensor provided by the distributed wrapper and returns the local Down projection partial. It does not create process groups or run AllGather, AllReduce, or ReduceScatter.

The distributed TP/CP/SP PRs can wrap this module and apply the required communication outside it.

For backward, the local module produces dX, dGate, dUp, dHidden, and the three local weight gradients. Cross-rank gradient reductions remain outside this PR.

dW reduces over the token dimension, so this PR checks repeat determinism for dW, but does not claim bitwise equality after changing token or CP reduction partitions.

Test plan

Static checks completed:

• Black
• isort
• flake8
• Python compileall
• git diff --check

SM90 validation:

• Build the CUDA extension with the deterministic SM90 GEMM path:

  KERNEL_ALIGN_DET_GEMM_SM90=1 \
    python -m pip install --no-build-isolation -e ".[dev]"

• Run the complete CUDA and Triton forward/backward, tolerance, real-shape, and batch-invariance tests:

  pytest -q tests/test_qwen3_ffn.py

• Compare the CUDA consistent path with the PyTorch fast path using the TP=2 rank-local Qwen3-8B shape:

  python benchmarks/benchmark_qwen3_ffn.py \
    --backend cuda \
    --profile-stages \
    --json-out /tmp/qwen3_ffn_tp2_cuda.json

• Compare the Triton consistent path with the same fast path and input data:

  python benchmarks/benchmark_qwen3_ffn.py \
    --backend triton \
    --profile-stages \
    --json-out /tmp/qwen3_ffn_tp2_triton.json

• Optional unsharded Qwen3-8B experiment:

  python benchmarks/benchmark_qwen3_ffn.py \
    --backend cuda \
    --intermediate-size 12288 \
    --json-out /tmp/qwen3_ffn_unsharded_cuda.json

SM90 results

Pending SM90 run. The benchmark tables and JSON metadata will be added here before marking the PR ready for review.

[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: a381398b-9854-4c43-a8f9-fcc217d25ad1

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.}