{"id":884,"date":"2026-05-11T16:36:00","date_gmt":"2026-05-11T08:36:00","guid":{"rendered":"https:\/\/connectword.dpdns.org\/?p=884"},"modified":"2026-05-11T16:36:00","modified_gmt":"2026-05-11T08:36:00","slug":"sakana-ai-and-nvidia-introduce-twell-with-cuda-kernels-for-20-5-inference-and-21-9-training-speedup-in-llms","status":"publish","type":"post","link":"https:\/\/connectword.dpdns.org\/?p=884","title":{"rendered":"Sakana AI and NVIDIA Introduce TwELL with CUDA Kernels for 20.5% Inference and 21.9% Training Speedup in LLMs"},"content":{"rendered":"

Scaling large language models (LLMs) is expensive. Every token processed during inference and every gradient computed during training flows through feedforward layers that account for over two-thirds of model parameters and more than 80% of total FLOPs in larger models. A team researchers from Sakana AI and NVIDIA have worked on a new research that directly targets this bottleneck \u2014 not by changing the architecture, but by making the computation inside feedforward layers significantly cheaper through unstructured sparsity.<\/p>\n

Sparsity Exists, But GPUs Ignore It<\/strong><\/h3>\n

Inside a transformer\u2019s feedforward block, for any given input token, only a small fraction of hidden neurons actually fire \u2014 the rest produce zero after passing through the activation function. This is called activation sparsity, and prior work has documented this phenomenon in models with ReLU activations.<\/p>\n

The frustrating reality is that this theoretical savings rarely translates into actual speedups. NVIDIA GPUs are heavily optimized for dense matrix multiplications using Tensor Cores, which operate on large contiguous tiles of data. Traditional sparse formats like ELLPACK (ELL) require a separate kernel pass to convert activations from dense to sparse representation, and that conversion overhead often cancels out what\u2019s saved by skipping the zeros.<\/p>\n

Critically, prior work on sparse LLM kernels (including TurboSparse, ProSparse, and Q-Sparse) has focused on memory-bound GEMV operations \u2014 the single- or few-token inference regime. The research team instead targets compute-bound GEMM operations in the batched setting with thousands of input tokens, where dense baselines on modern devices can execute orders-of-magnitude higher FLOP\/s with large tiles and Tensor Cores. That is a fundamentally harder problem, and the reason prior approaches didn\u2019t generalize to batched training or high-throughput inference.<\/p>\n

\n

<\/p>\n

\n
GUIDE<\/div>\n
\n
Sparser, Faster, Lighter LLMs \u2014 TwELL & Sparse CUDA Kernels<\/div>\n
Sakana AI \u00d7 NVIDIA \u00a0\u2014\u00a0 arXiv:2603.23198 \u00a0\u2014\u00a0 ICML 2026<\/div>\n<\/div>\n<\/div>\n

<\/p>\n

\n \n <\/div>\n

<\/p>\n

\n

<\/p>\n

\n
01 \u2014 The Problem<\/div>\n
Feedforward layers dominate LLM cost \u2014 and most of that work is wasted.<\/div>\n
\n
\n
>\u200a\u200a\u200a\u200a\u2154<\/div>\n
of all model parameters live in feedforward layers<\/div>\n<\/div>\n
\n
80%+<\/div>\n
of total FLOPs consumed by feedforward layers<\/div>\n<\/div>\n
\n
99%+<\/div>\n
of hidden activations can be zero with no accuracy drop<\/div>\n<\/div>\n<\/div>\n
\n

For any given token, only a tiny fraction of hidden neurons actually fire. The rest output zero after the activation function. This is called activation sparsity<\/strong> \u2014 and it has historically been impossible to exploit on modern GPUs because sparse operations ran slower than dense ones.<\/p>\n<\/div>\n

Prior sparse LLM kernels (TurboSparse, ProSparse, Q-Sparse) only targeted single-token GEMV operations<\/strong>. Sakana AI and NVIDIA tackle the harder problem: batched GEMM<\/strong> with thousands of tokens \u2014 the regime that covers both training and high-throughput inference.<\/div>\n<\/div>\n

<\/p>\n

\n
02 \u2014 The Innovation<\/div>\n
TwELL: a sparse format built around how GPU kernels actually work.<\/div>\n
\n
\n
Old Way \u2014 ELL<\/div>\n
Row-wide packing, costly to build<\/div>\n
Standard ELLPACK packs non-zeros row-by-row across the entire matrix. To construct it from a tiled matmul output you need a separate kernel launch, a full global memory read, and synchronization across all CTAs. Those overheads cancel out the savings from skipping zeros.<\/div>\n<\/div>\n
\n
New Way \u2014 TwELL<\/div>\n
Tile-wise packing, built in the epilogue<\/div>\n
TwELL partitions columns into horizontal tiles matching the matmul kernel\u2019s tile size T_n. Non-zeros are packed locally within each tile. By matching dimensions, TwELL is constructed inside the existing gate projection kernel epilogue<\/strong> \u2014 no extra kernel, no extra memory read, no synchronization overhead.<\/div>\n<\/div>\n<\/div>\n
\n

The inference pipeline uses one fused kernel<\/strong> that reads gate activations in TwELL format and performs up + down projections together. The intermediate hidden state is never written to global memory, cutting DRAM traffic at every forward pass.<\/p>\n<\/div>\n

\n

For training, a hybrid sparse format<\/strong> dynamically routes rows into a compact ELL matrix (sparse rows) or a dense backup (overflow rows). Sparsity during training is highly non-uniform \u2014 max non-zeros per row can be orders of magnitude above the average \u2014 so the hybrid design handles this without becoming brittle.<\/p>\n<\/div>\n<\/div>\n

<\/p>\n

\n
03 \u2014 Training Recipe<\/div>\n
Two changes to your training config. Nothing else.<\/div>\n
\n
\n
01<\/div>\n
Replace SiLU with ReLU<\/strong> as the gate activation function. ReLU produces exact zeros for negative inputs \u2014 this is what enables unstructured sparsity. No other architectural change is needed. (Unregularized ReLU sits slightly below SiLU on task accuracy: 46.4% vs 47.1% on the 1.5B model, offset by the efficiency gains.)<\/div>\n<\/div>\n
\n
02<\/div>\n
Add an L1 loss term<\/strong> on the hidden feedforward activations, averaged over all tokens and hidden dimensions across all layers. Recommended coefficient: L1 = 2\u00d710\u207b\u2075<\/code>. Add it to your standard cross-entropy loss. No changes to learning rate, weight decay, batch size, or optimizer.<\/div>\n<\/div>\n
\n
03<\/div>\n
Sparsity stabilizes fast.<\/strong> The non-zero count settles within ~1,000 training steps (~1B tokens). The training kernels deliver memory and throughput benefits for almost the entire training run, not just toward the end.<\/div>\n<\/div>\n<\/div>\n
\n
Watch Out<\/div>\n

At L1 = 2\u00d710\u207b\u2075, over 30% of neurons become permanently inactive (dead neurons)<\/strong> on average across layers. Downstream accuracy is not visibly affected at this level. The paper explores targeted gate weight reinitialization as a mitigation \u2014 yielding +19.1% speedup vs +17.9% baseline with no accuracy cost.<\/p>\n<\/div>\n<\/div>\n

<\/p>\n

\n
04 \u2014 Benchmark Results<\/div>\n
Accuracy preserved. Efficiency scales up with model size.<\/div>\n
\n\n\n\n\n\n\n\n\n
Model<\/th>\nAccuracy<\/th>\nInference<\/th>\nEnergy \/ tok<\/th>\nTraining<\/th>\nPeak Mem<\/th>\n<\/tr>\n<\/thead>\n
0.5B<\/td>\n40.4% \u2192 40.4%<\/td>\n+17.0%<\/span><\/td>\n\u221211.8%<\/span><\/td>\n\u22121.5%<\/span><\/td>\n\u221219.2%<\/span><\/td>\n<\/tr>\n
1B<\/td>\n44.6% \u2192 44.7%<\/td>\n+18.1%<\/span><\/td>\n\u221214.6%<\/span><\/td>\n+7.1%<\/span><\/td>\n\u221225.5%<\/span><\/td>\n<\/tr>\n
1.5B<\/td>\n46.4% \u2192 46.2%<\/td>\n+18.8%<\/span><\/td>\n\u221215.0%<\/span><\/td>\n+11.6%<\/span><\/td>\n\u221228.1%<\/span><\/td>\n<\/tr>\n
2B<\/td>\n49.1% \u2192 48.8%<\/td>\n+20.5%<\/span><\/td>\n\u221217.0%<\/span><\/td>\n+21.9%<\/span><\/td>\n+22.3%\u200a*<\/span><\/td>\n<\/tr>\n<\/tbody>\n<\/table><\/div>\n
All results at L1 = 2\u00d710\u207b\u2075 on a single node of eight H100 PCIe GPUs, sequence length 2048. Efficiency gains grow with scale<\/strong> \u2014 average non-zero activations drop from 39 (0.5B) to 24 (2B), giving the sparse kernels proportionally more computation to skip. * The 2B sparse model uses a larger micro-batch enabled by reduced activation memory, raising peak usage while improving throughput.<\/div>\n<\/div>\n

<\/p>\n

\n
05 \u2014 Key Findings<\/div>\n
What the paper reveals about where sparsity actually lives.<\/div>\n
\n
\n
\u25c6<\/div>\n
Early layers are least active.<\/strong> In a 28-layer 1.5B model, the first two layers have the fewest non-zero activations. Activity peaks in the early-to-middle layers \u2014 consistent with prior work showing LLM reasoning and knowledge retrieval concentrate there.<\/div>\n<\/div>\n
\n
\u25c6<\/div>\n
First tokens in a sequence fire far more neurons.<\/strong> The model allocates exponentially more computation to early sequence positions where contextual cues from prior tokens are absent. This non-uniformity is exactly what the sparse kernels exploit for speedups.<\/div>\n<\/div>\n
\n
\u25c6<\/div>\n
Strong inverse correlation between sparsity and speedup.<\/strong> The paper measures a Pearson correlation of \u22120.996 between each layer\u2019s average non-zero count and its inference speedup contribution. Sparser layers deliver proportionally larger gains.<\/div>\n<\/div>\n
\n
\u25c6<\/div>\n
Larger gains on less specialized hardware.<\/strong> On NVIDIA RTX PRO 6000 (188 SMs vs 114 on H100), training speedups are significantly higher. Dense GEMM is slower on the RTX 6000, while sparse ops run faster \u2014 widening the relative advantage of sparsity on accessible hardware.<\/div>\n<\/div>\n<\/div>\n<\/div>\n

<\/p>\n

\n
06 \u2014 Get Started<\/div>\n
Open-source. All kernels and training code released.<\/div>\n
\n
\n
\u25a0<\/div>\n
Architecture:<\/strong> Works with gated feedforward LLMs \u2014 Llama, Qwen, and any Transformer++ design. Non-gated (original transformer) variant also supported: 11.2% inference speedup vs 17.9% for gated at the same L1.<\/div>\n<\/div>\n
\n
\u25a0<\/div>\n
Hardware:<\/strong> CUDA kernels written for H100 GPUs using TMA-based pipelining and persistent cooperative design. Gains verified on RTX PRO 6000 with even larger speedups.<\/div>\n<\/div>\n
\n
\u25a0<\/div>\n
Existing models:<\/strong> Fine-tuning via sparsification approaches is flagged as a future direction for bringing these kernels to pretrained dense models \u2014 not yet demonstrated in this paper.<\/div>\n<\/div>\n<\/div>\n
\n \ud83d\udcc4\u00a0 GitHub \u2014 Code & Kernels<\/a>
\n \ud83d\udcd1\u00a0 arXiv Paper<\/a>
\n \ud83c\udf10\u00a0 Project Page<\/a>\n <\/div>\n<\/div>\n<\/div>\n

<\/p>\n

<\/p>\n

\n
\n