From 33212db2c2daa3bb0692864b1b8c93eb1c8e545d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 15 Oct 2026 22:57:37 +0200 Subject: [PATCH] [KV Connector][NIXL] Reject NVFP4 KV cache with heterogeneous TP head split (#61865) Signed-off-by: NickLucche Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- .../kv_connector/unit/test_nixl_connector.py | 90 +++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 17 ++++ 2 files changed, 107 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index db7dbdab09..f90de6297a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -73,6 +73,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, KVCacheLayout, + KVQuantMode, MLAAttentionSpec, compute_layer_kv_cache_shape_bytes, ) @@ -1551,6 +1552,95 @@ class TestNixlHandshake: with pytest.raises(AssertionError): worker.add_remote_agent(bad_meta, remote_tp_size=2) + @pytest.mark.parametrize( + ("kv_quant_mode", "local_tp", "remote_tp", "num_kv_heads", "rejected"), + [ + (KVQuantMode.NVFP4, 2, 1, 8, True), + (KVQuantMode.NVFP4, 1, 2, 8, True), + (KVQuantMode.NVFP4, 2, 2, 8, False), + # Both sides hold one replicated head per rank: no head split. + (KVQuantMode.NVFP4, 16, 8, 8, False), + (KVQuantMode.NONE, 2, 1, 8, False), + ], + ) + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_handshake_rejects_nvfp4_kv_head_split( + self, + default_vllm_config, + dist_init, + kv_quant_mode, + local_tp, + remote_tp, + num_kv_heads, + rejected, + ): + """NVFP4 pages keep all heads' data before all heads' scales, so + heterogeneous-TP head splitting would silently corrupt the KV. + """ + block_size = maybe_update_block_size(16) + vllm_config = create_vllm_config(block_size=block_size) + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=max(1, num_kv_heads // local_tp), + head_size=16, + dtype=torch.uint8, + kv_quant_mode=kv_quant_mode, + ), + ) + ], + ) + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 + return_value=local_tp, + ): + worker = FakeNixlConnectorWorker( + vllm_config, + "engine", + hand_shake_latency=0, + kv_cache_config=kv_cache_config, + ) + worker.transfer_topo.total_num_kv_heads = num_kv_heads + worker.transfer_topo.local_physical_heads = max(1, num_kv_heads // local_tp) + + local_len = 4096 * worker.block_size + remote_len = ( + local_len + * max(1, num_kv_heads // remote_tp) + // worker.transfer_topo.local_physical_heads + ) + worker.block_len_per_layer = [local_len] + worker.num_blocks = 1 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks + meta = NixlAgentMetadata( + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=[0], + device_id=0, + num_blocks=1, + block_lens=[remote_len], + block_strides=[remote_len], + kv_cache_layout=worker.kv_cache_layout, + block_size=worker.block_size, + ssm_sizes=(0, 0), + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=1, + ) + + if rejected: + with pytest.raises(RuntimeError, match="NVFP4"): + worker.add_remote_agent(meta, remote_tp_size=remote_tp) + else: + worker.add_remote_agent(meta, remote_tp_size=remote_tp) + # NOTE: resource cleanup in mp backend is a bit finicky, so the order in which # we put here is important. First run ray, it will clean up the resources, then 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 ed0e6eafe6..4e86000bb1 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 @@ -82,6 +82,7 @@ from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import make_zmq_path from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.kv_cache_interface import ( + AttentionSpec, CircularBufferSpec, FullAttentionSpec, KVCacheLayout, @@ -2671,6 +2672,22 @@ class NixlBaseConnectorWorker: "Heterogeneous TP head-dimension splitting requires contiguous heads. " "Use a block-contiguous layout (e.g. LBHNC) on the prefill side." ) + # NVFP4 pages store all heads' data before all heads' scales, so a + # head's bytes are not a contiguous slice of the page. + if ( + not self.use_mla + and self.transfer_topo.local_physical_heads + != max(1, self.transfer_topo.total_num_kv_heads // remote_tp_size) + and any( + isinstance(spec, AttentionSpec) and spec.kv_quant_mode.is_nvfp4 + for spec in self._layer_specs.values() + ) + ): + raise RuntimeError( + "NIXL KV-head re-sharding is not supported for NVFP4 KV cache " + f"(local TP {self.transfer_topo.tp_size}, remote TP " + f"{remote_tp_size}). Use matching TP sizes on prefill and decode." + ) # Per-region block_len validation enforcing the P/D invariant. # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer)