A curated TPU sparse-attention reading path: block-sparse contracts, Pallas kernel choices, SPMD sharding, and the runtime surfaces that keep long-context TPU work stable.
This hub is narrower than the general TPU/XLA archive. Start with the block-sparse and Pallas kernel notes, then move into sharding and planner surfaces, and finally the operational pieces that keep the TPU sparse-attention lane observable and stable.
Best if you care specifically about sparse attention, Pallas, and long-context TPU kernel work rather than the TPU stack as a whole.
These are the core sparse-attention and Pallas notes to read first.
How to frame block-sparse attention on TPU honestly: explicit mask contracts, MXU-aligned tile choices, and a preference for stable sparse layouts over data-dependent retracing.
The block-mask and MXU-friendly kernel contract that anchors the rest of this cluster.
Where Pallas beats the XLA lowering on TPU v6e, where it loses, the debugging workflow that keeps us sane, and the kernel deltas we kept versus the ones we reverted.
What the Pallas lane kept, what it deleted, and why the TPU kernel surface narrowed over time.
Why softcap attention on TPU needs a dedicated kernel surface: fuse the nonlinearity, keep masking contract-friendly, and avoid turning a stability trick into a second full pass over the score matrix.
The flash-attention and softcap companion piece once Pallas is already familiar.
Once the kernels are legible, these explain how the TPU layout stays stable.
How MegaCpp decomposes clustered sparse TPU attention into planner stages, legality checks, and backend dispatch rather than treating sparse attention as one giant kernel.
The planner-stage explanation for the clustered sparse TPU path.
Why explicit mark_sharding annotations matter on TPU XLA, what should be pinned explicitly, and why propagation is not a substitute for a stable sharding contract.
The concrete sharding annotations the TPU sparse-attention lane actually relies on.
Vocab-size constraints under XLA, the padding choices that keep the compile cache stable, sharded embedding init under SPMD, and the per-specialist platform vocab story.
The tokenizer and vocab decisions you have to lock down before the TPU sparse path stays compile-safe.
These are the nearby runtime notes that keep the sparse lane debuggable.
Why TPU telemetry has to be gated carefully: scalar reads can become host-device syncs, so sink and outlier tracking must be designed as explicit low-cadence instrumentation.
How to keep observability useful without turning the TPU lane into a measurement artifact.
Transformer Engine is an NVIDIA Hopper and Blackwell story. On TPU v6e it does not exist. This is the layer-spec abstraction and the XLA-friendly substitutes that let one model definition ship across both paths.
The TPU-side model-definition substitutions that keep the sparse or Pallas path compatible with the rest of the stack.
How the v6_enriched packed-rows pipeline feeds per-token structure IDs, chunk boundaries, and call edges into the XLA dataloader on TPU v6e without triggering compile cache misses, and how that contract lifts into the main path.
The packed-row data companion once long-context TPU execution has to stay compile-stable end to end.
Keep exploring
These hubs cover nearby parts of the blog without turning the archive into a giant taxonomy.
A curated GB10 and Blackwell reading path: consumer-versus-datacenter tensor paths, driver-visible false positives, arch-patch repros, and the serving or precision choices that survived contact with the hardware.
A curated Modal reading path: when the rented-GPU surface was useful, what broke on multi-GPU launches, how receipts were recorded, and how we kept the lane debuggable.
A curated Mamba3 reading path: why MegaCpp kept a hybrid stack, how the kernels evolved across CUDA, TileLang, and TPU, and where the runtime wins actually held.
A curated MLA reading path: the weight-absorption contract, Megatron-safe integration boundaries, dispatch and FP8 edges, and the adapter surfaces that keep MLA connected to the rest of the stack.
A curated evaluation reading path: verifier-first harnesses, ablation structure, benchmark receipts, and the evidence rules that keep comparisons from collapsing into anecdotes.
A curated Megatron reading path: the parallelism map, what actually splits, how NVIDIA and TPU wrappers differ, and the migration surfaces around NAM56R-style layouts.
A curated path through the H200 lane: operator bring-up, step-time anatomy, memory pressure, and the NVIDIA kernel surfaces that actually moved the stack.
A curated reading order for TPU work: bring-up, PJRT and Torch/XLA boundaries, SPMD sharding, and the kernel/runtime traps that made TPU performance non-obvious.
A curated archive for the C++ data path: corpus selection, semantic enrichment, packaging into training artifacts, and the file-level durability choices that keep the pipeline sane.
A curated path through the expert stack: what the specialist path changed, how routing works, and how the parallelism map constrains the model layout.