类别： 编译时：作用域 Vmem OOM
此错误表示程序需要的作用域向量内存 (Vmem) 超出了分配量。
示例出错提示：
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
XLA 后端： TPU
概览
TPU 具有向量内存 (VMEM)，这是一种仅供 TensorCore (TC) 使用的本地暂存内存。编译器会针对不同类型的分配管理 Vmem：
- 指令作用域分配： 在执行单个 HLO 指令时，Vmem 中的临时存储空间。这包括操作数跨度缓冲区（例如，用于双缓冲）和寄存器溢出。
- 程序作用域分配： 超出单个 HLO 指令范围的分配。这些通常是 HLO 临时变量和中间结果，它们是 HLO 指令的输入和/或输出。
当指令作用域分配超出该指令的分配限制时，就会发生编译时作用域 Vmem OOM。此限制由以下因素控制
- 通过标志 --xla_tpu_scoped_vmem_limit_kib 全局控制整个程序
和
- 通过 vmem_limit_bytes 参数控制每个自定义内核。
这些错误通常是由内部编译器 bug 或自定义内核超出其分配限制引起的。
调试
仔细分析错误消息，以确定错误是源于自定义内核还是标准 HLO。由自定义内核引起的错误应具有以下签名：
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- 自定义内核作用域 Vmem OOM：如果错误指向自定义内核 → 跳转到 重新调整内核。
- 非内核 Vmem 问题：如果 Vmem Omem 是由 非自定义内核操作引起的，则可能是内部编译器 bug。 请针对 XLA 提交 bug 报告，并附上 HLO 转储。
重新调整内核
如果错误源于自定义内核，请使用以下技术来降低内核的内存要求：
- 调整块大小： 减小内核配置中的块大小（平铺大小），以降低作用域 Vmem 使用量。
- 设置每个内核的作用域 Vmem 限制： 使用 vmem_limit_bytes 参数显式请求该特定内核所需的 内存量
- 修改内存着色： 使用pallas.tpu.with_memory_space_constraint将内核的 输入/输出显式着色/限制为 Vmem。 注意不要将过多的输入输出着色为 Vmem，因为这可能会导致整体 Vmem OOM。
- 调整 Vmem 限制： 如果难以进行内核特定的重新调整，或者问题 影响到多个内核，您可以使用标志 --xla_tpu_scoped_vmem_limit_kib 调整全局 Vmem 限制。