Skip to content

fix(server): gate tile-aligned prefill on the model, not on the hardware alone #1205

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
inureyes merged 1 commit into from
Aug 17, 2026
Merged

fix(server): gate tile-aligned prefill on the model, not on the hardware alone #1205

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 55 additions & 51 deletions src/server/batch/scheduler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]);
Expand Down Expand Up @@ -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]);
Expand Down Expand Up @@ -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]);
Expand Down
86 changes: 86 additions & 0 deletions tests/vlm_wrapper_capability_delegation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String> = 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 ")
);
}