hugo/DS4Server
1
0 Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Reject invalid Qwen logits and discard poisoned sessions

Browse Source
This commit is contained in:
Georg Bauer committed 2026-10-04 13:29:09 +02:00
1 parent 1de78fa3ef
commit 5ecf8ab4f3
31 files changed
+1812 -746
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/engine.rs
+192 -110
View File
@@ -844,54 +844,68 @@ impl Generator {
mut progress: impl FnMut(u32, u32, Option<f32>),
mut phase: impl FnMut(&'static str),
) -> Result<GenerationOutput, String> {
let checkpoint_present = checkpoint.is_file();
let selected = self.select_checkpoint(checkpoint, |tag| {
checkpoint_matches_prefix(
tag,
&settings.system_prompt,
settings.reasoning_mode,
messages,
)
})?;
if checkpoint_present && !selected.found {
phase(checkpoint_rebuild_activity(selected.incompatible));
}
if !selected.found
&& let Some(directory) = bootstrap
{
self.prepare_bootstrap(directory, settings, cancelled, &mut progress, &mut phase)?;
}
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
if prompt_complete {
self.executor.finalize_pending_sampled_token()?;
}
let mut completed = messages.to_vec();
completed.push(output.message.clone());
output.previous_checkpoint_bytes =
std::fs::metadata(checkpoint).ok().map(|item| item.len());
if matches!(self.executor, metal::Executor::QwenMtplx(_)) && !prompt_complete {
// The whole-turn owner retains only committed state. Do not stamp
// an old prefix as the partially emitted conversation or replace a
// valid file on cancellation/literal-stop.
output.checkpoint_bytes = output.previous_checkpoint_bytes.unwrap_or(0);
return Ok(output);
}
self.save_checkpoint(
checkpoint,
let result: Result<_, String> = (|| {
let checkpoint_present = checkpoint.is_file();
let selected = self.select_checkpoint(checkpoint, |tag| {
checkpoint_matches_prefix(
tag,
&settings.system_prompt,
settings.reasoning_mode,
messages,
)
})?;
if checkpoint_present && !selected.found {
phase(checkpoint_rebuild_activity(selected.incompatible));
}
if !selected.found
&& let Some(directory) = bootstrap
{
self.prepare_bootstrap(directory, settings, cancelled, &mut progress, &mut phase)?;
}
let result =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
if prompt_complete {
conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed)
} else {
[0; 32]
},
)?;
output.checkpoint_bytes = std::fs::metadata(checkpoint)
.map(|item| item.len())
.unwrap_or(0);
self.checkpoint = Some(checkpoint.to_owned());
self.resident_active = Some(checkpoint.to_owned());
Ok(output)
self.executor.finalize_pending_sampled_token()?;
}
let mut completed = messages.to_vec();
completed.push(output.message.clone());
output.previous_checkpoint_bytes =
std::fs::metadata(checkpoint).ok().map(|item| item.len());
if matches!(self.executor, metal::Executor::QwenMtplx(_)) && !prompt_complete {
// The whole-turn owner retains only committed state. Do not stamp
// an old prefix as the partially emitted conversation or replace a
// valid file on cancellation/literal-stop.
output.checkpoint_bytes = output.previous_checkpoint_bytes.unwrap_or(0);
return Ok(output);
}
self.save_checkpoint(
checkpoint,
if prompt_complete {
conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed)
} else {
[0; 32]
},
)?;
output.checkpoint_bytes = std::fs::metadata(checkpoint)
.map(|item| item.len())
.unwrap_or(0);
self.checkpoint = Some(checkpoint.to_owned());
self.resident_active = Some(checkpoint.to_owned());
Ok(output)
})();
if let Err(error) = &result {
let mut paths = vec![checkpoint.to_owned()];
if let Some(directory) = bootstrap {
let key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &[]);
let tag = Sha256::digest(&key).into();
paths.push(directory.join(format!("{}.bin", kvstore::key_hash(&key))));
paths.push(resident_key(directory, tag));
}
self.invalidate_qwen_failure(error, paths)?;
}
result
}
fn publish_execution_stats(&self) {
@@ -939,73 +953,140 @@ impl Generator {
mut emit: impl FnMut(bool, String),
mut progress: impl FnMut(u32, u32, Option<f32>),
) -> Result<GenerationOutput, String> {
let history = messages
.split_last()
.map_or(messages, |(_, history)| history);
let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
let history_key =
conversation_key(&settings.system_prompt, settings.reasoning_mode, history);
let history_tag: [u8; 32] = Sha256::digest(&history_key).into();
let previous_checkpoint =
self.restore_cached_prefix(directory, &store, &history_key, history_tag)?;
let result: Result<_, String> = (|| {
let history = messages
.split_last()
.map_or(messages, |(_, history)| history);
let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
let history_key =
conversation_key(&settings.system_prompt, settings.reasoning_mode, history);
let history_tag: [u8; 32] = Sha256::digest(&history_key).into();
let previous_checkpoint =
self.restore_cached_prefix(directory, &store, &history_key, history_tag)?;
let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
if prompt_complete {
self.executor.finalize_pending_sampled_token()?;
}
let mut completed = messages.to_vec();
completed.push(output.message.clone());
output.previous_checkpoint_bytes = previous_checkpoint
.as_deref()
.and_then(|path| std::fs::metadata(path).ok())
.map(|item| item.len());
if prompt_complete {
let completed_key =
conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed);
let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into();
if !settings.kv_cache.stores(
self.executor.position(),
history.is_empty(),
self.last_store_tokens,
) {
// A gated store still leaves the finished conversation in the
// live KV, so mark it and drop the stale file association.
self.executor.note_checkpoint_tag(completed_tag);
self.checkpoint = None;
self.resident_active = Some(resident_key(directory, completed_tag));
return Ok(output);
let result =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
self.publish_execution_stats();
let (mut output, prompt_complete) = result?;
if prompt_complete {
self.executor.finalize_pending_sampled_token()?;
}
let completed_checkpoint = store.checkpoint_path(&completed_key);
self.save_checkpoint(&completed_checkpoint, completed_tag)?;
let retained = store.record(
&completed_checkpoint,
&completed_key,
completed_tag,
self.executor.position(),
self.executor.context(),
if history.is_empty() {
StoreReason::Cold
let mut completed = messages.to_vec();
completed.push(output.message.clone());
output.previous_checkpoint_bytes = previous_checkpoint
.as_deref()
.and_then(|path| std::fs::metadata(path).ok())
.map(|item| item.len());
if prompt_complete {
let completed_key =
conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed);
let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into();
if !settings.kv_cache.stores(
self.executor.position(),
history.is_empty(),
self.last_store_tokens,
) {
// A gated store still leaves the finished conversation in the
// live KV, so mark it and drop the stale file association.
self.executor.note_checkpoint_tag(completed_tag);
self.checkpoint = None;
self.resident_active = Some(resident_key(directory, completed_tag));
return Ok(output);
}
let completed_checkpoint = store.checkpoint_path(&completed_key);
self.save_checkpoint(&completed_checkpoint, completed_tag)?;
let retained = store.record(
&completed_checkpoint,
&completed_key,
completed_tag,
self.executor.position(),
self.executor.context(),
if history.is_empty() {
StoreReason::Cold
} else {
StoreReason::Continued
},
)?;
output.checkpoint_bytes = retained
.then(|| std::fs::metadata(&completed_checkpoint).ok())
.flatten()
.map_or(0, |item| item.len());
if retained {
self.last_store_tokens = self.executor.position();
}
self.checkpoint = retained.then_some(completed_checkpoint);
if let Some(checkpoint) = &self.checkpoint {
self.resident_active = Some(checkpoint.clone());
} else {
StoreReason::Continued
},
)?;
output.checkpoint_bytes = retained
.then(|| std::fs::metadata(&completed_checkpoint).ok())
.flatten()
.map_or(0, |item| item.len());
if retained {
self.last_store_tokens = self.executor.position();
self.resident_active = Some(resident_key(directory, completed_tag));
}
}
self.checkpoint = retained.then_some(completed_checkpoint);
if let Some(checkpoint) = &self.checkpoint {
self.resident_active = Some(checkpoint.clone());
} else {
self.resident_active = Some(resident_key(directory, completed_tag));
Ok(output)
})();
if let Err(error) = &result {
// This server uses content-addressed conversation prefixes rather
// than MTPLX session IDs. Invalidate every ancestor of this turn,
// including prefixes no longer selected in memory or on disk.
let mut paths = Vec::new();
for count in 0..=messages.len() {
let key = conversation_key(
&settings.system_prompt,
settings.reasoning_mode,
&messages[..count],
);
let tag = Sha256::digest(&key).into();
paths.push(directory.join(format!("{}.bin", kvstore::key_hash(&key))));
paths.push(resident_key(directory, tag));
}
self.invalidate_qwen_failure(error, paths)?;
}
result
}
fn invalidate_qwen_failure(
&mut self,
error: &str,
mut paths: Vec<PathBuf>,
) -> Result<(), String> {
if !matches!(self.executor, metal::Executor::QwenMtplx(_))
|| !qwen::is_non_finite_logits(error)
{
return Ok(());
}
paths.extend(self.checkpoint.take());
paths.extend(self.resident_active.take());
self.last_store_tokens = 0;
paths.sort();
paths.dedup();
let mut failures = Vec::new();
for path in paths {
self.resident_sessions.remove(&path);
for file in [
path.clone(),
path.with_extension("meta"),
path.with_extension("tmp"),
] {
if let Err(failure) = std::fs::remove_file(&file)
&& failure.kind() != std::io::ErrorKind::NotFound
{
failures.push(format!("{}: {failure}", file.display()));
}
}
}
Ok(output)
// Persistence is synchronous on the serialized model worker. Clearing
// the owner before returning the error also forbids a later save or
// outgoing resident swap from republishing the failed turn.
if let Err(failure) = self.executor.reset() {
failures.push(failure);
}
if failures.is_empty() {
Ok(())
} else {
Err(format!(
"{error}; Qwen session invalidation failed: {}",
failures.join("; ")
))
}
}
fn prepare_bootstrap(
@@ -1119,7 +1200,7 @@ impl Generator {
let _ = std::fs::remove_file(checkpoint);
self.executor.reset()?;
self.checkpoint = None;
let result = (|| {
let result: Result<_, String> = (|| {
phase("Compacting durable task state…");
let mut private_messages = messages.to_vec();
private_messages.push(ChatTurn {
@@ -1224,7 +1305,8 @@ impl Generator {
checkpoint: checkpoint.to_owned(),
})
})();
if result.is_err() {
if let Err(error) = &result {
self.invalidate_qwen_failure(error, vec![checkpoint.to_owned()])?;
let _ = std::fs::remove_file(checkpoint);
self.executor.reset()?;
self.checkpoint = None;
src/engine/kvstore.rs
+1 -1
View File
@@ -305,7 +305,7 @@ fn index_size(key: &[u8]) -> Result<u64, String> {
.ok_or_else(|| "KV cache index is too large".to_owned())
}
fn key_hash(key: &[u8]) -> String {
pub(super) fn key_hash(key: &[u8]) -> String {
let digest: [u8; 32] = Sha256::digest(key).into();
const HEX: &[u8; 16] = b"0123456789abcdef";
let mut output = String::with_capacity(64);
src/engine/metal/qwen_mtplx.rs
+2
View File
@@ -31,6 +31,8 @@ mod indexing;
mod issue103_analysis;
#[cfg(test)]
mod issue104;
#[cfg(test)]
mod issue105;
mod linear;
mod load;
mod mlp;
src/engine/metal/qwen_mtplx/binary.rs
+2
View File
@@ -7,6 +7,7 @@ use super::{bytes, dispatch_geometry, ops, tensor, views};
#[derive(Clone, Copy, Debug)]
pub(super) enum Binary {
Equal,
NotEqual,
Greater,
GreaterEqual,
Less,
@@ -106,6 +107,7 @@ pub(super) fn binary(a: &Array, b: &Array, op: Binary, stream: Stream) -> Result
if matches!(
op,
Binary::Equal
| Binary::NotEqual
| Binary::Greater
| Binary::GreaterEqual
| Binary::Less
src/engine/metal/qwen_mtplx/execution.rs
+37 -6
View File
@@ -226,6 +226,10 @@ pub(super) struct Execution {
pub(super) prefill_throttle: Throttle,
pub(super) decode_throttle: Throttle,
runtime_overrides: super::policy::Overrides,
#[cfg(test)]
pub(super) fault: Option<super::issue105::Fault>,
#[cfg(test)]
pub(super) captured_prefixes: usize,
_context: Context,
}
@@ -335,6 +339,10 @@ impl Execution {
prefill_throttle,
decode_throttle,
runtime_overrides,
#[cfg(test)]
fault: None,
#[cfg(test)]
captured_prefixes: 0,
_context: context,
})
}
@@ -347,7 +355,7 @@ impl Execution {
) -> Result<Output, String> {
let ids = self.input(ids)?;
let options = self.graph_options(phase_prefill, capture)?;
self.model.with_dependent_mut(|_, model| {
let output = self.model.with_dependent_mut(|_, model| {
model.call(
&ids,
Some(&mut self.cache),
@@ -362,7 +370,10 @@ impl Execution {
&self.streams,
self.stream,
)
})
})?;
#[cfg(test)]
let output = self.inject_output(super::issue105::Stage::Target, output)?;
Ok(output)
}
/// Cold text _prefill. Session restore and optional owner-level cleanup /
@@ -775,7 +786,7 @@ impl Execution {
pub(super) fn draft(&mut self, hidden: &Array, ids: &[i32]) -> Result<Output, String> {
let ids = self.input(ids)?;
let options = self.graph_options(false, false)?;
self.model.with_dependent_mut(|_, model| {
let output = self.model.with_dependent_mut(|_, model| {
model.mtp_forward(
hidden,
&ids,
@@ -786,7 +797,10 @@ impl Execution {
&self.streams,
self.stream,
)
})
})?;
#[cfg(test)]
let output = self.inject_output(super::issue105::Stage::Draft, output)?;
Ok(output)
}
pub(super) fn append_mtp_history(
@@ -835,7 +849,7 @@ impl Execution {
self.stream,
&mut self.decode_throttle,
|ids, hidden| {
self.model.with_dependent_mut(|_, model| {
let output = self.model.with_dependent_mut(|_, model| {
model.call(
ids,
Some(&mut self.cache),
@@ -850,7 +864,24 @@ impl Execution {
&self.streams,
self.stream,
)
})
})?;
#[cfg(test)]
let output = Output {
logits: output
.logits
.map(|row| {
super::issue105::inject_logits(
&mut self.fault,
super::issue105::Stage::Target,
&row,
&self.streams,
self.stream,
)
})
.transpose()?,
..output
};
Ok(output)
},
emit,
);
src/engine/metal/qwen_mtplx/fused.rs
+1
View File
@@ -51,6 +51,7 @@ pub(super) fn operator(op: &Operation) -> Option<&'static str> {
},
Operation::Binary(op) => match op {
B::Equal => "Equal",
B::NotEqual => "NotEqual",
B::GreaterEqual => "GreaterEqual",
B::Greater => "Greater",
B::Less => "Less",
src/engine/metal/qwen_mtplx/issue103_analysis.rs
+23 -2
View File
@@ -91,6 +91,10 @@ fn issue103_sampling_receipts() {
serde_json::json!({"case":"sampling_matrix","temperature":temperature,"top_k":top_k,"top_p":top_p,"ordinary_probabilities":probabilities,"lazy_counts":counts.to_vec()})
);
}
let oracle = include_str!("../../../../docs/issue105-invalid-logits/oracle.jsonl")
.lines()
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
.collect::<Vec<_>>();
for (name, values) in [
("nan", vec![f32::NAN; 64]),
("positive_inf", [vec![f32::INFINITY], vec![0.; 63]].concat()),
@@ -123,10 +127,21 @@ fn issue103_sampling_receipts() {
"{}",
serde_json::json!({"case":"invalid_logits","row":name,"temperature":temperature,"ordinary":result.as_ref().map(|r|r.0).map_err(Clone::clone)})
);
let expected = oracle
.iter()
.find(|r| {
r["case"] == "ordinary" && r["row"] == name && r["temperature"] == temperature
})
.unwrap();
if name == "valid_mask" {
assert!([1, 3].contains(&result.unwrap().0));
assert_eq!(result.unwrap().0, expected["token"].as_i64().unwrap());
} else {
assert!(result.is_ok(), "diagnosis no longer reproduces");
let error = result.err().unwrap();
assert!(crate::engine::qwen::is_non_finite_logits(&error));
assert_eq!(
error,
format!("Qwen {}", expected["error"].as_str().unwrap())
);
}
}
let config = SamplerConfig {
@@ -139,6 +154,12 @@ fn issue103_sampling_receipts() {
"{}",
serde_json::json!({"case":"invalid_logits_lazy","row":name,"result":result})
);
let id = result.unwrap()[0];
if name == "valid_mask" {
assert!([1, 3].contains(&id));
} else {
assert_eq!(id, -1);
}
}
}
src/engine/metal/qwen_mtplx/issue105.rs
+141
View File
@@ -0,0 +1,141 @@
//! Test-only, owner-local injection at real model output boundaries.
use super::array::{Array, Dtype};
use super::execution::Execution;
use super::{model::Output, ops, sampling};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Stage {
Prefill,
Target,
Draft,
Capture,
}
#[derive(Clone, Copy, Debug)]
pub(super) enum Row {
Nan,
PositiveInf,
Masked,
LegalMask,
}
pub(super) struct Fault {
pub(super) stage: Stage,
pub(super) row: Row,
}
impl Execution {
pub(super) fn inject_logits(&mut self, stage: Stage, logits: &Array) -> Result<Array, String> {
inject_logits(&mut self.fault, stage, logits, &self.streams, self.stream)
}
pub(super) fn inject_output(&mut self, stage: Stage, output: Output) -> Result<Output, String> {
Ok(Output {
logits: output
.logits
.map(|row| self.inject_logits(stage, &row))
.transpose()?,
..output
})
}
pub(super) fn inject_capture_failure(&mut self, logits: &Array) -> Result<(), String> {
if self
.fault
.as_ref()
.is_some_and(|f| f.stage == Stage::Capture)
{
// Exercise final-capture error ownership, using the real numerical
// error produced by consuming a faulty row. Normal capture does not
// sample or add a validity pass to the oracle's execution route.
let row = self.inject_logits(Stage::Capture, logits)?;
sampling::checked_argmax(&row, &self.streams, self.stream)?;
}
Ok(())
}
}
pub(super) fn inject_logits(
fault: &mut Option<Fault>,
stage: Stage,
logits: &Array,
streams: &super::stream::Streams,
stream: super::stream::Stream,
) -> Result<Array, String> {
let Some(pending) = &*fault else {
return Ok(logits.clone());
};
if pending.stage != stage {
return Ok(logits.clone());
}
let fault = fault.take().unwrap();
// Settle the real computation before replacing its output. The normal
// caches, prefix capture and request/persistence owners all remain live.
ops::evaluate(streams, std::slice::from_ref(logits), stream, false)?;
let shape = logits.layout().shape().to_vec();
let vocab = *shape.last().unwrap() as usize;
let data = (0..logits.layout().size())
.flat_map(|i| {
let value = match fault.row {
Row::Nan => f32::NAN,
Row::PositiveInf if i % vocab == 0 => f32::INFINITY,
Row::PositiveInf => 0.,
Row::Masked => f32::NEG_INFINITY,
Row::LegalMask if i % vocab == 1 => 2.,
Row::LegalMask if i % vocab == 3 => 1.,
Row::LegalMask => f32::NEG_INFINITY,
};
value.to_le_bytes()
})
.collect::<Vec<_>>();
Array::new(&shape, Dtype::F32, super::scalar_buffer(&data)?)
}
#[test]
#[ignore = "Metal required; run serially under test-supervisor"]
fn issue105_batched_greedy_rejects_each_invalid_row() {
super::configure_sources().unwrap();
let _context = super::super::gpu::Context::open_qwen(0).unwrap();
let streams = super::stream::Streams::new(Some(super::allocator::Allocator::new().unwrap()));
let stream = streams.default_stream(super::stream::Device::Gpu).unwrap();
for width in [64, 4097] {
for row in [Row::Nan, Row::PositiveInf, Row::Masked, Row::LegalMask] {
for position in [0, 1] {
let healthy = Array::new(
&[1, width],
Dtype::F32,
super::scalar_buffer(&vec![0; width as usize * 4]).unwrap(),
)
.unwrap();
let injected = inject_logits(
&mut Some(Fault {
stage: Stage::Target,
row,
}),
Stage::Target,
&healthy,
&streams,
stream,
)
.unwrap();
let rows = if position == 0 {
[injected, healthy]
} else {
[healthy, injected]
};
let logits = ops::concatenate(&rows, 0, stream).unwrap();
let result = sampling::checked_argmax(&logits, &streams, stream);
if matches!(row, Row::LegalMask) {
let ids = ops::read_i64(&result.unwrap(), &streams, stream).unwrap();
assert_eq!(ids[position], 1);
assert_eq!(ids[1 - position], 0);
} else {
assert!(crate::engine::qwen::is_non_finite_logits(
&result.err().unwrap()
));
}
}
}
}
streams.clear_streams().unwrap();
}
src/engine/metal/qwen_mtplx/issue105_product.rs
+298
View File
@@ -0,0 +1,298 @@
//! Installed-model faults through Generator and both HTTP response adapters.
use super::super::issue105::{Fault, Row, Stage};
use super::*;
use crate::engine::{
Generator, conversation_key, conversation_tag,
kvstore::{KvStore, StoreReason},
metal::Executor,
};
use crate::runtime::{ActiveGeneration, GenerationEvent};
use crate::settings::{
ExecutionPreferences, GenerationPreferences, ReasoningMode, RuntimePreferences,
SpeculativePreferences, effective_settings,
};
use std::sync::{Arc, mpsc};
fn owner(generator: &mut Generator) -> &mut QwenExecutor {
match &mut generator.executor {
Executor::QwenMtplx(owner) => owner,
_ => unreachable!(),
}
}
fn message(text: &str) -> ChatTurn {
ChatTurn {
user: true,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: text.into(),
}
}
fn generate(
generator: &mut Generator,
path: &Path,
messages: &[ChatTurn],
settings: &TurnSettings,
) -> GenerationOutput {
generator
.generate(
path,
None,
messages,
settings,
&AtomicBool::new(false),
|_, _| {},
|_, _, _| {},
|_| {},
)
.unwrap()
}
#[test]
#[ignore = "installed Qwen model; serial test-supervisor with memory guard required"]
fn issue105_product_ar_cleanup_and_http_retry() {
product(false);
}
#[test]
#[ignore = "installed Qwen MTP model; serial test-supervisor with memory guard required"]
fn issue105_product_mtp_cleanup_and_http_retry() {
product(true);
}
fn product(mtp: bool) {
let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE105_DIR").unwrap());
std::fs::create_dir_all(&directory).unwrap();
let effective = effective_settings(
crate::model::ModelChoice::Qwen38FlashNext,
&GenerationPreferences {
context_tokens: 4096,
max_generated_tokens: 4,
system_prompt: "Be concise.".into(),
temperature: Some(0.6),
top_p: Some(0.95),
min_p: Some(0.),
reasoning_mode: ReasoningMode::Low,
..Default::default()
},
&RuntimePreferences {
execution: ExecutionPreferences {
power_percent: Some(100),
..Default::default()
},
speculative: SpeculativePreferences {
glm_mtp: mtp,
..Default::default()
},
..Default::default()
},
root.parent().unwrap(),
)
.unwrap();
let mut settings = effective.turn;
settings.seed = Some(7);
settings.kv_cache.min_tokens = 0;
settings.kv_cache.cold_max_tokens = 4096;
settings.kv_cache.continued_interval_tokens = 1;
let metrics = Arc::new(Metrics::new(&directory.join("metrics")));
eprintln!("#105 load product mtp={mtp}");
let mut generator = Generator::open(&effective.engine, metrics).unwrap();
generator.resident_limit = 8;
owner(&mut generator).policy.store_min_suffix = Some(0);
let baseline_path = directory.join("baseline.bin");
let mut messages = vec![message("Count from one.")];
let baseline = generate(&mut generator, &baseline_path, &messages, &settings);
messages.push(baseline.message);
let tag = conversation_tag(&settings.system_prompt, settings.reasoning_mode, &messages);
let key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &messages);
messages.push(message("Continue."));
let healthy_path = directory.join("healthy.bin");
generate(
&mut generator,
&healthy_path,
&[message("Say hello.")],
&settings,
);
let healthy_bytes = std::fs::read(&healthy_path).unwrap();
let bootstrap = directory.join("bootstrap");
let transient = directory.join("transient");
let mut cases = Vec::new();
for stage in [Stage::Prefill, Stage::Target, Stage::Capture]
.into_iter()
.chain(mtp.then_some(Stage::Draft))
{
for row in [Row::Nan, Row::PositiveInf, Row::Masked, Row::LegalMask] {
for streaming in [false, true] {
cases.push((stage, row, streaming));
}
}
}
for (index, (stage, row, streaming)) in cases.into_iter().enumerate() {
settings.temperature = if index % 4 < 2 { 0. } else { 0.6 };
// The real context-copy route can satisfy this short prompt without
// visiting the draft head. A supported penalty disables copy, ensuring
// draft injection is reached; target cases cover both routes.
owner(&mut generator).policy.sampler.presence_penalty =
if stage == Stage::Draft || (stage == Stage::Target && streaming) {
0.1
} else {
0.0
};
let path = if streaming {
let store = KvStore::open(&transient, settings.kv_cache.budget_bytes).unwrap();
let path = store.checkpoint_path(&key);
std::fs::copy(&baseline_path, &path).unwrap();
store
.record(&path, &key, tag, 1, 4096, StoreReason::Cold)
.unwrap();
path
} else {
let path = directory.join("failed.bin");
if index == 0 {
// Exercise a real bootstrap created by this failing cold turn.
assert!(!path.exists());
} else {
std::fs::copy(&baseline_path, &path).unwrap();
}
path
};
let before = owner(&mut generator).execution.captured_prefixes;
owner(&mut generator).execution.fault = Some(Fault { stage, row });
let (events, receiver) = mpsc::channel();
let cancelled = Arc::new(AtomicBool::new(false));
eprintln!(
"#105 case {index}: mtp={mtp} stage={stage:?} row={row:?} streaming={streaming} temperature={}",
settings.temperature
);
let chunks = Cell::new(0);
let emit = |reasoning, content| {
chunks.set(chunks.get() + 1);
events
.send(GenerationEvent::Chunk { reasoning, content })
.unwrap();
};
let progress = |used, limit, tokens_per_second| {
events
.send(GenerationEvent::Context {
used,
limit,
tokens_per_second,
})
.unwrap();
};
let result = if streaming {
generator
.generate_transient(&transient, &messages, &settings, &cancelled, emit, progress)
} else {
generator.generate(
&path,
Some(&bootstrap),
&messages,
&settings,
&cancelled,
emit,
progress,
|_| {},
)
};
assert!(
owner(&mut generator).execution.fault.is_none(),
"fault site must be reached: result={:?}",
result.as_ref().map(|o| o.completion_tokens)
);
assert!(
owner(&mut generator).execution.captured_prefixes > before,
"prefill snapshot before failure"
);
let valid = matches!(row, Row::LegalMask);
if !valid {
if stage == Stage::Prefill {
assert_eq!(chunks.get(), 0, "no fabricated token reaches the response");
}
assert!(crate::engine::qwen::is_non_finite_logits(
result.as_ref().err().unwrap()
));
assert!(owner(&mut generator).live.current.is_none());
assert!(owner(&mut generator).live.prefix.is_none());
assert!(owner(&mut generator).execution.boundaries.is_empty());
assert!(generator.checkpoint.is_none());
assert!(generator.resident_active.is_none());
assert_eq!(generator.last_store_tokens, 0);
assert!(!generator.resident_sessions.contains_key(&path));
assert!(!path.exists());
assert!(!path.with_extension("meta").exists());
assert!(!path.with_extension("tmp").exists());
let delayed = directory.join("delayed.bin");
assert!(generator.save_checkpoint(&delayed, tag).is_err());
assert!(!delayed.exists());
if !streaming {
let store = KvStore::open(&bootstrap, settings.kv_cache.budget_bytes).unwrap();
assert!(store.find(&key, 4096).is_none());
}
} else {
assert!(result.is_ok(), "legal masks: {:?}", result.as_ref().err());
}
events.send(GenerationEvent::Finished(result)).unwrap();
drop(events);
let wire = crate::server::issue105_wire(
ActiveGeneration {
events: receiver,
cancel: cancelled,
},
streaming,
);
if valid {
assert!(wire.starts_with("HTTP/1.1 200"));
assert!(!wire.contains("non-finite logits"));
} else if streaming {
assert!(wire.contains("event: error\n"));
assert!(wire.contains("Qwen non-finite logits"));
} else {
assert!(wire.starts_with("HTTP/1.1 500"));
assert!(wire.contains("Qwen non-finite logits"));
}
assert_eq!(std::fs::read(&healthy_path).unwrap(), healthy_bytes);
assert!(
generator.resident_sessions.contains_key(&healthy_path),
"unrelated healthy session remains resident"
);
if !valid {
let retried = if streaming {
generator
.generate_transient(
&transient,
&messages,
&settings,
&AtomicBool::new(false),
|_, _| {},
|_, _, _| {},
)
.unwrap()
} else {
generate(&mut generator, &path, &messages, &settings)
};
assert_eq!(retried.cached_tokens, 0, "failure retries cold");
assert!(retried.completion_tokens > 0);
}
println!(
"{}",
serde_json::json!({"case":index,"mtp":mtp,"stage":format!("{stage:?}"),"row":format!("{row:?}"),"streaming":streaming,"legal":valid,"healthy_preserved":true,"cold_retry":!valid})
);
// Do not let successful diagnostic requests evict the healthy control.
generator.executor.reset().unwrap();
generator.checkpoint = None;
generator.resident_active = None;
generator
.resident_sessions
.retain(|path, _| path == &healthy_path);
}
assert!(
generator
.select_checkpoint(&healthy_path, |_| true)
.unwrap()
.found
);
assert!(owner(&mut generator).live.current.is_some());
}
src/engine/metal/qwen_mtplx/ops.rs
+4
View File
@@ -80,6 +80,7 @@ pub(super) enum Operation {
Sum(Vec<usize>),
Min(Vec<usize>),
Max(Vec<usize>),
Any(Vec<usize>),
RandomBits,
SearchSorted(bool),
ArgMax(usize),
@@ -175,6 +176,9 @@ impl Operation {
Self::Min(axes) => {
super::reduce::evaluate(input, output, axes, super::reduce::Reduction::Min)
}
Self::Any(axes) => {
super::reduce::evaluate(input, output, axes, super::reduce::Reduction::Or)
}
Self::Max(axes) => {
super::reduce::evaluate(input, output, axes, super::reduce::Reduction::Max)
}
src/engine/metal/qwen_mtplx/owner.rs
+102 -77
View File
@@ -74,92 +74,113 @@ impl QwenExecutor {
progress: impl FnMut(u32, u32, Option<f32>),
emit: impl FnMut(bool, String),
) -> Result<(GenerationOutput, bool), String> {
if settings.context_tokens != self.identity.context as i32 {
return Err("Qwen turn context differs from the loaded engine".into());
}
self.execution.vision_prompt = None;
self.execution
.model
.with_dependent_mut(|_, model| model.vision = None);
if messages
.iter()
.any(|m| m.content.contains(crate::engine::VISION_DATA_START))
{
let root = self
.vision_path
.as_ref()
.ok_or("Qwen vision encoder is not installed; download it in Model Manager")?;
if self.vision_tower.is_none() {
self.vision_tower = Some(super::vision::Tower::load(
root,
let result = (|| {
if settings.context_tokens != self.identity.context as i32 {
return Err("Qwen turn context differs from the loaded engine".into());
}
self.execution.vision_prompt = None;
self.execution
.model
.with_dependent_mut(|_, model| model.vision = None);
if messages
.iter()
.any(|m| m.content.contains(crate::engine::VISION_DATA_START))
{
let root = self
.vision_path
.as_ref()
.ok_or("Qwen vision encoder is not installed; download it in Model Manager")?;
if self.vision_tower.is_none() {
self.vision_tower = Some(super::vision::Tower::load(
root,
&self.execution.streams,
self.execution.stream,
)?);
}
let prepared = super::vision::prepare(
cancelled,
&self.metadata,
messages,
settings,
self.vision_tower.as_ref().unwrap(),
&self.execution.streams,
self.execution.stream,
)?);
);
if !self.keep_vision_loaded {
self.vision_tower = None;
}
let (ids, input, identities) = prepared?;
self.execution
.model
.with_dependent_mut(|_, model| model.vision = Some(input));
self.execution.vision_prompt = Some((ids, identities));
}
let prepared = super::vision::prepare(
cancelled,
let baseline = self.stats;
let stats = &mut self.stats;
let result = self.live.generate(
&mut self.execution,
&self.metadata,
messages,
settings,
self.vision_tower.as_ref().unwrap(),
&self.execution.streams,
self.execution.stream,
);
if !self.keep_vision_loaded {
self.vision_tower = None;
&self.policy,
cancelled,
metrics,
progress,
emit,
|event| {
record_speculative_progress(stats, baseline, &event, metrics);
Ok(())
},
)?;
let complete = result
.session
.as_ref()
.is_some_and(|s| s.final_snapshot_error.is_none())
&& self.live.current.is_some();
if let Some(session) = &result.session {
let state_error = match &session.decoded {
TurnDecoded::Ar(decoded) => decoded.final_state_error.as_ref(),
TurnDecoded::Mtp(decoded) => decoded.final_state_error.as_ref(),
};
if let Some(error) = state_error {
return Err(if crate::engine::qwen::is_non_finite_logits(error) {
error.clone()
} else {
format!("Qwen final decode state failed: {error}")
});
}
if let Some(error) = &session.final_snapshot_error {
return Err(if crate::engine::qwen::is_non_finite_logits(error) {
error.clone()
} else {
format!("Qwen final state capture failed: {error}")
});
}
if let Some(error) = &session.prefix_snapshot_error {
eprintln!("Qwen prefix capture failed: {error}");
}
debug_assert_eq!(result.output.cached_tokens as usize, session.cached_tokens);
debug_assert!(session.prefill_seconds >= 0.0 && session.restore_seconds >= 0.0);
debug_assert!(session.decoded.decode_seconds() >= 0.0);
if let TurnDecoded::Mtp(decoded) = &session.decoded
&& decoded.pending.is_some()
{
return Err("Qwen final state still has a pending correction".into());
}
}
let (ids, input, identities) = prepared?;
self.execution
.model
.with_dependent_mut(|_, model| model.vision = Some(input));
self.execution.vision_prompt = Some((ids, identities));
}
let baseline = self.stats;
let stats = &mut self.stats;
let result = self.live.generate(
&mut self.execution,
&self.metadata,
messages,
settings,
&self.policy,
cancelled,
metrics,
progress,
emit,
|event| {
record_speculative_progress(stats, baseline, &event, metrics);
Ok(())
},
)?;
let complete = result
.session
.as_ref()
.is_some_and(|s| s.final_snapshot_error.is_none())
&& self.live.current.is_some();
if let Some(session) = &result.session {
let state_error = match &session.decoded {
TurnDecoded::Ar(decoded) => decoded.final_state_error.as_ref(),
TurnDecoded::Mtp(decoded) => decoded.final_state_error.as_ref(),
};
if let Some(error) = state_error {
return Err(format!("Qwen final decode state failed: {error}"));
}
if let Some(error) = &session.final_snapshot_error {
return Err(format!("Qwen final state capture failed: {error}"));
}
if let Some(error) = &session.prefix_snapshot_error {
eprintln!("Qwen prefix capture failed: {error}");
}
debug_assert_eq!(result.output.cached_tokens as usize, session.cached_tokens);
debug_assert!(session.prefill_seconds >= 0.0 && session.restore_seconds >= 0.0);
debug_assert!(session.decoded.decode_seconds() >= 0.0);
if let TurnDecoded::Mtp(decoded) = &session.decoded
&& decoded.pending.is_some()
{
return Err("Qwen final state still has a pending correction".into());
Ok((result.output, complete))
})();
if let Err(error) = &result
&& crate::engine::qwen::is_non_finite_logits(error)
{
// Drop all snapshots before synchronization, even if the device
// reports a second error while draining outstanding work.
self.live = LiveSession::default();
if let Err(cleanup) = self.clear_execution_cache() {
return Err(format!("{error}; Qwen cache cleanup failed: {cleanup}"));
}
}
Ok((result.output, complete))
result
}
pub(in crate::engine) fn prefill(
@@ -1040,3 +1061,7 @@ fn qwen_vision_chat_checkpoint_preserves_image_identity() {
assert_ne!(changed.message.content, first.message.content);
std::fs::remove_dir_all(directory).unwrap();
}
#[cfg(test)]
#[path = "issue105_product.rs"]
mod issue105_product;
src/engine/metal/qwen_mtplx/random.rs
+8 -1
View File
@@ -533,9 +533,16 @@ fn mtplx_sampler_resident_matches_reference() {
.iter()
.map(|t| ops::read_i64(t, &streams, gpu).unwrap()[0])
.collect::<Vec<_>>();
// Preserve the retired lane's finite-law/RNG fixture for #104.
// #105 replaces only its invalid IDs with the current checked sentinel.
let expected = if case["mode"] == "nonfinite" {
serde_json::json!(vec![-1; 4])
} else {
case["tokens"].clone()
};
assert_eq!(
serde_json::json!(actual),
case["tokens"],
expected,
"sampler case {index}: {case}"
);
assert_eq!(
src/engine/metal/qwen_mtplx/reduce.rs
+15 -2
View File
@@ -227,6 +227,7 @@ pub(super) enum Reduction {
Sum,
Min,
Max,
Or,
}
impl Reduction {
@@ -245,10 +246,21 @@ impl Reduction {
Self::Sum => "sum",
Self::Min => "min",
Self::Max => "max",
Self::Or => "or",
}
}
}
pub(super) fn any(
input: &Array,
axes: &[i32],
keepdims: bool,
stream: Stream,
) -> Result<Array, String> {
let input = ops::astype(input, Dtype::Bool, false, stream)?;
reduction(&input, axes, keepdims, stream, Reduction::Or)
}
pub(super) fn sum(
input: &Array,
axes: &[i32],
@@ -289,7 +301,7 @@ fn reduction(
let layout = input.layout();
let dtype = layout.dtype();
if !matches!(dtype, Dtype::BF16 | Dtype::F16 | Dtype::F32 | Dtype::I32)
&& !(dtype == Dtype::Bool && matches!(op, Reduction::Sum))
&& !(dtype == Dtype::Bool && matches!(op, Reduction::Sum | Reduction::Or))
{
return Err("Qwen reduction requires BF16/F16/F32/I32".into());
}
@@ -303,7 +315,7 @@ fn reduction(
for &axis in &axes {
shape[axis] = 1;
}
let out_dtype = if dtype == Dtype::Bool {
let out_dtype = if dtype == Dtype::Bool && matches!(op, Reduction::Sum) {
Dtype::I32
} else {
dtype
@@ -319,6 +331,7 @@ fn reduction(
Reduction::Sum => ops::Operation::Sum(axes.clone()),
Reduction::Min => ops::Operation::Min(axes.clone()),
Reduction::Max => ops::Operation::Max(axes.clone()),
Reduction::Or => ops::Operation::Any(axes.clone()),
},
)?
};
src/engine/metal/qwen_mtplx/request.rs
+7
View File
@@ -190,6 +190,13 @@ impl LiveSession {
Ok(())
},
);
if result
.as_ref()
.is_err_and(|error| crate::engine::qwen::is_non_finite_logits(error))
{
*self = Self::default();
return result;
}
// A valid prefill capture survives cancellation and unrelated callback
// errors. Never tag those partial turns as a completed conversation.
if let Some(prefix) = captured_prefix {
src/engine/metal/qwen_mtplx/sampling.rs
+118 -38
View File
@@ -3,7 +3,7 @@ use super::array::{Array, Dtype};
use super::binary::{Binary, binary};
use super::stream::{Stream, Streams};
use super::{indexing, normalization, ops, reduce, sort, unary};
use crate::engine::Rng;
use crate::engine::{Rng, qwen::non_finite_logits_error};
use std::borrow::Cow;
#[derive(Clone, Copy)]
@@ -238,7 +238,18 @@ pub(super) fn lazy_sample(
} else {
categorical_row(&top_vals, key, stream)?
};
indexing::take(&top_idx, &local, 0, stream)
let selected = ops::astype(
&indexing::take(&top_idx, &local, 0, stream)?,
Dtype::I32,
false,
stream,
)?;
let sentinel = Array::new(
&[],
Dtype::I32,
super::scalar_buffer(&(-1_i32).to_le_bytes())?,
)?;
indexing::select(&invalid_logits(row, stream)?, &sentinel, &selected, stream)
}
/// generation._distribution_from_mlx_logits, including the single-row GPU
@@ -300,7 +311,7 @@ pub(super) fn distribution_from_logits(
}
/// The ordinary generation._sample_from_logits path. Greedy evaluates model
/// logits before constructing/evaluating ArgMax and never advances PCG64.
/// logits and their validity flag together and never advances PCG64.
pub(super) fn sample_from_logits(
logits: &Array,
config: &SamplerConfig,
@@ -326,8 +337,7 @@ pub(super) fn sample_from_logits(
} else {
logits.clone()
};
ops::evaluate(streams, std::slice::from_ref(&logits), stream, false)?;
let token = reduce::argmax(&logits, -1, false, stream)?;
let token = checked_argmax(&logits, streams, stream)?;
if token.layout().size() != 1 {
return Err("greedy sampler requires a scalar token result".into());
}
@@ -338,6 +348,72 @@ pub(super) fn sample_from_logits(
}
}
/// The pinned greedy validity graph: any(isnan(row)) | isinf(max(row)).
/// A finite maximum permits negative-infinity masks. Read the full row only
/// on failure; valid turns evaluate the token and scalar fault flag together.
fn invalid_logits(logits: &Array, stream: Stream) -> Result<Array, String> {
let vocab = logits.layout().dim(-1)?;
// A valid row must not hide an all-masked neighbour in a batched verify.
// Preserve the ordinary scalar reduction when there is only one row.
let shape = if logits.layout().size() == vocab as usize {
vec![-1]
} else {
vec![-1, vocab]
};
let row = ops::reshape(logits, &shape, stream)?;
let nan = reduce::any(
&binary(&row, &row, Binary::NotEqual, stream)?,
&[-1],
false,
stream,
)?;
let maximum = reduce::max(&row, &[-1], false, stream)?;
let infinity = |value: f32| -> Result<Array, String> {
let scalar = ops::astype(
&super::random::scalar(value)?,
row.layout().dtype(),
false,
stream,
)?;
binary(&maximum, &scalar, Binary::Equal, stream)
};
binary(
&nan,
&binary(
&infinity(f32::INFINITY)?,
&infinity(f32::NEG_INFINITY)?,
Binary::LogicalOr,
stream,
)?,
Binary::LogicalOr,
stream,
)
}
pub(super) fn checked_argmax(
logits: &Array,
streams: &Streams,
stream: Stream,
) -> Result<Array, String> {
let token = reduce::argmax(logits, -1, false, stream)?;
let bad = invalid_logits(logits, stream)?;
ops::evaluate(streams, &[token.clone(), bad.clone()], stream, false)?;
if ops::read_i64(&bad, streams, stream)?
.iter()
.any(|&bad| bad != 0)
{
return Err(non_finite_logits_error(
&ops::read_f32_as_f64(
&ops::astype(logits, Dtype::F32, false, stream)?,
streams,
stream,
)?,
"greedy argmax",
));
}
Ok(token)
}
/// Sparse maps retain their iteration order. Counts cover generated tokens,
/// not prompt tokens. A disabled penalty is an identity, including storage.
pub(super) fn apply_penalties_host<'a>(
@@ -609,7 +685,7 @@ pub(super) struct SparseDistribution {
impl SparseDistribution {
pub(super) fn new(
mut ids: Vec<i64>,
ids: Vec<i64>,
mut probabilities: Vec<f64>,
vocabulary: usize,
) -> Result<Self, String> {
@@ -619,21 +695,15 @@ impl SparseDistribution {
if probabilities.iter().any(|p| p.is_finite() && *p < 0.0) {
return Err("sparse distribution probabilities must be nonnegative".into());
}
let original = probabilities.clone();
for p in &mut probabilities {
if !p.is_finite() || *p <= 0.0 {
*p = 0.0;
}
}
let mut total = pairwise_sum(&probabilities);
let total = pairwise_sum(&probabilities);
if !total.is_finite() || total <= 0.0 {
let token = ids
.iter()
.copied()
.find(|&i| i >= 0 && (i as u64) < vocabulary as u64)
.ok_or("sparse distribution has neither positive mass nor a valid token")?;
ids = vec![token];
probabilities = vec![1.0];
total = 1.0;
return Err(non_finite_logits_error(&original, "SparseDistribution"));
}
for p in &mut probabilities {
*p /= total;
@@ -978,7 +1048,7 @@ pub(super) fn dense_probabilities(
.collect::<Vec<_>>();
let total = pairwise_sum(&p);
if !total.is_finite() || total <= 0.0 {
return Err("cannot normalize logits into a probability distribution".into());
return Err(non_finite_logits_error(logits, "softmax"));
}
for value in &mut p {
*value /= total;
@@ -1053,18 +1123,14 @@ pub(super) fn dense_probabilities(
Ok(probabilities)
}
/// Original fast_sampling._host_sparse_distribution. Its invalid-mass result
/// is token zero; the ordinary dense generator path instead propagates error.
/// fast_sampling._host_sparse_distribution propagates invalid-mass failures.
fn host_sparse(
logits: &[f64],
temperature: f64,
top_p: f64,
top_k: i32,
) -> Result<SparseDistribution, String> {
let probabilities = match dense_probabilities(logits, temperature, top_p, top_k) {
Ok(p) => p,
Err(_) => return SparseDistribution::new(vec![0], vec![1.0], logits.len()),
};
let probabilities = dense_probabilities(logits, temperature, top_p, top_k)?;
let (ids, probabilities) = probabilities
.into_iter()
.enumerate()
@@ -1757,13 +1823,16 @@ fn mtplx_sampler_entry_matches_reference() {
.iter()
.filter(|v| **v == "eval")
.count();
let expected_commits = if profile == "greedy" {
explicit_evals + 1
} else if profile == "sparse" && mode == "nonfinite" {
explicit_evals - 1
} else {
explicit_evals
};
// NumPy's error census implicitly evaluates the F32 cast, which
// the mx.eval trace does not include.
let expected_commits =
if profile == "greedy" && mode == "nonfinite" && dtype != Dtype::F32 {
explicit_evals + 1
} else if profile == "sparse" && mode == "nonfinite" {
explicit_evals - 1
} else {
explicit_evals
};
assert_eq!(
commits() - before,
expected_commits,
@@ -2013,15 +2082,21 @@ fn mtplx_sampler_mixed_distributions_match_reference() {
.map(|i| (i * (vocab - 1) / (count - 1).max(1)) as i64)
.collect::<Vec<_>>();
let selected = ids.iter().map(|&i| p[i as usize]).collect();
Distribution::Sparse(SparseDistribution::new(ids, selected, vocab).unwrap())
SparseDistribution::new(ids, selected, vocab).map(Distribution::Sparse)
} else {
let total = pairwise_sum(&p);
Distribution::Dense(p.iter().map(|p| p / total).collect())
Ok(Distribution::Dense(p.iter().map(|p| p / total).collect()))
}
};
let kinds = case["kinds"].as_str().unwrap().as_bytes();
let target = make(p, kinds[0]);
let draft = make(q, kinds[1]);
if case["construction_error"] == true {
assert!(target.is_err() || draft.is_err());
continue;
}
let target = target.unwrap();
let draft = draft.unwrap();
let mut rng = Rng::new_qwen(12345);
let samples = (0..32)
.map(|_| target.sample(&mut rng))
@@ -2176,9 +2251,14 @@ fn mtplx_sampler_fallback_matches_reference() {
.unwrap(),
)
.unwrap();
let result = sparse_distributions(&logits, 0.6, top_p, 20, &streams, gpu)
.unwrap()
.unwrap();
let result = sparse_distributions(&logits, 0.6, top_p, 20, &streams, gpu);
if case["error"] == true {
assert!(crate::engine::qwen::is_non_finite_logits(
&result.err().unwrap()
));
continue;
}
let result = result.unwrap().unwrap();
assert_eq!(result.len(), 4);
for (row, distribution) in result.iter().enumerate() {
assert_eq!(
@@ -2304,10 +2384,10 @@ fn mtplx_sampler_distribution_and_rng_match_reference() {
assert!(SparseDistribution::new(vec![], vec![], 10).is_err());
assert!(SparseDistribution::new(vec![1], vec![-0.1], 10).is_err());
assert!(SparseDistribution::new(vec![-1, 10], vec![0.0, f64::NAN], 10).is_err());
let fallback =
SparseDistribution::new(vec![-1, 7, 3], vec![f64::NAN, f64::INFINITY, 0.0], 10).unwrap();
assert_eq!(fallback.ids, [7]);
assert_eq!(fallback.probabilities, [1.0]);
let failure = SparseDistribution::new(vec![-1, 7, 3], vec![f64::NAN, f64::INFINITY, 0.0], 10)
.err()
.unwrap();
assert!(crate::engine::qwen::is_non_finite_logits(&failure));
eprintln!(
"MTPLX distributions: {accepted} accepts and {rejected} corrections; exact probabilities, samples and subsequent RNG draws"
);
src/engine/metal/qwen_mtplx/session_cache.rs
+16 -1
View File
@@ -1158,6 +1158,11 @@ impl Execution {
} = self.prepare_session(prompt, saved, options, |event| {
observe(TurnProgress::Prefill(event))
})?;
#[cfg(test)]
let prepared = super::prefill::Prepared {
logits: self.inject_logits(super::issue105::Stage::Prefill, &prepared.logits)?,
..prepared
};
// Publish before decode (and before prefix storage) so a cancelled
// response retains the completed prompt timings, just like MTPLX.
observe(TurnProgress::PromptReady(crate::metrics::PromptTiming {
@@ -1183,6 +1188,10 @@ impl Execution {
.and_then(store_prefix)
})
.and_then(Result::err);
#[cfg(test)]
if options.store_min_suffix.is_some() && prefix_snapshot_error.is_none() {
self.captured_prefixes += 1;
}
observe(TurnProgress::Decode)?;
let decoded = match &options.decode {
DecodeOptions::Ar(decode) => TurnDecoded::Ar(self.decode_ar(
@@ -1216,7 +1225,13 @@ impl Execution {
for &id in decoded.tokens() {
committed.push(i32::try_from(id).map_err(|_| "generated token exceeds int32")?);
}
Some(self.capture_session(&committed, logits, hidden, history, options.lazy_kv))
#[cfg(test)]
let result = self.inject_capture_failure(logits).and_then(|()| {
self.capture_session(&committed, logits, hidden, history, options.lazy_kv)
});
#[cfg(not(test))]
let result = self.capture_session(&committed, logits, hidden, history, options.lazy_kv);
Some(result)
} else {
None
};
src/engine/metal/qwen_mtplx/speculative.rs
+1 -1
View File
@@ -440,7 +440,7 @@ impl Execution {
let logits = verified.logits.ok_or("copy verify omitted logits")?;
let hidden = verified.hidden.ok_or("copy verify omitted hidden")?;
let greedy = if sampler.temperature <= 0.0 {
let choices = super::reduce::argmax(&logits, -1, false, self.stream)?;
let choices = sampling::checked_argmax(&logits, &self.streams, self.stream)?;
Some(ops::read_i64(&choices, &self.streams, self.stream)?)
} else {
ops::evaluate(
src/engine/qwen.rs
+17
View File
@@ -17,6 +17,23 @@ use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::thread;
// The engine's public error channel is String. Keep this family-specific
// discriminator intact through request and persistence owners.
const NON_FINITE_LOGITS: &str = "Qwen non-finite logits in ";
pub(super) fn non_finite_logits_error(row: &[f64], location: &str) -> String {
format!(
"{NON_FINITE_LOGITS}{location} (nan={}, inf={}, vocab={})",
row.iter().filter(|x| x.is_nan()).count(),
row.iter().filter(|x| x.is_infinite()).count(),
row.len()
)
}
pub(super) fn is_non_finite_logits(error: &str) -> bool {
error.starts_with(NON_FINITE_LOGITS)
}
const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024;
const MANIFEST: &[u8] = include_bytes!("../../assets/models/qwen38-flash-next-bare-speed.json");
const INVENTORY: &str =
src/server.rs
+51
View File
@@ -578,6 +578,57 @@ fn empty_arguments() -> String {
"{}".into()
}
/// Feed actual Generator events through the production HTTP response adapters.
#[cfg(test)]
pub(crate) fn issue105_wire(active: crate::runtime::ActiveGeneration, streaming: bool) -> String {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let address = listener.local_addr().unwrap();
let client = thread::spawn(move || {
let mut stream = TcpStream::connect(address).unwrap();
let mut response = String::new();
stream.read_to_string(&mut response).unwrap();
response
});
let (mut stream, _) = listener.accept().unwrap();
let metrics = Arc::new(Metrics::new(std::path::Path::new("/nonexistent-issue105")));
let state = State {
generation: GenerationService::spawn(Arc::clone(&metrics)).unwrap(),
config: Arc::new(RwLock::new(Config::default())),
models_path: PathBuf::new(),
cache_path: PathBuf::new(),
sequence: AtomicU64::new(0),
connections: Arc::new(AtomicUsize::new(0)),
tool_memory: Mutex::new(HashMap::new()),
metrics,
cors: false,
};
let options = ResponseOptions {
protocol: Protocol::Chat,
reasoning_summary: false,
model_id: "qwen3.8-flash-next".into(),
stream: streaming,
include_usage: true,
has_tools: false,
cors: false,
};
let result = if streaming {
response::stream_response(&mut stream, &state, options, active, "issue105")
} else {
response::final_response(&mut stream, &state, options, active, "issue105")
};
if let Err((status, message)) = result {
send_json_with_cors(
&mut stream,
status,
&json!({"error":{"message":message,"type":"server_error"}}),
false,
)
.unwrap();
}
drop(stream);
client.join().unwrap()
}
#[cfg(test)]
mod tests {
use super::*;