执行摘要
- 一句话:修复 HiCache 启动时 ImportError 崩溃
- 推荐动作:值得精读,但属于中等篇幅即可读完的 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 落地」的案例收藏。
功能与动机
根因是发布顺序缺失: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 流程纪律上。
实现拆解
- 新增可测试、可缓存的算子解析入口 ——
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") 覆盖「符号在但算子未注册」的情形。
- 平台差异化的降级策略 —— 同一文件
- 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」。
- 调用方改造
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),其余逻辑不动。
- 测试配套 ——
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 下套件整体变红。
- 顺带修复 ——
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(模块 主机内存池;类别 source;类型 dependency-wiring;符号 _resolve_device_accessible_ptr_fn, make_kernel_ptr_table, _is_hip): 主线修复文件:新增 _resolve_device_accessible_ptr_fn() 承担算子存在性检查与平台差异化降级,并改造 make_kernel_ptr_table 的指针表构建路径,直接决定 HiCache registered host pool 能否在 pinned wheel 上启动。
test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py(模块 内核测试;类别 test;类型 test-coverage;符号 test_registered_mmap_kernel_ptr_table_fallback_matches_device_alias, test_registered_mmap_pointer_domains_and_all_layer_transfer): 测试配套:新增 CUDA 回退等价性测试并用 monkeypatch 覆盖回退分支,同时把既有 registered 内存指针测试改为走 resolver 并在缺 op 时跳过,保证 pinned wheel 下套件不会整体变红。
python/sglang/srt/mem_cache/unified_radix_cache.py(模块 前缀缓存;类别 source;类型 core-logic;符号 prefetch_from_storage): 顺带修复:prefetch_from_storage 的 rank_consensus 选择器由 req_id 改为 request.rid,与实参名对齐,属合入过程中被维护者发现的选择器失配。
关键符号:_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
主线修复文件:新增 _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
测试配套:新增 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
顺带修复: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 以及函数体均未改动
评论区精华
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”)。
- AOT 算子与调用方同 PR 合入导致 pinned wheel 版本错配 (design): 采取指南允许的例外:先检查 op 是否存在,保留等价回退等待已发布 wheel 追上,待 pin 更新后回退可移除或保留为版本错配保险。
- CUDA 回退与 ROCm fail-fast 的平台分工是否合理 (design): 平台差异化策略被接受,HaiShaw 以 APPROVED + LGTM 放行;维护者仅做了注释精简(commit “trim hicache ptr fallback comments”)。
- 缺 op 环境下测试该跳过还是失败 (testing): 选择 skip 以保证 pinned wheel 下套件保持绿色,代价是这些断言在旧 wheel 环境下并不真正执行,等 pin 更新后才会实际运行。
- 顺带修正 rank_consensus 选择器(req_id → request.rid) (correctness): 已随本 PR 合入,但没有独立测试、也没有对应讨论记录,属于被裹挟进修复 PR 的小范围正确性改动。
风险与影响
- 风险:
- 回退路径依赖 CUDA UVA 语义:
make_kernel_ptr_table 在缺 op 时用 tensor.data_ptr() 作为 device 可见地址,仅在 UVA 下成立;新增测试在 CUDA 上对此做了断言(aliased == raw),但在非 UVA 或未来引入不同 pinned 内存注册方式时,该假设需要重新验证。
@lru_cache(maxsize=1) 把解析结果固化在进程生命周期内:一旦在缺 op 的 wheel 上启动,即便后续 AOT 构建补齐也不会热生效,必须重启进程;这属于可接受的取舍,但排障时容易被忽略(尤其只有一条 warning)。
- ROCm 是刻意的 fail-fast:缺少 op 时
raise ImportError,把失败从 import 期推后到构建 registered host pool 时,错误信息更明确,但使用 pinned wheel 的 ROCm 集群会直接启动失败——这正是设计意图,不属于 bug,但会带来部署侧的硬性要求(必须从 python/sglang/kernels/aot 构建)。
- 测试覆盖有缺口:新测试只覆盖 CUDA 回退分支,ROCm 的 raise 分支和
torch.ops 二次检查都没有直接测试;此外缺 op 时测试走 pytest.skip,意味着在 pinned wheel 环境下这些断言实际上不会执行,回归保护存在真空期。
unified_radix_cache.py 的 rank_consensus(same_params=["request.rid", ...]) 属夹带修复:若 CacheRequestHandle 上没有 rid 字段或语义不符,多 rank 的 prefetch 共识校验会静默失效,而本 PR 没有为这行改动补测试。
- 回退在 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 固化解析结果
关联脉络
- PR #35233 [AMD] Fix registered HiCache host pointer aliases: 直接上游:该 PR 同时落地 AOT 算子 get_device_accessible_ptr 与调用方 make_kernel_ptr_table(含同一测试文件与 pool_host/common.py),跳过 sglang-kernel 发布顺序,是本 PR 要修的全部原因;正文明确说明回滚它会重新打断 MI355X。
- PR #39318 Scope prefetch cache state to the request attempt: 同样改动 python/sglang/srt/mem_cache/unified_radix_cache.py 与 prefetch/请求命名空间逻辑,很可能是 prefetch_from_storage 实参被称为 request 的来源;本 PR 顺带修正的 rank_consensus 选择器与之属同一功能线(需人工确认因果关系)。
参与讨论