JAX jit 中返回值顺序为何会显著影响性能？

2026-11-06数据小许

结论

仅凭这段伪代码和两组耗时，无法确定 638 微秒和 73 微秒的差异来自哪里。返回值顺序不会改变函数的数学计算，但会改变编译器看到的输出结构。启用 buffer_state donation，且函数涉及 buffer 和模型更新时，这种变化可能影响输入输出缓冲区的别名关系、内存规划或生成的设备程序。要定位原因，需要比较两种写法编译出的程序，并测量同步后的执行时间。

为什么返回顺序可能产生影响

jax.jit 会根据输入和输出结构编译函数。启用 donation 后，JAX/XLA 在满足条件时可以用被捐赠的输入缓冲区承载某个输出。输出位置也是输出结构的一部分；如果有多个兼容输出，调整顺序可能影响编译器选择的别名关系和内存规划。涉及动态更新、采样或较大状态时，这些差异可能影响执行时间。

这只是对现象的一种可能解释，尚未确认根因。伪代码没有说明 buffer 的实现、实际函数签名、JAX/XLA 版本或完整的基准测试方法。仅凭计时数字，也无法断定某个返回位置一定更快。

先确认计时测到的是设备执行

JAX 的设备调用可能是异步的。如果基准只测 Python 调用提交的时间，或只在一批调用结束后同步，结果反映的可能是排队和提交开销，而不是每次执行的实际耗时。两种写法应使用相同的预热、输入、迭代次数和同步方式，并在每次计时调用后等待结果就绪：

import time
import jax

# jitted_fn 是其中一种写法；args 是与其签名匹配的输入。
for _ in range(warmup_steps):
    out = jitted_fn(*args)
    jax.block_until_ready(out)

start = time.perf_counter()
for _ in range(num_steps):
    out = jitted_fn(*args)
    jax.block_until_ready(out)
elapsed = time.perf_counter() - start
print(f"平均耗时：{elapsed / num_steps * 1e6:.1f} 微秒")

如果 buffer_state 是被 donation 的输入，调用后原输入可能失效。循环中应将返回的状态作为下一次调用的输入，并按各函数的返回顺序正确解包，避免重复使用已捐赠的旧状态。两组测试也应使用等价的初始状态，避免某一组额外承担数据准备开销。

比较编译结果和设备执行

使用相同的输入结构分别降低两个 jitted 函数，检查编译器 IR 在缓冲区别名、复制、布局或执行算子方面是否存在差异。IR 的具体内容和可用接口可能随 JAX/XLA 版本变化：

lowered = jitted_fn.lower(*example_args)
print(lowered.compiler_ir(dialect="hlo").as_hlo_text())

也可以用 JAX profiler 记录有代表性的设备执行，查看差异来自设备计算、内存复制还是同步等待。比较时应覆盖完整计算路径。问题描述中，去掉模型更新或把 buffer 换成普通数组后，差异就消失了。这提示可以重点检查这些操作与状态缓冲区的交互，但还不能据此确定具体是哪一步造成差异。

建议的排查顺序

确保两种函数的运算、输入、预热、同步和计时区间完全一致，并重复多轮测试。
确认两种写法使用相同的 donate_argnums，且传入的 buffer_state 结构和输出结构符合预期。
比较两种写法的 HLO 或 profiler 结果，关注设备复制、缓冲区复用和执行时间的差异。
保留模型更新，逐步缩减 buffer 操作，找出差异首次出现的部分。每次只改动一处，并保持其他条件不变。
如果仍无法定位，记录实际的 JAX/XLA 版本、设备类型、完整的最小示例和同步后的重复计时数据，再分析编译与运行行为。

确认原因之前，不建议只根据返回顺序调整生产代码。如果两种写法生成的设备程序和稳定计时结果确实不同，可以把较快的写法作为当前版本和设备上的实测优化，并在目标运行环境中重新验证。