From dccf81971fab7165927592366a5a65e76b6486c1 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 4 Oct 2026 14:03:12 +0200 Subject: [PATCH] Preserve image-aware sparse QSA for Qwen vision --- docs/issue107-vision-qsa/README.md | 104 ++++++ .../baseline-capture.patch | 102 ++++++ docs/issue107-vision-qsa/before.jsonl | 1 + docs/issue107-vision-qsa/checks.json | 132 +++++++ docs/issue107-vision-qsa/indexer-native.jsonl | 15 + docs/issue107-vision-qsa/indexer-oracle.jsonl | 17 + docs/issue107-vision-qsa/model-oracle.jsonl | 7 + docs/issue107-vision-qsa/product.jsonl | 16 + .../issue107-vision-qsa/small-attention.jsonl | 4 + src/engine/metal/qwen_mtplx/binding.rs | 1 + src/engine/metal/qwen_mtplx/indexing.rs | 7 +- .../metal/qwen_mtplx/issue103_analysis.rs | 6 +- .../metal/qwen_mtplx/issue107_indexer.rs | 180 ++++++++++ .../metal/qwen_mtplx/issue107_product.rs | 328 ++++++++++++++++++ src/engine/metal/qwen_mtplx/owner.rs | 6 +- src/engine/metal/qwen_mtplx/qsa_attention.rs | 19 +- src/engine/metal/qwen_mtplx/qsa_indexer.rs | 99 ++++-- src/engine/metal/qwen_mtplx/qsa_select.rs | 1 + src/engine/metal/qwen_mtplx/request.rs | 9 + src/engine/metal/qwen_mtplx/rope.rs | 47 +++ .../metal/qwen_mtplx/session_checkpoint.rs | 6 + src/engine/metal/qwen_mtplx/vision.rs | 1 + src/engine/qwen_vision.rs | 26 +- .../fixtures/mtplx-qsa-attention-module.jsonl | 24 +- tools/issue107-indexer-reference.py | 53 +++ tools/issue107-model-reference.py | 48 +++ 26 files changed, 1215 insertions(+), 44 deletions(-) create mode 100644 docs/issue107-vision-qsa/README.md create mode 100644 docs/issue107-vision-qsa/baseline-capture.patch create mode 100644 docs/issue107-vision-qsa/before.jsonl create mode 100644 docs/issue107-vision-qsa/checks.json create mode 100644 docs/issue107-vision-qsa/indexer-native.jsonl create mode 100644 docs/issue107-vision-qsa/indexer-oracle.jsonl create mode 100644 docs/issue107-vision-qsa/model-oracle.jsonl create mode 100644 docs/issue107-vision-qsa/product.jsonl create mode 100644 docs/issue107-vision-qsa/small-attention.jsonl create mode 100644 src/engine/metal/qwen_mtplx/issue107_indexer.rs create mode 100644 src/engine/metal/qwen_mtplx/issue107_product.rs create mode 100644 tools/issue107-indexer-reference.py create mode 100644 tools/issue107-model-reference.py diff --git a/docs/issue107-vision-qsa/README.md b/docs/issue107-vision-qsa/README.md new file mode 100644 index 0000000..50561c9 --- /dev/null +++ b/docs/issue107-vision-qsa/README.md @@ -0,0 +1,104 @@ +# Issue 107: image-aware sparse QSA + +Qwen vision now retains the indexer's sparse selection. The indexer receives the +request's position table and delta, rotates queries row by row and pooled keys at +each block's first token, and uses the eager preparation route selected by pinned +MTPLX `9882703f3105363ddc37eca9f97aa09a1d387112`. Its compiled/fused preparation +is ineligible with image positions; the existing sparse selectors/consumers remain +eligible. Budgets, thresholds, text preparation and attention projection RoPE are +unchanged. Issue 108 still owns the attention projection's table-straddling case. + +The indexer's new helper implements `_vision_position_cos_sin` using the native +array operations, including I32 positions, clipped table gather, per-row selection, +axis gather, trigonometry and YaRN amplitude. Signed I32 Gather reuses the existing +U32 payload-copy kernel: it performs no value arithmetic or float conversion. +Tests include signed extrema on both front-axis and general gather routes. + +Image identities now include a numerical policy: 1 for image-aware sparse QSA, +0 for the old dense policy (also MTPLX's explicit diagnostic rollback flag). +Missing fields deserialize as 0. Old vision checkpoints are rejected before +tensor materialization, obsolete live/resident owners are discarded, and stale +owners cannot save a checkpoint. Identical image bytes and token spans do not +bypass this policy. Text-only identities remain empty and compatible. The V1 +signature rejection delivered in issue 106 remains in force; this issue also +rejects affected V2 vision files without invalidating unaffected V2 text files. + +Request prefix matching additionally rejects a restore ending strictly inside an +image span, matching MTPLX's image-safe restore boundary. Complete image spans and +preceding text prefixes remain eligible. Existing image digests continue to +distinguish different pixels with identical image-pad token IDs. + +## Acceptance and review + +- The owned `vision_sparse_qsa` regression now requires equal sparse output for + text and a vision context with equal text positions. Separate issue 108/109 + diagnostic assertions remain unchanged. +- Real indexer geometry: four heads, dim128, ratio4, top-k512, 2,072 raw key rows, + 518 pooled blocks. Query rows and block starts cross a 2,060-row position table. + Table, equal-axis, delta-only and empty-table cases match pinned MTPLX exactly + in F32 and BF16; independent NumPy math verifies the F32 values. Normal fused + selectors remain active while compiled preparation is correctly bypassed. +- A real 2,095-token, two-image request exceeds the sparse budget. Corrected final + model logits match MTPLX byte-for-byte: SHA256 + `8b97664316bec6113c85b7dfca9b09967391f019ea1064b83d81a3d9d5191354`. + Pre-fix logits differ (`a79aa02a98542adeea636f22c85fc5f929308dc5dc75658d4072e6f275d22abe`). +- The actual pre-fix dense-vision checkpoint is 412,628,836 bytes, SHA256 + `7ddc663d4f337ff272000a7237dac05a224262bf25bc88257e784547b2056326`. + Tests reject its normal disk load and import its original tensor payload only + through a test helper to verify stale live/resident rejection and save refusal. +- Product tests cover unchanged images, identical complete-state restore/reopen + and follow-up, text edits before/between/after two images, different pixels with + identical pad IDs, cancellation/resend, a prefix ending inside an image, and + unaffected text checkpoints. Comparisons cover all serialized tensors and state + metadata, not only response text. +- DS4/GLM do not call the modified Qwen modules. Their sampling, RNG, vision and + checkpoint formats are unchanged. This is MTPLX behavior plus DS4Server's + application-specific checkpoint policy integration. +- The 12 affected whole-attention fixture rows were regenerated from the pinned + oracle (profiles 7 and 9, TF32 on/off). All 56 fixture rows pass with both lazy + and prepared inputs, including output/state hashes, consumer routes, compiled + core lifecycle and counters. Existing text-only indexer route fixtures pass. + +All four commit gates passed. The initial general test run hit an unrelated +`dev_brain` fixture whose temporary short Git hash was all digits and therefore +parsed as a YAML integer. Its focused retry and the subsequent full suite passed; +no unrelated source change was included. `checks.json` records both outcomes. + +The tests and receipts are correctness evidence, not performance measurements. +See `checks.json` and the accompanying JSONL files for actual run outcomes. + +## Reproduction + +Use the same unchanged Bare-Speed artifact and environment documented in +`../issue104-ar/environment.json`, including the pinned vision tower. The input +image is `local-eval-results/qwen-vision-input/image.png` (1024 square pixels), +repeated twice; tests generate an independent 1024-square blue control. + +An isolated checkout of parent `a8cad88` plus `baseline-capture.patch` regenerates +the actual old snapshot with `issue107_vision_product`. Set `DS4_ISSUE107_DIR` to +its output directory. Do not replace a new file's header to fake old state. + +Generate fresh reference results in the isolated MLX 0.32.2 environment: + +```sh +python tools/issue107-model-reference.py PINNED_MTPLX MODEL IMAGE ORACLE_DIR +python tools/issue107-indexer-reference.py PINNED_MTPLX INDEXER_DIR +``` + +Build native tests with `cargo test --all-features --bin ds4-server --no-run`. +Run ignored tests one at a time via `test-supervisor MEMORY_MIB 60 120 --command +TEST_EXECUTABLE TEST_NAME --ignored --nocapture --test-threads=1`. Use 114688 MiB +for the product test and 4096 MiB for model-free tests. Set +`DS4_QWEN38_ARTIFACTS`, `DS4_QWEN_VISION_IMAGE`, `DS4_ISSUE107_DIR`, +`DS4_ISSUE107_BAD_CHECKPOINT`, `DS4_ISSUE107_ORACLE_DIR`, and +`DS4_ISSUE107_INDEXER` to the corresponding files/directories. + +Owned tests: `issue107_vision_product`, +`issue107_real_geometry_query_and_block_start_positions`, and the +`vision_sparse_qsa` case in `issue103_vision_position_receipts`. + +To refresh just the affected whole-attention receipts, set `PYTHONPATH` to the +pinned checkout and run `qsa_attention_module_fixtures(profile_filter=7)` and +`qsa_attention_module_fixtures(profile_filter=9)` from +`runpy.run_path("tools/mtplx-kernel-fixtures.py")` in separate processes with +`MLX_ENABLE_TF32=1` and `0`. Replace only those profile/TF32 rows in the fixture. diff --git a/docs/issue107-vision-qsa/baseline-capture.patch b/docs/issue107-vision-qsa/baseline-capture.patch new file mode 100644 index 0000000..56a1f6f --- /dev/null +++ b/docs/issue107-vision-qsa/baseline-capture.patch @@ -0,0 +1,102 @@ +--- a/src/engine/metal/qwen_mtplx/owner.rs ++++ b/src/engine/metal/qwen_mtplx/owner.rs +@@ -1068,3 +1068,7 @@ + #[cfg(test)] + #[path = "issue106_product.rs"] + mod issue106_product; ++ ++#[cfg(test)] ++#[path = "issue107_product.rs"] ++mod issue107_product; +--- /dev/null ++++ b/src/engine/metal/qwen_mtplx/issue107_product.rs +@@ -0,0 +1,89 @@ ++//! Long multi-image request and cache compatibility through the product owner. ++use super::*; ++use crate::settings::{ ++ GenerationPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, ++ effective_settings, ++}; ++use base64::Engine as _; ++ ++#[test] ++#[ignore = "installed model plus vision tower; serial external memory supervisor required"] ++fn issue107_vision_product() { ++ let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap()); ++ let image = std::fs::read(std::env::var_os("DS4_QWEN_VISION_IMAGE").unwrap()).unwrap(); ++ let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_DIR").unwrap()); ++ std::fs::create_dir_all(&directory).unwrap(); ++ let settings = effective_settings( ++ crate::model::ModelChoice::Qwen38FlashNext, ++ &GenerationPreferences { ++ context_tokens: 8192, ++ max_generated_tokens: 1, ++ system_prompt: String::new(), ++ temperature: Some(0.), ++ top_p: Some(0.95), ++ min_p: Some(0.), ++ seed: Some(1), ++ reasoning_mode: ReasoningMode::Low, ++ }, ++ &RuntimePreferences { ++ speculative: SpeculativePreferences { ++ glm_mtp: true, ++ ..Default::default() ++ }, ++ ..Default::default() ++ }, ++ root.parent().unwrap(), ++ ) ++ .unwrap(); ++ let marker = crate::engine::vision_data_marker(&format!( ++ "data:image/png;base64,{}", ++ base64::engine::general_purpose::STANDARD.encode(&image) ++ )); ++ let messages = [ChatTurn { ++ user: true, ++ tool: false, ++ system: false, ++ skip_previous_eos: false, ++ reasoning: None, ++ reasoning_complete: true, ++ content: format!("Describe these images{marker}{marker}"), ++ }]; ++ let mut owner = QwenExecutor::open(&settings.engine).unwrap(); ++ owner.policy.store_min_suffix = Some(0); ++ let metrics = Metrics::new(&directory.join("metrics")); ++ let (out, complete) = owner ++ .generate( ++ &messages, ++ &settings.turn, ++ &AtomicBool::new(false), ++ &metrics, ++ |p, _, _| { ++ eprintln!("#107 prefill {p}"); ++ }, ++ |_, _| {}, ++ ) ++ .unwrap(); ++ assert!(complete); ++ assert!(out.prompt_tokens > 2048); ++ assert_eq!(out.cached_tokens, 0); ++ let saved = owner.live.current.as_ref().unwrap(); ++ let logits = owner.execution.values_f32(&saved.logits).unwrap(); ++ std::fs::write( ++ directory.join("long-final.f32"), ++ logits ++ .iter() ++ .flat_map(|v| v.to_le_bytes()) ++ .collect::>(), ++ ) ++ .unwrap(); ++ std::fs::write(directory.join("long-input.json"),serde_json::to_vec(&serde_json::json!({"ids":saved.token_ids,"images":owner.live.current_images,"prompt_tokens":out.prompt_tokens})).unwrap()).unwrap(); ++ owner ++ .save_checkpoint(&directory.join("long.bin"), [107; 32], &mut |_| {}) ++ .unwrap(); ++ println!( ++ "{}", ++ serde_json::json!({"case":"vision_long","prompt_tokens":out.prompt_tokens,"cached_tokens":out.cached_tokens, ++ "tokens":owner.live.current.as_ref().unwrap().token_ids,"text":out.message.content,"reasoning":out.message.reasoning, ++ "logits_sha256":crate::engine::metal::qwen_mtplx_tests::hex_digest(&logits.iter().flat_map(|v|v.to_le_bytes()).collect::>())}) ++ ); ++} diff --git a/docs/issue107-vision-qsa/before.jsonl b/docs/issue107-vision-qsa/before.jsonl new file mode 100644 index 0000000..a41c80e --- /dev/null +++ b/docs/issue107-vision-qsa/before.jsonl @@ -0,0 +1 @@ +{"case":"vision_long","prompt_tokens":2095,"cached_tokens":0,"text":"","reasoning":"The","logits_sha256":"a79aa02a98542adeea636f22c85fc5f929308dc5dc75658d4072e6f275d22abe","bank_token_count":2096,"bank_token_ids_sha256":"a757dd27268f65ff6926bc0b2d6d81faec709b89512673068d3c536056c5105b","last_token":760} diff --git a/docs/issue107-vision-qsa/checks.json b/docs/issue107-vision-qsa/checks.json new file mode 100644 index 0000000..e0d2827 --- /dev/null +++ b/docs/issue107-vision-qsa/checks.json @@ -0,0 +1,132 @@ +{ + "before": { + "cases": 1, + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 524 filtered out; finished in 11.84s" + ], + "elapsed_ms": 12922, + "peak_physical_bytes": 81646134816, + "supervisor_error": null, + "log_sha256": "15b89366a4a1488df33e2ada91cc848b9ac6240bf7f8a44f3afe4ff2c2f0c757" + }, + "model-oracle": { + "cases": 7, + "test_result": [], + "elapsed_ms": 15413, + "peak_physical_bytes": 80281215320, + "supervisor_error": null, + "log_sha256": "efe9ea90eaf1d6a3d9e95669cd07f9318fdce6a3ae48d7b8dbecc545622511f2" + }, + "indexer-oracle": { + "cases": 17, + "test_result": [], + "elapsed_ms": 625, + "peak_physical_bytes": 215155408, + "supervisor_error": null, + "log_sha256": "aac8e9211ffc59569716960a5493420ddf9f235e0d54e38ae02a7d27de9cee84" + }, + "indexer-native-final": { + "cases": 15, + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 525 filtered out; finished in 0.28s" + ], + "elapsed_ms": 315, + "peak_physical_bytes": 181633888, + "supervisor_error": null, + "log_sha256": "1e634d173e8a5e4b6e50785ffcb31a70ac3438292293dddfa63316563d208b1e" + }, + "small": { + "cases": 4, + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 525 filtered out; finished in 0.14s" + ], + "elapsed_ms": 210, + "peak_physical_bytes": 37094192, + "supervisor_error": null, + "log_sha256": "afb6f32dc619c6e3ec70ca877e08aeae1fece96c64dc12019f0704354c0bfde6" + }, + "product": { + "cases": 16, + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 525 filtered out; finished in 165.40s" + ], + "elapsed_ms": 166771, + "peak_physical_bytes": 82617511280, + "supervisor_error": null, + "log_sha256": "844160c173fc4ba29b3475e871e940376b2ad553f0be75c17cf0fd049fe862de" + }, + "text-indexer": { + "cases": 0, + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 525 filtered out; finished in 38.29s" + ], + "elapsed_ms": 38368, + "peak_physical_bytes": 1108821168, + "supervisor_error": null, + "log_sha256": "0494120021c6accce4e4976465b6a7c6c2ddbe94af829bd0ca71dd1ef17f2809" + }, + "attention-native-1": { + "cases": 0, + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 525 filtered out; finished in 14.09s" + ], + "elapsed_ms": 14204, + "peak_physical_bytes": 1608435056, + "supervisor_error": null, + "log_sha256": "2ed6ca37f3700b0780cdf5c2262409ff16d2ed8c172ebd0206795c7a71fb4fea" + }, + "attention-native-0": { + "cases": 0, + "test_result": [ + "test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 525 filtered out; finished in 13.94s" + ], + "elapsed_ms": 14007, + "peak_physical_bytes": 1607861592, + "supervisor_error": null, + "log_sha256": "a32c72594ee7ce00b04553e65bca9de5b6cb0125b1d23b3abc574208fddb6e08" + }, + "pre_fix_checkpoint": { + "base": "a8cad8810060d22ee7ccac82c210ef40d4e08785", + "bytes": 412628836, + "sha256": "7ddc663d4f337ff272000a7237dac05a224262bf25bc88257e784547b2056326", + "magic": "QWMTKV02", + "vision_policy": "absent (legacy 0)", + "images": 2, + "prompt_tokens": 2095 + }, + "model_logits": { + "before": "a79aa02a98542adeea636f22c85fc5f929308dc5dc75658d4072e6f275d22abe", + "native": "8b97664316bec6113c85b7dfca9b09967391f019ea1064b83d81a3d9d5191354", + "oracle": "8b97664316bec6113c85b7dfca9b09967391f019ea1064b83d81a3d9d5191354" + }, + "input_image_sha256": "b7568d90f4df6180d9af14a824dd553cd995457561d29167354c1ba66b728347", + "commit_gates": { + "fmt": { + "status": "passed", + "log_sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + }, + "clippy": { + "status": "passed", + "log_sha256": "c437a2464fc3d0eb441259bf1126ad32e23230ebfde7e51a6c3fd06259d1bfc9" + }, + "bundle": { + "status": "passed", + "log_sha256": "0b9900a8392b683ce2804a420fbcbd9cac895658bbad67859d56a9b046110159" + }, + "tests-final": { + "status": "passed", + "log_sha256": "649b53e5b60e39c4a22a8f76246128b2024f76e6d9c9a9ec65b93e1b638cf261", + "results": [ + "test result: ok. 308 passed; 0 failed; 218 ignored; 0 measured; 0 filtered out; finished in 5.52s", + "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.29s" + ] + } + }, + "unrelated_initial_test_failure": { + "test": "dev_brain::tests::source_versions_follow_that_file_instead_of_the_whole_worktree", + "cause": "Temporary commit short hash 6496041 was parsed as an integer in unquoted YAML; test uses git2::Signature::now, so hash changes with time.", + "resolution": "Focused retry passed, followed by full suite passing. No unrelated source change." + }, + "review": "All issue 107 acceptance, oracle route/numerics, Qwen cache policy, and family boundaries reviewed; no remaining gap." +} diff --git a/docs/issue107-vision-qsa/indexer-native.jsonl b/docs/issue107-vision-qsa/indexer-native.jsonl new file mode 100644 index 0000000..caaf59e --- /dev/null +++ b/docs/issue107-vision-qsa/indexer-native.jsonl @@ -0,0 +1,15 @@ +{"case":"indexer","dtype":"f32","mode":"table","part":"pooled","exact_oracle_match":true} +{"case":"indexer","dtype":"f32","mode":"equal","part":"query","exact_oracle_match":true} +{"case":"indexer","dtype":"f32","mode":"equal","part":"pooled","exact_oracle_match":true} +{"case":"indexer","dtype":"f32","mode":"delta","part":"query","exact_oracle_match":true} +{"case":"indexer","dtype":"f32","mode":"delta","part":"pooled","exact_oracle_match":true} +{"case":"indexer","dtype":"f32","mode":"empty","part":"query","exact_oracle_match":true} +{"case":"indexer","dtype":"f32","mode":"empty","part":"pooled","exact_oracle_match":true} +{"case":"indexer","dtype":"bf16","mode":"table","part":"query","exact_oracle_match":true} +{"case":"indexer","dtype":"bf16","mode":"table","part":"pooled","exact_oracle_match":true} +{"case":"indexer","dtype":"bf16","mode":"equal","part":"query","exact_oracle_match":true} +{"case":"indexer","dtype":"bf16","mode":"equal","part":"pooled","exact_oracle_match":true} +{"case":"indexer","dtype":"bf16","mode":"delta","part":"query","exact_oracle_match":true} +{"case":"indexer","dtype":"bf16","mode":"delta","part":"pooled","exact_oracle_match":true} +{"case":"indexer","dtype":"bf16","mode":"empty","part":"query","exact_oracle_match":true} +{"case":"indexer","dtype":"bf16","mode":"empty","part":"pooled","exact_oracle_match":true} diff --git a/docs/issue107-vision-qsa/indexer-oracle.jsonl b/docs/issue107-vision-qsa/indexer-oracle.jsonl new file mode 100644 index 0000000..3cfce5e --- /dev/null +++ b/docs/issue107-vision-qsa/indexer-oracle.jsonl @@ -0,0 +1,17 @@ +{"case":"indexer","dtype":"f32","mode":"table","part":"query","shape":[1,16,4,128],"sha256":"00cf519652dde6ac163822403e7bd26c6f131fde0673633d71f9cd6523ebb730"} +{"case":"indexer","dtype":"f32","mode":"table","part":"pooled","shape":[1,518,128],"sha256":"d4c1f51750ebba1e83c2b5bed38a1962682630f39c0f3beae8455144444d8114"} +{"case":"indexer","dtype":"f32","mode":"equal","part":"query","shape":[1,16,4,128],"sha256":"6ed0da7d959383ae23041ff321283fd7fa99d4e4539c8a4bd0f22bafd8dd457a"} +{"case":"indexer","dtype":"f32","mode":"equal","part":"pooled","shape":[1,518,128],"sha256":"a6117d1182ed7abf52af125c9cfccba6bd1e416f3aa54d47b58d261c1b0f498d"} +{"case":"indexer","dtype":"f32","mode":"delta","part":"query","shape":[1,16,4,128],"sha256":"b7efbed01f410398d6bd2393f45c52071c1f1da109b244f5e9b851814ec7cebc"} +{"case":"indexer","dtype":"f32","mode":"delta","part":"pooled","shape":[1,518,128],"sha256":"f6533b66111c3230d4a6d1c0d32fbaccb97e9a25a42b5b0498671d83ffaa9bcf"} +{"case":"indexer","dtype":"f32","mode":"empty","part":"query","shape":[1,16,4,128],"sha256":"b7efbed01f410398d6bd2393f45c52071c1f1da109b244f5e9b851814ec7cebc"} +{"case":"indexer","dtype":"f32","mode":"empty","part":"pooled","shape":[1,518,128],"sha256":"f6533b66111c3230d4a6d1c0d32fbaccb97e9a25a42b5b0498671d83ffaa9bcf"} +{"case":"indexer","dtype":"bf16","mode":"table","part":"query","shape":[1,16,4,128],"sha256":"0e32ce79b902a3c1974a682486432f7e719dfe6c560490279948c629d7680a20"} +{"case":"indexer","dtype":"bf16","mode":"table","part":"pooled","shape":[1,518,128],"sha256":"a71f98ac2d9db6ad26a83daf9d8736c6a3a0ba0e4e03517a464718b4f72c1cad"} +{"case":"indexer","dtype":"bf16","mode":"equal","part":"query","shape":[1,16,4,128],"sha256":"50533b71c2e539a8f56b8fa5d6a4b4979dafe21b67e8975f8461792dab1f0b67"} +{"case":"indexer","dtype":"bf16","mode":"equal","part":"pooled","shape":[1,518,128],"sha256":"c57f4d215c9d4336c4250f1f5b494cefe156e43f88e712fcec53e789209b242b"} +{"case":"indexer","dtype":"bf16","mode":"delta","part":"query","shape":[1,16,4,128],"sha256":"e95114b8b5cf606c714c06c3d23bbadc1126c502527cbf5b84fc5e7d773e218f"} +{"case":"indexer","dtype":"bf16","mode":"delta","part":"pooled","shape":[1,518,128],"sha256":"b04fe5e8940fcd7cf15760692caa1951bcf7c338ba30d20fe84f31c1b6cca3cb"} +{"case":"indexer","dtype":"bf16","mode":"empty","part":"query","shape":[1,16,4,128],"sha256":"e95114b8b5cf606c714c06c3d23bbadc1126c502527cbf5b84fc5e7d773e218f"} +{"case":"indexer","dtype":"bf16","mode":"empty","part":"pooled","shape":[1,518,128],"sha256":"b04fe5e8940fcd7cf15760692caa1951bcf7c338ba30d20fe84f31c1b6cca3cb"} +{"case":"independent_numpy","passed":true} diff --git a/docs/issue107-vision-qsa/model-oracle.jsonl b/docs/issue107-vision-qsa/model-oracle.jsonl new file mode 100644 index 0000000..ba15336 --- /dev/null +++ b/docs/issue107-vision-qsa/model-oracle.jsonl @@ -0,0 +1,7 @@ +{"case":"environment","overrides":{"MTPLX_SKIP_VERIFY_SNAPSHOT":"0","MTPLX_FAMILY_CAPTURE_COMMIT":"1","MTPLX_FUSED_HC_V3":"1","MTPLX_FUSED_GDN_INPROJ":"1","MTPLX_FUSED_GATE_UP":"1","MTPLX_FUSED_GDN_CONVNORM":"1","MTPLX_FUSED_GDN_STEP":"1","MTPLX_FUSED_CONVNORM_VERIFY":"1","MTPLX_QSA_GATHER":"1","MTPLX_COMPILED_GDN":"1","MTPLX_QWEN4EXP_COMPILE":"1","MTPLX_QWEN4_BATCHED_TARGET_DISTRIBUTIONS":"1","MTPLX_QWEN4_FIXED_M4_VERIFY":"1","MTPLX_QWEN4_COMPILED_MTP_PREPARE":"1","MTPLX_QWEN4_RELAXED_DRAFT_TIES":"1","MTPLX_QWEN4_OPDIET":"1","MTPLX_QWEN4_PLE_PREFILL_LOOKAHEAD":"1","MTPLX_QWEN4_PLE_FIRST_GATHER_EARLY":"1","MTPLX_SESSION_BANK_SHED_BOUNDARIES":"1","MTPLX_SESSION_BANK_PROTECTED_TERMINAL":"1","MTPLX_NGRAM_PREWARM":"auto","MTPLX_QWEN4_VERIFY_GLUE":"1","MTPLX_QWEN4_VERIFY_GLUE_ITEMS":"qsa_rope,qsa_rope_idx","MTPLX_QSA_M4_FUSED_KV_GATHER":"1","MTPLX_QSA_GATHER_MAX_ROWS":"32","MTPLX_BATCH_TARGET_ARRAYS":"1","MTPLX_LAZY_TARGET_DISTRIBUTIONS":"0","MTPLX_COMPILED_VERIFY":"1","MTPLX_NAX_VERIFY":"0","MTPLX_QWEN4_PREFILL_WIDE_CHUNK":"4096","MTPLX_QSA_PREFILL_WIDE_MIN_CONTEXT":"16384","MTPLX_SUSTAINED_PREFILL":"1","MTPLX_PREFILL_CHUNK_SIZE":"2048"}} +{"case":"prefill","phase":"started","tokens_done":0,"tokens_total":2095,"cached_tokens":0,"new_prefill_tokens":2095,"elapsed_s":0.0,"started_s":174050.1008485} +{"case":"prefill","phase":"chunk","tokens_done":0,"tokens_total":2095,"cached_tokens":0,"new_prefill_tokens":2095,"elapsed_s":3.433300298638642e-05,"cache_hit":false,"reread":{"prompt_tokens":2095,"restore_point_tokens":0,"recompute_tokens":2095,"history_matched_tokens":0,"session_cached_tokens":0,"source":"none","cache_off":true,"image_changed_at":null,"image_limit_at":null}} +{"case":"prefill","phase":"chunk","tokens_done":2048,"tokens_total":2095,"cached_tokens":0,"elapsed_s":4.264302124996902,"prefill_tok_s":480.266158440049,"cumulative_prefill_tok_s":480.266158440049,"prefill_wall_tok_s":480.266158440049,"live_prefill_tok_s":480.6052339792707,"chunk_size":2048,"chunk_elapsed_s":4.26129358401522,"chunk_prefill_tok_s":480.6052339792707} +{"case":"prefill","phase":"chunk","tokens_done":2094,"tokens_total":2095,"cached_tokens":0,"elapsed_s":4.605880165996496,"prefill_tok_s":454.63623119403434,"cumulative_prefill_tok_s":454.63623119403434,"prefill_wall_tok_s":454.63623119403434,"live_prefill_tok_s":140.22877867288037,"chunk_size":46,"chunk_elapsed_s":0.32803537501604296,"chunk_prefill_tok_s":140.22877867288037} +{"case":"prefill","phase":"completed","tokens_total":2095,"cached_tokens":0,"new_prefill_tokens":2095,"elapsed_s":4.6907559579994995,"prompt_eval_time_s":4.686342919012532,"prefill_tok_s":447.04368335927103,"prefill_compute_tok_s":447.04368335927103,"prefill_wall_tok_s":446.6231069700479,"cache_hit":false} +{"case":"long_vision","prompt_tokens":2095,"text":"The","final_logits_sha256":"8b97664316bec6113c85b7dfca9b09967391f019ea1064b83d81a3d9d5191354","bank_token_count":1,"bank_token_ids_sha256":"a9cf0c5bfacb7118d245ee1a31149cbb264f7138ee585445dd76cc681d755554","last_token":760} diff --git a/docs/issue107-vision-qsa/product.jsonl b/docs/issue107-vision-qsa/product.jsonl new file mode 100644 index 0000000..c27318c --- /dev/null +++ b/docs/issue107-vision-qsa/product.jsonl @@ -0,0 +1,16 @@ +{"case":"vision_long","prompt_tokens":2095,"cached_tokens":0,"text":"","reasoning":"The","logits_sha256":"8b97664316bec6113c85b7dfca9b09967391f019ea1064b83d81a3d9d5191354","bank_token_count":2096,"bank_token_ids_sha256":"a757dd27268f65ff6926bc0b2d6d81faec709b89512673068d3c536056c5105b","last_token":760} +{"case":"unchanged","cached_tokens":2095,"prompt_tokens":2095,"state_sha256":"958ff3ac9ad31c986085a45d31a7424475660263371b5a509b9d0a4ccaf44ab3"} +{"case":"resident_followup","cached_tokens":2096,"prompt_tokens":2113,"state_sha256":"a5765c3da1ce1a75ce3b57e1b72a7a97b3f5874886c3eccd56c677127adebeef"} +{"case":"reopened_followup","cached_tokens":2096,"prompt_tokens":2113,"state_sha256":"a5765c3da1ce1a75ce3b57e1b72a7a97b3f5874886c3eccd56c677127adebeef"} +{"case":"text_before","cached_tokens":0,"prompt_tokens":2096,"state_sha256":"45a34b05a65378353e5c65df8f868dee7e75d38e381c897523058d37b64193d4"} +{"case":"text_before_cold_control","cached_tokens":0,"prompt_tokens":2096,"state_sha256":"45a34b05a65378353e5c65df8f868dee7e75d38e381c897523058d37b64193d4"} +{"case":"text_between","cached_tokens":0,"prompt_tokens":2098,"state_sha256":"69121fa827322e5308a065d61bbd9f53e6b85ced209214a0a7db4f711cc95e66"} +{"case":"text_between_cold_control","cached_tokens":0,"prompt_tokens":2098,"state_sha256":"69121fa827322e5308a065d61bbd9f53e6b85ced209214a0a7db4f711cc95e66"} +{"case":"text_after","cached_tokens":0,"prompt_tokens":2098,"state_sha256":"156daf37583f00996b960501612a77476e344eaead98726723958b79646ee67d"} +{"case":"text_after_cold_control","cached_tokens":0,"prompt_tokens":2098,"state_sha256":"156daf37583f00996b960501612a77476e344eaead98726723958b79646ee67d"} +{"case":"changed_pixels","cached_tokens":0,"prompt_tokens":2095,"state_sha256":"d1f3be2b52ca0314c9ec6c547babccb0847e29c733c9a31fdb4a55930f9af688"} +{"case":"changed_pixels_cold_control","cached_tokens":0,"prompt_tokens":2095,"state_sha256":"d1f3be2b52ca0314c9ec6c547babccb0847e29c733c9a31fdb4a55930f9af688"} +{"case":"cancel_resend","cached_tokens":0,"prompt_tokens":2095,"state_sha256":"958ff3ac9ad31c986085a45d31a7424475660263371b5a509b9d0a4ccaf44ab3"} +{"case":"inside_image_guard","cached_tokens":0,"prompt_tokens":2095,"state_sha256":"958ff3ac9ad31c986085a45d31a7424475660263371b5a509b9d0a4ccaf44ab3"} +{"case":"text_control","cached_tokens":0,"prompt_tokens":43,"state_sha256":"75bd53e84aa31a13851b02f13b025d5b7b76dddfec5236d2bd70294ab0548850"} +{"case":"compatibility","old_disk_rejected":true,"old_live_rejected":true,"old_resident_rejected":true,"stale_save_rejected":true,"new_vision_and_text_reopen_exactly":true} diff --git a/docs/issue107-vision-qsa/small-attention.jsonl b/docs/issue107-vision-qsa/small-attention.jsonl new file mode 100644 index 0000000..97845b7 --- /dev/null +++ b/docs/issue107-vision-qsa/small-attention.jsonl @@ -0,0 +1,4 @@ +{"case":"vision_positions","start":4,"rows":2,"table_end":4,"delta":-1,"different_rows":[],"actual_keys":[-0.3359375,0.02783203125,0.55078125,0.7890625,-0.94921875,1.25,1.3984375,1.5859375,0.3125,0.294921875,0.8203125,0.94140625,-1.203125,1.3203125,1.2109375,1.2578125],"expected_keys":[-0.3359375,0.02783203125,0.55078125,0.7890625,-0.94921875,1.25,1.3984375,1.5859375,0.3125,0.294921875,0.8203125,0.94140625,-1.203125,1.3203125,1.2109375,1.2578125]} +{"case":"vision_positions","start":2,"rows":4,"table_end":4,"delta":-1,"different_rows":[0,1],"actual_keys":[-0.72265625,0.275390625,0.578125,0.79296875,0.69921875,1.21875,1.390625,1.5859375,-1.2265625,0.55078125,0.84375,0.94140625,0.216796875,1.234375,1.1953125,1.2578125,-0.9609375,0.515625,0.88671875,0.96484375,-0.890625,1.28125,1.140625,1.1640625,0.1962890625,0.4296875,0.90234375,0.9765625,-1.328125,1.3203125,1.125,1.125],"expected_keys":[-0.98046875,0.02783203125,0.5390625,0.7890625,-0.2314453125,1.25,1.40625,1.5859375,-1.2265625,0.423828125,0.8203125,0.94140625,0.216796875,1.28125,1.2109375,1.2578125,-0.9609375,0.515625,0.88671875,0.96484375,-0.890625,1.28125,1.140625,1.1640625,0.1962890625,0.4296875,0.90234375,0.9765625,-1.328125,1.3203125,1.125,1.125]} +{"case":"vision_sparse_qsa","prefix_rows":12,"topk_blocks":1,"ratio":4,"position_delta":0,"text_output":[0.10986328125,-0.1923828125,0.08544921875,0.498046875,-0.0791015625,0.16796875,0.31640625,-0.03173828125],"vision_output":[0.10986328125,-0.1923828125,0.08544921875,0.498046875,-0.0791015625,0.16796875,0.31640625,-0.03173828125]} +{"case":"active_attention_gate","gate":-6.84375,"actual":[0.0003185272216796875,-0.00074005126953125,0.00021266937255859375,0.00095367431640625,-0.00021266937255859375,0.0004253387451171875,0.000637054443359375,-0.00010633468627929688],"expected_verify":[0.0003204345703125,-0.0007476806640625,0.000213623046875,0.0009613037109375,-0.000213623046875,0.00042724609375,0.000640869140625,-0.0001068115234375]} diff --git a/src/engine/metal/qwen_mtplx/binding.rs b/src/engine/metal/qwen_mtplx/binding.rs index 3ce91d5..ebd0bfd 100644 --- a/src/engine/metal/qwen_mtplx/binding.rs +++ b/src/engine/metal/qwen_mtplx/binding.rs @@ -847,6 +847,7 @@ impl<'a> Bind<'a> { q_norm: self.array(&format!("{ip}.q_layernorm.weight"))?, k_norm: self.array(&format!("{ip}.k_layernorm.weight"))?, inv_freq: &r.frequency, + mrope_axes: r.axes.as_ref(), eps: c.rms_norm_eps, scaling: r.scaling, scratch: 32 * 1024 * 1024, diff --git a/src/engine/metal/qwen_mtplx/indexing.rs b/src/engine/metal/qwen_mtplx/indexing.rs index f9907e5..4a0548f 100644 --- a/src/engine/metal/qwen_mtplx/indexing.rs +++ b/src/engine/metal/qwen_mtplx/indexing.rs @@ -1,5 +1,5 @@ //! Select, take/Gather, GatherAxis, Arange and Scan from the pinned runtime. -//! take exports BF16/F16/F32/U32 with I64 and I32 indices. +//! take exports BF16/F16/F32/I32/U32 with I64 and I32 indices. use super::super::gpu::Buffer; use super::array::{Array, Dtype, Layout}; use super::stream::Stream; @@ -503,7 +503,7 @@ pub(super) fn take(table: &Array, ids: &Array, axis: i32, stream: Stream) -> Res || !(il.dtype() == Dtype::I64 || il.dtype() == Dtype::U32 || il.dtype() == Dtype::I32) || !matches!( sl.dtype(), - Dtype::BF16 | Dtype::F16 | Dtype::F32 | Dtype::U32 + Dtype::BF16 | Dtype::F16 | Dtype::F32 | Dtype::I32 | Dtype::U32 ) || (sl.shape()[axis] == 0 && il.size() != 0) { @@ -551,7 +551,8 @@ fn gather_take_evaluate( Dtype::BF16 => "bf16", Dtype::F16 => "f16", Dtype::F32 => "f32", - Dtype::U32 => "u32", + // Gather copies 32-bit payloads without value arithmetic. + Dtype::I32 | Dtype::U32 => "u32", _ => unreachable!(), }; let large = sl.size() > i32::MAX as usize diff --git a/src/engine/metal/qwen_mtplx/issue103_analysis.rs b/src/engine/metal/qwen_mtplx/issue103_analysis.rs index 4223e12..ee64e4f 100644 --- a/src/engine/metal/qwen_mtplx/issue103_analysis.rs +++ b/src/engine/metal/qwen_mtplx/issue103_analysis.rs @@ -532,6 +532,7 @@ fn issue103_vision_position_receipts() { q_norm: &norm, k_norm: &norm, inv_freq: &freq, + mrope_axes: Some(&axes), eps: 1e-6, scaling: 1., scratch: 32 * 1024 * 1024, @@ -566,7 +567,10 @@ fn issue103_vision_position_receipts() { "{}", serde_json::json!({"case":"vision_sparse_qsa","prefix_rows":12,"topk_blocks":1,"ratio":4,"position_delta":0,"text_output":outputs[0],"vision_output":outputs[1]}) ); - assert_ne!(outputs[0], outputs[1], "diagnosis no longer reproduces"); + assert_eq!( + outputs[0], outputs[1], + "equal vision positions must preserve sparse QSA" + ); let qweight = bf( &[16, 8], &(0..128) diff --git a/src/engine/metal/qwen_mtplx/issue107_indexer.rs b/src/engine/metal/qwen_mtplx/issue107_indexer.rs new file mode 100644 index 0000000..c454d00 --- /dev/null +++ b/src/engine/metal/qwen_mtplx/issue107_indexer.rs @@ -0,0 +1,180 @@ +use super::*; + +#[test] +#[ignore = "Metal and independently verified pinned indexer fixtures; external supervisor required"] +fn issue107_real_geometry_query_and_block_start_positions() { + super::super::configure_sources().unwrap(); + let _context = crate::engine::metal::gpu::Context::open_qwen(0).unwrap(); + let streams = super::super::stream::Streams::new(Some( + super::super::allocator::Allocator::new().unwrap(), + )); + let stream = streams + .default_stream(super::super::stream::Device::Gpu) + .unwrap(); + let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_INDEXER").unwrap()); + let ints = |shape: &[i32], data: &[i32]| { + Array::new( + shape, + Dtype::I32, + super::super::scalar_buffer( + &data + .iter() + .flat_map(|i| i.to_le_bytes()) + .collect::>(), + ) + .unwrap(), + ) + .unwrap() + }; + let axes = ints(&[32], &rope::build_mrope_axes(&[11, 11, 10], true)); + // I32 Gather reuses the bit-preserving U32 copy kernel, including signed + // extrema that a detour through F32 could not preserve. + let signed = ints(&[2, 3], &[i32::MIN, -7, i32::MAX, 31, 0, -1]); + for (axis, indices, expected) in [ + (1, vec![2, 0], vec![i32::MAX, i32::MIN, -1, 31]), + (0, vec![1, 0], vec![31, 0, -1, i32::MIN, -7, i32::MAX]), + ] { + let copied = indexing::take(&signed, &ints(&[2], &indices), axis, stream).unwrap(); + assert_eq!( + ops::read_i64(&copied, &streams, stream).unwrap(), + expected.into_iter().map(i64::from).collect::>() + ); + } + let (freq, scaling) = rope::inv_freq_and_scaling(64, 10000000., None, stream).unwrap(); + for (name, dtype) in [("f32", Dtype::F32), ("bf16", Dtype::BF16)] { + let pattern = |shape: &[i32], seed: usize| { + let values = (0..shape.iter().map(|&n| n as usize).product()) + .map(|i| (((i * 17 + seed) % 97) as f32 - 48.) / 32.) + .collect::>(); + ops::astype( + &super::super::vision::f32_array(shape, &values).unwrap(), + dtype, + false, + stream, + ) + .unwrap() + }; + let norm = ops::astype( + &super::super::vision::f32_array(&[128], &[1.; 128]).unwrap(), + dtype, + false, + stream, + ) + .unwrap(); + let query = pattern(&[1, 16, 4, 128], 3); + let raw = pattern(&[1, 2072, 128], 11); + let mut indexer = Indexer { + projection: None, + heads: 4, + kv_heads: 1, + head_dim: 128, + block_topk: 512, + ratio: 4, + q_norm: &norm, + k_norm: &norm, + inv_freq: &freq, + mrope_axes: Some(&axes), + eps: 1e-6, + scaling, + scratch: 32 * 1024 * 1024, + compiled: None, + }; + for mode in ["table", "equal", "delta", "empty"] { + let delta = if mode == "equal" { 0 } else { -7 }; + let table = match mode { + "table" => Some(ints( + &[3, 2060], + &(0..3) + .flat_map(|axis| (0..2060).map(move |p| p / (axis + 1))) + .collect::>(), + )), + "equal" => Some(ints( + &[3, 2060], + &(0..3).flat_map(|_| 0..2060).collect::>(), + )), + "empty" => Some(ints(&[3, 0], &[])), + _ => None, + }; + let vision = Some((table.as_ref(), delta)); + let q = indexer + .prepare_queries(&query, 2056, true, vision, stream) + .unwrap(); + let k = indexer + .pool_keys(&raw, 0, 518, true, vision, stream) + .unwrap(); + for (part, array) in [("query", q), ("pooled", k)] { + let got = ops::read_f32_as_f64( + &ops::astype(&array, Dtype::F32, false, stream).unwrap(), + &streams, + stream, + ) + .unwrap(); + let expected = + std::fs::read(directory.join(format!("{name}-{mode}-{part}.f32"))).unwrap(); + let actual = got + .iter() + .flat_map(|v| (*v as f32).to_le_bytes()) + .collect::>(); + assert_eq!( + crate::engine::metal::qwen_mtplx_tests::hex_digest(&actual), + crate::engine::metal::qwen_mtplx_tests::hex_digest(&expected), + "{name}/{mode}/{part}" + ); + println!( + "{}", + serde_json::json!({"case":"indexer","dtype":name,"mode":mode,"part":part,"exact_oracle_match":true}) + ); + } + let mut cache = QsaCache::new(4); + cache.kv.offset = 2056; + cache.raw_keys = Some(super::super::cache::object( + ops::slice(&raw, &[0, 0, 0], &[1, 2056, 128], &[1; 3], stream).unwrap(), + )); + let qk = ops::concatenate( + &[ + ops::reshape(&query, &[1, 16, 512], stream).unwrap(), + ops::slice(&raw, &[0, 2056, 0], &[1, 2072, 128], &[1; 3], stream).unwrap(), + ], + -1, + stream, + ) + .unwrap(); + let options = super::super::QsaIndexerOptions { + fused: true, + compiled: true, + prefill: false, + phase_prefill: false, + prefill_min_rows: 128, + prefill_compile_rows: 2048, + prefill_min_context: 16384, + prefill_workspace: 0, + tile_rows: 64, + flash: false, + gather_decode: false, + gather: false, + gather_max_rows: 4, + gather_min_context: 0, + tf32: false, + }; + let (selection, lane) = indexer + .call( + &qk, + 2056, + &mut cache, + Some(&qk), + &options, + &streams, + vision, + stream, + ) + .unwrap(); + assert_ne!(lane, "compiled"); + assert_eq!(cache.pooled_len, 518); + let super::super::attention::EagerSelection::Dense(mask) = selection.unwrap() else { + panic!("expected sparse block selection mask") + }; + ops::evaluate(&streams, &[mask], stream, false).unwrap(); + } + } + streams.clear_streams().unwrap(); +} diff --git a/src/engine/metal/qwen_mtplx/issue107_product.rs b/src/engine/metal/qwen_mtplx/issue107_product.rs new file mode 100644 index 0000000..ef5612f --- /dev/null +++ b/src/engine/metal/qwen_mtplx/issue107_product.rs @@ -0,0 +1,328 @@ +//! Long multi-image request and cache compatibility through the product owner. +use super::*; +use crate::settings::{ + GenerationPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, + effective_settings, +}; +use base64::Engine as _; + +fn state(owner: &QwenExecutor) -> String { + 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, + |name, bytes| { + hashes.push(( + name, + crate::engine::metal::qwen_mtplx_tests::hex_digest(&bytes), + )); + Ok(()) + }, + || false, + ) + .unwrap(); + crate::engine::metal::qwen_mtplx_tests::hex_digest(&serde_json::to_vec(&serde_json::json!({"ids":saved.token_ids,"images":owner.live.current_images,"payload":payload,"hashes":hashes})).unwrap()) +} + +fn legacy_live(owner: &QwenExecutor, path: &Path) -> LiveSession { + // Import the actual baseline tensors solely to exercise an old live/resident + // object. The normal loader must reject this file, tested separately below. + let bytes = std::fs::read(path).unwrap(); + let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize; + let manifest: serde_json::Value = serde_json::from_slice(&bytes[offset..]).unwrap(); + let saved = super::super::session_codec::decode_payload( + &manifest["payload"], + serde_json::from_value(manifest["tokens"].clone()).unwrap(), + manifest["lazy_kv"].as_bool().unwrap(), + true, + &owner.execution.streams, + owner.execution.stream, + |name| { + let span = &manifest["blobs"][name]; + let start = span[0].as_u64().unwrap() as usize; + Ok(bytes[start..start + span[1].as_u64().unwrap() as usize].to_vec()) + }, + ) + .unwrap(); + LiveSession { + current: Some(saved), + current_images: serde_json::from_value(manifest["images"].clone()).unwrap(), + ..Default::default() + } +} + +#[test] +#[ignore = "installed model plus vision tower; serial external memory supervisor required"] +fn issue107_vision_product() { + let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap()); + let image = std::fs::read(std::env::var_os("DS4_QWEN_VISION_IMAGE").unwrap()).unwrap(); + let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_DIR").unwrap()); + std::fs::create_dir_all(&directory).unwrap(); + let settings = effective_settings( + crate::model::ModelChoice::Qwen38FlashNext, + &GenerationPreferences { + context_tokens: 8192, + max_generated_tokens: 1, + system_prompt: String::new(), + temperature: Some(0.), + top_p: Some(0.95), + min_p: Some(0.), + seed: Some(1), + reasoning_mode: ReasoningMode::Low, + }, + &RuntimePreferences { + speculative: SpeculativePreferences { + glm_mtp: true, + ..Default::default() + }, + ..Default::default() + }, + root.parent().unwrap(), + ) + .unwrap(); + let marker = crate::engine::vision_data_marker(&format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(&image) + )); + let messages = [ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: true, + content: format!("Describe these images{marker}{marker}"), + }]; + let mut owner = QwenExecutor::open(&settings.engine).unwrap(); + owner.policy.store_min_suffix = Some(0); + let metrics = Metrics::new(&directory.join("metrics")); + let old = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_BAD_CHECKPOINT").unwrap()); + assert!( + owner + .load_checkpoint(&old, &mut |_| {}) + .unwrap_err() + .contains("vision cache policy") + ); + let mut resident = Some(legacy_live(&owner, &old)); + owner.swap_resident_state(&mut resident).unwrap(); + assert!(owner.live.current.is_none()); + drop(resident); + owner.live = legacy_live(&owner, &old); + assert!( + owner + .save_checkpoint(&directory.join("stale-save.bin"), [0; 32], &mut |_| {}) + .unwrap_err() + .contains("vision cache policy") + ); + let (out, complete) = owner + .generate( + &messages, + &settings.turn, + &AtomicBool::new(false), + &metrics, + |p, _, _| { + eprintln!("#107 prefill {p}"); + }, + |_, _| {}, + ) + .unwrap(); + assert!(complete); + assert!(out.prompt_tokens > 2048); + assert_eq!(out.cached_tokens, 0); + let saved = owner.live.current.as_ref().unwrap(); + let logits = owner.execution.values_f32(&saved.logits).unwrap(); + let oracle = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_ORACLE_DIR").unwrap()); + assert_eq!( + logits + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + std::fs::read(oracle.join("final.f32")).unwrap() + ); + let oracle_ids: Vec = + serde_json::from_slice(&std::fs::read(oracle.join("input.json")).unwrap()).unwrap(); + assert_eq!(&saved.token_ids[..out.prompt_tokens as usize], oracle_ids); + std::fs::write( + directory.join("long-final.f32"), + logits + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + ) + .unwrap(); + std::fs::write(directory.join("long-input.json"),serde_json::to_vec(&serde_json::json!({"ids":saved.token_ids,"images":owner.live.current_images,"prompt_tokens":out.prompt_tokens})).unwrap()).unwrap(); + owner + .save_checkpoint(&directory.join("long.bin"), [107; 32], &mut |_| {}) + .unwrap(); + println!( + "{}", + serde_json::json!({"case":"vision_long","prompt_tokens":out.prompt_tokens,"cached_tokens":out.cached_tokens, + "tokens":owner.live.current.as_ref().unwrap().token_ids,"text":out.message.content,"reasoning":out.message.reasoning, + "logits_sha256":crate::engine::metal::qwen_mtplx_tests::hex_digest(&logits.iter().flat_map(|v|v.to_le_bytes()).collect::>())}) + ); + let original = state(&owner); + assert!( + owner + .load_checkpoint(&old, &mut |_| {}) + .unwrap_err() + .contains("vision cache policy") + ); + assert_eq!( + state(&owner), + original, + "rejected disk cache preserves healthy current state" + ); + let run = |owner: &mut QwenExecutor, messages: &[ChatTurn], label: &str| { + let (out, complete) = owner + .generate( + messages, + &settings.turn, + &AtomicBool::new(false), + &metrics, + |_, _, _| {}, + |_, _| {}, + ) + .unwrap(); + assert!(complete); + assert!( + !owner + .live + .current_images + .iter() + .any(|i| i.start < out.cached_tokens as usize + && (out.cached_tokens as usize) < i.end) + ); + println!( + "{}", + serde_json::json!({"case":label,"cached_tokens":out.cached_tokens,"prompt_tokens":out.prompt_tokens,"state_sha256":state(owner)}) + ); + out + }; + let repeated = run(&mut owner, &messages, "unchanged"); + assert_eq!(repeated.cached_tokens, out.prompt_tokens); + assert_eq!(state(&owner), original); + let mut followup = messages.to_vec(); + followup.push(out.message.clone()); + followup.push(ChatTurn { + content: "Continue.".into(), + ..messages[0].clone() + }); + let warm = run(&mut owner, &followup, "resident_followup"); + assert!(warm.cached_tokens >= out.prompt_tokens); + let warm_state = state(&owner); + owner.reset().unwrap(); + assert!( + owner + .load_checkpoint(&directory.join("long.bin"), &mut |_| {}) + .unwrap() + ); + assert_eq!(state(&owner), original); + let reopened = run(&mut owner, &followup, "reopened_followup"); + assert_eq!(reopened.cached_tokens, warm.cached_tokens); + assert_eq!(state(&owner), warm_state); + let mut other = std::io::Cursor::new(Vec::new()); + image::RgbImage::from_pixel(1024, 1024, image::Rgb([0, 0, 255])) + .write_to(&mut other, image::ImageFormat::Png) + .unwrap(); + let changed_marker = crate::engine::vision_data_marker(&format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(other.get_ref()) + )); + for (label, content) in [ + ( + "text_before", + format!("Please describe these images{marker}{marker}"), + ), + ( + "text_between", + format!("Describe these images{marker}Compare them.{marker}"), + ), + ( + "text_after", + format!("Describe these images{marker}{marker}Be brief."), + ), + ( + "changed_pixels", + format!("Describe these images{changed_marker}{changed_marker}"), + ), + ] { + owner.reset().unwrap(); + owner + .load_checkpoint(&directory.join("long.bin"), &mut |_| {}) + .unwrap(); + let edited = [ChatTurn { + content, + ..messages[0].clone() + }]; + let changed = run(&mut owner, &edited, label); + assert_eq!(changed.cached_tokens, 0); + if label == "changed_pixels" { + assert_eq!( + &owner.live.current.as_ref().unwrap().token_ids[..changed.prompt_tokens as usize], + oracle_ids + ); + assert_ne!(state(&owner), original); + } + let expected = state(&owner); + owner.reset().unwrap(); + run(&mut owner, &edited, &format!("{label}_cold_control")); + assert_eq!(state(&owner), expected); + } + owner.reset().unwrap(); + let cancelled = AtomicBool::new(false); + let interrupted = owner.generate( + &messages, + &settings.turn, + &cancelled, + &metrics, + |p, _, _| { + if p >= 2048 { + cancelled.store(true, std::sync::atomic::Ordering::Relaxed); + } + }, + |_, _| {}, + ); + assert!(cancelled.load(std::sync::atomic::Ordering::Relaxed)); + assert!(interrupted.is_err() || interrupted.is_ok_and(|(_, complete)| !complete)); + assert!(owner.live.current.is_none()); + let resent = run(&mut owner, &messages, "cancel_resend"); + assert_eq!(resent.cached_tokens, 0); + assert_eq!(state(&owner), original); + // An artificially supplied prefix ending within the second image must be + // refused before any of its tensors are restored. + owner + .live + .current + .as_mut() + .unwrap() + .token_ids + .truncate(2048); + owner.live.prefix = None; + let guarded = run(&mut owner, &messages, "inside_image_guard"); + assert_eq!(guarded.cached_tokens, 0); + assert_eq!(state(&owner), original); + let text = [ChatTurn { + content: "Say hello.".into(), + ..messages[0].clone() + }]; + run(&mut owner, &text, "text_control"); + let text_state = state(&owner); + assert!(owner.live.current_images.is_empty()); + owner + .save_checkpoint(&directory.join("text.bin"), [0; 32], &mut |_| {}) + .unwrap(); + owner.reset().unwrap(); + assert!( + owner + .load_checkpoint(&directory.join("text.bin"), &mut |_| {}) + .unwrap() + ); + assert_eq!(state(&owner), text_state); + println!( + "{}", + serde_json::json!({"case":"compatibility","old_disk_rejected":true,"old_live_rejected":true,"old_resident_rejected":true,"stale_save_rejected":true,"new_vision_and_text_reopen_exactly":true}) + ); +} diff --git a/src/engine/metal/qwen_mtplx/owner.rs b/src/engine/metal/qwen_mtplx/owner.rs index a5ce447..e2d3587 100644 --- a/src/engine/metal/qwen_mtplx/owner.rs +++ b/src/engine/metal/qwen_mtplx/owner.rs @@ -253,7 +253,8 @@ impl QwenExecutor { incoming: &mut Option, ) -> Result<(), String> { self.clear_execution_cache()?; - let next = incoming.take().unwrap_or_default(); + let mut next = incoming.take().unwrap_or_default(); + next.discard_incompatible_vision(); *incoming = Some(std::mem::replace(&mut self.live, next)); Ok(()) } @@ -1068,3 +1069,6 @@ mod issue105_product; #[cfg(test)] #[path = "issue106_product.rs"] mod issue106_product; +#[cfg(test)] +#[path = "issue107_product.rs"] +mod issue107_product; diff --git a/src/engine/metal/qwen_mtplx/qsa_attention.rs b/src/engine/metal/qwen_mtplx/qsa_attention.rs index 60bf7b7..3aaa252 100644 --- a/src/engine/metal/qwen_mtplx/qsa_attention.rs +++ b/src/engine/metal/qwen_mtplx/qsa_attention.rs @@ -106,6 +106,7 @@ impl Attention<'_> { idx.as_ref(), &options.indexer, streams, + vision, stream, )? .0 @@ -119,7 +120,16 @@ impl Attention<'_> { let selected = match &mut self.indexer { Some(indexer) => { indexer - .call(x, pos, cache, None, &options.indexer, streams, stream)? + .call( + x, + pos, + cache, + None, + &options.indexer, + streams, + vision, + stream, + )? .0 } None => None, @@ -192,8 +202,8 @@ impl Attention<'_> { let v = ops::transpose(&v, &[0, 2, 1, 3], stream)?; let (k, v) = cache.kv.update_and_fetch(&k, &v, stream)?; let total = k.borrow().layout().shape()[2]; - // Vision bypasses selection, not the indexer/cache work above. - if vision.is_some() { + // Match the oracle's explicit diagnostic rollback; normal vision keeps QSA. + if vision.is_some() && crate::engine::qwen_vision::cache_policy() == 0 { selected = None; } if let Some(EagerSelection::FlashBlocks(ids, tail)) = selected { @@ -492,6 +502,7 @@ fn mtplx_qsa_attention_module_matches_reference() { q_norm: &iqnorm, k_norm: &iknorm, inv_freq: &freq, + mrope_axes: (vision == "table").then_some(&axes), eps: 1e-6, scaling: 1., scratch: 32 * 1024 * 1024, @@ -512,7 +523,7 @@ fn mtplx_qsa_attention_module_matches_reference() { attn.indexer .as_ref() .unwrap() - .extend_pooled(&mut cache, prefix, options.indexer.fused, gpu) + .extend_pooled(&mut cache, prefix, options.indexer.fused, None, gpu) .unwrap(); let kv = ops::transpose(&source(&[1, prefix, 2, 256], 89), &[0, 2, 1, 3], gpu).unwrap(); diff --git a/src/engine/metal/qwen_mtplx/qsa_indexer.rs b/src/engine/metal/qwen_mtplx/qsa_indexer.rs index a64f881..094b499 100644 --- a/src/engine/metal/qwen_mtplx/qsa_indexer.rs +++ b/src/engine/metal/qwen_mtplx/qsa_indexer.rs @@ -18,6 +18,7 @@ pub(super) struct Indexer<'a> { pub(super) q_norm: &'a Array, pub(super) k_norm: &'a Array, pub(super) inv_freq: &'a Array, + pub(super) mrope_axes: Option<&'a Array>, pub(super) eps: f64, pub(super) scaling: f64, pub(super) scratch: u64, @@ -31,6 +32,10 @@ pub(super) struct PreparedRows { pub(super) total: i32, } +#[cfg(test)] +#[path = "issue107_indexer.rs"] +mod issue107; + impl<'a> Indexer<'a> { #[expect( clippy::too_many_arguments, @@ -195,6 +200,7 @@ impl<'a> Indexer<'a> { qk_rows: Option<&Array>, options: &super::QsaIndexerOptions, streams: &super::stream::Streams, + vision: Option<(Option<&Array>, i32)>, stream: Stream, ) -> Result<(Option, &'static str), String> { use super::{ @@ -221,16 +227,21 @@ impl<'a> Indexer<'a> { self.ratio as u32, self.block_topk as u32, ); - if self.compiled_supported( - qk_rows.unwrap_or(hidden), - cache, - pos, - qk_rows.is_some(), - decode, - mode, - options, - stream, - ) { + let vision = vision.filter(|_| { + self.mrope_axes.is_some() && crate::engine::qwen_vision::cache_policy() != 0 + }); + if vision.is_none() + && self.compiled_supported( + qk_rows.unwrap_or(hidden), + cache, + pos, + qk_rows.is_some(), + decode, + mode, + options, + stream, + ) + { return self .call_compiled( hidden, @@ -244,7 +255,7 @@ impl<'a> Indexer<'a> { .map(|a| (a, "compiled")); } let Some(prepared) = - self.prepare_rows(hidden, pos, cache, qk_rows, options.fused, stream)? + self.prepare_rows(hidden, pos, cache, qk_rows, options.fused, vision, stream)? else { return Ok((None, "update_only")); }; @@ -573,9 +584,13 @@ impl<'a> Indexer<'a> { q: &Array, pos: i32, fused: bool, + vision: Option<(Option<&Array>, i32)>, stream: Stream, ) -> Result { - if fused && qsa_prepare::supported(q, self.q_norm, self.inv_freq, 4, stream) { + if vision.is_none() + && fused + && qsa_prepare::supported(q, self.q_norm, self.inv_freq, 4, stream) + { return qsa_prepare::queries( q, self.q_norm, @@ -588,7 +603,18 @@ impl<'a> Indexer<'a> { } let q = ops::rms_norm(q, Some(self.q_norm), self.eps as f32, stream)?; let positions = indexing::arange_range_i32(pos, pos + q.layout().shape()[1], stream)?; - let (cos, sin) = rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?; + let (cos, sin) = match (vision, self.mrope_axes) { + (Some((table, delta)), Some(axes)) => rope::vision_position_cos_sin( + &positions, + table, + delta, + self.inv_freq, + axes, + self.scaling, + stream, + )?, + _ => rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?, + }; rope::apply_partial(&q, &cos, &sin, stream) } @@ -598,9 +624,13 @@ impl<'a> Indexer<'a> { old: i32, total: i32, fused: bool, + vision: Option<(Option<&Array>, i32)>, stream: Stream, ) -> Result { - if fused && qsa_prepare::supported(fresh, self.k_norm, self.inv_freq, 3, stream) { + if vision.is_none() + && fused + && qsa_prepare::supported(fresh, self.k_norm, self.inv_freq, 3, stream) + { return qsa_prepare::pool( fresh, self.k_norm, @@ -632,7 +662,18 @@ impl<'a> Indexer<'a> { Binary::Multiply, stream, )?; - let (cos, sin) = rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?; + let (cos, sin) = match (vision, self.mrope_axes) { + (Some((table, delta)), Some(axes)) => rope::vision_position_cos_sin( + &positions, + table, + delta, + self.inv_freq, + axes, + self.scaling, + stream, + )?, + _ => rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?, + }; ops::squeeze( &rope::apply_partial( &ops::expand_dims(&pooled, &[2], stream)?, @@ -650,6 +691,7 @@ impl<'a> Indexer<'a> { cache: &mut QsaCache, total: i32, fused: bool, + vision: Option<(Option<&Array>, i32)>, stream: Stream, ) -> Result, String> { let nb = total / self.ratio; @@ -669,7 +711,7 @@ impl<'a> Indexer<'a> { )?; drop(raw); cache.write_pooled( - &self.pool_keys(&fresh, old, nb, fused, stream)?, + &self.pool_keys(&fresh, old, nb, fused, vision, stream)?, old, nb, stream, @@ -689,6 +731,10 @@ impl<'a> Indexer<'a> { /// Common noncompiled _call_rows prefix. None means the complete-block /// budget fits: cache updates still happen, but lazy query work is unused. /// Cache.kv.offset is advanced by Attention, never by this indexer. + #[expect( + clippy::too_many_arguments, + reason = "original preparation inputs and request positions" + )] pub(super) fn prepare_rows( &self, hidden: &Array, @@ -696,6 +742,7 @@ impl<'a> Indexer<'a> { cache: &mut QsaCache, qk_rows: Option<&Array>, fused: bool, + vision: Option<(Option<&Array>, i32)>, stream: Stream, ) -> Result, String> { let shape = hidden.layout().shape().to_vec(); @@ -733,9 +780,9 @@ impl<'a> Indexer<'a> { let parts = ops::split(&qk, &[width], -1, stream)?; let query = ops::reshape(&parts[0], &[1, shape[1], self.heads, self.head_dim], stream)?; let keys = ops::reshape(&parts[1], &[1, shape[1], self.head_dim], stream)?; - let query = self.prepare_queries(&query, pos, fused, stream)?; + let query = self.prepare_queries(&query, pos, fused, vision, stream)?; cache.write_raw(&keys, stream)?; - let pooled = self.extend_pooled(cache, total, fused, stream)?; + let pooled = self.extend_pooled(cache, total, fused, vision, stream)?; if total / self.ratio <= self.block_topk { return Ok(None); } @@ -793,6 +840,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() { q_norm: &q_norm, k_norm: &k_norm, inv_freq: &inv_freq, + mrope_axes: None, eps: 1e-6, scaling: 1., scratch: 32 * 1024 * 1024, @@ -835,7 +883,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() { .write_raw(&source(&[1, prefix, 128], 88), gpu) .unwrap(); indexer - .extend_pooled(&mut cache, prefix, options.fused, gpu) + .extend_pooled(&mut cache, prefix, options.fused, None, gpu) .unwrap(); let kv = ops::transpose(&source(&[1, prefix, 2, 256], 89), &[0, 2, 1, 3], gpu) .unwrap(); @@ -875,6 +923,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() { qk.as_ref(), &options, &streams, + None, gpu, ) .unwrap() @@ -1111,6 +1160,7 @@ fn mtplx_qsa_indexer_compiled_commit_matches_reference() { q_norm: &q_norm[0], k_norm: &k_norm[0], inv_freq: &freq[0], + mrope_axes: None, eps: 1e-6, scaling: 1., scratch: 32 * 1024 * 1024, @@ -1399,6 +1449,7 @@ fn mtplx_qsa_indexer_prefix_matches_reference() { q_norm: &q_norm, k_norm: &k_norm, inv_freq: &inv_freq, + mrope_axes: None, eps: 1e-6, scaling: 1.125, scratch: 32 * 1024 * 1024, @@ -1454,7 +1505,15 @@ fn mtplx_qsa_indexer_prefix_matches_reference() { } let (prefix, calls) = capture_dispatches(|| { indexer - .prepare_rows(&hidden, pos, &mut cache, qk.as_ref(), fused, gpu) + .prepare_rows( + &hidden, + pos, + &mut cache, + qk.as_ref(), + fused, + None, + gpu, + ) .unwrap() }); assert!(calls.is_empty(), "QSA prefix must remain lazy"); diff --git a/src/engine/metal/qwen_mtplx/qsa_select.rs b/src/engine/metal/qwen_mtplx/qsa_select.rs index 0275b5b..8ac0de9 100644 --- a/src/engine/metal/qwen_mtplx/qsa_select.rs +++ b/src/engine/metal/qwen_mtplx/qsa_select.rs @@ -391,6 +391,7 @@ fn mtplx_qsa_fused_graph_matches_reference() { q_norm: &q, k_norm: &q, inv_freq: &q, + mrope_axes: None, eps: 1e-6, scaling: 1., scratch: 32 * 1024 * 1024, diff --git a/src/engine/metal/qwen_mtplx/request.rs b/src/engine/metal/qwen_mtplx/request.rs index 5d2e6b8..c10bcf2 100644 --- a/src/engine/metal/qwen_mtplx/request.rs +++ b/src/engine/metal/qwen_mtplx/request.rs @@ -100,6 +100,14 @@ pub(in crate::engine) struct LiveSession { } impl LiveSession { + pub(super) fn discard_incompatible_vision(&mut self) { + if !crate::engine::qwen_vision::compatible(&self.current_images) + || !crate::engine::qwen_vision::compatible(&self.prefix_images) + { + *self = Self::default(); + } + } + pub(super) fn prompt( &self, model: &QwenMetadata, @@ -131,6 +139,7 @@ impl LiveSession { emit: impl FnMut(bool, String), observe: impl FnMut(TurnProgress) -> Result<(), String>, ) -> Result { + self.discard_incompatible_vision(); let (prompt, images) = execution .vision_prompt .clone() diff --git a/src/engine/metal/qwen_mtplx/rope.rs b/src/engine/metal/qwen_mtplx/rope.rs index 05729ab..a56268a 100644 --- a/src/engine/metal/qwen_mtplx/rope.rs +++ b/src/engine/metal/qwen_mtplx/rope.rs @@ -246,6 +246,53 @@ pub(super) fn mrope_cos_sin( )) } +/// MTPLX _vision_position_cos_sin: choose table versus text+delta per row. +pub(super) fn vision_position_cos_sin( + positions: &Array, + table: Option<&Array>, + delta: i32, + inv_freq: &Array, + axes: &Array, + scaling: f64, + stream: Stream, +) -> Result<(Array, Array), String> { + let integer = |n: i32| Array::new(&[], Dtype::I32, super::scalar_buffer(&n.to_le_bytes())?); + let shifted = binary(positions, &integer(delta)?, Binary::Add, stream)?; + let mut positions3 = ops::broadcast_to( + &ops::expand_dims(&shifted, &[0], stream)?, + &[3, positions.layout().size() as i32], + stream, + )?; + if let Some(table) = table { + let shape = table.layout().shape().to_vec(); + if shape.len() != 2 || shape[0] != 3 { + return Err("vision positions require a three-axis table".into()); + } + if shape[1] > 0 { + let clipped = binary( + &binary(positions, &integer(0)?, Binary::Maximum, stream)?, + &integer(shape[1] - 1)?, + Binary::Minimum, + stream, + )?; + let prompt = indexing::take(table, &clipped, 1, stream)?; + let covered = binary(positions, &integer(shape[1])?, Binary::Less, stream)?; + positions3 = indexing::select( + &ops::expand_dims(&covered, &[0], stream)?, + &prompt, + &positions3, + stream, + )?; + } + } + let (mut cos, mut sin) = mrope_cos_sin(&positions3, inv_freq, axes, stream)?; + if scaling != 1.0 { + cos = binary(&cos, &scalar(scaling)?, Binary::Multiply, stream)?; + sin = binary(&sin, &scalar(scaling)?, Binary::Multiply, stream)?; + } + Ok((cos, sin)) +} + pub(super) fn apply_partial( x: &Array, cos: &Array, diff --git a/src/engine/metal/qwen_mtplx/session_checkpoint.rs b/src/engine/metal/qwen_mtplx/session_checkpoint.rs index 7ff6ca1..98e1887 100644 --- a/src/engine/metal/qwen_mtplx/session_checkpoint.rs +++ b/src/engine/metal/qwen_mtplx/session_checkpoint.rs @@ -78,6 +78,9 @@ impl LiveSession { .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()); } @@ -179,6 +182,9 @@ impl LiveSession { ); } 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 !manifest.payload["mtp_history_snapshot"].is_null() != identity.mtp { return Err("Qwen checkpoint MTP payload differs".into()); } diff --git a/src/engine/metal/qwen_mtplx/vision.rs b/src/engine/metal/qwen_mtplx/vision.rs index 6261739..4c0bfa8 100644 --- a/src/engine/metal/qwen_mtplx/vision.rs +++ b/src/engine/metal/qwen_mtplx/vision.rs @@ -510,6 +510,7 @@ pub(super) fn prepare( start, end: ids.len(), digest: *digest, + policy: crate::engine::qwen_vision::cache_policy(), }); index += 1; } else { diff --git a/src/engine/qwen_vision.rs b/src/engine/qwen_vision.rs index 8123480..1da810e 100644 --- a/src/engine/qwen_vision.rs +++ b/src/engine/qwen_vision.rs @@ -207,12 +207,27 @@ pub(super) struct Identity { pub(super) start: usize, pub(super) end: usize, pub(super) digest: [u8; 32], + /// Missing on checkpoints written before image-aware sparse QSA. + #[serde(default)] + pub(super) policy: u32, +} + +pub(super) fn cache_policy() -> u32 { + u32::from(std::env::var("MTPLX_QWEN4_VISION_QSA").as_deref() != Ok("0")) +} + +pub(super) fn compatible(images: &[Identity]) -> bool { + images.iter().all(|image| image.policy == cache_policy()) } pub(super) fn same_prefix(a: &[Identity], b: &[Identity], tokens: usize) -> bool { - a.iter() - .filter(|i| i.start < tokens) - .eq(b.iter().filter(|i| i.start < tokens)) + // Image pads do not identify partial image embeddings. MTPLX only restores + // before an image or after its complete span, never into the middle. + !b.iter() + .any(|image| image.start < tokens && tokens < image.end) + && a.iter() + .filter(|i| i.start < tokens) + .eq(b.iter().filter(|i| i.start < tokens)) } #[test] @@ -271,14 +286,19 @@ fn qwen_image_cache_keys_distinguish_pixels_and_preserve_prior_text() { start: 10, end: 14, digest: [1; 32], + policy: cache_policy(), }]; let changed = vec![Identity { start: 10, end: 14, digest: [2; 32], + policy: cache_policy(), }]; assert!(same_prefix(&old, &old, 20)); assert!(!same_prefix(&old, &changed, 20)); assert!(same_prefix(&old, &changed, 10)); assert!(!same_prefix(&old, &[], 20)); + assert!(!same_prefix(&old, &old, 11)); + assert!(!same_prefix(&old, &old, 13)); + assert!(same_prefix(&old, &old, 14)); } diff --git a/tests/fixtures/mtplx-qsa-attention-module.jsonl b/tests/fixtures/mtplx-qsa-attention-module.jsonl index 2393f1f..bb67d8c 100644 --- a/tests/fixtures/mtplx-qsa-attention-module.jsonl +++ b/tests/fixtures/mtplx-qsa-attention-module.jsonl @@ -17,15 +17,15 @@ {"kernel":"qsa_attention_p6_s0","sha256":["068422304f4191e74610562e007ffa7ec0d4b2ff5ed619997442ca1d254073b1","6f49058a22d7c62bbf4ffba4542b7aa32a130076b07537ac5605b2c85bb97162","158b3b53061c6d3d089049f9fdb47064a8c4f3d2f3fdb0ef5089afcaeba7cc5d","d4bc3f9184f206d1c86876a7ce98f086cf94f7912b460f7dc8c809b5b78d579b","011cc9bf7a9b2f86548e46728a9013e24273ca27ee50f9ee3f5509ecbc7c851f"],"profile":6,"step":0,"rows":2048,"pos":32768,"prefix":32768,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":32768,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["prefill_flash"],"counts":{"flash_kernel":1,"compiled_selector":1},"offset":34816,"pooled_len":8704,"state_present":[true,true,false,true,true],"shapes":[[1,2048,64],[1,65536,128],[1,16384,128],[1,2,34816,256],[1,2,34816,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":0,"qk_rows_calls":1,"traces":1,"selector_dispatches":1,"capacity_transitions":0,"prefill_score_producers":{"mpp":1,"mlx":0},"modes":{"blocks":0,"prefill_blocks":1,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw65536:pool16384:out0":1},"entry_count":1,"compiled_keys":[{"source":"qk_rows","mode":"prefill_blocks","input_shape":[1,2048,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"mpp","selector_chunk_rows":2048,"label":"qk_rows:prefill_blocks:s2048:raw65536:pool16384:out0:scorempp:chunk2048:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p6_s1","sha256":["bcbfcd4ff8d59b123b6ba9e84c2d621b9982c5e7430c84b82240f2b4416e48af","d50916defa6c9151ffeb706a7e0d6b8c04b861eaa94ed2b637f82fce79c8f2c9","158b3b53061c6d3d089049f9fdb47064a8c4f3d2f3fdb0ef5089afcaeba7cc5d","7c9ca3c2fd81b2f4749ce1f575a963b7054f004b91879a64212ecda5b0c34463","773fdcd29060aa9a94684db395aeccdbe1af803e191f3d0b938abc109d715fbb"],"profile":6,"step":1,"rows":1,"pos":34816,"prefix":32768,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":32768,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["flash"],"counts":{},"offset":34817,"pooled_len":8704,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,65536,128],[1,16384,128],[1,2,35072,256],[1,2,35072,256]],"core":{"calls":2,"compiled_calls":2,"hidden_calls":0,"qk_rows_calls":2,"traces":2,"selector_dispatches":2,"capacity_transitions":0,"prefill_score_producers":{"mpp":1,"mlx":0},"modes":{"blocks":1,"prefill_blocks":1,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw65536:pool16384:out0":2},"entry_count":2,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw65536:pool16384:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"prefill_blocks","input_shape":[1,2048,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"mpp","selector_chunk_rows":2048,"label":"qk_rows:prefill_blocks:s2048:raw65536:pool16384:out0:scorempp:chunk2048:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p6_s2","sha256":["e5f8b59d40e7e1c579d22d66af242e2882d1ab57727f17c7b9eeea894a5368e7","3b29542bd05bd5525466ad6c5951dd2d59759bed880564b8e07058d8feed1d42","9bc8c466e0a3164a6f4e22c0ffe23a895126f57945d4e7f0d5a07e79c63778ad","a41386c0de5defa3cc2eb26fe6005bab14865f573f9380f998d35e324abdbb43","270285c24ae2ea32e45240ada590a6cc1c9e7b14adfc23b270c7d7440c7372d4"],"profile":6,"step":2,"rows":4,"pos":34817,"prefix":32768,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":32768,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["rows_gather"],"counts":{},"offset":34821,"pooled_len":8705,"state_present":[true,true,false,true,true],"shapes":[[1,4,64],[1,65536,128],[1,16384,128],[1,2,35072,256],[1,2,35072,256]],"core":{"calls":3,"compiled_calls":3,"hidden_calls":0,"qk_rows_calls":3,"traces":3,"selector_dispatches":3,"capacity_transitions":0,"prefill_score_producers":{"mpp":1,"mlx":0},"modes":{"blocks":1,"prefill_blocks":1,"dense_mask":0,"row_tokens":1,"update_only":0},"buckets":{"raw65536:pool16384:out0":3},"entry_count":3,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw65536:pool16384:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"prefill_blocks","input_shape":[1,2048,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"mpp","selector_chunk_rows":2048,"label":"qk_rows:prefill_blocks:s2048:raw65536:pool16384:out0:scorempp:chunk2048:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,4,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"qk_rows:row_tokens:s4:raw65536:pool16384:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p7_s0","sha256":["f9c3c5622d5d7c346d6c2881634894e60f030654828e74d96041317d5d1b99e5","fa52f9dff396320af135c22c25b328674b3c44b6c768df9ed2d54d31bbecc71f","5385335d390b6950b06db2bd3a3e86efe8b4ca7d48d7458828f6891d0815be26","f583e84d01f4e3bfb152b8855f8520d25c5fd1639c20961343d44d3b494ee18f","ea3bda4b6f89e3552003b8afb307a5d04a67cb8d024ea18b7a7df90a172f7342"],"profile":7,"step":0,"rows":1,"pos":2051,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":2052,"pooled_len":513,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":0,"qk_rows_calls":1,"traces":1,"selector_dispatches":1,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw4096:pool1024:out0":1},"entry_count":1,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p7_s1","sha256":["96067951d1aa0a433bcc85d4c51f0dbbf8f142f2f97b0e4aa1f0396ce6af89de","c7e437185225097e4ce2bb012fb9dea74124314e3858db5b2d3f89d175628f8e","99fb41883cf58e2cec92897b8486a4b20b6817dc32c1b98d320280e2984756b4","4af7a650ea8717de842e55d3691237c8e2af5c5c8ac78868e1323773e1a792f1","f54917b64b4e33520dd53be4677edb0d92a41e62d868d44eff6fdbc030530df8"],"profile":7,"step":1,"rows":4,"pos":2052,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,4,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":2,"compiled_calls":2,"hidden_calls":0,"qk_rows_calls":2,"traces":2,"selector_dispatches":2,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":1,"update_only":0},"buckets":{"raw4096:pool1024:out0":2},"entry_count":2,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,4,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"qk_rows:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p7_s2","sha256":["5c48c9fcdb884b10fecce90489b58457a977528d3732d9ceb90d202f9a8bcabf","0962884629f0854f29ed23c253f65f1d44fcb8bfde142e948411c9e0ed65faf9","5b6ab0a7cc9a3a122481547531253bcba5f039fc507718fc85bef86dd8dc56ea","57eb1075e3ed3fe889a40c34bd1e59974119725a3a9ef009646a3427b3be9dd6","7fe9f3c5a939bb7dd64a147ea63354dc92c40828f2cc4d1cf8b3be331b86f4e5"],"profile":7,"step":2,"rows":2,"pos":2054,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,2,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":3,"compiled_calls":3,"hidden_calls":0,"qk_rows_calls":3,"traces":3,"selector_dispatches":3,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":2,"update_only":0},"buckets":{"raw4096:pool1024:out0":3},"entry_count":3,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,2,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":2,"label":"qk_rows:row_tokens:s2:raw4096:pool1024:out0:scorenone:chunk2:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,4,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"qk_rows:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p7_s3","sha256":["554527077cad0087f4b548f1cf28dfbf186762ce6e3e12982922b97eec865e1d","9424991712e6561268208093dc9c9df018591b376554ac208eab7f242bdba4e4","5b6ab0a7cc9a3a122481547531253bcba5f039fc507718fc85bef86dd8dc56ea","28b0c2bfcc8486f5c6daafec65e013f6d2efbf99b24cdf75ba25372716d828d9","6d3d7b11ac7cff92650bdc798b862cbd21b383cd0154f53eee439af721d3c55f"],"profile":7,"step":3,"rows":1,"pos":2056,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":2057,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":4,"compiled_calls":4,"hidden_calls":0,"qk_rows_calls":4,"traces":3,"selector_dispatches":4,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":2,"prefill_blocks":0,"dense_mask":0,"row_tokens":2,"update_only":0},"buckets":{"raw4096:pool1024:out0":4},"entry_count":3,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,2,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":2,"label":"qk_rows:row_tokens:s2:raw4096:pool1024:out0:scorenone:chunk2:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,4,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"qk_rows:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p7_s0","sha256":["9fc17259582838c33640717e4e21cdd556e3f5d03bd8dda03de124d43f7b55e8","48e48f96c54776ebd1a034cdfcf93ae5e027038346bff0b1cc4f6600a73bdbd4","56ff514413104108eea40eeb10e28ba89d4bcb5f273e7bfb327b650239b2608e","d1f4df60a631b8170515760665a9212d6375de4e6456f75d90eb31179c91596c","f583e84d01f4e3bfb152b8855f8520d25c5fd1639c20961343d44d3b494ee18f","ea3bda4b6f89e3552003b8afb307a5d04a67cb8d024ea18b7a7df90a172f7342"],"profile":7,"step":0,"rows":1,"pos":2051,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["flash"],"counts":{},"offset":2052,"pooled_len":513,"state_present":[true,true,true,true,true],"shapes":[[1,1,64],[1,2304,128],[1,1024,128],[1,1,128,1024],[1,2,2304,256],[1,2,2304,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p7_s1","sha256":["24e8ced9cf6525d3fb9499e8cb45e70960458b26278552ff5e3bdb825596f33a","1ceca1a52f5057c1aaee04689fc8026d3df9e1b657d718d75b9a72aeb8a88b25","8ee668bf4a7f6eca91380e543147f88eb73d8bd3384c1a375b869d4680d95753","9221cbebf9df5094116ab9269b2c0e543d03fcd7f7ec5e9c1c745fc7b2babd0e","4af7a650ea8717de842e55d3691237c8e2af5c5c8ac78868e1323773e1a792f1","f54917b64b4e33520dd53be4677edb0d92a41e62d868d44eff6fdbc030530df8"],"profile":7,"step":1,"rows":4,"pos":2052,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["rows_gather"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,true,true,true],"shapes":[[1,4,64],[1,2304,128],[1,1024,128],[1,1,128,1024],[1,2,2304,256],[1,2,2304,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p7_s2","sha256":["6cd50cb32aab449013ff1db077d2e4a2f480fb71f1ae11ed80499dfab78257d0","b49eaa7452b6018be70df323222b0624f230df95ea900e42a974eaba5e2ac975","a8f7c3b9add9760148121cf236dab8d6844d38d42ab5aa4a7de20bb14c342d97","d4abe49bd48fe5eeea175dd7102faf5cc4a0bcf4831d233e7521c51cf7a76932","57eb1075e3ed3fe889a40c34bd1e59974119725a3a9ef009646a3427b3be9dd6","7fe9f3c5a939bb7dd64a147ea63354dc92c40828f2cc4d1cf8b3be331b86f4e5"],"profile":7,"step":2,"rows":2,"pos":2054,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["rows_gather"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,true,true,true],"shapes":[[1,2,64],[1,2304,128],[1,1024,128],[1,1,128,1024],[1,2,2304,256],[1,2,2304,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p7_s3","sha256":["80cb2c6d1f34cc44dfbcd30192a17c8cb4a0c2a19017924058ece49e2f5c6486","59da7bcfd15c46c7c183a322559d2f455dba99f677bbef99b3626bd32160147f","a8f7c3b9add9760148121cf236dab8d6844d38d42ab5aa4a7de20bb14c342d97","d4abe49bd48fe5eeea175dd7102faf5cc4a0bcf4831d233e7521c51cf7a76932","28b0c2bfcc8486f5c6daafec65e013f6d2efbf99b24cdf75ba25372716d828d9","6d3d7b11ac7cff92650bdc798b862cbd21b383cd0154f53eee439af721d3c55f"],"profile":7,"step":3,"rows":1,"pos":2056,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["flash"],"counts":{},"offset":2057,"pooled_len":514,"state_present":[true,true,true,true,true],"shapes":[[1,1,64],[1,2304,128],[1,1024,128],[1,1,128,1024],[1,2,2304,256],[1,2,2304,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p8_s0","sha256":["2f2d1534efe9754560b2d2289f6c9756da611928c8c9e2321786ab03dd447111","84e64cfc1690832a3a6ed516b577d1001b89b382254c08b2663f493691330f2e","7da3f843c044b333d43259efd0a68f1c0a8f5230707898663f9ce096dd1f294a"],"profile":8,"step":0,"rows":3,"pos":0,"prefix":0,"fused":0,"vision":"none","no_indexer":true,"biased":true,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":3,"pooled_len":0,"state_present":[false,false,false,true,true],"shapes":[[1,3,64],[1,2,256,256],[1,2,256,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p8_s1","sha256":["dbe815a17715596ea030903a6800bba50c5ed74b9e03a3d29826e42acd57377b","50d2e91a257e44f7554a3ba6b1e0b0c0bef2a3f9ec4f1c4b6904e150f9b7b48c","a79deff710fa918973e19986ced87ae13347fbd66972c79fd4241d299f4472dc"],"profile":8,"step":1,"rows":1,"pos":3,"prefix":0,"fused":0,"vision":"none","no_indexer":true,"biased":true,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":4,"pooled_len":0,"state_present":[false,false,false,true,true],"shapes":[[1,1,64],[1,2,256,256],[1,2,256,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p8_s2","sha256":["faeacdecf2a8829ff7c7c5751858003e0931435cae65c610e5c362a3997ba050","d4c84ff24f73effb7d50f227eb1248f997cc7ea466b6a15e6fa2876cf9481ea0","af7554d1b0575d27146ecc0b2d3755a89a645e42231c214b62cbecfaf6d81402"],"profile":8,"step":2,"rows":4,"pos":4,"prefix":0,"fused":0,"vision":"none","no_indexer":true,"biased":true,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":8,"pooled_len":0,"state_present":[false,false,false,true,true],"shapes":[[1,4,64],[1,2,256,256],[1,2,256,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p9_s0","sha256":["91fd57e3f97be79622b57e34ed7e80b9934f65ce650592a6a619f0c88b322bd9","fa52f9dff396320af135c22c25b328674b3c44b6c768df9ed2d54d31bbecc71f","5385335d390b6950b06db2bd3a3e86efe8b4ca7d48d7458828f6891d0815be26","0823d3bc139c67f69911c8970fb2c1bcf5f757971e9191f8b98d56f6689840ae","ea3bda4b6f89e3552003b8afb307a5d04a67cb8d024ea18b7a7df90a172f7342"],"profile":9,"step":0,"rows":1,"pos":2051,"prefix":2051,"fused":3,"vision":"no_axes","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":2052,"pooled_len":513,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":1,"qk_rows_calls":0,"traces":1,"selector_dispatches":1,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw4096:pool1024:out0":1},"entry_count":1,"compiled_keys":[{"source":"hidden","mode":"blocks","input_shape":[1,1,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"hidden:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p9_s1","sha256":["893d7f0e21dc6a1c5099fa172b86630a0b694dc90ea7c6db4ae928c2151e9b00","c7e437185225097e4ce2bb012fb9dea74124314e3858db5b2d3f89d175628f8e","99fb41883cf58e2cec92897b8486a4b20b6817dc32c1b98d320280e2984756b4","e38b573ffab0799d27fcadc08d1e6a6aa31e40b4e6a61c5e9c98073a7012ba24","f54917b64b4e33520dd53be4677edb0d92a41e62d868d44eff6fdbc030530df8"],"profile":9,"step":1,"rows":4,"pos":2052,"prefix":2051,"fused":3,"vision":"no_axes","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["sdpa"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,4,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":2,"compiled_calls":2,"hidden_calls":2,"qk_rows_calls":0,"traces":2,"selector_dispatches":2,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":1,"update_only":0},"buckets":{"raw4096:pool1024:out0":2},"entry_count":2,"compiled_keys":[{"source":"hidden","mode":"blocks","input_shape":[1,1,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"hidden:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"hidden","mode":"row_tokens","input_shape":[1,4,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"hidden:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p9_s0","sha256":["6fd06b68b7059ae0220c03d7afada76bdadc86001123c4180f1a16380a89bece","fa52f9dff396320af135c22c25b328674b3c44b6c768df9ed2d54d31bbecc71f","5385335d390b6950b06db2bd3a3e86efe8b4ca7d48d7458828f6891d0815be26","0823d3bc139c67f69911c8970fb2c1bcf5f757971e9191f8b98d56f6689840ae","ea3bda4b6f89e3552003b8afb307a5d04a67cb8d024ea18b7a7df90a172f7342"],"profile":9,"step":0,"rows":1,"pos":2051,"prefix":2051,"fused":3,"vision":"no_axes","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["flash"],"counts":{},"offset":2052,"pooled_len":513,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":1,"qk_rows_calls":0,"traces":1,"selector_dispatches":1,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw4096:pool1024:out0":1},"entry_count":1,"compiled_keys":[{"source":"hidden","mode":"blocks","input_shape":[1,1,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"hidden:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p9_s1","sha256":["d8762135183d3a85dfd8713c2fda95cf6397f495329ee3f626a0fc56f3cb8e2d","c7e437185225097e4ce2bb012fb9dea74124314e3858db5b2d3f89d175628f8e","99fb41883cf58e2cec92897b8486a4b20b6817dc32c1b98d320280e2984756b4","e38b573ffab0799d27fcadc08d1e6a6aa31e40b4e6a61c5e9c98073a7012ba24","f54917b64b4e33520dd53be4677edb0d92a41e62d868d44eff6fdbc030530df8"],"profile":9,"step":1,"rows":4,"pos":2052,"prefix":2051,"fused":3,"vision":"no_axes","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":true},"consumers":["rows_gather"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,4,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":2,"compiled_calls":2,"hidden_calls":2,"qk_rows_calls":0,"traces":2,"selector_dispatches":2,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":1,"update_only":0},"buckets":{"raw4096:pool1024:out0":2},"entry_count":2,"compiled_keys":[{"source":"hidden","mode":"blocks","input_shape":[1,1,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"hidden:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"hidden","mode":"row_tokens","input_shape":[1,4,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"hidden:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p0_s0","sha256":["81cabd974019bcd738e4db77d9d7b089042d747910700154be3cc3da1ce986e7","316ed43e666fc6fb252dda7a8003ecb4be8da621d37db522985096db50465844","b4f1ff4375d0b1de7607bef982261bc0c7eba872a69ddc7b7dad6fbfdd3e1859","39b0f22d255932aefce0fff5850e71c764d11fd9e48c34b05fdbe22de04e7d44","b895415c8e478aff1ca87a13b4a553231b475db7a3969586a3b3d9b0d4291cd7"],"profile":0,"step":0,"rows":1,"pos":0,"prefix":0,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":1,"pooled_len":0,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,256,128],[1,256,128],[1,2,256,256],[1,2,256,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":0,"qk_rows_calls":1,"traces":1,"selector_dispatches":0,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":0,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":1},"buckets":{"raw256:pool256:out0":1},"entry_count":1,"compiled_keys":[{"source":"qk_rows","mode":"update_only","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,256,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,256,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":0,"label":"qk_rows:update_only:s1:raw256:pool256:out0:scorenone:chunk0:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p0_s1","sha256":["305e99a624ec1ef7144f43fad916352118cde06382377139cb606a09e5079bc9","8a6f35d3a8d1a938887f03baf9718fa9af582d121d1bed815dc68cb0482db7b1","46fe8ab534bd3d133ca7aa24d80a2f771907107ecac4f4f2496bc6501da6b855","847fb2445590d5b14c84e8a3afa037193a54e60aa9ab935f11815c0c237582c3","4152bdd42eacde7f856bace140740cc1b118fc0626c091daf2bbb7abaef2bc2d","cdb4f72a8a977ed2805ee125111ac4561a05997924610cb127d456aedb868553"],"profile":0,"step":1,"rows":4,"pos":1,"prefix":0,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":5,"pooled_len":1,"state_present":[true,true,true,true,true],"shapes":[[1,4,64],[1,256,128],[1,256,128],[1,1,128,256],[1,2,256,256],[1,2,256,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":0,"qk_rows_calls":1,"traces":1,"selector_dispatches":0,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":0,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":1},"buckets":{"raw256:pool256:out0":1},"entry_count":1,"compiled_keys":[{"source":"qk_rows","mode":"update_only","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,256,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,256,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":0,"label":"qk_rows:update_only:s1:raw256:pool256:out0:scorenone:chunk0:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p0_s2","sha256":["0db2e98676e803665a722901a32da71a553b2c589a9108e3b1f79a5be93c9992","75b526cbdc562be1855bd78936eb74132ac9e8c92186a06a517106e48544f312","d9b5b65fe3b0ac1e201e710f39f49c2d7a87f97dfb34cbdd71d1598d8b7d037d","fdfd8846b635f94a215516ccbca6b106a13d69c99a088893020fe5a1f27c110e","b456ba45bb69d6aec6338b8fc22e4bdb1e463a084966b8c33ed381c43d5579a3","ebc3501a599b08fc6c3f880049f7a5b4158343195f100e0afd4b769436043880"],"profile":0,"step":2,"rows":33,"pos":5,"prefix":0,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":38,"pooled_len":9,"state_present":[true,true,true,true,true],"shapes":[[1,33,64],[1,256,128],[1,256,128],[1,1,128,256],[1,2,256,256],[1,2,256,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":0,"qk_rows_calls":1,"traces":1,"selector_dispatches":0,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":0,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":1},"buckets":{"raw256:pool256:out0":1},"entry_count":1,"compiled_keys":[{"source":"qk_rows","mode":"update_only","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,256,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,256,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":0,"label":"qk_rows:update_only:s1:raw256:pool256:out0:scorenone:chunk0:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} @@ -45,12 +45,12 @@ {"kernel":"qsa_attention_p6_s0","sha256":["068422304f4191e74610562e007ffa7ec0d4b2ff5ed619997442ca1d254073b1","6f49058a22d7c62bbf4ffba4542b7aa32a130076b07537ac5605b2c85bb97162","158b3b53061c6d3d089049f9fdb47064a8c4f3d2f3fdb0ef5089afcaeba7cc5d","d4bc3f9184f206d1c86876a7ce98f086cf94f7912b460f7dc8c809b5b78d579b","011cc9bf7a9b2f86548e46728a9013e24273ca27ee50f9ee3f5509ecbc7c851f"],"profile":6,"step":0,"rows":2048,"pos":32768,"prefix":32768,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":32768,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["prefill_flash"],"counts":{"flash_kernel":1,"compiled_selector":1},"offset":34816,"pooled_len":8704,"state_present":[true,true,false,true,true],"shapes":[[1,2048,64],[1,65536,128],[1,16384,128],[1,2,34816,256],[1,2,34816,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":0,"qk_rows_calls":1,"traces":1,"selector_dispatches":1,"capacity_transitions":0,"prefill_score_producers":{"mpp":1,"mlx":0},"modes":{"blocks":0,"prefill_blocks":1,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw65536:pool16384:out0":1},"entry_count":1,"compiled_keys":[{"source":"qk_rows","mode":"prefill_blocks","input_shape":[1,2048,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"mpp","selector_chunk_rows":2048,"label":"qk_rows:prefill_blocks:s2048:raw65536:pool16384:out0:scorempp:chunk2048:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p6_s1","sha256":["bcbfcd4ff8d59b123b6ba9e84c2d621b9982c5e7430c84b82240f2b4416e48af","d50916defa6c9151ffeb706a7e0d6b8c04b861eaa94ed2b637f82fce79c8f2c9","158b3b53061c6d3d089049f9fdb47064a8c4f3d2f3fdb0ef5089afcaeba7cc5d","7c9ca3c2fd81b2f4749ce1f575a963b7054f004b91879a64212ecda5b0c34463","773fdcd29060aa9a94684db395aeccdbe1af803e191f3d0b938abc109d715fbb"],"profile":6,"step":1,"rows":1,"pos":34816,"prefix":32768,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":32768,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["flash"],"counts":{},"offset":34817,"pooled_len":8704,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,65536,128],[1,16384,128],[1,2,35072,256],[1,2,35072,256]],"core":{"calls":2,"compiled_calls":2,"hidden_calls":0,"qk_rows_calls":2,"traces":2,"selector_dispatches":2,"capacity_transitions":0,"prefill_score_producers":{"mpp":1,"mlx":0},"modes":{"blocks":1,"prefill_blocks":1,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw65536:pool16384:out0":2},"entry_count":2,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw65536:pool16384:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"prefill_blocks","input_shape":[1,2048,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"mpp","selector_chunk_rows":2048,"label":"qk_rows:prefill_blocks:s2048:raw65536:pool16384:out0:scorempp:chunk2048:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p6_s2","sha256":["e5f8b59d40e7e1c579d22d66af242e2882d1ab57727f17c7b9eeea894a5368e7","3b29542bd05bd5525466ad6c5951dd2d59759bed880564b8e07058d8feed1d42","9bc8c466e0a3164a6f4e22c0ffe23a895126f57945d4e7f0d5a07e79c63778ad","a41386c0de5defa3cc2eb26fe6005bab14865f573f9380f998d35e324abdbb43","270285c24ae2ea32e45240ada590a6cc1c9e7b14adfc23b270c7d7440c7372d4"],"profile":6,"step":2,"rows":4,"pos":34817,"prefix":32768,"fused":4,"vision":"none","no_indexer":false,"biased":false,"flash_min":32768,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["rows_gather"],"counts":{},"offset":34821,"pooled_len":8705,"state_present":[true,true,false,true,true],"shapes":[[1,4,64],[1,65536,128],[1,16384,128],[1,2,35072,256],[1,2,35072,256]],"core":{"calls":3,"compiled_calls":3,"hidden_calls":0,"qk_rows_calls":3,"traces":3,"selector_dispatches":3,"capacity_transitions":0,"prefill_score_producers":{"mpp":1,"mlx":0},"modes":{"blocks":1,"prefill_blocks":1,"dense_mask":0,"row_tokens":1,"update_only":0},"buckets":{"raw65536:pool16384:out0":3},"entry_count":3,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw65536:pool16384:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"prefill_blocks","input_shape":[1,2048,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"mpp","selector_chunk_rows":2048,"label":"qk_rows:prefill_blocks:s2048:raw65536:pool16384:out0:scorempp:chunk2048:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,4,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,65536,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,16384,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"qk_rows:row_tokens:s4:raw65536:pool16384:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p7_s0","sha256":["f9c3c5622d5d7c346d6c2881634894e60f030654828e74d96041317d5d1b99e5","fa52f9dff396320af135c22c25b328674b3c44b6c768df9ed2d54d31bbecc71f","5385335d390b6950b06db2bd3a3e86efe8b4ca7d48d7458828f6891d0815be26","f583e84d01f4e3bfb152b8855f8520d25c5fd1639c20961343d44d3b494ee18f","ea3bda4b6f89e3552003b8afb307a5d04a67cb8d024ea18b7a7df90a172f7342"],"profile":7,"step":0,"rows":1,"pos":2051,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":2052,"pooled_len":513,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":0,"qk_rows_calls":1,"traces":1,"selector_dispatches":1,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw4096:pool1024:out0":1},"entry_count":1,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p7_s1","sha256":["96067951d1aa0a433bcc85d4c51f0dbbf8f142f2f97b0e4aa1f0396ce6af89de","c7e437185225097e4ce2bb012fb9dea74124314e3858db5b2d3f89d175628f8e","99fb41883cf58e2cec92897b8486a4b20b6817dc32c1b98d320280e2984756b4","4af7a650ea8717de842e55d3691237c8e2af5c5c8ac78868e1323773e1a792f1","f54917b64b4e33520dd53be4677edb0d92a41e62d868d44eff6fdbc030530df8"],"profile":7,"step":1,"rows":4,"pos":2052,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,4,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":2,"compiled_calls":2,"hidden_calls":0,"qk_rows_calls":2,"traces":2,"selector_dispatches":2,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":1,"update_only":0},"buckets":{"raw4096:pool1024:out0":2},"entry_count":2,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,4,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"qk_rows:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p7_s2","sha256":["5c48c9fcdb884b10fecce90489b58457a977528d3732d9ceb90d202f9a8bcabf","0962884629f0854f29ed23c253f65f1d44fcb8bfde142e948411c9e0ed65faf9","5b6ab0a7cc9a3a122481547531253bcba5f039fc507718fc85bef86dd8dc56ea","57eb1075e3ed3fe889a40c34bd1e59974119725a3a9ef009646a3427b3be9dd6","7fe9f3c5a939bb7dd64a147ea63354dc92c40828f2cc4d1cf8b3be331b86f4e5"],"profile":7,"step":2,"rows":2,"pos":2054,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,2,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":3,"compiled_calls":3,"hidden_calls":0,"qk_rows_calls":3,"traces":3,"selector_dispatches":3,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":2,"update_only":0},"buckets":{"raw4096:pool1024:out0":3},"entry_count":3,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,2,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":2,"label":"qk_rows:row_tokens:s2:raw4096:pool1024:out0:scorenone:chunk2:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,4,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"qk_rows:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p7_s3","sha256":["554527077cad0087f4b548f1cf28dfbf186762ce6e3e12982922b97eec865e1d","9424991712e6561268208093dc9c9df018591b376554ac208eab7f242bdba4e4","5b6ab0a7cc9a3a122481547531253bcba5f039fc507718fc85bef86dd8dc56ea","28b0c2bfcc8486f5c6daafec65e013f6d2efbf99b24cdf75ba25372716d828d9","6d3d7b11ac7cff92650bdc798b862cbd21b383cd0154f53eee439af721d3c55f"],"profile":7,"step":3,"rows":1,"pos":2056,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":2057,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":4,"compiled_calls":4,"hidden_calls":0,"qk_rows_calls":4,"traces":3,"selector_dispatches":4,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":2,"prefill_blocks":0,"dense_mask":0,"row_tokens":2,"update_only":0},"buckets":{"raw4096:pool1024:out0":4},"entry_count":3,"compiled_keys":[{"source":"qk_rows","mode":"blocks","input_shape":[1,1,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"qk_rows:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,2,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":2,"label":"qk_rows:row_tokens:s2:raw4096:pool1024:out0:scorenone:chunk2:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"qk_rows","mode":"row_tokens","input_shape":[1,4,640],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"qk_rows:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p7_s0","sha256":["9fc17259582838c33640717e4e21cdd556e3f5d03bd8dda03de124d43f7b55e8","48e48f96c54776ebd1a034cdfcf93ae5e027038346bff0b1cc4f6600a73bdbd4","56ff514413104108eea40eeb10e28ba89d4bcb5f273e7bfb327b650239b2608e","d1f4df60a631b8170515760665a9212d6375de4e6456f75d90eb31179c91596c","f583e84d01f4e3bfb152b8855f8520d25c5fd1639c20961343d44d3b494ee18f","ea3bda4b6f89e3552003b8afb307a5d04a67cb8d024ea18b7a7df90a172f7342"],"profile":7,"step":0,"rows":1,"pos":2051,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["flash"],"counts":{},"offset":2052,"pooled_len":513,"state_present":[true,true,true,true,true],"shapes":[[1,1,64],[1,2304,128],[1,1024,128],[1,1,128,1024],[1,2,2304,256],[1,2,2304,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p7_s1","sha256":["24e8ced9cf6525d3fb9499e8cb45e70960458b26278552ff5e3bdb825596f33a","1ceca1a52f5057c1aaee04689fc8026d3df9e1b657d718d75b9a72aeb8a88b25","8ee668bf4a7f6eca91380e543147f88eb73d8bd3384c1a375b869d4680d95753","9221cbebf9df5094116ab9269b2c0e543d03fcd7f7ec5e9c1c745fc7b2babd0e","4af7a650ea8717de842e55d3691237c8e2af5c5c8ac78868e1323773e1a792f1","f54917b64b4e33520dd53be4677edb0d92a41e62d868d44eff6fdbc030530df8"],"profile":7,"step":1,"rows":4,"pos":2052,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["rows_gather"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,true,true,true],"shapes":[[1,4,64],[1,2304,128],[1,1024,128],[1,1,128,1024],[1,2,2304,256],[1,2,2304,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p7_s2","sha256":["6cd50cb32aab449013ff1db077d2e4a2f480fb71f1ae11ed80499dfab78257d0","b49eaa7452b6018be70df323222b0624f230df95ea900e42a974eaba5e2ac975","a8f7c3b9add9760148121cf236dab8d6844d38d42ab5aa4a7de20bb14c342d97","d4abe49bd48fe5eeea175dd7102faf5cc4a0bcf4831d233e7521c51cf7a76932","57eb1075e3ed3fe889a40c34bd1e59974119725a3a9ef009646a3427b3be9dd6","7fe9f3c5a939bb7dd64a147ea63354dc92c40828f2cc4d1cf8b3be331b86f4e5"],"profile":7,"step":2,"rows":2,"pos":2054,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["rows_gather"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,true,true,true],"shapes":[[1,2,64],[1,2304,128],[1,1024,128],[1,1,128,1024],[1,2,2304,256],[1,2,2304,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p7_s3","sha256":["80cb2c6d1f34cc44dfbcd30192a17c8cb4a0c2a19017924058ece49e2f5c6486","59da7bcfd15c46c7c183a322559d2f455dba99f677bbef99b3626bd32160147f","a8f7c3b9add9760148121cf236dab8d6844d38d42ab5aa4a7de20bb14c342d97","d4abe49bd48fe5eeea175dd7102faf5cc4a0bcf4831d233e7521c51cf7a76932","28b0c2bfcc8486f5c6daafec65e013f6d2efbf99b24cdf75ba25372716d828d9","6d3d7b11ac7cff92650bdc798b862cbd21b383cd0154f53eee439af721d3c55f"],"profile":7,"step":3,"rows":1,"pos":2056,"prefix":2051,"fused":4,"vision":"table","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["flash"],"counts":{},"offset":2057,"pooled_len":514,"state_present":[true,true,true,true,true],"shapes":[[1,1,64],[1,2304,128],[1,1024,128],[1,1,128,1024],[1,2,2304,256],[1,2,2304,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p8_s0","sha256":["2f2d1534efe9754560b2d2289f6c9756da611928c8c9e2321786ab03dd447111","84e64cfc1690832a3a6ed516b577d1001b89b382254c08b2663f493691330f2e","7da3f843c044b333d43259efd0a68f1c0a8f5230707898663f9ce096dd1f294a"],"profile":8,"step":0,"rows":3,"pos":0,"prefix":0,"fused":0,"vision":"none","no_indexer":true,"biased":true,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":3,"pooled_len":0,"state_present":[false,false,false,true,true],"shapes":[[1,3,64],[1,2,256,256],[1,2,256,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p8_s1","sha256":["dbe815a17715596ea030903a6800bba50c5ed74b9e03a3d29826e42acd57377b","50d2e91a257e44f7554a3ba6b1e0b0c0bef2a3f9ec4f1c4b6904e150f9b7b48c","a79deff710fa918973e19986ced87ae13347fbd66972c79fd4241d299f4472dc"],"profile":8,"step":1,"rows":1,"pos":3,"prefix":0,"fused":0,"vision":"none","no_indexer":true,"biased":true,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":4,"pooled_len":0,"state_present":[false,false,false,true,true],"shapes":[[1,1,64],[1,2,256,256],[1,2,256,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} {"kernel":"qsa_attention_p8_s2","sha256":["faeacdecf2a8829ff7c7c5751858003e0931435cae65c610e5c362a3997ba050","d4c84ff24f73effb7d50f227eb1248f997cc7ea466b6a15e6fa2876cf9481ea0","af7554d1b0575d27146ecc0b2d3755a89a645e42231c214b62cbecfaf6d81402"],"profile":8,"step":2,"rows":4,"pos":4,"prefix":0,"fused":0,"vision":"none","no_indexer":true,"biased":true,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":8,"pooled_len":0,"state_present":[false,false,false,true,true],"shapes":[[1,4,64],[1,2,256,256],[1,2,256,256]],"core":null,"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p9_s0","sha256":["91fd57e3f97be79622b57e34ed7e80b9934f65ce650592a6a619f0c88b322bd9","fa52f9dff396320af135c22c25b328674b3c44b6c768df9ed2d54d31bbecc71f","5385335d390b6950b06db2bd3a3e86efe8b4ca7d48d7458828f6891d0815be26","0823d3bc139c67f69911c8970fb2c1bcf5f757971e9191f8b98d56f6689840ae","ea3bda4b6f89e3552003b8afb307a5d04a67cb8d024ea18b7a7df90a172f7342"],"profile":9,"step":0,"rows":1,"pos":2051,"prefix":2051,"fused":3,"vision":"no_axes","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":2052,"pooled_len":513,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":1,"qk_rows_calls":0,"traces":1,"selector_dispatches":1,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw4096:pool1024:out0":1},"entry_count":1,"compiled_keys":[{"source":"hidden","mode":"blocks","input_shape":[1,1,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"hidden:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} -{"kernel":"qsa_attention_p9_s1","sha256":["893d7f0e21dc6a1c5099fa172b86630a0b694dc90ea7c6db4ae928c2151e9b00","c7e437185225097e4ce2bb012fb9dea74124314e3858db5b2d3f89d175628f8e","99fb41883cf58e2cec92897b8486a4b20b6817dc32c1b98d320280e2984756b4","e38b573ffab0799d27fcadc08d1e6a6aa31e40b4e6a61c5e9c98073a7012ba24","f54917b64b4e33520dd53be4677edb0d92a41e62d868d44eff6fdbc030530df8"],"profile":9,"step":1,"rows":4,"pos":2052,"prefix":2051,"fused":3,"vision":"no_axes","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["sdpa"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,4,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":2,"compiled_calls":2,"hidden_calls":2,"qk_rows_calls":0,"traces":2,"selector_dispatches":2,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":1,"update_only":0},"buckets":{"raw4096:pool1024:out0":2},"entry_count":2,"compiled_keys":[{"source":"hidden","mode":"blocks","input_shape":[1,1,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"hidden:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"hidden","mode":"row_tokens","input_shape":[1,4,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"hidden:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p9_s0","sha256":["6fd06b68b7059ae0220c03d7afada76bdadc86001123c4180f1a16380a89bece","fa52f9dff396320af135c22c25b328674b3c44b6c768df9ed2d54d31bbecc71f","5385335d390b6950b06db2bd3a3e86efe8b4ca7d48d7458828f6891d0815be26","0823d3bc139c67f69911c8970fb2c1bcf5f757971e9191f8b98d56f6689840ae","ea3bda4b6f89e3552003b8afb307a5d04a67cb8d024ea18b7a7df90a172f7342"],"profile":9,"step":0,"rows":1,"pos":2051,"prefix":2051,"fused":3,"vision":"no_axes","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["flash"],"counts":{},"offset":2052,"pooled_len":513,"state_present":[true,true,false,true,true],"shapes":[[1,1,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":1,"compiled_calls":1,"hidden_calls":1,"qk_rows_calls":0,"traces":1,"selector_dispatches":1,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":0,"update_only":0},"buckets":{"raw4096:pool1024:out0":1},"entry_count":1,"compiled_keys":[{"source":"hidden","mode":"blocks","input_shape":[1,1,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"hidden:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} +{"kernel":"qsa_attention_p9_s1","sha256":["d8762135183d3a85dfd8713c2fda95cf6397f495329ee3f626a0fc56f3cb8e2d","c7e437185225097e4ce2bb012fb9dea74124314e3858db5b2d3f89d175628f8e","99fb41883cf58e2cec92897b8486a4b20b6817dc32c1b98d320280e2984756b4","e38b573ffab0799d27fcadc08d1e6a6aa31e40b4e6a61c5e9c98073a7012ba24","f54917b64b4e33520dd53be4677edb0d92a41e62d868d44eff6fdbc030530df8"],"profile":9,"step":1,"rows":4,"pos":2052,"prefix":2051,"fused":3,"vision":"no_axes","no_indexer":false,"biased":false,"flash_min":2049,"prefill_gather":true,"options":{"fused":true,"compiled":true,"prefill":true,"phase_prefill":true,"prefill_min_rows":32,"prefill_compile_rows":2048,"prefill_min_context":2049,"prefill_workspace":134217728,"tile_rows":0,"flash":true,"gather_decode":false,"gather":true,"gather_max_rows":8,"gather_min_context":0,"tf32":false},"consumers":["rows_gather"],"counts":{},"offset":2056,"pooled_len":514,"state_present":[true,true,false,true,true],"shapes":[[1,4,64],[1,4096,128],[1,1024,128],[1,2,2304,256],[1,2,2304,256]],"core":{"calls":2,"compiled_calls":2,"hidden_calls":2,"qk_rows_calls":0,"traces":2,"selector_dispatches":2,"capacity_transitions":0,"prefill_score_producers":{"mpp":0,"mlx":0},"modes":{"blocks":1,"prefill_blocks":0,"dense_mask":0,"row_tokens":1,"update_only":0},"buckets":{"raw4096:pool1024:out0":2},"entry_count":2,"compiled_keys":[{"source":"hidden","mode":"blocks","input_shape":[1,1,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":1,"label":"hidden:blocks:s1:raw4096:pool1024:out0:scorenone:chunk1:mlx.core.bfloat16/mlx.core.bfloat16"},{"source":"hidden","mode":"row_tokens","input_shape":[1,4,64],"input_dtype":"mlx.core.bfloat16","raw_shape":[1,4096,128],"raw_dtype":"mlx.core.bfloat16","pooled_shape":[1,1024,128],"pooled_dtype":"mlx.core.bfloat16","n_heads":4,"kv_heads":1,"head_dim":128,"block_topk":512,"compress_ratio":4,"dense_output_tokens":0,"prefill_score_producer":"none","selector_chunk_rows":4,"label":"hidden:row_tokens:s4:raw4096:pool1024:out0:scorenone:chunk4:mlx.core.bfloat16/mlx.core.bfloat16"}],"require_bucketed_backing":true,"minimum_raw_capacity":256,"minimum_pooled_capacity":256,"selector_scratch_bytes":33554432,"prefill_score_workspace_bytes":134217728,"rope_attention_scaling":1.0},"frequencies":[1.0,0.6042963862419128,0.3651741147041321,0.22067339718341827,0.1333521455526352,0.08058421313762665,0.04869674891233444,0.02942727319896221,0.017782794311642647,0.010746077634394169,0.006493816152215004,0.003924189601093531,0.0023713738191872835,0.0014330126577988267,0.0008659642189741135,0.0005232990952208638,0.0003162277862429619,0.00019109529966954142,0.00011547819303814322,6.978305464144796e-05,4.2169649532297626e-05,2.5482966520939954e-05,1.5399264157167636e-05,9.305719686381053e-06,5.62341347176698e-06,3.398208264115965e-06,2.0535248950181995e-06,1.2409376495270408e-06,7.498942409256415e-07,4.5315832153391966e-07,2.738419482284371e-07,1.6548169412544667e-07]} diff --git a/tools/issue107-indexer-reference.py b/tools/issue107-indexer-reference.py new file mode 100644 index 0000000..9467691 --- /dev/null +++ b/tools/issue107-indexer-reference.py @@ -0,0 +1,53 @@ +"""Real-geometry image-aware QSA preparation, checked independently with NumPy.""" +import hashlib +import json +from pathlib import Path +import subprocess +import sys + +checkout, output = map(Path, sys.argv[1:3]) +assert subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112" +sys.path.insert(0, str(checkout)) +import mlx.core as mx +import numpy as np +from mtplx.models.qwen4_exp import TextArgs, QSAIndexer, _build_mrope_axes +from mtplx.attention_context import vision_rope + +output.mkdir(parents=True, exist_ok=True) +args = TextArgs(hidden_size=640, head_dim=256, indexer_head_dim=128, indexer_n_heads=4, + indexer_kv_heads=1, indexer_compress_ratio=4, indexer_budget=2048, + rope_parameters=dict(mrope_interleaved=True, mrope_section=[11,11,10], partial_rotary_factor=.25, + rope_theta=10000000, rope_type="default")) +idx = QSAIndexer(args) +axes = _build_mrope_axes([11,11,10],True) +table = np.stack([np.arange(2060, dtype=np.int32)//(axis+1) for axis in range(3)]) +def pattern(shape, seed): + return (((np.arange(np.prod(shape), dtype=np.int64)*17+seed)%97-48)/32).astype(np.float32).reshape(shape) +raw_q = pattern((1,16,4,128),3) +raw_k = pattern((1,2072,128),11) + +for dtype_name, dtype in [("f32",mx.float32),("bf16",mx.bfloat16)]: + idx.q_layernorm.weight = mx.ones((128,),dtype=dtype) + idx.k_layernorm.weight = mx.ones((128,),dtype=dtype) + q, k = mx.array(raw_q).astype(dtype), mx.array(raw_k).astype(dtype) + for mode in ["table", "equal", "delta", "empty"]: + delta = -7 if mode in ("table","delta","empty") else 0 + table_value = table if mode=="table" else np.broadcast_to(np.arange(2060,dtype=np.int32),(3,2060)) if mode=="equal" else np.zeros((3,0),dtype=np.int32) if mode=="empty" else None + with vision_rope(None if table_value is None else mx.array(table_value),delta): + query = idx._prepare_queries(q,2056) + pooled = idx._pool_keys_eager(k,0,518) + for name,value in [("query",query),("pooled",pooled)]: + data=np.asarray(value.astype(mx.float32)) + (output/f"{dtype_name}-{mode}-{name}.f32").write_bytes(data.tobytes()) + print(json.dumps(dict(case="indexer",dtype=dtype_name,mode=mode,part=name,shape=list(data.shape),sha256=hashlib.sha256(data.tobytes()).hexdigest())),flush=True) + if dtype_name=="f32": + def independent(values, positions): + values = values / np.sqrt(np.mean(values*values,axis=-1,keepdims=True)+args.rms_norm_eps) + pos3=np.stack([table_value[:,p] if table_value is not None and p<|image_pad|><|vision_end|>" * 2 +ids = _encode_messages_uncached(tokenizer, [ChatMessage(role="user", content=text)], enable_thinking=True, + reasoning_effort="low", preserve_reasoning_history=True, tools=None) +ids, splice = _materialize_vision_splice(argparse.Namespace(args=argparse.Namespace(model=str(root))), [image.read_bytes()]*2, ids) +assert len(ids) > 2048 +(output/"input.json").write_text(json.dumps(ids)) +out = generate_mtpk(runtime, ids, max_tokens=1, sampler=SamplerConfig(temperature=0., top_p=.95, top_k=20), + speculative_depth=3, seed=1, stop_token_ids=set(tokenizer.eos_token_ids), mtp_history_policy="committed", + verify_strategy="batched", vision_splice=splice, capture_final_state=True, + prefill_callback=lambda event: print(json.dumps(dict(case="prefill", **event)), flush=True)) +assert out.final_state.safe_to_commit +logits = np.asarray(out.final_state.final_logits.astype(mx.float32)).reshape(-1) +(output/"final.f32").write_bytes(logits.tobytes()) +print(json.dumps(dict(case="long_vision", prompt_tokens=len(ids), tokens=out.tokens, text=out.text, + final_logits_sha256=hashlib.sha256(logits.tobytes()).hexdigest())), flush=True)