From a8cad8810060d22ee7ccac82c210ef40d4e08785 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 4 Oct 2026 13:42:17 +0200 Subject: [PATCH] Commit the final Qwen MTP primary before saving state --- docs/issue106-final-primary/README.md | 73 ++++ .../baseline-save.patch | 22 + docs/issue106-final-primary/before.jsonl | 13 + docs/issue106-final-primary/checks.json | 108 +++++ .../native-matrix.jsonl | 8 + .../native-product.jsonl | 1 + .../issue106-final-primary/native-state.jsonl | 15 + docs/issue106-final-primary/oracle.jsonl | 9 + src/engine/metal/qwen_mtplx.rs | 2 + src/engine/metal/qwen_mtplx/execution.rs | 8 + .../metal/qwen_mtplx/issue103_analysis.rs | 126 ------ src/engine/metal/qwen_mtplx/issue106.rs | 399 ++++++++++++++++++ .../metal/qwen_mtplx/issue106_product.rs | 101 +++++ src/engine/metal/qwen_mtplx/owner.rs | 3 + .../metal/qwen_mtplx/session_checkpoint.rs | 4 +- src/engine/metal/qwen_mtplx/speculative.rs | 3 + tools/issue106-reference.py | 38 ++ 17 files changed, 806 insertions(+), 127 deletions(-) create mode 100644 docs/issue106-final-primary/README.md create mode 100644 docs/issue106-final-primary/baseline-save.patch create mode 100644 docs/issue106-final-primary/before.jsonl create mode 100644 docs/issue106-final-primary/checks.json create mode 100644 docs/issue106-final-primary/native-matrix.jsonl create mode 100644 docs/issue106-final-primary/native-product.jsonl create mode 100644 docs/issue106-final-primary/native-state.jsonl create mode 100644 docs/issue106-final-primary/oracle.jsonl create mode 100644 src/engine/metal/qwen_mtplx/issue106.rs create mode 100644 src/engine/metal/qwen_mtplx/issue106_product.rs create mode 100644 tools/issue106-reference.py diff --git a/docs/issue106-final-primary/README.md b/docs/issue106-final-primary/README.md new file mode 100644 index 0000000..4433c58 --- /dev/null +++ b/docs/issue106-final-primary/README.md @@ -0,0 +1,73 @@ +# Issue 106: complete MTP final-primary state + +Qwen's round-head stop/token-limit exit now re-arms its emitted primary for the +existing final commit, exactly as pinned MTPLX 2.12.2 +`9882703f3105363ddc37eca9f97aa09a1d387112`, `generation.py:13001-13019` does. +That commit updates committed MTP history and forwards the primary once. Sampling, +acceptance, RNG, snapshot ownership and other exit branches are unchanged. + +Qwen checkpoint magic is now `QWMTKV02`. Old `QWMTKV01` files are rejected before +tensor allocation; Generator's existing incompatible-cache path removes the file, +resets the owner and rebuilds cold. This deliberately invalidates all old Qwen +checkpoints, including AR files, because V1 does not certify complete MTP state. +It does not alter the DS4/GLM format or their owners. This is a DS4Server persistence +version change around MTPLX's existing inference behavior. + +There is no in-process executable hot-upgrade mechanism: resident sessions belong +to the process-local Generator. Installing a new executable takes effect on +restart, which drops all resident state. Runtime engine changes also drop the old +Generator (`runtime.rs`, `run_command`). New processes cannot import old resident +objects, and old processes' subsequently written V1 files are still rejected. + +## Acceptance review + +| Requirement | Evidence | +| --- | --- | +| Real pre-fix reproduction | Parent `5ecf8ab`, unchanged model, 12 first-primary cases: key 4, all 12 trunk offsets 3, MTP offset 2, pending absent and no error. `before.jsonl`. | +| Actually persisted bad state | A real `QWMTKV01` file was saved from the sampled depth-1 reproduction, with IDs `[9707,11,1075,314]`. File identity and hash in `checks.json`; the 116 MB cache is not checked into Git. | +| Full state, not offsets alone | All 12 cases now have trunk 4 and history 3. Every encoded tensor's raw bytes, dtype, shape and metadata, including GDN/conv, PLE, QSA/KV, MTP history, logits and hidden state, match independent teacher forcing. `native-state.jsonl`. | +| Exactly once and next turn | Teacher forcing restores the prompt, explicitly appends one MTP history row and forwards one primary. All corrected snapshots match it, persist/reopen exactly, restore four tokens, then match the full state after suffix `[1075]`. | +| All upstream exits | Eight cases compare token sequences, cache lengths and stop origin with the pinned upstream scripted-model test: bonus, deferred correction, accepted draft, sampled correction, first/fresh primary, and both max-token depths. Native tests script only logits; real model computation and caches still execute. Context copy is disabled in both oracle and native matrix so the intended draft/verify routes are reached. `oracle.jsonl`, `native-matrix.jsonl`. | +| No false final capture on failure | Cancellation at finalization returns an error; failed final forward leaves pending/error and no safe final arrays; disabling final capture leaves pending. `native-state.jsonl`. | +| Old cache recovery in product | The real old file is rejected and removed by Generator, and recovery reports zero reused tokens. Its entire newly persisted state and response equal an independent cold request. `native-product.jsonl`. | +| Existing persistence contracts | Model-free checkpoint roundtrip, corruption, failed restore, cancellation and atomic-write regression run under the Metal supervisor. | +| Other families | Production diff is confined to Qwen's MTP exit and checkpoint signature. DS4/GLM retain `DS4RKV01`/version 2 and their existing pending/state behavior; shared RNG and inference helpers are unchanged. Ordinary cargo tests retain DS4 token/RNG controls. | + +The former bug-asserting `issue103_mtp_final_receipts` diagnostic is replaced by +these positive-contract regressions. The only new Execution hooks are `cfg(test)` +scripted logits/failure injection, absent from release builds. No Python is called +by the product. These receipts establish correctness, not speed. + +## Reproduction + +Use the unchanged Bare-Speed revision +`74559cdf34fbfc0b593de72d17e93f37fd4f9ea7`, M5 Max 128 GiB, SSD PLE, +MLX 0.32.2 for the isolated oracle. Artifact/environment hashes were refreshed in +`../issue104-ar/environment.json`. Native raw prompt is `[9707,11,1075]`, +chunk 2048, seed 0, depths 1/2/3, temperatures 0/.6. + +To regenerate an actual old file, use an isolated checkout of `5ecf8ab`, apply +`baseline-save.patch`, build its tests and run `issue103_mtp_final_receipts` with +`DS4_ISSUE106_BAD_CHECKPOINT` pointing to the output file. The patch only adds +persistence to the old diagnostic; it does not manufacture a header or cache. + +Build current tests with `cargo test --all-features --bin ds4-server --no-run`. +Set `DS4_QWEN38_ARTIFACTS` to the installed model, `DS4_ISSUE106_BAD_CHECKPOINT` +to that real old file, and `DS4_ISSUE106_DIR` to a temporary output directory. +Run each ignored test separately and serially through +`target/debug/test-supervisor 114688 60 120 --command + --ignored --nocapture --test-threads=1`: + +- `issue106_complete_state_and_checkpoint_recovery` +- `issue106_product_discards_old_checkpoint_and_recovers_cold` +- `issue106_upstream_exit_matrix` (also `MTPLX_CONTEXT_COPY=0`, + `DS4_ISSUE106_ORACLE` pointing to fresh oracle JSONL). + +Generate the oracle with `MTPLX_CONTEXT_COPY=0 +tools/issue106-reference.py `, with pytest and the pinned +reference dependencies installed in that isolated environment. +Run `qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state` separately +with a 4096 MiB supervisor cap and `DS4_QWEN38_CHECKPOINT_DIR` set. + +All four mandatory commit gates and individual supervisor outcomes are recorded +in `checks.json`. diff --git a/docs/issue106-final-primary/baseline-save.patch b/docs/issue106-final-primary/baseline-save.patch new file mode 100644 index 0000000..42b1868 --- /dev/null +++ b/docs/issue106-final-primary/baseline-save.patch @@ -0,0 +1,22 @@ +--- a/src/engine/metal/qwen_mtplx/issue103_analysis.rs ++++ b/src/engine/metal/qwen_mtplx/issue103_analysis.rs +@@ -407,6 +407,19 @@ + "{}", + 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}) + ); ++ if depth == 1 && temperature == 0.6 && !stop_exit { ++ let path = std::env::var_os("DS4_ISSUE106_BAD_CHECKPOINT").unwrap(); ++ let metadata = crate::engine::qwen::QwenMetadata::open(&root, true).unwrap(); ++ let identity = super::session_checkpoint::Identity { ++ model: metadata.checkpoint_identity(), context: 4096, ++ vocabulary: metadata.tokenizer.vocab_size() as u32, mtp: true, ++ }; ++ let mut live = super::request::LiveSession { ++ current: Some(snapshot), ..Default::default() ++ }; ++ live.save_checkpoint(std::path::Path::new(&path), identity, ++ [106; 32], &execution.streams, execution.stream, |_| {}, || false).unwrap(); ++ } + assert_eq!(offsets, vec![3; 12]); + assert_eq!(reused, 4); + assert!(result.pending.is_none() && result.final_state_error.is_none()); diff --git a/docs/issue106-final-primary/before.jsonl b/docs/issue106-final-primary/before.jsonl new file mode 100644 index 0000000..e4a0781 --- /dev/null +++ b/docs/issue106-final-primary/before.jsonl @@ -0,0 +1,13 @@ +{"case":"mtp_final","exit":"max_tokens","depth":1,"temperature":0.0,"prompt_ids":[9707,11,1075],"tokens":[11],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"first_primary_stop","depth":1,"temperature":0.0,"prompt_ids":[9707,11,1075],"tokens":[11],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"max_tokens","depth":1,"temperature":0.6,"prompt_ids":[9707,11,1075],"tokens":[314],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"first_primary_stop","depth":1,"temperature":0.6,"prompt_ids":[9707,11,1075],"tokens":[314],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"max_tokens","depth":2,"temperature":0.0,"prompt_ids":[9707,11,1075],"tokens":[11],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"first_primary_stop","depth":2,"temperature":0.0,"prompt_ids":[9707,11,1075],"tokens":[11],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"max_tokens","depth":2,"temperature":0.6,"prompt_ids":[9707,11,1075],"tokens":[314],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"first_primary_stop","depth":2,"temperature":0.6,"prompt_ids":[9707,11,1075],"tokens":[314],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"max_tokens","depth":3,"temperature":0.0,"prompt_ids":[9707,11,1075],"tokens":[11],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"first_primary_stop","depth":3,"temperature":0.0,"prompt_ids":[9707,11,1075],"tokens":[11],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"max_tokens","depth":3,"temperature":0.6,"prompt_ids":[9707,11,1075],"tokens":[314],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_final","exit":"first_primary_stop","depth":3,"temperature":0.6,"prompt_ids":[9707,11,1075],"tokens":[314],"key_length":4,"restored_cached_tokens":4,"trunk_offsets":[3,3,3,3,3,3,3,3,3,3,3,3],"mtp_offset":2,"pending":null,"final_state_error":null} +{"case":"mtp_next_turn","missing_token":314,"suffix_ids":[1075],"different_logits":247643,"bad_sha256":"9c22c019eeb07fead0760dc06ef796e4fbf3f5b4946b02ad808f430ecc6cee5c","good_sha256":"9a556dd9d3063864b70c3338cb0cdf76cc578a01de1c7771b1da0abf3e776feb"} diff --git a/docs/issue106-final-primary/checks.json b/docs/issue106-final-primary/checks.json new file mode 100644 index 0000000..fa785af --- /dev/null +++ b/docs/issue106-final-primary/checks.json @@ -0,0 +1,108 @@ +{ + "before": { + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 521 filtered out; finished in 8.91s" + ], + "cases": 13, + "elapsed_ms": 11509, + "peak_physical_bytes": 73985465832, + "supervisor_error": null, + "log_sha256": "fd1c975e413b08c58ba486c6372e6d9d7e68bba87b71ca07633bfd5bc3efe070" + }, + "state": { + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 523 filtered out; finished in 101.17s" + ], + "cases": 15, + "elapsed_ms": 102581, + "peak_physical_bytes": 74587317128, + "supervisor_error": null, + "log_sha256": "cfcc068258fab9a2ec08632738ee4dcb148f9ff7b2f0a0b282e1ef8b026b3968" + }, + "matrix": { + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 523 filtered out; finished in 8.91s" + ], + "cases": 8, + "elapsed_ms": 9126, + "peak_physical_bytes": 74049414840, + "supervisor_error": null, + "log_sha256": "2d03f4fcb937abda28cfb00082aa0cefe9234d95645645e31021ca80da36fe36" + }, + "product": { + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 523 filtered out; finished in 12.86s" + ], + "cases": 1, + "elapsed_ms": 12972, + "peak_physical_bytes": 76933755256, + "supervisor_error": null, + "log_sha256": "f320e66e0a3355a588d07c6510b3d3449b4a3b2d7266cc18b5b1c82b0f02dc77" + }, + "codec": { + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 523 filtered out; finished in 0.09s" + ], + "cases": 0, + "elapsed_ms": 105, + "peak_physical_bytes": 213088, + "supervisor_error": null, + "log_sha256": "dbddf0b37a9600bdf3eec6868e9cbf5434296bb9998d68ebe1e51ba414ee94b8" + }, + "old_checkpoint": { + "base": "5ecf8ab4f31aeaef7a06b4f03fbb8d658f28c486", + "size": 116442383, + "sha256": "8cd38359515e9d30b191992cc477393ee59b5cac7b6d610aa5838e589ce02c64", + "magic": "QWMTKV01", + "tokens": [ + 9707, + 11, + 1075, + 314 + ], + "trunk_offsets": [ + 3, + 3, + 3, + 3, + 3, + 3, + 3, + 3, + 3, + 3, + 3, + 3 + ], + "mtp_offset": 2 + }, + "commit_gates": { + "fmt": { + "status": "passed", + "log_sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + }, + "clippy": { + "status": "passed", + "log_sha256": "d2d16d1a8d8e5adb2a590b9e1f424c52e7ed907a53b3ed35bbae82e7e858aa14" + }, + "bundle": { + "status": "passed", + "log_sha256": "abab34fc4488a041b6cdd2dd276957888c9e1994cceb1a88f4949c06e274effc" + }, + "tests": { + "status": "passed", + "log_sha256": "70e2bc94647ccdb67c1b3b6197590bed21f2bf90e7eaf2d4008a605a0ea74320", + "test_results": [ + "test result: ok. 308 passed; 0 failed; 216 ignored; 0 measured; 0 filtered out; finished in 5.57s", + "test result: ok. 4 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.01s", + "test result: ok. 4 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 2.39s" + ] + } + }, + "oracle": { + "revision": "9882703f3105363ddc37eca9f97aa09a1d387112", + "mlx": "0.32.2", + "mtp_cases": 8, + "ar_control": "passed" + } +} diff --git a/docs/issue106-final-primary/native-matrix.jsonl b/docs/issue106-final-primary/native-matrix.jsonl new file mode 100644 index 0000000..44ce865 --- /dev/null +++ b/docs/issue106-final-primary/native-matrix.jsonl @@ -0,0 +1,8 @@ +{"case":"mtp_exit","label":"bonus primary","tokens":[1,1,1,3],"key_length":9,"trunk_offsets":[9,9,9,9,9,9,9,9,9,9,9,9],"mtp_offset":8,"oracle_matches":true,"cycles":[{"primary":1,"primary_already_emitted":false,"drafts":[1,1],"accepted":2,"correction":null,"bonus":3,"capture_committed":false}]} +{"case":"mtp_exit","label":"deferred greedy correction","tokens":[1,1,1,3],"key_length":9,"trunk_offsets":[9,9,9,9,9,9,9,9,9,9,9,9],"mtp_offset":8,"oracle_matches":true,"cycles":[{"primary":1,"primary_already_emitted":false,"drafts":[1,1,1],"accepted":2,"correction":null,"bonus":null,"capture_committed":true},{"primary":3,"primary_already_emitted":false,"drafts":[],"accepted":0,"correction":null,"bonus":null,"capture_committed":false}]} +{"case":"mtp_exit","label":"accepted draft stop","tokens":[1,1,1,3],"key_length":9,"trunk_offsets":[9,9,9,9,9,9,9,9,9,9,9,9],"mtp_offset":8,"oracle_matches":true,"cycles":[{"primary":1,"primary_already_emitted":false,"drafts":[1,1,3],"accepted":3,"correction":null,"bonus":null,"capture_committed":false}]} +{"case":"mtp_exit","label":"sampled correction","tokens":[1,1,1,3],"key_length":9,"trunk_offsets":[9,9,9,9,9,9,9,9,9,9,9,9],"mtp_offset":8,"oracle_matches":true,"cycles":[{"primary":1,"primary_already_emitted":false,"drafts":[1,1,1],"accepted":2,"correction":3,"bonus":null,"capture_committed":true}]} +{"case":"mtp_exit","label":"first primary","tokens":[1,3],"key_length":7,"trunk_offsets":[7,7,7,7,7,7,7,7,7,7,7,7],"mtp_offset":6,"oracle_matches":true,"cycles":[{"primary":1,"primary_already_emitted":false,"drafts":[1,1,1],"accepted":0,"correction":null,"bonus":null,"capture_committed":true},{"primary":3,"primary_already_emitted":false,"drafts":[],"accepted":0,"correction":null,"bonus":null,"capture_committed":false}]} +{"case":"mtp_exit","label":"fresh primary depth 1","tokens":[1,1,3],"key_length":8,"trunk_offsets":[8,8,8,8,8,8,8,8,8,8,8,8],"mtp_offset":7,"oracle_matches":true,"cycles":[{"primary":1,"primary_already_emitted":false,"drafts":[1],"accepted":1,"correction":null,"bonus":3,"capture_committed":false}]} +{"case":"mtp_exit","label":"max tokens depth 3","tokens":[1,1,1,1,1,1,1],"key_length":12,"trunk_offsets":[12,12,12,12,12,12,12,12,12,12,12,12],"mtp_offset":11,"oracle_matches":true,"cycles":[{"primary":1,"primary_already_emitted":false,"drafts":[1,1,1],"accepted":3,"correction":null,"bonus":1,"capture_committed":false},{"primary":1,"primary_already_emitted":true,"drafts":[1,1],"accepted":2,"correction":null,"bonus":null,"capture_committed":false}]} +{"case":"mtp_exit","label":"max tokens depth 2","tokens":[1,1,1,1,1],"key_length":10,"trunk_offsets":[10,10,10,10,10,10,10,10,10,10,10,10],"mtp_offset":9,"oracle_matches":true,"cycles":[{"primary":1,"primary_already_emitted":false,"drafts":[1,1],"accepted":2,"correction":null,"bonus":1,"capture_committed":false},{"primary":1,"primary_already_emitted":true,"drafts":[1],"accepted":1,"correction":null,"bonus":null,"capture_committed":false}]} diff --git a/docs/issue106-final-primary/native-product.jsonl b/docs/issue106-final-primary/native-product.jsonl new file mode 100644 index 0000000..4cb8a15 --- /dev/null +++ b/docs/issue106-final-primary/native-product.jsonl @@ -0,0 +1 @@ +{"case":"product_checkpoint_recovery","old_removed":true,"cold_tokens":0,"new_magic":"QWMTKV02","fresh_state_equal":true} diff --git a/docs/issue106-final-primary/native-state.jsonl b/docs/issue106-final-primary/native-state.jsonl new file mode 100644 index 0000000..1e7bf8c --- /dev/null +++ b/docs/issue106-final-primary/native-state.jsonl @@ -0,0 +1,15 @@ +{"case":"complete_state","depth":1,"temperature":0.0,"exit":"max_tokens","tokens":[11],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"41fe15b0c3043a6ad54b55784288e2b2e51ff52106f2f80973be206319dd34db","next_sha256":"6232c8c27ec724f5b0bd585a1dcc8dc840f1f7f58e2a7cbf8050849dc4a002cd","old_checkpoint_rejected":true} +{"case":"complete_state","depth":1,"temperature":0.0,"exit":"stop","tokens":[11],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"41fe15b0c3043a6ad54b55784288e2b2e51ff52106f2f80973be206319dd34db","next_sha256":"6232c8c27ec724f5b0bd585a1dcc8dc840f1f7f58e2a7cbf8050849dc4a002cd","old_checkpoint_rejected":true} +{"case":"complete_state","depth":1,"temperature":0.6,"exit":"max_tokens","tokens":[314],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"06e36aaaee43822c7fc71fd1ed554fb4cd5ab92a94dd88218894f25d56f69b89","next_sha256":"e819026b44ad53bcee190e9121bf347e1f9825171ad5c7e97a051062f2cf91b4","old_checkpoint_rejected":true} +{"case":"complete_state","depth":1,"temperature":0.6,"exit":"stop","tokens":[314],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"06e36aaaee43822c7fc71fd1ed554fb4cd5ab92a94dd88218894f25d56f69b89","next_sha256":"e819026b44ad53bcee190e9121bf347e1f9825171ad5c7e97a051062f2cf91b4","old_checkpoint_rejected":true} +{"case":"complete_state","depth":2,"temperature":0.0,"exit":"max_tokens","tokens":[11],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"41fe15b0c3043a6ad54b55784288e2b2e51ff52106f2f80973be206319dd34db","next_sha256":"6232c8c27ec724f5b0bd585a1dcc8dc840f1f7f58e2a7cbf8050849dc4a002cd","old_checkpoint_rejected":true} +{"case":"complete_state","depth":2,"temperature":0.0,"exit":"stop","tokens":[11],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"41fe15b0c3043a6ad54b55784288e2b2e51ff52106f2f80973be206319dd34db","next_sha256":"6232c8c27ec724f5b0bd585a1dcc8dc840f1f7f58e2a7cbf8050849dc4a002cd","old_checkpoint_rejected":true} +{"case":"complete_state","depth":2,"temperature":0.6,"exit":"max_tokens","tokens":[314],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"06e36aaaee43822c7fc71fd1ed554fb4cd5ab92a94dd88218894f25d56f69b89","next_sha256":"e819026b44ad53bcee190e9121bf347e1f9825171ad5c7e97a051062f2cf91b4","old_checkpoint_rejected":true} +{"case":"complete_state","depth":2,"temperature":0.6,"exit":"stop","tokens":[314],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"06e36aaaee43822c7fc71fd1ed554fb4cd5ab92a94dd88218894f25d56f69b89","next_sha256":"e819026b44ad53bcee190e9121bf347e1f9825171ad5c7e97a051062f2cf91b4","old_checkpoint_rejected":true} +{"case":"complete_state","depth":3,"temperature":0.0,"exit":"max_tokens","tokens":[11],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"41fe15b0c3043a6ad54b55784288e2b2e51ff52106f2f80973be206319dd34db","next_sha256":"6232c8c27ec724f5b0bd585a1dcc8dc840f1f7f58e2a7cbf8050849dc4a002cd","old_checkpoint_rejected":true} +{"case":"complete_state","depth":3,"temperature":0.0,"exit":"stop","tokens":[11],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"41fe15b0c3043a6ad54b55784288e2b2e51ff52106f2f80973be206319dd34db","next_sha256":"6232c8c27ec724f5b0bd585a1dcc8dc840f1f7f58e2a7cbf8050849dc4a002cd","old_checkpoint_rejected":true} +{"case":"complete_state","depth":3,"temperature":0.6,"exit":"max_tokens","tokens":[314],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"06e36aaaee43822c7fc71fd1ed554fb4cd5ab92a94dd88218894f25d56f69b89","next_sha256":"e819026b44ad53bcee190e9121bf347e1f9825171ad5c7e97a051062f2cf91b4","old_checkpoint_rejected":true} +{"case":"complete_state","depth":3,"temperature":0.6,"exit":"stop","tokens":[314],"key_length":4,"trunk_offsets":[4,4,4,4,4,4,4,4,4,4,4,4],"mtp_offset":3,"reopened_cached_tokens":4,"state_sha256":"06e36aaaee43822c7fc71fd1ed554fb4cd5ab92a94dd88218894f25d56f69b89","next_sha256":"e819026b44ad53bcee190e9121bf347e1f9825171ad5c7e97a051062f2cf91b4","old_checkpoint_rejected":true} +{"case":"cancel","safe_final_state":false} +{"case":"forward","safe_final_state":false} +{"case":"capture_disabled","safe_final_state":false} diff --git a/docs/issue106-final-primary/oracle.jsonl b/docs/issue106-final-primary/oracle.jsonl new file mode 100644 index 0000000..267e363 --- /dev/null +++ b/docs/issue106-final-primary/oracle.jsonl @@ -0,0 +1,9 @@ +{"case": "mtp_exit", "label": "bonus primary", "stop_index": 8, "depth": 2, "omniscient": false, "temperature": 0.0, "max_tokens": 16, "tokens": [1, 1, 1, 3], "gap": 0, "origin": "bonus", "safe": true} +{"case": "mtp_exit", "label": "deferred greedy correction", "stop_index": 8, "depth": 3, "omniscient": false, "temperature": 0.0, "max_tokens": 16, "tokens": [1, 1, 1, 3], "gap": 0, "origin": "primary", "safe": true} +{"case": "mtp_exit", "label": "accepted draft stop", "stop_index": 8, "depth": 3, "omniscient": true, "temperature": 0.0, "max_tokens": 16, "tokens": [1, 1, 1, 3], "gap": 0, "origin": "accepted_draft", "safe": true} +{"case": "mtp_exit", "label": "sampled correction", "stop_index": 8, "depth": 3, "omniscient": false, "temperature": 0.7, "max_tokens": 16, "tokens": [1, 1, 1, 3], "gap": 0, "origin": "residual_correction", "safe": true} +{"case": "mtp_exit", "label": "first primary", "stop_index": 6, "depth": 3, "omniscient": false, "temperature": 0.0, "max_tokens": 16, "tokens": [1, 3], "gap": 0, "origin": "primary", "safe": true} +{"case": "mtp_exit", "label": "fresh primary depth 1", "stop_index": 7, "depth": 1, "omniscient": false, "temperature": 0.0, "max_tokens": 16, "tokens": [1, 1, 3], "gap": 0, "origin": "bonus", "safe": true} +{"case": "mtp_exit", "label": "max tokens depth 3", "stop_index": 10000, "depth": 3, "omniscient": false, "temperature": 0.0, "max_tokens": 7, "tokens": [1, 1, 1, 1, 1, 1, 1], "gap": 0, "origin": null, "safe": true} +{"case": "mtp_exit", "label": "max tokens depth 2", "stop_index": 10000, "depth": 2, "omniscient": false, "temperature": 0.0, "max_tokens": 5, "tokens": [1, 1, 1, 1, 1], "gap": 0, "origin": null, "safe": true} +{"case": "ar_control", "passed": true} diff --git a/src/engine/metal/qwen_mtplx.rs b/src/engine/metal/qwen_mtplx.rs index 58756c5..47bc1c4 100644 --- a/src/engine/metal/qwen_mtplx.rs +++ b/src/engine/metal/qwen_mtplx.rs @@ -33,6 +33,8 @@ mod issue103_analysis; mod issue104; #[cfg(test)] mod issue105; +#[cfg(test)] +mod issue106; mod linear; mod load; mod mlp; diff --git a/src/engine/metal/qwen_mtplx/execution.rs b/src/engine/metal/qwen_mtplx/execution.rs index fbeb9a8..57f7495 100644 --- a/src/engine/metal/qwen_mtplx/execution.rs +++ b/src/engine/metal/qwen_mtplx/execution.rs @@ -230,6 +230,8 @@ pub(super) struct Execution { pub(super) fault: Option, #[cfg(test)] pub(super) captured_prefixes: usize, + #[cfg(test)] + pub(super) script: Option, _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) } diff --git a/src/engine/metal/qwen_mtplx/issue103_analysis.rs b/src/engine/metal/qwen_mtplx/issue103_analysis.rs index 2e8c280..4223e12 100644 --- a/src/engine/metal/qwen_mtplx/issue103_analysis.rs +++ b/src/engine/metal/qwen_mtplx/issue103_analysis.rs @@ -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::>() - }; - 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::>()),"good_sha256":super::super::qwen_mtplx_tests::hex_digest(&good.iter().flat_map(|v|v.to_le_bytes()).collect::>())}) - ); - assert!(different > 0); -} - #[test] #[ignore = "#103 diagnostic; small active attention module on Metal; mixed position boundary"] fn issue103_vision_position_receipts() { diff --git a/src/engine/metal/qwen_mtplx/issue106.rs b/src/engine/metal/qwen_mtplx/issue106.rs new file mode 100644 index 0000000..b6717d3 --- /dev/null +++ b/src/engine/metal/qwen_mtplx/issue106.rs @@ -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, +} + +fn offsets(execution: &Execution) -> Vec { + 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 { + 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::>(); + 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::(l).unwrap()) + .collect::>(); + 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::>(), + ) + .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(); +} diff --git a/src/engine/metal/qwen_mtplx/issue106_product.rs b/src/engine/metal/qwen_mtplx/issue106_product.rs new file mode 100644 index 0000000..ea39626 --- /dev/null +++ b/src/engine/metal/qwen_mtplx/issue106_product.rs @@ -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}) + ); +} diff --git a/src/engine/metal/qwen_mtplx/owner.rs b/src/engine/metal/qwen_mtplx/owner.rs index 883d10f..a5ce447 100644 --- a/src/engine/metal/qwen_mtplx/owner.rs +++ b/src/engine/metal/qwen_mtplx/owner.rs @@ -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; diff --git a/src/engine/metal/qwen_mtplx/session_checkpoint.rs b/src/engine/metal/qwen_mtplx/session_checkpoint.rs index b83a04f..7ff6ca1 100644 --- a/src/engine/metal/qwen_mtplx/session_checkpoint.rs +++ b/src/engine/metal/qwen_mtplx/session_checkpoint.rs @@ -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; diff --git a/src/engine/metal/qwen_mtplx/speculative.rs b/src/engine/metal/qwen_mtplx/speculative.rs index ebc570e..557aa11 100644 --- a/src/engine/metal/qwen_mtplx/speculative.rs +++ b/src/engine/metal/qwen_mtplx/speculative.rs @@ -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; } diff --git a/tools/issue106-reference.py b/tools/issue106-reference.py new file mode 100644 index 0000000..8bcf7c9 --- /dev/null +++ b/tools/issue106-reference.py @@ -0,0 +1,38 @@ +"""Run the pinned MTPLX final-state exit matrix without loading model weights.""" +import importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys + +root = Path(sys.argv[1]).resolve() +assert subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112" +assert os.environ["MTPLX_CONTEXT_COPY"] == "0" +sys.path.insert(0, str(root)) +spec = importlib.util.spec_from_file_location("final_state", root / "tests/test_mtp_final_state_covers_committed_tokens.py") +m = importlib.util.module_from_spec(spec) +spec.loader.exec_module(m) +for label, stop_index, depth, omniscient, temperature, maximum in [ + ("bonus primary", 8, 2, False, 0.0, 16), + ("deferred greedy correction", 8, 3, False, 0.0, 16), + ("accepted draft stop", 8, 3, True, 0.0, 16), + ("sampled correction", 8, 3, False, 0.7, 16), + ("first primary", 6, 3, False, 0.0, 16), + ("fresh primary depth 1", 7, 1, False, 0.0, 16), + ("max tokens depth 3", 10000, 3, False, 0.0, 7), + ("max tokens depth 2", 10000, 2, False, 0.0, 5), +]: + sampler = m.SamplerConfig(temperature=temperature, top_k=1 if temperature else 0) + out = m.generate_mtpk(m._runtime(m._ScriptedModel(stop_index, omniscient)), m.PROMPT, + speculative_depth=depth, max_tokens=maximum, sampler=sampler, **m.COMMON) + gap, origin, safe = m._gap(m.PROMPT, out) + assert gap == 0 and safe + assert out.final_state.final_committed_mtp_cache[0].offset == len(m.PROMPT) + len(out.tokens) - 1 + if stop_index < 1000: + assert out.tokens[-1] == m.STOP + print(json.dumps(dict(case="mtp_exit", label=label, stop_index=stop_index, depth=depth, + omniscient=omniscient, temperature=temperature, max_tokens=maximum, + tokens=out.tokens, gap=gap, origin=origin, safe=safe)), flush=True) +m.test_ar_final_state_control() +print(json.dumps(dict(case="ar_control", passed=True)), flush=True)