Token-Centric Kernel Fusion
- Token-Centric Kernel Fusion is a family of methods that prioritize tokens for similarity computation, aggregation, and execution reorganization to optimize deep learning inference.
- The approach integrates token-space aggregation, kernel regression-based attention, and runtime fusion techniques to reduce memory traffic and improve throughput.
- Empirical studies show that these techniques yield improved accuracy, reduced latency, and lower GPU memory usage through effective token compression and fused computation.
Token-centric kernel fusion denotes a family of methods in which tokens are the primary objects of similarity computation, regression, scheduling, or execution fusion. In recent work, the term spans at least three levels: token-space aggregation that compresses or merges redundant tokens before or within a model; token mixers that interpret attention itself as a kernel regressor over tokens; and systems techniques that reorganize the per-token execution path so that communication, normalization, attention, and even entire decoder passes are fused into fewer synchronization domains (Pippi et al., 6 Mar 2025, Zheng et al., 7 May 2026, Gond et al., 16 May 2025, Jaber et al., 8 Jun 2026, Koeplinger et al., 2024). Across these levels, the recurring objective is to preserve task-relevant token information while reducing memory traffic, synchronization overhead, or context length.
1. Scope and historical framing
The literature uses kernel fusion in both a classical CUDA sense and a token-oriented inference sense. Classical kernel fusion arose from the observation that many GPU kernels are memory-bound, so sequences of map, reduce, or nested map/reduce kernels benefit when intermediate data remain in registers or shared memory instead of being written to and reread from global memory. In "Optimizing CUDA Code By Kernel Fusion—Application on BLAS" (Filipovič et al., 2013), a source-to-source compiler for such patterns generated implementations that were up to faster than comparable CUBLAS-based sequences. In "Operator Fusion in XLA: Analysis and Evaluation" (Snider et al., 2023), XLA’s HLO pipeline was shown to use instruction fusion, fusion merger, multi-output fusion, and horizontal fusion; on Cartpole, alternative fusion strategies produced up to speedup over the baseline implementation.
More recent work shifts the unit of analysis from generic tensors to tokens. "Token-Operations-Oriented Inference Optimization Techniques for Large Models" (Lian et al., 18 Jun 2026) makes this explicit by treating the token as the minimum unit of information processed, metered, billed and traded, and by organizing inference optimization around cost per token, tokens/sec, and token-supply stability. This framing places token-centric kernel fusion at the intersection of model structure, runtime scheduling, operator fusion, KV-cache layout, and distributed communication.
| Fusion level | Representative mechanism | Representative papers |
|---|---|---|
| Token-space compression | Similarity-based merging, pruning, norm-preserving interpolation | (Pippi et al., 6 Mar 2025, Kim et al., 2023) |
| Token mixing | Nadaraya-Watson or KRR over tokens | (Zheng et al., 7 May 2026) |
| Execution/runtime fusion | Token splitting, fused communication-normalization, megakernels, kernel looping | (Gond et al., 16 May 2025, Jaber et al., 8 Jun 2026, Koeplinger et al., 2024) |
This taxonomy suggests that token-centric kernel fusion is not a single algorithmic primitive but a layered design space. The common constraint is that fusion must respect the geometry, causality, and hardware footprint of token processing rather than merely collapse adjacent operators.
2. Token-space fusion: clustering, pruning, and merging
Two distinct methods named ToFu illustrate token-space fusion. "Token Fusion: Bridging the Gap between Token Pruning and Token Merging" (Kim et al., 2023) studies ViTs and argues that pruning is advantageous when the model is sensitive to input interpolations, whereas merging is effective when the model manifests close to linear responses to inputs. The method therefore uses pruned merging in early layers and average or MLERP merging in later layers. MLERP is presented as a variant of SLERP tailored to merge multiple tokens while maintaining the norm distribution, addressing the distributional shift caused by average merging. On ImageNet with ViT-B at , the baseline is $83.74$ Top-1 at $17.58$ GFLOPs and $323.61$ img/s, whereas ToFu MLERP reaches $83.22$ Top-1 at $13.12$ GFLOPs and $413.86$ img/s; at , ToFu MLERP reaches 0 Top-1 versus 1 for ToMe at the same 2 GFLOPs. The paper also reports better FID, LPIPS, and MS-SSIM than ToMe for Stable Diffusion 1.5 with essentially identical speed and memory.
The later multimodal "ToFu: Visual Tokens Reduction via Fusion for Multi-modal, Multi-patch, Multi-image Task" (Pippi et al., 6 Mar 2025) moves token fusion to the multimodal interface. It is visual encoder-agnostic, training-free, and operates after the adapter into LLM space. Given visual tokens
3
it maintains representatives
4
and weights
5
Its similarity function is cosine similarity,
6
and each incoming token is either preserved as a new representative if its maximum similarity is below a threshold 7, or fused into the nearest representative by the weighted update
8
This is explicitly described as a token-centric kernel aggregation scheme, or equivalently a sequential kernel clustering procedure whose representatives are cluster means.
The multimodal motivation is computationally specific. In multi-image settings, visual tokens can amount to 9 of the total context, and InternVL2 models can use 0–1 visual tokens per image. ToFu therefore sets 2 dynamically as a function of the total visual token count 3: 4 at 5, 6 at 7, and linear interpolation in between. On LLaVA-Interleave Bench, ToFu reduces the visual token count on average to about 8 of the original; for InternVL2-4B, overall accuracy improves from 9 to $83.74$0, max GPU memory drops from $83.74$1 GB to $83.74$2 GB, and runtime changes from $83.74$3 s to $83.74$4 s. On the ComPairs benchmark, with average $83.74$5K visual tokens per sample, ToFu improves accuracy for most tested models at similar token budgets, including InternVL2-8B from $83.74$6 to $83.74$7 with $83.74$8 tokens. The paper’s central claim is therefore that structured, similarity-based token compression preserves more semantic information than random dropping under comparable budgets.
3. Kernel-regression token mixers
A second meaning of token-centric kernel fusion treats token mixing itself as kernel regression. "Cubit: Token Mixer with Kernel Ridge Regression" (Zheng et al., 7 May 2026) starts from the observation that standard Transformer self-attention is exactly a Nadaraya-Watson regressor over tokens. For head $83.74$9,
$17.58$0
so the output for token $17.58$1 is a kernel-weighted average of values over all tokens. The paper defines the token kernel
$17.58$2
which makes attention a per-token Nadaraya-Watson smoother with row-wise normalization.
Cubit replaces this local normalization with a KRR-based token mixer. The conceptual KRR objective is
$17.58$3
with solution
$17.58$4
The resulting per-head token mixer is written as
$17.58$5
where $17.58$6 is a standard query-key similarity matrix, $17.58$7 implements the KRR-style global inverse-kernel normalization, and $17.58$8 is the Limited-Range Rescale (LRR) applied to values for conditioning and training stability. In implementation, Cubit constructs a token-token similarity matrix from a reference embedding $17.58$9, adds $323.61$0, and solves the corresponding linear system rather than using a low-rank or sparse approximation.
This formulation changes the semantics of token fusion. Standard attention fuses values by query-local kernel averaging; Cubit adds a global kernel coupling across all tokens before that averaging step. The paper therefore characterizes the architecture as a two-stage kernel fusion: first a global KRR solving step on values, then an attention-style Nadaraya-Watson aggregation. Empirically, Cubit reports lower validation loss than Transformer and DeltaFormer on Arxiv and Books3 at training lengths $323.61$1 and $323.61$2; at $323.61$3 tokens, the reported advantage stabilizes around $323.61$4–$323.61$5 loss. At $323.61$6M and $323.61$7B parameters, Cubit continues to outperform Transformer, and on ARC, HellaSwag, PIQA, SciQ, Winogrande, SocialIQA, and RACE it typically has the best or tied-best performance across the reported settings. The main algorithmic trade-off is complexity: without explicit low-rank approximation or sparsity, constructing similarity matrices is $323.61$8 and solving the linear system is $323.61$9 per head in the worst case.
4. Execution-level fusion and per-token decode pipelines
At systems level, token-centric kernel fusion reorganizes the execution path of a token so that fewer kernels, fewer synchronization boundaries, and less redundant communication occur during decode. The general systems view in (Lian et al., 18 Jun 2026) is that decode is dominated by repeated attention over KV cache, small GEMMs, kernel launch and scheduling overhead, and distributed communication; the most relevant fusion mechanisms therefore include fused attention, fused MLP and epilogues, Grouped GEMM for MoE, paged or radix KV layouts, continuous batching, CUDA Graph capture, and communication-compute overlap.
"TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference" (Gond et al., 16 May 2025) makes the token itself the decomposition unit. It splits the tokens in an inference batch into two wave-aware subsets, overlaps computation of one subset with communication of the other, reorders RMSNorm with respect to communication, and introduces a fused AllReduce-RMSNorm kernel using NVIDIA Hopper Multimem. The fused kernel combines ReduceScatter, residual add, RMSNorm, and AllGather, and uses only $83.22$0–$83.22$1 SMs while leaving the rest available for compute. The paper reports up to $83.22$2 latency reduction for the fused AllReduce-RMSNorm kernel relative to AllReduce + RMSNorm, and end-to-end gains of up to $83.22$3 latency and up to $83.22$4 throughput across models and workloads. In several settings, TokenWeave performs better than an equivalent model with all communication removed.
"AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis" (Jaber et al., 8 Jun 2026) pushes execution fusion further by compiling a HuggingFace Llama-family model into a single persistent cooperative CUDA kernel that runs the whole forward pass in one launch. Its central invariant is "one launch is one forward pass is one decoded token". The schedule IR is validated before launch for deadlock-freedom and race-freedom; across $83.22$5 adversarial schedules, including $83.22$6 unsafe ones, the validator had zero false-accepts and accepted all $83.22$7 real lowerings. On a real SmolLM2-135M checkpoint, the system reproduces HuggingFace greedy decode token-for-token and matches perplexity within $83.22$8. An unattended autoresearch loop improves the megakernel over its own baseline by $83.22$9–$13.12$0. A W8A16 megakernel beats CUDA-graphed cuBLAS bf16 at batch-1 decode on L4 by up to $13.12$1, on L40S by $13.12$2–$13.12$3, on A10G by up to $13.12$4, and on RTX 5090 by $13.12$5–$13.12$6, while trailing cuBLAS on A100 and H100.
On reconfigurable dataflow hardware, "Kernel Looping: Eliminating Synchronization Boundaries for Peak Inference Performance" (Koeplinger et al., 2024) treats decode as a synchronization-dominated token-wise workload and fuses consecutive calls to the same decoder kernel into one call to a modified kernel containing a pipelined outer loop. On TensorRT-LLM for H100, a single Llama 3.1-8B decoder is executed as $13.12$7 kernels and a full decode pass over $13.12$8 layers needs $13.12$9 kernel calls per token; on SambaNova SN40L-8, a baseline schedule uses one fused decoder kernel per layer, reducing this to $413.86$0 calls, but synchronization still accounts for more than $413.86$1 of TPOT. Kernel looping collapses the repeated decoder calls into a single looped kernel, yields decode speedups of up to $413.86$2 on SN40L, scaling speedups up to $413.86$3 on multiple sockets, and lets SN40L achieve over $413.86$4 of peak performance on $413.86$5 and $413.86$6 sockets, with up to $413.86$7 speedup over DGX H100.
5. Query-centric contrast and cache-fusion selection
A recurrent misconception is that token-centric fusion is necessarily local and token-intrinsic. "QCFuse: Query-Centric Cache Fusion for Efficient RAG Inference" (Yan et al., 30 Mar 2026) provides a counterexample by contrasting query-centric cache fusion with token-centric but locally driven methods such as CacheBlend, EPIC, and KVShare. Those methods select recomputation candidates from local signals such as KV deviation, positional heuristics, or early-layer attention, whereas QCFuse constructs a context-aware query representation from semantic summary anchors and then recomputes only the top-$413.86$8 query-related tokens, selected from attention at a critical middle layer.
The system workflow is precise. Offline, each chunk’s KV cache is stored on SSD and a small set of high-key-norm semantic anchors is stored in CPU memory. Online, anchors are prepended as a lightweight prefix during query forwarding to obtain a context-enhanced query representation without loading full context KV. Then only the K cache of a single critical middle layer is loaded from SSD, query-context attention is computed at that layer, and the top-$413.86$9 tokens by averaged multi-head attention are selected for recomputation. The recomputation ratio is configurable; the paper reports that QCFuse matches full computation accuracy at around 0 recomputation ratio, is two times faster than full computation, and reduces delay by an extra 1 when it reaches the same accuracy as the baseline. On HotpotQA, it is 2 points better than full computation due to an attention denoising effect.
This contrast clarifies an important boundary condition for token-centric kernel fusion. Token selection can be conditioned on token-token similarity, as in clustering and merging, or on token-query relevance, as in RAG cache reconstruction. The latter does not invalidate token-centric execution; it changes the scoring function that determines which tokens receive recomputation or fusion priority.
6. Design constraints, trade-offs, and open directions
The main limitations are method-specific but structurally consistent. In multimodal token fusion, ToFu assumes that average embeddings of similar patches retain task-relevant content; the paper explicitly notes loss of fine-grained details, sensitivity to 3, order dependence from sequential greedy clustering, a fixed cosine kernel, and the lack of an explicit exact-token-budget control (Pippi et al., 6 Mar 2025). In ViT token fusion, accuracy still drops at very high reduction ratios, BSM is layer-agnostic, and no explicit hardware-level kernel fusion is implemented (Kim et al., 2023). In Cubit, no low-rank or sparse approximation is used for the KRR branch, so complexity remains at least quadratic in sequence length and the Local Linear Regression variant is slow in naive form (Zheng et al., 7 May 2026).
Execution-level systems expose a different set of constraints. TokenWeave depends on coarse token splitting remaining large enough to avoid wave inefficiency and reports that Mixtral-8×22B can see smaller or negative gains at small sequence lengths because the FFN becomes memory-bound after expert splitting (Gond et al., 16 May 2025). AutoMegaKernel localizes a cross-SM synchronization bottleneck on A100 and H100, notes that long-context decode is not measured, and restricts current support to Llama-family models with specific architectural assumptions (Jaber et al., 8 Jun 2026). Kernel looping relies on regular repeated layer structure and sufficient on-chip capacity for promoted loop-carried buffers; its gains are therefore strongest on models whose decode can be expressed as long runs of repeated decoder kernels (Koeplinger et al., 2024).
The forward-looking directions are correspondingly diverse. The multimodal ToFu paper suggests richer kernels 4, task-conditioned kernels, learned fusion weights, and adaptive threshold or cluster-count optimization (Pippi et al., 6 Mar 2025). Cubit argues for a broader regression-theoretic design space including local linear regression, local KRR, and alternative kernel choices (Zheng et al., 7 May 2026). The ViT ToFu paper suggests combining token reduction with sparse or linear attention kernels and implementing fused kernels that combine attention, BSM, and merging (Kim et al., 2023). The systems literature emphasizes fusion-friendly model structures, paged or radix KV layouts, quantization with fused dequantization, Grouped GEMM for MoE, continuous batching, and overlap of communication with compute as the practical substrate for token-centric execution (Lian et al., 18 Jun 2026).
Taken together, these works indicate that token-centric kernel fusion is best understood as a unifying systems-and-algorithms principle rather than a single recipe. At one extreme, it compresses token sets by clustering in an embedding kernel space; at another, it redefines attention as a regularized kernel regressor over tokens; at the systems boundary, it turns the entire per-token execution path into the object of fusion. The shared criterion is that fusion decisions are made with respect to token identity, token geometry, or token service cost, rather than solely with respect to layers, operators, or tensor shapes.