Running DeepSeek V4 Flash on TPU v6

DeepSeek-V4-Flash has 256 routed experts (top-6) stored in MXFP4, with everything else in FP8 128×128 blocks: about 160 GB. vLLM’s TPU backend, tpu-inference, already had an experimental path for it, with Pallas kernels for sparse MLA, the KV compressor, the lightning indexer and mHC. But it had only been tested on TPU v7x (Ironwood) at TP=8 with prefix caching off. We brought it up on TPU v6e (Trillium) at TP=1, 2 and 4 with prefix caching on. With the kernel, weight-loading and scheduler fixes in place, the TPU matches DeepSeek’s fp32 reference as closely as the reference matches itself in bf16.

1. Four bits that v6e can’t hold

The MoE experts run through gmm_v2, tokamax’s grouped-matmul Pallas kernel. On v7x FP4 is a native type. On v6e the kernel doesn’t compile:

Unsupported type 'vector<8x128x8xf4E2M1FN>'

Mosaic, the compiler behind Pallas on TPU, can’t hold an f4 vector on v6e at all. Casting fp4 to bf16, f32 or fp8 all fail. But uint4 → int32 works, and decoding the bits by hand is bit-exact.

So the weights stay 4-bit and the kernel stops calling them floats. Once, when the model loads, each e2m1 weight is re-encoded as a signed 4-bit integer q = ±(2e + m) and stored as uint4: still 4 bits per weight in HBM.

The kernel’s loops don’t change. Tokens are sorted by expert, so the activations form one band of rows per expert. The outer loops walk N tiles and (expert, M tile) pairs; the inner loop walks K, streaming BK×BN weight tiles into VMEM and accumulating into an f32 tile that stays in VMEM. The one new step: each weight tile lands in VMEM as 4-bit codes, and we decode it to bf16 before anything else touches it.

Here is the kernel body as a schematic, without the pipelining and masking details, and the change we made to it. The weight tile is decoded right after it is pulled from HBM into VMEM, before the scales and the MXU see it:

 def gmm_kernel(lhs_hbm, rhs_hbm, scales_hbm, out_hbm, group_sizes):
     for n in range(num_n_tiles):                    # outer: N tiles
         for e, m in expert_m_tiles(group_sizes):    # outer: (expert, M)
             acc = zeros((BM, BN), f32)              # lives in VMEM
             for k in range(num_k_tiles):            # inner: walk K
                 a = load_tile(lhs_hbm, m, k)        # bf16 tokens
                 w = load_tile(rhs_hbm[e], k, n)     # 4-bit codes
+                w = decode_codes(w, bf16)           # -> bf16 in VMEM
                 w = w * load_tile(scales_hbm[e], k, n)
                 acc += dot(a, w)                    # bf16 MXU
             store_tile(out_hbm, acc, m, n)
+
+def decode_codes(codes, dtype):
+    q = bitcast(codes, int4).astype(bf16)           # q = +/-(2e + m)
+    v = 2 * q - clip(q, -6, 6) - 0.5 * clip(q, -4, 4)
+    return v.astype(dtype)                          # exact e2m1 value

Why that encoding? The eight e2m1 magnitudes 0, 0.5, 1, 1.5, 2, 3, 4, 6 are a piecewise-linear function of q = 2e + m: slope ½ up to 4, 1 up to 6, then 2. Two clips rebuild the kinks, so the decode is three bf16 vector ops, and every intermediate is a small multiple of ½, exact in bf16.

The existing dequantization paths then apply the per-32 block scales to the decoded tile, before the matmul. We also stopped quantizing the activations to fp8 on these chips: v6e reports the same fp8 throughput as bf16, so it only cost accuracy. The experts are now effectively W4A16 on the bf16 MXU.

Against an f32 reference with ragged groups, including empty ones, the kernel’s maximum relative error is 1.4e-7 with block 32 and 9.2e-8 with block 512.

Grouped matmul performance vs bf16

We benchmarked both the re-encoded MXFP4 kernel and a NVFP4 kernel that reads packed checkpoint bytes directly, decoding e2m1 weights and per-16 E4M3 scales in VMEM. Both use bf16 activations and bf16 MXU operations. On one TPU v6e chip, with 256 experts and random top-6 routing:

These are grouped-matmul timings, without the gate/up activation: w13 has K=N=4096; w2 has K=2048, N=4096. NVFP4 uses the best measured k-major tiles; MXFP4 and bf16 use gmm_v2’s default tiling. Timings exclude NVFP4’s activation permutation and one-time MXFP4 weight encoding. We use repeated JIT loops to reduce host overhead; for the noisy bf16 baseline we take the smaller of that estimate and the single-call time minus measured host overhead.

MXFP4 is about 1.5–2.6× faster than bf16 at the two decode sizes shown. Direct NVFP4 avoids load-time re-encoding, but its more expensive decoder makes it slower than MXFP4 at decode, and it does not always beat bf16. The 4,096-token w13 regression is a default-tiling problem: a separate fused-activation sweep reduced MXFP4 time from 10.1 ms to 3.2 ms by raising bucket_base from 64 to 128.

2. A prefix-caching crash

Every multi-chip config died at startup with prefix caching on (vLLM’s default). On more than one chip DSV4 uses attention data parallelism, whose DPScheduler forwarded hash_block_size only if it appeared in inspect.signature(AsyncScheduler). AsyncScheduler.__init__(*args, **kwargs) hides it, so it silently fell back to 1024, and vLLM’s hybrid KV coordinator asserts that every group’s block size (128, 128, 1024, 32, 16, 32) divides by it. Walking the MRO fixes it.

3. Results

We cut the checkpoint to its first five layers (sliding window, hash-routed CSA, HCA and CSA), a ~20 GB model that is bit-identical on TPU and GPU. The golden reference is DeepSeek’s own inference/model.py in float32 on CPU, with pure PyTorch in place of its TileLang kernels. A truncated model is rounding-sensitive: the reference in bf16 is itself ~0.13 mean |Δlogprob| from fp32, and that is the noise floor. Teacher-forced on a 7,600-token prompt:

The “after” row holds at TP=1, 2 and 4, with and without expert parallelism, on six prompts of 17 to 7,600 tokens; the error doesn’t grow with context. TP=2 and 4 agree with TP=1 on 98.5–99% of top-1 tokens (mean |Δlogprob| ~0.03), which is reduction-order noise.

• Decode: decode-time vs prefill-time logprobs of the same tokens differ by 0.038 on TPU (0.097 on H100), with no spikes at CSA/HCA compression boundaries.
• Prefix caching: the TPU is bit-deterministic for identical batches, so cache-on vs cache-off can be compared exactly: 21/21 requests bit-identical at TP=1, 2 and 4, with hits of 1,024 to 7,168 tokens.

Serving throughput

Coming soon: agentic-serving throughput on a 32-chip v6e slice.

If you try this yourself

• v6e VMs created through GCE need LIBTPU_INIT_ARGS=--enable_tpunetd_client=false.
• A 2-chip subslice needs TPU_CHIPS_PER_PROCESS_BOUNDS=1,2,1.