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

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 4, 13:45 UTC

[Bugfix] release KV blocks for skipped P-ranks to prevent invalid KV errors and timeouts when P_tp > D_tp and MLA (#40449)

Signed-off-by: yangruize <yangruize7@163.com> Co-authored-by: Roger Wang <hey@rogerw.io>

yangrz committed Apr 29, 2026 at 02:38 UTC de3fe8dc62f3d77eb8dab8125ca90436f606bccb

2 files changed +120 -1

tests/v1/kv_connector/unit/test_nixl_connector.py

+119

@@ -2479,3 +2479,122 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)

2479 remote_tp_size=1,

2480 expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

2481 )

2482 +

2483 + @patch(

2484 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",

2485 + FakeNixlWrapper,

2486 + )

2487 + def test_mla_broadcast_notif_uses_remote_request_id(

2488 + self, default_vllm_config, dist_init

2489 + ):

2490 + """MLA + remote TP > local TP: the broadcast notification sent to

2491 + non-read prefill ranks must be keyed by the prefill-side request

2492 + id (``meta.remote.request_id``), not the local decode request id.

2493 +

2494 + Prefill ranks key ``_reqs_to_send`` by their own request id, so a

2495 + broadcast keyed by the decode id is rejected in

2496 + ``_get_new_notifs`` with "Potentially invalid KV blocks for

2497 + unrecognized request" and the blocks only release via the abort

2498 + timeout. See ``_read_blocks_for_req`` in

2499 + ``vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py``.

2500 + """

2501 + decode_tp_size = 1

2502 + prefill_tp_size = 4

2503 +

2504 + vllm_config = create_vllm_config()

2505 + vllm_config.parallel_config.tensor_parallel_size = decode_tp_size

2506 +

2507 + connector = NixlConnector(

2508 + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

2509 + )

2510 + connector.connector_worker = FakeNixlConnectorWorker(

2511 + vllm_config, connector.engine_id, hand_shake_latency=0

2512 + )

2513 + worker = connector.connector_worker

2514 +

2515 + # Force the MLA path; only `self.use_mla` gates the branches we

2516 + # exercise inside `_read_blocks_for_req`.

2517 + worker.use_mla = True

2518 +

2519 + # Manually register the remote (P) engine and pre-populate the

2520 + # per-rank state the handshake would normally fill in. The real

2521 + # `_nixl_handshake` is unnecessary here — we only need

2522 + # `transfer_topo` to know `remote_tp_size`, and `_remote_agents`

2523 + # / `dst_xfer_side_handles` to be keyed by remote rank.

2524 + remote_engine_id = "remote_engine"

2525 + worker.transfer_topo.register_remote_engine(

2526 + remote_engine_id=remote_engine_id,

2527 + remote_tp_size=prefill_tp_size,

2528 + remote_block_size=worker.block_size,

2529 + remote_block_len=worker.block_size * 4096,

2530 + remote_physical_blocks_per_logical=1,

2531 + local_block_len=worker.block_size * 4096,

2532 + )

2533 + worker._remote_agents[remote_engine_id] = {

2534 + rank: f"agent_p{rank}" for rank in range(prefill_tp_size)

2535 + }

2536 + worker.dst_xfer_side_handles = {

2537 + remote_engine_id: {rank: 100 + rank for rank in range(prefill_tp_size)}

2538 + }

2539 + # Sanity: D TP=1, P TP=4 => tp_ratio = -4 (P > D).

2540 + assert worker.transfer_topo.tp_ratio(prefill_tp_size) == -prefill_tp_size

2541 +

2542 + # Distinct ids on each side — that's the whole point of the bug.

2543 + decode_req_id = "decode-req-AAAA"

2544 + prefill_req_id = "prefill-req-BBBB"

2545 + assert decode_req_id != prefill_req_id

2546 +

2547 + metadata = NixlConnectorMetadata()

2548 + metadata.add_new_req_to_recv(

2549 + request_id=decode_req_id,

2550 + local_block_ids=([0, 1, 2],),

2551 + kv_transfer_params={

2552 + "remote_block_ids": ([10, 11, 12],),

2553 + "remote_engine_id": remote_engine_id,

2554 + "remote_request_id": prefill_req_id,

2555 + "remote_host": "localhost",

2556 + "remote_port": 1234,

2557 + "remote_tp_size": prefill_tp_size,

2558 + },

2559 + )

2560 + meta = metadata.reqs_to_recv[decode_req_id]

2561 +

2562 + # Capture broadcast send_notif calls; stub `_read_blocks` so we

2563 + # don't need a working xfer path. Real `_read_blocks` emits its

2564 + # auto-notif via `make_prepped_xfer`, not via `send_notif`, so

2565 + # any captured `send_notif` here is a broadcast.

2566 + send_notif_calls: list[tuple[str, bytes]] = []

2567 + worker.nixl_wrapper.send_notif = ( # type: ignore[method-assign]

2568 + lambda agent_name, notif_msg: send_notif_calls.append(

2569 + (agent_name, notif_msg)

2570 + )

2571 + )

2572 + worker._read_blocks = MagicMock() # type: ignore[method-assign]

2573 +

2574 + worker._read_blocks_for_req(decode_req_id, meta)

2575 +

2576 + # MLA: read once from rank 0 and broadcast to the other ranks.

2577 + worker._read_blocks.assert_called_once()

2578 + assert worker._read_blocks.call_args.kwargs["remote_rank"] == 0

2579 + assert (

2580 + worker._read_blocks.call_args.kwargs["remote_request_id"] == prefill_req_id

2581 + )

2582 +

2583 + # Broadcast goes to ranks {1, 2, 3} only, never to the read target.

2584 + expected_recipients = {

2585 + worker._remote_agents[remote_engine_id][r]

2586 + for r in range(1, prefill_tp_size)

2587 + }

2588 + assert {agent for agent, _ in send_notif_calls} == expected_recipients

2589 +

2590 + # Every broadcast notif must be keyed by the prefill request id.

2591 + # Pre-fix this used the *decode* request id, which prefill ranks

2592 + # didn't recognize.

2593 + expected_notif = f"{prefill_req_id}:{decode_tp_size}".encode()

2594 + bad_notif = f"{decode_req_id}:{decode_tp_size}".encode()

2595 + for agent, notif in send_notif_calls:

2596 + assert notif == expected_notif, (

2597 + f"Broadcast notif to {agent!r} must use prefill_req_id; "

2598 + f"got {notif!r} (expected {expected_notif!r}, "

2599 + f"buggy form would be {bad_notif!r})"

2600 + )

vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py

+1 -1

@@ -1971,7 +1971,7 @@ class NixlConnectorWorker:

1971 if self.use_mla and tp_ratio < 0:

1972 # ..but we still need to notify the other remote ranks that we

1973 # have the blocks we need so they can update the request state.

1974 - notif_id = f"{req_id}:{self.world_size}".encode()

1974 + notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()

1975 remote_agents = self._remote_agents[meta.remote.engine_id]

1976 for rank_to_notify, agent in remote_agents.items():

1977 if rank_to_notify != remote_rank: