执行摘要

• 一句话：启用ROCm HiCache JIT分阶段写回并修复索引崩溃
• 推荐动作：该PR值得精读，尤其是JIT kernel平台适配和向量化技巧。对于维护多后端的团队，此PR展示了如何优雅处理PTX平台差异。建议关注后续是否回退#28473，以及是否需要在更多后端（如NPU）上推行类似模式。

功能与动机

ROCm上 page_first + kernel HiCache写回在第一次预填充时崩溃：RuntimeError: Destination indices must be a CUDA tensor。根因是HiCacheController.start_writing()将host_indices保留在CPU，但非JIT kernel需要设备索引。之前的workaround (#28473)强制使用layer_first布局，但关闭了JIT分阶段写回，导致ROCm性能降级并偏离CUDA路径。此PR从根源修复，使ROCm使用与CUDA相同的JIT写回kernel，保留带宽和向量化非临时load/store，同时消除维护两套路径的成本。

实现拆解

1. 扩展can_use_jit到HIP：在memory_pool_host.py和pool_host/mha.py中，将can_use_jit和can_use_write_back_jit的条件从_is_cuda扩展为_is_cuda or _is_hip，并添加详细注释，使得ROCm也能构建和执行JIT HiCache kernel。
2. ROCm非临时load/store适配：在hicache.cuh中，通过#ifdef USE_ROCM保护PTX指令，为load_nc/store_nc实现基于__builtin_nontemporal_load/store的ROCm版本。使用Clang ext_vector_type和__builtin_bit_cast确保向量化指令（global_{load,store}_dwordx{2,4}），避免编译器退化为标量操作导致带宽下降。
3. 设备匹配器统一：在hicache.cuh和staged_write_back.cuh中，将TensorMatcher的.with_device从硬编码的kDLCUDA/kDLCUDAHost替换为通过条件宏定义的kDLGPU/kDLGPUHost，兼容CUDA和ROCm。同样在utils.cuh中添加宏定义。
4. cache_controller条件修复：在start_writing()中，除了检查layout == "page_first"和io_backend == "kernel"外，额外检查can_use_write_back_jit属性（通过getattr安全访问）。仅在JIT kernel可用时保留host indices在CPU，否则移到设备，避免非JIT路径崩溃。
5. 新增单元测试：创建test_hicache_page_first_write_back.py，测试MHA和MLA场景的page_first写回路径，同时在CUDA和AMD CI中注册，验证JIT kernel在ROCm上的构建和执行。
6. 回退#28473：在PR描述中明确要求合并后回退#28473，移除对ROCm强制layer_first的fallback。

关键文件：

• python/sglang/srt/mem_cache/memory_pool_host.py（模块 缓存层；类别 source；类型 core-logic；符号 MLATokenToKVPoolHost.init, MLATokenToKVPoolHost._init_write_back_staging_buffers）: 核心开关：扩展can_use_jit到HIP，使ROCm使用JIT HiCache路径
• python/sglang/srt/mem_cache/pool_host/mha.py（模块 缓存层；类别 source；类型 core-logic；符号 MHATokenToKVPoolHost.init, MHATokenToKVPoolHost._init_write_back_staging_buffers）: MHA池对应启用JIT：与MLA池相同的HIP扩展
• python/sglang/srt/managers/cache_controller.py（模块 控制器；类别 source；类型 entrypoint；符号 HiCacheController.start_writing）: 修复崩溃的关键：在start_writing中检查can_use_write_back_jit，避免在JIT不可用时保留CPU索引
• python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh（模块 JIT内核；类别 other；类型 core-logic；符号 load_nc, store_nc, HiCacheKernel）: 平台适配核心：添加ROCm非 temporal load/store，使用Clang向量类型确保向量化，TensorMatcher兼容ROCm
• test/registered/jit/test_hicache_page_first_write_back.py（模块 测试套件；类别 test；类型 test-coverage；符号 _token_indices_for_pages, _pinned_host_pool, _fill_with_offset, _assert_pages_equal）: 新增单元测试覆盖page_first JIT写回路径，同时在CUDA和AMD CI中注册

关键符号：MLATokenToKVPoolHost.init, MLATokenToKVPoolHost._init_write_back_staging_buffers, MHATokenToKVPoolHost.init, MHATokenToKVPoolHost._init_write_back_staging_buffers, HiCacheController.start_writing, load_nc (uint1/uint2/uint4), store_nc (uint1/uint2/uint4), HiCacheStagedWriteBackKernel

关键源码片段

python/sglang/srt/mem_cache/memory_pool_host.py

核心开关：扩展can_use_jit到HIP，使ROCm使用JIT HiCache路径

# sglang/srt/mem_cache/memory_pool_host.py (partial)
​
class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
    def __init__(self, ...):
        # ... (super init omitted)
​
        # 关键变更：将 _is_cuda 扩展为 _is_cuda or _is_hip
        # 因为 JIT HiCache 内核也通过 hipcc 编译，并且 hicache.cuh 中
        # 的 PTX 辅助函数已被 USE_ROCM 保护，ROCm 使用非 temporal 内建函数作为替代
        self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel(
            element_size=self.kv_cache_dim * self.dtype.itemsize
        )
        # ... (data_refs, data_ptrs init omitted)
        self._init_write_back_staging_buffers()
​
    def _init_write_back_staging_buffers(self):
        self.staging_page_capacity = 0
        self.staging_token_capacity = 0
        self.staging_buffer = None
        self.can_use_write_back_jit = False
        if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
            return
​
        # 同样，ROCm 上也启用 staged write-back JIT 内核
        self.can_use_write_back_jit = (
            _is_cuda or _is_hip
        ) and can_use_write_back_jit_kernel(
            element_size=self.kv_cache_dim * self.dtype.itemsize,
        )
        if not self.can_use_write_back_jit:
            return
​
        # 分配 staging buffer（device memory）
        self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
        self.staging_token_capacity = self.staging_page_capacity * self.page_size
        self.staging_buffer = torch.empty(
            (self.staging_token_capacity, self.layer_num, 1, self.kv_cache_dim),
            dtype=self.dtype,
            device=self.device_pool.device,
        )

python/sglang/srt/managers/cache_controller.py

修复崩溃的关键：在start_writing中检查can_use_write_back_jit，避免在JIT不可用时保留CPU索引

# sglang/srt/managers/cache_controller.py (partial)
​
def start_writing(self) -> None:
    if len(self.write_queue) == 0:
        return
​
    op = CacheOperation.merge_ops(self.write_queue)
​
    # Kernel write-back 仅在 page_first 布局且分阶段 JIT 写回内核可用时
    # 才将 host indices 保留在 CPU（JIT 内核会将这些索引通过设备内存传递，
    # 并且接受 CPU 目标索引）。否则，普通的传输内核要求目标索引在设备上，
    # 因此必须将索引移动到设备。如果没有 can_use_write_back_jit 检查，
    # 在后端 JIT 内核不可用的情况下会崩溃，报错 "Destination indices must be a CUDA tensor"。
    if (
        self.io_backend == "kernel"
        and self.mem_pool_host.layout == "page_first"
        and getattr(self.mem_pool_host, "can_use_write_back_jit", False)
    ):
        host_indices, device_indices = op.host_indices, op.device_indices
    else:
        host_indices, device_indices = self.move_indices(
            op.host_indices, op.device_indices
        )
    self.write_queue.clear()
    # ...（stream 和实际写回调用省略）

python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh

平台适配核心：添加ROCm非 temporal load/store，使用Clang向量类型确保向量化，TensorMatcher兼容ROCm

// python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh (partial)
​
// NVIDIA 通过 PTX 指令 ld.global.L1::no_allocate 提供显式的一级缓存旁路。
// ROCm 没有等效 PTX，但非 temporal（流式）load/store 表达了相同的意图，
// 适用于 HiCache 写回这类不应污染缓存的一次性流量。
// 将 PTX 保护在 USE_ROCM 之后，使 JIT 模块也能通过 hipcc 编译。
​
#ifdef USE_ROCM
// 使用 Clang 原生向量类型，使得单个 __builtin_nontemporal_{load,store}
// 映射到一条向量化的 global_{load,store}_dwordx{2,4} 指令。
// 如果发出 N 个独立的 32 位非 temporal 操作，编译器可能不会合并，
// 甚至可能丢弃 non-temporal 提示，从而降低 HiCache 带宽。
// uint2/uint4 带有 8B/16B 对齐，因此指针 reinterpret_cast 保持正确对齐。
typedef uint32_t native_uint2 __attribute__((ext_vector_type(2)));
typedef uint32_t native_uint4 __attribute__((ext_vector_type(4)));
#endif
​
SGL_DEVICE uint1 load_nc(const uint1* __restrict__ src) {
#ifndef USE_ROCM
    uint32_t tmp;
    asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
    return uint1{tmp};
#else
    return uint1{__builtin_nontemporal_load(&src->x)};
#endif
}
​
SGL_DEVICE uint2 load_nc(const uint2* __restrict__ src) {
#ifndef USE_ROCM
    uint32_t tmp0, tmp1;
    asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
    return uint2{tmp0, tmp1};
#else
    native_uint2 tmp = __builtin_nontemporal_load(
        reinterpret_cast<const native_uint2*>(src));
    return __builtin_bit_cast(uint2, tmp);
#endif
}
​
SGL_DEVICE uint4 load_nc(const uint4* __restrict__ src) {
#ifndef USE_ROCM
    uint32_t tmp0, tmp1, tmp2, tmp3;
    asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src));
    return uint4{tmp0, tmp1, tmp2, tmp3};
#else
    native_uint4 tmp = __builtin_nontemporal_load(
        reinterpret_cast<const native_uint4*>(src));
    return __builtin_bit_cast(uint4, tmp);
#endif
}
​
SGL_DEVICE void store_nc(uint1* __restrict__ dst, const uint1& value) {
#ifndef USE_ROCM
    uint32_t tmp = value.x;
    asm volatile("st.global.L1::no_allocate.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
#else
    __builtin_nontemporal_store(value.x, &dst->x);
#endif
}
​
SGL_DEVICE void store_nc(uint2* __restrict__ dst, const uint2& value) {
#ifndef USE_ROCM
    uint32_t tmp0 = value.x; uint32_t tmp1 = value.y;
    asm volatile("st.global.L1::no_allocate.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
#else
    __builtin_nontemporal_store(
        __builtin_bit_cast(native_uint2, value),
        reinterpret_cast<native_uint2*>(dst));
#endif
}
​
SGL_DEVICE void store_nc(uint4* __restrict__ dst, const uint4& value) {
#ifndef USE_ROCM
    uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; uint32_t tmp2 = value.z; uint32_t tmp3 = value.w;
    asm volatile("st.global.L1::no_allocate.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
#else
    __builtin_nontemporal_store(
        __builtin_bit_cast(native_uint4, value),
        reinterpret_cast<native_uint4*>(dst));
#endif
}

评论区精华

主要review讨论集中在三点：

• 向量化非临时load/store：gemini-code-assist建议使用Clang ext_vector_type和__builtin_bit_cast来确保单个向量化指令，避免多个32位操作导致编译器不合并。作者采纳并修改。
• HostPoolGroup兼容：gemini提到当mem_pool_host是HostPoolGroup时可能没有can_use_jit属性。作者最终使用了getattr(self.mem_pool_host, "can_use_write_back_jit", False)来安全降级，但此问题未完全解决，仍存在静默降级风险。

• 使用kDLGPU宏：HaiShaw建议用统一的kDLGPU替代显式枚举kDLCUDA, kDLROCM，作者实施并修改了所有相关匹配器，减少了平台重复。

  • ROCm non-temporal load/store向量化 (performance): 作者采纳，在hicache.cuh中实现向量化版本
  • HostPoolGroup属性兼容 (design): 作者使用getattr + 默认False部分解决，但HostPoolGroup场景仍未完全覆盖
  • 使用kDLGPU宏统一设备类型 (design): 作者实施并在所有TensorMatcher中替换

风险与影响

• 风险：
  1. ROCm JIT kernel性能风险：非临时load/store的向量化质量依赖Clang编译器，可能仍有退化，需通过benchmark验证。
  2. HostPoolGroup覆盖不足：如果mem_pool_host是分组池，can_use_write_back_jit属性可能不存在，getattr返回False，导致JIT路径被静默关闭，用户可能未察觉。需要后续检查或统一添加属性。
  3. #28473回退依赖：若未按计划回退，ROCm仍会被强制layer_first，PR效果被抵消。需手动跟踪。
  4. 集成测试缺失：AMD CI仅运行单元测试，集成场景（PD分离、大模型）仅在nightly中覆盖，可能遗漏边界条件。
     - 影响：对用户：ROCm（AMD）用户现在可以使用page_first布局+ kernel写回，获得与CUDA相似的性能和稳定性。之前受影响的用户（被迫使用layer_first）可恢复正常配置。对系统：ROCm和CUDA在HiCache写回上保持一致，降低维护成本。对团队：需确保#28473被回退，否则布局强制fallback仍会生效。新测试在AMD CI中运行，增强信心。
     - 风险标记：ROCm JIT路径无集成测试, HostPoolGroup可能静默降级, #28473需手动回退

关联脉络

• PR #28473 [AMD] Fall back to layer_first layout for kernel write-back on ROCm: 先前的workaround强制ROCm使用layer_first，此PR修复根本原因后需要回退该PR
• PR #21631 [HiCache & JIT Kernel] Refactoring HiCache Write-Back Kernel: 该PR引入了page_first + JIT写回路径，使ROCm默认page_first后暴露问题