Making GLM-5.2 Sparse on TPUs

Sep 10, 2026TPU

A step-by-step look at how SGLang-JAX converted GLM-5.2's learned sparsity into real reductions in memory traffic, prefill time, and serving overhead on TPU.

Background

SemiAnalysis recently published a detailed overview of Google’s TPU inference stack (https://newsletter.semianalysis.com/p/tpu-inferencex-full-steam) , focusing on the work required to make TPUs first-class citizens for inference serving outside Google’s internal ecosystem. The article follows a series of optimizations in TPU-Inference, the public vLLM TPU backend, aimed at efficiently serving Qwen3.5-397B on TPU v7 Ironwood.

This write-up provides a complementary view from the SGLang ecosystem. Instead of TPU-Inference, we focus on SGLang-JAX, using DSA sparsity as a concrete example of the ongoing optimization work around GLM-5.2 inference.

What follows is a constructive walkthrough of a sequence of PRs that progressively improve GLM-5.2 serving performance, showing how the implementation evolved and where the gains came from. Several of these merged PRs include contributions from MoonMath, which is another reason we chose this case study.

While we use GLM-5.2 as the running example, the techniques discussed here apply equally to GLM-5.3.

Baseline and gameplan

GLM-5.2 already had a learned indexer and a layer-sharing schedule. SGLang-JAX initially loaded the model but still traversed the dense MLA cache, discarding the indexer’s output.

The first optimization was to consume that output; the next ones made sparse execution efficient in memory, compute and serving.^[1] Each step exposed the next bottleneck: dense reads gave way to page reads, then oversized indexer intermediates, duplicate page fetches and finally whole-cache movement around an append.

Separate the indexer from the attention it controls

Multi-head Latent Attention (MLA) compresses past-token KV representations; DeepSeek Sparse Attention (DSA) determines which positions attention reads. GLM-5.2 stores a 512-dimensional KV latent plus a 64-dimensional positional key component, shared across 64 attention heads. Its indexer has 32 heads of width 128 and a checkpoint budget of 2,048 tokens.^[2] Several prefill experiments instead use a 4,096-token budget: 32 pages of 128 tokens so their settings are not interchangeable with the checkpoint default.

The indexer scores candidate keys with a weighted sum of ReLU-transformed query–key dot products. These scores determine the selected positions, not the final attention probabilities. IndexShare reuses the result across layers: 21 full layers refresh it, while 57 shared layers reuse it until the next full layer. “Full” means indexer ownership, not dense attention.^[3]

The TPU path uses absorbed MLA: key projection moves onto the query, and value projection follows the weighted sum of latents. This avoids expanding cached latents into per-head keys and values. The target is therefore fewer reads from the existing compressed cache. Unselected entries must remain available to later queries, and the indexer retains its own key history. The attention scale remains tied to the original 192 + 64 query/key dimensions, not the absorbed 576-dimensional representation.^[4]

Make the sparsity unit match the memory unit

Sparse decode supplies the existing MLA kernel with a shorter physical-page list. The kernel reads directly from the original cache, avoiding an HBM gather; the new-token page is handled explicitly so its cache write reaches the correct destination.^[5]

Token top-k is not page top-k. Dispersed high-scoring tokens can touch most pages, and reading every token in those pages broadens attention support. A direct page budget instead ranks each page by its highest legal token score. Four top-scoring tokens in four four-token pages illustrate the difference: token top-k names four positions, page expansion reads sixteen, and a two-page budget reads eight. This bounds page reads but changes the approximation. Causal and valid-length masks must precede the page maximum, or an invisible token can promote its page and displace a valid one.^[6]

Benchmarking dense and page-sparse decode, with AOT dispatch enable for both, early 110k-context decode fell from 46.2 to 17.1 ms/token using 128 selected pages. At 1k, sparse was slower: 13.7 versus 11.8 ms/token. Indexer and metadata overhead need enough avoided attention work to pay off. The bounded page budget does not make decode context-independent: the indexer still scans candidate keys.^[7]

Carry sparse reads through prefill

Fused sparse prefill writes new latent rows into the paged cache, then traverses selected pages with online softmax: static work shapes, dynamic addresses. For T queries, Kp pages per query and b tokens per page, selected-attention work scales approximately as O(T × Kp × b), excluding head and feature dimensions.^[8]

At TP16, only four attention heads remain per device. The per-query kernel puts the abundant selected-key axis on the lane dimension and the small head axis on the sublane dimension, avoiding excessive padding. Three early full-indexer layers had taken a dense route, leaving three shared followers without a page list. Extending sparse routing through that six-layer region completed the path: matched 48k prefill moved from 155.7 s dense → 38.4 s sparse → 28.7 s with complete routing.^[9]

Pack requests and reuse prefixes

The single-request kernel becomes a serving path when each query carries its request ID, absolute position, valid KV length and page-table base. Cached prefixes and chunks are extend operations; decode can be a one-query extend. For prefix length P and local offset i, causality uses P + i: not an offset reset at each chunk.^[10]

This metadata lets variable-length requests share an execution grid without separate launches or padding every request to the longest. Prefix reuse avoids recomputation; chunking bounds active queries. An indexer page number must still be translated through the correct request’s page table. Neither packing nor prefix reuse changes which keys a query may read. The next bottleneck was upstream: the indexer still scored reserved capacity.

Stream indexer scores over live context

The reference prefill indexer scored the full configured KV extent and materialized token scores before reducing them to page scores. The streaming path instead reads valid KV blocks, computes scores, applies masks and pools page maxima in on-chip VMEM. Only page scores reach HBM; final page top-k remains a separate step.^[11]

For 8,192 queries and 135,168 configured positions, FP32 token scores occupy 4.125 GiB. Pooling 128-token pages reduces the logical score output to 33 MiB, which is 128× smaller. The indexer still scores tokens in visited blocks, and long-context prefill retains a quadratic component. For request length L, configured capacity Lmax and whole chunks of C tokens with no cached prefix, capacity-sized scoring scales with L × Lmax; scanning successive live prefixes scales with (L² + LC)/2. This work model explains the savings without claiming linear-time prefill.

Holding chunks at 4,096 tokens, enabling the streaming indexer reduced 76k TTFT from 62.1 to 44.5 s. Smaller temporaries also made 8,192-token chunks viable at roughly 135k capacity. In a separate concurrency-eight, 4,096-in/128-out workload, output throughput rose from 2.8 to 29.6 tokens/s by relieving serving stalls—not by accelerating decode arithmetic 10.6×.^[12]

Reuse pages across queries

Once the indexer stops writing huge intermediates, repeated sparse reads become visible. Nearby queries often select overlapping pages, yet the per-query kernel reloads them independently. Query blocking loads the union once, then restores each query’s own support with a membership mask before online softmax. Causality and valid-length masks still apply.^[13]

Unlike page-budget selection, this can preserve already-selected attention support. The union must not be truncated: its lossless capacity is min(QB × Kp, number of units). Membership is indexed by compact union slot rather than every possible page-table ID. The same support preserves the mathematical operation, but floating-point reduction order can still change; byte-identical generations are not guaranteed. For QB queries, Kp selected pages per query and union size U, the idealized trade-off is:

Page-load reduction = QB × Kp / U | Score-work expansion = U / Kp

At QB = 64, Kp = 32 and U ≈ 69, that is 29.7× fewer logical page loads for 2.16× more score work. These are analytical ratios, not latency predictions: they ignore cache reuse outside the kernel and do not capture DMA overlap, occupancy or preprocessing overhead. The kernel spends masked arithmetic to reduce discontiguous data movement.

The implementation sorts and compacts the union, builds membership without a multidimensional scatter, and overlaps fetches through a four-slot DMA ring. Blocking also improves matrix shape: four local heads at TP16 become 256 query–head positions for a 64-query block. The gain combines reuse, latency overlap and better compute geometry.

The paired 110k experiment reduced TTFT from 65.54 to 30.40 s, with decode essentially unchanged. A locality-shaped microbenchmark fell from 43.3 to 8.1 ms. Uniform-random selections, however, expanded the union to roughly 795 of 880 pages and made blocking about twice as slow. Query blocking depends on locality; larger blocks are not universally better.^[14]

Append new KV without moving the resident cache

After indexer and attention optimization, cache maintenance became the next visible cost. A flat-view scatter appeared to update only new rows, but compiler layout conversions copied the entire paged pool twice. Rewriting it as a four-dimensional scatter did not remove those transformations.^[15]

Direct paged write-back instead groups destination slots into contiguous runs and writes into aliased cache storage. Aligned interiors use direct copies; packed-word boundaries use read-modify-write to preserve neighbouring tokens. Excessive fragmentation takes a correctness-preserving fallback. The append should scale with new rows, not the resident allocation. Across many layers and chunks, an extra resident-capacity term can consume the gains from faster attention. This is cache maintenance, not eviction: previously stored rows remain intact and available to future queries.

The same follow-up raised the query block from 64 to 256. Its 110k sequence was 30.20 s at QB64 → 24.43 s with paged write-back → 21.89 s at QB256, a 27.5% reduction. This is a separate benchmark from the earlier 30.40 s result, not a continuous measured waterfall.^[16]

Test the contracts changed by each optimization

Indexer tests check legal page sets; blocked-attention tests check masked-softmax equivalence; serving tests check request permutation, isolation, prefix reuse, chunking and cache-write canaries. These checks separate whether the indexer selects acceptable support, whether the kernel preserves that support’s operation, and whether serving preserves the kernel’s inputs. Selecting every valid page should reproduce dense attention, but that only validates the execution path. Likewise, prompts shorter than the indexer budget barely test pruning. Long-context quality requires separate evaluation.^[17]

Use independent oracles and test the tests

Indexer-cache accounting provides a stronger resource test than comparing two copies of a sizing formula. Admission originally omitted indexer keys, leaving about 0.47 GiB unbudgeted above 8.75 GiB of latent cache. Tests allocate a real MLATokenToKVPool on a single-device CPU mesh, sum live arrays’ nbytes, and compare actual residency with both predicted and reported usage. Spare-page slack is explicit.^[18]

Discriminating shapes matter. GLM’s 512/64 widths cannot distinguish separate-segment alignment from concatenated alignment: both yield 640 features at 128-feature alignment. Synthetic 192/64 widths do distinguish them, 384 versus 256, and page size one exposes BF16 packing padding. Reverting the budget fix failed five of seven new cases; reverting reporting failed the live-allocation check. Redundant cases were removed according to which mutations they detected. These checks demonstrate sensitivity to plausible wrong implementations, rather than merely increasing the number of passing tests.^[19]

The open data-parallel proposal tests above the kernel: compare a request alone and inside DP2/DP4 batches at fixed attention TP, require rank isolation, and inject global page IDs where local IDs belong. Nontrivial TP reveals sharding mistakes hidden at TP1. The injected page-ID fault produced a reported output difference of 3.4 against FP32 tolerances of 1e−5. This is test-only work, not an unmerged speedup.^[20]

What the stack now delivers

Page-sparse reads bound expensive attention work; IndexShare reuses the page list across layers; packed metadata carries it through serving. The streaming indexer reduces padded work and intermediates, query blocking reuses page loads, and paged write-back keeps appends local.

DSA_PREFILL_SPARSE=1 and DSA_INDEXER_KERNEL_PREFILL=1 remain separate opt-ins; query blocking defaults to QB256 within sparse prefill. Decode’s page-budget and token-top-k indexer routes are alternatives, not additive stages.^[21] The latter regressed from 15.9 to 34.2 ms/token at single-request concurrency despite improving at higher concurrency.^[22] The objective is not simply to enable DSA, but to prevent the indexer, kernel and runtime from reintroducing the work sparsity removes.