Files
DS4Server/src/engine/metal/qwen_mtplx/session_checkpoint.rs
T

628 lines
23 KiB
Rust

//! A single atomic checkpoint for the existing DS4 KV index, not a second bank.
//! Tensor payloads keep the reference codec's dtypes, shapes and raw bytes.
use super::request::LiveSession;
use super::session_codec::decode_payload;
use super::stream::{Stream, Streams};
use crate::engine::metal::checkpoint::{read_u64, write_u64};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs::{self, File};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;
// 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;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub(super) struct Identity {
pub(super) model: [u8; 32],
pub(super) context: u32,
pub(super) vocabulary: u32,
pub(super) mtp: bool,
#[serde(default)]
pub(super) head_chunk: bool,
}
#[derive(Serialize, Deserialize)]
struct Manifest {
identity: Identity,
tag: [u8; 32],
#[serde(default)]
verify_gate: bool,
#[serde(default)]
sorted_gather_guard: bool,
#[serde(default)]
images: Vec<crate::engine::qwen_vision::Identity>,
tokens: Vec<i32>,
lazy_kv: bool,
payload: serde_json::Value,
// Offsets are absolute; blob names are codec-local, never filesystem paths.
blobs: BTreeMap<String, (u64, u64)>,
}
fn check_abort(abort: &impl Fn() -> bool) -> Result<(), String> {
if abort() {
Err("Qwen checkpoint interrupted".into())
} else {
Ok(())
}
}
fn validate_tokens(tokens: &[i32], identity: Identity) -> Result<(), String> {
if identity.context == 0
|| identity.vocabulary == 0
|| tokens.is_empty()
|| tokens.len() > identity.context as usize
|| tokens
.iter()
.any(|&token| token < 0 || token as u32 >= identity.vocabulary)
{
return Err("Qwen checkpoint token prefix is invalid".into());
}
Ok(())
}
impl LiveSession {
#[allow(clippy::too_many_arguments)]
pub(super) fn save_checkpoint(
&mut self,
path: &Path,
identity: Identity,
tag: [u8; 32],
streams: &Streams,
stream: Stream,
mut progress: impl FnMut(u64),
abort: impl Fn() -> bool,
) -> Result<(), String> {
check_abort(&abort)?;
let saved = self
.current
.as_ref()
.ok_or("Qwen has no committed checkpoint")?;
validate_tokens(&saved.token_ids, identity)?;
if !crate::engine::qwen_vision::compatible(&self.current_images) {
return Err("Qwen vision cache policy differs; rebuild the cache".into());
}
if saved.mtp.is_some() != identity.mtp {
return Err("Qwen checkpoint MTP configuration differs".into());
}
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|e| e.to_string())?;
}
// Same serialized-writer, staging-file and sync/rename contract as DS4.
let temporary = path.with_extension("tmp");
let result = (|| -> Result<(), String> {
let mut file = File::create(&temporary).map_err(|e| e.to_string())?;
file.write_all(MAGIC).map_err(|e| e.to_string())?;
write_u64(&mut file, 0)?;
progress(HEADER_BYTES);
let mut blobs = BTreeMap::new();
let payload = saved.encode_payload(
256,
streams,
stream,
|name, bytes| {
let offset = file.stream_position().map_err(|e| e.to_string())?;
for chunk in bytes.chunks(IO_CHUNK) {
check_abort(&abort)?;
file.write_all(chunk).map_err(|e| e.to_string())?;
progress(chunk.len() as u64);
}
if blobs.insert(name, (offset, bytes.len() as u64)).is_some() {
return Err("duplicate Qwen checkpoint blob".into());
}
Ok(())
},
&abort,
)?;
check_abort(&abort)?;
let offset = file.stream_position().map_err(|e| e.to_string())?;
let manifest = serde_json::to_vec(&Manifest {
verify_gate: true,
sorted_gather_guard: true,
identity,
tag,
images: self.current_images.clone(),
tokens: saved.token_ids.clone(),
lazy_kv: saved.lazy_kv,
payload,
blobs,
})
.map_err(|e| e.to_string())?;
if manifest.len() as u64 > MAX_MANIFEST_BYTES {
return Err("Qwen checkpoint metadata exceeds limit".into());
}
file.write_all(&manifest).map_err(|e| e.to_string())?;
progress(manifest.len() as u64);
file.seek(SeekFrom::Start(8)).map_err(|e| e.to_string())?;
write_u64(&mut file, offset)?;
file.sync_all().map_err(|e| e.to_string())?;
check_abort(&abort)?;
fs::rename(&temporary, path).map_err(|e| e.to_string())
})();
if result.is_err() {
let _ = fs::remove_file(&temporary);
} else {
self.tag = tag;
}
result
}
#[allow(clippy::too_many_arguments)]
pub(super) fn load_checkpoint(
&mut self,
path: &Path,
identity: Identity,
streams: &Streams,
stream: Stream,
mut progress: impl FnMut(u64),
abort: impl Fn() -> bool,
) -> Result<bool, String> {
check_abort(&abort)?;
let mut file = match File::open(path) {
Ok(file) => file,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(false),
Err(e) => return Err(e.to_string()),
};
let mut magic = [0; 8];
file.read_exact(&mut magic).map_err(|e| e.to_string())?;
if &magic != MAGIC {
return Err("Qwen checkpoint format differs; rebuild the cache".into());
}
let offset = read_u64(&mut file)?;
let length = file.metadata().map_err(|e| e.to_string())?.len();
if offset < HEADER_BYTES || offset >= length || length - offset > MAX_MANIFEST_BYTES {
return Err("Qwen checkpoint metadata range is invalid".into());
}
file.seek(SeekFrom::Start(offset))
.map_err(|e| e.to_string())?;
let mut raw = vec![0; (length - offset) as usize];
file.read_exact(&mut raw).map_err(|e| e.to_string())?;
let mut manifest: Manifest = serde_json::from_slice(&raw).map_err(|e| e.to_string())?;
progress(HEADER_BYTES + raw.len() as u64);
if (Identity {
head_chunk: identity.head_chunk,
..manifest.identity
}) != identity
{
return Err(
"Qwen checkpoint does not match the current model/context/MTP policy".into(),
);
}
validate_tokens(&manifest.tokens, identity)?;
if !crate::engine::qwen_vision::compatible(&manifest.images) {
return Err("Qwen vision cache policy differs; rebuild the cache".into());
}
if identity.mtp && !manifest.verify_gate {
return Err(
"Qwen checkpoint uses obsolete MTP attention-gate math; rebuild the cache".into(),
);
}
// The route also changes short dense prefills, not just MTP verify.
// A <=3-token AR prefix cannot have executed a >=3-row body, whereas
// MTP can retain a short frontier from a wider verification window.
if manifest.identity.head_chunk != identity.head_chunk
&& (identity.mtp || manifest.tokens.len() >= 4)
{
return Err("Qwen checkpoint SDPA head-chunk policy differs; rebuild the cache".into());
}
// Old files do not record prefill chunk sizes or producer hardware.
// At ten routes/token this length can contain an affected chunk;
// shorter sessions cannot. Reject before allocating/restoring state.
if !manifest.sorted_gather_guard && manifest.tokens.len() >= 3277 {
return Err(
"Qwen checkpoint predates the sorted-gather row guard; rebuild the cache".into(),
);
}
if !manifest.payload["mtp_history_snapshot"].is_null() != identity.mtp {
return Err("Qwen checkpoint MTP payload differs".into());
}
// Validate all extents before allocating any GPU state. Require exact
// coverage, so overlapping blobs and unindexed/trailing bytes fail closed.
let mut ranges = manifest.blobs.values().copied().collect::<Vec<_>>();
ranges.sort_unstable();
let mut end = HEADER_BYTES;
for (start, size) in ranges {
if start != end {
return Err("Qwen checkpoint blob range is invalid".into());
}
end = start
.checked_add(size)
.filter(|&n| n <= offset)
.ok_or("Qwen checkpoint blob exceeds payload")?;
}
if end != offset {
return Err("Qwen checkpoint payload coverage differs".into());
}
let saved = decode_payload(
&manifest.payload,
manifest.tokens,
manifest.lazy_kv,
true,
streams,
stream,
|name| {
check_abort(&abort)?;
let (start, size) = manifest
.blobs
.remove(name)
.ok_or("missing or repeated Qwen checkpoint blob")?;
file.seek(SeekFrom::Start(start))
.map_err(|e| e.to_string())?;
let mut bytes = Vec::new();
let size = usize::try_from(size).map_err(|_| "Qwen checkpoint blob too large")?;
bytes.try_reserve_exact(size).map_err(|e| e.to_string())?;
bytes.resize(size, 0);
for chunk in bytes.chunks_mut(IO_CHUNK) {
check_abort(&abort)?;
file.read_exact(chunk).map_err(|e| e.to_string())?;
progress(chunk.len() as u64);
}
Ok(bytes)
},
)?;
if !manifest.blobs.is_empty() {
return Err("unreferenced Qwen checkpoint blobs".into());
}
check_abort(&abort)?;
// Publish only after complete validation and materialization. Failure
// leaves the live conversation and its prefix untouched.
self.current = Some(saved);
self.prefix = None;
self.tag = manifest.tag;
self.current_images = manifest.images;
self.prefix_images.clear();
Ok(true)
}
}
#[test]
#[ignore = "model-free Metal checkpoint; external supervisor and DS4_QWEN38_CHECKPOINT_DIR required"]
fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() {
use super::super::gpu::Context;
use super::session_cache::{Boundary, SessionSnapshot, State};
use super::stream::Device;
use super::{
allocator::Allocator,
array::{Array, Dtype},
configure_sources, scalar_buffer,
};
use std::cell::Cell;
configure_sources().unwrap();
let _context = Context::open_qwen(0).unwrap();
let streams = Streams::new(Some(Allocator::new().unwrap()));
let stream = streams.default_stream(Device::Gpu).unwrap();
let leaf = Array::new(
&[1, 1, 4],
Dtype::BF16,
scalar_buffer(&[0, 0, 128, 63, 0, 64, 64, 64]).unwrap(),
)
.unwrap();
let state = Some(State::Recurrent(vec![Some(leaf.clone()), None]));
let mut live = LiveSession {
current_images: Vec::new(),
prefix_images: Vec::new(),
current: Some(SessionSnapshot {
token_ids: vec![1, 2, 3, 4],
logits: leaf.clone(),
hidden: Some(leaf.clone()),
trunk: vec![state.clone()],
mtp: Some(vec![None]),
lazy_kv: true,
boundaries: vec![Boundary {
tokens: 2,
state: vec![state],
hidden: Some(leaf),
}],
}),
..Default::default()
};
let identity = Identity {
model: [7; 32],
context: 64,
vocabulary: 256,
mtp: true,
head_chunk: super::policy::sdpa_head_chunk_enabled(|key| std::env::var(key).ok()),
};
let directory =
std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_CHECKPOINT_DIR").unwrap());
fs::create_dir_all(&directory).unwrap();
let path = directory.join("roundtrip.bin");
let fingerprint = |live: &LiveSession| {
let mut blobs = BTreeMap::new();
let payload = live
.current
.as_ref()
.unwrap()
.encode_payload(
256,
&streams,
stream,
|name, bytes| {
blobs.insert(name, bytes);
Ok(())
},
|| false,
)
.unwrap();
(payload, blobs)
};
let original = fingerprint(&live);
let mut written = 0;
live.save_checkpoint(
&path,
identity,
[9; 32],
&streams,
stream,
|n| written += n,
|| false,
)
.unwrap();
let valid = fs::read(&path).unwrap();
assert_eq!(written, valid.len() as u64);
assert_eq!(live.tag, [9; 32]);
live = LiveSession::default();
let mut read = 0;
assert!(
live.load_checkpoint(&path, identity, &streams, stream, |n| read += n, || false)
.unwrap()
);
assert_eq!(read, written);
assert_eq!(fingerprint(&live), original);
assert_eq!(live.tag, [9; 32]);
assert_eq!(live.current.as_ref().unwrap().token_ids, [1, 2, 3, 4]);
assert!(live.current.as_ref().unwrap().lazy_kv);
assert!(
!live
.load_checkpoint(
&directory.join("missing.bin"),
identity,
&streams,
stream,
|_| {},
|| false
)
.unwrap()
);
// Cancellation after IO has started must not replace a valid file/tag.
let interrupted = Cell::new(false);
assert!(
live.save_checkpoint(
&path,
identity,
[8; 32],
&streams,
stream,
|_| interrupted.set(true),
|| interrupted.get()
)
.is_err()
);
assert_eq!(fs::read(&path).unwrap(), valid);
assert!(!path.with_extension("tmp").exists());
interrupted.set(false);
assert!(
live.load_checkpoint(
&path,
identity,
&streams,
stream,
|_| interrupted.set(true),
|| interrupted.get()
)
.is_err()
);
for changed in [
Identity {
context: 65,
..identity
},
Identity {
mtp: false,
..identity
},
Identity {
model: [8; 32],
..identity
},
] {
assert!(
live.load_checkpoint(&path, changed, &streams, stream, |_| {}, || false)
.is_err()
);
}
let offset = u64::from_le_bytes(valid[8..16].try_into().unwrap()) as usize;
let broken = directory.join("broken.bin");
let mutations: [fn(&mut Manifest); 5] = [
|m| m.tokens[0] = -1,
|m| m.tokens.resize(65, 1),
|m| m.blobs.values_mut().next().unwrap().1 = u64::MAX,
|m| m.blobs.values_mut().next().unwrap().0 = 0,
|m| m.payload["logits"]["nbytes"] = serde_json::json!(99),
];
for mutate in mutations {
let mut manifest: Manifest = serde_json::from_slice(&valid[offset..]).unwrap();
mutate(&mut manifest);
let mut bytes = valid[..offset].to_vec();
bytes.extend(serde_json::to_vec(&manifest).unwrap());
fs::write(&broken, bytes).unwrap();
assert!(
live.load_checkpoint(&broken, identity, &streams, stream, |_| {}, || false)
.is_err()
);
}
for bytes in [&b"DS4QWN01"[..], &valid[..offset], &valid[..7]] {
fs::write(&broken, bytes).unwrap();
assert!(
live.load_checkpoint(&broken, identity, &streams, stream, |_| {}, || false)
.is_err()
);
}
assert_eq!(live.tag, [9; 32]);
assert_eq!(fingerprint(&live), original);
// Files predating the verify-phase law lack this field. Only MTP state
// changed numerically; old AR files must remain exactly restorable.
let without_gate_policy = |bytes: &[u8]| {
let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
let mut manifest: serde_json::Value = serde_json::from_slice(&bytes[offset..]).unwrap();
manifest.as_object_mut().unwrap().remove("verify_gate");
let mut result = bytes[..offset].to_vec();
result.extend(serde_json::to_vec(&manifest).unwrap());
result
};
fs::write(&broken, without_gate_policy(&valid)).unwrap();
assert!(
live.load_checkpoint(&broken, identity, &streams, stream, |_| {}, || false)
.unwrap_err()
.contains("obsolete MTP attention-gate")
);
assert_eq!(fingerprint(&live), original);
live.current.as_mut().unwrap().mtp = None;
let ar_identity = Identity {
mtp: false,
..identity
};
let ar_original = fingerprint(&live);
live.save_checkpoint(
&path,
ar_identity,
[10; 32],
&streams,
stream,
|_| {},
|| false,
)
.unwrap();
fs::write(&broken, without_gate_policy(&fs::read(&path).unwrap())).unwrap();
assert!(
live.load_checkpoint(&broken, ar_identity, &streams, stream, |_| {}, || false)
.unwrap()
);
assert_eq!(fingerprint(&live), ar_original);
let bytes = fs::read(&path).unwrap();
let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
let mut manifest: serde_json::Value = serde_json::from_slice(&bytes[offset..]).unwrap();
manifest
.as_object_mut()
.unwrap()
.remove("sorted_gather_guard");
let write_manifest = |manifest: &serde_json::Value| {
let mut old = bytes[..offset].to_vec();
old.extend(serde_json::to_vec(manifest).unwrap());
fs::write(&broken, old).unwrap();
};
write_manifest(&manifest);
assert!(
live.load_checkpoint(&broken, ar_identity, &streams, stream, |_| {}, || false)
.unwrap()
);
assert_eq!(fingerprint(&live), ar_original);
manifest["tokens"] = serde_json::json!(vec![1; 3277]);
manifest["identity"]["context"] = 4096.into();
write_manifest(&manifest);
let large_identity = Identity {
context: 4096,
..ar_identity
};
assert!(
live.load_checkpoint(&broken, large_identity, &streams, stream, |_| {}, || false)
.unwrap_err()
.contains("sorted-gather row guard")
);
assert_eq!(fingerprint(&live), ar_original);
// Policy changes in either direction invalidate affected checkpoints;
// missing metadata is the actual pre-port (unchunked) policy.
for mtp in [false, true] {
for count in [1, 3, 4] {
for produced in [false, true] {
let saved = live.current.as_mut().unwrap();
saved.token_ids = vec![1; count];
saved.mtp = mtp.then(|| vec![None]);
saved.boundaries.clear();
let producer = Identity {
mtp,
head_chunk: produced,
..identity
};
live.save_checkpoint(
&path,
producer,
[113; 32],
&streams,
stream,
|_| {},
|| false,
)
.unwrap();
let original = fingerprint(&live);
let bytes = fs::read(&path).unwrap();
let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
for legacy in [false, true] {
if legacy && produced {
continue;
}
let mut manifest: serde_json::Value =
serde_json::from_slice(&bytes[offset..]).unwrap();
if legacy {
manifest["identity"]
.as_object_mut()
.unwrap()
.remove("head_chunk");
}
let raw = serde_json::to_vec(&manifest).unwrap();
let mut file = bytes[..offset].to_vec();
file.extend(&raw);
fs::write(&broken, file).unwrap();
for consumed in [false, true] {
let consumer = Identity {
head_chunk: consumed,
..producer
};
let mut read = 0;
let result = live.load_checkpoint(
&broken,
consumer,
&streams,
stream,
|n| read += n,
|| false,
);
if produced != consumed && (mtp || count >= 4) {
assert!(
result
.unwrap_err()
.contains("SDPA head-chunk policy differs")
);
assert_eq!(
read,
HEADER_BYTES + raw.len() as u64,
"rejected before tensor reads"
);
} else {
assert!(result.unwrap());
}
assert_eq!(fingerprint(&live), original);
assert_eq!(live.current.as_ref().unwrap().token_ids, vec![1; count]);
}
}
}
}
}
println!(
"#113 SDPA policy: legacy/current, both directions, AR/MTP, 1/3/4-token frontiers; exact compatible restore and incompatible rejection before tensor reads"
);
println!(
"Qwen sorted-gather checkpoint policy: old short state exact, possibly affected long state rejected before restore"
);
println!(
"Qwen checkpoint: exact roundtrip, cancellation/corruption controls, old MTP rejected and old AR restored"
);
}