Designing Models from the Hardware Up

Designing Models from the Hardware Up

Speakers: Simran Arora

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Duration: 1636.0s

Overview

In this insightful MLSys 2025 talk, Simran Arora presents a compelling argument for designing AI models with hardware considerations from the ground up, rather than optimizing them post-hoc. The core premise addresses a critical inefficiency in the current machine learning landscape: while novel model architectures are emerging to rival the ubiquitous Transformer, their theoretical efficiency often fails to translate into tangible wall-clock speed in practice. Arora highlights that this gap stems from two primary issues: architectural choices that are incompatible with underlying hardware capabilities and the lack of flexible, high-performance programming abstractions for rapid hardware-aware development.

Arora introduces two significant contributions from her work. First, Thunder Kittens, a novel framework designed to simplify the process of mapping AI algorithms to hardware, bridging the chasm between asymptotic theoretical efficiency and real-world wall-clock performance. Second, the Based architecture, a new model design that explicitly incorporates hardware efficiency, demonstrating how combining local and global linearized attention can expand the Pareto frontier for sub-quadratic models. By focusing on hardware-aware building blocks and optimizing for GPU features like tensor cores and memory access patterns, Based achieves substantial improvements in hardware utilization and working memory compared to prior efficient architectures.

This work is pivotal for the future of AI/ML systems. As models grow larger and context lengths extend, the quadratic scaling of Transformers becomes increasingly untenable. Arora's research offers a dual solution: a set of tools to empower researchers to build hardware-efficient models more easily, and a concrete example of such a model that outperforms existing "efficient" alternatives in practical settings. It underscores a philosophical shift in ML system design, emphasizing that true innovation requires a deep understanding and co-design of both algorithms and the hardware they run on, ultimately unlocking orders of magnitude higher throughput and pushing the boundaries of what's possible with large-scale AI.

Background

▶ Watch: Transformer inefficiencies and sub-quadratic compute research goals (0:00)

The current era of AI is characterized by a "Cambrian explosion" of new model architectures, many of which aim to address the fundamental inefficiencies of the Transformer. While the Transformer has been the workhorse of deep learning since its popularization around 2015, its core attention mechanism scales quadratically with sequence length in terms of compute and linearly in terms of memory. This poses significant challenges for tasks requiring long context understanding, such as information extraction from lengthy documents, long-document question answering, code generation, and complex mathematical reasoning. Processing a "long textbook" with a Transformer demands prohibitive computational and memory resources.

The research community has actively pursued alternatives that offer sub-quadratic compute and constant memory usage. However, as Arora points out, there's "no free lunch" in this quest. The challenge lies in balancing the reduction of compute and memory with the ability to retain sufficient information about the sequence. Successive works have pushed this frontier, but often at the cost of information retention, meaning they cannot simply match Transformer performance across the board due to fundamental information bottlenecks.

A critical observation highlighted in the talk is the disconnect between theoretical efficiency and wall-clock efficiency. Architectures designed to be asymptotically more efficient, such as Mamba or RWKV, often fail to translate into faster real-world execution for practical sequence lengths and batch sizes. The "crossover point" where these models become faster than Transformers typically occurs at sequence lengths "way too late" to be beneficial in many common scenarios. This gap arises from the complexities of mapping AI workloads to hardware.

Two main reasons contribute to the difficulty of hardware-aware AI. Firstly, some architectural choices are inherently incompatible with hardware preferences. For instance, an algorithm might not effectively utilize the fastest compute units (like tensor cores on an NVIDIA GPU) or might access memory in a sub-optimal, non-coalesced fashion, leading to bottlenecks. Secondly, the available tooling and programming abstractions are often insufficient. AI researchers typically prototype in high-level Python frameworks like PyTorch, TensorFlow, or JAX. While these offer convenience and flexibility for rapid iteration, they often lack the fine-grained control necessary to achieve peak performance. On the other end of the spectrum, low-level tools like raw CUDA or Cutlass (templates over CUDA) offer complete hardware control but are notoriously complex, requiring deep expertise in GPU architecture and parallel programming. This complexity significantly hinders rapid research and development for the typical AI practitioner. The speaker cites the example of Flash Attention, which took years to optimize for the NVIDIA A100 GPU, only to see its performance degrade significantly (47% utilization drop) on the subsequent H100, requiring another two years to recover. This illustrates the arduous and time-consuming nature of achieving hardware-optimized performance with current methods.

Key Findings

▶ Watch: Talk structure: tools for hardware, then new architectures (3:59)

The talk presents two primary, interconnected key findings that address the challenges of designing efficient AI models for modern hardware.

First, the introduction of Thunder Kittens, a framework for rapid, hardware-aware AI algorithm mapping. This framework is designed to provide a concise and opinionated set of programming abstractions that are familiar to AI researchers (Pythonic, inspired by NumPy/PyTorch) while offering enough control to achieve peak hardware performance. The key achievement of Thunder Kittens is its ability to bridge the gap between high-level development ease and low-level hardware efficiency. Evaluations show that Thunder Kittens achieves performance comparable to highly optimized, low-level C++/CUDA tools like Cutlass but with drastically fewer lines of code. Crucially, it delivers up to 15x speedup compared to existing Pythonic interfaces like Triton while maintaining a comparable amount of code. This significantly lowers the barrier for AI researchers to develop custom, high-performance kernels, directly addressing the tooling problem in hardware-aware AI. Moreover, Thunder Kittens has demonstrated remarkable portability, with new hardware ports (e.g., Apple chips, Rivos, NVIDIA Blackwell) being completed in days or weeks by relatively junior engineers, showcasing its robust and adaptable design. Its catalytic impact is evident, inspiring new Pythonic frameworks like Cute 4, Cutlass 4, and Tiling.

Second, the development of the Based architecture, a novel model designed explicitly with hardware efficiency in mind, using the principles and insights gained from building tools like Thunder Kittens. Based aims to expand the Pareto frontier for efficient, sub-quadratic models by combining a local exact attention mechanism with a hardware-efficient global linearized attention. This hybrid approach capitalizes on the observation that nearby tokens require frequent, precise interaction, while far-apart tokens can rely on cheaper approximations. A key innovation in Based is its use of a Taylor approximation of the exp function in the softmax, which allows for reordering matrix multiplications to achieve constant memory and sub-quadratic compute. Unlike prior efficient architectures such as Mamba, Based is designed to fully leverage GPU tensor cores, which are significantly faster than other compute units (e.g., 16 times faster on H100). This hardware-aware design, coupled with optimized data movement and efficient state management using advanced H100 instructions, results in substantial performance gains. Specifically, Based achieves 5x higher hardware utilization and can use 8x the working memory or state size compared to prior architectures like Mamba, when implemented with these optimizations. This demonstrates that by co-designing models with hardware in mind, it is possible to achieve superior wall-clock efficiency, directly addressing the "theoretical vs. wall-clock efficiency" problem.

Technical Deep Dive

▶ Watch: Challenges in mapping AI workloads to hardware (4:00)

To understand the core innovations, it's essential to first grasp the underlying hardware architecture, specifically the NVIDIA H100 GPU. The H100 is composed of numerous streaming multiprocessors (SMs), which are the core compute units. Inside and close to these SMs are various memory hierarchies: tiny, extremely fast register memory; larger but slower L2 cache; and even larger, much slower High Bandwidth Memory (HBM). The SMs themselves contain specialized circuits for mathematical operations, load/store operations, and crucially, tensor cores for highly efficient matrix multiplications. The "golden rule" for AI workloads on these GPUs is to keep the tensor cores fully utilized, as they are dramatically faster (e.g., 16x faster than the next fastest compute on H100). The challenge lies in managing all other operations – nonlinearities, normalizations, and data movement – to minimize latency and ensure continuous tensor core engagement.

Thunder Kittens: Abstractions for Hardware-Aware AI

Thunder Kittens introduces a set of programming abstractions that mirror the GPU's hierarchical structure, aiming for both familiarity (Pythonic, NumPy/PyTorch-like) and control.

1. Thread Level (Warp):

• GPU execution involves tens of thousands of threads working in lock-step groups called warps (typically 32 threads).
• The challenge here is efficiently mapping logical data elements to physical threads and managing memory layouts – how data is arranged in memory for optimal access.
• Thunder Kittens' basic data abstraction is a 16x16 tile, analogous to a small tensor in PyTorch. Memory layouts for these tiles are automatically configured, relieving the user from this complex, error-prone task.
• Users interact with these tiles using familiar operations like X, Qumsum, and Matmols, inspired by high-level frameworks.

1. Thread Block Level (Occupancy):

• Work is scheduled onto SMs in units called thread blocks, which are groups of warps. Threads within a block can communicate very quickly.
• The goal is to maximize occupancy, meaning keeping different execution units within the SM (tensor cores, math units, load/store units) busy simultaneously. This requires co-scheduling multiple threads to handle different parts of the workload.
• However, increasing threads leads to contention for limited fast memory resources, creating a fundamental occupancy tradeoff.
• Thunder Kittens addresses this with a single unified kernel template called LCSF. This template predefines how different "workers" (groups of threads) specialize in various parts of the AI workload. Users simply fill in boilerplate functions (akin to a forward or init function in PyTorch) and can navigate the occupancy tradeoff space by toggling a few numerical parameters in the template, gaining fine-grained control without deep CUDA knowledge.

1. Grid Level (Launch Costs):

• The entire GPU (grid) consists of many SMs (e.g., 148 on B200s), each running multiple thread blocks.
• Launching each thread block incurs setup and teardown overheads (block launch costs), during which no useful AI work is performed.
• Thunder Kittens' template is designed to hide these costs from the user by design, abstracting away this low-level optimization.

Based Architecture: Hardware-Aware Model Design

The Based architecture is a direct application of the "hardware-up" philosophy, designed to overcome the wall-clock efficiency gap of prior sub-quadratic models.

1. Hybrid Attention Mechanism:

• Based combines two forms of attention: a local exact attention for nearby tokens (which are critical for fine-grained interactions and can be made very hardware-efficient due to their small scope) and a global linearized attention for far-apart tokens.
• The rationale is that distant interactions are less frequent and can tolerate cheaper approximations, reducing the quadratic scaling.

1. Linearized Attention via Taylor Approximation:

• The core of Transformer attention involves softmax(QK^T)V. The exp function within the softmax is problematic for linearization.
• Based uses a Taylor approximation of the exp function, transforming exp(Q K^T) into phi(Q) phi(K)^T, where phi is a polynomial approximation of exp. This allows for a more flexible reordering of matrix multiplications.
• Crucially, this linearization allows the computation to be reordered from (Q K^T) V (which is quadratic in sequence length) to Q (K^T V). The K^T V product can be computed first, resulting in constant memory usage and sub-quadratic compute.

1. Hardware Efficiency Optimizations:

• Tensor Core Utilization: Unlike prior architectures like Mamba, Based is explicitly designed to use tensor cores, leveraging their massive speedup for matrix multiplications.
• Optimized Data Movement: It efficiently loads Q (queries) and K (keys) from slower HBM to faster on-chip memory. Large intermediate computations, such as the Q outer product Q or K outer product K (required for certain Taylor terms), are performed entirely within fast GPU memory, minimizing costly data movement.
• Advanced Instruction Leveraging: The specific design of Based and its state management heavily utilizes advanced instructions available on the H100 GPU, enabling it to handle a very large working memory and state size extremely quickly. This includes fine-tuned control over memory access patterns and synchronization, which is exposed through the control mechanisms in Thunder Kittens.

These technical details underscore how Thunder Kittens provides the necessary granular control, while Based demonstrates the architectural principles for building models that are not just theoretically efficient but also optimized for the specific characteristics and capabilities of modern AI accelerators.

Experimental Setup & Results

▶ Watch: Current tooling landscape for hardware-aware AI code (5:59)

The evaluation of both Thunder Kittens and the Based architecture involved rigorous benchmarking against state-of-the-art alternatives across diverse hardware.

Thunder Kittens Evaluation:

The Thunder Kittens framework was evaluated by comparing its performance and expressiveness against two distinct classes of kernel development tools:

1. Low-level tools: These include raw CUDA and Cutlass, NVIDIA's highly optimized C++ template library for GPU programming. These tools offer maximum control but are known for their complexity.
2. Pythonic interfaces: Specifically, Triton from OpenAI, which aims to provide a more accessible, Python-like environment for writing GPU kernels. While simpler, Triton often sacrifices some aspects required for peak performance, such as fine-grained control over register memory or synchronization patterns.

The evaluation covered a breadth of AI workloads (not specified in detail, but implied to be common operations found in neural networks) and a breadth of hardware, including:

• NVIDIA H100 GPU
• NVIDIA RTX 4090 GPU
• Apple Silicon chips (e.g., M-series processors)

Metrics:

• Lines of Code (LoC): A proxy for development complexity and conciseness.
• Speed: Wall-clock execution time for a range of different input sizes to the kernel.

Results:

• Performance vs. Low-Level Tools: Thunder Kittens achieved comparable performance to highly optimized Cutlass and CUDA kernels. This is a significant finding, as it means researchers can achieve near-optimal hardware utilization without the steep learning curve and extensive development time associated with low-level programming.
• Expressiveness vs. Low-Level Tools: This comparable performance was achieved with drastically lower lines of code compared to Cutlass and CUDA, indicating a much simpler and more concise development experience.
• Performance vs. Pythonic Interfaces: Compared to Triton, Thunder Kittens maintained a comparable amount of code (suggesting similar ease of use) but delivered an impressive up to 15x speedup. This highlights Thunder Kittens' ability to provide both simplicity and peak performance, a critical advantage for rapid research.
• Portability: The framework demonstrated exceptional portability. Porting Thunder Kittens to Apple Silicon was accomplished by an undergraduate in just 2-3 weeks. The team in Arora's lab ported it to NVIDIA's new Blackwell architecture in a mere two days, achieving kernel performance on par with Cutlass's development rate. A port to AMD hardware is also underway.

Based Architecture Evaluation:

The Based architecture was evaluated for its efficiency and performance characteristics, particularly against prior "efficient" models.

Comparison:

• Triton-like implementations: This refers to implementations that might use more generic Pythonic kernel generation, which often cannot leverage advanced hardware features.
• Mamba: A prominent prior architecture known for its sub-quadratic scaling properties.

Hardware:

• NVIDIA H100 GPU

Metrics:

• Hardware utilization: A measure of how effectively the GPU's compute units (especially tensor cores) are being used.
• Working memory/state size: The amount of memory the model can effectively use for its internal state during processing.

Results:

• Hardware Utilization: Compared to Triton-like implementations, Based achieved drastically higher hardware utilization on H100s. This is attributed to Based's ability to leverage advanced H100 features and tensor cores, which generic implementations often cannot access.
• Memory and Utilization vs. Mamba: Based demonstrated a remarkable capability to use eight times the working memory or state size of Mamba, while simultaneously achieving five times the hardware utilization. This is a powerful testament to the "hardware-up" design philosophy, showing that a larger, richer internal state can be managed efficiently on hardware, breaking previous tradeoffs.
• Scaling and Quality: Through distillation techniques with labmates, a Transformer LLaMA 405 billion parameter model was successfully distilled into a Based model. This distilled Based model achieved "pretty interesting quality across tasks people care about" while maintaining "relatively low memory," indicating its potential for developing high-quality, efficient large language models.
• Community Impact: Since its release in 2023, the Based architecture has influenced subsequent developments in the field. Notably, Mamba B2 and RWKV V5/V6 have adopted some of the linear attention and large working memory principles demonstrated in Based, highlighting its impact on the design of new efficient models.

These experimental results provide strong evidence for the efficacy of both Thunder Kittens as a development tool and Based as a hardware-aware model architecture, validating the core tenets of the talk.

Practical Implications

▶ Watch: H100 GPU architecture and memory hierarchy overview (7:00)

The work presented by Simran Arora has profound practical implications for various stakeholders in the AI/ML ecosystem, from individual model builders to large infrastructure teams.

For practitioners and model builders, this research offers a pathway to break free from the performance limitations of the Transformer architecture without sacrificing wall-clock speed. Historically, exploring novel architectures beyond Transformers often meant accepting a performance penalty due to the lack of optimized kernels. Thunder Kittens directly addresses this by providing a high-level, Pythonic framework that empowers researchers to rapidly develop custom, high-performance kernels for their experimental models. This means new ideas can be prototyped and benchmarked against real-world hardware performance much faster, accelerating the pace of innovation for models that require sub-quadratic compute or constant memory. The Based architecture serves as a concrete example, demonstrating that it's possible to build models that are not only theoretically efficient but also achieve superior hardware utilization and larger working memory, translating directly into faster inference and training for long-context tasks. This empowers model builders to design more capable and context-aware AI systems.

For infrastructure teams and model deployers, the ability to quickly port and optimize kernels for new hardware generations is a game-changer. The example of Flash Attention taking years to optimize for A100 and then H100 highlights a significant bottleneck in the deployment pipeline. Thunder Kittens' demonstrated portability (e.g., Blackwell port in two days, Apple port in weeks) means that infra teams can adapt to new hardware much more rapidly, ensuring that their ML workloads always run on the most efficient available silicon. This reduces the total cost of ownership for AI infrastructure by maximizing hardware utilization and minimizing the engineering effort required for optimization. It also facilitates the adoption of diverse, specialized hardware (like Apple Silicon or custom accelerators) by making kernel development more accessible.

However, these innovations also come with tradeoffs and limitations. While Thunder Kittens simplifies hardware-aware development, it doesn't entirely abstract away the need for understanding hardware concepts. Users still need to make informed decisions by toggling parameters in the LCSF template, requiring some intuition about occupancy, memory hierarchies, and compute unit utilization. The "no free lunch" principle still holds: there's an inherent balance between reducing compute/memory and retaining information. The challenge, as highlighted by Based, is to strategically push this Pareto frontier by intelligently combining different architectural components and leveraging hardware features. Purely asymptotic-efficient models might still struggle to achieve competitive wall-clock speeds if their design doesn't align with hardware capabilities, emphasizing that the "crossover point" remains a critical consideration. The distillation of large Transformer models into Based models suggests a potential pathway for leveraging existing large models' knowledge while achieving the efficiency benefits of hardware-aware designs, but the fidelity of this distillation and its generalizability across all tasks would be an ongoing area of research.

In summary, this work provides both the tools and the blueprints for building the next generation of efficient AI models, promising faster research cycles, more performant deployments, and a more hardware-conscious approach to AI development.

Key Takeaways

• Theoretical vs. Wall-Clock Efficiency Gap: Novel AI architectures designed for asymptotic efficiency (sub-quadratic compute, constant memory) often fail to deliver practical wall-clock speed due to hardware incompatibility and a lack of flexible, high-performance programming tools.
• Thunder Kittens as a Solution: This framework offers a concise, Pythonic set of programming abstractions for hardware-aware kernel development. It achieves performance comparable to low-level CUDA/Cutlass with drastically less code and up to 15x speedup over Pythonic interfaces like Triton, making hardware optimization accessible to AI researchers.
• Hardware-Aware Model Design (Based): The Based architecture exemplifies how co-designing models with hardware in mind (combining local and global linearized attention, leveraging tensor cores, optimizing data movement) can achieve superior wall-clock efficiency, including 5x higher hardware utilization and 8x larger working memory compared to prior efficient models like Mamba.
• Portability and Rapid Adaptation: Thunder Kittens demonstrates exceptional portability, enabling new hardware ports (e.g., Apple, Blackwell) in days or weeks, significantly reducing the optimization bottleneck for new hardware generations.
• Catalytic Industry Impact: The philosophical approach and practical success of Thunder Kittens have catalyzed broader industry efforts towards developing more accessible, Pythonic frameworks for hardware-aware ML development (e.g., Cute 4, Cutlass 4, Tiling).
• The "Hardware-Up" Imperative: The talk emphasizes that designing models "from the hardware up" rather than attempting post-hoc optimization is crucial for truly expanding the Pareto frontier of efficient and capable AI systems, especially for long-context and large-scale applications.

About the Speaker(s)

Simran Arora is a researcher working at the intersection of AI and systems, with a focus on addressing the inefficiencies of current machine learning models and bridging the gap between theoretical algorithmic advancements and practical hardware performance. Her work, often conducted with labmates, centers on developing tools and architectures that facilitate hardware-aware AI model design. The mention of this talk originating from her "job talk" suggests a strong academic and research background, likely within a university or research institution setting, where she leads or contributes to projects exploring novel approaches to optimizing AI workloads for modern computing hardware.