feat(deepseek-v2-lite): add mixed-request serving gate with per-request KV ownership (#441)

* feat(deepseek-v2-lite): add mixed-request serving gate

* fix(deepseek-v2-lite): avoid decode cache clone in mixed batch path

Author: CAICAIIs

docs/index.md

@@ -67,7 +67,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l
 
 | Path | TL;DR |
 | --- | --- |
-| `models/deepseek-v2-lite/status.md` | DeepSeek-V2-Lite EP2 model status and benchmark ledger: HF/host-staged/NCCL use a committed small correctness case set; #278 adds covered-shape NCCL decode-step CUDA Graph capture/replay/verify evidence; direct batch, HTTP pressure, and vLLM rows remain diagnostic, not production serving parity. |
+| `models/deepseek-v2-lite/status.md` | DeepSeek-V2-Lite EP2 model status and benchmark ledger: HF/host-staged/NCCL use a committed correctness case set; #281 adds the first greedy mixed-request serving gate with per-request decode KV; direct batch, HTTP pressure, and vLLM rows remain diagnostic, not production serving parity. |
 | `models/deepseek-v2-lite/hf-accuracy-gate.md` | DeepSeek-V2-Lite EP2 HF accuracy gate after PR #149/#150/#274: HF `generate(use_cache=true)`, host-staged EP2, and NCCL EP2 are compared across the committed small case set. |
 | `models/deepseek-v2-lite/decode-attribution-gate.md` | DeepSeek-V2-Lite EP2 decode attribution gate for `Hello`/16-token batch sizes 1/4/8: structured JSON with accuracy hashes, timing/counters, separated NCCL all-reduce smoke, and fail-closed full-decode graph probe evidence for the retained batch-1 shape. |
 | `models/deepseek-v2-lite/source-layout.md` | DeepSeek-V2-Lite runtime layout refactor: `runtime.rs` split by responsibility, HF/host-staged/NCCL EP2 E2E exact on 2x RTX 5090; NCCL CUDA Graph smoke remains a diagnostic blocker on that host, independent of the passed correctness gate. |

docs/models/deepseek-v2-lite/status.md

@@ -1,6 +1,6 @@
 # DeepSeek-V2-Lite Status And Benchmark Ledger
 
-> **TL;DR:** DeepSeek-V2-Lite is a feature-gated EP2 correctness and attribution target. HF / host-staged / NCCL exactness is guarded by a committed case set. #278 adds probe-only CUDA Graph evidence for one NCCL batch-1 decode step: capture, instantiate, replay, and token verification all pass. Direct batch, HTTP pressure, and vLLM rows remain diagnostic and do not claim production serving parity.
+> **TL;DR:** DeepSeek-V2-Lite is a feature-gated EP2 correctness and serving-semantics target. HF / host-staged / NCCL exactness is guarded by a committed case set, and #281 adds the first greedy mixed-request serving gate with per-request decode KV ownership. CUDA Graph, direct batch, HTTP pressure, and vLLM rows remain diagnostic or unclaimed as production serving parity.
 
 Last touched: 2026-06
 
@@ -19,7 +19,7 @@ Last touched: 2026-06
 | Device-resident NCCL dense exchange | Available | Issue #276 reuses backend-owned bf16 dense-exchange scratch, clears rank1 zero-send every exchange, removes dense-exchange stream sync from the backend call, and preserves HF / host-staged / NCCL exactness on 2x RTX 5090. |
 | NCCL route-plan replay | Available | Issue #277 builds a token-major host route plan once after top-k routing, replays that plan for NCCL expert launches and device contribution accumulation, keeps route counters visible, and preserves HF / host-staged / NCCL exactness on 2x RTX 5090. This remains the eager NCCL oracle path. |
 | NCCL CUDA Graph readiness | Covered-shape diagnostic | Schema-2 `cuda_graph_readiness` now includes a fail-closed `full_decode_graph_probe`. The 2026-06-20 run reports capture, instantiate, replay, and verification success with `8/8` verified replays for the retained batch-1 NCCL decode step. |
-| Production continuous batching | Not available | The direct diagnostic batch path is not mixed-request HTTP serving. |
+| First mixed-request serving gate | Available | Issue #281 adds greedy-only request admission, FCFS deferral, explicit request-local rejection/error/finish events, and one owned `DecodeCache` per active request. The 2026-06-23 2x RTX 5090 run passed HF / host-staged / NCCL exactness and the mixed-serving E2E for host-staged and NCCL. |
 | vLLM production parity | Not claimed | The vLLM TP2 / TP2+EP2 snapshot below is a runnable comparison from a documented validation environment, not serving parity or a stock-install claim. |
 
 ## Correctness Contract
@@ -32,15 +32,17 @@ The retained correctness gate is deliberately narrow:
 - generation mode: greedy;
 - backends: host-staged and `OPENINFER_DSV2_LITE_EP_BACKEND=nccl`.
 
-The comparison gate must be run on the same model snapshot for HF, host-staged, and NCCL outputs. Same-host comparison remains strict: HF, host-staged, and NCCL must be token-exact and text-exact for every committed case and every diagnostic batch row. Host-staged remains the baseline oracle for NCCL transport changes. The latest retained evidence is the 2026-06-20 2x RTX 5090 case-set run with `case_count=5`, top-level `classification=all_token_text_exact`, and no comparison warnings.
+The comparison gate must be run on the same model snapshot for HF, host-staged, and NCCL outputs. Same-host comparison remains strict: HF, host-staged, and NCCL must be token-exact and text-exact for every committed case and every diagnostic batch row. Host-staged remains the baseline oracle for NCCL transport changes. The latest retained evidence is the 2026-06-23 2x RTX 5090 case-set run with `case_count=5`, top-level `classification=all_token_text_exact`, and no comparison warnings.
+
+The mixed-request serving E2E computes sequential greedy token-id oracles with `DeepSeekV2LiteEp2Generator::generate_greedy`, then submits concurrent requests through `start_engine`. The retained 2026-06-23 run covers same-length mixed prompts for same-position batch decode, different-length mixed prompts for single-row decode fallback, and a valid request submitted beside an invalid `logprobs` request to prove explicit rejection does not poison the valid stream. Host-staged and NCCL both passed the mixed-serving E2E.
 
 The Rust E2E accepts the known HF-confirmed RTX 5090 and A800 hash pairs for this narrow shape, because the same model snapshot has produced different exact greedy text on those hosts while still matching HF on each host. Do not use the static hash pair list as a substitute for the same-host HF comparison when changing accuracy-sensitive code.
 
 ## Benchmark Ledger
 
 ### Direct Same-Prompt Diagnostic Batch
 
-This path is useful for attribution and for avoiding the earlier row-loop TPOT measurement. It is not production continuous batching:
+This path is useful for attribution and for avoiding the earlier row-loop TPOT measurement. It is separate from the first mixed-request serving gate and is not production continuous batching:
 
 - every row uses the same prompt;
 - prefill remains conservative;
@@ -119,6 +121,7 @@ Use these labels consistently:
 | --- | --- | --- |
 | `direct single-row` | In-process batch `1` decode. | HTTP serving throughput. |
 | `direct same-prompt diagnostic batch` | Fixed same-prompt direct batch sizes `1/4/8`. | Production continuous batching or mixed-request scheduling. |
+| `first mixed-request serving gate` | Greedy-only EP2 scheduler path with explicit admission/rejection/deferral, per-request host-side decode `DecodeCache`, active cap `8`, and exact sequential-oracle E2E. | vLLM parity, sparse dispatch, production EP readiness, HTTP throughput scaling, non-greedy sampling, or logprobs support. |
 | `covered NCCL decode graph probe` | Probe-only batch-1 `Hello` decode step captured, instantiated, replayed, and token-verified under CUDA Graph. | Default serving graph coverage, multi-step graph replay, batch `4/8` graph coverage, or performance improvement. |
 | `HTTP concurrency pressure` | `vllm bench serve --max-concurrency N` against an HTTP endpoint. | True OpenInfer batch size unless the engine path proves it. |
 | `vLLM comparison from documented environment` | vLLM TP2 / TP2+EP2 after target-environment package/toolchain fixes. | Stock vLLM install support, OpenInfer serving parity, or production readiness. |
@@ -159,11 +162,11 @@ The next implementation should be chosen from measured evidence:
    - vLLM TP2+EP2 when supported.
    - default vLLM configuration plus a controlled configuration with cache/flag choices recorded.
 
-4. Add real request batching / serving semantics before broader throughput claims.
-   - request admission;
-   - per-request KV ownership;
-   - mixed request state;
-   - decode iterations that carry multiple live `/v1/completions` requests.
+4. Widen the first mixed-request serving gate before broader throughput claims.
+   - keep the fixed EP2 path and exact sequential oracle until a wider oracle replaces it;
+   - keep greedy-only admission explicit until sampling/logprobs have their own gate;
+   - keep direct same-prompt batch labeled diagnostic;
+   - add HTTP-serving evidence before claiming `/v1/completions` parity or production continuous batching.
 
 5. Keep MoE internals readable.
    - routing, dispatch, expert execution, and combine should remain distinguishable in code and attribution;

openinfer-deepseek-v2-lite/src/engine.rs

@@ -1,164 +1,27 @@
-use std::{
-    path::Path,
-    time::{Instant, SystemTime, UNIX_EPOCH},
-};
+use std::{path::Path, time::Instant};
 
 use anyhow::{Context, Result};
 use log::info;
-use openinfer_engine::engine::{
-    EngineHandle, EngineLoadOptions, FinishReason, GenerateRequest, TokenEvent, TokenSink,
-};
+use openinfer_engine::engine::{EngineHandle, EngineLoadOptions};
 use tokio::sync::mpsc;
 
-use crate::runtime::{DeepSeekV2LiteEp2Generator, GenerationResult};
+use crate::{runtime::DeepSeekV2LiteEp2Generator, scheduler::MixedRequestScheduler};
 
 pub(crate) fn start_engine(model_path: &Path, options: EngineLoadOptions) -> Result<EngineHandle> {
     let started = Instant::now();
     info!("starting DeepSeek-V2-Lite EP2 engine");
-    let mut generator = DeepSeekV2LiteEp2Generator::load(model_path, options)?;
-    let (submit_tx, mut submit_rx) = mpsc::unbounded_channel();
+    let generator = DeepSeekV2LiteEp2Generator::load(model_path, options)?;
+    let servable_len = generator.config().supported_plain_rope_context() as u32;
+    let (submit_tx, submit_rx) = mpsc::unbounded_channel();
 
     let join_handle = std::thread::Builder::new()
         .name("deepseek-v2-lite-ep2".to_string())
-        .spawn(move || {
-            while let Some(req) = submit_rx.blocking_recv() {
-                handle_request(&mut generator, &req);
-            }
-        })
+        .spawn(move || MixedRequestScheduler::new(generator, submit_rx).run())
         .context("spawn DeepSeek-V2-Lite EP=2 engine thread")?;
 
     info!(
         "DeepSeek-V2-Lite EP2 engine started cost {:.2}s",
         started.elapsed().as_secs_f64()
     );
-    Ok(EngineHandle::new_with_join_handle(submit_tx, join_handle))
-}
-
-fn handle_request(generator: &mut DeepSeekV2LiteEp2Generator, req: &GenerateRequest) {
-    let prompt_tokens = req.prompt_tokens.len();
-    let now = unix_time_secs();
-    let _ = req.token_tx.send(TokenEvent::Scheduled {
-        queued_at_unix_s: req.queued_at_unix_s.unwrap_or(now),
-        scheduled_at_unix_s: now,
-        prompt_tokens,
-        cached_tokens: 0,
-    });
-    if req.echo {
-        let _ = req.token_tx.send(TokenEvent::PromptTokens {
-            ids: req.prompt_tokens.clone(),
-            logprobs: vec![None; prompt_tokens],
-        });
-    }
-    if !req.params.is_greedy() {
-        reject_request(
-            req,
-            prompt_tokens,
-            format!(
-                "DeepSeek-V2-Lite EP=2 first gate serves greedy decoding only; requested temperature={}, top_k={}, top_p={}",
-                req.params.temperature, req.params.top_k, req.params.top_p
-            ),
-        );
-        return;
-    }
-    if req.logprobs > 0 {
-        reject_request(
-            req,
-            prompt_tokens,
-            "DeepSeek-V2-Lite EP=2 first gate does not return logprobs yet".to_string(),
-        );
-        return;
-    }
-    if req.max_tokens == 0 {
-        let _ = req.token_tx.send(TokenEvent::Finished {
-            finish_reason: FinishReason::Length,
-            prompt_tokens,
-            completion_tokens: 0,
-        });
-        return;
-    }
-
-    match generator.generate_greedy(&req.prompt_tokens, req.max_tokens, req.params.ignore_eos) {
-        Ok(result) => {
-            emit_generation_result(&req.token_tx, prompt_tokens, &result);
-        }
-        Err(err) => {
-            let _ = req.token_tx.send(TokenEvent::Error {
-                message: err.to_string(),
-                prompt_tokens,
-                completion_tokens: 0,
-            });
-        }
-    }
-}
-
-fn reject_request(req: &GenerateRequest, prompt_tokens: usize, message: String) {
-    let _ = req.token_tx.send(TokenEvent::Rejected {
-        message,
-        prompt_tokens,
-        completion_tokens: 0,
-    });
-}
-
-fn emit_generation_result(token_tx: &TokenSink, prompt_tokens: usize, result: &GenerationResult) {
-    for token in &result.tokens {
-        let _ = token_tx.send(TokenEvent::Token {
-            id: *token,
-            logprob: None,
-        });
-    }
-    let _ = token_tx.send(TokenEvent::Finished {
-        finish_reason: result.finish_reason,
-        prompt_tokens,
-        completion_tokens: result.tokens.len(),
-    });
-}
-
-fn unix_time_secs() -> f64 {
-    SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .map_or(0.0, |duration| duration.as_secs_f64())
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::runtime::GenerationStats;
-
-    #[test]
-    fn stop_generation_streams_tokens_and_stop_finish() {
-        let (tx, mut rx) = TokenSink::standalone();
-
-        emit_generation_result(
-            &tx,
-            3,
-            &GenerationResult {
-                tokens: vec![10, 11],
-                finish_reason: FinishReason::Stop,
-                stats: GenerationStats::default(),
-            },
-        );
-        drop(tx);
-
-        match rx.try_recv().expect("expected first token").1 {
-            TokenEvent::Token { id, .. } => assert_eq!(id, 10),
-            _ => panic!("expected first token event"),
-        }
-        match rx.try_recv().expect("expected second token").1 {
-            TokenEvent::Token { id, .. } => assert_eq!(id, 11),
-            _ => panic!("expected second token event"),
-        }
-        match rx.try_recv().expect("expected finished event").1 {
-            TokenEvent::Finished {
-                finish_reason,
-                prompt_tokens,
-                completion_tokens,
-            } => {
-                assert_eq!(finish_reason, FinishReason::Stop);
-                assert_eq!(prompt_tokens, 3);
-                assert_eq!(completion_tokens, 2);
-            }
-            _ => panic!("expected finished event"),
-        }
-        assert!(rx.try_recv().is_err());
-    }
+    Ok(EngineHandle::new_with_join_handle(submit_tx, join_handle).with_servable_len(servable_len))
 }

openinfer-deepseek-v2-lite/src/host_ops.rs

@@ -4,12 +4,12 @@ use openinfer_core::tensor::{DeviceContext, HiddenStates};
 
 use crate::{Config, device::activate};
 
-#[derive(Clone, Default)]
+#[derive(Default)]
 pub(crate) struct DecodeCache {
     pub(crate) layers: Vec<LayerCache>,
 }
 
-#[derive(Clone, Default)]
+#[derive(Default)]
 pub(crate) struct LayerCache {
     keys: Vec<f32>,
     values: Vec<f32>,
@@ -41,6 +41,27 @@ impl DecodeCache {
                 .collect(),
         }
     }
+
+    pub(crate) fn position(&self, config: &Config) -> Result<usize> {
+        ensure!(
+            self.layers.len() == config.num_hidden_layers,
+            "decode cache layer count mismatch: cache={}, expected={}",
+            self.layers.len(),
+            config.num_hidden_layers
+        );
+        let Some(first) = self.layers.first() else {
+            return Ok(0);
+        };
+        let position = first.len(config);
+        for (layer_idx, layer) in self.layers.iter().enumerate().skip(1) {
+            ensure!(
+                layer.len(config) == position,
+                "decode cache layer {layer_idx} position mismatch: cache_len={}, expected={position}",
+                layer.len(config)
+            );
+        }
+        Ok(position)
+    }
 }
 
 pub(crate) fn normalize_compressed_kv(

openinfer-deepseek-v2-lite/src/lib.rs

@@ -15,6 +15,8 @@ mod nccl_backend;
 #[cfg(feature = "deepseek-v2-lite")]
 mod runtime;
 #[cfg(feature = "deepseek-v2-lite")]
+mod scheduler;
+#[cfg(feature = "deepseek-v2-lite")]
 mod weights;
 
 use std::path::Path;

openinfer-deepseek-v2-lite/src/runtime/generation.rs

@@ -165,6 +165,21 @@ impl DeepSeekV2LiteEp2Generator {
         Ok((result, attribution))
     }
 
+    pub(crate) fn config(&self) -> &Config {
+        &self.config
+    }
+
+    pub(crate) fn new_generation_stats(&self, prompt_tokens: usize) -> GenerationStats {
+        GenerationStats {
+            model_path: self.model_path.clone(),
+            device_ordinals: self.device_ordinals.clone(),
+            ep_backend: self.backend.kind().as_str().to_string(),
+            ep_size: 2,
+            prompt_tokens,
+            ..GenerationStats::default()
+        }
+    }
+
     pub fn generate_greedy_batch_same_prompt_with_timings(
         &mut self,
         prompt_tokens: &[u32],
@@ -230,14 +245,7 @@ impl DeepSeekV2LiteEp2Generator {
         );
 
         let generation_start = Instant::now();
-        let mut stats = GenerationStats {
-            model_path: self.model_path.clone(),
-            device_ordinals: self.device_ordinals.clone(),
-            ep_backend: self.backend.kind().as_str().to_string(),
-            ep_size: 2,
-            prompt_tokens: prompt_tokens.len() * batch_size,
-            ..GenerationStats::default()
-        };
+        let mut stats = self.new_generation_stats(prompt_tokens.len() * batch_size);
         let mut caches: Vec<_> = (0..batch_size)
             .map(|_| DecodeCache::new(&self.config))
             .collect();
@@ -319,14 +327,7 @@ impl DeepSeekV2LiteEp2Generator {
             max_new_tokens
         );
 
-        let mut stats = GenerationStats {
-            model_path: self.model_path.clone(),
-            device_ordinals: self.device_ordinals.clone(),
-            ep_backend: self.backend.kind().as_str().to_string(),
-            ep_size: 2,
-            prompt_tokens: prompt_tokens.len(),
-            ..GenerationStats::default()
-        };
+        let mut stats = self.new_generation_stats(prompt_tokens.len());
 
         let mut cache = DecodeCache::new(&self.config);
         let mut generated = Vec::with_capacity(max_new_tokens);
@@ -370,7 +371,7 @@ impl DeepSeekV2LiteEp2Generator {
         })
     }
 
-    pub(super) fn prefill_next_token(
+    pub(crate) fn prefill_next_token(
         &mut self,
         prompt_tokens: &[u32],
         cache: &mut DecodeCache,
@@ -396,7 +397,7 @@ impl DeepSeekV2LiteEp2Generator {
         )
     }
 
-    pub(super) fn decode_next_token(
+    pub(crate) fn decode_next_token(
         &mut self,
         token: u32,
         position: usize,
@@ -432,7 +433,7 @@ impl DeepSeekV2LiteEp2Generator {
         )
     }
 
-    fn decode_next_tokens_batch(
+    pub(crate) fn decode_next_tokens_batch(
         &mut self,
         tokens: &[u32],
         position: usize,