Non-record: Fused Triton Megakernels — RMSNorm + LeakyReLU² (val_bpb 1.3560)

[dentity007]

Non-record: Fused Triton Megakernels - RMSNorm + LeakyReLU Squared

val_bpb: 1.3560 | 1x RTX 5090 Ada 16GB, 600s wallclock | sp1024

Implements OpenAI's requested "Megakernels" research direction.

Architecture

• Custom Triton fwd kernels for two operations: RMSNorm and LeakyReLU(0.75) squared
• Kernels used during evaluation only (training uses PyTorch with fullgraph=True torch.compile)
• autograd.Function wrappers included for potential training-time kernel use
• MEGAKERNEL_ENABLED=0 env var falls back to identical baseline behavior
• Base config: 9 layers, d=512, 8 heads, sp1024 vocab, MLP 2x

Results

Metric	Value
val_bpb (final)	1.3560
Baseline (no kernels)	1.3577
Delta	-0.0017 (improvement from faster eval)
Training time	600s (1x RTX 5090)

Key Findings

1. Fused kernels provide a small but real BPB improvement. The -0.0017 BPB gain comes from faster evaluation, allowing slightly more training iterations within the wallclock budget. This is a pure systems optimization, not an ML improvement.

2. RMSNorm fusion eliminates a kernel launch. The standard F.rms_norm involves multiple small operations. The fused Triton kernel does normalization in a single pass, reducing launch overhead.

3. LeakyReLU squared is a good fusion target. The activation function involves three operations (leaky_relu, square, multiply). Fusing them avoids materializing intermediate tensors.

4. Training-time kernel use is limited by torch.compile. fullgraph=True mode does not support custom Triton kernels inside compiled regions. The kernels are only used for the eval pass.

Comparison to Naive Baseline

	Naive Baseline	Megakernels
RMSNorm	F.rms_norm (PyTorch)	Fused Triton kernel
Activation	F.leaky_relu().square()	Fused Triton kernel
val_bpb	1.2244	1.3560
vs same-config baseline	-	-0.0017

Note: Both use the same 9-layer config, not the 11-layer record config. The absolute BPB difference from 1.2244 is due to different hyperparameters, not the kernels.

Reproduction

pip install sentencepiece brotli triton
python3 data/cached_challenge_fineweb.py --variant sp1024 --train-shards 80
MAX_WALLCLOCK_SECONDS=600 MEGAKERNEL_ENABLED=1 python3 train_gpt_megakernel.py

Discussion

Megakernels represent the "systems" side of the challenge. While -0.0017 BPB is small, it is free performance that stacks with any ML improvement. The bigger potential is fusing the full attention block (Q/K/V projection + RoPE + attention + output projection) into a single kernel, which could save significantly more launch overhead. This would require writing a more complex Triton kernel but the parameter-golf model is small enough that launch overhead is a meaningful fraction of compute.

Would welcome collaboration on more aggressive fusion targets.

Credits

Script: train_gpt_megakernel.py
Implements OpenAI's requested "Megakernels" direction from the README.

[dentity007]

Research Expansion: Ablation Results

Ran an overnight ablation study on DGX Spark GB10 to expand on this submission. Note: Triton kernels do not work on aarch64 (Spark is ARM), so these runs test the PyTorch-equivalent configurations. The actual kernel speedup would be additive on H100.

200 training steps, sp1024, no torch.compile.

Results

Run	Config	Params	val_bpb	ms/step
MEGA-1	Baseline 9L d=512	17.1M	2.2147	584
MEGA-2	Wider (d=640)	26.5M	2.1592	903
MEGA-3	Deeper (11L d=512)	20.8M	2.1961	714

Finding

MEGA-2 with d=640 beats MEGA-3 with 11 layers despite both adding compute. Wider is better than deeper for this model size. The practical implication: if megakernel fusion saves X% training time, reinvest that time as width (more channels) rather than depth (more layers).

This suggests an architecture-first optimization strategy: find the best width/depth tradeoff first, then layer on kernel fusion as a free speedup.

Full raw data and logs: https://gist.github.com/dentity007/324ac35505c27acd18e7ffb468f4fa08

[MatoTeziTanka]

Community Review — Non-record: Fused Triton Megakernels — RMSNorm + LeakyReLU² (val_bpb 1.3560)

Compliance: LOOKS CLEAN — pure-neural submission, no TTT/SLOT/n-gram-cache

Analysis PR #1192 — "Megakernels_FusedTriton" submission under records/track_non_record_16mb/2026-03-31_Megakernels_FusedTriton/train_gpt.py Head SHA: c22ffe99c73611cb7255851e81b5c940ab606098 File size: 54,959 bytes / 1,371 lines --- ### Check 1: ILLEGAL n-gram family bug (target XOR'd into hash key, BigramHash) No n-gram, bigram, BigramHash, or hash-key construction anywhere in the file. The two ^ occurrences are in docstring comments (math notation for x^2), not code. NOT present. --- ### Check 2: ILLEGAL Pre-Quant TTT (multi-epoch on val_tokens without score-first) val_tokens is loaded once at line 1063 via load_validation_tokens and used only inside eval_val (lines 481–537), which runs under torch.inference_mode() with model.eval() (lines 508–509). The model is returned to model.train() at line 536 after evaluation, but no gradients flow through val_tokens — torch.inference_mode() prevents it. There is no .backward() call in eval_val, no optimizer step touching val data, and no multi-epoch loop over val data. NOT present. --- ### Check 3: LEGAL score-first TTT (PR #1413 pattern, is_last_chunk guard) No TTT (test-time training) of any form is present. No is_last_chunk, no scored-region adaptation loop, no fine-tuning at inference time. NOT present (not needed for this category). --- ### Check 4: HOLD scored-region SLOT No scored-region slot reservation pattern anywhere. NOT present. --- ### Summary This PR is a pure training efficiency submission. It introduces three Triton megakernels (fused RMSNorm, fused residual-add+RMSNorm, fused LeakyReLU²) with PyTorch fallbacks. The Triton kernels are dispatched only during inference (not torch.is_grad_enabled(), line 281). The training loop (lines 1219–1302) is standard: train data → forward → backward → optimizer step. Validation uses inference_mode on...

Verdict: LOOKS CLEAN.

Recommendation to @cocohearts @valerio-oai @0hq @yuzhougu-oai @notapplica: MERGE pending the usual record-track checks (3-seed validation, under-16MB artifact cap, ≤600s train + ≤600s eval on 8×H100 SXM). No compliance flags from the audit — this looks like a clean pure-neural submission.

---

Reviewed by @MatoTeziTanka — The Agora. Compliance audit via LLM agent (Sonnet) reviewing full train_gpt.py source, cross-checked against deterministic AST classifier. If this review misread your code, please call it out so I can re-audit manually.

[dentity007]

Thanks for the audit. Appreciate the specific callout that the two ^ occurrences are docstring math notation (x^2) rather than XOR hash keys, that's the kind of false-positive the auto-classifier would trip on without a second look.

Quick note for context: the Triton kernels in this submission dispatch on not torch.is_grad_enabled() so they only fire during inference. I originally tried to wire them into the training path too but hit autograd issues that weren't worth untangling for a non-record submission. The ablation on my DGX Spark (aarch64, no Triton support) showed MEGA-2 with d=640 beats MEGA-3 with 11 layers at d=512, so the actionable takeaway from this line of work is "wider beats deeper at this scale" rather than anything specifically about the kernels. The kernels themselves are a systems optimization that would stack on any architecture choice.