Reject invalid Qwen logits and discard poisoned sessions

src/engine/metal/qwen_mtplx/binary.rs

@@ -7,6 +7,7 @@ use super::{bytes, dispatch_geometry, ops, tensor, views};
 #[derive(Clone, Copy, Debug)]
 pub(super) enum Binary {
     Equal,
+    NotEqual,
     Greater,
     GreaterEqual,
     Less,
@@ -106,6 +107,7 @@ pub(super) fn binary(a: &Array, b: &Array, op: Binary, stream: Stream) -> Result
             if matches!(
                 op,
                 Binary::Equal
+                    | Binary::NotEqual
                     | Binary::Greater
                     | Binary::GreaterEqual
                     | Binary::Less

src/engine/metal/qwen_mtplx/execution.rs

@@ -226,6 +226,10 @@ pub(super) struct Execution {
     pub(super) prefill_throttle: Throttle,
     pub(super) decode_throttle: Throttle,
     runtime_overrides: super::policy::Overrides,
+    #[cfg(test)]
+    pub(super) fault: Option<super::issue105::Fault>,
+    #[cfg(test)]
+    pub(super) captured_prefixes: usize,
     _context: Context,
 }
 
@@ -335,6 +339,10 @@ impl Execution {
             prefill_throttle,
             decode_throttle,
             runtime_overrides,
+            #[cfg(test)]
+            fault: None,
+            #[cfg(test)]
+            captured_prefixes: 0,
             _context: context,
         })
     }
@@ -347,7 +355,7 @@ impl Execution {
     ) -> Result<Output, String> {
         let ids = self.input(ids)?;
         let options = self.graph_options(phase_prefill, capture)?;
-        self.model.with_dependent_mut(|_, model| {
+        let output = self.model.with_dependent_mut(|_, model| {
             model.call(
                 &ids,
                 Some(&mut self.cache),
@@ -362,7 +370,10 @@ impl Execution {
                 &self.streams,
                 self.stream,
             )
-        })
+        })?;
+        #[cfg(test)]
+        let output = self.inject_output(super::issue105::Stage::Target, output)?;
+        Ok(output)
     }
 
     /// Cold text _prefill. Session restore and optional owner-level cleanup /
@@ -775,7 +786,7 @@ impl Execution {
     pub(super) fn draft(&mut self, hidden: &Array, ids: &[i32]) -> Result<Output, String> {
         let ids = self.input(ids)?;
         let options = self.graph_options(false, false)?;
-        self.model.with_dependent_mut(|_, model| {
+        let output = self.model.with_dependent_mut(|_, model| {
             model.mtp_forward(
                 hidden,
                 &ids,
@@ -786,7 +797,10 @@ impl Execution {
                 &self.streams,
                 self.stream,
             )
-        })
+        })?;
+        #[cfg(test)]
+        let output = self.inject_output(super::issue105::Stage::Draft, output)?;
+        Ok(output)
     }
 
     pub(super) fn append_mtp_history(
@@ -835,7 +849,7 @@ impl Execution {
             self.stream,
             &mut self.decode_throttle,
             |ids, hidden| {
-                self.model.with_dependent_mut(|_, model| {
+                let output = self.model.with_dependent_mut(|_, model| {
                     model.call(
                         ids,
                         Some(&mut self.cache),
@@ -850,7 +864,24 @@ impl Execution {
                         &self.streams,
                         self.stream,
                     )
-                })
+                })?;
+                #[cfg(test)]
+                let output = Output {
+                    logits: output
+                        .logits
+                        .map(|row| {
+                            super::issue105::inject_logits(
+                                &mut self.fault,
+                                super::issue105::Stage::Target,
+                                &row,
+                                &self.streams,
+                                self.stream,
+                            )
+                        })
+                        .transpose()?,
+                    ..output
+                };
+                Ok(output)
             },
             emit,
         );

src/engine/metal/qwen_mtplx/fused.rs

@@ -51,6 +51,7 @@ pub(super) fn operator(op: &Operation) -> Option<&'static str> {
         },
         Operation::Binary(op) => match op {
             B::Equal => "Equal",
+            B::NotEqual => "NotEqual",
             B::GreaterEqual => "GreaterEqual",
             B::Greater => "Greater",
             B::Less => "Less",

src/engine/metal/qwen_mtplx/issue103_analysis.rs

@@ -91,6 +91,10 @@ fn issue103_sampling_receipts() {
             serde_json::json!({"case":"sampling_matrix","temperature":temperature,"top_k":top_k,"top_p":top_p,"ordinary_probabilities":probabilities,"lazy_counts":counts.to_vec()})
         );
     }
+    let oracle = include_str!("../../../../docs/issue105-invalid-logits/oracle.jsonl")
+        .lines()
+        .map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
+        .collect::<Vec<_>>();
     for (name, values) in [
         ("nan", vec![f32::NAN; 64]),
         ("positive_inf", [vec![f32::INFINITY], vec![0.; 63]].concat()),
@@ -123,10 +127,21 @@ fn issue103_sampling_receipts() {
                 "{}",
                 serde_json::json!({"case":"invalid_logits","row":name,"temperature":temperature,"ordinary":result.as_ref().map(|r|r.0).map_err(Clone::clone)})
             );
+            let expected = oracle
+                .iter()
+                .find(|r| {
+                    r["case"] == "ordinary" && r["row"] == name && r["temperature"] == temperature
+                })
+                .unwrap();
             if name == "valid_mask" {
-                assert!([1, 3].contains(&result.unwrap().0));
+                assert_eq!(result.unwrap().0, expected["token"].as_i64().unwrap());
             } else {
-                assert!(result.is_ok(), "diagnosis no longer reproduces");
+                let error = result.err().unwrap();
+                assert!(crate::engine::qwen::is_non_finite_logits(&error));
+                assert_eq!(
+                    error,
+                    format!("Qwen {}", expected["error"].as_str().unwrap())
+                );
             }
         }
         let config = SamplerConfig {
@@ -139,6 +154,12 @@ fn issue103_sampling_receipts() {
             "{}",
             serde_json::json!({"case":"invalid_logits_lazy","row":name,"result":result})
         );
+        let id = result.unwrap()[0];
+        if name == "valid_mask" {
+            assert!([1, 3].contains(&id));
+        } else {
+            assert_eq!(id, -1);
+        }
     }
 }
 

src/engine/metal/qwen_mtplx/issue105.rs

@@ -0,0 +1,141 @@
+//! Test-only, owner-local injection at real model output boundaries.
+use super::array::{Array, Dtype};
+use super::execution::Execution;
+use super::{model::Output, ops, sampling};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(super) enum Stage {
+    Prefill,
+    Target,
+    Draft,
+    Capture,
+}
+
+#[derive(Clone, Copy, Debug)]
+pub(super) enum Row {
+    Nan,
+    PositiveInf,
+    Masked,
+    LegalMask,
+}
+
+pub(super) struct Fault {
+    pub(super) stage: Stage,
+    pub(super) row: Row,
+}
+
+impl Execution {
+    pub(super) fn inject_logits(&mut self, stage: Stage, logits: &Array) -> Result<Array, String> {
+        inject_logits(&mut self.fault, stage, logits, &self.streams, self.stream)
+    }
+
+    pub(super) fn inject_output(&mut self, stage: Stage, output: Output) -> Result<Output, String> {
+        Ok(Output {
+            logits: output
+                .logits
+                .map(|row| self.inject_logits(stage, &row))
+                .transpose()?,
+            ..output
+        })
+    }
+
+    pub(super) fn inject_capture_failure(&mut self, logits: &Array) -> Result<(), String> {
+        if self
+            .fault
+            .as_ref()
+            .is_some_and(|f| f.stage == Stage::Capture)
+        {
+            // Exercise final-capture error ownership, using the real numerical
+            // error produced by consuming a faulty row. Normal capture does not
+            // sample or add a validity pass to the oracle's execution route.
+            let row = self.inject_logits(Stage::Capture, logits)?;
+            sampling::checked_argmax(&row, &self.streams, self.stream)?;
+        }
+        Ok(())
+    }
+}
+
+pub(super) fn inject_logits(
+    fault: &mut Option<Fault>,
+    stage: Stage,
+    logits: &Array,
+    streams: &super::stream::Streams,
+    stream: super::stream::Stream,
+) -> Result<Array, String> {
+    let Some(pending) = &*fault else {
+        return Ok(logits.clone());
+    };
+    if pending.stage != stage {
+        return Ok(logits.clone());
+    }
+    let fault = fault.take().unwrap();
+    // Settle the real computation before replacing its output. The normal
+    // caches, prefix capture and request/persistence owners all remain live.
+    ops::evaluate(streams, std::slice::from_ref(logits), stream, false)?;
+    let shape = logits.layout().shape().to_vec();
+    let vocab = *shape.last().unwrap() as usize;
+    let data = (0..logits.layout().size())
+        .flat_map(|i| {
+            let value = match fault.row {
+                Row::Nan => f32::NAN,
+                Row::PositiveInf if i % vocab == 0 => f32::INFINITY,
+                Row::PositiveInf => 0.,
+                Row::Masked => f32::NEG_INFINITY,
+                Row::LegalMask if i % vocab == 1 => 2.,
+                Row::LegalMask if i % vocab == 3 => 1.,
+                Row::LegalMask => f32::NEG_INFINITY,
+            };
+            value.to_le_bytes()
+        })
+        .collect::<Vec<_>>();
+    Array::new(&shape, Dtype::F32, super::scalar_buffer(&data)?)
+}
+
+#[test]
+#[ignore = "Metal required; run serially under test-supervisor"]
+fn issue105_batched_greedy_rejects_each_invalid_row() {
+    super::configure_sources().unwrap();
+    let _context = super::super::gpu::Context::open_qwen(0).unwrap();
+    let streams = super::stream::Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let stream = streams.default_stream(super::stream::Device::Gpu).unwrap();
+    for width in [64, 4097] {
+        for row in [Row::Nan, Row::PositiveInf, Row::Masked, Row::LegalMask] {
+            for position in [0, 1] {
+                let healthy = Array::new(
+                    &[1, width],
+                    Dtype::F32,
+                    super::scalar_buffer(&vec![0; width as usize * 4]).unwrap(),
+                )
+                .unwrap();
+                let injected = inject_logits(
+                    &mut Some(Fault {
+                        stage: Stage::Target,
+                        row,
+                    }),
+                    Stage::Target,
+                    &healthy,
+                    &streams,
+                    stream,
+                )
+                .unwrap();
+                let rows = if position == 0 {
+                    [injected, healthy]
+                } else {
+                    [healthy, injected]
+                };
+                let logits = ops::concatenate(&rows, 0, stream).unwrap();
+                let result = sampling::checked_argmax(&logits, &streams, stream);
+                if matches!(row, Row::LegalMask) {
+                    let ids = ops::read_i64(&result.unwrap(), &streams, stream).unwrap();
+                    assert_eq!(ids[position], 1);
+                    assert_eq!(ids[1 - position], 0);
+                } else {
+                    assert!(crate::engine::qwen::is_non_finite_logits(
+                        &result.err().unwrap()
+                    ));
+                }
+            }
+        }
+    }
+    streams.clear_streams().unwrap();
+}

src/engine/metal/qwen_mtplx/issue105_product.rs

@@ -0,0 +1,298 @@
+//! Installed-model faults through Generator and both HTTP response adapters.
+use super::super::issue105::{Fault, Row, Stage};
+use super::*;
+use crate::engine::{
+    Generator, conversation_key, conversation_tag,
+    kvstore::{KvStore, StoreReason},
+    metal::Executor,
+};
+use crate::runtime::{ActiveGeneration, GenerationEvent};
+use crate::settings::{
+    ExecutionPreferences, GenerationPreferences, ReasoningMode, RuntimePreferences,
+    SpeculativePreferences, effective_settings,
+};
+use std::sync::{Arc, mpsc};
+
+fn owner(generator: &mut Generator) -> &mut QwenExecutor {
+    match &mut generator.executor {
+        Executor::QwenMtplx(owner) => owner,
+        _ => unreachable!(),
+    }
+}
+fn message(text: &str) -> ChatTurn {
+    ChatTurn {
+        user: true,
+        tool: false,
+        system: false,
+        skip_previous_eos: false,
+        reasoning: None,
+        reasoning_complete: true,
+        content: text.into(),
+    }
+}
+fn generate(
+    generator: &mut Generator,
+    path: &Path,
+    messages: &[ChatTurn],
+    settings: &TurnSettings,
+) -> GenerationOutput {
+    generator
+        .generate(
+            path,
+            None,
+            messages,
+            settings,
+            &AtomicBool::new(false),
+            |_, _| {},
+            |_, _, _| {},
+            |_| {},
+        )
+        .unwrap()
+}
+
+#[test]
+#[ignore = "installed Qwen model; serial test-supervisor with memory guard required"]
+fn issue105_product_ar_cleanup_and_http_retry() {
+    product(false);
+}
+#[test]
+#[ignore = "installed Qwen MTP model; serial test-supervisor with memory guard required"]
+fn issue105_product_mtp_cleanup_and_http_retry() {
+    product(true);
+}
+
+fn product(mtp: bool) {
+    let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
+    let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE105_DIR").unwrap());
+    std::fs::create_dir_all(&directory).unwrap();
+    let effective = effective_settings(
+        crate::model::ModelChoice::Qwen38FlashNext,
+        &GenerationPreferences {
+            context_tokens: 4096,
+            max_generated_tokens: 4,
+            system_prompt: "Be concise.".into(),
+            temperature: Some(0.6),
+            top_p: Some(0.95),
+            min_p: Some(0.),
+            reasoning_mode: ReasoningMode::Low,
+            ..Default::default()
+        },
+        &RuntimePreferences {
+            execution: ExecutionPreferences {
+                power_percent: Some(100),
+                ..Default::default()
+            },
+            speculative: SpeculativePreferences {
+                glm_mtp: mtp,
+                ..Default::default()
+            },
+            ..Default::default()
+        },
+        root.parent().unwrap(),
+    )
+    .unwrap();
+    let mut settings = effective.turn;
+    settings.seed = Some(7);
+    settings.kv_cache.min_tokens = 0;
+    settings.kv_cache.cold_max_tokens = 4096;
+    settings.kv_cache.continued_interval_tokens = 1;
+    let metrics = Arc::new(Metrics::new(&directory.join("metrics")));
+    eprintln!("#105 load product mtp={mtp}");
+    let mut generator = Generator::open(&effective.engine, metrics).unwrap();
+    generator.resident_limit = 8;
+    owner(&mut generator).policy.store_min_suffix = Some(0);
+    let baseline_path = directory.join("baseline.bin");
+    let mut messages = vec![message("Count from one.")];
+    let baseline = generate(&mut generator, &baseline_path, &messages, &settings);
+    messages.push(baseline.message);
+    let tag = conversation_tag(&settings.system_prompt, settings.reasoning_mode, &messages);
+    let key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &messages);
+    messages.push(message("Continue."));
+    let healthy_path = directory.join("healthy.bin");
+    generate(
+        &mut generator,
+        &healthy_path,
+        &[message("Say hello.")],
+        &settings,
+    );
+    let healthy_bytes = std::fs::read(&healthy_path).unwrap();
+    let bootstrap = directory.join("bootstrap");
+    let transient = directory.join("transient");
+    let mut cases = Vec::new();
+    for stage in [Stage::Prefill, Stage::Target, Stage::Capture]
+        .into_iter()
+        .chain(mtp.then_some(Stage::Draft))
+    {
+        for row in [Row::Nan, Row::PositiveInf, Row::Masked, Row::LegalMask] {
+            for streaming in [false, true] {
+                cases.push((stage, row, streaming));
+            }
+        }
+    }
+    for (index, (stage, row, streaming)) in cases.into_iter().enumerate() {
+        settings.temperature = if index % 4 < 2 { 0. } else { 0.6 };
+        // The real context-copy route can satisfy this short prompt without
+        // visiting the draft head. A supported penalty disables copy, ensuring
+        // draft injection is reached; target cases cover both routes.
+        owner(&mut generator).policy.sampler.presence_penalty =
+            if stage == Stage::Draft || (stage == Stage::Target && streaming) {
+                0.1
+            } else {
+                0.0
+            };
+        let path = if streaming {
+            let store = KvStore::open(&transient, settings.kv_cache.budget_bytes).unwrap();
+            let path = store.checkpoint_path(&key);
+            std::fs::copy(&baseline_path, &path).unwrap();
+            store
+                .record(&path, &key, tag, 1, 4096, StoreReason::Cold)
+                .unwrap();
+            path
+        } else {
+            let path = directory.join("failed.bin");
+            if index == 0 {
+                // Exercise a real bootstrap created by this failing cold turn.
+                assert!(!path.exists());
+            } else {
+                std::fs::copy(&baseline_path, &path).unwrap();
+            }
+            path
+        };
+        let before = owner(&mut generator).execution.captured_prefixes;
+        owner(&mut generator).execution.fault = Some(Fault { stage, row });
+        let (events, receiver) = mpsc::channel();
+        let cancelled = Arc::new(AtomicBool::new(false));
+        eprintln!(
+            "#105 case {index}: mtp={mtp} stage={stage:?} row={row:?} streaming={streaming} temperature={}",
+            settings.temperature
+        );
+        let chunks = Cell::new(0);
+        let emit = |reasoning, content| {
+            chunks.set(chunks.get() + 1);
+            events
+                .send(GenerationEvent::Chunk { reasoning, content })
+                .unwrap();
+        };
+        let progress = |used, limit, tokens_per_second| {
+            events
+                .send(GenerationEvent::Context {
+                    used,
+                    limit,
+                    tokens_per_second,
+                })
+                .unwrap();
+        };
+        let result = if streaming {
+            generator
+                .generate_transient(&transient, &messages, &settings, &cancelled, emit, progress)
+        } else {
+            generator.generate(
+                &path,
+                Some(&bootstrap),
+                &messages,
+                &settings,
+                &cancelled,
+                emit,
+                progress,
+                |_| {},
+            )
+        };
+        assert!(
+            owner(&mut generator).execution.fault.is_none(),
+            "fault site must be reached: result={:?}",
+            result.as_ref().map(|o| o.completion_tokens)
+        );
+        assert!(
+            owner(&mut generator).execution.captured_prefixes > before,
+            "prefill snapshot before failure"
+        );
+        let valid = matches!(row, Row::LegalMask);
+        if !valid {
+            if stage == Stage::Prefill {
+                assert_eq!(chunks.get(), 0, "no fabricated token reaches the response");
+            }
+            assert!(crate::engine::qwen::is_non_finite_logits(
+                result.as_ref().err().unwrap()
+            ));
+            assert!(owner(&mut generator).live.current.is_none());
+            assert!(owner(&mut generator).live.prefix.is_none());
+            assert!(owner(&mut generator).execution.boundaries.is_empty());
+            assert!(generator.checkpoint.is_none());
+            assert!(generator.resident_active.is_none());
+            assert_eq!(generator.last_store_tokens, 0);
+            assert!(!generator.resident_sessions.contains_key(&path));
+            assert!(!path.exists());
+            assert!(!path.with_extension("meta").exists());
+            assert!(!path.with_extension("tmp").exists());
+            let delayed = directory.join("delayed.bin");
+            assert!(generator.save_checkpoint(&delayed, tag).is_err());
+            assert!(!delayed.exists());
+            if !streaming {
+                let store = KvStore::open(&bootstrap, settings.kv_cache.budget_bytes).unwrap();
+                assert!(store.find(&key, 4096).is_none());
+            }
+        } else {
+            assert!(result.is_ok(), "legal masks: {:?}", result.as_ref().err());
+        }
+        events.send(GenerationEvent::Finished(result)).unwrap();
+        drop(events);
+        let wire = crate::server::issue105_wire(
+            ActiveGeneration {
+                events: receiver,
+                cancel: cancelled,
+            },
+            streaming,
+        );
+        if valid {
+            assert!(wire.starts_with("HTTP/1.1 200"));
+            assert!(!wire.contains("non-finite logits"));
+        } else if streaming {
+            assert!(wire.contains("event: error\n"));
+            assert!(wire.contains("Qwen non-finite logits"));
+        } else {
+            assert!(wire.starts_with("HTTP/1.1 500"));
+            assert!(wire.contains("Qwen non-finite logits"));
+        }
+        assert_eq!(std::fs::read(&healthy_path).unwrap(), healthy_bytes);
+        assert!(
+            generator.resident_sessions.contains_key(&healthy_path),
+            "unrelated healthy session remains resident"
+        );
+        if !valid {
+            let retried = if streaming {
+                generator
+                    .generate_transient(
+                        &transient,
+                        &messages,
+                        &settings,
+                        &AtomicBool::new(false),
+                        |_, _| {},
+                        |_, _, _| {},
+                    )
+                    .unwrap()
+            } else {
+                generate(&mut generator, &path, &messages, &settings)
+            };
+            assert_eq!(retried.cached_tokens, 0, "failure retries cold");
+            assert!(retried.completion_tokens > 0);
+        }
+        println!(
+            "{}",
+            serde_json::json!({"case":index,"mtp":mtp,"stage":format!("{stage:?}"),"row":format!("{row:?}"),"streaming":streaming,"legal":valid,"healthy_preserved":true,"cold_retry":!valid})
+        );
+        // Do not let successful diagnostic requests evict the healthy control.
+        generator.executor.reset().unwrap();
+        generator.checkpoint = None;
+        generator.resident_active = None;
+        generator
+            .resident_sessions
+            .retain(|path, _| path == &healthy_path);
+    }
+    assert!(
+        generator
+            .select_checkpoint(&healthy_path, |_| true)
+            .unwrap()
+            .found
+    );
+    assert!(owner(&mut generator).live.current.is_some());
+}

src/engine/metal/qwen_mtplx/ops.rs

@@ -80,6 +80,7 @@ pub(super) enum Operation {
     Sum(Vec<usize>),
     Min(Vec<usize>),
     Max(Vec<usize>),
+    Any(Vec<usize>),
     RandomBits,
     SearchSorted(bool),
     ArgMax(usize),
@@ -175,6 +176,9 @@ impl Operation {
             Self::Min(axes) => {
                 super::reduce::evaluate(input, output, axes, super::reduce::Reduction::Min)
             }
+            Self::Any(axes) => {
+                super::reduce::evaluate(input, output, axes, super::reduce::Reduction::Or)
+            }
             Self::Max(axes) => {
                 super::reduce::evaluate(input, output, axes, super::reduce::Reduction::Max)
             }

src/engine/metal/qwen_mtplx/owner.rs

@@ -74,92 +74,113 @@ impl QwenExecutor {
         progress: impl FnMut(u32, u32, Option<f32>),
         emit: impl FnMut(bool, String),
     ) -> Result<(GenerationOutput, bool), String> {
-        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,
+        let result = (|| {
+            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 prepared = super::vision::prepare(
-                cancelled,
+            let baseline = self.stats;
+            let stats = &mut self.stats;
+            let result = self.live.generate(
+                &mut self.execution,
                 &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;
+                &self.policy,
+                cancelled,
+                metrics,
+                progress,
+                emit,
+                |event| {
+                    record_speculative_progress(stats, baseline, &event, metrics);
+                    Ok(())
+                },
+            )?;
+            let complete = result
+                .session
+                .as_ref()
+                .is_some_and(|s| s.final_snapshot_error.is_none())
+                && self.live.current.is_some();
+            if let Some(session) = &result.session {
+                let state_error = match &session.decoded {
+                    TurnDecoded::Ar(decoded) => decoded.final_state_error.as_ref(),
+                    TurnDecoded::Mtp(decoded) => decoded.final_state_error.as_ref(),
+                };
+                if let Some(error) = state_error {
+                    return Err(if crate::engine::qwen::is_non_finite_logits(error) {
+                        error.clone()
+                    } else {
+                        format!("Qwen final decode state failed: {error}")
+                    });
+                }
+                if let Some(error) = &session.final_snapshot_error {
+                    return Err(if crate::engine::qwen::is_non_finite_logits(error) {
+                        error.clone()
+                    } else {
+                        format!("Qwen final state capture failed: {error}")
+                    });
+                }
+                if let Some(error) = &session.prefix_snapshot_error {
+                    eprintln!("Qwen prefix capture failed: {error}");
+                }
+                debug_assert_eq!(result.output.cached_tokens as usize, session.cached_tokens);
+                debug_assert!(session.prefill_seconds >= 0.0 && session.restore_seconds >= 0.0);
+                debug_assert!(session.decoded.decode_seconds() >= 0.0);
+                if let TurnDecoded::Mtp(decoded) = &session.decoded
+                    && decoded.pending.is_some()
+                {
+                    return Err("Qwen final state still has a pending correction".into());
+                }
             }
-            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(
-            &mut self.execution,
-            &self.metadata,
-            messages,
-            settings,
-            &self.policy,
-            cancelled,
-            metrics,
-            progress,
-            emit,
-            |event| {
-                record_speculative_progress(stats, baseline, &event, metrics);
-                Ok(())
-            },
-        )?;
-        let complete = result
-            .session
-            .as_ref()
-            .is_some_and(|s| s.final_snapshot_error.is_none())
-            && self.live.current.is_some();
-        if let Some(session) = &result.session {
-            let state_error = match &session.decoded {
-                TurnDecoded::Ar(decoded) => decoded.final_state_error.as_ref(),
-                TurnDecoded::Mtp(decoded) => decoded.final_state_error.as_ref(),
-            };
-            if let Some(error) = state_error {
-                return Err(format!("Qwen final decode state failed: {error}"));
-            }
-            if let Some(error) = &session.final_snapshot_error {
-                return Err(format!("Qwen final state capture failed: {error}"));
-            }
-            if let Some(error) = &session.prefix_snapshot_error {
-                eprintln!("Qwen prefix capture failed: {error}");
-            }
-            debug_assert_eq!(result.output.cached_tokens as usize, session.cached_tokens);
-            debug_assert!(session.prefill_seconds >= 0.0 && session.restore_seconds >= 0.0);
-            debug_assert!(session.decoded.decode_seconds() >= 0.0);
-            if let TurnDecoded::Mtp(decoded) = &session.decoded
-                && decoded.pending.is_some()
-            {
-                return Err("Qwen final state still has a pending correction".into());
+            Ok((result.output, complete))
+        })();
+        if let Err(error) = &result
+            && crate::engine::qwen::is_non_finite_logits(error)
+        {
+            // Drop all snapshots before synchronization, even if the device
+            // reports a second error while draining outstanding work.
+            self.live = LiveSession::default();
+            if let Err(cleanup) = self.clear_execution_cache() {
+                return Err(format!("{error}; Qwen cache cleanup failed: {cleanup}"));
             }
         }
-        Ok((result.output, complete))
+        result
     }
 
     pub(in crate::engine) fn prefill(
@@ -1040,3 +1061,7 @@ fn qwen_vision_chat_checkpoint_preserves_image_identity() {
     assert_ne!(changed.message.content, first.message.content);
     std::fs::remove_dir_all(directory).unwrap();
 }
+
+#[cfg(test)]
+#[path = "issue105_product.rs"]
+mod issue105_product;

src/engine/metal/qwen_mtplx/random.rs

@@ -533,9 +533,16 @@ fn mtplx_sampler_resident_matches_reference() {
             .iter()
             .map(|t| ops::read_i64(t, &streams, gpu).unwrap()[0])
             .collect::<Vec<_>>();
+        // Preserve the retired lane's finite-law/RNG fixture for #104.
+        // #105 replaces only its invalid IDs with the current checked sentinel.
+        let expected = if case["mode"] == "nonfinite" {
+            serde_json::json!(vec![-1; 4])
+        } else {
+            case["tokens"].clone()
+        };
         assert_eq!(
             serde_json::json!(actual),
-            case["tokens"],
+            expected,
             "sampler case {index}: {case}"
         );
         assert_eq!(

src/engine/metal/qwen_mtplx/reduce.rs

@@ -227,6 +227,7 @@ pub(super) enum Reduction {
     Sum,
     Min,
     Max,
+    Or,
 }
 
 impl Reduction {
@@ -245,10 +246,21 @@ impl Reduction {
             Self::Sum => "sum",
             Self::Min => "min",
             Self::Max => "max",
+            Self::Or => "or",
         }
     }
 }
 
+pub(super) fn any(
+    input: &Array,
+    axes: &[i32],
+    keepdims: bool,
+    stream: Stream,
+) -> Result<Array, String> {
+    let input = ops::astype(input, Dtype::Bool, false, stream)?;
+    reduction(&input, axes, keepdims, stream, Reduction::Or)
+}
+
 pub(super) fn sum(
     input: &Array,
     axes: &[i32],
@@ -289,7 +301,7 @@ fn reduction(
     let layout = input.layout();
     let dtype = layout.dtype();
     if !matches!(dtype, Dtype::BF16 | Dtype::F16 | Dtype::F32 | Dtype::I32)
-        && !(dtype == Dtype::Bool && matches!(op, Reduction::Sum))
+        && !(dtype == Dtype::Bool && matches!(op, Reduction::Sum | Reduction::Or))
     {
         return Err("Qwen reduction requires BF16/F16/F32/I32".into());
     }
@@ -303,7 +315,7 @@ fn reduction(
     for &axis in &axes {
         shape[axis] = 1;
     }
-    let out_dtype = if dtype == Dtype::Bool {
+    let out_dtype = if dtype == Dtype::Bool && matches!(op, Reduction::Sum) {
         Dtype::I32
     } else {
         dtype
@@ -319,6 +331,7 @@ fn reduction(
                 Reduction::Sum => ops::Operation::Sum(axes.clone()),
                 Reduction::Min => ops::Operation::Min(axes.clone()),
                 Reduction::Max => ops::Operation::Max(axes.clone()),
+                Reduction::Or => ops::Operation::Any(axes.clone()),
             },
         )?
     };

src/engine/metal/qwen_mtplx/request.rs

@@ -190,6 +190,13 @@ impl LiveSession {
                 Ok(())
             },
         );
+        if result
+            .as_ref()
+            .is_err_and(|error| crate::engine::qwen::is_non_finite_logits(error))
+        {
+            *self = Self::default();
+            return result;
+        }
         // A valid prefill capture survives cancellation and unrelated callback
         // errors. Never tag those partial turns as a completed conversation.
         if let Some(prefix) = captured_prefix {

src/engine/metal/qwen_mtplx/sampling.rs

@@ -3,7 +3,7 @@ use super::array::{Array, Dtype};
 use super::binary::{Binary, binary};
 use super::stream::{Stream, Streams};
 use super::{indexing, normalization, ops, reduce, sort, unary};
-use crate::engine::Rng;
+use crate::engine::{Rng, qwen::non_finite_logits_error};
 use std::borrow::Cow;
 
 #[derive(Clone, Copy)]
@@ -238,7 +238,18 @@ pub(super) fn lazy_sample(
     } else {
         categorical_row(&top_vals, key, stream)?
     };
-    indexing::take(&top_idx, &local, 0, stream)
+    let selected = ops::astype(
+        &indexing::take(&top_idx, &local, 0, stream)?,
+        Dtype::I32,
+        false,
+        stream,
+    )?;
+    let sentinel = Array::new(
+        &[],
+        Dtype::I32,
+        super::scalar_buffer(&(-1_i32).to_le_bytes())?,
+    )?;
+    indexing::select(&invalid_logits(row, stream)?, &sentinel, &selected, stream)
 }
 
 /// generation._distribution_from_mlx_logits, including the single-row GPU
@@ -300,7 +311,7 @@ pub(super) fn distribution_from_logits(
 }
 
 /// The ordinary generation._sample_from_logits path. Greedy evaluates model
-/// logits before constructing/evaluating ArgMax and never advances PCG64.
+/// logits and their validity flag together and never advances PCG64.
 pub(super) fn sample_from_logits(
     logits: &Array,
     config: &SamplerConfig,
@@ -326,8 +337,7 @@ pub(super) fn sample_from_logits(
         } else {
             logits.clone()
         };
-        ops::evaluate(streams, std::slice::from_ref(&logits), stream, false)?;
-        let token = reduce::argmax(&logits, -1, false, stream)?;
+        let token = checked_argmax(&logits, streams, stream)?;
         if token.layout().size() != 1 {
             return Err("greedy sampler requires a scalar token result".into());
         }
@@ -338,6 +348,72 @@ pub(super) fn sample_from_logits(
     }
 }
 
+/// The pinned greedy validity graph: any(isnan(row)) | isinf(max(row)).
+/// A finite maximum permits negative-infinity masks. Read the full row only
+/// on failure; valid turns evaluate the token and scalar fault flag together.
+fn invalid_logits(logits: &Array, stream: Stream) -> Result<Array, String> {
+    let vocab = logits.layout().dim(-1)?;
+    // A valid row must not hide an all-masked neighbour in a batched verify.
+    // Preserve the ordinary scalar reduction when there is only one row.
+    let shape = if logits.layout().size() == vocab as usize {
+        vec![-1]
+    } else {
+        vec![-1, vocab]
+    };
+    let row = ops::reshape(logits, &shape, stream)?;
+    let nan = reduce::any(
+        &binary(&row, &row, Binary::NotEqual, stream)?,
+        &[-1],
+        false,
+        stream,
+    )?;
+    let maximum = reduce::max(&row, &[-1], false, stream)?;
+    let infinity = |value: f32| -> Result<Array, String> {
+        let scalar = ops::astype(
+            &super::random::scalar(value)?,
+            row.layout().dtype(),
+            false,
+            stream,
+        )?;
+        binary(&maximum, &scalar, Binary::Equal, stream)
+    };
+    binary(
+        &nan,
+        &binary(
+            &infinity(f32::INFINITY)?,
+            &infinity(f32::NEG_INFINITY)?,
+            Binary::LogicalOr,
+            stream,
+        )?,
+        Binary::LogicalOr,
+        stream,
+    )
+}
+
+pub(super) fn checked_argmax(
+    logits: &Array,
+    streams: &Streams,
+    stream: Stream,
+) -> Result<Array, String> {
+    let token = reduce::argmax(logits, -1, false, stream)?;
+    let bad = invalid_logits(logits, stream)?;
+    ops::evaluate(streams, &[token.clone(), bad.clone()], stream, false)?;
+    if ops::read_i64(&bad, streams, stream)?
+        .iter()
+        .any(|&bad| bad != 0)
+    {
+        return Err(non_finite_logits_error(
+            &ops::read_f32_as_f64(
+                &ops::astype(logits, Dtype::F32, false, stream)?,
+                streams,
+                stream,
+            )?,
+            "greedy argmax",
+        ));
+    }
+    Ok(token)
+}
+
 /// Sparse maps retain their iteration order. Counts cover generated tokens,
 /// not prompt tokens. A disabled penalty is an identity, including storage.
 pub(super) fn apply_penalties_host<'a>(
@@ -609,7 +685,7 @@ pub(super) struct SparseDistribution {
 
 impl SparseDistribution {
     pub(super) fn new(
-        mut ids: Vec<i64>,
+        ids: Vec<i64>,
         mut probabilities: Vec<f64>,
         vocabulary: usize,
     ) -> Result<Self, String> {
@@ -619,21 +695,15 @@ impl SparseDistribution {
         if probabilities.iter().any(|p| p.is_finite() && *p < 0.0) {
             return Err("sparse distribution probabilities must be nonnegative".into());
         }
+        let original = probabilities.clone();
         for p in &mut probabilities {
             if !p.is_finite() || *p <= 0.0 {
                 *p = 0.0;
             }
         }
-        let mut total = pairwise_sum(&probabilities);
+        let total = pairwise_sum(&probabilities);
         if !total.is_finite() || total <= 0.0 {
-            let token = ids
-                .iter()
-                .copied()
-                .find(|&i| i >= 0 && (i as u64) < vocabulary as u64)
-                .ok_or("sparse distribution has neither positive mass nor a valid token")?;
-            ids = vec![token];
-            probabilities = vec![1.0];
-            total = 1.0;
+            return Err(non_finite_logits_error(&original, "SparseDistribution"));
         }
         for p in &mut probabilities {
             *p /= total;
@@ -978,7 +1048,7 @@ pub(super) fn dense_probabilities(
             .collect::<Vec<_>>();
         let total = pairwise_sum(&p);
         if !total.is_finite() || total <= 0.0 {
-            return Err("cannot normalize logits into a probability distribution".into());
+            return Err(non_finite_logits_error(logits, "softmax"));
         }
         for value in &mut p {
             *value /= total;
@@ -1053,18 +1123,14 @@ pub(super) fn dense_probabilities(
     Ok(probabilities)
 }
 
-/// Original fast_sampling._host_sparse_distribution. Its invalid-mass result
-/// is token zero; the ordinary dense generator path instead propagates error.
+/// fast_sampling._host_sparse_distribution propagates invalid-mass failures.
 fn host_sparse(
     logits: &[f64],
     temperature: f64,
     top_p: f64,
     top_k: i32,
 ) -> Result<SparseDistribution, String> {
-    let probabilities = match dense_probabilities(logits, temperature, top_p, top_k) {
-        Ok(p) => p,
-        Err(_) => return SparseDistribution::new(vec![0], vec![1.0], logits.len()),
-    };
+    let probabilities = dense_probabilities(logits, temperature, top_p, top_k)?;
     let (ids, probabilities) = probabilities
         .into_iter()
         .enumerate()
@@ -1757,13 +1823,16 @@ fn mtplx_sampler_entry_matches_reference() {
                 .iter()
                 .filter(|v| **v == "eval")
                 .count();
-            let expected_commits = if profile == "greedy" {
-                explicit_evals + 1
-            } else if profile == "sparse" && mode == "nonfinite" {
-                explicit_evals - 1
-            } else {
-                explicit_evals
-            };
+            // NumPy's error census implicitly evaluates the F32 cast, which
+            // the mx.eval trace does not include.
+            let expected_commits =
+                if profile == "greedy" && mode == "nonfinite" && dtype != Dtype::F32 {
+                    explicit_evals + 1
+                } else if profile == "sparse" && mode == "nonfinite" {
+                    explicit_evals - 1
+                } else {
+                    explicit_evals
+                };
             assert_eq!(
                 commits() - before,
                 expected_commits,
@@ -2013,15 +2082,21 @@ fn mtplx_sampler_mixed_distributions_match_reference() {
                     .map(|i| (i * (vocab - 1) / (count - 1).max(1)) as i64)
                     .collect::<Vec<_>>();
                 let selected = ids.iter().map(|&i| p[i as usize]).collect();
-                Distribution::Sparse(SparseDistribution::new(ids, selected, vocab).unwrap())
+                SparseDistribution::new(ids, selected, vocab).map(Distribution::Sparse)
             } else {
                 let total = pairwise_sum(&p);
-                Distribution::Dense(p.iter().map(|p| p / total).collect())
+                Ok(Distribution::Dense(p.iter().map(|p| p / total).collect()))
             }
         };
         let kinds = case["kinds"].as_str().unwrap().as_bytes();
         let target = make(p, kinds[0]);
         let draft = make(q, kinds[1]);
+        if case["construction_error"] == true {
+            assert!(target.is_err() || draft.is_err());
+            continue;
+        }
+        let target = target.unwrap();
+        let draft = draft.unwrap();
         let mut rng = Rng::new_qwen(12345);
         let samples = (0..32)
             .map(|_| target.sample(&mut rng))
@@ -2176,9 +2251,14 @@ fn mtplx_sampler_fallback_matches_reference() {
                 .unwrap(),
             )
             .unwrap();
-            let result = sparse_distributions(&logits, 0.6, top_p, 20, &streams, gpu)
-                .unwrap()
-                .unwrap();
+            let result = sparse_distributions(&logits, 0.6, top_p, 20, &streams, gpu);
+            if case["error"] == true {
+                assert!(crate::engine::qwen::is_non_finite_logits(
+                    &result.err().unwrap()
+                ));
+                continue;
+            }
+            let result = result.unwrap().unwrap();
             assert_eq!(result.len(), 4);
             for (row, distribution) in result.iter().enumerate() {
                 assert_eq!(
@@ -2304,10 +2384,10 @@ fn mtplx_sampler_distribution_and_rng_match_reference() {
     assert!(SparseDistribution::new(vec![], vec![], 10).is_err());
     assert!(SparseDistribution::new(vec![1], vec![-0.1], 10).is_err());
     assert!(SparseDistribution::new(vec![-1, 10], vec![0.0, f64::NAN], 10).is_err());
-    let fallback =
-        SparseDistribution::new(vec![-1, 7, 3], vec![f64::NAN, f64::INFINITY, 0.0], 10).unwrap();
-    assert_eq!(fallback.ids, [7]);
-    assert_eq!(fallback.probabilities, [1.0]);
+    let failure = SparseDistribution::new(vec![-1, 7, 3], vec![f64::NAN, f64::INFINITY, 0.0], 10)
+        .err()
+        .unwrap();
+    assert!(crate::engine::qwen::is_non_finite_logits(&failure));
     eprintln!(
         "MTPLX distributions: {accepted} accepts and {rejected} corrections; exact probabilities, samples and subsequent RNG draws"
     );

src/engine/metal/qwen_mtplx/session_cache.rs

@@ -1158,6 +1158,11 @@ impl Execution {
         } = self.prepare_session(prompt, saved, options, |event| {
             observe(TurnProgress::Prefill(event))
         })?;
+        #[cfg(test)]
+        let prepared = super::prefill::Prepared {
+            logits: self.inject_logits(super::issue105::Stage::Prefill, &prepared.logits)?,
+            ..prepared
+        };
         // Publish before decode (and before prefix storage) so a cancelled
         // response retains the completed prompt timings, just like MTPLX.
         observe(TurnProgress::PromptReady(crate::metrics::PromptTiming {
@@ -1183,6 +1188,10 @@ impl Execution {
                 .and_then(store_prefix)
             })
             .and_then(Result::err);
+        #[cfg(test)]
+        if options.store_min_suffix.is_some() && prefix_snapshot_error.is_none() {
+            self.captured_prefixes += 1;
+        }
         observe(TurnProgress::Decode)?;
         let decoded = match &options.decode {
             DecodeOptions::Ar(decode) => TurnDecoded::Ar(self.decode_ar(
@@ -1216,7 +1225,13 @@ impl Execution {
             for &id in decoded.tokens() {
                 committed.push(i32::try_from(id).map_err(|_| "generated token exceeds int32")?);
             }
-            Some(self.capture_session(&committed, logits, hidden, history, options.lazy_kv))
+            #[cfg(test)]
+            let result = self.inject_capture_failure(logits).and_then(|()| {
+                self.capture_session(&committed, logits, hidden, history, options.lazy_kv)
+            });
+            #[cfg(not(test))]
+            let result = self.capture_session(&committed, logits, hidden, history, options.lazy_kv);
+            Some(result)
         } else {
             None
         };

src/engine/metal/qwen_mtplx/speculative.rs

@@ -440,7 +440,7 @@ impl Execution {
         let logits = verified.logits.ok_or("copy verify omitted logits")?;
         let hidden = verified.hidden.ok_or("copy verify omitted hidden")?;
         let greedy = if sampler.temperature <= 0.0 {
-            let choices = super::reduce::argmax(&logits, -1, false, self.stream)?;
+            let choices = sampling::checked_argmax(&logits, &self.streams, self.stream)?;
             Some(ops::read_i64(&choices, &self.streams, self.stream)?)
         } else {
             ops::evaluate(