Merge branch 'main' into wentao-fix-es-v2-bug

Signed-off-by: yewentao256 <zhyanwentao@126.com>

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -202,3 +202,79 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
     output = mrv2.GPUModelRunner.sample_tokens(runner, None)
     assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None)
     assert events == ["receive", "postprocess_num_computed_tokens", "eplb"]
+
+
+def test_v2_save_serving_state_round_trips_block_table_state():
+    class FakeStagedWriteTensor:
+        def __init__(self, values: torch.Tensor):
+            self.gpu = values.clone()
+            self.clear_calls = 0
+
+        def clear_staged_writes(self) -> None:
+            self.clear_calls += 1
+
+    class FakeNumBlocks:
+        def __init__(self, values: torch.Tensor):
+            self.cpu = values.clone()
+            self.gpu = values.clone()
+            self.copy_to_uva_calls = 0
+
+        def copy_to_uva(self) -> torch.Tensor:
+            self.copy_to_uva_calls += 1
+            self.gpu = self.cpu.clone()
+            return self.gpu
+
+    block_tables = SimpleNamespace(
+        block_tables=[
+            FakeStagedWriteTensor(
+                torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
+            ),
+            FakeStagedWriteTensor(torch.tensor([[7, 8], [9, 10]], dtype=torch.int32)),
+        ],
+        input_block_tables=[
+            torch.tensor([[11, 12, 13], [14, 15, 16]], dtype=torch.int32),
+            torch.tensor([[17, 18], [19, 20]], dtype=torch.int32),
+        ],
+        num_blocks=FakeNumBlocks(torch.tensor([[3, 2], [1, 2]], dtype=torch.int32)),
+    )
+    runner = _make_runner(block_tables=block_tables)
+
+    original_rows = [bt.gpu.clone() for bt in block_tables.block_tables]
+    original_inputs = [bt.clone() for bt in block_tables.input_block_tables]
+    original_num_blocks = block_tables.num_blocks.cpu.clone()
+
+    saved_state = mrv2.GPUModelRunner.save_serving_state(runner)
+
+    # test states saved correctly
+    assert [bt.clear_calls for bt in block_tables.block_tables] == [1, 1]
+    assert all(torch.count_nonzero(bt.gpu) == 0 for bt in block_tables.block_tables)
+    assert all(
+        torch.count_nonzero(input_bt) == 0
+        for input_bt in block_tables.input_block_tables
+    )
+    assert torch.count_nonzero(block_tables.num_blocks.cpu) == 0
+    assert torch.equal(saved_state.block_table_rows[0], original_rows[0])
+    assert torch.equal(saved_state.block_table_rows[1], original_rows[1])
+    assert torch.equal(saved_state.input_block_tables[0], original_inputs[0])
+    assert torch.equal(saved_state.input_block_tables[1], original_inputs[1])
+    assert torch.equal(saved_state.num_blocks, original_num_blocks)
+    assert block_tables.num_blocks.copy_to_uva_calls == 1
+
+    # fill with new values
+    block_tables.block_tables[0].gpu.fill_(99)
+    block_tables.block_tables[1].gpu.fill_(98)
+    block_tables.input_block_tables[0].fill_(97)
+    block_tables.input_block_tables[1].fill_(96)
+    block_tables.num_blocks.cpu.fill_(95)
+
+    mrv2.GPUModelRunner.restore_serving_state(runner, saved_state)
+
+    # test states restored correctly
+    assert [bt.clear_calls for bt in block_tables.block_tables] == [2, 2]
+    assert torch.equal(block_tables.block_tables[0].gpu, original_rows[0])
+    assert torch.equal(block_tables.block_tables[1].gpu, original_rows[1])
+    assert torch.equal(block_tables.input_block_tables[0], original_inputs[0])
+    assert torch.equal(block_tables.input_block_tables[1], original_inputs[1])
+    assert torch.equal(block_tables.num_blocks.cpu, original_num_blocks)
+    assert torch.equal(block_tables.num_blocks.gpu, original_num_blocks)
+    assert block_tables.num_blocks.copy_to_uva_calls == 2

vllm/distributed/elastic_ep/elastic_execute.py

@@ -500,23 +500,11 @@ class ElasticEPScalingExecutor:
             compilation_counter.stock_torch_compile_count += 1
             self.worker.model_runner.model.compile(fullgraph=True, backend=backend)
 
-        multi_block_table = self.worker.model_runner.input_batch.block_table
-        saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
-        for bt in multi_block_table.block_tables:
-            saved_block_tables.append(
-                (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
-            )
-        multi_block_table.clear()
-
+        saved_state = self.worker.model_runner.save_serving_state()
         unlock_workspace()
         self.worker.compile_or_warm_up_model()
         lock_workspace()
-
-        for bt, (saved_gpu, saved_cpu) in zip(
-            multi_block_table.block_tables, saved_block_tables
-        ):
-            bt.block_table.gpu.copy_(saved_gpu)
-            bt.block_table.cpu.copy_(saved_cpu)
+        self.worker.model_runner.restore_serving_state(saved_state)
         if new_dp_size < old_dp_size:
             self._set_eplb_suppressed(False)
 
@@ -651,13 +639,7 @@ class ElasticEPScalingExecutor:
         # Save and clear block tables so profile_run/compile_or_warm_up_model
         # don't write dummy slot mappings into real KV-cache blocks (mirrors
         # switch_and_prepare's pattern).
-        multi_block_table = self.worker.model_runner.input_batch.block_table
-        saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
-        for bt in multi_block_table.block_tables:
-            saved_block_tables.append(
-                (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
-            )
-        multi_block_table.clear()
+        saved_state = self.worker.model_runner.save_serving_state()
 
         # _ensure_workspace_size allocates a fresh tensor on grow, leaving
         # captured CUDA graphs with stale data pointers; drop graphs before
@@ -678,9 +660,4 @@ class ElasticEPScalingExecutor:
         self.worker.compile_or_warm_up_model()
 
         lock_workspace()
-
-        for bt, (saved_gpu, saved_cpu) in zip(
-            multi_block_table.block_tables, saved_block_tables
-        ):
-            bt.block_table.gpu.copy_(saved_gpu)
-            bt.block_table.cpu.copy_(saved_cpu)
+        self.worker.model_runner.restore_serving_state(saved_state)

vllm/v1/worker/gpu/block_table.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable
+from dataclasses import dataclass
 
 import torch
 
@@ -196,6 +197,49 @@ class BlockTables:
         return self.slot_mappings[:, :num_tokens]
 
 
+@dataclass
+class BlockTableServingState:
+    block_table_rows: list[torch.Tensor]
+    input_block_tables: list[torch.Tensor]
+    num_blocks: torch.Tensor
+
+
+def save_block_table_serving_state(
+    block_tables: BlockTables,
+) -> BlockTableServingState:
+    saved_state = BlockTableServingState(
+        block_table_rows=[bt.gpu.clone() for bt in block_tables.block_tables],
+        input_block_tables=[
+            input_bt.clone() for input_bt in block_tables.input_block_tables
+        ],
+        num_blocks=block_tables.num_blocks.cpu.clone(),
+    )
+
+    for bt in block_tables.block_tables:
+        bt.gpu.zero_()
+        bt.clear_staged_writes()
+    for input_bt in block_tables.input_block_tables:
+        input_bt.zero_()
+    block_tables.num_blocks.cpu.zero_()
+    block_tables.num_blocks.copy_to_uva()
+    return saved_state
+
+
+def restore_block_table_serving_state(
+    block_tables: BlockTables,
+    saved_state: BlockTableServingState,
+) -> None:
+    for bt, saved_gpu in zip(block_tables.block_tables, saved_state.block_table_rows):
+        bt.gpu.copy_(saved_gpu)
+        bt.clear_staged_writes()
+    for input_bt, saved_input in zip(
+        block_tables.input_block_tables, saved_state.input_block_tables
+    ):
+        input_bt.copy_(saved_input)
+    block_tables.num_blocks.cpu.copy_(saved_state.num_blocks)
+    block_tables.num_blocks.copy_to_uva()
+
+
 @triton.jit(do_not_specialize=["num_reqs"])
 def _gather_block_tables_kernel(
     batch_idx_to_req_idx,  # [batch_size]

vllm/v1/worker/gpu/model_runner.py

@@ -59,7 +59,12 @@ from vllm.v1.worker.gpu.attn_utils import (
     init_attn_backend,
     init_kv_cache,
 )
-from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.block_table import (
+    BlockTables,
+    BlockTableServingState,
+    restore_block_table_serving_state,
+    save_block_table_serving_state,
+)
 from vllm.v1.worker.gpu.buffer_utils import (
     async_copy_to_gpu,
     set_default_max_concurrency,
@@ -676,6 +681,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
     def post_kv_cache_wake_up(self) -> None:
         self.block_tables.init_block_table_layout_tensors()
 
+    def save_serving_state(self) -> BlockTableServingState:
+        return save_block_table_serving_state(self.block_tables)
+
+    def restore_serving_state(self, saved_state: BlockTableServingState) -> None:
+        restore_block_table_serving_state(self.block_tables, saved_state)
+
     def reset_mm_cache(self) -> None:
         if self.encoder_cache is not None:
             self.encoder_cache.reset_mm_cache()

vllm/v1/worker/gpu_model_runner.py

@@ -969,6 +969,26 @@ class GPUModelRunner(
     def post_kv_cache_wake_up(self) -> None:
         self.init_fp8_kv_scales()
 
+    def save_serving_state(self) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        multi_block_table = self.input_batch.block_table
+        saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for bt in multi_block_table.block_tables:
+            saved_block_tables.append(
+                (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
+            )
+        multi_block_table.clear()
+        return saved_block_tables
+
+    def restore_serving_state(
+        self, saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]]
+    ) -> None:
+        multi_block_table = self.input_batch.block_table
+        for bt, (saved_gpu, saved_cpu) in zip(
+            multi_block_table.block_tables, saved_block_tables
+        ):
+            bt.block_table.gpu.copy_(saved_gpu)
+            bt.block_table.cpu.copy_(saved_cpu)
+
     @torch.inference_mode()
     def init_fp8_kv_scales(self) -> None:
         """