Handcraft Kernel Dispatch for a Fixed Transformer Layer

Inspiration

The disclosed test matrix for a fixed Transformer layer spans a batch size from 1 to 10,000, a head dimension from 8 to 256, and a sequence length from 32 to 100,000. No single kernel is best across that range, and one shape cannot execute at all under the reference formulation: its attention-score tensor alone would require 10.24 TB.

The disclosed shapes are not perturbations of one workload. They are launch-bound, memory-bound, projection-bound, capacity-bound, or infeasible as written. An optimization that wins decisively in one regime is frequently neutral or harmful in another.

That permission to select implementations by shape check became the whole design. We submitted a strict shape-aware dispatcher rather than one universal kernel, and treated the choice of route as a measurement problem instead of a design preference.

What It Does

The project implements a fixed causal Transformer layer faster while remaining numerically indistinguishable from the immutable PyTorch reference under the executable per-element criterion:

$$ |c-r| \le 0.002 \quad \lor \quad |c-r| \le 0.02|r| $$

where (c) is the candidate output and (r) is the reference output.

The dispatcher routes on the complete disclosed tuple:

$$ (B,\, N,\, d_{\mathrm{model}},\, H,\, \mathrm{FFN},\, \mathrm{causal},\, L) $$

and on the runtime numerical contract: input shape, device, dtype, device capability, mask contract, inference state, float32 matmul precision, and TF32 setting.

The final implementation combines:

Compiled float32 PyTorch SDPA.
Strided head views that avoid unnecessary contiguous layout copies.
Causal padding-key mask elimination when the mask is verified to be right-padded.
Packed QKV projection for the two shapes where one combined GEMM is beneficial.
Batch streaming for the extreme-batch case.
Prefix and sample streaming for the extreme long-context case.
Guarded fused Triton polynomial attention for Case 14.

Anything unrecognized or outside the validated runtime contract falls back to reference arithmetic where safe. The two extreme cases reject rather than fall back, because dense fallback would be memory-unsafe.

Results

On an NVIDIA GeForce RTX 5080:

Cases 1-13 passed 65/65 accuracy trials with zero failed elements.
The directly comparable cases achieved a 3.611x geometric-mean speedup over the immutable dense reference.
Speedups range from 1.116x on the projection-bound Case 8 to 8.577x on long-sequence Case 13.
Case 6 reduced peak allocation from 10,672.719 MiB to 2,312.512 MiB, a 78.3% reduction, while improving latency by 2.394x.
Case 13 reduced peak allocation from 2,372.229 MiB to 227.104 MiB, a 90.4% reduction.

Every accepted route was gated against the immutable reference, timed with a paired interleaved schedule that cancels GPU clock drift, and admitted only when its gain exceeded the run's own noise floor.

Case 14: Feasibility With Proven Accuracy

Case 14 is:

$$ B=32,\quad H=16,\quad N=100{,}000,\quad d_{\mathrm{model}}=1024,\quad L=2 $$

At this shape, the score tensor contains:

$$ 5.12 \times 10^{12} $$

values, or 10.24 TB in FP16. The FP32 input and output alone are about 24.4 GiB. Dense attention cannot execute at the target shape.

Three mechanisms make it run:

Prefix streaming: process one valid prefix at a time, bounding the working set.
FP32 interface with FP16 compute: parameters and the interface remain FP32; internal attention compute is FP16.
Polynomial feature-map attention: a degree-2 feature map converts causal attention into a chunked linear scan.

The approximation is guarded at runtime. We estimate the score standard deviation from a 512-row sample, and enable polynomial attention only when:

$$ \sigma \le 0.40 $$

Otherwise, the dispatcher forces exact Flash SDPA, never a quadratic fallback that would OOM.

The first 4,096 tokens are computed exactly because approximation error concentrates at the beginning of the sequence. The running state remains FP32: an FP16 state was faster, but failed at long sequence length.

The full dense reference cannot run for Case 14, so we first validated a linear-memory FP32 oracle against the immutable dense model at (B=1, N=4096), with 0 failures across 4,194,304 elements. At the full target shape, the final candidate achieved:

PASS 5/5 trials
0 failures across 16,384,000,000 output elements
3,643.988 MiB peak allocation
67.095 seconds candidate time
601.975 seconds streamed FP32-oracle time
8.972x diagnostic oracle-to-candidate ratio

The 8.972x ratio compares two linear-memory paths and is explicitly not an immutable dense-baseline speedup. The defensible claim for Case 14 is feasibility with proven accuracy.

How We Built It

We started from the immutable reference benchmark and characterized every shape family analytically before coding. This identified the plausible optimization lever for each regime.

For memory-bound attention, the goal was to avoid materializing the (N \times N) score tensor and to remove unnecessary layout copies. For projection-bound Case 8, attention optimization has a strict Amdahl-law ceiling because dense GEMMs dominate the work. For Case 6, the problem is activation capacity rather than arithmetic. For Case 14, only FLOP reduction can help.

We used PyTorch profiler traces, Perfetto, and NVTX ranges to attribute runtime by operation. We used Codex agents from VS Code over WSL2 for repository inspection, upstream-source inspection, task decomposition, profiler interpretation, hypothesis generation, implementation, conflict review, benchmark orchestration, and documentation.

How We Used AI Tools

AGENTS.md defined the process we followed around that work. We defined this file to ensure that agents worked well, with the latest info and full context of the rapid, ongoing research to prevent conflicts and duplicated effort.

TASK.md defined the scope of the work: the fixed Transformer computation, the 14 disclosed shapes, the executable numerical criterion, and the immutable torch_transformer_benchmark.py reference. We also include the exact split of the tasks documenting each person's responsibility, with a git ignored markdown that stores each contributors role and current progress.

The governing rule was simple: an AI suggestion was a hypothesis, not a result. Nothing entered the dispatcher until it passed the immutable benchmark's numerical checker, multi-seed stress tests, three-reference comparison, and paired timing gate.

This discipline corrected several initially plausible directions. Universal packed QKV was narrowed to Cases 2 and 3 after Case 8 measured only 1.003x end-to-end. A blanket FP16 conclusion was refined after we separated FP16 internal attention compute from a fully FP16 residual path. Case-14 kernels tuned for the wrong GPU architecture were rejected on the target device. An attention-core microbenchmark was rejected after full-model integration measured more than an order of magnitude slower.

Challenges We Faced

The main challenge was distinguishing real gains from plausible but misleading ones.

Universal packed QKV was narrowed to Cases 2 and 3 after the end-to-end Case-8 screen measured only 1.003x. Whole-model FP16 failed on four-layer shapes because repeated residual-stream rounding produced rare errors near zero-valued outputs. max-autotune produced no meaningful gain over reduce-overhead and failed numerical validation on Case 8.

Case-14 kernel configurations tuned for the wrong GPU architecture were rejected on the target device. An attention-core microbenchmark that looked faster was rejected after full-model integration measured more than an order of magnitude slower. A polynomial-attention guard threshold was lowered from 0.60, then 0.45, to 0.40 after kernel changes moved the accuracy boundary.

We also encountered GPU clock variation large enough to fabricate speedups. Identical code measured in separate sessions drifted by 17.5%, and repeated profiles spread by 2.17x. We addressed this with fixed settling periods, alternating AB/BA paired timing, CUDA events, batched forward passes for launch-bound cases, and a noise floor derived from within-pair disagreement.

What We Learned

The project answers the task's question, can AI-assisted development produce efficient shape-specific implementations in limited time, with a qualified yes.

AI gave us breadth and speed of hypothesis generation. The real leverage came from building an evidence gate strong enough to reject incorrect, noisy, architecture-mismatched, or whole-model-neutral optimizations.

We learned that:

Shape-aware dispatch is more valuable than a universal kernel when the workload matrix spans different hardware regimes.
Memory traffic and layout transformations can dominate attention arithmetic.
A kernel-level speedup is not evidence of a whole-model speedup.
Mixed precision must be evaluated across the complete residual path.
Long-context attention requires algorithmic complexity reduction, not only better scheduling.
Negative results are valuable engineering output and save future work.
Performance claims must remain pinned to the GPU, software build, precision mode, and timing methodology that validated them.