最近朋友甩了张截图给我问我在 JAX 里跑同一个矩阵乘法有时候几十微秒μs就跑完了有时候却要几百毫秒ms中间差了整整三个数量级是不是 JAX 出 bug 了。这已经不是我第一次被人用 JAX 的时间单位问懵了。说实话ms毫秒和 μs微秒这两个单位放在 JAX 的语境里对应的根本不是同一类开销不理解这一点你测出来的性能数字多半也是自欺欺人。如果你正在用 JAX 做数值计算、深度学习训练或者正准备把代码从 NumPy / PyTorch 迁到 JAX这篇就是给你写的。我会从 JAX 的执行机制讲起把 ms 和 μs 分别从哪里来、怎么测才准确、哪些慢能优化、哪些慢是白费力气全部拆开讲清楚最后再列一遍我踩过的测量坑。看完你会发现JAX 并不玄学性能数字背后全是编译、调度和数据搬移这几件事在作祟。1. 先搞清楚ms 和 μs 在 JAX 里各自代表什么1.1 两个时间单位的体感差异先复习一下物理常识1 毫秒是 10 的负三次方秒1 微秒是 10 的负六次方秒1 毫秒等于 1000 微秒。正常人眨一次眼大约要 300 毫秒也就是说在 1 毫秒里你能完成的事情仅仅是眨眼的千分之三这已经是人类完全无法感知的尺度了。而微秒更夸张一个 3GHz 的 CPU 在一个微秒里大约能执行几千条指令光在真空中跑 1 微秒大约能走 300 米。落到 JAX 场景里这两个单位的含义差别很大。一次 jitted 函数调用里的小算子比如两个向量相加在 GPU 或 CPU 上稳态执行可能只要 5 到 50 微秒一个 1024 乘 1024 的矩阵乘通常在几十微秒到几百微秒之间而一次典型的深度学习训练 step常见耗时在 1 到 10 毫秒如果要从 CPU 往 GPU 拷贝几百 MB 的数据那轻松就是十几毫秒甚至更多。所以当你听到JAX 很快这句口号时必须先问一句他说的快是指编译后的稳态执行只要几微秒还是指整个流程包含编译、传输也只要几毫秒这两个数字背后的优化手段完全不同。1.2 一个容易混淆的搜索词JAX 与 JAX-RS搜jax这个词的时候经常同时出现jax rs很多新手会误以为这是同一个东西的两种写法。其实两者八竿子打不着。本文说的 JAX是 Google 和 DeepMind 生态里那个基于 Python 的高性能数值计算库主打 NumPy 风格的 API加上 grad、jit、vmap、pmap 这些可组合变换底层用 XLA 编译到 CPU、GPU、TPU而 JAX-RS 是 Java 社区里定义 RESTful Web 服务的一套规范就是写 Path、GET、POST 注解那套东西跟 Python 数值计算没有任何关系。顺带说一句JAX 这个名字本身并不是某个词组的官方缩写官方文档里也没有解释它的全称社区里有各种玩笑式的说法但你别把它和 JAX-RS 搞混就行。如果你搜jax是想查 REST 接口规范那这篇内容对你不适用如果你是想搞明白 JAX 的性能为什么忽快忽慢那咱们继续往下看。2. 为什么 JAX 的运行时间会横跨 ms 和 μs 两个量级2.1 JIT 编译毫秒级等待的真正来源JAX 最核心的机制是 jit即 just-in-time 编译。当你第一次调用一个被 jax.jit 装饰的函数时JAX 并不会直接像普通 Python 函数那样逐行执行而是先用 tracer 把函数的执行路径录下来生成一份中间表示 JAXPR再转换成 XLA 的 HLO 计算图经过算子融合、内存规划、指令调度等一系列优化最后编译成针对当前硬件CPU、GPU 或 TPU的原生内核。这一步非常重。一个中等复杂度的函数首次编译花几百毫秒到几秒都是常见现象复杂度高的模型编译几分钟也不稀奇。这就是毫秒级耗时的第一个主要来源它是一次性成本不是每次调用都要付的。编译产物会按输入数组的形状和 dtype 缓存下来第二次用相同形状调用时JAX 直接拿出缓存的 executable 执行单次开销瞬间掉到微秒级。我在实际项目里见过太多人被这个首次编译吓到第一次调用慢得离谱就以为 JAX 性能不行。其实只要把 warmup 做掉后面每次调用才是真正的稳态性能。记住这个结论ms 通常对应一次性开销μs 才是每次调用的稳态开销。2.2 异步派发微秒级数字测谎的关键JAX 在 GPU 和 TPU 上默认采用异步派发机制。调用一个 jitted 函数后函数会立刻返回一个数组对象但这个数组只是个占位符或者说 future真正的 kernel 还在设备上排队等待执行。如果你用 time.perf_counter 把函数调用包起来计时却不等它执行完测到的时间其实只是 Python 到运行时之间的派发时间根本不是计算时间。这会导致一个非常经典的陷阱测出来的数字虚低低到让你以为自己写出了性能怪兽。正确做法是在计时结束点调用 block_until_ready() 方法它会阻塞当前线程直到这个数组依赖的所有 kernel 都执行完毕。在 CPU 后端行为可能略有不同但养成无论什么后端都加同步的习惯能让你少踩很多坑。我经常跟朋友说一句话在 JAX 里说这个算子只要 5 微秒之前先问自己一句——我到底等它跑完了没有很多人说 JAX 跑得飞快其实是根本没等它真的跑完。2.3 算子融合为什么整段代码编译后反而更快XLA 在编译阶段会做算子融合把多个可以合并的计算合并成单个内核。比如矩阵乘 激活函数 缩放在底层可能被融合成一个 kernel中间结果根本不落内存。这意味着整个函数只有一次 kernel 启动、一次数据读写。而不做 jit 时jnp.dot 一次调度、jnp.tanh 又一次调度、逐元素乘法再一次调度每次都至少伴随一次 Python 层的调用开销和一次内核启动。在 CPU 上一次 Python 调度的开销大约几微秒GPU 上内核启动也要几微秒到几十微秒。如果一个函数里有 50 个逐元素步骤未 jit 就是 50 次调度jit 之后可能只有一个融合内核。这就是为什么看起来一样的计算逻辑有没有 jit单次耗时能从几百微秒差到几十微秒。3. 手把手实测把单次 JAX 调用时间量到微秒级3.1 用对计时工具perf_counter 与 Profiler测量 JAX 性能第一步是选对计时工具。time.time() 虽然简单但它取的是墙钟时间分辨率不够高还可能受系统时间调整影响测微秒级耗时基本是开玩笑。time.perf_counter() 是 Python 标准库里推荐的单调时钟高分辨率做短时测量就用它。timeit 模块封装了重复测定的逻辑也可以直接用但要注意 JAX 的计算必须配合 block_until_ready否则 timeit 只能测到派发时间。如果你需要看算子级别的详细耗时而不是整段函数的整体耗时那就该上 jax.profiler。通过 jax.profiler.start_trace 和 jax.profiler.stop_trace 可以导出 XLA 和 kernel 级别的 trace 文件在里面能看到每个算子的耗时、内存占用和依赖关系。这个工具适合你已经确认整体耗时异常、需要往下钻取的时候用日常基准测试用 perf_counter 就够了。3.2 一组可以复用的微基准模板下面这个 benchmarks 函数是我自己一直在用的模板逻辑很简单先 warmup再重复测若干次分别记录最小值、中位数和平均值。为什么要 warmup因为首次调用会触发编译必须让 JAX 把编译和缓存这一步先走完warmup 循环里的 block_until_ready 就是用来干这个的。import time import jax import jax.numpy as jnp def bench_jax(fn, *args, warmup10, repeat100): # 热身穿触发编译并等待执行完成 for _ in range(warmup): out fn(*args) out.block_until_ready() samples_us [] for _ in range(repeat): t0 time.perf_counter() out fn(*args) out.block_until_ready() t1 time.perf_counter() samples_us.append((t1 - t0) * 1e6) # 秒 - 微秒 samples_us.sort() min_us samples_us[0] med_us samples_us[repeat // 2] mean_us sum(samples_us) / repeat return min_us, med_us, mean_us注意这个模板假设你传入的 args 已经是设备上的数组。如果传入的是普通 NumPy 数组JAX 会在每次调用时先做 host 到 device 的拷贝那测出来的就不是纯计算时间而是传输加计算的总时间这在分析时是两码事。用这个模板去测一个 1024 乘 1024 矩阵乘可以看到非常典型的结果第一次 jitted 调用因为包含编译耗时可能在几百毫秒甚至几秒第二次开始同形状调用直接走缓存耗时降到几十微秒。这中间的差距就是 ms 和 μs 最直观的对比。3.3 实测数据怎么读我把常见场景的耗时段位整理成了一张表方便你建立量级直觉。注意这些数字不是标准答案只是量级参考不同硬件、不同 JAX 版本、不同线程数下会有明显差异但段位关系是不会变的。场景典型耗时说明第一次调用 jitted 函数100 ms ~ 数秒XLA 编译 优化 代码生成同形状再次调用小算子如向量加5 ~ 30 μs内核执行 派发开销同形状再次调用 1024x1024 矩阵乘30 ~ 200 μs视 CPU/GPU 而异未 jit 的单个 jnp.dot100 ~ 500 μs包含 Python 调度与内核启动大数组 host 到 device 拷贝10 ms 以上PCIe / NVLink 带宽限制一次典型训练 step1 ~ 10 ms取决于模型规模这张表最大的价值在于当你看到某个耗时是毫秒级时先判断它属于哪个环节——是编译、是传输、还是真正的计算。判断错了优化方向就会跑偏。4. 从数字到优化怎么把延迟从 ms 压到 μs4.1 能用 jit 就别裸跑最常见的优化就是把计算逻辑尽量放进一个 jitted 函数里让 XLA 把整段计算融合成少量内核。比如下面的例子如果拆开在 Python 里逐行跑jnp.tanh 一次、矩阵乘一次、逐元素乘法一次、求和又一次至少四次调度用 jit 包起来之后这些操作会被 XLA 融合规划调度次数大幅减少。jax.jit def fused_fn(w, x, b): h jnp.tanh(w x) return jnp.sum(h * b) 1.0实测下来一个由十几个逐元素操作组成的函数未 jit 时单次调用可能上百微秒jit 之后能压到十几微秒。这就是为什么我一直强调JAX 的推荐用法是把函数当整体编译而不是像 NumPy 那样一行一行地调用算子。4.2 减少 host 和 device 之间的数据搬移如果每次调用 jitted 函数时都传 NumPy 数组JAX 会先执行 device_put 把数据搬到设备上。在 GPU 后端一次传几十 MB 数据就是毫秒级开销这个耗时会被算进你的计时里。正确的做法是把固定数据提前放到设备上循环里复用同一个设备数组而不是反复从 CPU 侧往里塞数据。x_dev jax.device_put(x_np) # 只搬一次 for step in range(1000): loss train_step(x_dev) # 后续不再触发传输这条规则在很多新手代码里被忽略。数据搬移是毫秒级成本混进微秒级路径的典型元凶优化收益往往比调算子还明显。4.3 把 Python 循环换成 JAX 原语Python 循环里反复调用 jitted 函数每一轮都有 Python 到运行时的调度开销通常是几十微秒。如果循环 1000 轮光调度就吃掉几十毫秒。正确的思路是让循环发生在 jit 的内部用 jax.lax.fori_loop、jax.lax.scan 这类原语替代 Python 的 for 循环或者用 jnp 的向量化操作、vmap 来替代显式循环。举个最简单的例子如果你要对一个序列做累乘用 Python 循环逐次调用 jitted 函数和用 lax.scan 在编译后的内核里做循环单次迭代的耗时差距能从几十微秒降到几微秒。这个优化对序列模型、RNN、迭代算法这类场景尤其重要。4.4 别让形状抖动触发重复编译JAX 的编译缓存是按输入数组的形状和 dtype 做 key 的。batch size 从 32 变成 33哪怕只差 1JAX 也会认为这是不同的输入签名触发布局全新编译然后又来一次几百毫秒的开销。如果你观察到一个函数耗时在极低和几百毫秒之间反复横跳十有八九就是形状或 dtype 在变。对策也很朴素能固定形状就固定形状需要变长输入就统一 padding 到固定长度。另外注意 JAX 默认是 float32如果你开了 jax.enable_x64很多消费级 GPU 的 fp64 计算会慢到令人怀疑人生这也是一个容易被忽略的毫秒级陷阱。5. 常见测量翻车现场与排查策略5.1 翻车一没做 warmup数字虚高这是新手最容易犯的错误。直接拿 time.perf_counter 包住 jitted 函数的第一次调用测出来的几百毫秒其实是编译时间跟函数本身的性能没有关系。如果 repeat 次数又很少平均值会被这个编译耗时严重污染得到一个毫无参考价值的数字。对策是 warmup 至少 5 到 10 次再进入正式计时循环统计时优先看最小值和中位数不要被平均值骗了。最小值代表系统在最优状态下能达到的延迟中位数代表常见状态平均值则容易被偶然的 GC、后台任务拖高。5.2 翻车二忘了 block_until_ready数字虚低虚高容易发现虚低才可怕。因为 JAX 是异步派发的如果你这样写计时t0 time.perf_counter() out jitted_dot(x, y) # 错误不等计算完成 t1 time.perf_counter() print(f{(t1 - t0) * 1e6:.1f} μs) # 数字很漂亮但它是假的测到的只是派发时间kernel 可能还没开始执行。正确版本是在计时结束点加一行 out.block_until_ready()。判断自己有没有中招很简单如果你的耗时数字小得离谱而且不管数据量从 1024 涨到 4096耗时几乎不变那基本就是没做同步。5.3 翻车三拿平均值当真相被环境噪声骗了笔记本的省电模式、CPU 降频、后台进程抢占、GPU 共享都会让单次测量的抖动非常大。我见过有人拿 mean 值做基准结果一次系统更新后所有数字都变差了其实只是噪声。正确处理是取多次采样后的最小值或者中位数并且尽量在机器负载稳定的时候测。GPU 场景还要注意第一次调用某个形状的 kernel 可能触发 cuDNN 的 autotune所以 warmup 次数要足够多。5.4 问题速查表下面这张表是我在实际排查中总结出来的遇到异常的 ms/μs 数字可以先对着表自查一遍。症状可能原因对策每次调用都几百 ms且伴随进程重启首次编译固定形状 warmup耗时在极低和插几百 ms 之间横跳输入形状或 dtype 变化padding 固定形状数字小得离谱且不随数据量变化没 block_until_ready计时结束点加同步平均耗时远大于最小耗时后台任务、GC、降频取 min / medianGPU 上整体偏慢反复传 NumPy 数组device_put 后复用Python 循环内多次 jit 调用调度开销叠加用 lax.scan / fori_loop6. 一点个人心得我在 JAX 上被快这个字骗过很长时间。刚上手那阵子我用一个没加 block_until_ready 的计时脚本测出来的漂亮数字兴奋地以为自己写出了一个微秒级神算子直到后来把同步加上真实耗时露出了真面目。踩过这些坑之后我对 ms 和 μs 建立了比较踏实的直觉它们不是对立的量级而是分别对应两类完全不同的开销——毫秒级通常是编译、传输、初始化这类低频高成本的一次性投入微秒级才是内核执行的高频低成本稳态开销。优化 JAX 程序核心不是把 μs 压成更小的 μs而是别让毫秒级的事混进每一轮的执行路径里。最后分享一个小技巧把 bench_jax 这种模板存成自己工具库里的公共函数顺手把硬件型号、JAX 版本、x64 开关都记在输出里。这样你的性能数字才是可复现的而不是每次换个环境就变成另一个故事。真正会优化 JAX 的人都是从先把时间量准开始的。