diff --git a/TECHNICAL_REPORTS/2133-youtu-vl-image-pad-expansion-20261006.en.md b/TECHNICAL_REPORTS/2133-youtu-vl-image-pad-expansion-20261006.en.md new file mode 100644 index 000000000..da4548f5d --- /dev/null +++ b/TECHNICAL_REPORTS/2133-youtu-vl-image-pad-expansion-20261006.en.md @@ -0,0 +1,90 @@ +# Technical Report: PR #2133 - fix(youtu_vl): expand the template image placeholder per merged feature + +**Date**: 2026-10-06 +**Author**: mlxcel maintainers +**Reviewer**: implementation review cycle +**Status**: Completed (validated on the real checkpoint on CUDA through the CLI and mlxcel-server) +**Languages**: Rust +**Risk Level**: Low (one family's prompt preparation; the new mismatch check turns a silent truncation into a request error) + +--- + +## Executive Summary + +After #1610 fixed the patch emission order, Youtu-VL still described both three-shape fixtures as "a single circle" while the solid orange control read correctly. Issue #1618 pointed at the windowed-attention path because the only correct fixture was the only single-window one. The real cause was upstream of the vision tower: the chat template emits one `<|image_pad|>` per image, mlxcel never expanded it to one token per merged feature, and the masked scatter therefore put only the first feature into the prompt. PR #2133 expands the placeholder the way the checkpoint's processor does and makes a count mismatch an error. CLI and server output now match the transformers greedy output word for word on all three fixtures. + +--- + +## 1. Problem Statement + +### 1.1 Background + +`YoutuVLProcessor.__call__` in the checkpoint's `processing_youtu_vl.py` replaces each `<|image_pad|>` in the rendered text with `h * w / 4` copies, one per merged vision feature. mlxcel's equivalent, `insert_youtu_vl_image_tokens`, was written for callers that pass no placeholder at all and splices a framed run after BOS. Its guard returned early whenever the prompt already contained an image token, on the assumption that such a prompt was already expanded. + +### 1.2 Existing Issues + +The chat template always renders the placeholder, so the guard fired on every templated request. The prompt reached the tower with one image token (`128262, 128264, 128263`) against 49, 121 or 196 features. `merge_llava`'s masked scatter fills placeholders in order and drops what does not fit, so the language model saw exactly one feature: the top-left merged patch. + +### 1.3 Why it looked like a window defect + +A uniform image produces the same content in every merged token, so the 224 orange fixture read correctly from its one surviving token. The two shape fixtures both have light grey corners, which the model turned into "a single circle on a plain background" in two polarities. Grid size and window count correlated with correctness only because the uniform fixture happened to be the small one. + +--- + +## 2. Technical Review + +### 2.1 Stage-by-stage diff + +A transformers 4.56.0 oracle (the checkpoint's remote code, f32 on CPU) produced the reference processor output and the tower plus merger output. mlxcel's processor matched to 3e-8 at 224 and 448 (336 differs only by Lanczos versus bilinear resampling). Feeding the reference pixels through mlxcel's bf16 tower gave per-merged-token cosine mean 0.9999, 0.998 and 0.997 at 14x14, 22x22 and 28x28 patch grids. That rules out `get_window_index`, the `cu_window_seqlens` padding and deduplication, the full-attention boundaries, and the window inverse, which were the issue's candidates. + +### 2.2 Locating the defect + +The reference generated the correct answers in bf16, and the vision features matched, so the divergence was after the tower. Dumping the prompt ids showed one image token where the reference has 196. + +--- + +## 3. Technical Decisions + +### 3.1 Mirror the processor's replace loop + +One placeholder per image is expanded in place, keeping the template's own start/end framing. A prompt that already carries one token per feature is left alone, checked first so a grid whose runs are a single token is not expanded twice. No placeholder keeps the existing splice-after-BOS fallback. Any other count is an error rather than a guess. + +### 3.2 Fail on a placeholder/feature mismatch + +Upstream raises when image tokens and features disagree. `YoutuVLModel::get_input_embeddings` now does the same, so this class of defect surfaces as a request error instead of a fluent description of one patch. + +--- + +## 4. Validation + +| fixture | before | after (CLI and server) | +|---|---|---| +| 224 orange | "a solid, uniform orange color" | unchanged | +| 336 shapes | "a single white circle on a black background" | red square, blue circle, green triangle (identical to transformers) | +| 448 shapes | "a single, solid black circle on a plain white background" | "- Red square - Blue circle - Green triangle" (identical to transformers) | + +The two formerly known-failing tests in `tests/youtu_vl_parity.rs` pass on the real checkpoint. Server `prompt_tokens` are 82, 154 and 229, equal to the reference processor's counts. New unit tests cover expansion on a 28x28 grid, multi-image ordering, the mismatch error, and `get_window_index` against reference values on multi-window grids. + +--- + +## 5. Change Summary + +- `src/multimodal/youtu_vl_prompt.rs`, `youtu_vl_prompt_tests.rs`: placeholder expansion and tests. +- `src/vision/youtu_vl.rs`: count check before the scatter. +- `src/multimodal/vlm_runtime.rs`: both errors propagate as request errors. +- `src/vision/encoders/youtu_vl_tests.rs`: window-index reference test. +- `tests/youtu_vl_parity.rs`: un-gated, comments corrected. + +Related: #1618 (closed by this PR), #1610, #1600, #1611. + +--- + +## 6. Follow-up Actions + +### Transferable lesson + +A solid-color fixture cannot detect a defect in how features reach prompt slots: any single feature carries the whole image. Content tests need a spatially varied fixture, and a VLM port should assert placeholder count equals feature count rather than relying on a scatter that truncates. + +### Open + +The processor's `smart_resize` rounds edges to a multiple of 32 where the reference rounds up, so some input sizes other than these fixtures get a different patch grid. #1611's premise (cap at `max_num_patches` = 256) is contradicted by `YoutuVLProcessor.__call__`, which passes `max_image_patches=36864` and overrides the preprocessor config. diff --git a/TECHNICAL_REPORTS/2133-youtu-vl-image-pad-expansion-20261006.ko.md b/TECHNICAL_REPORTS/2133-youtu-vl-image-pad-expansion-20261006.ko.md new file mode 100644 index 000000000..c3217e169 --- /dev/null +++ b/TECHNICAL_REPORTS/2133-youtu-vl-image-pad-expansion-20261006.ko.md @@ -0,0 +1,90 @@ +# 기술 보고서: PR #2133 - fix(youtu_vl): 템플릿 이미지 placeholder를 병합 feature 수만큼 확장 + +**날짜**: 2026-10-06 +**작성자**: mlxcel maintainers +**리뷰어**: implementation review cycle +**상태**: 완료 (CUDA에서 실제 체크포인트로 CLI와 mlxcel-server 모두 검증) +**언어**: Rust +**위험도**: 낮음 (한 모델 계열의 프롬프트 준비 경로만 변경, 새 개수 검사는 조용한 잘림을 요청 오류로 바꿈) + +--- + +## 요약 + +#1610이 패치 방출 순서를 고친 뒤에도 Youtu-VL은 세 도형 fixture 두 개를 모두 "원 하나"로 설명했고, 단색 주황 대조 이미지만 올바르게 읽었습니다. 올바른 fixture가 유일하게 단일 attention window 크기였기 때문에 #1618은 windowed attention 경로를 의심했습니다. 실제 원인은 vision tower 앞단에 있었습니다. 채팅 템플릿은 이미지마다 `<|image_pad|>` 하나를 렌더링하는데 mlxcel은 이를 병합 feature 수만큼 확장하지 않았고, masked scatter가 첫 번째 feature 하나만 프롬프트에 넣었습니다. PR #2133은 체크포인트의 processor와 같은 방식으로 placeholder를 확장하고, 개수가 맞지 않으면 오류로 처리합니다. 세 fixture 모두에서 CLI와 서버 출력이 transformers greedy 출력과 단어 단위로 일치합니다. + +--- + +## 1. 문제 정의 + +### 1.1 배경 + +체크포인트의 `processing_youtu_vl.py`에 있는 `YoutuVLProcessor.__call__`은 렌더링된 텍스트의 각 `<|image_pad|>`를 병합 vision feature 하나당 하나씩, 즉 `h * w / 4`개로 바꿉니다. mlxcel의 대응 함수 `insert_youtu_vl_image_tokens`는 placeholder 없이 이미지만 넘기는 호출자를 위해 작성되어 BOS 뒤에 framing된 토큰 열을 삽입합니다. 이 함수는 프롬프트에 이미지 토큰이 하나라도 있으면 이미 확장된 것으로 보고 바로 반환했습니다. + +### 1.2 기존 문제 + +채팅 템플릿은 항상 placeholder를 렌더링하므로 이 조건이 모든 템플릿 요청에서 걸렸습니다. 프롬프트에는 이미지 토큰이 하나(`128262, 128264, 128263`)뿐이었고 feature는 49, 121, 196개였습니다. `merge_llava`의 masked scatter는 placeholder를 순서대로 채우고 남는 것은 버리므로, 언어 모델은 왼쪽 위 병합 패치 하나만 보았습니다. + +### 1.3 window 결함처럼 보인 이유 + +균일한 이미지는 모든 병합 토큰이 같은 내용을 담으므로 224 주황 fixture는 남은 토큰 하나로도 올바르게 읽혔습니다. 두 도형 fixture는 모서리가 모두 밝은 회색이라 모델은 "단색 배경 위의 원 하나"를 두 가지 극성으로 답했습니다. 격자 크기와 window 수가 정답 여부와 상관관계를 보인 것은 균일한 fixture가 마침 작은 이미지였기 때문입니다. + +--- + +## 2. 기술 검토 + +### 2.1 단계별 비교 + +transformers 4.56.0 oracle(체크포인트 원격 코드, CPU f32)로 참조 processor 출력과 tower+merger 출력을 만들었습니다. mlxcel processor는 224와 448에서 3e-8 이내로 일치했고 336은 Lanczos와 bilinear 리샘플링 차이뿐이었습니다. 참조 픽셀을 mlxcel bf16 tower에 넣었을 때 병합 토큰별 cosine 평균은 14x14, 22x22, 28x28 패치 격자에서 각각 0.9999, 0.998, 0.997이었습니다. 이로써 이슈가 지목한 `get_window_index`, `cu_window_seqlens` 패딩과 중복 제거, full-attention 경계, window 역순열이 모두 배제됩니다. + +### 2.2 결함 위치 특정 + +참조 구현은 bf16에서도 정답을 생성했고 vision feature도 일치했으므로 차이는 tower 이후에 있었습니다. 프롬프트 id를 덤프하자 참조 구현에서 196개인 이미지 토큰이 mlxcel에는 하나뿐이었습니다. + +--- + +## 3. 기술적 결정 + +### 3.1 processor의 치환 루프를 그대로 따름 + +이미지당 placeholder 하나는 템플릿의 시작/끝 framing을 유지한 채 그 자리에서 확장합니다. feature당 토큰이 이미 있는 프롬프트는 그대로 두며, 토큰 열이 하나씩인 격자를 두 번 확장하지 않도록 이 조건을 먼저 검사합니다. placeholder가 없으면 기존처럼 BOS 뒤에 삽입합니다. 그 밖의 개수는 추측하지 않고 오류로 처리합니다. + +### 3.2 placeholder와 feature 개수 불일치 시 실패 + +upstream은 이미지 토큰과 feature 수가 다르면 예외를 냅니다. 이제 `YoutuVLModel::get_input_embeddings`도 같은 검사를 하므로 이런 결함은 패치 하나에 대한 유창한 설명 대신 요청 오류로 드러납니다. + +--- + +## 4. 검증 + +| fixture | 이전 | 이후 (CLI와 서버) | +|---|---|---| +| 224 주황 | "a solid, uniform orange color" | 변화 없음 | +| 336 도형 | "a single white circle on a black background" | 빨간 사각형, 파란 원, 초록 삼각형 (transformers와 동일) | +| 448 도형 | "a single, solid black circle on a plain white background" | "- Red square - Blue circle - Green triangle" (transformers와 동일) | + +`tests/youtu_vl_parity.rs`에서 실패로 알려졌던 두 테스트가 실제 체크포인트에서 통과합니다. 서버의 `prompt_tokens`는 82, 154, 229로 참조 processor의 토큰 수와 같습니다. 새 단위 테스트는 28x28 격자 확장, 다중 이미지 순서, 불일치 오류, 다중 window 격자에서 참조값과의 `get_window_index` 일치를 다룹니다. + +--- + +## 5. 변경 요약 + +- `src/multimodal/youtu_vl_prompt.rs`, `youtu_vl_prompt_tests.rs`: placeholder 확장과 테스트. +- `src/vision/youtu_vl.rs`: scatter 전 개수 검사. +- `src/multimodal/vlm_runtime.rs`: 두 오류를 요청 오류로 전파. +- `src/vision/encoders/youtu_vl_tests.rs`: window index 참조 테스트. +- `tests/youtu_vl_parity.rs`: ignore 사유 정리, 주석 수정. + +관련: #1618 (이 PR로 종료), #1610, #1600, #1611. + +--- + +## 6. 후속 조치 + +### 일반화할 교훈 + +단색 fixture는 feature가 프롬프트 슬롯에 도달하는 방식의 결함을 잡지 못합니다. feature 하나가 이미지 전체를 대표하기 때문입니다. 내용 테스트에는 공간적으로 다양한 fixture가 필요하고, VLM 포팅은 잘림을 허용하는 scatter에 기대지 말고 placeholder 수와 feature 수가 같은지 확인해야 합니다. + +### 남은 항목 + +processor의 `smart_resize`는 가장자리를 32의 배수로 반올림하지만 참조 구현은 올림하므로, 이번 fixture가 아닌 일부 입력 크기에서는 패치 격자가 달라집니다. #1611의 전제(`max_num_patches` = 256으로 제한)는 `YoutuVLProcessor.__call__`이 `max_image_patches=36864`를 넘겨 preprocessor 설정을 덮어쓴다는 사실과 모순됩니다. diff --git a/src/multimodal/vlm_runtime.rs b/src/multimodal/vlm_runtime.rs index 6710789de..c490cb7f5 100644 --- a/src/multimodal/vlm_runtime.rs +++ b/src/multimodal/vlm_runtime.rs @@ -1850,9 +1850,9 @@ where .try_preprocess_with_spatial(images) .map_err(|err| anyhow::anyhow!("Youtu-VL image preprocessing failed: {err}"))?; - // Splice image-token runs into the prompt when they aren't already - // present. We mirror the upstream ` + image*N + - // ` framing. + // Expand the template's `<|image_pad|>` placeholders (or splice + // framed runs when there are none) so the prompt carries one image + // token per merged feature, as the checkpoint's processor does. let preparation = insert_youtu_vl_image_tokens( prompt_tokens, &spatial_shapes, @@ -1861,6 +1861,7 @@ where youtu.vision_end_token_id, youtu.image_token_id, ) + .map_err(|err| anyhow::anyhow!("Youtu-VL prompt preparation failed: {err}"))? .map(|stats| VlmPreparationSummary::YoutuVL { image_blocks: stats.image_blocks, total_image_tokens: stats.total_image_tokens, @@ -1875,8 +1876,9 @@ where let _ = image_cache_keys; let input_ids_arr = prompt_ids_array(prompt_tokens); - let embeddings = - youtu.get_input_embeddings(&input_ids_arr, &pixel_values, &spatial_shapes); + let embeddings = youtu + .get_input_embeddings(&input_ids_arr, &pixel_values, &spatial_shapes) + .map_err(|err| anyhow::anyhow!(err))?; Ok(Some(PreparedVlmEmbeddings { embeddings, diff --git a/src/multimodal/youtu_vl_prompt.rs b/src/multimodal/youtu_vl_prompt.rs index 3b13de8e5..eb6664b6a 100644 --- a/src/multimodal/youtu_vl_prompt.rs +++ b/src/multimodal/youtu_vl_prompt.rs @@ -14,12 +14,28 @@ //! Youtu-VL prompt token insertion. //! -//! Mirrors upstream's prompt-format expectation: each image consumes -//! `(h_patches / merge) * (w_patches / merge)` `image_token_id` tokens, -//! framed by `vision_start_token_id` and `vision_end_token_id`. When the -//! prompt does not already carry an `image_token_id`, the run is spliced in -//! after the BOS token (mirroring the Qwen-VL fallback behaviour for callers -//! that pass only an image without an explicit placeholder). +//! Each image consumes `(h_patches / merge) * (w_patches / merge)` +//! `image_token_id` tokens, one per merged vision feature. The checkpoint's +//! chat template renders a single `<|vision_start|><|image_pad|><|vision_end|>` +//! per image, and the checkpoint's own processor (`YoutuVLProcessor.__call__` +//! in `processing_youtu_vl.py`) then replaces each `<|image_pad|>`, in order, +//! with that image's full run. This module performs the same rewrite on token +//! ids: +//! +//! - one placeholder per image: each is expanded in place to its image's run, +//! keeping the template's own start/end framing; +//! - already expanded (one token per feature): left as is; +//! - no placeholder at all (a caller that passed only an image): one framed run +//! per image is spliced in after the BOS token, mirroring the Qwen-VL +//! fallback; +//! - any other placeholder count is an error. +//! +//! Issue #1618: the first case used to be read as "already expanded", so the +//! prompt kept one image token while the tower produced one feature per merged +//! patch. The embedding scatter then placed only the first (top-left) feature, +//! and every multi-object image was described from its corner. + +use thiserror::Error; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct InsertedYoutuVlmTokens { @@ -27,10 +43,24 @@ pub struct InsertedYoutuVlmTokens { pub total_image_tokens: i32, } -/// Insert image-placeholder runs into `prompt_tokens` based on each image's -/// `(h_patches, w_patches)`. Returns `None` when the inputs are empty or -/// the prompt already contains expanded image tokens (in which case the -/// caller should not re-expand). +/// The prompt's image placeholders cannot be matched to the images supplied. +#[derive(Debug, Clone, PartialEq, Eq, Error)] +#[error( + "Youtu-VL prompt carries {found} image placeholder token(s), but {images} image(s) need either one placeholder each or {expected} expanded tokens" +)] +pub struct YoutuVlPlaceholderMismatch { + pub found: usize, + pub images: usize, + pub expected: usize, +} + +/// Rewrite `prompt_tokens` so it carries exactly one `image_token_id` per +/// merged vision feature, given each image's `(h_patches, w_patches)`. +/// +/// Returns `Ok(Some(stats))` when the prompt was rewritten, `Ok(None)` when +/// there was nothing to do (empty inputs, or the prompt was already expanded), +/// and an error when the placeholder count matches neither the image count nor +/// the expanded token count. pub fn insert_youtu_vl_image_tokens( prompt_tokens: &mut Vec, spatial_shapes: &[(i32, i32)], @@ -38,26 +68,62 @@ pub fn insert_youtu_vl_image_tokens( vision_start_token_id: i32, vision_end_token_id: i32, image_token_id: i32, -) -> Option { - if prompt_tokens.is_empty() - || spatial_shapes.is_empty() - || prompt_tokens.contains(&image_token_id) - || spatial_merge_size == 0 - { - return None; +) -> Result, YoutuVlPlaceholderMismatch> { + if prompt_tokens.is_empty() || spatial_shapes.is_empty() || spatial_merge_size == 0 { + return Ok(None); } let merge = spatial_merge_size as i32; - let mut image_tokens: Vec = Vec::new(); - let mut total_image_tokens = 0i32; + let per_image: Vec = spatial_shapes + .iter() + .map(|&(h, w)| (h / merge) * (w / merge)) + .collect(); + let total_image_tokens: i32 = per_image.iter().sum(); + let expected = usize::try_from(total_image_tokens).unwrap_or(0); + let stats = InsertedYoutuVlmTokens { + image_blocks: spatial_shapes.len(), + total_image_tokens, + }; - for &(h, w) in spatial_shapes { - let tokens_per_image = (h / merge) * (w / merge); - total_image_tokens += tokens_per_image; - image_tokens.push(vision_start_token_id); - for _ in 0..tokens_per_image { - image_tokens.push(image_token_id); + let found = prompt_tokens + .iter() + .filter(|&&t| t == image_token_id) + .count(); + + // Already one token per feature. Checked before the per-image case so a + // grid whose runs are a single token each is not expanded twice. + if found == expected { + return Ok(None); + } + + if found == spatial_shapes.len() { + let mut expanded = Vec::with_capacity(prompt_tokens.len() + expected); + let mut runs = per_image.iter(); + for &token in prompt_tokens.iter() { + // `found == runs.len()`, so every placeholder has a run. + match (token == image_token_id).then(|| runs.next()).flatten() { + Some(&count) => { + expanded.extend(std::iter::repeat_n(image_token_id, count.max(0) as usize)) + } + None => expanded.push(token), + } } + *prompt_tokens = expanded; + return Ok(Some(stats)); + } + + if found > 0 { + return Err(YoutuVlPlaceholderMismatch { + found, + images: spatial_shapes.len(), + expected, + }); + } + + let mut image_tokens: Vec = Vec::with_capacity(expected + 2 * per_image.len()); + for &count in &per_image { + image_tokens.push(vision_start_token_id); + image_tokens.extend(std::iter::repeat_n(image_token_id, count.max(0) as usize)); image_tokens.push(vision_end_token_id); } @@ -67,58 +133,9 @@ pub fn insert_youtu_vl_image_tokens( prompt_tokens.extend(image_tokens); prompt_tokens.extend(rest); - Some(InsertedYoutuVlmTokens { - image_blocks: spatial_shapes.len(), - total_image_tokens, - }) + Ok(Some(stats)) } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn skips_when_prompt_already_has_image_tokens() { - let mut prompt = vec![1, 100, 100, 200]; - let result = insert_youtu_vl_image_tokens(&mut prompt, &[(4, 4)], 2, 128_262, 128_263, 100); - assert!(result.is_none()); - // Prompt unchanged. - assert_eq!(prompt, vec![1, 100, 100, 200]); - } - - #[test] - fn splices_image_tokens_after_bos() { - let mut prompt = vec![1i32, 200, 300]; - let result = - insert_youtu_vl_image_tokens(&mut prompt, &[(4, 4)], 2, 128_262, 128_263, 128_264) - .unwrap(); - - // 4 patches at merge=2 → (4/2)*(4/2) = 4 image tokens. - assert_eq!(result.image_blocks, 1); - assert_eq!(result.total_image_tokens, 4); - - // Layout: [BOS, vstart, img × 4, vend, original tail...] - assert_eq!(prompt[0], 1); - assert_eq!(prompt[1], 128_262); - assert_eq!(&prompt[2..6], &[128_264; 4]); - assert_eq!(prompt[6], 128_263); - assert_eq!(&prompt[7..], &[200, 300]); - } - - #[test] - fn handles_multiple_images() { - let mut prompt = vec![1i32, 200]; - let result = insert_youtu_vl_image_tokens( - &mut prompt, - &[(4, 4), (2, 2)], - 2, - 128_262, - 128_263, - 128_264, - ) - .unwrap(); - // First image: 4 tokens, second image: 1 token, total = 5. - assert_eq!(result.image_blocks, 2); - assert_eq!(result.total_image_tokens, 5); - } -} +#[path = "youtu_vl_prompt_tests.rs"] +mod tests; diff --git a/src/multimodal/youtu_vl_prompt_tests.rs b/src/multimodal/youtu_vl_prompt_tests.rs new file mode 100644 index 000000000..cfc594858 --- /dev/null +++ b/src/multimodal/youtu_vl_prompt_tests.rs @@ -0,0 +1,128 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; + +const VSTART: i32 = 128_262; +const VEND: i32 = 128_263; +const IMG: i32 = 128_264; + +/// The prompt `tencent/Youtu-VL-4B-Instruct`'s chat template renders for one +/// image and a question, as token ids: system turn, then a user turn that +/// opens with `<|vision_start|><|image_pad|><|vision_end|>`. +fn templated_prompt() -> Vec { + vec![ + 128_000, 12_864, 198, 2_836, 128_001, 198, 128_000, 1_178, 198, VSTART, IMG, VEND, 3_753, + 21_314, 13, 128_001, 198, 128_000, 122_596, 198, + ] +} + +#[test] +fn leaves_an_already_expanded_prompt_unchanged() { + let mut prompt = vec![1, VSTART, IMG, IMG, IMG, IMG, VEND, 200]; + let before = prompt.clone(); + let result = insert_youtu_vl_image_tokens(&mut prompt, &[(4, 4)], 2, VSTART, VEND, IMG); + assert_eq!(result, Ok(None)); + assert_eq!(prompt, before); +} + +#[test] +fn splices_image_tokens_after_bos_when_the_prompt_has_no_placeholder() { + let mut prompt = vec![1i32, 200, 300]; + let result = insert_youtu_vl_image_tokens(&mut prompt, &[(4, 4)], 2, VSTART, VEND, IMG) + .unwrap() + .unwrap(); + + // 4 patches at merge=2 → (4/2)*(4/2) = 4 image tokens. + assert_eq!(result.image_blocks, 1); + assert_eq!(result.total_image_tokens, 4); + assert_eq!(prompt, vec![1, VSTART, IMG, IMG, IMG, IMG, VEND, 200, 300]); +} + +#[test] +fn handles_multiple_images_without_placeholders() { + let mut prompt = vec![1i32, 200]; + let result = insert_youtu_vl_image_tokens(&mut prompt, &[(4, 4), (2, 2)], 2, VSTART, VEND, IMG) + .unwrap() + .unwrap(); + // First image: 4 tokens, second image: 1 token, total = 5. + assert_eq!(result.image_blocks, 2); + assert_eq!(result.total_image_tokens, 5); + assert_eq!(prompt.iter().filter(|&&t| t == IMG).count(), 5); +} + +/// Issue #1618. The chat template leaves one `<|image_pad|>` per image; the +/// checkpoint's processor replaces it with one token per merged feature. A +/// 28x28 patch grid is the `test_image_shapes.png` fixture: a 14x14 merged +/// grid that spans 2x2 attention windows of 8x8, so 196 features must each +/// get their own slot, not just the first. +#[test] +fn expands_the_template_placeholder_to_one_token_per_merged_feature() { + let mut prompt = templated_prompt(); + let pad_at = prompt.iter().position(|&t| t == IMG).unwrap(); + let head = prompt[..pad_at].to_vec(); + let tail = prompt[pad_at + 1..].to_vec(); + + let result = insert_youtu_vl_image_tokens(&mut prompt, &[(28, 28)], 2, VSTART, VEND, IMG) + .unwrap() + .unwrap(); + + assert_eq!(result.image_blocks, 1); + assert_eq!(result.total_image_tokens, 196); + assert_eq!(prompt.len(), head.len() + 196 + tail.len()); + // The run sits exactly where the placeholder was, inside the template's + // own framing; no second start/end pair and no text moved. + assert_eq!(&prompt[..head.len()], head.as_slice()); + assert!( + prompt[head.len()..head.len() + 196] + .iter() + .all(|&t| t == IMG) + ); + assert_eq!(&prompt[head.len() + 196..], tail.as_slice()); + assert_eq!(prompt.iter().filter(|&&t| t == VSTART).count(), 1); + assert_eq!(prompt.iter().filter(|&&t| t == VEND).count(), 1); +} + +#[test] +fn expands_each_placeholder_with_its_own_image_run_in_order() { + let mut prompt = vec![1, VSTART, IMG, VEND, 50, VSTART, IMG, VEND, 60]; + let result = + insert_youtu_vl_image_tokens(&mut prompt, &[(22, 22), (4, 2)], 2, VSTART, VEND, IMG) + .unwrap() + .unwrap(); + assert_eq!(result.total_image_tokens, 121 + 2); + + let mut expected = vec![1, VSTART]; + expected.extend(std::iter::repeat_n(IMG, 121)); + expected.extend([VEND, 50, VSTART]); + expected.extend(std::iter::repeat_n(IMG, 2)); + expected.extend([VEND, 60]); + assert_eq!(prompt, expected); +} + +#[test] +fn rejects_a_placeholder_count_that_matches_neither_images_nor_features() { + let mut prompt = vec![1, IMG, IMG, 200]; + let before = prompt.clone(); + let result = insert_youtu_vl_image_tokens(&mut prompt, &[(28, 28)], 2, VSTART, VEND, IMG); + assert_eq!( + result, + Err(YoutuVlPlaceholderMismatch { + found: 2, + images: 1, + expected: 196, + }) + ); + assert_eq!(prompt, before); +} diff --git a/src/vision/encoders/youtu_vl_tests.rs b/src/vision/encoders/youtu_vl_tests.rs index ea2f86813..4bd3f5e10 100644 --- a/src/vision/encoders/youtu_vl_tests.rs +++ b/src/vision/encoders/youtu_vl_tests.rs @@ -394,3 +394,60 @@ fn reverse_window_indices_round_trips_a_multi_image_batch() { assert_eq!(misplaced_when_applied_twice(&window_index), 192); assert!(is_identity(&round_trip(&window_index))); } + +/// Issue #1618 named `get_window_index` as the first suspect for the +/// multi-window misdescriptions. These values come from running the +/// checkpoint's own `Siglip2Encoder.get_window_index` (`modeling_siglip2.py`, +/// including `torch.unique_consecutive`) on the two shape fixtures' grids, +/// one image each and both batched, plus a non-square grid with a partial +/// last window column. The real defect was the unexpanded prompt placeholder +/// (`multimodal::youtu_vl_prompt`); this pins that the window path, which the +/// stage-by-stage diff found matching, stays matching. +#[test] +fn window_index_matches_the_reference_on_multi_window_grids() { + let cases: &[(&[(i32, i32)], &[i32], [i32; 10], i32, usize)] = &[ + ( + &[(22, 22)], + &[0, 256, 352, 448, 484], + [81, 82, 83, 84, 8, 9, 10, 19, 20, 21], + 7260, + 121, + ), + ( + &[(28, 28)], + &[0, 256, 448, 640, 784], + [102, 103, 104, 105, 8, 9, 10, 11, 12, 13], + 19110, + 196, + ), + ( + &[(22, 22), (28, 28)], + &[0, 256, 352, 448, 484, 740, 932, 1124, 1268], + [81, 82, 83, 84, 8, 9, 10, 19, 20, 21], + 50086, + 317, + ), + ( + &[(16, 40)], + &[0, 256, 512, 640], + [144, 145, 146, 147, 8, 9, 10, 11, 12, 13], + 12720, + 160, + ), + ]; + for &(shapes, cu, slice, sum, len) in cases { + let (window_index, cu_window_seqlens) = window::get_window_index(shapes, 2, 256, 16, 4); + assert_eq!(cu_window_seqlens, cu, "cu_window_seqlens for {shapes:?}"); + assert_eq!(window_index.len(), len, "length for {shapes:?}"); + assert_eq!( + &window_index[60..70], + &slice, + "window_index[60..70] for {shapes:?}" + ); + assert_eq!( + window_index.iter().sum::(), + sum, + "index sum for {shapes:?}" + ); + } +} diff --git a/src/vision/youtu_vl.rs b/src/vision/youtu_vl.rs index ec80795de..6f7f8f49b 100644 --- a/src/vision/youtu_vl.rs +++ b/src/vision/youtu_vl.rs @@ -65,7 +65,7 @@ impl YoutuVLModel { input_ids: &MlxArray, pixel_values: &MlxArray, spatial_shapes: &[(i32, i32)], - ) -> InputEmbeddings { + ) -> Result { // Text embeddings. let inputs_embeds = self.text_model.get_embed_tokens(input_ids); @@ -82,30 +82,52 @@ impl YoutuVLModel { // Replace image tokens with vision features. Upstream falls back to // the video token when there are no image tokens; we do the same by // first checking which token id is present in the prompt. - let target_token_id = self.choose_target_token_id(input_ids); + let (target_token_id, placeholders) = self.choose_target_token_id(input_ids); + + // The scatter below fills placeholders with features in order and + // silently drops whatever does not fit. Upstream raises when the two + // counts differ; do the same so a prompt that was not expanded to one + // token per merged feature (#1618) fails instead of describing only + // the image's first patch. + let features = vision_output.hidden_states.as_ref().unwrap(); + let feature_rows = mlxcel_core::array_shape(features)[0]; + if placeholders != feature_rows { + return Err(format!( + "Youtu-VL image features and image tokens do not match: {placeholders} placeholder token(s) for {feature_rows} feature(s)" + )); + } - merge::merge_llava( + Ok(merge::merge_llava( target_token_id, - vision_output.hidden_states.as_ref().unwrap(), + features, &inputs_embeds, input_ids, - ) + )) } /// Pick the placeholder token id to merge against. Returns /// `image_token_id` when any image token appears in the prompt; otherwise - /// falls back to `video_token_id`. This mirrors the upstream behaviour - /// where the merge silently switches placeholder kinds when the caller - /// used a video token instead of an image token. - fn choose_target_token_id(&self, input_ids: &MlxArray) -> i32 { - // Compare the int-typed input ids to a scalar of the same dtype, then - // count matches via `sum_axis`. We materialize a single scalar count - // back to the host so the dispatch decision stays a plain Rust if/else. - let target = mlxcel_core::full_f32( - &[1], - self.image_token_id as f32, - mlxcel_core::array_dtype(input_ids), - ); + /// falls back to `video_token_id`, together with how many positions carry + /// it. This mirrors the upstream behaviour where the merge silently + /// switches placeholder kinds when the caller used a video token instead + /// of an image token. + fn choose_target_token_id(&self, input_ids: &MlxArray) -> (i32, i32) { + let images = Self::count_token(input_ids, self.image_token_id); + if images > 0 { + (self.image_token_id, images) + } else { + ( + self.video_token_id, + Self::count_token(input_ids, self.video_token_id), + ) + } + } + + /// Number of positions in `input_ids` equal to `token_id`, read back to + /// the host so dispatch stays a plain Rust branch. + fn count_token(input_ids: &MlxArray, token_id: i32) -> i32 { + let target = + mlxcel_core::full_f32(&[1], token_id as f32, mlxcel_core::array_dtype(input_ids)); let cmp = mlxcel_core::equal(input_ids, &target); // `equal` returns a bool array; cast to int32 so `sum_axis` accumulates // a numeric count (mlx does not currently sum bool tensors directly). @@ -113,11 +135,7 @@ impl YoutuVLModel { let flat = mlxcel_core::flatten(&cmp_int); let count = mlxcel_core::sum_axis(&flat, 0, false); mlxcel_core::eval(&count); - if mlxcel_core::item_i32(&count) > 0 { - self.image_token_id - } else { - self.video_token_id - } + mlxcel_core::item_i32(&count) } } diff --git a/tests/youtu_vl_parity.rs b/tests/youtu_vl_parity.rs index bffe8378c..4accdcc45 100644 --- a/tests/youtu_vl_parity.rs +++ b/tests/youtu_vl_parity.rs @@ -26,12 +26,15 @@ //! //! `models/mlx/youtu-vl-4b-instruct` has `patch_size=16`, `spatial_merge_size=2` //! and `window_size=256`, so `smart_resize` snaps every edge to a multiple of 32 -//! and the merged grid is one eighth of the pixel edge. The three fixture sizes -//! therefore sit on three different points of the window path: 224 gives a 7x7 -//! merged grid, a single 8x8 window where `get_window_index` is the identity and -//! the window inverse from #1600 / #1603 provably cannot change the result; 336 -//! resizes to 352 for an 11x11 merged grid and 448 gives 14x14, both of which -//! are 2x2 windows and do exercise it. +//! and the merged grid is one eighth of the pixel edge. 224 gives a 7x7 merged +//! grid, a single 8x8 attention window; 336 resizes to 352 for an 11x11 merged +//! grid and 448 gives 14x14, both of which are 2x2 windows. +//! +//! The solid orange fixture is also the one image that cannot catch a defect +//! in how features reach their prompt slots: every merged token of a uniform +//! image carries the same content, so a run that places only one feature +//! (#1618, where the template's single `<|image_pad|>` was never expanded) +//! still reads as orange. The two three-shape fixtures are what catch that. //! //! To run them: //! ```text @@ -126,10 +129,9 @@ fn assert_mentions(body: &str, groups: &[&[&str]]) { ); } -/// The single-window control, and the test that pins issue #1610. +/// The solid-color control, and the test that pins issue #1610. /// -/// At 224x224 the merged grid is 7x7, one 8x8 window, so `get_window_index` is -/// the identity and no window-permutation defect can reach this fixture. Before +/// At 224x224 the merged grid is 7x7, one 8x8 window. Before /// the patch-order fix the model answered "The image is completely black and /// contains no visible content, objects, text, or details. It is a solid black /// square or rectangle with no variation in color or texture." for a solid @@ -157,17 +159,13 @@ fn youtu_vl_describes_the_solid_color_fixture_as_a_solid_color() { /// Every object and every color at 448x448, a 14x14 merged grid. /// -/// KNOWN FAILING, tracked by #1618. The patch-order fix in #1610 -/// changed this answer from "The image contains a single, solid black circle on -/// a white background." to "The image contains a single, solid black circle on a -/// plain white background.", which is still wrong: the fixture is a red square, -/// a blue circle and a green triangle on light grey. At least one further defect -/// remains in this family's vision path and it is out of scope for #1610. The -/// assertion is deliberately left at the full correct answer rather than -/// weakened to something this build can satisfy, so that fixing #1618 -/// is what turns it green. +/// Pins issue #1618. Before it, the chat template's single `<|image_pad|>` was +/// never expanded to one token per merged feature, so only the top-left feature +/// reached the language model and the answer was "The image contains a single, +/// solid black circle on a plain white background." The fixture is a red +/// square, a blue circle and a green triangle on light grey. #[test] -#[ignore = "known failing, blocked on #1618 (second Youtu-VL vision defect); also requires the local youtu-vl-4b-instruct checkpoint"] +#[ignore = "requires the local youtu-vl-4b-instruct checkpoint"] fn youtu_vl_names_every_object_in_the_three_shape_image() { let Some(model_dir) = checkpoint_dir() else { return; @@ -189,13 +187,10 @@ fn youtu_vl_names_every_object_in_the_three_shape_image() { /// The same scene at 336x336, which `smart_resize` lifts to 352 for an 11x11 /// merged grid. /// -/// KNOWN FAILING, tracked by #1618, for the same reason as the 448 -/// case. The patch-order fix moved this answer from "The image contains a -/// single, solid black circle on a white background." to "The image contains a -/// single white circle on a black background.": the polarity flipped, so the -/// tower's output did change, but the scene is still not the one in the file. +/// Pins issue #1618 for the same reason as the 448 case. Before it, the answer +/// was "The image contains a single white circle on a black background." #[test] -#[ignore = "known failing, blocked on #1618 (second Youtu-VL vision defect); also requires the local youtu-vl-4b-instruct checkpoint"] +#[ignore = "requires the local youtu-vl-4b-instruct checkpoint"] fn youtu_vl_names_every_object_at_a_multi_window_grid() { let Some(model_dir) = checkpoint_dir() else { return;