Learning-Guided Transformer Program Search

cover
shape14_capacity
workload_landscape
shape14_streaming
workload_sensitivity
performance_summary
useful_throughput
architecture_overview
search_evidence
deployed_program_matrix

Inspiration

The official benchmark evaluates the same Transformer semantics across 14 very different workload Shapes. Batch size ranges from 1 to 10,000, model width from 32 to 1,024, and sequence length from 32 to 100,000.

These Shapes stress different parts of a GPU system: kernel-launch overhead, matrix throughput, memory traffic, tensor layouts, attention working sets, and device capacity. This means that one universal implementation is unlikely to perform well everywhere.

We were inspired to build a system that could learn from real GPU measurements and discover a specialized Transformer execution program for each workload Shape.

What it does

Learning-Guided Transformer Program Search generates, validates, benchmarks, and deploys complete Transformer execution programs for a target GPU.

Instead of searching for only one kernel, it composes programs from:

PyTorch operators and scaled-dot-product attention backends
Custom Triton kernels
Tensor layouts and materialization strategies
FP32 and internal mixed-precision paths
Operator and cross-operator fusion
Eager execution, TorchInductor, and CUDA Graphs
Kernel launch parameters, batch tiling, and microbatching

Every candidate must pass static legality checks, execution-path validation, and numerical correctness checks before its performance can be considered.

Promising candidates move through Screen, Enhanced, and Formal measurement stages. A challenger is deployed only when paired GPU measurements show a repeatable improvement over the current program.

How we built it

We built the project in Python using PyTorch, Triton, Optuna, CUDA Events, SQLite, and structured JSON/JSONL evidence.

Each candidate is represented by a typed ConfigSpec containing both structural program choices and scheduling parameters. PlanBuilder either compiles it into an immutable execution plan or rejects it with explicit violations. This prevents an invalid candidate or hidden fallback from being reported as a successful optimization.

For Shapes 1–13, we use constraint-aware, branch-local TPE combined with fixed-budget survivor racing. Screen results guide the sampler, selected candidates receive Enhanced measurements, and one locked challenger enters a paired Formal comparison.

Shape 14 has a sequence length of 100,000 and cannot materialize a dense attention matrix on a 16 GB-class GPU. We therefore created a separate streamed causal-attention runtime that processes ordered attention blocks and microbatches without storing a global S × S tensor.

GPU operations are protected by an exclusive device lease, and each Shape is evaluated in a fresh process to reduce interference from allocator state, compilation caches, CUDA Graphs, and earlier model executions.

OpenAI Codex was used for implementation, refactoring, testing, technical research, and multi-agent review. Human guidance defined the optimization goals, numerical constraints, hardware boundaries, and promotion decisions. Correctness checking, benchmarking, comparison, and deployment remain deterministic program operations.

Challenges we ran into

One major challenge was the size and conditional nature of the search space. Many combinations of precision, layout, fusion, kernel, and runtime were invalid or meaningless. Typed program construction and explicit execution-trace validation allowed us to reject these combinations before spending GPU time on them.

Reliable benchmarking was another challenge. Sub-millisecond workloads were highly sensitive to background GPU activity, measurement order, compilation state, and cached allocations. We addressed this with CUDA-event timing, fixed comparison inputs, fresh processes, exclusive GPU access, and interleaved incumbent-versus-challenger measurements.

Shape 14 required a fundamentally different approach. Treating it as a larger resident Shape produced infeasible memory requirements and extremely slow reference execution. We separated its search space, measurement protocol, runtime, and deployment lifecycle, then implemented streamed attention and bounded microbatching.

We also had to balance exploration with a strict time budget. Broad search discovered new program families, but repeatedly measuring weak branches wasted valuable GPU time. Survivor racing and persistent studies allowed later optimization cycles to reuse evidence and focus on more promising structures.

Accomplishments that we're proud of

On the declared NVIDIA GeForce RTX 4080 snapshot:

Shapes 1–13 passed the official-compatible correctness comparator.
Shapes 1–13 achieved a 14.49x geometric-mean speedup over the measured PyTorch baseline.
Individual resident-Shape speedups ranged from 3.41x to 36.09x.
The streamed Shape 14 path completed the full logical 100,000-token workload in 17.24 seconds.
Shape 14 used 6.56 GiB peak allocated memory for each inner B=2 streamed forward.

We are also proud that the result is supported by a complete optimization loop rather than a collection of manually selected kernels. Across four resident optimization cycles, the system reduced 3,933 Screen-stage entries to 381 Enhanced entries, 50 Formal comparisons, and six evidence-supported deployment updates.

The final deployment contains different generated programs for different Shapes, combining mature PyTorch operations with focused Triton kernels, fusion strategies, precision plans, and runtime choices.

These results are local engineering measurements on the disclosed RTX 4080 environment. They are not presented as an official competition score or as a cross-hardware performance claim.

What we learned

We learned that GPU optimization is an end-to-end program problem, not simply a search for one universally fast kernel. A mature library operator may be best for one Shape, while a custom Triton kernel, fused boundary, compiled graph, or streamed runtime wins another.

We also learned that measurement infrastructure is part of the optimization system. A fast result is not useful unless the candidate is correct, the executed path matches the intended path, the environment is identified, and the improvement survives a controlled comparison.

Finally, extreme workloads should not always be forced into the same optimization lifecycle. Shape 14 became tractable only after we treated device capacity as a first-class constraint and designed a dedicated streamed execution regime.

What's next for Learning-Guided Transformer Program Search

Next, we plan to evaluate the system on additional GPU architectures and build a cross-device performance model that can warm-start search on new hardware.

We would also expand the program vocabulary with additional attention, projection, normalization, and cross-operator fusion primitives; improve profiling-guided candidate generation; and evaluate the deployment registry under concurrent serving workloads.

For Shape 14, we plan to strengthen full-scale validation when complete official large-shape input/output evidence becomes available and explore new streamed fusion strategies beyond the current attention implementation.