[Rebase] Rebase to vllm 0.26.0 (#5443)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: tzhouam <tzhouam@gmail.com>

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -780,6 +780,10 @@ def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_ada
 class _HashableRequest(SimpleNamespace):
     """SimpleNamespace that can be added to a set (needed by scheduler internals)."""
 
+    # vLLM 0.26: update_from_output settles this counter for every scheduled
+    # request; real Requests initialise it to 0 (vllm/v1/request.py).
+    num_in_flight_tokens = 0
+
     def __hash__(self):
         return hash(self.request_id)
 
@@ -833,7 +837,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     scheduler.requests = {"req-s1": request}
 
     scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
-    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
     scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
     scheduler.running = [request]
     scheduler.waiting = mocker.MagicMock()
@@ -918,7 +922,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     scheduler._update_request_with_output = mocker.MagicMock(return_value=([], True))
     scheduler._process_kv_transfer_trigger = mocker.MagicMock(return_value=False)
     scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
-    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
     scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
     scheduler.running = [request]
     scheduler.waiting = mocker.MagicMock()
@@ -1047,7 +1051,7 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
     scheduler._pending_finish_reqs = list(pending_finish_reqs)
 
     scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
-    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
     scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
     scheduler.running = list(running)
     scheduler.waiting = mocker.MagicMock()
@@ -1103,7 +1107,7 @@ def test_deferred_finish_emits_finished_output(mocker: MockerFixture):
         running=[request],
         pending_finish_reqs=[request],
     )
-    scheduler._free_request.return_value = {"mock": "kv_params"}
+    scheduler._free_request.return_value = ({"mock": "kv_params"}, None)
 
     result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
 

tests/distributed/omni_connectors/test_kv_flow.py

@@ -330,6 +330,31 @@ def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants,
     assert normalized is None
 
 
+def test_normalize_layer_kv_unpacks_packed_layout(common_constants):
+    """The vLLM >= 0.23 packed layout (num_blocks, n_heads, block_size,
+    2*head_dim) is unpacked into (num_blocks, block_size, n_heads, head_dim)
+    key/value blocks when block_size is provided."""
+    block_size = common_constants["block_size"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    req_id = common_constants["req_id"]
+    num_blocks = 10
+
+    layer_kv = torch.randn(num_blocks, num_heads, block_size, 2 * head_dim)
+
+    # Without block_size the 4-D packed layout is not recognized.
+    assert normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0) is None
+
+    normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0, block_size=block_size)
+    assert normalized is not None
+    key_blocks, value_blocks = normalized
+    assert key_blocks.shape == (num_blocks, block_size, num_heads, head_dim)
+    assert value_blocks.shape == (num_blocks, block_size, num_heads, head_dim)
+    reference = layer_kv.transpose(1, 2)
+    assert torch.equal(key_blocks, reference[..., :head_dim])
+    assert torch.equal(value_blocks, reference[..., head_dim:])
+
+
 def test_manager_reception(kv_config, mock_connector, common_constants):
     """Test reception and injection logic in OmniKVTransferManager."""
     num_layers = common_constants["num_layers"]