2026年3月25日
TPU 架构与 Pallas Kernel 编程入门：从内存层次结构到 FlashAttention
摘要： 做过 GPU kernel 优化的人对以下编程模型肯定不会陌生：写一个 CUDA kernel分发到流式多处理器（SM）上执行，缓存层次结构自行负责数据搬运。而TPU 则完全不同，除非明确告诉编译器要把哪些数据块搬到哪里，否则kernel 根本无法编译。实际操作确实和听起来一样繁琐，所以JAX 的P 阅读全文
posted @ 2026-03-25 20:32 deephub 阅读(12) 评论(0) 推荐(0)