Project Image :)

Shape-Aware Transformer Kernels for NVIDIA T4

Short description

A shape-aware PyTorch and CUDA optimization for a causal Transformer layer that combines packed QKV projections, mask fast paths, and a custom fused residual-add/LayerNorm kernel to achieve up to 1.537x speedup on an NVIDIA Tesla T4 while passing the benchmark's numerical accuracy requirements.

Project overview

Transformers are highly parallel, but a straightforward PyTorch implementation still launches many separate GPU operations. Query, key, and value projections, masking, residual additions, normalization, attention, and feed-forward operations each introduce memory traffic and kernel-launch overhead. These costs are particularly noticeable on smaller shapes, where the GPU spends a significant portion of its time coordinating operations rather than performing useful arithmetic.

This project optimizes the Transformer implementation supplied by the hackathon for an NVIDIA Tesla T4. The aim was not merely to replace the reference with a numerically different fast implementation, but to improve latency while preserving its behavior closely enough for every output element to satisfy:

absolute_error <= 0.002 OR relative_error <= 2%

The final solution passes all 13 runnable published test cases. It achieves a 1.305x average per-case speedup, a 1.297x geometric-mean speedup, and a 1.230x aggregate speedup when the median latencies are combined. This reduces combined median latency by approximately 18.7%, with the best individual case improving from 3.9075 ms to 2.5415 ms, or 1.537x.

How the solution addresses the problem

The solution keeps the benchmark's model interface and parameters compatible with the reference implementation, but replaces unnecessary work with several carefully validated fast paths.

1. Packed QKV projection

The baseline launches three separate linear projections to produce query, key, and value tensors. The optimized implementation concatenates their weights and biases and computes all three with one torch.nn.functional.linear call before splitting the result. This reduces kernel-launch overhead and gives the underlying CUDA libraries a larger operation to process.

FP16 is sensitive to operation ordering and rounding. Stress testing showed that packed QKV could occasionally exceed the accuracy threshold for a small subset of shapes. The implementation therefore uses a shape-aware fallback to the original three-projection path wherever exact baseline behavior is safer.

2. Mask fast paths and caching

The published cases use causal attention and generate masks repeatedly. The optimized implementation caches reusable causal masks and detects the common all-valid-token case. When every token is valid, it avoids redundant key masking and output masking while preserving the result.

3. Custom fused residual-add and LayerNorm CUDA kernel

The central custom GPU kernel fuses:

Residual addition
Mean and variance calculation using Welford's algorithm
LayerNorm normalization
Learned scale and bias application

Fusing these stages avoids writing the residual tensor to global memory and then reading it back in a separate LayerNorm launch. The kernel uses one CUDA block per row and selects a block size based on the hidden dimension. FP16 and FP32 inputs are supported, while LayerNorm statistics are accumulated in FP32 for stability.

Matching the reference implementation required more than using the same formula. The kernel explicitly rounds the residual sum to the output dtype before calculating normalization statistics, reproducing the baseline's FP16 operation boundary. This detail was essential for meeting the strict element-wise accuracy requirement.

4. Shape-aware dispatch with a safe fallback

The fastest implementation was not automatically the most reliable for every shape. Instead of applying the custom kernel universally, I stress-tested each path using multiple random seeds and additional accuracy trials. The fused residual/LayerNorm kernel is enabled only for published cases where it was consistently accurate and beneficial; other inputs use native PyTorch operations.

This approach trades some theoretical generality for dependable results on the disclosed target workload. It also ensures that unseen shapes remain functional through the native fallback.

Development process

I first established FP32 and FP16 baselines using CUDA-event timing. I then introduced one optimization at a time and ran the supplied accuracy checker before benchmarking it. Some implementations appeared correct over five trials but failed on a single element during longer stress tests, so they were removed from dispatch rather than accepted based on average error.

The custom CUDA extension was tested independently across:

Hidden dimensions of 32, 128, and 1,024
Row counts of 128, 2,048, and 8,192
FP16 and FP32
Multiple random seeds
Exact residual comparisons and tolerance-based LayerNorm comparisons

After standalone testing, the extension was integrated into the full Transformer and evaluated across the published shape matrix using Slurm job arrays. The final 13-case run passed every accuracy check and produced speedups ranging from 1.114x to 1.537x.

Development tools used

OpenAI ChatGPT/Codex — assisted with benchmark interpretation, optimization brainstorming, CUDA integration review, numerical debugging, test planning, and documentation. AI suggestions were treated as hypotheses and accepted only after testing on the target GPU.
Git and GitHub — source control and public project hosting.
SSH and Linux command-line tools — remote development and job management on the compute cluster.
Slurm (sbatch, srun, and job arrays) — GPU allocation, reproducible builds, standalone testing, and execution of the published test matrix.
NVIDIA CUDA Compiler (nvcc) — compilation of the custom CUDA kernel.
GCC/G++ and Ninja — C++ compilation and accelerated extension builds.
nvidia-smi and PyTorch CUDA events — environment inspection and GPU latency measurement.

APIs used

No external web or data API is required at runtime.

The implementation uses the following programming interfaces:

PyTorch Python API for model construction, tensor operations, weight management, correctness checking, and benchmarking.
PyTorch C++/CUDA Extension API to expose the custom CUDA kernel as an importable Python module.
ATen tensor interfaces for validating and passing PyTorch tensors through the C++ binding.
CUDA Runtime API for kernel launches, device execution, and error handling.

OpenAI ChatGPT/Codex was used interactively as a development tool; the submitted Transformer does not call an OpenAI API during execution.

Libraries and frameworks used

PyTorch 2.13.0+cu130
CUDA 13.3 toolkit with PyTorch's CUDA 13.0 runtime build
PyTorch torch.utils.cpp_extension
Ninja build system
Python 3.12
C++17 and CUDA C++

The final measurements were collected on an NVIDIA Tesla T4 with compute capability 7.5 and 15 GB of GPU memory.

Datasets and assets used

This project does not train a model and does not use an external dataset.

The main provided asset was the hackathon's torch_transformer_benchmark.py reference benchmark. It supplies the Transformer architecture, accuracy rule, timing procedure, and published test configurations. Test inputs and model weights are synthetic tensors generated by PyTorch using controlled random seeds. The published shape table was used as the workload specification for optimization and validation.

The repository also contains the custom CUDA/C++ source files, standalone kernel tests, Slurm build and benchmark scripts, and summarized performance results produced during this project.

Limitations and future work

The current kernel and dispatch policy are tuned for the Tesla T4 (sm_75). Other GPUs can compile the code by changing TORCH_CUDA_ARCH_LIST, but should be profiled and retuned. The extension implements an inference forward path and does not include a custom backward pass for training. It also fuses only residual addition and LayerNorm; attention, GEMMs, softmax, and the feed-forward network still depend on PyTorch and CUDA library kernels.

The supplied baseline materializes the full attention score matrix, so its memory usage grows quadratically with sequence length. Published case 14 uses a sequence length of 100,000. Its score tensor would contain 5.12 trillion FP16 values and require about 10.24 TB before temporary allocations, far beyond the T4's memory. Because the reference itself cannot run that case, it was not possible to produce a fair accuracy or latency comparison without changing the benchmark methodology.

With more time, I would implement a tiled or FlashAttention-style causal kernel that never materializes the full score matrix, extend fusion to QKV reshaping and the feed-forward path, add architecture-specific autotuning, profile the kernels with Nsight Systems and Nsight Compute, and add backward support and automated multi-GPU testing.