Thank you for your positive review of our work! We are glad that you find our approach "fresh" and "valuable" and results "impressive". We carefully address your outstanding questions in our response.

Addressing weakness 1. TK’s flexibility for complex AI workloads.

We focus on a research question: What are the tradeoffs between programming complexity and the accessibility of peak hardware performance? We view the fact that TK is uses minimal primitives, yet provides competitive or higher performance as prior kernels, as a strength rather than a limitation. We provide the summary table below (and in Appendix B.3), showing that:

1. TK competes with or outperforms the baselines across settings.
2. TK uses a similar amount of lines of code per kernel as Triton, and fewer lines than the other CUDA reference kernels

Table: Lines of code across TK H100 kernels, state of the art non TK kernels, and the TK speed up over the reference kernels.

Overall, our core abstractions handle the optimization patterns needed for most AI kernels (as demonstrated in our case studies), TK is explicitly designed to be extensible when specialized optimizations are needed. Users can seamlessly mix TK's high-level primitives with custom CUDA code - as we show in our Mamba-2 implementation. This hybrid approach means users get the benefits of TK's abstractions for the bulk of their kernel (typically 95% of the effort), while retaining the flexibility to implement specialized operations where needed.

Philosophically, TK focuses on providing well-chosen defaults that compose well with both each other and custom code. This design ensures TK remains practical and maintainable (ThunderKittens’ source is just a few thousand lines of code) while supporting the full spectrum of AI workloads.

Addressing weakness 2. TK’s portability across hardware platforms.

To address this concern, we provide 2 new kernels for NVIDIA 4090 GPUs and 3 new kernels for Apple M2 chips, to demonstrate that the TK framework ports across hardware platforms. We have added results and discussion for these kernels in Appendix B.

Table: 4090 Attention FWD Performance (non-causal, batch=16, heads=16)

Table: Apple M2 Standard Attention FWD vs Apple MLX (non-causal, batch=16, heads=16)

Table: Apple M2 GEMM TK vs Apple MLX Performance

We originally highlighted NVIDIA H100s since they represent the trend of AI hardware growing more complex, with an increasing number of specialized instructions, levels of parallel execution, and opportunities for asynchronous execution.

Addressing weakness 5. Highlighting TK’s value to the ML audience

While TK advances systems capabilities, its impact lies in enabling the next wave of ML innovation. As recognized in the ICLR 2025 call for papers, implementation efficiency and hardware optimization are no longer secondary concerns - they are fundamental bottlenecks in advancing AI research.

Our work offers three key contributions that directly address these bottlenecks:

1. Democratizing Hardware-Efficient ML Research. The complexity of hardware optimization has become a critical barrier to ML architecture innovation. TK dramatically lowers this barrier, enabling researchers to rapidly prototype and evaluate novel neural architectures that would previously have required months of specialized CUDA engineering. For instance, our framework makes tensor core optimization accessible through simple interfaces (stated on lines 88-91, 270-280 of our submission).

2. Making New ML Architectures Practical. History shows that major ML advances come from scaling - but scaling requires efficiency. Our framework has already revealed that architectures previously dismissed as impractical can be transformed into state-of-the-art approaches through proper optimization. For example, our implementation shows that linear attention can outperform standard attention in wall-clock time at much shorter sequences than previously thought possible, fundamentally changing the calculus of architecture design decisions. This is just one example of how TK can unlock new families of architectures that were previously considered computationally infeasible.

3. Quantitative Impact on ML Workloads. We demonstrate substantial improvements across core ML operations:

• $10-40\%$ faster attention computation
• $8\times$ acceleration for long-range convolutions
• $6\times-14\times$ speedup for linear attention variants These improvements directly enable research into larger contexts, alternative attention mechanisms, and novel architecture designs.

In our view, these optimizations represent the difference between architectures that can scale and those that cannot. Just as FlashAttention's impact went far beyond its speedup numbers to enable the current wave of large language models, TK aims to unlock the next generation of ML architectures by removing critical implementation barriers that currently constrain innovation.

Addressing question 1. Bank conflicts in FlashAttention-3.

CUTLASS and TK support the exact same capabilities, which is stated on lines 200-202 of our submission. We did not state that CUTLASS faces fundamental challenges causing bank conflicts. We simply observe that bank conflicts are avoidable and the user does not need to be burdened in thinking about layouts, as stated on lines 206-208 of our submission. TK handles the layout selection for the user to help minimize such issues.

Addressing question 2. CUTLASS ping pong and cooperative kernels.

ThunderKittens supports ping-pong kernels like CUTLASS -- our work has not emphasized ping-pong style kernels, as we have found that it adds considerable complexity for marginal performance improvements. Our GEMM kernel is also cooperative in nature; we also have threads exchange data through shared memory before performing coalesced stores to global memory. We have added comparison to CUTLASS GEMM in Figure 7, and also provide an example ThunderKittens ping-pong GEMM kernel in the appendix which makes further use of the asynchronous tensor core instructions. We find it achieves an additional 0.4% performance above the kernel proposed in the main paper.

Addressing question 3. Extensibility of TK’s ideas across hardware platforms

The review mentions that our original set of kernels are for H100 GPUs. We provide 2 new kernels for NVIDIA 4090 GPUs and 3 new kernels for Apple M2 chips, to demonstrate that the TK framework ports across hardware platforms. We have added results and discussion for these kernels in Appendix B3.

Addressing question 4. Extensibility to future GPU generations

Thank you for this thoughtful question! While it is difficult to exactly know what changes will be required, since we do not have Blackwell GPUs, we are optimistic that TK’s approach will extend. Our primitives are designed around fundamental hardware properties – e.g., for memory, shared memory is banked, we need coalesced HBM loads, register utilization needs to be carefully managed; e.g., for compute, hardware has multiple execution units that benefit from occupancy, latencies need to be hidden via asynchronous execution. These are not platform specific properties; these fundamentals have already cleanly transferred across hardware platforms and vendors (NVIDIA 4090; NVIDIA H100; Apple M2). We expect the trend to continue.

Thank you for your thorough review of our work! We are grateful for your detailed feedback, which significantly helped us improve our work. We provide our response to your feedback below.

Addressing weakness 1: Simplicity of programming with TK, and Addressing weakness 3: Lines of code comparisons.

The review mentions that TK is too low-level, which may not simplify the programming burden. Our results and analysis suggest that TK is quite helpful. We provide the summary table below (and in Appendix B.3), showing that:

1. TK competes with or outperforms the baselines across settings.
2. TK uses a similar amount of lines of code per kernel as Triton, and fewer lines than the other CUDA reference kernels

Table: Lines of code across TK H100 kernels, state of the art non TK kernels, and the TK speed up over the reference kernels

Additionally, we reiterate that TK manages the following for the user:

1. Warp-level: We automatically manage memory layouts for the user. We provide optimized library functions inspired by PyTorch (mma, exp, cumsum) to uplevel the programming experience. We provide tile data structures for both shared and register memory (and we manage data type conversions, broadcasts, templating functions around the tile sizes, loads and stores across the memory hierarchy) to help users manage their memory resource utilization.
2. Block-level: We provide a general kernel template that helps users utilize two important asynchronous execution patterns that are helpful in AI workloads: (1) multi-stage pipeline and (2) producer consumer worker specialization. We set these patterns up for the user. We think that giving users the option to set barriers is important for enabling peak performance. Triton does not offer this, and we show in our paper that asynchrony is important. Our goal is to simplify the primitives, without sacrificing performance.
3. Grid-level: We provide persistent grid support and a function to help users set the thread block launch pattern, which influences L2 reuse.

Addressing weaknesses 2: Novelty of TK compared to prior frameworks

We focus on a research question: What are the tradeoffs between programming complexity and the accessibility of peak hardware performance? Our novelty is in showing a new point on the tradeoff space exists. We can use fewer primitives to achieve higher performance on AI workloads, compared to what’s been demonstrated in prior frameworks. Identifying the primitives is non-trivial given the complexity of GPUs, and this is core to our contribution.

To address your feedback, we have included new analysis in Appendix B.3, we believe it has helped the paper and thank you for your suggestion. Please also see this new content below and in our common response. We compare the features supported by different frameworks:

We recognize that the following is an imperfect metric. However, we include the sizes of various CUDA libraries below as a rough proxy. For CUTLASS and TK we report the size of the “include/” directory, and for Triton we report the combined size of the “include/” directories in Triton plus the “include/” in the core MLIR compiler dependency.

Addressing weakness 4. Additional kernels highlighting the scope of TK's features

The review notes that our experiments miss some useful kernels. We provide new results for the requested kernels: quantized FP8 GEMM, causal attention. Please find these results in Appendix B2 and in the common response. We hope that these results address any remaining concerns on experimental validation.

Thank you for your review, we are very grateful that you took the time to provide the thoughtful detailed feedback, which has helped us significantly strengthen our work. We carefully address your feedback in our response and hope that it helps.

Demonstrating the broad scope of ThunderKittens

Thank you for the question! We provide new results beyond the submission for FP8 GEMM, causal attention, NVIDIA 4090 kernels, and Apple M2 kernels to further emphasize the scope of our framework. We also show that TK supports tile dimensions that are not multiples of 16; we provide demo attention kernels that use this feature. Please find these results in the common response and in Appendix B of our revision. We also added new discussion of our implementations for these features in Appendix D.1.

Connections to the computer architecture literature

Thank you for raising this concern! We have contextualized both the cost model in Section 2.2 and the results in Table 3 to highlight the existing literature. We did not aim for Section 2.2 to appear as our contribution and appreciate the suggestions to further establish that this is fundamental background knowledge. We correspondingly added citations to computer architecture works in this section. Please let us know if any other changes are helpful!

Extending to Fully Shared Data Parallel (FSDP)

We appreciate this suggestion. While multi-GPU operations like FSDP are indeed important for large-scale training, our current focus is on optimizing single-GPU kernel performance. Users of our framework can currently handle multi-GPU communication through established libraries like NCCL, which provides highly optimized primitives for inter-GPU collective operations. We agree that extending the producer-consumer model to multi-GPU contexts is an interesting direction. We are excited to explore such directions in future work.

Presentation improvements

We have made several presentation improvements to improve readability. We apologize for the presentation errors. Additionally:

• We removed the mention of external endorsements from the paper, per your recommendation.

• The review notes that Figures 2 and 5 are too simplified compared to the full implementation. We have updated the paper to include more precise captions, to highlight that the figures represent specific aspects of the library: Figure 2 is only meant to show that the library functions and tiles are Pytorch-like. Figure 5 is only meant to improve intuition on separation of responsibility across workers.

Thank you for your feedback on presentation, which has helped improve the paper.

Technical details

We have updated the paper to reflect your feedback, and appreciate your help in tightening the technical writing. We also provide specific responses below:

• With respect to the TC:ALU ratio: We were specifically referring to the BF16 tensor core to ALU compute ratio, while your citation refers to FP16 FMA performance. Nonetheless, there was a slight error in our writing for the H100; we've adjusted the number to match NVIDIA's official 15x figure. (The A100 GPU officially has 16x.)
• Regarding kernel memory access: Kernels can operate without memory access. We’ve updated the wording to "typically loads" for accuracy, as it describes the typical pattern for deep learning kernels, which is our focus.
• On warp execution units: We've revised this section to be more precise. You're correct that individual warps cannot simultaneously occupy different functional units. We've clarified that we were referring to instruction-level parallelism (e.g., overlapping memory loads with computation) and concurrent execution across different warps.
• Concerning register spills: We've updated this to be more precise. Register spills occur to local memory, which may be cached in L1 depending on hardware policies. Thank you for helping us make this clearer.
• On SM thread block scheduling: We've added explicit mention of resource constraints as a factor in thread block collocation.
• As to kernel tail effects: While concurrent kernel execution is possible across multiple asynchronous CUDA streams, our discussion focuses on the common case in deep learning where kernels execute sequentially within a single CUDA stream, and for these workloads tail effects are an important consideration. We've added a note acknowledging advanced multi-stream techniques while maintaining our emphasis on typical AI workloads. We are excited to explore overlapping kernel execution in future work, such as in the recently-released async TP workload.
• We've also added a brief note about instruction cache considerations, though this wasn't a significant factor in our evaluated kernels.

Please let us know if anything else would be helpful in your review!

Technical Details

Thank you for addressing these corrections. However, a few issues remain:

5. Regarding the TC:ALU ratio, where are you finding this number? According to the H100 datasheet, NVIDIA reports a 7.4x increase, not 16x. This is the case for both FP16 and BF16, as they utilize the same hardware units, likely to support TF32. Are you perhaps considering the FLOPs with sparse matmul?

6. The discussion of register spilling to L1 still seems imprecise. It would be more accurate to describe this as “spills into global memory HBM.” The subsequent sentence effectively explains how SMEM and L1 can be traded to provide additional L1 capacity for these spilled registers, where I believe the local memory L1 policy is write-back but evictions can occur during global memory load / stores. However, it is worth implying that, in the most general case, spilling registers may incur a full hierarchy cache miss, potentially causing significant performance degradation. As currently written, this performance issue may be understated.

7. There are a few minor technical inaccuracies in Section 2.1 that I’d like to highlight:

   • The ALU operations of min/max and FMA are computed on the same units, but the text seems to suggest they are distinct. While speculative, these operations likely share significant overlap within the datapaths of the ALU/CUDA Core.
   • Could you confirm whether the statement “throughput differs across pipelines” is accurate? Or is the observed throughput difference due to resource contention rather than variations in individual unit throughput? My understanding is that all compute units, including the XFU, have a consistent throughput of 1 op/clk.
   • Regarding thread stalling, the term “fixed instruction latencies” may not be sufficiently precise. I suggest replacing it with “resource contention,” which provides a broader and more accurate explanation.

Additional Comments

• Regarding the instruction cache, this is particularly relevant to your discussion on reducing the number of registers per thread. An increase in operations could lead to more instructions, potentially resulting in I$ misses.
• Thank you for your response regarding kernel tail effects. I had not considered the limitation of sequential kernel launches, and this would indeed be an excellent direction for follow-up TK!

Thank you for your positive and thoughtful review! We are glad that you appreciate TK’s simplicity and performance, including providing kernels that far outperform the prior state-of-the-art kernels by up to 14x. Here, we address your remaining questions.

Simplicity of tuning kernels within TK

In our work, we de-emphasize autotuning, as we find that tuning kernels is usually straightforward once the right primitives are in place. This is, in a sense, a limitation of our framework compared to compiler-based frameworks like Triton, which automatically block and tile. However, our core argument is that for most AI workloads, sophisticated ML compilers are not needed to concisely, performantly, and portably express kernels.

To the extent that tuning is necessary, it can usually be done by twiddling occupancy and pipeline depth, for which there are usually only a few values that make sense, anyways. TK makes it easy to twiddle these two values by changing a single number in the kernel, and recompiling.

For tile size optimizations, users may need to consider hardware constraints such as shared memory and register limitations, requiring more careful analysis. However, TK's templated operation design minimizes the code changes needed to implement these adjustments.

TK supports dynamic shapes beyond multiples of 16

Thank you for the insightful question with respect to the potential limitations of 16x16 tiles. Yes, TK does support workloads where the dimensions are not multiples of 16. To address your feedback, we include new results in Appendix B.2, where we report that TK gives high performance on attention kernels that require tiles that are not multiples of 16.

We also describe how we support these dimensions in Appendix D.1. The common solution to support any dimension is to use padding. ThunderKittens supports padding within the kernel, and other frameworks, such as CUTLASS and Triton also pad data into tiles. Padding is the solution the hardware prefers, too, since tensor cores generally operate in multiples of some base size (usually 16).

We hope these responses help address your feedback, and please let us know if you have any additional questions!

Performance Benchmarks

Please also find these results discussed in Appendix B of our revised submission.

TK extends across hardware platforms

To further highlight the extensibility of TK, we provide new kernels for the NVIDIA 4090 and Apple M2 platforms.

Table: 4090 Attention FWD (non-causal, batch=16, heads=16)

Table: Apple M2 Pro Attention FWD vs Apple MLX (non-causal, batch=16, heads=16)

Table: Apple M2 Pro GEMM TK vs Apple MLX Performance

Comparing TK to CUTLASS and Triton:

We focus on a research question: What are the tradeoffs between programming complexity and the accessibility of peak hardware performance? Our novelty is in showing a new point on the tradeoff space exists. We can use fewer primitives to achieve higher performance on AI workloads, compared to what’s been demonstrated in prior frameworks.

Overview: CUTLASS provides many layers of templated C++ primitives. Compiler approaches like Triton expose fewer optimizations, but simplify the programming experience. Researchers often turn to Triton, pointing to the relative simplicity of Triton. Our novelty is showing a small set of C++ primitives that captures most of what we need and and consistently unlocks state-of-the-art kernels across hardware platforms. Identifying this small set is not trivial given the complexity of GPUs.

Below, we summarize how Triton, CUTLASS, and TK support key features needed for high performance:

Table: Framework Feature Comparison

As proxies for the simplicity versus efficiency resulting from different frameworks, we measure: (1) the size of various frameworks and (2) the lines of code used across kernels. We validate two claims:

1. TK compares to Triton in lines of code, and both outperform CUTLASS
2. TK competes with or outperforms Triton and CUTLASS/raw CUDA baselines

Table: Library Size Comparison

Table: Lines of code across TK H100 kernels, state of the art non TK kernels, and the TK speed up over the reference kernel.