M3 FP4 Indexer

这是一个面向 NVIDIA Blackwell 的 FP4 prefill/decode indexer。它接收已经量化并 打包的 MXFP4/NVFP4 Q、paged K 及 block scale，输出每个 128-token K block 的 最大 QK score。

当前接口支持两条自动 dispatch 的计算路径：

• 短 Q、较小 GQA ratio 使用 packed-Q decode kernel；
• 其余 shape 使用 staged prefill kernel。

公共约束为：

• max_seqlen_q > 0
• Hq % Hkv == 0
• D = 128
• K page / score block 大小固定为 128

本仓库只包含 indexer block-score 实现。TopK、CSR/schedule 构造和下游 attention 均不在这里实现，也没有引入 Quack 依赖。

环境与安装

需要：

• NVIDIA SM100/SM103 GPU；当前验证机器为 NVIDIA B200（SM100）
• 支持对应 GPU 的 CUDA、PyTorch
• nvidia-cutlass-dsl >= 4.5.2
• 精度测试额外需要 pytest

安装声明的 Python 依赖：

python -m pip install -r requirements-dev.txt

第一次调用可能触发 JIT 编译，性能采样前应先完成预热。

Public API

入口为 fp4_indexer_interface.py 中的 fp4_indexer_block_scores：

from fp4_indexer_interface import fp4_indexer_block_scores

scores = fp4_indexer_block_scores(
    q_fp4,
    k_fp4,
    q_scale,
    k_scale,
    cu_seqlens_q,
    cu_seqlens_k,
    cu_page_offsets,
    max_seqlen_q=max_seqlen_q,
    max_seqlen_k=max_seqlen_k,
    kv_indices=kv_indices,
    fp4_format="nvfp4",          # 或 "mxfp4"
    causal=True,
    qo_offset=None,
    scale_layout="public",       # 生产路径建议使用 "preordered_mma"
)

自动 dispatch

令 R = Hq / Hkv，interface 只根据调用方传入的 max_seqlen_q 和 tensor shape 中的 R 选择 score kernel：

条件	Dispatch
max_seqlen_q <= 8 且 R <= 16	packed-Q decode kernel
max_seqlen_q > 8 或 R > 16	staged prefill kernel

因此，即使实际 Q 很短，只要 R > 16 也会走 staged kernel；反过来，只有两个 decode 条件同时满足才走 packed-Q kernel。dispatch 不依赖 max_seqlen_k、batch、 causal 或 scale_layout。

wrapper 不会读取 cu_seqlens_q 的内容来决定路径，以避免 CUDA→CPU 同步；调用方 必须保证 max_seqlen_q 是所有 batch item 实际 q_len 的有效上界。staged 路径中， max_seqlen_q >= 128 会选择 Q-scale TMA 变体，较短 Q 使用普通 Q-scale staging； 这不会改变接口或数学语义。

符号

• B：batch size
• Tq = sum_b q_len[b]：varlen Q 的 token 总数
• P = sum_b ceil(k_len[b] / 128)：逻辑 K page 总数
• Hq / Hkv：Q head 数和 K/V head 数
• G：每个 D=128 向量的 scale group 数；MXFP4 为 4，NVFP4 为 8
• Ktiles = ceil(max_seqlen_k / 128)

输入 shape 和类型

参数	Shape	Dtype / 存储	含义
q_fp4	[Tq, Hq, 64]	CUDA uint8、int8 或可用时的 float4_e2m1fn_x2	逻辑 [Tq,Hq,128]，每 byte 打包两个 E2M1 FP4
k_fp4	[P, Hkv, 128, 64]	同 q_fp4	paged K；第三维是 page 内 token
q_scale（public）	[Tq, Hq, G]	MX: float8_e8m0fnu；NV: float8_e4m3fn	Q block scale
k_scale（public）	[P, Hkv, 128, G]	同上	K block scale
cu_seqlens_q	[B+1]	CUDA int32	Q 长度的 exclusive prefix sum
cu_seqlens_k	[B+1]	CUDA int32	K token 长度的 exclusive prefix sum
cu_page_offsets	[B+1]	CUDA int32	每个 batch 逻辑 K page 数的 exclusive prefix sum
kv_indices	[P]	CUDA int32	逻辑 page 到 k_fp4 physical page 的映射
qo_offset	[B] 或 None	CUDA int32	可选 causal offset；仅在 causal=True 时允许
max_seqlen_q	host scalar	Python int	所有 batch item 的最大 Q 长度，必须为正；同时参与 dispatch
max_seqlen_k	host scalar	Python int	最大 K token 长度，用于确定输出 K tile 数

所有 tensor 必须 contiguous 且位于同一 CUDA device。q_fp4/k_fp4 的 data pointer 必须 128-byte aligned；MMA scale storage 必须 32-byte aligned。

batch b 的第 k_block 个逻辑 page 通过下式定位：

physical_page = kv_indices[cu_page_offsets[b] + k_block]

cu_seqlens_k 不用于 physical page 寻址，但用于最后一个 K page 的 tail mask 和 默认 causal offset 计算。

为避免在调用前引入 CUDA→CPU 同步，wrapper 只校验这些 metadata tensor 的 shape、dtype、device 和 contiguous 属性，不读取其内容。调用方必须保证：

• cu_seqlens_q[0] == cu_seqlens_k[0] == cu_page_offsets[0] == 0
• 三个 prefix-sum tensor 单调不减
• cu_seqlens_q[-1] == Tq，且每个 q_len[b] 在 [1,max_seqlen_q] 内
• 每个 k_len[b] <= max_seqlen_k
• cu_page_offsets[b+1]-cu_page_offsets[b] == ceil(k_len[b]/128)，且末值为 P
• kv_indices 长度为 P，每个 physical page id 都在 [0,P) 内

输出 shape 和类型

scores: [Hq, Ktiles, Tq], CUDA torch.float32

其中第二维是 batch-local 的逻辑 K block 编号。属于较短 batch 的越界 block 和 被 causal mask 完全遮住的 block 为 -inf；最后一页的 tail token 在 max 前被 mask，不参与该 block 的归约。

该接口只返回 block scores，不执行 TopK。调用方如需 TopK，应沿 Ktiles 维 （dim=1）处理。

FP4 和 scale layout

两种 format 的 FP4 value 都是 E2M1，区别是 block scale：

Format	每个 scale 覆盖的 FP4 value 数	G	Scale dtype
MXFP4	32	4	torch.float8_e8m0fnu
NVFP4	16	8	torch.float8_e4m3fn

NVFP4 路径不接收、也不乘 tensor-level FP32 global scale；输入 scale 是 E4M3 block scale。

scale_layout="public"

这是易于构造和验证的 layout。wrapper 会在计算前将它转换为内部需要的 storage， 因此调用耗时包含该转换开销。

scale_layout="preordered_mma"

该 layout 期望调用方提前提供 contiguous scale storage：

shape  = (L, ceil(mn/128), ceil(G/4), 32, 4, 4)
stride = (
    512 * ceil(mn/128) * ceil(G/4),
    512 * ceil(G/4),
    512,
    16,
    4,
    1,
)

• Q scale：mn=Tq, L=Hq
• K scale：mn=128, L=P*Hkv

可以在初始化阶段调用 fp4_indexer_reorder_scales_for_mma_cute：

from fp4_indexer_interface import fp4_indexer_reorder_scales_for_mma_cute

q_mma_view, k_mma_view = fp4_indexer_reorder_scales_for_mma_cute(
    q_scale_public,
    k_scale_public,
    fp4_format="nvfp4",
)

# helper 返回 semantic view；permute 回接口需要的 storage。
q_scale_storage = q_mma_view.permute(5, 2, 4, 0, 1, 3)
k_scale_storage = k_mma_view.permute(5, 2, 4, 0, 1, 3)

之后把两个 storage tensor 传给 fp4_indexer_block_scores(..., scale_layout="preordered_mma")。benchmark 使用这条路径，并把一次性的转换排除在 计时外。

数学语义

令 scale group 大小 S=32（MXFP4）或 S=16（NVFP4），并定义：

q_global = cu_seqlens_q[b] + q_local
hk = hq // (Hq / Hkv)
physical_page = kv_indices[cu_page_offsets[b] + kb]
page_row = k_local - 128 * kb
g = floor(d / S)

Q_real(q_global, hq, d) =
    E2M1_unpack(q_fp4, q_global, hq, d) * q_scale[q_global, hq, g]

K_real(b, k_local, hk, d) =
    E2M1_unpack(k_fp4, physical_page, hk, page_row, d)
    * k_scale[physical_page, hk, page_row, g]

scores[hq, kb, q_global] =
    max over valid k_local in block kb (
        sum over d=0..127 (
            Q_real(q_global, hq, d) * K_real(b, k_local, hk, d)
        )
    )

每个 K block 覆盖 [128*kb, 128*(kb+1))。超过真实 k_len[b] 的位置不参与 max；如果一个 block 没有任何有效 k_local，结果为 -inf。

上式以 public scale layout 表示；preordered_mma 只改变 scale 的存储方式，不改变 数值语义。

causal=True 默认使用 bottom-right 对齐：

offset[b] = k_len[b] - q_len[b]
visible iff k_local <= q_local + offset[b]

传入 qo_offset[b] 时用它替代默认 offset。causal=False 时传 qo_offset 会报错。

精度测试

测试同时覆盖 staged prefill 和 packed-Q decode：

• MXFP4 / NVFP4
• staged：Q dispatch 8/9、GQA ratio 16/17、Q-scale TMA 128-row 对齐与 ragged fallback、causal compact schedule、显式 qo_offset
• decode：q_len in {1,4,8}，K full page 和 tail：k_len in {128,257}
• causal / non-causal
• 默认 bottom-right causal offset 和显式 per-batch qo_offset
• public / preordered MMA scale
• batch 1/2、随机 physical page permutation、两个 seed
• CPU E2M1 LUT + block scale dequant + FP32 matmul/block-max reference
• score tolerance：atol=0.1, rtol=0.01

运行代表性用例：

make test-smoke

运行完整 prefill/decode 精度矩阵：

make test
# 等价于：python -m pytest -q test_fp4_indexer.py

首次运行需要 JIT 编译，完整矩阵耗时会明显长于后续已缓存运行。

性能测试

列出内置 prefill/decode workload：

python test_fp4_indexer.py benchmark --list-cases

运行一个 preordered_mma scale-layout workload：

python test_fp4_indexer.py benchmark \
  --case decode_uniform \
  --format nvfp4 \
  --scale-layout preordered_mma \
  --warmup 10 --iters 100 --repeats 3

运行所有 prefill/decode workload 和两种 FP4 format：

make benchmark-all WARMUP=10 ITERS=100 REPEATS=3

自定义 prefill workload 示例：

python test_fp4_indexer.py benchmark \
  --b 1 --sq 4096 --skv 4096 \
  --head-kv 4 --qhead-per-kv 16 \
  --format nvfp4 --scale-layout preordered_mma \
  --causal

benchmark 使用 CUDA Event；JIT、首次调用、显式 warmup 和预先执行的 scale layout 转换不计入采样。Time ms 是 public API 的调用耗时。Eff TFLOPS 的 causal workload 只统计可见 QK pair，计算式为：

2 * Hq * D * sum_b(visible_qk_pairs[b]) / time

性能数据

Staged prefill：B200 CUDA Event

测试环境为 NVIDIA B200（SM100），配置为 causal、B=1、Sq=Skv=8192、 Hq/Hkv=64/4、D=128、preordered_mma，5 次 warmup、每组 20 次 launch、3 组均值：

Format	Public API time (ms)	Effective TFLOPS	Effective GB/s
MXFP4	0.8449	650.721	203.7
NVFP4	0.8756	627.938	199.1

复现命令：

python test_fp4_indexer.py benchmark \
  --case prefill_q8k_k8k --format both \
  --scale-layout preordered_mma \
  --warmup 5 --iters 20 --repeats 3

Packed-Q decode：B200 Nsight Systems

测试环境为 NVIDIA B200，配置为 NVFP4、preordered_mma、non-causal、 B=128、Sq=4、Hq/Hkv=4/1、D=128，每个 shape 统计 30 次 launch：

K length	K blocks	Core score avg (ms)	Median (ms)	Effective TFLOPS	Effective GB/s
8,192	64	0.023715	0.023679	181.108	3,213.305
16,384	128	0.053906	0.053887	159.350	2,824.510
32,768	256	0.105015	0.104991	163.594	2,898.320
65,536	512	0.199114	0.199165	172.563	3,056.466

Effective GB/s 与 benchmark CLI 使用相同的逻辑 I/O 口径：packed Q/K、 preordered Q/K scale、metadata 和 scores 各计一次，再除以 core score 时间；它不是 Nsight 采集的物理 HBM transaction 带宽。

这组数据来自 Nsight Systems 对核心 block-score 计算的单独统计，不是 public API 端到端耗时，也不包含输入 layout 转换、Python/API 开销和 TopK。因此它不能与上面 benchmark CLI 输出的 Time ms 直接混用。

性能依赖 GPU clocks、driver、CUDA/PyTorch/CUTLASS DSL 版本。代码或环境变化后应在 目标机器重新运行 CLI；需要拆分内部耗时时可使用 Nsight Systems 复测。

已知限制

• 只支持 D=128、K page/block=128。
• staged 路径要求 B * Hq <= 65535，以满足 CUDA grid Y 上限。
• 只返回 block max scores；不包含 TopK 或稀疏 metadata 构造。
• 输入必须由上游完成 FP4 packing 和 block-scale 量化。
• prefill 输出大小为 4 * Hq * ceil(max_seqlen_k/128) * Tq bytes；长序列、 大 head 数时应提前核算显存。
• 第一次调用可能触发 JIT 编译，性能采样前需要预热。