工程 · 2026 年 9 月 23 日

Kimi K3 达到 700 TPS：TPU Megakernel 的实践案例

作者：George Novack、Xuting Liu、Jeff Ma 和 Woosuk Kwon

今天，我们发布了 inferact/tpu-megakernels，这是一个面向 TPU v7 的 megakernel 集合。我们的 Kimi K3 实现在采用推测解码时可实现超过 700 tokens/s，而 GB200 上的对应成绩为 452 tokens/s。在不使用推测解码的情况下，我们针对 K3 和 Qwen 3.8 27B 的 megakernel，在 batch size 1 至 8 时可达到 GB200 基线解码吞吐的约 1.4 至 2 倍。

模型	TPU v7	GB200
Kimi K3（tokens/s）	709（16× TPU v7）	452（16× GB200）
Qwen 3.8 27B（tokens/s）	1,515（4× TPU v7）	695（4× GB200）

• TPU megakernel
• vLLM 基线

在接受长度（acceptance length）为 6 的推测解码条件下，每秒输出的 token 数。

解码速度受硬件峰值内存带宽的限制。我们的 megakernel 通过有效利用 TPU 大容量片上 VMEM、顺序编程模型，以及用于跨层预取权重的显式异步流水线，使端到端性能更接近这一理论峰值。本文将介绍我们的设计，以及 megakernel 与 TPU 为何十分契合。

背景：TPU 与 GPU

如果只看顶层规格，Google 的 TPU v7 Ironwood 芯片与 NVIDIA 的 GB200 GPU 大致相当。

规格	Google TPU v7	NVIDIA GB200（每张 GPU）
峰值 BF16 算力	2.31 PFLOPS	2.5 PFLOPS
峰值 FP8 算力	4.61 PFLOPS	5 PFLOPS
HBM 容量	206 GB	186 GB
HBM 带宽	7,380 GB/s	8,000 GB/s
互连	1,200 GB/s（ICI）	1,800 GB/s（NVLink 5）

来源：Google Cloud TPU7x 文档、NVIDIA GB200 NVL72。

但这一苹果对苹果的比较掩盖了 TPU 最有趣的特性之一：其片上内存的形态和容量。解码速度主要取决于能否高效地将数据从 HBM 移入高速片上内存，因此这些维度是关键设计考量。

与 GPU 的多级内存层级不同——GPU 包含硬件管理的 L1/L2 缓存，以及由程序管理的 Shared Memory 和 Tensor Memory 空间——TPU 的内存层级更为简单。来自 HBM 的数据会被传输到一个称为 Vector Memory（VMEM）的大型片上 SRAM 池；TPU 的向量寄存器（VREG）则直接由 VMEM 填充。重要的是，VMEM 中所有数据的生命周期完全由我们的 TPU 程序管理：当我们将数据移入 VMEM 后，它会一直保留在那里，直到我们明确用其他内容将其覆盖。

NVIDIA GB200（每张 GPU）	TPU v7
HBM：186 GB	HBM：206 GB
程序发起的 DMA
L2 缓存：硬件管理
SM：152 个	VMEM：每个 TensorCore 64 MiB
共享内存：228 KB；Tensor Memory：256 KB	单一内存池，软件管理；程序决定其中保存什么内容及其保留时长
一个 block 仅能访问自身 SM 的共享内存
Tensor Cores / CUDA Cores	MXU / VPU / XLU
约 111 MiB SRAM，被分割到 152 个 SM	每颗芯片 128 MiB，作为两个内存池寻址

每个 TensorCore 拥有 64 MiB VMEM，因此我们可以预取大多数模型权重张量，确保在任意操作需要使用它们时，这些数据已经被置于高速片上 SRAM 中。GPU 的 Tensor Memory（TMEM）组织方式则不同：整张 GPU 约有 38 MiB TMEM，但其容量分配给全部 152 个 SM；每个 SM 仅有 256 KB，几乎没有余量提前暂存权重。

为什么选择 megakernel？

大多数 kernel 都有大致相同的结构：从 HBM 加载权重和激活值，执行计算，再将结果写回。只有加载阶段会从 HBM 拉取大量数据，在其他阶段带宽会处于未充分利用状态。设计良好的 kernel 会通过细分并重叠各个阶段来缓解这些带宽气泡，但改善程度有限。单个解码步骤可能会依次调度数百个这样的 kernel，这些气泡会不断累积。

LOAD：加载 n、加载 n+1、加载 n+2

COMPUTE：设置、等待加载 n、计算 n、计算 n+1、计算 n+2

STORE：写入 n、写入 n+1、写入 n+2

预热阶段：尚未发起加载；排空阶段：最后一次加载已发起，计算和写入完成；稳态：加载 n+1 与计算 n 重叠。

时间 →（不按比例）

• HBM 流量
• 计算
• 空闲（气泡）

GPU 已经提供了应对这一问题的方法。CUDA Graphs 可降低主机端启动开销，stream 可让独立分支并发执行，而程序化依赖启动（programmatic dependent launch，PDL）则允许依赖 kernel 在其生产者完成前就开始准备和加载权重。这些机制在保留独立 kernel 的同时恢复了部分重叠。

megakernel 则移除了这个边界。当整个解码步骤都位于同一个程序中时，权重加载不再绑定于使用该权重的 kernel——只要片上存储允许，就可以尽早发起加载。此时约束变成缓冲区的生命周期和容量。在 TPU v7 上，大容量 VMEM 和显式异步传输使我们能够在当前层仍在计算或通信时，暂存下一层的权重。目标是在一个以搬运权重为主要成本的解码步骤中，让 HBM 接口始终保持忙碌。

现有 TPU 实现依赖 XLA 执行融合和调度，但寻找最优的跨层流水线，需要对存储复用、传输时机和同步作出联合决策。自动优化未必能发现我们希望的调度方案。Pallas 让我们可以直接表达这些决策，同时将指令调度留给编译器。

使用 Pallas 编写 TPU megakernel

Pallas 是 TPU 的 kernel 语言。据我们所知，这是首个使用 Pallas 编写的开源推理 megakernel，我们发现该语言非常适合这项工作。以下是本项目中最重要的 Pallas 特性。

无 Grid 的 kernel 启动

Pallas 语言可用于为 TPU 和 GPU 编写 kernel。因此，可以以任何 GPU kernel 工程师都熟悉的方式编写 Pallas kernel：将 kernel 作为一个程序网格（grid）启动，每个程序接收并处理输入的一个特定 tile。

但 TPU 遵循一种本质上单线程的执行模型。即使我们将 kernel 作为 grid 启动，每个单独的程序仍会一次一个地顺序运行。尽管基于 grid 的编程模型有其优势，我们发现，对于 megakernel 开发而言，拥抱 TPU 的单线程特性能够同时提升可读性和细粒度控制能力。我们的 megakernel 作为一个没有显式 grid 的单一程序启动。

有作用域的 VMEM 分配

我们的 megakernel 必须在激进预取模型权重与每个 TensorCore 有限的 64 MiB VMEM 池之间取得平衡。一些数据，例如激活值，会在整个前向传播过程中一直驻留在 VMEM 中；另一些数据，例如 MLA 使用的 KV cache，则具有很短的生命周期。Pallas 的 run_scoped 原语可对特定 VMEM 分配的生命周期进行细粒度控制。

这简化了程序员的工作：我们不必显式管理任意时间点 VMEM 中哪些字节范围可以安全读写，而是可以编写自包含的“小型 kernel”，精确声明其所需 VMEM 的形状和大小。

异步 DMA

HBM 与 VMEM 之间的异步 DMA 是 megakernel 成立的关键。我们的实现大量使用 make_async_copy(...).start() 和 make_async_copy(...).wait()，以指定应在何时启动权重预取，以及何时暂停执行直至预取完成。在 JAX 中管理操作间的异步依赖关系向来颇为棘手：XLA 会将操作链作为一个整体进行调度，并保留重排操作或回收其缓冲区的权利。

在单个 kernel 内，这个问题不会出现：我们启动并等待的复制操作永远不会跨越 XLA 可以自由调度的边界。

我们的 Kimi K3 megakernel

Kimi K3 megakernel 将 K3 的 92 个 Mixture-of-Experts（MoE）层组合为一次 Pallas 调用，并根据层索引有条件地选择 Kimi Delta Attention（KDA）或 Multi-head Latent Attention（MLA）操作。在 kernel 内，每个权重张量都会通过在更早阶段发起的传输被暂存到 VMEM 中，因此理想情况下，当其消费者运行时，它已经驻留于此。由于层循环位于 kernel 内部，异步复制可以跨越层边界。例如，在第 N 层仍处于 MoE 阶段时，第 N+1 层的注意力投影权重就开始传输。下方展示了跨层预取策略的简化视图。

COMPUTE：第 N 层通信、注意力、路由器、共享专家、路由专家、通信、注意力。

HBM → VMEM：第 N 层注意力、共享专家、路由专家、第 N+1 层注意力。

第 N+1 层注意力传输在第 N 层期间发起，并在第 N+1 层需要它之前驻留于 VMEM。

时间 →（不按比例）

• HBM → VMEM 传输
• 矩阵计算
• 向量计算
• 集合通信

该 kernel 将模型切分到 16 颗 TPU v7 芯片上，合计 32 个 TensorCore。注意力头分布在全部 32 个 rank 上，而路由专家采用 TP4 × EP8。MoE 的 latent-to-hidden 投影在每台主机内部（8 个 TensorCore）切分，并跨主机复制，以额外读取权重为代价换取成本更低的集合通信。

为特定目标请求并发度调优 megakernel，需要探索更具定制性的跨芯片并行方式，超越人们熟悉的张量并行、专家并行和数据并行策略。例如，即使解码 batch size 从 1 看似仅小幅增加到 8，也暴露出以残差流向量算术为主的操作中的新瓶颈，例如注意力残差和 RMSNorm。为此，我们将张量并行与残差流序列并行相结合，除了在 TPU 间切分矩阵乘法外，也对向量算术进行切分。

megakernel 方法还有一项我们未曾预见的优势：编译速度快。事后看来这很合理：kernel 作者手工且只需一次性支付排序和重叠的成本；否则，XLA 必须在每次新的编译中、跨每一个操作重新发现这些关系。在 TPU 上，编译由数百或数千个 XLA op 构成的大模型通常需要 30 分钟以上；相比之下，我们完整的 megakernel 从零开始编译不到 90 秒，支持更快速的开发和实验周期。

DSpark

随着支持 batch size 8，该 megakernel 还可在使用 DSpark 推测解码时充当验证器。每次 kernel 启动会在单一步骤中评估 1 个锚点（anchor）和 7 个提议续写。根据具体工作负载，草稿模型通常每步会产生 3 至 6 个被接受的 token。每个解码步骤约为 8.5 ms 时，该 megakernel 可为单个用户提供超过 700 tokens/s 的吞吐。

接受长度	Kimi K3 megakernel（16× TPU v7）	vLLM 基线（16× GB200）	加速比
3	350 accepted output tokens/s	229 accepted output tokens/s	1.53× 更快
6	709 accepted output tokens/s	452 accepted output tokens/s	1.57× 更快

• Kimi K3 megakernel — 16× TPU v7
• vLLM 基线 — 16× GB200

使用 RedHatAI/Kimi-K3-speculator.dspark speculator 时，不同接受长度下的推测解码吞吐。

结果

我们在 16 颗 TPU v7 Ironwood 芯片上评估该 megakernel，并采用 vLLM 已发布的、运行于 16 张 GB200 GPU 上的 Kimi K3 配方作为基线。在不使用推测解码时，megakernel 在 batch size 1 至 8 的每一个测试点都领先，并且在 batch size 1 时实现近两倍吞吐。

Batch size	Kimi K3 megakernel（16× TPU v7）	vLLM 基线（16× GB200）
1	249 tok/s	127 tok/s
2	392 tok/s	227 tok/s
4	515 tok/s	373 tok/s
8	865 tok/s	636 tok/s

• Kimi K3 megakernel — 16× TPU v7
• vLLM 基线 — 16× GB200

我们还通过一个简单的 API server 对外提供服务，对实现进行了端到端准确性验证。在贪心解码和最高推理强度（max reasoning effort）下，该 megakernel 在 GPQA-Diamond 上得到 0.944 分，在 GSM8K 上得到 0.972 分。用于复现这些结果的脚本与实现一同提供。

后续工作

这个 kernel 只是第一步。其构建理念——全模型可见性、显式 VMEM 暂存、在消费阶段之前数层就发起传输——并不局限于 Kimi K3。接下来的大部分工作，是让这些理念可复用并将其进一步推进。我们正在开展的若干方向如下：

更多工作负载。 本研究表明，megakernel 对小 batch 解码是有效的，但生产流量可能呈现更多形态。更高并发会改变瓶颈，因为向量计算和跨设备流量会随 batch size 增长；而 Agent 工作负载则使 KV cache 的移动和管理成为决定性能的关键。

更多拓扑。 我们的实现针对 2×2×4 TPU 芯片拓扑进行了定制。支持其他形状意味着要编写新的集合通信实现，因为我们的通信和并行化策略围绕这 16 颗芯片的连接方式构建。在 NVIDIA 的机架级系统中，一个 NVLink 域内的每张 GPU 到其他 GPU 的距离相同，因此唯一的决策是使用多少张 GPU；而在 TPU 上，我们还需要选择拓扑，这为搜索增加了一个维度。

更好的工具。 Pallas 在编写 kernel 时表现良好。不过，一些瓶颈只有通过直接分析编译后的指令包才能发现。随着这项工作持续扩展，能够为开发者以及未来会编写更多此类 kernel 的编程 Agent 指出这些低效之处的工具，将非常重要。

Inferact 的业务

这项工作源自 Inferact 所做的事情：我们在客户的基础设施上优化推理，这意味着针对眼前的特定硬件、模型和工作负载形态进行优化。Kernel 工作是我们通过平台及面向客户提供的推理优化服务的一个例子。由于该平台构建于开源 vLLM 之上，客户无需改变现有服务方式即可获得这些优化。

Engineering · September 23, 2026

700 TPS on Kimi K3: A Case for TPU Megakernels

Written byGeorge Novack, Xuting Liu, Jeff Ma, and Woosuk Kwon

Today we released inferact/tpu-megakernels, a collection of megakernels for TPU v7. Our Kimi K3 implementation delivers over 700 tokens/s with speculative decoding, compared with 452 tokens/s on GB200. Without speculative decoding, our megakernels for K3 and Qwen 3.8 27B deliver roughly 1.4 to 2× the decode throughput of the GB200 baseline at batch sizes 1 through 8.

Kimi K3 tokens / second Qwen 3.8 27B tokens / second 709 16× TPU v7 452 16× GB200 1,515 4× TPU v7 695 4× GB200

• TPU megakernel
• vLLM baseline

Output tokens per second with speculative decoding at an acceptance length of 6.

Decode speed is bound by the hardware's peak memory bandwidth. Our megakernels push end-to-end performance closer to this theoretical peak through effective use of TPU's large on-chip VMEM, sequential programming model, and explicit asynchronous pipelines to prefetch weights across layer boundaries. This post explains our design and why megakernels and TPUs are a good fit.

Background — TPUs vs. GPUs

Considering only the top-level specs, Google's TPU v7 Ironwood chip is comparable to NVIDIA's GB200 GPU.

Spec	Google TPU v7	NVIDIA GB200 (per GPU)
Peak BF16 Compute	2.31 PFLOPS	2.5 PFLOPS
Peak FP8 Compute	4.61 PFLOPS	5 PFLOPS
HBM capacity	206 GB	186 GB
HBM bandwidth	7,380 GB/s	8,000 GB/s
Interconnect	1,200 GB/s (ICI)	1,800 GB/s (NVLink 5)

Sources: Google Cloud TPU7x docs, NVIDIA GB200 NVL72.

This apples-to-apples comparison glosses over one of the TPU's most interesting features: the shape and size of its on-chip memory. Decode speed depends mostly on moving data efficiently from HBM to fast on-chip memory, which makes those dimensions a key design consideration.

Compared to a GPU's multi-level memory hierarchy of hardware-managed L1/L2 caches and program-managed Shared Memory and Tensor Memory spaces, a TPU's memory hierarchy is simple. Data from HBM is transferred to a large pool of on-chip SRAM called Vector Memory (VMEM); the TPU's vector registers (VREGs) are populated directly from VMEM. Importantly, the lifespan of all data in VMEM is managed entirely by our TPU program: when we move some data to VMEM, it stays there until we explicitly overwrite it with something else.

NVIDIA GB200 (per GPU) TPU v7 HBM — 186 GB HBM — 206 GB program-issued DMA L2 cache — hardware managed SM — ×152 shared mem 228 KB tensor mem 256 KB a block addresses only its own SM's shared memory VMEM — 64 MiB per TensorCore one pool, software managed the program decides what lives here, and for how long tensor cores / CUDA cores MXU / VPU / XLU ~111 MiB of SRAM, split 152 ways 128 MiB per chip, addressed as two pools

With 64 MiB of VMEM per TensorCore, we can prefetch most model weight tensors, ensuring they are staged in fast on-chip SRAM by the time any operation needs to use them. The GPU's Tensor Memory (TMEM) is organized differently. An entire GPU holds ~38 MiB of it, but that capacity is divided among all 152 SMs, and 256 KB per SM leaves little room to stage weights ahead of time.

Why megakernels?

Most kernels have the same rough outline: load weights and activations from HBM, perform some computation, then store the result back. Only the load phase pulls a large amount of data from HBM, leaving bandwidth underutilized during the other phases. Well-designed kernels mitigate these bandwidth bubbles by subdividing and overlapping each phase, but only to a degree. A single decode step might dispatch hundreds of kernels like this, one after another, and these bubbles add up.

LOAD Load n Load n+1 Load n+2 no loads in flight COMPUTE setup wait for Load n Compute n Compute n+1 Compute n+2 STORE Store n Store n+1 Store n+2 ramp: no loads issued yet drain: last load issued, compute and stores finish steady state — Load n+1 overlaps Compute n Time → (not to scale)

• HBM traffic
• Compute
• Idle (bubble)

GPUs already provide ways to address this. CUDA Graphs reduce host launch overhead, streams allow independent branches to execute concurrently, and programmatic dependent launch (PDL) lets a dependent kernel begin preparation and weight loading before its producer finishes. These mechanisms recover some overlap while retaining separate kernels.

A megakernel removes the boundary. With the whole decoder step inside one program, a weight load is no longer tied to the kernel that uses it — it can be issued as early as on-chip storage allows. The binding constraints become buffer lifetimes and capacity. On TPU v7, large VMEM and explicit asynchronous transfers let us stage the next layer's weights while the current layer is still computing or communicating. The goal is to keep the HBM interface busy throughout a decode step whose cost is dominated by moving weights.

Existing TPU implementations rely on XLA to perform fusion and scheduling, but finding an optimal cross-layer pipeline requires joint decisions about storage reuse, transfer timing, and synchronization. Automatic optimization may not discover the schedule we want. Pallas lets us express those decisions directly while leaving instruction scheduling to the compiler.

Writing TPU megakernels with Pallas

Pallas is the kernel language for TPUs. To our knowledge this is the first open-source inference megakernel written in Pallas, and we found the language surprisingly well suited to it. Below are the Pallas features that mattered most in this work.

Grid-less kernel launch

The Pallas language can be used to write kernels for both TPU and GPU. As such, it is possible to write Pallas kernels in a way that will feel familiar to any GPU kernel engineer: we can launch our kernels as a grid of programs, with each program receiving and operating on a specific tile of the input.

TPUs, though, follow an inherently single-threaded execution model. Even if we launch our kernel as a grid, each individual program is still run sequentially, one at a time. While the grid-based programming model has its advantages, we found that for megakernel development, embracing TPU's single-threaded nature is the way to go for both readability and fine-grained control. Our megakernel is launched as a single program with no explicit grid.

Scoped VMEM allocations

Our megakernel must balance aggressive prefetching of model weights against the limited pool of 64 MiB VMEM per TensorCore. Some data, like activations, remain resident in VMEM throughout the entire forward pass, while others, like the KV cache used in MLA, have a very short lifetime. The Pallas run_scoped primitive enables fine-grained control over the lifetime of specific VMEM allocations. This simplifies the task of the programmer: rather than explicitly managing which byte ranges in VMEM are safe to read and write at any point in time, we can write self-contained "mini-kernels" that define exactly the shape and size of the VMEM they need.

Asynchronous DMAs

Asynchronous DMAs between HBM and VMEM are what make the megakernel possible. Our implementation makes heavy use of make_async_copy(...).start() and make_async_copy(...).wait() to specify when weight prefetching should be initiated and when to pause execution until it completes. Managing asynchronous dependencies between operations is notoriously awkward in JAX: XLA schedules chains of operations as a unit and reserves the right to reorder them or reclaim their buffers. Inside a single kernel the problem does not arise: a copy we start and wait on never crosses a boundary XLA is free to schedule across.

Our Kimi K3 megakernel

The Kimi K3 megakernel combines K3's 92 Mixture-of-Experts (MoE) layers into a single Pallas call, conditionally choosing between Kimi Delta Attention (KDA) and Multi-head Latent Attention (MLA) operations based on the layer index. Inside the kernel, every weight tensor is staged into VMEM by a transfer issued during an earlier phase so that it is, ideally, already resident when its consumer runs. Because the layer loop lives inside the kernel, asynchronous copies can reach across layer boundaries. For example, layer N+1's attention projection weights start moving while layer N is still in its MoE phase. A simplified view of the cross-layer prefetching strategy is shown below.

COMPUTE layer N communications attention router shared expert routed experts communications attention layer N−1 layer N+1 HBM → VMEM layer N attn shared experts routed experts layer N+1 attention issued during layer N, resident before layer N+1 needs it Time → (not to scale)

• HBM → VMEM transfer
• Matrix work
• Vector work
• Collective

The kernel shards the model across 16 TPU v7 chips, a total of 32 TensorCores. Attention heads are split across all 32 ranks, while routed experts use TP4 × EP8. The MoE latent-to-hidden projection is sharded within each host (8 TensorCores) and replicated across hosts, trading additional weight reads for cheaper collective communications.

Tuning megakernels for a specific target request concurrency requires exploring more bespoke means of parallelizing work across chips beyond the familiar tensor, expert, and data parallel strategies. For example, even the seemingly modest increase in decode batch size from 1 to 8 exposed new bottlenecks in operations that primarily perform vector arithmetic on the residual stream (e.g., attention residuals and RMSNorm); to address this, we combine tensor parallelism with residual stream sequence parallelism to split vector arithmetic in addition to matrix multiplication across TPUs.

One benefit of the megakernel approach that we did not foresee is fast compilation. In hindsight it makes sense: the kernel author pays by hand, and only once, for the ordering and overlap that XLA must otherwise rediscover across every operation on every new compile. On TPU, compiling a large model made up of hundreds or thousands of XLA ops can regularly take 30+ minutes; our entire megakernel, on the other hand, compiles from scratch in less than 90 seconds, enabling much more rapid development and experimentation cycles.

DSpark

With support for batch size 8, the megakernel can also serve as the verifier when using DSpark speculative decoding. Each kernel launch evaluates 1 anchor and 7 proposed continuations in a single step. Depending on the exact workload, the draft model typically yields between 3 and 6 accepted tokens per step. At ~8.5 ms per decode step, the megakernel can deliver over 700 tokens per second to a single user.

0 200 400 600 800 accepted output tokens / second 350 229 acceptance length 3 1.53× faster 709 452 acceptance length 6 1.57× faster

• Kimi K3 megakernel — 16× TPU v7
• vLLM baseline — 16× GB200

Speculative decode throughput by acceptance length using the RedHatAI/Kimi-K3-speculator.dspark speculator.

Results

We evaluate the megakernel on 16 TPU v7 Ironwood chips and take vLLM's published Kimi K3 recipe on 16 GB200 GPUs as our baseline. With no speculative decoding, the megakernel leads at every batch size from 1 through 8, delivering nearly twice the throughput at batch size 1.

0 200 400 600 800 tok/s 249 127 1 392 227 2 515 373 4 865 636 8 batch size

• Kimi K3 megakernel — 16× TPU v7
• vLLM baseline — 16× GB200

We validated the accuracy of our implementation end to end by exposing it through a simple API server. With greedy decoding and max reasoning effort, the megakernel scores 0.944 on GPQA-Diamond and 0.972 on GSM8K. Scripts to reproduce these results are included alongside the implementation.

Future work

This kernel is a first step. The ideas it is built from — whole-model visibility, explicit VMEM staging, transfers issued layers ahead of the phase that consumes them — are not specific to Kimi K3, and most of the work ahead is about making them reusable and pushing them further. A few directions we are taking from here:

More workloads. This work shows that megakernels can be effective for small-batch decode, but production traffic can take many more shapes. Higher concurrency moves the bottleneck, since vector work and cross-device traffic grow with the batch size, and agentic workloads make KV cache movement and management crucial for performance.

More topologies. Our implementation is tailored to a 2×2×4 TPU chip topology. Supporting other shapes means writing new collectives, since our communication and parallelization strategies are built around how those 16 chips are wired. On NVIDIA's rack-scale systems every GPU in an NVLink domain is equally close, so the only decision is how many to use; on TPU we choose the topology as well, which adds a dimension to the search.

Better tooling. Pallas worked well for writing the kernel. Some bottlenecks, though, only become visible through direct analysis of compiled instruction bundles. Tooling that points at those inefficiencies, for people and for the coding agents that will write more of these kernels, would matter a lot as this work grows.

What Inferact does

This work comes out of what we do at Inferact. We optimize inference on our customers' infrastructure, which means optimizing against the specific hardware, model, and workload shape in front of us. Kernel work is one example of the inference optimizations we provide in the platform and to our customers, and because the platform builds on open-source vLLM, customers pick these up without having to change how they already serve.