Top-k

Top-k 从每行中选出最大的 k 个值及其下标。它把第 12 节的跨通道归约推进了一步：不仅要知道最大值是多少，还要知道它在哪里，并且要重复 k 次。

本节是本章的一次综合。前面各节给每个硬件单元测出了延迟和发射间隔，这里用它们回答一个完整的问题：一种 top-k 的写法要多少周期，时间花在哪里，换一种写法能快多少。顺序是：先从 argmax 出发，看 XLU 怎样同时给出最大值和下标；再看 k = 8 时的重复选取，量出它的周期数，并归纳决定周期数的三类资源；然后看当前 Pallas 实现的错误，以及修正它要多少周期；最后换一类不串行的算法。

这个问题在 top-k 论文中有完整的研究：正确性的严格定义、25 个形状上的对照、几种新算法，以及让编译器生成所需指令的办法。本节只取其中与“怎样用周期数分析一个 kernel”有关的部分，实验都是自包含的，数字与论文一致；更深入的内容在相应的位置指向论文。

本节实验：指令与正确性 源码、输出；原生 XLA 的清单 源码、输出；vmax.index.xlane 的规则与计算段的模型 源码、输出；TC VMEM 边界内的周期数 源码、输出。几种写法在 top_k_common.py 中。

写法：两个输出

本小节实验源码、输出。

top-k 有两个输出：值与下标。相对于第 1 节的最小 kernel，out_type 是两个输出组成的 tuple，scratch 中多两个 TC VMEM buffer，计算一次写两个 Ref：

out_type=(pltpu.HBM(values.shape, values.dtype), pltpu.HBM(indices.shape, indices.dtype)),
...
def kernel(x_hbm: Ref, v_hbm: Ref, i_hbm: Ref, x_vmem: Ref, v_vmem: Ref, i_vmem: Ref, sem: Ref) -> None:
    pltpu.async_copy(x_hbm, x_vmem, sem).wait()
    v_vmem[...], i_vmem[...] = f(x_vmem[...])
    pltpu.async_copy(v_vmem, v_hbm, sem).wait()
    pltpu.async_copy(i_vmem, i_hbm, sem).wait()

输出 Ref 按 out_type 的顺序排在输入之后。输出明确放在 HBM 的原因见第 12 节。

基本形式：argmax 是一次 XLU 操作

jnp.max(x, axis=1, keepdims=True) 与 jnp.argmax(x, axis=1, keepdims=True) 一起，输入 DMA 之后的清单是：

{ va0: vlaneseq.8x128.u32 v1 ; vld: vld.8x128 v0, [vmem:0x0] }
{ va0: vand.8x128.u32 v2, 0x7f, v1 ; vx0: vmax.xlane.2.8x128.f32 trf0, v0 }   # v2 = lane 号；提交求最大值
{ va0: veq.8x128.s32 vm0, v2, 0 }                                            # vm0 = 第 0 个 lane
{ vx0: vmax.index.xlane.0.8x128.f32 trf0, v0 }                               # 提交求最大值的下标
{ vr0: vpop.8x128 v3, trf0 }                                                 # 取回最大值
{ vst: vst.msk.8x128 [vmem:0x8], vm0, v3 }                                   # 只写第 0 个 lane
{ vr0: vpop.8x128 v4, trf0 }                                                 # 取回下标
{ vst: vst.msk.8x128 [vmem:0x10], vm0, v4 }

XLU 有专门给出最大值下标的指令 vmax.index.xlane，不需要先求最大值、再比较、再找位置。两次提交进入同一个队列 trf0，按提交的顺序取回。与第 12 节的 vadd.xlane 一样，结果广播到整行的 128 个通道；输出是 [8,1]，只占每行的第 0 个通道，所以编译器用 vlaneseq 与 0x7f 相与得到通道号，比较出“第 0 个通道”的掩码，再用带掩码的 vst.msk 只写这一个通道。下标直接作为 int32 写出：vmax.index.xlane 助记符中的 f32 指输入的类型，结果是整数下标。值与下标都与 CPU 一致。

vmax.index.xlane 的规则

本小节实验源码、输出。

上面的输入每行互不相同。并列、NaN 时选哪个下标，要用 tpuasm 直接构造输入来看。实验用第 3 节的 LccProbe 执行手写片段：先用 vlaneseq、vand、vcvt、veq、vsel 在 v13 中构造一种输入（每个子通道相同），再依次提交 vmax.index.xlane、vmax.xlane 和 vmin.index.xlane，取回写出。实验只改构造 v13 的那几条指令：

由此得到三条规则：

• 并列时取通道号最大的一个。 最大值或最小值出现在多个通道时，vmax.index 与 vmin.index 都返回其中最大的通道号：lane & 15 的最大值 15 返回 127，最小值 0 返回 112。全部相等时返回 127。这与 jnp.argmax 的约定（取第一个）相反，而 Mosaic 直接用这条指令实现 jnp.argmax，没有做任何修正。第一个实验加了一组每行最大值在 lane 15、31、…、127 并列的输入：

  ## max 与 argmax，每行的最大值在 lane 15、31、…、127 并列
    值与 CPU 结果一致：True；下标与 CPU 结果一致：False
    第 0 行下标：[127]；CPU：[15]

  所以在 Pallas kernel 中，一行不超过 128 个元素时，jnp.argmax 遇到并列返回最后一个下标，与 JAX 的语义不同（更宽的行见“只改宽度”一小节）。这是 JAX issue #34620 记录的问题，jnp.argmin 也一样。需要第一个下标时，可以用下文手写 top-k 的办法：先求最大值，再在等于它的位置中用 vmin.xlane 取最小的下标。

• NaN 优先。 有 NaN 时，最大值与最小值都是 NaN，两个下标都指向它。

• −0 与 +0 相等。 它们按并列处理。

还测了从提交到取回的延迟：vmax.xlane、vmax.index.xlane、vmin.xlane 都是 79 个周期，换用 vx1 槽与 trf1 队列也一样，与第 12 节测得的 vadd.xlane 相同。

写法的限制：is_stable 与 bf16

jax.lax.top_k(x, k) 在 kernel 中无法编译：

is_stable=True is not supported in Pallas top_k. For efficiency, only is_stable=False is supported

lax.top_k 默认要求稳定：相等的值按原来的顺序排列。Pallas 只支持 jax.lax.top_k(x, k, is_stable=False)，即相等值之间的顺序不确定。实验的每行都是互不相同的值，两者结果相同。

bf16 输入无法编译：bfloat16 top_k is not supported on TPUv5 or older。需要时先转换成 f32。

只改 k：重复选取 8 次

jax.lax.top_k(x, 8, is_stable=False) 的值和下标都正确。清单中是 8 轮选取：

清单的前两轮（只保留计算指令）：

{ va0: vand.8x128.u32 v2, 0x7f, v1 ; vx0: vmax.index.xlane.2.8x128.f32 trf0, v0 }   # 第 1 轮：原始输入的下标
{ va0: veq.8x128.s32 vm6, v2, 0 ; va1: vlt.8x128.s32 vm7, v2, 2 }                    # 拼结果用的掩码
{ vr0: vpop.8x128 v3, trf0 }                                                         # 第 1 个下标
{ va0: veq.8x128.s32 vm0, v2, v3 }                                                   # 选中的位置
{ va0: vsel.8x128 v4, vm0, 0xff800000, v0 ; va1: vlt.8x128.s32 vm0, v2, 3 }          # 把它改成 −inf
{ vx0: vmax.index.xlane.0.8x128.f32 trf0, v4 }                                       # 第 2 轮
{ vr0: vpop.8x128 v5, trf0 }
{ va0: veq.8x128.s32 vm1, v2, v5 ; va1: vsel.8x128 v15, vm6, v3, v5 }                # lane 0 放第 1 个下标，其余放第 2 个
...

每一轮是一条依赖链：提交 vmax.index.xlane，等 79 个周期取回下标，veq 比较出选中的 lane，vsel 把它改成 0xff800000（−inf），下一轮对改过的数组再求下标。下标链在关键路径上；最大值则不在：编译器在后面对第 1–8 轮的数组 v0、v4、v6… 各提交一次 vmax.xlane，它们之间没有依赖，每 8 个周期可以发射一条。

8 个结果要拼成一行的第 0–7 个通道。编译器事先算出“lane < 2”“lane < 3”…“lane < 8”的掩码（7 条 vlt，加上“lane = 0”的一条 veq），每取回一个结果就做一次 vsel：小于 j 的 lane 保留之前的结果，其余放第 j 个结果。值和下标各拼一次。

助记符中的 .0–.3 与队列的对应是固定的：.0、.2 的结果进入 trf0，.1、.3 进入 trf1（与第 10 节 vxpose 的变体相同），编译器轮流使用它们，并同时用 vx0、vx1 两个槽提交。

周期数：一条每轮 82 个周期的链

本小节实验源码、输出。

把 f32[8,128] 的 top-k 放进第 3 节的 TC VMEM 边界：输入和两个结果都在 TC VMEM，两个结果共同经过 optimization_barrier。每行是 0 到 127 的随机排列，4 份输入各运行一次，值和下标全部核对。只改 k：

Pallas 的周期数与 k 成正比，每轮 82 个周期。 从 k = 8 到 16 多 656 个周期，从 16 到 32 多 1312 个，都是每轮 82 个。这个数可以直接从清单读出来：每一轮提交 vmax.index.xlane，等 79 个周期取回下标（第 12 节），再用 vpop、veq、vsel 各一个周期把选中的位置改掉，下一轮才能提交。k 轮串成一条链，79 个周期的等待一轮也省不掉。

求最大值的那 k 次 vmax.xlane 不在这条链上：它们只依赖各轮的工作数组，不被后面的轮次等待。编译器把它们排在链的等待里，XLU 每 8 个周期可以接受一条（第 9 节），所以归约的次数是 2k，周期数却只按 k 次往返增长。

把计算部分的 bundle 原样插进 LccProbe（第 13 节的办法），用发射模型预测，结果与实测逐个相同：

  argmax（max 与 argmax）：9 个 bundle，2 次 XLU 归约；实测 R2 − R0 = 103，模型 103
  lax.top_k(x, 8, is_stable=False)：54 个 bundle，16 次 XLU 归约；实测 R2 − R0 = 671，模型 671

54 个 bundle 要 658 个周期（671 − 13），周期数与 bundle 数几乎无关，只与链上往返的次数有关。

原生 XLA 的周期数与 k 无关。 XLA 对 k = 1 调用一个通用的 TopK 实现，对其余的 k 把整行连同下标一起排序，再取前 k 个；区间内还有一次 CMEM → TC VMEM 的 DMA（它生成的下标数组放在 CMEM），所以读数有几个周期的波动。排序整行大约 2565 个周期，取几名都一样。

两条规律相交：k = 8 时 Pallas 是 XLA 的四分之一，k = 32 时已经比 XLA 慢。一种写法是快是慢，要看它的周期数随什么增长。

三类资源

到这里，本章遇到的周期数可以归成三类。一段不含 DMA 的计算，时间大致由其中一类决定：

判断属于哪一类，决定了改哪里才有用。lax.top_k 属于第一类：减少指令没有用（清单里本来就没有多少指令在链上），只能缩短链、让链不经过 XLU，或者换成不串行的算法。下面几个小节依次改变形状，看瓶颈怎样在这三类之间移动。

这只是用来判断方向的框架，不是精确的预测：精确的周期数要逐个 bundle 推算（第三章第 4 节的发射模型）并实测。论文的 8.2 节把这种估计给出的下限与实测放在一起，链的估计最准（79 对 82），不经过 XLU 的写法偏差最大。

只改宽度：跨两个 lane tile

f32[8,256] 每行跨两个 TC VREG，结果仍然正确。清单多出了 8 条 vmax.8x128.f32、8 条 vge.8x128.f32 和 8 条 vperm：每一轮先在两个 TC VREG 之间逐元素比较、取较大者，再对较大者做跨通道的选取，下标也要换算回 0–255 的范围。行宽每多一个 TC VREG，每一轮就要多一层这样的合并。

这种合并改变了并列时的选择。只看 argmax：两个 TC VREG 先用 vge 逐通道比较，再用 vsel 各取较大的值和它的下标（来自后一个 TC VREG 的下标由 vadd 加上 128），vmax.index.xlane 只在合并后的 128 个通道上执行一次，最后 vperm 按它给出的通道取出下标。于是并列时的选择分两级：vge 在相等时取前一个 TC VREG，vmax.index.xlane 取通道号最大的一个。实验加了三组并列的输入：

结果既不总是第一个下标，也不总是最后一个：选中的是“下标除以 128 的余数”最大的那个，余数相同时取靠前的 TC VREG。所以 256 个元素全部相等时得到 127，而不是 0 或 255。

周期数。 同一个实验测了更宽的行，k = 8：

行宽超过 128 之后，每轮从 82 个周期变成 160 多个。原因就是上面清单里多出的那条 vperm：vmax.index.xlane 给出的只是通道号，还要再经过一次 XLU（69 个周期）查出这个通道来自哪个 TC VREG，下一轮才知道该把哪个元素标为已选。链上从一次往返变成了两次。行宽从 256 加到 1024，每轮只多十几个周期：多出来的 TC VREG 由逐元素的 vmax、vge、vsel 合并，属于第三类资源，比 XLU 的往返便宜得多。

f32[8,256] 上原生 XLA 在这个外壳里编译失败（Fused computation shape ... is not equal to the fusion shape）。这是 XLA 把排序和取前 k 列融合之后，Pin 的内存空间没有传到融合内部造成的，论文的附录 C.2 分析了它并给出了进程内的修复，修复之后是 2325 个周期。

论文的 5.1 节给出了去掉第二次往返的写法：先在每个通道内把各个 TC VREG 的元素排好，每轮只对最上面一层做一次 argmax，选中的元素来自哪个 TC VREG 由随排序移动的位置数组查出，不在链上。f32[8,1024] 取前 8 名由此从 1445 降到 943 个周期。

只改行数：XLU 排不下

行宽保持 128，增加行数，k = 8：

行数增加 8 倍，Pallas 只多 81 个周期：每轮 79 个周期的等待是所有行共用的，8 个 TC VREG 每轮各提交两条归约，两个队列在一轮的等待里排得下（16 条 × 8 个周期 ÷ 2 个队列 = 64 个周期）。再增加一倍就排不下了：16 个 TC VREG 每轮 32 条归约，两个队列要 128 个周期才能提交完，超过了一轮的等待，周期数开始按归约的总数增长。瓶颈从第一类（往返延迟）换成了第二类（发射间隔）。

这时有用的改动也随之不同：减少每轮的归约条数（每轮只求下标，循环结束后用一次 gather 取回全部的值），或者把数据转置之后沿子通道处理，完全不经过 XLU。论文的 5.4 节与第 6 章给出了这两种写法，f32[128,128] 取前 8 名从 1235 降到 937（只用 Pallas 源码）和 802 个周期（用到 Pallas 目前编译不出的按索引读，第 9 节）。

一个边界错误：有效值不足 k 个

输入的每行只有 3 个有限值，其余都是 -inf，求 top-8：

第 0 行下标：[0, 2, 1, 127, 127, 127, 127, 127]
CPU：        [0, 2, 1, 3, 4, 5, 6, 7]

值是对的，但后 5 个下标全都是 127，重复了。is_stable=False 允许相等值以任意顺序出现，但不允许同一个下标出现多次。“vmax.index.xlane 的规则”一小节解释了它：前 3 轮选完有限值后，整行都是 −inf，已选位置被改成的 −inf 与原有的 −inf 无法区分；全部并列时 vmax.index.xlane 返回 lane 127，第 4 轮选中它、把它改成 −inf（值不变），第 5 轮又全部并列，再次返回 127。

这类“有效值少于 k”的情况在实际中并不少见，例如对已经做过掩码的分数求 top-k。

错误不止这一处。 论文的第 4 章系统地检查了 Pallas 的这个实现，以 IEEE 754 的 totalOrder（原生 XLA 与 CPU 的 lax.top_k 实际遵守的顺序）为正确的定义，归纳出三类错误：

• 已选的位置被重新选中，即上面的重复下标。后果可以比重复更重：被反复选中的 lane 127 如果本来是一个有限值（论文 4.1 节的例子），它先作为真实的值返回一次，之后又以 −inf 的值重复出现，值与下标不再配对。
• argmax 的并列与 NaN 规则随数据的摆放变化。“只改宽度”一小节已经看到跨两个 lane tile 时并列的选择不同；NaN 更严重：合并两个 TC VREG 时值用 vmax 取（传播 NaN），下标却由 vge 的比较结果决定（与 NaN 的比较总是假），只有落在最后一个 lane tile 中的 NaN 才被找到，top-k 的每一轮都返回同一个 NaN。
• 不保持完整位型的顺序。负号的 NaN 应当排在 −inf 之后，XLU 把任何 NaN 都当作最大值；+0 应当排在 −0 之前，XLU 认为两者相等。

在 197592 条覆盖特殊值的输入上，这个实现在 70706 条上出错，原生 XLA 没有错误。它在正态随机数、小整数并列和只含有限值的输入上全部正确，这正是上游测试覆盖的范围。写 kernel 时要记住：lax.top_k 在 Pallas 中只对“没有 −inf、没有 NaN、行宽不超过 128”的输入可靠。

修正要多少周期

JAX 仓库中这个实现的注释说，保留重复下标是为了避免昂贵的防御性掩码。修正到底要多少周期，取决于修正的工作放在哪里。

用掩码排除已选位置：两次归约都在链上

问题出在“把已选位置的值改成 -inf”：改完之后，已选位置与原本就是 -inf 的位置无法区分。最直接的修正是把“是否已选”与值分开，记在掩码寄存器里（第 7 节）：

def f(x):
    lane = jax.lax.broadcasted_iota(jnp.int32, x.shape, 1).astype(jnp.float32)
    taken = jnp.zeros(x.shape, jnp.bool_)
    for _ in range(k):
        best = jnp.max(jnp.where(taken, -jnp.inf, x), axis=1, keepdims=True)
        index = jnp.min(jnp.where(~taken & (x == best), lane, float(x.shape[1])), axis=1, keepdims=True)
        taken = taken | (lane == index)
        ...

每一轮先求未选位置中的最大值，再在“未选且等于最大值”的位置中取最小的下标，最后把这个下标标为已选。只有 3 个有限值的输入上，值和下标都与 CPU 完全一致；“相等时取最小的下标”恰好就是稳定排序的顺序。归约的次数与 lax.top_k 一样是 2k 次，周期数却是两倍：

每轮 166 个周期。区别在依赖关系：lax.top_k 每轮的两次归约中，只有求下标的一次在链上；这里每轮先求最大值，再用它求下标，两次 79 个周期的归约都在链上。注释所说的“昂贵”指的大概是这一类写法。

这个写法还不能用于含 NaN 的输入。最大值是 NaN 时，x == best 处处为假，没有一个位置参与取最小下标，归约得到的是填充值 128：

## 手写 top-8，用掩码排除已选位置，lane 7 为 NaN
  值与 CPU 结果一致：False；下标与 CPU 结果一致：False
  第 0 行下标：[128, 128, 128, 128, 128, 128, 128, 128]；CPU：[7, 115, 95, 0, 21, 97, 91, 4]

抬高：修正不在链上

上面的周期数表明，lax.top_k 的时间几乎全是在等那条 argmax 的链：每轮 79 个周期的等待中，XLU 每 8 个周期就可以再接受一条归约，向量 ALU 更是空闲。只要不加长这条链，多做一些工作几乎不花周期。

论文 4.7 节的“抬高”就是这样安排的。设一行有 m 个不等于 −inf 的元素。前 m 轮里总有未选的、不等于 −inf 的元素，argmax 不会选到原有的 −inf，这段时间用 −inf 表示已选是安全的。选完这 m 个之后，把原有的 −inf 抬到 f32 的最小有限值：真实的值都已选走，不会冲突；已选的位置仍是 −inf，排在它们后面。

def f(x):
    lane = jax.lax.broadcasted_iota(jnp.int32, x.shape, 1)
    lowest = jnp.finfo(x.dtype).min
    original = x == -jnp.inf
    m = jnp.sum(jnp.where(original, 0.0, 1.0), axis=1, keepdims=True)
    lift_round = jnp.where(original, jnp.maximum(m, 1.0), -1.0)     # 每个原有的 -inf 在第几轮之后被抬高
    current = x
    for j in range(k):
        index = jnp.argmax(current, axis=1, keepdims=True).astype(jnp.int32)
        values.append(jnp.where(m <= j, -jnp.inf, jnp.max(current, axis=1, keepdims=True)))
        indices.append(index)
        current = jnp.where(lane == index, -jnp.inf, jnp.where(lift_round == j + 1, lowest, current))

与 lax.top_k 相比多了三样东西，都不在链上：m 是一次归约，只依赖输入，与第 0 轮的 argmax 同时提交；每个位置在哪一轮被抬高是预先算好的 lift_round；每轮多一条 vsel 完成抬高，它与标记已选的那条 vsel 合在一起，在等待本轮 argmax 的 79 个周期里早已就绪。

## 手写 top-8，抬高，每行只有 3 个有限值，其余为 -inf
  值与 CPU 结果一致：True；下标与 CPU 结果一致：False
  第 0 行下标：[0, 2, 1, 127, 126, 125, 124, 123]；CPU：[0, 2, 1, 3, 4, 5, 6, 7]

后 5 个下标与 CPU 不同，但互不相同，对应的输入都是 −inf，这是 is_stable=False 允许的。全部为 −inf 的输入、含 NaN 的输入、没有 −inf 的输入，结果也都正确。周期数：

归约多了一次，周期数只多 3 到 10 个，与 k 无关。两种修正做的事情差不多，一个慢一倍，一个几乎不花周期，决定周期数的是工作有没有落在链上，而不是工作的多少。“修正昂贵”这个前提不成立。

抬高只解决重复下标，不解决前面列出的后两类错误。论文的 4.8、4.9 节继续往下做：把修正完全移出循环（用 MXU 算出每个 −inf 的名次，685 个周期，与不修正相同），以及做到与原生 XLA 逐位一致（717 个周期，比不修正多 4.7%）。

换一类算法：秩计数

前面的写法都保留了那条长度为 k 的链，k = 32 时已经不如原生 XLA 的排序。k 大时应当换成第二类：往返很多但互不依赖，让时间由发射间隔而不是延迟决定。

秩计数就是这样的算法。一个元素的秩是排在它前面的元素个数，秩小于 k 的元素就是结果，秩正好是它在输出中的位置。把整行循环移动 d 个通道（第 9 节的 vrot），与原来的行逐元素比较，就完成了所有相距 d 的元素对；d 从 1 到 127 各做一次，这 127 次移位互不依赖：

def f(x):
    n = x.shape[1]
    lane = jax.lax.broadcasted_iota(jnp.int32, x.shape, 1)
    bits = jax.lax.bitcast_convert_type(x, jnp.int32)
    key = jnp.where(bits < 0, bits ^ 0x7fffffff, bits)
    count = jnp.where(key == -2**31, lane, 0)
    lower = key - 1
    for distance in range(1, n):
        other = pltpu.roll(key, distance, axis=1)
        count += jnp.where(other > jnp.where(lane >= distance, lower, key), 1, 0)
    indices = jnp.concatenate([jnp.max(jnp.where(count == rank, lane.astype(jnp.float32), -1.0), axis=1, keepdims=True) for rank in range(k)], axis=1).astype(jnp.int32)
    return jnp.take_along_axis(x, indices, axis=1), indices

• 比较用整数键。 key 把 f32 的位型变成一个有符号整数：非负的位型原样保留，负号的位型翻转低 31 位。按整数比较 key 的顺序正好是浮点数的全序，NaN、−inf、±0 都只是键的一个取值，不需要任何特例。比较在向量 ALU 中进行，不经过只接受 f32 的 XLU 归约。
• 相等的元素规定下标小的在前。 与下标更小的元素比较时用 key − 1，于是所有元素的秩互不相同，结果是稳定的。键等于最小整数时减一会回绕，这些位置的秩用自己的下标预先补上。
• 结果。 秩等于 r 的那个位置的通道号由一次 vmax.xlane 取出，k 个名次 k 次归约，互不依赖；值再用一次沿通道的 gather 取回。

它的结果在全部测试输入上与 CPU 逐个相同，包括并列、只有 3 个有限值和含 NaN 的输入。周期数：

秩计数的周期数几乎与 k 无关：每多取一名只多一次归约，两个队列平均每 4 个周期完成一次，从 k = 16 到 32 正好多 64 个周期。它的时间由提交的条数决定：127 次移位分给两个队列，光是提交就要 64 × 8 = 512 个周期，再加上最后一次移位、取下标的归约和取值的 gather 各一次往返。所以 k 小时它不如链（k = 8 时 798 对 695），k = 16 起反超，k = 32 时是 lax.top_k 的三分之一，也是原生 XLA 的三分之一。

三种瓶颈在这张表里都出现了：lax.top_k 受往返延迟限制，与 k 成正比；秩计数受发射间隔限制，与 k 几乎无关；原生 XLA 排序整行，区间里有上千个 bundle 的逐元素指令。按 k 选择写法，就是比较这几条随 k 变化的直线。

更进一步

本节的写法都只用了 Pallas 现有的接口。论文在 25 个形状上按瓶颈逐一给出了写法，相对 Pallas 的 lax.top_k 周期数的几何平均加速比是 1.67×，并且全部快于原生 XLA，同时做到与原生 XLA 逐位一致。几条与本章内容直接相关的结论：

最后两行用到了第 9 节的按索引读写，以及本教程没有介绍的分段转置 vsxpose，这些指令 Pallas 目前编译不出，论文先用 tpuasm 把手写的片段插进编译出的程序，再在进程内给编译器打补丁让它生成这些指令。其中有一个与本章主题直接有关的数：同一个算法、同一组指令，编译器排出的顺序是 714 个周期，按实测的延迟重排之后是 595 个周期，一条指令也没有改。周期数取决于每种资源被谁占用、占多久，以及这些占用怎样错开，第三章第 4 节接着讲这件事。

对照：原生 XLA

本小节实验源码、输出。

XLA 对 top-1 调用一个通用的 TopK 实现，对 top-8 则把整行连同下标一起排序再取前 8 个。两者的结果都正确，包括 -inf 的边界情况（XLA 的 top_k 默认 is_stable=True）。但对于“每行 128 个数中取最大的几个”，Pallas 中一次 vmax.xlane 加一次 vmax.index.xlane 的 argmax，比 XLA 的上千个 bundle 少得多；只是要自己处理前面几小节的错误。

bundle 数不是周期数：上面“周期数”各小节在同一个 TC VMEM 边界内测得，top-8 时 XLA 是 2565 个周期，Pallas 的 lax.top_k 是 685 个；但 Pallas 的周期数随 k 增长，k = 32 时已经超过 XLA，而且 XLA 在全部输入上都正确。越过 XLA 这条底线，既要更快，也要同样正确。