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

[NPU] Fix disaggregation metadata buffer bootstrap_room_dtype for npu backend (#19423)

Browse Source
This commit is contained in:
Shangming Cai
2026-02-26 21:10:50 +08:00
committed by GitHub
parent 86eb80007e
commit 97f1fa5e6b
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/utils.py
+4 -1
View File
@@ -98,10 +98,13 @@ class MetadataBuffers:
custom_mem_pool: torch.cuda.MemPool = None, custom_mem_pool: torch.cuda.MemPool = None,
): ):
self.custom_mem_pool = custom_mem_pool self.custom_mem_pool = custom_mem_pool
bootstrap_room_dtype = torch.uint64
device = "cpu" device = "cpu"
if is_npu(): if is_npu():
# For ascend backend, output tokens are placed in the NPU and will be transferred by D2D channel. # For ascend backend, output tokens are placed in the NPU and will be transferred by D2D channel.
device = "npu" device = "npu"
# TODO: Fix me when npu backend supports torch.uint64
bootstrap_room_dtype = torch.int64
elif self.custom_mem_pool: elif self.custom_mem_pool:
# TODO(shangming): Fix me (use 'cuda') when nvlink_transport of Mooncake is bug-free # TODO(shangming): Fix me (use 'cuda') when nvlink_transport of Mooncake is bug-free
device = "cpu" device = "cpu"
@@ -144,7 +147,7 @@ class MetadataBuffers:
) )
# Request validation: store bootstrap_room to detect metadata corruption # Request validation: store bootstrap_room to detect metadata corruption
self.bootstrap_room = torch.zeros( self.bootstrap_room = torch.zeros(
(size, 8), dtype=torch.uint64, device=device (size, 8), dtype=bootstrap_room_dtype, device=device
) )
def get_buf_infos(self): def get_buf_infos(self):