Pallas：JAX 内核语言

Pallas 是 JAX 的一个扩展，它支持为 GPU 和 TPU 编写自定义内核。其目标是提供对生成代码的细粒度控制，同时结合 JAX 追踪的高级人机工程学和 jax.numpy API。

本节包含使用 Pallas 的教程、指南和示例。另请参阅 jax.experimental.pallas 模块 API 文档。

警告

Pallas 仍处于实验阶段，并且频繁更新。请参阅 Pallas 变更日志 查看最近的更改。

您可能会遇到错误和未实现的情况，例如，当降低需要仿真的高级 JAX 概念时，或者仅仅因为 Pallas 仍在开发中。

指南

• Pallas 快速入门
  • Pallas 中的 Hello world
  • Pallas 编程模型
• 软件流水线
  • 内存层次结构
  • 流水线基础
  • Pallas 流水线 API
  • 棘手之处
  • 性能分析
• 网格和 BlockSpecs
  • grid，即循环中的内核
  • BlockSpec，即如何分块输入

TPU 后端指南

• Pallas TPU
  • 使用 Pallas 编写 TPU 内核
  • TPU 流水线
  • 矩阵乘法
  • 标量预取和块稀疏计算
  • Pallas 在 TPU 上的分布式计算

Mosaic GPU 后端指南

• Pallas: Mosaic GPU
  • 使用 Pallas 编写 Mosaic GPU 内核
  • Mosaic GPU 流水线

设计说明

• Pallas 设计说明
  • Pallas 设计
  • Pallas 异步操作

其他

• Pallas 变更日志