diff --git a/docs/qwen-vision-20260911.md b/docs/qwen-vision-20260911.md new file mode 100644 index 0000000..68c9fd3 --- /dev/null +++ b/docs/qwen-vision-20260911.md @@ -0,0 +1,89 @@ +# Qwen vision verification — 2026-09-11 + +Qwen3.8 Flash Next now accepts the existing PNG/JPEG attachments. Its optional +vision encoder has a separate Model Manager entry; downloading, validating or +deleting it does not alter the text-model artifact set. The existing keep-vision- +weights-loaded preference also applies to Qwen. Application preprocessing, +loading and inference are Rust, using the existing Metal runtime kernels. + +The four pinned vision files (897,900,287 bytes total) were downloaded and SHA-256 +verified in `~/Library/Application Support/de.rfc1437.ds4server/models/qwen3.8-flash-next`. +The artifact revision is `74559cdf34fbfc0b593de72d17e93f37fd4f9ea7` of +`Youssofal/Qwen3.8-Flash-Next-MTPLX-Bare-Speed`; the text manifest remains unchanged. +The video processor configuration is part of that artifact set; this change +implements still images. + +## Grounding check + +Every primary description run used exactly `Describe this image`. The supplied +1024×1024 PNG was copied unchanged to +`local-eval-results/qwen-vision-input/image.png`. Its SHA-256 is +`b7568d90f4df6180d9af14a824dd553cd995457561d29167354c1ba66b728347`. +Only the image bytes and prompt enter the model; neither the source filename nor +the neutral filename is included in its text input. + +Rust/Metal generation succeeded both with MTP and with ordinary autoregressive +decoding. Qwen described a golden-tan cartoon llama/alpaca, large eyes, upright +ears, an open smiling mouth, mountains, a sunset and a grainy poster texture. +These details are visible in the supplied image. The direct cold-session result +begins: + +> This is a stylized, cartoon-style illustration of a llama (or alpaca) shown from the neck up, set against a sunset landscape. + +Controls used the same prompt: + +| Input / session | Observed result | +| --- | --- | +| No image, Rust and MTPLX | Reports no attached image and requests one | +| Solid blue image, same dimensions, after the animal image | Describes a uniform blue field; zero cached prompt tokens | +| Same image repeated | Same description; all 1,069 prompt tokens reused | +| Saved checkpoint, reset, restore, follow-up | Correct animal description; 1,446 cached tokens out of 1,460 | + +This demonstrates image-dependent descriptions for these inputs, not a general +guarantee against hallucinations. + +## Oracle and reproducibility + +The oracle is local MTPLX reference `e652d55` with MLX 0.32.2. Python scripts under +`tools/qwen-vision*-reference.py` run only that reference, never the application. +The Rust tower matches its exported values exactly at patch embedding, position +embedding, rotary positions, blocks 0 and 26, and the final merger. Both the +1024×1024 input (2,621,440 final values) and a small non-square fixture (168,960 +final values) had zero differing values. CPU resize/preprocessing also matches +three Pillow/MTPLX golden hashes. This is exact encoder agreement; full generated +token-sequence parity is not claimed. + +Local evidence is retained under `local-eval-results/`: + +- `qwen-vision-rust-mtp.jsonl`, `qwen-vision-rust-ar.jsonl`, + `qwen-vision-rust-no-image.jsonl`: complete application runs. +- `qwen-vision-lifecycle.jsonl`: cold, repeat, restored and changed-image runs. +- `qwen-vision-chat-reference.jsonl`: oracle image/no-image runs. +- `qwen-vision-image/`, `qwen-vision-small/`: exported oracle arrays. +- `qwen-vision-small-rust.log`: small-fixture exact comparison. + +Example application invocation (empty YAML config avoids an unrelated system +prompt): + +```sh +target/release/ds4-server model-eval \ + --model qwen3.8-flash-next --config /tmp/qwen-vision-config.yaml \ + --prompt 'Describe this image' \ + --image-file local-eval-results/qwen-vision-input/image.png \ + --context 8192 --max-tokens 1024 --reasoning low \ + --temperature 0 --top-p 0.95 --seed 1 --acceleration on \ + --prefill-chunk 2048 --warmup off --canary on --max-memory-gib 108 +``` + +The ignored GPU tests `qwen_vision_tower_matches_mtplx_image` and +`qwen_vision_chat_checkpoint_preserves_image_identity` are runnable with +`DS4_QWEN38_ARTIFACTS` pointing to the model directory and respectively +`DS4_QWEN_VISION_REFERENCE` pointing to exported arrays or +`DS4_QWEN_VISION_IMAGE` pointing to the neutral input. Run one GPU model process +at a time under `test-supervisor` with an appropriate memory limit. + +Final checks: `cargo fmt --all -- --check`, Clippy with all targets/features and +warnings denied, `RUST_TEST_THREADS=1 cargo test --all-features` (315 passed, +204 opt-in tests ignored), and `make bundle` all succeeded. The two encoder +comparisons and image lifecycle test were additionally executed explicitly with +GPU access. The updated, signed application is `target/release/DS4Server.app`. diff --git a/src/app.rs b/src/app.rs index aed98e9..7927edc 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1402,7 +1402,7 @@ impl App { } Message::ComposerAction(action) => self.composer.perform(action), Message::ChooseVisionImage => { - if self.config.model != ModelChoice::Glm53Flash || self.generating { + if !self.config.model.supports_vision() || self.generating { return Task::none(); } return Task::perform( diff --git a/src/app/preferences.rs b/src/app/preferences.rs index ec5e603..ea52fab 100644 --- a/src/app/preferences.rs +++ b/src/app/preferences.rs @@ -604,7 +604,7 @@ impl App { return; } self.config = config; - if self.config.model != ModelChoice::Glm53Flash { + if !self.config.model.supports_vision() { self.pending_vision_image = None; } self.context_limit = self.config.active_generation().context_tokens.max(0) as u32; @@ -889,7 +889,7 @@ impl App { } Message::PreferenceKeepVisionLoadedChanged(value) => { self.preference_draft.keep_vision_loaded = - self.preference_draft.acceleration_model == ModelChoice::Glm53Flash && value; + self.preference_draft.acceleration_model.supports_vision() && value; self.preference_error = None; } Message::PreferenceDsparkConfidenceChanged(value) => { diff --git a/src/app/projects.rs b/src/app/projects.rs index 7f3a761..26c7c6e 100644 --- a/src/app/projects.rs +++ b/src/app/projects.rs @@ -11,7 +11,7 @@ impl App { self.config = config; self.preference_draft = PreferenceDraft::from_saved(&self.config); self.context_limit = self.config.active_generation().context_tokens.max(0) as u32; - if model != ModelChoice::Glm53Flash { + if !model.supports_vision() { self.pending_vision_image = None; } #[cfg(target_os = "macos")] diff --git a/src/app/view/chat.rs b/src/app/view/chat.rs index f0a7e0e..b0db497 100644 --- a/src/app/view/chat.rs +++ b/src/app/view/chat.rs @@ -337,8 +337,8 @@ impl App { .padding([2, 6]) .into() }; - let vision_ready = self.config.model == ModelChoice::Glm53Flash - && model::engine_artifacts(ModelChoice::Glm53Flash, false, &models_path()) + let vision_ready = self.config.model.supports_vision() + && model::engine_artifacts(self.config.model, false, &models_path()) .vision .is_some(); let attach = if vision_ready && !self.generating && self.pending_vision_image.is_none() @@ -356,8 +356,10 @@ impl App { attach, container(text(if vision_ready { "Attach PNG or JPEG" + } else if self.config.model.supports_vision() { + "Download the vision encoder in Model Manager" } else { - "GLM 5.3 Flash vision sidecar is not ready" + "This model does not support image attachments" })) .padding(8) .style(preference_group_style), diff --git a/src/app/view/preferences.rs b/src/app/view/preferences.rs index 937ae9d..4100fe1 100644 --- a/src/app/view/preferences.rs +++ b/src/app/view/preferences.rs @@ -25,7 +25,7 @@ impl App { .supports_integrated_mtp() .then_some(Message::PreferenceGlmMtpTimingChanged); let keep_vision_loaded_toggle: Option Message> = - (self.preference_draft.acceleration_model == ModelChoice::Glm53Flash) + (self.preference_draft.acceleration_model.supports_vision()) .then_some(Message::PreferenceKeepVisionLoadedChanged); let dspark_strict_toggle: Option Message> = self .preference_draft @@ -623,7 +623,7 @@ impl App { ), hint( toggle(self.preference_draft.keep_vision_loaded) - .label("Keep GLM 5.3 vision weights loaded") + .label("Keep vision weights loaded") .on_toggle_maybe(keep_vision_loaded_toggle), "Keeps the vision encoder mapped between image turns for lower image latency. Off releases it after encoding all images in a turn, leaving more memory for long text contexts.", ), diff --git a/src/engine.rs b/src/engine.rs index 0d13620..ceb3cf0 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -6,6 +6,7 @@ mod kvstore; #[cfg(target_os = "macos")] mod metal; mod qwen; +mod qwen_vision; #[cfg(test)] #[path = "../tools/sampler-replay-benchmark.rs"] mod sampler_replay_benchmark; @@ -1504,9 +1505,6 @@ impl Generator { .iter() .any(|message| message.content.contains(VISION_DATA_START)); if let metal::Executor::QwenMtplx(executor) = &mut self.executor { - if has_vision { - return Err("vision input requires GLM 5.3 Flash".into()); - } return executor.generate(messages, settings, cancelled, &self.metrics, progress, emit); } let (tokens, overlays) = if has_vision { diff --git a/src/engine/metal/qwen_mtplx.rs b/src/engine/metal/qwen_mtplx.rs index 3a8efba..b2b22c4 100644 --- a/src/engine/metal/qwen_mtplx.rs +++ b/src/engine/metal/qwen_mtplx.rs @@ -70,6 +70,8 @@ mod stream; mod submission; mod unary; mod views; +mod vision; +mod vision_kernels; mod weights; #[cfg(test)] pub(super) use encoder::Scope; diff --git a/src/engine/metal/qwen_mtplx/compile_graph.rs b/src/engine/metal/qwen_mtplx/compile_graph.rs index fb67710..ee74f91 100644 --- a/src/engine/metal/qwen_mtplx/compile_graph.rs +++ b/src/engine/metal/qwen_mtplx/compile_graph.rs @@ -549,6 +549,7 @@ fn equivalent_operation(a: &Operation, b: &Operation, x: &Array, y: &Array) -> b | (Softmax(a), Softmax(b)) | (SearchSorted(a), SearchSorted(b)) | (CumsumLastF32(a), CumsumLastF32(b)) => a == b, + (Vision(a), Vision(b)) => a == b, (RmsNorm(a), RmsNorm(b)) => a == b, ( SdpaVector { diff --git a/src/engine/metal/qwen_mtplx/dense.rs b/src/engine/metal/qwen_mtplx/dense.rs index a0c4dbe..235870e 100644 --- a/src/engine/metal/qwen_mtplx/dense.rs +++ b/src/engine/metal/qwen_mtplx/dense.rs @@ -181,6 +181,15 @@ fn dot(a: &Array, b: &Array, out: &Array, k: u32) -> Result<(), String> { } pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String> { + evaluate_bias(a, b, None, out) +} + +pub(super) fn evaluate_bias( + a: &Array, + b: &Array, + bias: Option<&Array>, + out: &Array, +) -> Result<(), String> { let dtype = out.layout().dtype(); let dtype_name = dtype.kernel_name(); if a.layout().size() == 0 || b.layout().size() == 0 { @@ -247,6 +256,7 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String> } let ab = a.buffer(); let bb = b.buffer(); + let cb = bias.map(Array::buffer); let ob = out.buffer(); let ai = |index| { tensor(index, &ab) @@ -480,15 +490,33 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String> }, [32, wn, wm], )?; + let zero = 0_i32; + let one = 1_i32; + let scale = 1_f32; + let mut accum = vec![ + tensor(0, &intermediate).input(parts as usize * stride as usize), + output(1), + bytes(2, &parts), + bytes(3, &stride), + bytes(4, &n), + ]; + if let (Some(c), Some(cb)) = (bias, cb.as_ref()) { + accum.extend([ + tensor(5, cb) + .at_byte_offset(c.offset()) + .input(c.data_size()), + bytes(6, &zero), + bytes(7, &one), + bytes(8, &scale), + bytes(9, &scale), + ]); + } dispatch_geometry( - &format!("steel_gemm_splitk_accum_{dtype_name}_float32"), - &[ - tensor(0, &intermediate).input(parts as usize * stride as usize), - output(1), - bytes(2, &parts), - bytes(3, &stride), - bytes(4, &n), - ], + &format!( + "steel_gemm_splitk_accum_{dtype_name}_float32{}", + if bias.is_some() { "_axbpy" } else { "" } + ), + &accum, &[], [n, m, 1], super::block_dims([n, m, 1]), @@ -520,6 +548,29 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String> padding: 0, }; let mut bindings = vec![ai(0), bi(1), output(3), bytes(4, ¶ms)]; + #[repr(C)] + struct Add { + ldc: i32, + fdc: i32, + stride: i64, + alpha: f32, + beta: f32, + } + let add = Add { + ldc: 0, + fdc: 1, + stride: 0, + alpha: 1., + beta: 1., + }; + if let (Some(c), Some(cb)) = (bias, cb.as_ref()) { + bindings.extend([ + tensor(2, cb) + .at_byte_offset(c.offset()) + .input(c.data_size()), + bytes(5, &add), + ]); + } let batch_strides = a_batch.iter().chain(&b_batch).copied().collect::>(); if batch_ndim > 1 { bindings.extend([ @@ -535,7 +586,7 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String> &bindings, &constants(&[ (10, batch_ndim > 1), - (100, false), + (100, bias.is_some()), (110, false), (200, m.is_multiple_of(64)), (201, n.is_multiple_of(bn)), diff --git a/src/engine/metal/qwen_mtplx/execution.rs b/src/engine/metal/qwen_mtplx/execution.rs index 19fc5b1..e275eb7 100644 --- a/src/engine/metal/qwen_mtplx/execution.rs +++ b/src/engine/metal/qwen_mtplx/execution.rs @@ -216,6 +216,7 @@ impl Options { pub(super) struct Execution { pub(super) model: BoundModel, + pub(super) vision_prompt: Option<(Vec, Vec)>, pub(super) cache: Vec>, pub(super) mtp_cache: [Option; 1], pub(super) streams: Streams, @@ -324,6 +325,7 @@ impl Execution { let mtp_cache = model.borrow_dependent().make_mtp_cache(4); Ok(Self { model, + vision_prompt: None, cache, mtp_cache, streams, diff --git a/src/engine/metal/qwen_mtplx/fused.rs b/src/engine/metal/qwen_mtplx/fused.rs index 7966649..04f537a 100644 --- a/src/engine/metal/qwen_mtplx/fused.rs +++ b/src/engine/metal/qwen_mtplx/fused.rs @@ -47,6 +47,7 @@ pub(super) fn operator(op: &Operation) -> Option<&'static str> { U::Rsqrt => "Rsqrt", U::Sin => "Sin", U::Cos => "Cos", + U::Erf => "Erf", }, Operation::Binary(op) => match op { B::Equal => "Equal", diff --git a/src/engine/metal/qwen_mtplx/model.rs b/src/engine/metal/qwen_mtplx/model.rs index 4f392c5..247ee98 100644 --- a/src/engine/metal/qwen_mtplx/model.rs +++ b/src/engine/metal/qwen_mtplx/model.rs @@ -89,6 +89,7 @@ pub(super) struct TextModel<'a> { pub(super) mixer: hyper::GatedResidual<'a>, pub(super) head: Option>, // None means tied embedding weights. pub(super) mtp: Option>, + pub(super) vision: Option, pub(super) last_widened: Option, policy: CompilePolicy, decode_runs: Option>, @@ -111,7 +112,7 @@ pub(super) struct Output { pub(super) hidden: Option, } -fn embedding(layer: Linear<'_>, ids: &Array, stream: Stream) -> Result { +pub(super) fn embedding(layer: Linear<'_>, ids: &Array, stream: Stream) -> Result { match layer { Linear::Dense(weight) => indexing::take_rows(weight, ids, stream), Linear::Quantized(layer) => linear::dequantize( @@ -171,6 +172,7 @@ impl<'a> TextModel<'a> { head, parameter_order: Vec::new(), mtp: None, + vision: None, last_widened: None, policy, decode_runs: None, @@ -213,6 +215,8 @@ impl<'a> TextModel<'a> { streams: &Streams, stream: Stream, ) -> Result { + let image = self.vision.clone(); + let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta))); let emb = embedding(self.embedding, ids, stream)?; let mtp = self .mtp @@ -242,7 +246,18 @@ impl<'a> TextModel<'a> { streams: &Streams, stream: Stream, ) -> Result { - let emb = match input_embeddings { + let image = self.vision.clone(); + let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta))); + let spliced = match (&image, input_embeddings) { + (Some(v), None) => v.embeddings( + self.embedding, + ids, + super::vision::cache_position(cache.as_deref()) + 1, + stream, + )?, + _ => None, + }; + let emb = match input_embeddings.or(spliced.as_ref()) { Some(emb) => emb.clone(), None => embedding(self.embedding, ids, stream)?, }; @@ -440,7 +455,18 @@ impl<'a> TextModel<'a> { streams: &Streams, stream: Stream, ) -> Result { - let h = match input_embeddings { + let image = self.vision.clone(); + let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta))); + let spliced = match (&image, input_embeddings) { + (Some(v), None) => v.embeddings( + self.embedding, + ids, + super::vision::cache_position(cache.as_deref()), + stream, + )?, + _ => None, + }; + let h = match input_embeddings.or(spliced.as_ref()) { Some(h) => h.clone(), None => embedding(self.embedding, ids, stream)?, }; diff --git a/src/engine/metal/qwen_mtplx/ops.rs b/src/engine/metal/qwen_mtplx/ops.rs index d435aae..e43b863 100644 --- a/src/engine/metal/qwen_mtplx/ops.rs +++ b/src/engine/metal/qwen_mtplx/ops.rs @@ -6,6 +6,7 @@ use super::stream::{Device, Stream, Streams}; use super::{eval, views}; pub(super) enum Operation { + Vision(super::vision_kernels::Kernel), Compiled(super::fused::Expression), Load(super::load::Load), Reshape, @@ -208,7 +209,7 @@ impl Operation { | Self::ArangeU32 | Self::ArangeF32 { .. } | Self::ArangeI64 => Err("invalid indexing arity".into()), - Self::RmsNorm(_) => Err("RMSNorm requires two inputs".into()), + Self::Vision(_) | Self::RmsNorm(_) => Err("RMSNorm requires two inputs".into()), Self::QuantizedLinearBf16 { .. } => Err("quantized linear requires four inputs".into()), Self::GatherQmmBf16 { .. } => Err("GatherQMM requires six inputs".into()), Self::AffineQuantize { .. } => Err("quantize requires three sibling outputs".into()), @@ -257,6 +258,9 @@ pub(super) fn evaluate( .operation() .ok_or("array has no executable primitive")?; match &*operation { + Operation::Vision(k) => { + return super::vision_kernels::evaluate(*k, inputs, &outputs[0]); + } Operation::Compiled(expression) => { if stream.device() != Device::Gpu { return Err("compiled expression requires Metal".into()); diff --git a/src/engine/metal/qwen_mtplx/owner.rs b/src/engine/metal/qwen_mtplx/owner.rs index 195c68b..880d8be 100644 --- a/src/engine/metal/qwen_mtplx/owner.rs +++ b/src/engine/metal/qwen_mtplx/owner.rs @@ -15,6 +15,9 @@ use std::sync::atomic::AtomicBool; pub(in crate::engine) struct QwenExecutor { // All snapshot arrays must die before the executor's native Metal context. live: LiveSession, + vision_tower: Option, + vision_path: Option, + keep_vision_loaded: bool, metadata: QwenMetadata, identity: Identity, policy: TurnOptions<'static>, @@ -43,6 +46,13 @@ impl QwenExecutor { let (execution, policy) = Execution::open(settings)?; Ok(Self { live: LiveSession::default(), + vision_tower: None, + vision_path: settings + .artifacts + .vision + .as_ref() + .map(|_| settings.artifacts.model.clone()), + keep_vision_loaded: settings.speculative.keep_vision_loaded, metadata, identity, policy, @@ -67,6 +77,43 @@ impl QwenExecutor { if settings.context_tokens != self.identity.context as i32 { return Err("Qwen turn context differs from the loaded engine".into()); } + self.execution.vision_prompt = None; + self.execution + .model + .with_dependent_mut(|_, model| model.vision = None); + if messages + .iter() + .any(|m| m.content.contains(crate::engine::VISION_DATA_START)) + { + let root = self + .vision_path + .as_ref() + .ok_or("Qwen vision encoder is not installed; download it in Model Manager")?; + if self.vision_tower.is_none() { + self.vision_tower = Some(super::vision::Tower::load( + root, + &self.execution.streams, + self.execution.stream, + )?); + } + let prepared = super::vision::prepare( + cancelled, + &self.metadata, + messages, + settings, + self.vision_tower.as_ref().unwrap(), + &self.execution.streams, + self.execution.stream, + ); + if !self.keep_vision_loaded { + self.vision_tower = None; + } + let (ids, input, identities) = prepared?; + self.execution + .model + .with_dependent_mut(|_, model| model.vision = Some(input)); + self.execution.vision_prompt = Some((ids, identities)); + } let baseline = self.stats; let stats = &mut self.stats; let result = self.live.generate( @@ -163,6 +210,10 @@ impl QwenExecutor { fn clear_execution_cache(&mut self) -> Result<(), String> { self.execution.streams.synchronize(self.execution.stream)?; + self.execution.vision_prompt = None; + self.execution + .model + .with_dependent_mut(|_, model| model.vision = None); let model = self.execution.model.borrow_dependent(); self.execution.cache = model.make_cache(4); self.execution.mtp_cache = model.make_mtp_cache(4); @@ -866,3 +917,126 @@ fn qwen_product_generator_chat_matches_reference() { std::panic::resume_unwind(error); } } + +#[test] +#[ignore = "requires installed Qwen, Metal, and DS4_QWEN_VISION_IMAGE; run under test-supervisor"] +fn qwen_vision_chat_checkpoint_preserves_image_identity() { + use crate::settings::{ + GenerationPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, + }; + use base64::Engine as _; + let root = std::path::PathBuf::from(std::env::var("DS4_QWEN38_ARTIFACTS").unwrap()); + let image = std::fs::read(std::env::var("DS4_QWEN_VISION_IMAGE").unwrap()).unwrap(); + let directory = std::env::temp_dir().join(format!("qwen-vision-chat-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + let settings = crate::settings::effective_settings( + crate::model::ModelChoice::Qwen38FlashNext, + &GenerationPreferences { + context_tokens: 8192, + max_generated_tokens: 1024, + 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(); + eprintln!("vision lifecycle: loading Qwen"); + let mut executor = QwenExecutor::open(&settings.engine).unwrap(); + let metrics = Metrics::new(&directory); + let cancelled = AtomicBool::new(false); + let user = |content: String| ChatTurn { + user: true, + tool: false, + system: false, + skip_previous_eos: false, + reasoning: None, + reasoning_complete: false, + content, + }; + let marker = crate::engine::vision_data_marker(&format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(image) + )); + let mut messages = vec![user(format!("Describe this image{marker}"))]; + let run = |executor: &mut QwenExecutor, messages: &[ChatTurn], label: &str| { + let (output, complete) = executor + .generate( + messages, + &settings.turn, + &cancelled, + &metrics, + |p, _, _| { + if p % 128 == 0 { + eprintln!("{label}: position {p}"); + } + }, + |_, _| {}, + ) + .unwrap(); + assert!(complete); + assert_eq!(output.finish_reason, "stop"); + println!( + "{}", + serde_json::json!({"event":"vision_lifecycle","phase":label,"content":output.message.content,"reasoning":output.message.reasoning,"prompt":output.prompt_tokens,"cached":output.cached_tokens}) + ); + output + }; + let first = run(&mut executor, &messages, "cold"); + assert_eq!(first.prompt_tokens, 1069); + assert_eq!(first.cached_tokens, 0); + let content = first.message.content.to_lowercase(); + assert!(content.contains("llama") || content.contains("alpaca")); + assert!(content.contains("mountain")); + assert!( + executor.vision_tower.is_none(), + "default policy must release vision weights" + ); + assert_eq!(executor.live.current_images.len(), 1); + let identities = executor.live.current_images.clone(); + let repeat = run(&mut executor, &messages, "same_image"); + assert_eq!(repeat.cached_tokens, first.prompt_tokens); + assert_eq!(repeat.message.content, first.message.content); + messages.push(repeat.message); + let tag = crate::engine::conversation_tag("", ReasoningMode::Low, &messages); + let checkpoint = directory.join("vision.kv"); + executor + .save_checkpoint(&checkpoint, tag, &mut |_| {}) + .unwrap(); + executor.reset().unwrap(); + assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap()); + assert_eq!(executor.live.current_images, identities); + messages.push(user("Describe this image".into())); + let followup = run(&mut executor, &messages, "restored_followup"); + assert!(followup.cached_tokens >= first.prompt_tokens); + let text = followup.message.content.to_lowercase(); + assert!(text.contains("llama") || text.contains("alpaca")); + // Same text and same image-pad span, different pixels must miss the old KV. + let blue = image::RgbImage::from_pixel(1024, 1024, image::Rgb([0, 0, 255])); + let mut encoded = std::io::Cursor::new(Vec::new()); + blue.write_to(&mut encoded, image::ImageFormat::Png) + .unwrap(); + let marker = crate::engine::vision_data_marker(&format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(encoded.get_ref()) + )); + let changed = run( + &mut executor, + &[user(format!("Describe this image{marker}"))], + "different_image", + ); + assert_eq!(changed.cached_tokens, 0); + assert!(changed.message.content.to_lowercase().contains("blue")); + assert_ne!(changed.message.content, first.message.content); + std::fs::remove_dir_all(directory).unwrap(); +} diff --git a/src/engine/metal/qwen_mtplx/request.rs b/src/engine/metal/qwen_mtplx/request.rs index b3408d9..762b702 100644 --- a/src/engine/metal/qwen_mtplx/request.rs +++ b/src/engine/metal/qwen_mtplx/request.rs @@ -95,6 +95,8 @@ pub(in crate::engine) struct LiveSession { pub(super) current: Option, pub(super) prefix: Option, pub(super) tag: [u8; 32], + pub(super) current_images: Vec, + pub(super) prefix_images: Vec, } impl LiveSession { @@ -129,7 +131,10 @@ impl LiveSession { emit: impl FnMut(bool, String), observe: impl FnMut(TurnProgress) -> Result<(), String>, ) -> Result { - let prompt = self.prompt(model, messages, settings); + let (prompt, images) = execution + .vision_prompt + .clone() + .unwrap_or_else(|| (self.prompt(model, messages, settings), Vec::new())); let context = u32::try_from(settings.context_tokens) .ok() .filter(|&n| n > 0) @@ -154,11 +159,18 @@ impl LiveSession { o.stops = &stop_ids; } } - let saved = [&self.current, &self.prefix] - .into_iter() - .flatten() - .filter(|s| prompt.starts_with(&s.token_ids)) - .max_by_key(|s| s.token_ids.len()); + let saved = [ + (&self.current, &self.current_images), + (&self.prefix, &self.prefix_images), + ] + .into_iter() + .filter_map(|(s, old)| { + s.as_ref().filter(|s| { + crate::engine::qwen_vision::same_prefix(old, &images, s.token_ids.len()) + }) + }) + .filter(|s| prompt.starts_with(&s.token_ids)) + .max_by_key(|s| s.token_ids.len()); let mut captured_prefix = None; let result = execution.generate_text( &prompt, @@ -182,6 +194,7 @@ impl LiveSession { // errors. Never tag those partial turns as a completed conversation. if let Some(prefix) = captured_prefix { self.prefix = Some(prefix); + self.prefix_images = images.clone(); } let mut result = result?; if let Some(snapshot) = result @@ -194,6 +207,7 @@ impl LiveSession { self.tag = conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed); self.current = Some(snapshot); + self.current_images = images; } Ok(result) } diff --git a/src/engine/metal/qwen_mtplx/session_checkpoint.rs b/src/engine/metal/qwen_mtplx/session_checkpoint.rs index bafacc4..b83a04f 100644 --- a/src/engine/metal/qwen_mtplx/session_checkpoint.rs +++ b/src/engine/metal/qwen_mtplx/session_checkpoint.rs @@ -27,6 +27,8 @@ pub(super) struct Identity { struct Manifest { identity: Identity, tag: [u8; 32], + #[serde(default)] + images: Vec, tokens: Vec, lazy_kv: bool, payload: serde_json::Value, @@ -111,6 +113,7 @@ impl LiveSession { let manifest = serde_json::to_vec(&Manifest { identity, tag, + images: self.current_images.clone(), tokens: saved.token_ids.clone(), lazy_kv: saved.lazy_kv, payload, @@ -230,6 +233,8 @@ impl LiveSession { self.current = Some(saved); self.prefix = None; self.tag = manifest.tag; + self.current_images = manifest.images; + self.prefix_images.clear(); Ok(true) } } @@ -258,6 +263,8 @@ fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() { .unwrap(); let state = Some(State::Recurrent(vec![Some(leaf.clone()), None])); let mut live = LiveSession { + current_images: Vec::new(), + prefix_images: Vec::new(), current: Some(SessionSnapshot { token_ids: vec![1, 2, 3, 4], logits: leaf.clone(), diff --git a/src/engine/metal/qwen_mtplx/unary.rs b/src/engine/metal/qwen_mtplx/unary.rs index ff17b7f..e5cee1a 100644 --- a/src/engine/metal/qwen_mtplx/unary.rs +++ b/src/engine/metal/qwen_mtplx/unary.rs @@ -15,12 +15,12 @@ pub(super) enum Unary { #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))] Log, Sigmoid, - #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))] Tanh, Sqrt, Rsqrt, Sin, Cos, + Erf, } pub(super) fn unary(input: &Array, op: Unary, stream: Stream) -> Result { diff --git a/src/engine/metal/qwen_mtplx/vision.rs b/src/engine/metal/qwen_mtplx/vision.rs new file mode 100644 index 0000000..6261739 --- /dev/null +++ b/src/engine/metal/qwen_mtplx/vision.rs @@ -0,0 +1,544 @@ +//! Qwen3-VL tower, following MTPLX e652d55 vision/qwen3_vl_tower.py. +use super::array::{Array, Dtype}; +use super::binary::{Binary, binary}; +use super::stream::{Stream, Streams}; +use super::{ + indexing, ops, + unary::{Unary, unary}, + vision_kernels::{self, Kernel}, +}; +use std::path::Path; + +pub(super) fn f32_array(shape: &[i32], data: &[f32]) -> Result { + Array::new( + shape, + Dtype::F32, + super::scalar_buffer( + &data + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect::>(), + )?, + ) +} +fn i32_array(shape: &[i32], data: &[i32]) -> Result { + Array::new( + shape, + Dtype::I32, + super::scalar_buffer( + &data + .iter() + .flat_map(|x| x.to_le_bytes()) + .collect::>(), + )?, + ) +} +fn scalar(x: f32, dtype: Dtype, s: Stream) -> Result { + ops::astype(&f32_array(&[], &[x])?, dtype, false, s) +} +fn add(a: &Array, b: &Array, s: Stream) -> Result { + binary(a, b, Binary::Add, s) +} +fn mul(a: &Array, b: &Array, s: Stream) -> Result { + binary(a, b, Binary::Multiply, s) +} +fn pad_last(a: &Array, extra: i32, s: Stream) -> Result { + let mut shape = a.layout().shape().to_vec(); + *shape.last_mut().unwrap() = extra; + ops::concatenate( + &[a.clone(), ops::zeros(&shape, a.layout().dtype(), s)?], + -1, + s, + ) +} +fn contiguous(a: &Array, s: Stream) -> Result { + ops::contiguous(a, false, s) +} + +pub(super) struct Tower { + weights: super::weights::Parameters, +} +impl Tower { + pub(super) fn load(root: &Path, streams: &Streams, s: Stream) -> Result { + let weights = super::load::safetensors(&root.join("model-vision.safetensors"), streams)?; + if weights.len() != 333 || weights.keys().any(|k| !k.starts_with("vision_tower.")) { + return Err("Qwen vision tower must contain its 333 pinned tensors".into()); + } + let mut tower = Self { weights }; + let key = "vision_tower.patch_embed.proj.weight"; + let w = tower + .weights + .get(key) + .ok_or("missing Qwen vision patch weight")?; + let w = match w.layout().shape() { + [1152, 3, 2, 16, 16] => ops::transpose(w, &[0, 2, 3, 4, 1], s)?, + [1152, 2, 16, 16, 3] => w.clone(), + _ => return Err("unexpected Qwen vision patch layout".into()), + }; + tower + .weights + .insert(key.into(), contiguous(&pad_last(&w, 13, s)?, s)?); + ops::evaluate( + streams, + &tower.weights.values().cloned().collect::>(), + s, + false, + )?; + Ok(tower) + } + fn weight(&self, name: &str) -> Result<&Array, String> { + self.weights + .get(&format!("vision_tower.{name}")) + .ok_or_else(|| format!("missing Qwen vision tensor {name}")) + } + fn linear(&self, x: &Array, name: &str, s: Stream) -> Result { + let w = self.weight(&format!("{name}.weight"))?; + let b = self.weight(&format!("{name}.bias"))?; + let rows = x.layout().dim(0)?; + if rows < 16 + || x.layout().dtype() != Dtype::BF16 + || x.layout().shape().len() != 2 + || x.layout().dim(1)? != w.layout().dim(1)? + { + return Err("invalid Qwen vision projection layout".into()); + } + vision_kernels::call( + Kernel::Linear, + &[contiguous(x, s)?, ops::transpose(w, &[1, 0], s)?, b.clone()], + &[rows, w.layout().dim(0)?], + s, + ) + } + fn norm(&self, x: &Array, name: &str, s: Stream) -> Result { + vision_kernels::call( + Kernel::LayerNorm, + &[ + contiguous(x, s)?, + self.weight(&format!("{name}.weight"))?.clone(), + self.weight(&format!("{name}.bias"))?.clone(), + ], + x.layout().shape(), + s, + ) + } + fn positions(&self, h: i32, w: i32, s: Stream) -> Result<(Array, Array), String> { + let mut indices = [Vec::new(), Vec::new(), Vec::new(), Vec::new()]; + let mut weights = [Vec::new(), Vec::new(), Vec::new(), Vec::new()]; + let mut rotary = Vec::new(); + for br in 0..h / 2 { + for bc in 0..w / 2 { + for ir in 0..2 { + for ic in 0..2 { + let (r, c) = (br * 2 + ir, bc * 2 + ic); + // mx.linspace(0,47,n) uses a double step cast to F32. + let y = (r as f64 * 47.0 / f64::from(h - 1)) as f32; + let x = (c as f64 * 47.0 / f64::from(w - 1)) as f32; + let (yf, xf) = (y as i32, x as i32); + let (yc, xc) = ((yf + 1).min(47), (xf + 1).min(47)); + let (dy, dx) = (y - yf as f32, x - xf as f32); + for (i, (idx, v)) in [ + (yf * 48 + xf, (1. - dy) * (1. - dx)), + (yf * 48 + xc, (1. - dy) * dx), + (yc * 48 + xf, dy * (1. - dx)), + (yc * 48 + xc, dy * dx), + ] + .into_iter() + .enumerate() + { + indices[i].push(idx); + weights[i].push(v) + } + rotary.extend([r as f32, c as f32]); + } + } + } + } + let mut parts = Vec::new(); + for i in 0..4 { + let p = indexing::take_rows( + self.weight("pos_embed.weight")?, + &i32_array(&[h * w], &indices[i])?, + s, + )?; + let weights = + ops::astype(&f32_array(&[h * w, 1], &weights[i])?, Dtype::BF16, false, s)?; + parts.push(mul(&p, &weights, s)?); + } + let position = add( + &add(&add(&parts[0], &parts[1], s)?, &parts[2], s)?, + &parts[3], + s, + )?; + let exponent = binary( + &f32_array(&[18], &(0..18).map(|i| (i * 2) as f32).collect::>())?, + &f32_array(&[], &[36.])?, + Binary::Divide, + s, + )?; + let powers = binary(&f32_array(&[], &[10000.])?, &exponent, Binary::Power, s)?; + let inv = binary(&f32_array(&[], &[1.])?, &powers, Binary::Divide, s)?; + let rotary = ops::reshape( + &mul(&f32_array(&[h * w, 2, 1], &rotary)?, &inv, s)?, + &[h * w, 36], + s, + )?; + Ok((position, rotary)) + } + fn attention(&self, x: &Array, name: &str, rotary: &Array, s: Stream) -> Result { + let n = x.layout().dim(0)?; + let qkv = self.linear(x, &format!("{name}.qkv"), s)?; + let qkv = ops::transpose(&ops::reshape(&qkv, &[n, 3, 16, 72], s)?, &[1, 0, 2, 3], s)?; + let mut parts = ops::split(&qkv, &[1, 2], 0, s)?; + let angles = ops::reshape(rotary, &[1, n, 1, 36], s)?; + let cos = unary(&angles, Unary::Cos, s)?; + let sin = unary(&angles, Unary::Sin, s)?; + let cos = ops::concatenate(&[cos.clone(), cos], -1, s)?; + let sin = ops::concatenate(&[sin.clone(), sin], -1, s)?; + for p in &mut parts[..2] { + let halves = ops::split(p, &[36], -1, s)?; + let rotated = ops::concatenate( + &[unary(&halves[1], Unary::Negative, s)?, halves[0].clone()], + -1, + s, + )?; + *p = ops::astype( + &add(&mul(p, &cos, s)?, &mul(&rotated, &sin, s)?, s)?, + Dtype::BF16, + false, + s, + )?; + } + let inputs = parts + .iter() + .map(|p| contiguous(&pad_last(&ops::transpose(p, &[0, 2, 1, 3], s)?, 8, s)?, s)) + .collect::, _>>()?; + let output = vision_kernels::call(Kernel::Attention, &inputs, &[1, 16, n, 80], s)?; + let output = ops::slice(&output, &[0, 0, 0, 0], &[1, 16, n, 72], &[1; 4], s)?; + let output = ops::reshape(&ops::transpose(&output, &[0, 2, 1, 3], s)?, &[n, 1152], s)?; + self.linear(&output, &format!("{name}.proj"), s) + } + pub(super) fn encode( + &self, + pixels: &Array, + h: i32, + w: i32, + streams: &Streams, + s: Stream, + mut observe: impl FnMut(&str, &Array) -> Result<(), String>, + ) -> Result { + if h < 2 || w < 2 || h % 2 != 0 || w % 2 != 0 || pixels.layout().shape() != [h * w, 1536] { + return Err("invalid Qwen vision patch grid".into()); + } + let pixels = ops::astype(pixels, Dtype::BF16, false, s)?; + let patches = ops::transpose( + &ops::reshape(&pixels, &[h * w, 3, 2, 16, 16], s)?, + &[0, 2, 3, 4, 1], + s, + )?; + let patches = contiguous(&pad_last(&patches, 13, s)?, s)?; + let mut x = vision_kernels::call( + Kernel::PatchConv, + &[patches, self.weight("patch_embed.proj.weight")?.clone()], + &[h * w, 1152], + s, + )?; + x = add(&x, self.weight("patch_embed.proj.bias")?, s)?; + observe("patch", &x)?; + let (position, rotary) = self.positions(h, w, s)?; + observe("position", &position)?; + observe("rotary", &rotary)?; + x = add(&x, &position, s)?; + for i in 0..27 { + observe("progress", &x)?; + let name = format!("blocks.{i}"); + let a = self.attention( + &self.norm(&x, &format!("{name}.norm1"), s)?, + &format!("{name}.attn"), + &rotary, + s, + )?; + x = add(&x, &a, s)?; + let mlp = self.linear( + &self.norm(&x, &format!("{name}.norm2"), s)?, + &format!("{name}.mlp.linear_fc1"), + s, + )?; + let mlp = gelu(&mlp, true, s)?; + x = add( + &x, + &self.linear(&mlp, &format!("{name}.mlp.linear_fc2"), s)?, + s, + )?; + ops::evaluate(streams, std::slice::from_ref(&x), s, false)?; + if i == 0 || i == 26 { + observe(&format!("block{i}"), &x)?; + } + } + let x = ops::reshape(&self.norm(&x, "merger.norm", s)?, &[h * w / 4, 4608], s)?; + let x = gelu(&self.linear(&x, "merger.linear_fc1", s)?, false, s)?; + let x = self.linear(&x, "merger.linear_fc2", s)?; + ops::evaluate(streams, std::slice::from_ref(&x), s, false)?; + observe("embeddings", &x)?; + Ok(x) + } +} +fn gelu(x: &Array, approx: bool, s: Stream) -> Result { + let inputs = std::slice::from_ref(x); + let graph = super::compiled::Replay::trace_compiled(inputs, |args| { + let x = &args[0]; + let dtype = x.layout().dtype(); + let c = |v| scalar(v, dtype, s); + let result = if approx { + let cubic = binary(x, &c(3.)?, Binary::Power, s)?; + let inner = mul( + &c((2_f64 / std::f64::consts::PI).sqrt() as f32)?, + &add(x, &mul(&c(0.044715)?, &cubic, s)?, s)?, + s, + )?; + mul( + &mul(&c(0.5)?, x, s)?, + &add(&c(1.)?, &unary(&inner, Unary::Tanh, s)?, s)?, + s, + )? + } else { + let erf = unary( + &binary(x, &c(2_f64.sqrt() as f32)?, Binary::Divide, s)?, + Unary::Erf, + s, + )?; + binary( + &mul(x, &add(&c(1.)?, &erf, s)?, s)?, + &c(2.)?, + Binary::Divide, + s, + )? + }; + Ok(vec![result]) + })?; + Ok(graph.run(inputs)?.remove(0)) +} + +#[test] +#[ignore = "requires Metal and pinned Qwen vision artifacts/reference values"] +fn qwen_vision_tower_matches_mtplx_image() { + super::configure_sources().unwrap(); + let _context = super::super::gpu::Context::open_qwen(0).unwrap(); + let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap())); + let s = streams.default_stream(super::stream::Device::Gpu).unwrap(); + let root = std::path::PathBuf::from(std::env::var("DS4_QWEN38_ARTIFACTS").unwrap()); + let dir = std::path::PathBuf::from(std::env::var("DS4_QWEN_VISION_REFERENCE").unwrap()); + let read = |name: &str| { + std::fs::read(dir.join(format!("{name}.f32"))) + .unwrap() + .chunks_exact(4) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect::>() + }; + let grid: Vec = + serde_json::from_slice(&std::fs::read(dir.join("grid.json")).unwrap()).unwrap(); + let (h, w) = (grid[1], grid[2]); + let pixels = f32_array(&[h * w, 1536], &read("pixels")).unwrap(); + let tower = Tower::load(&root, &streams, s).unwrap(); + tower + .encode(&pixels, h, w, &streams, s, |name, a| { + if name == "progress" { + return Ok(()); + } + let a = ops::astype(a, Dtype::F32, false, s)?; + ops::evaluate(&streams, std::slice::from_ref(&a), s, false)?; + let got = ops::read_f32_as_f64(&a, &streams, s)?; + let expected = read(name); + let max = got + .iter() + .zip(&expected) + .map(|(a, b)| (a - f64::from(*b)).abs()) + .fold(0., f64::max); + let different = got + .iter() + .zip(&expected) + .filter(|(a, b)| **a != f64::from(**b)) + .count(); + eprintln!( + "{name}: values={} different={different} max_error={max}", + got.len() + ); + assert_eq!(got.len(), expected.len()); + assert_eq!(different, 0, "{name} differs from MTPLX (max error {max})"); + Ok(()) + }) + .unwrap(); +} + +#[derive(Clone)] +pub(super) struct Input { + pub(super) table: Array, + pub(super) delta: i32, + pub(super) images: Vec<(std::ops::Range, Array)>, +} +impl Input { + pub(super) fn embeddings( + &self, + layer: super::mlp::Linear<'_>, + ids: &Array, + start: usize, + s: Stream, + ) -> Result, String> { + let n = ids.layout().dim(1)? as usize; + let end = start.checked_add(n).ok_or("Qwen image window overflow")?; + if !self + .images + .iter() + .any(|(r, _)| start < r.end && end > r.start) + { + return Ok(None); + } + let mut embedded = super::model::embedding(layer, ids, s)?; + for (span, rows) in &self.images { + let a = start.max(span.start); + let b = end.min(span.end); + if a >= b { + continue; + } + let rows = ops::slice( + rows, + &[(a - span.start) as i32, 0], + &[(b - span.start) as i32, 2560], + &[1, 1], + s, + )?; + let rows = ops::reshape(&rows, &[1, (b - a) as i32, 2560], s)?; + embedded = ops::slice_update( + &embedded, + &rows, + &[0, (a - start) as i32, 0], + &[1, (b - start) as i32, 2560], + &[1; 3], + s, + )?; + } + Ok(Some(embedded)) + } +} +pub(super) fn cache_position(cache: Option<&[Option]>) -> usize { + cache + .into_iter() + .flatten() + .flatten() + .find_map(|c| match &c.attention { + super::decoder::AttentionCache::Qsa(c) => Some(c.kv.offset as usize), + _ => None, + }) + .unwrap_or(0) +} + +pub(super) fn prepare( + cancelled: &std::sync::atomic::AtomicBool, + model: &crate::engine::qwen::QwenMetadata, + messages: &[crate::engine::ChatTurn], + settings: &crate::settings::TurnSettings, + tower: &Tower, + streams: &Streams, + s: Stream, +) -> Result<(Vec, Input, Vec), String> { + use crate::engine::{VISION_DATA_END, VISION_DATA_START}; + use base64::Engine as _; + use sha2::{Digest, Sha256}; + let check = || { + if cancelled.load(std::sync::atomic::Ordering::Relaxed) { + Err("Qwen vision input cancelled".to_owned()) + } else { + Ok(()) + } + }; + check()?; + let mut rendered = messages.to_vec(); + let mut decoded = Vec::new(); + let mut total = 0; + for message in &mut rendered { + if !message.content.contains(VISION_DATA_START) { + continue; + } + if !message.user && !message.tool { + return Err("vision input is allowed only in user or tool messages".into()); + } + let mut content = String::new(); + let mut rest = message.content.as_str(); + while let Some(start) = rest.find(VISION_DATA_START) { + content.push_str(&rest[..start]); + let image = &rest[start + VISION_DATA_START.len()..]; + let end = image + .find(VISION_DATA_END) + .ok_or("unterminated image input")?; + let uri = &image[..end]; + let payload = uri + .strip_prefix("data:image/png;base64,") + .or_else(|| uri.strip_prefix("data:image/jpeg;base64,")) + .ok_or("image input must be an inline PNG or JPEG data URI")?; + if decoded.len() >= 16 || payload.len() > 90 * 1024 * 1024 { + return Err("image request exceeds limits".into()); + } + check()?; + let raw = base64::engine::general_purpose::STANDARD + .decode(payload) + .map_err(|_| "invalid image base64")?; + total += raw.len(); + if total > 64 * 1024 * 1024 { + return Err("image inputs exceed the 64 MiB request limit".into()); + } + let digest: [u8; 32] = Sha256::digest(&raw).into(); + let image = crate::engine::qwen_vision::decode(&raw)?; + decoded.push((image, digest)); + content.push_str("<|vision_start|><|image_pad|><|vision_end|>"); + rest = &rest[start + VISION_DATA_START.len() + end + VISION_DATA_END.len()..]; + } + content.push_str(rest); + message.content = content; + } + let plain = + model.render_conversation(&settings.system_prompt, &rendered, settings.reasoning_mode); + let mut ids = Vec::new(); + let mut spans = Vec::new(); + let mut index = 0; + for id in plain { + if id == 248056 { + let (image, digest) = decoded + .get(index) + .ok_or("image pad has no supplied image")?; + let start = ids.len(); + ids.extend(std::iter::repeat_n(id, (image.h * image.w / 4) as usize)); + spans.push(crate::engine::qwen_vision::Identity { + start, + end: ids.len(), + digest: *digest, + }); + index += 1; + } else { + ids.push(id) + } + if ids.len() >= settings.context_tokens as usize { + return Err("image prompt exceeds configured context".into()); + } + } + if index != decoded.len() { + return Err("rendered Qwen prompt lost an image placeholder".into()); + } + let grids = decoded.iter().map(|(i, _)| (i.h, i.w)).collect::>(); + let (table, delta) = crate::engine::qwen_vision::positions(&ids, &grids)?; + let table = i32_array(&[3, ids.len() as i32], &table)?; + let mut images = Vec::new(); + for ((image, _), span) in decoded.iter().zip(&spans) { + check()?; + let pixels = f32_array(&[image.h * image.w, 1536], &image.pixels)?; + let rows = tower.encode(&pixels, image.h, image.w, streams, s, |_, _| check())?; + images.push((span.start..span.end, rows)); + } + Ok(( + ids, + Input { + table, + delta, + images, + }, + spans, + )) +} diff --git a/src/engine/metal/qwen_mtplx/vision_kernels.rs b/src/engine/metal/qwen_mtplx/vision_kernels.rs new file mode 100644 index 0000000..32ff886 --- /dev/null +++ b/src/engine/metal/qwen_mtplx/vision_kernels.rs @@ -0,0 +1,191 @@ +//! Vision-only dispatches from pinned MLX 0.32.2 normalization, conv and SDPA. +//! Rust supplies the same layouts and parameters to the existing Metal library. +use super::super::gpu::{Buffer, MetalConstant}; +use super::array::{Array, Dtype, Layout}; +use super::{bytes, dispatch_geometry, dispatch_specialized, ops, stream::Stream, tensor}; + +#[derive(Clone, Copy, PartialEq)] +pub(super) enum Kernel { + LayerNorm, + PatchConv, + Attention, + Linear, +} + +pub(super) fn call( + kernel: Kernel, + inputs: &[Array], + shape: &[i32], + stream: Stream, +) -> Result { + Array::make_operation_with_inputs( + stream, + inputs, + Layout::new(shape, Dtype::BF16)?, + ops::Operation::Vision(kernel), + ) +} + +pub(super) fn evaluate(kernel: Kernel, inputs: &[Array], out: &Array) -> Result<(), String> { + if kernel == Kernel::Linear { + return super::dense::evaluate_bias(&inputs[0], &inputs[1], Some(&inputs[2]), out); + } + let nbytes = out.layout().nbytes() as u64; + out.set_data(Buffer::mtplx_bytes(nbytes)?)?; + let buffers = inputs.iter().map(Array::buffer).collect::>(); + let ob = out.buffer(); + let mut bindings = inputs + .iter() + .zip(&buffers) + .enumerate() + .map(|(i, (a, b))| { + tensor(i as u32, b) + .at_byte_offset(a.offset()) + .input(a.data_size()) + }) + .collect::>(); + bindings.push(tensor(inputs.len() as u32, &ob).output(out.data_size())); + match kernel { + Kernel::Linear => { + super::dense::evaluate_bias(&inputs[0], &inputs[1], Some(&inputs[2]), out) + } + Kernel::LayerNorm => { + let axis = out.layout().dim(-1)? as u32; + let rows = out.layout().size() as u32 / axis; + let threads = axis.div_ceil(8).div_ceil(32) * 32; + let eps = 1e-6_f32; + let stride = 1_u32; + bindings.extend([ + bytes(4, &eps), + bytes(5, &axis), + bytes(6, &stride), + bytes(7, &stride), + ]); + dispatch_geometry( + "layer_normbfloat16", + &bindings, + &[], + [rows * threads, 1, 1], + [threads, 1, 1], + true, + ) + } + Kernel::PatchConv => { + // Inputs are contiguous [patches,2,16,16,16], weights + // [1152,2,16,16,16]. Original Conv3d pads RGB channels to 16. + #[repr(C)] + struct Params { + n: i32, + c: i32, + o: i32, + input: [i32; 3], + weight: [i32; 3], + output: [i32; 3], + stride: [i32; 3], + pad: [i32; 3], + kdil: [i32; 3], + idil: [i32; 3], + is: [i64; 5], + ws: [i64; 5], + os: [i64; 5], + groups: i32, + flip: u8, + padding: [u8; 3], + } + let n = out.layout().dim(0)?; + let params = Params { + n, + c: 16, + o: 1152, + input: [2, 16, 16], + weight: [2, 16, 16], + output: [1, 1, 1], + stride: [2, 16, 16], + pad: [0; 3], + kdil: [1; 3], + idil: [1; 3], + is: [8192, 4096, 256, 16, 1], + ws: [8192, 4096, 256, 16, 1], + os: [1152, 1152, 1152, 1152, 1], + groups: 1, + flip: 0, + padding: [0; 3], + }; + let gemm = [n, 1152, 8192, 512, 16, 16, 16, -8160, 18, (n + 31) / 32, 0]; + bindings.extend([bytes(3, ¶ms), bytes(4, &gemm)]); + dispatch_specialized( + "implicit_gemm_conv_3d_bfloat16_bm32_bn64_bk16_wm2_wn2_filter_s", + &bindings, + &[], + [18, ((n + 31) / 32) as u32, 1], + [32, 2, 2], + ) + } + Kernel::Attention => { + #[repr(C)] + struct Params { + b: i32, + h: i32, + d: i32, + ql: i32, + kl: i32, + gqa: i32, + scale: f32, + nq: i32, + nk: i32, + aq: i32, + ak: i32, + rq: i32, + rk: i32, + off: i32, + strides: [[i64; 3]; 4], + } + let shape = out.layout().shape().to_vec(); + let n = shape[2]; + let mut strides = [[0; 3]; 4]; + for (s, a) in strides + .iter_mut() + .zip(inputs.iter().chain(std::iter::once(out))) + { + s.copy_from_slice(&a.layout().strides()[..3]); + } + let params = Params { + b: 1, + h: 16, + d: 80, + ql: n, + kl: n, + gqa: 1, + scale: (72_f64.powf(-0.5)) as f32, + nq: ((n + 31) / 32), + nk: ((n + 31) / 32), + aq: n / 32, + ak: n / 32, + rq: n % 32, + rk: n % 32, + off: 0, + strides, + }; + let constants = [ + (200, n % 32 == 0), + (201, n % 32 == 0), + (300, false), + (301, false), + (302, false), + ] + .map(|(index, v)| MetalConstant { + index, + value: u32::from(v), + kind: 0, + }); + bindings.push(bytes(4, ¶ms)); + dispatch_specialized( + "steel_attention_bfloat16_bq32_bk32_bd80_wm4_wn1_maskbfloat16", + &bindings, + &constants, + [((n + 31) / 32) as u32, 16, 1], + [32, 4, 1], + ) + } + } +} diff --git a/src/engine/metal/qwen_mtplx/weights.rs b/src/engine/metal/qwen_mtplx/weights.rs index 4236c23..1bd0541 100644 --- a/src/engine/metal/qwen_mtplx/weights.rs +++ b/src/engine/metal/qwen_mtplx/weights.rs @@ -34,7 +34,10 @@ pub(super) fn load_shards( let entry = entry.map_err(|e| e.to_string())?; let name = entry.file_name(); let Some(name) = name.to_str() else { continue }; - if name.starts_with("model") && name.ends_with(".safetensors") { + if name.starts_with("model") + && name.ends_with(".safetensors") + && name != "model-vision.safetensors" + { parameters.extend(super::load::safetensors(&entry.path(), streams)?); files += 1; } diff --git a/src/engine/qwen.rs b/src/engine/qwen.rs index bdc6390..8d1626c 100644 --- a/src/engine/qwen.rs +++ b/src/engine/qwen.rs @@ -200,7 +200,9 @@ impl QwenMetadata { ) .ok_or("Qwen artifact size overflows")?; } - Ok(Self::new(tokenizer, mapped_bytes, tensor_count)) + let mut metadata = Self::new(tokenizer, mapped_bytes, tensor_count); + metadata.summary.vision_loaded = root.join("model-vision.safetensors").is_file(); + Ok(metadata) } } diff --git a/src/engine/qwen_vision.rs b/src/engine/qwen_vision.rs new file mode 100644 index 0000000..8123480 --- /dev/null +++ b/src/engine/qwen_vision.rs @@ -0,0 +1,284 @@ +//! PNG/JPEG preprocessing and image positions from MTPLX's Qwen3-VL path. +use image::{DynamicImage, ImageDecoder, ImageReader, RgbImage}; +use std::io::Cursor; + +pub(super) struct Image { + pub(super) pixels: Vec, + pub(super) h: i32, + pub(super) w: i32, +} + +pub(super) fn decode(bytes: &[u8]) -> Result { + if bytes.is_empty() || bytes.len() > 50 * 1024 * 1024 { + return Err("Qwen image exceeds the 50 MiB encoded limit".into()); + } + let reader = ImageReader::new(Cursor::new(bytes)) + .with_guessed_format() + .map_err(|e| e.to_string())?; + if !matches!( + reader.format(), + Some(image::ImageFormat::Png | image::ImageFormat::Jpeg) + ) { + return Err("Qwen images must be PNG or JPEG".into()); + } + let mut decoder = reader.into_decoder().map_err(|e| e.to_string())?; + let (w, h) = decoder.dimensions(); + if w == 0 || h == 0 || w > 8000 || h > 8000 { + return Err("Qwen images must have dimensions between 1 and 8000".into()); + } + let orientation = decoder.orientation().map_err(|e| e.to_string())?; + let mut image = DynamicImage::from_decoder(decoder).map_err(|e| e.to_string())?; + image.apply_orientation(orientation); + let mut image = image.into_rgb8(); + let (w, h) = image.dimensions(); + let (rh, rw) = smart_resize(h, w)?; + if (w, h) != (rw, rh) { + image = resize(&image, rw, rh); + } + let (h, w) = ((rh / 16) as i32, (rw / 16) as i32); + let mut pixels = Vec::with_capacity(h as usize * w as usize * 1536); + for br in 0..h / 2 { + for bc in 0..w / 2 { + for ir in 0..2 { + for ic in 0..2 { + for c in 0..3 { + for _t in 0..2 { + for y in 0..16 { + for x in 0..16 { + let v = image.get_pixel( + ((bc * 2 + ic) * 16 + x) as u32, + ((br * 2 + ir) * 16 + y) as u32, + )[c]; + pixels.push((f32::from(v) * (1_f32 / 255.) - 0.5) / 0.5); + } + } + } + } + } + } + } + } + Ok(Image { pixels, h, w }) +} +fn smart_resize(h: u32, w: u32) -> Result<(u32, u32), String> { + if f64::from(h.max(w)) / f64::from(h.min(w)) > 200. { + return Err("Qwen image aspect ratio exceeds 200".into()); + } + let (hf, wf) = (f64::from(h), f64::from(w)); + let (mut rh, mut rw) = ( + (hf / 32.).round_ties_even() as u32 * 32, + (wf / 32.).round_ties_even() as u32 * 32, + ); + if u64::from(rh) * u64::from(rw) > 16777216 { + let beta = (hf * wf / 16777216.).sqrt(); + rh = ((hf / beta / 32.).floor() as u32 * 32).max(32); + rw = ((wf / beta / 32.).floor() as u32 * 32).max(32); + } else if u64::from(rh) * u64::from(rw) < 65536 { + let beta = (65536. / (hf * wf)).sqrt(); + rh = (hf * beta / 32.).ceil() as u32 * 32; + rw = (wf * beta / 32.).ceil() as u32 * 32; + } + Ok((rh, rw)) +} +// Pillow Resample.c: bicubic a=-0.5, normalized 22-bit coefficients, +// horizontal then vertical passes with rounding/clipping after each pass. +fn coefficients(input: u32, output: u32) -> Vec<(usize, Vec)> { + let scale = f64::from(input) / f64::from(output); + let filter_scale = scale.max(1.); + let support = 2. * filter_scale; + (0..output) + .map(|i| { + let center = (f64::from(i) + 0.5) * scale; + let first = ((center - support + 0.5) as i64).max(0) as usize; + let last = ((center + support + 0.5) as i64).min(i64::from(input)) as usize; + let ws = (first..last) + .map(|j| { + let x = ((j as f64 - center + 0.5) / filter_scale).abs(); + if x < 1. { + ((1.5 * x - 2.5) * x) * x + 1. + } else if x < 2. { + (((x - 5.) * x + 8.) * x - 4.) * (-0.5) + } else { + 0. + } + }) + .collect::>(); + let sum = ws.iter().sum::(); + ( + first, + ws.iter() + .map(|w| (w / sum * 4194304.).round() as i32) + .collect(), + ) + }) + .collect() +} +fn resize(input: &RgbImage, w: u32, h: u32) -> RgbImage { + let horizontal = if w == input.width() { + input.clone() + } else { + let coeff = coefficients(input.width(), w); + RgbImage::from_fn(w, input.height(), |x, y| { + let (start, ws) = &coeff[x as usize]; + image::Rgb(std::array::from_fn(|c| { + let sum = ws.iter().enumerate().fold(1_i64 << 21, |sum, (i, k)| { + sum + i64::from(*k) * i64::from(input.get_pixel((*start + i) as u32, y)[c]) + }); + (sum >> 22).clamp(0, 255) as u8 + })) + }) + }; + if h == horizontal.height() { + return horizontal; + } + let coeff = coefficients(horizontal.height(), h); + RgbImage::from_fn(w, h, |x, y| { + let (start, ws) = &coeff[y as usize]; + image::Rgb(std::array::from_fn(|c| { + let sum = ws.iter().enumerate().fold(1_i64 << 21, |sum, (i, k)| { + sum + i64::from(*k) * i64::from(horizontal.get_pixel(x, (*start + i) as u32)[c]) + }); + (sum >> 22).clamp(0, 255) as u8 + })) + }) +} + +pub(super) fn positions(ids: &[i32], grids: &[(i32, i32)]) -> Result<(Vec, i32), String> { + let n = ids.len(); + let mut axes = vec![0; n * 3]; + let (mut start, mut next) = (0, 0_i32); + for &(h, w) in grids { + let h = h / 2; + let w = w / 2; + if h <= 0 || w <= 0 { + return Err("invalid Qwen image grid".into()); + } + let end = start + + ids[start..] + .iter() + .position(|&t| t == 248056) + .ok_or("missing Qwen image pad")?; + for i in start..end { + for axis in 0..3 { + axes[axis * n + i] = next + (i - start) as i32; + } + } + next += (end - start) as i32; + let count = (h * w) as usize; + if end + count > n || ids[end..end + count].iter().any(|&t| t != 248056) { + return Err("Qwen image pad count differs from grid".into()); + } + for i in 0..count { + axes[end + i] = next; + axes[n + end + i] = next + i as i32 / w; + axes[n * 2 + end + i] = next + i as i32 % w; + } + next += h.max(w).max(1); + start = end + count; + } + if ids[start..].contains(&248056) { + return Err("Qwen image pads exceed image grids".into()); + } + for i in start..n { + for axis in 0..3 { + axes[axis * n + i] = next + (i - start) as i32; + } + } + next += (n - start) as i32; + Ok((axes, next - n as i32)) +} + +#[test] +fn qwen_vision_positions_and_resize_contract() { + assert_eq!(smart_resize(1024, 1024).unwrap(), (1024, 1024)); + assert_eq!(smart_resize(16, 16).unwrap(), (256, 256)); + assert!(smart_resize(1, 201).is_err()); + let (table, delta) = positions(&[1, 248056, 248056, 248056, 248056, 2], &[(4, 4)]).unwrap(); + assert_eq!( + table, + vec![0, 1, 1, 1, 1, 3, 0, 1, 1, 2, 2, 3, 0, 1, 2, 1, 2, 3] + ); + assert_eq!(delta, -2); + assert!(positions(&[248056], &[(4, 4)]).is_err()); +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub(super) struct Identity { + pub(super) start: usize, + pub(super) end: usize, + pub(super) digest: [u8; 32], +} + +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)) +} + +#[test] +fn qwen_image_preprocessing_matches_pillow_mtplx() { + use sha2::{Digest, Sha256}; + for (w, h, gh, gw, expected) in [ + ( + 17, + 31, + 22, + 12, + "339e4ac4e02bb8fcab6d521194200191e4c77fbcce224de98b40d8d48393930f", + ), + ( + 333, + 177, + 12, + 22, + "3fc0890cb34e0c37bf32c5ce2a8bfa9b178ae3564c8e97b8e6496e71a66b2644", + ), + ( + 1537, + 257, + 16, + 96, + "96bcea68014063c4edc51b422af82930e1bd43984d57a41e12f44865eed45368", + ), + ] { + let image = RgbImage::from_fn(w, h, |x, y| { + image::Rgb(std::array::from_fn(|c| { + ((x * 13 + y * 7 + c as u32 * 31) % 256) as u8 + })) + }); + let mut bytes = Cursor::new(Vec::new()); + image.write_to(&mut bytes, image::ImageFormat::Png).unwrap(); + let input = decode(bytes.get_ref()).unwrap(); + assert_eq!((input.h, input.w), (gh, gw)); + let hash = Sha256::digest( + input + .pixels + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect::>(), + ); + assert_eq!( + hash.iter().map(|b| format!("{b:02x}")).collect::(), + expected, + "{w}x{h}" + ); + } +} + +#[test] +fn qwen_image_cache_keys_distinguish_pixels_and_preserve_prior_text() { + let old = vec![Identity { + start: 10, + end: 14, + digest: [1; 32], + }]; + let changed = vec![Identity { + start: 10, + end: 14, + digest: [2; 32], + }]; + assert!(same_prefix(&old, &old, 20)); + assert!(!same_prefix(&old, &changed, 20)); + assert!(same_prefix(&old, &changed, 10)); + assert!(!same_prefix(&old, &[], 20)); +} diff --git a/src/model.rs b/src/model.rs index 4470bab..8d7ae1c 100644 --- a/src/model.rs +++ b/src/model.rs @@ -17,7 +17,7 @@ pub(crate) const MODEL_CHOICES: [ModelChoice; 5] = [ ModelChoice::Glm53Flash, ModelChoice::Qwen38FlashNext, ]; -pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [ +pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 8] = [ ManagedArtifactId::DeepSeekV4Flash0731, ManagedArtifactId::DeepSeekV4Flash0731Dspark, ManagedArtifactId::DeepSeekV4Pro, @@ -25,6 +25,7 @@ pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [ ManagedArtifactId::Glm53Flash, ManagedArtifactId::Glm53FlashVision, ManagedArtifactId::Qwen38FlashNext, + ManagedArtifactId::Qwen38FlashNextVision, ]; const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf"; @@ -128,6 +129,10 @@ impl ModelChoice { self == Self::Qwen38FlashNext } + pub(crate) fn supports_vision(self) -> bool { + matches!(self, Self::Glm53Flash | Self::Qwen38FlashNext) + } + pub(crate) fn supports_integrated_mtp(self) -> bool { self.is_glm() || self.is_qwen38() } @@ -189,7 +194,8 @@ pub(crate) fn engine_artifacts( return EngineArtifacts { model: qwen::root(models_path), support: None, - vision: None, + vision: qwen::vision_is_installed(&qwen::root(models_path)) + .then(|| qwen::root(models_path).join("model-vision.safetensors")), }; } EngineArtifacts { @@ -218,8 +224,14 @@ pub(crate) fn validate_engine_artifacts( return Err(format!("DSpark is not compatible with {model}")); } if model.is_qwen38() { - if artifacts.support.is_some() || artifacts.vision.is_some() { - return Err("Qwen3.8 does not accept GGUF support or vision artifacts".into()); + if artifacts.support.is_some() { + return Err("Qwen3.8 does not accept GGUF support artifacts".into()); + } + if let Some(path) = &artifacts.vision + && (path != &artifacts.model.join("model-vision.safetensors") + || !qwen::vision_is_installed(&artifacts.model)) + { + return Err("Qwen vision artifacts have not passed verification".into()); } return qwen::validate_installed(&artifacts.model); } @@ -258,9 +270,14 @@ pub(crate) enum ManagedArtifactId { Glm53Flash, Glm53FlashVision, Qwen38FlashNext, + Qwen38FlashNextVision, } impl ManagedArtifactId { + fn is_qwen(self) -> bool { + self.model().is_qwen38() + } + pub(crate) fn model(self) -> ModelChoice { match self { Self::DeepSeekV4Flash0731 | Self::DeepSeekV4Flash0731Dspark => { @@ -269,7 +286,7 @@ impl ManagedArtifactId { Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro, Self::Glm52 => ModelChoice::Glm52, Self::Glm53Flash | Self::Glm53FlashVision => ModelChoice::Glm53Flash, - Self::Qwen38FlashNext => ModelChoice::Qwen38FlashNext, + Self::Qwen38FlashNext | Self::Qwen38FlashNextVision => ModelChoice::Qwen38FlashNext, } } @@ -281,15 +298,17 @@ impl ManagedArtifactId { Self::Glm52 => &GLM, Self::Glm53Flash => &GLM53_FLASH, Self::Glm53FlashVision => &GLM53_FLASH_VISION, - Self::Qwen38FlashNext => unreachable!("Qwen uses a pinned artifact set"), + Self::Qwen38FlashNext | Self::Qwen38FlashNextVision => { + unreachable!("Qwen uses a pinned artifact set") + } } } } impl fmt::Display for ManagedArtifactId { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - if *self == Self::Qwen38FlashNext { - formatter.write_str(qwen::LABEL) + if self.is_qwen() { + formatter.write_str(qwen::label(*self)) } else { formatter.write_str(self.artifact().label) } @@ -480,8 +499,8 @@ pub(crate) fn managed_artifacts(models_path: &Path) -> Vec { MANAGED_ARTIFACTS .into_iter() .map(|id| { - if id == ManagedArtifactId::Qwen38FlashNext { - return qwen::managed_artifact(models_path); + if id.is_qwen() { + return qwen::managed_artifact(id, models_path); } let model = id.model(); let artifact = id.artifact(); @@ -527,8 +546,8 @@ pub(crate) fn artifact_download_progress( id: ManagedArtifactId, models_path: &Path, ) -> DownloadProgress { - if id == ManagedArtifactId::Qwen38FlashNext { - return qwen::download_progress(models_path); + if id.is_qwen() { + return qwen::download_progress(id, models_path); } let model = id.model(); let artifact = id.artifact(); @@ -560,8 +579,8 @@ pub(crate) fn artifact_verification_progress( id: ManagedArtifactId, verified: u64, ) -> DownloadProgress { - if id == ManagedArtifactId::Qwen38FlashNext { - return qwen::verification_progress(verified); + if id.is_qwen() { + return qwen::verification_progress(id, verified); } let artifact = id.artifact(); DownloadProgress { diff --git a/src/model/qwen.rs b/src/model/qwen.rs index aec32f2..5cff06b 100644 --- a/src/model/qwen.rs +++ b/src/model/qwen.rs @@ -22,6 +22,7 @@ struct Manifest { format: u32, source: Source, files: Vec, + excluded: Vec, } #[derive(Deserialize)] @@ -33,6 +34,7 @@ struct Source { #[derive(Deserialize)] pub(super) struct Artifact { pub(super) path: String, + #[serde(default)] pub(super) role: String, pub(super) size: u64, pub(super) sha256: String, @@ -61,6 +63,23 @@ pub(super) fn artifacts() -> &'static [Artifact] { &manifest().files } +pub(super) fn selected(id: ManagedArtifactId) -> &'static [Artifact] { + match id { + ManagedArtifactId::Qwen38FlashNext => artifacts(), + ManagedArtifactId::Qwen38FlashNextVision => &manifest().excluded, + _ => unreachable!("not a Qwen artifact set"), + } +} + +pub(super) fn label(id: ManagedArtifactId) -> &'static str { + match id { + ManagedArtifactId::Qwen38FlashNext => LABEL, + ManagedArtifactId::Qwen38FlashNextVision => "Qwen3.8 Flash Next vision encoder", + _ => unreachable!("not a Qwen artifact set"), + } +} + +#[cfg(test)] pub(super) fn total_bytes() -> u64 { artifacts().iter().map(|artifact| artifact.size).sum() } @@ -106,6 +125,12 @@ pub(super) fn is_installed(root: &Path) -> bool { .all(|artifact| artifact_is_installed(root, artifact)) } +pub(super) fn vision_is_installed(root: &Path) -> bool { + selected(ManagedArtifactId::Qwen38FlashNextVision) + .iter() + .all(|a| artifact_is_installed(root, a)) +} + pub(super) fn validate_installed(root: &Path) -> Result<(), String> { for artifact in artifacts() { let artifact_path = path(root, artifact); @@ -136,8 +161,8 @@ pub(super) fn validate_installed(root: &Path) -> Result<(), String> { Ok(()) } -fn stored_bytes(root: &Path) -> u64 { - artifacts() +fn stored_bytes(id: ManagedArtifactId, root: &Path) -> u64 { + selected(id) .iter() .flat_map(|artifact| [path(root, artifact), partial_path(root, artifact)]) .filter_map(|path| path.metadata().ok()) @@ -145,8 +170,8 @@ fn stored_bytes(root: &Path) -> u64 { .sum() } -fn downloaded_bytes(root: &Path) -> u64 { - artifacts() +fn downloaded_bytes(id: ManagedArtifactId, root: &Path) -> u64 { + selected(id) .iter() .map(|artifact| { if artifact_is_installed(root, artifact) { @@ -161,8 +186,8 @@ fn downloaded_bytes(root: &Path) -> u64 { .sum() } -fn complete_payloads(root: &Path) -> bool { - artifacts().iter().all(|artifact| { +fn complete_payloads(id: ManagedArtifactId, root: &Path) -> bool { + selected(id).iter().all(|artifact| { path(root, artifact) .metadata() .or_else(|_| partial_path(root, artifact).metadata()) @@ -170,12 +195,12 @@ fn complete_payloads(root: &Path) -> bool { }) } -pub(super) fn managed_artifact(models_path: &Path) -> ManagedArtifact { +pub(super) fn managed_artifact(id: ManagedArtifactId, models_path: &Path) -> ManagedArtifact { let root = root(models_path); - let stored = stored_bytes(&root); - let state = if is_installed(&root) { + let stored = stored_bytes(id, &root); + let state = if selected(id).iter().all(|a| artifact_is_installed(&root, a)) { ManagedArtifactState::Ready - } else if complete_payloads(&root) { + } else if complete_payloads(id, &root) { ManagedArtifactState::NeedsVerification } else if stored > 0 { ManagedArtifactState::Partial @@ -183,44 +208,45 @@ pub(super) fn managed_artifact(models_path: &Path) -> ManagedArtifact { ManagedArtifactState::Missing }; ManagedArtifact { - id: ManagedArtifactId::Qwen38FlashNext, + id, stored, - expected: total_bytes(), + expected: selected(id).iter().map(|a| a.size).sum(), state, } } -pub(super) fn download_progress(models_path: &Path) -> DownloadProgress { +pub(super) fn download_progress(id: ManagedArtifactId, models_path: &Path) -> DownloadProgress { let root = root(models_path); - let downloaded = downloaded_bytes(&root); - let installed = is_installed(&root); - let verification = (!installed && complete_payloads(&root)).then_some(VerificationProgress { - verified: 0, - total: total_bytes(), - }); + let downloaded = downloaded_bytes(id, &root); + let installed = selected(id).iter().all(|a| artifact_is_installed(&root, a)); + let verification = + (!installed && complete_payloads(id, &root)).then_some(VerificationProgress { + verified: 0, + total: selected(id).iter().map(|a| a.size).sum(), + }); let phase = if installed { DownloadPhase::Complete } else if verification.is_some() { - DownloadPhase::Verifying(LABEL) + DownloadPhase::Verifying(label(id)) } else if downloaded > 0 { - DownloadPhase::Downloading(LABEL) + DownloadPhase::Downloading(label(id)) } else { - DownloadPhase::Pending(LABEL) + DownloadPhase::Pending(label(id)) }; DownloadProgress { downloaded, - total: total_bytes(), + total: selected(id).iter().map(|a| a.size).sum(), phase, verification, } } -pub(super) fn verification_progress(verified: u64) -> DownloadProgress { - let total = total_bytes(); +pub(super) fn verification_progress(id: ManagedArtifactId, verified: u64) -> DownloadProgress { + let total = selected(id).iter().map(|a| a.size).sum(); DownloadProgress { downloaded: total, total, - phase: DownloadPhase::Verifying(LABEL), + phase: DownloadPhase::Verifying(label(id)), verification: Some(VerificationProgress { verified: verified.min(total), total, @@ -336,3 +362,41 @@ mod tests { fs::remove_dir_all(root).unwrap(); } } + +#[test] +fn qwen_vision_is_an_independent_managed_artifact_set() { + let id = ManagedArtifactId::Qwen38FlashNextVision; + let root = std::env::temp_dir().join(format!("qwen-vision-catalog-{}", std::process::id())); + let model_root = self::root(&root); + fs::create_dir_all(&model_root).unwrap(); + assert_eq!(selected(id).len(), 4); + assert_eq!( + selected(id).iter().map(|a| a.size).sum::(), + 897_900_287 + ); + assert!( + selected(id) + .iter() + .all(|v| artifacts().iter().all(|a| a.path != v.path)) + ); + for artifact in selected(id) { + File::create(path(&model_root, artifact)) + .unwrap() + .set_len(artifact.size) + .unwrap(); + fs::write(verification_path(&model_root, artifact), &artifact.sha256).unwrap(); + } + let text = model_root.join("model-00001-of-00017.safetensors"); + fs::write(&text, b"preserve text model").unwrap(); + assert_eq!( + managed_artifact(id, &root).state, + ManagedArtifactState::Ready + ); + assert_eq!(download_progress(id, &root).phase, DownloadPhase::Complete); + assert!(!is_installed(&model_root)); + assert!(vision_is_installed(&model_root)); + super::delete_managed_artifact(id, &root).unwrap(); + assert_eq!(fs::read(&text).unwrap(), b"preserve text model"); + assert!(!vision_is_installed(&model_root)); + fs::remove_dir_all(root).unwrap(); +} diff --git a/src/model/transfer.rs b/src/model/transfer.rs index 7c9d31c..5b609e8 100644 --- a/src/model/transfer.rs +++ b/src/model/transfer.rs @@ -14,8 +14,8 @@ pub(crate) fn download_managed_artifact( cancel: &AtomicBool, verified_bytes: &AtomicU64, ) -> Result { - if id == ManagedArtifactId::Qwen38FlashNext { - return download_qwen(models_path, cancel, verified_bytes); + if id.is_qwen() { + return download_qwen(id, models_path, cancel, verified_bytes); } download_artifact_with_cancel( id.model(), @@ -32,8 +32,8 @@ pub(crate) fn validate_managed_artifact( cancel: &AtomicBool, verified_bytes: &AtomicU64, ) -> Result { - if id == ManagedArtifactId::Qwen38FlashNext { - return validate_qwen(models_path, cancel, verified_bytes); + if id.is_qwen() { + return validate_qwen(id, models_path, cancel, verified_bytes); } let model = id.model(); let artifact = id.artifact(); @@ -73,8 +73,8 @@ pub(crate) fn delete_managed_artifact( id: ManagedArtifactId, models_path: &Path, ) -> Result<(), String> { - if id == ManagedArtifactId::Qwen38FlashNext { - return delete_qwen(models_path); + if id.is_qwen() { + return delete_qwen(id, models_path); } let model = id.model(); let artifact = id.artifact(); @@ -93,6 +93,7 @@ pub(crate) fn delete_managed_artifact( } fn download_qwen( + id: ManagedArtifactId, models_path: &Path, cancel: &AtomicBool, verified_bytes: &AtomicU64, @@ -100,7 +101,7 @@ fn download_qwen( let root = qwen::root(models_path); fs::create_dir_all(&root).map_err(|error| format!("{}: {error}", root.display()))?; verified_bytes.store(0, Ordering::Relaxed); - for artifact in qwen::artifacts() { + for artifact in qwen::selected(id) { if cancel.load(Ordering::Relaxed) { return Ok(DownloadOutcome::Stopped); } @@ -145,26 +146,29 @@ fn download_qwen( fs::write(qwen::verification_path(&root, artifact), &artifact.sha256) .map_err(|error| error.to_string())?; } - if let Err(error) = crate::engine::validate_qwen_artifacts(&root) { - clear_qwen_markers(&root)?; + if id == ManagedArtifactId::Qwen38FlashNext + && let Err(error) = crate::engine::validate_qwen_artifacts(&root) + { + clear_qwen_markers(id, &root)?; return Err(error); } Ok(DownloadOutcome::Complete) } fn validate_qwen( + id: ManagedArtifactId, models_path: &Path, cancel: &AtomicBool, verified_bytes: &AtomicU64, ) -> Result { let root = qwen::root(models_path); verified_bytes.store(0, Ordering::Relaxed); - for artifact in qwen::artifacts() { + for artifact in qwen::selected(id) { if !qwen::path(&root, artifact).exists() && !qwen::partial_path(&root, artifact).exists() { return Err(format!("{} is not downloaded", artifact.path)); } } - for artifact in qwen::artifacts() { + for artifact in qwen::selected(id) { let destination = qwen::path(&root, artifact); let partial = qwen::partial_path(&root, artifact); let (source, promote) = if destination.exists() { @@ -189,15 +193,17 @@ fn validate_qwen( fs::write(qwen::verification_path(&root, artifact), &artifact.sha256) .map_err(|error| error.to_string())?; } - if let Err(error) = crate::engine::validate_qwen_artifacts(&root) { - clear_qwen_markers(&root)?; + if id == ManagedArtifactId::Qwen38FlashNext + && let Err(error) = crate::engine::validate_qwen_artifacts(&root) + { + clear_qwen_markers(id, &root)?; return Err(error); } Ok(DownloadOutcome::Complete) } -fn clear_qwen_markers(root: &Path) -> Result<(), String> { - for artifact in qwen::artifacts() { +fn clear_qwen_markers(id: ManagedArtifactId, root: &Path) -> Result<(), String> { + for artifact in qwen::selected(id) { let path = qwen::verification_path(root, artifact); match fs::remove_file(&path) { Ok(()) => {} @@ -208,9 +214,9 @@ fn clear_qwen_markers(root: &Path) -> Result<(), String> { Ok(()) } -fn delete_qwen(models_path: &Path) -> Result<(), String> { +fn delete_qwen(id: ManagedArtifactId, models_path: &Path) -> Result<(), String> { let root = qwen::root(models_path); - for artifact in qwen::artifacts() { + for artifact in qwen::selected(id) { for path in [ qwen::path(&root, artifact), qwen::partial_path(&root, artifact), @@ -225,7 +231,14 @@ fn delete_qwen(models_path: &Path) -> Result<(), String> { } match fs::remove_dir(&root) { Ok(()) => Ok(()), - Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::NotFound | std::io::ErrorKind::DirectoryNotEmpty + ) => + { + Ok(()) + } Err(error) => Err(format!("{}: {error}", root.display())), } } @@ -480,7 +493,7 @@ mod tests { ); assert!(ModelChoice::from_id("unknown").is_none()); assert_eq!(MODEL_CHOICES.len(), 5); - assert_eq!(MANAGED_ARTIFACTS.len(), 7); + assert_eq!(MANAGED_ARTIFACTS.len(), 8); assert_eq!( ModelChoice::from_id("qwen3.8-flash-next"), Some(ModelChoice::Qwen38FlashNext) @@ -600,20 +613,31 @@ mod tests { )) .unwrap(); assert!( - validate_qwen(&models, &AtomicBool::new(false), &AtomicU64::new(0)) - .unwrap_err() - .contains("mtp.safetensors is not downloaded") + validate_qwen( + ManagedArtifactId::Qwen38FlashNext, + &models, + &AtomicBool::new(false), + &AtomicU64::new(0) + ) + .unwrap_err() + .contains("mtp.safetensors is not downloaded") ); fs::hard_link(source.join("mtp.safetensors"), root.join("mtp.safetensors")).unwrap(); let verified = AtomicU64::new(0); assert_eq!( - validate_qwen(&models, &AtomicBool::new(false), &verified).unwrap(), + validate_qwen( + ManagedArtifactId::Qwen38FlashNext, + &models, + &AtomicBool::new(false), + &verified + ) + .unwrap(), DownloadOutcome::Complete ); assert_eq!(verified.load(Ordering::Relaxed), qwen::total_bytes()); assert!(qwen::is_installed(&root)); assert!(crate::engine::validate_qwen_artifacts(&root).is_ok()); - delete_qwen(&models).unwrap(); + delete_qwen(ManagedArtifactId::Qwen38FlashNext, &models).unwrap(); assert!(!root.exists()); fs::remove_dir_all(models).unwrap(); } diff --git a/src/model_eval.rs b/src/model_eval.rs index 0353789..73c004b 100644 --- a/src/model_eval.rs +++ b/src/model_eval.rs @@ -38,6 +38,7 @@ struct Options { prompt: String, additional_prompts: Vec, input_file: Option, + image_file: Option, models_path: PathBuf, config_path: PathBuf, plain_chat: bool, @@ -68,6 +69,7 @@ impl Options { let mut prompt_overridden = false; let mut additional_prompts = Vec::new(); let mut input_file = None; + let mut image_file = None; let mut models_path = crate::app::models_path(); let mut config_path = crate::app::config_path(); let mut context = None; @@ -120,6 +122,7 @@ impl Options { prompt_overridden = true; } } + "--image-file" => image_file = Some(value()?.into()), "--input-file" => input_file = Some(value()?.into()), "--models-dir" => models_path = value()?.into(), "--config" => config_path = value()?.into(), @@ -190,6 +193,7 @@ impl Options { prompt, additional_prompts, input_file, + image_file, models_path, config_path, plain_chat, @@ -249,6 +253,7 @@ Models: deepseek-v4-flash-0731, deepseek-v4-pro, glm-5.2, glm-5.3-flash, qwen3.8 Options:\n\ --prompt TEXT Chat turn; repeat for one ongoing conversation\n\ (default: autonomous short story)\n\ + --image-file PATH Attach PNG/JPEG pixels to the first turn (no filename sent)\n\ --input-file PATH Append a UTF-8 file to the first turn (maximum: 16 MiB)\n\ --context TOKENS Context size; otherwise use the model profile\n\ --max-tokens TOKENS Override the model profile generation limit\n\ @@ -937,6 +942,30 @@ fn stop_worker( } fn prompts_with_input_file(options: &Options) -> Result, String> { + let mut prompts = text_prompts_with_input_file(options)?; + if let Some(path) = &options.image_file { + use base64::Engine as _; + if !options.model.supports_vision() { + return Err("selected model does not support images".into()); + } + if fs::metadata(path).map_err(|e| e.to_string())?.len() > 50 * 1024 * 1024 { + return Err("image exceeds 50 MiB".into()); + } + let bytes = fs::read(path).map_err(|e| e.to_string())?; + let mime = match image::guess_format(&bytes).map_err(|e| e.to_string())? { + image::ImageFormat::Png => "image/png", + image::ImageFormat::Jpeg => "image/jpeg", + _ => return Err("image must be PNG or JPEG".into()), + }; + prompts[0].push_str(&crate::engine::vision_data_marker(&format!( + "data:{mime};base64,{}", + base64::engine::general_purpose::STANDARD.encode(bytes) + ))); + } + Ok(prompts) +} + +fn text_prompts_with_input_file(options: &Options) -> Result, String> { let mut prompts = Vec::with_capacity(1 + options.additional_prompts.len()); let Some(path) = &options.input_file else { prompts.push(options.prompt.clone()); diff --git a/src/settings.rs b/src/settings.rs index 562cb18..60c58a9 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -51,8 +51,8 @@ impl SpeculativePreferences { if !model.supports_integrated_mtp() && (self.glm_mtp || self.glm_mtp_timing) { return Err("Integrated MTP is unavailable for the selected model.".into()); } - if self.keep_vision_loaded && model != ModelChoice::Glm53Flash { - return Err("Persistent vision weights are available only for GLM 5.3 Flash.".into()); + if self.keep_vision_loaded && !model.supports_vision() { + return Err("Persistent vision weights are available only for GLM 5.3 Flash and Qwen3.8 Flash Next.".into()); } if self.dspark_enabled && !model.supports_dspark() { return Err("DSpark is not available for the selected model.".into()); diff --git a/tools/qwen-vision-chat-reference.py b/tools/qwen-vision-chat-reference.py new file mode 100644 index 0000000..9764908 --- /dev/null +++ b/tools/qwen-vision-chat-reference.py @@ -0,0 +1,32 @@ +"""Describe a neutral image and run a no-image control through pinned MTPLX. +Usage: test-supervisor ... --command python tools/qwen-vision-chat-reference.py MODEL_DIR IMAGE +Uses the reference runtime only as an oracle; application inference is Rust/Metal. +""" +import os,sys,json,time,argparse +from pathlib import Path +root=Path(sys.argv[1]);image=Path(sys.argv[2]) +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 +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 +print(json.dumps(dict(event='loading')),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) +rt=MTPLXRuntime(model,tokenizer,root,True,MTPContract()) +for with_image in [True,False]: + text='Describe this image'+('<|vision_start|><|image_pad|><|vision_end|>' if with_image else '') + ids=_encode_messages_uncached(tokenizer,[ChatMessage(role='user',content=text)],enable_thinking=True,reasoning_effort='low',preserve_reasoning_history=True,tools=None) + splice=None + if with_image:ids,splice=_materialize_vision_splice(argparse.Namespace(args=argparse.Namespace(model=str(root))),[image.read_bytes()],ids) + print(json.dumps(dict(event='input',with_image=with_image,prompt='Describe this image',ids=ids)),flush=True) + out=generate_mtpk(rt,ids,max_tokens=1024,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 d:print(json.dumps(dict(event='prefill',**d)),flush=True),token_callback=lambda ids:print(json.dumps(dict(event='tokens',ids=ids)),flush=True)) + print(json.dumps(dict(event='result',with_image=with_image,text=out.text,tokens=out.tokens,finish_reason=out.finish_reason)),flush=True) diff --git a/tools/qwen-vision-reference.py b/tools/qwen-vision-reference.py new file mode 100644 index 0000000..9f7007f --- /dev/null +++ b/tools/qwen-vision-reference.py @@ -0,0 +1,24 @@ +"""Export pinned MTPLX vision stages; oracle only, never application code. +Usage: python tools/qwen-vision-reference.py MODEL_DIR IMAGE OUTPUT_DIR +Run in the pinned MTPLX reference environment with GPU access. +""" +import sys,json,hashlib,pathlib +import mlx.core as mx +import numpy as np +from mtplx.vision import load_vision_tower +from mtplx.vision.processing import decode_image,preprocess_images +root=pathlib.Path(sys.argv[1]); image=pathlib.Path(sys.argv[2]); out=pathlib.Path(sys.argv[3]);out.mkdir(exist_ok=True) +def save(name,x): + mx.eval(x); a=np.asarray(x.astype(mx.float32));a.tofile(out/(name+'.f32')); print(json.dumps(dict(stage=name,shape=list(x.shape),dtype=str(x.dtype),min=float(a.min()),max=float(a.max()),sha256=hashlib.sha256(a.tobytes()).hexdigest())),flush=True) +pixels,grids=preprocess_images([decode_image(image.read_bytes())],json.loads((root/'preprocessor_config.json').read_text())) +(out/'grid.json').write_text(json.dumps(grids[0]));save('pixels',pixels); print(json.dumps(dict(grids=grids)),flush=True) +tower=load_vision_tower(root) +h=tower.patch_embed(pixels.astype(tower.patch_embed.proj.weight.dtype));save('patch',h) +p=tower.fast_pos_embed_interpolate(grids);save('position',p) +h=h+p;r=tower.rot_pos_emb(grids);save('rotary',r) +for i,b in enumerate(tower.blocks): + h=b(h,[],r) + if i in [0,26]:save('block'+str(i),h) + else:mx.eval(h) + print(json.dumps(dict(block=i)),flush=True) +h=tower.merger(h);save('embeddings',h)