类别：编译时：范围限定的 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 OOM 是由非自定义内核操作引起的，则很可能是内部编译器 bug。请提交包含 HLO 转储的 XLA bug。
重新调整内核
如果错误源自自定义内核，请使用以下方法来降低内核的内存要求：
- 调整块大小：减小内核配置中的块大小（平铺大小），以降低 Scoped Vmem 用量。
- 设置每个内核范围的 Vmem 限制：使用 vmem_limit_bytes 参数为特定内核明确请求所需的内存量
- 修改内存着色：使用 pallas.tpu.with_memory_space_constraint 将内核的输入/输出显式着色/限制为 VMEM。但请注意，不要为过多的输入/输出着色到 Vmem，因为这可能会导致整体 VMEM OOM。
- 如果难以进行内核特定的重新调整，或者问题影响了许多内核，您可以使用标志 --xla_tpu_scoped_vmem_limit_kib 调整全局 Vmem 限制。