Inferact / tpu-megakernels

# TPU megakernels for Kimi K3, Qwen3.8-27B, and Muse-Spark-Open
Companion code for the blog post
700 TPS on Kimi K3: A Case for TPU Megakernels.
This repository contains fused decode implementations for three models:
- **Kimi K3 + DSpark** runs on 32 TPU devices across four hosts. One megakernel
 executes the KDA, MLA, MoE, residual, normalization, and collective work for
 the target model; DSpark provides fused speculative decoding with one anchor
 and seven draft rows.
- **Qwen3.8-27B + DFlash2** runs on eight TPU devices on one host. Its dense
 BF16 megakernel covers the hybrid gated-delta-net and grouped-query-attention
 stack, with an optional fused DFlash2 speculative-decoding loop.
- **Muse-Spark-Open + DSpark** runs the NVFP4 model on 8, 16, or 32 TPU devices across one, two, or four hosts, with optional speculative decoding and non-speculative batching.
The demos support terminal interaction and an OpenAI-compatible HTTP server.
## Results
With speculative decoding at an acceptance length of 6, the megakernels reach
709 accepted tokens per second on Kimi K3 (16× TPU v7 vs 16× GB200, 1.57×) and
1,515 on Qwen3.8-27B (4× TPU v7 vs 4× GB200, 2.18×):
DSpark keeps a large lead at shorter acceptance lengths as well:
Without speculative decoding, the megakernels deliver roughly 1.4–2× the
decode throughput of the GB200 baseline at batch sizes 1 through 8:
See the blog post for the full
analysis and measurement setup.
## Setup
Create the environment from the repository root:
 [code] 
The launchers run Hugging Face in offline mode, so the model snapshots must
already be present in the cache selected by `HF_HOME`, or supplied as explicit
local paths. All scripts automatically load an optional, untracked `.env` file
from the repository root. A typical configuration is:
 [code] 
The Kimi pre-sharded variables are optional. Without them, the demo prepares
its kernel layouts from `moonshotai/Kimi-K3` and
`RedHatAI/Kimi-K3-speculator.dspark`. The Qwen demo loads
`Qwen/Qwen3.8-27B` and `z-lab/Qwen3.8-27B-DFlash2`; packed weight containers
under `checkpoints/qwen38-tp8` (or `$QWEN_PACKED_WEIGHTS` when set) are used
instead when present. The optional `KIMI_CHECKPOINT_REVISION`,
`KIMI_DRAFT_REVISION`, `QWEN_CHECKPOINT_REVISION`, and
`QWEN_DRAFT_REVISION` variables pin the commits resolved from the `HF_HOME`
cache for the default repositories. Explicit `--checkpoint`,
`--draft`, `--weights`, and `--draft-weights` arguments override these
defaults.
Muse-Spark-Open loads `meta-models/Muse-Spark-Open-816B-A42B-NVFP4` and `Inferact/Muse-Spark-Open-DSpark`. Set the matching `MPK_PACKED_WEIGHTS_TP8`, `MPK_PACKED_WEIGHTS_TP16`, or `MPK_PACKED_WEIGHTS_TP32` to an existing target pack, or pass `--weights`. To create a pack from an existing checkpoint on a high-memory compute node:
 [code] 
Use `--tp 8`, `16`, or `32` and a separate output directory on model-weight storage for each. The DSpark draft loads directly from its checkpoint. `MPK_CHECKPOINT_REVISION` and `MP