[Docs] JAX, Pallas, and TPU execution architecture guide

[Tar-ive]

Adds a developer guide that explains why this repository is written in JAX and Pallas, what mathematics those languages encode, how they compile, and how the resulting programs run on a TPU.

Why

JAX (often heard as "Jacks") and Pallas (often heard as "Palace") are easy to treat as "just Python." In tpu-inference they are the whole lowering path: models and layers are JAX array programs; attention, MoE, and collectives are Pallas kernels; XLA and Mosaic (C++ compilers) emit TPU machine code. There was no single document that walked that path from transformer math down to the MXU.

What landed

• New guide: docs/developer_guides/jax_pallas_architecture.md
  • Why these languages exist (two-language problem, systolic arrays, ragged inference)
  • The math they compute (einsum, RMSNorm, RoPE, online softmax / FlashAttention, MoE, GSPMD)
  • Where JAX stops and Pallas starts in this repo
  • Compilation: Python → jaxpr → StableHLO/HLO → XLA/Mosaic C++ → TPU VLIW (not generated C++ source)
  • Architecture-level execution: HBM/VMEM/SMEM, MXU, VPU, SparseCore, ICI, VLIW issue slots
  • A request walk through TPUWorker → CompilationManager → RPA
  • 15 diagrams
• Wired into MkDocs nav (with Mermaid fences) and the docs home / JAX model-dev pages

Tests

Docs-only. No runtime tests. Mermaid is GitHub-native and enabled in MkDocs via pymdownx.superfences custom fences.

  

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}