@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7e6edb1469

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[NIXL][HeteroTP] Enable KV transfer from HND prefill to NHD decode (#26556)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

Chendi.Xue committed Oct 14, 2025 at 04:46 UTC 7e6edb14698c1a760272dd44363a288aeb07b571

6 files changed +150 -5

docs/features/nixl_connector_usage.md

+10

@@ -156,6 +156,16 @@ python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \

156 NixlConnector currently does not distinguish `kv_role`; the actual prefiller/decoder roles are determined by the upper-level proxy (e.g., `toy_proxy_server.py` using `--prefiller-hosts` and `--decoder-hosts`).

157 Therefore, `kv_role` in `--kv-transfer-config` is effectively a placeholder and does not affect NixlConnector's behavior.

158

159 +## Experimental Feature

160 +

161 +### Heterogenuous KV Layout support

162 +

163 +Support use case: Prefill with 'HND' and decode with 'NHD' with experimental configuration

164 +

165 +```bash

166 +--kv-transfer-config '{..., "enable_permute_local_kv":"True"}'

167 +```

168 +

169 ## Example Scripts/Code

170

171 Refer to these example scripts in the vLLM repository:

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

+11 -2

@@ -19,11 +19,18 @@ done

20 echo "Running accuracy tests with kv_buffer_device=$KV_BUFFER_DEVICE"

22 +DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"HND"} # Default to HND, optional NHD

23 +if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then

24 + KV_CONFIG_HETERO_LAYOUT=',"enable_permute_local_kv":"True"'

25 +else

26 + KV_CONFIG_HETERO_LAYOUT=''

27 +fi

28 +

29 # Build the kv-transfer-config once

30 if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then

24 - KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"}'

31 + KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}'

32 else

26 - KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\"}"

33 + KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"

34 fi

36 # Models to run

@@ -117,6 +124,7 @@ run_tests_for_model() {

124

125 # Build the command with or without model-specific args

126 BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \

127 + VLLM_KV_CACHE_LAYOUT='HND' \

128 UCX_NET_DEVICES=all \

129 VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \

130 vllm serve $model_name \

@@ -157,6 +165,7 @@ run_tests_for_model() {

165

166 # Build the command with or without model-specific args

167 BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \

168 + VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \

169 UCX_NET_DEVICES=all \

170 VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \

171 vllm serve $model_name \

tests/v1/kv_connector/unit/test_nixl_connector.py

+58 -2

@@ -286,9 +286,12 @@ def test_prompt_less_than_block_size():

286 class FakeNixlConnectorWorker(NixlConnectorWorker):

287 REMOTE_ENGINE_ID = "remote_engine"

288

289 - def __init__(self, *args, hand_shake_latency: float = 1.8, **kwargs):

289 + def __init__(

290 + self, *args, hand_shake_latency: float = 1.8, kv_cache_layout="HND", **kwargs

291 + ):

292 super().__init__(*args, **kwargs)

293 self._hand_shake_latency = hand_shake_latency

294 + self.kv_cache_layout = kv_cache_layout

295

296 def _nixl_handshake(

297 self, host: str, port: int, remote_tp_size: int, expected_engine_id: str

@@ -564,10 +567,63 @@ class TestNixlHandshake:

567

568 # We don't check layout for homogeneous TP and MLA for now, as the

569 # whole block is moved.

567 - worker.add_remote_agent(meta, remote_tp_size=2)

570 + with pytest.raises(RuntimeError):

571 + # mismatched layout is expected to fail

572 + worker.add_remote_agent(meta, remote_tp_size=2)

573 with pytest.raises(AssertionError):

574 worker.add_remote_agent(meta, remote_tp_size=1)

575

576 + @patch(

577 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

578 + FakeNixlWrapper,

579 + )

580 + def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental(

581 + self, dist_init

582 + ):

583 + """

584 + Verify that adding a remote agent fails if kv_cache_layout differs.

585 + This test is only relevant for heterogeneous TP.

586 + """

587 + vllm_config = create_vllm_config(enable_permute_local_kv=True)

588 +

589 + # Mock TP world size to 2 to force heterogeneous TP when

590 + # remote_tp_size=1

591 + with patch(

592 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size", # noqa: E501

593 + return_value=2,

594 + ):

595 + # Initialize connector and worker (with fake NIXL wrapper)

596 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

597 + connector.connector_worker = FakeNixlConnectorWorker(

598 + vllm_config,

599 + connector.engine_id,

600 + hand_shake_latency=0,

601 + kv_cache_layout="NHD",

602 + )

603 + worker = connector.connector_worker

604 +

605 + # Minimal local registration params used by add_remote_agent

606 + worker.slot_size_per_layer = [2048]

607 + worker.block_len_per_layer = [2048 * worker.block_size]

608 + worker.num_blocks = 1

609 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks

610 +

611 + # Metadata with different kv_cache_layout than local worker

612 + meta = NixlAgentMetadata(

613 + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

614 + agent_metadata=FakeNixlWrapper.AGENT_METADATA,

615 + kv_caches_base_addr=[0],

616 + num_blocks=1,

617 + # prefill TP=1, decode TP=2, remote block_lens is double to local

618 + block_lens=[i * 2 for i in worker.block_len_per_layer],

619 + attn_backend_name=worker.backend_name,

620 + kv_cache_layout="HND",

621 + )

622 +

623 + # We don't check layout for homogeneous TP and MLA for now, as the

624 + # whole block is moved.

625 + worker.add_remote_agent(meta, remote_tp_size=1)

626 +

627

628 # NOTE: resource cleanup in mp backend is a bit finicky, so the order in which

629 # we put here is important. First run ray, it will clean up the resources, then

tests/v1/kv_connector/unit/utils.py

@@ -83,6 +83,7 @@ def create_vllm_config(

83 block_size: int = 16,

84 max_model_len: int = 10000,

85 enable_chunked_prefill: bool = True,

86 + enable_permute_local_kv: bool = False,

87 ) -> VllmConfig:

88 """Initialize VllmConfig For Testing."""

89 scheduler_config = SchedulerConfig(

@@ -108,6 +109,7 @@ def create_vllm_config(

109 kv_transfer_config = KVTransferConfig(

110 kv_connector="NixlConnector",

111 kv_role="kv_both",

112 + enable_permute_local_kv=enable_permute_local_kv,

113 )

114 return VllmConfig(

115 scheduler_config=scheduler_config,

vllm/config/kv_transfer.py

@@ -61,6 +61,9 @@ class KVTransferConfig:

61 """The Python module path to dynamically load the KV connector from.

62 Only supported in V1."""

64 + enable_permute_local_kv: bool = False

65 + """Experiment feature flag to enable HND to NHD KV Transfer"""

66 +

67 def compute_hash(self) -> str:

68 """

69 WARNING: Whenever a new field is added to this config,

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+66 -1

@@ -563,6 +563,7 @@ class NixlConnectorWorker:

563 self.world_size = get_tensor_model_parallel_world_size()

564 self.tp_group = get_tp_group()

565 self.num_blocks = 0

566 + self.enable_permute_local_kv = False

567

568 # KV Caches and nixl tracking data.

569 self.device_type = current_platform.device_type

@@ -1094,6 +1095,23 @@ class NixlConnectorWorker:

1095 is_kv_replicated = self._tp_size[engine_id] // total_num_kv_heads >= 1

1096

1097 remote_block_len = nixl_agent_meta.block_lens[0]

1098 + if nixl_agent_meta.kv_cache_layout != self.kv_cache_layout:

1099 + if (

1100 + self.vllm_config.kv_transfer_config is not None

1101 + and self.vllm_config.kv_transfer_config.enable_permute_local_kv

1102 + and nixl_agent_meta.kv_cache_layout == "HND"

1103 + ):

1104 + logger.info(

1105 + "Remote is HND and local is NHD, enabled additional permute "

1106 + "on local device KV."

1107 + )

1108 + self.enable_permute_local_kv = True

1109 + else:

1110 + raise RuntimeError(

1111 + "Heterogeneous TP expects same kv_cache_layout. "

1112 + "Or enable experimental feature to use HND to NHD support by "

1113 + "setting 'enable_permute_local_kv'=True in --kv-transfer-config."

1114 + )

1115 if self.use_mla or is_kv_replicated:

1116 # With replicated KV cache, only the number of blocks can differ.

1117 assert self.block_len_per_layer == nixl_agent_meta.block_lens, (

@@ -1114,7 +1132,10 @@ class NixlConnectorWorker:

1132 remote_block_size //= 2

1133 if tp_ratio > 1:

1134 # Heterogeneous TP expects same kv_cache_layout.

1117 - assert nixl_agent_meta.kv_cache_layout == self.kv_cache_layout

1135 + if nixl_agent_meta.kv_cache_layout == "NHD":

1136 + raise ValueError(

1137 + "Heterogeneous TP is not supported for remote with NHD."

1138 + )

1139 if self.device_type == "xpu":

1140 raise ValueError("Heterogeneous TP is not supported on XPU")

1141

@@ -1226,6 +1247,41 @@ class NixlConnectorWorker:

1247 "d2h",

1248 )

1249

1250 + def permute_device_kv(self, block_ids: list[int]):

1251 + """Transforms the layout of received KV cache blocks to the local format.

1252 +

1253 + This method corrects layout mismatches from direct memory copies by

1254 + permuting the tensor dimensions.

1255 +

1256 + - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]`

1257 + - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]`

1258 +

1259 + Args:

1260 + block_ids: A list of block IDs to update and permute.

1261 +

1262 + Implementation:

1263 + - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout

1264 + - permuted_blocks = x.permute(*inv_order) # transpose n_kv_heads, block_size

1265 + - cache.index_copy_(0, indices, permuted_blocks) # copy permuted kv back

1266 +

1267 + """

1268 + split_k_and_v = not (self.use_mla or self._use_pallas or self._use_flashinfer)

1269 + inv_order = [0, 2, 1, 3]

1270 + sample_cache = list(self.device_kv_caches.values())[0][0]

1271 + target_shape = list(sample_cache.shape)

1272 + target_shape[0] = -1

1273 + src_shape = tuple(target_shape[i] for i in inv_order)

1274 + indices = torch.tensor(block_ids, device=sample_cache.device)

1275 +

1276 + for _, cache_or_caches in self.device_kv_caches.items():

1277 + cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]

1278 + for cache in cache_list:

1279 + blocks_to_update = cache.index_select(0, indices)

1280 + permuted_blocks = blocks_to_update.reshape(src_shape).permute(

1281 + *inv_order

1282 + )

1283 + cache.index_copy_(0, indices, permuted_blocks)

1284 +

1285 def get_finished(self) -> tuple[set[str], set[str]]:

1286 """

1287 Get requests that are done sending or recving on this specific worker.

@@ -1273,6 +1329,15 @@ class NixlConnectorWorker:

1329 del self._reqs_to_send[req_id]

1330 done_sending.add(req_id)

1331

1332 + if self.enable_permute_local_kv and len(done_recving) > 0:

1333 + block_ids = []

1334 + for req_id in done_recving:

1335 + meta = self._recving_metadata.pop(req_id)

1336 + assert meta, f"{req_id} not found in recving_metadata list"

1337 + block_ids += meta.local_block_ids

1338 +

1339 + self.permute_device_kv(block_ids)

1340 +

1341 return done_sending, done_recving

1342

1343 def _get_new_notifs(self) -> set[str]: