Skip to content

perf: add opt-in fused BF16 SwiGLU - #54

Closed
0z5a wants to merge 1 commit into
ViperEkura:mainfrom
0z5a:codex/fuse-small-batch-swiglu-main-0z5a
Closed

perf: add opt-in fused BF16 SwiGLU#54
0z5a wants to merge 1 commit into
ViperEkura:mainfrom
0z5a:codex/fuse-small-batch-swiglu-main-0z5a

Conversation

@0z5a

@0z5a 0z5a commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add a directly callable CUDA BF16 SwiGLU primitive for contiguous M=1..8 decode inputs
  • fuse the up/gate projections, BF16 rounding boundaries, SiLU, and multiply into one launch
  • use CTA weight reuse generally and a measured warp-per-row tiling for the native AstrAI 1B (N,K)=(6912,1536) M=2/4/8 cells
  • route dense MLPs through a safe backend: ASTRAI_SWIGLU=0 is unfused, 1 is explicit opt-in, and default auto currently enables no shape
  • check in the reproducible benchmark script, raw L20 operator/engine/checkpoint evidence, and an explicit InfraSWE repository Draft

The unfused fallback still calls AstrAI's existing linear backend, so independently qualified GEMV bands remain available.

L20 benchmark

CUDA Graph medians for the AstrAI 1B MLP shape:

M torch (ms) GEMV chain (ms) fused (ms) vs best unfused
1 0.02564 0.02298 0.01375 +67.13%
2 0.02484 0.02628 0.01416 +75.40%
4 0.02507 0.03839 0.01806 +38.82%
8 0.02563 0.07007 0.03339 -23.24%

The five-shape matrix also covers LLaMA 2 7B/13B, LLaMA 3 8B, and GPT-NeoX 20B. The wide shapes stay within about -1.2% to +0.9% of the best unfused path, so they are not automatically selected.

Real 24-layer AstrAI checkpoint results through InferenceEngine (scheduler, sampling, and CUDA Graph; A-B-B-A):

Batch unfused (ms/step) forced fused (ms/step) gain
1 4.125 3.925 +5.10%
2 4.245 4.055 +4.69%
4 4.475 4.305 +3.95%

The fused operator uses one CUDA kernel. Direct comparisons stayed within max_abs <= 2.4e-4 with cosine similarity approximately 1.0.

Dispatch decision

Deterministic greedy checkpoint generations changed for M=1, M=2, and M=4 because the fused FP32 reduction order differs. The before/after hashes and outputs are checked in rather than omitted.

Therefore this PR deliberately leaves the auto table empty. Default behavior is unchanged; only ASTRAI_SWIGLU=1 opts into the experimental primitive. A future automatic band must pass both performance and checkpoint-output gates.

Validation

  • Ruff format and import-order checks: pass
  • local full suite on current main: 623 passed, 123 skipped
  • L20 focused SwiGLU suite: 23 passed
  • L20 full built-extension suite on current main: 746 passed
  • InfraSWE Draft engine: 53 passed; Draft validate/resolve: pass
  • SM89 build: CTA variants 36-64 registers / one barrier; warp-row variants 42 registers / zero barriers; zero spills

InfraSWE's frozen formula gives diagnostic ProjectFit 90.95/100 and BenchmarkTrust 95.87/100. These are explicitly non-official: official ProjectFit remains unresolved because the D3 Draft is unsealed and lacks the required fresh-process count, system traces, hidden probes, and verified manifest.

Environment disclosure

Benchmarks ran on NVIDIA L20 (sm_89), PyTorch 2.11.0+cu128, and CUDA 12.8. An existing GPU5 AstrAI service remained resident at about 15.4 GiB but was idle at sampling boundaries. No process or container was stopped.

Related: #53 (independent common-shape GEMV tuning).

@ViperEkura

Copy link
Copy Markdown
Owner

Landed on main as d4a292b (squashed with #53)

@ViperEkura ViperEkura closed this Sep 2, 2026
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.

2 participants