From e7425557bdca3953d7392e8e92270e4921e54ba3 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 17 Aug 2026 23:53:23 +0900 Subject: [PATCH] fix(server): gate tile-aligned prefill on the model, not on the hardware alone execute_full_prefill and both chunked-prefill paths padded the prompt to a 32-token Neural Accelerator tile whenever the hardware had one, without asking whether the model tolerates it. Every hybrid and recurrent family answers supports_padded_prefill() == false, because pad positions are absorbed by a conv / SSM / GatedDeltaNet state that trimming the KV caches afterwards does not rewind. So every such model served on an M5-class host had its recurrent state corrupted whenever a prompt or chunk was not already tile-aligned. The same file knows the hazard. The boundary-snapshot path declines to pad and says why in a comment, and the batched path reads the predicate into can_pad_prefill. These three sites simply never asked. Measured on M5 Max with qwen3.8-27b-4bit, whose backbone answers false: the same /v1/completions request at temperature 0 returns different text before and after this change, diverging at character 74 of the completion. The after side is the correct one. This is the server-side twin of #1201, and wider: that one needed a VLM wrapper to lose the predicate by not delegating it, while this one reaches every hybrid model directly, VLM-wrapped or not. The guard test added with #1202 grows a second case for the second way to lose a guard. The first covers a type that answers the trait default instead of forwarding; this covers a call site that never asks. It accepts a local bound to the predicate as the guard, since the batched path hoists it fourteen lines up, and skips the definition and its own unit tests. Verified to fail by removing one of the three guards. --- src/server/batch/scheduler.rs | 106 +++++++++++---------- tests/vlm_wrapper_capability_delegation.rs | 86 +++++++++++++++++ 2 files changed, 141 insertions(+), 51 deletions(-) diff --git a/src/server/batch/scheduler.rs b/src/server/batch/scheduler.rs index 9ff1b75fd..4278662e2 100644 --- a/src/server/batch/scheduler.rs +++ b/src/server/batch/scheduler.rs @@ -5851,27 +5851,29 @@ impl BatchScheduler { // text backbone then builds a causal mask sized to the embeddings, // matching the CLI generate path. Token-id (text-only) prefill — for // VLMs and plain text models alike — is unaffected. - let (effective_tokens, pad_mask_opt) = - if should_align_prefill() && seq.vlm_embeddings.is_none() { - let padded_len = align_to_na_tile(actual_len); - if padded_len > actual_len { - let mut padded = suffix_tokens.clone(); - padded.resize(padded_len, 0); - // The padding mask anchors to the adopted cache offset so - // the newly-prefilled positions see the correct KV-history - // positions on M5+ hardware. - let mask = create_padded_prefill_mask( - actual_len as i32, - padded_len as i32, - seq.prefill_start_offset as i32, - ); - (padded, Some(mask)) - } else { - (suffix_tokens.clone(), None) - } + let (effective_tokens, pad_mask_opt) = if self.model.supports_padded_prefill() + && should_align_prefill() + && seq.vlm_embeddings.is_none() + { + let padded_len = align_to_na_tile(actual_len); + if padded_len > actual_len { + let mut padded = suffix_tokens.clone(); + padded.resize(padded_len, 0); + // The padding mask anchors to the adopted cache offset so + // the newly-prefilled positions see the correct KV-history + // positions on M5+ hardware. + let mask = create_padded_prefill_mask( + actual_len as i32, + padded_len as i32, + seq.prefill_start_offset as i32, + ); + (padded, Some(mask)) } else { (suffix_tokens.clone(), None) - }; + } + } else { + (suffix_tokens.clone(), None) + }; let eff_len = effective_tokens.len() as i32; let input = mlxcel_core::from_slice_i32(&effective_tokens, &[1, eff_len]); @@ -6024,25 +6026,26 @@ impl BatchScheduler { // Align the first chunk to a 32-token tile boundary on M5+ hardware. let actual_chunk_len = chunk.len(); - let (eff_chunk, pad_mask_opt) = if should_align_prefill() { - let padded_len = align_to_na_tile(actual_chunk_len); - if padded_len > actual_chunk_len { - let mut padded = chunk.to_vec(); - padded.resize(padded_len, 0); - // Mask anchored to the KV offset the adopted prefix already - // installed (starts at zero for cold prefills). - let mask = create_padded_prefill_mask( - actual_chunk_len as i32, - padded_len as i32, - start as i32, - ); - (padded, Some(mask)) + let (eff_chunk, pad_mask_opt) = + if self.model.supports_padded_prefill() && should_align_prefill() { + let padded_len = align_to_na_tile(actual_chunk_len); + if padded_len > actual_chunk_len { + let mut padded = chunk.to_vec(); + padded.resize(padded_len, 0); + // Mask anchored to the KV offset the adopted prefix already + // installed (starts at zero for cold prefills). + let mask = create_padded_prefill_mask( + actual_chunk_len as i32, + padded_len as i32, + start as i32, + ); + (padded, Some(mask)) + } else { + (chunk.to_vec(), None) + } } else { (chunk.to_vec(), None) - } - } else { - (chunk.to_vec(), None) - }; + }; let eff_len = eff_chunk.len() as i32; let input = mlxcel_core::from_slice_i32(&eff_chunk, &[1, eff_len]); @@ -6241,23 +6244,24 @@ impl BatchScheduler { offset as i32 } }; - let (eff_chunk, pad_mask_opt) = if should_align_prefill() { - let padded_len = align_to_na_tile(actual_chunk_len); - if padded_len > actual_chunk_len { - let mut padded = chunk.to_vec(); - padded.resize(padded_len, 0); - let mask = create_padded_prefill_mask( - actual_chunk_len as i32, - padded_len as i32, - kv_offset, - ); - (padded, Some(mask)) + let (eff_chunk, pad_mask_opt) = + if self.model.supports_padded_prefill() && should_align_prefill() { + let padded_len = align_to_na_tile(actual_chunk_len); + if padded_len > actual_chunk_len { + let mut padded = chunk.to_vec(); + padded.resize(padded_len, 0); + let mask = create_padded_prefill_mask( + actual_chunk_len as i32, + padded_len as i32, + kv_offset, + ); + (padded, Some(mask)) + } else { + (chunk.to_vec(), None) + } } else { (chunk.to_vec(), None) - } - } else { - (chunk.to_vec(), None) - }; + }; let eff_len = eff_chunk.len() as i32; let input = mlxcel_core::from_slice_i32(&eff_chunk, &[1, eff_len]); diff --git a/tests/vlm_wrapper_capability_delegation.rs b/tests/vlm_wrapper_capability_delegation.rs index 3264719bb..65ef129e7 100644 --- a/tests/vlm_wrapper_capability_delegation.rs +++ b/tests/vlm_wrapper_capability_delegation.rs @@ -193,3 +193,89 @@ fn a_vision_wrapper_delegates_padded_prefill_when_its_backbone_refuses_it() { violations.join("\n ") ); } + +/// Every tile-aligned padded prefill must be guarded on the model predicate. +/// +/// The wrapper test above covers one way to lose the guard: a type that +/// answers the trait default instead of forwarding. This covers the other: +/// a call site that never asks. `src/server/batch/scheduler.rs` had three, +/// `execute_full_prefill` and both chunked-prefill paths, each gated on +/// hardware alone, so every hybrid model served on Neural Accelerator +/// hardware had its recurrent state corrupted by pad positions. The same +/// file's boundary-snapshot path documents the hazard in a comment and +/// declines to pad, which is what makes the omission a slip rather than a +/// disagreement. +/// +/// The rule is deliberately shallow: the guard must appear within the dozen +/// lines before the call, either as the predicate itself or as a local bound +/// to it (`let can_pad_prefill = self.model.supports_padded_prefill()`, which +/// the batched path does read fourteen lines up). A guard further away than +/// that is one a reader cannot see either. +#[test] +fn every_tile_aligned_prefill_is_guarded_on_the_model_predicate() { + const WINDOW: usize = 12; + let mut files = Vec::new(); + rust_sources(&repo_root().join("src"), &mut files); + files.sort(); + + let mut unguarded = Vec::new(); + for file in &files { + let Ok(src) = fs::read_to_string(file) else { + continue; + }; + let lines: Vec<&str> = src.lines().collect(); + // Locals bound to the predicate count as the guard, so a site that + // hoists it out of a loop is not reported as unguarded. + let mut aliases: Vec = Vec::new(); + for line in &lines { + let Some(rest) = line.trim().strip_prefix("let ") else { + continue; + }; + if !line.contains("supports_padded_prefill") { + continue; + } + let name: String = rest + .chars() + .take_while(|c| c.is_alphanumeric() || *c == '_') + .collect(); + if !name.is_empty() { + aliases.push(name); + } + } + for (i, line) in lines.iter().enumerate() { + let trimmed = line.trim_start(); + // The definition, its doc examples and its own unit tests are not + // prefill sites and have nothing to guard. + if !line.contains("align_to_na_tile(") + || trimmed.starts_with("//") + || trimmed.starts_with("fn ") + || trimmed.starts_with("pub fn ") + || trimmed.starts_with("assert") + { + continue; + } + let start = i.saturating_sub(WINDOW); + if lines[start..=i].iter().any(|l| { + l.contains("supports_padded_prefill") + || aliases.iter().any(|a| l.contains(a.as_str())) + }) { + continue; + } + unguarded.push(format!( + "{}:{}: {}", + file.strip_prefix(repo_root()).unwrap_or(file).display(), + i + 1, + line.trim() + )); + } + } + + assert!( + unguarded.is_empty(), + "a tile-aligned prefill must be gated on the model's own \ + `supports_padded_prefill()`, not on hardware alone: padding a hybrid \ + model corrupts its conv / SSM / GatedDeltaNet recurrent state, which \ + trimming the KV caches afterwards does not undo. Unguarded sites:\n {}", + unguarded.join("\n ") + ); +}