[MRV2] Miscellaneous code structure updates

tests/v1/spec_decode/test_ngram_gpu.py

@@ -147,7 +147,7 @@ def _propose(
         next_prefill_tokens=torch.zeros(B, dtype=torch.int32, device=DEVICE),
         temperature=torch.zeros(B, dtype=torch.float32, device=DEVICE),
         seeds=torch.zeros(B, dtype=torch.int64, device=DEVICE),
-        dp_sync=None,
+        dp_sync_state=None,
     )
     return drafts.cpu().tolist()
 

tests/v1/worker/test_gpu_kv_connector.py

@@ -63,13 +63,11 @@ def test_load_start_phase(
 
     request_indices = torch.tensor([3, 1])
     request_ids = ["first", "second"]
-    attn_metadata = {"layer": object()}
-    connector.pre_forward(
-        output,
-        request_state_indices=request_indices,
-        request_ids=request_ids,
-        attn_metadata=attn_metadata,
+    input_batch = SimpleNamespace(
+        idx_mapping=request_indices, req_ids=request_ids, num_tokens=5
     )
+    attn_metadata = {"layer": object()}
+    connector.pre_forward(output, input_batch, attn_metadata=attn_metadata)
     assert events == (
         ["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"]
     )
@@ -80,6 +78,8 @@ def test_load_start_phase(
     kwargs = connector.kv_connector.start_load_kv.call_args.kwargs
     assert kwargs["request_state_indices"] is request_indices
     assert kwargs["request_ids"] is request_ids
+    assert kwargs["num_tokens"] == 5
+    assert kwargs["scheduler_output"] is output
     assert kwargs["attn_metadata"] is attn_metadata
 
     # A subsequent step without a forward must not reuse the prior batch.

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -21,6 +21,7 @@
 from vllm.v1.worker.gpu.async_utils import async_copy_to_np
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner
+from vllm.v1.worker.gpu.spec_decode.utils import get_drafter_hidden_states
 
 
 def test_non_last_pp_rank_uses_global_batch_for_sample_feedback():
@@ -43,7 +44,7 @@ def test_non_last_pp_rank_uses_global_batch_for_sample_feedback():
         slot_mappings_by_layer=None,
         hidden_states=None,
         aux_hidden_states=None,
-        dp_sync=None,
+        dp_sync_state=None,
         finished_req_ids=set(),
         ec_connector_output=None,
         cudagraph_stats=None,
@@ -342,16 +343,15 @@ def test_capture_model_profile_only_skips_lock(monkeypatch):
 @pytest.mark.parametrize("target_buffer", ["absent", "none", "tensor"])
 def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer):
     """Targets allocate the MTP hidden buffer only for hidden-state drafters."""
-    runner = GPUModelRunner.__new__(GPUModelRunner)
     hidden_states = torch.zeros(4, 8)
     buffer = torch.arange(16 * 8, dtype=torch.float32).view(16, 8)
     if target_buffer == "absent":
-        runner.model = SimpleNamespace()
+        model = SimpleNamespace()
     else:
         returned = buffer if target_buffer == "tensor" else None
-        runner.model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned)
+        model = SimpleNamespace(get_mtp_target_hidden_states=lambda: returned)
 
-    out = runner._get_drafter_hidden_states(hidden_states)
+    out = get_drafter_hidden_states(model, hidden_states)
 
     if target_buffer == "tensor":
         assert torch.equal(out, buffer[:4])

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -185,7 +185,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
         slot_mappings_by_layer=None,
         hidden_states=None,
         aux_hidden_states=None,
-        dp_sync=None,
+        dp_sync_state=None,
         finished_req_ids=set(),
         ec_connector_output=None,
         routed_experts=None,

tests/v1/worker/test_gpu_ubatch_slicing.py

@@ -350,10 +350,10 @@ def test_decode_threshold_only_applies_when_every_rank_is_a_uniform_decode():
 
 def test_microbatching_pads_all_ranks_to_the_largest():
     """Ranks must run the same token count so microbatch sizes line up."""
-    desc, dp_sync = _sync_dp([200, 256])
+    desc, dp_sync_state = _sync_dp([200, 256])
     assert desc.num_tokens == 256
-    assert dp_sync is not None
-    assert dp_sync.num_tokens_across_dp.tolist() == [256, 256]
+    assert dp_sync_state is not None
+    assert dp_sync_state.num_tokens_across_dp.tolist() == [256, 256]
 
 
 def test_microbatching_survives_a_rank_that_cannot_fill_it():
@@ -959,7 +959,7 @@ def test_microbatched_graph_needs_every_rank_to_reach_the_split():
     assert desc.num_ubatches == 2
 
     # Rank 1 pads the group up to a 128-token graph, but rank 0 only reaches 33.
-    desc, dp_sync = _sync_dp([33, 128], uniform, cudagraph_manager=manager)
+    desc, dp_sync_state = _sync_dp([33, 128], uniform, cudagraph_manager=manager)
     assert desc.cg_mode == CUDAGraphMode.NONE
     assert desc.num_ubatches == 2
-    assert dp_sync is not None and dp_sync.eager
+    assert dp_sync_state is not None and dp_sync_state.eager

vllm/compilation/backends.py

@@ -1091,12 +1091,12 @@ def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any:
         # TODO(patchy): the V1 torch.compile ngram-gpu kernel causes vllm
         # torch compile cache errors. The V2 implementation is pure Triton and
         # does not need the cache disabled.
-        is_ngram_gpu_enabled = (
+        is_v1_ngram_gpu_enabled = (
             vllm_config.speculative_config is not None
             and vllm_config.speculative_config.use_ngram_gpu()
             and not vllm_config.use_v2_model_runner
         )
-        disable_cache = disable_cache or is_ngram_gpu_enabled
+        disable_cache = disable_cache or is_v1_ngram_gpu_enabled
 
         if disable_cache:
             logger.info("%svLLM's torch.compile cache is disabled.", self._log_prefix())

vllm/v1/worker/gpu/dp_utils.py

@@ -152,7 +152,7 @@ def sync_cudagraph_and_dp_padding(
                 num_reqs=num_reqs,
             )
 
-    synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item()))
+    synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min()))
 
     # If any rank wants to run eager, all ranks run eager
     if synced_cg_mode == CUDAGraphMode.NONE:
@@ -175,13 +175,13 @@ def sync_cudagraph_and_dp_padding(
         "cudagraph_manager should only be None during profile run, "
         "where synced_cg_mode must be NONE across all DP ranks"
     )
-    synced_num_tokens = int(num_tokens_across_dp.max().item())
+    synced_num_tokens = int(num_tokens_across_dp.max())
 
     # Varlen decode graphs are selected by the query-length bound, so ranks must agree
     # on it or they pad to different token counts below.
     synced_max_query_len: int | None = None
-    if bool(torch.all(max_query_lens_across_dp != -1).item()):
-        synced_max_query_len = int(max_query_lens_across_dp.max().item())
+    if bool(torch.all(max_query_lens_across_dp != -1)):
+        synced_max_query_len = int(max_query_lens_across_dp.max())
 
     # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs
     # so we don't perform request padding for PIECEWISE graphs.
@@ -223,12 +223,12 @@ def dispatch_cg_and_sync_dp(
     parallel_config: ParallelConfig | None = None,
     allow_ubatching: bool = False,
     uniform_decode: bool = False,
-    dp_sync: DPSyncState | None = None,
+    dp_sync_state: DPSyncState | None = None,
 ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]:
     """Pick a cudagraph descriptor for this batch, agreeing it across DP ranks.
 
     Runs a collective when dp_size > 1 so every rank dispatches to the same
-    shape. Pass `dp_sync` from a dispatch already made over this same batch (a
+    shape. Pass `dp_sync_state` from a dispatch already made over this same batch (a
     drafter's prefill runs the target's batch shape) to reuse that agreement
     instead, with no collective.
 
@@ -238,7 +238,7 @@ def dispatch_cg_and_sync_dp(
         num_reqs: Requests in this rank's batch.
         num_tokens: Tokens in this rank's batch, already padded by the caller.
         uniform_token_count: Per-request token count if this rank's batch is a
-            uniform decode, else None. `dp_sync.uniform_token_count` takes its
+            uniform decode, else None. `dp_sync_state.uniform_token_count` takes its
             place when a sync is reused, since that one is agreed across ranks.
         dp_size: Data-parallel world size. 1 skips all cross-rank work.
         dp_rank: This rank's index in the DP group.
@@ -254,7 +254,7 @@ def dispatch_cg_and_sync_dp(
             microbatches. Agreed across ranks before it takes effect.
         uniform_decode: Whether this rank's batch is a uniform decode, used
             to pick the microbatching split.
-        dp_sync: Agreement from a prior dispatch over this same batch, to reuse.
+        dp_sync_state: Agreement from a prior dispatch over this same batch, to reuse.
             Must come from a batch with this same padded `num_tokens` and the
             same `uniform_token_count`; `num_reqs` may differ, as neither
             depends on it. Passing a sync from a different batch is a caller
@@ -265,7 +265,7 @@ def dispatch_cg_and_sync_dp(
         dispatch to reuse. It is None when `dp_size` is 1 or no rank has work.
 
     """
-    reuse_eager = dp_sync is not None and dp_sync.eager
+    reuse_eager = dp_sync_state is not None and dp_sync_state.eager
 
     if need_eager or reuse_eager:
         batch_desc = BatchExecutionDescriptor(
@@ -282,32 +282,32 @@ def dispatch_cg_and_sync_dp(
         batch_desc = cudagraph_manager.dispatch(
             num_reqs,
             num_tokens,
-            dp_sync.uniform_token_count if dp_sync is not None else uniform_token_count,
+            dp_sync_state.uniform_token_count if dp_sync_state else uniform_token_count,
             num_active_loras=num_active_loras,
             max_query_len=max_query_len,
         )
 
     if dp_size == 1:
         return batch_desc, None
 
-    if dp_sync is not None:
-        assert dp_sync.num_tokens_across_dp[dp_rank] == num_tokens, (
+    if dp_sync_state is not None:
+        assert dp_sync_state.num_tokens_across_dp[dp_rank] == num_tokens, (
             "reusing a DP sync taken over a different batch"
         )
         assert (
-            dp_sync.uniform_token_count is None
-            or uniform_token_count == dp_sync.uniform_token_count
+            dp_sync_state.uniform_token_count is None
+            or uniform_token_count == dp_sync_state.uniform_token_count
         ), "reusing a DP sync taken over a different batch"
-        if not dp_sync.eager and batch_desc.num_tokens != num_tokens:
+        if not dp_sync_state.eager and batch_desc.num_tokens != num_tokens:
             # Capture sizes can differ between managers, so this one may
             # pad further. Every rank pads alike, so report what will run.
-            dp_sync = replace(
-                dp_sync,
+            dp_sync_state = replace(
+                dp_sync_state,
                 num_tokens_across_dp=torch.full_like(
-                    dp_sync.num_tokens_across_dp, batch_desc.num_tokens
+                    dp_sync_state.num_tokens_across_dp, batch_desc.num_tokens
                 ),
             )
-        return batch_desc, dp_sync
+        return batch_desc, dp_sync_state
 
     return sync_cudagraph_and_dp_padding(
         cudagraph_manager,

vllm/v1/worker/gpu/kv_connector.py

@@ -24,12 +24,18 @@
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.worker.gpu.input_batch import InputBatch
 
 
 class KVConnector:
     """KVConnector interface used by GPUModelRunner."""
 
-    def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None:
+    def pre_forward(
+        self,
+        scheduler_output: "SchedulerOutput",
+        input_batch: "InputBatch | None" = None,
+        attn_metadata: dict[str, Any] | None = None,
+    ) -> None:
         pass
 
     def finish_forward(self) -> None:
@@ -61,22 +67,49 @@ def __init__(
         self._pending_load_kwargs: dict[str, Any] | None = None
         self._disabled = False
 
-    def pre_forward(self, scheduler_output: "SchedulerOutput", **kwargs: Any) -> None:
+    def pre_forward(
+        self,
+        scheduler_output: "SchedulerOutput",
+        input_batch: "InputBatch | None" = None,
+        attn_metadata: dict[str, Any] | None = None,
+    ) -> None:
         if self._disabled:
             return
 
         kv_connector_metadata = scheduler_output.kv_connector_metadata
         assert kv_connector_metadata is not None
         self.kv_connector.handle_preemptions(kv_connector_metadata)
         self.kv_connector.bind_connector_metadata(kv_connector_metadata)
-        self._pending_load_kwargs = kwargs
+        self._pending_load_kwargs = self._build_load_kwargs(
+            scheduler_output, input_batch, attn_metadata
+        )
 
         if scheduler_output.has_sync_kv_loads:
             # Sync loads need to run before this step's forward.
             self._start_load_kv()
         # Otherwise start the async load after forward to keep submission
         # off the critical path.
 
+    @staticmethod
+    def _build_load_kwargs(
+        scheduler_output: "SchedulerOutput",
+        input_batch: "InputBatch | None",
+        attn_metadata: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        """Keyword args for the connector's start_load_kv()."""
+        if input_batch is None:
+            # No forward this step.
+            return {}
+        load_kwargs: dict[str, Any] = dict(
+            scheduler_output=scheduler_output,
+            request_state_indices=input_batch.idx_mapping,
+            request_ids=input_batch.req_ids,
+            num_tokens=input_batch.num_tokens,
+        )
+        if attn_metadata is not None:
+            load_kwargs["attn_metadata"] = attn_metadata
+        return load_kwargs
+
     def _start_load_kv(self) -> None:
         load_kwargs = self._pending_load_kwargs
         assert load_kwargs is not None