minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[HiCache]: Check return code of cudaHostRegister (#26301)

Browse Source
This commit is contained in:
Chao Shi
2026-05-26 17:44:15 +08:00
committed by GitHub
parent d25a220fdb
commit 48f3264807
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool_host.py
+3 -1
View File
@@ -187,9 +187,11 @@ def alloc_with_host_register(
""" """
buffer = allocator.allocate(dims, dtype=dtype, device=device) buffer = allocator.allocate(dims, dtype=dtype, device=device)
if pin_memory: if pin_memory:
torch.cuda.cudart().cudaHostRegister( ret = torch.cuda.cudart().cudaHostRegister(
buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0 buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0
) )
if ret != 0:
raise RuntimeError(f"cudaHostRegister failed with error code {ret}")
return buffer return buffer