Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e739507803
|
||
|
|
ffafe3f4dc
|
||
|
|
91255619b6
|
||
|
|
d453d5deb8
|
||
|
|
01c604195d
|
||
|
|
ea1e7f1997
|
||
|
|
e94633c3a3
|
||
|
|
7cc37ebf37
|
||
|
|
98d63cc3fd
|
@@ -20,6 +20,11 @@ pub struct MessagesRequest {
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
/// Anthropic's Messages API carries `top_k` natively, so it maps
|
||||
/// straight onto the OpenAI-side field rather than being dropped in
|
||||
/// translation (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<usize>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub stream: Option<bool>,
|
||||
#[serde(flatten)]
|
||||
|
||||
@@ -18,6 +18,31 @@ pub struct ChatCompletionRequest {
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
/// Truncate sampling to the `k` most likely tokens (#272).
|
||||
///
|
||||
/// Not in the OpenAI schema but universal among serving stacks
|
||||
/// (vLLM, SGLang, llama.cpp, Ollama), and required by models that
|
||||
/// publish it — Qwen3 specifies `top_k = 20`. A named field rather
|
||||
/// than `extra`, because a sampling parameter that lands in `extra`
|
||||
/// is accepted and silently discarded.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<usize>,
|
||||
/// Pin the sampler's RNG for reproducible output (OpenAI core).
|
||||
///
|
||||
/// Previously ignored on every text path, which called
|
||||
/// `unix_subsec_nanos()` regardless — a caller asking for
|
||||
/// determinism got fresh randomness and a `200`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub seed: Option<u64>,
|
||||
/// Penalty applied to recently-generated tokens; `1.0` disables.
|
||||
/// Defaults to 1.1 when absent (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub repetition_penalty: Option<f32>,
|
||||
/// How many recent tokens `repetition_penalty` considers. Defaults
|
||||
/// to 64 — fine for chat, short for a long reasoning block, which is
|
||||
/// why it is a knob rather than a constant (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub repeat_last_n: Option<usize>,
|
||||
/// The deprecated spelling of the output cap. Still what most
|
||||
/// clients send; see [`ChatCompletionRequest::effective_max_tokens`].
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
@@ -149,8 +174,9 @@ pub struct Usage {
|
||||
/// `completion_tokens`, never added into `total_tokens`).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub completion_tokens_details: Option<CompletionTokensDetails>,
|
||||
/// OpenAI-standard breakdown of `prompt_tokens`. Populated once
|
||||
/// prompt caching lands (#11); `None` until then.
|
||||
/// OpenAI-standard breakdown of `prompt_tokens`, carrying
|
||||
/// `cached_tokens` (#269). Omitted when nothing was reused, so
|
||||
/// "absent" stays distinguishable from "measured as zero".
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_tokens_details: Option<PromptTokensDetails>,
|
||||
/// helexa extension (non-OpenAI): server-measured prefill/decode
|
||||
@@ -182,8 +208,9 @@ pub struct CompletionTokensDetails {
|
||||
/// Sub-counts of `Usage::prompt_tokens`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PromptTokensDetails {
|
||||
/// Prompt tokens served from cache (cache-read rate). Populated
|
||||
/// once prompt caching lands (#11).
|
||||
/// Prompt tokens served from neuron's prefix KV cache (#11), so a
|
||||
/// caller can see the saving rather than assume none (#269). A
|
||||
/// sub-count of `prompt_tokens`, never added into `total_tokens`.
|
||||
pub cached_tokens: u64,
|
||||
}
|
||||
|
||||
|
||||
@@ -55,6 +55,21 @@ pub struct ResponsesRequest {
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
/// Truncate sampling to the `k` most likely tokens (#272). Not in
|
||||
/// the Responses schema, but a sampling parameter that lands in
|
||||
/// `extra` is accepted and silently discarded, which is the defect
|
||||
/// this closes.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<usize>,
|
||||
/// Pin the sampler's RNG for reproducible output (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub seed: Option<u64>,
|
||||
/// Penalty on recently-generated tokens; `1.0` disables (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub repetition_penalty: Option<f32>,
|
||||
/// Window `repetition_penalty` considers (#272).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub repeat_last_n: Option<usize>,
|
||||
/// Chained-conversation identifier. We don't store responses
|
||||
/// server-side yet; if this is `Some`, the handler returns 400.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
@@ -249,8 +264,8 @@ pub struct ResponsesUsage {
|
||||
/// sub-count of `output_tokens`, never added into `total_tokens`).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub output_tokens_details: Option<OutputTokensDetails>,
|
||||
/// OpenAI-standard breakdown of `input_tokens`. Populated once
|
||||
/// prompt caching lands (#11); `None` until then.
|
||||
/// OpenAI-standard breakdown of `input_tokens`, carrying
|
||||
/// `cached_tokens` (#269). Omitted when nothing was reused.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub input_tokens_details: Option<InputTokensDetails>,
|
||||
}
|
||||
@@ -265,8 +280,8 @@ pub struct OutputTokensDetails {
|
||||
/// Sub-counts of `ResponsesUsage::input_tokens`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct InputTokensDetails {
|
||||
/// Input tokens served from cache (cache-read rate). Populated
|
||||
/// once prompt caching lands (#11).
|
||||
/// Input tokens served from neuron's prefix KV cache (#11), so a
|
||||
/// caller can see the saving rather than assume none (#269).
|
||||
pub cached_tokens: u64,
|
||||
}
|
||||
|
||||
|
||||
@@ -101,6 +101,13 @@ pub fn anthropic_to_openai(req: MessagesRequest) -> ChatCompletionRequest {
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
top_p: req.top_p,
|
||||
// Anthropic models top_k too, so it survives the hop (#272).
|
||||
top_k: req.top_k,
|
||||
// Anthropic's Messages API has no seed or penalty knobs; absent
|
||||
// means the model's own generation_config decides.
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
max_tokens: Some(req.max_tokens),
|
||||
// Anthropic has one spelling; it maps onto the legacy field.
|
||||
max_completion_tokens: None,
|
||||
|
||||
@@ -141,6 +141,15 @@ pub fn usage_sink(
|
||||
/// returns a guard the caller settles with actual usage; on refusal returns
|
||||
/// the #63 envelope (`rate_limit_exceeded` + `Retry-After` for a resetting
|
||||
/// window, `insufficient_quota` for a hard balance — never `402`).
|
||||
// The `Err` variant is the #63 envelope itself, which is a data-carrying
|
||||
// struct (status, type, code, message, param, retry hint, diagnostics) and
|
||||
// is passed by value everywhere it is built, matched and rendered. Boxing it
|
||||
// at this one boundary would push `Box` conversions through three handlers
|
||||
// and both HTTP adapters to save an allocation on a path taken once per
|
||||
// rejected request. clippy 1.98 lowered what it considers large enough to
|
||||
// flag; the shape is deliberate. Same call as the existing allows on the
|
||||
// figment-returning config loaders.
|
||||
#[allow(clippy::result_large_err)]
|
||||
pub async fn reserve_or_reject(
|
||||
provider: Arc<dyn EntitlementProvider>,
|
||||
principal: &Principal,
|
||||
|
||||
@@ -121,6 +121,10 @@ fn describe_metrics() {
|
||||
"cortex_request_errors_total",
|
||||
"Total number of failed proxy requests"
|
||||
);
|
||||
metrics::describe_counter!(
|
||||
"cortex_cached_prompt_tokens_total",
|
||||
"Prompt tokens served from neuron's prefix KV cache; over cortex_prompt_tokens_total this is the cache-hit rate (#269)"
|
||||
);
|
||||
metrics::describe_counter!("cortex_evictions_total", "Total number of model evictions");
|
||||
metrics::describe_counter!(
|
||||
"cortex_cold_starts_total",
|
||||
|
||||
@@ -171,6 +171,9 @@ impl IntoResponse for ProxyError {
|
||||
// over the decode window (first→last chunk); falls back to the
|
||||
// full request duration for single-chunk (non-streaming) bodies
|
||||
// cortex_prompt_tokens_total / cortex_completion_tokens_total (counters)
|
||||
// cortex_cached_prompt_tokens_total (counter) — prompt tokens served
|
||||
// from neuron's prefix cache; over prompt_tokens_total this is the
|
||||
// fleet's cache-hit rate (#269)
|
||||
|
||||
/// Cap on the retained body tail. The usage object rides on the final
|
||||
/// chunk, so a generous tail is plenty; the cap bounds memory on huge
|
||||
@@ -243,6 +246,15 @@ impl ChunkObserver for CortexMetrics {
|
||||
if let Some(prompt) = prompt {
|
||||
metrics::counter!("cortex_prompt_tokens_total", &self.labels).increment(prompt);
|
||||
}
|
||||
// Prefix-cache reuse (#269), read from the same usage object.
|
||||
// Against `cortex_prompt_tokens_total` this gives fleet-wide
|
||||
// cache-hit rate as a ratio, which was previously only
|
||||
// obtainable by grepping `reused=` out of a neuron journal —
|
||||
// per-host, unaggregated, and gone on log rotation.
|
||||
if let Some(cached) = last_count_for(self.tail.as_str(), "cached_tokens") {
|
||||
metrics::counter!("cortex_cached_prompt_tokens_total", &self.labels)
|
||||
.increment(cached);
|
||||
}
|
||||
if let Some(completion) = completion.filter(|c| *c > 0) {
|
||||
metrics::counter!("cortex_completion_tokens_total", &self.labels)
|
||||
.increment(completion);
|
||||
|
||||
+402
-250
File diff suppressed because it is too large
Load Diff
@@ -297,14 +297,23 @@ async fn op_drop_snap(cfg: &EngineConfig, session: &mut ActiveSession, id: u64)
|
||||
|
||||
/// Prefill one joining request at B=1 through the backend's existing
|
||||
/// chunked-prefill + prefix-cache paths. Returns the prefill logits.
|
||||
/// Prefill, returning the logits **and** how many leading tokens the
|
||||
/// prefix cache let us skip (#269).
|
||||
///
|
||||
/// The reuse count is computed here and was previously only logged, so
|
||||
/// no caller could report it — every client saw a 0% cache hit rate on
|
||||
/// a cache that was working. Returning it is what lets `usage` carry
|
||||
/// `cached_tokens`.
|
||||
async fn op_prefill(
|
||||
cfg: &EngineConfig,
|
||||
session: &mut ActiveSession,
|
||||
prompt_tokens: &[u32],
|
||||
) -> Result<Vec<f32>> {
|
||||
) -> Result<(Vec<f32>, usize)> {
|
||||
let prompt_len = prompt_tokens.len();
|
||||
let prefill_start = std::time::Instant::now();
|
||||
let logits = match (&cfg.backend, session) {
|
||||
// Each arm yields its logits plus how many leading tokens the
|
||||
// prefix cache served, so neither can be forgotten.
|
||||
let (logits, cached_tokens) = match (&cfg.backend, session) {
|
||||
(
|
||||
BackendConfig::Single {
|
||||
worker,
|
||||
@@ -340,7 +349,7 @@ async fn op_prefill(
|
||||
None => chunked_prefill_via_worker(worker, *handle, prompt_tokens, reused).await?,
|
||||
};
|
||||
prefill_rate.record(prompt_len, prefill_start.elapsed());
|
||||
logits
|
||||
(logits, reused)
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
(BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => {
|
||||
@@ -384,12 +393,12 @@ async fn op_prefill(
|
||||
}
|
||||
};
|
||||
tp.prefill_rate.record(prompt_len, prefill_start.elapsed());
|
||||
logits
|
||||
(logits, reused)
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
_ => anyhow::bail!("engine backend/session mismatch"),
|
||||
};
|
||||
Ok(logits)
|
||||
Ok((logits, cached_tokens))
|
||||
}
|
||||
|
||||
/// One queued request. Admission has already been passed — the permit
|
||||
@@ -397,9 +406,10 @@ async fn op_prefill(
|
||||
pub struct EngineRequest {
|
||||
pub prompt_tokens: Vec<u32>,
|
||||
pub max_new: usize,
|
||||
pub temperature: f64,
|
||||
pub top_p: Option<f64>,
|
||||
pub seed: u64,
|
||||
/// Sampling resolved once from request → model `generation_config`
|
||||
/// → fallback (#272), rather than loose scalars each caller could
|
||||
/// default differently.
|
||||
pub sampling: crate::harness::sampling::SamplingParams,
|
||||
pub eos_id: Option<u32>,
|
||||
pub tool_schemas: ToolSchemas,
|
||||
pub tx: mpsc::Sender<InferenceEvent>,
|
||||
@@ -439,6 +449,8 @@ enum RouterMsg {
|
||||
reason: FinishReason,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
/// Leading prompt tokens the prefix cache let us skip (#269).
|
||||
cached_tokens: u32,
|
||||
timing: FinishTiming,
|
||||
},
|
||||
}
|
||||
@@ -455,6 +467,10 @@ struct Slot {
|
||||
max_new: usize,
|
||||
eos_id: Option<u32>,
|
||||
lp: LogitsProcessor,
|
||||
/// The slot's resolved sampling (#272). Held per-slot because the
|
||||
/// repetition penalty and its window are per-request knobs, and
|
||||
/// slots in one batch can carry different ones.
|
||||
sampling: crate::harness::sampling::SamplingParams,
|
||||
router: mpsc::Sender<RouterMsg>,
|
||||
/// Set by the router when the consumer hangs up; the engine stops
|
||||
/// feeding the slot and compacts it out.
|
||||
@@ -462,6 +478,9 @@ struct Slot {
|
||||
finished: Option<FinishReason>,
|
||||
prefill_ms: u32,
|
||||
prefill_tokens: u32,
|
||||
/// Prompt tokens served from the prefix cache (#269), reported to
|
||||
/// the caller as `cached_tokens` so a working cache is visible.
|
||||
cached_tokens: u32,
|
||||
decode_start: std::time::Instant,
|
||||
_admit: AdmissionPermit,
|
||||
}
|
||||
@@ -559,22 +578,23 @@ async fn run_engine(cfg: EngineConfig, mut rx: mpsc::Receiver<EngineRequest>) {
|
||||
break;
|
||||
}
|
||||
};
|
||||
let nt = match sample_with_penalty(&logits, &slot.generated, &mut slot.lp) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
let health = logits_health_slice(&logits_vec);
|
||||
tracing::warn!(
|
||||
?health,
|
||||
error = %e,
|
||||
"batch engine: sample failed; logits unhealthy"
|
||||
);
|
||||
// Unhealthy logits are a device-level problem —
|
||||
// fail the whole engine, mirroring the B=1 path's
|
||||
// poison classification.
|
||||
fatal = Some(e);
|
||||
break;
|
||||
}
|
||||
};
|
||||
let nt =
|
||||
match sample_with_penalty(&logits, &slot.generated, &mut slot.lp, &slot.sampling) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
let health = logits_health_slice(&logits_vec);
|
||||
tracing::warn!(
|
||||
?health,
|
||||
error = %e,
|
||||
"batch engine: sample failed; logits unhealthy"
|
||||
);
|
||||
// Unhealthy logits are a device-level problem —
|
||||
// fail the whole engine, mirroring the B=1 path's
|
||||
// poison classification.
|
||||
fatal = Some(e);
|
||||
break;
|
||||
}
|
||||
};
|
||||
if Some(nt) == slot.eos_id {
|
||||
finish_slot(slot, FinishReason::Stop, active_rate(&cfg, sess)).await;
|
||||
continue;
|
||||
@@ -627,6 +647,7 @@ async fn finish_slot(slot: &mut Slot, reason: FinishReason, rate: &PrefillRateEm
|
||||
reason,
|
||||
prompt_tokens: slot.prompt_len as u32,
|
||||
completion_tokens: slot.generated.len() as u32,
|
||||
cached_tokens: slot.cached_tokens,
|
||||
timing: FinishTiming {
|
||||
prefill_ms: slot.prefill_ms,
|
||||
decode_ms: slot.decode_start.elapsed().as_millis() as u32,
|
||||
@@ -759,14 +780,10 @@ async fn prefill_join(
|
||||
session: &mut ActiveSession,
|
||||
req: EngineRequest,
|
||||
) -> Result<Option<(Slot, u64)>> {
|
||||
use candle_transformers::generation::Sampling;
|
||||
|
||||
let EngineRequest {
|
||||
prompt_tokens,
|
||||
max_new,
|
||||
temperature,
|
||||
top_p,
|
||||
seed,
|
||||
sampling,
|
||||
eos_id,
|
||||
tool_schemas,
|
||||
tx,
|
||||
@@ -774,27 +791,17 @@ async fn prefill_join(
|
||||
span,
|
||||
} = req;
|
||||
|
||||
let mut lp = {
|
||||
let sampling = if temperature <= 0.0 {
|
||||
Sampling::ArgMax
|
||||
} else {
|
||||
match top_p {
|
||||
Some(p) => Sampling::TopP { p, temperature },
|
||||
None => Sampling::All { temperature },
|
||||
}
|
||||
};
|
||||
LogitsProcessor::from_sampling(seed, sampling)
|
||||
};
|
||||
let mut lp = LogitsProcessor::from_sampling(sampling.seed, sampling.to_sampling());
|
||||
|
||||
let prompt_len = prompt_tokens.len();
|
||||
let prefill_start = std::time::Instant::now();
|
||||
let logits_vec = op_prefill(cfg, session, &prompt_tokens).await?;
|
||||
let (logits_vec, cached_tokens) = op_prefill(cfg, session, &prompt_tokens).await?;
|
||||
let prefill_elapsed = prefill_start.elapsed();
|
||||
|
||||
// First token from the prefill logits.
|
||||
let generated: Vec<u32> = Vec::new();
|
||||
let logits = Tensor::new(logits_vec.as_slice(), &Device::Cpu)?;
|
||||
let first = match sample_with_penalty(&logits, &generated, &mut lp) {
|
||||
let first = match sample_with_penalty(&logits, &generated, &mut lp, &sampling) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
let health = logits_health_slice(&logits_vec);
|
||||
@@ -827,11 +834,13 @@ async fn prefill_join(
|
||||
let mut slot = Slot {
|
||||
prefix_len: prompt_len,
|
||||
prompt_len,
|
||||
cached_tokens: cached_tokens as u32,
|
||||
generated,
|
||||
next_token: first,
|
||||
max_new,
|
||||
eos_id,
|
||||
lp,
|
||||
sampling,
|
||||
router: router_tx,
|
||||
hangup,
|
||||
finished: None,
|
||||
@@ -890,6 +899,9 @@ async fn run_router(
|
||||
let mut tool_call_idx: usize = 0;
|
||||
let mut emitted_tool_call = false;
|
||||
let mut consumer_alive = true;
|
||||
// Whether a Finish was seen, so the fall-through below can tell a
|
||||
// clean end from a dropped slot.
|
||||
let mut finished = false;
|
||||
|
||||
while let Some(msg) = rx.recv().await {
|
||||
match msg {
|
||||
@@ -978,17 +990,43 @@ async fn run_router(
|
||||
mut reason,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
cached_tokens,
|
||||
timing,
|
||||
} => {
|
||||
if emitted_tool_call && reason == FinishReason::Stop {
|
||||
reason = FinishReason::ToolCalls;
|
||||
}
|
||||
// Close the journal trail for this turn (#268). The
|
||||
// non-streaming path has always logged this; streamed
|
||||
// turns logged nothing at all, so whether one finished,
|
||||
// hit its output cap, or was abandoned mid-flight was
|
||||
// unanswerable afterwards — and an absence of lines was
|
||||
// repeatedly misread as an absence of work.
|
||||
//
|
||||
// `run_router` is instrumented with the request's span,
|
||||
// so req_id and model come along without being passed.
|
||||
tracing::info!(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens = reasoning_token_count,
|
||||
cached_tokens,
|
||||
finish_reason = ?reason,
|
||||
prefill_ms = timing.prefill_ms,
|
||||
decode_ms = timing.decode_ms,
|
||||
// The turn ran to completion with nobody reading it:
|
||||
// it held its admission permit and the pool lock the
|
||||
// whole time. Invisible until now.
|
||||
client_gone = !consumer_alive,
|
||||
"chat_completion (stream): done"
|
||||
);
|
||||
finished = true;
|
||||
let _ = tx
|
||||
.send(InferenceEvent::Finish {
|
||||
reason,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens: reasoning_token_count,
|
||||
cached_tokens,
|
||||
timing: Some(timing),
|
||||
})
|
||||
.await;
|
||||
@@ -996,6 +1034,17 @@ async fn run_router(
|
||||
}
|
||||
}
|
||||
}
|
||||
// Fell out of the loop without a Finish: the engine dropped this
|
||||
// slot's channel (fatal worker error, shutdown). The client's stream
|
||||
// ends either way, but silently — so say so rather than leaving the
|
||||
// turn's last trace to be a prefill line (#268).
|
||||
if !finished {
|
||||
tracing::warn!(
|
||||
reasoning_tokens = reasoning_token_count,
|
||||
client_gone = !consumer_alive,
|
||||
"chat_completion (stream): ended without a finish event — engine dropped the slot"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// `tracing::Instrument` without importing the trait at every use
|
||||
@@ -1013,6 +1062,92 @@ mod tests {
|
||||
use crate::config::AdmissionConfig;
|
||||
use crate::harness::admission::AdmissionController;
|
||||
|
||||
/// A streamed turn must close its own journal trail (#268).
|
||||
///
|
||||
/// Streaming used to end silently: no token counts, no
|
||||
/// `finish_reason`, no duration. That made a finished turn, a turn
|
||||
/// that hit its output cap, and a turn abandoned by a disconnected
|
||||
/// client indistinguishable after the fact — and during the #267
|
||||
/// investigation an absence of log lines was twice mistaken for an
|
||||
/// absence of work on a daemon that was healthy and decoding.
|
||||
///
|
||||
/// Asserts the emitted line rather than the struct, because a
|
||||
/// `Finish` event that reaches the client but leaves no server-side
|
||||
/// record is exactly the state this fixes.
|
||||
#[tokio::test]
|
||||
async fn a_streamed_turn_logs_how_it_ended() {
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tracing_subscriber::prelude::*;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct Capture(Arc<Mutex<Vec<u8>>>);
|
||||
impl std::io::Write for Capture {
|
||||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||
self.0
|
||||
.lock()
|
||||
.expect("capture poisoned")
|
||||
.extend_from_slice(buf);
|
||||
Ok(buf.len())
|
||||
}
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for Capture {
|
||||
type Writer = Self;
|
||||
fn make_writer(&'a self) -> Self::Writer {
|
||||
self.clone()
|
||||
}
|
||||
}
|
||||
|
||||
let capture = Capture::default();
|
||||
let subscriber = tracing_subscriber::registry()
|
||||
.with(tracing_subscriber::fmt::layer().with_writer(capture.clone()));
|
||||
let _guard = tracing::subscriber::set_default(subscriber);
|
||||
|
||||
let (router_tx, router_rx) = mpsc::channel::<RouterMsg>(8);
|
||||
let (tx, mut rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let handle = tokio::spawn(run_router(
|
||||
tiny_tokenizer(8),
|
||||
None,
|
||||
None,
|
||||
ToolSchemas::default(),
|
||||
false,
|
||||
tx,
|
||||
Arc::new(AtomicBool::new(false)),
|
||||
router_rx,
|
||||
));
|
||||
|
||||
router_tx
|
||||
.send(RouterMsg::Finish {
|
||||
reason: FinishReason::Length,
|
||||
prompt_tokens: 7413,
|
||||
completion_tokens: 24576,
|
||||
cached_tokens: 2068,
|
||||
timing: FinishTiming {
|
||||
prefill_ms: 6000,
|
||||
decode_ms: 900_000,
|
||||
prefill_tokens: 7413,
|
||||
},
|
||||
})
|
||||
.await
|
||||
.expect("router accepts finish");
|
||||
drop(router_tx);
|
||||
handle.await.expect("router task joins");
|
||||
while rx.recv().await.is_some() {}
|
||||
|
||||
let logged = String::from_utf8(capture.0.lock().expect("capture poisoned").clone())
|
||||
.expect("utf8 log output");
|
||||
assert!(
|
||||
logged.contains("chat_completion (stream): done"),
|
||||
"streamed turn left no completion line:\n{logged}"
|
||||
);
|
||||
// The counts are the point — "it ended" without them cannot
|
||||
// distinguish a complete answer from an exhausted budget.
|
||||
assert!(logged.contains("24576"), "no completion_tokens:\n{logged}");
|
||||
assert!(logged.contains("Length"), "no finish_reason:\n{logged}");
|
||||
}
|
||||
|
||||
/// A WordLevel tokenizer whose vocab covers the whole fixture
|
||||
/// vocab (`w0`..`w511`), so every decoded token maps to a unique
|
||||
/// word and the emitted text uniquely encodes the token sequence.
|
||||
@@ -1046,9 +1181,17 @@ mod tests {
|
||||
.submit(EngineRequest {
|
||||
prompt_tokens: prompt,
|
||||
max_new,
|
||||
temperature: 0.0, // greedy — deterministic
|
||||
top_p: None,
|
||||
seed: 0,
|
||||
// Greedy, deterministic: the fixture asserts exact
|
||||
// token sequences.
|
||||
sampling: crate::harness::sampling::SamplingParams::resolve(
|
||||
&crate::harness::sampling::RequestedSampling {
|
||||
temperature: Some(0.0),
|
||||
seed: Some(0),
|
||||
..Default::default()
|
||||
},
|
||||
&crate::harness::sampling::ModelGenerationDefaults::default(),
|
||||
0,
|
||||
),
|
||||
eos_id: None,
|
||||
tool_schemas: ToolSchemas::new(),
|
||||
tx,
|
||||
|
||||
@@ -12,6 +12,7 @@ pub mod image;
|
||||
pub mod prefix_cache;
|
||||
pub mod preflight;
|
||||
pub mod preprocess;
|
||||
pub mod sampling;
|
||||
pub mod speculative;
|
||||
pub mod tp;
|
||||
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
//! Sampling parameter resolution (#272).
|
||||
//!
|
||||
//! Three sources, in priority order:
|
||||
//!
|
||||
//! 1. what the **request** asked for,
|
||||
//! 2. what the **model** published in `generation_config.json`,
|
||||
//! 3. a built-in fallback.
|
||||
//!
|
||||
//! Before this, (2) did not exist and (3) was a hardcoded
|
||||
//! `temperature = 0.7` with `top_p = None` — and `None` selects
|
||||
//! `Sampling::All`, meaning no truncation at all. So a caller that
|
||||
//! omitted sampling fields (the common case: the DeepSeek Harness sends
|
||||
//! only `temperature`) got untruncated sampling across the full ~250k
|
||||
//! vocabulary on a model whose authors published `top_k = 20,
|
||||
//! top_p = 0.95, temperature = 1.0`.
|
||||
//!
|
||||
//! Picking "no truncation" for an absent field is not a neutral default,
|
||||
//! it is the widest possible one. The neutral default is what the model
|
||||
//! shipped with.
|
||||
//!
|
||||
//! Resolved **once** per request and threaded whole. The eight
|
||||
//! `LogitsProcessor` construction sites this replaces could each have
|
||||
//! disagreed about a default — the same shape of bug as #252, where two
|
||||
//! config files described one model and drifted.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::Path;
|
||||
|
||||
/// Fallback temperature when neither the request nor the model says.
|
||||
/// Retained from the previous hardcoded value so a model without a
|
||||
/// `generation_config.json` behaves exactly as it did before.
|
||||
const FALLBACK_TEMPERATURE: f64 = 0.7;
|
||||
|
||||
/// Repetition penalty applied to recently-generated tokens before
|
||||
/// sampling. 1.0 disables it; >1.0 makes recently-emitted tokens less
|
||||
/// likely. mistral.rs and llama.cpp default to 1.1, which is enough to
|
||||
/// stop small quantized models from degenerating into "Wait, no, no..."
|
||||
/// loops without distorting normal output.
|
||||
pub const DEFAULT_REPEAT_PENALTY: f32 = 1.1;
|
||||
|
||||
/// Number of recently-generated tokens fed into the repetition penalty.
|
||||
/// Matches the candle quantized-qwen3 example default.
|
||||
///
|
||||
/// Defensible for chat and poorly suited to a 24k-token reasoning block,
|
||||
/// where a derivation repeated 500 tokens later is well outside the
|
||||
/// window — which is why it is now a per-request knob rather than a
|
||||
/// constant.
|
||||
pub const DEFAULT_REPEAT_LAST_N: usize = 64;
|
||||
|
||||
/// What a model published in its `generation_config.json`.
|
||||
///
|
||||
/// Every field optional: the file is a HuggingFace convention, not a
|
||||
/// schema, and models populate different subsets of it.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct ModelGenerationDefaults {
|
||||
pub temperature: Option<f64>,
|
||||
pub top_p: Option<f64>,
|
||||
pub top_k: Option<usize>,
|
||||
}
|
||||
|
||||
impl ModelGenerationDefaults {
|
||||
/// Read `generation_config.json` from a model snapshot directory.
|
||||
///
|
||||
/// Absent or unparseable is not an error — we fall back exactly as
|
||||
/// before — but it is logged, because silently guessing the model's
|
||||
/// sampling is how the untruncated default went unnoticed.
|
||||
pub fn load_from_dir(dir: &Path) -> Self {
|
||||
let path = dir.join("generation_config.json");
|
||||
let text = match std::fs::read_to_string(&path) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"sampling: no generation_config.json; using built-in fallbacks"
|
||||
);
|
||||
return Self::default();
|
||||
}
|
||||
};
|
||||
match serde_json::from_str::<Self>(&text) {
|
||||
Ok(d) => {
|
||||
tracing::info!(
|
||||
path = %path.display(),
|
||||
temperature = ?d.temperature,
|
||||
top_p = ?d.top_p,
|
||||
top_k = ?d.top_k,
|
||||
"sampling: model generation defaults loaded"
|
||||
);
|
||||
d
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"sampling: generation_config.json did not parse; using built-in fallbacks"
|
||||
);
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the caller asked for. All optional — absent means "no opinion",
|
||||
/// which defers to the model's published value.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct RequestedSampling {
|
||||
pub temperature: Option<f64>,
|
||||
pub top_p: Option<f64>,
|
||||
pub top_k: Option<usize>,
|
||||
pub seed: Option<u64>,
|
||||
pub repeat_penalty: Option<f32>,
|
||||
pub repeat_last_n: Option<usize>,
|
||||
}
|
||||
|
||||
/// Fully-resolved sampling for one request.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct SamplingParams {
|
||||
pub temperature: f64,
|
||||
pub top_p: Option<f64>,
|
||||
pub top_k: Option<usize>,
|
||||
pub seed: u64,
|
||||
pub repeat_penalty: f32,
|
||||
pub repeat_last_n: usize,
|
||||
}
|
||||
|
||||
impl SamplingParams {
|
||||
/// Resolve request over model defaults over fallback.
|
||||
///
|
||||
/// `random_seed` is the per-request seed used when the caller did
|
||||
/// not pin one; it is passed in rather than generated here so this
|
||||
/// stays a pure function and the "same seed reproduces" property is
|
||||
/// testable without a clock.
|
||||
pub fn resolve(
|
||||
requested: &RequestedSampling,
|
||||
model: &ModelGenerationDefaults,
|
||||
random_seed: u64,
|
||||
) -> Self {
|
||||
Self {
|
||||
temperature: requested
|
||||
.temperature
|
||||
.or(model.temperature)
|
||||
.unwrap_or(FALLBACK_TEMPERATURE),
|
||||
// No `unwrap_or` here on purpose: `None` still means "do not
|
||||
// truncate", but it is now only reachable when neither the
|
||||
// caller nor the model expressed a preference, rather than
|
||||
// whenever the caller stayed quiet.
|
||||
top_p: requested.top_p.or(model.top_p),
|
||||
top_k: requested.top_k.or(model.top_k),
|
||||
seed: requested.seed.unwrap_or(random_seed),
|
||||
repeat_penalty: requested.repeat_penalty.unwrap_or(DEFAULT_REPEAT_PENALTY),
|
||||
repeat_last_n: requested.repeat_last_n.unwrap_or(DEFAULT_REPEAT_LAST_N),
|
||||
}
|
||||
}
|
||||
|
||||
/// The candle sampling strategy these parameters select.
|
||||
///
|
||||
/// `temperature <= 0` is greedy regardless of the rest — a caller
|
||||
/// asking for determinism that way should not have it silently
|
||||
/// softened by a `top_k` the model happened to publish.
|
||||
pub fn to_sampling(&self) -> candle_transformers::generation::Sampling {
|
||||
use candle_transformers::generation::Sampling;
|
||||
if self.temperature <= 0.0 {
|
||||
return Sampling::ArgMax;
|
||||
}
|
||||
let temperature = self.temperature;
|
||||
match (self.top_k, self.top_p) {
|
||||
(Some(k), Some(p)) => Sampling::TopKThenTopP { k, p, temperature },
|
||||
(Some(k), None) => Sampling::TopK { k, temperature },
|
||||
(None, Some(p)) => Sampling::TopP { p, temperature },
|
||||
(None, None) => Sampling::All { temperature },
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_transformers::generation::Sampling;
|
||||
|
||||
fn qwen3_defaults() -> ModelGenerationDefaults {
|
||||
ModelGenerationDefaults {
|
||||
temperature: Some(1.0),
|
||||
top_p: Some(0.95),
|
||||
top_k: Some(20),
|
||||
}
|
||||
}
|
||||
|
||||
/// The case that motivated #271: a caller that names no sampling
|
||||
/// parameters must get the model's published ones, not untruncated
|
||||
/// sampling at a hardcoded 0.7.
|
||||
#[test]
|
||||
fn a_silent_request_inherits_the_models_published_sampling() {
|
||||
let p = SamplingParams::resolve(&RequestedSampling::default(), &qwen3_defaults(), 42);
|
||||
assert_eq!(p.temperature, 1.0);
|
||||
assert_eq!(p.top_p, Some(0.95));
|
||||
assert_eq!(p.top_k, Some(20));
|
||||
assert_eq!(
|
||||
p.to_sampling(),
|
||||
Sampling::TopKThenTopP {
|
||||
k: 20,
|
||||
p: 0.95,
|
||||
temperature: 1.0
|
||||
},
|
||||
"the combination Qwen3 asks for must be expressible"
|
||||
);
|
||||
}
|
||||
|
||||
/// The caller still wins. A harness that deliberately sets a low
|
||||
/// temperature must not have it overridden by the model's file.
|
||||
#[test]
|
||||
fn the_request_overrides_the_model() {
|
||||
let requested = RequestedSampling {
|
||||
temperature: Some(0.2),
|
||||
top_k: Some(5),
|
||||
..Default::default()
|
||||
};
|
||||
let p = SamplingParams::resolve(&requested, &qwen3_defaults(), 42);
|
||||
assert_eq!(p.temperature, 0.2);
|
||||
assert_eq!(p.top_k, Some(5));
|
||||
// Untouched fields still fall through to the model.
|
||||
assert_eq!(p.top_p, Some(0.95));
|
||||
}
|
||||
|
||||
/// A model with no `generation_config.json` must behave exactly as
|
||||
/// it did before this landed, so nothing regresses on the models
|
||||
/// that do not ship one.
|
||||
#[test]
|
||||
fn without_model_defaults_the_old_behaviour_is_preserved() {
|
||||
let p = SamplingParams::resolve(
|
||||
&RequestedSampling::default(),
|
||||
&ModelGenerationDefaults::default(),
|
||||
42,
|
||||
);
|
||||
assert_eq!(p.temperature, FALLBACK_TEMPERATURE);
|
||||
assert_eq!(p.top_p, None);
|
||||
assert_eq!(p.top_k, None);
|
||||
assert_eq!(p.to_sampling(), Sampling::All { temperature: 0.7 });
|
||||
}
|
||||
|
||||
/// A pinned seed is honoured; an absent one takes the per-request
|
||||
/// random value. Reproducibility is the whole point — the text paths
|
||||
/// previously called `unix_subsec_nanos()` unconditionally, so
|
||||
/// `seed` was accepted and discarded.
|
||||
#[test]
|
||||
fn seed_is_pinned_when_given_and_random_otherwise() {
|
||||
let pinned = SamplingParams::resolve(
|
||||
&RequestedSampling {
|
||||
seed: Some(7),
|
||||
..Default::default()
|
||||
},
|
||||
&ModelGenerationDefaults::default(),
|
||||
42,
|
||||
);
|
||||
assert_eq!(pinned.seed, 7);
|
||||
|
||||
let unpinned = SamplingParams::resolve(
|
||||
&RequestedSampling::default(),
|
||||
&ModelGenerationDefaults::default(),
|
||||
42,
|
||||
);
|
||||
assert_eq!(unpinned.seed, 42);
|
||||
}
|
||||
|
||||
/// Greedy beats every truncation setting: a caller asking for
|
||||
/// `temperature = 0` wants determinism, not a softened version of it.
|
||||
#[test]
|
||||
fn zero_temperature_is_greedy_even_with_model_top_k() {
|
||||
let requested = RequestedSampling {
|
||||
temperature: Some(0.0),
|
||||
..Default::default()
|
||||
};
|
||||
let p = SamplingParams::resolve(&requested, &qwen3_defaults(), 42);
|
||||
assert_eq!(p.to_sampling(), Sampling::ArgMax);
|
||||
}
|
||||
|
||||
/// Each truncation knob alone selects the matching strategy.
|
||||
#[test]
|
||||
fn single_knobs_select_their_own_strategy() {
|
||||
let only_k = SamplingParams::resolve(
|
||||
&RequestedSampling {
|
||||
top_k: Some(20),
|
||||
..Default::default()
|
||||
},
|
||||
&ModelGenerationDefaults::default(),
|
||||
1,
|
||||
);
|
||||
assert_eq!(
|
||||
only_k.to_sampling(),
|
||||
Sampling::TopK {
|
||||
k: 20,
|
||||
temperature: 0.7
|
||||
}
|
||||
);
|
||||
|
||||
let only_p = SamplingParams::resolve(
|
||||
&RequestedSampling {
|
||||
top_p: Some(0.9),
|
||||
..Default::default()
|
||||
},
|
||||
&ModelGenerationDefaults::default(),
|
||||
1,
|
||||
);
|
||||
assert_eq!(
|
||||
only_p.to_sampling(),
|
||||
Sampling::TopP {
|
||||
p: 0.9,
|
||||
temperature: 0.7
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/// Repetition penalty and its window are now per-request, defaulting
|
||||
/// to the values that were previously constants.
|
||||
#[test]
|
||||
fn repetition_knobs_default_to_the_former_constants() {
|
||||
let d = SamplingParams::resolve(
|
||||
&RequestedSampling::default(),
|
||||
&ModelGenerationDefaults::default(),
|
||||
1,
|
||||
);
|
||||
assert_eq!(d.repeat_penalty, DEFAULT_REPEAT_PENALTY);
|
||||
assert_eq!(d.repeat_last_n, DEFAULT_REPEAT_LAST_N);
|
||||
|
||||
let tuned = SamplingParams::resolve(
|
||||
&RequestedSampling {
|
||||
repeat_penalty: Some(1.05),
|
||||
repeat_last_n: Some(512),
|
||||
..Default::default()
|
||||
},
|
||||
&ModelGenerationDefaults::default(),
|
||||
1,
|
||||
);
|
||||
assert_eq!(tuned.repeat_penalty, 1.05);
|
||||
assert_eq!(tuned.repeat_last_n, 512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_generation_config_is_not_an_error() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let d = ModelGenerationDefaults::load_from_dir(dir.path());
|
||||
assert!(d.temperature.is_none() && d.top_p.is_none() && d.top_k.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_real_generation_config_parses() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
// Verbatim from Qwen/Qwen3.8-27B, extra fields included — the
|
||||
// reader must tolerate the ones it does not model.
|
||||
std::fs::write(
|
||||
dir.path().join("generation_config.json"),
|
||||
r#"{"bos_token_id":248044,"do_sample":true,"eos_token_id":[248046,248044],
|
||||
"pad_token_id":248044,"temperature":1.0,"top_k":20,"top_p":0.95}"#,
|
||||
)
|
||||
.expect("write");
|
||||
let d = ModelGenerationDefaults::load_from_dir(dir.path());
|
||||
assert_eq!(d.temperature, Some(1.0));
|
||||
assert_eq!(d.top_p, Some(0.95));
|
||||
assert_eq!(d.top_k, Some(20));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_generation_config_falls_back_rather_than_failing() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
std::fs::write(dir.path().join("generation_config.json"), "{not json").expect("write");
|
||||
let d = ModelGenerationDefaults::load_from_dir(dir.path());
|
||||
assert!(d.temperature.is_none());
|
||||
}
|
||||
}
|
||||
@@ -84,6 +84,13 @@ pub enum InferenceEvent {
|
||||
/// `output_tokens_details.reasoning_tokens` (responses).
|
||||
/// Zero for non-reasoning models.
|
||||
reasoning_tokens: u32,
|
||||
/// Leading prompt tokens served from the prefix cache (#269) — a
|
||||
/// sub-count of `prompt_tokens`, never added to `total_tokens`.
|
||||
/// Surfaced as `prompt_tokens_details.cached_tokens` (chat) /
|
||||
/// `input_tokens_details.cached_tokens` (responses). The cache
|
||||
/// has worked since #11; without this the saving was invisible
|
||||
/// and every client reported a 0% hit rate.
|
||||
cached_tokens: u32,
|
||||
/// Server-measured prefill/decode timing for the request, or
|
||||
/// `None` on paths that don't measure it (CPU fallback that
|
||||
/// doesn't instrument, tests). Streaming projectors surface
|
||||
|
||||
@@ -27,7 +27,8 @@
|
||||
//! propagate without us writing any logic.
|
||||
|
||||
use cortex_core::openai::{
|
||||
ChatCompletionChunk, ChunkChoice, CompletionTokensDetails, HelexaTiming, Usage,
|
||||
ChatCompletionChunk, ChunkChoice, CompletionTokensDetails, HelexaTiming, PromptTokensDetails,
|
||||
Usage,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tokio::sync::mpsc;
|
||||
@@ -219,6 +220,7 @@ pub fn project_chat_stream_with(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
cached_tokens,
|
||||
timing,
|
||||
} => {
|
||||
// The finish_reason chunk, then an OpenAI-style
|
||||
@@ -233,9 +235,12 @@ pub fn project_chat_stream_with(
|
||||
&id,
|
||||
created,
|
||||
&model_id,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
UsageCounts {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
cached_tokens,
|
||||
},
|
||||
timing,
|
||||
),
|
||||
]
|
||||
@@ -385,15 +390,32 @@ fn final_chunk(
|
||||
/// ignore the empty-choices chunk; clients that do (opencode, and
|
||||
/// cortex's Anthropic translator) get the token counts they need to
|
||||
/// track context.
|
||||
/// The token counts a finished turn reports. Grouped rather than passed
|
||||
/// loose: adding `cached_tokens` (#269) put the list past the point
|
||||
/// where four adjacent bare `u32`s are safe to read at a call site.
|
||||
struct UsageCounts {
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
/// Sub-count of `completion_tokens` — tokens spent reasoning.
|
||||
reasoning_tokens: u32,
|
||||
/// Sub-count of `prompt_tokens` — tokens served from the prefix
|
||||
/// cache. Neither sub-count is added into `total_tokens`.
|
||||
cached_tokens: u32,
|
||||
}
|
||||
|
||||
fn usage_chunk(
|
||||
id: &str,
|
||||
created: u64,
|
||||
model_id: &str,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
reasoning_tokens: u32,
|
||||
counts: UsageCounts,
|
||||
timing: Option<FinishTiming>,
|
||||
) -> ChatCompletionChunk {
|
||||
let UsageCounts {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
cached_tokens,
|
||||
} = counts;
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
@@ -409,7 +431,12 @@ fn usage_chunk(
|
||||
completion_tokens_details: (reasoning_tokens > 0).then_some(CompletionTokensDetails {
|
||||
reasoning_tokens: reasoning_tokens as u64,
|
||||
}),
|
||||
prompt_tokens_details: None,
|
||||
// Prefix-cache reuse (#269). Omitted at zero so a client
|
||||
// sees unchanged JSON when nothing was cached, and so
|
||||
// "absent" stays distinguishable from "measured as none".
|
||||
prompt_tokens_details: (cached_tokens > 0).then_some(PromptTokensDetails {
|
||||
cached_tokens: cached_tokens as u64,
|
||||
}),
|
||||
// helexa extension (#85): server-measured prefill/decode
|
||||
// timing for the bench harness. Omitted on paths that don't
|
||||
// measure it so standard clients see unchanged JSON.
|
||||
@@ -458,6 +485,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -496,6 +524,7 @@ mod tests {
|
||||
prompt_tokens: 128,
|
||||
completion_tokens: 64,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: Some(FinishTiming {
|
||||
prefill_ms: 200,
|
||||
decode_ms: 1500,
|
||||
@@ -520,6 +549,73 @@ mod tests {
|
||||
assert_eq!(timing.prefill_tokens, 128);
|
||||
}
|
||||
|
||||
/// A cache hit must reach the caller (#269).
|
||||
///
|
||||
/// Prefix caching has worked since #11 — neuron logs `reused=2068`
|
||||
/// on prefills — but nothing carried the number onto the wire, so
|
||||
/// every client reported a 0% hit rate on a cache that was saving
|
||||
/// real prefill. dsh showed `Cache hit 0%` against a 42.1K-token
|
||||
/// context, and an agentic harness that believes nothing is cached
|
||||
/// may restructure the very prefix that was being reused.
|
||||
#[tokio::test]
|
||||
async fn cached_tokens_reach_the_usage_chunk() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 7413,
|
||||
completion_tokens: 12,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 2068,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let out = collect(out_rx).await;
|
||||
let usage = out
|
||||
.iter()
|
||||
.find_map(|c| c.usage.as_ref())
|
||||
.expect("usage chunk present");
|
||||
let details = usage
|
||||
.prompt_tokens_details
|
||||
.as_ref()
|
||||
.expect("prompt_tokens_details populated when the cache was hit");
|
||||
assert_eq!(details.cached_tokens, 2068);
|
||||
// A sub-count of prompt_tokens, never added into the total —
|
||||
// otherwise a cache hit would inflate what a caller is billed.
|
||||
assert_eq!(usage.prompt_tokens, 7413);
|
||||
assert_eq!(usage.total_tokens, 7425);
|
||||
}
|
||||
|
||||
/// No reuse means no `prompt_tokens_details` at all, so a client
|
||||
/// sees byte-identical JSON to before this landed — and "absent"
|
||||
/// stays distinguishable from "measured as zero".
|
||||
#[tokio::test]
|
||||
async fn no_cache_hit_omits_the_details_object() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 1,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let usage = collect(out_rx)
|
||||
.await
|
||||
.iter()
|
||||
.find_map(|c| c.usage.clone())
|
||||
.expect("usage chunk present");
|
||||
assert!(usage.prompt_tokens_details.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_text_delta_is_dropped() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
@@ -541,6 +637,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -637,6 +734,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -689,6 +787,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -735,6 +834,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -781,6 +881,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -809,6 +910,7 @@ mod tests {
|
||||
prompt_tokens: 42,
|
||||
completion_tokens: 5,
|
||||
reasoning_tokens: 2,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -843,6 +945,7 @@ mod tests {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 7,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
|
||||
@@ -33,9 +33,9 @@
|
||||
|
||||
use cortex_core::openai::{ChatCompletionRequest, ChatMessage, MessageContent};
|
||||
use cortex_core::responses::{
|
||||
OutputTokensDetails, ResponsesContentPart, ResponsesInput, ResponsesInputElement,
|
||||
ResponsesInputItem, ResponsesMessageContent, ResponsesOutputContent, ResponsesOutputItem,
|
||||
ResponsesRequest, ResponsesResponse, ResponsesUsage, events,
|
||||
InputTokensDetails, OutputTokensDetails, ResponsesContentPart, ResponsesInput,
|
||||
ResponsesInputElement, ResponsesInputItem, ResponsesMessageContent, ResponsesOutputContent,
|
||||
ResponsesOutputItem, ResponsesRequest, ResponsesResponse, ResponsesUsage, events,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc;
|
||||
@@ -156,6 +156,10 @@ pub fn request_to_chat(req: ResponsesRequest) -> Result<ChatCompletionRequest, T
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
top_p: req.top_p,
|
||||
top_k: req.top_k,
|
||||
seed: req.seed,
|
||||
repetition_penalty: req.repetition_penalty,
|
||||
repeat_last_n: req.repeat_last_n,
|
||||
max_tokens: req.max_output_tokens,
|
||||
// Responses spells the cap `max_output_tokens`, already
|
||||
// resolved onto the legacy field above.
|
||||
@@ -540,6 +544,7 @@ async fn run_projection(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
reasoning_tokens,
|
||||
cached_tokens,
|
||||
// Responses-side `helexa_timing` surfacing not wired yet;
|
||||
// the bench harness reads timing off the chat path (#85).
|
||||
timing: _,
|
||||
@@ -556,7 +561,12 @@ async fn run_projection(
|
||||
output_tokens_details: (reasoning_tokens > 0).then_some(OutputTokensDetails {
|
||||
reasoning_tokens: reasoning_tokens as u64,
|
||||
}),
|
||||
input_tokens_details: None,
|
||||
// Prefix-cache reuse (#269). Omitted at zero so a
|
||||
// client sees unchanged JSON when nothing was
|
||||
// cached.
|
||||
input_tokens_details: (cached_tokens > 0).then_some(InputTokensDetails {
|
||||
cached_tokens: cached_tokens as u64,
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -1002,6 +1012,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1024,6 +1038,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1047,6 +1065,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: Some("resp_prev".into()),
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1079,6 +1101,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1109,6 +1135,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1154,6 +1184,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1191,6 +1225,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1217,6 +1255,10 @@ mod tests {
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
seed: None,
|
||||
repetition_penalty: None,
|
||||
repeat_last_n: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
@@ -1351,6 +1393,41 @@ mod tests {
|
||||
/// timeout could never finish. Its watchdog resets on any parsed
|
||||
/// event, which is why emitting these fixes it and why SSE comment
|
||||
/// keep-alives could not: idle timers count events, not comments.
|
||||
/// The Responses surface reports cache reuse too (#269), in its own
|
||||
/// spelling: `input_tokens_details.cached_tokens`.
|
||||
///
|
||||
/// This is the surface dsh uses, and the one that showed
|
||||
/// `Cache hit 0%` against a 42.1K-token context while neuron was
|
||||
/// reusing thousands of tokens per prefill.
|
||||
#[tokio::test]
|
||||
async fn cached_tokens_reach_the_completed_usage() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("ok".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
prompt_tokens: 7413,
|
||||
completion_tokens: 12,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 2068,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let usage = &frames.last().unwrap().data["response"]["usage"];
|
||||
assert_eq!(usage["input_tokens_details"]["cached_tokens"], 2068);
|
||||
assert_eq!(usage["input_tokens"], 7413);
|
||||
// Sub-count, not additive — a cache hit must not inflate the total.
|
||||
assert_eq!(usage["total_tokens"], 7425);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_projects_a_full_item_lifecycle_before_the_message() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(16);
|
||||
@@ -1371,6 +1448,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 2,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1444,6 +1522,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 1,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1487,6 +1566,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 1,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1522,6 +1602,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1583,6 +1664,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1655,6 +1737,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1684,6 +1767,7 @@ mod tests {
|
||||
prompt_tokens: 30,
|
||||
completion_tokens: 12,
|
||||
reasoning_tokens: 4,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1715,6 +1799,7 @@ mod tests {
|
||||
prompt_tokens: 8,
|
||||
completion_tokens: 3,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1740,6 +1825,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
@@ -1787,6 +1873,7 @@ mod tests {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
timing: None,
|
||||
})
|
||||
.await
|
||||
|
||||
Reference in New Issue
Block a user