--- 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: