perf: tile wide GeGLU rows on SM103

[heiheiha798]

Summary

• Add a 2D column-tiled GeGLU forward/backward path for exact SM103 (blackwell_ultra).
• Dispatch from the input tensor's CUDA device only when next_power_of_2(n_cols) >= 16384; other architectures and narrower rows keep the existing one-row kernels and launches.
• Add forced-path correctness and device-routing coverage plus a paired JSONL kernel benchmark with immutable backward inputs.

Related fork tracker: heiheiha798#1

Engineering record: heiheiha798#2

Details

On an NVIDIA B300 SXM6 AC, the legacy BF16 backward kernel at 8192x14336 used 128 registers per thread. Registers limited it to one block per SM, with 25% theoretical and 21.75% achieved occupancy. NCU also reported 46.02% of cycles with no eligible warp and 29.07% DRAM throughput.

The fixed 1024-element, 4-warp tile was selected from {256, 512, 1024, 2048, 4096} x {4, 8} and checked against 2048x4 over widths 11008, 14336, 18944, and 29568. It was the robust choice across the measured row counts and widths.

The optimized path is exact-SM103 only. B200, consumer Blackwell, Hopper, AMD/NPU, and narrow widths retain the legacy path. This intentionally duplicates the small elementwise kernel body so the legacy launch and compilation surface remain unchanged.

The optional cuTile GeGLU backend in #1250 is separate. This change follows the default-Triton Blackwell column-tiling direction established for SwiGLU in #1271, with a narrower architecture guard because these measurements are B300-specific.

Testing Done

The GPU correctness, performance, and profile results below are from Slurm job 4777 on an NVIDIA B300 SXM6 AC, using the clean pushed head 02697305d0870f4fb574da24fa09382b4b979014, driver 580.126.09, CUDA 13.0, PyTorch 2.13.0+cu130, and Triton 3.7.1.

Targeted correctness:

.venv/bin/pytest -q test/transformers/test_geglu.py
31 passed, 1 warning in 11.04s

The added tests force the legacy and tiled paths on the same GPU and require exact equality for forward, da, and db across FP32, FP16, and BF16; aligned and ragged widths; the 8193-column dispatch boundary; and mocked nonzero CUDA-device routing.

Static checks:

.venv/bin/ruff check --no-cache benchmark/tools/compare_geglu_sm103_tiling.py src/liger_kernel/ops/geglu.py test/transformers/test_geglu.py
.venv/bin/ruff format --check --no-cache benchmark/tools/compare_geglu_sm103_tiling.py src/liger_kernel/ops/geglu.py test/transformers/test_geglu.py
.venv/bin/python -m py_compile benchmark/tools/compare_geglu_sm103_tiling.py src/liger_kernel/ops/geglu.py test/transformers/test_geglu.py
git diff --check
PATH="$PWD/.cache/b300-geglu/checkstyle-env/bin:$PWD/.venv/bin:$PATH" make checkstyle

All passed. make checkstyle checked all 371 files and left them unchanged.

B300 kernel performance

Each result is the median across five alternating legacy/dispatched pairs. Each inner measurement uses 100 ms warmup and 500 ms repetition time. The tracked benchmark restores working a and b from immutable originals before every backward or full launch; restoration is outside the timed region. Earlier dirty-worktree jobs 4732 and 4738 and their mutated-operand results are superseded by job 4777.

python benchmark/tools/compare_geglu_sm103_tiling.py \
  --rows 1024 8192 \
  --cols 4096 11008 14336 18944 29568 \
  --dtype bfloat16 \
  --warmup 100 \
  --rep 500 \
  --rounds 5

rows	cols	legacy fwd ms	tiled fwd ms	legacy bwd ms	tiled bwd ms	full speedup
1024	4096	0.012288	0.012288	0.014336	0.014336	1.000x
1024	11008	0.030720	0.018432	0.036864	0.024576	1.665x
1024	14336	0.038912	0.022528	0.049120	0.028672	1.738x
1024	18944	0.040960	0.028672	0.059392	0.036864	1.567x
1024	29568	0.065504	0.038944	0.104416	0.053248	1.883x
8192	4096	0.043008	0.043008	0.059392	0.059392	1.000x
8192	11008	0.186432	0.102432	0.225344	0.147456	1.685x
8192	14336	0.229376	0.129056	0.303136	0.184224	1.720x
8192	18944	0.255968	0.172032	0.394240	0.253952	1.539x
8192	29568	0.430048	0.262144	0.731136	0.374752	1.824x

At the dispatch boundary, 8192 columns retained the legacy path at 1.000x full speed, while 8193 columns reached 1.371x at 1024 rows and 1.380x at 8192 rows.

NCU

Paired legacy/tiled BF16 backward profiles at 8192x14336:

metric	legacy	tiled
registers/thread	128	38
theoretical occupancy	25.00%	75.00%
achieved occupancy	21.75%	68.55%
no eligible warp	46.02%	11.48%
issue slots busy	50.39%	87.12%
DRAM throughput	29.07%	50.15%
local spill requests	1,638,400	0
profiled duration	504.224 us	291.968 us

Full MLP check

A supplemental paired BF16 LigerGEGLUMLP check used hidden size 4096 and intermediate size 14336. Identical weights and inputs were measured over three alternating rounds with 100 ms warmup and 500 ms repetition time.

tokens	legacy full ms	tiled full ms	speedup
512	0.551872	0.535456	1.031x
2048	1.764256	1.721344	1.025x

The repository-standard python benchmark/scripts/benchmark_geglu.py --model llama_3_8b end-to-end benchmark was not rerun.

• Hardware Type: NVIDIA B300 SXM6 AC (SM103)
• run make test to ensure correctness (targeted GeGLU suite passed as reported above)
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence