diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
index 0cc78f3371..b542388952 100644
--- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
@@ -166,9 +166,7 @@ def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
         out[num_reqs:].zero_()
         return out
 
-    with patch.object(
-        base_speculator, "maybe_prepare_dcp_local_seq_lens", fake_prepare
-    ):
+    with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
         captured = _run_build(
             fake,
             num_reqs=3,
diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py
index 13a92d2c0f..920277cb81 100644
--- a/tests/v1/worker/test_gpu_input_batch_v2.py
+++ b/tests/v1/worker/test_gpu_input_batch_v2.py
@@ -81,7 +81,7 @@ def test_make_dummy_padding_controls_moe_routing(monkeypatch, is_padding: bool):
     assert bool((topk_ids >= 0).all()) is not is_padding
 
 
-def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
+def test_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
     """The batch must view the caller-owned buffer, sliced to padded length.
 
     Runtime (and capture) paths all funnel through this helper; attention
@@ -115,7 +115,7 @@ def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
             return fake_kernel
 
     monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", FakeKernel())
-    batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
+    batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens(
         buffers.dcp_local_seq_lens,
         batch.seq_lens,
         batch.num_reqs,
@@ -130,37 +130,10 @@ def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
     assert batch.dcp_local_seq_lens.tolist() == [1, 2, 0, 0]
 
 
-def test_maybe_prepare_dcp_local_seq_lens_clears_stale_metadata(monkeypatch):
-    """dcp_size == 1 resets the field to None instead of leaving a leftover.
-
-    DCP is toggled per deployment, and batches are recycled; a value from an
-    earlier DCP batch would otherwise be consumed as if it were current.
-    """
-    buffers = InputBuffers(max_num_reqs=2, max_num_tokens=2, device=torch.device("cpu"))
-    batch = InputBatch.make_dummy(1, 1, buffers)
-    batch.dcp_local_seq_lens = buffers.dcp_local_seq_lens[:1]
-
-    def fail_if_called(*args, **kwargs):
-        raise AssertionError("kernel must not run with dcp_size == 1")
-
-    monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", fail_if_called)
-    batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
-        buffers.dcp_local_seq_lens,
-        batch.seq_lens,
-        batch.num_reqs,
-        1,
-        0,
-        1,
-        num_reqs_padded=batch.num_reqs_after_padding,
-    )
-
-    assert batch.dcp_local_seq_lens is None
-
-
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
 @pytest.mark.parametrize("dcp_size", [2, 4])
 @pytest.mark.parametrize("cp_interleave", [1, 16])
-def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
+def test_prepare_dcp_local_seq_lens_matches_reference(
     dcp_size: int, cp_interleave: int
 ):
     """Every rank's local lengths must equal the reference torch formula."""
@@ -172,7 +145,7 @@ def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
 
     for dcp_rank in range(dcp_size):
         buffers.dcp_local_seq_lens.fill_(-1)
-        batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
+        batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens(
             buffers.dcp_local_seq_lens,
             batch.seq_lens,
             batch.num_reqs,
diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py
index 2bd7be2cfa..547e985c17 100644
--- a/tests/v1/worker/test_gpu_pcp_manager.py
+++ b/tests/v1/worker/test_gpu_pcp_manager.py
@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import replace
 from types import SimpleNamespace
-from unittest.mock import MagicMock
 
 import numpy as np
 import pytest
@@ -170,38 +169,44 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
         )
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels")
 @pytest.mark.parametrize(
     ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"),
     [
-        (CUDAGraphMode.NONE, 4, None, None),
-        (CUDAGraphMode.PIECEWISE, None, 8, None),
+        (CUDAGraphMode.NONE, 4, 2, 2),
+        (CUDAGraphMode.PIECEWISE, None, 8, 2),
         (CUDAGraphMode.FULL, 4, 8, 4),
     ],
 )
 def test_partition_padding_is_derived_from_batch_descriptor(
     cg_mode, num_reqs, expected_tokens, expected_reqs
 ):
-    manager = MagicMock()
-    input_batch = MagicMock()
-    manager.partition_batch.return_value = input_batch
-    batch_desc = BatchExecutionDescriptor(
-        cg_mode=cg_mode,
-        num_tokens=8,
-        num_reqs=num_reqs,
+    """partition_batch derives graph padding from the batch descriptor: no
+    padding without graphs, token padding only for PIECEWISE, and request
+    padding on top for FULL."""
+    device = torch.device("cuda:0")
+    manager = PCPManager(
+        pcp_world_size=2,
+        pcp_rank=0,
+        device=device,
+        max_num_reqs=4,
+        max_num_tokens=8,
+    )
+    global_batch = _make_global_decode_batch(
+        [16, 24], InputBuffers(4, 8, device), device
     )
 
-    result = pcp_manager_module.maybe_partition_pcp_batch(
-        manager,
-        input_batch,
-        batch_desc,
+    local_batch = manager.partition_batch(
+        global_batch,
+        BatchExecutionDescriptor(
+            cg_mode=cg_mode,
+            num_tokens=8,
+            num_reqs=num_reqs,
+        ),
     )
 
-    assert result is input_batch
-    manager.partition_batch.assert_called_once_with(
-        input_batch,
-        padded_num_tokens=expected_tokens,
-        padded_num_reqs=expected_reqs,
-    )
+    assert local_batch.num_tokens_after_padding == expected_tokens
+    assert local_batch.num_reqs_after_padding == expected_reqs
 
 
 def test_capture_uses_pcp_persistent_inputs():
@@ -472,8 +477,11 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
 
     local_batch = manager.partition_batch(
         global_batch,
-        padded_num_tokens=4,
-        padded_num_reqs=4,
+        BatchExecutionDescriptor(
+            cg_mode=CUDAGraphMode.FULL,
+            num_tokens=4,
+            num_reqs=4,
+        ),
     )
 
     assert local_batch.dcp_local_seq_lens is None
@@ -489,7 +497,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
 
     # What execute_model does next: derive DCP metadata from the final batch
     # on the PCP-owned buffers.
-    local_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
+    local_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens(
         manager.input_buffers.dcp_local_seq_lens,
         local_batch.seq_lens,
         local_batch.num_reqs,
diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py
index 9d3a92e623..ebf334e45d 100644
--- a/tests/v1/worker/test_gpu_ubatch_slicing.py
+++ b/tests/v1/worker/test_gpu_ubatch_slicing.py
@@ -756,7 +756,7 @@ def test_microbatches_recompute_dcp_lens_from_truncated_seq_lens(dcp_rank: int):
         [1, 1, 10, 1, 1], [64, 96, 512, 32, 48]
     )
     # execute_model has populated the merged batch's DCP metadata already.
-    input_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
+    input_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens(
         buffers.dcp_local_seq_lens,
         input_batch.seq_lens,
         input_batch.num_reqs,
diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py
index 6563e30c34..096d9b99aa 100644
--- a/vllm/v1/worker/gpu/block_table.py
+++ b/vllm/v1/worker/gpu/block_table.py
@@ -59,8 +59,7 @@ class BlockTables:
             self.block_tables.append(block_table)
 
         self.num_blocks = UvaBackedTensor(
-            (self.num_kv_cache_groups, self.max_num_reqs),
-            dtype=torch.int32,
+            (self.num_kv_cache_groups, self.max_num_reqs), dtype=torch.int32
         )
         self.fused_writer: FusedStagedWriter | None = None
         if self.num_kv_cache_groups > 1:
diff --git a/vllm/v1/worker/gpu/cp_utils.py b/vllm/v1/worker/gpu/cp_utils.py
index 470271a95b..1c906f4214 100644
--- a/vllm/v1/worker/gpu/cp_utils.py
+++ b/vllm/v1/worker/gpu/cp_utils.py
@@ -5,7 +5,7 @@ import torch
 from vllm.triton_utils import tl, triton
 
 
-def maybe_prepare_dcp_local_seq_lens(
+def prepare_dcp_local_seq_lens(
     dcp_local_seq_lens: torch.Tensor,
     seq_lens: torch.Tensor,
     num_reqs: int,
@@ -16,8 +16,7 @@ def maybe_prepare_dcp_local_seq_lens(
     num_reqs_padded: int | None = None,
 ) -> torch.Tensor | None:
     """Populate caller-owned storage and return its padded view, or None without DCP."""
-    if dcp_size == 1:
-        return None
+    assert dcp_size > 1
 
     max_num_reqs = dcp_local_seq_lens.shape[0]
     BLOCK_SIZE = 128
diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py
index ddab7799d4..6febb1e15b 100644
--- a/vllm/v1/worker/gpu/cudagraph_utils.py
+++ b/vllm/v1/worker/gpu/cudagraph_utils.py
@@ -40,7 +40,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches
@@ -643,8 +643,7 @@ class ModelCudaGraphManager(CudaGraphManager):
                     self.intermediate_tensors[k][:num_tokens] = v
 
         def create_forward_fn(
-            desc: BatchExecutionDescriptor,
-            warmup: bool,
+            desc: BatchExecutionDescriptor, warmup: bool
         ) -> Callable[[CUDAGraphMode], None]:
             num_tokens = desc.num_tokens
             num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
@@ -789,15 +788,16 @@ def prepare_inputs_to_capture(
         slot_mappings, kv_cache_config
     )
 
-    input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
-        input_buffers.dcp_local_seq_lens,
-        input_batch.seq_lens,
-        input_batch.num_reqs,
-        block_tables.cp_size,
-        block_tables.cp_rank,
-        block_tables.cp_interleave,
-        num_reqs_padded=input_batch.num_reqs_after_padding,
-    )
+    if block_tables.cp_size > 1:
+        input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
+            input_buffers.dcp_local_seq_lens,
+            input_batch.seq_lens,
+            input_batch.num_reqs,
+            block_tables.cp_size,
+            block_tables.cp_rank,
+            block_tables.cp_interleave,
+            num_reqs_padded=input_batch.num_reqs_after_padding,
+        )
 
     # NOTE(woosuk): Attention metadata is required not just by standard attention
     # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,
diff --git a/vllm/v1/worker/gpu/ec_connector.py b/vllm/v1/worker/gpu/ec_connector.py
index c81fd7c11e..72f09f74e7 100644
--- a/vllm/v1/worker/gpu/ec_connector.py
+++ b/vllm/v1/worker/gpu/ec_connector.py
@@ -29,10 +29,7 @@ class ECConnector:
     ) -> Generator[ECConnectorOutput | None, None, None]:
         yield None
 
-    def no_forward(
-        self,
-        scheduler_output: "SchedulerOutput",
-    ) -> ModelRunnerOutput:
+    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
         return EMPTY_MODEL_RUNNER_OUTPUT
 
 
@@ -84,10 +81,7 @@ class ActiveECConnector(ECConnector):
             output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta()
             ec_connector.clear_connector_metadata()
 
-    def no_forward(
-        self,
-        scheduler_output: "SchedulerOutput",
-    ) -> ModelRunnerOutput:
+    def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
         # EC send/recv even if no work to do.
         with self.maybe_get_output(scheduler_output) as ec_connector_output:
             pass
@@ -99,8 +93,7 @@ NO_OP_EC_CONNECTOR = ECConnector()
 
 
 def get_ec_connector(
-    vllm_config: VllmConfig,
-    encoder_cache: "EncoderCache | None",
+    vllm_config: VllmConfig, encoder_cache: "EncoderCache | None"
 ) -> ECConnector:
     if (
         not has_ec_transfer()
diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py
index 24c0f0c4e1..d1bbd3338e 100644
--- a/vllm/v1/worker/gpu/eplb_utils.py
+++ b/vllm/v1/worker/gpu/eplb_utils.py
@@ -77,19 +77,13 @@ class EPLBController:
         assert speculative_config is not None
         assert speculative_config.draft_model_config is not None
         assert self.state is not None
-        self.state.add_model(
-            draft_moe_model,
-            speculative_config.draft_model_config,
-        )
+        self.state.add_model(draft_moe_model, speculative_config.draft_model_config)
         speculator.set_eplb_state(self.state)
         self._has_registered_models = True
         return True
 
     def maybe_register_model(
-        self,
-        model: nn.Module,
-        model_config: Any,
-        load_dummy_weights: bool,
+        self, model: nn.Module, model_config: Any, load_dummy_weights: bool
     ) -> bool:
         if not self.parallel_config.enable_eplb or load_dummy_weights:
             return False
@@ -110,11 +104,7 @@ class EPLBController:
         if eplb_models_added and self.state is not None and self.state.is_async:
             self.state.start_async_loop()
 
-    def step(
-        self,
-        is_dummy: bool = False,
-        is_profile: bool = False,
-    ) -> None:
+    def step(self, is_dummy: bool = False, is_profile: bool = False) -> None:
         if (
             not self.parallel_config.enable_eplb
             or self.suppressed
@@ -123,11 +113,8 @@ class EPLBController:
         ):
             return
 
-        self.state.step(
-            is_dummy,
-            is_profile,
-            log_stats=self.parallel_config.eplb_config.log_balancedness,
-        )
+        log_stats = self.parallel_config.eplb_config.log_balancedness
+        self.state.step(is_dummy, is_profile, log_stats=log_stats)
 
     def prepare_forward(
         self,
diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py
index 32c083e459..66b183f192 100644
--- a/vllm/v1/worker/gpu/input_batch.py
+++ b/vllm/v1/worker/gpu/input_batch.py
@@ -17,12 +17,7 @@ if TYPE_CHECKING:
 
 
 class InputBuffers:
-    def __init__(
-        self,
-        max_num_reqs: int,
-        max_num_tokens: int,
-        device: torch.device,
-    ):
+    def __init__(self, max_num_reqs: int, max_num_tokens: int, device: torch.device):
         self.max_num_reqs = max_num_reqs
         self.max_num_tokens = max_num_tokens
         self.device = device
@@ -474,11 +469,7 @@ def combine_sampled_and_draft_tokens(
     num_reqs = idx_mapping.shape[0]
     num_speculative_steps = draft_tokens.shape[-1]
 
-    logits_indices = torch.empty(
-        num_logits,
-        dtype=torch.int64,
-        device=input_ids.device,
-    )
+    logits_indices = torch.empty(num_logits, dtype=torch.int64, device=input_ids.device)
     _combine_sampled_and_draft_tokens_kernel[(num_reqs,)](
         input_ids,
         idx_mapping,
@@ -711,9 +702,7 @@ def expand_idx_mapping(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     num_reqs = idx_mapping.shape[0]
     expanded_idx_mapping = idx_mapping.new_empty(total_num_logits)
-    expanded_local_pos = torch.empty(
-        total_num_logits, dtype=torch.int32, device=idx_mapping.device
-    )
+    expanded_local_pos = idx_mapping.new_empty(total_num_logits)
     _expand_idx_mapping_kernel[(num_reqs,)](
         idx_mapping,
         expanded_idx_mapping,
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index eb2a85b168..70f601a4c0 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -36,10 +36,7 @@ from vllm.distributed.aux_output_connector.worker import (
     AuxOutputWorkerConnector,
     get_aux_output_connector,
 )
-from vllm.distributed.parallel_state import (
-    get_dcp_group,
-    get_pp_group,
-)
+from vllm.distributed.parallel_state import get_dcp_group, get_pp_group
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
@@ -97,7 +94,7 @@ from vllm.v1.worker.gpu.attn_utils import (
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import (
     BatchExecutionDescriptor,
     ModelCudaGraphManager,
@@ -494,8 +491,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     ),
                 )
             self.prompt_logprobs_worker = PromptLogprobsWorker(
-                self.max_num_reqs,
-                logprobs_mode=self.model_config.logprobs_mode,
+                self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode
             )
             self.structured_outputs_worker = StructuredOutputsWorker(
                 max_num_logits=self.max_num_reqs * self.decode_query_len,
@@ -508,9 +504,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         if self.is_pooling_model and self.is_last_pp_rank:
             self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
         eplb_models_added |= self.eplb.maybe_register_model(
-            self.model,
-            self.model_config,
-            load_dummy_weights,
+            self.model, self.model_config, load_dummy_weights
         )
         self.eplb.maybe_start_async_loop(eplb_models_added)
 
@@ -869,14 +863,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.step_timing.drafter_start()
             mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
             if self.speculator.supports_mm_inputs:
-                mm_inputs = (
-                    [],
-                    torch.zeros(
-                        input_batch.num_tokens,
-                        dtype=torch.bool,
-                        device="cpu",
-                    ),
+                all_false = torch.zeros(
+                    input_batch.num_tokens, dtype=torch.bool, device="cpu"
                 )
+                mm_inputs = [], all_false
 
             # Let the target override the hidden state fed to the drafter
             # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
@@ -945,9 +935,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             mm_config = self.model_config.multimodal_config
             if mm_config is not None and not mm_config.skip_mm_profiling:
                 mm_budget = MultiModalBudget(
-                    self.vllm_config,
-                    self.mm_registry,
-                    enable_cache=False,
+                    self.vllm_config, self.mm_registry, enable_cache=False
                 )
                 dummy_mm_inputs = mm_budget.get_dummy_encoder_profile_inputs()
                 self.model_state.encoder_runner.profile_encoder_cache(
@@ -1294,7 +1282,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         req_ids = batch_req_state.req_ids
         num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
         idx_mapping_np = batch_req_state.idx_mapping_np
-        idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device)
+        idx_mapping = async_tensor_h2d(
+            idx_mapping_np, device=self.device, dtype=torch.int32
+        )
         num_reqs = len(req_ids)
 
         # Get the number of draft tokens for each request.
@@ -1471,11 +1461,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 else None
             ),
         )
-        input_batch = pcp.maybe_partition_pcp_batch(
-            self.pcp_manager,
-            input_batch,
-            batch_desc,
-        )
+        if self.pcp_manager is not None:
+            input_batch = self.pcp_manager.partition_batch(input_batch, batch_desc)
         return input_batch
 
     def prepare_attn(
@@ -1486,8 +1473,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
         block_tables = self.block_tables.gather_block_tables(
-            input_batch.idx_mapping,
-            num_reqs_padded=input_batch.num_reqs_after_padding,
+            input_batch.idx_mapping, num_reqs_padded=input_batch.num_reqs_after_padding
         )
         # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
         # Kernel pads beyond num_tokens with PAD_SLOT_ID.
@@ -1506,15 +1492,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs)
         if valid_state_slots:
             state_slots = torch.arange(
-                1,
-                input_batch.num_reqs + 1,
-                dtype=torch.int32,
-                device=self.device,
+                1, input_batch.num_reqs + 1, dtype=torch.int32, device=self.device
             )
             for block_table in block_tables:
                 block_table[:, 0].copy_(state_slots)
-        slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(
-            self.pcp_manager, self.block_tables, input_batch.num_tokens
+        slot_mappings = block_table_provider.get_dummy_slot_mappings(
+            input_batch.num_tokens
         )
         return block_tables, slot_mappings
 
@@ -1607,11 +1590,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         query_start_loc: torch.Tensor | None = None,
     ) -> None:
         # Update the number of computed tokens.
+        output_bin_counts = None
         if self.is_last_pp_rank:
             assert self.sampler is not None
             output_bin_counts = self.sampler.penalties_state.output_bin_counts
-        else:
-            output_bin_counts = None
         post_update(
             idx_mapping,
             self.req_states.num_computed_tokens.gpu,
@@ -1769,32 +1751,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     "Attention metadata must be prepared for dummy runs when using "
                     "FULL cudagraph mode."
                 )
-                block_tables = None
-                slot_mappings = None
+                block_tables, slot_mappings = None, None
             else:
                 block_tables, slot_mappings = self.prepare_dummy_attn(
                     input_batch, valid_dummy_state_slots
                 )
-            if not skip_attn_for_dummy_run and context_len:
-                assert block_tables is not None
-                set_dummy_context(
-                    input_batch,
-                    self.block_tables,
-                    context_len,
-                    self.kv_cache_config.num_blocks,
-                    self.max_model_len,
-                    input_block_tables=block_tables,
-                )
+                if context_len:
+                    set_dummy_context(
+                        input_batch,
+                        self.block_tables,
+                        context_len,
+                        self.kv_cache_config.num_blocks,
+                        self.max_model_len,
+                        input_block_tables=block_tables,
+                    )
 
-        attn_metadata = None
-        slot_mappings_by_layer = None
-        if not (dummy_run and skip_attn_for_dummy_run):
-            input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
-                (
-                    self.pcp_manager.input_buffers
-                    if self.pcp_manager is not None
-                    else self.input_buffers
-                ).dcp_local_seq_lens,
+        if self.dcp_size > 1 and not (dummy_run and skip_attn_for_dummy_run):
+            input_buffers = self.input_buffers
+            if self.pcp_manager is not None:
+                input_buffers = self.pcp_manager.input_buffers
+            input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
+                input_buffers.dcp_local_seq_lens,
                 input_batch.seq_lens,
                 input_batch.num_reqs,
                 self.dcp_size,
@@ -1802,6 +1779,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 self.cp_interleave,
                 num_reqs_padded=input_batch.num_reqs_after_padding,
             )
+
+        attn_metadata = None
+        slot_mappings_by_layer = None
         ubatch_state: UBatchState | None = None
         if batch_desc.num_ubatches > 1:
             assert self.ubatch_runner is not None
@@ -2046,16 +2026,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
 
         # Last rank: sample tokens
+        assert hidden_states is not None
         draft_hidden_states = hidden_states
-        assert draft_hidden_states is not None
-        hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(
-            self.pcp_manager, hidden_states, input_batch
-        )
-        if self.pcp_manager is not None and aux_hidden_states is not None:
-            aux_hidden_states = [
-                self.pcp_manager.restore_hidden_states(states)
-                for states in aux_hidden_states
-            ]
+        if self.pcp_manager is not None:
+            hidden_states, aux_hidden_states, input_batch = (
+                self.pcp_manager.restore_for_sampling(hidden_states, aux_hidden_states)
+            )
 
         sampler_output, num_sampled, num_rejected = self.sample(
             hidden_states, input_batch, grammar_output
@@ -2305,13 +2281,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.eplb.suppressed = suppressed
 
     def setup_eplb_from_mapping(
-        self,
-        expanded_physical_to_logical: torch.Tensor,
+        self, expanded_physical_to_logical: torch.Tensor
     ) -> None:
         self.eplb.setup_from_mapping(
-            self.model,
-            self.model_config,
-            expanded_physical_to_logical,
+            self.model, self.model_config, expanded_physical_to_logical
         )
 
     ########### EPLB methods end ###########
diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py
index a61384ed7f..bebb9cd41b 100644
--- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py
+++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py
@@ -27,9 +27,7 @@ class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata):
     encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]]
 
     def get_extra_common_attn_kwargs(
-        self,
-        kv_cache_group_id: int,
-        num_reqs: int,
+        self, kv_cache_group_id: int, num_reqs: int
     ) -> dict[str, Any]:
         encoder_seq_lens = self.encoder_seq_lens.get(kv_cache_group_id)
         if encoder_seq_lens is None:
diff --git a/vllm/v1/worker/gpu/model_states/encoder_only.py b/vllm/v1/worker/gpu/model_states/encoder_only.py
index 19c97e2658..1b23723fdc 100644
--- a/vllm/v1/worker/gpu/model_states/encoder_only.py
+++ b/vllm/v1/worker/gpu/model_states/encoder_only.py
@@ -188,10 +188,7 @@ class EncoderOnlyModelState(DefaultModelState):
         return attn_metadata
 
     def _build_encoder_attn_metadata(
-        self,
-        input_batch: InputBatch,
-        cudagraph_mode: CUDAGraphMode,
-        for_capture: bool,
+        self, input_batch: InputBatch, cudagraph_mode: CUDAGraphMode, for_capture: bool
     ) -> dict[str, Any]:
         if cudagraph_mode == CUDAGraphMode.FULL:
             num_reqs = input_batch.num_reqs_after_padding
diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py
index b8cdd76940..ee3e519a6e 100644
--- a/vllm/v1/worker/gpu/model_states/interface.py
+++ b/vllm/v1/worker/gpu/model_states/interface.py
@@ -28,16 +28,12 @@ class ModelSpecificAttnMetadata:
     """Base class for model-specific attention metadata."""
 
     def get_extra_common_attn_kwargs(
-        self,
-        kv_cache_group_id: int,
-        num_reqs: int,
+        self, kv_cache_group_id: int, num_reqs: int
     ) -> dict[str, Any]:
         return {}
 
     def get_extra_attn_kwargs(
-        self,
-        attn_metadata_builder: Any,
-        num_reqs: int,
+        self, attn_metadata_builder: Any, num_reqs: int
     ) -> dict[str, Any]:
         return {}
 
diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
index d079991907..df527bd891 100644
--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
@@ -43,16 +43,12 @@ class MambaHybridAttnMetadata(ModelSpecificAttnMetadata):
     num_decode_draft_tokens_cpu: torch.Tensor | None = None
 
     def get_extra_common_attn_kwargs(
-        self,
-        kv_cache_group_id: int,
-        num_reqs: int,
+        self, kv_cache_group_id: int, num_reqs: int
     ) -> dict[str, Any]:
         return {"is_prefilling": self.is_prefilling[:num_reqs]}
 
     def get_extra_attn_kwargs(
-        self,
-        attn_metadata_builder: Any,
-        num_reqs: int,
+        self, attn_metadata_builder: Any, num_reqs: int
     ) -> dict[str, Any]:
         if not isinstance(
             attn_metadata_builder,
@@ -338,9 +334,7 @@ class MambaHybridModelState(DefaultModelState):
         )
         if self.recoverssm is not None:
             self.recoverssm.record_step(
-                attn_metadata,
-                attn_groups,
-                for_capture=for_capture,
+                attn_metadata, attn_groups, for_capture=for_capture
             )
         return attn_metadata
 
diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py
index 376e0f2510..cf3941813d 100644
--- a/vllm/v1/worker/gpu/pcp_manager.py
+++ b/vllm/v1/worker/gpu/pcp_manager.py
@@ -192,8 +192,7 @@ class PCPManager:
 
     @staticmethod
     def _reorder_segments(
-        segments: list[RankSegment],
-        is_prefilling: np.ndarray,
+        segments: list[RankSegment], is_prefilling: np.ndarray
     ) -> list[RankSegment]:
         """Order this rank's rows decodes-first, then prefills, canonically."""
 
@@ -218,9 +217,7 @@ class PCPManager:
         return segments
 
     def replicated_requests(
-        self,
-        num_scheduled_tokens: np.ndarray,
-        is_prefilling: np.ndarray,
+        self, num_scheduled_tokens: np.ndarray, is_prefilling: np.ndarray
     ) -> np.ndarray:
         """Per global request, whether every PCP rank gets the whole query."""
         num_chunks = 2 * self.pcp_world_size
@@ -407,10 +404,7 @@ class PCPManager:
         return self._input_buffers
 
     def partition_batch(
-        self,
-        input_batch: InputBatch,
-        padded_num_tokens: int | None = None,
-        padded_num_reqs: int | None = None,
+        self, input_batch: InputBatch, batch_desc: "BatchExecutionDescriptor"
     ) -> InputBatch:
         assert self._input_buffers is not None
         input_buffers = self._input_buffers
@@ -422,6 +416,13 @@ class PCPManager:
         num_computed_tokens = global_batch.num_computed_tokens_np
         is_prefilling = global_batch.is_prefilling_np
 
+        padded_num_tokens = None
+        padded_num_reqs = None
+        if batch_desc.cg_mode != CUDAGraphMode.NONE:
+            padded_num_tokens = batch_desc.num_tokens
+            if batch_desc.cg_mode == CUDAGraphMode.FULL:
+                padded_num_reqs = batch_desc.num_reqs
+
         segments_by_rank, per_rank_num_tokens = self._build_batch_layout(
             num_scheduled_tokens,
             num_computed_tokens,
@@ -442,9 +443,7 @@ class PCPManager:
 
         num_local_reqs = len(local_segments)
         num_reqs_after_padding = self._resolve_num_reqs_after_padding(
-            global_batch,
-            padded_num_reqs,
-            num_local_reqs,
+            global_batch, padded_num_reqs, num_local_reqs
         )
         if num_reqs_after_padding > input_buffers.max_num_reqs:
             raise RuntimeError(
@@ -714,8 +713,7 @@ class PCPManager:
         return self._gathered_kv_slot_mappings[:, : num_tokens * self.pcp_world_size]
 
     def _convert_to_gathered_slot_mappings(
-        self,
-        global_batch_slot_mappings: torch.Tensor,
+        self, global_batch_slot_mappings: torch.Tensor
     ) -> torch.Tensor:
         assert self._padded_gather_idx is not None
         assert self._gathered_kv_write_mask is not None
@@ -754,9 +752,7 @@ class PCPManager:
         return self.draft_prefill_batch or input_buffers
 
     def prepare_draft_prefill(
-        self,
-        input_batch: InputBatch,
-        input_ids: torch.Tensor,
+        self, input_batch: InputBatch, input_ids: torch.Tensor
     ) -> None:
         if input_batch is not self._global_batch or self._local_batch is None:
             return
@@ -789,54 +785,15 @@ class PCPManager:
         return last_hidden_states, hidden_states
 
     def restore_for_sampling(
-        self,
-        hidden_states: torch.Tensor,
-    ) -> tuple[torch.Tensor, InputBatch]:
+        self, hidden_states: torch.Tensor, aux_hidden_states: list[torch.Tensor] | None
+    ) -> tuple[torch.Tensor, list[torch.Tensor] | None, InputBatch]:
         assert self._global_batch is not None
-        return self.restore_hidden_states(hidden_states), self._global_batch
-
-
-def maybe_partition_pcp_batch(
-    manager: PCPManager | None,
-    input_batch: InputBatch,
-    batch_desc: "BatchExecutionDescriptor",
-) -> InputBatch:
-    if manager is None:
-        return input_batch
-
-    padded_num_tokens = None
-    padded_num_reqs = None
-    if batch_desc.cg_mode != CUDAGraphMode.NONE:
-        padded_num_tokens = batch_desc.num_tokens
-        if batch_desc.cg_mode == CUDAGraphMode.FULL:
-            padded_num_reqs = batch_desc.num_reqs
-
-    return manager.partition_batch(
-        input_batch,
-        padded_num_tokens=padded_num_tokens,
-        padded_num_reqs=padded_num_reqs,
-    )
-
-
-def maybe_get_pcp_dummy_slot_mappings(
-    manager: PCPManager | None,
-    block_tables: BlockTables,
-    num_tokens: int,
-) -> torch.Tensor:
-    if manager is None:
-        return block_tables.get_dummy_slot_mappings(num_tokens)
-    return manager.get_dummy_slot_mappings(num_tokens)
-
-
-def maybe_restore_pcp_for_sampling(
-    manager: PCPManager | None,
-    hidden_states: torch.Tensor | None,
-    input_batch: InputBatch,
-) -> tuple[torch.Tensor, InputBatch]:
-    assert hidden_states is not None
-    if manager is None:
-        return hidden_states, input_batch
-    return manager.restore_for_sampling(hidden_states)
+        hidden_states = self.restore_hidden_states(hidden_states)
+        if aux_hidden_states is not None:
+            aux_hidden_states = [
+                self.restore_hidden_states(states) for states in aux_hidden_states
+            ]
+        return hidden_states, aux_hidden_states, self._global_batch
 
 
 def maybe_build_pcp_manager(
diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
index 7e829cf309..f47b4b7cbf 100644
--- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
+++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
@@ -255,9 +255,7 @@ class AdaptiveVerificationManager:
         logger.debug("DSpark cost tables: %s", self.cost_tables)
 
     def record_confidences(
-        self,
-        confidence_probs: torch.Tensor,
-        input_batch: "InputBatch",
+        self, confidence_probs: torch.Tensor, input_batch: "InputBatch"
     ) -> None:
         """Publish this step's raw confidences for the ranking kernel and start
         copying them to the CPU, where a later step's budget reads them."""
@@ -305,9 +303,9 @@ class AdaptiveVerificationManager:
         )
         num_non_draft_tokens = scheduled_tokens - scheduled_drafts
         slots = np.fromiter(
-            (self.req_states.req_id_to_index[req_id] for req_id in req_ids),
+            map(self.req_states.req_id_to_index.__getitem__, req_ids),
             dtype=np.int32,
-            count=len(req_ids),
+            count=num_reqs,
         )
         stale_confidences = self._stale_confidences[self._stale_idx].np[slots]
         survival_probability = np.cumprod(stale_confidences.astype(np.float64), axis=1)
@@ -469,8 +467,7 @@ def maybe_create_adaptive_verification_manager(
     # The selector rejects unsupported backends, but models that
     # hard-wire theirs (e.g. DeepSeek-V4) never go through it.
     backend = get_query_lens_mismatch_unsupported_backend(
-        attn_groups,
-        checked_layer_names=target_layer_names,
+        attn_groups, checked_layer_names=target_layer_names
     )
     if backend is not None:
         raise ValueError(
diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
index c787efa07f..ca0662076b 100644
--- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
@@ -53,10 +53,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             return
 
         self.inputs_embeds = torch.zeros(
-            self.max_num_tokens,
-            self.hidden_size,
-            dtype=self.dtype,
-            device=self.device,
+            self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=self.device
         )
 
     # Lifecycle hooks for model-specific optimizations. Subclasses override
@@ -140,10 +137,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
 
         # Initialize cudagraph manager for draft decodes (draft positions > 0).
         self.decode_cudagraph_manager = SpeculatorCudaGraphManager(
-            self.vllm_config,
-            self.device,
-            cudagraph_mode,
-            decode_query_len=1,
+            self.vllm_config, self.device, cudagraph_mode, decode_query_len=1
         )
 
     def capture(self) -> None:
@@ -272,8 +266,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
 
         if self.pcp_manager is not None:
             self.pcp_manager.prepare_draft_prefill(
-                input_batch,
-                self.input_buffers.input_ids[:num_tokens_padded],
+                input_batch, self.input_buffers.input_ids[:num_tokens_padded]
             )
             prefill = self.pcp_manager.draft_prefill_batch
             if prefill is not None:
diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
index 725714fcdb..e8f8be7b64 100644
--- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
+++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
@@ -11,7 +11,7 @@ from vllm.v1.worker.gpu.attn_utils import (
     build_slot_mappings_by_layer,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import (
     AttentionState,
     BatchExecutionDescriptor,
@@ -42,15 +42,16 @@ def _prepare_dflash_inputs_to_capture(
     attn_metadata = None
     if not skip_attn:
         query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
-        input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
-            input_buffers.dcp_local_seq_lens,
-            input_batch.seq_lens,
-            input_batch.num_reqs,
-            block_tables.cp_size,
-            block_tables.cp_rank,
-            block_tables.cp_interleave,
-            num_reqs_padded=input_batch.num_reqs_after_padding,
-        )
+        if block_tables.cp_size > 1:
+            input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
+                input_buffers.dcp_local_seq_lens,
+                input_batch.seq_lens,
+                input_batch.num_reqs,
+                block_tables.cp_size,
+                block_tables.cp_rank,
+                block_tables.cp_interleave,
+                num_reqs_padded=input_batch.num_reqs_after_padding,
+            )
         attn_metadata = build_attn_metadata(
             attn_groups=attn_groups,
             num_reqs=num_reqs,
diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
index 411aa7f041..b86f58a450 100644
--- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
@@ -149,10 +149,7 @@ class DSparkSpeculator(DFlashSpeculator):
         )
         if self.draft_watermarker is not None:
             sampled = self.draft_watermarker.sample(
-                logits,
-                sampled,
-                idx_map,
-                self.temperature,
+                logits, sampled, idx_map, self.temperature
             )
         return sampled
 
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
index 427c93b73c..743d66010d 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
@@ -134,8 +134,7 @@ class RejectionSampler:
         if expanded_logits:
             if self.enable_adaptive_verification:
                 # Adaptive verification keeps the true per-request boundaries
-                # on device only; cu_num_logits_np holds the pre-compacted
-                # layout.
+                # on device only; cu_num_logits_np holds the pre-compacted layout.
                 cu_num_generated_tokens = cu_num_logits.clone()
             else:
                 cu_num_generated_tokens = cu_num_logits_np.tolist()
diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py
index eef1704ac1..90e02395e0 100644
--- a/vllm/v1/worker/gpu/spec_decode/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/speculator.py
@@ -26,7 +26,7 @@ from vllm.v1.worker.gpu.attn_utils import (
     init_attn_backend,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.dp_utils import DPSyncState
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
@@ -202,9 +202,7 @@ class DraftModelSpeculator(BaseSpeculator):
 
     @abstractmethod
     def load_draft_model(
-        self,
-        target_model: nn.Module,
-        target_attn_layer_names: set[str],
+        self, target_model: nn.Module, target_attn_layer_names: set[str]
     ) -> nn.Module:
         pass
 
@@ -337,7 +335,7 @@ class DraftModelSpeculator(BaseSpeculator):
         if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
             # Draft steps advance and rewind their own global sequence lengths,
             # so the target model's DCP-local lengths may already be stale.
-            dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
+            dcp_local_seq_lens = prepare_dcp_local_seq_lens(
                 self.input_buffers.dcp_local_seq_lens,
                 self.input_buffers.seq_lens,
                 num_reqs,
@@ -422,10 +420,7 @@ class DraftModelSpeculator(BaseSpeculator):
             )
             if self.draft_watermarker is not None:
                 sampled = self.draft_watermarker.sample(
-                    logits,
-                    sampled,
-                    idx_mapping,
-                    temperature,
+                    logits, sampled, idx_mapping, temperature
                 )
         elif self.use_local_argmax_reduction:
             return self.model.get_top_tokens(hidden_states)
@@ -436,10 +431,7 @@ class DraftModelSpeculator(BaseSpeculator):
         return sampled
 
     def _maybe_predict_acceptance(
-        self,
-        logits: torch.Tensor,
-        idx_mapping: torch.Tensor,
-        draft_step: torch.Tensor,
+        self, logits: torch.Tensor, idx_mapping: torch.Tensor, draft_step: torch.Tensor
     ) -> None:
         if self.acceptance_estimator is not None:
             self.acceptance_estimator.predict(
diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py
index 025c8300fc..db27bf013f 100644
--- a/vllm/v1/worker/gpu/states.py
+++ b/vllm/v1/worker/gpu/states.py
@@ -78,10 +78,7 @@ class RequestState:
         )
 
         self.next_prefill_tokens = torch.zeros(
-            num_prefill_lookahead,
-            self.max_num_reqs,
-            dtype=torch.int32,
-            device=device,
+            num_prefill_lookahead, self.max_num_reqs, dtype=torch.int32, device=device
         )
 
     @property
diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py
index e3cade2d0b..d6f24fcc69 100644
--- a/vllm/v1/worker/gpu/ubatch_utils.py
+++ b/vllm/v1/worker/gpu/ubatch_utils.py
@@ -24,7 +24,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils.torch_utils import current_stream
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
-from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.ubatch_utils import (
@@ -150,7 +150,7 @@ def _slice_input_batch(
     dcp_local_seq_lens = None
     if dcp_size > 1:
         assert dcp_local_seq_lens_buf is not None
-        dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
+        dcp_local_seq_lens = prepare_dcp_local_seq_lens(
             dcp_local_seq_lens_buf,
             seq_lens,
             num_reqs,