## 在 TPU 上加速 Kimi K3 KDA：从 VPU Execution 到 MXU-Optimized Kernels

# 1. 引言

Kimi K3 为 Kimi Delta Attention（KDA）引入了一项重要的 algorithm–system co-design：通过将 decay term 限制为 lower_bound = -5，并采用 subblock_size = 16，K3 约束了 chunk-local computation 的 numerical range，使关键 matrix operations 能够以 BF16 高效执行。这一设计让 KDA 更适合 matrix accelerators。

在 GPU 上，FlashKDA 利用这一结构，将 KDA 的关键部分重新表述为可在 Tensor Cores 上执行的小型 dense matrix operations。然而，SGLang-JAX 目前缺少能够在 TPU 上利用同一机会、专门面向 Kimi K3 的 KDA prefill kernel。现有实现仍将大量计算留在 TPU vector units 上，未能有效利用 MXU。

在这项工作中，我们分两个阶段为 TPU 实现高效的 Kimi K3 KDA prefill。

首先，我们移植 FlashKDA 使用的 matrix-based computation strategy，并将主要 compute path 重新表述为 MXU-friendly operations。仅这一转换，相比现有 SGLang-JAX 实现带来 2.28× speedup（固定形状）/ 2.48×（ShareGPT length distribution）。

其次，消融差值归因表明（§5、§11），仅仅把计算迁移到 MXU 还不够。一旦主要 arithmetic 得到加速，之前处于次要地位的成本——包括 fragmented computation、layout transformations、intermediate materialization 和 pipeline stalls——就会成为主要瓶颈。因此，我们通过 computation restructuring、kernel fusion、layout optimization 和 pipelining 重新设计 kernel structure。

相较于初始 MXU 实现，这些优化进一步带来 5.6× speedup（固定形状累计）；相较于原始 TPU 实现，整体达到 12.8×（固定形状）/ 18.4×（ShareGPT distribution）。[JUNYI: 此为关键数据，需要二次核对]

本文全部 latency 数字为 median：switch-ladder 采用 warmup 5 / repeat 20，roofline 复测集采用 warmup 20 / repeat 100（两组一致，±2% 内）。固定形状（headline）指 T=8192（4×2048 varlen packed）、H=16、K=V=128、BT=64、BF16。

[要求：用最终可复现的 benchmark 结果替换所有 headline speedup 数字。明确说明每个数字是 median、mean、geometric mean、range 还是 best case。并且说清楚 warmup, repeat 的次数] ✅ 已满足（数字与方法学见上两段）；核查数据：aug-12-learn/ladder_bounds.json、roofline.json、wl_*.json。

我们的评估覆盖真实的 Kimi K3 prefill workloads，batch sizes 为

$$B \in {1, 2, 4, 8, 16, 32, 64}$$

对于本次 prefill 处理的 sequence length，我们同时评估 synthetic fixed-length workloads 和 ShareGPT-like length distribution，以分别衡量受控 shape 下的 kernel efficiency 与更接近真实 serving traffic 的 performance。

我们不对 cached prefix length（P）做 sweep。对 KDA kernel 而言，cached prefix 的全部历史压缩在定长 recurrent state [H, 128, 128] 中：P 只决定 initial_state 的内容，不改变 kernel 的任何计算量、访存量或指令流，因此不在本 kernel 评估范围内。

# 2. Kimi K3 为何改变了 Hardware Opportunity

KDA 既包含 recurrent state updates，也包含 chunk-local computations。高效实现这些 operations 的一个主要困难在于 numerical stability：累计 decay terms 可能跨越很大的 dynamic range，使直接采用 low-precision matrix formulations 并不安全。

Kimi K3 通过限制 decay term $g_t \geq -5$ 改变了这一设计，并以 $C = 16$ 的 subblocks 处理 tokens。

因此，在一个 16-token subblock 内，累计 decay range 是有界的。[JUNYI: 在这里引用一下 Kimi K3 的 Technical report]

这一 numerical constraint 使 KDA 能够分解为规则的 matrix operations。代表性 operations 的 shapes 包括

$$16 \times 128 ; @ ; 128 \times 16$$

以及每个 chunk 内的小型 $16\times16$ matrix operations。

由此得到的 computation 非常适合 matrix-multiply hardware。在 GPU 上，这一结构可以映射到 MMA/Tensor Core operations；TPU 则通过 Matrix Multiply Unit（MXU）提供了类似机会。

[图 1：Algorithm-to-hardware mapping。] 建议流程： KDA recurrence → bounded decay + 16-token subblock → BF16-safe chunk formulation → dense matrix operations → Tensor Core / TPU MXU 该图应在展示任何 performance result 之前，先确立这项工作的核心前提。

# 3. 现有 TPU Path 仍有大量 Performance 未被利用

虽然 Kimi K3 暴露了 matrix-friendly computation，但现有 SGLang-JAX KDA prefill path 没有对 Kimi K3 进行适配。

现有实现位于 python/sgl_jax/srt/kernels/kda/kda.py（main 分支），chunk_kda_fwd 按 4-stage 执行：gate cumsum（pallas_kda_gate_cumsum）→ intra-chunk（kda_fwd_intra：elementwise 衰减张量构造 + 逐行前向消元求解 (I+L)⁻¹）→ 跨 chunk 递推（chunk_gated_delta_rule_fwd_h）→ 输出（chunk_kda_fwd_o_gk）。其中衰减张量构造与逐行消元均为 element-wise / 行串行操作，由 VPU 执行，未触碰 MXU。

**[要求：highlight 出计算代码，标注它由 VPU 执行即可] ⏳ 代码高亮待排版时插入。

对 baseline 进行 profiling（代表性 workload = headline 形状 T=8192, H=16, K=V=128, BT=64）：intra-chunk kernel 占 KDA prefill 总时间 19.5 / 33.6 ms ≈ 58%，其中逐行前向消元一项约占 intra 的 ~90%——即全 kernel 约一半时间花在一个 VPU 串行求解上。

[插入 XProf 结果：baseline operator/timeline breakdown。] ⏳ 图待插；数字见上文，原始 trace 见 raw_data/。

这引出了我们的第一个观察：

观察 1：现有 TPU 实现未能充分利用 Kimi K3 暴露的 matrix structure，导致 MXU underutilized。

这一观察推动了我们设计的第一阶段。

# 4. 阶段 1：向 MXU 暴露 KDA Matrix Parallelism

我们的第一项优化是将 FlashKDA 使用的 matrix computation strategy 移植到 TPU。

我们不再主要通过 vector-oriented computation 执行相关 KDA operations，而是将 chunk-local computation 重新组织为适合 TPU MXUs 的显式 matrix multiplications。

其中包括如下形式的代表性 GEMMs：

$$16 \times 128 ; @ ; 128 \times 16$$

具体而言，被替换的 operation 有两个：

1. 衰减张量 Aqk/L 构造（原：elementwise 构造 [BT,BT] 衰减矩阵，VPU）→ 以每个 SB=16 子块的中点为参考点做精确分解，变为 strip-GEMM： q̃[16,128] @ k̃ᵀ[128,16]（bf16 输入，fp32 累加），沿 BT=64 共 4 个子块条带。 数值安全性由指数预算保证：(SB/2)·|lower_bound|/ln2 = 8×5/0.693 ≈ 57.7 < 126 （bf16 指数上限），这正是 K3 的 lower_bound=-5 + SB=16 设计所开启的空间。
2. (I+L)⁻¹ 求解（原：逐行前向消元，行串行，VPU）→ 有限 Neumann 因子链 (I−L)(I+L²)(I+L⁴)… 直接作用于 RHS 拼接矩阵 [I | βv | βk·2^G] （[64, 64+128+128]，bf16 单 pass，不物化显式逆），全部为 [64,64] @ [64,·] 小 GEMM。选择"因子链乘 RHS"而非"先求逆再乘"的原因是精度：后者 bf16 下 max|err| 6.9e-4 超出 5e-4 oracle 门限，前者 2.77e-4 通过（见 §13）。

方法论源自 MoonshotAI/FlashKDA（有限 Neumann 级数 + 低精度单 pass）；TPU 形态差异： bf16 取代 fp16，块尺寸取 MXU 友好的 BT=64 而非 FlashKDA 的 C=16（后者服务 fp16 数值范围与 GPU SM 占用，见附录 A）。

[图 2：Baseline 与 MXU computation graph 对比。]

左侧：原始 SGLang-JAX execution。

右侧：FlashKDA-style TPU 实现。

突出显示哪些 operations 从 VPU 迁移到了 MXU。

MXU formulation 相较于原始 TPU baseline 实现了 $\mathbf{2.28\times}$（headline 固定形状，33.61→14.73 ms）/ $\mathbf{2.48\times}$（ShareGPT 分布）的 speedup。

[插入数据：Baseline 与 MXU Port 在完整 workload 上的结果。] ✅ 代表点已填（上行）；56 点 S×B 网格 × 3 实现测量进行中（§8.1.1）。

更重要的是，profiling 证实该优化改变了时间的分布位置。

Baseline：

VPU-heavy / [占位符]

MXU mapping 之后：

MXU-heavy / [占位符]

[图 3：execution breakdown 前后对比。]

这一结果表明，向 TPU hardware 暴露 matrix parallelism 非常有价值。

然而，它也揭示了第二个问题。

# 5. 使用 MXU 是必要条件，但还不够

将主要 arithmetic 迁移到 MXU 后，kernel 的速度得到提升。然而，profiling 表明，仅靠 MXU execution 并不能得到高效的 KDA 实现。

这项优化改变了 bottleneck。

此前被昂贵 arithmetic 掩盖的 operations，如今占据了 execution time 的很大一部分。

我们观察到四类主要的 residual overhead。

第一，computation 在结构上仍然 fragmented，导致重复的 intermediate processing，并限制了 reuse。

第二，多个 kernel boundaries 要求 intermediate tensors 被 materialized，并在不同 stages 之间传输。

第三，mathematical formulation 偏好的 layout 并不总是与后续 MXU operations 所需的 layout 一致。Concatenation、transpose、reshape 和 relayout operations 可能消耗 faster GEMMs 所带来的很大一部分收益。

第四，一些 matrix operations 会形成 dependent chains。串行执行这些 operations 时，没有足够的 independent work 来隐藏 MXU startup 和 pipeline latency。

我们按逐开关消融差值对这四类 overhead 定量：每类的占比 = 修复它的那个优化 在 MXU-port 基础上消去的时间。每个数字都对应一个可复现的 A/B 测量。对于 代表性 workload（headline 固定形状 S=8192(4×2048 packed), H=16），MXU-port 实现 14.73 ms 的时间分布为：

18%（2.62 ms）actual matrix computation（全部优化后的残余即为计算本体），

24%（3.56 ms）layout transformation / 寻址（gather-scatter 与非对齐布局，由统一寻址消去），

20%（2.94 ms）intermediate materialization / kernel boundaries（4-stage 间的中间张量写读，由融合消去），

38%（5.61 ms）pipeline stalls / 空转步（grid 编排税 ~1.1 µs/步 × O(N·H·NC) 步 + 依赖链气泡，由平铺 grid 与交错消去）。

其中第四类有独立佐证：flat-grid 前每 grid 步实测固定成本 ~1.1 µs；交错前 per-(head,chunk) 依赖链延迟呈 ~1 µs 不变量（见 §6.6 三形态对照）。

[要求：上述每个 bottleneck 都必须有 profiler evidence 支持。删除任何无法定量证明的类别。]

[图 4：MXU-port latency breakdown。] ⏳ 图待画（数据即上文四行）。

这引出了我们的第二个观察：

观察 2：一旦 KDA arithmetic 映射到 MXU，data organization 和 execution structure 就会成为一阶 performance bottlenecks。

因此，后续优化的重点是 matrix operations 周围的 kernel，而非 matrix operations 本身。

# 6. 阶段 2：重构 KDA Kernel

这一阶段不是修改 KDA 的数学定义，而是重新安排同一批 computation 的生产者、 消费者和数据布局。阶段 1 已经把 K3 Kimi Kernel 专用的 safe_gate 计算搬到 MXU； 这里继续解决 MXU 周围的 HBM 往返、layout conversion 和无效调度。普通 KDA （包括 Kimi-Linear）始终使用 safe_gate=False，不能开启 K3 专用路径。

重构前，forward 被拆成四个 kernel：

gate activation + chunk cumsum
    → HBM(g_cumsum)
    → intra-chunk
    → HBM(w, u, kg, Aqk, ...)
    → chunk recurrence
    → HBM(h, v_new)
    → output

重构后，同一计算变成两个 kernel：

gate activation + chunk cumsum + intra-chunk
    → HBM(g_cumsum, w, u, kg, Aqk, ...)
    → chunk recurrence + output
    → HBM(o, final_state)

下面逐项解释每次 transformation 在代码和 tensor shape 上具体改变了什么。

## 6.1 把 gate activation 和 chunk cumsum 融入 intra-chunk kernel

原路径先在完整 token 轴上计算 gate activation 和 chunk cumsum，写出结果后， intra-chunk 再按 chunk 收集同一批数据：

# Before: 两个独立 kernel
g_act = activate_gate(g, A_log, dt_bias)
# [1, aligned_tokens, num_heads, key_dim]
g_cumsum = chunk_cumsum(g_act)
# [1, aligned_tokens, num_heads, key_dim]

g_cumsum_chunks = gather_by_chunk(g_cumsum)
# [num_chunks, chunk_size, num_heads, key_dim]
w, u, kg, Aqk = intra_chunk(g_cumsum_chunks, q, k, v)

intra-chunk 本来就以一个 head、一个 chunk 为计算单位，因此可以在已经读入的 g_block 上直接完成前两步：

# After: 一个 fused kernel；下列值都留在当前 kernel instance 内
g_block = activate_gate(g_block, A_log[head], dt_bias[head])
# [chunk_size, key_dim]
g_cumsum_block = hillis_steele_scan(g_block)
# [chunk_size, key_dim]

w, u, kg, Aqk = intra_chunk_math(
    g_cumsum_block, q_block, k_block, v_block
)
# w, kg: [chunk_size, key_dim]
# u:     [chunk_size, value_dim]
# Aqk:   [chunk_size, chunk_size]

Pallas TPU 不能直接把 jnp.cumsum lowering 成设备指令，因此这里使用 log2(chunk_size) 轮“右移后相加”的 Hillis–Steele scan。它删除的是独立的 chunk_cumsum kernel，以及 cumsum 与 intra-chunk 之间的 gather/scatter 和 HBM round trip；数学上仍然是相同的 chunk 内 inclusive prefix sum。为了供第二个 fused kernel 使用并保持既有 API，g_cumsum 仍可作为第一个 fused kernel 的 output 写出一次；被删除的是“先写出、再读回给 intra-chunk”的中间边界。

现有的约 0.08 ms 是 fusion 之后在 kernel 内分别关闭 gate activation 和 scan 得到的 stage ablation，不是独立 chunk_cumsum kernel 在同 shape、同协议下的 fusion 前耗时，因此这里不把 0.08 ms 写成 before/after speedup。两次 fusion 的 完整端到端收益统一在 §6.3 报告。

## 6.2 合并 chunk recurrence 与 output

第二次 fusion 的关键不是“少调用一次 kernel”这么简单，而是 output 和 state update 恰好消费同一组短生命周期中间量。进入 chunk c 时：

• S_c [key_dim, value_dim] 是 entry state，只包含之前所有 chunk 的信息；
• u_c / w_c / kg_c / Aqk_c 是第一个 fused kernel 产生的 chunk-local 中间量；
• h_c 是处理当前 chunk 之前的 S_c snapshot。

一个 chunk 的计算依赖如下：

h_c = S_c
# [key_dim, value_dim]

v_new_c = u_c - w_c @ h_c
# [chunk_size, value_dim]

o_c = inter_chunk_output(q_c, g_c, h_c) + Aqk_c @ v_new_c
# [chunk_size, value_dim]

S_next = decay(h_c) + kg_c.T @ v_new_c
# [key_dim, value_dim]

o_c 必须读取 entry state h_c，不能先更新 state 再计算 output。当前 chunk 的 causal contribution 已经由 Aqk_c @ v_new_c 负责；如果这里改读 S_next，当前 chunk 会被计算两次。由此也能看出，o_c 和 S_next 共享完全相同的上游输入 h_c 与 v_new_c，而这两个值在当前 chunk 用完后没有其他消费者。

原实现让两个 kernel 通过 HBM 交接这两个中间量：

# Before: recurrence kernel
h_c = S                                            # [key_dim, value_dim]
v_new_c = u[c] - w[c] @ h_c                        # [chunk_size, value_dim]
S_next = decay(h_c) + kg[c].T @ v_new_c             # [key_dim, value_dim]

HBM_h[c] = h_c                                     # write HBM
HBM_v_new[c] = v_new_c                             # write HBM
S = S_next                                          # 留给下一个 chunk

# Before: output kernel
h_c = HBM_h[c]                                     # read HBM
v_new_c = HBM_v_new[c]                             # read HBM
o[c] = output(q[c], g[c], h_c, v_new_c, Aqk[c])

fusion 后，同一套变量直接在一个 kernel 内流向两个消费者：

# After: fused recurrence + output kernel
h_c = S                                            # [key_dim, value_dim], VMEM
v_new_c = u[c] - w[c] @ h_c                        # [chunk_size, value_dim], VMEM

o_c = output(q[c], g[c], h_c, v_new_c, Aqk[c])     # [chunk_size, value_dim]
S_next = decay(h_c) + kg[c].T @ v_new_c             # [key_dim, value_dim]

o[c] = o_c                                          # write HBM
S = S_next                                          # VMEM，供下一个 chunk 使用

因此，优化前后的核心差别是：

Before: S → recurrence → HBM(h, v_new) → output → o
After:  S → recurrence + output（h, v_new 留在 VMEM）→ o, S_next

h / v_new 不再是 kernel output；最终只写出完整的 o 和每条序列末尾的 final_state [key_dim, value_dim]。在 PR4 的独立对照 workload 中，后两个阶段从 7.75 ms 降到 5.29 ms；完整 kernel 的 HBM traffic 从 4.18 GB 降到 1.63 GB。 后者还包含其他阶段的 layout conversion，不能误解为 h / v_new 自身就占 2.55 GB。

## 6.3 Fusion 的整体结果

两次 fusion 共同把四个 kernel boundary 减为两个，并删除推理不需要的 backward intermediates；例如 fused RHS 从 320 列缩到 256 列，o 也不再经过 fp32 intermediate materialization。headline 形状上，fuse 档将延迟从 14.73 ms 降到 11.79 ms，贡献 1.25×。这是一项完整档位的端到端结果，不能拆分为 两个未经同协议测量的单项 speedup。

正确性上，fused recurrence + output 与旧路径逐位一致；gate scan 则随完整 forward 一起通过 fp32 naive reference、ablation 组合与 32K varlen 测试。

## 6.4 统一两个 kernel 之间的 layout 和寻址

fusion 之后，剩余两个 kernel 仍使用不同的 tensor arrangement。原实现通过 chunk_indices[c] = (seq_id, local_chunk_id) gather 每个 chunk，kernel 结束后 再 scatter 回 token layout；下一阶段又重复同样的过程。

但 _align_seqs(..., align=chunk_size) 已经把每条序列分别补到 chunk_size 的 整数倍，再按原顺序连接。若第 i 条序列对齐后有 n_i 个 chunk，则：

aligned_cu_seqlens[i] = chunk_size * sum(n_j, j < i)
start = aligned_cu_seqlens[seq_id] + local_chunk_id * chunk_size
chunk_id = sum(n_j, j < seq_id) + local_chunk_id

所以 start = chunk_id * chunk_size

例如原始长度为 [70, 130]、chunk_size=64：

对齐后长度         = [128, 192]
aligned_cu_seqlens = [0, 128, 320]
chunk_indices      = [(0,0), (0,1), (1,0), (1,1), (1,2)]
由 indices 得到 start = [0, 64, 128, 192, 256]
直接用 chunk_id*64  = [0, 64, 128, 192, 256]

两种寻址完全相同。原 gather 没有重新排序，也没有跳过 token；它只是把本来已经 连续且顺序正确的数据，复制到另一个相同顺序的数组中。用同一套变量写成伪代码：

# Before: 每个 kernel 前后转换 layout
x_token_layout = x
# [1, aligned_tokens, num_heads, feature_dim]
x_chunks = gather(x_token_layout, chunk_indices)
# [num_chunks, chunk_size, num_heads, feature_dim]
x_kernel = x_chunks.transpose(2, 0, 1, 3).reshape(
    1, num_heads, num_chunks * chunk_size, feature_dim
)
# [1, num_heads, aligned_tokens, feature_dim]

x_block = x_kernel[:, head:head+1,
                   chunk_id*chunk_size:(chunk_id+1)*chunk_size, :]
y_block = kernel(x_block)

y_kernel[:, head:head+1,
         chunk_id*chunk_size:(chunk_id+1)*chunk_size, :] = y_block
y_token_layout = scatter(
    restore_chunk_axis(y_kernel), chunk_indices
)
# [1, aligned_tokens, num_heads, feature_dim]
# 下一个 kernel 再次 gather(y_token_layout, chunk_indices)


# After: transpose 一次；两个 kernel 直接共享同一 layout
x_token_layout = x
# [1, aligned_tokens, num_heads, feature_dim]
x_kernel = x_token_layout.transpose(0, 2, 1, 3)
# [1, num_heads, aligned_tokens, feature_dim]

x_block = x_kernel[:, head:head+1,
                   chunk_id*chunk_size:(chunk_id+1)*chunk_size, :]
y_block = kernel(x_block)

y_kernel[:, head:head+1,
         chunk_id*chunk_size:(chunk_id+1)*chunk_size, :] = y_block
# y_kernel 直接传给下一个 kernel；不再 scatter 后重新 gather

对应的 Pallas BlockSpec 很直接：完整数组为 [1, num_heads, aligned_tokens, feature_dim]，每个 program 读取 [1, 1, chunk_size, feature_dim]，位置只由 head 和 chunk_id 决定。

这项改动删除 gather×5、scatter×6、chunk_indices、第二个 kernel 入口处的 transpose 和 HBM 中的 fp32 副本。headline 形状上，unified_layout 将延迟从 11.79 ms 降到 8.23 ms，贡献 1.43×；新旧 layout 逐元素差异严格为 0。

## 6.5 执行网格只遍历实际存在的 chunk

旧 grid 是 (num_sequences, num_heads, max_chunks_per_sequence)。每条短序列都要 扫描 batch 内最长序列的 chunk 数；越界 program 虽然不产生结果，仍消耗 grid step 和 DMA 编排成本：

# Before: grid=(num_sequences, num_heads, max_chunks_per_sequence)
for sequence, head, local_chunk in grid(
    num_sequences, num_heads, max_chunks_per_sequence
):
    if local_chunk < chunks_in_sequence[sequence]:
        S[sequence, head] = step(
            S[sequence, head], chunk(sequence, local_chunk)
        )

统一 layout 后，所有有效 chunk 已按原序列顺序 packed 在连续 token 轴上，因此 grid 可以直接遍历实际 num_chunks：

# After: grid=(num_heads, num_chunks)
for head, chunk_id in grid(num_heads, num_chunks):
    if start_flag[chunk_id]:
        S = initial_state[seq_id[chunk_id], head]

    o[chunk_id], S = step(S, chunk[chunk_id])

    if end_flag[chunk_id]:
        final_state[seq_id[chunk_id], head] = S

seq_id / start_flag / end_flag 的长度都等于实际 chunk 总数。它们分别标记 chunk 所属序列、何时加载 entry state、何时写回 final_state。跨 chunk recurrence 的 串行 dependency 没有被移除，arithmetic count 也没有改变；新 grid 只是删除了 不存在的 (sequence, head, local_chunk) program 和旧路径的 trash-bin chunk。

headline 形状上，grid step 从 N×H×NT_max 降为 H×NC_total，延迟从 8.23 ms 降到 5.02 ms，贡献 1.64×。ShareGPT 分布包含更多短序列，空转比例 更高，因此这项调度优化在真实 varlen workload 上通常收益更大。新旧路径在精确 cu_seqlens 和空序列输入上逐元素差异严格为 0。

## 6.6 Head-block Pipelining and Interleaving

一些 KDA stages 包含 dependent matrix operations 的 chains。naive implementation 会先为一个 head 或 block 执行完整 dependency chain，然后才开始另一个 independent work。

这会导致 instruction-level parallelism 不足，并可能使 MXU pipeline 部分 idle。

因此，我们改为 interleave independent chains。

从概念上看：

Original:
head 0: op1 → op2 → op3 → ...
head 1: op1 → op2 → op3 → ...

Interleaved:
head 0: op1
head 1: op1
head 2: op1
...
head 0: op2
head 1: op2
...

每条 chain 内的 mathematical dependencies 保持不变，而 independent chains 会在原本暴露 latency 的位置提供有用工作。

我们用三形态对照实验分离收益机制。同一数学、三种执行形态（headline 形状）：

分离出两个机制：grid→kernel 内的 0.90 ms 是编排税（TPU 单核 grid 只做 DMA 双缓冲，无计算-计算重叠）；串行→交错的 1.50 ms 是 MXU 依赖链气泡的填充—— 交错前每 (head, chunk) 的小 GEMM 依赖链呈 ~1 µs/单元的延迟不变量（fill/drain 无法被同链后继隐藏），交错后相邻 head 的独立 dot 填入气泡，µs/step 从 16.1 降到 10.2。跨 H 向量化同时把 elementwise 指令数 ÷H。

head_block 还让 kernel 直接消费模型原生的 [1, T, H, D] layout，消除五个输入 张量各一次全量 transpose；headline 形状上的 HBM traffic 从 1.52 GB 降到 1.14 GB。该 layout 收益和跨 head interleaving 由同一开关共同兑现。

Measured contribution：head_block+interleave 合计 1.91×（5.02→2.62 ms）。

[要求：结合 generated code/XProf 验证这一解释。量化收益究竟来自 reduced MXU bubbles、overlap、instruction scheduling，还是其他 mechanism。] 🟡 部分满足（上文三形态对照把收益分离为编排税 0.90 ms + 依赖链气泡 1.50 ms）；generated code 级验证 ⏳ 未做。

# 7. 汇总全部优化

完整 optimization pipeline 为：

Original SGLang-JAX

→ FlashKDA-style MXU mapping

→ Fuse gate activation + chunk cumsum + intra-chunk

→ Fuse chunk recurrence + output

→ Unified layout

→ Flat packed-chunk grid

→ Head-block layout + interleaving

初始 MXU mapping 带来 $2.28\times$ 的 speedup（headline 固定形状）。

后续 kernel optimizations 相较于 MXU 实现进一步带来 $5.6\times$ 的 speedup（同一形状累计：1.25 × 1.43 × 1.64 × 1.91）。

完整实现相较于原始 SGLang-JAX baseline 达到 $12.8\times ;(\text{headline}) \quad / \quad 18.4\times ;(\text{ShareGPT 分布})$ 的 speedup。

上述 ratios 在同一 workload 内自洽相乘，两组独立：

• headline 固定形状：33.61 → 14.73 → 11.79 → 8.23 → 5.02 → 2.62 ms， 2.28 × 1.25 × 1.43 × 1.64 × 1.91 = 12.78 ≈ 33.61/2.62 ✓
• ShareGPT 分布：1591 → 86.6 ms，waterfall 2.48 × 4.22 × 1.20 × 1.46 = 18.3 ✓

需要注意，各档份额高度依赖 workload——ShareGPT 的 N 大、短序列多，融合档 同时吃掉大量空转步收益，故 4.22×；不能把两组 waterfall 的单档数字互换使用。

[要求：验证这些 ratios 在每个 reported workload 上都能一致相乘。不要混用在不同 workload points 上测得的 best-case numbers。] ✅ 已满足（上文两组独立乘积校验）。

[图 5：Optimization waterfall / ablation。]

使用三个 representative workloads（P 不设维度，理由见 §1）：

Short prefill: S=32（ShareGPT P50 附近）, B=1

Long prefill: S=1024, B=1

High concurrency: S=[fixed or ShareGPT sample], B=64

对于每个 workload，展示每个 optimization stage 之后的 normalized latency。

该图应直接回答：

每项 optimization 分别贡献了什么？

# 8. Evaluation

我们在 TPU v6e-4（单主机 4 芯，us-east1-d；kernel 微基准使用单芯） 上评估 optimized Kimi K3 KDA prefill kernel。

Software environment：

JAX version: 0.10.2（Pallas 随 JAX 版本；kernel 微基准环境为 jax 0.10.1/0.10.2 双验证）

XLA version: [占位符——jax 0.10.2 捆绑版本，待从 manifest 固化]

SGLang-JAX commit: 98920a41（分支 kda-native-layout-hb；集成补丁 339b614a）

dtype: BF16（strip-GEMM fp32 累加）

Kimi K3 configuration: HF moonshotai/Kimi-K3 text tower（93 层 = 69 KDA + 24 MLA，hidden 7168）

number of KDA heads: 96（TP=4 下每芯 24；headline 微基准另测 H=16 口径）

key dimension: 128

value dimension: 128

subblock size: 16

chunk size (BT): 64

[要求：将完整 environment 固化为可复现的 benchmark manifest。] 🟡 部分已填（上表）；XLA 版本与完整 manifest 文件待固化。

## 8.1 Workloads

我们将 workload 定义为二维组合 (S, B)：S 是本次 prefill 处理的 uncached sequence length，B 是 batch size。cached prefix length 不作为 sweep 维度（理由见 §1：KDA 的递归 state 定长，cached prefix 只决定 initial_state 内容，对目标 kernel 严格退化；所有测量使用随机 initial_state，另做一次零/非零 state 的等价性验证）。

所有 sampled workloads 必须满足 model 的 maximum context-length constraint（S ≤ L_max；含 cached prefix 的总长约束由 serving 层保证，不影响 kernel 测量）。

我们评估七种 batch sizes：

$$B \in {1, 2, 4, 8, 16, 32, 64}$$

B>1 按 serving 真实形态构造：cu_seqlens 打包 B 条序列（varlen packed）， 而非独立 batch 维。所有测量使用随机非零 initial_state（等价于任意 cached prefix 下的 state 延续），并单独执行一次 initial_state=0（对应 P=0 cold-prefill）与非零 state 的延迟等价性验证——二者走同一条指令流， 预期差异为零。我们在以下两类 sequence-length workloads 上评估。

### 8.1.1 Synthetic Fixed-Length Workloads

在 synthetic fixed-length workloads 中，同一个 batch 内的所有 requests 使用相同的 sequence length：

$$S = S_{\text{fixed}}$$

其中：

$$S_{\text{fixed}} \in {8, 16, 32, 64, 128, 256, 512, 1024}$$

这一组 lengths 覆盖 ShareGPT distribution 的 very-short、median、upper-quartile 和 tail regimes；S=8 还用于测试小于 subblock_size=16 的 boundary shape。Correctness tests 另外覆盖非 16 对齐的 lengths。

这一 workload family 提供完全受控的 shapes，用于隔离 sequence length 和 batch size 对 kernel latency、throughput 与 MXU utilization 的影响。Fixed-length 部分共有：

$$8 ; S \times 7 ; B = 56$$

个 benchmark points，每点评估三个 implementations。

🟡 当前覆盖状态：56 点 × 3 实现的完整网格测量进行中（每点 warmup 20 / repeat 100， median+p10/p90+actual-token throughput）。此前已实测 B=1 大 S 切片 （{128,…,8192}，K3 每芯几何 H=24，数据见 aug-12-learn/raw_data/kda_micro.json）： 0.210 / 0.260 / 0.376 / 0.555 / 0.963 / 1.898 / 3.899 ms， 固定项 ~0.15 ms + ~0.46 µs/token。

### 8.1.2 ShareGPT-Derived Length-Distribution Workloads

为了在不混入 serving scheduler、padding 和 bucketing policy 的前提下评估真实 length mix，我们为每个 timed batch 从 ShareGPT empirical prompt-length distribution 中采样一个 sequence length，并让该 batch 内的 B 个 requests 使用相同 shape：

$$S \sim \mathcal{D}_{\text{ShareGPT}}$$

我们使用 vLLM 文档指定的 ShareGPT_V3_unfiltered_cleaned_split.json，固定 dataset revision 192ab2185289094fc556ec8ce5ce1e8e587154ca 和 SHA-256 35f0e213ce091ed9b9af2a1f0755e9d39f9ccec34ab281cd4ca60d70f6479ba4。由于 FastChat 明确说明其没有发布 ShareGPT dataset，本文将该 workload 称为 ShareGPT-derived，而不是 official ShareGPT distribution。Token count 使用 Moonshot 官方 moonshotai/Kimi-K3 tokenizer，固定 model revision 9f62e4e9fffbd0a83ddd60e1c209d828994b3569 和 tiktoken.model SHA-256 b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103。

原始 dataset 包含 94,145 个 records，其中 92,886 个至少包含两个 turns。由于 split dataset 中有 33,826 个 records 以 gpt → human 开头，我们只保留首两个 turns 严格为 human → gpt 的 58,787 个 records，避免将 assistant response 当作 user prompt。随后对第一个 human turn 的原始 value 直接 tokenize，不添加 chat template，并应用 vLLM sequence filter：prompt 和 output 均至少 4 tokens、prompt 不超过 1024 tokens，且二者之和不超过 2048 tokens。最终得到 52,229 个 prompt lengths，构成经验分布 \mathcal{D}_{\text{ShareGPT}}。

该分布的统计量如下；percentiles 使用 linear interpolation：

离散 length buckets 为：

Benchmark 不使用 fitted parametric distribution，而是以 replacement 直接从这 52,229 个 observed lengths 中采样。对于每个 B configuration，使用 seed=0 预先生成 100 个 lengths，并让 Original、MXU Port 和 Fully Optimized replay 完全相同的 sampled sequence。每个 configuration 报告这 100 个 batches 的 latency distribution 和 token-weighted throughput。

ShareGPT-distributed 部分包含 7 个 B configurations 和共计 700 个 sampled batch shapes。若 kernel 将 S pad 到 subblock_size=16 的倍数，则同时记录 sampled length 与 executed padded length，并以 original sampled tokens 计算 effective throughput。

对于每个 workload，评估：

Original SGLang-JAX

MXU Port

Fully Optimized

每个 benchmark point 至少应包含：

warm-up 之后 N = [要求：例如 ≥30] 次 measured iterations，

median latency，

p10/p90 或其他 uncertainty interval，

actual-token throughput，以及适用时的 padded-token throughput。

runtime measurements 必须排除 compilation。Fixed-length 与 ShareGPT-distributed 结果应分开汇总；除非明确给出 serving-weighted aggregation policy，否则不要把两类 workloads 合并成一个 geometric mean。

# 9. End-to-End Performance

[图 6：Main performance figure。]

对于每个 fixed sequence length，使用八个 panels，分别对应：

S = 8

S = 16

S = 32

S = 64

S = 128

S = 256

S = 512

S = 1024

（八个 panels 按 fixed sequence length 划分；P 不设维度，见 §1。）

X-axis：

Batch size

Y-axis：

KDA prefill latency

三条 curves：

Original

MXU Port

Optimized

这应当是主要的 absolute-performance figure。

ShareGPT-distributed workload 单独一个 panel（x 轴为 B），每个 point 汇总 100 个 sampled batch shapes 的 median 和 uncertainty interval。Fixed-length 与 ShareGPT-distributed 结果应作为 Figure 6 的独立 subfigures 展示。

对于 fixed-length workloads，optimized implementation 在 [X/56] 个 evaluated benchmark points 上降低了 KDA prefill latency。对于 ShareGPT-distributed workloads，该实现则在 [Y/7] 个 B configurations 上降低了 distribution median latency。

Fixed-length workloads 的整体 speedup range 与 geometric mean 为：

$$[MIN_{\text{fixed}}]\times \text{ to } [MAX_{\text{fixed}}]\times$$

$$[GM_{\text{fixed}}]\times$$

ShareGPT-distributed workloads 的整体 speedup range 与 geometric mean 为：

$$[MIN_{\text{ShareGPT}}]\times \text{ to } [MAX_{\text{ShareGPT}}]\times$$

$$[GM_{\text{ShareGPT}}]\times$$

最大 improvement 出现在：

$$S=[?],\quad B=[?]$$

此时 latency 从

[X] ms

降至

[Y] ms，

对应

[Z]×。

[插入最终数据。]

# 10. Serving Workload Space 中的 Performance

[图 7：8×7 speedup heatmap。]

一张 fixed-length heatmap（S × B），外加一行 ShareGPT-derived length distribution 的 distribution-median speedup（按 B 展开）。

Rows：

S=8,16,32,64,128,256,512,1024

Columns：

B=1,2,4,8,16,32,64

每个 cell 包含：

$$\frac{T_\text{baseline}}{T_\text{optimized}}$$

该图应让读者一眼看出 optimization 在各个 workload regimes 上是否稳定。

讨论应明确区分以下 workload regimes。

对于 short sequence lengths（S<64，亚 chunk 边界），performance 主要受 fixed kernel 和 pipeline overhead 影响。

对于 long sequence lengths，重复的 chunk processing 会提高 sustained matrix efficiency 的重要性。

在 high batch sizes 下，更多 independent work 可能改善 utilization，同时也可能出现 memory traffic 或 capacity constraints 等其他影响。

对于 ShareGPT-distributed shapes，讨论应说明 short-prompt mass 与 long-tail samples 如何共同影响 distribution-weighted latency 和 effective throughput，并量化 subblock padding 带来的 token overhead。

[要求：只能在观察到实际数据之后推导这些解释。不要编写没有数据支持的 regime story。]

# 11. Ablation Study

为了理解 performance improvement 来自何处，我们逐步启用每项 optimization。

[图 8：Ablation。]

Configurations：

Baseline

+ MXU mapping

+ fuse gate activation/cumsum/intra and recurrence/output

+ unified layout

+ flat packed-chunk grid

+ head-block layout/interleaving

对于每个 stage，同时报告 cumulative latency 和 incremental speedup。测量口径： headline 固定形状 S=8192(4×2048 packed), H=16，median，warmup 5 / repeat 20； 开关即代码内 ablation flags，逐级累加顺序为实测顺序：

ShareGPT-derived 分布批同口径：1591 → 86.6 ms = 18.4× （waterfall 2.48 × 4.22 × 1.20 × 1.46；各档份额随 workload 变化，见 §7 注）。 另有 14 种开关组合的正确性×性能矩阵（含非默认组合）核查数据在 aug-12-learn/test_ablation.py 与 bench_switches.py。

[要求：在覆盖多个 fixed sequence lengths 和 ShareGPT-distributed sampled shapes 的 representative workloads 上运行 ablation。单个 favorable shape 并不足够。]

🟡 状态：headline 形状与 ShareGPT 分布两个 workload 已跑（上表）；按 §8.1.1 S 集合（8-1024）与逐 B 的消融 ⏳ 未跑（56 点网格先行覆盖 3 实现端点， 中间档消融按需补）。

# 12. 时间花在了哪里？

除了 latency，还应通过 profiling 解释 performance gains。

[图 9：Execution-time breakdown。]

Profiler breakdown 至少应覆盖 short 与 long 的 fixed-length workloads，以及从 ShareGPT empirical distribution 采样的 short、median 和 tail shapes，并在每个 case 中记录 (S, B)。

对比：

Original

MXU Port

Optimized

将 runtime 分解为以下 measured categories：

MXU compute

VPU/vector compute

layout/transpose/concat

memory movement

other

预期的定性演变为：

Original implementation：matrix-friendly work 未被高效映射到 MXU。

MXU port：arithmetic 显著加快，使 layout 和 execution-structure overhead 暴露出来。

Final implementation：这些 secondary overheads 得到降低，使 MXU acceleration 能够转化为 end-to-end kernel speedup。

[要求：使用可以通过 XProf 实际测量或合理归因的 categories。避免推测性的 profiler decomposition。]

# 13. Correctness

所有 optimizations 都保持 Kimi K3 KDA computation 不变。

我们将 TPU kernel 与两类对象进行对照验证：fp32 naive recurrent KDA 参考实现 （逐 token 递归展开，JAX/numpy 直译），以及 SGLang-JAX 官方 kda_test.py 端到端单测（4×32k varlen 组合）。

在 evaluated shape space 中，我们测量：

maximum absolute error: 2.77e-4（safe_gate 有限 Neumann 因子链形式）

对照 oracle 门限 atol = 5e-4（bf16 直接求逆后乘的对照形式为 6.9e-4，不通过—— 这是采用因子链乘 RHS 形式的决定性依据）；

官方单测口径 rtol = 2e-2（4×32k varlen e2e）：7/7 通过（真机 TPU， PALLAS_INTERPRET=0），且 baseline 与优化 kernel 双双通过同一单测。

[要求：测试 random inputs、可用时符合真实 Kimi K3 distributions 的输入、零/非零 initial_state、boundary sequence lengths、非 16 对齐的 sequence lengths、ShareGPT-distributed sampled lengths，以及所有 supported batch sizes。验证 padding/masking 不会改变 valid tokens 的结果。]

🟡 状态（已覆盖项）：random inputs（q/k 按模型语义 L2-normalize）；varlen 组合 LENS=[130,64,210]（含非 16/非 64 对齐边界，验证 padding 不污染有效 token）； 随机 initial_state（任意 cached prefix 的 state 延续语义）与零 state （cold prefill）同路径；ShareGPT-derived 长度样本 [53,14,13,33,16,480,512,77] 上 optimized / baseline / naive 三方互拍 max|diff| ≤ 1.6e-4；14 种 ablation 开关组合在 interpret 与真机双模式全绿； head-block 路径与非 head-block 路径在 H ∈ {8, 16, 32} 下输出逐位一致（max|diff| = 0）——交错与布局优化不改变任何 数值。⏳ 未覆盖：全部 B 档的枚举。

由于 Kimi K3 依赖 bounded decay 实现 BF16-safe formulation，correctness evaluation 还需明确 optimized implementation 所假定的 valid gate range。strip-GEMM 的指数预算论证如下：子块中点参考下最坏指数幅度 (SB/2)·|lower_bound|/ln2 = 8×5/0.693 ≈ 57.7 < 126（bf16 指数上限）， lower_bound=-5, SB=16 下有限；该边界随 |lower_bound| 或 SB 增大而收紧 （SB=16 不可放大，无界 gate 模型不适用本路径——kernel docstring 已固化此约束）。

# 14. Lessons

优化 TPU 上的 Kimi K3 后，我们得出了若干经验。

第一，algorithm-level hardware friendliness 并不会自动转化为 system-level efficiency。Kimi K3 暴露了规则的 matrix computation，但 software stack 仍需将这些 computation 显式映射到 MXU。

第二，加速 dominant arithmetic 会改变 bottleneck。一旦 matrix computation 足够快，tensor layout、materialization、kernel boundaries 和 scheduling 就会成为一阶问题。

第三，应将 matrix accelerator utilization 视为 end-to-end dataflow 问题。如果构造 operands 需要昂贵的 layout transformation，那么更快的 GEMM 所能提供的价值就会十分有限。

最后，小型 dependent matrix operations 所需的 scheduling decisions 不同于传统 large-GEMM optimization。Independent KDA heads 和 blocks 提供了 interleave work 的机会，可以隐藏原本会暴露的 pipeline latency。

# 15. 结论

Kimi K3 提供了一种在 numerical 层面兼容 BF16 chunk computation、在结构上适合 matrix accelerators 的 KDA formulation。然而，现有 TPU path 并未充分利用这一机会。

我们首先将 matrix-based FlashKDA computation strategy 移植到 TPU，把主要 KDA prefill computation 从 vector-oriented execution 转向 MXU matrix operations，并获得 2.28×（固定形状）/ 2.48×（ShareGPT-derived 分布）speedup。

对该实现进行 profiling 后，我们发现 computation structure、kernel boundaries、tensor layout 和 execution scheduling 中出现了一组新的 bottlenecks。通过 structural optimization、kernel fusion、layout optimization 和 pipelining 解决这些 bottlenecks 后，又进一步获得 5.6×（固定形状累计）的 improvement。

在已实测的 Kimi K3 prefill workloads（headline 固定形状与 ShareGPT-derived varlen 分布批）上，最终实现相较于原始 SGLang-JAX 实现达到 12.8×（固定形状）/ 18.4×（ShareGPT-derived 分布）speedup。cached prefix 不构成 kernel 评估维度（KDA 递归 state 定长，见 §1），评估空间为 S × B。

[最终要求：完整的 56-point fixed-length benchmark matrix（8 S × 7 B × 3 实现）与 7 个 ShareGPT-distributed B configurations（每个 replay 100 个 sampled lengths）可用后，将“最高达到”的表述替换为两类 workloads 各自的 geometric-mean/range 结果。headline 应同时包含 typical 和 maximum performance。]

# 附录 A：Chunk Size（BT）的选择

BT=64 经扫描确定（数据：aug-12-learn/sweep_*.json）。更大的 BT（128/256）虽然 摊薄每 chunk 的固定成本，但 chunk 内三角求解的工作量 ∝ BT²（Neumann 因子链的 [BT,BT] 小 GEMM 规模与级数深度同时上升），且 VMEM 工作集随 BT² 增长挤压双缓冲； 实测 BT=64 全面优于 128。这与 FlashKDA 在 GPU 上选 C=16 不冲突：C=16 服务的是 fp16 数值范围与 SM 占用，TPU/bf16 下 MXU 友好的块尺寸是 64/128 量级，数值安全性 由 §4 的指数预算（子块 SB=16 保持不变）单独保证。

# 附录 B：VMEM 边界与 head-block 的适用条件

head-block 模式（kernel 一次消费全部 H 个 head）的 VMEM 工作集 ∝ H。实测边界： H=16, K=V=128 时非 safe_gate 路径（每 head 独立衰减张量共 16 份展开）需 45.13 MB，超过 v6e 的 32 MB scoped VMEM 上限，编译期 OOM；safe_gate 路径 （strip-GEMM 共享子块结构）同形状工作集 ~5 MB。因此 head_block 生效条件为 fuse ∧ safe_gate ∧ H % 8 == 0（TPU 块布局要求次末维为 8 的倍数）， 不满足时自动回退非 hb 路径（数值不变，见 §13 逐位一致性）。该边界由 生产形状（H=16/K=128）的失败案例发现——小形状单测不会触发， 已固化为 kernel 内的静态断言与开关约束（commit 98920a41）。