[NIXL] Use int32 array for indices to avoid intermediate conversion (#51952)
Signed-off-by: Ilia Yastrebov <iyastrebov@nvidia.com>
Ilia Yastrebov committed
Sep 2, 2026 at 16:30 UTC
41848caa63b131f931618055d3ece0b487ce8edb
9 files changed
+29
-22
tests/v1/kv_offload/tiering/test_obj_tier.py
+1
-1
@@ -113,7 +113,7 @@ def make_job(
113
return TransferJob(
114
job_id=job_id,
115
keys=keys,
116
- block_ids=np.array(block_ids, dtype=np.int64),
116
+ block_ids=np.array(block_ids, dtype=np.int32),
117
is_promotion=False,
118
req_context=_CTX,
119
)
vllm/distributed/eplb/eplb_communicator.py
+1
-1
@@ -527,7 +527,7 @@ class NixlEplbCommunicator(EplbCommunicator):
527
remote_desc,
528
)
529
530
- indices = list(range(len(local_descs)))
530
+ indices = np.arange(len(local_descs), dtype=np.int32)
531
xfer_handle = self._nixl_wrapper.make_prepped_xfer(
532
"READ",
533
local_handle,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+12
-10
@@ -141,9 +141,11 @@ class NixlBaseConnectorWorker:
141
# read across all regions, same for [3], but group0-group1 blocks will
142
# always differ (different areas). Therefore we can just flatten the
143
# block_ids and compute the descs ids for all groups at once.
144
- block_arr = np.concatenate(block_ids)[None, :]
145
- region_ids = np.arange(self.num_regions)[:, None]
146
- return (region_ids * num_blocks + block_arr).flatten()
144
+ block_arr = np.concatenate(
145
+ [np.asarray(g, dtype=np.int32) for g in block_ids]
146
+ )[None, :]
147
+ region_ids = np.arange(self.num_regions, dtype=np.int32)[:, None]
148
+ return (region_ids * num_blocks + block_arr).ravel()
149
150
# Compute desc ids per group using the right stride: FA descs have
151
# num_blocks entries per region (kernel granularity, expanded by
@@ -153,18 +155,18 @@ class NixlBaseConnectorWorker:
155
logical_blocks = dst_num_blocks // physical_blocks_per_logical
156
all_descs: list[np.ndarray] = []
157
for i, group in enumerate(block_ids):
156
- group_arr = np.asarray(group)
158
+ group_arr = np.asarray(group, dtype=np.int32)
159
spec_type = self._group_spec_types[i]
160
if _is_attention_spec(spec_type):
161
# A scratch cache lives only in its own regions; every other
162
# attention group spans all of them.
163
fa_region_ids = (
162
- np.asarray(self._scratch_region_indices, dtype=np.int64)
164
+ np.asarray(self._scratch_region_indices, dtype=np.int32)
165
if spec_type is CircularBufferSpec
164
- else np.arange(self.num_regions)
166
+ else np.arange(self.num_regions, dtype=np.int32)
167
)[:, None]
168
all_descs.append(
167
- (fa_region_ids * num_blocks + group_arr[None, :]).flatten()
169
+ (fa_region_ids * num_blocks + group_arr[None, :]).ravel()
170
)
171
elif _is_ssm_spec(spec_type):
172
# NOTE (NickLucche) SSM and Attention block regions can
@@ -175,16 +177,16 @@ class NixlBaseConnectorWorker:
177
# P and D can have different num_blocks (and thus
178
# different FA desc counts).
179
ssm_region_ids = (
178
- np.arange(num_ssm_regions, num_ssm_regions + 1)
180
+ np.arange(num_ssm_regions, num_ssm_regions + 1, dtype=np.int32)
181
if i == self._ple_group_index
180
- else np.arange(num_ssm_regions)
182
+ else np.arange(num_ssm_regions, dtype=np.int32)
183
)[:, None]
184
all_descs.append(
185
(
186
ssm_region_ids * logical_blocks
187
+ group_arr[None, :]
188
+ num_fa_descs
187
- ).flatten()
189
+ ).ravel()
190
)
191
else:
192
raise ValueError(
vllm/v1/kv_offload/base.py
+1
-1
@@ -403,7 +403,7 @@ class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
403
"""
404
405
def __init__(self, block_ids: list[int]):
406
- self.block_ids = np.array(block_ids, dtype=np.int64)
406
+ self.block_ids = np.array(block_ids, dtype=np.int32)
407
408
def __repr__(self) -> str:
409
return repr(self.block_ids)
vllm/v1/kv_offload/tiering/manager.py
+1
-1
@@ -483,7 +483,7 @@ class TieringOffloadingManager(OffloadingManager):
483
job_metadata = TransferJob(
484
job_id=job_id,
485
keys=entry.keys,
486
- block_ids=np.array(entry.block_ids, dtype=np.int64),
486
+ block_ids=np.array(entry.block_ids, dtype=np.int32),
487
is_promotion=True,
488
req_context=entry.req_context,
489
)
vllm/v1/kv_offload/tiering/obj/manager.py
+7
-4
@@ -7,6 +7,8 @@ import time
7
from collections.abc import Iterable
8
from typing import TYPE_CHECKING, ClassVar, NamedTuple
9
10
+import numpy as np
11
+
12
from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent
13
from vllm.distributed.nixl_utils import nixl_agent_config
14
from vllm.logger import init_logger
@@ -217,12 +219,13 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
219
def _submit_transfer(
220
self,
221
job_id: int,
220
- block_ids: Iterable[int],
222
+ block_ids: np.ndarray,
223
obj_keys: Iterable[str],
224
op: str,
225
) -> None:
226
"""Submit an async transfer. op is 'WRITE' (store) or 'READ' (load)."""
225
- block_ids_list = [int(bid) for bid in block_ids]
227
+ # NIXL takes indices as int32; a wider dtype silently costs a conversion.
228
+ assert block_ids.dtype == np.int32
229
# The OBJ backend maps devId -> obj_key. All descriptors must have
230
# unique devIds or later registrations overwrite earlier ones.
231
nixl_files = [
@@ -247,9 +250,9 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
250
xfer_handle = self._agent.make_prepped_xfer(
251
op,
252
self._dram_prepped_handle,
250
- block_ids_list,
253
+ block_ids,
254
obj_handle,
252
- list(range(len(nixl_files))),
255
+ np.arange(len(nixl_files), dtype=np.int32),
256
)
257
if not xfer_handle:
258
logger.warning("make_prepped_xfer failed for job %d", job_id)
vllm/v1/kv_offload/tiering/p2p/data/nixl.py
+4
-2
@@ -10,6 +10,8 @@ import itertools
10
from collections.abc import Iterable
11
from typing import Any, NamedTuple
12
13
+import numpy as np
14
+
15
from vllm.distributed.nixl_utils import NixlWrapper as _NixlAgent
16
from vllm.distributed.nixl_utils import nixl_agent_config as _NixlAgentConfig
17
from vllm.logger import init_logger
@@ -178,9 +180,9 @@ class NixlTransport(DataTransport):
180
handle = self._agent.make_prepped_xfer(
181
"WRITE",
182
self._local_dlist,
181
- local_idxs,
183
+ np.asarray(local_idxs, dtype=np.int32),
184
remote_dlist,
183
- remote_idxs,
185
+ np.asarray(remote_idxs, dtype=np.int32),
186
)
187
self._agent.transfer(handle)
188
transfer_id = next(self._next_id)
vllm/v1/kv_offload/tiering/p2p/manager.py
+1
-1
@@ -417,7 +417,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
417
def submit_store(self, job_metadata: TransferJob) -> None:
418
job_id = job_metadata.job_id
419
keys = list(job_metadata.keys)
420
- block_ids = job_metadata.block_ids
420
+ block_ids = job_metadata.block_ids.tolist()
421
422
assert len(keys) == len(block_ids)
423
vllm/v1/kv_offload/tiering/p2p/session/server.py
+1
-1
@@ -593,7 +593,7 @@ class ServerRole:
593
self.add_stored_blocks(
594
lookup.kv_request_id,
595
list(meta.keys),
596
- list(meta.block_ids),
596
+ meta.block_ids.tolist(),
597
meta.job_id,
598
round_seq=lookup.round_seq,
599
from_lookup=True,