Optimizing Muse Spark Open on TPU v7 with SGLang-TorchTPU

Featured projects We serve Meta’s Muse Spark Open, an 816B-parameter mixture-of-experts model, on a single TPU v7 host of four chips, with SGLang running on the TorchTPU PyTorch backend. At the center is a megakernel: all 62 layers of a decode or verify step run as one Pallas program, and every weight load, KV read and cross-device transfer is pipelined behind compute. The rest of the work falls into three groups: faster decode steps (collective folding, a max-free attention loop that pipelines its KV chunks, an expert loop that keeps register spills off its common path, and DFlash2 speculative decoding), more KV cache per chip (sequence-on-lane KV layout in the megakernel, and L2 HiCache on TPU with session-aware eviction), and faster prefill (a fused MoE kernel and sequence parallelism). Decode speed. The megakernel halves the time per output token for one or two concurrent requests, where it helps the most; the gain narrows as concurrency grows. Its verify step runs within 1.5–1.7× of the time it takes just to stream its weights from HBM. AgentX. On the 1M-context version of SemiAnalysis’s agentic-coding benchmark, our curve runs from 433 output tokens/s per request at 6,101 tokens/s per chip to 25,906 tokens/s per chip at 95 tokens/s per request (Figure 1). Figure 1. InferenceX AgentX, 1M-context trace set: tokens per dollar against p90 interactivity. TPU v7: one host, four chips, FP8 compute with NVFP4 expert weights. GB200 (FP4) is shown as a baseline from InferenceX’s published results as of 2026-11-02. Costs are InferenceX’s external owning costs: $1.21 per TPU v7 chip-hour and $1.86 per GB200 GPU-hour. AgentX is built from 393 opt-in Claude Code sessions. It replays their request lengths, prefix sharing between turns, tool pauses and subagent structure, with synthetic tokens in place of the original content. We run its 1M-context trace set, as the GB200 baseline does, where prompts reach about 990K tokens; models with shorter context windows run a version capped at 256K. Throughput counts input and output tokens. Every TPU point in Figure 1 finishes requests within 47 s at p90. Background The model Muse Spark Open is Meta’s open-weights mixture-of-experts (MoE) [1] model, with 816B parameters in total and about 42B active per token. 62 decoder layers, all of them MoE. Each layer routes every token to 8 of 256 experts. Hybrid attention. 46 layers use sliding-window attention [2] over the last 2,048 tokens; the other 16 attend over the full context, up to 1M tokens. Attention has 128 query heads and 16 KV heads, each 64 wide. NVFP4 experts. Expert weights are 4-bit, with one scale per 16 values, a microscaling format [3]. The hardware One TPU v7 (Ironwood) host has four chips. Each chip packs two chiplets, each with one TensorCore and its own HBM, and frameworks expose each chiplet as a device, so a host has 8 devices. We shard the model 8 ways with tensor parallelism [4], one shard per device. Figure 2. One TPU v7 chip, simplified from the architecture diagram in Google Cloud’s TPU7x documentation. The numbers are Google’s published figures. Each device is a single-threaded processor: one program runs at a time, so speed comes from keeping the matrix units fed and overlapping memory transfers with compute inside that program. These are the units that matter in this post, as described in Google Cloud’s TPU7x documentation and JAX’s Pallas TPU guide: Scalar unit. One per device, with its own small scalar memory (SMEM). It runs the control flow: loops, branches and address math. MXU, the matrix unit. Two per device. They do the matmuls; the chip’s peak is 2,307 TFLOPs in bf16 and 4,614 in fp8. At small batch they mostly wait for weights to arrive from HBM. Vector unit (VPU). Elementwise work such as norms, softmax exponentials and format conversions. On recent TPUs it computes on 2D vector registers of 8 × 128 32-bit values, 8 sublanes by 128 lanes, so a data layout that leaves lanes empty wastes both vector work and memory (section 7). A cross-lane unit (XLU) next to it does transposes and permutes. VMEM. On-chip SRAM next to the MXUs and the vector unit, with much higher bandwidth than HBM. Kernels stage their tiles here, so its size bounds how much a kernel can keep in flight. DMA engines. They move data between HBM, VMEM and other devices in the background: a kernel starts a copy and waits for it later, while it computes. The megakernel’s pipelining (section 2) is built on them. HBM. 192 GiB per chip at 7.37 TB/s, 96 GiB of it on each device. Die-to-die link. It joins the two devices of a chip; Google documents it as six times faster than a 1D ICI link. Section 8’s pair exchange leans on it. ICI. Six links per chip at 200 GB/s each, counting both directions, for 1,200 GB/s per chip. The four chips of a host talk to each other over it; slices of 64 chips or more form 3D tori. Host. Each chip connects to the host’s CPUs over PCIe Gen5 ×16. A 4-chip host has 960 GB of DRAM, which L2 HiCache uses as a second cache tier (section 7). SparseCores. Four per chip. This work does not use them. The TorchTPU post has a longer primer. The software The server is SGLang-TorchTPU: SGLang’s [5] scheduler, KV cache and model code in PyTorch, with Pallas kernels on the hot paths. We decode with speculative decoding [6, 7]: a small draft model guesses the next few tokens, and the full model checks them all in one forward pass, the verify step. Each (request, position) pair in a verify step is a row, so two requests verifying 8 tokens each make a 16-row step. 1. A megakernel At small batch, generating a token does little math per weight, so the step is as fast as HBM can stream the weights, as long as nothing interrupts the stream. A decode or verify step should therefore be one uninterrupted stream of weights from HBM into the MXU. Our previous path (“native” below) ran each layer as several kernels plus collectives, and every boundary stopped the stream. The megakernel runs all 62 layers of a step as one Pallas program per device, so the next weights are always in flight. Whole-step megakernels have been built for GPUs [8, 9] and, on TPU v7, for Kimi K3 and Qwen 3.8 [10]. Ours runs inside SGLang’s serving path, on its paged KV cache, for the decode and DFlash2 verify steps of many requests at once. The native path was already built from fused Pallas blocks, yet its verify step took 2.9× the time needed to stream its weights. Profiling our first one-kernel prototype showed why: at batch 1 the matmuls and attention cost almost nothing, while collectives cost about 45 µs per layer. Nothing overlaps across a program boundary, so the step was bound by overhead, not by bandwidth. Figure 3. One step. Inside the megakernel every layer runs the same flow (white); the shaded boxes are data movement the kernel issues itself. For each step, every device runs one pallas_call that covers all 62 layers and the final norm. SGLang still runs the embedding, the draft model, a small KV-write kernel, the LM head and sampling around it. Figure 4. Time per verify step, both paths with DFlash2 speculative decoding (GSM8K, measured on an earlier Muse Spark checkpoint). The floor is the time to stream that step’s distinct experts and dense weights from HBM. Per output token, the megakernel is 1.94× faster than the native path for a single request (1.95 ms against 3.78 ms) and 1.85× for two concurrent requests. The next sections show what made one program this fast. 2. Pipelining everything inside the megakernel A step is mostly data movement: weights, KV cache and partial sums between devices. Once the whole step is one program, the kernel can schedule every transfer so that it lands just before the compute that needs it [11], including across layer boundaries [10]. Figure 5 sketches one layer. Figure 5. One layer of a step, schematic. Each DMA stream runs ahead of, or alongside, the compute that consumes it. Weights arrive before they are needed A ring for dense weights. Attention, o-proj and router weights stream through a ring of tiles that prefetches across layer boundaries. About ten tiles in flight sustain 3.34 TB/s per device, while a single DMA stream reaches only 1.13 TB/s. Experts as soon as routing finishes. The router’s choice is known only mid-layer, so the kernel requests each chosen expert’s weights the moment routing ends, with three experts in flight. The next layer during the expert loop. Refills for the next layer’s attention weights are held until the expert loop runs, so they never compete with the expert weights the current layer is waiting for. Collectives leave as soon as their data exists Tensor parallelism exchanges partial sums between the 8 devices several times per layer. Instead of a separate collective after each matmul, the kernel writes partial results straight into the owning device’s memory with remote DMAs issued from inside the kernel; a small round costs about 3 µs. Streamed o-proj. Each o-proj output tile leaves for its owner as soon as it is computed, while the next tile is still being multiplied. Pipelined KV writes Each step writes its new K and V rows into the cache. In our fp8 sequence-on-lane cache (section 7), a token is one lane column of a 256-token page, so writing a token means reading the whole page block into VMEM, merging the new columns, and writing the page back. This runs as a small kernel right after the megakernel, because XLA copies the aliased operands of any kernel that issues remote DMA. It used to do every page in turn: read, wait, merge, write, wait, for each request and each of 62 layers. Figure 6. Schematic. The pipelined kernel issues all of a layer’s page reads at once, merges each page as it lands, and starts its write; the next layer’s reads start immediately on a second VMEM bank. The new kernel loads all layers’ new rows once, then keeps two VMEM banks: layer l reads into one while layer l−1’s writes drain from the other, and a bank is reused only after the writes of two layers earlier have landed. The write cost per step fell from 810 to 155 µs at 16 rows and from 2,286 to 336 µs at 64 rows, with bit-identical results. 3. Collective folding Because each layer is split across 8 devices, a device often has to collect the others’ results before it can continue. After the MoE and after attention, each device holds one slice of the layer’s output along the hidden dimension, and must all-gather the other seven slices before the next matvec can run. Collective folding lets that matvec start on the slices as they arrive instead of waiting for the whole gather, as decomposed collective matmuls do [12], here across the RMS norm that sits between the two. Figure 7. Schematic. The folded matvec multiplies each gathered slice as soon as it lands, starting with the device’s own, and applies the per-row normalization at the end. What makes this possible is where Muse Spark normalizes. The next layer’s attention input is an RMS norm [13] of the residual stream α·x + β·s·f, where x is the residual this device already holds, f is the gathered MoE output and s is f’s own norm factor. The two norm factors, s and the final r, are each one scalar per row. Scalars per row commute with the matmul, so the projection splits into a part the device can compute right away and a part that adds up slice by slice: Here g is the norm’s gain vector, f_d is the slice from device d, and W_d is the matching block of weight rows. The kernel stacks the local part and the arriving slice into one product per block, consumes the slices in a fixed peer order, and reorders the weight ring so that W’s blocks stream in the same order. Once the last slice is in, it computes s and r and scales the result. The kernel folds two all-gathers per layer this way: the MoE output into the next layer’s qkv projection, and the attention output into the pooler and router. The residual stream itself is unchanged; only the point where the matmul input is rounded to bf16 moves, so results are not bit-identical, and our accuracy checks cover it. 4. Enabling better pipelining in attention At long context, attention is the largest part of a step. The kernel walks the KV cache chunk by chunk, and the key bottleneck was the running max of online softmax [14]: it chained every chunk to the one before, so chunks could not be pipelined. The running max chains every chunk Online softmax, as used by FlashAttention [15], keeps the maximum score seen so far. Each new chunk computes its scores, updates that maximum, exponentiates against it, and rescales everything accumulated so far whenever the maximum grows. So each chunk is a serial chain (wait for its DMA, QK, max, exp, sum, P·V, rescale) that cannot start its exponentials until the previous chunk’s maximum is known. In a microbenchmark this chain costs about 0.17 µs per chunk, whether the chunk is fp8 or bf16. Muse Spark makes the max unnecessary Muse Spark normalizes q and k (QK-norm [16, 17]) with an RMS norm that has no learned weight, and gives q a fixed gain. No attention score can exceed 43.8 in magnitude, far from where exp overflows in f32. The running max exists only to prevent that overflow, so the kernel drops it, along with the rescale. It checks the bound from the model’s constants when it is built. FlashDecoding++ [18] breaks the same chain on GPUs with a unified max value and a recomputation path for the rare overflow; here the bound is exact, so the kernel needs no fallback. Figure 8. Schematic. With a running max, each chunk’s chain waits on the previous chunk. Without it, chunks are independent and their stages overlap. Independent chunks pipeline Without the max, chunks no longer depend on each other. The kernel issues all DMA waits at the top of each loop iteration, splits the chunks across several independent softmax streams that merge at the end, as Flash-Decoding does across the KV sequence [19], and software-pipelines the loop [11]: it scores the next chunks while it multiplies the previous chunks’ probabilities by V. Three smaller changes shorten each chain further: Chunk size by layer type. Global layers read the cache in long chunks of 1,024 to 2,048 tokens, so a 100K context needs far fewer chains. Sliding-window layers use 512-token chunks, since their window spans only a few. fp8 keys straight into the MXU. Widening each fp8 chunk to bf16 on the vector unit cost more vector work than the MXU work it fed. Mixed bf16 × fp8 dots are exact, so the keys now go in as fp8. The attention scale in the queries. Muse Spark’s attention scale is 1/8, a power of two, so multiplying the queries once instead of every score is exact. Figure 9. Microbenchmark at production shapes. The memory floor is the time to read the KV bytes. In the microbenchmark, longer chunks with independent streams cut attention time 2.4×, and dropping the running max took it to 3.9×. In the server, the pipelined loop raised per-request decode speed by 18% for two requests at 500K context. 5. Keeping register pressure off the common path A processor computes on values held in its registers, a small, fixed set of storage slots next to the arithmetic units. When code needs more live values than the registers hold, the compiler spills some of them to memory and loads them back later [20], and those extra stores and loads cost time on every pass through the code. In a loop with several branches, the compiler has to fit the live values of all branches into the same registers, so one heavy branch can make every pass through the loop spill, even when that branch never runs. The megakernel’s expert loop had exactly this problem. A step loops over the experts its rows were routed to. For each one, the kernel waits for the expert’s weights and runs a path chosen by how many rows the expert serves: a few-row path, with small variants for one to eight rows, and a dequantize-then-multiply path for busier experts. As the batch grows, a step visits more experts, but most of them still serve only one or two rows, so the light path is the one that matters. One heavy branch slows the light path The 8-row variant keeps large f32 accumulators that do not fit in the vector registers even on their own. Because it shared one loop body with the one- and two-row variants, the common path spilled registers to memory and loaded them back on every expert: about 0.37 µs per expert, against roughly 1.2 µs to stream that expert’s weights. That overhead had nothing to do with the work the light path itself did. Two loops instead of one The fix keeps the heavy code out of the light path’s loop. The kernel sorts a step’s experts by row count and runs two loops: one for experts with at most four rows, and one for the rest. Each loop body now holds only its own variants, so the few-row loop carries 61–64% fewer spill instructions. Every expert still takes the same path as before, so the output is bit-identical. Figure 10. Schematic. In one shared loop body, every few-row expert is followed by a spill and refill, and the loop falls behind its weight stream. In their own loop, few-row experts keep pace with the weights. In a kernel test with random routing, decode steps with many requests got 13–19% faster, and the expert loop now keeps pace with its weight stream. 6. Speculative decoding with DFlash2 Without speculation, every output token costs one full pass over the active weights. DFlash [21] drafts a whole block of tokens in one forward pass of a small model; we use Muse Spark’s DFlash2 draft model. Figure 11. One speculative step. The draft model reads hidden states captured from the full model. The token strip is an illustration. The draft attends bidirectionally within its block and proposes all of the block’s tokens at once. The full model checks the whole block in one megakernel step and keeps the longest prefix of drafts that matches its own choices, plus one token of its own. This suits TPU well: MXU time grows with the number of row and output tiles, not with multiply-adds, and a block’s few rows fit in the same row tile as one token, so verifying a block costs little more than verifying one token. Measured on GSM8K with real acceptance, DFlash2 cut the megakernel’s time per output token for a single request from 3.93 to 1.95 ms, 2.0×. 7. More context per device: KV layout and L2 HiCache On agent traffic, KV capacity is throughput: every prefix that falls out of the cache comes back as a 100K-token re-prefill. After the weights, each device has only about 20 GiB left for KV, split between a full-attention pool for the 16 global layers and a sliding-window pool for the other 46. fp8 KV in a sequence-on-lane layout Each device holds 2 of the 16 KV heads, each 64 values wide: half of a vector register’s 128 lanes. A token-major cache packs a head’s K and V side by side to fill the lanes. That works in bf16, but fp8 [22] packs four values per 32-bit word down the sublanes, so a K-and-V pair is padded to four slots and switching to fp8 alone saves nothing. The sequence-on-lane layout puts a page’s tokens along the 128 lanes and the head dimension down the sublanes, so each word holds four fp8 values of one token and nothing is padded. It halves both the KV memory and the KV bytes each step reads. Google engineers added this layout to SGLang-TorchTPU’s ragged paged attention kernel (#344), which serves prefill here, and the TorchTPU post describes it for Qwen 3.5. We rewrote the megakernel’s attention reads and its KV writes for it. The draft model’s KV in the sliding-window pool The draft model only attends to the last 2,048 tokens, but the server allocated its KV for the full context, almost a third of all KV memory per cached token. Serving it from the sliding-window pool lets the same HBM hold 38% more tokens. L2 HiCache on TPU We enabled SGLang’s hierarchical cache (HiCache) [23] on TPU, with host DRAM as an L2 behind every HBM pool: full-attention, sliding-window and draft. When HBM evicts a session’s prefix during a tool pause, its pages stay in host memory and are copied back on the next turn instead of being recomputed, as in multi-turn KV caching systems such as CachedAttention [24]. Making it work alongside speculative decoding and the overlap scheduler needed a new synchronization fence. On agent traffic, the L2 cache raised throughput by up to 77%. Session-aware eviction Agent traffic brought a second problem: the cache did not know which old prefixes would be needed again. An agent session is a tree of conversations: a main agent can sit idle for a quarter of an hour while its sub-agents run, then come back with its whole context. To plain LRU eviction, that idle conversation looks like the coldest data in the cache, so it goes first, and its next turn recomputes a long prefix. A hybrid model makes the loss worse: reusing a prefix needs both its full-attention KV and the sliding-window KV of its last 2,048 tokens, so losing only the small window makes the large full-attention KV useless. On this trace, prefill computed about 1.2× the tokens that the best possible cache would. Figure 12. Schematic. A main agent idles while its sub-agents run. Plain LRU evicts its KV first; session-aware eviction keeps it in host memory until everything else has gone. SGLang’s session-aware radix cache [5] (sgl-project/sglang#29173) gives the cache that knowledge. Requests from the same conversation share a session id, and KV that a live session will reuse is evicted only after everything else. We adapted it to agent traffic: the session id comes from the per-conversation id the client already sends, and the protection applies only in host memory, so HBM stays plain LRU for the requests that are running while host memory holds the idle conversations. In the highest-throughput run of Figure 1, prefill now computes within 1–5% of what the best possible cache would. 8. Prefill: a fused MoE kernel and fair batching Before a request can start decoding, the server runs its new prompt tokens through the model in chunks of 4,096 tokens [25, 26]: the prefill. A chunk has enough rows to keep the matrix unit busy, so unlike decode, prefill time goes to compute and to moving results between devices. On agent traffic a single turn can bring hundreds of thousands of new tokens, so prefill speed decides how long requests wait before their first token. In the MoE layers of prefill, every device computes a partial result for every (token, expert) pair, and each token’s owner must add up 8 of them. Moving those partials, one small write at a time, cost more than the matrix math. The fused MoE kernel and the pair exchange Experts are sharded with tensor parallelism: each device holds 1/8 of every expert. Each device therefore runs every (token, expert) row of a prefill chunk against its slice, producing a full-width partial per row; Muse Spark normalizes each expert’s output before the weighted sum, so a row’s 8 partials must all be added up before its owner can apply the norm and combine experts. The kernel streams one expert’s weights while the next prefetches, runs gate/up, SwiGLU [27] and down, and writes each finished row straight into its owner’s HBM, overlapping later tiles. Figure 13. The two devices on a chip share a die-to-die link, which Google documents as six times faster than a 1D ICI link to another chip. Each row’s write took about 50 ns of scalar work on the same unit that issues the matmuls, so writes serialized with compute. The pair exchange uses the chip’s topology instead. Each device sends its whole tile to its sibling on the same chip over the fast on-chip link, and the pair’s sum then goes to each remote owner once, as one DMA per contiguous run of rows. Cross-chip bytes halve, each owner adds 5 blocks instead of 8, and the MoE kernel went from 3,247 to 2,368 µs per layer. Smaller changes trimmed the rest of the path: Weight ping-pong. Alternating two weight buffers lets each expert’s fetch start while the previous expert computes. A one-pass combine. The owner’s sum of partials, per-expert norm and weighted sum ran as three XLA passes; one fused Pallas pass reads the partials once. No zero fill, unrolled writes. Every output row is written exactly once, so the zero fill of the output buffer was dropped, and the per-row DMA loop issues 8 rows per iteration. Together with these smaller changes, the pair exchange cut a 4,096-token prefill chunk from 440 to 384 ms on a cold 100K-token prompt. Sequence-parallel prefill Tensor parallelism splits every matmul across the 8 devices, but each device still holds the activations of the whole chunk, so the per-token work between the matmuls, such as norms, residual updates and routing, is repeated on every device. Sequence parallelism [28] splits the tokens instead: between layers, each device keeps only its own 512 of the chunk’s 4,096 tokens. The norms and residual updates then run on an eighth of the tokens, and attention’s all-reduce becomes a reduce-scatter, which delivers each device only its own tokens’ sums. Routing packs each token’s expert ids and weights into one small all-gather. The MoE output reaches its tokens’ devices through an all-to-all: each device receives only its own tokens’ columns from the other seven, an eighth of the bytes of the all-gather it replaces, which sent every device the columns of all 4,096 tokens. This cut a 4,096-token forward at 131K context by 51 ms, about 11%. It changes the order of some reductions, so results are not bit-identical, and our accuracy checks cover it. Chunk sharing A cold prompt of hundreds of thousands of tokens is prefilled in 4,096-token chunks of about half a second each. The scheduler used to give every chunk to that prompt until it finished, so other turns’ few thousand new tokens waited 15 to 95 s behind it. Figure 14. Schematic. The long prompt keeps at least 1,024 tokens of each batch, waiting turns fill the rest, and the long prompt finishes a little later. Now a prompt in the middle of a chunked prefill keeps at least 1,024 tokens of each batch, and waiting turns fill the rest. In an AgentX test, this cut median end-to-end latency by 58% and p90 by 33%, with throughput within 5%. Future work Attention closer to the memory floor, by packing the 64-wide heads to fill the MXU and shortening the per-chunk chain further. A topology-aware MoE layout, with expert parallelism across chips and tensor parallelism inside each chip, to make most of the MoE exchange local. Faster prefill, now the limit at high concurrency, through larger prefill graphs, collectives overlapped with compute and a reworked MoE prefill kernel. Prefill/decode (P/D) disaggregation, optimizing AgentX performance with prefill and decode on separate hosts. Upstreaming this work into SGLang-TorchTPU’s main branch. Everything here is open source in SGLang-TorchTPU, branch haolin/ms-v11. The megakernel lives in sglang_tpu/muse_megakernel/, the fused MoE prefill kernel in sglang_tpu/moe/pallas/fused_rs/. Acknowledgments We thank Google and Meta for their kernel contributions to SGLang-TorchTPU, which this work builds on, and give special thanks to Qi Zhou, Yunlong Liu, Charles Chen, Yin Lin, Jiaxin Cao and Anne Stern for their work on the kernels and AgentX optimization described here. We also thank SemiAnalysis, Google, Meta and the SGLang community for InferenceX, TPU v7 and TorchTPU, Muse Spark Open, and SGLang. References N. Shazeer, A. Mirhoseini, K. Maziarz, et al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. arXiv:1701.06538 I. Beltagy, M. E. Peters, A. Cohan. Longformer: The Long-Document Transformer. 2020. arXiv:2004.05150 B. D. Rouhani, R. Zhao, A. More, et al. Microscaling Data Formats for Deep Learning. 2023. arXiv:2310.10537 M. Shoeybi, M. Patwary, R. Puri, et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. arXiv:1909.08053 L. Zheng, L. Yin, Z. Xie, et al. SGLang: Efficient Execution of Structured Language Model Programs. NeurIPS 2024. arXiv:2312.07104 Y. Leviathan, M. Kalman, Y. Matias. Fast Inference from Transformers via Speculative Decoding. ICML 2023. arXiv:2211.17192 C. Chen, S. Borgeaud, G. Irving, et al. Accelerating Large Language Model Decoding with Speculative Sampling. 2023. arXiv:2302.01318 B. Spector, J. Juravsky, et al. Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B. Hazy Research, 2025. Blog post X. Cheng, Z. Zhang, Y. Zhou, et al. MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs. OSDI 2026. Paper G. Novack, X. Liu, J. Ma, W. Kwon. TPU megakernels for Kimi K3 and Qwen3.8-27B. 2026. Code M. Lam. Software Pipelining: An Effective Scheduling Technique for VLIW Machines. PLDI 1988. doi:10.1145/53990.54022 S. Wang, J. Wei, A. Sabne, et al. Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models. ASPLOS 2023. doi:10.1145/3567955.3567959 B. Zhang, R. Sennrich. Root Mean Square Layer Normalization. NeurIPS 2019. arXiv:1910.07467 M. Milakov, N. Gimelshein. Online Normalizer Calculation for Softmax. 2018. arXiv:1805.02867 T. Dao, D. Y. Fu, S. Ermon, et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135 A. Henry, P. R. Dachapally, S. Pawar, et al. Query-Key Normalization for Transformers. Findings of EMNLP 2020. arXiv:2010.04245 M. Dehghani, J. Djolonga, B. Mustafa, et al. Scaling Vision Transformers to 22 Billion Parameters. ICML 2023. arXiv:2302.05442 K. Hong, G. Dai, J. Xu, et al. FlashDecoding++: Faster Large Language Model Inference on GPUs. MLSys 2024. arXiv:2311.01282 T. Dao, D. Haziza, F. Massa, G. Sizov. Flash-Decoding for Long-Context Inference. Stanford CRFM, 2023. Blog post G. J. Chaitin. Register Allocation & Spilling via Graph Coloring. SIGPLAN Symposium on Compiler Construction, 1982. doi:10.1145/800230.806984 J. Chen, Y. Liang, Z. Liu, et al. DFlash: Block Diffusion for Flash Speculative Decoding. 2026. arXiv:2602.06036 P. Micikevicius, D. Stosic, N. Burgess, et al. FP8 Formats for Deep Learning. 2022. arXiv:2209.05433 SGLang Team. SGLang HiCache: Fast Hierarchical KV Caching with Your Favorite Storage Backends. LMSYS Org, 2025. Blog post B. Gao, Z. He, P. Sharma, et al. Cost-Efficient Large Language Model Serving for Multi-turn Conversations with CachedAttention. USENIX ATC 2024. arXiv:2403.19708 A. Agrawal, A. Panwar, J. Mohan, et al. SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills. 2023. arXiv:2308.16369 A. Agrawal, N. Kedia, A. Panwar, et al. Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve. OSDI 2024. arXiv:2403.02310 N. Shazeer. GLU Variants Improve Transformer. 2020. arXiv:2002.05202 V. Korthikanti, J. Casper, S. Lym, et al. Reducing Activation Recomputation in Large Transformer Models. MLSys 2023. arXiv:2205.05198