launcher/dcp2/diffs/22-v1_worker_gpu_model_runner.py.diff

--- baked-ours/v1/worker/gpu/model_runner.py	2026-09-10 19:56:14.000000000 +0200
+++ baked/v1/worker/gpu/model_runner.py	2026-09-11 10:13:47.000000000 +0200
@@ -70,6 +70,7 @@
     KVCacheConfig,
     MambaSpec,
     UniformTypeKVCacheSpecs,
+    is_dcp_replicated_spec,
 )
 from vllm.v1.outputs import (
     DraftTokenIds,
@@ -170,6 +171,82 @@
 
 logger = init_logger(__name__)
 
+# Stage 7 (DSV41_PREFILL_TIMING=1): per-phase budget of eager prefill steps.
+import os as _dsv41_os  # noqa: E402
+
+_DSV41_PREFILL_TIMING = _dsv41_os.environ.get("DSV41_PREFILL_TIMING", "0") == "1"
+
+
+def _dsv41_prefill_timer_start(input_batch, batch_desc, dummy_run):
+    if (
+        not _DSV41_PREFILL_TIMING
+        or dummy_run
+        or input_batch.num_tokens < 1024
+        or batch_desc.cg_mode == CUDAGraphMode.FULL
+        or torch.cuda.is_current_stream_capturing()
+    ):
+        return None
+    try:  # drop phase time accumulated outside this forward (e.g. the DSpark drafter of the previous step)
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import prefill_acc_take
+
+        prefill_acc_take()
+    except Exception:  # noqa: BLE001
+        pass
+    e0 = torch.cuda.Event(enable_timing=True)
+    e0.record()
+    return e0
+
+
+def _dsv41_prefill_timer_end(e0, input_batch, model_state):
+    if e0 is None:
+        return
+    e1 = torch.cuda.Event(enable_timing=True)
+    e1.record()
+    e1.synchronize()
+    fwd = e0.elapsed_time(e1)
+    try:
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
+            prefill_acc_take,
+        )
+
+        acc = prefill_acc_take()
+    except Exception:  # noqa: BLE001
+        acc = {}
+    g = acc.get("attn_gather", 0.0)
+    k = acc.get("attn_kernel", 0.0)
+    c = acc.get("attn_combine", 0.0)
+    idx = acc.get("indexer", 0.0)
+    cand = acc.get("candidates", 0.0)
+    stager = getattr(model_state, "engram_stager", None)
+    eng = float(getattr(stager, "_t_last_total_ms", 0.0) or 0.0) if stager is not None else 0.0
+    n = int(input_batch.num_tokens)
+    moe = acc.get("moe", 0.0)
+    ar = acc.get("allreduce", 0.0)
+    egf = acc.get("engram_fwd", 0.0)
+    rest = fwd - (g + k + c) - idx
+    logger.info(
+        "Prefill step timing: %d tokens / %d reqs: forward %.0f ms (%.0f tok/s) = "
+        "attention %.0f (gather %.0f, kernel %.0f, combine %.0f) + indexer %.0f "
+        "(incl. candidates %.0f) + rest[MoE/dense/other] %.0f [moe %.0f, allreduce %.0f, "
+        "engram_fwd %.0f, other %.0f]; Engram staging before the step %.0f ms",
+        n,
+        int(input_batch.num_reqs),
+        fwd,
+        1e3 * n / max(fwd, 1e-6),
+        g + k + c,
+        g,
+        k,
+        c,
+        idx,
+        cand,
+        rest,
+        moe,
+        ar,
+        egf,
+        rest - moe - ar - egf,
+        eng,
+    )
+
 
 class GPUModelRunner(LoRAModelRunnerMixin):
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
@@ -558,9 +635,13 @@
         block_sizes = []
         max_num_blocks_per_group = []
         slot_mapping_enabled = []
+        cp_sharded = []
         for kv_cache_group in kv_cache_config.kv_cache_groups:
             spec = kv_cache_group.kv_cache_spec
             block_sizes.append(spec.block_size)
+            # DCP-replicated groups (V4.1 SWA cache, compressor ring) store
+            # every position on every rank; only sharded groups round-robin.
+            cp_sharded.append(not is_dcp_replicated_spec(spec))
             layer_spec = (
                 spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
             )
@@ -623,7 +704,15 @@
             cp_size=self.dcp_size,
             cp_rank=self.dcp_rank,
             cp_interleave=self.cp_interleave,
+            cp_sharded=cp_sharded,
         )
+        if self.use_dcp:
+            logger.info(
+                "DCP%d KV cache groups: sharded=%s (block sizes %s)",
+                self.dcp_size,
+                cp_sharded,
+                block_sizes,
+            )
         self.pcp_manager = pcp.maybe_build_pcp_manager(
             self.vllm_config,
             self.device,
@@ -1752,6 +1841,7 @@
             input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL
         )
         self.step_timing.forward_start()
+        _dsv41_pf = _dsv41_prefill_timer_start(input_batch, batch_desc, dummy_run)
 
         # Run model.
         if batch_desc.cg_mode == CUDAGraphMode.FULL:
@@ -1794,6 +1884,7 @@
                 else:
                     # Eager (NONE): call the raw model directly.
                     model_output = self.model(**model_inputs)
+        _dsv41_prefill_timer_end(_dsv41_pf, input_batch, self.model_state)
 
         if self.is_last_pp_rank:
             if self.use_aux_hidden_state_outputs: