700 tok/s on a 2.8 Trillion Model: The 92-Layer TPU Megakernel

Описание: Generating tokens on a 2.8-trillion parameter model for a single user usually crawls at memory bandwidth limits[cite: 10]. Every single token requires streaming roughly 208 GB of active weights across high-bandwidth memory (HBM), and stock inference engines chop that decode step into hundreds of separate kernel launches—leaving massive execution bubbles where the memory bus sits idle[cite: 10].

In this episode of Skildcore, we dissect the landmark megakernel release from the Inferact team on Google TPU v7 (Ironwood)[cite: 10]. By fusing an entire 92-layer forward pass into a single persistent Pallas program, the machine pre-fetches layer n+1 weights while layer n computes, amortizing memory traffic with speculative decoding to hit a staggering 707 tokens per second on Kimi K3—a 4.4x leap over stock vLLM[cite: 10].

Here is how software-managed vector memory (VMEM), XLA compiler contracts, and cross-chip interconnects make 92 layers in one kernel possible[cite: 10].

---
TIMESTAMPS:
0:00 - The 700 tok/s Reality: 2.8 Trillion Parameters, Single Stream[cite: 10]
0:40 - The Baseline: Why LLM Decode is a Pure Memory Bandwidth Problem[cite: 10]
1:20 - Stock Engines: Hundreds of Launches, Hundreds of Idle Bubbles[cite: 10]
2:08 - The NVIDIA Stack: CUDA Graphs, Streams, and the Shared Ceiling[cite: 10]
2:45 - The Megakernel Concept: One Program for the Whole Decode Step[cite: 10]
3:26 - Part I: The Machine (TPU v7 Ironwood vs. NVIDIA GB200)[cite: 10]
4:04 - 64 MiB VMEM: The Cache You Command (Software-Managed SRAM)[cite: 10]
4:22 - How It Is Written: Pallas (Python That Speaks TPU)[cite: 10]
5:02 - The Compiler Contract: You Place the Moves, XLA Plays Them[cite: 10]
5:37 - Part II: Inside the Kernel (Kimi K3 Across 8 TPU v7 Chips)[cite: 10]
6:30 - The Choreography: Layer n+1 Loads While Layer n Computes[cite: 10]
7:20 - Part III: The Honest Number (Speculative Decoding Breaks the Floor)[cite: 10]
8:20 - The Scoreboard: 161 tok/s Stock vs 707 tok/s Megakernel[cite: 10]
9:08 - What It Costs: Static Shapes, Cold Compiles, and Tradeoffs[cite: 10]
9:40 - Architectural Verdict: The Megakernel Did Not Break Physics[cite: 10]

---
KEY TECHNICAL TAKEAWAYS:
• The memory wall: Decoding Kimi K3 activates 104B parameters per token, demanding ~208 GB of memory traffic per step[cite: 10]. Even across 7,380 GB/s HBM, standard serving hits a hardware floor of ~28 ms per token[cite: 10].
• Deleting the launch tax: In standard engines, kernels constantly ramp up and drain[cite: 10]. In a 92-layer megakernel, layer n+1 weights pre-fetch over the memory bus while layer n tensor cores crunch math, eliminating memory silence[cite: 10].
• GPU hardware caches vs. TPU VMEM: While GPUs rely on hardware-managed L1/L2 caches and thread blocks, TPU v7 provides 64 MiB of software-managed Vector Memory (VMEM) per tensor core[cite: 10]. You don't ask the cache for a favor; your program dictates exactly what stays in SRAM[cite: 10].
• Pallas and XLA: Compiling a 92-layer model usually takes 30+ minutes in XLA[cite: 10]. Pallas bypasses automatic fusion search by explicitly encoding buffer lifetimes and direct memory copies, compiling in under 90 seconds[cite: 10].
• Speculative amortization: Combining persistent execution with DSpark speculative decoding amortizes weight reads over multiple verified tokens, breaking the single-token memory ceiling[cite: 10].

CITATIONS & OPEN SOURCE CODE:
• Inferact Team: TPU Megakernels for Kimi K3 (github.com/inferact/tpu-megakernels)[cite: 10]
• Hardware: Google TPU v7 Ironwood (2.31 PF BF16, 7,380 GB/s HBM, 1,200 GB/s ICI)[cite: 10]
• Models: Moonshot AI Kimi K3 (2.8T total, 104B active, 92 layers MoE)[cite: 10]

#TPU #GoogleCloud #LLMInference #MachineLearning #KimiK3 #Megakernel #Pallas #GPUArchitecture #DeepSystems #Skildcore