diff --git a/docs/index.md b/docs/index.md index 8f41a7f65..e9a8e5ff6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -66,7 +66,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. | diff --git a/docs/models/deepseek-v2-lite/status.md b/docs/models/deepseek-v2-lite/status.md index c6ad639e9..2aa320bcd 100644 --- a/docs/models/deepseek-v2-lite/status.md +++ b/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,7 +32,9 @@ 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. @@ -40,7 +42,7 @@ The Rust E2E accepts the known HF-confirmed RTX 5090 and A800 hash pairs for thi ### 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; diff --git a/openinfer-deepseek-v2-lite/src/engine.rs b/openinfer-deepseek-v2-lite/src/engine.rs index 8e374eb58..2fd2d17f5 100644 --- a/openinfer-deepseek-v2-lite/src/engine.rs +++ b/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 { 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)) } diff --git a/openinfer-deepseek-v2-lite/src/host_ops.rs b/openinfer-deepseek-v2-lite/src/host_ops.rs index 4fc3c09e0..a1da74401 100644 --- a/openinfer-deepseek-v2-lite/src/host_ops.rs +++ b/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, } -#[derive(Clone, Default)] +#[derive(Default)] pub(crate) struct LayerCache { keys: Vec, values: Vec, @@ -41,6 +41,27 @@ impl DecodeCache { .collect(), } } + + pub(crate) fn position(&self, config: &Config) -> Result { + 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( diff --git a/openinfer-deepseek-v2-lite/src/lib.rs b/openinfer-deepseek-v2-lite/src/lib.rs index cd33d6600..6a9224493 100644 --- a/openinfer-deepseek-v2-lite/src/lib.rs +++ b/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; diff --git a/openinfer-deepseek-v2-lite/src/runtime/generation.rs b/openinfer-deepseek-v2-lite/src/runtime/generation.rs index 68f6a90c0..b10034348 100644 --- a/openinfer-deepseek-v2-lite/src/runtime/generation.rs +++ b/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, diff --git a/openinfer-deepseek-v2-lite/src/scheduler.rs b/openinfer-deepseek-v2-lite/src/scheduler.rs new file mode 100644 index 000000000..360638538 --- /dev/null +++ b/openinfer-deepseek-v2-lite/src/scheduler.rs @@ -0,0 +1,927 @@ +//! Mixed-request greedy serving for the DeepSeek-V2-Lite EP2 gate. +//! +//! This is the first serving-semantics gate for the model. It keeps one +//! `DecodeCache` per active request, admits only shapes the current runtime can +//! honor exactly, and retires each request independently when validation, +//! disconnect, EOS, length, or request-local decode errors occur. + +use std::{collections::VecDeque, mem}; + +use anyhow::{Result, ensure}; +use openinfer_engine::{ + engine::{FinishReason, GenerateRequest, TokenEvent, TokenSink, unix_now_s}, + sampler::SamplingParams, +}; +use tokio::sync::mpsc; + +use crate::{ + Config, + attribution::DecodeAttributionProfile, + host_ops::DecodeCache, + runtime::{DeepSeekV2LiteEp2Generator, GenerationStats}, +}; + +pub(crate) const DEFAULT_MAX_ACTIVE_REQUESTS: usize = 8; + +pub(crate) struct MixedRequestScheduler { + generator: DeepSeekV2LiteEp2Generator, + submit_rx: mpsc::UnboundedReceiver, + pending: VecDeque, + active: Vec, + max_active_requests: usize, +} + +struct PendingRequest { + request_id: Option, + queued_at_unix_s: Option, + prompt_tokens: Vec, + params: SamplingParams, + max_tokens: usize, + lora_adapter: Option, + token_tx: TokenSink, + logprobs: usize, + echo: bool, +} + +struct ActiveRequestState { + request_id: Option, + token_tx: TokenSink, + prompt_len: usize, + max_tokens: usize, + generated: usize, + last_token: u32, + finish_policy: FinishPolicy, + cache: DecodeCache, + stats: GenerationStats, +} + +#[derive(Clone, Copy)] +struct FinishPolicy { + eos_token_id: u32, + ignore_eos: bool, +} + +struct AdmissionBatch { + admitted: Vec, + rejected: Vec<(PendingRequest, String)>, + finished: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +enum AdmissionDecision { + Admit, + Reject(String), + Finish(FinishReason), +} + +#[derive(Clone, Debug, Eq, PartialEq)] +enum DecodeGrouping { + Empty, + BatchSamePosition { position: usize, rows: usize }, + SingleRows, +} + +impl MixedRequestScheduler { + pub(crate) fn new( + generator: DeepSeekV2LiteEp2Generator, + submit_rx: mpsc::UnboundedReceiver, + ) -> Self { + Self { + generator, + submit_rx, + pending: VecDeque::new(), + active: Vec::new(), + max_active_requests: DEFAULT_MAX_ACTIVE_REQUESTS, + } + } + + pub(crate) fn run(mut self) { + while self.block_until_work() { + self.drain_pending_submissions(); + self.admit_ready_requests(); + if !self.active.is_empty() { + self.decode_round(); + } + } + } + + fn block_until_work(&mut self) -> bool { + if !self.pending.is_empty() || !self.active.is_empty() { + return true; + } + + match self.submit_rx.blocking_recv() { + Some(req) => { + self.pending.push_back(PendingRequest::from(req)); + true + } + None => false, + } + } + + fn drain_pending_submissions(&mut self) { + while let Ok(req) = self.submit_rx.try_recv() { + self.pending.push_back(PendingRequest::from(req)); + } + } + + fn admit_ready_requests(&mut self) { + let supported_context = self.generator.config().supported_plain_rope_context(); + let batch = take_admission_batch( + &mut self.pending, + self.active.len(), + self.max_active_requests, + supported_context, + ); + + for (pending, message) in batch.rejected { + if send_scheduled(&pending) { + let _ = send_prompt_echo(&pending); + let _ = pending.token_tx.send(TokenEvent::Rejected { + message, + prompt_tokens: pending.prompt_tokens.len(), + completion_tokens: 0, + }); + } + } + + for pending in batch.finished { + if send_scheduled(&pending) { + let _ = send_prompt_echo(&pending); + let _ = pending.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: pending.prompt_tokens.len(), + completion_tokens: 0, + }); + } + } + + for pending in batch.admitted { + if self.active.len() >= self.max_active_requests { + self.pending.push_front(pending); + break; + } + if let Some(active) = self.prefill_request(pending) { + self.active.push(active); + } + } + } + + fn prefill_request(&mut self, pending: PendingRequest) -> Option { + let prompt_len = pending.prompt_tokens.len(); + if !send_scheduled(&pending) { + return None; + } + + if !send_prompt_echo(&pending) { + return None; + } + + let mut cache = DecodeCache::new(self.generator.config()); + let mut stats = self.generator.new_generation_stats(prompt_len); + let mut attribution = DecodeAttributionProfile::disabled(); + let next = match self.generator.prefill_next_token( + &pending.prompt_tokens, + &mut cache, + &mut stats, + &mut attribution, + ) { + Ok(token) => token, + Err(err) => { + let _ = pending.token_tx.send(TokenEvent::Error { + message: err.to_string(), + prompt_tokens: prompt_len, + completion_tokens: 0, + }); + return None; + } + }; + + let mut active = ActiveRequestState { + request_id: pending.request_id, + token_tx: pending.token_tx, + prompt_len, + max_tokens: pending.max_tokens, + generated: 0, + last_token: next, + finish_policy: FinishPolicy { + eos_token_id: self.generator.config().eos_token_id, + ignore_eos: pending.params.ignore_eos, + }, + cache, + stats, + }; + + if active.emit_token_or_finish(next) { + return None; + } + Some(active) + } + + fn decode_round(&mut self) { + self.retire_bad_cache_positions(); + let positions: Vec<_> = self + .active + .iter() + .map(ActiveRequestState::next_decode_position) + .collect(); + match decode_grouping_for_positions(&positions) { + DecodeGrouping::Empty => {} + DecodeGrouping::BatchSamePosition { position, rows } if rows > 1 => { + self.decode_batch_round(position); + } + DecodeGrouping::BatchSamePosition { .. } | DecodeGrouping::SingleRows => { + self.decode_single_rows(); + } + } + } + + fn retire_bad_cache_positions(&mut self) { + let config = self.generator.config(); + let mut survivors = Vec::with_capacity(self.active.len()); + for state in self.active.drain(..) { + match state.cache_position(config) { + Ok(()) => survivors.push(state), + Err(message) => state.emit_error(message.to_string()), + } + } + self.active = survivors; + } + + fn decode_batch_round(&mut self, position: usize) { + let tokens: Vec<_> = self.active.iter().map(|state| state.last_token).collect(); + let token_index = self + .active + .iter() + .map(|state| state.generated) + .min() + .unwrap_or(0); + let prompt_tokens = self.active.iter().map(|state| state.prompt_len).sum(); + let mut stats = self.generator.new_generation_stats(prompt_tokens); + let mut attribution = DecodeAttributionProfile::disabled(); + let mut caches: Vec<_> = self + .active + .iter_mut() + .map(|state| mem::take(&mut state.cache)) + .collect(); + let result = self.generator.decode_next_tokens_batch( + &tokens, + position, + &mut caches, + &mut stats, + &mut attribution, + token_index, + ); + + match result { + Ok(next_tokens) if next_tokens.len() == self.active.len() => { + for (state, cache) in self.active.iter_mut().zip(caches) { + state.cache = cache; + } + self.apply_decoded_tokens(next_tokens); + } + // The batched path mutates per-row caches as it advances through the + // model. This gate avoids full-cache rollback clones; a batch decode + // failure is therefore a shared runtime error for the active rows. + Ok(next_tokens) => self.retire_active_batch_error(format!( + "DeepSeek-V2-Lite batched decode returned {} rows for {} active requests", + next_tokens.len(), + self.active.len() + )), + Err(err) => self.retire_active_batch_error(format!( + "DeepSeek-V2-Lite batched decode failed for {} active requests: {err}", + self.active.len() + )), + } + } + + fn decode_single_rows(&mut self) { + let mut survivors = Vec::with_capacity(self.active.len()); + for mut state in self.active.drain(..) { + let token = state.last_token; + let position = state.next_decode_position(); + let token_index = state.generated; + let result = self.generator.decode_next_token( + token, + position, + &mut state.cache, + &mut state.stats, + &mut DecodeAttributionProfile::disabled(), + token_index, + ); + match result { + Ok(next) => { + if !state.emit_token_or_finish(next) { + survivors.push(state); + } + } + Err(err) => state.emit_error(err.to_string()), + } + } + self.active = survivors; + } + + fn apply_decoded_tokens(&mut self, next_tokens: Vec) { + let mut survivors = Vec::with_capacity(self.active.len()); + for (mut state, token) in self.active.drain(..).zip(next_tokens) { + if !state.emit_token_or_finish(token) { + survivors.push(state); + } + } + self.active = survivors; + } + + fn retire_active_batch_error(&mut self, message: String) { + retire_active_requests_with_error(&mut self.active, message); + } +} + +fn retire_active_requests_with_error(active: &mut Vec, message: String) { + for state in active.drain(..) { + state.emit_error(message.clone()); + } +} + +impl From for PendingRequest { + fn from(req: GenerateRequest) -> Self { + Self { + request_id: req.request_id, + queued_at_unix_s: req.queued_at_unix_s, + prompt_tokens: req.prompt_tokens, + params: req.params, + max_tokens: req.max_tokens, + lora_adapter: req.lora_adapter, + token_tx: req.token_tx, + logprobs: req.logprobs, + echo: req.echo, + } + } +} + +impl ActiveRequestState { + fn next_decode_position(&self) -> usize { + self.prompt_len + self.generated - 1 + } + + fn cache_position(&self, config: &Config) -> Result<()> { + let expected = self.next_decode_position(); + let actual = self.cache.position(config)?; + ensure!( + actual == expected, + "DeepSeek-V2-Lite request {:?} cache position mismatch: cache_len={}, expected={expected}", + self.request_id, + actual + ); + Ok(()) + } + + fn emit_token_or_finish(&mut self, token: u32) -> bool { + self.last_token = token; + if !self.finish_policy.ignore_eos && token == self.finish_policy.eos_token_id { + let _ = self.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Stop, + prompt_tokens: self.prompt_len, + completion_tokens: self.generated, + }); + return true; + } + + if self + .token_tx + .send(TokenEvent::Token { + id: token, + logprob: None, + }) + .is_err() + { + return true; + } + self.generated += 1; + + if self.generated == self.max_tokens { + let _ = self.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: self.prompt_len, + completion_tokens: self.generated, + }); + return true; + } + false + } + + fn emit_error(self, message: String) { + let _ = self.token_tx.send(TokenEvent::Error { + message, + prompt_tokens: self.prompt_len, + completion_tokens: self.generated, + }); + } +} + +fn send_scheduled(pending: &PendingRequest) -> bool { + let now = unix_now_s(); + pending + .token_tx + .send(TokenEvent::Scheduled { + queued_at_unix_s: pending.queued_at_unix_s.unwrap_or(now), + scheduled_at_unix_s: now, + prompt_tokens: pending.prompt_tokens.len(), + cached_tokens: 0, + }) + .is_ok() +} + +fn send_prompt_echo(pending: &PendingRequest) -> bool { + if !pending.echo { + return true; + } + pending + .token_tx + .send(TokenEvent::PromptTokens { + ids: pending.prompt_tokens.clone(), + logprobs: vec![None; pending.prompt_tokens.len()], + }) + .is_ok() +} + +fn take_admission_batch( + pending: &mut VecDeque, + active_len: usize, + max_active_requests: usize, + supported_context: usize, +) -> AdmissionBatch { + let mut batch = AdmissionBatch { + admitted: Vec::new(), + rejected: Vec::new(), + finished: Vec::new(), + }; + + while let Some(pending_req) = pending.pop_front() { + let can_admit = active_len + batch.admitted.len() < max_active_requests; + match admission_decision(&pending_req, supported_context) { + AdmissionDecision::Admit if can_admit => batch.admitted.push(pending_req), + AdmissionDecision::Admit => { + pending.push_front(pending_req); + break; + } + AdmissionDecision::Reject(message) => batch.rejected.push((pending_req, message)), + AdmissionDecision::Finish(FinishReason::Length) => batch.finished.push(pending_req), + AdmissionDecision::Finish(reason) => { + batch.rejected.push(( + pending_req, + format!("DeepSeek-V2-Lite unsupported admission finish reason: {reason:?}"), + )); + } + } + } + + batch +} + +fn admission_decision(req: &PendingRequest, supported_context: usize) -> AdmissionDecision { + let prompt_tokens = req.prompt_tokens.len(); + if !req.params.is_greedy() { + return AdmissionDecision::Reject(format!( + "DeepSeek-V2-Lite EP=2 mixed serving gate supports greedy decoding only; requested temperature={}, top_k={}, top_p={}", + req.params.temperature, req.params.top_k, req.params.top_p + )); + } + if req.logprobs > 0 { + return AdmissionDecision::Reject( + "DeepSeek-V2-Lite EP=2 mixed serving gate does not return logprobs yet".to_string(), + ); + } + if req.lora_adapter.is_some() { + return AdmissionDecision::Reject( + "DeepSeek-V2-Lite EP=2 mixed serving gate does not support LoRA adapters".to_string(), + ); + } + if req.prompt_tokens.is_empty() { + return AdmissionDecision::Reject( + "DeepSeek-V2-Lite EP=2 mixed serving gate requires a non-empty prompt".to_string(), + ); + } + if req.max_tokens == 0 { + return AdmissionDecision::Finish(FinishReason::Length); + } + + let Some(requested_context) = prompt_tokens.checked_add(req.max_tokens) else { + return AdmissionDecision::Reject(format!( + "DeepSeek-V2-Lite EP=2 mixed serving gate context length overflow: prompt_tokens={prompt_tokens} max_new_tokens={}", + req.max_tokens + )); + }; + if requested_context > supported_context { + return AdmissionDecision::Reject(format!( + "DeepSeek-V2-Lite EP=2 mixed serving gate supports plain RoPE context <= {supported_context} tokens; requested prompt_tokens={prompt_tokens} max_new_tokens={} total={requested_context}. YaRN rope_scaling long context is not implemented yet.", + req.max_tokens + )); + } + + AdmissionDecision::Admit +} + +fn decode_grouping_for_positions(positions: &[usize]) -> DecodeGrouping { + let Some((&first, rest)) = positions.split_first() else { + return DecodeGrouping::Empty; + }; + if positions.len() > 1 && rest.iter().all(|position| *position == first) { + return DecodeGrouping::BatchSamePosition { + position: first, + rows: positions.len(), + }; + } + DecodeGrouping::SingleRows +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, atomic::AtomicBool}; + + use openinfer_engine::engine::RequestTag; + use openinfer_engine::sampler::SamplingParams; + use tokio::sync::mpsc; + + use super::*; + use crate::config::test_lite_config; + + fn request( + id: &str, + prompt_len: usize, + max_tokens: usize, + ) -> ( + PendingRequest, + openinfer_engine::engine::TokenStreamReceiver, + ) { + let (token_tx, token_rx) = TokenSink::standalone(); + ( + PendingRequest { + request_id: Some(id.to_string()), + queued_at_unix_s: None, + prompt_tokens: vec![1; prompt_len], + params: SamplingParams::default(), + max_tokens, + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + }, + token_rx, + ) + } + + fn recv_event(rx: &mut openinfer_engine::engine::TokenStreamReceiver) -> TokenEvent { + rx.try_recv().expect("expected event").1 + } + + #[test] + fn admission_rejects_unsupported_shapes() { + let context = 16; + + let (mut sampling, _rx) = request("sampling", 1, 1); + sampling.params.temperature = 0.8; + assert!(matches!( + admission_decision(&sampling, context), + AdmissionDecision::Reject(message) if message.contains("greedy") + )); + + let (mut logprobs, _rx) = request("logprobs", 1, 1); + logprobs.logprobs = 1; + assert!(matches!( + admission_decision(&logprobs, context), + AdmissionDecision::Reject(message) if message.contains("logprobs") + )); + + let (mut lora, _rx) = request("lora", 1, 1); + lora.lora_adapter = Some("adapter-a".to_string()); + assert!(matches!( + admission_decision(&lora, context), + AdmissionDecision::Reject(message) if message.contains("LoRA") + )); + + let (empty, _rx) = request("empty", 0, 1); + assert!(matches!( + admission_decision(&empty, context), + AdmissionDecision::Reject(message) if message.contains("non-empty prompt") + )); + + let (zero, _rx) = request("zero", 1, 0); + assert_eq!( + admission_decision(&zero, context), + AdmissionDecision::Finish(FinishReason::Length) + ); + } + + #[test] + fn context_overflow_is_rejected() { + let (req, _rx) = request("too-long", 12, 5); + + assert!(matches!( + admission_decision(&req, 16), + AdmissionDecision::Reject(message) + if message.contains("context") && message.contains("total=17") + )); + } + + #[test] + fn active_cap_defers_in_fcfs_order() { + let mut pending = VecDeque::new(); + pending.push_back(request("first", 2, 1).0); + pending.push_back(request("second", 2, 1).0); + pending.push_back(request("third", 2, 1).0); + + let batch = take_admission_batch(&mut pending, 1, 3, 16); + + assert_eq!( + batch + .admitted + .iter() + .map(|req| req.request_id.as_deref()) + .collect::>(), + vec![Some("first"), Some("second")] + ); + assert!(batch.rejected.is_empty()); + assert!(batch.finished.is_empty()); + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].request_id.as_deref(), Some("third")); + } + + #[test] + fn terminal_requests_do_not_wait_for_active_capacity() { + let mut pending = VecDeque::new(); + pending.push_back(request("zero", 2, 0).0); + let (mut invalid, _rx) = request("invalid", 2, 1); + invalid.logprobs = 1; + pending.push_back(invalid); + pending.push_back(request("valid", 2, 1).0); + + let batch = take_admission_batch(&mut pending, 8, 8, 16); + + assert!(batch.admitted.is_empty()); + assert_eq!(batch.finished.len(), 1); + assert_eq!(batch.finished[0].request_id.as_deref(), Some("zero")); + assert_eq!(batch.rejected.len(), 1); + assert_eq!(batch.rejected[0].0.request_id.as_deref(), Some("invalid")); + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].request_id.as_deref(), Some("valid")); + } + + #[test] + fn invalid_request_does_not_block_later_admission_when_cap_has_room() { + let mut pending = VecDeque::new(); + let (mut invalid, _rx) = request("invalid", 2, 1); + invalid.logprobs = 1; + pending.push_back(invalid); + pending.push_back(request("valid", 2, 1).0); + + let batch = take_admission_batch(&mut pending, 0, 2, 16); + + assert_eq!(batch.rejected.len(), 1); + assert_eq!(batch.rejected[0].0.request_id.as_deref(), Some("invalid")); + assert_eq!(batch.admitted.len(), 1); + assert_eq!(batch.admitted[0].request_id.as_deref(), Some("valid")); + assert!(pending.is_empty()); + } + + #[test] + fn terminal_admission_events_keep_scheduler_contract() { + let (mut zero, mut zero_rx) = request("zero", 2, 0); + zero.echo = true; + assert!(send_scheduled(&zero)); + assert!(send_prompt_echo(&zero)); + let _ = zero.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: zero.prompt_tokens.len(), + completion_tokens: 0, + }); + + assert!(matches!( + recv_event(&mut zero_rx), + TokenEvent::Scheduled { .. } + )); + assert!(matches!( + recv_event(&mut zero_rx), + TokenEvent::PromptTokens { ids, .. } if ids == vec![1, 1] + )); + assert!(matches!( + recv_event(&mut zero_rx), + TokenEvent::Finished { + finish_reason: FinishReason::Length, + completion_tokens: 0, + .. + } + )); + + let (rejected, mut rejected_rx) = request("rejected", 2, 1); + assert!(send_scheduled(&rejected)); + let _ = rejected.token_tx.send(TokenEvent::Rejected { + message: "nope".to_string(), + prompt_tokens: rejected.prompt_tokens.len(), + completion_tokens: 0, + }); + + assert!(matches!( + recv_event(&mut rejected_rx), + TokenEvent::Scheduled { .. } + )); + assert!(matches!( + recv_event(&mut rejected_rx), + TokenEvent::Rejected { + completion_tokens: 0, + .. + } + )); + } + + #[test] + fn eos_retirement_is_independent_per_request() { + let config = test_lite_config(); + let (tx_stop, mut rx_stop) = TokenSink::standalone(); + let (tx_live, mut rx_live) = TokenSink::standalone(); + let mut stop_state = ActiveRequestState { + request_id: Some("stop".to_string()), + token_tx: tx_stop, + prompt_len: 3, + max_tokens: 4, + generated: 1, + last_token: 10, + finish_policy: FinishPolicy { + eos_token_id: config.eos_token_id, + ignore_eos: false, + }, + cache: DecodeCache::new(&config), + stats: GenerationStats::default(), + }; + let mut live_state = ActiveRequestState { + request_id: Some("live".to_string()), + token_tx: tx_live, + prompt_len: 2, + max_tokens: 4, + generated: 1, + last_token: 11, + finish_policy: FinishPolicy { + eos_token_id: config.eos_token_id, + ignore_eos: false, + }, + cache: DecodeCache::new(&config), + stats: GenerationStats::default(), + }; + + assert!(stop_state.emit_token_or_finish(config.eos_token_id)); + assert!(!live_state.emit_token_or_finish(12)); + + match recv_event(&mut rx_stop) { + TokenEvent::Finished { + finish_reason, + completion_tokens, + .. + } => { + assert_eq!(finish_reason, FinishReason::Stop); + assert_eq!(completion_tokens, 1); + } + _ => panic!("EOS request should finish without emitting EOS"), + } + match recv_event(&mut rx_live) { + TokenEvent::Token { id, .. } => assert_eq!(id, 12), + _ => panic!("live request should receive its own token"), + } + assert!(rx_live.try_recv().is_err()); + } + + #[test] + fn cancelled_token_sink_retires_request() { + let config = test_lite_config(); + let (stream_tx, mut stream_rx) = mpsc::unbounded_channel(); + let cancelled = Arc::new(AtomicBool::new(true)); + let sink = TokenSink::new( + RequestTag::from("cancelled"), + stream_tx, + Arc::clone(&cancelled), + ); + let mut state = ActiveRequestState { + request_id: Some("cancelled".to_string()), + token_tx: sink, + prompt_len: 2, + max_tokens: 4, + generated: 1, + last_token: 11, + finish_policy: FinishPolicy { + eos_token_id: config.eos_token_id, + ignore_eos: false, + }, + cache: DecodeCache::new(&config), + stats: GenerationStats::default(), + }; + + assert!(state.emit_token_or_finish(12)); + assert!(stream_rx.try_recv().is_err()); + } + + #[test] + fn closed_token_sink_retires_request() { + let config = test_lite_config(); + let (sink, rx) = TokenSink::standalone(); + drop(rx); + let mut state = ActiveRequestState { + request_id: Some("closed".to_string()), + token_tx: sink, + prompt_len: 2, + max_tokens: 4, + generated: 1, + last_token: 11, + finish_policy: FinishPolicy { + eos_token_id: config.eos_token_id, + ignore_eos: false, + }, + cache: DecodeCache::new(&config), + stats: GenerationStats::default(), + }; + + assert!(state.emit_token_or_finish(12)); + } + + #[test] + fn batch_decode_error_retires_all_active_requests() { + let config = test_lite_config(); + let (first_tx, mut first_rx) = TokenSink::standalone(); + let (second_tx, mut second_rx) = TokenSink::standalone(); + let mut active = vec![ + ActiveRequestState { + request_id: Some("first".to_string()), + token_tx: first_tx, + prompt_len: 3, + max_tokens: 8, + generated: 2, + last_token: 11, + finish_policy: FinishPolicy { + eos_token_id: config.eos_token_id, + ignore_eos: false, + }, + cache: DecodeCache::new(&config), + stats: GenerationStats::default(), + }, + ActiveRequestState { + request_id: Some("second".to_string()), + token_tx: second_tx, + prompt_len: 4, + max_tokens: 8, + generated: 1, + last_token: 12, + finish_policy: FinishPolicy { + eos_token_id: config.eos_token_id, + ignore_eos: false, + }, + cache: DecodeCache::new(&config), + stats: GenerationStats::default(), + }, + ]; + + retire_active_requests_with_error(&mut active, "batch failed".to_string()); + + assert!(active.is_empty()); + match recv_event(&mut first_rx) { + TokenEvent::Error { + message, + prompt_tokens, + completion_tokens, + } => { + assert_eq!(message, "batch failed"); + assert_eq!(prompt_tokens, 3); + assert_eq!(completion_tokens, 2); + } + _ => panic!("first active request should receive batch error"), + } + match recv_event(&mut second_rx) { + TokenEvent::Error { + message, + prompt_tokens, + completion_tokens, + } => { + assert_eq!(message, "batch failed"); + assert_eq!(prompt_tokens, 4); + assert_eq!(completion_tokens, 1); + } + _ => panic!("second active request should receive batch error"), + } + } + + #[test] + fn decode_grouping_batches_only_uniform_positions() { + assert_eq!(decode_grouping_for_positions(&[]), DecodeGrouping::Empty); + assert_eq!( + decode_grouping_for_positions(&[5]), + DecodeGrouping::SingleRows + ); + assert_eq!( + decode_grouping_for_positions(&[7, 7, 7]), + DecodeGrouping::BatchSamePosition { + position: 7, + rows: 3 + } + ); + assert_eq!( + decode_grouping_for_positions(&[7, 8, 7]), + DecodeGrouping::SingleRows + ); + } +} diff --git a/openinfer-deepseek-v2-lite/tests/e2e_ep2.rs b/openinfer-deepseek-v2-lite/tests/e2e_ep2.rs index f5f12963d..f3a09e7f8 100644 --- a/openinfer-deepseek-v2-lite/tests/e2e_ep2.rs +++ b/openinfer-deepseek-v2-lite/tests/e2e_ep2.rs @@ -1,11 +1,19 @@ use std::{ env, fs, path::{Path, PathBuf}, + sync::{Arc, Barrier}, + thread, }; use anyhow::{Context, Result, ensure}; use openinfer_deepseek_v2_lite::DeepSeekV2LiteEp2Generator; -use openinfer_engine::engine::{EngineLoadOptions, FinishReason}; +use openinfer_engine::{ + engine::{ + EngineLoadOptions, FinishReason, GenerateRequest, TokenEvent, TokenSink, + TokenStreamReceiver, + }, + sampler::SamplingParams, +}; use serde::Deserialize; use sha2::{Digest, Sha256}; use vllm_text::tokenizer::{HuggingFaceTokenizer, Tokenizer}; @@ -72,7 +80,8 @@ fn test_deepseek_v2_lite_ep2_rust_generation() -> Result<()> { "duplicate CUDA ordinal error should mention distinct devices, got {duplicate_ordinal_err:#}" ); - run_rust_generation(&model_path_label, &model_path) + run_rust_generation(&model_path_label, &model_path)?; + run_mixed_serving_generation(&model_path, &model_path_label) } fn run_rust_generation(model_path_label: &str, model_path: &Path) -> Result<()> { @@ -473,6 +482,363 @@ fn run_case_set_case( } } +fn run_mixed_serving_generation(model_path: &Path, model_path_label: &str) -> Result<()> { + let tokenizer_path = model_path.join("tokenizer.json"); + let tokenizer = HuggingFaceTokenizer::new(&tokenizer_path).map_err(|err| { + anyhow::anyhow!( + "failed to load tokenizer {}: {err:?}", + tokenizer_path.display() + ) + })?; + let cases = [ + ("mixed-hello", "Hello", 8usize, false), + ("mixed-world", "World", 8usize, false), + ("mixed-paris", "Paris", 8usize, false), + ("mixed-china", "China", 8usize, false), + ]; + + let mut sequential = DeepSeekV2LiteEp2Generator::load( + model_path, + EngineLoadOptions { + enable_cuda_graph: false, + enable_prefill_profile: false, + device_ordinals: vec![0, 1], + seed: 42, + ..EngineLoadOptions::default() + }, + )?; + let mut encoded_cases = Vec::with_capacity(cases.len()); + let mut expected = Vec::with_capacity(cases.len()); + for (id, prompt, max_tokens, ignore_eos) in cases { + let prompt_tokens = tokenizer + .encode(prompt, false) + .map_err(|err| anyhow::anyhow!("encode prompt for {id} failed: {err:?}"))?; + ensure!( + !prompt_tokens.is_empty(), + "tokenizer returned empty prompt for mixed-serving case {id}" + ); + let result = sequential.generate_greedy(&prompt_tokens, max_tokens, ignore_eos)?; + expected.push((id.to_string(), result.tokens, result.finish_reason)); + encoded_cases.push(( + id.to_string(), + prompt.to_string(), + prompt_tokens, + max_tokens, + ignore_eos, + )); + } + let hello_len = encoded_cases + .iter() + .find(|(id, _, _, _, _)| id == "mixed-hello") + .map(|(_, _, prompt_tokens, _, _)| prompt_tokens.len()) + .context("mixed-hello case missing")?; + let world_len = encoded_cases + .iter() + .find(|(id, _, _, _, _)| id == "mixed-world") + .map(|(_, _, prompt_tokens, _, _)| prompt_tokens.len()) + .context("mixed-world case missing")?; + ensure!( + hello_len == world_len, + "mixed-serving e2e requires Hello and World to tokenize to the same length so the batch decode path is exercised; got Hello={hello_len}, World={world_len}" + ); + for (id, _, prompt_tokens, _, _) in &encoded_cases { + ensure!( + prompt_tokens.len() == hello_len, + "mixed-serving e2e requires every main prompt to share one tokenized length so the batch decode path is exercised; case {id} has {}, expected {hello_len}", + prompt_tokens.len() + ); + } + let isolation_id = "mixed-valid-beside-invalid"; + let isolation_prompt = "A valid request beside a rejected request"; + let isolation_prompt_tokens = tokenizer + .encode(isolation_prompt, false) + .map_err(|err| anyhow::anyhow!("encode prompt for {isolation_id} failed: {err:?}"))?; + ensure!( + !isolation_prompt_tokens.is_empty(), + "tokenizer returned empty prompt for mixed-serving case {isolation_id}" + ); + let isolation_expected = sequential.generate_greedy(&isolation_prompt_tokens, 6, false)?; + + let fallback_cases = [ + ("mixed-short-single-row", "Hello", 6usize, false), + ( + "mixed-long-single-row", + "A valid request beside a rejected request", + 6usize, + false, + ), + ]; + let mut fallback_encoded = Vec::with_capacity(fallback_cases.len()); + let mut fallback_expected = Vec::with_capacity(fallback_cases.len()); + for (id, prompt, max_tokens, ignore_eos) in fallback_cases { + let prompt_tokens = tokenizer + .encode(prompt, false) + .map_err(|err| anyhow::anyhow!("encode prompt for {id} failed: {err:?}"))?; + ensure!( + !prompt_tokens.is_empty(), + "tokenizer returned empty prompt for mixed-serving case {id}" + ); + let result = sequential.generate_greedy(&prompt_tokens, max_tokens, ignore_eos)?; + fallback_expected.push((id.to_string(), result.tokens, result.finish_reason)); + fallback_encoded.push((id.to_string(), prompt_tokens, max_tokens, ignore_eos)); + } + let fallback_prompt_lengths: Vec<_> = fallback_encoded + .iter() + .map(|(_, prompt_tokens, _, _)| prompt_tokens.len()) + .collect(); + ensure!( + fallback_prompt_lengths + .windows(2) + .any(|pair| pair[0] != pair[1]), + "mixed-serving fallback e2e requires at least two prompt lengths so single-row decode is exercised; got {fallback_prompt_lengths:?}" + ); + drop(sequential); + + let handle = openinfer_deepseek_v2_lite::start_engine( + model_path, + EngineLoadOptions { + enable_cuda_graph: false, + enable_prefill_profile: false, + device_ordinals: vec![0, 1], + seed: 42, + ..EngineLoadOptions::default() + }, + )?; + + let mut requests = Vec::with_capacity(encoded_cases.len()); + let mut receivers = Vec::with_capacity(encoded_cases.len()); + for (id, _prompt, prompt_tokens, max_tokens, ignore_eos) in encoded_cases { + let (token_tx, token_rx) = TokenSink::standalone(); + let req = GenerateRequest { + request_id: Some(id.clone()), + queued_at_unix_s: None, + prompt_tokens, + params: SamplingParams { + ignore_eos, + ..SamplingParams::default() + }, + max_tokens, + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + }; + receivers.push((id, token_rx)); + requests.push(req); + } + submit_concurrently(&handle, requests)?; + + let mut actual = Vec::with_capacity(receivers.len()); + for (id, mut token_rx) in receivers { + let (tokens, finish_reason) = drain_engine_stream(&id, &mut token_rx)?; + actual.push((id, tokens, finish_reason)); + } + + ensure!( + actual == expected, + "mixed-serving output drift on {model_path_label}: actual={actual:?} expected={expected:?}" + ); + + run_mixed_serving_position_fallback(&handle, fallback_encoded, fallback_expected, &mut actual)?; + + run_mixed_serving_rejection_isolation( + &handle, + isolation_id, + isolation_prompt_tokens, + isolation_expected.tokens, + isolation_expected.finish_reason, + &mut actual, + )?; + + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema": 1, + "report_type": "deepseek-v2-lite-ep2-mixed-serving-e2e", + "model_path": model_path_label, + "ep_backend": current_backend(), + "case_count": actual.len(), + "cases": actual + .iter() + .map(|(id, tokens, finish_reason)| { + serde_json::json!({ + "id": id, + "generated_tokens": tokens.len(), + "output_token_sha256": token_sha256(tokens), + "finish_reason": format!("{finish_reason:?}"), + }) + }) + .collect::>(), + }))? + ); + Ok(()) +} + +fn run_mixed_serving_position_fallback( + handle: &openinfer_engine::engine::EngineHandle, + encoded_cases: Vec<(String, Vec, usize, bool)>, + expected: Vec<(String, Vec, FinishReason)>, + actual: &mut Vec<(String, Vec, FinishReason)>, +) -> Result<()> { + let mut requests = Vec::with_capacity(encoded_cases.len()); + let mut receivers = Vec::with_capacity(encoded_cases.len()); + for (id, prompt_tokens, max_tokens, ignore_eos) in encoded_cases { + let (token_tx, token_rx) = TokenSink::standalone(); + let req = GenerateRequest { + request_id: Some(id.clone()), + queued_at_unix_s: None, + prompt_tokens, + params: SamplingParams { + ignore_eos, + ..SamplingParams::default() + }, + max_tokens, + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + }; + receivers.push((id, token_rx)); + requests.push(req); + } + submit_concurrently(handle, requests)?; + + let mut fallback_actual = Vec::with_capacity(receivers.len()); + for (id, mut token_rx) in receivers { + let (tokens, finish_reason) = drain_engine_stream(&id, &mut token_rx)?; + fallback_actual.push((id, tokens, finish_reason)); + } + + ensure!( + fallback_actual == expected, + "mixed-serving single-row fallback output drift: actual={fallback_actual:?} expected={expected:?}" + ); + actual.extend(fallback_actual); + Ok(()) +} + +fn submit_concurrently( + handle: &openinfer_engine::engine::EngineHandle, + requests: Vec, +) -> Result<()> { + let barrier = Arc::new(Barrier::new(requests.len() + 1)); + let mut threads = Vec::with_capacity(requests.len()); + for req in requests { + let handle = handle.clone(); + let barrier = Arc::clone(&barrier); + threads.push(thread::spawn(move || -> Result<()> { + barrier.wait(); + handle + .submit(req) + .map_err(|_| anyhow::anyhow!("DeepSeek-V2-Lite mixed-serving engine closed"))?; + Ok(()) + })); + } + barrier.wait(); + for submit_thread in threads { + submit_thread + .join() + .map_err(|_| anyhow::anyhow!("mixed-serving submit thread panicked"))??; + } + Ok(()) +} + +fn run_mixed_serving_rejection_isolation( + handle: &openinfer_engine::engine::EngineHandle, + valid_id: &str, + valid_prompt_tokens: Vec, + expected_tokens: Vec, + expected_finish_reason: FinishReason, + actual: &mut Vec<(String, Vec, FinishReason)>, +) -> Result<()> { + let (invalid_tx, mut invalid_rx) = TokenSink::standalone(); + let invalid_req = GenerateRequest { + request_id: Some("mixed-invalid-logprobs".to_string()), + queued_at_unix_s: None, + prompt_tokens: vec![1, 2, 3], + params: SamplingParams::default(), + max_tokens: 4, + lora_adapter: None, + token_tx: invalid_tx, + logprobs: 1, + echo: false, + }; + + let (valid_tx, mut valid_rx) = TokenSink::standalone(); + let valid_req = GenerateRequest { + request_id: Some(valid_id.to_string()), + queued_at_unix_s: None, + prompt_tokens: valid_prompt_tokens, + params: SamplingParams::default(), + max_tokens: 6, + lora_adapter: None, + token_tx: valid_tx, + logprobs: 0, + echo: false, + }; + submit_concurrently(handle, vec![invalid_req, valid_req])?; + + let mut saw_rejection = false; + while let Some((_tag, event)) = invalid_rx.blocking_recv() { + match event { + TokenEvent::Scheduled { .. } => {} + TokenEvent::Rejected { message, .. } => { + ensure!( + message.contains("logprobs"), + "invalid request rejection should mention logprobs, got {message}" + ); + saw_rejection = true; + break; + } + TokenEvent::Token { .. } + | TokenEvent::PromptTokens { .. } + | TokenEvent::Finished { .. } + | TokenEvent::Error { .. } => { + anyhow::bail!("invalid mixed-serving request reached unexpected event") + } + } + } + ensure!( + saw_rejection, + "invalid mixed-serving request stream closed without Rejected" + ); + + let (tokens, finish_reason) = drain_engine_stream(valid_id, &mut valid_rx)?; + ensure!( + tokens == expected_tokens && finish_reason == expected_finish_reason, + "valid mixed-serving request drifted next to invalid request: actual_tokens={tokens:?} expected_tokens={expected_tokens:?} actual_finish={finish_reason:?} expected_finish={expected_finish_reason:?}" + ); + ensure!( + !tokens.is_empty(), + "valid mixed-serving request produced no tokens after neighboring rejection" + ); + actual.push((valid_id.to_string(), tokens, finish_reason)); + Ok(()) +} + +fn drain_engine_stream( + case_id: &str, + token_rx: &mut TokenStreamReceiver, +) -> Result<(Vec, FinishReason)> { + let mut tokens = Vec::new(); + loop { + match token_rx.blocking_recv() { + Some((_tag, TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. })) => {} + Some((_tag, TokenEvent::Token { id, .. })) => tokens.push(id), + Some((_tag, TokenEvent::Finished { finish_reason, .. })) => { + return Ok((tokens, finish_reason)); + } + Some((_tag, TokenEvent::Error { message, .. })) => { + anyhow::bail!("mixed-serving case {case_id} failed: {message}") + } + Some((_tag, TokenEvent::Rejected { message, .. })) => { + anyhow::bail!("mixed-serving case {case_id} rejected: {message}") + } + None => anyhow::bail!("mixed-serving case {case_id} stream closed before Finished"), + } + } +} + fn load_case_set_from_env() -> Result)>> { let Ok(raw_path) = env::var(E2E_CASE_SET_ENV) else { return Ok(None);