arXiv Science⌕ Search
Home/Archive/arXiv

arXiv · 2609.12379

ForgeMegakernel: A General Framework for Efficient Auto-Regressive Model Decode Megakernels

Leshan Li·Zhui Zhu·Xianglong Deng·Yaojian Chen·Qingfeng He·Yuxuan Li·Rong Zhao·Xu Han·Zhiyuan Liu
View original source ↗Read source document ↗Export metadata ↓

Abstract

Auto-regressive model decode is bandwidth-bound, since every weight and key/value-cache byte crosses high-bandwidth memory once per token. A megakernel is an ideal solution, but existing automatic megakernel generation approaches cannot achieve both generalization across models and correctness guarantees. We present ForgeMegakernel, which generates a per-model high-performance decode megakernel using coding agents. ForgeMegakernel pairs a universal knowledge base of ten progressive milestones with an independent mid-state test oracle. The milestones provide the megakernel's structural properties: a fine-grained instruction stream for each SM, dependency counters replacing the global synchronization, and a shared-memory buffer pool for workload balance across SMs and greater parallelism. The test oracle derives the mid-states of the megakernel and checks the performance, error and precision during the generation process, guaranteeing a correct and trustworthy forged megakernel. We evaluated ForgeMegakernel on 14 representative decoding operations across eight model families spanning 0.6B-13B parameters. The generated megakernels achieved 50.5-85.9% MBU and geometric mean speedups of 1.21x over SGLang 0.5.18 and 1.54x over a megakernel compiler under identical configurations. Inside SGLang, evaluated on GSM8K with ragged prompts, all 14 megakernels decoded faster than the SGLang engine at comparable answer accuracy.

Explore related subjects

cs.DC

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Leshan Li, Zhui Zhu, Xianglong Deng, Yaojian Chen, Qingfeng He, Yuxuan Li, Rong Zhao, Xu Han, Zhiyuan Liu. 2026-09-11. ForgeMegakernel: A General Framework for Efficient Auto-Regressive Model Decode Megakernels. https://arxiv.org/abs/2609.12379

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Multimmit: Extending Blocks for Faster Finality

To meet the throughput demands of modern blockchain systems, protocols for State Machine Replication (SMR) increasingly have many processors disseminate blocks of transactions in parallel, with consensus then establishing a total ordering on the blocks of all producers. Such designs face a choice as to when a block may enter the ordering. Certified approaches wait for a quorum to attest a block's availability, which is robust but adds message delays to every transaction. Uncertified approaches let proposals reference blocks immediately, which is fast but degrades rapidly when referenced data must be fetched on the critical path. Raptr, the state of the art, takes a middle course, finalising the longest prefix of the leader's proposal that a quorum holds, so that no processor ever blocks or fetches. The remaining weakness is sensitivity to order: if the data behind a single early batch is withheld, the proposal finalises little or nothing, so individual faulty producers can still deny the system its optimistic path. We present Multimmit, a protocol for $n \ge 5f+1$ processors combining a consensus layer requiring one round of voting per view with multi-chain data dissemination. Votes are cast relative to the leader's proposal, reporting per chain how far the voter can support it, and may themselves attest fresh blocks beyond it. A transaction block disseminated at time $t$ is ordered by $t+3δ$ in expectation and $t+2δ$ at best, measured from the block's dissemination rather than the leader's proposal. Degradation under faults is graceful: a faulty producer delays only its own chain's blocks, costing other chains at most a one-view wait for placement. No leader can both finalise its leader block and exclude a fresh, well-circulated block of an honest chain. Consensus traffic is tens of kilobytes per view, independent of transaction volume.

cs.DC↗
arXiv2026

Comparison of Algebraic Block Multi-Coloring and Leiden Methods for Parallel Preconditioning in the ICCG Method

In the application of incomplete Cholesky preconditioning to the incomplete Cholesky-conjugate gradient (ICCG) method, forward and backward substitutions exhibit sequential dependencies that constitute a major bottleneck for parallelization in multicore environments. To alleviate this bottleneck, the algebraic block multi-coloring (ABMC) method achieves both parallelism and data locality through block-wise coloring. However, ABMC requires the number of blocks to be specified as an input parameter in advance. This study evaluates the Leiden method as an alternative blocking approach for parallel preconditioning in the ICCG method. As a community detection technique that maximizes a quality function for graph partitioning, the Leiden method automatically generates blocks that reflect the matrix structure without requiring the number of blocks a priori. We partition the adjacency graphs of sparse matrices using the Leiden method and utilize the resulting blocks for parallel preconditioning. We implement the Leiden method using modularity and the constant Potts model as quality functions and compare its performance with that of the ABMC method in terms of the number of iterations, execution time, and L2 cache efficiency across eight symmetric positive definite matrices. The experimental results demonstrate that the Leiden method with the constant Potts model achieves performance comparable to that of the ABMC method configured with an optimized number of blocks.

cs.DC↗
arXiv2026

ForgeStencil: Automating Per-Case Stencil Specialization from Kernels to 100+ Real Applications

Industrial and scientific computing rests on a few core kernels, and the stencil is among the most widely used: weather and climate models, seismic imaging, fluid dynamics, and image processing all run on it. No single stencil implementation is fastest: the optimal kernel changes qualitatively with stencil shape, grid shape, precision, and host application. For two decades the field has answered with general methods (DSLs, code generators, autotuners), because specialized solutions were too expensive to build per case, so all reuse one human-authored recipe. That reuse costs performance; we call the cost the generality tax. This premise no longer holds: code-synthesis agents now build a correct, specialized solution per case at acceptable cost. ForgeStencil automates this. A Kernel Agent synthesizes CUDA and forges a per-configuration map of specialized operators, removing the tax case by case. On an A100 the map beats the strongest public baseline in 37 of 37 cases: geometric mean 2.35x against same-precision f32 baselines and 1.95x for fp16, each reported under its own precision. The same change reaches end-to-end application performance. A generic operator library is tuned once for its own general case and reused across applications, so its shapes, layouts, and launch boundaries are optimal for none of them: using it is the application-level form of the tax. An App Agent instead forges a specialized solution per application, locating hotspots, rewriting application structure, and validating and integrating each change. Across 100 real industrial and scientific codes the end-to-end median speedup is 1.41x against each application's own GPU baseline. To our knowledge this is the first demonstration that per-case synthesis carries from a kernel library to complete applications at this breadth, and evidence that reuse is no longer the default in a domain built on it for two decades.

cs.DC↗