ML ENGINES · Low-bit decode attention · arXiv + Under Review
FOSCARAccelerating Decode Attention over 2- and 4-Bit KV Caches through Affine Folding
Donglin Zhuang1,2,Pragaash Ponnusamy1,Xiaoxia Wu1,Shuaiwen Leon Song2,Zhongzhu Zhou1,2
1 Together AI 2 University of Sydney
FOSCAR moves scale and offset out of per-value KV reconstruction. Hopper uses BF16 Tensor Core operands, and Blackwell can feed packed INT2 codes into FP4 Tensor Cores after remapping.
Every decode step reads the KV cache. Packing it into 2 or 4 bits cuts memory traffic, but reconstructing each value still takes GPU instructions. FOSCAR builds Tensor Core operands from packed codes and applies each group's scale to a matrix input or partial result. The offset enters as an output correction.
Step 1/3 · Find the shared group and the reduction axis.第 1/3 步 · 找到共享分组与求和方向。Step 2/3 · Follow the gold box to see where the scale is applied.第 2/3 步 · 沿着金色方框,看缩放在哪一步发生。Step 3/3 · Add the offset correction shown in blue.第 3/3 步 · 加上蓝色标出的偏移修正。
Scale the QK partial sum缩放 QK 部分和
One token’s scale is shared across the channels that QK sums over.QK 沿通道求和,同一个 token 的所有通道共享缩放。
Scale Q before QKᵀ先缩放 Q,再计算 QKᵀ
Each channel has its own scale. Apply it to the matching Q component.每个通道有自己的缩放系数,先把它乘到 Q 对应的分量上。
Scale P before PV先缩放 P,再计算 PV
Each token has its own scale. Apply it to the matching attention weight P.每个 token 有自己的缩放系数,先把它乘到对应的注意力权重 P 上。
Scale the PV partial sum缩放 PV 部分和
One channel’s scale is shared across the tokens that PV sums over.PV 沿 token 求和,同一个通道的所有 token 共享缩放。
D · channels通道Σ QKᵀ reduces hereΣ QKᵀ 沿此求和
L · cached tokens缓存 tokenΣ PV reduces hereΣ PV 沿此求和
qqqqs1 z1qqqqs2 z2qqqqs3 z3s1 z1s2 z2s3 z3s4 z4
Illustrative KV tile of quantized values q. Gold marks one group sharing a scale s and offset z.示意 KV 分块,q 为量化值。金色标出共享同一缩放 s 与偏移 z 的一组。
QPother matrix input另一侧矩阵输入× s · scale before multiply乘法前缩放
Q · KᵀP · Vmatrix product on code-derived operands对编码构造出的操作数做矩阵乘
Group-local partial result组内部分和× s · scale after multiply乘法后缩放
+ Σ(z · Q)+ Σ(z · P)offset correction before softmaxsoftmax 前加上偏移修正offset correction before softmax normalizationsoftmax 归一化前加上偏移修正
softmax÷ softmax denominator÷ softmax 分母corrected result continues修正后的结果继续计算
Scaled product缩放后的乘积sl · Σd Qd ql,d
+
Offset correction偏移修正zl · Σd Qd
Scaled product缩放后的乘积Σd (sd Qd) ql,d
+
Offset correction偏移修正Σd zd Qd
Scaled product缩放后的乘积Σl (sl Pl) ql,d
+
Offset correction偏移修正Σl zl Pl
Scaled product缩放后的乘积sd · Σl Pl ql,d
+
Offset correction偏移修正zd · Σl Pl
Interactive adaptation of Figure 4.根据论文图 4 改编。Try the arithmetic on four keys用四个键看折叠怎么算
Use query Q = [1, 2, 1, 1], INT2 codes q = [0, 1, 2, 3], scale s = 0.5 and offset z = −1. Reconstructing the keys gives [−1, −0.5, 0, 0.5]. Its dot product with Q is −1.5.
The folded route gets the same result in exact arithmetic. It applies scale and offset once to this group's contributions.
在精确算术下,折叠路径得到同一个结果。缩放和偏移作用在这一组对结果的贡献上。
From packed bits to Tensor Cores从打包位到 Tensor Cores
TMA moves packed tiles into shared memory. On Hopper, consumer warpgroups construct BF16 operands and accumulate in registers. On Blackwell, the INT2 path zero-extends codes into FP4 operands in shared memory, and the matrix products accumulate in tensor memory (TMEM).
TMA copies the packed KV tile. For channel-wise K, WG0 also prepares scaled Q.TMA 搬运打包的 KV 分块。若 K 按通道分组,WG0 还负责准备缩放后的 Q。HBM → SMEM
Copies later tiles ahead.提前搬运后续分块。
—
WG1 / WG2 · consumers消费者
The two consumers alternate even and odd tiles; one takes this tile.两个消费者交替处理奇偶分块,其中一个接手本块。
Builds BF16 operands from packed bits in registers and issues wgmma for QKᵀ.在寄存器中由打包比特构造 BF16 操作数,并发出 wgmma 计算 QKᵀ。BF16 · REG
QK correction → online softmax → PV → V correction → output, all in registers. At the end, one consumer merges both partial outputs.QK 修正 → 在线 softmax → PV → V 修正 → 输出,全程在寄存器中完成。最后由一个消费者合并两路部分结果。REG
Tensor Cores
—
wgmma accumulates QKᵀ in the issuing warpgroup's registers.wgmma 把 QKᵀ 累加到发出指令的 warpgroup 的寄存器中。REG
PV accumulates in the same consumer's registers.PV 同样累加在该消费者的寄存器中。REG
WG0 · TMA & tcgen05 issueTMA 与 tcgen05 发射
Starts the TMA copy of the packed KV tile.启动 TMA,搬运打包的 KV 分块。HBM → SMEM
Issues tcgen05 for QKᵀ on the prepared operand tile.对准备好的操作数分块发出 tcgen05,计算 QKᵀ。
Issues tcgen05 for PV once P is ready.P 就绪后发出 tcgen05,计算 PV。
WG1–WG3 · CUDA cores
For channel-wise K, the output warpgroup prepares scaled Q.若 K 按通道分组,由输出 warpgroup 准备缩放后的 Q。FP8 Q
The preparation warpgroup builds the KV operand tile in shared memory; INT2 bits become FP4 directly.准备 warpgroup 在共享内存中构造 KV 操作数分块,INT2 比特直接成为 FP4。FP4 KV · SMEMFP8 Q/P
Softmax warpgroup: reads QK, applies K corrections in FP32, computes P. Output warpgroup: reads PV, applies V corrections, keeps the running output.softmax warpgroup:读取 QK,以 FP32 完成 K 修正并计算 P。输出 warpgroup:读取 PV,完成 V 修正并维护累计输出。
Tensor Cores · TMEM
—
tcgen05 accumulates QKᵀ in TMEM.tcgen05 把 QKᵀ 累加到 TMEM。TMEM
PV accumulates in TMEM, readable by other warpgroups.PV 累加在 TMEM 中,其他 warpgroup 可以读取。TMEM
Interactive adaptation of Figure 7. Stages show the dependencies for one tile. The original schedules below show how tiles overlap.论文图 7 的交互式改编。各阶段展示单个分块的依赖关系。下方原图展示不同分块如何重叠执行。Original register packing and dataflow figures查看原始寄存器打包与数据流图
Figure 6a · Register packing and bit-level operand construction图 6a · 寄存器打包与操作数位构造Figure 6b · How an INT2 K tile reaches the Tensor Cores图 6b · INT2 K 分块如何到达 Tensor Cores
Original H100 and B200 pipeline schedules查看 H100 和 B200 原始流水线调度图Figure 7a · H100, alternating consumer warpgroups图 7a · H100 的两个消费者 warpgroup 交替处理分块Figure 7b · B200, separate issue, operand preparation, and correction lanes图 7b · B200 分别调度指令发射、操作数准备与修正
Decode kernel speedup解码内核加速
Geometric mean latency speedup over FlashInfer BF16 on each GPU. Bars share a 0–4× scale.
Replace OSCAR's INT2 decode kernel with FOSCAR while keeping its rotations and BF16 sink and recent windows. The benchmark uses Qwen3-8B, 100K input tokens, 1K output tokens, and a 99% prefix-cache hit rate. Values are peak aggregate output tokens/s across tested concurrency levels.
将 OSCAR 的 INT2 解码内核换成 FOSCAR,沿用 OSCAR 的旋转和 BF16 起始与近期窗口。测试使用 Qwen3-8B、100K 输入、1K 输出和 99% 前缀缓存命中率。表中为各并发设置下测得的峰值总输出 token/s。
GPU
OSCAR INT2
With FOSCAR换用 FOSCAR
Gain倍数
H100
269
487
1.81×
B200
333
868
2.61×
See the B200 throughput trade-off查看 B200 的单用户与总吞吐对比Figure 12c · Each point is one tested concurrency setting. Higher on the vertical axis means more total output tokens per second.图 12c · 每个点对应一个测试并发设置。纵轴越高,每秒输出的 token 总数越多。
Original figures from the FOSCAR paper (arXiv:2611.09722, CC BY 4.0).原图来自 FOSCAR 论文(arXiv:2611.09722,CC BY 4.0)。