Inspiration
Transformer workloads that look similar at the model level can behave very differently on a GPU. A small-batch workload may be dominated by kernel-launch overhead, a long-sequence workload may be memory-bound, while a very large workload may simply exceed available VRAM.
For TikTok TechJam Track 3, we wanted to go beyond finding one fast kernel and instead ask:
Can we automatically match each workload to an optimization strategy based on its actual hardware bottleneck?
What we built
We built a hardware-aware Transformer kernel optimization system for the 14 official Track 3 workloads on an NVIDIA RTX 3060 Laptop GPU.
The system follows this loop:
- Characterize the workload and GPU.
- Estimate the likely bottleneck: compute, memory bandwidth, launch overhead, occupancy, or capacity.
- Select optimization techniques appropriate for that bottleneck.
- Search implementations and kernel parameters.
- Run deterministic correctness checks.
- Benchmark surviving candidates.
- Promote the best measured implementation as the champion for that workload.
The key principle is:
The agent proposes; the deterministic evaluator decides.
Coding agents were used to help propose and implement optimization ideas, but they never decide whether a candidate is accepted. Every implementation must pass the benchmark's correctness requirement before performance is considered.
Optimizations
Different workloads ended up using different strategies, including:
- Triton tiled / streaming attention
- online-softmax attention
- causal tile skipping
- head-dimension blocking
- query partitioning
- FlashAttention-style memory reduction
- microbatching for VRAM-constrained workloads
- shape-specific kernel parameter tuning
Rather than forcing one implementation across all shapes, the system specializes execution based on workload characteristics.
Results
All 14 / 14 official workloads pass correctness.
On our RTX 3060 Laptop GPU:
- 1.532× geometric-mean speedup over the SDPA comparison
- 3.600× speedup on the long-sequence Shape 13
- 5.276× speedup on the capacity-bound Shape 14
- Peak measured hardware utilization of approximately 73.9% MFU
For Shape 13, latency dropped from approximately 117.2 ms to 32.5 ms.
For Shape 14, where the ordinary full workload exceeds available VRAM, the optimized implementation reduced the capacity-safe SDPA runtime from approximately 13.3 s to 2.52 s.
Full per-shape results and benchmarking methodology are available in the technical report in the repository.
Challenges
One major challenge was that there was no universal optimization.
An optimization that helped a long-sequence workload could be useless or even slower for a small workload. This required separating the problem into workload diagnosis, optimization selection, parameter search, correctness verification, and measurement.
Memory capacity was another major constraint. The largest workloads could not run using the ordinary reference execution path on a 6 GB GPU, so we built exact capacity-safe reference and comparison paths without changing the required computation.
We also had to be careful about benchmarking methodology. Compilation is excluded from timing, GPU execution is measured with CUDA events, and correctness is treated as a hard constraint rather than a performance trade-off.
What we learned
The biggest lesson was that GPU optimization is not just about writing a faster kernel.
Performance depends on the interaction between:
workload shape × algorithm × implementation × kernel parameters × hardware
A useful optimization system therefore needs to understand both the computation and the machine, form hypotheses about the bottleneck, and then let actual measurements determine whether those hypotheses were correct.
The result is not one "magic kernel", but an evidence-driven system for specializing Transformer execution to the workload and hardware.
Log in or sign up for Devpost to join the conversation.