fix(mocker): use logical KV tokens for decode timing [DYN-3851]

lib/mocker/src/common/perf_model.rs

@@ -82,7 +82,7 @@ pub enum PerfModel {
     #[default]
     Polynomial,
     /// Interpolation-based model using profiler data
-    /// Decode axes: (active_kv_tokens, context_length)
+    /// Decode axes: (scheduled logical KV tokens, mean context length)
     Interpolated {
         prefill_interp: Arc<dyn PrefillInterpolator>,
         decode_interp: Arc<dyn DecodeInterpolator>,
@@ -126,7 +126,7 @@ impl PerfModel {
     /// Expected arrays in NPZ file:
     /// - prefill_isl: 1D array of input sequence lengths
     /// - prefill_ttft_ms: 1D array of time to first token in milliseconds
-    /// - decode_active_kv_tokens: 1D array of active KV token counts
+    /// - decode_active_kv_tokens: 1D array of scheduled logical KV token counts
     /// - decode_context_length: 1D array of context lengths
     /// - decode_itl: 2D array of inter-token latencies in milliseconds
     pub fn from_npz(path: &Path) -> Result<Self> {
@@ -248,8 +248,12 @@ impl PerfModel {
 
     /// Predict decode time in milliseconds.
     ///
+    /// `active_kv_tokens` is the sum of logical context lengths in the scheduled
+    /// batch, not the number of distinct physically resident tokens.
+    ///
     /// Callers always pass all parameters; each variant uses what it needs:
-    /// - Polynomial: uses (active_kv_tokens, total_kv_tokens) as utilization
+    /// - Polynomial: uses logical active KV tokens relative to total capacity,
+    ///   clamped to full utilization
     /// - Interpolated: uses (active_kv_tokens, context_length)
     /// - Aiconfigurator: uses (batch_size, context_length)
     pub fn predict_decode_time(
@@ -288,7 +292,7 @@ fn polynomial_prefill_time(batch_size: usize, new_tokens_per_request: usize) ->
 
 fn polynomial_decode_time(active_kv_tokens: usize, total_kv_tokens: usize) -> f64 {
     let active_perc = if total_kv_tokens > 0 {
-        active_kv_tokens as f64 / total_kv_tokens as f64
+        (active_kv_tokens as f64 / total_kv_tokens as f64).min(1.0)
     } else {
         tracing::warn!("Total KV tokens is 0, using 1.0 as capacity");
         1.0

lib/mocker/src/kv_manager/g1_manager.rs

@@ -59,15 +59,6 @@ enum DecodeBlockReservationBackend {
     Native(NativeDecodeBlockReservation),
 }
 
-impl DecodeBlockReservation {
-    pub(crate) fn len(&self) -> usize {
-        match &self.inner {
-            DecodeBlockReservationBackend::Kvbm(reservation) => reservation.len(),
-            DecodeBlockReservationBackend::Native(reservation) => reservation.len(),
-        }
-    }
-}
-
 pub(crate) struct DestinationReservation {
     inner: DestinationReservationBackend,
 }

lib/mocker/src/kv_manager/vllm_backend.rs

@@ -143,12 +143,6 @@ pub(crate) struct NativeDecodeBlockReservation {
     pool: BlockReservation,
 }
 
-impl NativeDecodeBlockReservation {
-    pub(crate) fn len(&self) -> usize {
-        self.pool.fresh_len()
-    }
-}
-
 pub(crate) struct NativeDestinationReservation {
     request_id: Uuid,
     pool: BlockReservation,

lib/mocker/src/scheduler/sglang/decode.rs

@@ -247,7 +247,7 @@ pub(super) fn simulate_decode_step_with_sampler(
         .map(SglangRequest::current_sequence_len)
         .sum();
     let avg_context = total_context / running.len();
-    let active_kv_tokens = total_context.min(config.total_kv_tokens);
+    let active_kv_tokens = total_context;
     let decode_time = config.perf_model.predict_decode_time(
         running.len(),
         active_kv_tokens,

lib/mocker/src/scheduler/vllm/core.rs

@@ -2200,8 +2200,8 @@ impl VllmCore {
         let (decode_time, decode_end_ms) = if self.args.worker_type == WorkerType::Prefill {
             (Duration::ZERO, decode_start_ms)
         } else {
-            let active_kv_tokens = self.kv_manager.num_active_blocks() * self.args.block_size;
             let total_kv_tokens = self.args.num_gpu_blocks * self.args.block_size;
+            let active_kv_tokens = total_length;
             let context_length = total_length / ready.len();
             let decode_ms = self.args.perf_model.predict_decode_time(
                 ready.len(),
@@ -2435,12 +2435,8 @@ impl VllmCore {
         let (decode_time, decode_end_ms) = if self.args.worker_type == WorkerType::Prefill {
             (Duration::ZERO, decode_start_ms)
         } else {
-            let active_kv_tokens = self
-                .kv_manager
-                .num_active_blocks()
-                .saturating_sub(reservation.len())
-                * self.args.block_size;
             let total_kv_tokens = self.args.num_gpu_blocks * self.args.block_size;
+            let active_kv_tokens = total_length;
             let context_length = total_length / ready.len();
             let decode_ms = self.args.perf_model.predict_decode_time(
                 ready.len(),