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 6, 13:15 UTC
[NIXL][Bugfix] metrics & testing minor bug (#36051)
Signed-off-by: Andy Lo <andy@mistral.ai>
Andy Lo committed
Mar 18, 2026 at 13:39 UTC
98b09ddc2761545f3164d930b143f84737b1ab43
2 files changed
+9
-8
tests/v1/kv_connector/unit/test_nixl_connector.py
+8
-7
@@ -694,16 +694,18 @@ class TestNixlHandshake:
694
)
695
@pytest.mark.parametrize("local_tp_size", [1, 2])
696
def test_prefill_tp_size_greater_than_decode_tp_size(
697
- self, local_tp_size: int, default_vllm_config, dist_init
697
+ self, local_tp_size: int, default_vllm_config, dist_init, monkeypatch
698
):
699
"""
700
Verify remote TP > local TP handshake succeeds with different
701
remote configurations.
702
"""
703
+ monkeypatch.setattr(
704
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size",
705
+ lambda: local_tp_size,
706
+ )
707
708
vllm_config = create_vllm_config()
705
- local_tp_size = 1
706
- vllm_config.parallel_config.tensor_parallel_size = local_tp_size
709
710
connector = NixlConnector(
711
vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
@@ -738,10 +740,10 @@ class TestNixlHandshake:
740
remote_agents = worker._nixl_handshake(
741
host="localhost",
742
port=1234,
741
- remote_tp_size=2,
743
+ remote_tp_size=4,
744
expected_engine_id=worker.REMOTE_ENGINE_ID,
745
)
744
- check_handshake(2)
746
+ check_handshake(4)
747
748
# NOTE flexibility: a second remote with higher number of ranks is
749
# discovered. This is not a scenario we actively support right now, but
@@ -759,9 +761,8 @@ class TestNixlHandshake:
761
"vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
762
FakeNixlWrapper,
763
)
762
- @pytest.mark.parametrize("local_tp_size", [1, 2])
764
def test_prefill_tp_size_greater_than_decode_tp_size_mla(
764
- self, local_tp_size: int, default_vllm_config, dist_init
765
+ self, default_vllm_config, dist_init
766
):
767
"""
768
Verify remote TP > local TP handshake succeeds with different
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+1
-1
@@ -1318,12 +1318,12 @@ class NixlConnectorWorker:
1318
f"Expected {expected_engine_id},"
1319
f"received {metadata.engine_id}."
1320
)
1321
- setup_agent_time = time.perf_counter()
1321
1322
# Register Remote agent.
1323
remote_agent_name = self.add_remote_agent(
1324
metadata, remote_rank, remote_tp_size
1325
)
1326
+ setup_agent_time = time.perf_counter()
1327
logger.debug(
1328
"NIXL handshake: add agent took: %s",
1329
setup_agent_time - got_metadata_time,