[perf][gemm]: optimize long-token FFN deterministic path

[maxiaosong1124]

Summary

Optimize the long-token CUDA deterministic GEMM path used by the Qwen3 FFN, building on PR #351's short-token dW optimization.

The PR351 short-token path is intentionally unchanged: when tokens < 32, det_gemm_db_small_k is still used. This change targets the tokens >= 32 SM90 path.

The long-token path now:

• packs the deterministic tree state into BF16x2 values;
• uses vectorized BF16x2 conversion and __hadd2 tree merges;
• computes the mid-split merge count directly instead of carrying interval metadata;
• limits the TMA barrier wait to the issuing thread while retaining the CTA synchronization;
• preserves the fixed reduction order and deterministic output behavior.

Validation

• tests/test_det_gemm.py tests/test_qwen_ffn.py: 69 passed
• Short-token dispatch remains on the PR351 det_gemm_db_small_k kernel.
• Benchmark and profile artifacts are included in benchmarks/results/qwen_ffn_h100_trace/.

H100 benchmark results

Hardware: NVIDIA H100 80GB HBM3, BF16, Qwen3-8B FFN (hidden=4096, intermediate=12288).

CUDA-event latency uses 5 warmups, 20 forward iterations, and 10 forward+backward iterations. Values are the median of three trial medians. Baseline is origin/pr-351 @ adb575d.

Tokens	Det forward before (ms)	Det forward after (ms)	Speedup	Det fwd+bwd before (ms)	Det fwd+bwd after (ms)	Speedup
1024	6.5133	4.8531	1.34x	19.9793	15.1710	1.32x
2048	12.8660	10.0312	1.28x	39.8176	31.3641	1.27x
4096	25.4928	20.0340	1.27x	78.7306	62.3420	1.26x
8192	50.0523	40.4205	1.24x	155.3476	126.4405	1.23x

Nsight Systems kernel attribution

The cuda_gpu_kern_sum aggregates show the deterministic GEMM reduction remains positive across all long-token scenarios:

Tokens	Det GEMM before (ms)	Det GEMM after (ms)	Reduction
1024	212.642	155.672	26.8%
2048	431.163	329.599	23.6%
4096	868.660	674.838	22.3%
8192	1713.346	1378.469	19.5%

Deterministic GEMM accounts for approximately 85%–91% of GPU kernel time in these profiles and remains the primary FFN bottleneck. The percentage benefit decreases with token count as the long-token GEMM work dominates fixed launch/reduction overheads.

Reproduction

PYTHONPATH=$PWD python benchmarks/results/qwen_ffn_h100_trace/bench_long_token_std.py \
  1024 2048 4096 8192

The raw .nsys-rep files are retained on the profiling node; the SVGs and summary are included as review artifacts.

[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: d9943cf5-81f2-4099-8e45-6755e04a2a62

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

---

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

[maxiaosong1124]

Superseded by #354, which uses the correctly named perf/ffn-long-token-backward branch.