观潮 · 模式预报与气象 AI 观察
Fortran 物理过程搬上 GPU：LAFT 让四个大模型翻译 Kessler 暖雨方案，最快版本比串行 Fortran 快 357 倍
> CESM、WRF、E3SM 里的物理过程包大多是几十年迭代下来的 Fortran，而把它们搬到 GPU 上要靠既懂加速器编程、又懂内存布局和编译器行为的人手工重写。这篇提出的 LAFT（LLM-Assisted Fortran Translation）不把翻译当"一次提问出代码"，而是拆成五段：确定性的 Fortran 前端解析、桥接层生成、策略引导的 LLM 翻译、对 Fortran 参考实现的数值校验、以及 XLA 层面的性能剖析反馈。在 Kessler 暖雨微物理这个算例上，Claude Sonnet 4.6、GPT-5.4 Thinking、Gemini 3.1 Pro、Qwen2.5-32B 四个模型生成的 JAX 版本都复现到了接近浮点精度，但 GPU 性能差得很远：Gemini 那份在 10⁶ 个大气柱的缓存执行下比串行 Fortran 快最多 357 倍，Qwen 那份数值也对，却因为保留了对列的顺序循环，五次针对性修复都没改掉。作者的落点是：数值正确是必要条件，不是充分条件。
缺的是哪一块
自动代码翻译在软件工程里做了几十年，从 TXL 这类规则变换系统，到基于静态分析、中间表示和领域重写规则的方案。它们对翻译行为控制得很死，代价是要为每个源语言—目标语言组合手写变换逻辑。大模型换了个思路：能吃更宽的上下文，写出的目标语言代码更像人写的。已有研究也证明，在提示里给足上下文能改善翻译质量。
问题在于，大模型生成的代码不保证能编译、不保证语义等价、更不保证跑得快。已有的实证研究报告过编译错误、运行时错误、功能错误、依赖缺失、API 映射错误、数据类型错误等各类语义缺陷。这类风险在科学计算代码上尤其致命——一个很小的语义漂移就能产出科学上无效的结果。
更难的一点在于本文想说的那层：就算数值完全对，翻译出来的代码在 GPU 上也可能慢得没法用。原因不在模型"翻译错"，而在生成代码没有把并行性暴露给 XLA。所以真正缺的不是更好的提示词，而是一套把上下文、约束、校验和执行反馈都管起来的工作流。
LAFT 的五段流程
LAFT 的出发点是不直接采信模型输出：翻译结果必须对着原始 Fortran 参考实现验证，再用 GPU 与 XLA 的剖析结果去改。五个阶段如下：
阶段 · 做什么 · 为什么不由模型做
1 · 确定性元数据抽取 · 用 Tree-sitter 的 Fortran 语法解析源码，还原过程边界、声明、控制流、模块结构、源码范围；抽取带类型的实参签名（数据类型、数组秩、符号化维度表达式、intent），模块层面记录 USE 依赖、模块级变量、可访问性 · 程序结构应该是显式、可复现的，不能交给模型"临场理解"
2 · 依赖分析 · 建模块依赖图，识别循环依赖并分层，无前置依赖的模块先翻 · 同上
3 · 桥接层生成 · 从前端元数据确定性生成桥接代码，把列主序的输入在接口处一次性转成 JAX 期望的行主序；二维数组在 `jnp.transpose` 后进核函数，返回前再转回来，两次转置都在 JIT 边界内，XLA 可以把布局置换和核计算一起融合 · 布局转换是可以被明确定义并单独验证的事，不该混进每个核函数的翻译里
4 · 策略化 LLM 翻译 · 每个过程走四遍提示：初始生成、JIT 安全性检查、面向 XLA 的向量化重构、最终一致性检查 · —
5 · 校验与剖析反馈 · 用 MAE 对 Fortran 参考实现比对；再用 XLA trace 查过度核函数启动、意外重编译、阻止并行的顺序循环结构 · —
拆成四遍提示是有依据的：上下文窗口大不等于模型能均衡用到提示的每个部分，已有工作表明相关信息出现在长上下文中部时模型表现会退化。科学代码翻译里，变量定义、依赖关系、数组维度、可变性、内存布局假设、控制流约束这些关键细节很容易被埋掉；模型漏掉任何一条，生成的代码看着语法正确，数值或性能却已经错了。所以拆分的目的是让每个子任务更窄、更可检视、更容易单独验证。
桥接层还留了几条明确的边界：一维数组原样通过，标量不转换，整数或逻辑标量在 JAX 追踪需要时标记为 static，实数数组统一转 64 位浮点以对齐 Fortran 双精度；字符参数留在主机侧，因为字符串过不了 JIT 边界；模块变量显式建模成 INOUT 实参，由驱动方显式管理模块状态，不依赖隐藏的全局状态。
四个模型翻译出了四份不同的执行结构
算例用 Kessler 暖雨微物理，来自 ESCOMP 大气物理代码库。它只有水汽、云水、雨水三个水成物、不含冰相，每根大气柱调用一次、每个模式时间步更新温度和湿度。选它的理由是"麻雀虽小"：既小到可以做逐项校验和剖析，又包含垂直循环、列结构、阈值条件更新、时间子循环、多个相互作用的预报与诊断变量——这些正是大参数化里也有的难点。
四份翻译都过了静态的 JAX 质量检查，也都在数值上通过了校验，但生成的执行结构不一样：
模型 · 列方向组织 · 垂直层更新方式 · XLA 操作数 · 唯一操作数
Claude Sonnet 4.6 · `jax.vmap` 跨列 · 对选定层的 `jnp.arange` 子集表达式，独立的边界索引逻辑，gather–compute–scatter · 292 · 39
GPT-5.4 Thinking · `jax.vmap` 跨列 · 同上 · 292 · 39
Gemini 3.1 Pro · `jax.vmap` 跨列 · 对整层数组用 `jnp.arange(nz)`，边界用 `jnp.clip` 与 `jnp.where` 处理 · 138 · 25
Qwen2.5-32B · `lax.fori_loop` 顺序扫列 · 列内向量化，`jnp.arange` 加 `.at[klevs].set` 索引更新 · 远多于前三者 · —
差别集中在怎么写垂直层。Claude 和 GPT 只对选定的层做子集运算，于是需要单独的标量边界更新和 scatter 式索引写入，XLA 操作数因此翻倍。Gemini 改成对整个垂直数组做全数组掩膜，边界交给裁剪和条件选择，不需要单独的标量边界更新，也不需要 scatter 式写入，操作数大约只有前两者的一半。Qwen 在列内做了向量化，但把子循环嵌进了列循环，于是每根柱各进入一次子循环，列级并行被彻底放弃。
数值对得上，性能差多少
性能口径是缓存后的 JAX 执行时间，含派发、设备执行与同步，不含编译和主机—设备传输；桥接开销单独测，分主机到设备的传输、输入输出的置换、设备回主机三段，按 C 连续和 Fortran 序两种调用方分别报总和。这个区分是有意的：核函数一旦数据上了 GPU 可以很快，但真实部署还得算布局转换和传输。
串行 Fortran 基线随列数近线性增长，从 2.42 ms 一路涨到 54,546.96 ms。两个 GPU 版本在最小算例下都低于 1 ms，且随规模增长平缓得多。Gemini 那份在 10⁶ 个大气柱的缓存执行下比串行 Fortran 快最多 357 倍。
与手工加速的 OpenACC GPU 版本对照，结论不是"谁更好"，而是看规模：
列数区间 · OpenACC · Gemini JAX · 谁快
小规模 · 0.218 ms · 0.74 ms · OpenACC
中等规模 · 2.96 ms · 3.05 ms · 基本持平
100,000 列 · 54.38 ms · 16.30 ms · Gemini JAX
最大规模 · 751.16 ms · 152.64 ms · Gemini JAX
OpenACC 开销低、在小列数上占优；Gemini 那份全数组公式在大批量执行下更能吃到融合的红利。一旦规模大到能把 JAX/XLA 的固定开销摊薄，它的扩展性就超过 OpenACC。作者明确写了：这不意味着某个 GPU 编程模型全面更优，优势取决于问题规模以及代码把并行性和融合机会暴露得多好。
作者承认的限制
- Qwen 那份的顺序列循环没修掉。它在数值上通过了校验，但在连续五次针对该问题的性能修复后仍然保留着顺序列循环。作者把它留在结果里，就是为了记录这种执行设计及其性能后果。
- 性能规模区间是分段给的。文中按列数从 50 起、垂直层数固定的设定给出上述各档耗时，没有给单一拟合曲线。
- 目前只做了 Kessler 一个方案。作者明确把它定位为初始算例；P3 微物理和 Droplet 方案的实验还在进行，会带来更多过程、更丰富的控制流、更强的跨过程依赖、查表、沉降逻辑，以及更复杂的桥接与校验需求。
- 桥接层还不完整。当前的生成器对所有非字符数组套用同一套转换，尚未保留整数与逻辑数组类型——只是 Kessler 接口用不到所以没暴露问题。另外，按列主序存放的数组（真实 Fortran 宿主模式给的就是这种）目前在上机前还要多一次主机端拷贝。
- 性能口径不含编译与传输。报告的是缓存执行时间；桥接开销单独列表给出，但不在同一张性能图里。
- 可微性没做。JAX 本身支持自动微分，但本文只谈数值正确性和 GPU 性能。
- 接入真实模式框架仍是进行中工作。初步研究表明，LAFT 翻译出的 Kessler 版本桥接进 EAMxx 大气过程后，复杂度和性能与原 Fortran 实现相当——作者用的是"初步研究"这个措辞，不是已完成的结论。
- 提示策略还没吃到剖析反馈。目前剖析结果只用于翻译后的修复；让剖析结论反过来改写提示策略（显式鼓励全数组掩膜、抑制 scatter 密集更新和顺序列迭代）被列为下一步。
一句话判断
把 Fortran 物理过程交给大模型翻译这件事，难点已经不在"译得对不对"——四个模型在校验驱动的工作流里都对到了接近浮点精度——而在"译出来的代码把并行性交给编译器了没有"。LAFT 的价值在于把确定性的部分（解析、桥接、校验、剖析）从模型里拿出来自己管，只把真正需要生成的核函数留给模型；357 倍那个数字来自 Gemini 生成的全数组写法，而不是来自"用了 JAX"这件事本身。