执行摘要

• 一句话：DeepSeek V4 DSpark 支持 AMD HIP，打通 unified KV ring 注入
• 推荐动作：值得精读。核心看点：_unified_inject_loc 的 ring 寻址与双重 skip 设计、HIP 后端 ragged verify 图 key 改造（token-tier graph key）、以及围绕 nightly-only 覆盖缺口的 CI 流程讨论。若要在 AMD 上部署 DSV4 DSpark，需同时跟进 #34147 保证夜间覆盖，并在后续改动中保持 target_verify_num_draft_tokens 与 server args 的 gamma 约定一致。

功能与动机

PR body 开篇明确说明动机："Follow https://github.com/sgl-project/sglang/pull/30261, support dspark for deepseek v4 on AMD platform"。此前 HIP 后端（DeepseekV4HipRadixBackend）在 ragged verify 的 CUDA graph 路径直接抛 NotImplementedError，DSpark 在 ROCm 上完全不可用；同时 DSpark draft 的 KV 存储与 target-hidden 注入仅支持非 unified 的 fp8 swa_kv_pool，无法利用 unified_kv_triton 的共享 bf16 ring，吞吐与显存都受限制。本 PR 需要同时补齐 HIP ragged verify 图支持与 unified KV 写入通道，并为 ROCm 构建未注册的 sgl_kernel 算子提供 torch 回退。

实现拆解

1. HIP 后端打通 ragged verify CUDA graph

涉及 python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py：

• 新增类属性 supports_ragged_verify_graph = True，移除 TARGET_VERIFY bucket 中针对 ragged verify 的 NotImplementedError 守卫。
• __init__ 新增 is_dspark_draft 标志与 target_verify_num_draft_tokens = speculative_num_draft_tokens - 1：DSpark draft worker 只验证 gamma 行，服务端参数沿用 CUDA 侧 "gamma + 1" 约定，避免原地改属性导致下游 5 处用法隐式消费 gamma（对应 review 意见）。
• init_forward_metadata_target_verify_old 接受 ragged_layout：通过 ExpandPrefillCausally 在 GPU 上做 prefill 展开，新增 compress_gpu_plan / extend_start_loc 支持 GPU 侧压缩计划；eager 路径回退到 sum(verify_lens)。
• init_forward_metadata_out_graph 的 TARGET_VERIFY bucket 用 ragged_layout.padded_to_bucket() 后的 graph_num_tokens 作为 token-tier 图 key（graph_key = num_tokens_v）；make_forward_metadata_from_raw_verify 同步改用 target_verify_num_draft_tokens。

2. unified KV ring 写入通道

涉及 python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py 与 python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py：

• 新增 Triton _scatter_loc_kernel 与包装函数 scatter_bf16_into_unified：逐行把已 norm + RoPE 的 bf16 K scatter 进 unified_kv[loc]，loc < 0 的行跳过（即按行 commit 掩码）。
• DeepSeekV4TokenToKVPool 新增 set_unified_key_buffer_radix_fused_norm_rope：先 fused_norm_rope_inplace（与主模型 _compute_kv_bf16 同一 freqs_cis 路径），再 scatter 进 get_unified_kv(layer_id)，签名与非 unified 版本完全一致。

3. 模型与注入/验证路径路由

• python/sglang/srt/models/deepseek_v4_dspark.py：_store_block_kv 在 is_unified_kv_triton() 时改用 get_unified_swa_loc（逐位置实时重算，支持多步 draft decode）+ unified 写入；write_target_hidden_kv 选择 unified store 函数，swa_loc/positions 契约不变。
• python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py：新增 _unified_inject_loc，地址公式 loc = state_slot * ring + pos % ring，叠加两层 skip（SWA 窗口裁剪旧 token 防 ring 槽位竞争、commit 门控丢未验收 token）；inject_ragged 分流到 build_unified_commit_inject_layout。
• python/sglang/srt/speculative/dspark_components/dspark_verify.py：commit_hidden 为 unified 模式组装逐 token state_slot（由 batch.req_pool_indices 展开）；_commit_inject 分流到 unified 布局构建。
• python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py：_forward_prefill 用 repeat_interleave 为每个 prefill token 计算 state_slot 与 final_pos，使注入器能裁剪窗口外 token——长 prefill chunk 下这些 token 会共享 ring 槽位造成覆盖。

4. ROCm 采样回退与测试配套

• python/sglang/srt/speculative/dflash_utils.py：为 top-k / top-p 概率重归一化增加 HIP torch 回退，CUDA/MUSA 行为不变。
• 新增 test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py：两个 kernel 单测用参考实现逐元素比对（test_build_unified_commit_inject_layout、test_scatter_bf16_into_unified）；8 卡 MI35x 全量 GSM8K（1319 题，5-shot）e2e 注册为 nightly 套件，门禁 accuracy > 0.92 且 avg_spec_accept_length > 3.0。

关键文件：

• test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py（模块 端到端测试；类别 test；类型 test-coverage；符号 TestDSparkUnifiedKVKernelsAMD, test_build_unified_commit_inject_layout, test_scatter_bf16_into_unified, TestDeepseekV4DSparkUnifiedKVGSM8K）: 本 PR 唯一的专用验证：两个 unified KV kernel 单测（参考实现逐元素比对）+ 8 卡 MI35x 全量 GSM8K e2e，注册为 nightly 套件；也是 PR CI 覆盖缺口讨论的中心文件。
• python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py（模块 投机注入；类别 source；类型 dependency-wiring；符号 _unified_inject_loc, inject_target_hidden, inject_ragged, _inject_mla）: DSpark target-hidden 注入的核心入口：新增 _unified_inject_loc 实现 ring 寻址与 SWA/commit 双重 skip，inject_ragged 分流到 unified 布局构建，是 unified KV 数据契约落地的关键。
• python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py（模块 注意力后端；类别 source；类型 dependency-wiring；符号 DeepseekV4HipRadixBackend, init_forward_metadata_prefill, init_forward_metadata_target_verify_old, init_forward_metadata_out_graph）: HIP 后端 DSpark 支持的核心：supports_ragged_verify_graph 开启、显式 gamma 语义（target_verify_num_draft_tokens）、GPU prefill 展开与压缩计划、token-tier 图 key，是本 PR 改动量最大的源码文件。
• python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py（模块 KV 缓存；类别 source；类型 core-logic；符号 set_unified_key_buffer_radix_fused_norm_rope）: unified KV 写入通道的核心：新增 set_unified_key_buffer_radix_fused_norm_rope，把 DSpark draft 的 norm+RoPE 后 bf16 K 写入 bf16 ring，替代非 unified 的 fp8 swa_kv_pool 路径。
• python/sglang/srt/models/deepseek_v4_dspark.py（模块 模型层；类别 source；类型 data-contract；符号 _store_block_kv, write_target_hidden_kv）: DSpark 模型层的数据契约改造：_store_block_kv 与 write_target_hidden_kv 在 unified 模式下改写 ring，保持 swa_loc/positions 外部契约不变。
• python/sglang/srt/speculative/dspark_components/dspark_verify.py（模块 投机验证；类别 source；类型 dependency-wiring；符号 commit_hidden, _commit_inject）: 验证提交注入的适配点：commit_hidden 为 unified 模式传递逐 token state_slot，_commit_inject 选择 unified commit-inject 布局。
• python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py（模块 投机工作器；类别 source；类型 dependency-wiring；符号 _forward_prefill）: prefill 注入的数据准备：为每个 token 计算 state_slot 与 final_pos，保证注入器能裁剪 SWA 窗口外的旧 token，避免 ring 槽位竞争。
• python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py（模块 内核层；类别 infra；类型 infrastructure；符号 _scatter_loc_kernel, scatter_bf16_into_unified）: 新增 Triton scatter 内核，是 unified ring 写入的底层算子，被 memory pool 与注入路径共同依赖。
• python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py（模块 内核层；类别 infra；类型 infrastructure；符号 build_unified_commit_inject_layout）: 新增 build_unified_commit_inject_layout：unified 模式下 commit-inject 布局的静态 shape 构建器，CUDA graph 安全。
• python/sglang/srt/speculative/dflash_utils.py（模块 投机采样；类别 source；类型 compatibility）: ROCm 兼容性回退：top-k / top-p 概率重归一化在 HIP 上无法使用 sgl_kernel 算子时退回 torch 实现，保证采样正确性。

关键符号：scatter_bf16_into_unified, _scatter_loc_kernel, set_unified_key_buffer_radix_fused_norm_rope, _unified_inject_loc, inject_target_hidden, inject_ragged, build_unified_commit_inject_layout, init_forward_metadata_target_verify_old, init_forward_metadata_out_graph, make_forward_metadata_from_raw_verify, commit_hidden, _commit_inject, _store_block_kv, write_target_hidden_kv

关键源码片段

python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py

DSpark target-hidden 注入的核心入口：新增 _unified_inject_loc 实现 ring 寻址与 SWA/commit 双重 skip，inject_ragged 分流到 unified 布局构建，是 unified KV 数据契约落地的关键。

# dspark_kv_inject.py —— unified_kv 下 target-hidden 注入的 ring 行地址计算
​
def _unified_inject_loc(
    self,
    *,
    pool,
    positions: torch.Tensor,
    cache_loc_2d: Optional[torch.Tensor],
    commit_lens: Optional[torch.Tensor],
    state_slot: Optional[torch.Tensor],
    final_pos: Optional[torch.Tensor],
) -> torch.Tensor:
    """unified_kv 模式下 target-hidden 注入的 ring 行地址计算。
​
    地址公式为 loc = state_slot * ring + pos % ring，并带两层 skip（-1）规则：
      * SWA 窗口裁剪：每个请求只有最后 win 个 token 落在 ring 内；更早的 token
        会与窗口内 token 共享 ring 槽位（pos % ring 相同）造成互相覆盖，所以
        直接丢弃——长 prefill chunk 场景必须依赖此规则。
      * commit 门控：未通过验收的 verify token（col >= commit_len）丢弃。
    """
    if state_slot is None:
        raise RuntimeError(
            "unified_kv target-hidden injection requires state_slot "
            "(per-token draft req_pool_indices)."
        )
    ring = pool.unified_swa_ring_size
    win = pool.unified_swa_window
    pos = positions.to(torch.int64)
    loc = state_slot.to(torch.int64) * ring + pos % ring
    if final_pos is not None:
        # 只保留窗口内的最新 token，其余置为 -1 由 scatter 跳过
        keep = pos > (final_pos.to(torch.int64) - win)
        loc = torch.where(keep, loc, torch.full_like(loc, -1))
    if commit_lens is not None and cache_loc_2d is not None:
        # 按请求展开列索引，未 commit 的列整体置 -1
        bs, verify_len = cache_loc_2d.shape
        col = torch.arange(verify_len, device=positions.device).view(1, -1)
        committed = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1)
        loc = torch.where(committed, loc, torch.full_like(loc, -1))
    return loc.to(torch.int32)

python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py

新增 Triton scatter 内核，是 unified ring 写入的底层算子，被 memory pool 与注入路径共同依赖。

# unified_kv_kernels/runtime.py —— bf16 K 的 ring scatter 内核与包装函数
​
@triton.jit
def _scatter_loc_kernel(
    kv_ptr, # [T, D] bf16，已 norm + rope 的 draft K
    loc_ptr, # [T] int，unified ring 行号；< 0 表示跳过
    unified_ptr, # [pages, D] bf16，目标 unified KV buffer
    n_rows,
    D: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    if row >= n_rows:
        return
    loc = tl.load(loc_ptr + row).to(tl.int64)
    if loc < 0:
        return # 未 commit 或被窗口裁剪的 token，跳过写入
    offs = tl.arange(0, BLOCK_D)
    mask = offs < D
    vals = tl.load(kv_ptr + row * D + offs, mask=mask, other=0.0)
    tl.store(unified_ptr + loc * D + offs, vals, mask=mask)
​
​
def scatter_bf16_into_unified(
    *,
    kv: torch.Tensor, # [T, head_dim] bf16，已做 norm + RoPE
    loc: torch.Tensor, # [T] int32/int64，ring 行号；< 0 跳过
    unified_kv: torch.Tensor, # [pages, head_dim] bf16
) -> None:
    """把已 norm + rope 的 bf16 K scatter 进 unified_kv[loc]，loc < 0 跳过。
​
    与 store_swa_into_unified 配套：本函数面向已持有预计算 ring 行号的调用方
    （DSpark draft 的 get_unified_swa_loc、或 commit-inject 布局），
    用 loc == -1 表达按行的 commit 掩码，避免单独传 mask。
    """
    n_rows, D = kv.shape
    if n_rows == 0:
        return
    assert kv.is_contiguous() and kv.dtype == unified_kv.dtype
    assert loc.is_contiguous()
    assert unified_kv.is_contiguous()
    _scatter_loc_kernel[(n_rows,)](
        kv,
        loc,
        unified_kv,
        n_rows,
        D=D,
        BLOCK_D=triton.next_power_of_2(D),
        num_warps=8,
    )

评论区精华

核心 review 交锋集中在三处：

1. kkHuang-amd 指出 __init__ 中原地修改 speculative_num_draft_tokens -= 1 是 action-at-a-distance，建议镜像 CUDA 后端的局部变量模式；head 版本已改为显式 is_dspark_draft + target_verify_num_draft_tokens，并在注释中说明 CUDA 约定 gamma+1，同时提醒核对 graph buffer sizing。
2. 1am9trash 提醒新测试应复用 PR#28920 清理后的默认环境变量，避免硬编码导致未来默认值变更时逐文件同步；作者已按建议更新并附上 CI 结果。
3. amd-bot 多次给出同一结论：本 PR 唯一专用测试是 nightly-only 且路径被 SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton 门控，PR CI 完全不验证变更代码，合入前必须由作者在 MI35x 上跑通 nightly 套件；michaelzhang-ai 随后手动调度该套件通过（GSM8K accuracy 0.9515 > 0.92，avg_spec_accept_length 3.730 > 3.0），并指出套件未挂入任何 nightly workflow、已提交 #34147 补排程。

• speculative_num_draft_tokens 原地减一（gamma off-by-one 风险） (design): head 版本已改为显式 is_dspark_draft 标志 + target_verify_num_draft_tokens 局部语义，与 CUDA 模式一致。
• commit_hidden 中的死变量 (style): 已清理，reviewer 最终 APPROVED（LGTM）。
• 测试环境变量应默认化而非硬编码 (design): 作者已按建议更新环境变量并附上 CI 结果。
• nightly-only 覆盖缺口与合入门禁 (testing): michaelzhang-ai 手动调度 nightly 套件（run 31291406552）通过：GSM8K accuracy 0.9515 > 0.92、avg_spec_accept_length 3.730 > 3.0，满足门禁。
• nightly 套件未挂入排程（coverage dormant） (other): 待 #34147 合入后闭环；本 PR 合入时该缺口仍存在。

风险与影响

• 风险：
  • 测试覆盖空白（高风险）：PR CI 不执行任何本 PR 的代码路径；合入时 nightly 排程也缺失（#34147 未合入），覆盖实际处于休眠状态，回归只能靠夜间任务兜底。
  • 门控依赖环境变量：is_unified_kv_triton() 依赖 SGLANG_HACK_FLASHMLA_BACKEND，unified / 非 unified 双路径增加组合数，新布局代码只在 unified 模式激活，普通 CI 无法触达。
  • gamma 语义耦合：target_verify_num_draft_tokens 与 graph runner 读取的 server_args.speculative_num_draft_tokens（未减一）必须始终对齐，kkHuang 也提醒检查 graph buffer sizing。
  • 分支卫生：64 个 commit、约 20+ 次 main merge、多次 revert/re-apply（如 CUTLASS FP8 相关 revert），存在冲突解决引入无关变更的风险。
  • 性能：scatter 内核每行一个 program、num_warps = 8，极端 shape 下可能成为瓶颈；但实测吞吐提升明显，风险可控。
• 影响：
  • 对 AMD 用户：MI35x 等 ROCm 平台首次获得 DSV4 Pro DSpark 投机解码支持，TP8 单并发 TTT 最高 2.93x、TP8DP8 256 并发 unified 后端 1.37x 且 TTFT 改善 2.4x 以上；实现为 HIP-only，不影响 CUDA/MUSA 路径。
  • 对系统：改动横跨 attention backend、mem cache、模型层、speculative 三组件与两个 kernel 模块，"state_slot * ring + pos % ring" 成为 DSV4 DSpark 注入的新数据契约。
  • 对团队流程：暴露并推动了 nightly 套件排程补齐（#34147），也示范了 env-gated HIP 特性的合入门禁做法。
  • 风险标记：仅 nightly 覆盖，PR CI 不验证, nightly 排程缺口待 #34147 补齐, HIP 专属 + 环境变量门控路径, 投机解码核心路径变更, 64 commits 长分支，多次 main 合并

关联脉络

• PR #30261 DeepSeek V4 DSpark support (CUDA baseline): 本 PR 的移植基准：PR body 明确写 Follow #30261，将 CUDA 侧 DSpark 逻辑镜像到 HIP，并沿用其 gamma+1 参数约定。
• PR #34147 Add nightly-8-gpu-mi35x-deepseek-v4-pro-dspark-rocm720 job: michaelzhang-ai 在评论中提交的补排程 PR：把本 PR 的 nightly 套件真正挂入 nightly-test-amd 工作流，否则覆盖处于休眠状态。
• PR #31705 Fix DeepSeek-V4 CUDA sparse-prefill crash: amd-bot CI 报告指出 PR CI 唯一真实失败是既存 DSV4 CUDA sparse-prefill 崩溃，修复在 #31705 进行，与本 PR 无代码交集但影响合入判定。
• PR #31531 [Refactor] Separate ROCm-specific DeepSeek MHA and MLA forward paths: 同属 AMD + DeepSeek 功能线，为 ROCm 侧 DeepSeek 注意力路径的独立性打下基础，本 PR 的 HIP 后端改造与其同一演进脉络。
• PR #33892 [Fix] Speculative decoding crashes with DP-Attention: 投机解码与 DP-Attention 组合的稳定性修复；本 PR 的 TP8DP8 压测数据同样验证了 DSpark 与 DP-Attention 在 AMD 上的组合可用性。