Prhub

#37413 [AMD][DSV4] feat: enable fp8 two-pool unified_kv on gfx950

原始 PR 作者 amd-danli103 合并时间 2026-09-14 17:49 文件变更 21 提交数 13 评论 4 代码增减 +3594 / -104

执行摘要

DSV4 在 gfx950 启用 fp8 双池 KV,容量提升 1.5×

PR body 明确指出问题:unified_kv 下 DSv4 把整个 latent 都以 bf16 保存,qk_nope_head_dim 448 + qk_rope_head_dim 64 意味着“1024 B per token per layer”,而“The nope half actually does not need that precision”。把 nope 按 1x64 tile 量化成 fp8(带 E8M0 scale)、rope 保持 bf16,行宽降到 640 B,“which measures out at 1.51× the KV capacity at identical memory on gfx950”。目标是在固定 HBM 下让更多 token 常驻,支撑更长上下文与更高并发;同时刻意把变更限定为默认关闭、gfx950 必需、其他架构 warn + 回退 bf16,以保证既有路径行为不变(“so nothing on the existing path changes”)。

值得精读,尤其是正在做 KV 量化、异构内存池或新硬件后端的人。建议按以下顺序看:先读 python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py 与 unified_kv_kernels/layout.py 建立行宽/两池的单一真相,再读 models/deepseek_v4.py 的 _forward_prepare 看“用 caller buffer 决定布局分支”的设计,最后看 runtime.py 与 fused_qk_norm_rope_store.py 理解写入/读取两侧的形态约束。

值得关注的设计决策有三个:(1) 把 E8M0 scale 在行内重复写两份、并用 512 B 2 的幂 stride 对齐 asm reader,属于“为内核取数模式服务”的布局设计;(2) 所有未就绪组合一律显式拒绝并附原因与后续 PR 号,而不是靠默认值静默降级;(3) 判据挂在 caller 传入的 buffer 上,从结构上消除 prepare 与 backend 对布局的分歧。

如果要把它用于生产,请先确认三件事:目标场景的并发区间(c=8 附近有 7.6% 解码损失)、是否需要 PD/HiCache/CP(当前均不可用)、以及是否有独立复测(现有精度证据只有 GSM8K 的四组配对,性能证据全部来自作者自测)。若只是评估设计而非上线,本 PR 的价值在于它展示了一次“KV 行宽级重构”如何做到默认零影响、边界显式化。

讨论亮点

本 PR 没有行内 review 评论,评审信号只有一条 APPROVED: LGTM(HaiShaw)与 4 条 issue 级评论;真正有交锋价值的内容几乎全部由作者写在 PR body 里,属于“作者自我 review + 一条 LGTM 放行”的评审形态。

值得记录的讨论要点(引自 PR body / 评论):

HaiShaw:/tag-and-rerun-ci;随后 @ispobock @Fridge003 @hnyls2002 @xiezhq-hermann @hzh0425 @hebiao064 @BBuf @ShangmingCai Please have a review.

xiezhq-hermann:@DarkSharpness can you help take a look?

HaiShaw:PR Test Base & PR Test Extra all pass.(review 结论:LGTM)

作者主动标注最差格子:“The cell to flag is 8192/1024 at c=8 under TP=8: 7.6% less throughput, ITL +1.18 ms. That is the worst of the 14 TP=8 cells … The leftover gap is the decode reader, not anything this PR still owns.” 结论是 aiter 的 mla_a8w8 asm 内核按 64 q head 调优,TP=8 只剩 16 个本地 head 时受占用率限制,“Closing the TP=8 gap needs an mla_a8w8 variant for 16 q heads on the aiter side; that is a separate PR.”

能力边界的处理方式:“PD disaggregation and HiCache offload are refused for now because both assume one pool per layer.” 代码里落实为 get_buf_infos / get_contiguous_buf_infos / get_unified_swa_ring_buf_infos 抛 NotImplementedError,并带 TODO(danli103) 与后续 PR 号。

被放弃的方案:“A Triton rewrite was measured and was slower at every shape, so it was dropped.”

测试落位也经过一轮修正(commit f28857c):“#37436 rejects new files under test/registered/amd/. These stay on the mi35x stage-b suite, so e2e/dsv4 is the kind that keeps that registration.”

未在讨论中解决、只能靠后续 PR 收口的问题:TP8 低并发解码回退、PD(#39166)、HiCache(#37778)、DSpark MTP(#38901)、HiSparse 与 PCP。

实现拆解

本 PR 是“布局 + 写入 + 读取 + 接线 + 调优”一体化改造,核心思路是把一个 bf16 latent 拆成行数一致、row index 语义共享的两个池,让所有读写方按同一份行宽定义工作。

  1. 布局与 sizing(python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py、python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/layout.py):新增 DSV4_FP8_NOPE_ROW_BYTES = 512、DSV4_FP8_QUANT_TILE = 64 与 check_two_pool_pair 守卫;DeepSeekV4UnifiedKVPool.__init__ 增加 fp8 形参,原 kv_buffer 改为 fp8 nope 行缓冲并新增 kv_buffer_rope(bf16 rope,行数与 nope 池严格相同);dsv4_unified_row_bytes() 成为唯一行宽公式(bf16 1024 B / fp8 640 B),越界时抛 ValueError 而非 assert,理由写在代码注释里——“sizing runs under -O too, and a silently skipped check here overreports capacity”。pool_configurator.py 把行宽接入显存预算(bytes_per_full_token 由 8918.00 降到 5945.00),使容量报告与实际分配同源,这是 1.50× 数字可信的前提。
  2. 写入侧三处 store:SWA scatter store_swa_into_unified 增加可选的 rope tensor / rope pool 参数,两个池按同一个 ring row index 落行,单字节行按 uint8 搬运以避开 E8M0 scale 的浮点转换;fused_qk_norm_rope_store.py 新增 _fp8_2buff_store,交给 aiter group-quant 融合内核,一次 launch 同时产出打包 fp8 nope 行与 bf16 rope 行,既有 packed 与 bf16 路径及其 module key 不变;compressor 侧新增 fp8_2buff 包装(dsv4/compress.py、compressor_v2.py,仅对 512 宽 latent 构建);JIT 模板 fused_norm_rope_v2.cuh 是 common code 的唯一改动点——新增一个 elif 分支,旧分支不动。
  3. 读取侧两个内核:decode 走 aiter mla_a8w8 v4 nm asm(runtime.decode_fp8_2buff),prefill 走 aiter opus(runtime.prefill_fp8_2buff,把 paged prefix 与本 chunk 的 flat extend 区用同一行布局寻址);runtime.py 里补齐两池/行宽的启动前断言,并新增 decode_qo_indptr(qo_indptr 只依赖 token 数,改为每次 forward 建一次)。
  4. 模型与后端接线(python/sglang/srt/models/deepseek_v4.py、layers/attention/deepseek_v4_backend_hip_radix.py):MQALayer._forward_prepare 新增 q_rope_out / k_nope_out / k_rope_out 三个 buffer 形参,并新增 fuse_prefill 与 fuse_verify_fp8 两个分支;所有判据都挂在 caller 传入的 buffer 上(“Keyed off those buffers the same way fuse_prefill is, so the two arms cannot disagree about the layout”),从结构上排除 prepare 与 backend 对布局判断不一致的可能。fp8 下 prefill/verify 的 kv 由 k_nope_out 顶替并保持 save_kv_cache,ring 写的时机由 backend 决定(prefill 在 attention 之后、verify 在之前)。
  5. 性能调优与测试配套:decode 路径删掉对空 segment 的 masked_fill_(builders 产不出该形态)、跳过 group-quant 内核不需要的 kv.contiguous()、>40 token 时把 num_kv_splits 上限压到 4;prefill 分支同样去掉一次白拷贝。测试侧新增 8 个文件,覆盖行宽公式与两池分配(test_dsv4_unified_fp8_pool.py)、SWA scatter、fused QK norm+RoPE、compressor store、decode/prefill 读数、backend prefill 接线与模型打包契约;其中两个纯 CPU 用例注册到 base-a-test-cpu,其余注册到 stage-b-test-1-gpu-small-amd-mi35x(非 gfx950 运行器会 skip)。
文件 模块 状态 重要度
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py KV 缓存池 modified 8.17
python/sglang/srt/models/deepseek_v4.py DSv4 模型 modified 7.97
python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py 内核运行时 modified 7.51
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py 注意力后端 modified 7.09
python/sglang/kernels/ops/attention/fused_qk_norm_rope_store.py 融合归一化 modified 6.67
python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/layout.py 行布局 added 6.16
python/sglang/srt/model_executor/pool_configurator.py 显存配置 modified 7.0
test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py 打包契约 added 7.97
test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py 解码读取 added 7.56
test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py 预填写入 added 7.56
test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py 融合归一化 added 7.56
test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py 行宽计算 added 7.25
python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh JIT 内核 modified 5.67
python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py 特性开关 modified 5.27
python/sglang/srt/layers/attention/dsv4/compressor_v2.py 压缩器 modified 5.88
python/sglang/srt/environ.py 环境变量 modified 5.35

关键符号

dsv4_unified_row_bytes DeepSeekV4UnifiedKVPool.get_unified_kv_rope DeepSeekV4TokenToKVPool.get_unified_kv_rope MQALayer._forward_prepare DeepseekV4HipRadixBackend._forward_unified_kv decode_fp8_2buff prefill_fp8_2buff decode_qo_indptr _fp8_2buff_store fused_qk_norm_rope_swa_store is_unified_kv_fp8 check_two_pool_pair store_swa_into_unified

关键源码片段

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py data-contract

双池布局与行宽公式的落点:新增 fp8 nope 池 + bf16 rope 池、`dsv4_unified_row_bytes()` 统一 sizing,并对单池假设的传输接口显式拒绝,是整条链路的地基。

# python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
# 行宽与 tile 尺寸由 kernels 侧的 layout 模块定义,这里只做重导出:
# 分配行的池与消费行的内核必须共用同一份 stride 真相。
DSV4_FP8_NOPE_ROW_BYTES = layout.DSV4_FP8_NOPE_ROW_BYTES # 512 B
DSV4_FP8_QUANT_TILE = layout.DSV4_FP8_QUANT_TILE # 每 64 个元素一个 E8M0 scale
​
​
def dsv4_unified_row_bytes(
    qk_nope_head_dim: int, qk_rope_head_dim: int, fp8: bool
) -> int:
    """一个 token 在 unified_kv 中占用的字节数(两个池相加)。"""
    if not fp8:
        # bf16 单池:448 + 64 = 512 个 bf16 元素 = 1024 B / token / layer
        return (qk_nope_head_dim + qk_rope_head_dim) * 2
    # fp8 双池:nope 每元素 1 B,外加每个 1x64 tile 一个 E8M0 scale;
    # scale 在行内写两份,因为 asm reader 每个 tile 会读两次。
    num_tiles = -(-qk_nope_head_dim // DSV4_FP8_QUANT_TILE) # 向上取整:448 -> 7
    scale_bytes = 2 * num_tiles
    # 这里必须是 raise 而不是 assert:sizing 在 python -O 下同样要生效,
    # 静默跳过会把可用显存算多,进而按错误的行数去分配。
    if qk_nope_head_dim + scale_bytes > DSV4_FP8_NOPE_ROW_BYTES:
        raise ValueError(
            f"fp8 nope row overflows: {qk_nope_head_dim} latent values at 1 B "
            f"{scale_bytes} B scales > {DSV4_FP8_NOPE_ROW_BYTES} B stride"
        )
    # 512 B 定长 nope 行 + 128 B bf16 rope 行 = 640 B / token / layer(0.625x)
    return DSV4_FP8_NOPE_ROW_BYTES + qk_rope_head_dim * 2
​
​
class DeepSeekV4UnifiedKVPool:
    """bf16 是单池 ``[rows, head_dim]``;fp8 是行数完全相同的两个池。
​
    row index 在两池中指向同一个 token(SWA ring 用
    ``req_pool_indices * swa_window + pos % swa_window``,压缩区用
    ``swa_pages + page_index``),所以两边行数必须一一对应。
    """
​
    def __init__(self, ..., swa_ring_size: int, fp8: bool = False): # 其余形参略
        self.fp8 = fp8
        self.rope_head_dim = qk_rope_head_dim
        self.head_dim = qk_nope_head_dim + qk_rope_head_dim
        bufs, rope_bufs = [], []
        # rows = swa_pages + padded_compress_rows,两个池共用同一个 rows
        if self.fp8:
            bufs.append(
                torch.zeros(
                    rows, DSV4_FP8_NOPE_ROW_BYTES,
                    dtype=torch.float8_e4m3fn, device=device,
                )
            )
            rope_bufs.append(
                torch.zeros(
                    rows, self.rope_head_dim,
                    dtype=torch.bfloat16, device=device,
                )
            )
        else:
            bufs.append(
                torch.zeros(rows, self.head_dim,
                            dtype=torch.bfloat16, device=device)
            )
            rope_bufs.append(None) # bf16 布局没有第二个池
        self.kv_buffer = bufs
        self.kv_buffer_rope = rope_bufs
​
    def get_unified_kv_rope(self, local_layer_id: int) -> torch.Tensor:
        # rope 池只存在于 fp8 布局,调用方须先用 is_unified_kv_fp8() 分流
        assert self.fp8, "rope pool only exists under SGLANG_DSV4_UNIFIED_KV_FP8"
        return self.kv_buffer_rope[local_layer_id]
​
    def get_buf_infos(self):
        if self.fp8:
            # 单池假设:一层一个指针 + 一个行宽只描述了 nope 池,
            # 消费者会搬走半行且毫无察觉,因此显式拒绝而不是给出错数据。
            # TODO(danli103): 等下游确实需要时再上报两个池。
            raise NotImplementedError(
                "get_buf_infos describes one pool per layer; the fp8 rope pool "
                "would be dropped (SGLANG_DSV4_UNIFIED_KV_FP8=1)."
            )
        ...
python/sglang/srt/models/deepseek_v4.py core-logic

模型侧接线:`_forward_prepare` 新增三个 buffer 形参与 `fuse_prefill` / `fuse_verify_fp8` 分支,决定打包对在 decode/prefill/verify 三种模式下的流向与 `save_kv_cache` 语义。

# python/sglang/srt/models/deepseek_v4.py —— MQALayer._forward_prepare 的 fp8 分流
unified = is_unified_kv_triton()
fp8_2buff = is_unified_kv_fp8()
is_decode = forward_batch.forward_mode.is_decode_or_idle()
​
fuse_verify = (
    envs.SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY.get()
    and forward_batch.forward_mode.is_target_verify()
)
# verify 在 fp8 下像 prefill 一样打包 K,但保留 verify 自己的时序:
# ring 写在 attention 之前。判据全部挂在 caller 提供的 buffer 上,
# prepare 与 backend 因此不可能对布局产生分歧。
fuse_verify_fp8 = (
    fuse_verify and unified and fp8_2buff
    and k_nope_out is not None and k_rope_out is not None
)
# prefill 复用同一个 fused store:本 chunk 的 K 既是 attention 的 extend 区,
# 也是 attention 之后 ring 写入的源,所以只量化一次、两边共用这对 buffer。
fuse_prefill = (
    unified and fp8_2buff
    and k_nope_out is not None and k_rope_out is not None
    and not is_decode
    and not forward_batch.forward_mode.is_target_verify()
)
do_fused_qk_norm_rope = (
    unified and (is_decode or fuse_verify or fuse_prefill)
) or (not unified and self.use_fused_qk_norm_rope)
​
# ……(中间省略 call kernel 的部分,调用处会带上 fp8_2buff / swa_rope_cache /
# k_nope_out / k_rope_out / q_rope_out)
​
swa_rope_cache = None
if unified and fuse_verify:
    # verify 读 ring 的姿势与 decode 相同,但自带 per-token slot map;
    # kv 是 qkv_a 的 strided slice,未融合路径要在这里显式 contiguous() 出来。
    kv = kv.contiguous()
    swa_cache, swa_loc = None, None
    swa_page_size, bf16_store = 1, not fuse_verify_fp8
elif unified and fuse_prefill:
    # 这次不写 ring:这些行是本 fwd 的 extend 区,prefix 池必须保持
    # attention 期望的形态,ring 由 backend 在 attention 之后写入。
    # kv 继续保持 qkv_a 的 stride —— fp8 下内核只读它,并且按 stride
    # 参数取行,所以原来的 contiguous() 是纯粹的白拷贝(本 PR 把它删掉)。
    swa_cache, swa_loc = None, None
    swa_page_size, bf16_store = 1, False
elif unified:
    swa_cache = token_to_kv_pool.get_unified_kv(self.layer_id)
    # swa_loc 与 layer 无关:backend 每次 forward 计算一次并缓存在 metadata
    swa_loc = attn_backend.get_unified_swa_loc(forward_batch)
    swa_page_size, bf16_store = 1, not fp8_2buff
    if fp8_2buff:
        # decode 的第二个写入目标:bf16 rope 池
        swa_rope_cache = token_to_kv_pool.get_unified_kv_rope(self.layer_id)
​
if unified and (fuse_prefill or fuse_verify_fp8):
    # 打包后的 nope 半行沿 kv 槽返回,attention 把它当 attn_k,
    # save_kv_cache 保持打开让 backend 完成 ring 写;rope 半行去了
    # caller 的 buffer,这里没有第二个返回槽。prefill 在 attention 之后写、
    # verify 在之前写,但两者读的是同一对 buffer。
    kv = k_nope_out
elif not (unified and fuse_verify):
    kv = None
test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py test-coverage

CPU 级契约测试:pin 住模型交给 backend 的“打包对”在 decode/prefill/verify/CP 各种模式下的形态,以及未融合 store 与 CP 的拒绝语义。

# test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py
# 这些用例 pin 的是“交给 backend 的是不是一个打包对”,而不是内核算术:
# decode 只需要打包的 Q(K 已经在 ring 里);prefill 自己就是 KV 源,
# 因此还要一对打包的 K;verify 两者都要——像 decode 一样读、像 prefill 一样写。
​
​
class _Harness(deepseek_v4.MQALayer):
    def _forward_prepare(
        self, x, positions, forward_batch, attn_backend,
        q_out=None, x_quant=None, q_rope_out=None,
        k_nope_out=None, k_rope_out=None,
    ):
        # 记下 caller 传入的 buffer,用于断言“store 填的”就是“attention 读的”
        self.prepare_kwargs = dict(
            q_out=q_out, q_rope_out=q_rope_out,
            k_nope_out=k_nope_out, k_rope_out=k_rope_out,
        )
        q_out.zero_()
        # 与 prefill 分支一致:打包后的 nope 半行沿 kv 槽返回,
        # 这正是调用方把 save_kv_cache 打开的原因
        return q_out, k_nope_out
​
​
class TestUnifiedFp8QPair(unittest.TestCase):
    def test_fp8_decode_hands_the_backend_a_packed_pair(self):
        layer, call = _run(fp8=True)
​
        q, q_rope = call["q"], call["q_rope"]
        self.assertEqual(q.dtype, torch.float8_e4m3fn)
        # 行宽取自 pool 而不是 head_dim:fixture 里两者被刻意设成不同值,
        # 否则“误读 self.head_dim”的回归会照样通过
        self.assertEqual(tuple(q.shape), (TOKENS, N_LOCAL_HEADS, NOPE_ROW_BYTES))
        self.assertEqual(q_rope.dtype, torch.bfloat16)
        # asm reader 把两者当扁平 buffer 走,没有任何 stride 参数
        self.assertTrue(q.is_contiguous())
        self.assertTrue(q_rope.is_contiguous())
        # 同一对 buffer 也送进了 store,否则根本没有东西把它们写出来
        self.assertIs(layer.prepare_kwargs["q_out"], q)
        self.assertIs(layer.prepare_kwargs["q_rope_out"], q_rope)
​
    def test_fp8_target_verify_gets_the_packed_pair(self):
        """verify 像 decode 一样读 ring,又像 prefill 一样喂 ring"""
        layer, call = _run(fp8=True, mode=ForwardMode.TARGET_VERIFY)
​
        # 打包 Q 是选中 decode reader 而不是 Triton reader 的开关
        self.assertEqual(call["q"].dtype, torch.float8_e4m3fn)
        k, k_rope = call["k"], call["k_rope"]
        self.assertEqual(k.dtype, torch.float8_e4m3fn)
        self.assertEqual(tuple(k.shape), (TOKENS, NOPE_ROW_BYTES))
        self.assertIs(layer.prepare_kwargs["k_nope_out"], k)
        self.assertIs(layer.prepare_kwargs["k_rope_out"], k_rope)
        # 与 prefill 不同,ring 写在 attention 之前,但标志位与 buffer 是同一对
        self.assertTrue(call["save_kv_cache"])

评论区精华

1.50× 容量收益与 GSM8K 无回归的验证口径 设计

PR body 给出四组“CLI 逐字节相同、只差一个 env 开关”的配对:TP8 135.59 GB 下 full_token 16,126,208 → 24,223,744(1.5021×),bytes_per_full_token 8918.00 → 5945.00(0.6666×),并解释为什么不是行宽的 0.625——“the indexer term in _get_bytes_per_full_token does not scale with the KV row width”。精度侧给 1319 题 GSM8K 的四格配对(0.940/0.945/0.952/0.947 对 0.949/0.948/0.947/0.946)。

结论:评审以一条 LGTM 接受作者自证的数据,没有对测量方法提出独立质询。 · 已解决(单条 LGTM 放行)

TP8 低并发解码回退的归因与归属 性能

作者主动标注 8192/1024、c=8 的 TP8 格子 ttt 比值 0.924、ITL +1.18 ms,并给出根因:aiter `mla_a8w8` asm 内核按 64 q head 调优,TP=8 只剩 16 个本地 head,读得更少(640 B vs 1024 B)却更慢;Triton 重写“was slower at every shape, so it was dropped”。

结论:不在本 PR 内修复,结论是“Closing the TP=8 gap needs an mla_a8w8 variant for 16 q heads on the aiter side; that is a separate PR”,风险被显式转嫁给外部依赖。 · 未解决(转后续 aiter 侧改动)

PD/HiCache 单池假设的显式拒绝 正确性

`get_buf_infos`、`get_contiguous_buf_infos`、`get_unified_swa_ring_buf_infos` 在 fp8 下抛 NotImplementedError,注释写明“one pointer and one row size per layer describes the nope pool only, so whoever picks this up next would move half a row and not notice”,并留下 TODO(danli103)。

结论:以“宁可报错不可静默丢 rope”为原则冻结能力,PD 与 HiCache 分别由 #39166、#37778 补齐。 · 已知限制,后续 PR 处理

未支持组合用运行时错误而非静默回退 测试

测试用 assertRaisesRegex 固化三类拒绝:fp8 target-verify 缺少融合 store 时抛 NotImplementedError("SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY")、fp8 prefill CP 抛 cp_size(“the gather hands kv back in global token order after norm+RoPE, so packing would have to move ahead of it”)、两池行数 / 行宽不一致抛 AssertionError。

结论:边界以可测试的错误语义固化,避免组合爆炸变成静默错算。 · 已解决

测试目录规约与 CI 桩对象的兼容性修补 测试

commit f28857c 说明“#37436 rejects new files under test/registered/amd/”,因此双池用例落到 test/registered/e2e/dsv4/ 并注册 mi35x stage-b 套件;commit 11cddd2 与 ecc80c8 则修补 `object.__new__` 桩对象(SWA pool-floor fixture、CPU CI 构造 DeepSeekV4TokenToKVPool 时跳过 __init__)导致的 `_fixed_swa_bytes` / `_unified_kv_fp8` AttributeError。

结论:通过类级默认 `_unified_kv_fp8 = False` 与新目录规约对齐,代价是部分测试在非 gfx950 运行器上整体 skip。 · 已解决

风险与影响

兼容性与回归

  • 默认关闭 + gfx95 门控(env_gate.is_unified_kv_fp8() 要求 SGLANG_DSV4_UNIFIED_KV_FP8 且 gfx95),bf16 路径在所有写入/读取点都保留了独立分支,测试也逐条 pin 了“bf16 不变”断言(如 test_bf16_decode_still_gets_one_plain_tensor、test_the_bf16_arm_never_touches_the_rope_pool、test_bf16_single_pool_unchanged),因此存量用户风险低。风险集中在开启开关之后的组合空间。
  • 未支持组合(PD、HiCache offload、DSA prefill CP、HiSparse、PCP)靠 NotImplementedError / AssertionError 显式拒绝,不会静默退化;但这依赖调用方先经 is_unified_kv_fp8() 分流——get_unified_kv_rope() 内部是 assert self.fp8,门控不一致时会以断言形式炸在运行时,而非给出明确语义错误。

性能

  • 作者自测最差格子为 TP8、ISL/OSL 8k/1k、concurrency=8:ttt 比值 0.924(bf16 4837 → fp8 4470),ITL +1.18 ms;14 个 TP8 格子中其余都在 4.3% 以内,concurrency ≥ 16 全部在 3.2% 以内,DP8 表基本落进噪声。根因是 aiter mla_a8w8 asm 内核按 64 q head 调优、TP=8 下只剩 16 个本地 head,属于外部依赖问题,本 PR 无法消除。
  • 本 PR 已顺手做了若干调优(删 masked_fill_、decode_qo_indptr 复用、去掉两处 kv.contiguous()、num_kv_splits 上限 4),但这些都是作者自述,未见独立 benchmark 复核。

显存与数据正确性

  • 容量收益完全建立在 dsv4_unified_row_bytes() 与两池实际分配一致之上;test_fp8_pool_bytes_match_the_sizing_row_width 做了交叉校验,test_fp8_saves_exactly_three_eighths 钉住 640/1024 = 0.625。若未来 head_dim 或 tile 尺寸变化,ValueError(而非 assert)保证 -O 下也能拦住。
  • 两池行数必须相等,row index 才在两池中语义一致;check_two_pool_pair 与测试里的“pool rows differ”断言提供了启动前保护。
  • get_buf_infos 在 fp8 下直接抛错,说明有人用时会立刻失败而不是搬半行——这是好的防御,但也意味着 PD 场景在 #39166 落地前完全不可用。

精度

  • nope 走 1x64 tile 的 E8M0 scale,误差上限约为该 tile absmax 的 2^-4(测试用 0.07 * tile_absmax 作为 dequant 判据);PR 只给出 GSM8K 1319 题的四组配对结果(0.940/0.945/0.952/0.947 对 0.949/0.948/0.947/0.946),长上下文端到端精度未做量化评测。

测试与 CI 覆盖

  • 所有双池硬件用例都注册在 stage-b-test-1-gpu-small-amd-mi35x,在 mi300 等非 gfx950 运行器上会被整体 skip,等于覆盖强绑特定硬件池。
  • PR body 的 CI 状态里 “Latest PR Test (AMD ROCm 10)” 标记为失败(:x:),而 issue 评论只声明 PR Test Base & PR Test Extra 通过;合并依据没有覆盖 AMD ROCm 10 这条线。
  • 3594 行新增、21 个文件、13 个 commit,但只有一条 LGTM 评审、零行内评论,评审深度实质等同于作者自证 + CI 绿灯。

过程风险

  • PR 自带 checklist 全部未勾选(格式化、单测、文档、benchmark),但 body 里实际附了完整可复现实验步骤与基准数据,说明流程痕迹与实际证据不一致。
  • 代码里留有两处 TODO(danli103) 与显式拒绝路径,属于“已知未完成”,若后续 PR 未落地会成为长期分支。

对用户

  • 在 MI355X / gfx950 上跑 DSv4 的用户,打开 SGLANG_DSV4_UNIFIED_KV_FP8 即可在字节级相同的 CLI 下获得约 1.50× 的 KV 容量(TP8:16,126,208 → 24,223,744 full_token;TP8+DP8+MTP 为 13,612,288 → 20,424,448),代价是低并发解码 1%~7.6% 吞吐;默认关闭,存量部署零影响。
  • 需要长上下文/高并发但可以接受轻微解码损失的线上场景(例如大并发 68k 输入)收益明显:作者测的 68k/200 各并发点 ttt 比值 0.980~1.002,基本无损。

对系统

  • KV 行宽从 1024 B 变为 640 B,改变了显存预算公式(bytes_per_full_token 8918.00 → 5945.00,实际比值 0.6666 而非行宽的 0.625,因为 indexer 项不随行宽缩放),直接影响 pool sizing、full_token 报告与容量规划;容量计算与分配同源是这套收益成立的前提。
  • row index 语义从“一个池的行”变为“两个池共享的行”,所有池消费者(PD 传输、HiCache offload、region buffer)都必须重新审视单池假设;本 PR 以显式拒绝换来安全性,把工作留给 #39166 / #37778。
  • 引入了对 aiter(v4 nm asm decode、opus prefill、group-quant 融合 store)与 gfx950 的硬依赖,非 gfx950 平台只有警告回退路径。

对团队

  • 功能线被拆成清晰的后续清单(DSpark #38901、HiCache #37778、PD #39166、HiSparse 与 PCP 待定),短期对外能力是“仅在 gfx950 且关闭 PD/HiCache/CP 的前提下可用”。
  • 测试资产(8 个新文件,覆盖写入/读取/接线/拒绝路径)为后续在每个新能力上扩展提供了模版,尤其是用 sentinel 值区分“没写”和“写了 0”、用反量化参考隔离“接线错误 vs fp8 舍入”的做法。
核心 KV 布局变更 默认关闭的硬件专属特性 TP8 低并发吞吐回退 PD/HiCache 仅拒绝未支持 测试强绑 mi35x(非 gfx950 全 skip) 评审覆盖薄(单条 LGTM)

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论