Prhub

#39516 [Fix] HiCache startup ImportError on the pinned kernel wheel

原始 PR 作者 mmangkad 合并时间 2026-09-15 13:31 文件变更 3 提交数 8 评论 1 代码增减 +83 / -5

执行摘要

修复 HiCache 启动时 ImportError 崩溃

根因是发布顺序缺失:PR body 明确写到 #35233 把 AOT op 和调用方放在一个 PR 里落地,跳过了贡献指南 Update sglang-kernel 一节要求的顺序。PR CI 构建 PR-local wheel,所以当时是绿的;而 main 与定时 CI 跑的是 pinned sglang-kernel==0.4.7,那里没有这个符号,于是「any server that builds a registered host pool dies at startup on ImportError: cannot import name 'get_device_accessible_ptr'. That was 32 jobs in the last scheduled run.」作者同时排除了直接回滚:ROCm 从本仓库构建 kernel、已经在跑这个修复,回滚会重新打断 MI355X,而且这个算子终究要发布。因此采取指南允许的例外——先做存在性检查,保留等价回退,直到已发布 wheel 追上来。Issue 评论中 HaiShaw 对 AMD 同学补充了「FYI. Let's do more JIT, and be careful with AOT process.」,把这次事故归因到 AOT 流程纪律上。

值得精读,但属于中等篇幅即可读完的 PR(3 文件、+83/-5)。重点看三处:(1) _resolve_device_accessible_ptr_fn 的双重存在性检查(import + torch.ops)与 lru_cache(maxsize=1) 的收敛设计,这是跨仓库发布顺序错配时的通用处置模式;(2) HIP 上 fail-fast、CUDA 上 fail-soft 的平台分工,背后的依据是 UVA 下 host/device 地址别名关系;(3) 测试如何用 alloc_with_host_register + monkeypatch resolver 直接验证回退等价性,以及为什么在缺 op 时选择 skip 而不是 fail。阅读时请额外确认夹带的 unified_radix_cache.py 选择器改动(req_id → request.rid)在本仓库当前签名下确实必要,这行改动没有配套测试。若你是部署/CI 维护者,可作为「AOT op 与调用方必须分两个 PR 落地」的案例收藏。

讨论亮点

Review 侧只有一条 LGTM,实质论证集中在 PR body 与 commit 记录里:

HaiShaw(APPROVED):“LGTM”
HaiShaw(Issue 评论,@AMD-yanfeiwang @amd-danli103 @1am9trash @zijiecode):“FYI. Let's do more JIT, and be careful with AOT process.”

PR body 主动交代了取舍:“That was 32 jobs in the last scheduled run. Reverting re-breaks MI355X ... So this takes the exception that section allows: check the op is there, keep an equivalent fallback until a released wheel catches up.” 并明确了 CUDA 回退等价(UVA 即 device alias)、ROCm 必须 raise 的平台分工。

commits 显示维护者 hnyls2002 在合入前做了最后打磨(“trim hicache ptr fallback comments”),并亲自提交 “fix prefetch consensus selector”,把 req_id 选择器改成 request.rid。值得注意的是:没有任何逐行 review 评论,也没有人讨论“长期停留在回退分支会不会掩盖 wheel 未更新”这类 observability 问题,回退的去留被留给了将来的 pin 更新(“the fallback can go or stay as skew insurance”)。

实现拆解

  1. 新增可测试、可缓存的算子解析入口 —— python/sglang/srt/mem_cache/pool_host/common.py
    - 新增 from functools import lru_cache、from sglang.srt.utils import is_hip,并在模块层缓存 _is_hip = is_hip()。
    - 新增 _resolve_device_accessible_ptr_fn() 并加 @lru_cache(maxsize=1):进程内只解析一次,把「op 是否存在」的判断从调用点内联 import 收敛成单一入口,便于测试 monkeypatch,也便于将来 pin 更新后一处收口。
    - 双重存在性检查:from sgl_kernel.kvcacheio import get_device_accessible_ptr 捕获 ImportError 置 None;import 成功后再用 hasattr(torch.ops.sgl_kernel, "get_device_accessible_ptr") 覆盖「符号在但算子未注册」的情形。
  2. 平台差异化的降级策略 —— 同一文件
    - HIP(_is_hip 为真):直接 raise ImportError,错误信息给出重建指引 python/sglang/kernels/aot (setup_rocm.py)。因为 ROCm 上 registered host memory 携带独立的 device 地址,静默回退会写错指针,缺 op 即意味着构建是坏的。
    - CUDA:只 logger.warning 一次并返回 None,让调用方走原始地址。因为 UVA 下 host 地址本身就是 device alias,「nothing changes past one warning」。
  3. 调用方改造 make_kernel_ptr_table —— 同一文件
    - 原实现在 host_memory_registered and device.type == "cuda" 分支里内联 from sgl_kernel.kvcacheio import ...;现在先调用 _resolve_device_accessible_ptr_fn(),返回 None 就落到 tensor.data_ptr() 分支。
    - 效果:缺少 op 的 CUDA wheel 上,行为回到 #35233 之前的语义,不会在启动/建表阶段抛异常;地址计算仍是 torch.tensor(pointers, dtype=torch.uint64, device=device),其余逻辑不动。
  4. 测试配套 —— test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py
    - 新增 test_registered_mmap_kernel_ptr_table_fallback_matches_device_alias:HIP 上 skipif;resolver 返回 None 时 pytest.skip;否则用 alloc_with_host_register((PAGE_SIZE * 4, 128), torch.bfloat16, "cpu", True, HostTensorAllocator()) 建真实 pinned 内存,比较别名表与 monkeypatch 掉 resolver 后的 raw 表,finally 中 _cuda_host_unregister 释放。
    - 已有 test_registered_mmap_pointer_domains_and_all_layer_transfer 由「直接 import」改为走 _resolve_device_accessible_ptr_fn(),拿不到就 skip 并提示从 AOT 构建,避免 pinned wheel 下套件整体变红。
  5. 顺带修复 —— python/sglang/srt/mem_cache/unified_radix_cache.py
    - prefetch_from_storage 的 @rank_consensus(same_params=["req_id", "len(new_input_tokens)"]) 改为 same_params=["request.rid", "len(new_input_tokens)"],与实参 request: CacheRequestHandle 对齐;该改动来自维护者 hnyls2002 的 commit「fix prefetch consensus selector」,说明它是合入过程中被发现并顺手修掉的选择器失配,而不是本 PR 的主线目标。
文件 类别 关键符号 作用
python/sglang/srt/mem_cache/pool_host/common.py source _resolve_device_accessible_ptr_fn、make_kernel_ptr_table 主线修复:解析 + 平台差异化降级
test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py test test_registered_mmap_kernel_ptr_table_fallback_matches_device_alias 校验 CUDA 回退与别名表等价,并让原测试在缺 op 时跳过
python/sglang/srt/mem_cache/unified_radix_cache.py source prefetch_from_storage 顺带修正 rank_consensus 选择器字段
文件 模块 状态 重要度
python/sglang/srt/mem_cache/pool_host/common.py 主机内存池 modified 7.19
test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py 内核测试 modified 5.57
python/sglang/srt/mem_cache/unified_radix_cache.py 前缀缓存 modified 4.52

关键符号

_resolve_device_accessible_ptr_fn make_kernel_ptr_table test_registered_mmap_kernel_ptr_table_fallback_matches_device_alias test_registered_mmap_pointer_domains_and_all_layer_transfer prefetch_from_storage

关键源码片段

python/sglang/srt/mem_cache/pool_host/common.py dependency-wiring

主线修复文件:新增 _resolve_device_accessible_ptr_fn() 承担算子存在性检查与平台差异化降级,并改造 make_kernel_ptr_table 的指针表构建路径,直接决定 HiCache registered host pool 能否在 pinned wheel 上启动。

# python/sglang/srt/mem_cache/pool_host/common.py
from functools import lru_cache
​
from sglang.srt.utils import is_hip
​
# 平台判定只在模块加载时做一次,避免在解析函数里重复调用 is_hip()
_is_hip = is_hip()
​
​
@lru_cache(maxsize=1)
def _resolve_device_accessible_ptr_fn():
    """解析 sgl_kernel 的 device-accessible 指针算子;缺失时按平台降级或报错。
​
    maxsize=1 的缓存确保每个进程只解析一次:wheel 内容在进程内不会变化。
    """
    try:
        # 第一层:pinned wheel 里可能没有这个模块级符号(#35233 新增的 AOT op)
        from sgl_kernel.kvcacheio import get_device_accessible_ptr
    except ImportError:
        get_device_accessible_ptr = None
    else:
        # 第二层:符号存在但 torch.ops 命名空间未注册时,调用依然会失败
        if not hasattr(torch.ops.sgl_kernel, "get_device_accessible_ptr"):
            get_device_accessible_ptr = None
​
    if get_device_accessible_ptr is None:
        # CUDA 的 UVA 让 host 地址本身就是 device 可见地址,回退等价;
        # ROCm 上两个地址确实不同,静默回退会写错指针,因此直接 fail fast。
        if _is_hip:
            raise ImportError(
                "sgl_kernel.kvcacheio.get_device_accessible_ptr is missing from the "
                "installed sglang-kernel. It is required on ROCm, where registered "
                "host memory carries a distinct device address. Rebuild sglang-kernel "
                "from python/sglang/kernels/aot (setup_rocm.py)."
            )
        logger.warning(
            "sgl_kernel.kvcacheio.get_device_accessible_ptr is missing from the "
            "installed sglang-kernel; using raw host addresses for kernel pointer "
            "tables. Build sglang-kernel from python/sglang/kernels/aot to enable it."
        )
    return get_device_accessible_ptr
​
​
def make_kernel_ptr_table(
    tensors: list[torch.Tensor],
    target_device: torch.device | str,
    *,
    host_memory_registered: bool,
) -> torch.Tensor:
    # 只有 registered 的 CUDA host 内存才需要地址转换,其他情况一律用原始地址
    device = torch.device(target_device)
    get_device_accessible_ptr = (
        _resolve_device_accessible_ptr_fn()
        if host_memory_registered and device.type == "cuda"
        else None
    )
    if get_device_accessible_ptr is not None:
        if device.index is None:
            device_index = torch.cuda.current_device()
        else:
            device_index = device.index
        pointers = [
            get_device_accessible_ptr(tensor, device_index) for tensor in tensors
        ]
    else:
        # 回退分支:缺 op 的 CUDA wheel 上,UVA 让 data_ptr() 即为 device 别名
        pointers = [tensor.data_ptr() for tensor in tensors]
    return torch.tensor(
        pointers,
        dtype=torch.uint64,
        device=device,
    )
test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py test-coverage

测试配套:新增 CUDA 回退等价性测试并用 monkeypatch 覆盖回退分支,同时把既有 registered 内存指针测试改为走 resolver 并在缺 op 时跳过,保证 pinned wheel 下套件不会整体变红。

# test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py
@pytest.mark.skipif(
    is_hip(),
    reason="ROCm maps registered host memory at a distinct device address.",
)
def test_registered_mmap_kernel_ptr_table_fallback_matches_device_alias(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """CUDA maps registered host memory at the host address itself;
    ``make_kernel_ptr_table``'s raw-host-address fallback depends on it."""
    # pinned wheel 没有该 op 时别名表本身无法构建,测试跳过而不是失败
    if _resolve_device_accessible_ptr_fn() is None:
        pytest.skip(
            "installed sglang-kernel has no get_device_accessible_ptr; "
            "build it from python/sglang/kernels/aot to run this test"
        )
​
    # 使用真实注册的 pinned host 内存,避免 UVA 别名关系是伪造出来的
    buffer = alloc_with_host_register(
        (PAGE_SIZE * 4, 128),
        torch.bfloat16,
        "cpu",
        True,
        HostTensorAllocator(),
    )
    try:
        aliased = make_kernel_ptr_table([buffer], DEVICE, host_memory_registered=True)
        # 强行让 resolver 返回 None,强制走 raw data_ptr() 回退路径
        monkeypatch.setattr(
            pool_host_common, "_resolve_device_accessible_ptr_fn", lambda: None
        )
        raw = make_kernel_ptr_table([buffer], DEVICE, host_memory_registered=True)
        # 两者不等说明 UVA 回退假设被打破,回退路径会写错指针
        assert torch.equal(aliased, raw)
    finally:
        _cuda_host_unregister(buffer)
​
​
def test_registered_mmap_pointer_domains_and_all_layer_transfer() -> None:
    # 统一走 resolver:缺 op 的 wheel 上跳过,而不是在 import 阶段就 ImportError
    get_device_accessible_ptr = _resolve_device_accessible_ptr_fn()
    if get_device_accessible_ptr is None:
        pytest.skip(
            "installed sglang-kernel has no get_device_accessible_ptr; "
            "build it from python/sglang/kernels/aot to run this test"
        )
python/sglang/srt/mem_cache/unified_radix_cache.py core-logic

顺带修复:prefetch_from_storage 的 rank_consensus 选择器由 req_id 改为 request.rid,与实参名对齐,属合入过程中被维护者发现的选择器失配。

# python/sglang/srt/mem_cache/unified_radix_cache.py
# rank_consensus 的 same_params 选择器此前写的是 "req_id",与实参名 request 不匹配,
# 会让多 rank 的 prefetch 共识校验取不到值;此处顺带改成 request.rid。
@rank_consensus(same_params=["request.rid", "len(new_input_tokens)"])
def prefetch_from_storage(
    self,
    request: CacheRequestHandle,
    # 其余参数 last_host_node_id / new_input_tokens / last_hash / prefix_keys /
    # matched_prefix_tokens / extra_key / cache_salt 以及函数体均未改动

评论区精华

AOT 算子与调用方同 PR 合入导致 pinned wheel 版本错配 设计

PR body 说明 #35233 跳过了贡献指南 Update sglang-kernel 要求的顺序,PR CI 用 PR-local wheel 所以是绿的,而 main 与定时 CI 用 pinned sglang-kernel==0.4.7,缺 'get_device_accessible_ptr' 导致 32 个 job 启动即崩溃;Issue 评论中 HaiShaw 对 AMD 同学补充「Let's do more JIT, and be careful with AOT process.」

结论:采取指南允许的例外:先检查 op 是否存在,保留等价回退等待已发布 wheel 追上,待 pin 更新后回退可移除或保留为版本错配保险。 · 已解决

CUDA 回退与 ROCm fail-fast 的平台分工是否合理 设计

PR body 论证 CUDA 的 raw host address 在 UVA 下就是 device alias,因此回退除了多一条 warning 外行为不变;ROCm 上两个地址确实不同,缺 op 意味着构建是坏的,所以直接 raise。新增测试在 CUDA 上断言别名表与回退表相等,HIP 上以 skipif 排除。

结论:平台差异化策略被接受,HaiShaw 以 APPROVED + LGTM 放行;维护者仅做了注释精简(commit “trim hicache ptr fallback comments”)。 · 已解决

缺 op 环境下测试该跳过还是失败 测试

test_registered_mmap_pointer_domains_and_all_layer_transfer 由直接 import 改为通过 _resolve_device_accessible_ptr_fn() 取函数,返回 None 即 pytest.skip 并提示从 python/sglang/kernels/aot 构建;新增的回退等价性测试同样先 skip。

结论:选择 skip 以保证 pinned wheel 下套件保持绿色,代价是这些断言在旧 wheel 环境下并不真正执行,等 pin 更新后才会实际运行。 · 已解决

顺带修正 rank_consensus 选择器(req_id → request.rid) 正确性

维护者 hnyls2002 在合入过程中提交 “fix prefetch consensus selector”,把 prefetch_from_storage 的 same_params 从 ["req_id", ...] 改为 ["request.rid", ...],与实参 request: CacheRequestHandle 对齐。

结论:已随本 PR 合入,但没有独立测试、也没有对应讨论记录,属于被裹挟进修复 PR 的小范围正确性改动。 · 已解决

风险与影响

  1. 回退路径依赖 CUDA UVA 语义:make_kernel_ptr_table 在缺 op 时用 tensor.data_ptr() 作为 device 可见地址,仅在 UVA 下成立;新增测试在 CUDA 上对此做了断言(aliased == raw),但在非 UVA 或未来引入不同 pinned 内存注册方式时,该假设需要重新验证。
  2. @lru_cache(maxsize=1) 把解析结果固化在进程生命周期内:一旦在缺 op 的 wheel 上启动,即便后续 AOT 构建补齐也不会热生效,必须重启进程;这属于可接受的取舍,但排障时容易被忽略(尤其只有一条 warning)。
  3. ROCm 是刻意的 fail-fast:缺少 op 时 raise ImportError,把失败从 import 期推后到构建 registered host pool 时,错误信息更明确,但使用 pinned wheel 的 ROCm 集群会直接启动失败——这正是设计意图,不属于 bug,但会带来部署侧的硬性要求(必须从 python/sglang/kernels/aot 构建)。
  4. 测试覆盖有缺口:新测试只覆盖 CUDA 回退分支,ROCm 的 raise 分支和 torch.ops 二次检查都没有直接测试;此外缺 op 时测试走 pytest.skip,意味着在 pinned wheel 环境下这些断言实际上不会执行,回归保护存在真空期。
  5. unified_radix_cache.py 的 rank_consensus(same_params=["request.rid", ...]) 属夹带修复:若 CacheRequestHandle 上没有 rid 字段或语义不符,多 rank 的 prefetch 共识校验会静默失效,而本 PR 没有为这行改动补测试。
  6. 回退在 CUDA 上只打一次 warning,缺少指标/计数暴露;长期跑在回退路径上不会在监控里体现。

用户侧:使用官方 pinned sglang-kernel wheel(0.4.7)的 CUDA 用户,此前只要开启 registered host pool(HiCache page_first + kernel 模式的写回/加载路径)就会在启动时 ImportError 崩溃,现在可正常启动,代价仅是日志里一条 warning;ROCm 用户不受影响,仍要求从源码构建 kernel(否则明确报错而不是算错指针)。系统与 CI 侧:恢复 main 与定时 CI 的绿色(此前 32 个 job 因该 ImportError 失败),并让 PR CI 与 main CI 的行为差异(PR-local wheel vs pinned wheel)显式化。团队流程侧:本 PR 是「AOT 算子需先发布 wheel、再合入调用方」这条纪律缺口的直接产物,Issue 评论中 AMD 相关同学被点名提醒“do more JIT, be careful with AOT process”,预计会推动后续同类改动优先走 JIT 路径或补上 fallback。影响范围集中在 HiCache host 内存池的建表入口,不触及正常 AOT wheel 的服务路径。

启动路径降级 依赖 CUDA UVA 平台假设 ROCm 缺 op 直接 raise 缺 op 时测试跳过 夹带 rank_consensus 修复 lru_cache 固化解析结果

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论