Commit the final Qwen MTP primary before saving state

This commit is contained in:
Georg Bauer committed 2026-10-04 13:42:17 +02:00
1 parent 5ecf8ab4f3
commit a8cad88100
17 files changed
+806 -127

No files matched your search

+2
View File
@@ -33,6 +33,8 @@ mod issue103_analysis;
mod issue104;
#[cfg(test)]
mod issue105;
#[cfg(test)]
mod issue106;
mod linear;
mod load;
mod mlp;
+8
View File
@@ -230,6 +230,8 @@ pub(super) struct Execution {
pub(super) fault: Option<super::issue105::Fault>,
#[cfg(test)]
pub(super) captured_prefixes: usize,
#[cfg(test)]
pub(super) script: Option<super::issue106::Script>,
_context: Context,
}
@@ -343,6 +345,8 @@ impl Execution {
fault: None,
#[cfg(test)]
captured_prefixes: 0,
#[cfg(test)]
script: None,
_context: context,
})
}
@@ -373,6 +377,8 @@ impl Execution {
})?;
#[cfg(test)]
let output = self.inject_output(super::issue105::Stage::Target, output)?;
#[cfg(test)]
let output = self.script_output(false, output)?;
Ok(output)
}
@@ -800,6 +806,8 @@ impl Execution {
})?;
#[cfg(test)]
let output = self.inject_output(super::issue105::Stage::Draft, output)?;
#[cfg(test)]
let output = self.script_output(true, output)?;
Ok(output)
}
@@ -316,132 +316,6 @@ fn issue103_tokenizer_receipts() {
assert_eq!(mismatches, 5, "diagnosis changed; inspect exact token IDs");
}
#[test]
#[ignore = "#103 diagnostic; full pinned model, SSD PLE and external memory supervisor"]
fn issue103_mtp_final_receipts() {
use super::decoder::AttentionCache;
use super::execution::Execution;
use super::speculative::Options;
let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
assert!(!super::ngram::resident_policy(), "use SSD PLE");
eprintln!("#103 loading unchanged Bare-Speed model");
let mut execution = Execution::load(&root, true, false, 100).unwrap();
let ids = [9707, 11, 1075]; // fixed valid IDs; no template is involved in this model-level repro
let prepared = execution
.prefill_mtp_history(
&ids,
&super::prefill::Options {
chunk_size: Some(2048),
final_logits_only: true,
external_cache_only: false,
external_emit_logits: true,
return_hidden: true,
},
None,
false,
|_| Ok(()),
|_| Ok(()),
)
.unwrap();
let baseline = execution
.capture_session(
&ids,
&prepared.target.logits,
prepared.target.hidden.as_ref(),
true,
false,
)
.unwrap();
let offset = |e: &Execution| {
e.cache
.iter()
.filter_map(|c| match &c.as_ref()?.attention {
AttentionCache::Qsa(c) => Some(c.kv.offset),
_ => None,
})
.collect::<Vec<_>>()
};
let mut final_token = 0;
for depth in [1, 2, 3] {
for temperature in [0., 0.6] {
for stop_exit in [false, true] {
let (logits, hidden, _) = execution.restore_session(&baseline, &ids).unwrap();
let sampler = SamplerConfig {
temperature,
..Default::default()
};
let result = execution
.decode_mtp(
&logits,
hidden.as_ref().unwrap(),
&sampler,
&sampler,
&Options {
prompt_ids: &ids,
depth,
max_tokens: if stop_exit { 8 } else { 1 },
seed: 0,
stops: if stop_exit {
if temperature == 0. { &[11] } else { &[314] }
} else {
&[]
},
committed_history: true,
family_capture: true,
lazy_history: false,
defer_repair_eval: false,
capture_final_state: true,
},
|_| Ok(()),
|_| Ok(()),
)
.unwrap();
let mut key = ids.to_vec();
key.push(result.tokens[0] as i32);
let snapshot = execution
.capture_session(&key, &result.logits, Some(&result.hidden), true, false)
.unwrap();
let (.., reused) = execution.restore_session(&snapshot, &key).unwrap();
let offsets = offset(&execution);
println!(
"{}",
serde_json::json!({"case":"mtp_final","exit":if stop_exit {"first_primary_stop"} else {"max_tokens"},"depth":depth,"temperature":temperature,"prompt_ids":ids,"tokens":result.tokens,"key_length":key.len(),"restored_cached_tokens":reused,"trunk_offsets":offsets,"mtp_offset":execution.mtp_offset().unwrap(),"pending":result.pending,"final_state_error":result.final_state_error})
);
assert_eq!(offsets, vec![3; 12]);
assert_eq!(reused, 4);
assert!(result.pending.is_none() && result.final_state_error.is_none());
final_token = result.tokens[0] as i32;
}
}
}
execution.restore_session(&baseline, &ids).unwrap();
let bad = execution
.forward(&[1075], false, false)
.unwrap()
.logits
.unwrap();
let bad = execution.values_f32(&bad).unwrap();
execution.restore_session(&baseline, &ids).unwrap();
let good = execution
.forward(&[final_token, 1075], false, false)
.unwrap()
.logits
.unwrap();
let good = execution
.values_f32(&super::ar::last_row(&good, execution.stream).unwrap())
.unwrap();
let different = bad
.iter()
.zip(&good)
.filter(|(a, b)| a.to_bits() != b.to_bits())
.count();
println!(
"{}",
serde_json::json!({"case":"mtp_next_turn","missing_token":final_token,"suffix_ids":[1075],"different_logits":different,"bad_sha256":super::super::qwen_mtplx_tests::hex_digest(&bad.iter().flat_map(|v|v.to_le_bytes()).collect::<Vec<_>>()),"good_sha256":super::super::qwen_mtplx_tests::hex_digest(&good.iter().flat_map(|v|v.to_le_bytes()).collect::<Vec<_>>())})
);
assert!(different > 0);
}
#[test]
#[ignore = "#103 diagnostic; small active attention module on Metal; mixed position boundary"]
fn issue103_vision_position_receipts() {
+399
View File
@@ -0,0 +1,399 @@
//! Final-primary regression on real caches, plus the pinned oracle's exit matrix.
use super::array::{Array, Dtype};
use super::decoder::AttentionCache;
use super::execution::Execution;
use super::model::Output;
use super::request::LiveSession;
use super::sampling::SamplerConfig;
use super::session_cache::{SessionSnapshot, TurnDecoded, TurnProgress};
use super::session_checkpoint::Identity;
use super::speculative::Options;
use super::{ar, ops};
use serde_json::{Value, json};
use std::path::PathBuf;
pub(super) struct Script {
stop_index: i32,
omniscient: bool,
fail_offset: Option<i32>,
}
fn offsets(execution: &Execution) -> Vec<i32> {
execution
.cache
.iter()
.filter_map(|c| match &c.as_ref()?.attention {
AttentionCache::Qsa(c) => Some(c.kv.offset),
_ => None,
})
.collect()
}
impl Execution {
pub(super) fn script_output(&self, draft: bool, output: Output) -> Result<Output, String> {
let Some(script) = &self.script else {
return Ok(output);
};
let end = if draft {
self.mtp_offset()?
} else {
offsets(self)[0]
};
if !draft && script.fail_offset == Some(end) {
return Err("#106 final forward failure".into());
}
let logits = output.logits.unwrap();
ops::evaluate(
&self.streams,
std::slice::from_ref(&logits),
self.stream,
false,
)?;
let rows = logits.layout().dim(1)?;
let values = (0..rows)
.flat_map(|j| {
let index = end - rows + j + if draft { 2 } else { 1 };
let token = if (!draft || script.omniscient) && index == script.stop_index {
3
} else {
1
};
(0..4).flat_map(move |i| (if i == token { 1f32 } else { 0f32 }).to_le_bytes())
})
.collect::<Vec<_>>();
Ok(Output {
logits: Some(Array::new(
&[1, rows, 4],
Dtype::F32,
super::scalar_buffer(&values)?,
)?),
..output
})
}
}
fn fingerprint(execution: &Execution, saved: &SessionSnapshot) -> Value {
let mut blobs = std::collections::BTreeMap::new();
let payload = saved
.encode_payload(
256,
&execution.streams,
execution.stream,
|name, bytes| {
blobs.insert(name, super::super::qwen_mtplx_tests::hex_digest(&bytes));
Ok(())
},
|| false,
)
.unwrap();
json!({"ids":saved.token_ids,"lazy_kv":saved.lazy_kv,"payload":payload,"blobs":blobs})
}
fn digest(value: &Value) -> String {
super::super::qwen_mtplx_tests::hex_digest(&serde_json::to_vec(value).unwrap())
}
fn prefill(execution: &mut Execution, ids: &[i32]) -> SessionSnapshot {
let prepared = execution
.prefill_mtp_history(
ids,
&super::prefill::Options {
chunk_size: Some(2048),
final_logits_only: true,
external_cache_only: false,
external_emit_logits: true,
return_hidden: true,
},
None,
false,
|_| Ok(()),
|_| Ok(()),
)
.unwrap();
execution
.capture_session(
ids,
&prepared.target.logits,
prepared.target.hidden.as_ref(),
true,
false,
)
.unwrap()
}
fn options(ids: &[i32]) -> Options<'_> {
Options {
prompt_ids: ids,
depth: 1,
max_tokens: 1,
seed: 0,
stops: &[],
committed_history: true,
family_capture: true,
lazy_history: false,
defer_repair_eval: false,
capture_final_state: true,
}
}
fn advance(execution: &mut Execution, saved: &SessionSnapshot, token: i32) -> SessionSnapshot {
let (_, hidden, _) = execution.restore_session(saved, &saved.token_ids).unwrap();
let updated = execution
.append_mtp_history(hidden.as_ref().unwrap(), &[token], false)
.unwrap();
ops::evaluate(&execution.streams, &[updated], execution.stream, false).unwrap();
let output = execution.forward(&[token], false, false).unwrap();
let logits = ar::last_row(&output.logits.unwrap(), execution.stream).unwrap();
let mut key = saved.token_ids.clone();
key.push(token);
execution
.capture_session(&key, &logits, output.hidden.as_ref(), true, false)
.unwrap()
}
#[test]
#[ignore = "installed model and actual pre-fix checkpoint; external memory guard required"]
fn issue106_complete_state_and_checkpoint_recovery() {
let root = PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
let old = PathBuf::from(std::env::var_os("DS4_ISSUE106_BAD_CHECKPOINT").unwrap());
let directory = PathBuf::from(std::env::var_os("DS4_ISSUE106_DIR").unwrap());
std::fs::create_dir_all(&directory).unwrap();
assert!(!super::ngram::resident_policy());
let mut execution = Execution::load(&root, true, false, 100).unwrap();
let metadata = crate::engine::qwen::QwenMetadata::open(&root, true).unwrap();
let identity = Identity {
model: metadata.checkpoint_identity(),
context: 4096,
vocabulary: metadata.tokenizer.vocab_size() as u32,
mtp: true,
};
let mut live = LiveSession::default();
assert_eq!(&std::fs::read(&old).unwrap()[..8], b"QWMTKV01");
let error = live
.load_checkpoint(
&old,
identity,
&execution.streams,
execution.stream,
|_| {},
|| false,
)
.unwrap_err();
assert!(error.contains("format differs"));
assert!(live.current.is_none());
let ids = [9707, 11, 1075];
let baseline = prefill(&mut execution, &ids);
for depth in [1, 2, 3] {
for temperature in [0., 0.6] {
for stop in [false, true] {
let (logits, hidden, reused) = execution.restore_session(&baseline, &ids).unwrap();
assert_eq!(reused, ids.len());
let sampler = SamplerConfig {
temperature,
..Default::default()
};
let mut settings = options(&ids);
settings.depth = depth;
settings.max_tokens = if stop { 8 } else { 1 };
settings.stops = if !stop {
&[]
} else if temperature == 0. {
&[11]
} else {
&[314]
};
let result = execution
.decode_mtp(
&logits,
hidden.as_ref().unwrap(),
&sampler,
&sampler,
&settings,
|_| Ok(()),
|_| Ok(()),
)
.unwrap();
assert_eq!(result.tokens, [if temperature == 0. { 11 } else { 314 }]);
assert!(result.pending.is_none());
assert!(result.final_state_error.is_none());
assert_eq!(offsets(&execution), vec![4; 12]);
assert_eq!(execution.mtp_offset().unwrap(), 3);
let mut key = ids.to_vec();
key.push(result.tokens[0] as i32);
let actual = execution
.capture_session(&key, &result.logits, Some(&result.hidden), true, false)
.unwrap();
let actual_state = fingerprint(&execution, &actual);
let expected = advance(&mut execution, &baseline, key[3]);
assert_eq!(
actual_state,
fingerprint(&execution, &expected),
"all trunk/history/logit/hidden tensors must match teacher forcing"
);
let path = directory.join(format!("{depth}-{temperature}-{stop}.bin"));
live.current = Some(actual);
live.save_checkpoint(
&path,
identity,
[106; 32],
&execution.streams,
execution.stream,
|_| {},
|| false,
)
.unwrap();
live = LiveSession::default();
assert!(
live.load_checkpoint(
&path,
identity,
&execution.streams,
execution.stream,
|_| {},
|| false
)
.unwrap()
);
let reopened = live.current.take().unwrap();
assert_eq!(actual_state, fingerprint(&execution, &reopened));
let (.., reused) = execution.restore_session(&reopened, &key).unwrap();
assert_eq!(reused, 4);
let warm_next = advance(&mut execution, &reopened, 1075);
let expected_next = advance(&mut execution, &expected, 1075);
let next = fingerprint(&execution, &warm_next);
assert_eq!(next, fingerprint(&execution, &expected_next));
println!(
"{}",
json!({"case":"complete_state","depth":depth,"temperature":temperature,"exit":if stop {"stop"} else {"max_tokens"},"tokens":result.tokens,"key_length":4,"trunk_offsets":vec![4;12],"mtp_offset":3,"reopened_cached_tokens":reused,"state_sha256":digest(&actual_state),"next_sha256":digest(&next),"old_checkpoint_rejected":true})
);
}
}
}
// A cancelled finalization and a failed final forward cannot expose safe arrays.
let sampler = SamplerConfig::default();
for failure in ["cancel", "forward", "capture_disabled"] {
let (logits, hidden, _) = execution.restore_session(&baseline, &ids).unwrap();
let mut settings = options(&ids);
settings.capture_final_state = failure != "capture_disabled";
if failure == "forward" {
execution.script = Some(Script {
stop_index: 10000,
omniscient: false,
fail_offset: Some(4),
});
}
let result = execution.decode_mtp(
&logits,
hidden.as_ref().unwrap(),
&sampler,
&sampler,
&settings,
|_| Ok(()),
|progress| {
if failure == "cancel" && matches!(progress, TurnProgress::Finalize { .. }) {
Err("cancelled".into())
} else {
Ok(())
}
},
);
execution.script = None;
if failure == "cancel" {
assert!(result.is_err());
} else {
let result = result.unwrap();
assert!(result.pending.is_some());
assert_eq!(result.final_state_error.is_some(), failure == "forward");
assert!(TurnDecoded::Mtp(result).final_arrays(true).is_none());
}
println!("{}", json!({"case":failure,"safe_final_state":false}));
}
execution.streams.clear_streams().unwrap();
}
#[test]
#[ignore = "installed model, scripted logits on real caches; external memory guard required"]
fn issue106_upstream_exit_matrix() {
assert_eq!(std::env::var("MTPLX_CONTEXT_COPY").unwrap(), "0");
let root = PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
let reference =
std::fs::read_to_string(std::env::var_os("DS4_ISSUE106_ORACLE").unwrap()).unwrap();
let reference = reference
.lines()
.map(|l| serde_json::from_str::<Value>(l).unwrap())
.collect::<Vec<_>>();
let mut execution = Execution::load(&root, true, false, 100).unwrap();
let ids = [0; 5];
let baseline = prefill(&mut execution, &ids);
for expected in reference.iter().filter(|r| r["case"] == "mtp_exit") {
let (_, hidden, _) = execution.restore_session(&baseline, &ids).unwrap();
let sampler = SamplerConfig {
temperature: expected["temperature"].as_f64().unwrap(),
top_k: if expected["temperature"] == 0.7 { 1 } else { 0 },
..Default::default()
};
execution.script = Some(Script {
stop_index: expected["stop_index"].as_i64().unwrap() as i32,
omniscient: expected["omniscient"].as_bool().unwrap(),
fail_offset: None,
});
let input = Array::new(
&[1, 1, 4],
Dtype::F32,
super::scalar_buffer(
&[0f32, 1., 0., 0.]
.into_iter()
.flat_map(f32::to_le_bytes)
.collect::<Vec<_>>(),
)
.unwrap(),
)
.unwrap();
let mut settings = options(&ids);
settings.depth = expected["depth"].as_u64().unwrap() as usize;
settings.max_tokens = expected["max_tokens"].as_u64().unwrap() as usize;
settings.stops = &[3];
let result = execution
.decode_mtp(
&input,
hidden.as_ref().unwrap(),
&sampler,
&sampler,
&settings,
|_| Ok(()),
|_| Ok(()),
)
.unwrap();
assert_eq!(
json!(result.tokens),
expected["tokens"],
"{}",
expected["label"]
);
match expected["origin"].as_str() {
Some("bonus") => assert!(result.cycles.iter().any(|c| c.bonus == Some(3))),
Some("accepted_draft") => assert!(
result
.cycles
.iter()
.any(|c| { c.drafts[..c.accepted].contains(&3) })
),
Some("residual_correction") => {
assert!(result.cycles.iter().any(|c| c.correction == Some(3)))
}
Some("primary") => assert_eq!(result.cycles.last().unwrap().primary, 3),
None => assert_eq!(result.tokens.len(), settings.max_tokens),
other => panic!("uncovered oracle exit: {other:?}"),
}
let length = ids.len() + result.tokens.len();
assert_eq!(offsets(&execution), vec![length as i32; 12]);
assert_eq!(execution.mtp_offset().unwrap(), length as i32 - 1);
assert!(result.pending.is_none() && result.final_state_error.is_none());
println!(
"{}",
json!({"case":"mtp_exit","label":expected["label"],"tokens":result.tokens,"key_length":length,"trunk_offsets":offsets(&execution),"mtp_offset":execution.mtp_offset().unwrap(),"oracle_matches":true,"cycles":result.cycles})
);
execution.script = None;
}
execution.streams.clear_streams().unwrap();
}
@@ -0,0 +1,101 @@
//! Existing Generator rejection/rebuild path for a real pre-fix checkpoint.
use super::*;
use crate::engine::{Generator, metal::Executor};
use crate::settings::{
GenerationPreferences, RuntimePreferences, SpeculativePreferences, effective_settings,
};
use std::sync::Arc;
#[test]
#[ignore = "installed model and actual pre-fix checkpoint; external memory guard required"]
fn issue106_product_discards_old_checkpoint_and_recovers_cold() {
let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
let old = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE106_BAD_CHECKPOINT").unwrap());
let directory =
std::path::PathBuf::from(std::env::var_os("DS4_ISSUE106_DIR").unwrap()).join("product");
std::fs::create_dir_all(&directory).unwrap();
let effective = effective_settings(
crate::model::ModelChoice::Qwen38FlashNext,
&GenerationPreferences {
context_tokens: 4096,
max_generated_tokens: 1,
temperature: Some(0.6),
..Default::default()
},
&RuntimePreferences {
speculative: SpeculativePreferences {
glm_mtp: true,
..Default::default()
},
..Default::default()
},
root.parent().unwrap(),
)
.unwrap();
let mut settings = effective.turn;
settings.seed = Some(0);
settings.kv_cache.min_tokens = 0;
let mut generator = Generator::open(
&effective.engine,
Arc::new(Metrics::new(&directory.join("metrics"))),
)
.unwrap();
let path = directory.join("old.bin");
std::fs::copy(&old, &path).unwrap();
let selection = generator.select_checkpoint(&path, |_| true).unwrap();
assert!(!selection.found && selection.incompatible);
assert!(!path.exists());
let Executor::QwenMtplx(owner) = &mut generator.executor else {
unreachable!()
};
assert!(owner.live.current.is_none() && owner.live.prefix.is_none());
let messages = [ChatTurn {
user: true,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Say hello.".into(),
}];
let mut states = Vec::new();
for target in [&path, &directory.join("fresh.bin")] {
let out = generator
.generate(
target,
None,
&messages,
&settings,
&AtomicBool::new(false),
|_, _| {},
|_, _, _| {},
|_| {},
)
.unwrap();
assert_eq!(out.cached_tokens, 0);
assert_eq!(&std::fs::read(target).unwrap()[..8], b"QWMTKV02");
let Executor::QwenMtplx(owner) = &mut generator.executor else {
unreachable!()
};
let saved = owner.live.current.as_ref().unwrap();
let mut hashes = Vec::new();
let payload = saved
.encode_payload(
256,
&owner.execution.streams,
owner.execution.stream,
|_, bytes| {
hashes.push(crate::engine::metal::qwen_mtplx_tests::hex_digest(&bytes));
Ok(())
},
|| false,
)
.unwrap();
states.push(serde_json::json!({"tokens":saved.token_ids,"payload":payload,"hashes":hashes,"text":out.message.content,"reasoning":out.message.reasoning}));
}
assert_eq!(states[0], states[1]);
println!(
"{}",
serde_json::json!({"case":"product_checkpoint_recovery","old_removed":true,"cold_tokens":0,"new_magic":"QWMTKV02","fresh_state_equal":true})
);
}
+3
View File
@@ -1065,3 +1065,6 @@ fn qwen_vision_chat_checkpoint_preserves_image_identity() {
#[cfg(test)]
#[path = "issue105_product.rs"]
mod issue105_product;
#[cfg(test)]
#[path = "issue106_product.rs"]
mod issue106_product;
@@ -10,7 +10,9 @@ use std::fs::{self, File};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;
const MAGIC: &[u8; 8] = b"QWMTKV01";
// V1 may bank an emitted MTP primary without its trunk/history state (#106).
// Reject it before decoding tensors; the owner rebuilds the selected cache.
const MAGIC: &[u8; 8] = b"QWMTKV02";
const HEADER_BYTES: u64 = 16;
const MAX_MANIFEST_BYTES: u64 = 64 * 1024 * 1024;
const IO_CHUNK: usize = 8 * 1024 * 1024;
@@ -734,6 +734,9 @@ impl Execution {
..Default::default()
};
if result.tokens.len() >= options.max_tokens || stopped(cycle.primary) {
// MTPLX keeps an emitted round-head primary pending until the
// final commit advances both trunk and committed MTP history.
result.pending = Some(cycle.primary);
finish_cycle(&mut result, cycle)?;
break;
}