FastTransformer: Breaking the Compiler Ceiling via Architectural Co-Design

1. I Hit the Compiler Ceiling

For weeks, I had been building a from-scratch pure CUDA engine for a standard GPT-2 decoder ($B=32, T=256, C=256, L=6, H=8$). I went deep into hardware-level micro-optimizations:

• Allocated all weights and gradients into a single 1D contiguous GPU buffer to eliminate runtime heap allocation overhead.
• Wrote vectorized 128-bit (float4) LayerNorm kernels with in-register warp shuffle reductions (__shfl_down_sync).
• Fused causal softmax directly into warp registers to eliminate global DRAM round-trips.
• Wrote a single-pass fused AdamW kernel updating first and second moments in registers.

The result? My hand-written engine reached 107.21 ms ($76{,}411\text{ tok/s}$). PyTorch Eager sat at 125.20 ms. Hand-written CUDA was outrunning PyTorch by a wide margin.

Then I turned on torch.compile(mode="reduce-overhead"). TorchInductor and Triton lowered the autograd graph, fused point-wise backward operations, and clocked in at 101.30 ms ($80{,}869\text{ tok/s}$).

I was within ~5.9 ms of beating the compiler. But closing that final gap was like hitting concrete. I was optimizing kernels for fractions of a millisecond, wrestling with shared memory bank conflicts and cuBLAS stream bubbles.

Figure 1: Full-step training latency on NVIDIA Tesla T4 ($B=32, T=256, L=6, d=256$). Pure CUDA handily beats PyTorch Eager and CUDA Graphs, but collides with torch.compile's automated graph lowering and memory fusions at ~101 ms.

That’s when I stepped back and asked a different question:

2. Profiling the Workload

Instead of profiling kernel launch overhead, I calculated the exact arithmetic operations (FLOPs) and memory transfers for every single layer in the block.

At batch size $B=32$ and sequence length $T=256$, each training step processes $N_{\text{tok}} = 8{,}192$ tokens. Here is where the FLOPs actually go in a 6-layer GPT block:

Figure 2: Full-step phase latency breakdown across frameworks on Tesla T4. Notice how backward and MLP phases dominate execution time.

3. The Surprising Bottleneck

Look at that breakdown. At sequence length $T = 256$:

• The attention matrix multiplication ($\mathbf{Q}\mathbf{K}^T$ and $\mathbf{A}\mathbf{V}$) accounts for only 4.0% of total block compute!
• The MLP and linear projections account for 96.0% of the block compute and memory traffic!

All that community obsession with optimizing attention dot products? At context length 256, it was tackling a component that represented 4% of the runtime.

Meanwhile, standard GPT-2 had two massive structural flaws:

1. DRAM Traffic from Attention Materialization: Standard attention insists on computing and writing the full $(B, H, T, T)$ probability tensor to global VRAM so autograd can read it during backward propagation. That’s $805\text{ MB}$ of memory traffic per step across 6 layers!
2. The 4× MLP Compute Monster: Expanding the hidden dimension by $4\times$ ($256 \to 1024 \to 256$) required $8.59\text{ GFLOPs}$ per layer. The GPU was spending $64\%$ of its life multiplying FFN weights.

Changing the architecture produces a fundamentally larger performance gain than further micro-optimizing the original kernels.

4. Architectural Co-Design

Instead of accepting the 2017 Transformer graph as immutable sacred text, what if we co-designed the architecture for modern GPU memory hierarchies?

I set four explicit hardware engineering requirements:

1. Slash Key-Value memory traffic without dropping multi-head representation capacity.
2. Eliminate the $(B, H, T, T)$ attention score tensor from global DRAM entirely.
3. Cut the MLP compute cost in half without sacrificing representational capacity.
4. Eliminate two-pass normalization barriers that cause warp synchronization bubbles.

5. FastTransformer Architecture

To solve these four bottlenecks, I designed FastTransformer:

5.1 Multi-Query Attention (MQA)

Instead of having 8 separate Key heads and 8 Value heads, FastTransformer uses 8 Query heads ($d=32$) and shares 1 Key head and 1 Value head across all queries:

This immediately slashed KV projection weights and activation storage by $58.3\%$! In cuBLAS, this is executed by setting $\text{Stride } B = 0$, broadcasting the single Key/Value head across all Query heads with zero memory replication.

5.2 Native Hardware-Fused SDPA

Instead of materializing $\mathbf{P} = \operatorname{Softmax}(\mathbf{Q}\mathbf{K}^T / \sqrt{d})$ in global memory, FastTransformer calls fused Scaled Dot-Product Attention (SDPA). On our Tesla T4, PyTorch routes this directly into cuDNN / FlashAttention hardware kernels. The attention matrix is computed in fast on-chip SRAM; zero bytes of $(B, H, T, T)$ scratchpad touch global VRAM.

5.3 Lean 2× Fused MLP

We replaced the bloated $4\times$ MLP ($d_{\text{ff}} = 1024$) with a lean $2\times$ expansion ($d_{\text{ff}} = 512$):

This halved the computation of the most expensive phase ($64\%$ of block FLOPs) from $8.59\text{ GFLOPs} \to 4.30\text{ GFLOPs}$, while cutting intermediate activation storage in half.

5.4 Vectorized Pre-RMSNorm

Replaced LayerNorm with Root Mean Square Normalization (RMSNorm). RMSNorm removes the mean-centering calculation:

This allows a single-pass 128-bit vectorized CUDA kernel that updates activations without warp-synchronization pipeline bubbles.

import torch
import torch.nn as nn
import torch.nn.functional as F

class RMSNorm(nn.Module):
    """Vectorized Root Mean Square Normalization (removes mean centering)."""
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm_x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm_x * self.weight

class MultiQueryAttention(nn.Module):
    """MQA with 8 Query heads and 1 shared Key/Value head for zero-copy broadcasting."""
    def __init__(self, d_model: int = 256, n_heads: int = 8):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads  # 32
        
        # Projection slashes weights from 256x768 to 256x320 (-58.3%)
        self.qkv_proj = nn.Linear(d_model, d_model + 2 * self.d_head, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        qkv = self.qkv_proj(x)
        q, k, v = torch.split(qkv, [self.d_model, self.d_head, self.d_head], dim=-1)
        
        q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)  # (B, H, T, d)
        k = k.view(B, T, 1, self.d_head).transpose(1, 2)            # (B, 1, T, d)
        v = v.view(B, T, 1, self.d_head).transpose(1, 2)            # (B, 1, T, d)

        # Fused SDPA in SRAM: zero (B, H, T, T) DRAM allocation
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out)

class FastTransformerBlock(nn.Module):
    """Hardware co-designed block: Pre-RMSNorm + MQA + Lean 2x Fused MLP."""
    def __init__(self, d_model: int = 256, n_heads: int = 8):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn = MultiQueryAttention(d_model, n_heads)
        self.norm2 = RMSNorm(d_model)
        
        # Lean 2x MLP cuts 50% of block FLOPs
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 2 * d_model, bias=False),
            nn.GELU(),
            nn.Linear(2 * d_model, d_model, bias=False)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

6. The Results: 1.45× Speedup

I ran the benchmarks on an NVIDIA Tesla T4 GPU (FP32 compute, 50 measured iterations). Here is what happened:

Figure 3: Sustained training throughput on Tesla T4. FastTransformer compiled reaches 116,080 tok/s, a +44.5% speedup over compiled standard GPT-2.

Did We Destroy Model Convergence?

Cutting parameters by $47\%$ ($4.84\text{M} \to 2.56\text{M}$) sounds risky. So I trained both models on TinyShakespeare for 200 optimization steps under identical hyper-parameters:

Figure 4: 200-step cross-entropy loss trajectory on TinyShakespeare. FastTransformer converged faster and reached a lower perplexity (11.98 vs 12.45).

7. Why It Works: FLOPs & Memory Co-Design

The explanation is simple:

7.1 Architectural Co-Design vs. Mere Model Pruning

We did not randomly drop layers or shrink embedding dimensions ($C=256$). Instead, we surgically restructured the two components that starved GPU memory bandwidth:

• Halving the Compute Giant (The $2\times$ vs $4\times$ MLP): Standard GPT-2 expands $C \to 4C \to C$ ($256 \to 1024 \to 256$), demanding $8.59\text{ GFLOPs}$ per block ($64\%$ of total block compute). Co-designing to a lean $2\times$ expansion ($512$) eliminated $4.29\text{ GFLOPs}$ of matrix multiplications without harming gradient propagation or perplexity.
• Shedding Redundant Key/Value Heads (MQA): Standard Multi-Head Attention maintains 8 separate Key heads and 8 Value heads ($\mathbf{W}_{qkv} \in \mathbb{R}^{256 \times 768}$). Multi-Query Attention shares 1 Key and 1 Value head across all 8 Query heads ($\mathbf{W}_{qkv} \in \mathbb{R}^{256 \times 320}$). This slashed projection weights by $58.3\%$ and eliminated gradient accumulation round-trips for 7 redundant Key/Value pairs.
• Keeping Activations in On-Chip SRAM (Zero Register Spills): Tesla T4 provides $320\text{ GB/s}$ of GDDR6 memory bandwidth. When intermediate activations are cut in half, the backward autograd tape fits within hardware registers and L2 cache, eliminating DRAM writebacks and warp register spills.
• Industry Convergence: This is the exact architectural path modern frontier LLMs followed: migrating from GPT-2/3 (Multi-Head Attention + $4\times$ MLP + LayerNorm) to state-of-the-art models like Llama 3, Mistral, and Gemma (Multi/Grouped-Query Attention + tuned MLP ratios + Pre-RMSNorm).

When you give torch.compile a standard Transformer block, it is forced to compile around quadratic attention and a fat 4× MLP. It does an admirable job fusing elementwise ops, but it cannot fundamentally alter the math you gave it.

When you feed the compiler FastTransformer:

• The compiler doesn't have to generate memory-heavy backward passes for the attention matrix because native SDPA keeps it in SRAM.
• The backward pass for the MLP processes half the activations, meaning registers don't spill into local memory.
• MQA cuts parameter gradients for keys and values by $87.5\%$.

The compiler had an easier, leaner workload. And that’s why Step 1 compilation latency dropped from 4.31 seconds down to 1.22 seconds ($3.5\times$ faster JIT tracing)!

Figure 5: Throughput scaling across sequence lengths $T \in [128, 1024]$. At $T=128$, FastTransformer hits 123,177 tokens/sec (1.58x speedup).

8. What I Learned

Here are the three engineering principles I took away from this discovery:

1. Profile arithmetic intensity before you write a single line of CUDA.
   I spent two weeks writing vectorized LayerNorm and warp-reduction softmax kernels to speed up operators that only accounted for 4% of the runtime. Five minutes with a FLOPs calculator would have pointed me straight to the MLP.
2. The compiler is your partner, not your adversary.
   I was trying to beat torch.compile on its home turf. But compilers are exceptional at lowering clean, lean graphs. When you change the workload to be hardware-friendly, compiler auto-tuners work with you instead of fighting memory bandwidth limits.
3. The Eager Inversion is the gold standard of architectural efficiency.
   If your new architecture in raw, uncompiled PyTorch eager mode runs faster than the baseline under full JIT compilation ($98\text{ ms}$ vs $102\text{ ms}$), you know your performance gain isn't a compiler trick. It’s fundamental algorithmic efficiency.

9. Where This Goes Next

Now that FastTransformer has established a new ceiling ($116{,}080\text{ tok/s}$), here is what I’m building next:

• Custom Pure CUDA FastTransformer Engine: Implementing the MQA zero-copy cuBLAS striding and vectorized RMSNorm directly in my from-scratch C++ engine to target $> 140{,}000\text{ tok/s}$ with hardware CUDA Graph capture.
• Long-Context Evaluation ($T \ge 2048$): As sequence length grows into full documents, MQA's memory savings scale exponentially.
• Mixed-Precision FP16 / BF16 Tensor Cores: Moving from FP32 CUDA cores to Tensor Core MMA instructions.

All code, models, benchmark scripts, and publication figures are open-source in the GitHub repository:
github.com/sarimahsan/transformer-cuda

Read More Systems Research & Engineering