bench(ffn): add H100 cuBLAS and vLLM comparisons

[frank-2077]

Summary

This PR removes avoidable tensor-layout materializations from the deterministic
SM90 GEMMs used by the Qwen3 FFN while preserving its bitwise-consistency
contract. The optimized path consumes canonical [out, in] weights directly
and writes canonical contiguous weight gradients directly, without changing
the fixed MMA/K-tree reduction or TP/CP/SP collective order.

On an NVIDIA H100 80GB, the layout change improves the stripped Qwen3-8B FFN
by 2.36–2.40x in forward and 1.22–1.70x in forward + backward versus an
in-process replay of the previous layout contract. All deterministic outputs,
input gradients, and weight gradients remain bitwise identical. The PR also
adds reproducible cuBLAS-preferred and vLLM batch-invariant reference paths,
raw H100 results, and separate forward and forward + backward figures.

Note

These are operator-only, single-GPU measurements of a stripped TP=1,
three-GEMM FFN core. They do not load or benchmark a model checkpoint,
tokenizer, dataset, scheduler, communication path, or serving engine. The
vLLM reference is forward-only and is not an end-to-end vLLM benchmark.

Comparison contract

Question	Baseline / reference	Candidate	Acceptance / metric
Does the layout change preserve deterministic results?	Replayed deterministic layout contract from 7207ebd and the public deterministic qwen3_ffn path	Optimized canonical-layout path	Zero raw BF16 mismatches for inference output, training output, dHidden, and all three weight gradients
How much layout overhead was removed?	Preserved legacy CUDA APIs plus the removed weight and weight-gradient transpose/copy materializations	New layout-aware CUDA APIs	Median full-FFN latency and direct_copy_kernel launch counts
How far is the deterministic operator from production-oriented GEMMs?	PyTorch CUDA BLAS (cuBLAS preferred) and vLLM's batch-invariant persistent Triton matmul	Optimized deterministic path	Median latency plus numerical sanity checks; vLLM additionally requires identical-first-row batch invariance

cuBLAS and vLLM use different reduction orders and are numerical/performance
references, not bitwise references. The legacy comparison is replayed in the
same process and build so that it isolates layout materialization cost from
cross-build noise.

Design

• Add a RHS-transposed GEMM entry point that consumes canonical [out, in]
  weights directly, avoiding the previous weight.t().contiguous() round trip.
• Add a transposed-output weight-gradient entry point that writes the logical
  GEMM result directly as canonical contiguous [out, in], removing three
  large post-GEMM transposes.
• Reuse the existing physical [out, in] weight representation for
  input-gradient GEMMs.
• Support both normal and transposed output indexing in the SM90 and scalar
  fallback kernels.
• Preserve the existing M-padding/cropping behavior. For already-contiguous
  TMA operands, clone only when the base pointer is not 16-byte aligned.
• Keep the legacy GEMM entry points available and preserve the original MMA
  instruction order, fixed K-tree, no-Split-K behavior, and TP/CP/SP
  communication/reduction order.
• Vendor vLLM's batch-invariant persistent matmul and complete configuration
  table under benchmarks/, pinned to vLLM PR
  #53247 merge commit
  7797b602.

Operator test matrix

Surface	Cases	Contract	Result
Layout-aware CUDA GEMMs	Aligned SM90, M-padding/cropping, scalar fallback, and unaligned contiguous TMA views	New entry points match the legacy APIs bitwise; transposed weight gradients are canonical and contiguous	Pass
Qwen3 FFN, TP=1	BF16, H=4096, I=12288, M=1/8/32/128	Inference, training, dHidden, and all weight gradients match the replayed legacy path bitwise	Pass
vLLM forward reference	Identical first input row at M=1/8/32/128	Identical first output row in raw BF16 bits	Pass, 0 mismatches
Distributed FFN	2-GPU TP, TP+SP, and TP/CP/SP cross-configuration checks	Existing consistency contract remains unchanged	Pass, 3 tests
Single-GPU suite	GEMM, simulated TP, benchmark artifact, and Qwen FFN tests	Functional, layout, report, and chart checks	Pass, 67 tests; 14 skipped

H100 environment

Field	Value
GPU	NVIDIA H100 80GB HBM3
Compute capability	SM90 (9.0)
Driver	595.71.05
Python	3.11.15
PyTorch	2.13.0+cu130
CUDA	13.0
Triton	3.7.1
Dtype / shape	BF16; H=4096, I=12288, M=1/8/32/128
Matmul controls	TF32 disabled; float32 matmul precision highest; preferred BLAS library Cublas
Sampling	3 warmups; 20 forward samples; 10 forward + backward samples per path
Measurement time	2026-08-27T13:36:55Z
Kernel commit at measurement	15ca2627f7b1843901dc59877d6db623399dd86a
Benchmark/results commit	814a99a

results.json records git_dirty=true because the benchmark, vendored vLLM
reference, and generated artifacts had not yet been committed when the run was
captured. It also records SHA-256 hashes for the benchmark, vLLM kernel,
configuration table, and built CUDA extension.

Correctness results

All deterministic entries below are raw BF16 mismatch counts. Timing is
aborted if any count is non-zero.

M	Inference output	Training output	dHidden	dGateW	dUpW	dDownW
1	0	0	0	0	0	0
8	0	0	0	0	0	0
32	0	0	0	0	0	0
128	0	0	0	0	0	0

The optimized timing core also matches the public deterministic qwen3_ffn
path bitwise for every value above, and every returned weight gradient has its
parameter-matching shape, stride, and contiguous layout.

CUDA BLAS is not expected to match bitwise. Before timing, every output and
gradient must remain below the benchmark gates of 2% relative-L2 error and 5%
normalized maximum error.

M	Output rel. L2	dHidden rel. L2	Max dWeight rel. L2	Max normalized error
1	0.955%	1.001%	0.807%	1.500%
8	0.960%	1.000%	0.809%	1.242%
32	0.965%	0.997%	0.811%	1.379%
128	0.965%	0.998%	0.848%	1.235%

The vLLM path is checked numerically against the optimized deterministic path
and separately checked for identical-first-row batch invariance.

M	Relative L2	Normalized max error	Full-output BF16 mismatch fraction	First-row mismatch vs. M=1
1	0.954%	1.000%	81.567%	0
8	0.960%	0.990%	82.443%	0
32	0.965%	0.837%	82.053%	0
128	0.965%	0.948%	82.207%	0

Reproduction commands (the recorded benchmark exposed physical GPU 5 as
logical cuda:0; select any target H100 through CUDA_VISIBLE_DEVICES):

KERNEL_ALIGN_DET_GEMM_SM90=1 MAX_JOBS=2 \
  python setup.py build_ext --inplace

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

CUDA_VISIBLE_DEVICES=0 python -m pytest -q \
  tests/test_det_gemm.py \
  tests/distributed/test_det_gemm_simulated_tp.py \
  tests/test_benchmark_qwen_ffn_layout.py \
  tests/test_qwen_ffn.py
# 67 passed, 14 skipped

CUDA_VISIBLE_DEVICES=0,1 python -m pytest -q \
  tests/test_qwen_ffn.py::test_qwen_ffn_tp_correctness_and_batch_invariance \
  tests/test_qwen_ffn.py::test_qwen_ffn_tp_sp_correctness_and_batch_invariance \
  tests/test_qwen_ffn.py::test_qwen_ffn_world2_tp_sp_and_cp_match_tp1_cp1_bitwise
# 3 passed

Performance results

CUDA events measure complete stripped-FFN calls. Forward samples use a
prefix-balanced cycle over all 24 execution-order permutations of legacy,
optimized, CUDA BLAS, and vLLM. Forward + backward uses the balanced six-order
cycle over the three training-capable paths. Tables and figures report median
latency; the JSON artifact retains p95/min/max and every raw sample.

Forward

M	Replayed legacy	Optimized deterministic	vLLM batch-invariant	CUDA BLAS	Layout speedup	Optimized / vLLM	Optimized / CUDA BLAS
1	4.0067 ms	1.6985 ms	0.4206 ms	0.1201 ms	2.36x	4.04x	14.14x
8	3.9980 ms	1.6922 ms	0.3890 ms	0.1223 ms	2.36x	4.35x	13.84x
32	3.9976 ms	1.6896 ms	0.3544 ms	0.1300 ms	2.37x	4.77x	13.00x
128	3.9509 ms	1.6476 ms	0.2043 ms	0.1312 ms	2.40x	8.06x	12.56x

Forward + backward

M	Replayed legacy	Optimized deterministic	CUDA BLAS	Layout speedup	Optimized / CUDA BLAS
1	12.5711 ms	10.2401 ms	0.4988 ms	1.23x	20.53x
8	12.6107 ms	10.2956 ms	0.3841 ms	1.22x	26.80x
32	9.0592 ms	5.3159 ms	0.3966 ms	1.70x	13.40x
128	9.4869 ms	5.7568 ms	0.4147 ms	1.65x	13.88x

Layout speedup = replayed legacy / optimized. Ratios above 1 in the two
Optimized / reference columns mean that the reference in the denominator is
faster; they are not speedup claims for the deterministic path.

Layout-copy profile

One warmed M=128 call was captured with torch.profiler:

Direction	Legacy direct copies	Optimized direct copies	Legacy GEMMs	Optimized GEMMs
Forward	6	0	3	3
Forward + backward	18	9	9	9

Full H100 report ·
Raw JSON ·
Summary CSV

Interpretation and limitations

• The vLLM reference uses the same stripped TP=1 three-GEMM topology, weight
  transpose views, seeded BF16 inputs, and RL-Kernel SwiGLU as the other paths.
  It measures the vendored matmul kernels, not vLLM serving.
• The Qwen3-8B projection shapes (N,K)=(12288,4096) and (4096,12288) do
  not match PR #53247's BF16 tuned table. All vLLM rows therefore use the
  upstream default BM=128, BN=128, BK=64, 8 warps, 3 stages; these are not
  shape-tuned PR #53247 results.
• PyTorch reports Cublas as the preferred BLAS library in this environment,
  but torch.matmul does not guarantee a specific cuBLAS algorithm.
• The remaining gap to both production references is visible in the tables and
  charts. Further GEMM tiling/scheduling work is separate from the layout-only
  optimization in this PR.

[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: 2c780e8c-ce73-4a32-a084-c6fd89cef36b

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