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

This commit is contained in:
Shangming Cai
2026-02-26 21:10:50 +08:00
committed by GitHub
parent 86eb80007e
commit 97f1fa5e6b

View File

@@ -98,10 +98,13 @@ class MetadataBuffers:
custom_mem_pool: torch.cuda.MemPool = None,
):
self.custom_mem_pool = custom_mem_pool
bootstrap_room_dtype = torch.uint64
device = "cpu"
if is_npu():
# For ascend backend, output tokens are placed in the NPU and will be transferred by D2D channel.
device = "npu"
# TODO: Fix me when npu backend supports torch.uint64
bootstrap_room_dtype = torch.int64
elif self.custom_mem_pool:
# TODO(shangming): Fix me (use 'cuda') when nvlink_transport of Mooncake is bug-free
device = "cpu"
@@ -144,7 +147,7 @@ class MetadataBuffers:
)
# Request validation: store bootstrap_room to detect metadata corruption
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):