Skip to content
Bud's Wiki

IO-Aware Kernel Design

TL;DR

Section titled “TL;DR”

A recurring GPU-kernel design template, originating with FlashAttention, that treats HBM round-trips of intermediate tensors as the binding bottleneck — not arithmetic intensity. The recipe: identify an N × K intermediate that’s materialized in HBM only to be reduced along K; fuse the produce/consume pair into a streaming kernel with on-chip state; exploit an algebraic decomposition (softmax over blocks, argmax over partitions, segmented sums) that lets the reduction commute with tiling. The result is mathematically exact, not an approximation. Recent filed work generalizes the template off attention onto Lloyd’s k-means and LM-head categorical sampling.

Key claims

Section titled “Key claims”

Recent contributions

Section titled “Recent contributions”
  • [2026-08-27] Redwood: A Frontier AI Accelerator Designed, Verified, and Deployed from Scratch in 2 Weeks by AI: Redwood (Architect Labs) is the first filed instance where the hardware is co-designed with IO-aware kernel primitives from the ground up: the SIMD engine reuses existing FP resources to run FlashAttention-4’s emulated softmax (avoiding a dedicated softmax unit), each tile pairs a matrix engine that streams directly into the vector engine over wide banked scratchpad memory, and a CTM-to-CTM messaging fabric expresses double-buffering / prefetch / out-of-order compute as message flows rather than centralized arbitration. A concrete datapoint on how much of the IO-aware-kernel discipline transfers into silicon when compiler, kernel, and RTL are co-generated by a single AI-for-AI-research system.
  • [2026-08-20] Binh ports PhysX to AMD Radeon and wins the global AMD robotics hackathon: Hackathon-scale port of NVIDIA PhysX 5’s GPU rigid-body solver from CUDA to ROCm/AMD Radeon; W7900 vs RTX5090 gap on ManiSkill Franka PushCube (67,320 vs 110,705 steps/s at 4096 envs) tracks the 2× memory-bandwidth ratio between the two cards — a nice in-the-wild confirmation that the physics-sim step is memory-bound and porting is bottlenecked by bandwidth, not by algorithmic differences between CUDA and HIP.
  • [2026-08-18] Mojo🔥 is now open source!: Modular open-sources the Mojo compiler + tooling + standard library under Apache 2.0 + LLVM exception — a third-alternative substrate for GPU-kernel work alongside CuTeDSL-in-Python (FA4, Gram-NS, Attn-QAT) and the “agents write kernels directly” position (Retire the Abstractions — CUDA DSLs Heading Toward Retirement). Mojo is a standalone systems language, not a Python-embedded DSL; its four-year cadence (stdlib open 2024 → kernels open → tooling → compiler now) and its 1.0 with source-stability shipping the week prior turn it into a permissively-licensed test of the “compile-time as the bottleneck” open question on this page. Compiler-side contributions not yet accepted (planned end of 2026); stdlib contributions have been open since 2024.
  • [2026-08-17] Training Dyna-2 at million-hour scale, repeatably: Dyna Robotics’s 1M-hour training-infrastructure retrospective extends the “don’t materialize the intermediate” template from GPU-kernel IO up to the dataset-delivery layer: MCAP topic-group chunking (group topics that share a read pattern, write each group time-major) plus Alluxio page-granularity caching on node NVMe with consistent-hash ownership. A sample fetch is O(read-groups) instead of O(topics); page-caching means an NVMe covers far more distinct episodes and eviction drops a page not a whole file. Reports ~68% episode-size reduction, ~2.9× faster sample reads, ~2 GB/s per node from cluster-local cache, and 98% GPU utilization on warm multi-node runs. Sibling to KnapFormer’s compute-bags at the sequence-parallel layer — both treat per-sample shape as a first-class primitive rather than a fixed container property, applied at different points in the stack.
  • [2026-08-13] Dion3: Full-Stack Orthogonal Updates: Dion3 packages the Gram Newton-Schulz template (structured-matmul exploitation via symmetric-GEMM CuteDSL kernels, from Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon) together with a megabatching strategy that specifically targets the communication overhead of orthogonalization on sharded weights — extending the IO-aware kernel template with a distributed-communication axis beyond the single-device HBM-traffic focus of FA4/Flash-KMeans/FlashSampling. Claims up to 6× optimizer step time reduction, the strongest end-to-end wall-clock result yet for the structured-matmul branch of this concept.
  • [2026-08-06] Retire the Abstractions — CUDA DSLs Heading Toward Retirement: Hazy Research (creators of ThunderKittens) argue that CUDA DSLs are on the retirement list because coding agents now do the cognitive-offloading job C++ template abstractions used to do — reports building this year’s Cursor Mixture of Kittens MoE megakernel without an intermediate C++ abstraction layer, with the abstraction expressed informally in the prompt instead. Directly counters the CuTeDSL-in-Python trajectory established by FA4 / Gram Newton-Schulz / Attn-QAT on this page, and pairs with CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation‘s KernelBench evidence as the manifesto version of the same thesis.
  • [2026-07-28] Online KL Shampoo (OKLS) — Tilde Research's KL-optimal full-matrix AdaGrad approximation optimizer with 1.45× parameter efficiency over Muon: Tilde’s OKLS optimizer relies on Scaled CANS Coupled Newton–Schulz for zero-staleness inverse-square-root preconditioning: 10 iterations, 27 FP16 GEMMs per step, FP32 accumulation. Reports 98% of Muon’s training throughput despite recomputing preconditioners every step — a datapoint that per-step matrix-inverse-square-root is now a practical optimizer primitive alongside the Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon Gram-Newton–Schulz Muon kernel.
  • [2026-07-28] Sol-Attn: Accelerating Video Generation Inference via On-the-Fly Attention Sparsification: Sol-Attn (NVIDIA/HAN Lab) adds a third IO-aware design point to the sparse-attention family: never materialize the proxy map at all. FlashAttention’s contribution was softmax-over-blocks with running statistics; Sparse VideoGen2: Accelerate Video Generation with Sparse Attention via Semantic-Aware Permutation added tile-shape variability as a first-class concern; Sol-Attn folds routing (on-the-fly block thresholding against proxy scores) and approximate correction (proxy-score reuse for unselected blocks) into the same online-softmax pass. The user-facing knob is a threshold, not a fixed K or p — per-query budget becomes a data-dependent outcome rather than an imposed constraint. Reported 2.1× (generation) / 2.3× (editing) end-to-end speedups on B200 vs FlashAttention-3.
  • [2026-07-27] Kimi K3 Technical Report — Open Frontier Intelligence: Kimi K3 tech report — KDA training kernel redesigned to hold the MLA attention output in FP32 (to correct flash attention’s biased rounding error) by overlapping the FP32 output tile with KV staging buffers instead of the query tile, freeing shared memory for a deeper KV pipeline.
  • [2026-07-23] Lucy 2.5 Real-Time AI Video Editing: Decart's 40ms Live Stream Revolution: Lucy 2.5 (Decart, community writeup) — “deep kernel fusion” merging multi-operator diffusion graphs into single CUDA kernels is cited as a load-bearing latency lever for the sub-40 ms budget, alongside NVFP4 quantization and dynamic sparse attention.
  • [2026-07-22] Sol Video Inference Engine: Agent-Native Full-Stack Acceleration Framework for Efficient Video Generation: Sol Video Inference Engine treats kernel fusion as a late-stage local-tuning problem on the partially accelerated stack, driven by the current bottlenecks (which shift after cache/sparsity/pruning are applied). The kernel-fusion agent profiles per-kernel time, decides custom-vs-compiler-generated, and picks operator sequences (GEMM+GELU, GEMM+residual, RoPE+norm) to match the current tensor shapes and precision formats. New design point: kernel fusion scheduling as part of an automated full-stack acceleration pipeline rather than as hand-coded epilogues.
  • [2026-07-16] Kimi K3 — Open Frontier Intelligence (2.8T MoE with KDA + AttnRes): Kimi K3 contributes a KDA-compatible prefix-cache implementation to vLLM — a serving-side kernel dependency for the hybrid-linear-attention backbone — and recommends supernode deployment (≥64 accelerators) so expert-parallel communication stays inside a single high-bandwidth domain.
  • [2026-07-14] Inside TPU and GPU Clusters: The Anatomy of Collective Communication: Documents the multi-GPU-side hardware constraints that the compute-communication kernel work on this page (ThunderKittens ParallelKittens, FA4, SonicMoE, NVSHMEM) is designed to exploit: TPU ICI torus with per-hop latency and bandwidth hierarchy vs GPU fat-tree with NVSwitch SHARP in-network reduction and multicast. Concretely notes SHARP’s ~2× ideal All-Reduce speedup collapses to ~1.3× in practice on NCCL 2.29.7 / H100 IB — the specific gap that motivates kernel-level comm/compute overlap in Loads and Loads of Fluffy Kittens and NVSHMEM device-initiated communication.
  • [2026-07-10] The 4-bitter Lesson: Balancing Stability and Performance in NVFP4 RL: Concrete Blackwell/PTX kernel work for end-to-end NVFP4 RL: (a) per-token FP32 activation scaling fused into the quant kernel to avoid batch-shared scales; (b) a Four-Over-Six fast path using cvt.rn.f16x2.e2m1x2, cvt.rn.f16x2.e4m3x2, and mul.rn.f16x2 to keep FP4·FP8 multiplies lossless in FP16 registers (BF16 lacks mantissa bits), ~2.8× faster than the FP32-scalar reference at >99.97% scale-choice agreement; (c) TransformerEngine PR eliminating duplicate quantized tensor storage, cutting NVFP4 linear peak memory 150 MB → 45 MB (–70%) on a 2k×2k→8k proxy.
  • [2026-06-30] waterloo_intern (ali / Baseten) — Worklog: retrofitting sparsity for the world's fastest video generation kernel (54×): Baseten engineer ali (@waterloo_intern) teases a worklog on retrofitting sparsity into a frontier OSS video generation model at inference time, claiming the world’s fastest video generation kernel and 54× cumulative speedup. Worklog details (sparsity pattern, kernel implementation, baseline) not retrievable at filing time — captured as a tweet-as-pointer until the full writeup is accessible.
  • [2026-06-29] Demystifying NVSHMEM: A System-Level Analysis on Symmetric Memory and Device-Initiated Operations in GPU Communication: Consolidates the NVSHMEM design — symmetric memory + device-initiated put/get + IBGDA — that the compute-communication kernel work on this page (ThunderKittens, FA4, SonicMoE) sits one layer above; identifies where current runtime defaults (completion semantics, fence/quiet cost) leak into upstream kernel performance, complementing Hazy Research’s earlier observation that off-the-shelf NVSHMEM defaults are suboptimal.
  • [2026-05-26] Laguna M.1/XS.2 Technical Report: Laguna fuses MoE all-to-all dispatch/combine into the CUTLASS grouped-GEMM kernel: 8 of 132 SMs per H200 copy expert tokens over NVLink in expert order and set HBM ready-flags, the tile scheduler waits on flags before processing, and the GEMM epilogue stores results directly to remote HBM via shared-memory vectorized stores — concrete production instance of the compute-communication-fusion thesis.
  • [2026-05-24] VideoNSA: Native Sparse Attention Scales Video Understanding: Locates the wall-clock bottleneck inside NSA’s three-branch sparse attention — at 128K video-text context, the compression branch dominates runtime (linear-in-context wall clock), while selection and sliding-window remain cheap. Concrete kernel target for the IO-aware template applied to learnable sparse attention: fuse NSA’s per-block average-pooling + MLP projection with downstream gated combination so the block-level KV tensor is never materialized to HBM. Companion observation (Finding 6) that learnable sparse attention produces branch-specific attention sinks — compression manufactures them, selection eliminates them — adds a numerical-stability dimension worth tracking when designing the fused kernel.
  • [2026-05-24] Attn-QAT: 4-Bit Attention With Quantization-Aware Training: extends the template to training-side co-design with low-precision fused attention. Identifies that FlashAttention’s memory-efficient backward identity (Pi⊤dPi=dOi⊤Oi\mathbf{P}_i^\top \mathbf{dP}_i = \mathbf{dO}_i^\top \mathbf{O}_i) implicitly assumes forward and backward use matching precision; under fake-quantized FP4 forward the identity breaks and gradients become inconsistent. Two fixes: (1) materialize an auxiliary O′=PVF\mathbf{O}' = \mathbf{P}\mathbf{V}^F during the forward (25% extra storage) and use it in the backward; (2) fake-quantize the recomputed P\mathbf{P} in the backward to match the forward precision. Result: FP4 attention matches BF16 on Wan-2.1-14B video diffusion (human 2AFC) and Qwen3-14B / Llama-3.1-70B continued pretraining + SFT, without any outlier-mitigation heuristics. Paired CuTeDSL kernel (FlashAttention-4 FP4 on B200) runs block-scaled NVFP4 on QK\mathbf{QK} + BF16 on PV\mathbf{PV} — the “FP4 only where it’s free” design point dictated by Blackwell’s softmax bottleneck — reaching 1801 TFLOPs/s and 1.39× over FA4 BF16 at long context.
  • [2026-05-23] Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon: extends the template to optimizer-state computation. Gram Newton-Schulz rewrites Muon’s Newton-Schulz iteration so the dominant computation runs on the small symmetric Gram matrix XX⊤\mathbf X \mathbf X^\top instead of the rectangular momentum matrix, then exploits custom CuTeDSL symmetric-GEMM kernels (Hopper + Blackwell) that write only the lower triangle. 40–50% wall-clock reduction in the Muon orthogonalization step on trillion-parameter MoEs (Kimi K2-scale), with training quality preserved within 0.01 validation perplexity. Adds “structured-matmul exploitation” as a distinct mechanism alongside HBM-traffic elimination, and introduces a half-precision stability investigation (spurious negative eigenvalues amplified by (15/8)2(15/8)^2 per step, plus eigenvector drift) fixed by periodic restart.
  • [2026-05-23] FlashSampling: Fast and Memory-Efficient Exact Sampling: extends the template to LM-head categorical sampling. Argmax-over-partitions plays the same role for sampling that softmax-over-blocks plays for attention; the categorical distribution’s hierarchical factorization enables a tensor-parallel variant that replaces all_gather(logits) with streaming per-rank (max, argmax) exchanges.
  • [2026-05-23] Flash-KMeans: Fast and Memory-Efficient Exact K-Means: applies the template to Lloyd’s k-means. FlashAssign fuses streaming distance computation with online argmin (no N×K distance matrix in HBM); Sort-Inverse Update converts scatter-atomic centroid updates into argsort + segmented partial sums. Up to 200× over FAISS at iteration latency.
  • [2026-05-23] FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling: re-tunes the original FlashAttention design for Blackwell. New roofline accounts for SMEM traffic and MUFU exp throughput as co-dominant with matmul; introduces software-emulated 2^x on FMA units, conditional softmax rescaling, and a backward pass that uses 2-CTA MMA to halve operand-B SMEM staging.

Open questions

Section titled “Open questions”
  • Where else does the template apply on the wiki? Any kernel that materializes an N × K intermediate solely to reduce along one axis is a candidate: top-k logits, MoE router gating, perceptual losses, RoPE applied before attention, contrastive logits. None of these have been filed yet as IO-aware rewrites. Partial answer: Gram Newton-Schulz is an adjacent variant — instead of fusing a produce/reduce pair, it rewrites the iteration so the dominant matmul is structured, and ships a kernel that exploits the structure. This suggests an enlarged template: “find an intermediate whose algebraic structure (symmetry, low rank, banded) is implicit but not exploited by standard libraries, restructure the algorithm to make the structure dominant, write the structured kernel.” VideoNSA adds a learnable sparse candidate: NSA’s compression branch is N × K (per-block mean-pool + MLP) → reduced by the per-head gate; fusing the pool + projection + gate combination is the obvious next target.
  • Backward-pass coverage. FA4 has a backward pass; Flash-KMeans, FlashSampling, and Gram Newton-Schulz are inference / online primitives only. Attn-QAT is the first filed paper that targets the backward pass head-on, showing that the FlashAttention backward identity implicitly assumes forward/backward precision match — and that under low-precision training this assumption must be made explicit by carrying a small high-precision auxiliary output. Whether the template applies to backward passes that don’t have an obvious partition decomposition (e.g. CE loss + LM-head backward, which is exactly the locus that Lost in Backpropagation: The LM Head is a Gradient Bottleneck flags as a bottleneck) is still unaddressed.
  • Compile-time as the bottleneck. All filed papers either rebuild compile infrastructure (FA4’s, Gram-NS’s, and Attn-QAT’s CuTeDSL moves) or replace autotune sweeps with analytical heuristics (Flash-KMeans’s cache-aware config selector). The implicit claim is that iteration speed on the kernel-design loop itself is now load-bearing for architecture-specific tuning at the rate Blackwell/Hopper/B300 ship; whether this means Python-DSL becomes the dominant paradigm or whether C++-template + cache modeling wins is open. Dao-AILab’s stack (FA4 + Gram-NS + quack/gemm_symmetric + Attn-QAT FA4-FP4) is the strongest existing endorsement of the CuTeDSL-in-Python direction.
  • Approximation-as-extension. Three of the original four papers emphasize exactness. The natural next axis is controlled approximation — block-sparse attention, approximate k-means with bounded error, top-k sampling — built on the same kernel skeleton. The wiki has MonarchRT: Efficient Attention for Real-Time Video Generation as one such datapoint on the attention side; Attn-QAT is the precision-side version of approximation — quantization is a controlled approximation, and the training-time fix lets it match BF16 in practice. VideoNSA / NSA is the learnable sparse version of the same axis: sparsity pattern + gating are trained, not statically chosen.
  • Numerical stability in half precision. Gram Newton-Schulz’s restart trick is the explicit instance of this: bfloat16 spurious negatives blow up exponentially via rt<(15/8)2rt−1r_t < (15/8)^2 r_{t-1}. FA4’s conditional softmax rescaling addresses a milder version of the same problem. Attn-QAT adds a third instance: the FlashAttention backward identity silently breaks under precision mismatch, and only the auxiliary O′\mathbf{O}' trick + matched backward recomputation restores it. Whether other applications of the template (Flash-KMeans, FlashSampling) have undiagnosed half-precision pathologies is open.
  • Where does the FP4 attention story go next? Attn-QAT’s design point on B200 — NVFP4 QK\mathbf{QK} + BF16 PV\mathbf{PV} — is dictated by the softmax bottleneck; B300 doubles MUFU exp throughput and Rubin quadruples it, which should make NVFP4/MXFP8 QK\mathbf{QK} + FP8 PV\mathbf{PV} a net win. The open meta-question is whether the QAT recipe needs to be re-derived per hardware generation, or whether the precision-consistency fixes (auxiliary output + matched backward recomputation) transfer.

Papers

Section titled “Papers”

2026-05

Section titled “2026-05”

2026-04

Section titled “2026-04”

2026-03

Section titled “2026-03”

2026-02

Section titled “2026-02”

2026-01

Section titled “2026-01”

2025-12

Section titled “2025-12”

2025-11

Section titled “2025-11”

2025-10

Section titled “2025-10”

2025-09

Section titled “2025-09”

2025-07

Section titled “2025-07”

2025-06

Section titled “2025-06”

2025-05

Section titled “2025-05”

2025-04

Section titled “2025-04”

Unknown date

Section titled “Unknown date”