[NIXL] Use int32 array for indices to avoid intermediate conversion (#51952)

Signed-off-by: Ilia Yastrebov <iyastrebov@nvidia.com>

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,
     )

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,

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(

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)

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,
                 )

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)

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)

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)
 

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,