给一条工单选部门、判断是否需要退款、给紧急程度打分,这些任务需要的是几个结构化判断。模型可以直接输出选项和概率,不必逐 token 生成一段回答。
Laya 就是这样一个 Decision Model:用一次前向计算处理选择、评分和是非问题。我把它的推理路径移植到了 JAX,做成了 deciax,让这类模型可以在 TPU 上运行,再用 Pallas 优化其中的计算。
这次测试里1,单芯片 TPU v6e 的英文模型单问题调用 p50 是 2.18 ms,多语言模型是 1.55 ms。作为对比,A100 SXM4、H100 PCIe、H200 NVL 的最快实测结果分别是 6.4 / 4.0 ms、4.8 / 3.2 ms、5.8 / 4.1 ms(英文 / 多语言)。
每次 50 个问题时,v6e Pallas 的 p50 是 33.69 / 17.83 ms,相较普通 JAX 实现提速约 2.3 倍;A100、H100、H200 的最快实测结果分别是 68.8 / 36.2 ms、44.0 / 23.2 ms、39.8 / 21.2 ms。

