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