Skip to content

[WS2][GEMM][3th merge] feat: add wrapper to ablation matrix - #322

Merged
Flink-ddd merged 14 commits into
RL-Align:testfrom
frank-2077:wrapper_ablation_matrix
Aug 24, 2026
Merged

[WS2][GEMM][3th merge] feat: add wrapper to ablation matrix#322
Flink-ddd merged 14 commits into
RL-Align:testfrom
frank-2077:wrapper_ablation_matrix

Conversation

@frank-2077

Copy link
Copy Markdown
Collaborator

Summary

Expose a Qwen3 dense FFN wrapper so the ablation matrix can switch GEMM backends.

  • Default disable_split_k=True: deterministic det_gemm (no split-K), TP/CP/SP with fixed-tree collectives.
  • disable_split_k=False: torch.matmul (cuBLASLt / CUTLASS, split-K allowed).
qwen3_ffn(x, gate_w, up_w, down_w)                       # det_gemm
qwen3_ffn(x, gate_w, up_w, down_w, disable_split_k=False)  # vendor GEMM

Keep the FP32 running sum for tensor-core speed, and only quantize once at store time.
Keep FP32 leaves for tensor-core speed and BF16-add internal nodes. TP=2 is one commutative add; TP=8 left-fold is a different tree and is expected to diverge.
Replace NCCL with fixed-tree collectives and AllGather CP tokens before det_gemm_db so weight grads match CP=1. Fold FFN coverage into tests/test_qwen_ffn.py.
Default True keeps det_gemm with no split-K. False uses torch.matmul so callers can take the cuBLASLt/CUTLASS path.
@coderabbitai

coderabbitai Bot commented Aug 19, 2026

Copy link
Copy Markdown

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: 921f92e8-dbb8-4cc1-a9bc-150968e7f604

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

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

@Flink-ddd Flink-ddd left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@maxiaosong1124
maxiaosong1124 self-requested a review August 24, 2026 15:58
@Flink-ddd
Flink-ddd merged commit 856e98a into RL-Align:test Aug 24, 2026
5 of 6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants