diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index d0c4d4cb04..1d631e7782 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -113,7 +113,7 @@ def make_job( return TransferJob( job_id=job_id, keys=keys, - block_ids=np.array(block_ids, dtype=np.int64), + block_ids=np.array(block_ids, dtype=np.int32), is_promotion=False, req_context=_CTX, ) diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 14efd16177..77362e593c 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -527,7 +527,7 @@ class NixlEplbCommunicator(EplbCommunicator): remote_desc, ) - indices = list(range(len(local_descs))) + indices = np.arange(len(local_descs), dtype=np.int32) xfer_handle = self._nixl_wrapper.make_prepped_xfer( "READ", local_handle, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 19b35b8253..46d0ec3092 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -141,9 +141,11 @@ class NixlBaseConnectorWorker: # read across all regions, same for [3], but group0-group1 blocks will # always differ (different areas). Therefore we can just flatten the # block_ids and compute the descs ids for all groups at once. - block_arr = np.concatenate(block_ids)[None, :] - region_ids = np.arange(self.num_regions)[:, None] - return (region_ids * num_blocks + block_arr).flatten() + block_arr = np.concatenate( + [np.asarray(g, dtype=np.int32) for g in block_ids] + )[None, :] + region_ids = np.arange(self.num_regions, dtype=np.int32)[:, None] + return (region_ids * num_blocks + block_arr).ravel() # Compute desc ids per group using the right stride: FA descs have # num_blocks entries per region (kernel granularity, expanded by @@ -153,18 +155,18 @@ class NixlBaseConnectorWorker: logical_blocks = dst_num_blocks // physical_blocks_per_logical all_descs: list[np.ndarray] = [] for i, group in enumerate(block_ids): - group_arr = np.asarray(group) + group_arr = np.asarray(group, dtype=np.int32) spec_type = self._group_spec_types[i] if _is_attention_spec(spec_type): # A scratch cache lives only in its own regions; every other # attention group spans all of them. fa_region_ids = ( - np.asarray(self._scratch_region_indices, dtype=np.int64) + np.asarray(self._scratch_region_indices, dtype=np.int32) if spec_type is CircularBufferSpec - else np.arange(self.num_regions) + else np.arange(self.num_regions, dtype=np.int32) )[:, None] all_descs.append( - (fa_region_ids * num_blocks + group_arr[None, :]).flatten() + (fa_region_ids * num_blocks + group_arr[None, :]).ravel() ) elif _is_ssm_spec(spec_type): # NOTE (NickLucche) SSM and Attention block regions can @@ -175,16 +177,16 @@ class NixlBaseConnectorWorker: # P and D can have different num_blocks (and thus # different FA desc counts). ssm_region_ids = ( - np.arange(num_ssm_regions, num_ssm_regions + 1) + np.arange(num_ssm_regions, num_ssm_regions + 1, dtype=np.int32) if i == self._ple_group_index - else np.arange(num_ssm_regions) + else np.arange(num_ssm_regions, dtype=np.int32) )[:, None] all_descs.append( ( ssm_region_ids * logical_blocks + group_arr[None, :] + num_fa_descs - ).flatten() + ).ravel() ) else: raise ValueError( diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index dede311446..5ad89eef57 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -403,7 +403,7 @@ class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC): """ def __init__(self, block_ids: list[int]): - self.block_ids = np.array(block_ids, dtype=np.int64) + self.block_ids = np.array(block_ids, dtype=np.int32) def __repr__(self) -> str: return repr(self.block_ids) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 8ba23cd83c..d2499859f6 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -483,7 +483,7 @@ class TieringOffloadingManager(OffloadingManager): job_metadata = TransferJob( job_id=job_id, keys=entry.keys, - block_ids=np.array(entry.block_ids, dtype=np.int64), + block_ids=np.array(entry.block_ids, dtype=np.int32), is_promotion=True, req_context=entry.req_context, ) diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index fa055c2f0a..ed3c552459 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -7,6 +7,8 @@ import time from collections.abc import Iterable from typing import TYPE_CHECKING, ClassVar, NamedTuple +import numpy as np + from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent from vllm.distributed.nixl_utils import nixl_agent_config from vllm.logger import init_logger @@ -217,12 +219,13 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): def _submit_transfer( self, job_id: int, - block_ids: Iterable[int], + block_ids: np.ndarray, obj_keys: Iterable[str], op: str, ) -> None: """Submit an async transfer. op is 'WRITE' (store) or 'READ' (load).""" - block_ids_list = [int(bid) for bid in block_ids] + # NIXL takes indices as int32; a wider dtype silently costs a conversion. + assert block_ids.dtype == np.int32 # The OBJ backend maps devId -> obj_key. All descriptors must have # unique devIds or later registrations overwrite earlier ones. nixl_files = [ @@ -247,9 +250,9 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): xfer_handle = self._agent.make_prepped_xfer( op, self._dram_prepped_handle, - block_ids_list, + block_ids, obj_handle, - list(range(len(nixl_files))), + np.arange(len(nixl_files), dtype=np.int32), ) if not xfer_handle: logger.warning("make_prepped_xfer failed for job %d", job_id) diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index 8d53b923be..bf3adf4b60 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -10,6 +10,8 @@ import itertools from collections.abc import Iterable from typing import Any, NamedTuple +import numpy as np + from vllm.distributed.nixl_utils import NixlWrapper as _NixlAgent from vllm.distributed.nixl_utils import nixl_agent_config as _NixlAgentConfig from vllm.logger import init_logger @@ -178,9 +180,9 @@ class NixlTransport(DataTransport): handle = self._agent.make_prepped_xfer( "WRITE", self._local_dlist, - local_idxs, + np.asarray(local_idxs, dtype=np.int32), remote_dlist, - remote_idxs, + np.asarray(remote_idxs, dtype=np.int32), ) self._agent.transfer(handle) transfer_id = next(self._next_id) diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 9d8ab38c65..70ed3da456 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -417,7 +417,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): def submit_store(self, job_metadata: TransferJob) -> None: job_id = job_metadata.job_id keys = list(job_metadata.keys) - block_ids = job_metadata.block_ids + block_ids = job_metadata.block_ids.tolist() assert len(keys) == len(block_ids) diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 8558acdf91..7d24f3773f 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -593,7 +593,7 @@ class ServerRole: self.add_stored_blocks( lookup.kv_request_id, list(meta.keys), - list(meta.block_ids), + meta.block_ids.tolist(), meta.job_id, round_seq=lookup.round_seq, from_lookup=True,