XLA 如何调整性能 flags 提升 TPU 负载的执行性能【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow如果你的 TensorFlow 程序跑在 TPU 上、已经走 XLA 编译路径但执行速度还没有达到预期可以尝试调整 XLA 的性能 flags。TensorFlow 仓库自带的 XLA Flags Guidance 维护了一份经过整理的 flag 清单其中明确列出了哪些 flag 影响运行时性能、默认值是什么、建议值是什么、候选值范围是什么以及哪些 flag 必须成组使用。这篇文章按照这份文档给出的分组说明如何为你的 TPU 负载挑选 flags、如何通过环境变量生效以及如何核对编译结果。前提确认负载走的是 XLA 编译路径flags 只对经过 XLA 编译的部分起作用所以先确认你的程序确实在用 XLA。XLA for TensorFlow 文档给出两种启用方式显式编译在tf.function上标注jit_compileTrue。注意jit_compile是 must-compile 语义——要么整个函数被 XLA 编译要么抛出errors.InvalidArgumentErrorAuto-clustering不改代码通过环境变量自动聚类可编译子图$ TF_XLA_FLAGS--tf_xla_auto_jit2 path/to/your/tf/program两种方式的文档说明都以 GPU 为主要优化对象TPU 负载则依赖运行时选择的 TPU backend。如果你的程序根本没有进入 XLA 编译下面的 flag 调整不会有任何效果。通过 XLA_FLAGS 环境变量设置性能 flagsXLA flags 的通用传递方式是XLA_FLAGS环境变量--flagvalue形式拼接多个 flag。HLO Dumps 文档确认该变量对 JAX、TensorFlow、PyTorch/XLA 都生效。应用 flags 的基本命令形态是$ XLA_FLAGS--flag_avalue_a --flag_bvalue_b path/to/your/tf/program按文档分组选择 flagsflags_guidance.md 的 Performance Flags 一节把 flag 分成几类每类都标明了适用条件。选择依据如下通信重叠类Pipelining 与异步集合通信分组flags需成组使用默认值建议值Pipeliningxla_should_allow_loop_variant_parameter_in_chain、xla_should_add_loop_invariant_op_in_chain、xla_tpu_enable_ici_ag_pipelining前两个为kDisabled第三个为false前两个为kEnabled第三个为truev5e/Asyncall-gatherxla_enable_async_all_gather、xla_tpu_enable_async_collective_fusion、xla_tpu_enable_async_collective_fusion_fuse_all_gatherkAuto、true、true同默认值v5e/Asyncall-reducexla_tpu_enable_async_collective_fusion、xla_tpu_enable_async_collective_fusion_fuse_all_reducetrue、falsetrue、trueAsyncall-to-allxla_tpu_enable_async_all_to_allfalsetrue文档明确说明第一组 3 个 flag 应同时使用用于启用 ICIInterchip-Interconnectall-gather 的 collective pipelining制造更多执行重叠的机会第二、三组标注 v5e/Async即针对 v5e 平台的异步 all-gather / all-reduce。如果你的负载不是 v5e这两组的收益没有文档依据不要默认开启。推理负载Latency-bound 阈值如果你的负载以推理为主、存在小尺寸latency-bound的集合通信文档给出 4 个阈值 flag默认值都是-1即未启用xla_all_gather_latency_bound_threshold_in_bytesxla_all_reduce_latency_bound_threshold_in_bytesxla_collective_permute_latency_bound_threshold_in_bytesxla_all_to_all_latency_bound_threshold_in_bytes文档说明它们是intended for latency-bound (i.e., small-sized)对应类型的集合操作启用后触发针对性优化以降低执行时间typically its used in inference workloads。建议值区间为4~16Mb文档同时给出了换算4~16 * 1024 * 1024字节候选值范围是[0, 9223372036854775807]。计算密集型Compute centric 类这组 flag 文档给出的默认值与建议值基本一致适合作为核对项而非调优起点flag默认值 建议值候选值xla_tpu_enable_dot_strength_reductiontruetrue/falsexla_tpu_dot_dot_fusiontruetrue/falsexla_jf_enable_multi_output_fusiontruetrue/falsexla_tpu_scoped_vmem_limit_kib16384[4096, VMEM size of the architecture - 1024]xla_tpu_async_copy_bandwidth_scaling_factor1(0, 1]xla_msa_enable_cross_program_prefetch_freeingenabledenabled/disabledxla_tpu_msa_inefficient_use_to_copy_ratio0.5[0, 1]文档对这几个 flag 的用途说明分别是将非计算密集 dot 重写为 multiply reduce、执行 dot-dot fusion 以减少慢速/主存上的中间输出、启用多输出 fusion、设置每个算子可用的 scratchpad VMEM 大小单位 KiB、缩放 async copy 的有效带宽用于 prefetch 决策、启用跨程序 prefetch buffer 的释放优化、以及 VMEM 放置决策中判断分配点是否低效的 use/copy 字节比0 视为所有点高效1 要求使用字节数至少等于 async copy 字节数。一条最短的主路径示例以文档第一组 Pipelining flags 为例把三个 flag 设为建议值并运行程序$ XLA_FLAGS--xla_should_allow_loop_variant_parameter_in_chainkEnabled \ --xla_should_add_loop_invariant_op_in_chainkEnabled \ --xla_tpu_enable_ici_ag_pipeliningtrue \ path/to/your/tf/program如果同时是 v5e 上以推理为主的负载可以在同一XLA_FLAGS中追加一个 latency-bound 阈值例如取建议区间内的8 * 1024 * 1024字节$ XLA_FLAGS--xla_should_allow_loop_variant_parameter_in_chainkEnabled \ --xla_should_add_loop_invariant_op_in_chainkEnabled \ --xla_tpu_enable_ici_ag_pipeliningtrue \ --xla_all_gather_latency_bound_threshold_in_bytes8388608 \ path/to/your/tf/program这里8388608即8 * 1024 * 1024落在文档建议的4~16Mb区间内具体取哪个值需要你在自己的负载上实验确定——文档只给出区间没有给出固定推荐点。验证 flags 是否生效Inspect compiled programs 和 XLA Tooling 两节都说明把--xla_dump_to放进XLA_FLAGS即可把编译产物导出到指定目录。在上面的命令里加上该 flag$ XLA_FLAGS--xla_dump_to/tmp/generated \ --xla_should_allow_loop_variant_parameter_in_chainkEnabled \ --xla_should_add_loop_invariant_op_in_chainkEnabled \ --xla_tpu_enable_ici_ag_pipeliningtrue \ path/to/your/tf/program运行结束后/tmp/generated目录下会生成每个编译 cluster 的产物包括module_XXXX.*_optimizations.txt优化后的 XLA 程序文本每个 cluster 一份以及中间表示文件。对比调整 flag 前后两次 dump 中 HLO 指令序列的变化是文档提供的核对手段这些 dump 在提交 XLA bug report 时也是被要求附带的材料。文档给出的成功/失败判据很直接调整的是can significantly impact runtime performance的 flag收益需要在你的负载上自行测量而如果启用某个 flag 后出现 crashflags_guidance.md 的建议是回退到默认设置并创建 GitHub issue——也就是说flag 组合出问题时的官方动作不是继续尝试新组合。边界与不支持项以下限制来自文档原文调整 flags 时需要一并遵守v5e/Async 分组只针对 v5e 平台xla_enable_async_all_gather等 flag 的描述明确写着 activate asynchronous all-gather operations on v5e其他 TPU 版本上开启没有文档依据。Latency-bound 阈值主要针对推理文档反复注明 Typically its used in inference workloads训练负载是否适用文档没有说明。Memory flags 只在 HBM 报 out of memory 时才动xla_latency_hiding_scheduler_rerun、xla_memory_scheduler等 memory 类 flag 文档明确警告should only be adjusted if you encounter HBM out of memory errors during model compilation其他场景应保持默认因为改变它们可能反过来损害性能。成组 flag 不要拆开调Pipelining 三件套、v5e 两组 flags文档都用 should be used in conjunction 描述单独改其中一个不符合文档用法。xla_tpu_scoped_vmem_limit_kib的候选值上界与硬件相关候选范围写为[4096, VMEM size of the architecture - 1024]即上界取决于你使用的 TPU 架构 VMEM 大小文档没有给出统一数值。另外flags_guidance.md 的 TPU XLA flags 小节还列出了 5 个 TPU 专项 boolean flag如xla_tpu_enable_data_parallel_all_reduce_opt用于增加数据并行 all-reduce 的重叠机会、xla_tpu_spmd_rng_bit_generator_unsafe用于分区执行 RngBitGenerator 但会牺牲确定性它们属于按负载特征按需启用的选项默认文档未列建议值不建议与上面的性能组合一起盲开。下一步若需要理解编译器在不同优化等级下会启用哪些 pass参考 Effort LevelsEFFORT_O2是文档标注的生产负载默认档位。若 dump 出的 HLO 需要脱离原程序单独编译或定位某个 pass 的影响XLA Tooling 中的hlo-opt、run_hlo_module是文档给出的工具路径。【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/9 21:40:48 · 阅读 1.2k XLA 如何调整性能 flags 提升 TPU 负载的执行性能？
XLA 如何调整性能 flags 提升 TPU 负载的执行性能【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow如果你的 TensorFlow 程序跑在 TPU 上、已经走 XLA 编译路径但执行速度还没有达到预期可以尝试调整 XLA 的性能 flags。TensorFlow 仓库自带的 XLA Flags Guidance 维护了一份经过整理的 flag 清单其中明确列出了哪些 flag 影响运行时性能、默认值是什么、建议值是什么、候选值范围是什么以及哪些 flag 必须成组使用。这篇文章按照这份文档给出的分组说明如何为你的 TPU 负载挑选 flags、如何通过环境变量生效以及如何核对编译结果。前提确认负载走的是 XLA 编译路径flags 只对经过 XLA 编译的部分起作用所以先确认你的程序确实在用 XLA。XLA for TensorFlow 文档给出两种启用方式显式编译在tf.function上标注jit_compileTrue。注意jit_compile是 must-compile 语义——要么整个函数被 XLA 编译要么抛出errors.InvalidArgumentErrorAuto-clustering不改代码通过环境变量自动聚类可编译子图$ TF_XLA_FLAGS--tf_xla_auto_jit2 path/to/your/tf/program两种方式的文档说明都以 GPU 为主要优化对象TPU 负载则依赖运行时选择的 TPU backend。如果你的程序根本没有进入 XLA 编译下面的 flag 调整不会有任何效果。通过 XLA_FLAGS 环境变量设置性能 flagsXLA flags 的通用传递方式是XLA_FLAGS环境变量--flagvalue形式拼接多个 flag。HLO Dumps 文档确认该变量对 JAX、TensorFlow、PyTorch/XLA 都生效。应用 flags 的基本命令形态是$ XLA_FLAGS--flag_avalue_a --flag_bvalue_b path/to/your/tf/program按文档分组选择 flagsflags_guidance.md 的 Performance Flags 一节把 flag 分成几类每类都标明了适用条件。选择依据如下通信重叠类Pipelining 与异步集合通信分组flags需成组使用默认值建议值Pipeliningxla_should_allow_loop_variant_parameter_in_chain、xla_should_add_loop_invariant_op_in_chain、xla_tpu_enable_ici_ag_pipelining前两个为kDisabled第三个为false前两个为kEnabled第三个为truev5e/Asyncall-gatherxla_enable_async_all_gather、xla_tpu_enable_async_collective_fusion、xla_tpu_enable_async_collective_fusion_fuse_all_gatherkAuto、true、true同默认值v5e/Asyncall-reducexla_tpu_enable_async_collective_fusion、xla_tpu_enable_async_collective_fusion_fuse_all_reducetrue、falsetrue、trueAsyncall-to-allxla_tpu_enable_async_all_to_allfalsetrue文档明确说明第一组 3 个 flag 应同时使用用于启用 ICIInterchip-Interconnectall-gather 的 collective pipelining制造更多执行重叠的机会第二、三组标注 v5e/Async即针对 v5e 平台的异步 all-gather / all-reduce。如果你的负载不是 v5e这两组的收益没有文档依据不要默认开启。推理负载Latency-bound 阈值如果你的负载以推理为主、存在小尺寸latency-bound的集合通信文档给出 4 个阈值 flag默认值都是-1即未启用xla_all_gather_latency_bound_threshold_in_bytesxla_all_reduce_latency_bound_threshold_in_bytesxla_collective_permute_latency_bound_threshold_in_bytesxla_all_to_all_latency_bound_threshold_in_bytes文档说明它们是intended for latency-bound (i.e., small-sized)对应类型的集合操作启用后触发针对性优化以降低执行时间typically its used in inference workloads。建议值区间为4~16Mb文档同时给出了换算4~16 * 1024 * 1024字节候选值范围是[0, 9223372036854775807]。计算密集型Compute centric 类这组 flag 文档给出的默认值与建议值基本一致适合作为核对项而非调优起点flag默认值 建议值候选值xla_tpu_enable_dot_strength_reductiontruetrue/falsexla_tpu_dot_dot_fusiontruetrue/falsexla_jf_enable_multi_output_fusiontruetrue/falsexla_tpu_scoped_vmem_limit_kib16384[4096, VMEM size of the architecture - 1024]xla_tpu_async_copy_bandwidth_scaling_factor1(0, 1]xla_msa_enable_cross_program_prefetch_freeingenabledenabled/disabledxla_tpu_msa_inefficient_use_to_copy_ratio0.5[0, 1]文档对这几个 flag 的用途说明分别是将非计算密集 dot 重写为 multiply reduce、执行 dot-dot fusion 以减少慢速/主存上的中间输出、启用多输出 fusion、设置每个算子可用的 scratchpad VMEM 大小单位 KiB、缩放 async copy 的有效带宽用于 prefetch 决策、启用跨程序 prefetch buffer 的释放优化、以及 VMEM 放置决策中判断分配点是否低效的 use/copy 字节比0 视为所有点高效1 要求使用字节数至少等于 async copy 字节数。一条最短的主路径示例以文档第一组 Pipelining flags 为例把三个 flag 设为建议值并运行程序$ XLA_FLAGS--xla_should_allow_loop_variant_parameter_in_chainkEnabled \ --xla_should_add_loop_invariant_op_in_chainkEnabled \ --xla_tpu_enable_ici_ag_pipeliningtrue \ path/to/your/tf/program如果同时是 v5e 上以推理为主的负载可以在同一XLA_FLAGS中追加一个 latency-bound 阈值例如取建议区间内的8 * 1024 * 1024字节$ XLA_FLAGS--xla_should_allow_loop_variant_parameter_in_chainkEnabled \ --xla_should_add_loop_invariant_op_in_chainkEnabled \ --xla_tpu_enable_ici_ag_pipeliningtrue \ --xla_all_gather_latency_bound_threshold_in_bytes8388608 \ path/to/your/tf/program这里8388608即8 * 1024 * 1024落在文档建议的4~16Mb区间内具体取哪个值需要你在自己的负载上实验确定——文档只给出区间没有给出固定推荐点。验证 flags 是否生效Inspect compiled programs 和 XLA Tooling 两节都说明把--xla_dump_to放进XLA_FLAGS即可把编译产物导出到指定目录。在上面的命令里加上该 flag$ XLA_FLAGS--xla_dump_to/tmp/generated \ --xla_should_allow_loop_variant_parameter_in_chainkEnabled \ --xla_should_add_loop_invariant_op_in_chainkEnabled \ --xla_tpu_enable_ici_ag_pipeliningtrue \ path/to/your/tf/program运行结束后/tmp/generated目录下会生成每个编译 cluster 的产物包括module_XXXX.*_optimizations.txt优化后的 XLA 程序文本每个 cluster 一份以及中间表示文件。对比调整 flag 前后两次 dump 中 HLO 指令序列的变化是文档提供的核对手段这些 dump 在提交 XLA bug report 时也是被要求附带的材料。文档给出的成功/失败判据很直接调整的是can significantly impact runtime performance的 flag收益需要在你的负载上自行测量而如果启用某个 flag 后出现 crashflags_guidance.md 的建议是回退到默认设置并创建 GitHub issue——也就是说flag 组合出问题时的官方动作不是继续尝试新组合。边界与不支持项以下限制来自文档原文调整 flags 时需要一并遵守v5e/Async 分组只针对 v5e 平台xla_enable_async_all_gather等 flag 的描述明确写着 activate asynchronous all-gather operations on v5e其他 TPU 版本上开启没有文档依据。Latency-bound 阈值主要针对推理文档反复注明 Typically its used in inference workloads训练负载是否适用文档没有说明。Memory flags 只在 HBM 报 out of memory 时才动xla_latency_hiding_scheduler_rerun、xla_memory_scheduler等 memory 类 flag 文档明确警告should only be adjusted if you encounter HBM out of memory errors during model compilation其他场景应保持默认因为改变它们可能反过来损害性能。成组 flag 不要拆开调Pipelining 三件套、v5e 两组 flags文档都用 should be used in conjunction 描述单独改其中一个不符合文档用法。xla_tpu_scoped_vmem_limit_kib的候选值上界与硬件相关候选范围写为[4096, VMEM size of the architecture - 1024]即上界取决于你使用的 TPU 架构 VMEM 大小文档没有给出统一数值。另外flags_guidance.md 的 TPU XLA flags 小节还列出了 5 个 TPU 专项 boolean flag如xla_tpu_enable_data_parallel_all_reduce_opt用于增加数据并行 all-reduce 的重叠机会、xla_tpu_spmd_rng_bit_generator_unsafe用于分区执行 RngBitGenerator 但会牺牲确定性它们属于按负载特征按需启用的选项默认文档未列建议值不建议与上面的性能组合一起盲开。下一步若需要理解编译器在不同优化等级下会启用哪些 pass参考 Effort LevelsEFFORT_O2是文档标注的生产负载默认档位。若 dump 出的 HLO 需要脱离原程序单独编译或定位某个 pass 的影响XLA Tooling 中的hlo-opt、run_hlo_module是文档给出的工具路径。【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考