执行摘要

• 一句话：修复 ROCm 融合 KV 和 KDA 路径的同步与池类型问题
• 推荐动作：值得精读的部分在于处理统一池与混合池差异的设计决策，以及 Triton 内核中同步屏障的添加方式。对于维护 ROCm 后端的工程师有直接参考价值。

功能与动机

ROCm 后端在 fused KV 缓存和 KDA 路径上存在数据一致性问题：KDA 内核中跨 warp 重载 tile 缺少同步屏障；融合 KV 缓存设置中未区分统一池和混合池，在非混合池（--disable-hybrid-swa-memory）下访问了不存在的 full_to_swa_index_mapping，导致错误。

实现拆解

1. python/sglang/kernels/ops/attention/fla/chunk_intra.py: 在 chunk_kda_fwd_kernel_inter_solve_fused 内核中，于全局 tile 存储（tl.store）之后、跨 warp 进行 forward substitution 重载之前，插入 tl.debug_barrier()，确保所有 warp 的数据可见性。
2. python/sglang/srt/models/utils.py: 在 create_fused_set_kv_buffer_arg 函数的 ROCm 分支中，先通过 isinstance(token_to_kv_pool, SWAKVPool) 判断池类型，再决定是否获取 full_to_swa_index_mapping；仅当池为混合池且层有滑动窗口时，才设置 slot_mapping_swa，否则设为 None。同时将 FusedSetKVBufferArg 的导入移到文件级，并添加注释说明 CUDA/ROCm 的差异。

关键文件：

• python/sglang/srt/models/utils.py（模块 模型工具；类别 source；类型 data-contract；符号 create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer）: 核心修复文件：重写 ROCm 分支中 fused KV 缓存参数的构建逻辑，区分统一池和混合池，并调整导入方式。
• python/sglang/kernels/ops/attention/fla/chunk_intra.py（模块 注意力内核；类别 infra；类型 infrastructure；符号 chunk_kda_fwd_kernel_inter_solve_fused）: 在内核内添加 tl.debug_barrier()，解决 KDA fused store 后跨 warp 重载的同步问题。

关键符号：create_fused_set_kv_buffer_arg, chunk_kda_fwd_kernel_inter_solve_fused

关键源码片段

python/sglang/srt/models/utils.py

核心修复文件：重写 ROCm 分支中 fused KV 缓存参数的构建逻辑，区分统一池和混合池，并调整导入方式。

def create_fused_set_kv_buffer_arg(
    value: torch.Tensor,
    layer: RadixAttention,
    forward_batch: ForwardBatch,
):
    layer_id = layer.layer_id
    token_to_kv_pool = get_token_to_kv_pool()
​
    k_buffer = token_to_kv_pool.get_key_buffer(layer_id)
    v_buffer = token_to_kv_pool.get_value_buffer(layer_id)
​
    if not _is_hip:
        # CUDA path: only bf16 KV cache, no scaling support.
        assert layer.k_scale is None and layer.v_scale is None, "scale not supported"
        return FusedSetKVBufferArg(
            value=value,
            k_buffer=k_buffer.view(k_buffer.shape[0], -1),
            v_buffer=v_buffer.view(v_buffer.shape[0], -1),
            cache_loc=forward_batch.out_cache_loc,
        )
    else:
        # ROCm path: support bf16/fp16/fp8 KV cache.
        page_size = token_to_kv_pool.page_size
        # A non-hybrid pool has no full->SWA remap: SWA and full layers
        # share one slot space indexed directly by out_cache_loc (as
        # --disable-hybrid-swa-memory gives). Leaving swa_slot_mapping=None
        # makes the fused store write at out_cache_loc, matching CUDA path.
        full_to_swa = (
            token_to_kv_pool.full_to_swa_index_mapping
            if isinstance(token_to_kv_pool, SWAKVPool)
            else None
        )
        slot_mapping_swa = (
            full_to_swa.long()
            if layer.sliding_window_size > 0 and full_to_swa is not None
            else None
        )
        # ... rest of the function uses slot_mapping_swa ...

python/sglang/kernels/ops/attention/fla/chunk_intra.py

在内核内添加 tl.debug_barrier()，解决 KDA fused store 后跨 warp 重载的同步问题。

@triton.jit
def chunk_kda_fwd_kernel_inter_solve_fused(...):
    # ... kernel computation ...
    # Store results to global memory
    tl.store(p_Akkd11, b_Akk_d1.to(Akkd.dtype.element_ty), boundary_check=(0, 1))
    tl.store(p_Akkd22, b_Akk_d2.to(Akkd.dtype.element_ty), boundary_check=(0, 1))
    tl.store(p_Akkd33, b_Akk_d3.to(Akkd.dtype.element_ty), boundary_check=(0, 1))
    # Forward substitution reloads these global tiles across warps below.
    tl.debug_barrier()
    b_Ai00 = b_Akk_d0
    b_Ai11 = b_Akk_d1
    # ... continue with forward substitution ...

评论区精华

Gemini Code Assist 的自动审查未提出具体问题，仅概述了变更内容，并提示该工具即将停用。无人工 review 评论。

• 暂无高价值评论线程

风险与影响

• 风险：风险极低：KDA 屏障仅影响 Triton kernel 内跨 warp 同步，语义安全；池类型检查通过 isinstance 显式判断，避免动态属性访问的隐式错误。未引入新依赖或核心逻辑变更。
• 影响：影响范围窄：仅修复 ROCm 后端的两个特定路径。用户层面，AMD GPU 用户可能在某些配置下（如 --disable-hybrid-swa-memory 或 KDA 注意力）遇到之前的错误，本 PR 修正后不再发生。系统层面，非 ROCm 后端不受影响。
• 风险标记：低风险, 仅影响 ROCm 后端

关联脉络

• PR #31687 [Scheduler] Move the WAR barrier to right after each run_batch launch: 同为同步屏障的修复，但位于调度器层级，讨论 WAR barrier 的正确放置。
• PR #29353 [Scheduler] Add SGLANG_FORCE_COARSE_WAR_BARRIER opt-in for a whole-forward WAR barrier: 引入粗粒度 WAR 屏障环境变量，与 KDA 内核的内联屏障属于同一类问题（跨 warp/stream 同步）。
• PR #31474 Fix KDA prefix caching under mamba extra_buffer and enable it for kimi_linear: 之前对 KDA 后端的修复，本 PR 在此基础上补充了 ROCm 特定同步修复。