The Case for Co-Designing Model Architectures with Hardware

ArXiv: 2401.14489

🎯 Pitch

A 2.7B parameter transformer can run 39% faster with zero accuracy loss simply by reshaping its dimensions to satisfy GPU Tensor Core alignment rules—like dropping attention heads from 32 to 20 so the per-head dimension becomes a multiple of 64. These gains come entirely from eliminating wasted computation in the matrix multiply kernels that dominate transformer latency, yet most practitioners blindly copy hyperparameters from prior work without considering hardware.

---

1. Executive Summary

This paper analyzes how the shape of transformer model architectures—specifically hyperparameters controlling model width, depth, and attention head count—interacts with the underlying GPU hardware to affect training and inference throughput. Working with decoder-only transformer models on NVIDIA V100, A100, and H100 GPUs, the authors trace performance bottlenecks back to General Matrix Multiplication (GEMM) kernels, which account for 68–95% of total model latency depending on scale. The paper identifies three primary hardware-aware optimization mechanisms: satisfying Tensor Core requirements (ensuring GEMM dimensions are multiples of 64 FP16 elements—e.g., setting the per-head dimension $h/a$ to a multiple of 64), mitigating wave quantization effects (aligning the number of thread blocks with the GPU's streaming multiprocessor count—e.g., 108 on A100), and minimizing tile quantization waste (ensuring output matrices divide evenly into hardware tile sizes—e.g., $128 \times 256$ blocks). By applying these principles, the authors achieve up to a 39% throughput improvement while preserving accuracy on a widely-adopted 2.7B parameter architecture, establishing that seemingly minor hyperparameter choices—such as reducing attention heads from 32 to 20 to make $h/a$ a power-of-two multiple—can yield substantial efficiency gains without any algorithmic changes, though these benefits are contingent on the specific GPU architecture used at both training and inference time.

2. Context and Motivation

The Core Problem: Model Architecture Design Ignores Hardware Realities

The fundamental issue this paper tackles is a persistent disconnect in deep learning practice: researchers and engineers routinely design transformer model architectures with little to no consideration for the GPU hardware those models will run on, despite the fact that hardware characteristics profoundly affect training and inference throughput. This is not a hypothetical concern—the paper demonstrates that models with nearly identical parameter counts but different architectural "shapes" can exhibit radically different runtimes, with throughput variations approaching 39% (Figure 1). At the scale of modern LLMs where training runs cost millions of dollars, this gap represents enormous wasted resources.

The paper traces this disconnect to three specific root causes (Section I):

1. Opacity of the software stack: Few resources trace performance impacts from the high-level transformer implementation all the way down to the underlying GEMM kernels executing on GPU hardware. Practitioners work with frameworks like PyTorch that abstract away hardware details, making it non-obvious why one set of hyperparameters runs faster than another even when the total FLOP count is identical.

2. Fragmented and inaccessible documentation: The existing knowledge about how transformer hyperparameters map to GPU kernel performance exists in scattered, non-standard formats—tweets from practitioners (Andrej Karpathy noting that padding the vocabulary from 50,257 to 50,304 tokens yields a 25% speedup, Reference [19]), footnotes in systems papers, and in-code comments within training libraries like Megatron and GPT-NeoX. There is no single reference that systematically compiles these insights or explains the underlying hardware first principles.

3. Institutional inertia from benchmark-driven research: Researchers frequently copy architectures wholesale from prior papers to enable cleaner comparisons, inadvertently "locking in" suboptimal design choices. The paper identifies the 2.7B parameter architecture from GPT-3 (Brown et al., 2020) as a canonical example—this specific model shape was replicated by OPT (Zhang et al., 2022), GPT-Neo (Black et al., 2021), Cerebras-GPT (Dey et al., 2023), RedPajama-INCITE (Together AI, 2023), and Pythia (Biderman et al., 2023), all inheriting the same hardware inefficiency despite the models being developed years apart and by different organizations. The paper shows that simply reducing the number of attention heads from 32 to 20 in this architecture—preserving total parameter count—yields nearly a 20% speedup (Figure 1).

This is not merely a pedagogical problem. The authors hypothesize that the lack of accessible, hardware-first documentation has caused multiple independent research groups to rediscover similar optimization principles through trial and error, wasting collective effort that could have been avoided with a consolidated reference.

Why This Problem Matters: The Scale of Waste and the Permanence of Design Choices

The practical significance of this gap operates on multiple levels:

Economic waste at scale. Training runs for large language models now cost tens of millions of dollars in compute. The paper's central finding—that model shape choices can cause up to 39% throughput differences at identical parameter counts—means that a poorly-shaped architecture effectively inflates the cost of training by nearly 40% compared to an optimized alternative with the same total FLOPs. Since GEMM kernels account for 68.3% of latency in medium-sized models and 94.9% in large models (Figure 2), these inefficiencies compound with scale: as models grow, the fraction of time spent in the operations most sensitive to shape choices actually increases, making the problem more severe for the largest and most expensive training runs.

Lock-in across the model lifecycle. A model's architectural hyperparameters—once chosen at design time—are essentially permanent. They affect not only the initial pretraining run but every subsequent fine-tuning run, every inference query, and every downstream application built on top of the model. The paper emphasizes this explicitly: "Optimizing model shapes for efficient GEMMs will increase throughput for the entire lifetime of the model, decreasing training time and inference costs for production models." A suboptimal choice made at architecture design time therefore amortizes its cost across potentially billions of inference calls over the model's deployment lifetime.

The inference dimension. While much of the paper's analysis focuses on training throughput, the implications extend to inference as well. The underlying forward-pass GEMMs are identical in both training and inference, so shape optimizations that improve training throughput also improve inference latency. The paper demonstrates this concretely in Section VII-C using the Pythia model suite: Pythia-1B, with its carefully chosen hidden dimension and fewer attention heads, achieves significantly higher inference throughput than Pythia-410M despite being 2.4× larger, because its shape is more hardware-efficient. The test loss remains on-trend with the rest of the Pythia family, confirming that the throughput gains do not come at the cost of model quality.

Cross-hardware portability challenges. The paper surfaces a subtle but important practical issue: model shapes optimized for one GPU architecture may not be optimal—or even usable—on another. Section VII-A describes the case of Oak Ridge National Lab's Summit supercomputer, which uses 6-GPU nodes rather than the more common 8-GPU configuration. When tensor parallelism is set to the number of GPUs per node (which is commonly the most efficient 3D-parallelism scheme, per Narayanan et al., 2021), the hidden size per GPU $h/t$ may no longer be a multiple of a power of two—breaking the alignment needed for efficient Tensor Core utilization. This creates a three-way tension: architectures efficient on 8-GPU nodes may not run efficiently (or at all) on 6-GPU nodes; adapting to 6-GPU nodes may require compromises that hurt performance; and any such compromises will affect downstream users who deploy the model on different hardware configurations. This is a concrete illustration of the paper's thesis that model design and hardware choices are deeply intertwined in ways that current practice largely ignores.

Prior Approaches and Their Limitations

The paper positions itself against several categories of existing work, each of which partially addresses the hardware-model co-design problem but falls short of providing a practical, unified guide for transformer practitioners.

GPU kernel characterization studies (Li et al., 2020; Mittal and Vaishay, 2019): A substantial body of work characterizes and profiles GPU kernels for deep learning workloads, measuring throughput, utilization, and bottlenecks across different operations and hardware generations. While valuable, these studies tend to focus on the kernel level in isolation—they characterize what kernels are slow, but do not connect these findings back to the architectural hyperparameters that practitioners actually control when designing a model. A researcher deciding whether to use 32 or 40 attention heads needs guidance on which choice will produce more efficient GEMMs, not just a profile showing that certain BMM dimensions are faster than others.

Kernel optimization frameworks (Aminabadi et al., 2022; Fang et al., 2021; NVIDIA FasterTransformer; Dao et al., 2022; Dao, 2023): A parallel line of work develops optimized implementations of specific transformer operations—FlashAttention being the most prominent recent example. These efforts improve performance by rewriting the implementation of attention or other operators, often achieving dramatic speedups without changing the model architecture. However, this approach is orthogonal to the paper's concern: kernel optimization addresses how an operation is implemented, while the paper addresses what operation dimensions are chosen. Even with an optimally-implemented attention kernel, the underlying GEMM will run faster if its dimensions align with hardware tile sizes and Tensor Core requirements. The paper's recommendations are therefore complementary to kernel optimization work, not competitive with it.

Cross-accelerator benchmarking (Emani et al., 2022; Wang et al., 2019; MLPerf): Several studies compare deep learning performance across different hardware platforms—GPUs, TPUs, and wafer-scale accelerators—to help users choose appropriate hardware for their workloads. While informative for procurement decisions, these comparisons typically treat the model architecture as fixed and vary only the hardware, missing the opportunity to co-optimize both dimensions simultaneously. The paper's contribution is orthogonal: it shows that within a given hardware platform (specifically NVIDIA datacenter GPUs), significant gains are available by tuning the model architecture to the hardware, regardless of which hardware is ultimately chosen.

Distributed training systems (Narayanan et al., 2021; Shoeybi et al., 2019): The Megatron-LM line of work provides detailed performance analysis and optimization for distributed transformer training across GPU clusters, including recommendations about parallelism strategies and their interaction with model dimensions. This paper cites this work positively and builds on it—many of the paper's recommendations about tensor parallelism degree $t$ and its relationship to $h/t$ draw directly from Megatron's insights. However, Megatron's performance guide focuses primarily on the distribution strategy (how to split the model across GPUs) rather than the intrinsic architecture (what shape the model should have in the first place). The paper's thesis is that these two concerns are inseparable—the optimal tensor parallel degree depends on the hidden dimension, which depends on hardware tile sizes—and that a unified treatment is needed.

Scattered practitioner knowledge (Karpathy, He, Shoeybi et al., GPT-NeoX comments): Perhaps most tellingly, the paper documents that the key insights about transformer shape optimization have been independently rediscovered multiple times by different groups and shared in ephemeral formats: a tweet about vocabulary size padding yielding 25% speedups (Karpathy, Reference [19]), a tweet about PyTorch tiling details (He, Reference [18]), comments in the Megatron codebase, and inline documentation in GPT-NeoX's training library. The paper's contribution is not novel discovery per se but rather systematization and explanation: collecting these scattered insights, grounding them in GPU first principles (Tensor Core requirements, wave quantization, tile quantization), and providing a unified framework that allows practitioners to reason about the performance implications of their architectural choices without needing to become GPU kernel experts.

The specific gap: no actionable, principle-based guide for model architects. Prior work either (a) profiles kernels without connecting to hyperparameter choices, (b) optimizes kernel implementations without changing model architecture, (c) compares hardware platforms with fixed models, or (d) exists as informal, scattered practitioner lore. None provides what the paper aims to deliver: a concise, principled set of rules that a model architect can apply at design time to ensure their transformer model will run efficiently on target GPU hardware. The paper explicitly frames itself as filling this gap: "We seek to provide explanations for these takeaways from the perspective of fundamental GPU first-principles, and to aggregate these explanations into a concise set of takeaways for efficient transformer training and inference."

How This Paper Positions Itself

The paper's framing is explicitly pedagogical rather than methodological: it is not proposing a new algorithm, architecture, or optimization technique. Instead, it acts as a translation layer between two communities that typically operate at different levels of abstraction—GPU hardware architects and kernel implementers on one side, and transformer model designers on the other—by mapping the performance-relevant hyperparameters of the transformer architecture onto the performance characteristics of the underlying GEMM kernels.

This positioning is evident in several structural choices:

First-principles derivation. Rather than simply reporting empirical speedups from specific hyperparameter choices, the paper walks through why those choices matter. Section III provides background on GEMM tiling (how output matrices are divided into tiles and scheduled onto streaming multiprocessors), wave quantization (how thread blocks are dispatched in waves, and how a final "tail wave" can waste compute), and Tensor Core requirements (the byte-alignment constraints needed to use the fastest hardware units). Section V then demonstrates these effects empirically through controlled GEMM microbenchmarks, showing how throughput varies with matrix dimensions in predictable, hardware-explainable ways. Only then does Section VI map these GEMM-level phenomena onto the specific matrix multiplications that comprise a transformer layer—attention QKV transforms, attention score computation, attention over value, post-attention projection, and MLP blocks—showing for each operation which dimension sizes will trigger which hardware effects.

Model architecture as the independent variable. The paper consistently treats model hyperparameters (hidden size $h$, number of attention heads $a$, vocabulary size $v$, tensor parallel degree $t$) as the levers that practitioners can adjust, and GPU throughput as the dependent variable that responds to those choices. This is a deliberate inversion of the typical hardware optimization framing, which treats the model as fixed and asks how to optimize kernel implementations or scheduling for it. The paper's thesis is that for transformer models specifically, there is enough flexibility in architectural hyperparameters (within the constraint of preserving total parameter count and model quality) that significant gains are available without any systems-level optimization at all.

Hardware specificity as a feature, not a limitation. The paper explicitly notes that its recommendations are hardware-dependent—what works for an A100 may not be optimal for a V100 or H100 due to different Tensor Core alignment requirements (64 elements on V100, 64 elements on A100 as well in FP16) and different SM counts (80 on V100, 108 on A100, 144 on H100). Rather than viewing this as a weakness, the paper frames it as evidence for the broader thesis: model architectures should be co-designed with specific hardware targets in mind. The paper covers multiple GPU generations (V100, A100, H100) and even includes results on AMD MI250X GPUs to demonstrate that the principles generalize across vendors, though the specific numerical thresholds change.

Practical, prescriptive recommendations. The paper culminates in a concrete set of actionable rules (Section VI-B): ensure vocabulary size is divisible by 64; make $b \cdot s$, $h/a$, and $h/t$ multiples of a power of two (up to 64); ensure $(b \cdot a)/t$ is an integer; minimize tensor parallel degree $t$; and prefer larger hidden dimensions to saturate the MLP GEMMs. These rules are presented as directly applicable by a practitioner without requiring deep GPU architecture knowledge—the preceding sections provide the explanation, but the rules themselves function as a self-contained checklist.

Reconciliation with modern architectural variants. The paper explicitly addresses how its analysis interacts with popular architectural modifications: parallel attention/MLP layers (Wang and Komatsuzaki, 2021), alternative positional embeddings (Su et al., 2021; Press et al., 2021), FlashAttention (Dao et al., 2022), and SwiGLU activations with their non-standard MLP ratios (Shazeer, 2020). For each, it explains whether and how the shape optimization principles change—for example, FlashAttention simplifies the attention sizing guidelines (only requiring $h$ to be as large as possible, since FlashAttention follows a simpler roofline model per Figure 12), while SwiGLU significantly complicates MLP block sizing because the $8/3$ intermediate dimension factor breaks alignment with hardware tile sizes (Section VII-B). This engagement with contemporary practice signals that the paper is intended as a living guide, not a historical artifact.

The Pythia-1B case study as proof of concept. The paper concludes with a concrete demonstration using the Pythia model suite (Section VII-C, Figure 13): Pythia-1B was designed with the hardware-aware principles described in the paper (fewer attention heads, larger hidden dimension relative to its size), and as a result achieves substantially higher inference throughput than Pythia-410M while maintaining on-trend test loss. The fact that the Pythia authors (including one of this paper's co-authors, Stella Biderman) applied these principles in practice and achieved the predicted benefits serves as validation that the recommendations are not merely theoretical.

In summary, the paper positions itself as filling a specific documentation and knowledge-transfer gap that has caused measurable waste in the transformer training ecosystem. It is not competing with kernel optimization research or hardware benchmarking studies; rather, it synthesizes insights from those communities into a form directly actionable by the model architect who controls hyperparameters but may never look at a GPU kernel profile. The thesis is that model shape is a first-class performance consideration that deserves as much attention as hardware choice or distributed training strategy—and that with relatively minor changes to standard architectures, practitioners can recover substantial throughput gains that compound across the entire model lifecycle.

3. Technical Approach

3.1 Reader Orientation

This is a hardware-model co-design analysis paper—not a new algorithm or model architecture, but a systematic framework for choosing transformer hyperparameters to maximize GPU throughput. The problem it addresses is that model architects routinely copy hyperparameters from prior work without considering how those choices interact with the underlying GPU hardware, causing wasted compute at every scale. The proposed solution takes the form of a concrete checklist: a set of rules mapping transformer dimensions to GPU GEMM requirements that practitioners can apply at design time without becoming hardware experts.

3.2 Big-Picture Architecture (Diagram in Words)

The paper's analytical pipeline operates in four sequential stages:

1. GPU Hardware Profiling (Sections III, V): Characterize how raw GEMM throughput varies with matrix dimensions on V100, A100, and H100 GPUs. This isolates the fundamental performance factors—Tensor Core alignment, wave quantization, tile quantization—in controlled microbenchmarks independent of any model architecture.

2. Transformer Operator Decomposition (Section III-C, VI-A): Break the decoder-only transformer into its constituent matrix multiplications. Each operator (QKV transform, attention score BMM, attention-over-value BMM, post-attention projection, MLP up/down projections, logit layer) is assigned its exact GEMM signature in terms of model hyperparameters: hidden size $h$, number of attention heads $a$, sequence length $s$, batch size $b$, and tensor parallel degree $t$.

3. Hyperparameter-to-Hardware Mapping (Section VI): Overlay the GEMM profiling results onto the transformer decomposition. For each operator, this mapping reveals which hyperparameter choices trigger which hardware phenomena. The key mappings are:

   • $h/a$ (the per-head dimension) controls Tensor Core efficiency in attention BMMs
   • $b \cdot s$ (the batch-times-sequence dimension) controls wave quantization in the QKV, output projection, and MLP GEMMs
   • $h$ and $4h$ (the MLP dimensions) control whether the MLP GEMMs reach the compute-bound saturation regime

4. Prescriptive Rule Set (Section VI-B): Condense the mapping into actionable guidelines (e.g., "vocabulary size divisible by 64," "$h/a$ a multiple of a power of two up to 64," "minimize tensor parallel degree") that ensure good hardware utilization across all GEMMs in the transformer layer.

Information flows linearly: hardware profiling $\rightarrow$ operator decomposition $\rightarrow$ mapping $\rightarrow$ rules. The rules are validated both on controlled transformer layer benchmarks (Figure 1, achieving up to 39% speedup) and on real training workloads (the Pythia case study in Section VII-C).

3.3 Roadmap for the Deep Dive

• First, the paper's formal abstraction of a transformer as a series of GEMM signatures (the operator decomposition table), because every subsequent mapping depends on knowing exactly which matrix multiplications exist and what dimensions they involve.
• Second, the hardware performance model—Tensor Core requirements, tile quantization, and wave quantization—with the paper's diagnostic microbenchmarks that isolate each effect, because this is the "why" behind every rule.
• Third, the hyperparameter-to-hardware mapping that connects the two: how each transformer dimension (h, a, s, b, t) flows through to GEMM dimensions (m, n, k) and triggers specific hardware phenomena. This is the paper's core analytical contribution.
• Fourth, the prescriptive rule set derived from the mapping, with the paper's concrete demonstration on the GPT-3 2.7B architecture (reducing attention heads from 32 to 20, achieving 1.18× speedup).
• Fifth, the treatment of modern architectural variants (parallel layers, alternative position embeddings, FlashAttention, SwiGLU) and how they interact with or modify the base recommendations.
• Finally, the paper's systematic approach to handling hardware heterogeneity—how rules change across V100/A100/H100 and across 6-GPU vs. 8-GPU node configurations—since this directly validates the co-design thesis.

3.4 Detailed, Sentence-Based Technical Breakdown

This is an empirical systems analysis paper whose core idea is that transformer model hyperparameters should be chosen to satisfy specific numerical constraints derived from GPU GEMM kernel properties. The paper's method is to decompose the transformer into its constituent matrix multiplications, characterize how each matrix multiplication's throughput depends on its dimensions through GPU microbenchmarks, and then derive a set of concrete constraints on hyperparameters that keep all GEMMs in high-throughput regimes.

---

3.4.1 The Transformer as a Formal Series of GEMMs

The paper's analytical foundation is a precise accounting of every matrix multiplication inside a standard decoder-only transformer layer (following GPT-2, Radford et al., 2019). This decomposition serves as the shared vocabulary between the model architecture world (where practitioners control hyperparameters like $h$, $a$, $L$) and the GPU kernel world (where performance depends on matrix dimensions $m$, $n$, $k$).

Variable definitions. Table I defines the core hyperparameters:

• $b$: microbatch size (number of independent sequences processed simultaneously)
• $s$: sequence length (number of tokens per sequence)
• $h$: hidden dimension (the width of the residual stream and primary representational dimension)
• $a$: number of attention heads
• $t$: tensor parallel degree (number of GPUs across which the model is split, each GPU processing $h/t$ of the hidden dimension)
• $v$: vocabulary size (number of output tokens in the embedding and logit layers)
• $L$: number of transformer layers (depth)

The relationship $h/a$ gives the per-head dimension, which the paper identifies as a critical efficiency variable because it appears as an inner dimension in the attention score and attention-over-value batched matrix multiplications.

The five GEMM classes inside one transformer layer. The paper enumerates exactly which matrix multiplications occur each time a transformer layer processes a batch of sequences (Section III-C, Table II), making explicit the mapping from hyperparameters to GEMM dimensions. For each, the input activations are a matrix whose dimensions are given by the current tensor shapes, and the weight matrices (learned during training) have fixed sizes determined by the architecture:

1. Attention key, value, query (QKV) transform. This fuses the three attention projections into a single matrix multiplication:

   $\text{GEMM size: } (b \cdot s, h) \times \left(h, \frac{3h}{t}\right)$

   where $(b \cdot s, h)$ is the input activation matrix—each of the $b \cdot s$ tokens has an $h$-dimensional hidden state—and $(h, 3h/t)$ is the concatenated Q, K, V weight matrix, already split across $t$ GPUs via tensor parallelism so that each GPU sees only $h/t$ of the output width per head and 3 copies (query, key, value). The output is of size $(b \cdot s, 3h/t)$.

2. Attention score computation (the $KQ^T$ batched matrix multiply). After the QKV outputs are reshaped into per-head tensors, the attention logits are computed as batched matrix multiplications:

   $\text{BMM count: } \frac{b \cdot a}{t} \text{ batched multiplications, each of size } \left(s, \frac{h}{a}\right) \times \left(\frac{h}{a}, s\right)$

   The batch size for the BMM is $b \cdot a / t$ because there are $a$ attention heads per layer, each sequence in the batch has all heads computed independently, and tensor parallelism splits heads across GPUs (so each GPU sees $a/t$ heads). Each individual multiplication takes a query matrix of shape $(s, h/a)$ and a transposed key matrix of shape $(h/a, s)$, producing an $(s, s)$ attention score matrix.

3. Attention over value computation (the score-times-value BMM). Once the attention scores are computed and softmax-normalized, they are multiplied with the value vectors:

   $\text{BMM count: } \frac{b \cdot a}{t} \text{ batched multiplications, each of size } (s, s) \times \left(s, \frac{h}{a}\right)$

   Each batched multiplication multiplies the $(s, s)$ attention probability matrix by the $(s, h/a)$ value matrix, producing a $(s, h/a)$ context vector for that head.

4. Post-attention linear projection. The per-head outputs are concatenated and projected back to the hidden dimension:

   $\text{GEMM size: } \left(b \cdot s, \frac{h}{t}\right) \times \left(\frac{h}{t}, h\right)$

   The input $(b \cdot s, h/t)$ is the concatenated multi-head output after tensor-parallel splitting, and the weight matrix $(h/t, h)$ projects it back to the full $h$-dimensional hidden space. The output is $(b \cdot s, h)$.

5. MLP block (two matrix multiplications). The standard feed-forward network expands and then contracts the hidden dimension:

   $\text{GEMM 1 (up-projection): } (b \cdot s, h) \times \left(h, \frac{4h}{t}\right)$ $\text{GEMM 2 (down-projection): } \left(b \cdot s, \frac{4h}{t}\right) \times \left(\frac{4h}{t}, h\right)$

   The standard expansion factor is 4, so the intermediate dimension is $4h$ (split to $4h/t$ per GPU under tensor parallelism). The second projection contracts back to $h$.

Vocabulary embedding and logit layers (not per-layer, but model-wide). Outside the transformer layers, the embedding lookup at the input and the logit projection at the output involve matrix multiplications of size $(b \cdot s, h) \times (h, v)$ and its transpose, where $v$ is the vocabulary size.

Parameter counting and compute estimation. The paper provides the standard approximations that link hyperparameters to total model size and FLOPs. The total parameter count is:

$P = 12h^2L + 13hL + (v + s)h \approx 12h^2L$

where the first term $12h^2L$ captures the dominant contribution from attention and MLP weight matrices (each layer has roughly $12h^2$ parameters when accounting for Q, K, V, output projection, and the two MLP matrices), and the lower-order terms $13hL + (v+s)h$ capture bias terms and embedding parameters. The paper notes this is commonly approximated as $P = 12h^2L$, omitting lower-order terms.

What this decomposition enables. By mapping every transformer operation to its exact GEMM signature, the paper can now apply GEMM-level performance analysis directly to model hyperparameter choices. If a particular GEMM size—say, $(s, h/a) \times (h/a, s)$—is known to run inefficiently on a target GPU because $h/a$ does not satisfy Tensor Core alignment, the model architect can adjust either $h$ or $a$ to fix the problem at the source. This mapping is the paper's core analytical contribution: it closes the loop between high-level architectural design and low-level kernel performance without requiring the architect to run any profiling tools.

---

3.4.2 The GPU GEMM Performance Model: Tensor Cores, Wave Quantization, and Tile Quantization

The paper's rules derive from three distinct hardware phenomena that govern GEMM throughput on NVIDIA GPUs. The paper explains each phenomenon from first principles and then demonstrates it empirically through controlled microbenchmarks (Figures 5 and 6) that sweep matrix dimensions in isolation.

Tensor Core requirements: the fundamental alignment constraint. NVIDIA Tensor Cores are specialized hardware units within each streaming multiprocessor (SM) that perform mixed-precision matrix multiply-accumulate operations at much higher throughput than standard CUDA cores. However, Tensor Cores can only operate on GEMMs whose dimensions satisfy specific alignment constraints because they process data in fixed-size "warps" of 32 threads operating on structured tile sizes.

The exact constraint is stated in Section III-B: "Tensor Cores can be fully utilized when GEMM dimensions $m$, $k$, and $n$ are multiples of 16 bytes and 128 bytes for V100 and A100 GPUs, respectively." Since the paper assumes FP16 computation (2 bytes per element), this translates to:

• On V100 GPUs: dimensions must be multiples of 8 elements ($16 \text{ bytes} / 2 \text{ bytes per element} = 8$)
• On A100 GPUs: dimensions must be multiples of 64 elements ($128 \text{ bytes} / 2 \text{ bytes per element} = 64$)

If dimensions do not satisfy these constraints, the computation falls back to slower CUDA cores or uses Tensor Cores with padding (wasting compute on zero-padded elements). The paper notes that "if these dimension sizes are not possible, Tensor Cores perform better with larger multiples of 2 bytes," meaning that even when full alignment is impossible, having GEMM dimensions that are multiples of larger powers of 2 (e.g., 16, 32 out of the ideal 64) reduces the amount of wasted computation relative to dimensions that share no common factors with the tile size.

Tile quantization: wasted compute from mismatched output grids. NVIDIA GPUs divide the output matrix of a GEMM into fixed-size rectangular regions called "tiles" or "thread blocks" (Figure 3 illustrates this tiling). Each tile is assigned to one SM for execution. The tile size is determined by the kernel implementation—typically $128 \times 256$ elements for the most efficient kernel variant—and cannot be changed by the user.

If the output matrix dimensions $X$ and $Y$ do not divide evenly into the tile dimensions $t_1 \times t_2$, some tiles along the boundary will cover elements that extend beyond the actual output matrix. These "partial tiles" still require full execution time on the SM (because the hardware processes the entire tile regardless of how many output elements are actually needed), but the excess computation is discarded. This is called tile quantization waste.

The paper notes that tile quantization is "hard to observe by the user" because the partial tiles execute concurrently with full tiles in the same scheduling wave. The latency penalty is therefore bounded—the kernel runs with "the same latency as a kernel with a larger problem size"—but throughput per FLOP decreases because the hardware is computing results that are never used.

Wave quantization: the dominant observable effect. The most visible and impactful performance phenomenon is wave quantization. The mechanism is:

1. The output matrix is divided into tiles.
2. Tiles (thread blocks) are scheduled onto the GPU's streaming multiprocessors (SMs). An A100 has 108 SMs, a V100 has 80, and an H100 has 144.
3. Only as many thread blocks as there are SMs can execute simultaneously (one "wave").
4. If the total number of thread blocks is not an exact multiple of the SM count, the final wave contains fewer thread blocks than available SMs, leaving some SMs idle.
5. This final "tail wave" has nearly the same latency as a full wave but performs useful computation on only a fraction of the SMs.

For example, if a GEMM produces 109 thread blocks, the A100 schedules: wave 1 with 108 blocks (full utilization), then wave 2 with 1 block (1 of 108 SMs active, 107 idle). The tail wave's latency is essentially identical to the first wave (since they execute the same computation), but its throughput per unit time drops to approximately $1/108$ of peak.

The paper formalizes this in Section VI-B. Assuming the most efficient tile size of $128 \times 256$, a matrix of size $(X, Y)$ avoids wave quantization waste when:

$\lceil \frac{X}{128} \rceil \cdot \lceil \frac{Y}{256} \rceil \equiv 0 \pmod{\text{\#SMs}}$

or equivalently with the tile dimensions swapped. In words: the number of thread blocks in the output grid (computed as the ceiling of $X$ divided by the tile width, times the ceiling of $Y$ divided by the tile height) must be evenly divisible by the number of streaming multiprocessors on the GPU.

The paper's key practical insight is that wave quantization produces a characteristic "sawtooth" pattern in throughput benchmarks: as one GEMM dimension increases, throughput rises (because larger matrices are more compute-dense) until a new tail wave is triggered, at which point throughput drops sharply, then rises again as the tail wave fills up. This pattern is clearly visible in Figure 5b, where sweeping the $k$ dimension of a $(27648, 4096) \times (4096, k)$ GEMM shows periodic throughput drops at thresholds where the number of thread blocks crosses a multiple of the SM count.

Demonstrating the effects through microbenchmarks. Figures 5 and 6 provide the empirical foundation. Figure 5a sweeps $m$ in a $(m, 4096) \times (4096, m)$ GEMM, showing throughput increasing with matrix size as the kernel transitions from memory-bound to compute-bound. Figure 5b sweeps $k$ in a $(27648, 4096) \times (4096, k)$ GEMM, showing clear wave quantization sawtooth patterns. Figure 5c sweeps $k$ in a smaller $(2304, 4096) \times (4096, k)$ GEMM where PyTorch "is able to better balance the improvements from GEMM parallelization and inefficiencies from wave quantization," showing a smoother curve. Figures 6a-d show BMM throughput sweeping batch size and inner dimension on V100 and A100 GPUs, demonstrating that BMMs inherit the same patterns from their constituent GEMMs.

Why these three effects are the complete model. The paper argues that GPU GEMM performance reduces to these three factors for the purpose of model architecture design. Memory bandwidth limitations (whether the kernel is memory-bound vs. compute-bound) are addressed implicitly by ensuring dimensions are large enough to reach the compute-bound regime. The key controllable variable for the model architect is the alignment between GEMM dimensions and hardware constants (tile sizes, SM counts, Tensor Core requirements)—and these three effects capture all the relevant alignment constraints.

---

3.4.3 The Hyperparameter-to-Hardware Mapping

This section is the paper's analytical core: it takes each class of GEMM identified in the transformer decomposition and overlays the GPU performance model to determine which hyperparameter values lead to efficient GEMMs.

The attention QKV transform: optimizing $b \cdot s$ and $h/t$. The QKV GEMM has signature $(b \cdot s, h) \times (h, 3h/t)$. The paper observes that Figure 2 shows this GEMM growing in relative importance as model size increases (accounting for the largest share of latency in the largest models, along with the MLP GEMMs). The relevant dimensions for hardware efficiency are:

• $b \cdot s$ (the batch-times-sequence dimension, which becomes the $m$ dimension in the GEMM output): should be as large as possible to move the kernel into the compute-bound regime. The paper cites Nado et al. (2021) for the recommendation that microbatch size should be as large as possible (within memory constraints).
• $h/t$ (the hidden dimension per GPU after tensor parallelism): should be divisible by a power of two up to 64 to satisfy Tensor Core alignment on A100 GPUs. Since $t$ is typically chosen as the number of GPUs per node, this creates a coupling between hardware topology and model architecture: if the node has 6 GPUs, $h$ must be divisible by both 6 (for $t=6$) and 64 (for Tensor Cores), making $h$ a multiple of 192. This constraint may be impossible to satisfy for small models, creating the tension described in the 6-GPU node case study (Section VII-A).

Attention score computation ($KQ^T$): the $h/a$ constraint. This is a BMM with $b \cdot a/t$ independent multiplications, each of size:

$(s, \frac{h}{a}) \times (\frac{h}{a}, s)$

Here, $h/a$ is the per-head dimension, and it appears as both an inner dimension (the contracted $k$ dimension in the GEMM) and an outer dimension (one side of each input matrix). The paper's key finding (Figures 7a-b, 8-9, and the extensive Appendix B) is that the throughput of this BMM depends strongly on the largest power of two that divides $h/a$:

• When $h/a$ is divisible by 64 (the maximum relevant alignment on A100), Tensor Cores are fully utilized and throughput is maximized.
• When $h/a$ is divisible by smaller powers of two (32, 16, 8, 4, 2), throughput degrades in steps because the alignment with the 128-byte Tensor Core tile size becomes progressively worse.
• When $h/a$ is not divisible by any power of two greater than 1 (i.e., it is odd, or its highest power-of-two factor is $2^0$), throughput drops substantially.
• Going beyond 64 provides no additional benefit because the Tensor Core tile size requires only 64 FP16 elements for full alignment.

The paper visualizes this in Figures 7a and 7b by "splitting" the single throughput series by the largest power of two that divides $h/a$, showing clearly stratified performance bands. For example, at $a=32$ heads (Figure 7a), when hidden size is swept such that $h/32$ is divisible by 64 (top band, blue), throughput is substantially higher than when $h/32$ is only divisible by 8 (lower band).

This explains the paper's repeated emphasis on the ratio $h/a$ as the single most important hyperparameter for attention efficiency. It also explains why decreasing the number of attention heads (increasing $h/a$) improves throughput: Figures 8 and 9 show that "a decrease in $a$ is an increase in $h/a$ and these two GEMMs are memory bound, [so] an increase in component matrices size creates much more efficient GEMMs." More precisely, larger per-head dimensions make each BMM sub-GEMM larger and more compute-dense, pulling it out of the memory-bound regime where kernel launch overhead and memory bandwidth dominate.

Attention over value computation: same $h/a$ constraint, different tail wave pattern. The attention-over-value BMM has signature:

$(s, s) \times (s, \frac{h}{a})$

with $b \cdot a/t$ batched copies. The same $h/a$ power-of-two constraint applies because $h/a$ again appears as a GEMM dimension. Figure 9 additionally shows wave quantization effects in the attention-over-value computation: "Since each line moves in steps of $64h/a$, the BMMs corresponding to each line grow at different rates. This causes the period of the wave quantization effect to appear different for each $a$ value." In plain language: as hidden size increases, the GEMM dimensions for different head counts increase at different rates (because $h/a$ changes differently for each $a$), causing each head count to hit its wave quantization thresholds at different $h$ values.

MLP GEMMs: saturating the compute-bound regime. The MLP up-projection $(b \cdot s, h) \times (h, 4h/t)$ and down-projection $(b \cdot s, 4h/t) \times (4h/t, h)$ are typically the second-largest GEMMs in the transformer layer (along with the QKV transform, per Figure 11). Figures 10a and 10b show the paper's sweep of MLP throughput as a function of hidden dimension $h$ with $a=128$ heads fixed. The plots show that throughput increases with $h$ and begins to saturate at large $h$, as the GEMMs become fully compute-bound. The paper's recommendation is therefore to make $h$ as large as possible (subject to the total parameter budget and other constraints) to push the MLP GEMMs into saturation, where they achieve the highest teraFLOP/s.

Figure 11 provides a complementary view: it shows the proportion of total layer latency spent in each GEMM as model size grows. For the largest models, the QKV transform and MLP blocks dominate (together accounting for the majority of latency), while attention score computation and attention-over-value shrink in relative importance. This justifies the paper's strategic recommendation: "for the largest models, the QKV transformation in the attention block along with the MLP block are the most prevalent GEMMs. Therefore, the overall latency of the model would benefit most from optimizing these kernels." Meanwhile, the fact that attention score and AOV computations are a small fraction of total latency means that even if those kernels are somewhat suboptimal (because $h/a$ cannot be perfectly aligned), the end-to-end impact is limited.

The logit layer: vocabulary size alignment. The final logit projection maps the hidden state back to the vocabulary dimension: $(b \cdot s, h) \times (h, v)$. Figures 20a and 20b show the empirical effect: throughput is maximized when $v$ is a multiple of 64, and also benefits from $h$ being a multiple of 64. The paper notes that padding the vocabulary size to the next multiple of 64 (e.g., from 50,257 to 50,304 in the GPT-2 tokenizer, as noted in Karpathy's tweet, Reference [19]) is a nearly cost-free optimization that improves logit layer efficiency at negligible parameter cost.

Why the paper focuses on single-GPU computations despite distributed training. The paper explicitly states that since it "focuses on the computations being done on a single GPU," the hidden size should be understood as "the hidden size per GPU"—that is, $h/t$ where $t$ is the tensor parallel degree. The coupling between $t$ and $h$ is critical: if $h$ is chosen to satisfy Tensor Core alignment on a single GPU, but then tensor parallelism splits it across GPUs, the per-GPU dimension $h/t$ may break that alignment. The paper includes $h/t$ in its rule set and discusses the 6-GPU node case as a concrete example of this tension, but leaves a full analysis of pipeline and sequence parallelism interactions to future work.

---

3.4.4 The Prescriptive Rule Set

Section VI-B condenses the entire preceding analysis into five concrete rules for model architects. Each rule is directly traceable to one or more of the hardware phenomena documented in Sections V and VI-A.

Rule 1: The vocabulary size should be divisible by 64. This ensures the embedding and logit layer GEMMs satisfy Tensor Core alignment on A100 GPUs. The paper demonstrates this empirically in Figures 20a-b, showing clear throughput peaks at multiples of 64. This rule is essentially zero-cost: padding $v$ from 50,257 to 50,304 adds only 47 unused token embeddings, representing a negligible fraction of the total parameter count.

Rule 2: The microbatch size $b$ should be as large as possible. This pushes all GEMMs whose $m$ dimension involves $b \cdot s$ further into the compute-bound regime, increasing utilization. The paper cites Nado et al. (2021) for the theoretical justification. Importantly, $b$ itself does not need to be divisible by a power of two because $s$ "is a large power of two" (standard sequence lengths are 512, 1024, 2048, etc.), so the product $b \cdot s$ naturally satisfies alignment as long as $s$ does. This is a subtle but practical distinction: the architect does not need to constrain batch size to particular values beyond "as large as memory allows."

Rule 3: $b \cdot s$, $h/a$, and $h/t$ should be divisible by a power of two, with no further benefit to going beyond 64. This is the central alignment rule, bundling three related constraints:

• $b \cdot s$ (via $s$ being a power of two) ensures the $m$ dimension in QKV, output projection, and MLP GEMMs satisfies Tensor Core alignment.
• $h/a$ ensures the attention score and attention-over-value BMMs satisfy Tensor Core alignment. Since this dimension appears as both an inner and outer dimension in attention GEMMs, it is the single most impactful per-head hyperparameter.
• $h/t$ ensures that tensor parallelism does not break alignment: the hidden dimension per GPU must still satisfy Tensor Core requirements.

The "no benefit beyond 64" clause comes from the A100 Tensor Core requirement of 128 bytes = 64 FP16 elements. Making $h/a = 128$ does not further improve Tensor Core utilization (though it may improve overall throughput by making the GEMM larger and more compute-dense).

Rule 4: $(b \cdot a)/t$ should be an integer. This is a correctness constraint, not a performance optimization. The BMM batch size in attention score and AOV computations is $b \cdot a/t$ (sequences × attention heads, split across $t$ GPUs). If this is not an integer, the tensor operations would require ragged batch sizes, which standard BMM implementations do not support. This rule ensures the model can actually execute.

Rule 5: $t$ should be as small as possible. Higher tensor parallelism reduces the per-GPU hidden dimension $h/t$, which can push GEMMs toward the memory-bound regime (smaller matrices are less compute-dense) and can break alignment if $h/t$ loses its power-of-two divisibility. The paper cites Narayanan et al. (2021) for this recommendation. The practical constraint is that $t$ is often fixed by the number of GPUs per node (commonly 8, sometimes 6 as in the Summit case study), so this rule is really a recommendation to use the minimum tensor parallelism necessary to fit the model in GPU memory, and to rely on other parallelism strategies (pipeline parallelism, data parallelism) for additional scaling.

Additional rule from Section VI-B preamble: The number of layers $L$ should be divisible by the number of pipeline parallel stages. This is a pipeline parallelism constraint, not a GEMM constraint: if the number of layers does not divide evenly into pipeline stages, some stages will have more layers than others, creating load imbalance. The paper notes this is "further evidence for our thesis that model dimensions should be chosen with hardware details in mind."

Demonstration on GPT-3 2.7B. The paper applies these rules to the canonical GPT-3 2.7B architecture (hidden dimension $h = 2560$, attention heads $a = 32$, layers $L = 32$). The problem is immediately visible: $h/a = 2560 / 32 = 80$, which is not a multiple of 64 (it is $64 + 16$). The Tensor Cores on an A100 operate at partial efficiency with this head dimension, wasting approximately 25% of peak throughput on attention score and attention-over-value computations.

The paper proposes two fixes: increase $h$ to 4096 (which would make $h/a = 4096/32 = 128$, satisfying alignment), or decrease $a$ to 20 (making $h/a = 2560/20 = 128$, satisfying alignment). The first option doubles the parameter count to 6.7B, violating the constraint of preserving model size. The second option preserves the 2.7B parameter count while fixing the alignment.

Figure 1 shows the results: configuration C1 ($h=2560, a=64$) and C2 ($h=2560, a=40$) represent intermediate improvements, with the fully-optimized architecture achieving a $1.18\times$ speedup (the paper states "almost 20% faster" in Section I and "1.18× speed-up" in Section VI-B—these are consistent, as 1.18× means 18% faster).

The tradeoff between attention head count and attention efficiency. The paper explicitly addresses a potential objection: decreasing $a$ increases $h/a$, which reduces the number of independent BMMs but makes each BMM larger. Larger BMMs are more efficient (more compute-dense, better Tensor Core utilization), so throughput improves despite the reduced parallelism. The paper further notes that attention score and AOV computations are a small fraction of total latency for large models (Figure 11), so even if there were some overhead from reduced head parallelism, the impact on end-to-end throughput would be minimal.

The paper's recommended strategy when $h/a$ cannot be fully aligned. For models where the ideal $h/a$ cannot be achieved (e.g., due to accuracy constraints on the minimum number of attention heads, or total parameter budget constraints), the paper recommends two alternatives:

1. Use FlashAttention v2 (Dao, 2023) to replace the attention BMMs with a fused kernel that follows a simpler roofline model (Figure 12), effectively bypassing the $h/a$ alignment constraint. The paper notes this is "for small models to mitigate these effects."
2. Increase $h$ as much as possible to saturate the MLP GEMMs (Figures 10a-b), which improves overall throughput even if the attention GEMMs remain suboptimal.

---

3.4.5 Interactions with Modern Architectural Variants

Section VI-C addresses how four popular modifications to the standard transformer architecture interact with the paper's shape optimization principles.

Parallel layers (Wang and Komatsuzaki, 2021). The standard transformer sequentially computes attention and MLP:

$y = x + \text{MLP}(\text{Norm}(x + \text{Attn}(\text{Norm}(x))))$

The parallel formulation computes them independently and sums:

$y = x + \text{MLP}(\text{Norm}(x)) + \text{Attn}(\text{Norm}(x))$

The paper notes that "in practice the two branches are not computed simultaneously" because the GPUs execute them sequentially. The speedup comes from fusing the attention and MLP kernels into a single kernel, reducing kernel launch overhead and memory traffic. The paper recommends parallel attention as "the default best practice" and notes that "it does not impact our analysis at all"—the GEMMs remain the same shape and the alignment constraints are unchanged.

Alternative positional embeddings (Su et al., 2021; Press et al., 2021). Rotary and ALiBi embeddings replace the original sinusoidal or learned positional embeddings. While these are "slightly faster than the GEMM necessary for Rotary and ALiBi embeddings," the paper considers the model quality improvements "well worth it" and notes that custom kernels have reduced the cost further. Again, these do not change the GEMM shape analysis.

FlashAttention (Dao et al., 2022; Dao, 2023). FlashAttention v1 and v2 rewrite the attention computation to be IO-aware, fusing the attention score computation, softmax, and attention-over-value steps into a single kernel that never materializes the full $(s, s)$ attention matrix in GPU memory. The paper evaluates FlashAttention v2 in Figure 12, sweeping hidden dimension with $a=128$ fixed. The finding: FlashAttention "follows a roofline model," meaning its throughput scales smoothly with problem size without the wave quantization sawtooth pattern seen in standard attention BMMs. This simplifies the attention sizing guidelines: "our attention takeaways [simplify] to only require that $h$ be as large as possible; the takeaways for MLPs remain unchanged." In other words, when using FlashAttention, the $h/a$ power-of-two constraint becomes irrelevant for the attention computation itself—though it may still matter for other operations that use $h/a$ as a dimension.

SwiGLU and $8h/3$ MLPs (Shazeer, 2020). This is the most consequential architectural variant for the paper's analysis. The SwiGLU activation function contains an additional learned matrix (the "gate" projection), so the MLP block now has 3 weight matrices instead of the standard 2. To preserve the total parameter count ratio between attention and MLP blocks, the common practice is to reduce the MLP expansion factor from $4h$ to $8h/3$:

$\text{Standard: } d_{\text{ff}} = 4h \quad \rightarrow \quad \text{SwiGLU: } d_{\text{ff}} = \frac{8}{3}h$

The paper identifies a critical problem: $8/3$ is not an integer, so $8h/3$ will almost never be divisible by 64 or any useful power of two. This breaks all the alignment constraints derived in the previous sections, producing MLP GEMMs that run substantially slower than they could with a properly-aligned expansion factor.

The practical solution: the $8/3$ coefficient is negotiable. The paper makes a crucial observation: "the $8/3$ coefficient is only a suggestion and thus it's possible to find other coefficients that would lead to better-shaped MLP matrices." The Llama-2 models provide a concrete example:

• Llama-2-7B uses $d_{\text{ff}} = 11008$ with $h = 4096$, giving $11008/4096 = 2.6875$, close to $8/3 = 2.667$.
• Llama-2-70B uses $d_{\text{ff}} = 28672$ with $h = 8192$, giving $28672/8192 = 3.5$, far from $8/3$.

Section VII-B describes how to operationalize this insight: "now that we know the recommended coefficient isn't exact and since a good $h$ has already been chosen, one can now search for a good nearby number that still leads to high-performance GEMMs in the MLP. Running a brute-force search reveals that Llama-2-7B's intermediate size is indeed one of the best performing sizes in its range." The procedure is: (1) choose $h$ to satisfy Tensor Core alignment; (2) search for an $d_{\text{ff}}$ near $8h/3$ that is also divisible by 64; (3) accept the small deviation in parameter count ratio as a worthwhile tradeoff for hardware efficiency.

---

3.4.6 Handling Hardware Heterogeneity

The paper's thesis is that model architectures should be co-designed with hardware. This inherently means that rules are hardware-specific, and the paper explicitly addresses how recommendations vary across GPU generations and node configurations.

GPU generation differences. The three NVIDIA GPU generations covered (V100, A100, H100) differ in two critical parameters:

• Tensor Core alignment requirement: V100 requires 16 bytes = 8 FP16 elements for full Tensor Core utilization. A100 and H100 require 128 bytes = 64 FP16 elements. This means the "multiple of 64" rule for $h/a$, $h/t$, and $b \cdot s$ applies to A100/H100 but can be relaxed to "multiple of 8" for V100.
• Streaming multiprocessor count: V100 has 80 SMs, A100 has 108 SMs, H100 has 144 SMs. This changes the wave quantization thresholds: the condition for avoiding tail-wave waste becomes divisibility by 80, 108, or 144 respectively.

The MI250X results. The paper includes results on AMD MI250X GPUs to demonstrate cross-vendor generalizability. The same GEMM-level performance factors apply, but the specific thresholds differ based on AMD's hardware architecture. This supports the broader thesis that the methodology (decompose model into GEMMs, characterize hardware GEMM performance, derive alignment constraints) is portable even though the specific numerical rules are not.

The 6-GPU node challenge (Section VII-A). When a supercomputer node has 6 GPUs instead of 8, tensor parallelism with $t=6$ (one GPU per node) creates three specific problems:

1. Compatibility: Architectures designed for 8-GPU nodes ($t=8$) have $h/t$ dimensions that assume $h$ is divisible by 8. With $t=6$, $h$ must be divisible by 6, which may conflict with original $h$ values.
2. Efficiency: Even if the model fits, $h/t$ with $t=6$ may not be a power-of-two multiple, breaking Tensor Core alignment. For A100 GPUs, $h/t$ would need to be divisible by both 6 (for the tensor parallel split) and 64 (for Tensor Cores), requiring $h$ to be a multiple of $lcm(6, 64) = 192$.
3. Portability: If the model is optimized for $t=6$ (e.g., $h$ chosen as a multiple of 192), it may not run efficiently on 2-GPU, 4-GPU, or 8-GPU nodes used for fine-tuning or inference, because the $h/t$ alignment would break for different $t$ values.

The paper frames this as a design choice: "Does one choose the most efficient hyperparameters for pretraining only (which would involve a tensor-parallel degree of 6 and therefore a hidden dimension divisible by 6 and 64), or should the pretraining team choose a set of hyperparameters that are more amenable to the node architectures commonly used for finetuning or inference?" This question has no universal answer; it depends on the relative importance of pretraining throughput vs. downstream deployment flexibility. The paper presents it as evidence for the co-design thesis—hardware topology fundamentally constrains model architecture in ways that are rarely considered at design time.

The Pythia inference validation (Section VII-C). To demonstrate that training-time shape optimizations also benefit inference (since "the underlying forward-pass GEMMs are the same"), the paper benchmarks the Pythia model suite using DeepSpeed-MII (Figure 13). Pythia-1B, with its hardware-optimized shape (fewer attention heads, larger hidden dimension), achieves substantially higher inference throughput than Pythia-410M despite being 2.4× larger, while maintaining on-trend test loss. This validates that the paper's shape optimization principles apply throughout the model lifecycle and do not degrade model quality.

---

3.4.7 Summary of the Paper's Methodological Approach

The paper's technical contribution is a reductionist framework: reduce the complex problem of "how do I make my transformer run fast on GPUs?" to a small set of numerical constraints on a handful of hyperparameters, each constraint traceable to a specific GPU hardware phenomenon. The framework's power comes from its completeness—it covers every GEMM in the transformer layer, accounts for the dominant performance effects (Tensor Core alignment, wave quantization, tile quantization), and extends to modern architectural variants (FlashAttention, SwiGLU) and hardware configurations (multi-GPU nodes, different GPU generations, cross-vendor GPUs).

The cost of this reduction is that the rules are hardware-dependent. A model optimized for A100 will not automatically be optimal on H100 or V100. The paper treats this as a feature, not a bug: it is precisely the argument for co-design. The framework provides the analysis tools to re-derive the rules for any GPU architecture, given its SM count and Tensor Core alignment requirements.

4. Key Insights and Innovations

Innovation 1: The Transformer as a Formal GEMM Decomposition — A Complete Mapping from Hyperparameters to Kernel Signatures

Prior work on GPU performance optimization for deep learning typically operated at one of two disconnected levels: (1) kernel-level profiling and optimization (Li et al., 2020; Mittal and Vaishay, 2019) that characterized which GEMM shapes were fast or slow but did not connect these findings back to the architect's levers, or (2) distributed training system design (Narayanan et al., 2021; Shoeybi et al., 2019) that focused on parallelism strategies with implicit—but never fully enumerated—assumptions about model dimensions. Between these two lived the model architect, who controlled hyperparameters like $h$, $a$, and $L$ but had no systematic way to predict how those choices would translate into GPU throughput.

The paper's intellectual contribution is the construction of this missing translation layer. By exhaustively enumerating every matrix multiplication in a standard decoder-only transformer layer and writing each GEMM signature in terms of the architect's hyperparameters—e.g., the attention score BMM is $(s, h/a) \times (h/a, s)$ with batch count $b \cdot a/t$—the paper creates a complete, closed-form mapping from the design space to the hardware performance space. This mapping is what enables all subsequent analysis: it turns the vague intuition that "wider models might run faster" into precise constraints like "$h/a$ must be divisible by 64 on A100 GPUs."

Why is this a fundamental contribution rather than just careful bookkeeping? Because it changes who can do hardware-aware model design. Before this paper, optimizing a transformer architecture for GPU throughput required either deep systems expertise (understanding CUDA kernel launch configurations, profiling individual GEMMs, interpreting nsight reports) or copying hyperparameters from other papers and hoping they were efficient. After this paper, the architect applies a checklist: is $h/a$ a multiple of 64? Is $b \cdot s$ large? Is $v$ padded to the nearest multiple of 64? The intellectual work of translating between the two abstraction levels is done once, in this paper, and subsequent architects inherit the result. This is analogous to what Hoffmann et al. (2022) did for pretraining compute allocation—not inventing a new method, but providing the systematic framework that turns ad-hoc practice into principled design.

The empirical validation of the mapping's completeness comes from Figure 2: as model size increases from medium to large, the fraction of latency attributable to GEMMs rises from 68.3% to 94.9%. This means the mapping becomes more accurate—not less—for the most expensive models, which is exactly where the economic stakes are highest.

Innovation 2: Wave Quantization as a First-Class Performance Model for Transformer GEMMs — and Why Tile Quantization Isn't

The GPU computing literature has long recognized wave quantization as a performance factor, but prior work treating it as one of many second-order effects to be managed through kernel autotuning. The paper elevates wave quantization to the primary observable performance phenomenon that model architects must design around, while simultaneously arguing that tile quantization—often discussed in the same breath—is largely irrelevant from the architect's perspective.

The key conceptual move is the distinction between observability and controllability. The paper argues that wave quantization is "more easily observable" because it produces characteristic sawtooth patterns in throughput benchmarks (Figure 5b): as a GEMM dimension increases, throughput rises until a new tail wave is triggered, at which point it drops sharply. This pattern is visible to anyone running a sweep, and the underlying mechanism (thread blocks schedule in waves equal to the SM count) is both intuitive and quantifiable—the formal condition for avoiding tail-wave waste is that the number of thread blocks be divisible by the SM count (108 for A100, 80 for V100).

Tile quantization, by contrast, "is hard to observe by the user." The reason, which the paper explains clearly, is that partial tiles execute concurrently with full tiles in the same scheduling wave, so the latency penalty is bounded and hidden. A GEMM with tile waste runs with essentially the same latency as a slightly larger GEMM without waste—the throughput per FLOP decreases, but the absolute time does not spike. This means the architect cannot easily diagnose tile quantization from benchmarks and, more importantly, cannot easily control it because PyTorch's linear algebra backend can choose different tile sizes for different GEMMs, making the effective tile size unpredictable.

This diagnostic contribution—distinguishing which hardware effects actually matter for architectural decisions—is significant because it prevents practitioners from optimizing the wrong thing. The paper formalizes this by providing the exact wave quantization condition (Section VI-B) while noting that there is "not a transformer configuration with GEMMs that fill tensor core requirements without wave quantization inefficiency," implying some degree of wave quantization is inevitable and the goal is to manage it, not eliminate it. The paper's rules ($b \cdot s$ as large as possible, $h/a$ a power-of-two multiple) are specifically designed to minimize wave quantization waste within the constraints of standard transformer shapes.

Innovation 3: The $8/3$ Coefficient Is Negotiable — Unfreezing MLP Expansion Factors for Hardware Alignment

When the SwiGLU activation function (Shazeer, 2020) was introduced, it came with a specific prescription for preserving parameter count parity: reduce the MLP expansion factor from $4h$ to $8h/3$. Subsequent architectures—PaLM, LLaMA, Mistral—largely adopted this prescription, treating $8/3$ as an architectural constant on par with the factor of 4 in standard MLPs. The paper identifies a critical problem with this inherited practice: $8/3$ times any integer $h$ that satisfies Tensor Core alignment will almost never produce an integer $d_{ff}$ that itself satisfies Tensor Core alignment, because $8/3$ is not an integer. The MLP intermediate dimension—which, per Figure 11, accounts for a substantial fraction of total transformer latency—will systematically fall on misaligned GEMM dimensions.

This is an institutional design critique disguised as a performance optimization. The paper argues that the field has been treating $8/3$ as a hard constraint when it is in fact "only a suggestion," and that this misconception has caused a generation of SwiGLU-based models to ship with avoidable MLP inefficiencies. The evidence: Llama-2-7B uses $d_{ff} = 11008$ with $h = 4096$, giving a ratio of 2.6875 (close to 2.667), while Llama-2-70B uses $d_{ff} = 28672$ with $h = 8192$, giving a ratio of 3.5 (far from 2.667). These are not random deviations—the paper's brute-force search shows that "Llama-2-7B's intermediate size is indeed one of the best performing sizes in its range." The Llama team likely performed the same kind of hardware-aware search the paper advocates, but without documenting the principle, so each subsequent team must independently rediscover that the $8/3$ rule is breakable.

The broader intellectual contribution is the identification of a class of "suggested constants" in the architecture design literature that practitioners treat as immutable. The paper provides the intellectual justification—and the procedural recipe—for questioning such constants: (1) identify the function the constant serves (preserving parameter count ratios); (2) determine how much flexibility exists (ratios like 2.6875 and 3.5 both work, as evidenced by Llama-2); (3) search within the flexible range for values that satisfy hardware alignment constraints. This is not just about SwiGLU—it is a generalizable methodology for hardware-aware architecture design that treats "standard" hyperparameters as negotiable within empirically validated bounds.

Innovation 4: Hardware Specificity as a Feature — The Co-Design Thesis and Its Boundary Conditions

The paper's most provocative intellectual move is reframing hardware dependence from a limitation to a design principle. Standard practice treats hardware independence as desirable: a good model architecture should run efficiently on all GPUs. The paper argues almost the opposite: efficiency requires embracing hardware specificity. The optimal $h/a$ on A100 is a multiple of 64; on V100, it is a multiple of 8. A model optimized for V100 will leave A100 Tensor Cores partially underutilized. A model optimized for 8-GPU nodes may not run on 6-GPU nodes at all.

This reframing has concrete consequences. The paper's Section VII-A case study reveals a genuine tension: the most efficient pretraining configuration for Oak Ridge's Summit supercomputer uses $t=6$ and therefore requires $h$ divisible by 192 (to satisfy both the tensor-parallel split and A100 Tensor Core alignment). But this model will run inefficiently on the 8-GPU nodes common in cloud deployments. There is no universal solution—the architect must choose whether to optimize for pretraining throughput or downstream portability, and this choice is fundamentally a hardware-model co-design decision, not an architecture choice alone.

The strength of this framing is that it is empirically grounded and bounded. The paper does not claim that every hyperparameter must be hardware-tuned. It identifies which hyperparameters interact with hardware ($h/a$, $h/t$, $b \cdot s$, $v$) and which do not ($L$, the choice of positional embedding, the use of parallel layers). It shows that the hardware-specific rules follow the same methodology across GPU generations and even across vendors (MI250X results follow the same pattern with different thresholds). And it validates that shape-optimized models do not sacrifice model quality: Pythia-1B achieves on-trend test loss while being substantially more hardware-efficient than its smaller Pythia-410M sibling (Figure 13).

The boundary condition is equally important: the paper's rules are for GEMM-dominated transformer workloads on NVIDIA-style tensor-core GPUs. On CPUs, TPUs, or architectures without tensor cores, different rules would apply. But rather than weakening the contribution, this boundedness is what makes it actionable. The paper is not claiming universal truths about deep learning hardware; it is providing a complete, validated framework for the specific (and highly prevalent) case of training transformers on NVIDIA datacenter GPUs.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper operates entirely at the level of GPU kernel throughput and single-layer transformer latency, not end-to-end model training or benchmark accuracy. There is no evaluation dataset in the conventional sense. The "data" are GPU microbenchmarks that sweep GEMM dimensions and transformer hyperparameter configurations on fixed hardware. The only model-level quality validation appears in Section VII-C, where the Pythia model suite test losses are referenced to confirm that shape-optimized models remain on-trend with their non-optimized counterparts, but no accuracy numbers are reported.

• Base model(s). The analysis targets decoder-only transformer models following the GPT-2 architecture (Radford et al., 2019), with specific attention to the 2.7B parameter configuration from GPT-3 (Brown et al., 2020): hidden dimension $h = 2560$, attention heads $a = 32$, layers $L = 32$. This architecture was chosen because it has been directly replicated by OPT (Zhang et al., 2022), GPT-Neo (Black et al., 2021), Cerebras-GPT (Dey et al., 2023), RedPajama-INCITE (Together AI, 2023), and Pythia (Biderman et al., 2023), making it the most widely-copied inefficient architecture in the open-source LLM ecosystem. Additional coverage of Pythia-410M and Pythia-1B in Section VII-C provides inference validation.

• Metrics. All performance is measured in throughput, defined in two complementary ways: (a) teraFLOP/s for raw GEMM and BMM microbenchmarks—this is the absolute computational throughput achieved by the kernel, computed as (2 × FLOPs in the operation) / (kernel execution time), and reported directly by PyTorch's benchmarking infrastructure; (b) single-layer transformer latency for end-to-end operator analysis, where the paper reports the proportion of time spent in each GEMM component (Figure 2, Figure 11) and the relative speedup of optimized vs. standard architectures (Figure 1). The paper does not report absolute wall-clock time or tokens/second. Model quality is referenced only in passing via Pythia test loss trends (Section VII-C) and the claim that shape optimization does not degrade accuracy, but no quantitative accuracy comparison is presented.

• Baselines. The primary baseline throughout the paper is the standard GPT-3 2.7B architecture ($h=2560$, $a=32$, $L=32$), which the authors identify as representative of the status quo due to its replication across multiple model families. In Figure 1, this baseline is compared against two optimized configurations: C1 ($h=2560$, $a=64$—which makes $h/a = 40$) and C2 ($h=2560$, $a=40$—which makes $h/a = 64$, satisfying A100 alignment). The optimal configuration (reducing $a$ to 20, achieving $h/a = 128$) is described in the text but the exact throughput number relative to baseline is reported as 1.18× speedup. For the Pythia inference analysis (Section VII-C, Figure 13), the baselines are the other models in the Pythia suite (particularly Pythia-410M), with the claim that Pythia-1B is "off-trend" in its throughput-to-parameter relationship due to its hardware-efficient shape.

• Generation budget / compute accounting. There is no concept of a "generation budget" in this paper—the analysis is entirely about throughput per unit time for fixed computational work. Every GEMM microbenchmark executes the specified matrix multiplication and measures achieved teraFLOP/s. Every transformer layer benchmark executes a single forward pass through one layer and measures the latency breakdown. The efficiency gains are reported as multiplicative speedups (e.g., 1.18×, "up to 39%") on fixed hardware, not as FLOP reductions to achieve equivalent output. The "39%" figure in the abstract corresponds to the throughput difference between the most and least efficient model shapes shown in Figure 1 at the same total parameter count (approximately 2.7B).

• Cross-validation / statistical protocol. None. This is a deterministic systems benchmarking paper: given fixed hardware, fixed software stack, and fixed GEMM dimensions, the measured throughput is essentially single-valued (modulo minor system noise). The paper does not report error bars, confidence intervals, or multiple trials. The breadth of the analysis comes from sweeping across dimensions (various $h$, $a$, $b$, $s$, $t$ values, multiple GPU architectures) rather than from statistical resampling. The validation strategy is cross-hardware replication: results are shown on V100, A100, and H100 GPUs from NVIDIA, and MI250X from AMD, with the paper arguing that consistent patterns across these platforms validate the underlying first-principles explanation.

---

Main Quantitative Results

The paper's experimental results are organized into three tiers: (1) raw GEMM microbenchmarks that isolate individual hardware effects, (2) transformer-component benchmarks that map hyperparameters to throughput, and (3) end-to-end demonstrations on real model architectures.

Raw GEMM and BMM Microbenchmarks (Section V)

GEMM throughput scales with size but suffers periodic wave quantization drops. Figure 5 presents three controlled sweeps of standalone GEMM operations on V100 and A100 GPUs:

• Figure 5a, sweeping $m$ in $(m, 4096) \times (4096, m)$: Throughput increases monotonically with $m$ as the GEMM transitions from memory-bound (small $m$, where data movement dominates) to compute-bound (large $m$, where arithmetic dominates). On A100, throughput rises from approximately 50 teraFLOP/s at $m=256$ to approximately 300 teraFLOP/s at $m=8192$. No wave quantization effects are visible in this sweep because the tile grid changes smoothly with $m$.

• Figure 5b, sweeping $k$ in $(27648, 4096) \times (4096, k)$: This is the paper's clearest demonstration of wave quantization sawtooth patterns. On the A100 curve, throughput rises with $k$ until sharply dropping at specific thresholds, then rising again. These drops occur when the number of thread blocks crosses a multiple of 108 (the A100 SM count), triggering a new tail wave. The magnitude of the drops is substantial—peak-to-trough variations of approximately 20–30 teraFLOP/s are visible. On V100, the pattern is similar but with thresholds at multiples of 80 SMs.

• Figure 5c, sweeping $k$ in $(2304, 4096) \times (4096, k)$: With a smaller $m$ dimension (2304 vs. 27648), the wave quantization effects are "lessened, as PyTorch is able to better balance the improvements from GEMM parallelization and inefficiencies from wave quantization." The curve is smoother, with smaller peak-to-trough variations, because the smaller grid size gives the autotuner more flexibility in tile size selection. This is a subtle finding: wave quantization exists for all GEMM sizes, but its observability depends on how much freedom the kernel autotuner has to select alternative tile sizes.

Batched matrix multiplications follow the same patterns as individual GEMMs. Figure 6 shows BMM throughput sweeps on V100 (a, c) and A100 (b, d):

• Figures 6a-b, sweeping batch size $b$ in $(b, m, m) \times (b, m, m)$: Throughput increases with batch size as the total computational work grows, saturating when the GPU is fully utilized. The batch size at which saturation occurs depends on $m$—for $m=512$, saturation occurs at batch size ~32 on A100, while for $m=4096$, it occurs at batch size ~4.

• Figures 6c-d, sweeping $m$ in $(b, m, 4096) \times (b, 4096, m)$ for various batch sizes: Throughput increases with $m$ as each sub-GEMM becomes more compute-dense. The paper notes that "the same wave quantization effects would apply (though they do not for these BMM sizes and on these GPU architectures)," meaning the specific dimensions swept happen to avoid triggering tail waves, not that BMMs are immune to wave quantization.

These microbenchmarks serve as the paper's calibration data: they establish the baseline throughput curves for each GEMM shape, which are then referenced when analyzing transformer operators that share those same shape signatures. The paper's entire subsequent analysis is, in effect, an exercise in mapping transformer hyperparameters onto the $m$, $n$, $k$ dimensions that appear in these microbenchmarks.

Transformer Component Benchmarks (Section VI-A)

The attention score and attention-over-value BMMs are highly sensitive to $h/a$ power-of-two alignment. The extensive sweep data in Figures 7a-b through 9 (main text) and Figures 21-47 (Appendix B) constitute the most detailed experimental contribution. For each number of attention heads $a$ (ranging from 8 to 512), the paper sweeps hidden dimension $h$ and measures the throughput of the attention key-query score computation (the $KQ^T$ BMM) and the attention-over-value computation.

Figures 7a and 7b illustrate the key pattern for $a=32$ heads on A100:

"Each plot is a single series... but split by the largest power of two that divides $h/a$ to demonstrate that more powers of two leads to better performance up to $h/a=64$."

When $h/a$ is divisible by 64 (the top band in each plot), throughput is substantially higher than when $h/a$ is divisible by only 32, 16, 8, or smaller powers of two. The throughput stratification is structured and systematic: within a given power-of-two band, throughput increases smoothly with hidden size; between bands, there are discrete jumps upward as $h/a$ crosses thresholds that make it divisible by higher powers of two.

Decreasing attention heads increases throughput, not decreases it—counter to the intuition that more heads mean more parallelism. Figures 8 and 9 show the attention computations with $h/a$ fixed at 64. The key finding:

"a decrease in $a$ is an increase in $h/a$ and these two GEMMs are memory bound, an increase in component matrices size creates much more efficient GEMMs."

At any given hidden size, the configuration with fewer attention heads (larger per-head dimension) achieves higher throughput because each individual BMM sub-operation is larger and more compute-dense, moving it out of the memory-bound regime. The paper notes this is a case where hardware efficiency and naive parallelism expectations conflict: more heads means more independent BMMs, but each BMM becomes so small that kernel launch overhead and memory bandwidth dominate, yielding lower aggregate throughput.

Figure 9 additionally shows wave quantization effects in the attention-over-value computation: "since each line moves in steps of $64h/a$, the BMMs corresponding to each line grow at different rates. This causes the period of the wave quantization effect to appear different for each $a$ value." In practice, this means that the hidden dimensions that avoid tail waves are different for each attention head count—a model architect optimizing for $a=32$ heads might choose a different $h$ than one optimizing for $a=40$, even at the same total parameter budget.

The MLP GEMMs saturate at large hidden dimensions, confirming the compute-bound transition. Figures 10a and 10b show MLP throughput as a function of hidden dimension for $a=128$ fixed. Throughput rises with $h$ and begins to plateau at approximately $h = 4096$ on A100, indicating the MLP GEMMs have become compute-bound. The paper's strategic implication: make $h$ as large as possible (within the parameter budget) to push these GEMMs into saturation, since Figures 2 and 11 show the MLP and QKV GEMMs dominate total latency for large models.

The QKV transform and MLP block are the high-value optimization targets for large models. Figure 11 breaks down the proportion of latency spent in each GEMM component for one transformer layer at various model sizes. For the largest configurations (right side of the plot), the QKV transform and the two MLP GEMMs (up-projection and down-projection) collectively account for the vast majority of total latency, while attention score computation and attention-over-value shrink to negligible fractions. This finding grounds the paper's strategic advice: "for the largest models, the QKV transformation in the attention block along with the MLP block are the most prevalent GEMMs. Therefore, the overall latency of the model would benefit most from optimizing these kernels." Since $h/a$ primarily affects the attention score and AOV computations (which become less significant at scale), the paper argues that "only a small portion of the latency of large models is the attention score computation and attention over value computation GEMMs, so an increase in the latency of these components will have only a small effect on the end-to-end model performance."

FlashAttention simplifies the attention sizing requirements. Figure 12 sweeps hidden dimension with FlashAttention v2 (Dao, 2023) and $a=128$ fixed. The paper reports that "FlashAttention follows a roofline model," meaning its throughput scales smoothly with $h$ without the wave quantization sawtooth pattern or power-of-two stratification seen in standard attention BMMs. This means that when using FlashAttention, the $h/a$ alignment constraint can be relaxed for the attention computation—only the requirement that $h$ be "as large as possible" remains. The paper explicitly frames this as a simplification for small models: "we recommend either using FlashAttention v2 for small models to mitigate these effects, or increasing $h$ as much as possible to reach the saturation point."

The vocabulary embedding transformation aligns at multiples of 64. Figures 20a and 20b sweep vocabulary size $v$ and hidden dimension $h$ for the logit layer GEMM. The paper reports: "The performance of the logit layer is maximized when $v$ is a multiple of 64, therefore it is best to pad the vocab size to the nearest multiple of 64. Likewise, the layer also performs best with a hidden size that is a multiple of 64." This corroborates the practitioner observation (Karpathy, Reference [19]) that padding GPT-2's vocabulary from 50,257 to 50,304 tokens yields a 25% speedup—the paper provides the hardware explanation (Tensor Core alignment at 64 FP16 elements) and generalizes it beyond a single model.

End-to-End Demonstrations on Real Architectures

The GPT-3 2.7B architecture achieves 1.18× speedup through head count reduction. Figure 1 provides the headline result: the standard GPT-3 2.7B configuration ($h=2560$, $a=32$, giving $h/a=80$) achieves a certain single-layer throughput (shown as the leftmost bar). Configuration C1 ($h=2560$, $a=64$, giving $h/a=40$) achieves slightly lower throughput because the per-head dimension is even further from optimal alignment. Configuration C2 ($h=2560$, $a=40$, giving $h/a=64$) achieves higher throughput—the per-head dimension now satisfies A100 Tensor Core alignment. The optimal configuration (reducing $a$ to 20, giving $h/a=128$) is described in the text as achieving a 1.18× speedup over the baseline, meaning the layer runs 18% faster while preserving the total parameter count of approximately 2.7 billion.

The paper explicitly contrasts two options for fixing $h/a$: increasing $h$ from 2560 to 4096 (which would give $h/a = 128$ with $a=32$, satisfying alignment) or decreasing $a$ from 32 to 20 (also giving $h/a = 128$). The first option doubles the parameter count to approximately 6.7 billion. The second preserves the 2.7B parameter count. The paper chooses the latter for the 2.7B comparison, but notes that the former demonstrates how the alignment constraint interacts with model scaling decisions.

Pythia-1B validates that shape-optimized models maintain quality at higher throughput. Figure 13 shows inference latency for the Pythia model suite using DeepSpeed-MII. The paper identifies Pythia-410M and Pythia-1B as "off-trend" in their throughput relative to the rest of the suite: Pythia-1B achieves higher throughput than would be expected from its parameter count, while Pythia-410M achieves lower throughput than expected. The explanation: Pythia-1B was designed with hardware-aware principles (fewer attention heads, larger hidden dimension relative to its size), while Pythia-410M uses a less efficient shape. The critical validation: "Despite these architectural changes, the test loss of Pythia-1B is on-trend with the rest of the suite while having significantly higher training and inference throughput." This is the paper's only empirical evidence that shape optimization does not come at a model quality cost.

---

Ablation Studies and Robustness Checks

This paper does not contain ablation studies in the conventional machine learning sense (systematically removing or varying components to isolate their contribution). Instead, its robustness is established through dimensional sweeps across multiple axes and cross-platform replication. The following are the closest analogues to ablation experiments:

Power-of-two alignment sweep for attention GEMMs (Appendix B, Figures 21-47): For each attention head count from $a=8$ to $a=512$, the paper provides a full sweep of hidden dimension and color-codes the results by the largest power of two dividing $h/a$. This is effectively an ablation across alignment quality: it shows that throughput within each power-of-two band is similar, and that bands are cleanly stratified with higher powers of two achieving strictly higher throughput. The key quantitative finding: the throughput drop from $h/a$ divisible by 64 to $h/a$ divisible by 32 is substantial; the drop from 32 to 16 is smaller; and below 8, throughput varies erratically because Tensor Core utilization becomes severely degraded. This systematic sweep validates that the "multiple of 64" rule is not an arbitrary threshold but reflects a genuine hardware discontinuity.

Tensor parallel degree sweep for QKV transform (Figure 16): The paper sweeps tensor parallel degree $t$ in the QKV transform GEMM, showing that as $t$ increases (making $h/t$ smaller), throughput decreases. This validates the Rule 5 recommendation that $t$ "should be as small as possible." The underlying mechanism: smaller $h/t$ pushes the GEMM toward the memory-bound regime, reducing teraFLOP/s.

3D vs. 2D tensor dimension ordering for GEMMs (Figure 14): The paper benchmarks GEMMs with batched input tensors in different dimension orderings: $(2048, 4, n) \times (n, 3n)$, $(4, 2048, n) \times (n, 3n)$, and the flattened $(8192, n) \times (n, 3n)$. The finding: "the ordering of the batched dimension does not affect performance. The batched implementation is also the same speed as a 2-dimensional GEMM." This is a robustness check on the paper's modeling choice: representing transformer GEMMs as 2-dimensional matrix multiplications (as done throughout Sections III-C and VI-A) is valid because PyTorch's torch.nn.functional.linear handles batched inputs with the same efficiency as unbatched GEMMs.

GPU architecture comparison (V100 vs. A100 vs. H100; NVIDIA vs. AMD): Throughout the paper, the same sweeps and patterns are demonstrated across V100, A100, and H100 GPUs. Figure 5 shows GEMM throughput on both V100 and A100, with both exhibiting wave quantization but at different SM-count thresholds (80 vs. 108). Figure 6 shows BMM throughput on both platforms. Appendix B provides A100-specific sweep data. The MI250X results (mentioned in the experimental setup) are included to show that the principles extend to non-NVIDIA hardware. The key cross-platform finding: the qualitative behavior (wave quantization, Tensor Core alignment at power-of-two boundaries) is consistent across platforms, but the quantitative thresholds (alignment requirement of 8 vs. 64 FP16 elements, SM counts of 80/108/144) are hardware-specific. This is precisely what the paper's co-design thesis predicts.

FlashAttention vs. standard attention comparison (Figure 12 vs. Figures 7-9): Figure 12's FlashAttention sweep can be viewed as an ablation of the kernel implementation: replacing the standard attention BMMs with FlashAttention eliminates the power-of-two stratification and wave quantization patterns visible in Figures 7-9, replacing them with a smooth roofline curve. This validates that the complex patterns in the standard attention benchmarks are indeed kernel-implementation artifacts and not properties of the attention computation itself. It also establishes the boundary of applicability for the paper's $h/a$ rules: they apply to standard attention implementations but can be relaxed when FlashAttention is used.

SwiGLU intermediate dimension search (Section VII-B, narrative only): While not presented as a formal ablation table, the paper describes performing "a brute-force search" over intermediate dimensions near $8h/3$ to find values that satisfy Tensor Core alignment while preserving approximate parameter count parity. The finding that "Llama-2-7B's intermediate size is indeed one of the best performing sizes in its range" serves as a post-hoc validation of the paper's claim that the $8/3$ coefficient is negotiable. However, the paper does not present the sweep data—no figure or table shows throughput as a function of MLP expansion factor. This is a notable omission: the reader is asked to trust that the brute-force search produces gains without being shown the tradeoff curve.

---

Critical Assessment

Does the 39% throughput improvement claim hold? Yes, but only for the specific configuration shown in Figure 1.

The abstract claims "the throughput of models with 'efficient' model shapes is up to 39% higher while preserving accuracy compared to models with a similar number of parameters but with unoptimized shapes." Figure 1 supports this for the specific case of 2.7B parameter transformers with different shapes. However, the 39% figure represents the maximum gap between any two configurations in Figure 1, not an average improvement or an improvement at the optimal configuration. The optimal configuration (C2, $a=40$, $h/a=64$) achieves a speedup closer to the ~18% figure cited for the $a=20$ optimization. The 39% figure should be understood as an upper bound on the possible variation in throughput due to shape choices at this parameter scale, not as a guarantee that any model can be improved by 39%.

More importantly, the paper does not systematically characterize how the throughput gap varies with model scale. Figure 1 covers 2.7B parameters. Does the gap grow, shrink, or stay constant at 1B, 7B, 13B, 70B? Figure 11 suggests that for larger models, the GEMMs most sensitive to shape (attention score and AOV) shrink in relative importance, while the less shape-sensitive GEMMs (QKV, MLP) dominate. This implies the throughput gap might narrow at larger scales—the very models where the paper's claimed economic stakes are highest. Without sweep data across model sizes, this remains speculative.

Do the rules generalize beyond the models tested? Partially, but the evidence is thin.

The paper presents the prescriptive rule set (vocabulary divisible by 64, $h/a$ multiple of power of two, etc.) as universal guidelines for transformer architects. The experimental support for this claim comes entirely from:

• GEMM-level microbenchmarks (Figures 5-6), which validate that the underlying hardware phenomena (Tensor Core alignment, wave quantization) are real.
• Single-layer transformer component sweeps (Figures 7-12, Appendix B), which validate that the mapping from hyperparameters to GEMM dimensions is correct for a standard decoder-only architecture.
• Two end-to-end demonstrations: the GPT-3 2.7B reconfiguration (Figure 1) and the Pythia inference comparison (Figure 13).

The gap between "these dimension constraints produce efficient GEMMs" and "these dimension constraints produce efficient end-to-end training of models at arbitrary scale" is substantial. Multi-layer training involves pipeline parallelism, activation checkpointing, gradient accumulation, optimizer overhead, and communication—none of which is captured by single-layer throughput measurements. The paper explicitly acknowledges this: "We leave an analysis of the implications of pipeline and sequence parallelism on optimal model shapes to future work." This is not a minor caveat; it means the paper's measured speedups are upper bounds on what would be observed in actual distributed training, where overhead from communication and scheduling may dilute the GEMM-level gains.

Is model quality preservation adequately demonstrated? No.

The paper's claim that shape optimization "preserves accuracy" rests on a single sentence in Section VII-C: "the test loss of Pythia-1B is on-trend with the rest of the suite." No accuracy numbers, no downstream evaluation, no comparison of the standard 2.7B configuration against the optimized 2.7B configuration on any benchmark. This is a significant gap because the paper recommends reducing attention heads—a change that alters the model's inductive biases. Fewer heads means each head must capture a larger fraction of the representation, which could affect the model's ability to attend to multiple distinct patterns simultaneously. The fact that Pythia-1B's test loss is "on-trend" is weak evidence: Pythia-1B and Pythia-410M are different sizes, so the comparison does not isolate the effect of shape. A proper accuracy evaluation would train the standard and optimized 2.7B architectures from scratch on identical data and compare benchmark scores.

This is not a fatal flaw for a systems paper, but it means the paper's core claim—that shape optimization improves throughput "while preserving accuracy"—is asserted rather than demonstrated. A reader who wants to apply these rules must decide, without guidance from the paper, whether the specific hyperparameter changes recommended (fewer heads, padded vocabulary, non-standard MLP expansion factors) will affect their model's quality.

Is the SwiGLU brute-force search result credible? Yes, but the paper provides no evidence for it.

The Section VII-B narrative about Llama-2's MLP intermediate dimensions is compelling as a case study, but it describes post-hoc validation rather than a controlled experiment. The paper states that "running a brute-force search reveals that Llama-2-7B's intermediate size is indeed one of the best performing sizes in its range" without showing the search results, stating the search space, or reporting the throughput of alternative expansion factors. This is particularly problematic because the SwiGLU extension is one of the paper's most actionable contributions—many modern architectures use SwiGLU—yet the paper provides no quantitative data to support its specific recommendation.

Are the experiments sufficient to distinguish wave quantization from other effects? Mostly yes, but with a key ambiguity.

The sawtooth pattern in Figure 5b is the paper's primary empirical evidence for wave quantization, and it is convincing: the periodic drops at thresholds corresponding to the SM count are a clean diagnostic. However, the paper's discussion of Figure 5c introduces an ambiguity: "when the size of the GEMM is sufficiently large, PyTorch may automatically choose a tile size that decreases quantization effects. In Figure 5c, the effects of wave quantization are lessened, as PyTorch is able to better balance the improvements from GEMM parallelization and inefficiencies from wave quantization to improve throughput." This means wave quantization is not a fixed property of a GEMM dimension but depends on the kernel autotuner's choice of tile size—which the user does not directly control. A dimension that triggers wave quantization with the default tile size might avoid it if the autotuner selects a different tile size. This makes the paper's formal wave quantization condition (Section VI-B) less practically useful than it appears: it describes the condition under a specific tile size assumption, but the actual tile size is chosen by PyTorch at runtime.

The paper acknowledges this complexity obliquely: "there is not a transformer configuration with GEMMs that fill tensor core requirements without wave quantization inefficiency. Further, PyTorch's linear algebra backend can use different tile sizes for each GEMM. Therefore, PyTorch is unable to efficiently overcome the effects of wave quantization." The reasoning here is somewhat circular—PyTorch cannot overcome the effects, but also the effects are lessened when PyTorch selects different tile sizes. A reader seeking to apply the wave quantization condition to their own architecture would need more guidance on when the condition holds in practice versus when the autotuner compensates.

What experiments would strengthen the paper?

Several missing experiments are notable:

1. End-to-end training throughput for the optimized 2.7B architecture. The paper reports single-layer throughput (Figure 1) but does not train the full 32-layer model to convergence and compare wall-clock time against the standard architecture. Single-layer measurements omit multi-GPU communication, gradient synchronization, pipeline bubbles, and optimizer overhead—all of which could dilute or amplify the GEMM-level gains.

2. Accuracy benchmarks for shape-optimized vs. standard architectures at matched parameter count. Training the standard ($h=2560$, $a=32$) and optimized ($h=2560$, $a=20$ or $h=2560$, $a=40$) 2.7B architectures on identical data and comparing perplexity and downstream task performance would directly validate the "while preserving accuracy" claim.

3. Throughput sweep across model scales. Measuring the throughput gap between "efficient" and "inefficient" shapes at 125M, 350M, 1.3B, 2.7B, 6.7B, and 13B parameters would reveal whether the 39% gap is scale-dependent. Figure 11 implies it narrows at larger scales (because attention GEMMs shrink in relative importance), but this needs direct measurement.

4. Ablation of individual rules. The paper bundles multiple recommendations ($h/a$ alignment, $v$ padding, $b \cdot s$ large, $t$ small) into a single rule set. An ablation that applies each rule independently to the same base architecture and measures the marginal throughput gain from each would show which rules matter most at which scales.

5. SwiGLU intermediate dimension sweep with quantitative data. A figure showing MLP throughput as a function of expansion factor near $8/3$, with the standard $4h$ MLP as a baseline, would substantiate one of the paper's most actionable recommendations.

Summary of claims and support

Claim: Throughput of models with efficient shapes is up to 39% higher. Supported by Figure 1 for 2.7B parameters, but this represents the maximum gap rather than the typical or optimal improvement, and may not generalize to other scales.

Claim: Shape optimization preserves accuracy. Weakly supported by Pythia-1B test loss trend (Section VII-C, Figure 13). No direct comparison of shape-optimized vs. standard architectures at matched parameter count is provided.

Claim: Rules derive from GPU first principles (Tensor Cores, wave quantization, tile quantization). Well supported by the GEMM microbenchmarks (Figures 5-6) and the systematic stratification by power-of-two alignment (Figures 7-9, Appendix B). The hardware explanation is the paper's strongest contribution.

Claim: Rules generalize across GPU architectures. Supported by V100/A100/H100 comparisons, with the caveat that quantitative thresholds change. The MI250X results are mentioned but not shown, weakening this claim for non-NVIDIA hardware.

Claim: The $8/3$ SwiGLU coefficient is negotiable. Supported by the post-hoc observation that Llama-2 models deviate from it, and the assertion that a brute-force search validates these deviations. However, no sweep data is presented, making this the paper's least empirically-grounded recommendation.

Overall assessment. The paper's empirical work is strongest at the level it focuses on: characterizing how GEMM dimensions map to throughput, and demonstrating that transformer hyperparameters control those dimensions. It is weakest at connecting these kernel-level insights to end-to-end training outcomes—both throughput in realistic multi-GPU training and model quality. The paper is best read as establishing why shape optimization should work, based on first-principles GPU behavior, rather than as conclusive proof that the specific rules produce universal gains at all scales. For a practitioner, the rules are high-confidence for the attention and vocabulary operations they directly address, lower-confidence for their interaction with distributed training overhead, and essentially unvalidated for their effect on model quality beyond the single Pythia data point.

6. Limitations and Trade-offs

Single-GPU Analysis Without Distributed Training Overhead

The assumption or constraint. The paper explicitly restricts its analysis to single-GPU computations: "Since this paper focuses on the computations being done on a single GPU, we will largely ignore parallelism" (Section III-C). The prescriptive rules (Section VI-B) and all throughput measurements (Figures 1, 7–12) are derived from single-layer, single-GPU benchmarks. The paper acknowledges this gap directly: "We leave an analysis of the implications of pipeline and sequence parallelism on optimal model shapes to future work" (Section III-C).

The consequence. In realistic large-scale training, multi-GPU communication—all-reduce operations for tensor parallelism, point-to-point transfers for pipeline parallelism, and gradient synchronization for data parallelism—can dominate end-to-end latency. The paper's measured GEMM-level speedups (18–39% on single-layer throughput) are therefore upper bounds on what would be observed in distributed training. A shape optimization that improves single-GPU GEMM throughput by 20% might yield only a 5–10% improvement in end-to-end distributed training throughput if the model is communication-bound. Moreover, the paper's recommendations interact with parallelism strategies in ways that could be counterproductive: making $h$ larger (to saturate MLP GEMMs) increases the total parameter count per layer, which may require reducing pipeline micro-batch size or increasing tensor parallelism degree—both of which introduce overheads the paper does not model. The recommendation to minimize tensor parallel degree $t$ (Rule 5) is sensible for single-GPU efficiency but conflicts with the memory constraints that force high tensor parallelism in the first place for large models.

What evidence exists in the paper. None. The paper contains no multi-GPU training benchmarks, no communication overhead measurements, and no end-to-end training time comparisons for a full model at scale. The Pythia inference validation (Figure 13) uses DeepSpeed-MII for single-GPU inference, not multi-GPU training. The entire evidence base consists of (a) isolated GEMM microbenchmarks (Figures 5–6), (b) single-layer transformer component sweeps (Figures 7–12), and (c) single-layer comparisons of alternative 2.7B architectures (Figure 1). The critical gap—how these single-layer gains translate to multi-GPU, multi-layer training throughput—is unmeasured.

Mitigation status. The paper acknowledges the limitation (Section III-C, Section VIII) but does not attempt to address it. The suggestion to study "pipeline and sequence parallelism" interactions in future work is stated but no methodology or preliminary results are provided. A practitioner deploying these recommendations at scale must therefore treat the quoted speedup figures as optimistic estimates and should benchmark their specific distributed training configuration rather than relying on the paper's single-GPU numbers.

---

Model Quality Preservation Is Asserted, Not Demonstrated

The assumption or constraint. The paper's thesis is that shape optimization improves throughput "while preserving accuracy" (abstract, Section IX), "without any algorithmic changes" (Section 1 of prior summary), and with model quality remaining "on-trend" (Section VII-C). The only evidence offered for this claim is a single sentence about the Pythia model suite: "Despite these architectural changes, the test loss of Pythia-1B is on-trend with the rest of the suite while having significantly higher training and inference throughput" (Section VII-C, Figure 13).

The consequence. The paper recommends concrete hyperparameter changes that alter model inductive biases: reducing the number of attention heads (from 32 to 20 in the 2.7B case study), padding vocabulary size, and adjusting MLP expansion factors (in the SwiGLU case). Each of these changes could affect model quality in ways the paper does not measure:

• Fewer attention heads: Each head has a larger per-head dimension ($h/a$ increases), meaning each head captures a broader slice of the representation. This could reduce the model's ability to attend to multiple distinct syntactic or semantic patterns simultaneously—a well-documented tradeoff in multi-head attention (Michel et al., 2019; Voita et al., 2019). The paper does not evaluate whether the optimized 2.7B architecture with 20 heads achieves comparable perplexity, downstream task performance, or scaling behavior to the standard 32-head architecture when trained from scratch on identical data.
• Padded vocabulary: Adding 47 unused token embeddings (from 50,257 to 50,304) is a trivial change unlikely to affect quality, but the paper does not verify this.
• Non-standard MLP expansion factors: For SwiGLU architectures, the paper recommends deviating from the $8/3$ coefficient to find hardware-aligned intermediate dimensions (Section VII-B). The effect of these deviations on model quality is completely unmeasured—the paper simply observes that Llama-2 models deviate from $8/3$ and asserts that a brute-force search validates the choice, without presenting any quality data.

The Pythia-1B comparison is fundamentally insufficient as evidence for quality preservation: Pythia-1B and Pythia-410M are different model sizes (1B vs. 410M parameters), so the comparison does not isolate the effect of shape from the effect of scale. A model that is 2.4× larger should achieve lower test loss than a smaller model, regardless of shape. Showing that the 1B model's loss is "on-trend" with the rest of the suite tells us nothing about whether a differently-shaped 1B model would perform better or worse.

What evidence exists in the paper. A single sentence in Section VII-C and Figure 13, which plots inference latency (not model quality) against a qualitative claim about test loss trends. The paper provides no perplexity numbers, no downstream task evaluations, no training curves, and no head-to-head comparison of the standard and optimized 2.7B architectures trained on identical data. The paper does not even state what "on-trend" means quantitatively—is the test loss within 0.01? 0.1? 1.0 nats of the trend line?

Mitigation status. Not attempted. The paper treats accuracy preservation as a self-evident consequence of preserving total parameter count, but this conflates capacity (parameters) with inductive bias (architecture). The paper's only suggestion of future work in this direction is implicit in the broader call for co-design, but there is no proposal to validate the quality implications of the specific hyperparameter changes recommended.

---

Difficulty Estimation for Wave Quantization Is Left to the Autotuner, Making the Formal Condition Less Actionable Than It Appears

The assumption or constraint. Section VI-B provides a formal condition for avoiding wave quantization inefficiency: the number of thread blocks must be divisible by the SM count. This condition assumes a specific tile size ("assuming a tile size of $128 \times 256$ which is the most efficient"). However, the paper also observes that PyTorch's linear algebra backend can dynamically select different tile sizes, and that this flexibility can mitigate wave quantization: "when the size of the GEMM is sufficiently large, PyTorch may automatically choose a tile size that decreases quantization effects" (Section V, discussing Figure 5c). The paper further acknowledges that "there is not a transformer configuration with GEMMs that fill tensor core requirements without wave quantization inefficiency," and that "PyTorch is unable to efficiently overcome the effects of wave quantization" (Section VI-B).

The consequence. There is a tension between two statements: (1) wave quantization is the dominant observable performance effect that architects must design around, and (2) the autotuner can select tile sizes that mitigate it, but not eliminate it. This ambiguity makes it unclear when the formal condition should guide hyperparameter choices versus when the autotuner will compensate. A model architect who follows the paper's wave quantization condition may be optimizing for a tile size that PyTorch never actually uses, while ignoring a different tile size configuration that would produce better actual throughput. Conversely, the paper provides no guidance on how to predict or influence which tile size the autotuner will select, making the wave quantization condition a theoretical ideal rather than a reliable practical constraint.

More specifically, the paper's claim that PyTorch "is unable to efficiently overcome the effects of wave quantization" is followed by the prescriptive rule to make $b \cdot s$ "as large as possible" (Rule 2)—but $b \cdot s$ is the dimension that controls the number of thread blocks in the output grid, and making it larger changes the wave quantization pattern rather than eliminating it. The paper does not provide a method for choosing $b \cdot s$ to land at a "good" point on the sawtooth curve (before the drop, rather than after), which would be more useful than the blanket "as large as possible" recommendation. This is analogous to the difference between knowing that a function is periodic and knowing how to evaluate it at its peaks—the paper provides the periodicity condition but not the peak-finding procedure.

What evidence exists in the paper. The sawtooth pattern in Figure 5b clearly demonstrates wave quantization at work, and the formal condition (Section VI-B) correctly describes the mechanism. However, Figure 5c shows that the effect is "lessened" when PyTorch has more autotuning flexibility, and the paper does not characterize when this mitigation occurs versus when it does not. The paper provides no sweep of $b \cdot s$ values for a fixed transformer architecture showing where the throughput peaks and valleys actually fall, which would be the natural way to operationalize the wave quantization condition.

Mitigation status. Partial. The paper acknowledges that wave quantization cannot be fully eliminated ("there is not a transformer configuration... without wave quantization inefficiency") and provides the formal condition as a diagnostic tool. However, the gap between the condition (which assumes a specific tile size) and the autotuner's actual behavior (which may select different tile sizes) is not addressed beyond the observation that the autotuner can sometimes help. The paper's primary mitigation strategy is implicit: choose dimensions that are large (to reduce the relative impact of tail waves) and aligned (to get the best possible Tensor Core utilization in the threads that do execute). This is sensible but leaves the architect without a concrete procedure for predicting or optimizing wave quantization in their specific configuration.

---

Single Model Family, Single Task Domain, Single Hardware Vendor (Effectively)

The assumption or constraint. All experimental results use decoder-only transformer architectures (GPT-2 style, Section III-C) on NVIDIA datacenter GPUs (V100, A100, H100). The paper notes that "most of our conclusions also apply to encoder-only models" but that the analysis "will largely not apply to encoder-decoder models" (Section III-C). No non-transformer architectures are considered. No non-NVIDIA GPU results are shown in the main experimental figures (the MI250X is mentioned in the experimental setup, Table III, but no MI250X data appears in any figure or table).

The consequence. The paper's prescriptive rules are architecture-specific and hardware-specific in ways that are not fully bounded. The core mapping from hyperparameters to GEMM signatures (Table II) assumes a specific decoder-only design with fused QKV projection, per-head dimension $h/a$, and 4× MLP expansion. For encoder-decoder models (e.g., T5, BART), the cross-attention mechanism introduces additional GEMMs with different dimension signatures—the encoder output sequence length becomes an additional variable, and the key/value projections come from a different source than the queries—all of which changes which dimensions control Tensor Core alignment and wave quantization. For mixture-of-experts architectures, the conditional computation introduces sparse GEMMs that may have entirely different performance characteristics. For non-transformer architectures (state-space models, recurrent models, convolutional models), the paper's entire GEMM decomposition would need to be reconstructed from scratch.

On the hardware side, the paper's rules are calibrated to NVIDIA's Tensor Core architecture, which imposes specific alignment constraints (8 FP16 elements on V100, 64 on A100/H100). AMD's MI250X has different matrix core requirements, and Google's TPUs have a completely different systolic array architecture with different dimension sensitivity. The paper mentions MI250X results (Section IV-A) but does not present them, so the reader cannot verify whether the same methodology produces useful rules on non-NVIDIA hardware or whether the specific thresholds (multiples of 64) are NVIDIA-specific.

What evidence exists in the paper. The restriction to decoder-only transformers is stated explicitly (Section III-C). The MI250X is mentioned only in the hardware table (Table III) and software setup (Section IV-B). No cross-architecture experiments (encoder-only vs. decoder-only, transformer vs. non-transformer) are performed. No non-NVIDIA GPU benchmarking results are shown in any figure.

Mitigation status. The paper is transparent about the architecture restriction ("our analysis will largely not apply to encoder-decoder models") but does not discuss the implications for other architecture families beyond this sentence. The hardware specificity is framed as a feature of the co-design thesis (Section 4, prior summary) rather than a limitation, but the absence of non-NVIDIA data means the generalization claim—that the methodology transfers even if the specific thresholds differ—is asserted rather than demonstrated. A practitioner using AMD GPUs or TPUs cannot apply the paper's rules directly and receives no guidance on how to adapt them.

---

Latency vs. Throughput: The Paper Optimizes for a Metric That May Not Match Deployment Constraints

The assumption or constraint. The paper measures only throughput (teraFLOP/s or single-layer latency for a fixed computational unit). It does not measure wall-clock latency for processing a user request, which depends on the serial dependencies within and across layers. The revision model example from the previous paper illustrates the issue: sequential revisions maximize throughput per generation but serial execution makes wall-clock time far higher than parallel sampling.

The consequence. The paper's recommendations can create a throughput-latency conflict that is never discussed. The recommendation to make $b \cdot s$ as large as possible (Rule 2) improves GEMM utilization but may increase latency for interactive applications where the user waits for a single sequence. The recommendation to minimize tensor parallel degree $t$ (Rule 5) improves per-GPU GEMM efficiency but may force the use of pipeline parallelism (to fit the model), which introduces pipeline bubbles that increase latency for individual requests. More subtly, the recommendation to use fewer attention heads (to increase $h/a$) reduces the parallelism available within the attention computation—each head processes independently, so fewer heads means less opportunity for concurrent execution of attention operations. The paper argues this is acceptable because attention GEMMs are "a small portion of the latency of large models" (Figure 11), but Figure 11 measures proportion of total compute, not wall-clock latency. If attention is on the critical path, then reducing head parallelism could increase end-to-end latency even if it improves throughput.

The distinction matters for deployment: a throughput-maximizing configuration is appropriate for offline batch processing (training, bulk inference), while a latency-minimizing configuration is appropriate for interactive serving (chatbots, real-time applications). The paper's rules are optimized exclusively for the former, and the paper does not acknowledge the existence of the latter use case or how its recommendations might differ.

What evidence exists in the paper. None. Every metric in the paper—teraFLOP/s, single-layer latency proportion—measures throughput. Wall-clock latency for end-to-end inference is mentioned only in Figure 13 (DeepSpeed-MII inference), which also reports latency but does not discuss the throughput-latency tradeoff. The paper does not measure how changing $b$, $a$, or $t$ affects the latency of processing a single sequence from input to output.

Mitigation status. Not addressed. The paper does not discuss latency, interactive serving, or real-time constraints. A practitioner serving models in production cannot determine from the paper whether the recommended shape optimizations will improve or degrade p50/p99 latency for their specific workload.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper does not introduce a new algorithm, architecture, or training technique. Its contribution is methodological reframing: it demonstrates that model architecture hyperparameters—specifically those controlling the shape of matrix multiplications—should be treated as first-class performance levers in the GPU computing stack, not as arbitrary inherited constants. This is not a paradigm shift (the underlying GPU phenomena were known), but it is a reclassification of responsibility: the paper moves shape optimization from the domain of GPU kernel engineers (who cannot change model architecture) into the domain of model architects (who can), by providing the translation layer between the two.

The reframing has several concrete consequences for how the field thinks about transformer design:

Model architecture becomes a hardware-dependent optimization problem. Prior to this work, the standard practice for choosing transformer hyperparameters was what the paper calls "convenient borrowing"—copying $h$, $a$, and $L$ from a prior paper's architecture table, occasionally adjusting for total parameter budget, but rarely considering hardware at all. The paper's demonstration that a 2.7B parameter model's single-layer throughput can vary by up to 39% based purely on shape choices (Figure 1) makes this practice empirically indefensible. More importantly, by providing the complete GEMM decomposition (Table II) and the prescriptive rule set (Section VI-B), the paper makes it easy to stop borrowing blindly—the architect now has a checklist that takes minutes to evaluate against any proposed architecture. The economic argument is simple: at the scale of modern LLM training runs, a 20% throughput difference represents millions of dollars of compute, and the cost of checking whether $h/a$ is divisible by 64 is zero.

The $8/3$ coefficient is revealed as a design convention, not a constraint. The SwiGLU case study (Section VII-B) identifies a specific instance where the field has been treating a suggested constant as immutable. The observation that Llama-2-7B and Llama-2-70B use different MLP expansion factors (2.6875 and 3.5, respectively, versus the recommended 2.667) is presented as evidence that the Llama team independently discovered—but did not document—what the paper now articulates as a general principle: hardware alignment constraints should override parameter-count parity conventions when the two conflict, because the throughput cost of misalignment exceeds the quality cost (if any) of a slight parameter count deviation. This is a design norm the paper seeks to establish: any "standard" architectural constant that does not derive from a mathematical necessity (e.g., the requirement that $h/a$ be integer) should be treated as negotiable when hardware efficiency is at stake.

Wave quantization is elevated to a first-order concern for model architects. While wave quantization was previously understood by GPU kernel implementers—and occasionally surfaced in practitioner tweets and code comments (References [18, 19, 33])—the paper provides the first formal condition (Section VI-B) and the first systematic demonstration (Figures 5b, 7–9, Appendix B) of how it affects transformer throughput. More importantly, the paper distinguishes wave quantization (observable, partially controllable via dimension sizing) from tile quantization (hard to observe, not controllable by the architect), giving practitioners a clear diagnostic target. This is a knowledge consolidation that takes scattered, informal practitioner lore and grounds it in first-principles GPU architecture, making it teachable and reproducible.

Hardware specificity becomes a design principle rather than a limitation. The paper's co-design thesis—that model architectures should be tuned to specific GPU targets—challenges the prevailing norm of hardware-agnostic architecture design. This is a normative shift: the paper argues not just that hardware-aware design can improve throughput, but that it should be standard practice, and that the field's habit of designing architectures in hardware-agnostic terms is leaving substantial performance on the table. The tension this creates—models optimized for A100 may run inefficiently on H100, and vice versa—is not avoided but embraced as a necessary tradeoff that architects should explicitly manage rather than ignore.

Which research directions become more attractive, and which become less so. The paper's analysis suggests that:

• More attractive: Work on automated architecture search where the objective function includes hardware throughput. The paper's rule set provides a differentiable or at least evaluable proxy for GPU efficiency that could be incorporated into neural architecture search (NAS) pipelines. Work on hardware-dependent scaling laws—extending Hoffmann et al. (2022) to include model shape as a dimension—becomes natural, since the paper shows that shape can trade off against parameter count for fixed throughput. Work on compiler-level optimization that automatically rewrites model architectures to satisfy the paper's constraints (padding dimensions, adjusting head counts) without architect intervention.

• Less attractive: Manual, trial-and-error shape optimization based on profiling individual kernels. The paper's rule set largely automates this analysis. Research that assumes model architecture hyperparameters are free variables independent of hardware—the paper's core finding is that this assumption is costly. Work that focuses exclusively on kernel implementation optimization (rewriting attention, optimizing GEMM kernels) without considering whether the model dimensions are aligned—these are complementary but the paper shows that shape optimization alone captures a large fraction of the available gain without implementation changes.

---

Follow-Up Research This Work Enables

Automated shape-hardware co-optimization via differentiable throughput proxies. The paper's rule set ($h/a$ multiple of 64, $b \cdot s$ large, $v$ padded) is currently applied manually. A natural next step is to encode these constraints—or more ambitiously, to learn a differentiable proxy for GEMM throughput as a function of dimensions—into neural architecture search pipelines that jointly optimize for model quality and hardware efficiency. The paper's extensive microbenchmark data (Figures 5–6, Appendix B) provides the training signal for such a proxy: given a dataset of (GEMM dimensions, measured teraFLOP/s) pairs across multiple GPU architectures, one could train a small predictor network that estimates throughput for any proposed architecture without running benchmarks. The key experiment: compare architectures discovered by this proxy-guided search against both the paper's hand-derived rules and standard borrowed architectures, measuring both training throughput (on the target GPU) and downstream task quality. A strong result would show that the search discovers non-obvious shape choices (e.g., attention head counts that are not powers of two, or MLP expansion factors far from 4× or $8/3\times$) that outperform the paper's rules on specific hardware configurations.

End-to-end distributed training validation of the single-GPU rules. The paper's most significant limitation is that all throughput measurements are single-layer, single-GPU (Section 6, Limitations). A direct follow-up would train the standard GPT-3 2.7B architecture ($h=2560$, $a=32$) and the paper's optimized variant ($h=2560$, $a=20$) from scratch on identical data using a realistic multi-GPU training setup (e.g., 8×A100 nodes with tensor parallelism, pipeline parallelism, and data parallelism as appropriate for the model scale). The experiment would measure three quantities: (a) total wall-clock training time to a fixed number of tokens, capturing communication and scheduling overhead; (b) final validation perplexity and downstream benchmark scores, validating or refuting the paper's implicit quality-preservation claim; (c) the ratio of single-layer speedup to end-to-end speedup, which would calibrate how much of the paper's reported gain survives in distributed training. A negative result—where the 18% single-layer speedup translates to only 2–3% end-to-end gain due to communication bottlenecks—would be highly informative, as it would bound the practical importance of shape optimization at scale and motivate research on shape-aware parallelism strategies.

Quality-shape Pareto frontier: how far can $h/a$ be pushed before accuracy degrades? The paper recommends reducing attention heads to increase $h/a$, citing Pythia-1B's "on-trend" test loss as the only evidence that quality is preserved. This is insufficient. A rigorous follow-up would train a family of models at fixed parameter count (e.g., 1.3B or 2.7B) with varying $a$ and correspondingly adjusted $h$ to hold total parameters constant, creating a sweep from many narrow heads (large $a$, small $h/a$) to few wide heads (small $a$, large $h/a$, following the paper's alignment rules). For each configuration, measure both training throughput and a battery of eval metrics (perplexity, MMLU, HellaSwag, code generation, reasoning benchmarks). The output is a quality-shape Pareto frontier showing exactly how much throughput gain is available before measurable quality degradation sets in, and whether that threshold depends on model scale. The paper's case study suggests $a=20$ is safe for a 2.7B model, but provides no evidence for whether $a=16$, $a=12$, or even $a=8$ would cross a quality cliff. This experiment would directly address the paper's most impactful unanswered question for practitioners: how aggressively can I apply these rules before my model gets worse?

Cross-architecture generalization: do the rules hold for encoder-decoder and mixture-of-experts models? The paper explicitly restricts its analysis to decoder-only transformers and notes that "our analysis will largely not apply to encoder-decoder models" (Section III-C). A natural extension would reconstruct the GEMM decomposition for encoder-decoder architectures (T5, BART) and mixture-of-experts (MoE) transformers, deriving the analog of Table II for each architecture family. For encoder-decoder models, the cross-attention mechanism introduces GEMMs where key and value projections depend on encoder output sequence length rather than decoder sequence length, changing which dimensions control wave quantization. For MoE models, the expert routing introduces sparse GEMMs whose performance characteristics may differ from dense GEMMs, and the expert capacity parameter constrains effective batch sizes in ways that interact with the paper's Rule 2 ("make $b \cdot s$ as large as possible"). The key experiment: train small encoder-decoder and MoE models with both "standard" and "shape-optimized" configurations (applying the paper's methodology, not its literal rules), and measure whether comparable throughput gains are achievable. A null result for encoder-decoder models would be valuable—it would precisely bound the scope of the paper's framework and motivate new analysis for architectures where cross-attention GEMMs dominate.

Hardware-adaptive model compilation: can shape optimization be automated at deployment time? The paper's rules are hardware-specific (multiples of 64 for A100, multiples of 8 for V100), but the model architect typically must fix hyperparameters before knowing all deployment targets. A more ambitious follow-up would develop a model compiler that takes a trained model checkpoint and automatically adjusts its effective shape for a target GPU without retraining. For example: given a model trained with $a=32$ heads, the compiler could fuse pairs of heads post-hoc (averaging their weight matrices) to produce an equivalent model with $a=16$ heads and $2\times$ larger per-head dimension, satisfying Tensor Core alignment on A100 without any quality loss. Alternatively, the compiler could pad weight matrices with zeros to make dimensions divisible by 64, exploiting the fact that zero-padded GEMM dimensions are essentially free (the zeros contribute nothing to the output). The key experiment: apply this compiler to a suite of publicly available models (Pythia, OPT, GPT-Neo, Llama) and measure the throughput improvement on target hardware against the paper's predicted gains. A successful compiler would decouple architecture design from hardware deployment, allowing architects to choose hyperparameters for quality and let the compiler handle hardware alignment.

Stress-test: do the rules survive future GPU architectures? The paper's rules are calibrated to NVIDIA V100/A100/H100 architectures. As new GPU generations introduce different Tensor Core tile sizes, SM counts, or memory hierarchies, the specific numerical constraints (multiple of 64, 108 SMs) will change. A forward-looking experiment would apply the paper's methodology—GEMM microbenchmarks, transformer decomposition, rule derivation—to a next-generation GPU (e.g., NVIDIA B100 or AMD MI300X) and measure whether the structure of the rules (power-of-two alignment of $h/a$, $b \cdot s$ as large as possible, vocabulary padding) persists even as the thresholds change. A finding that the rules are fundamentally tied to specific tile sizes—and that entirely different phenomena dominate on a new architecture—would refine the paper's co-design thesis: perhaps architecture design should target a distribution of expected hardware rather than a specific GPU, requiring robustness to different alignment constraints. This experiment would also provide an early signal to the community about whether models optimized for H100 will need architectural changes to run efficiently on the next generation.

---

Practical Applications and Downstream Use Cases

Batch inference cost reduction for API providers. The paper's rules are directly applicable to organizations running large-scale batch inference pipelines—API providers serving completions at scale, companies generating synthetic training data, or research labs evaluating models on benchmark suites. For such workloads, throughput is the dominant concern (latency per individual request matters less than total tokens processed per dollar). The paper's recommendation to pad vocabulary size to the nearest multiple of 64 (e.g., 50,257 → 50,304) provides a near-zero-cost improvement to the embedding and logit layer GEMMs that are on the critical path for every token generated. The paper's demonstration that the 2.7B architecture can be improved by 18% purely through head count adjustment generalizes: any model serving provider can audit their deployed architectures against the paper's checklist and apply compatible fixes (padding dimensions, adjusting $h/a$) without retraining, using weight padding or head fusion where applicable. At the scale of millions of inference calls per day, an 18% throughput improvement translates to an ~15% reduction in GPU-hours and therefore cost.

Architecture design for cost-efficient pretraining from scratch. For organizations planning to pretrain a new LLM, the paper's framework should be applied at architecture design time, before a single GPU-hour is spent. The procedure: (1) choose the target total parameter budget $P$; (2) select a candidate $(h, L)$ pair using $P \approx 12h^2L$; (3) adjust $a$ so that $h/a$ is divisible by 64 (on A100/H100) or 8 (on V100); (4) ensure $h/t$ is divisible by the same power-of-two threshold for the planned tensor parallel degree $t$; (5) if using SwiGLU, brute-force search for an intermediate MLP dimension near $8h/3$ that satisfies the alignment constraint; (6) pad $v$ to the nearest multiple of 64. The paper's claim that the 2.7B GPT-3 architecture can be trained "almost 20% faster" through these steps means that a $1M pretraining run could have been a$1.2M pretraining run with the same outcome simply by choosing different hyperparameters—or, equivalently, that the $1M budget could have trained a model ~20% larger (more layers or wider) for the same wall-clock time.

Hardware procurement and architecture co-planning. The paper's Section VII-A case study on 6-GPU nodes reveals a hidden coupling between hardware topology and model architecture that procurement teams rarely consider. Organizations planning to acquire GPU clusters for LLM training should evaluate candidate node configurations (4-GPU, 6-GPU, 8-GPU) not just in terms of raw FLOPs and interconnect bandwidth, but in terms of whether their intended model architectures can be efficiently mapped to those nodes. For example, a 6-GPU node requires $h$ to be divisible by 192 (least common multiple of 6, for tensor parallelism, and 64, for A100 Tensor Cores), which may constrain architecture choices for small-to-medium models in ways that an 8-GPU node (requiring $h$ divisible by 64, since $h/8$ automatically satisfies alignment if $h$ is a multiple of 64) does not. The paper's framework allows procurement decisions and architecture decisions to be evaluated jointly before either is finalized, potentially avoiding situations like Summit's where models must be designed around an unusual GPU count.

Post-deployment throughput auditing for open-source models. Organizations that fine-tune and deploy open-source models (LLaMA, Mistral, Pythia, OPT, GPT-Neo) inherit the original architects' hardware efficiency choices. The paper's framework provides a simple audit: for any downloaded model, check whether $h/a$, $h/t$, and $v$ satisfy the alignment constraints for your deployment GPU. If not, the model is leaving throughput on the table, and the paper's Section VII-B analysis suggests simple mitigations—padding the vocabulary embedding, adjusting tensor parallelism to change $h/t$, or in some cases fusing attention heads—that can be applied to the checkpoint without retraining. The paper's finding that Pythia-1B is substantially more efficient than Pythia-410M (Figure 13) despite being larger demonstrates that model size alone is a poor predictor of inference cost; a shape audit provides a complementary metric. This is immediately actionable for any organization serving fine-tuned open-weight models in production.