Pallas：JAX 内核语言完整指南

原文档：https://docs.jax.dev/en/latest/pallas/

本中文翻译覆盖 Pallas 文档的所有页面，保留完整代码和图片。

快速入门

#

• Pallas 快速入门 — Hello World、编程模型、Grid、BlockSpec

核心概念

#

• Grid 和 BlockSpec — Grid 迭代空间、BlockSpec 切分输入、索引模式
• 软件流水线 — 内存层次、双缓冲、流水线 API、性能分析

Pallas TPU

#

• 使用 Pallas 编写 TPU 内核 — TPU 架构、数组布局、支持的操作
• TPU 流水线 — 内存空间、多重缓冲、emit_pipeline、Megacore
• 矩阵乘法 — 分块矩阵乘法、性能分析、bf16、融合激活
• 标量预取与块稀疏计算 — 动态块索引、稀疏矩阵乘法、预取映射
• TPU 分布式计算 — RDMA、信号量、双缓冲、ppermute/all_gather/psum/psum_scatter
• 核心特定编程 — core_map、逐核心内核、SparseCore 映射
• SparseCore 内核编写 — SparseCore 硬件、流水线、Gather/Scatter
• 伪随机数生成 — jax.random API、硬件 PRNG、块不变采样

Pallas:Mosaic GPU

#

• 使用 Pallas 编写 Mosaic GPU 内核 — SM 架构、内存空间、MMA、core_map、同步、异步拷贝
• Mosaic GPU 流水线 — GPU 流水线、Warp 特化、矩阵乘法示例
• Blackwell 高性能矩阵乘法 — 7 步渐进优化、2CTA MMA、持久化内核
• 集合矩阵乘法 — Ring All-Gather、设备间通信、JAX 集成

设计笔记

#

• Pallas 设计 — 设计理念、JAX 内核扩展、Lowering 流程
• Pallas 异步操作 — 异步 DMA、信号量状态、优化屏障