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-11-06] cuRoboV2-Fast: Fast Torque-Limited Collision-Free Motion Planning: cuRoboV2-Fast (NVIDIA LPR; arXiv:2611.06913) transplants the template onto classical robot trajectory optimization — TimeOpt represents squared path speed as a B-spline and exploits the spline’s local support to produce banded Newton systems solved efficiently on GPU, the same “exploit algebraic structure of the intermediate” lever Gram Newton-Schulz uses for symmetric GEMMs (Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon) but applied to the Newton-system matrix of a trajectory-optimization problem rather than the Gram matrix of an optimizer step. The geometry-refinement stage adds a fused cost-reduction + gradient-accumulation kernel that removes separate pointwise and gradient-add launches (the same produce/reduce-pair fusion pattern FlashAttention and Flash-KMeans use), and a spline-aware backward gather where each control point gathers its gradients with one warp reduction without atomics. First filed instance of the template operating on motion-planning kernels rather than LLM / attention / sampling / clustering / optimizer-state primitives; the per-cycle result (planning at 1% of robot cycle time on real hardware) is a direct datapoint on how much latency the template unlocks on a non-LLM GPU workload.
  • [2026-11-05] ERNIE 5.5 Flash Base: FlashMask-4 (ERNIE 5.5 Flash infrastructure co-release) is a CuTe-DSL attention kernel for large head dimensions (qk_head_dim up to 576, v_head_dim up to 512) — beyond what FlashAttention-4 supports — via 2-CTA tiled parallel computation and (on SM100) double-buffer rotation in tensor memory, while preserving FlashMask’s compact O(N) column-wise mask representation and fully-masked-block skipping. Reported up to 175% faster forward and 74% faster backward than FlashAttention-4 MaskMod on SM100 at the model’s head dimensions. The head-dimension axis matters because ERNIE 5.5 Flash’s HCA layers deliberately use wide (512-dim) heads to spend compute where the sequence is cheap (128× compressed), and MLA layers use narrow (256-dim) heads where attention is expensive — a per-layer attention-cost-aware head-dimension schedule that off-the-shelf kernels under-support. Combined on the distributed side with Muon over the expert-parallel ring + inter-node balanced optimizer-state sharding.
  • [2026-11-04] GPU-accelerated Weight Teleportation: humans& ai apply the IO-aware-kernel template to the weight-sync path: the Blackwell Decompression Engine (fixed-function ASIC for supported codecs) runs concurrently with SM compute, so a double-buffered pipeline streams one model layer through DE decompression while SMs apply bitwise XOR via a single Triton kernel launch, hiding Triton work entirely under DE setup/decompress loops and bounding the rollout pause by DE throughput (0.2196 s on NVFP4 GLM-5.3). Novel sub-template elements: pinned-CPU-RAM input to DE via CU_MEM_LOCATION_TYPE_HOST_NUMA + CU_MEM_CREATE_USAGE_HW_DECOMPRESS allocations (not ordinary pinned); two layer-sized HBM slots fit in existing activation memory so KV cache survives; fused-kernel-layout-aware mask remapping (paired gate-up weights in groups of 64) so bitwise XOR aligns under non-canonical physical weight layouts. Design-point selection (Snappy inner + Zstd outer, CPU-Zstd + DE-Snappy split) is driven explicitly by the DE support matrix, extending the “fuse produce/consume + exploit algebraic structure” template with “exploit fixed-function decompression hardware.”
  • [2026-11-03] Latent Sparse Attention: Towards Subquadratic Ultra-Long Context Attention: Latent Sparse Attention (Prefrontal AI / ETH Zurich) is the long-context sparse-attention sibling of the “attack the dominant kernel cost at the target regime” template this page tracks: at 10M tokens the indexer, not the attention or the KV read, is the bottleneck (~98% of attention-layer time). LSA clusters keys hierarchically in a learned latent space so queries score a tree of cluster centroids; the hierarchical variant D-LSA (Tree) is approximately O(T log T) in prefill indexing cost. Released kernels ship as a custom SGLang fork + CUTLASS-style flat and tree search variants under src/lsa/core, baselines include Full Attention / SWA / NSA / PISA. Natural next step on the thread VideoNSA opened — identifying NSA’s compression branch as the kernel bottleneck — by moving the index space itself into a learned, hierarchical representation rather than keeping a fixed block grid.
  • [2026-10-28] Herman Brunborg — 83% NanoGPT speedrun improvement via retrieval-augmented n-gram hints: Herman Brunborg’s 83% NanoGPT speedrun improvement (39.9s → 6.7s) ports the “treat the bottleneck as kernel work” discipline onto the CPU side of the pretraining job: a branch-free SIMD n-gram-hint-corpus builder that fits in L2 cache runs on otherwise-idle host CPUs during the timed GPU run, reportedly 500× faster than Python and 11× faster than the best auto-researched Rust code — a direct datapoint on the “agents write kernels directly” question that Retire the Abstractions — CUDA DSLs Heading Toward Retirement opens. Unverified at filing (no code, no writeup, no ablations), but if the claim stands it is the host-side analog of this concept’s GPU-side template applied to a data-pipeline bottleneck.
  • [2026-10-25] CroissAttn — fixing FlashAttention backward's zero-sum bug for end-to-end FP4 training: CroissAttn (Johansen, ICLR’27 submission) adds a per-draw correction rule to the FA backward: the standard shortcut uses δ = ⟨dO, O⟩ with the stored O, which is correct in expectation but wrong on every single draw once dP is produced by FP4 stochastic rounding. The fix computes δ from the actually-rounded dP row (two sweeps, one fewer high-precision product than the shortcut), restoring the exact zero-row-sum identity per draw. Distinct from Attn-QAT’s approach to the same precision-mismatch failure (Attn-QAT: 4-Bit Attention With Quantization-Aware Training materializes an auxiliary high-precision O’ and uses it in the backward identity) — CroissAttn operates on dP directly, no auxiliary tensor. Pathology concentrates on attention-sink rows where the shortcut’s error is unbounded. Author notes the submitted kernels still used the shortcut via a high-precision auxiliary product; whether the per-draw rule’s “delete the auxiliary product” variant is a net wall-clock win is still being profiled — a direct contest between the two kernel implementations of the same backward-identity fix.
  • [2026-10-16] Introducing Muse Spark Open — The Most Capable US Open-Weight Model to Date: Muse Spark Open ships auto-tunable Helion GPU kernels inside SGLang as part of the release package — claimed 1.03–3.95× kernel-level speedup on NVIDIA and AMD GPUs, 3% end-to-end speedup on AMD MI350. First filed instance of a frontier open-weights release co-shipping its own kernel implementation into a third-party serving framework at launch (rather than deferring to community-tuned kernels post-release).
  • [2026-10-15] Frontier Pre-training at Microsoft AI (MAI-Base-2): MAI-Base-2 catalogs four frontier-scale IO/kernel optimizations for MoE training on 16K GB200: (1) grouped GEMM + NVLink per expert group to amortize kernel-launch overhead and tail effects; (2) device-resident token metadata + GPU-initiated all-to-all to remove host-device sync; (3) symmetric-memory all-to-all over NVLink without staging copies, improving throughput and reducing HBM contention under concurrent GEMMs; (4) batched load-balancing-loss all-reduce across all MoE layers (replacing per-layer all-reduces). Also: pipelined dispatch/compute/collect rounds overlap with weight quantization and shared-expert compute (forward) and with input-dispatch/weight-gradient computations (backward). Together these keep per-expert overhead low enough that sparser MoE configurations don’t degrade MFU.
  • [2026-10-13] The Will to Act — Natura-1: Natura-1 (Nature Will) reports custom attention kernels delivering ~1,200 TFLOP/s average on forward + backward benchmarks — >6× the throughput of a tuned FlexAttention baseline on their workloads — translating to >1.7× end-to-end training speedup and 2× inference speedup in the evaluated configurations. Shared GPU kernel layer sits across both training and real-time inference, so optimizations compose: fused operator pairs, reused intermediate results, dynamically adapted execution plans, reduced memory traffic + kernel-launch overhead. The shared-operator discipline across training and inference is pitched as load-bearing for numerical consistency between large-scale learning and real-time 100 ms control-loop deployment on physical robots — the “same op, both regimes” pattern makes per-op kernel work pay twice. Companion in framing to InstinctFlash — High-Performance Serving Runtime for Robotics Models InstinctFlash and Embodied.cpp: A Portable Inference Runtime of Embodied AI Models on Heterogeneous Robots Embodied.cpp but at the kernel-layer rather than runtime-layer; the 6× FlexAttention margin is one of the sharpest reported attention-kernel speedups filed on this page.
  • [2026-10-02] Veda: Scalable Video Diffusion via Distilled Sparse Attention: Veda ships a ThunderKittens sparse-attention kernel that fetches only the top-K selected K/V tiles via TMA and overlaps sparse gathers with WGMMA compute, with no proxy-score map materialized to HBM — the exploit-algebraic-structure template applied to sparse-attention gathers rather than fused softmax, co-designed with per-head learned tile shapes (pt, ph, pw) chosen to minimize sparse-output error. 5.1× E2E on Waver-T2V-12B 720P / 241 frames at 95% sparsity.
  • [2026-10-01] Improving Video Sparse Attention with Fine-grained Router and Sparse Rebasing: VSA2 extends the trainable-sparse-attention branch (alongside VideoNSA) by making the router fine-grained and allowing per-query variable KV count — the per-query budget becomes a trained, data-dependent property rather than a hyperparameter. The Hard-to-Easy training curriculum (train at higher sparsity than inference) is a new axis on the “trainable sparsity pattern + gating” sub-template VideoNSA opened: here the training-side discipline is what unlocks the quality parity, not the sparsity pattern alone. 8.9× attention speedup over FlashAttention-3 at 720p.
  • [2026-09-30] Beyond kernels: Letting agents optimize the whole AI job: Whole-job counterpart to kernel-level IO-aware design: the FastVideo case study (unnecessary post-decode GPU→CPU return of 10.4 GB VAE weights, 5.87s of 8.87s decode time) is a canonical IO-inefficiency at the job level rather than the kernel level, and the fix (pin CPU copies, skip return transfer) is IO-aware discipline applied one layer up. The plugin’s profiling scripts explicitly surface CPU–GPU timelines / transfers / synchronization as first-class diagnostic views.
  • [2026-09-30] Sol-H3: Recursive Self-Improvement for MiniMax-H3 Inference Acceleration on Sol-Engine across Cloud and Edge: Sol-H3 replaces hand-crafted kernel-tuning heuristics with a Recursive Self-Improvement loop that searches over kernel fusions and memory layouts, scoring candidates on latency together with numerical agreement with a reference. Datapoint on the “compile-time as the bottleneck” thread: instead of a compiler rebuild (FastAttention-4, Gram-NS, Attn-QAT) or analytical autotune replacement (FlashSampling), the tuning search itself becomes an autoresearch loop with a correctness gate. Applied here to accelerate MiniMax H3 across an 8×GB200 cloud node and a single DGX Spark edge workstation from the same searched kernel library.
  • [2026-09-12] AI Performance Engineering — curated resource list for GPU perf engineering and production inference: Wafer AI’s curated resource list; the Kernel Optimization + Programming Models sections cover FlashAttention 1–4, DeepGEMM, CUTLASS/CuTe, CUDA Tile IR, Triton, and the ROCm Composable Kernel / HipKittens AMD stack, ordered as a reading path. Explicit source policy requires baseline + hardware + workload + correctness for any performance claim — the curation-side counterpart to the “numerical-stability-in-target-precision” and “compile-time-as-bottleneck” open questions this concept tracks.
  • [2026-09-12] All About Transformer Inference — Chapter 7 of How To Scale Your Model: Chapter 7 of “How To Scale Your Model” — establishes the roofline framing this concept’s kernel work is designed to exploit. Derives the attention-intensity formula ST/(S+T)ST/(S+T) that motivates every KV-reduction kernel (GMQA/MQA, sliding-window, KV quantization, PagedAttention/ragged HBM reads) the concept’s papers instantiate; explicitly cites FlashAttention variants and PagedAttention as the load-bearing primitives for the generation-phase memory-bandwidth wall.
  • [2026-09-03] FastVideo FastH3 V1 — 4-Step Sparse Distilled MiniMax H3 for 14x Blackwell Speedup: FastH3 Preview v1 distillation is calibrated to FastVideo’s VSA-H3 (Video Sparse Attention) kernel backend — the distilled weights explicitly require the custom kernel and do not fall back to dense attention. Continuation of the FastWan / FastWan-QAD pattern of shipping checkpoints tightly coupled to a specific kernel implementation, and a datapoint for the “distillation + kernel” co-design axis that most other diffusion-serving entries treat as independent.
  • [2026-09-03] VideoDeltaNet on MiniMax H3 — Near-Lossless Hybrid Attention for Video Models: VideoDeltaNet-H3 reports a per-block latency waterfall isolating each hardware-facing lever on B200: architectural change (VDA + Sliding Window Softmax) cuts 332.5 ms → 192.1 ms; adding optimized attention kernels + FP8 linears drops another 66.8 ms to 125.3 ms (2.65× total per-layer). The 5+3 branch-parallel all-to-all sequence-parallelism split (softmax branch on 5 GPUs, linear branch on 3, chosen by profiling) buys 13.3% on top of standard Ulysses. Illustrates that in hybrid diffusion transformers, kernel-level and communication-level optimizations remain roughly stackable with architectural change.
  • [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”