Signal Brief

JAXBench与Pallas内核优化

Google、哈佛和UC Berkeley联合发布JAXBench,包含50个基于真实MaxText架构的JAX工作负载。使用Gemini 3 Flash并配合文档条件化,将正确率从5.8%提升至37.3%,解决48/50个基准,速度提升1.28倍,波束搜索进一步优化至1.36倍。

twitter关注列表 elvis (@omarsar0) 发布 2026-07-27 收录 2026-07-27 观察

一句话判断

首次提出面向TPU的JAX内核优化基准JAXBench,并展示了文档条件化带来的显著提升,值得TPU优化领域从业者阅读原文。

核心信息

Google、哈佛和UC Berkeley联合发布JAXBench,包含50个基于真实MaxText架构的JAX工作负载。使用Gemini 3 Flash并配合文档条件化,将正确率从5.8%提升至37.3%,解决48/50个基准,速度提升1.28倍,波束搜索进一步优化至1.36倍。

原始内容

Great technical paper from Google. Great read on why context beats scale for agents working against unfamiliar APIs. (bookmark it) GPU kernel optimization has KernelBench to hillclimb on. TPUs had nothing, and the Pallas DSL is documented thinly enough that models mostly guess. New research from Google, Harvard, and UC Berkeley introduces JAXBench, 50 JAX workloads built from real MaxText architectures like Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, and AlphaFold2. Eight operators ship with hand-tuned Pallas kernels from Tokamax, so agent output gets measured against expert work instead of a naive baseline. With Gemini 3 Flash, conditioning on curated TPU documentation lifts per-sample correctness from 5.8% to 37.3%, solving 48 of 50 benchmarks at a 1.28x geomean speedup. Beam search then pushes it to 1.36x. Correctness turned out to be a documentation problem and speed turned out to be a search problem. Paper: https://t.co/8lOLtFHlfq Learn to build effective AI agents in our academy: https://t.co/1e8RZKs4uX ![photo](https://pbs.twimg.com/media/HOMUhixbQAA0KQ0.jpg)

相关动态

02

NVIDIA 研究:AdamW 在大批量训练中存在规模上限

NVIDIA 研究发现,AdamW 在批量大小达到 1 亿 token 时训练稳定性下降,而 SOAP 和 Muon 保持稳定。团队通过每步 QR 正交化和改进预处理策略消除了 SOAP 在大批量下的 loss spikes,并在多十亿参数模型上训练数万亿 token 后,SOAP 和 Muon 一致优于 AdamW。

twitter关注列表2026-07-26#模型#技术#研究
观察
03

Kimi K3: 2.8万亿参数开源模型发布

Kimi.ai 发布 Kimi K3 模型,拥有 2.8 万亿参数、100 万上下文、原生多模态能力。采用 Kimi Delta Attention 技术,在百万 token 上下文中实现高达 6.3 倍解码加速;Attention Residuals 以低于 2% 的额外成本带来约 25% 的训练效率提升。模型已上线官网、Kimi Wo...

twitter关注列表2026-07-26#大模型#模型发布#技术突破
值得跟进
04

From Memory to Skills: Evidence-Grounded Co-Evolution Governance for Long-Horizon LLM Agents

论文指出大多数代理记忆系统只保存事件并在上下文召回,导致重复推理错误。作者提出MSCE方法,将经验组织为声明性事实、诱导程序化策略和基础步骤轨迹,仅将重复模式提升为可调用的过程(如pip安装失败时运行规则),每个策略需具备触发条件、边界、正向收益估计和证据支持。

twitter关注列表2026-07-26#AI#研究#技术
观察
05

Understanding Reasoning from Pretraining to Post-Training

论文通过象棋实验(5M-1B参数模型,36种预训练-RL组合)发现:预训练验证损失能高精度预测后RL的pass@1(相关性从0.93到0.99),RL计算量每10倍增益与预训练token数对数正相关(r=+0.84);RL最优计算量份额在50M参数时为20%,680M时为28%,且RL会放大模型在困难问题上的错误模式。

twitter关注列表2026-07-26#研究#模型#技术
观察