JAX
JAX自定义算子开发
与优化实战
🔥 30章
从入门到高阶 · 系统化课程
📚
30
章节
⚡
Pallas
+
Triton
🎯
GPU/TPU
优化
01
JAX生态概览
入门
JAX与NumPy关系 · 自动微分/JIT/向量化 · 安装验证
→
02
JAX基础入门
入门
jnp.array vs np.array · device_buffer · jax.random · 纯函数
→
03
JIT编译机制
核心
jit工作原理 · 静态/动态参数 · trace/compile · 陷阱调试
→
04
自动微分深入
核心
grad/vmap/pmap协同 · 高阶梯度 · Hessian · 自定义VJP
→
05
向量化与并行
进阶
vmap语义 · pmap多设备 · sharded_map · 数据并行策略
→
06
Pallas内核入门
Pallas
Pallas是什么 · 与Triton对比 · 第一个内核：向量加法
→
07
Pallas内核进阶
Pallas
块状编程 · 内存层级/tiling · pallas_call · 调试
→
08
自定义算子基础
核心
为什么需要 · JAX原生 vs 自定义 · 应用场景 · 性能预期
→
09
XLA与HLO
底层
XLA编译栈 · HLO中间表示 · StableHLO · 查看HLO图
→
10
自定义VJP规则
进阶
implicit/explicit VJP · custom_vjp · 前反向实现 · 梯度验证
→
11
自定义primitive
高阶
primitive注册 · 前向/反向规则 · jit兼容性
→
12
lowering规则
高阶
lower到XLA · HLO lowering · 多后端(CPU/GPU/TPU) · 调试
→
13
Pallas自定义算子实战
Pallas
高性能Softmax · Flash Attention · 内存优化 · 性能对比
→
14
自定义算子性能分析
工具
jax.profiler · Nsight集成 · 带宽/吞吐分析 · 瓶颈定位
→
15
算子融合策略
优化
手动/自动融合 · jax.lax组合 · 自定义融合pass
→
16
GPU内核优化
GPU
warp调度/occupancy · 共享内存 · 向量化加载 · bank conflict
→
17
TPU内核优化
TPU
脉动阵列 · MXU · TPU内存层级 · 特有优化技巧
→
18
混合精度计算
精度
fp16/bf16 · jax.dtype · 精度管理 · 梯度缩放
→
19
自动微分优化
进阶
checkpointing · 梯度图优化 · 重计算权衡 · jax.checkpoint
→
20
分布式自定义算子
分布式
pmap自定义 · ppermute/all-reduce · SPMD · 分布式语义
→
21
编译缓存
工程
JAX缓存机制 · 命中/未命中 · 缓存策略 · 调试清理
→
22
测试与验证
工程
单元测试 · 数值正确性 · 梯度测试 · 性能回归
→
23
文档与发布
工程
API文档 · 类型注解 · 示例代码 · PyPI发布
→
24
JAX与Triton集成
生态
Triton内核基础 · JAX调用Triton · 与Pallas对比 · 混合策略
→
25
安全与稳定性
工程
内存安全 · 数值稳定性 · 错误处理 · 优雅降级
→
26
JAX生态工具链
生态
Optax · Flax · Haiku · Orbax检查点
→
27
LLM推理应用
实战
KV Cache自定义 · GQA/MQA优化 · PageAttention · 推理优化
→
28
科学计算应用
实战
ODE求解器 · FFT优化 · 稀疏矩阵 · 物理模拟
→
29
前沿趋势
展望
动态shape · 可微分编程 · JAX与MLIR · 未来方向
→
30
综合实战项目
🚀 毕业
从零构建高性能算子库 · 架构设计 · 基准测试 · 社区贡献
→