[perf][gemm]: tile short-token FFN weight gradients

[Flink-ddd]

Summary

Optimize the short-token CUDA deterministic GEMM backward path used by the Qwen3 FFN benchmark in PR #348.

When tokens < 32, the three weight-gradient GEMMs currently transpose the input and fall back to a 16x16 scalar launch. For Qwen3-8B shapes this creates roughly 200k CTAs per GEMM. This change adds a deterministic 128x128 tiled outer-product kernel that:

• reads [tokens, in] directly (no X.T.contiguous() materialization);
• stages the short token dimension in shared memory;
• preserves ascending-token FP32 accumulation and BF16 rounding;
• leaves the existing SM90 path unchanged for tokens >= 32.

Validation

• CPU/mock FFN tests pass.
• Added CUDA regression cases for tokens=1 and tokens=8 to verify canonical contiguous gradients and bitwise parity with the legacy API.
• The current environment has ROCm but no CUDA device, so the H100 benchmark must be run on the target branch as requested. The expected largest gain is in PR bench(ffn): add H100 cuBLAS and vLLM comparisons #348 M=1/8 forward+backward rows.

Benchmark command from PR #348:

CUDA_VISIBLE_DEVICES=0 python benchmarks/benchmark_qwen_ffn_layout.py \
  --tokens 1,8,32,128 --hidden 4096 --intermediate 12288 \
  --seed 20260825 --warmup 3 --samples 20 --training-samples 10 \
  --device-index 0 --output-dir benchmarks/results/qwen_ffn_layout_h100

H100 benchmark results

A single-run benchmark was collected on an NVIDIA H100 80GB HBM3 with CUDA 12.4 and PyTorch 2.4.1+cu124. The comparison uses the benchmark's optimized.median_ms values; both branches used the same Triton compatibility workaround (tl.range(..., flatten=True) removed because the installed Triton does not accept that keyword).

Tokens	Path	Before (ms)	After (ms)	Speedup	Latency change
1	forward	4.9174	4.8701	1.010x	-1.0%
1	forward + backward	20.6728	14.1171	1.464x	-31.7%
8	forward	4.9226	4.8480	1.015x	-1.5%
8	forward + backward	20.5909	14.1231	1.458x	-31.4%
32	forward	4.9083	4.9467	0.992x	+0.8%
32	forward + backward	20.1835	19.6997	1.025x	-2.4%
128	forward	4.8917	5.0628	0.966x	+3.5%
128	forward + backward	19.4240	19.4565	0.998x	+0.2%

All correctness checks passed for tokens 1, 8, 32, and 128. The largest gain is in the intended short-token forward + backward cases: 1.464x at M=1 and 1.458x at M=8. A higher-sample repeat is pending because the GPU is currently busy.

The CUDA 12.4 build-compatibility follow-up is included in commit 8fbd363.

Nsight Systems kernel profile

The H100 trace confirms that the optimization is localized to the intended short-token deterministic dW path. These are aggregated kernel-time breakdown charts from cuda_gpu_kern_sum for the same tokens=1,8 profiling run (CUDA 12.4, H100 80GB HBM3). They are not synthetic timings: the source data came from the before/after .nsys-rep reports on the target node. These are kernel-time breakdown charts, not standard call-stack flame graphs. A true Nsight flame graph requires CPU/Python stack sampling and will be added from a dedicated trace when the H100 is available.

Before

After

The replaced kernel is det_gemm_naive<bf16,true>: 90 launches and 157.657 ms before, versus 90 det_gemm_db_small_k launches and 12.006 ms after (13.13x kernel-time speedup, 92.38% reduction). The existing det_gemm_sm90_kernel row remains effectively unchanged (392.263 ms → 391.452 ms), which is a control for the short-token-specific change.

Detailed source data and calculation: kernel_summary.md.
The kernel-breakdown artifacts are included in commits 949b088 and adb575d.

Nsight Systems event-level details

The screenshots below are direct views from the same before and after.nsys-rep traces.

Before

The selected event is det_gemm_naive<bf16, bool(1)>, with approximately 1.75 ms per launch.

After

The selected event is det_gemm_db_small_k, with approximately 94–178 us per launch.

[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: 5cdbafa0-f2c0-4876-bb84-ca74c5187f16

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