Preserve image-aware sparse QSA for Qwen vision

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.

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::<Vec<_>>(),
++    )
++    .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::<Vec<_>>())})
++    );
++}

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}

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."
+}

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}

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}

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}

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}

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]}

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,

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

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)

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::<Vec<_>>(),
+            )
+            .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::<Vec<_>>()
+        );
+    }
+    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::<Vec<_>>();
+            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::<Vec<_>>(),
+                )),
+                "equal" => Some(ints(
+                    &[3, 2060],
+                    &(0..3).flat_map(|_| 0..2060).collect::<Vec<_>>(),
+                )),
+                "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::<Vec<_>>();
+                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();
+}

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::<Vec<_>>(),
+        std::fs::read(oracle.join("final.f32")).unwrap()
+    );
+    let oracle_ids: Vec<i32> =
+        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::<Vec<_>>(),
+    )
+    .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::<Vec<_>>())})
+    );
+    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})
+    );
+}

src/engine/metal/qwen_mtplx/owner.rs

@@ -253,7 +253,8 @@ impl QwenExecutor {
         incoming: &mut Option<LiveSession>,
     ) -> 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;

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();

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<super::attention::EagerSelection>, &'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<Array, String> {
-        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<Array, String> {
-        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<Option<Array>, 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<Option<PreparedRows>, 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");

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,

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<TextTurn, String> {
+        self.discard_incompatible_vision();
         let (prompt, images) = execution
             .vision_prompt
             .clone()

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,

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());
         }

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 {

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));
 }

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]}

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<table_value.shape[1] else np.full(3,p+delta) for p in positions],axis=1)
+                angles=pos3[axes].T.astype(np.float32)*np.asarray(idx._inv_freq)[None,:]
+                angles=np.concatenate([angles,angles],axis=-1)[None,:,None,:]
+                first=values[...,:64];rotated=np.concatenate([-first[...,32:],first[...,:32]],axis=-1)
+                return np.concatenate([first*np.cos(angles)+rotated*np.sin(angles),values[...,64:]],axis=-1)
+            np.testing.assert_allclose(np.asarray(query),independent(raw_q,np.arange(2056,2072)),rtol=3e-6,atol=3e-6)
+            np.testing.assert_allclose(np.asarray(pooled),independent(raw_k.reshape(1,518,4,128).mean(axis=2)[:,:,None],np.arange(0,2072,4))[:,:,0],rtol=3e-6,atol=3e-6)
+print(json.dumps(dict(case="independent_numpy",passed=True)),flush=True)

tools/issue107-model-reference.py

@@ -0,0 +1,48 @@
+"""Pinned MTPLX oracle for issue 107's real long two-image product request."""
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+checkout, root, image, output = map(Path, sys.argv[1:5])
+assert subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
+sys.path.insert(0, str(checkout))
+from mtplx.server.openai import _server_runtime_env_overrides
+os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root), generation_mode="mtp"), None))
+os.environ["MTPLX_SUSTAINED_PREFILL"] = "1"
+os.environ["MTPLX_PREFILL_CHUNK_SIZE"] = "2048"
+import mlx.core as mx
+import numpy as np
+from mlx_lm.utils import load_model
+from mtplx.models import qwen4_exp as qwen
+from mtplx.runtime import _load_tokenizer_resilient, MTPLXRuntime
+from mtplx.mtp_patch import MTPContract
+from mtplx.sampling import SamplerConfig
+from mtplx.generation import generate_mtpk
+from mtplx.server.openai import ChatMessage, _encode_messages_uncached, _materialize_vision_splice
+
+output.mkdir(parents=True, exist_ok=True)
+print(json.dumps(dict(case="environment", overrides={k:v for k,v in os.environ.items() if k.startswith("MTPLX_")})), flush=True)
+tokenizer = _load_tokenizer_resilient(root, json.loads((root/"config.json").read_text()))
+model, _ = load_model(root, lazy=False, strict=True, get_model_classes=lambda **_: (qwen.Model,qwen.ModelArgs))
+model.post_weight_load(root)
+assert model.attach_mtp(root)
+runtime = MTPLXRuntime(model, tokenizer, root, True, MTPContract())
+text = "Describe these images" + "<|vision_start|><|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)