我们写了5个新算子，让vLLM支持DeepSeek-V4.1-Flash跑在A100上

DeepSeek-V4.1-Flash 发布后很快进了 vLLM 主干，但官方验证过的平台只有 H200、GB200、GB300、MI350X。翻一遍代码就明白为什么 —— 稀疏 MLA 的每一个注意力后端都把 Ampere 挡在门外：

flashmla_sparse.py:171          capability.major in [9, 10]
flashinfer_mla_sparse.py:98     capability.major == 10
flashinfer_mla_sparse_sm90.py   capability.major == 9
flashattn_mla_sparse.py:72      capability.major == 9

A100/A800 是 major == 8，一个都不匹配。再往下，indexer 依赖 DeepGEMM，而 DeepGEMM 在 Ampere 上直接抛 Unsupported architecture。

也就是说：A100 / A800 这一代卡，一行都跑不起来。

我们把它跑起来了。本文讲清楚两件事：拦路的到底是什么，以及为了绕过去我们写了哪些算子。

所有数据来自一台 8×A800 80GB PCIe 的机器，模型是完整的 DeepSeek-V4.1-Flash（475.3 GiB 权重，含视觉塔）。

先把话说在前面：这些改动没有合进 vLLM 主干，我们也不指望它合进去（原因见文末第八节）。交付形式是一个可直接用的 fork 和一个 docker 镜像。另外，我们只在 sm_80 上实测过 —— 标题里的"Ampere 卡"请按第八节的适用范围理解。

一、先说结果

相对最初"能跑起来"的版本，单流优化了 **+48%**。

关于那两列：真实权重是最终交付性能，但它在实例之间有约 7% 的抖动（Triton autotune 每个进程选到的配置不同），做 A/B 时会淹没一个 8% 的改进。所以所有优化对比都在一个 40 层齐全、只缩小 engram 表的 dummy 权重配置上做，读数确定性到 0.01 ms。两者绝对值差 6–10%，趋势一致。方法细节见文末附录 —— 我们在这一点上吃过亏。

一个容易被忽略的数字：每卡只占 36.59 GiB。475.3 GiB 的权重按 8 卡均摊该是 59.4 GiB，对不上 —— 差额的来源可以在 checkpoint 里直接查到：

model-00047-of-00048.safetensors    95 GiB   layers.1.engram.*
model-00048-of-00048.safetensors    95 GiB   layers.14.engram.*
其余 46 个分片                     285.3 GiB

48 个分片里的最后两个各 95 GiB，装的全是 engram 的 n-gram 表，而且只属于第 1 层和第 14 层（每层 6 个张量）。这 190 GiB 不进显存，是在主机共享内存里 mmap 后注册给 CUDA、经 UVA 访问的。剩下的 285.3 GiB 按 8 卡分是 35.7 GiB —— 和实测的 36.59 GiB 正好对上。

这也解释了为什么容器必须给 --ipc=host 或 --shm-size=256g：默认的 64MB /dev/shm 连一张表的零头都放不下，会直接启动失败。

二、拦路的到底是什么

我们没有一上来就写代码，而是先做了三个实测，用来判断哪些墙是真的、哪些是绕得过去的。

实测一：DeepSeek 官方的参考算子，一半能在 Ampere 上跑

模型目录里带了 DeepSeek 自己的参考实现（TileLang 写的），而且它显式设了 TL_DISABLE_WARP_SPECIALIZED 和 TL_DISABLE_TMA_LOWER —— 关掉了 Hopper 专有特性。我们直接在 A800 上编译运行：

结论比预想的乐观得多：稀疏注意力的数学本身是可移植的，真正过不去的只有窄精度 MMA —— Ampere 没有 FP8/FP4 张量核指令。而这一块 vLLM 本来就有 Marlin 和 dequant-to-BF16 的替代路径。

实测二：vLLM 自己就有一个平台无关的 BF16 稀疏 MLA 算子，只是名字起错了

vllm/v1/attention/ops/xpu_mla_sparse.py 里的 triton_bf16_mla_sparse_interface，零 XPU 专有代码、不涉及 fp8。我们按 DSv4.1 的真实几何（h_q=64，dim_qk=576，d_v=512，topk=512）在 A800 上一跑就过，rel err 2.5e-3。

它还返回 lse（log-sum-exp），这一点后面会变得非常关键。

实测三：Triton 在 sm_80 上拒绝 fp8e4nv —— 但有个后门

这是整个移植里最硬的一堵墙。DSv4.1 的 KV cache 是 OCP E4M3 格式，而 Triton 的 CUDA 后端在 sm_80/sm_86 上连这个类型本身都不接受，连 bitcast=True 都不行：

ValueError: type fp8e4nv not supported in this architecture.
            The supported fp8 dtypes are ('fp8e4b15', 'fp8e5')

注意报错里那句"支持的类型"。我们把所有 fp8 类型扫了一遍：

fp8e4b15 和 OCP E4M3 只差一个指数偏置（15 vs 7）。这意味着把 E4M3 的位模式当作 e4b15 读出来，得到的值恰好小 2⁸ 倍 —— 乘个 256 就还原了。次正规数也对得上，因为两种格式的次正规缩放因子相同。

这就是我们的第一个算子。

三、我们写的 5 个算子

算子 1 & 2：可移植的 E4M3 编解码

@triton.jit
defe4m3_bytes_to_float(x_u8, NATIVE_FP8: tl.constexpr):
    if NATIVE_FP8:
        return x_u8.to(tl.float8e4nv, bitcast=True).to(tl.float32)
    # E4M3 偏置是 7，e4b15 是 15，重解释出来的位读小了 2^-8
    return x_u8.to(tl.float8e4b15, bitcast=True).to(tl.float32) * 256.0


@triton.jit
deffloat_to_e4m3_bytes(x, NATIVE_FP8: tl.constexpr):
    if NATIVE_FP8:
        return x.to(tl.float8e4nv).to(tl.uint8, bitcast=True)
    # 0.00390625 == 2^-8，先抵消偏置差再转换
    return (x * 0.00390625).to(tl.float8e4b15).to(tl.uint8, bitcast=True)

两行代码，但精度要交代清楚：

• 解码方向：254 个有限字节模式全部逐位精确。 只有 0x7F / 0xFF（E4M3 的两个 NaN 编码）不同，而任何格式良好的编码器都不会产生它们。这个方向可以无保留地直接用。
• 编码方向：在"恰好中点"上与 IEEE 不同。 我们是 round-half-away-from-zero，IEEE 是 round-half-to-even —— 252 个中点里 126 个不同（全部偏向远离零）。超范围饱和、次正规、负数全部与 torch 一致；随机数据上逐字节 100% 吻合，因为恰好中点是零测集。NaN 输入会饱和成 448 而不是传播。对 KV cache 量化这是亚 ULP 差异，但必须如实标注。

为什么值得单独拿出来讲：在此之前，这个问题的已知解法是手写整数位操作（V4 时代的 fork 就是这么干的），或者像上游某个 PR 那样引入 654 行 CUDA 加一个预编译的 .bc bitcode 文件。而位偏置这个性质让它塌缩成一次 bitcast 加一次乘法，无查找表、无额外依赖、无主机侧内存。

这一个改动就打通了整条 fp8 KV cache 的读写路径：cache 读写、indexer 的 K store、engram 的 fp8 嵌入 gather。

算子 3：按槽位 gather + 反量化分页 KV

DSv4.1 的 KV cache 是分页的，每 token 一条 584 字节记录：448 维 fp8（64 维一组 UE8M0 scale）+ 64 维 bf16 RoPE，一个 block 的 scale 排在它所有 token 数据之后。

解码时需要按任意槽位列表把这些行取出来、反量化成 BF16。上游有按序列连续 gather 的算子，但没有按任意索引 gather 的。我们写了一个：

row = tl.program_id(0).to(tl.int64)
slot = tl.load(indices_ptr + row).to(tl.int64)
...
block_idx = slot // CACHE_BLOCK_SIZE
pos_in_block = slot % CACHE_BLOCK_SIZE
# 全程 int64：block_idx * block_stride 在大 cache 上会溢出 int32
cache_block_ptr = cache_ptr + block_idx * block_stride

那句 int64 的注释不是防御性编程 —— V4 时代我们在这里踩过一模一样的坑：block_idx 按 int32 乘以 pool 的 block stride，在 block 序号超过约 2065 时回绕成负偏移，表现为长上下文下的非法内存访问，或者更恶劣的静默数据损坏。

算子 4 & 5：split-KV 稀疏注意力 + LSE 合并

这是收益最大的一个，也是最能说明"为什么通用算子在小 batch 解码下会垮"的例子。

问题：实测二里那个能用的 BF16 稀疏 MLA 算子，只在 token 和 head 两个维度并行。batch 1 解码时 h_q=64、BLOCK_H=16，于是：

grid = (num_tokens, cdiv(h_q, BLOCK_H)) = (1, 4)

只有 4 个 CTA。 A800 有 108 个 SM，96% 的算力在空转，而这 4 个 CTA 里每一个都要串行走完全部 40 个 key block。

诊断方法很直接：喂 8 个 token 和喂 1 个 token，耗时都是 79 µs。八倍的活、同样的时间 —— 这是启动受限的教科书特征，不是吞吐受限。

解法是 flash-decoding：把 key 维度也切开并行，每个 CTA 维护自己的 running max 和分母，最后用 log-sum-exp 合并。实测二里提到的那个"返回 lse"在这里派上了用场 —— 合并所需的量本来就在。

顺带把 per-head 的 attention sink 折进合并步骤：sink 是一个附加到 softmax 分母上的裸 logit，既然合并本来就在算分母，多加一项等于免费，还省掉了原先 wrapper 里一整趟 elementwise。

结果：batch 1 的 77 µs → 13.5 µs，5.7×，而且在实测的每一个 batch 尺寸上都更快。精度还更好了（3.3e-3 vs 6.3e-3，因为合并在 fp32 里做）。

但这里有个坑，值得单独写一段。

第一版的合并算子是按 head-block 循环 split，每次迭代加载一个 [BLOCK_H, BLOCK_D] 的 fp32 tile。分开计时才发现：

split 算子标度得非常漂亮（77 → 9.2 µs，8.4×），但合并算子在超过 4 个 split 后爆炸 —— 比它要合并的对象还慢 10 倍。 原因是 static_range 展开后每次迭代 32KB 的 tile，8 个 split 就开始寄存器溢出。这一版整体反而比原算子还慢。

改成一个 CTA 只管一个 head（64 个 CTA 而不是 4 个），一次加载 [n_splits, BLOCK_D] 整块沿 split 轴规约，合并稳定在 ~4 µs，与 split 数无关。这才有了上面那个 5.7×。

教训很朴素：split 和合并必须分开计时。 只看端到端的话，你只会看到"split-KV 没用"，然后放弃一个能给 5.7× 的方向。

split 数的启发式瞄准约 256 个总 CTA，在 7 个实测形状上全部命中最优点，测试里把这些配对钉死了，以后要重调必须是有意识的。

四、性能：延迟、吞吐、并发

单流的优化是分步叠加出来的（下表为完整 40 层架构 + dummy 权重，读数确定性到 0.01 ms）：

并发曲线（同上，dummy 权重，--max-num-seqs 32）则揭示了一个没有单一最优解的事实，单位 tok/s：

交叉点在并发 2 和 4 之间。

PP 为什么在低并发赢：这台机器的拓扑是 GPU 0–3 和 4–7 各自 PIX 互联（对应两个 NUMA 节点），跨组走 SYS（CPU 互联）。TP=8 的每次 allreduce 都要跨这一跳；而 PP=2 下每个 TP=4 组正好落在一个岛内，跨 socket 那一跳被并行策略本身消掉了。所以 PP 配置不需要也不应该再开分层 allreduce —— 4 卡组本来就没有两个岛可分。

PP 为什么在高并发输：流水线气泡。请求少时凑不满两个阶段，TP=8 的利用率优势就压过来了。

还有一个细节值得记：我们本来预期 PP=8 × TP=1 会是大赢家（allreduce 完全消失），但它被模型自身的结构挡死。DSv4.1 的 kv/index 源层是 [2, 8, 14, 20]，每个源层的消费者一直延伸到下一个源层，所以合法的流水线边界只能落在这几个位置。均分 10 层或 5 层都会把共享组切断。20/20 是唯一既均衡又合法的切法 —— 也就是 PP=2。 这不是调参调出来的，是架构决定的。

上下文长度

模型 config 声明 max_position_embeddings = 1048576（YaRN 外推到 1M）。我们能实际拉到多长，取决于 KV cache 放不放得下 —— 而 DSv4.1 的 KV 是压缩 MLA + fp8，单 token 成本极低：

注意 262144 那一行：cache 总量反而更大，因为更长的上下文让 vLLM 把 block 表调得更省。即使按 256k 上下文开，也还能塞下 34 路并发 —— 长上下文在这个模型上不是"只能跑一条"的奢侈品。

实测我们喂进去 150,068 个 prompt token 做 needle-in-a-haystack，埋在中段的信息准确取回。更长的没测，不是因为跑不动，而是构造和验证一个可信的 500k token 测试本身要花掉数小时的机时 —— 所以我们只敢说"实测到 150k"，1M 是模型的声明值，不是我们的实测值。

五、我们试过但否掉的方向

技术文章只写成功的部分是不诚实的。这几条我们都完整实现并实测过，结论是负面的：

1. 打开 torch.compile / inductor。VLLM_USE_BREAKABLE_CUDAGRAPH=0 确实能解锁 CompilationMode.VLLM_COMPILE，模型也能跑。但每步 kernel 启动数完全没变（1850 vs 1845）—— DSv4.1 的计算几乎全在 custom op 和 TileLang/Triton 算子里，inductor 根本看不进去，没东西可融。而且更慢（16.03 vs 13.45 ms/step）。

2. 把 gather-反量化融进注意力算子。 写完了，数值也对（3–5e-3，边界情况覆盖），但更慢：标准化对比下 batch 1 到 32 全线 0.83–0.95×。分散的 584 字节记录读取加上注意力状态的寄存器压力，在内层循环里比单独一趟连续反量化更贵，省下的 HBM 往返和 78 次 kernel 启动补不回来。已回退。

3. 手写 MXFP8 GEMV 替代 Marlin。 batch 1 时所有 GEMM 都是 M=1 的瘦长形状，Marlin 实测比权重访存下限差 3–15 倍，看起来有很大空间。但原型做出来只在 5 个真实形状里赢了 2 个（1.65× 和 1.39×），其余的反而更慢，端到端只值 3.1% —— 不值得为此维护两套权重布局。

4. 继续优化 allreduce。 10 KB 消息在 8 路 PCIe 上要 44 µs（NCCL 自己是 63 µs），换算成带宽只有 190 MB/s。传输本身只要 0.5 µs，99% 是延迟。这已经接近物理下限，压不动了。

顺带一个观察：这台机器上每个 kernel 有 ~8–10 µs 的固有下限，我们手写的干净 GEMV 也撞同一堵墙。每步 1845 次启动 —— 系统性的出路是融合成更少更大的算子，但如上所述，可融合的面比想象中小得多。

六、路上修的两个上游 bug

Triton 里 constexpr 守卫的提前 return 不会裁剪后续代码。

ifnot QUANTIZE:      # QUANTIZE 是 tl.constexpr
    tl.store(out, x)
    return
x_quant = ...to(tl.float8e4nv)   # 仍然会被编译！

改成显式 if/else 才会正确裁剪。这个性质导致 fused_inv_rope_fp8_quant(quantize=False) 在任何 fp8 编译不了的架构上都是坏的 —— 在 SM90 上没人发现，因为那里 fp8 能编译。

多模态模型的 PP 拿不到 input_ids。 DSv4.1 的 MoE gate 每一层都要靠 input_ids 识别图像 token，但 vLLM 在非首流水线阶段无条件把它置空。而 VL 模型明明声明了 requires_raw_input_tokens = True，runner 在多模态 embeds 那条路上尊重了这个协议，PP 这条路上没有。要在两处修（真实 forward 和 cudagraph 捕获）—— 只修前者的话，捕获阶段会再挂一次。

七、怎么用

docker run --gpus all --init --ipc=host --network host \