[ZenCPU] Changes with respect to relevant cpu tests for ZenDNN

Signed-off-by: Chinmay Kulkarni <Chinmay.Kulkarni@amd.com>

.buildkite/hardware_tests/zen.yaml

@@ -0,0 +1,25 @@
+group: Zen CPU
+depends_on: []
+steps:
+- label: Zen-CPU-Kernel Tests
+  depends_on: []
+  soft_fail: false
+  device: zen5
+  no_plugin: true
+  source_file_dependencies:
+  - setup.py
+  - docker/Dockerfile.cpu
+  - requirements/cpu.txt
+  - requirements/build/cpu.txt
+  - csrc/cpu/
+  - cmake/cpu_extension.cmake
+  - CMakeLists.txt
+  - vllm/model_executor/layers/utils.py
+  - vllm/platforms/cpu.py
+  - vllm/platforms/zen_cpu.py
+  - vllm/platforms/__init__.py
+  - tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-zen-cpu-test.sh 20m "
+      pytest -x -v -s tests/model_executor/test_cpu_unquantized_gemm_dispatch.py"

.buildkite/scripts/hardware_ci/run-zen-cpu-test.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# This script build the CPU docker image and run the offline inference inside the container.
+# It serves a sanity check for compilation and basic model usage.
+set -euox pipefail
+
+# allow to bind to different cores
+CORE_RANGE=${CORE_RANGE:-48-95}
+NUMA_NODE=${NUMA_NODE:-1}
+IMAGE_NAME="zen-cpu-test-$NUMA_NODE"
+TIMEOUT_VAL=$1
+TEST_COMMAND=$2
+
+# building the docker image
+echo "--- :docker: Building Docker image"
+docker build --progress plain --tag "$IMAGE_NAME" --target vllm-zen-test -f docker/Dockerfile.cpu .
+
+# Run the image, setting --shm-size=4g for tensor parallel.
+docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
+        timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"

.buildkite/test-amd.yaml

@@ -2828,6 +2828,7 @@ steps:
   - rocm-smi
   - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
   - pytest -v -s tests/kernels/attention/test_attention_selector.py
+  - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
 
 - label: Kernels Attention Test %N # TBD
   timeout_in_minutes: 60

docker/Dockerfile.cpu

@@ -318,3 +318,30 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install "vllm[zen]"
 
 ENTRYPOINT ["vllm", "serve"]
+
+######################### ZEN CPU TEST IMAGE #########################
+FROM vllm-openai-zen AS vllm-zen-test
+
+WORKDIR /vllm-workspace
+
+COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install -r requirements/test/cpu.txt
+
+ADD ./tests/ ./tests/
+ADD ./examples/ ./examples/
+ADD ./benchmarks/ ./benchmarks/
+ADD ./vllm/collect_env.py .
+ADD ./docker/ ./docker/
+ADD ./.buildkite/ ./.buildkite/
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install -e tests/vllm_test_utils
+
+# enable fast downloads from hf (for testing)
+ENV HF_XET_HIGH_PERFORMANCE=1
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT=60
+
+ENTRYPOINT []

docs/models/supported_models.md

@@ -626,8 +626,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Step3VLForConditionalGeneration` | Step3-VL | T + I<sup>+</sup> | `stepfun-ai/step3` | | ✅︎ |
 | `StepVLForConditionalGeneration` | Step3-VL-10B | T + I<sup>+</sup> | `stepfun-ai/Step3-VL-10B` | | ✅︎ |
 | `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I<sup>+</sup> | `stepfun-ai/Step-3.7-Flash` | | ✅︎ |
-| `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
-| `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
 | `UltravoxModel` | Ultravox | T + A<sup>E+</sup> | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |
 | `UnlimitedOCRForCausalLM` | Unlimited-OCR | T + I<sup>+</sup> | `baidu/Unlimited-OCR`, etc. | ✅︎ | ✅︎ |
 

examples/generate/multimodal/vision_language_multi_image_offline.py

@@ -1275,55 +1275,6 @@ def load_step_vl(question: str, image_urls: list[str]) -> ModelRequestData:
     )
 
 
-def load_tarsier(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "omni-research/Tarsier-7b"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=4096,
-        limit_mm_per_prompt={"image": len(image_urls)},
-    )
-
-    prompt = f"USER: {'<image>' * len(image_urls)}\n{question}\n ASSISTANT:"
-    image_data = [fetch_image(url) for url in image_urls]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image_data=image_data,
-    )
-
-
-def load_tarsier2(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "omni-research/Tarsier2-Recap-7b"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=32768,
-        limit_mm_per_prompt={"image": len(image_urls)},
-        hf_overrides={
-            "architectures": ["Tarsier2ForConditionalGeneration"],
-            "model_type": "tarsier2",
-        },
-    )
-
-    prompt = (
-        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-        f"<|im_start|>user\n<|vision_start|>{'<|image_pad|>' * len(image_urls)}"
-        f"<|vision_end|>{question}<|im_end|>\n"
-        "<|im_start|>assistant\n"
-    )
-    image_data = [fetch_image(url) for url in image_urls]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image_data=image_data,
-    )
-
-
 # GLM-4.1V
 def load_glm4_1v(question: str, image_urls: list[str]) -> ModelRequestData:
     model_name = "zai-org/GLM-4.1V-9B-Thinking"
@@ -1507,8 +1458,6 @@ model_example_map = {
     "smolvlm": load_smolvlm,
     "step3": load_step3,
     "stepvl": load_step_vl,
-    "tarsier": load_tarsier,
-    "tarsier2": load_tarsier2,
     "glm4_1v": load_glm4_1v,
     "glm4_5v": load_glm4_5v,
     "glm4_5v_fp8": load_glm4_5v_fp8,

examples/generate/multimodal/vision_language_offline.py

@@ -2347,65 +2347,6 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
     )
 
 
-# omni-research/Tarsier-7b
-def run_tarsier(questions: list[str], modality: str) -> ModelRequestData:
-    assert modality == "image"
-    model_name = "omni-research/Tarsier-7b"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=4096,
-        limit_mm_per_prompt={modality: 1},
-    )
-    prompts = [(f"USER: <image>\n{question} ASSISTANT:") for question in questions]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
-def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
-    model_name = "omni-research/Tarsier2-Recap-7b"
-
-    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
-    engine_args = EngineArgs(
-        model=model_name,
-        max_model_len=4096,
-        hf_overrides={
-            "architectures": ["Tarsier2ForConditionalGeneration"],
-            "model_type": "tarsier2",
-        },
-        limit_mm_per_prompt=mm_limit,
-    )
-
-    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
-
-    if modality == "image":
-        placeholder = image_placeholder
-    elif modality == "video":
-        placeholder = video_placeholder
-    elif modality == "image+video":
-        placeholder = image_placeholder + video_placeholder
-
-    prompts = [
-        (
-            "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-            f"<|im_start|>user\n{placeholder}"
-            f"{question}<|im_end|>\n"
-            "<|im_start|>assistant\n"
-        )
-        for question in questions
-    ]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
 model_example_map = {
     "aria": run_aria,
     "aya_vision": run_aya_vision,
@@ -2479,8 +2420,6 @@ model_example_map = {
     "smolvlm": run_smolvlm,
     "step3": run_step3,
     "stepvl": run_step_vl,
-    "tarsier": run_tarsier,
-    "tarsier2": run_tarsier2,
 }
 
 

requirements/cuda.txt

@@ -8,7 +8,7 @@ torch==2.11.0
 torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
-PyNvVideoCodec==2.0.4
+PyNvVideoCodec==2.1.0
 # FlashInfer should be updated together with the Dockerfile
 flashinfer-python==0.6.13
 flashinfer-cubin==0.6.13

rust/Cargo.lock

@@ -5310,6 +5310,7 @@ dependencies = [
  "vllm-llm",
  "vllm-metrics",
  "vllm-text",
+ "vllm-tokenizer",
  "zeromq",
 ]
 

rust/src/chat/Cargo.toml

@@ -50,6 +50,7 @@ tokio.workspace = true
 tracing-subscriber.workspace = true
 uuid.workspace = true
 vllm-engine-core-client = { workspace = true, features = ["test-util"] }
+vllm-tokenizer = { workspace = true, features = ["test-utils"] }
 zeromq.workspace = true
 
 [lints]

rust/src/chat/src/backend/hf.rs

@@ -154,7 +154,8 @@ mod tests {
     use thiserror_ext::AsReport as _;
     use vllm_text::Prompt;
     use vllm_text::backend::hf::TokenizerSource;
-    use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
+    use vllm_text::tokenizer::DynTokenizer;
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::HfChatBackend;
     use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions};
@@ -196,32 +197,8 @@ mod tests {
         }
     }
 
-    struct TestTokenizer;
-
-    impl Tokenizer for TestTokenizer {
-        fn encode(
-            &self,
-            _text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_text::tokenizer::Result<Vec<u32>> {
-            Ok(Vec::new())
-        }
-
-        fn decode(
-            &self,
-            _token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_text::tokenizer::Result<String> {
-            Ok(String::new())
-        }
-
-        fn token_to_id(&self, _token: &str) -> Option<u32> {
-            None
-        }
-    }
-
     fn test_tokenizer() -> DynTokenizer {
-        Arc::new(TestTokenizer)
+        Arc::new(TestTokenizer::new())
     }
 
     fn backend_for_selection(

rust/src/chat/src/multimodal.rs

@@ -563,7 +563,7 @@ mod tests {
 
     use llm_multimodal::TokenId;
     use vllm_engine_core_client::protocol::tensor::WireArrayData;
-    use vllm_text::tokenizer::{IncrementalDecoder, Tokenizer, TokenizerError};
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::*;
 
@@ -574,60 +574,14 @@ mod tests {
     const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093;
     const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094;
 
-    struct TestTokenizer;
-
-    impl Tokenizer for TestTokenizer {
-        fn encode(
-            &self,
-            text: &str,
-            _add_special_tokens: bool,
-        ) -> std::result::Result<Vec<u32>, TokenizerError> {
-            Ok(match text {
-                "<|image|>" => vec![LLAMA4_IMAGE_ID],
-                text => text.bytes().map(u32::from).collect(),
-            })
-        }
-
-        fn decode(
-            &self,
-            _token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> std::result::Result<String, TokenizerError> {
-            Ok(String::new())
-        }
-
-        fn token_to_id(&self, token: &str) -> Option<u32> {
-            match token {
-                "<|image_start|>" => Some(LLAMA4_IMAGE_START_ID),
-                "<|image_end|>" => Some(LLAMA4_IMAGE_END_ID),
-                "<|image|>" => Some(LLAMA4_IMAGE_ID),
-                "<|patch|>" => Some(LLAMA4_PATCH_ID),
-                "<|tile_x_separator|>" => Some(LLAMA4_TILE_X_SEPARATOR_ID),
-                "<|tile_y_separator|>" => Some(LLAMA4_TILE_Y_SEPARATOR_ID),
-                _ => None,
-            }
-        }
-
-        fn id_to_token(&self, id: u32) -> Option<String> {
-            match id {
-                LLAMA4_IMAGE_START_ID => Some("<|image_start|>".to_string()),
-                LLAMA4_IMAGE_END_ID => Some("<|image_end|>".to_string()),
-                LLAMA4_IMAGE_ID => Some("<|image|>".to_string()),
-                LLAMA4_PATCH_ID => Some("<|patch|>".to_string()),
-                LLAMA4_TILE_X_SEPARATOR_ID => Some("<|tile_x_separator|>".to_string()),
-                LLAMA4_TILE_Y_SEPARATOR_ID => Some("<|tile_y_separator|>".to_string()),
-                _ => None,
-            }
-        }
-
-        fn create_decode_stream(
-            &self,
-            _prompt_token_ids: &[u32],
-            _skip_special_tokens: bool,
-            _min_bytes_to_buffer: usize,
-        ) -> Box<dyn IncrementalDecoder + '_> {
-            unreachable!("not used")
-        }
+    fn llama4_tokenizer() -> TestTokenizer {
+        TestTokenizer::new()
+            .with_regular_token("<|image_start|>", LLAMA4_IMAGE_START_ID)
+            .with_regular_token("<|image_end|>", LLAMA4_IMAGE_END_ID)
+            .with_regular_token("<|image|>", LLAMA4_IMAGE_ID)
+            .with_regular_token("<|patch|>", LLAMA4_PATCH_ID)
+            .with_regular_token("<|tile_x_separator|>", LLAMA4_TILE_X_SEPARATOR_ID)
+            .with_regular_token("<|tile_y_separator|>", LLAMA4_TILE_Y_SEPARATOR_ID)
     }
 
     fn test_info(model_type: &str, config: serde_json::Value) -> MultimodalModelInfo {
@@ -635,7 +589,7 @@ mod tests {
             model_id: format!("{model_type}-test"),
             model_type: Some(model_type.to_string()),
             config,
-            tokenizer: TokenizerResolver(Arc::new(TestTokenizer)),
+            tokenizer: TokenizerResolver(Arc::new(llama4_tokenizer())),
         };
         let spec = context
             .resolve_model_spec()

rust/src/chat/src/output/default/mod.rs

@@ -189,46 +189,19 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor {
 mod tests {
     use std::sync::Arc;
 
-    use vllm_tokenizer::Tokenizer;
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::DefaultChatOutputProcessor;
     use crate::Error;
     use crate::parser::ParserSelection;
     use crate::request::ChatRequest;
 
-    struct FakeTokenizer;
-
-    impl Tokenizer for FakeTokenizer {
-        fn encode(
-            &self,
-            text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            Ok(text.chars().map(u32::from).collect())
-        }
-
-        fn decode(
-            &self,
-            token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            Ok(token_ids
-                .iter()
-                .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
-                .collect())
-        }
-
-        fn token_to_id(&self, token: &str) -> Option<u32> {
-            match token {
-                "<|channel>" => Some(1),
-                "<channel|>" => Some(2),
-                _ => None,
-            }
-        }
-    }
-
-    fn tokenizer() -> Arc<FakeTokenizer> {
-        Arc::new(FakeTokenizer)
+    fn tokenizer() -> Arc<TestTokenizer> {
+        Arc::new(
+            TestTokenizer::new()
+                .with_regular_token("<|channel>", 256)
+                .with_regular_token("<channel|>", 257),
+        )
     }
 
     #[test]

rust/src/chat/src/output/harmony/tests.rs

@@ -1,14 +1,8 @@
-//! Harmony output tests share the upstream `openai-harmony` tiktoken cache.
-//!
-//! Use a file lock for tests that load the encoding so `cargo nextest` cannot
-//! start multiple processes that concurrently populate the same cache file.
-
 use std::sync::Arc;
 
 use futures::executor::block_on;
 use futures::{TryStreamExt as _, stream};
 use openai_harmony::chat::{Message, Role};
-use serial_test::file_serial;
 use vllm_text::output::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, Finished};
 
 use super::*;
@@ -91,7 +85,6 @@ fn request_with_tools() -> ChatRequest {
 }
 
 #[test]
-#[file_serial(harmony_tiktoken_cache)]
 fn interrupted_final_message_is_preserved() {
     let tokens = completion_tokens(&[text_message("final", "hello")]);
     let events = block_on(collect_events(
@@ -127,7 +120,6 @@ fn interrupted_final_message_is_preserved() {
 }
 
 #[test]
-#[file_serial(harmony_tiktoken_cache)]
 fn eos_flush_preserves_trailing_replacement_text() {
     let mut tokens = completion_tokens(&[text_message("final", "Hi")]);
     tokens.pop();
@@ -153,7 +145,6 @@ fn eos_flush_preserves_trailing_replacement_text() {
 }
 
 #[test]
-#[file_serial(harmony_tiktoken_cache)]
 fn interrupted_analysis_message_is_preserved() {
     let tokens = completion_tokens(&[text_message("analysis", "think")]);
     let events = block_on(collect_events(
@@ -189,7 +180,6 @@ fn interrupted_analysis_message_is_preserved() {
 }
 
 #[test]
-#[file_serial(harmony_tiktoken_cache)]
 fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
     let tokens = completion_tokens(&[
         text_message("commentary", "Let me check."),
@@ -217,7 +207,6 @@ fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
 }
 
 #[test]
-#[file_serial(harmony_tiktoken_cache)]
 fn multiple_messages_get_newline_separators() {
     let tokens = completion_tokens(&[
         text_message("analysis", "first think"),
@@ -249,7 +238,6 @@ fn multiple_messages_get_newline_separators() {
 }
 
 #[test]
-#[file_serial(harmony_tiktoken_cache)]
 fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
     let tokens = completion_tokens(&[tool_message(
         "get_weather",
@@ -302,7 +290,6 @@ fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
 }
 
 #[test]
-#[file_serial(harmony_tiktoken_cache)]
 fn semantic_events_precede_same_update_logprobs() {
     let tokens = completion_tokens(&[text_message("final", "hello")]);
     let events = block_on(collect_events(
@@ -353,7 +340,6 @@ fn rejects_generic_parser_overrides() {
 }
 
 #[test]
-#[file_serial(harmony_tiktoken_cache)]
 fn allows_auto_auto_only() {
     validate_harmony_parser_overrides(&ParserSelection::Auto, &ParserSelection::Auto).unwrap();
     let _ = HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap();

rust/src/chat/src/parser/reasoning/tests.rs

@@ -1,32 +1,9 @@
 use std::sync::Arc;
 
-use vllm_tokenizer::Tokenizer;
+use vllm_tokenizer::test_utils::TestTokenizer;
 
 use super::{ReasoningParserFactory, names};
 
-struct FakeTokenizer;
-
-impl Tokenizer for FakeTokenizer {
-    fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
-        Ok(text.chars().map(u32::from).collect())
-    }
-
-    fn decode(
-        &self,
-        token_ids: &[u32],
-        _skip_special_tokens: bool,
-    ) -> vllm_tokenizer::Result<String> {
-        Ok(token_ids
-            .iter()
-            .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
-            .collect())
-    }
-
-    fn token_to_id(&self, _token: &str) -> Option<u32> {
-        None
-    }
-}
-
 #[test]
 fn factory_contains_and_lists_registered_parsers() {
     let factory = ReasoningParserFactory::new();
@@ -107,7 +84,7 @@ fn factory_resolves_minimax_m3_before_generic_minimax() {
 
 #[test]
 fn factory_rejects_unknown_parser_names() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(TestTokenizer::new());
     let factory = ReasoningParserFactory::new();
     let error = match factory.create("missing", tokenizer) {
         Ok(_) => panic!("expected parser lookup to fail"),

rust/src/chat/src/parser/unified.rs

@@ -75,39 +75,14 @@ impl UnifiedParserFactory {
 mod tests {
     use std::sync::Arc;
 
-    use vllm_tokenizer::Tokenizer;
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::{UnifiedParserFactory, names};
 
-    struct FakeTokenizer;
-
-    impl Tokenizer for FakeTokenizer {
-        fn encode(
-            &self,
-            text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            Ok(text.chars().map(u32::from).collect())
-        }
-
-        fn decode(
-            &self,
-            token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            Ok(token_ids
-                .iter()
-                .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
-                .collect())
-        }
-
-        fn token_to_id(&self, token: &str) -> Option<u32> {
-            match token {
-                "<|channel>" => Some(1),
-                "<channel|>" => Some(2),
-                _ => None,
-            }
-        }
+    fn tokenizer() -> TestTokenizer {
+        TestTokenizer::new()
+            .with_regular_token("<|channel>", 256)
+            .with_regular_token("<channel|>", 257)
     }
 
     #[test]
@@ -119,6 +94,6 @@ mod tests {
             factory.resolve_name_for_model("google/gemma-4-27b-it"),
             Some(names::GEMMA4)
         );
-        factory.create(names::GEMMA4, &[], Arc::new(FakeTokenizer)).unwrap();
+        factory.create(names::GEMMA4, &[], Arc::new(tokenizer())).unwrap();
     }
 }

rust/src/chat/tests/chat.rs

@@ -21,15 +21,17 @@ use vllm_engine_core_client::protocol::{
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
 use vllm_llm::Llm;
-use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
+use vllm_text::tokenizer::DynTokenizer;
 use vllm_text::{
     DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTokenLogprob, Prompt,
     TextBackend,
 };
+use vllm_tokenizer::test_utils::TestTokenizer;
 use zeromq::prelude::{SocketRecv, SocketSend};
 use zeromq::{DealerSocket, PushSocket, ZmqMessage};
 
 const SPECIAL_STOP_TOKEN_ID: u32 = 256;
+const UNKNOWN_DECODE_TOKEN_ID: u32 = 10_000;
 
 fn request_output(
     request_id: &str,
@@ -158,45 +160,18 @@ async fn connect_chat_llm_with_ipc(
 struct FakeChatBackend {
     has_template: bool,
     model_id: String,
+    tokenizer: DynTokenizer,
 }
 
-#[derive(Debug)]
-struct FakeChatTokenizer;
-
-impl Tokenizer for FakeChatTokenizer {
-    fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
-        Ok(text.bytes().map(u32::from).collect())
-    }
-
-    fn decode(
-        &self,
-        token_ids: &[u32],
-        skip_special_tokens: bool,
-    ) -> vllm_tokenizer::Result<String> {
-        let bytes = token_ids
-            .iter()
-            .filter_map(|id| {
-                if skip_special_tokens && *id == SPECIAL_STOP_TOKEN_ID {
-                    None
-                } else {
-                    Some(*id as u8)
-                }
-            })
-            .collect::<Vec<_>>();
-        Ok(String::from_utf8_lossy(&bytes).into_owned())
-    }
-
-    fn token_to_id(&self, token: &str) -> Option<u32> {
-        match token {
-            "<think>" => Some(0xF001),
-            "</think>" => Some(0xF002),
-            "<|START_THINKING|>" => Some(0xF003),
-            "<|END_THINKING|>" => Some(0xF004),
-            "◁think▷" => Some(0xF005),
-            "◁/think▷" => Some(0xF006),
-            _ => None,
-        }
-    }
+fn fake_chat_tokenizer() -> TestTokenizer {
+    TestTokenizer::new()
+        .with_special_token("<stop>", SPECIAL_STOP_TOKEN_ID)
+        .with_regular_token("<think>", 0xF001)
+        .with_regular_token("</think>", 0xF002)
+        .with_regular_token("<|START_THINKING|>", 0xF003)
+        .with_regular_token("<|END_THINKING|>", 0xF004)
+        .with_regular_token("◁think▷", 0xF005)
+        .with_regular_token("◁/think▷", 0xF006)
 }
 
 impl fmt::Debug for FakeChatBackend {
@@ -210,6 +185,7 @@ impl FakeChatBackend {
         Self {
             has_template: true,
             model_id: "test-model".to_string(),
+            tokenizer: Arc::new(fake_chat_tokenizer()),
         }
     }
 
@@ -217,6 +193,7 @@ impl FakeChatBackend {
         Self {
             has_template: false,
             model_id: "test-model".to_string(),
+            tokenizer: Arc::new(fake_chat_tokenizer()),
         }
     }
 
@@ -224,13 +201,19 @@ impl FakeChatBackend {
         Self {
             has_template: true,
             model_id: model_id.into(),
+            tokenizer: Arc::new(fake_chat_tokenizer()),
         }
     }
+
+    fn with_tokenizer(mut self, tokenizer: DynTokenizer) -> Self {
+        self.tokenizer = tokenizer;
+        self
+    }
 }
 
 impl TextBackend for FakeChatBackend {
     fn tokenizer(&self) -> DynTokenizer {
-        Arc::new(FakeChatTokenizer)
+        Arc::clone(&self.tokenizer)
     }
 
     fn model_id(&self) -> &str {
@@ -282,65 +265,6 @@ impl ChatRenderer for FakeChatBackend {
     }
 }
 
-#[derive(Clone, Debug)]
-struct FailingDecodeBackend {
-    inner: FakeChatBackend,
-}
-
-#[derive(Debug)]
-struct FailingDecodeTokenizer;
-
-impl Tokenizer for FailingDecodeTokenizer {
-    fn encode(&self, text: &str, add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
-        FakeChatTokenizer.encode(text, add_special_tokens)
-    }
-
-    fn decode(
-        &self,
-        token_ids: &[u32],
-        skip_special_tokens: bool,
-    ) -> vllm_tokenizer::Result<String> {
-        if token_ids.contains(&(b'i' as u32)) {
-            return Err(vllm_tokenizer::TokenizerError("decode failed".to_string()));
-        }
-        FakeChatTokenizer.decode(token_ids, skip_special_tokens)
-    }
-
-    fn token_to_id(&self, token: &str) -> Option<u32> {
-        FakeChatTokenizer.token_to_id(token)
-    }
-}
-
-impl TextBackend for FailingDecodeBackend {
-    fn tokenizer(&self) -> DynTokenizer {
-        Arc::new(FailingDecodeTokenizer)
-    }
-
-    fn model_id(&self) -> &str {
-        self.inner.model_id()
-    }
-}
-
-impl ChatBackend for FailingDecodeBackend {
-    fn chat_renderer(&self) -> DynChatRenderer {
-        Arc::new(self.clone())
-    }
-
-    fn new_chat_output_processor(
-        &self,
-        _request: &mut ChatRequest,
-        _options: NewChatOutputProcessorOptions<'_>,
-    ) -> vllm_chat::Result<DynChatOutputProcessor> {
-        Ok(Box::new(DefaultChatOutputProcessor::plain_text_only()))
-    }
-}
-
-impl ChatRenderer for FailingDecodeBackend {
-    fn render(&self, request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
-        self.inner.render(request)
-    }
-}
-
 /// Skip `LogprobsDelta` events that carry only token_ids (no logprobs),
 /// returning the next semantically interesting event.
 async fn next_semantic<S>(stream: &mut S) -> Option<Result<ChatEvent, vllm_chat::Error>>
@@ -738,7 +662,12 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
                 send_outputs(
                     push,
                     EngineCoreOutputs {
-                        outputs: vec![request_output("chat-4", vec![b'i' as u32], None, None)],
+                        outputs: vec![request_output(
+                            "chat-4",
+                            vec![UNKNOWN_DECODE_TOKEN_ID],
+                            None,
+                            None,
+                        )],
                         ..Default::default()
                     },
                 )
@@ -747,9 +676,8 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
         },
     );
 
-    let backend: Arc<dyn ChatTextBackend> = Arc::new(FailingDecodeBackend {
-        inner: FakeChatBackend::new(),
-    });
+    let backend: Arc<dyn ChatTextBackend> =
+        Arc::new(FakeChatBackend::new().with_tokenizer(Arc::new(TestTokenizer::new())));
     let chat = connect_chat_llm_with_ipc(
         EngineCoreClientConfig::new_single(handshake_address),
         &ipc,
@@ -769,7 +697,10 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
 
     match timeout(Duration::from_secs(2), stream.next()).await.unwrap() {
         Some(Err(vllm_chat::Error::Text(vllm_text::Error::Tokenizer(message)))) => {
-            assert_eq!(message, "decode failed");
+            assert_eq!(
+                message,
+                format!("test tokenizer cannot decode unknown token id {UNKNOWN_DECODE_TOKEN_ID}")
+            );
         }
         other => panic!("unexpected event after close: {other:?}"),
     }

rust/src/chat/tests/roundtrip.rs

@@ -1,8 +1,8 @@
-//! Text-level roundtrip tests for the real chat-template and output-processor pairing.
+//! Roundtrip tests for the real chat-template and output-processor pairing.
 //!
 //! The invariant under test is that a structured assistant message rendered as history can be
 //! parsed from the generated assistant completion and then rendered back to the exact same
-//! assistant-completion text.
+//! assistant completion.
 
 use std::pin::Pin;
 use std::sync::Arc;
@@ -18,6 +18,10 @@ use vllm_chat::{
     RendererSelection, load_model_backends,
 };
 use vllm_text::{DecodedTextEvent, Finished, Prompt};
+use vllm_tokenizer::Tokenizer;
+
+const TEXT_COMPLETION_CHUNK_CHARS: usize = 7;
+const TOKEN_COMPLETION_CHUNK_TOKENS: usize = 1;
 
 /// One model/parser configuration used to run the fixed roundtrip fixtures.
 #[derive(Clone)]
@@ -191,14 +195,28 @@ impl RoundtripCase {
             sort_json_keys: false,
         }
     }
+
+    /// GPT-OSS Harmony token-id renderer and native Harmony output processor.
+    fn gpt_oss() -> Self {
+        Self {
+            model_id: "openai/gpt-oss-20b",
+            assistant_stop_suffix: "", // not applicable for token-id cases
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            thinking_behavior: ThinkingBehavior::Always { value: true },
+            json_fmt: compact_json_fmt(),
+            sort_json_keys: false,
+        }
+    }
 }
 
 macro_rules! roundtrip_tests {
-    ($($case:ident => [$($fixture:ident),* $(,)?]),+ $(,)?) => {
+    ($($case:ident => [$($(#[$fixture_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => {
         paste::paste! {
             $(
                 $(
                     #[tokio::test]
+                    $(#[$fixture_attr])*
                     #[file_serial([<hf_ $case>])]
                     async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
                         [<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
@@ -217,9 +235,9 @@ roundtrip_tests! {
     glm47 => [reasoning_and_content, tool_call_mix],
     seed_oss => [reasoning_and_content],
     step3p5 => [reasoning_and_content],
-
     gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
     kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
+    gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call
 }
 
 /// Run the fixed reasoning+content fixture for one model/parser case.
@@ -421,10 +439,10 @@ struct RoundtripResult {
     parsed_message: AssistantMessage,
     /// Assistant-completion suffix cut from rendering the expected assistant as
     /// history.
-    closed_completion: String,
+    closed_completion: Prompt,
     /// Assistant-completion suffix cut after rendering the parsed assistant
     /// back as history.
-    rerendered_closed_completion: String,
+    rerendered_closed_completion: Prompt,
 }
 
 /// Render, parse, and rerender one assistant turn through the production
@@ -436,60 +454,59 @@ async fn run_roundtrip(
     assistant: AssistantMessage,
 ) -> Result<RoundtripResult> {
     let renderer = backends.chat_backend.chat_renderer();
-    let (prompt, closed_completion_text) =
-        render_closed_completion(renderer.as_ref(), request, &assistant)?;
-    let completion_body = closed_completion_text
-        .strip_suffix(case.assistant_stop_suffix)
-        .with_context(|| {
-            format!(
-                "closed assistant completion did not end with {:?}: {:?}",
-                case.assistant_stop_suffix, closed_completion_text
-            )
-        })?;
+    let rendered = render_closed_completion(renderer.as_ref(), request, &assistant)?;
 
-    let parsed_message =
-        parse_completion(case, backends, request, &prompt, completion_body).await?;
-    let (_, rerendered_closed_completion) =
-        render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
+    let parsed_message = parse_completion(case, backends, request, &rendered).await?;
+    let rerendered = render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
 
     Ok(RoundtripResult {
         parsed_message,
-        closed_completion: closed_completion_text,
-        rerendered_closed_completion,
+        closed_completion: rendered.completion,
+        rerendered_closed_completion: rerendered.completion,
     })
 }
 
+/// Rendered prompt/completion artifacts at the renderer boundary.
+struct RenderedTurn {
+    prompt: Prompt,
+    completion: Prompt,
+}
+
 /// Render `history` as a production prompt and `history + assistant` as closed
 /// history, then return the production prompt and assistant-completion suffix.
 fn render_closed_completion(
     renderer: &dyn vllm_chat::ChatRenderer,
     base_request: &ChatRequest,
     assistant: &AssistantMessage,
-) -> Result<(String, String)> {
+) -> Result<RenderedTurn> {
     let mut prompt_request = base_request.clone();
     prompt_request.chat_options.generation_prompt_mode = GenerationPromptMode::StartNewAssistant;
-    let prompt = render_text(renderer, &prompt_request).context("failed to render prompt")?;
+    let prompt = renderer.render(&prompt_request).context("failed to render prompt")?.prompt;
 
     let mut full_request = base_request.clone();
     full_request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
     full_request.messages.push(ChatMessage::from(assistant.clone()));
-    let full = render_text(renderer, &full_request).context("failed to render full prompt")?;
+    let full = renderer.render(&full_request).context("failed to render full prompt")?.prompt;
 
-    ensure!(
-        full.starts_with(&prompt),
-        "full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
-    );
-    let completion = full[prompt.len()..].to_string();
+    let completion = match (&prompt, full) {
+        (Prompt::Text(prompt), Prompt::Text(full)) => {
+            ensure!(
+                full.starts_with(prompt),
+                "full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
+            );
+            Prompt::Text(full[prompt.len()..].to_string())
+        }
+        (Prompt::TokenIds(prompt), Prompt::TokenIds(full)) => {
+            ensure!(
+                full.starts_with(prompt),
+                "full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
+            );
+            Prompt::TokenIds(full[prompt.len()..].to_vec())
+        }
+        (prompt, full) => bail!("prompt kind changed between renders: {prompt:?} vs {full:?}"),
+    };
 
-    Ok((prompt, completion))
-}
-
-/// Render one chat request and require a text prompt.
-fn render_text(renderer: &dyn vllm_chat::ChatRenderer, request: &ChatRequest) -> Result<String> {
-    match renderer.render(request)?.prompt {
-        Prompt::Text(text) => Ok(text),
-        other => bail!("roundtrip tests expect text prompts, got {other:?}"),
-    }
+    Ok(RenderedTurn { prompt, completion })
 }
 
 /// Feed one rendered assistant completion body into the real output processor
@@ -498,13 +515,15 @@ async fn parse_completion(
     case: &RoundtripCase,
     backends: &vllm_chat::LoadedModelBackends,
     base_request: &ChatRequest,
-    prompt: &str,
-    completion_body: &str,
+    rendered: &RenderedTurn,
 ) -> Result<AssistantMessage> {
     let tokenizer = backends.text_backend.tokenizer();
-    let prompt_token_ids = tokenizer
-        .encode(prompt, base_request.add_special_tokens)
-        .context("failed to encode rendered prompt")?;
+    let prompt_token_ids = match &rendered.prompt {
+        Prompt::Text(prompt) => tokenizer
+            .encode(prompt, base_request.add_special_tokens)
+            .context("failed to encode rendered prompt")?,
+        Prompt::TokenIds(token_ids) => token_ids.clone(),
+    };
 
     let mut request = base_request.clone();
     let processor = backends.chat_backend.new_chat_output_processor(
@@ -515,7 +534,12 @@ async fn parse_completion(
         },
     )?;
 
-    let decoded = decoded_completion_stream(prompt_token_ids, completion_body);
+    let decoded = decoded_completion_stream(
+        tokenizer.as_ref(),
+        prompt_token_ids,
+        &rendered.completion,
+        case.assistant_stop_suffix,
+    )?;
     let mut events = processor.process(decoded)?;
 
     while let Some(event) = events.next().await {
@@ -538,16 +562,46 @@ async fn parse_completion(
 /// split into small chunks to exercise streaming parser state across marker
 /// and JSON boundaries.
 fn decoded_completion_stream(
+    tokenizer: &dyn Tokenizer,
     prompt_token_ids: Vec<u32>,
-    completion_body: &str,
-) -> Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>> {
-    let prompt_token_count = prompt_token_ids.len();
+    completion: &Prompt,
+    assistant_stop_suffix: &str,
+) -> Result<Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>>> {
     let mut events = vec![DecodedTextEvent::Start {
-        prompt_token_ids: Arc::from(prompt_token_ids.into_boxed_slice()),
+        prompt_token_ids: Arc::from(prompt_token_ids.clone().into_boxed_slice()),
         prompt_logprobs: None,
     }];
 
-    let chunks = split_by_chars(completion_body, 7);
+    let chunks = match completion {
+        Prompt::Text(text) => {
+            let body = text.strip_suffix(assistant_stop_suffix).with_context(|| {
+                format!(
+                    "closed assistant completion did not end with {:?}: {:?}",
+                    assistant_stop_suffix, text
+                )
+            })?;
+            split_by_chars(body, TEXT_COMPLETION_CHUNK_CHARS)
+                .into_iter()
+                .map(|delta| DecodedCompletionChunk {
+                    delta,
+                    token_ids: Vec::new(), // unused for text-level roundtrip cases
+                })
+                .collect()
+        }
+        Prompt::TokenIds(token_ids) => {
+            ensure!(
+                assistant_stop_suffix.is_empty(),
+                "token-id roundtrip cases do not support text stop suffixes"
+            );
+            incremental_decode_chunks(
+                tokenizer,
+                &prompt_token_ids,
+                token_ids,
+                TOKEN_COMPLETION_CHUNK_TOKENS,
+            )?
+        }
+    };
+
     if chunks.is_empty() {
         events.push({
             DecodedTextEvent::TextDelta {
@@ -555,11 +609,7 @@ fn decoded_completion_stream(
                 token_ids: Vec::new(),
                 logprobs: None,
                 finished: Some(Finished {
-                    usage: vllm_llm::TokenUsage {
-                        prompt_token_count: 0,
-                        output_token_count: 0,
-                        cached_token_count: 0,
-                    },
+                    usage: Default::default(),
                     finish_reason: FinishReason::stop_eos(),
                     kv_transfer_params: None,
                 }),
@@ -569,24 +619,26 @@ fn decoded_completion_stream(
         let last_index = chunks.len() - 1;
         for (index, chunk) in chunks.into_iter().enumerate() {
             let finished = (index == last_index).then(|| Finished {
-                usage: vllm_llm::TokenUsage {
-                    prompt_token_count,
-                    output_token_count: completion_body.chars().count(),
-                    cached_token_count: 0,
-                },
+                usage: Default::default(),
                 finish_reason: FinishReason::stop_eos(),
                 kv_transfer_params: None,
             });
             events.push(DecodedTextEvent::TextDelta {
-                delta: chunk,
-                token_ids: Vec::new(),
+                delta: chunk.delta,
+                token_ids: chunk.token_ids,
                 logprobs: None,
                 finished,
             });
         }
     }
 
-    stream::iter(events).map(Ok).boxed()
+    Ok(stream::iter(events).map(Ok).boxed())
+}
+
+/// One decoded completion chunk fed into the output processor.
+struct DecodedCompletionChunk {
+    delta: String,
+    token_ids: Vec<u32>,
 }
 
 /// Split text into chunks containing at most `chunk_chars` Unicode scalar
@@ -612,6 +664,49 @@ fn split_by_chars(text: &str, chunk_chars: usize) -> Vec<String> {
     chunks
 }
 
+/// Split token ids into chunks containing at most `chunk_size` ids.
+fn split_by_count(token_ids: &[u32], chunk_size: usize) -> Vec<Vec<u32>> {
+    token_ids.chunks(chunk_size).map(<[u32]>::to_vec).collect()
+}
+
+/// Decode token ids incrementally using the production tokenizer stream.
+fn incremental_decode_chunks(
+    tokenizer: &dyn Tokenizer,
+    prompt_token_ids: &[u32],
+    token_ids: &[u32],
+    chunk_size: usize,
+) -> Result<Vec<DecodedCompletionChunk>> {
+    let mut decoder = tokenizer.create_decode_stream(prompt_token_ids, false, 0);
+    let mut chunks = Vec::new();
+    for chunk_token_ids in split_by_count(token_ids, chunk_size) {
+        let mut delta = String::new();
+        for token_id in chunk_token_ids.iter().copied() {
+            decoder.push_token(token_id)?;
+            while let Some(chunk) = decoder.next_chunk() {
+                delta.push_str(&chunk);
+            }
+        }
+        chunks.push(DecodedCompletionChunk {
+            delta,
+            token_ids: chunk_token_ids,
+        });
+    }
+
+    let (last_chunk, _) = decoder.flush(None)?;
+    if let Some(last_chunk) = last_chunk {
+        if let Some(delta) = chunks.last_mut() {
+            delta.delta.push_str(&last_chunk);
+        } else {
+            chunks.push(DecodedCompletionChunk {
+                delta: last_chunk,
+                token_ids: Vec::new(),
+            });
+        }
+    }
+
+    Ok(chunks)
+}
+
 /// Build a chat request fixture with parser-enabling tool-choice semantics.
 fn roundtrip_request(
     request_id: impl Into<String>,

rust/src/parser/Cargo.toml

@@ -23,6 +23,7 @@ expect-test.workspace = true
 futures.workspace = true
 openai-protocol.workspace = true
 tool-parser.workspace = true
+vllm-tokenizer = { workspace = true, features = ["test-utils"] }
 
 [[bench]]
 name = "deepseek_v3"

rust/src/parser/benches/utils/adapter.rs

@@ -27,6 +27,10 @@ impl Tokenizer for BenchTokenizer {
     fn token_to_id(&self, _token: &str) -> Option<u32> {
         Some(u32::MAX)
     }
+
+    fn id_to_token(&self, _id: u32) -> Option<String> {
+        Some("\u{FFFD}".to_string())
+    }
 }
 
 /// Bench-only adapter that exposes a unified parser through the tool-parser

rust/src/parser/src/reasoning/seed_oss.rs

@@ -49,11 +49,14 @@ mod tests {
     use std::sync::Arc;
 
     use super::SeedOssReasoningParser;
-    use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
+    use crate::reasoning::{
+        ReasoningParser,
+        tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer},
+    };
 
     #[test]
     fn without_prompt_markers_expects_start_token() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
 
         let delta = parser.push("implicit reasoning</seed:think>answer").unwrap();
@@ -66,10 +69,10 @@ mod tests {
 
     #[test]
     fn picks_up_prompt_start_boundary() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
-        // Prompt prefills `<seed:think>` (id 10), opening reasoning before the stream.
-        parser.initialize(&[10]).unwrap();
+        // Prompt prefills `<seed:think>`, opening reasoning before the stream.
+        parser.initialize(&[SEED_THINK_START_ID]).unwrap();
 
         let delta = parser.push("reason</seed:think>answer").unwrap();
         assert_eq!(delta.reasoning.as_deref(), Some("reason"));
@@ -78,10 +81,10 @@ mod tests {
 
     #[test]
     fn respects_prompt_end_boundary() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
-        // Prompt already closed reasoning with `</seed:think>` (id 11).
-        parser.initialize(&[11]).unwrap();
+        // Prompt already closed reasoning with `</seed:think>`.
+        parser.initialize(&[SEED_THINK_END_ID]).unwrap();
 
         let delta = parser.push("answer").unwrap();
         assert_eq!(delta.reasoning, None);
@@ -91,7 +94,7 @@ mod tests {
     #[test]
     fn handles_explicit_start_token() {
         // An explicit start delimiter must not leak into reasoning text.
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
 
         let delta = parser.push("<seed:think>reason</seed:think>answer").unwrap();
@@ -103,7 +106,7 @@ mod tests {
     fn streams_explicit_start_token_across_pushes() {
         // Start token, reasoning body, end token, and content arrive in separate
         // streaming deltas.
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
 
         let mut reasoning = String::new();
@@ -131,9 +134,9 @@ mod tests {
 
     #[test]
     fn handles_partial_delimiters_across_pushes() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
-        parser.initialize(&[10]).unwrap();
+        parser.initialize(&[SEED_THINK_START_ID]).unwrap();
 
         // Closing delimiter `</seed:think>` arrives in two halves.
         let first = parser.push("reason</seed:").unwrap();

rust/src/parser/src/reasoning/step3p5.rs

@@ -127,14 +127,17 @@ mod tests {
     use std::sync::Arc;
 
     use super::Step3p5ReasoningParser;
-    use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
+    use crate::reasoning::{
+        ReasoningParser,
+        tests::{THINK_START_ID, fake_tokenizer},
+    };
 
     #[test]
     fn picks_up_prompt_start_boundary() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
-        // Prompt prefills `<think>` (id 1), opening reasoning before the stream.
-        parser.initialize(&[1]).unwrap();
+        // Prompt prefills `<think>`, opening reasoning before the stream.
+        parser.initialize(&[THINK_START_ID]).unwrap();
 
         let delta = parser.push("This is a reasoning section</think>This is the rest").unwrap();
         assert_eq!(
@@ -146,7 +149,7 @@ mod tests {
 
     #[test]
     fn handles_unterminated_reasoning() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let pushed = parser.push("<think>reason without end").unwrap();
@@ -159,7 +162,7 @@ mod tests {
 
     #[test]
     fn handles_empty_input() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let pushed = parser.push("").unwrap();
@@ -172,9 +175,9 @@ mod tests {
     fn complex_newline_pattern_trims_only_single_framing_newline_each_side() {
         // Only the immediately-adjacent framing `\n` is dropped on each side of
         // `</think>`; surrounding newlines remain part of reasoning/content.
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
-        parser.initialize(&[1]).unwrap();
+        parser.initialize(&[THINK_START_ID]).unwrap();
 
         let delta = parser
             .push("\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest")
@@ -188,7 +191,7 @@ mod tests {
 
     #[test]
     fn drops_framing_newlines_in_single_push() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let delta = parser.push("<think>reason\n</think>\nanswer").unwrap();
@@ -198,7 +201,7 @@ mod tests {
 
     #[test]
     fn drops_framing_newlines_across_pushes() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         // The trailing `\n` from the first push is held until we know whether
@@ -219,7 +222,7 @@ mod tests {
 
     #[test]
     fn replays_held_newline_when_more_reasoning_follows() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let first = parser.push("<think>reason\n").unwrap();
@@ -232,7 +235,7 @@ mod tests {
 
     #[test]
     fn finish_flushes_held_newline_in_unterminated_stream() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let first = parser.push("<think>reason\n").unwrap();
@@ -245,7 +248,7 @@ mod tests {
 
     #[test]
     fn preserves_inner_newlines_in_reasoning() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let delta = parser.push("<think>line1\nline2</think>tail").unwrap();
@@ -257,7 +260,7 @@ mod tests {
     fn trims_only_one_trailing_reasoning_newline() {
         // Only the single framing newline immediately before `</think>` is
         // dropped; earlier newlines in the reasoning body are preserved.
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let delta = parser.push("<think>reason\n\n</think>answer").unwrap();
@@ -269,7 +272,7 @@ mod tests {
     fn drops_only_first_content_newline_after_transition() {
         // The leading-`\n` drop applies only to the first content delta after
         // `</think>`; later deltas pass through untouched.
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let first = parser.push("<think>reason</think>").unwrap();
@@ -288,7 +291,7 @@ mod tests {
 
     #[test]
     fn passes_through_clean_boundary_without_framing_newlines() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let delta = parser.push("<think>reason</think>tail").unwrap();
@@ -298,7 +301,7 @@ mod tests {
 
     #[test]
     fn handles_empty_reasoning_section() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(fake_tokenizer());
         let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
 
         let delta = parser.push("<think></think>answer").unwrap();

rust/src/parser/src/reasoning/tests.rs

@@ -1,54 +1,42 @@
 use std::sync::Arc;
 
-use vllm_tokenizer::Tokenizer;
+use vllm_tokenizer::test_utils::TestTokenizer;
 
 use super::{
     DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
     Qwen3ReasoningParser, ReasoningParser,
 };
 
-pub(crate) struct FakeTokenizer;
+pub(crate) const THINK_START_ID: u32 = 256;
+pub(crate) const THINK_END_ID: u32 = 257;
+pub(crate) const START_THINKING_ID: u32 = 258;
+pub(crate) const END_THINKING_ID: u32 = 259;
+pub(crate) const MINIMAX_THINK_START_ID: u32 = 260;
+pub(crate) const MINIMAX_THINK_END_ID: u32 = 261;
+pub(crate) const SPECIAL_BOUNDARY_ID: u32 = 262;
+pub(crate) const MM_THINK_START_ID: u32 = 263;
+pub(crate) const MM_THINK_END_ID: u32 = 264;
+pub(crate) const SEED_THINK_START_ID: u32 = 265;
+pub(crate) const SEED_THINK_END_ID: u32 = 266;
 
-impl Tokenizer for FakeTokenizer {
-    fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
-        Ok(text.chars().map(u32::from).collect())
-    }
-
-    fn decode(
-        &self,
-        token_ids: &[u32],
-        _skip_special_tokens: bool,
-    ) -> vllm_tokenizer::Result<String> {
-        Ok(token_ids
-            .iter()
-            .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
-            .collect())
-    }
-
-    fn token_to_id(&self, token: &str) -> Option<u32> {
-        match token {
-            "<think>" => Some(1),
-            "</think>" => Some(2),
-            "<|START_THINKING|>" => Some(3),
-            "<|END_THINKING|>" => Some(4),
-            "◁think▷" => Some(5),
-            "◁/think▷" => Some(6),
-            "<mm:think>" => Some(8),
-            "</mm:think>" => Some(9),
-            "<seed:think>" => Some(10),
-            "</seed:think>" => Some(11),
-            _ => None,
-        }
-    }
-
-    fn is_special_id(&self, token_id: u32) -> bool {
-        token_id == 7
-    }
+pub(crate) fn fake_tokenizer() -> TestTokenizer {
+    TestTokenizer::new()
+        .with_regular_token("<think>", THINK_START_ID)
+        .with_regular_token("</think>", THINK_END_ID)
+        .with_regular_token("<|START_THINKING|>", START_THINKING_ID)
+        .with_regular_token("<|END_THINKING|>", END_THINKING_ID)
+        .with_regular_token("◁think▷", MINIMAX_THINK_START_ID)
+        .with_regular_token("◁/think▷", MINIMAX_THINK_END_ID)
+        .with_special_token("<special-boundary>", SPECIAL_BOUNDARY_ID)
+        .with_regular_token("<mm:think>", MM_THINK_START_ID)
+        .with_regular_token("</mm:think>", MM_THINK_END_ID)
+        .with_regular_token("<seed:think>", SEED_THINK_START_ID)
+        .with_regular_token("</seed:think>", SEED_THINK_END_ID)
 }
 
 #[test]
 fn delimited_content_only_stream() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser =
         DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
 
@@ -60,7 +48,7 @@ fn delimited_content_only_stream() {
 
 #[test]
 fn delimited_single_chunk_with_reasoning_and_content() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser =
         DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
 
@@ -71,7 +59,7 @@ fn delimited_single_chunk_with_reasoning_and_content() {
 
 #[test]
 fn delimited_partial_tokens_across_chunks() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser =
         DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
 
@@ -83,10 +71,10 @@ fn delimited_partial_tokens_across_chunks() {
 
 #[test]
 fn delimited_finish_flushes_buffer() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser =
         DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
-    parser.initialize(&[1]);
+    parser.initialize(&[THINK_START_ID]);
 
     let delta = parser.push("unfinished</thi");
     assert_eq!(delta.reasoning.as_deref(), Some("unfinished"));
@@ -96,7 +84,7 @@ fn delimited_finish_flushes_buffer() {
 
 #[test]
 fn qwen3_without_prompt_markers_expects_start_token() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
 
     let delta = parser.push("reason</think>answer").unwrap();
@@ -106,9 +94,9 @@ fn qwen3_without_prompt_markers_expects_start_token() {
 
 #[test]
 fn qwen3_prompt_end_marker_starts_in_content() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
-    parser.initialize(&[2]).unwrap();
+    parser.initialize(&[THINK_END_ID]).unwrap();
 
     let delta = parser.push("answer").unwrap();
     assert_eq!(delta.reasoning, None);
@@ -117,7 +105,7 @@ fn qwen3_prompt_end_marker_starts_in_content() {
 
 #[test]
 fn qwen3_tolerates_old_and_new_formats() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
 
     let mut old_parser = Qwen3ReasoningParser::new(tokenizer.clone()).unwrap();
     let old = old_parser.push("<think>reason</think>answer").unwrap();
@@ -125,7 +113,7 @@ fn qwen3_tolerates_old_and_new_formats() {
     assert_eq!(old.content.as_deref(), Some("answer"));
 
     let mut new_parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
-    new_parser.initialize(&[1]).unwrap();
+    new_parser.initialize(&[THINK_START_ID]).unwrap();
     let new = new_parser.push("reason</think>answer").unwrap();
     assert_eq!(new.reasoning.as_deref(), Some("reason"));
     assert_eq!(new.content.as_deref(), Some("answer"));
@@ -133,10 +121,10 @@ fn qwen3_tolerates_old_and_new_formats() {
 
 #[test]
 fn qwen3_stops_scanning_at_last_special_token() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
 
-    parser.initialize(&[1, 7]).unwrap();
+    parser.initialize(&[THINK_START_ID, SPECIAL_BOUNDARY_ID]).unwrap();
 
     let delta = parser.push("answer").unwrap();
     assert_eq!(delta.reasoning, None);
@@ -145,7 +133,7 @@ fn qwen3_stops_scanning_at_last_special_token() {
 
 #[test]
 fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();
 
     let delta = parser.push("reason</think>answer").unwrap();
@@ -155,10 +143,10 @@ fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() {
 
 #[test]
 fn deepseek_r1_stops_scanning_at_last_special_token() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();
 
-    parser.initialize(&[2, 7]).unwrap();
+    parser.initialize(&[THINK_END_ID, SPECIAL_BOUNDARY_ID]).unwrap();
 
     let delta = parser.push("reason</think>answer").unwrap();
     assert_eq!(delta.reasoning.as_deref(), Some("reason"));
@@ -167,7 +155,7 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
 
 #[test]
 fn minimax_m3_handles_explicit_think_delimiters() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
 
     let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
@@ -177,7 +165,7 @@ fn minimax_m3_handles_explicit_think_delimiters() {
 
 #[test]
 fn minimax_m3_drops_leading_end_marker() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
 
     let delta = parser.push("</mm:think>answer").unwrap();
@@ -187,7 +175,7 @@ fn minimax_m3_drops_leading_end_marker() {
 
 #[test]
 fn minimax_m3_preserves_non_leading_end_marker() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
 
     let delta = parser.push("XXX</mm:think>YYY").unwrap();
@@ -197,7 +185,7 @@ fn minimax_m3_preserves_non_leading_end_marker() {
 
 #[test]
 fn minimax_m3_drops_split_leading_end_marker() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
 
     assert!(parser.push("</mm").unwrap().is_empty());
@@ -208,9 +196,9 @@ fn minimax_m3_drops_split_leading_end_marker() {
 
 #[test]
 fn minimax_m3_uses_prompt_prefilled_start_marker() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
-    parser.initialize(&[8]).unwrap();
+    parser.initialize(&[MM_THINK_START_ID]).unwrap();
 
     let delta = parser.push("reason</mm:think>answer").unwrap();
     assert_eq!(delta.reasoning.as_deref(), Some("reason"));
@@ -219,9 +207,9 @@ fn minimax_m3_uses_prompt_prefilled_start_marker() {
 
 #[test]
 fn minimax_m3_uses_prompt_prefilled_end_marker() {
-    let tokenizer = Arc::new(FakeTokenizer);
+    let tokenizer = Arc::new(fake_tokenizer());
     let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
-    parser.initialize(&[9]).unwrap();
+    parser.initialize(&[MM_THINK_END_ID]).unwrap();
 
     let delta = parser.push("answer").unwrap();
     assert_eq!(delta.reasoning, None);

rust/src/parser/src/unified/combined.rs

@@ -124,42 +124,17 @@ impl UnifiedParser for CombinedParser {
 mod tests {
     use std::sync::Arc;
 
-    use vllm_tokenizer::Tokenizer;
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::CombinedParser;
     use crate::reasoning::{Qwen3ReasoningParser, ReasoningDelta, ReasoningParser};
     use crate::tool::{Qwen3XmlToolParser, Tool, ToolParser};
     use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
 
-    struct FakeTokenizer;
-
-    impl Tokenizer for FakeTokenizer {
-        fn encode(
-            &self,
-            text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            Ok(text.chars().map(u32::from).collect())
-        }
-
-        fn decode(
-            &self,
-            token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            Ok(token_ids
-                .iter()
-                .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
-                .collect())
-        }
-
-        fn token_to_id(&self, token: &str) -> Option<u32> {
-            match token {
-                "<think>" => Some(1),
-                "</think>" => Some(2),
-                _ => None,
-            }
-        }
+    fn tokenizer() -> TestTokenizer {
+        TestTokenizer::new()
+            .with_regular_token("<think>", 256)
+            .with_regular_token("</think>", 257)
     }
 
     fn test_tools() -> Vec<Tool> {
@@ -273,7 +248,7 @@ mod tests {
 
     #[test]
     fn combined_parser_emits_reasoning_and_text() {
-        let tokenizer = Arc::new(FakeTokenizer);
+        let tokenizer = Arc::new(tokenizer());
         let reasoning = Qwen3ReasoningParser::create(tokenizer).unwrap();
         let mut parser = CombinedParser::new(Some(reasoning), None);
 

rust/src/parser/src/unified/gemma4.rs

@@ -515,7 +515,7 @@ mod tests {
 
     use serde_json::{Value, json};
     use thiserror_ext::AsReport;
-    use vllm_tokenizer::Tokenizer;
+    use vllm_tokenizer::test_utils::TestTokenizer;
     use winnow::combinator::{eof, terminated};
     use winnow::error::ErrMode;
     use winnow::prelude::*;
@@ -527,66 +527,15 @@ mod tests {
     use crate::tool::Tool;
     use crate::unified::{UnifiedParserEvent, parsing_failed};
 
-    struct FakeTokenizer;
+    const CHANNEL_START_ID: u32 = 256;
+    const CHANNEL_END_ID: u32 = 257;
+    const TURN_BOUNDARY_ID: u32 = 258;
 
-    impl Tokenizer for FakeTokenizer {
-        fn encode(
-            &self,
-            text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            Ok(text.chars().map(u32::from).collect())
-        }
-
-        fn decode(
-            &self,
-            token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            Ok(token_ids
-                .iter()
-                .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
-                .collect())
-        }
-
-        fn token_to_id(&self, token: &str) -> Option<u32> {
-            match token {
-                CHANNEL_START => Some(100),
-                CHANNEL_END => Some(101),
-                _ => None,
-            }
-        }
-
-        fn is_special_id(&self, token_id: u32) -> bool {
-            matches!(token_id, 100..=105)
-        }
-    }
-
-    struct MissingTokenTokenizer;
-
-    impl Tokenizer for MissingTokenTokenizer {
-        fn encode(
-            &self,
-            text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            Ok(text.chars().map(u32::from).collect())
-        }
-
-        fn decode(
-            &self,
-            token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            Ok(token_ids
-                .iter()
-                .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
-                .collect())
-        }
-
-        fn token_to_id(&self, _token: &str) -> Option<u32> {
-            None
-        }
+    fn tokenizer() -> TestTokenizer {
+        TestTokenizer::new()
+            .with_special_token(CHANNEL_START, CHANNEL_START_ID)
+            .with_special_token(CHANNEL_END, CHANNEL_END_ID)
+            .with_special_token("<turn-boundary>", TURN_BOUNDARY_ID)
     }
 
     trait UnifiedParserTestExt {
@@ -716,12 +665,12 @@ mod tests {
     }
 
     fn test_parser() -> Gemma4UnifiedParser {
-        Gemma4UnifiedParser::new(&test_tools(), Arc::new(FakeTokenizer)).unwrap()
+        Gemma4UnifiedParser::new(&test_tools(), Arc::new(tokenizer())).unwrap()
     }
 
     #[test]
     fn gemma4_create_requires_channel_start_token() {
-        let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(MissingTokenTokenizer)) {
+        let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(TestTokenizer::new())) {
             Ok(_) => panic!("expected missing token error"),
             Err(error) => error,
         };
@@ -1046,7 +995,7 @@ mod tests {
     #[test]
     fn gemma4_initialize_open_channel_prompt_starts_in_reasoning() {
         let mut parser = test_parser();
-        parser.initialize(&[100, 3000, 3001]).unwrap();
+        parser.initialize(&[CHANNEL_START_ID, 3000, 3001]).unwrap();
 
         let output = parser.parse_complete("reason<channel|>answer").unwrap();
 
@@ -1057,7 +1006,7 @@ mod tests {
     #[test]
     fn gemma4_initialize_turn_prompt_starts_in_text() {
         let mut parser = test_parser();
-        parser.initialize(&[104, 3000, 3001]).unwrap();
+        parser.initialize(&[TURN_BOUNDARY_ID, 3000, 3001]).unwrap();
 
         let output = parser.parse_complete("<|channel>thought\nreason<channel|>answer").unwrap();
 
@@ -1068,7 +1017,7 @@ mod tests {
     #[test]
     fn gemma4_initialize_special_token_caps_boundary_scan() {
         let mut parser = test_parser();
-        parser.initialize(&[100, 3000, 104, 3001]).unwrap();
+        parser.initialize(&[CHANNEL_START_ID, 3000, TURN_BOUNDARY_ID, 3001]).unwrap();
 
         let output = parser.parse_complete("answer").unwrap();
 
@@ -1079,7 +1028,7 @@ mod tests {
     #[test]
     fn gemma4_initialize_closed_channel_prompt_starts_in_text() {
         let mut parser = test_parser();
-        parser.initialize(&[100, 3000, 3001, 101]).unwrap();
+        parser.initialize(&[CHANNEL_START_ID, 3000, 3001, CHANNEL_END_ID]).unwrap();
 
         let output = parser.parse_complete("answer").unwrap();
 

rust/src/server/Cargo.toml

@@ -61,8 +61,10 @@ expect-test.workspace = true
 rmp-serde.workspace = true
 serial_test.workspace = true
 tempfile.workspace = true
+tokio = { workspace = true, features = ["test-util"] }
 tower.workspace = true
 vllm-engine-core-client = { workspace = true, features = ["test-util"] }
+vllm-tokenizer = { workspace = true, features = ["test-utils"] }
 zeromq.workspace = true
 
 [lints]

rust/src/server/src/grpc/mod.rs

@@ -4,21 +4,16 @@ mod convert;
 
 use std::pin::Pin;
 use std::sync::Arc;
-use std::task::{Context, Poll};
 
-use futures::{Stream, StreamExt as _, stream};
+use futures::{Stream, StreamExt as _};
 use thiserror_ext::AsReport as _;
-use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio::sync::mpsc;
-use tokio_openssl::SslStream;
 use tokio_stream::wrappers::ReceiverStream;
-use tonic::transport::server::{Connected, TcpConnectInfo};
 use tonic::{Request, Response, Status};
 use tracing::info;
 use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _};
 
 use self::convert::ResponseOpts;
-use crate::listener::{Listener, ListenerIo};
 use crate::state::AppState;
 
 /// Generated protobuf/gRPC types for the `vllm` package.
@@ -31,78 +26,6 @@ pub use pb::generate_server::GenerateServer;
 #[cfg(test)]
 mod tests;
 
-/// Newtype over `tokio-openssl`'s `SslStream` so we can implement tonic's
-/// [`Connected`] on it (the orphan rule blocks doing so on the foreign type).
-pub(crate) struct GrpcTlsStream {
-    inner: SslStream<ListenerIo>,
-}
-
-impl GrpcTlsStream {
-    pub(crate) fn new(inner: SslStream<ListenerIo>) -> Self {
-        Self { inner }
-    }
-}
-
-impl AsyncRead for GrpcTlsStream {
-    fn poll_read(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &mut ReadBuf<'_>,
-    ) -> Poll<std::io::Result<()>> {
-        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
-    }
-}
-
-impl AsyncWrite for GrpcTlsStream {
-    fn poll_write(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &[u8],
-    ) -> Poll<std::io::Result<usize>> {
-        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
-    }
-
-    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
-    }
-
-    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
-    }
-}
-
-impl Connected for GrpcTlsStream {
-    type ConnectInfo = TcpConnectInfo;
-
-    fn connect_info(&self) -> TcpConnectInfo {
-        self.inner.get_ref().connect_info()
-    }
-}
-
-/// Adapt the shared server listener into tonic's incoming stream shape.
-pub(crate) fn incoming(listener: Listener) -> impl Stream<Item = std::io::Result<ListenerIo>> {
-    stream::unfold(listener, |mut listener| async move {
-        let (io, _) = axum::serve::Listener::accept(&mut listener).await;
-        Some((Ok(io), listener))
-    })
-}
-
-/// Wrap the gRPC listener so each accepted connection completes a TLS handshake
-/// before tonic serves it.
-pub(crate) fn tls_incoming(
-    listener: Listener,
-    context: openssl::ssl::SslContext,
-    handshake_timeout: std::time::Duration,
-) -> impl Stream<Item = std::io::Result<GrpcTlsStream>> {
-    tls_listener::builder(context)
-        .handshake_timeout(handshake_timeout)
-        .listen(listener)
-        .map(|res| {
-            res.map(|(inner, _addr)| GrpcTlsStream::new(inner))
-                .map_err(std::io::Error::other)
-        })
-}
-
 /// gRPC Generate service implementation backed by the shared application state.
 pub struct GenerateServiceImpl {
     state: Arc<AppState>,

rust/src/server/src/grpc/tests.rs

@@ -24,14 +24,15 @@ use vllm_engine_core_client::protocol::{
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
 use vllm_llm::Llm;
-use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
+use vllm_text::tokenizer::DynTokenizer;
 use vllm_text::{Prompt, TextBackend};
+use vllm_tokenizer::test_utils::TestTokenizer;
 use zeromq::prelude::{SocketRecv, SocketSend};
 use zeromq::{DealerSocket, PushSocket, ZmqMessage};
 
 use super::pb::generate_client::GenerateClient;
-use super::{GenerateServer, GenerateServiceImpl, incoming, pb, tls_incoming};
-use crate::listener::Listener;
+use super::{GenerateServer, GenerateServiceImpl, pb};
+use crate::listener::{Listener, MaybeTlsListener};
 use crate::state::AppState;
 use crate::tls;
 use crate::tls_tests::{TestCerts, server_tls};
@@ -155,37 +156,9 @@ fn test_llm(client: EngineCoreClient) -> Llm {
 #[derive(Clone, Debug)]
 struct FakeTextBackend;
 
-#[derive(Debug)]
-struct FakeTokenizer;
-
-impl Tokenizer for FakeTokenizer {
-    fn encode(
-        &self,
-        text: &str,
-        _add_special_tokens: bool,
-    ) -> vllm_text::tokenizer::Result<Vec<u32>> {
-        Ok(text.bytes().map(u32::from).collect())
-    }
-
-    fn decode(
-        &self,
-        token_ids: &[u32],
-        _skip_special_tokens: bool,
-    ) -> vllm_text::tokenizer::Result<String> {
-        Ok(
-            String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
-                .into_owned(),
-        )
-    }
-
-    fn token_to_id(&self, token: &str) -> Option<u32> {
-        token.bytes().next().map(u32::from)
-    }
-}
-
 impl TextBackend for FakeTextBackend {
     fn tokenizer(&self) -> DynTokenizer {
-        Arc::new(FakeTokenizer)
+        Arc::new(TestTokenizer::new())
     }
 
     fn model_id(&self) -> &str {
@@ -287,7 +260,7 @@ async fn grpc_test_server(
     let addr = listener.local_addr().expect("local addr");
 
     let server_task = tokio::spawn(async move {
-        let incoming = incoming(Listener::Tcp(listener));
+        let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
         TonicServer::builder()
             .add_service(svc)
             .serve_with_incoming(incoming)
@@ -318,7 +291,7 @@ async fn grpc_tls_test_server(
     let addr = listener.local_addr().expect("local addr").to_string();
 
     let server_task = tokio::spawn(async move {
-        let incoming = tls_incoming(Listener::Tcp(listener), context, tls::TLS_HANDSHAKE_TIMEOUT);
+        let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
         TonicServer::builder()
             .add_service(svc)
             .serve_with_incoming(incoming)
@@ -413,7 +386,7 @@ async fn grpc_server_with_keepalive(
     }
 
     let server_task = tokio::spawn(async move {
-        let incoming = incoming(Listener::Tcp(listener));
+        let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
         builder
             .add_service(svc)
             .serve_with_incoming(incoming)

rust/src/server/src/lib.rs

@@ -27,7 +27,6 @@ pub use config::{
     ApiServerOptions, Config, CoordinatorMode, CorsConfig, DEFAULT_KEEP_ALIVE_TIMEOUT,
     HttpListenerMode, TlsConfig,
 };
-use futures::FutureExt as _;
 use hyper::body::Incoming;
 use hyper::server::conn::http1;
 use hyper_util::rt::{TokioIo, TokioTimer};
@@ -45,7 +44,7 @@ use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
 use vllm_llm::Llm;
 use vllm_text::TextLlm;
 
-use crate::listener::Listener;
+use crate::listener::{Listener, MaybeTlsListener};
 use crate::routes::build_router;
 use crate::server_info::ServerInfoSnapshot;
 use crate::state::AppState;
@@ -179,7 +178,7 @@ where
     let listener = Listener::bind(&config.listener_mode)
         .await
         .context("failed to bind listener for OpenAI server")?;
-    let bind_address = listener.local_addr()?;
+    let bind_address = listener.local_addr_display()?;
     let model = state.primary_model_name().to_owned();
     let app = extend_router(build_router(state.clone()));
 
@@ -252,7 +251,6 @@ where
     // silent client cannot hold the connection open.
     let keep_alive_timeout = config.keep_alive_timeout;
     let timeouts = ConnectionTimeouts {
-        handshake: tls::TLS_HANDSHAKE_TIMEOUT,
         header_read: if keep_alive_timeout.is_zero() {
             DEFAULT_KEEP_ALIVE_TIMEOUT
         } else {
@@ -266,9 +264,15 @@ where
         let server_shutdown = server_shutdown.clone();
         let force_shutdown = force_shutdown.clone();
         async move {
+            let listener = match tls_config {
+                Some(context) => MaybeTlsListener::tls(listener, context),
+                None => MaybeTlsListener::plain(listener),
+            };
+            let server = serve_connections(listener, app, shutdown.cancelled_owned(), timeouts);
+
             let result = tokio::select! {
-                result = serve_listener(listener, tls_config, app, shutdown.cancelled_owned(), timeouts) => {
-                    result
+                result = server => {
+                    result.context("HTTP server failed")
                 }
                 _ = force_shutdown.cancelled() => {
                     warn!("HTTP graceful shutdown deadline elapsed; aborting server");
@@ -292,18 +296,11 @@ where
                 shutdown.cancelled().await;
                 return Ok(());
             };
-            // Box to unify the TLS and plaintext arms' different stream types.
-            let server = match grpc_tls {
-                Some(context) => {
-                    let incoming =
-                        grpc::tls_incoming(grpc_listener, context, tls::TLS_HANDSHAKE_TIMEOUT);
-                    svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()).boxed()
-                }
-                None => {
-                    let incoming = grpc::incoming(grpc_listener);
-                    svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()).boxed()
-                }
+            let incoming = match grpc_tls {
+                Some(context) => MaybeTlsListener::tls(grpc_listener, context),
+                None => MaybeTlsListener::plain(grpc_listener),
             };
+            let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
 
             let result = tokio::select! {
                 result = server => {
@@ -333,61 +330,19 @@ where
 /// Per-connection timeouts applied while serving HTTP/HTTPS.
 #[derive(Clone, Copy)]
 pub(crate) struct ConnectionTimeouts {
-    /// Max time for a client to complete the TLS handshake (TLS path only).
-    pub(crate) handshake: Duration,
     /// HTTP/1 header-read timeout (bounds idle keep-alive and the head read).
     pub(crate) header_read: Duration,
     /// Whether HTTP/1 keep-alive is enabled; `false` closes after each response.
     pub(crate) keep_alive_enabled: bool,
 }
 
-/// Apply optional TLS termination and per-connection HTTP timeouts, then serve
-/// `app`. Shared by [`serve_with_router_extension`] and the TLS tests.
-async fn serve_listener(
-    listener: Listener,
-    tls: Option<openssl::ssl::SslContext>,
-    app: Router,
-    shutdown: impl Future<Output = ()> + Send + 'static,
-    timeouts: ConnectionTimeouts,
-) -> Result<()> {
-    match tls {
-        Some(context) => {
-            // tls-listener terminates TLS (handshake + timeout); serve_connections
-            // owns the HTTP keep-alive/idle bound that axum::serve cannot express.
-            // Failed handshakes (incl. timeouts) log at ERROR via tls-listener.
-            let listener = tls_listener::builder(context)
-                .handshake_timeout(timeouts.handshake)
-                .listen(listener);
-            serve_connections(
-                listener,
-                app,
-                shutdown,
-                timeouts.header_read,
-                timeouts.keep_alive_enabled,
-            )
-            .await
-            .context("HTTPS server failed")
-        }
-        None => serve_connections(
-            listener,
-            app,
-            shutdown,
-            timeouts.header_read,
-            timeouts.keep_alive_enabled,
-        )
-        .await
-        .context("HTTP server failed"),
-    }
-}
-
 /// Serve `app` per connection (HTTP/1) with a keep-alive idle timeout and
 /// graceful drain. Hand-rolled on hyper because [`axum::serve()`] takes no config.
 async fn serve_connections<L>(
     mut listener: L,
     app: Router,
     shutdown: impl Future<Output = ()> + Send,
-    header_read: Duration,
-    keep_alive_enabled: bool,
+    timeouts: ConnectionTimeouts,
 ) -> Result<()>
 where
     L: axum::serve::Listener,
@@ -404,8 +359,8 @@ where
             app.clone().map_request(|req: Request<Incoming>| req.map(Body::new)),
         );
         let mut builder = http1::Builder::new();
-        builder.timer(TokioTimer::new()).header_read_timeout(header_read);
-        if !keep_alive_enabled {
+        builder.timer(TokioTimer::new()).header_read_timeout(timeouts.header_read);
+        if !timeouts.keep_alive_enabled {
             builder.keep_alive(false);
         }
         let connection = builder.serve_connection(TokioIo::new(io), service);

rust/src/server/src/listener.rs

@@ -12,13 +12,14 @@ use std::pin::Pin;
 use std::task::{Context, Poll, ready};
 
 use auto_enums::enum_derive;
+use openssl::ssl::SslContext;
 use socket2::Socket;
 use tls_listener::{AsyncAccept, AsyncListener};
 use tokio::net::{TcpListener, TcpStream, UnixListener, UnixStream};
 use tonic::transport::server::{Connected, TcpConnectInfo};
 use tracing::trace;
 
-use crate::HttpListenerMode;
+use crate::{HttpListenerMode, tls};
 
 /// Runtime listener type used by the OpenAI-compatible HTTP or gRPC server,
 /// which is either a TCP listener or a Unix-domain listener.
@@ -61,7 +62,7 @@ impl Listener {
 
     /// Return a log-friendly local address string for either TCP or Unix
     /// sockets.
-    pub fn local_addr(&self) -> Result<String> {
+    pub fn local_addr_display(&self) -> Result<String> {
         match self {
             Self::Tcp(listener) => Ok(listener.local_addr()?.to_string()),
             Self::Unix(listener) => Ok(match listener.local_addr()?.as_pathname() {
@@ -92,7 +93,7 @@ impl Listener {
         }
     }
 
-    fn listener_addr(&self) -> Result<ListenerAddr> {
+    fn local_addr(&self) -> Result<ListenerAddr> {
         match self {
             Self::Tcp(listener) => listener.local_addr().map(ListenerAddr::Tcp),
             Self::Unix(listener) => listener.local_addr().map(ListenerAddr::Unix),
@@ -100,6 +101,7 @@ impl Listener {
     }
 }
 
+/// Allow the unified listener to plug directly into tonic's gRPC server.
 impl Connected for ListenerIo {
     type ConnectInfo = TcpConnectInfo;
 
@@ -144,7 +146,7 @@ impl axum::serve::Listener for Listener {
     }
 
     fn local_addr(&self) -> Result<Self::Addr> {
-        self.listener_addr()
+        self.local_addr()
     }
 }
 
@@ -173,10 +175,98 @@ impl AsyncAccept for Listener {
         }
     }
 }
-
 impl AsyncListener for Listener {
     fn local_addr(&self) -> Result<Self::Address> {
-        self.listener_addr()
+        self.local_addr()
+    }
+}
+
+/// A listener that may be either a plain TCP/UDS listener or a TLS listener over it.
+pub enum MaybeTlsListener {
+    Plain(Listener),
+    Tls(tls_listener::TlsListener<Listener, SslContext>),
+}
+
+impl MaybeTlsListener {
+    /// Create a plain listener without TLS.
+    pub fn plain(listener: Listener) -> Self {
+        Self::Plain(listener)
+    }
+
+    /// Create a TLS listener over the given plain listener.
+    pub fn tls(listener: Listener, context: SslContext) -> Self {
+        Self::Tls(
+            tls_listener::builder(context)
+                .handshake_timeout(tls::TLS_HANDSHAKE_TIMEOUT)
+                .listen(listener),
+        )
+    }
+}
+
+/// Listener I/O type that may be either a plain TCP/UDS stream or a TLS stream over it.
+#[derive(Debug)]
+#[enum_derive(tokio1::AsyncRead, tokio1::AsyncWrite)]
+pub enum MaybeTlsStream {
+    Plain(ListenerIo),
+    Tls(tokio_openssl::SslStream<ListenerIo>),
+}
+
+/// Allow the maybe-TLS listener to plug directly into `axum::serve(...)`.
+impl axum::serve::Listener for MaybeTlsListener {
+    type Addr = ListenerAddr;
+    type Io = MaybeTlsStream;
+
+    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
+        match self {
+            Self::Plain(listener) => {
+                let (io, addr) = axum::serve::Listener::accept(listener).await;
+                (MaybeTlsStream::Plain(io), addr)
+            }
+            Self::Tls(tls_listener) => {
+                let (io, addr) = axum::serve::Listener::accept(tls_listener).await;
+                (MaybeTlsStream::Tls(io), addr)
+            }
+        }
+    }
+
+    fn local_addr(&self) -> tokio::io::Result<Self::Addr> {
+        match self {
+            Self::Plain(listener) => listener.local_addr(),
+            Self::Tls(tls_listener) => tls_listener.local_addr(),
+        }
+    }
+}
+
+/// Allow the maybe-TLS listener to plug directly into tonic's gRPC server.
+impl Connected for MaybeTlsStream {
+    type ConnectInfo = TcpConnectInfo;
+
+    fn connect_info(&self) -> TcpConnectInfo {
+        match self {
+            Self::Plain(stream) => stream.connect_info(),
+            Self::Tls(stream) => stream.get_ref().connect_info(),
+        }
+    }
+}
+
+/// Allow the maybe-TLS listener to be adaptable to tonic's incoming stream shape.
+impl futures::Stream for MaybeTlsListener {
+    type Item = std::io::Result<MaybeTlsStream>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match self.get_mut() {
+            Self::Plain(listener) => {
+                let listener = Pin::new(listener);
+                let (io, _) = ready!(listener.poll_accept(cx))?;
+                Poll::Ready(Some(Ok(MaybeTlsStream::Plain(io))))
+            }
+            Self::Tls(tls_listener) => {
+                let tls_listener = Pin::new(tls_listener);
+                let (io, _) =
+                    ready!(tls_listener.poll_accept(cx)).map_err(std::io::Error::other)?;
+                Poll::Ready(Some(Ok(MaybeTlsStream::Tls(io))))
+            }
+        }
     }
 }
 

rust/src/server/src/routes/http_client_tests.rs

@@ -24,8 +24,9 @@ use vllm_engine_core_client::protocol::{
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
 use vllm_llm::Llm;
-use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
+use vllm_text::tokenizer::DynTokenizer;
 use vllm_text::{Prompt, TextBackend};
+use vllm_tokenizer::test_utils::TestTokenizer;
 use zeromq::prelude::{SocketRecv, SocketSend};
 use zeromq::{DealerSocket, PushSocket, ZmqMessage};
 
@@ -151,37 +152,9 @@ fn test_llm(client: EngineCoreClient) -> Llm {
 #[derive(Clone, Debug)]
 struct FakeChatBackend;
 
-#[derive(Debug)]
-struct FakeChatTokenizer;
-
-impl Tokenizer for FakeChatTokenizer {
-    fn encode(
-        &self,
-        text: &str,
-        _add_special_tokens: bool,
-    ) -> vllm_text::tokenizer::Result<Vec<u32>> {
-        Ok(text.bytes().map(u32::from).collect())
-    }
-
-    fn decode(
-        &self,
-        token_ids: &[u32],
-        _skip_special_tokens: bool,
-    ) -> vllm_text::tokenizer::Result<String> {
-        Ok(
-            String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
-                .into_owned(),
-        )
-    }
-
-    fn token_to_id(&self, token: &str) -> Option<u32> {
-        token.bytes().next().map(u32::from)
-    }
-}
-
 impl TextBackend for FakeChatBackend {
     fn tokenizer(&self) -> DynTokenizer {
-        Arc::new(FakeChatTokenizer)
+        Arc::new(TestTokenizer::new())
     }
 
     fn model_id(&self) -> &str {

rust/src/server/src/routes/openai/completions/convert.rs

@@ -194,7 +194,7 @@ mod tests {
     use axum::http::HeaderMap;
     use serde_json::json;
     use vllm_text::Prompt;
-    use vllm_text::tokenizer::Tokenizer;
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::prepare_completion_request;
     use crate::lora::LoraModelResolution;
@@ -212,32 +212,8 @@ mod tests {
         }
     }
 
-    #[derive(Debug)]
-    struct TestTokenizer;
-
-    impl Tokenizer for TestTokenizer {
-        fn encode(
-            &self,
-            text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_text::tokenizer::Result<Vec<u32>> {
-            Ok(text.bytes().map(u32::from).collect())
-        }
-
-        fn decode(
-            &self,
-            token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_text::tokenizer::Result<String> {
-            Ok(
-                String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
-                    .into_owned(),
-            )
-        }
-
-        fn token_to_id(&self, _token: &str) -> Option<u32> {
-            None
-        }
+    fn test_tokenizer() -> TestTokenizer {
+        TestTokenizer::new()
     }
 
     fn base_request_json() -> serde_json::Value {
@@ -297,7 +273,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
 
@@ -340,7 +316,7 @@ mod tests {
                 request,
                 &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
                 ResolvedRequestContext::default(),
-                &TestTokenizer,
+                &test_tokenizer(),
             )
             .expect("prepare")
             .text_request
@@ -374,7 +350,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
 
@@ -399,7 +375,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
 
@@ -422,7 +398,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
 
@@ -446,7 +422,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
 
@@ -471,7 +447,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
 
@@ -498,7 +474,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
 
@@ -523,7 +499,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
 
@@ -549,7 +525,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
         assert_eq!(prepared.text_request.sampling_params.logprobs, Some(1));
@@ -574,7 +550,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             request_context(&headers, None),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
         assert_eq!(prepared.text_request.data_parallel_rank, Some(3));
@@ -593,7 +569,7 @@ mod tests {
             request,
             &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
             ResolvedRequestContext::default(),
-            &TestTokenizer,
+            &test_tokenizer(),
         )
         .expect("prepare");
         assert_eq!(prepared.text_request.data_parallel_rank, None);

rust/src/server/src/routes/tests.rs

@@ -40,8 +40,9 @@ use vllm_engine_core_client::{
 };
 use vllm_llm::Llm;
 use vllm_metrics::METRICS;
-use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
+use vllm_text::tokenizer::DynTokenizer;
 use vllm_text::{Prompt, TextBackend};
+use vllm_tokenizer::test_utils::TestTokenizer;
 use zeromq::prelude::{SocketRecv, SocketSend};
 use zeromq::{DealerSocket, PushSocket, ZmqMessage};
 
@@ -417,82 +418,20 @@ struct FakeChatBackend {
 }
 
 /// Synthetic BOS id used when `add_special_tokens` is true in tests.
-const FAKE_BOS_TOKEN_ID: u32 = 1;
+const FAKE_BOS_TOKEN_ID: u32 = 256;
+const UNKNOWN_DECODE_TOKEN_ID: u32 = 10_000;
 
-#[derive(Debug)]
-struct FakeChatTokenizer;
-
-impl Tokenizer for FakeChatTokenizer {
-    fn encode(
-        &self,
-        text: &str,
-        add_special_tokens: bool,
-    ) -> vllm_text::tokenizer::Result<Vec<u32>> {
-        let mut token_ids = Vec::new();
-        if add_special_tokens {
-            token_ids.push(FAKE_BOS_TOKEN_ID);
-        }
-        let mut rest = text;
-        while !rest.is_empty() {
-            if let Some(stripped) = rest.strip_prefix("<image>") {
-                token_ids.push(999);
-                rest = stripped;
-                continue;
-            }
-            if let Some(stripped) = rest.strip_prefix("<|image_pad|>") {
-                token_ids.push(151655);
-                rest = stripped;
-                continue;
-            }
-
-            let ch = rest.chars().next().expect("rest is not empty");
-            let mut buf = [0; 4];
-            token_ids.extend(ch.encode_utf8(&mut buf).bytes().map(u32::from));
-            rest = &rest[ch.len_utf8()..];
-        }
-        Ok(token_ids)
-    }
-
-    fn decode(
-        &self,
-        token_ids: &[u32],
-        _skip_special_tokens: bool,
-    ) -> vllm_text::tokenizer::Result<String> {
-        Ok(
-            String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
-                .into_owned(),
-        )
-    }
-
-    fn token_to_id(&self, token: &str) -> Option<u32> {
-        match token {
-            "<image>" => Some(999),
-            "<|image_pad|>" => Some(151655),
-            "<think>" => Some(0xF001),
-            "</think>" => Some(0xF002),
-            "<|START_THINKING|>" => Some(0xF003),
-            "<|END_THINKING|>" => Some(0xF004),
-            "◁think▷" => Some(0xF005),
-            "◁/think▷" => Some(0xF006),
-            _ => None,
-        }
-    }
-
-    fn id_to_token(&self, id: u32) -> Option<String> {
-        match id {
-            FAKE_BOS_TOKEN_ID => Some("<bos>".to_string()),
-            999 => Some("<image>".to_string()),
-            151655 => Some("<|image_pad|>".to_string()),
-            0xF001 => Some("<think>".to_string()),
-            0xF002 => Some("</think>".to_string()),
-            0xF003 => Some("<|START_THINKING|>".to_string()),
-            0xF004 => Some("<|END_THINKING|>".to_string()),
-            0xF005 => Some("◁think▷".to_string()),
-            0xF006 => Some("◁/think▷".to_string()),
-            id if id < 128 => char::from_u32(id).map(|ch| ch.to_string()),
-            _ => None,
-        }
-    }
+fn fake_chat_tokenizer() -> TestTokenizer {
+    TestTokenizer::new()
+        .with_bos_token("<bos>", FAKE_BOS_TOKEN_ID)
+        .with_regular_token("<image>", 999)
+        .with_regular_token("<|image_pad|>", 151655)
+        .with_regular_token("<think>", 0xF001)
+        .with_regular_token("</think>", 0xF002)
+        .with_regular_token("<|START_THINKING|>", 0xF003)
+        .with_regular_token("<|END_THINKING|>", 0xF004)
+        .with_regular_token("◁think▷", 0xF005)
+        .with_regular_token("◁/think▷", 0xF006)
 }
 
 impl FakeChatBackend {
@@ -530,7 +469,7 @@ impl fmt::Debug for FakeChatBackend {
 
 impl TextBackend for FakeChatBackend {
     fn tokenizer(&self) -> DynTokenizer {
-        Arc::new(FakeChatTokenizer)
+        Arc::new(fake_chat_tokenizer())
     }
 
     fn model_id(&self) -> &str {
@@ -630,7 +569,7 @@ fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
         Some("qwen2_vl".to_string()),
         Some(&config_path),
         None,
-        Arc::new(FakeChatTokenizer),
+        Arc::new(fake_chat_tokenizer()),
     )
     .expect("load multimodal info")
     .expect("qwen multimodal info is registered");
@@ -638,70 +577,6 @@ fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
     info
 }
 
-#[derive(Clone, Debug)]
-struct FailingDecodeChatBackend;
-
-#[derive(Debug)]
-struct FailingDecodeTokenizer;
-
-impl Tokenizer for FailingDecodeTokenizer {
-    fn encode(
-        &self,
-        text: &str,
-        add_special_tokens: bool,
-    ) -> vllm_text::tokenizer::Result<Vec<u32>> {
-        FakeChatTokenizer.encode(text, add_special_tokens)
-    }
-
-    fn decode(
-        &self,
-        token_ids: &[u32],
-        skip_special_tokens: bool,
-    ) -> vllm_text::tokenizer::Result<String> {
-        if token_ids.contains(&(b'i' as u32)) {
-            return Err(vllm_text::tokenizer::TokenizerError(
-                "forced decode failure for streaming test".to_string(),
-            ));
-        }
-
-        FakeChatTokenizer.decode(token_ids, skip_special_tokens)
-    }
-
-    fn token_to_id(&self, token: &str) -> Option<u32> {
-        FakeChatTokenizer.token_to_id(token)
-    }
-}
-
-impl TextBackend for FailingDecodeChatBackend {
-    fn tokenizer(&self) -> DynTokenizer {
-        Arc::new(FailingDecodeTokenizer)
-    }
-
-    fn model_id(&self) -> &str {
-        "test-model"
-    }
-}
-
-impl ChatBackend for FailingDecodeChatBackend {
-    fn chat_renderer(&self) -> DynChatRenderer {
-        Arc::new(self.clone())
-    }
-
-    fn new_chat_output_processor(
-        &self,
-        _request: &mut ChatRequest,
-        _options: NewChatOutputProcessorOptions<'_>,
-    ) -> vllm_chat::Result<DynChatOutputProcessor> {
-        Ok(Box::new(DefaultChatOutputProcessor::plain_text_only()))
-    }
-}
-
-impl ChatRenderer for FailingDecodeChatBackend {
-    fn render(&self, request: &ChatRequest) -> vllm_chat::Result<vllm_chat::RenderedPrompt> {
-        FakeChatBackend::new().render(request)
-    }
-}
-
 async fn test_models_with_engine_outputs_and_backend_inner(
     engine_id: impl Into<EngineId>,
     output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
@@ -2755,8 +2630,8 @@ async fn load_endpoint_resets_when_stream_response_is_dropped() {
 #[serial]
 async fn stream_error_is_returned_as_openai_error_sse() {
     let (app, engine_task) = test_app_with_backend_and_stream_output_specs(
-        Arc::new(FailingDecodeChatBackend),
-        default_stream_output_specs(),
+        Arc::new(FakeChatBackend::new()),
+        vec![(vec![UNKNOWN_DECODE_TOKEN_ID], None)],
     )
     .await;
     let response = app
@@ -2789,7 +2664,9 @@ async fn stream_error_is_returned_as_openai_error_sse() {
     assert!(text.contains("\"role\":\"assistant\""), "{text}");
     assert!(text.contains("\"type\":\"server_error\""), "{text}");
     assert!(
-        text.contains("forced decode failure for streaming test"),
+        text.contains(&format!(
+            "test tokenizer cannot decode unknown token id {UNKNOWN_DECODE_TOKEN_ID}"
+        )),
         "{text}"
     );
     assert!(!text.contains("\"usage\":"), "{text}");

rust/src/server/src/tls_tests.rs

@@ -1,5 +1,5 @@
 //! TLS tests: `build_server_config` unit checks plus end-to-end OpenSSL handshakes
-//! through the production `serve_listener` path, with a trivial router since TLS
+//! through the production listener/connection path, with a trivial router since TLS
 //! terminates below the app.
 
 use std::pin::Pin;
@@ -23,8 +23,8 @@ use tokio_openssl::SslStream;
 use tokio_util::sync::CancellationToken;
 
 use crate::config::{HttpListenerMode, TlsConfig};
-use crate::listener::Listener;
-use crate::{ConnectionTimeouts, serve_listener, tls};
+use crate::listener::{Listener, MaybeTlsListener};
+use crate::{ConnectionTimeouts, serve_connections, tls};
 
 // ============================================================================
 // Test infrastructure
@@ -251,7 +251,6 @@ fn build_tls(certs: &TestCerts, cert: &str, key: Option<&str>) -> TlsConfig {
 
 /// Generous per-connection timeouts that never fire during the fast tests.
 const TEST_TIMEOUTS: ConnectionTimeouts = ConnectionTimeouts {
-    handshake: Duration::from_secs(60),
     header_read: Duration::from_secs(5),
     keep_alive_enabled: true,
 };
@@ -261,9 +260,8 @@ async fn spawn_server(tls_config: Option<TlsConfig>) -> (String, CancellationTok
 }
 
 /// Bind an ephemeral listener and serve a trivial router via the production
-/// `serve_listener`, optionally with TLS. The listener is bound (and thus
-/// accepting into the backlog) before returning, so a client may connect
-/// immediately without a sleep.
+/// listener/connection path. The listener is bound (and thus accepting into the
+/// backlog) before returning, so a client may connect immediately without a sleep.
 async fn spawn_server_with_timeouts(
     tls_config: Option<TlsConfig>,
     timeouts: ConnectionTimeouts,
@@ -274,7 +272,7 @@ async fn spawn_server_with_timeouts(
     })
     .await
     .expect("bind listener");
-    let addr = listener.local_addr().expect("local addr");
+    let addr = listener.local_addr_display().expect("local addr");
 
     let server_config =
         tls_config.map(|cfg| tls::build_server_config(&cfg).expect("build server config"));
@@ -282,14 +280,11 @@ async fn spawn_server_with_timeouts(
     let shutdown = CancellationToken::new();
     let server_shutdown = shutdown.clone();
     tokio::spawn(async move {
-        let _ = serve_listener(
-            listener,
-            server_config,
-            app,
-            server_shutdown.cancelled_owned(),
-            timeouts,
-        )
-        .await;
+        let listener = match server_config {
+            Some(context) => MaybeTlsListener::tls(listener, context),
+            None => MaybeTlsListener::plain(listener),
+        };
+        let _ = serve_connections(listener, app, server_shutdown.cancelled_owned(), timeouts).await;
     });
     (addr, shutdown)
 }
@@ -521,20 +516,19 @@ async fn plain_http_serves_when_tls_is_disabled() {
     shutdown.cancel();
 }
 
-#[tokio::test]
+#[tokio::test(start_paused = true)]
 async fn tls_handshake_timeout_drops_silent_client() {
     // Silent client (no ClientHello) must be dropped at the handshake deadline.
     let certs = TestCerts::generate();
-    let timeouts = ConnectionTimeouts {
-        handshake: Duration::from_millis(150),
-        header_read: Duration::from_secs(5),
-        keep_alive_enabled: true,
-    };
-    let (addr, shutdown) = spawn_server_with_timeouts(Some(server_tls(&certs, 0)), timeouts).await;
+    let (addr, shutdown) = spawn_server(Some(server_tls(&certs, 0))).await;
 
     let mut tcp = TcpStream::connect(&addr).await.expect("connect");
+    tokio::task::yield_now().await;
+    tokio::time::advance(tls::TLS_HANDSHAKE_TIMEOUT + Duration::from_millis(1)).await;
+    tokio::task::yield_now().await;
+
     let mut buf = [0u8; 1];
-    let read = tokio::time::timeout(Duration::from_secs(5), tcp.read(&mut buf)).await;
+    let read = tokio::time::timeout(Duration::from_secs(1), tcp.read(&mut buf)).await;
     assert!(
         matches!(read, Ok(Ok(0)) | Ok(Err(_))),
         "server must drop a stalled TLS handshake (expected close, got {read:?})"
@@ -546,7 +540,6 @@ async fn tls_handshake_timeout_drops_silent_client() {
 async fn keep_alive_timeout_closes_idle_connection() {
     // Idle keep-alive connection must be closed at the deadline.
     let timeouts = ConnectionTimeouts {
-        handshake: Duration::from_secs(60),
         header_read: Duration::from_millis(150),
         keep_alive_enabled: true,
     };
@@ -582,7 +575,6 @@ async fn keep_alive_timeout_closes_idle_tls_connection() {
     // still fires through tls-listener's post-handshake SslStream, not just plaintext.
     let certs = TestCerts::generate();
     let timeouts = ConnectionTimeouts {
-        handshake: Duration::from_secs(60),
         header_read: Duration::from_millis(150),
         keep_alive_enabled: true,
     };
@@ -618,7 +610,6 @@ async fn keep_alive_timeout_closes_idle_tls_connection() {
 async fn idle_timeout_closes_silent_client() {
     // Silent client closed by the header-read timeout (http1-only arms it from byte 0).
     let timeouts = ConnectionTimeouts {
-        handshake: Duration::from_secs(60),
         header_read: Duration::from_millis(150),
         keep_alive_enabled: true,
     };
@@ -638,7 +629,6 @@ async fn idle_timeout_closes_silent_client() {
 async fn keep_alive_zero_disables_keep_alive() {
     // 0 disables keep-alive (serve, then close), like uvicorn's timeout_keep_alive=0.
     let timeouts = ConnectionTimeouts {
-        handshake: Duration::from_secs(60),
         header_read: Duration::from_secs(5),
         keep_alive_enabled: false,
     };
@@ -671,7 +661,6 @@ async fn disabled_keep_alive_still_closes_silent_client() {
     // Even with keep-alive off, the head read stays bounded, so a silent client
     // is dropped rather than held open.
     let timeouts = ConnectionTimeouts {
-        handshake: Duration::from_secs(60),
         header_read: Duration::from_millis(150),
         keep_alive_enabled: false,
     };

rust/src/text/Cargo.toml

@@ -31,6 +31,7 @@ serial_test.workspace = true
 tempfile.workspace = true
 tokio.workspace = true
 vllm-llm = { workspace = true, features = ["test-util"] }
+vllm-tokenizer = { workspace = true, features = ["test-utils"] }
 
 [lints]
 workspace = true

rust/src/text/src/lower.rs

@@ -275,6 +275,7 @@ mod tests {
     use std::collections::{BTreeSet, HashMap};
 
     use serial_test::file_serial;
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::*;
     use crate::backend::hf::HfTextBackend;
@@ -282,60 +283,8 @@ mod tests {
     use crate::error::{LogprobsError, TokenIdsError};
     use crate::request::{Prompt, TextRequest};
 
-    /// Stub tokenizer that returns empty token IDs — sufficient for tests that
-    /// don't exercise bad-words tokenization.
-    struct StubTokenizer;
-
-    impl Tokenizer for StubTokenizer {
-        fn encode(
-            &self,
-            _text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            Ok(vec![])
-        }
-
-        fn decode(
-            &self,
-            _token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            Ok(String::new())
-        }
-
-        fn token_to_id(&self, _token: &str) -> Option<u32> {
-            None
-        }
-    }
-
-    fn stub_tokenizer() -> StubTokenizer {
-        StubTokenizer
-    }
-
-    struct FixedTokenizer {
-        token_ids: Vec<u32>,
-    }
-
-    impl Tokenizer for FixedTokenizer {
-        fn encode(
-            &self,
-            _text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            Ok(self.token_ids.clone())
-        }
-
-        fn decode(
-            &self,
-            _token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            Ok(String::new())
-        }
-
-        fn token_to_id(&self, _token: &str) -> Option<u32> {
-            None
-        }
+    fn stub_tokenizer() -> TestTokenizer {
+        TestTokenizer::new()
     }
 
     fn sample_request() -> TextRequest {
@@ -952,9 +901,7 @@ mod tests {
 
     #[test]
     fn lower_sampling_params_rejects_out_of_vocab_bad_words() {
-        let tokenizer = FixedTokenizer {
-            token_ids: vec![1999, 2000],
-        };
+        let tokenizer = TestTokenizer::new().with_regular_token("blocked", 2000);
         let error = lower_sampling_params(
             SamplingParams {
                 bad_words: Some(vec!["blocked".to_string()]),

rust/src/text/src/output/decoded.rs

@@ -323,37 +323,11 @@ mod tests {
     use futures::{Stream, stream};
     use vllm_engine_core_client::AbortCause;
     use vllm_llm::GenerateOutput;
-    use vllm_tokenizer::Tokenizer;
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::*;
     use crate::output::TextOutputStreamExt as _;
 
-    /// Backend that treats each token ID as a raw byte, producing lossy UTF-8.
-    struct ByteTokenizer;
-
-    impl Tokenizer for ByteTokenizer {
-        fn encode(
-            &self,
-            _text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            unreachable!()
-        }
-
-        fn decode(
-            &self,
-            token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            let bytes = token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>();
-            Ok(String::from_utf8_lossy(&bytes).into_owned())
-        }
-
-        fn token_to_id(&self, _token: &str) -> Option<u32> {
-            unreachable!()
-        }
-    }
-
     /// Helper: run `decoded_text_event_stream` to completion and return the
     /// collected output.
     async fn run_to_completion(
@@ -366,7 +340,7 @@ mod tests {
             token_ids,
             Some(FinishReason::Length),
         ))]);
-        let tokenizer: DynTokenizer = Arc::new(ByteTokenizer);
+        let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
         decoded_text_event_stream("test".into(), tokenizer, raw_stream, decode_options, false)
             .collect_output()
             .await
@@ -419,7 +393,7 @@ mod tests {
             ))),
             dropped_cause: Arc::clone(&dropped_cause),
         };
-        let tokenizer: DynTokenizer = Arc::new(ByteTokenizer);
+        let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
 
         let output = decoded_text_event_stream(
             "test".into(),

rust/src/text/src/output/logprobs.rs

@@ -129,40 +129,13 @@ fn decode_position_logprobs<T: Tokenizer + ?Sized>(
 #[cfg(test)]
 mod tests {
     use vllm_llm::{Logprobs, PositionLogprobs, TokenLogprob};
+    use vllm_tokenizer::test_utils::TestTokenizer;
 
     use super::*;
 
-    #[derive(Debug)]
-    struct ByteTokenizer;
-
-    impl vllm_tokenizer::Tokenizer for ByteTokenizer {
-        fn encode(
-            &self,
-            _text: &str,
-            _add_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<Vec<u32>> {
-            unreachable!()
-        }
-
-        fn decode(
-            &self,
-            token_ids: &[u32],
-            _skip_special_tokens: bool,
-        ) -> vllm_tokenizer::Result<String> {
-            Ok(String::from_utf8_lossy(
-                &token_ids.iter().map(|token_id| *token_id as u8).collect::<Vec<_>>(),
-            )
-            .into_owned())
-        }
-
-        fn token_to_id(&self, _token: &str) -> Option<u32> {
-            unreachable!()
-        }
-    }
-
     #[test]
     fn decode_logprobs_decodes_every_candidate_token() {
-        let tokenizer = ByteTokenizer;
+        let tokenizer = TestTokenizer::new();
         let logprobs = Logprobs {
             positions: vec![PositionLogprobs {
                 entries: vec![
@@ -205,7 +178,7 @@ mod tests {
 
     #[test]
     fn decode_prompt_logprobs_separates_first_prompt_token() {
-        let tokenizer = ByteTokenizer;
+        let tokenizer = TestTokenizer::new();
         let logprobs = Logprobs {
             positions: vec![PositionLogprobs {
                 entries: vec![TokenLogprob {

rust/src/tokenizer/Cargo.toml

@@ -4,6 +4,9 @@ version.workspace = true
 edition.workspace = true
 license.workspace = true
 
+[features]
+test-utils = []
+
 [dependencies]
 base64.workspace = true
 fastokens.workspace = true

rust/src/tokenizer/src/incremental.rs

@@ -170,6 +170,10 @@ mod tests {
         fn token_to_id(&self, _token: &str) -> Option<u32> {
             unreachable!()
         }
+
+        fn id_to_token(&self, _id: u32) -> Option<String> {
+            unreachable!()
+        }
     }
 
     #[test]
@@ -248,6 +252,10 @@ mod tests {
         fn token_to_id(&self, _token: &str) -> Option<u32> {
             unreachable!()
         }
+
+        fn id_to_token(&self, _id: u32) -> Option<String> {
+            unreachable!()
+        }
     }
 
     #[test]
@@ -320,6 +328,10 @@ mod tests {
         fn token_to_id(&self, _token: &str) -> Option<u32> {
             unreachable!()
         }
+
+        fn id_to_token(&self, _id: u32) -> Option<String> {
+            unreachable!()
+        }
     }
 
     /// Without the char-boundary fix, this panics slicing mid-emoji.

rust/src/tokenizer/src/lib.rs

@@ -8,6 +8,8 @@ mod error;
 mod hf;
 mod incremental;
 mod tekken;
+#[cfg(any(test, feature = "test-utils"))]
+pub mod test_utils;
 mod tiktoken;
 
 pub use error::{Result, TokenizerError};
@@ -28,11 +30,7 @@ pub trait Tokenizer: Send + Sync {
     fn token_to_id(&self, token: &str) -> Option<u32>;
 
     /// Convert one token ID into the tokenizer's raw token string.
-    fn id_to_token(&self, _id: u32) -> Option<String> {
-        // TODO: remove default impl and require this to be implemented by all
-        // tokenizers
-        None
-    }
+    fn id_to_token(&self, id: u32) -> Option<String>;
 
     /// Return the vocabulary size. Backends that cannot report it fall back to
     /// `usize::MAX`, an effectively unbounded value used only by test stubs.

rust/src/tokenizer/src/test_utils.rs

@@ -0,0 +1,434 @@
+use std::collections::BTreeMap;
+
+use crate::{Result, Tokenizer, TokenizerError};
+
+const FIRST_CONFIGURED_TOKEN_ID: u32 = 256;
+
+/// Whether a configured test token should be treated as special.
+///
+/// Special tokens are skipped by [`Tokenizer::decode`] when
+/// `skip_special_tokens` is set. Regular configured tokens are always emitted.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum TestTokenKind {
+    /// Token is skipped when `skip_special_tokens = true`.
+    Special,
+    /// Token is emitted regardless of `skip_special_tokens`.
+    Regular,
+}
+
+impl TestTokenKind {
+    fn is_special(self) -> bool {
+        matches!(self, Self::Special)
+    }
+}
+
+/// Decode behavior for token ids that are neither configured tokens nor byte ids.
+///
+/// The default is [`UnknownDecode::Error`] so tests notice missing tokenizer
+/// fixtures instead of silently accepting impossible ids. Individual tests can
+/// opt into empty or replacement output when they are explicitly modeling a
+/// lenient detokenization path.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum UnknownDecode {
+    /// Return a tokenizer error on the first unknown id.
+    Error,
+    /// Drop unknown ids from decoded output.
+    Empty,
+    /// Emit U+FFFD for each unknown id.
+    Replacement,
+}
+
+#[derive(Debug, Clone)]
+struct TestToken {
+    text: String,
+    kind: TestTokenKind,
+}
+
+/// Configurable tokenizer for Rust frontend tests.
+///
+/// `TestTokenizer` is intentionally small, but its methods obey the same basic
+/// contract as production tokenizers:
+///
+/// - ordinary text encodes as UTF-8 byte ids;
+/// - configured token ids start at 256, leaving `0..=255` for byte fallback;
+/// - configured token ids and token text are unique;
+/// - configured tokens are matched before ordinary bytes, using longest-prefix matching so
+///   multi-character markers such as `<think>` work naturally;
+/// - `token_to_id` and `id_to_token` are consistent for configured tokens;
+/// - `decode` is strict by default for ids outside the byte range and the configured token table;
+/// - `vocab_size` is an exclusive upper bound covering byte ids and configured token ids unless a
+///   test sets it explicitly.
+///
+/// Prefer this helper over ad-hoc fake tokenizers for tests that rely on
+/// tokenizer semantics. Keep dedicated tiny fakes for error injection or for
+/// tests that deliberately need a degenerate tokenizer.
+#[derive(Debug, Clone)]
+pub struct TestTokenizer {
+    token_to_id: BTreeMap<String, u32>,
+    id_to_token: BTreeMap<u32, TestToken>,
+    unknown_decode: UnknownDecode,
+    vocab_size: Option<usize>,
+    bos_token_id: Option<u32>,
+}
+
+impl Default for TestTokenizer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl TestTokenizer {
+    /// Create a byte-level test tokenizer with strict unknown-id decode.
+    pub fn new() -> Self {
+        Self {
+            token_to_id: BTreeMap::new(),
+            id_to_token: BTreeMap::new(),
+            unknown_decode: UnknownDecode::Error,
+            vocab_size: None,
+            bos_token_id: None,
+        }
+    }
+
+    /// Add a configured token and return the updated tokenizer.
+    ///
+    /// Configured tokens must use ids outside the byte range and may be marked
+    /// special or regular.
+    pub fn with_token(mut self, token: impl Into<String>, id: u32, kind: TestTokenKind) -> Self {
+        self.insert_token(token, id, kind);
+        self
+    }
+
+    /// Add a special configured token and return the updated tokenizer.
+    pub fn with_special_token(self, token: impl Into<String>, id: u32) -> Self {
+        self.with_token(token, id, TestTokenKind::Special)
+    }
+
+    /// Add a regular configured token and return the updated tokenizer.
+    pub fn with_regular_token(self, token: impl Into<String>, id: u32) -> Self {
+        self.with_token(token, id, TestTokenKind::Regular)
+    }
+
+    /// Add a special BOS token inserted by `encode(..., true)`.
+    ///
+    /// This also registers the token in the normal token/id maps so
+    /// `token_to_id`, `id_to_token`, `decode`, and `is_special_id` stay
+    /// consistent for the inserted id.
+    pub fn with_bos_token(mut self, token: impl Into<String>, id: u32) -> Self {
+        self.insert_token(token, id, TestTokenKind::Special);
+        self.bos_token_id = Some(id);
+        self
+    }
+
+    /// Set decode behavior for unknown non-byte ids.
+    pub fn with_unknown_decode(mut self, behavior: UnknownDecode) -> Self {
+        self.unknown_decode = behavior;
+        self
+    }
+
+    /// Set an explicit vocabulary size.
+    ///
+    /// Use this when a test needs a model-like vocabulary bound that differs
+    /// from the highest configured token id plus one.
+    pub fn with_vocab_size(mut self, vocab_size: usize) -> Self {
+        self.vocab_size = Some(vocab_size);
+        self
+    }
+
+    fn insert_token(&mut self, token: impl Into<String>, id: u32, kind: TestTokenKind) {
+        let token = token.into();
+        assert!(
+            !token.is_empty(),
+            "configured test token text must be non-empty"
+        );
+        assert!(
+            id >= FIRST_CONFIGURED_TOKEN_ID,
+            "configured test token id {id} overlaps byte fallback range 0..=255"
+        );
+        assert!(
+            token.len() > 1,
+            "configured test token text {token:?} overlaps byte fallback token text"
+        );
+        if self.token_to_id.insert(token.clone(), id).is_some() {
+            panic!("configured test token text {token:?} was registered more than once");
+        }
+        if self.id_to_token.insert(id, TestToken { text: token, kind }).is_some() {
+            panic!("configured test token id {id} was registered more than once");
+        }
+    }
+
+    fn byte_to_token(id: u32) -> Option<String> {
+        u8::try_from(id).ok().map(|byte| String::from_utf8_lossy(&[byte]).into_owned())
+    }
+
+    fn flush_bytes(bytes: &mut Vec<u8>, output: &mut String) {
+        if !bytes.is_empty() {
+            output.push_str(&String::from_utf8_lossy(bytes));
+            bytes.clear();
+        }
+    }
+
+    fn configured_token_prefix(&self, text: &str) -> Option<(&str, u32)> {
+        self.token_to_id
+            .iter()
+            .filter_map(|(token, &id)| text.starts_with(token).then_some((token.as_str(), id)))
+            .max_by_key(|(token, _)| token.len())
+    }
+
+    fn inferred_vocab_size(&self) -> usize {
+        let max_configured =
+            self.id_to_token.last_key_value().map(|(&id, _)| id as usize + 1).unwrap_or(0);
+        256.max(max_configured)
+    }
+}
+
+impl Tokenizer for TestTokenizer {
+    fn encode(&self, text: &str, add_special_tokens: bool) -> Result<Vec<u32>> {
+        let mut ids = Vec::new();
+        if add_special_tokens && let Some(bos_token_id) = self.bos_token_id {
+            ids.push(bos_token_id);
+        }
+
+        let mut rest = text;
+        while !rest.is_empty() {
+            if let Some((token, id)) = self.configured_token_prefix(rest) {
+                ids.push(id);
+                rest = &rest[token.len()..];
+                continue;
+            }
+
+            let ch = rest.chars().next().expect("rest is not empty");
+            let mut buf = [0_u8; 4];
+            ids.extend(ch.encode_utf8(&mut buf).bytes().map(u32::from));
+            rest = &rest[ch.len_utf8()..];
+        }
+
+        Ok(ids)
+    }
+
+    fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
+        let mut output = String::new();
+        let mut pending_bytes = Vec::new();
+        for &id in token_ids {
+            if let Some(token) = self.id_to_token.get(&id) {
+                Self::flush_bytes(&mut pending_bytes, &mut output);
+                if !(skip_special_tokens && token.kind.is_special()) {
+                    output.push_str(&token.text);
+                }
+            } else if let Ok(byte) = u8::try_from(id) {
+                pending_bytes.push(byte);
+            } else {
+                Self::flush_bytes(&mut pending_bytes, &mut output);
+                match self.unknown_decode {
+                    UnknownDecode::Error => {
+                        return Err(TokenizerError(format!(
+                            "test tokenizer cannot decode unknown token id {id}"
+                        )));
+                    }
+                    UnknownDecode::Empty => {}
+                    UnknownDecode::Replacement => output.push('\u{FFFD}'),
+                }
+            }
+        }
+        Self::flush_bytes(&mut pending_bytes, &mut output);
+        Ok(output)
+    }
+
+    fn token_to_id(&self, token: &str) -> Option<u32> {
+        self.token_to_id.get(token).copied().or_else(|| {
+            let bytes = token.as_bytes();
+            (bytes.len() == 1).then(|| u32::from(bytes[0]))
+        })
+    }
+
+    fn id_to_token(&self, id: u32) -> Option<String> {
+        self.id_to_token
+            .get(&id)
+            .map(|token| token.text.clone())
+            .or_else(|| Self::byte_to_token(id))
+    }
+
+    fn vocab_size(&self) -> usize {
+        self.vocab_size.unwrap_or_else(|| self.inferred_vocab_size())
+    }
+
+    fn is_special_id(&self, token_id: u32) -> bool {
+        self.id_to_token.get(&token_id).is_some_and(|token| token.kind.is_special())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn byte_text_roundtrips_and_reports_byte_ids() {
+        let tokenizer = TestTokenizer::new();
+
+        let ids = tokenizer.encode("hi", false).unwrap();
+        assert_eq!(ids, vec![b'h' as u32, b'i' as u32]);
+        assert_eq!(tokenizer.decode(&ids, false).unwrap(), "hi");
+        assert_eq!(tokenizer.token_to_id("h"), Some(b'h' as u32));
+        assert_eq!(tokenizer.id_to_token(b'h' as u32).as_deref(), Some("h"));
+        assert_eq!(tokenizer.vocab_size(), 256);
+    }
+
+    #[test]
+    fn configured_tokens_use_longest_prefix_matching() {
+        let tokenizer = TestTokenizer::new()
+            .with_regular_token("<image>", 999)
+            .with_regular_token("<image></image>", 1000);
+
+        assert_eq!(
+            tokenizer.encode("a<image></image>b", false).unwrap(),
+            vec![b'a' as u32, 1000, b'b' as u32,]
+        );
+        assert_eq!(
+            tokenizer.decode(&[b'a' as u32, 1000, b'b' as u32], false).unwrap(),
+            "a<image></image>b"
+        );
+        assert_eq!(tokenizer.token_to_id("<image>"), Some(999));
+        assert_eq!(
+            tokenizer.id_to_token(1000).as_deref(),
+            Some("<image></image>")
+        );
+        assert_eq!(tokenizer.vocab_size(), 1001);
+    }
+
+    #[test]
+    fn non_ascii_text_roundtrips_through_buffered_byte_decode() {
+        let tokenizer = TestTokenizer::new();
+        let text = "你好, café, 🚀";
+
+        let ids = tokenizer.encode(text, false).unwrap();
+        assert_eq!(
+            ids,
+            text.as_bytes().iter().copied().map(u32::from).collect::<Vec<_>>()
+        );
+        assert_eq!(tokenizer.decode(&ids, false).unwrap(), text);
+    }
+
+    #[test]
+    fn buffered_byte_decode_flushes_around_configured_tokens() {
+        let tokenizer = TestTokenizer::new()
+            .with_regular_token("<image>", 999)
+            .with_special_token("<skip>", 1000);
+
+        assert_eq!(
+            tokenizer.encode("你<image>好<skip>🚀", false).unwrap(),
+            vec![228, 189, 160, 999, 229, 165, 189, 1000, 240, 159, 154, 128]
+        );
+        assert_eq!(
+            tokenizer
+                .decode(
+                    &[228, 189, 160, 999, 229, 165, 189, 1000, 240, 159, 154, 128],
+                    false
+                )
+                .unwrap(),
+            "你<image>好<skip>🚀"
+        );
+        assert_eq!(
+            tokenizer
+                .decode(
+                    &[228, 189, 160, 999, 229, 165, 189, 1000, 240, 159, 154, 128],
+                    true
+                )
+                .unwrap(),
+            "你<image>好🚀"
+        );
+    }
+
+    #[test]
+    fn invalid_utf8_bytes_decode_lossily_as_a_sequence() {
+        let tokenizer = TestTokenizer::new();
+
+        assert_eq!(tokenizer.decode(&[0xE4, 0xBD], false).unwrap(), "\u{FFFD}");
+        assert_eq!(
+            tokenizer.decode(&[0xFF, b'a' as u32], false).unwrap(),
+            "\u{FFFD}a"
+        );
+    }
+
+    #[test]
+    fn special_tokens_respect_skip_special_tokens() {
+        let tokenizer = TestTokenizer::new()
+            .with_bos_token("<bos>", 256)
+            .with_special_token("<think>", 0xF001)
+            .with_regular_token("</think>", 0xF002);
+
+        assert_eq!(
+            tokenizer.encode("<think>x</think>", true).unwrap(),
+            vec![256, 0xF001, b'x' as u32, 0xF002,]
+        );
+        assert_eq!(
+            tokenizer.decode(&[256, 0xF001, b'x' as u32, 0xF002], false).unwrap(),
+            "<bos><think>x</think>"
+        );
+        assert_eq!(
+            tokenizer.decode(&[256, 0xF001, b'x' as u32, 0xF002], true).unwrap(),
+            "x</think>"
+        );
+        assert!(tokenizer.is_special_id(0xF001));
+        assert!(!tokenizer.is_special_id(0xF002));
+    }
+
+    #[test]
+    #[should_panic(expected = "configured test token id 255 overlaps byte fallback range 0..=255")]
+    fn configured_token_id_must_stay_outside_byte_range() {
+        let _ = TestTokenizer::new().with_regular_token("<token>", 255);
+    }
+
+    #[test]
+    #[should_panic(expected = "configured test token text \"a\" overlaps byte fallback token text")]
+    fn configured_token_text_must_not_shadow_byte_tokens() {
+        let _ = TestTokenizer::new().with_regular_token("a", 256);
+    }
+
+    #[test]
+    #[should_panic(
+        expected = "configured test token text \"<token>\" was registered more than once"
+    )]
+    fn configured_token_text_must_be_unique() {
+        let _ = TestTokenizer::new()
+            .with_regular_token("<token>", 256)
+            .with_regular_token("<token>", 257);
+    }
+
+    #[test]
+    #[should_panic(expected = "configured test token id 256 was registered more than once")]
+    fn configured_token_id_must_be_unique() {
+        let _ = TestTokenizer::new()
+            .with_regular_token("<token-a>", 256)
+            .with_regular_token("<token-b>", 256);
+    }
+
+    #[test]
+    fn unknown_decode_is_strict_by_default_and_configurable() {
+        let strict = TestTokenizer::new();
+        assert!(strict.decode(&[300], false).is_err());
+        assert_eq!(
+            TestTokenizer::new()
+                .with_unknown_decode(UnknownDecode::Empty)
+                .decode(&[b'a' as u32, 300, b'b' as u32], false)
+                .unwrap(),
+            "ab"
+        );
+        assert_eq!(
+            TestTokenizer::new()
+                .with_unknown_decode(UnknownDecode::Replacement)
+                .decode(&[300], false)
+                .unwrap(),
+            "\u{FFFD}"
+        );
+        assert_eq!(strict.id_to_token(300), None);
+    }
+
+    #[test]
+    fn explicit_vocab_size_overrides_inferred_bound() {
+        let tokenizer = TestTokenizer::new()
+            .with_regular_token("<high>", 10_000)
+            .with_vocab_size(20_000);
+
+        assert_eq!(tokenizer.vocab_size(), 20_000);
+        assert_eq!(tokenizer.id_to_token(10_000).as_deref(), Some("<high>"));
+    }
+}

tests/distributed/test_ray_v2_executor.py

@@ -17,6 +17,7 @@ import ray
 from vllm import LLM
 from vllm.config import VllmConfig
 from vllm.engine.arg_utils import EngineArgs
+from vllm.v1.executor import ray_executor_v2
 from vllm.v1.executor.ray_executor_v2 import RayExecutorV2
 
 pytestmark = pytest.mark.usefixtures("enable_ray_v2_backend")
@@ -95,6 +96,43 @@ def assert_executor(executor, tp_size, pp_size):
         assert handle.node_id is not None
 
 
+def test_select_tcpstore_port_seeds_disjoint_windows(monkeypatch):
+    """Co-located DP engines scan distinct, adjacent port windows, so two
+    engines on a node cannot pick the same TCPStore port."""
+    requested = []
+
+    def fake_get_open_port(start_port, max_attempts):
+        requested.append((start_port, max_attempts))
+        return start_port
+
+    monkeypatch.setattr(ray_executor_v2, "_get_open_port", fake_get_open_port)
+
+    ports = [
+        RayExecutorV2._select_tcpstore_port(rank, master_port=29500)
+        for rank in range(4)
+    ]
+
+    assert requested == [(29600, 32), (29632, 32), (29664, 32), (29696, 32)]
+    assert len(set(ports)) == 4
+
+
+def test_select_tcpstore_port_non_dp_uses_random(monkeypatch):
+    """A non-DP engine has no local rank and uses a random port."""
+    monkeypatch.setattr(ray_executor_v2, "get_open_port", lambda: 54321)
+    assert RayExecutorV2._select_tcpstore_port(None, master_port=29500) == 54321
+
+
+def test_select_tcpstore_port_full_window_uses_random(monkeypatch):
+    """A fully occupied window falls back to a random port."""
+
+    def raise_full(start_port, max_attempts):
+        raise RuntimeError("no open port")
+
+    monkeypatch.setattr(ray_executor_v2, "_get_open_port", raise_full)
+    monkeypatch.setattr(ray_executor_v2, "get_open_port", lambda: 54321)
+    assert RayExecutorV2._select_tcpstore_port(0, master_port=29500) == 54321
+
+
 @pytest.mark.parametrize("tp_size, pp_size", [(1, 1), (2, 1), (4, 1), (2, 2)])
 def test_ray_v2_executor(tp_size, pp_size):
     """Validate RayExecutorV2 with various TP/PP configs."""

tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression test for AITER MLA persistent decode metadata dtypes.
+
+For the gfx950 fp8/fp8 nhead=32 qlen=1 fold path, the split/reduce metadata
+layout depends on the q/kv element size. The builder must forward dtype_q/dtype_kv
+to ``get_mla_metadata_v1``; omitting them lays out the work for the wrong dtype
+and corrupts decode output. The test pins the builder's metadata to a golden
+recomputed at runtime with the explicit correct dtypes.
+"""
+
+import types
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm._aiter_ops import is_aiter_found
+from vllm.platforms import current_platform
+
+
+def _on_gfx950() -> bool:
+    if not (current_platform.is_rocm() and is_aiter_found()):
+        return False
+    from vllm.platforms.rocm import on_gfx950
+
+    return on_gfx950()
+
+
+pytestmark = pytest.mark.skipif(
+    not _on_gfx950(),
+    reason="AITER MLA fp8 persistent decode metadata is gfx950-only",
+)
+
+# The fold path that the bug corrupted: fp8 query + fp8 KV-cache, 32 query
+# heads, single-token decode, batch 128, context 8192, page_size 1.
+NUM_QUERY_HEADS = 32
+DECODE_QLEN = 1
+BATCH_SIZE = 128
+CONTEXT_LEN = 8192
+PAGE_SIZE = 1
+
+# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8
+# KV-cache -> fp8_e4m3 kv.
+EXPECTED_Q_DTYPE = torch.bfloat16
+EXPECTED_KV_DTYPE = torch.float8_e4m3fn
+
+# The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
+# is excluded: it holds raw device pointers, never equal across allocations.
+_CONTENT_METADATA_FIELDS = (
+    "work_indptr",
+    "work_info_set",
+    "reduce_indptr",
+    "reduce_final_map",
+    "reduce_partial_map",
+)
+
+# The builder's get_mla_metadata_v1 call passes 6 input args then 6 output
+# buffers (see AiterMLAMetadataBuilder._build_decode). Output order -> field.
+_NUM_INPUT_ARGS = 6
+_OUTPUT_ARG_FIELDS = (
+    "work_meta_data",
+    "work_info_set",
+    "work_indptr",
+    "reduce_indptr",
+    "reduce_final_map",
+    "reduce_partial_map",
+)
+
+
+def _build_decode_metadata():
+    """Build AITER MLA decode metadata for the fp8/fp8 nhead=32 fold path.
+
+    Returns ``(metadata, captured)`` where ``captured`` records the positional
+    args/kwargs the builder passed to ``get_mla_metadata_v1``, so the golden can
+    be recomputed from the identical inputs.
+    """
+    from tests.v1.attention.utils import (
+        BatchSpec,
+        create_common_attn_metadata,
+        create_vllm_config,
+    )
+    from vllm.config.vllm import set_current_vllm_config
+    from vllm.v1.attention.backends.registry import AttentionBackendEnum
+    from vllm.v1.kv_cache_interface import MLAAttentionSpec
+    from vllm.v1.worker.workspace import init_workspace_manager
+
+    device = torch.device("cuda:0")
+
+    vllm_config = create_vllm_config(
+        model_name="deepseek-ai/DeepSeek-R1",
+        max_model_len=CONTEXT_LEN,
+        # One flat page per token (page_size=1); +buffer for the null block.
+        num_gpu_blocks=BATCH_SIZE * CONTEXT_LEN + 200,
+        block_size=PAGE_SIZE,
+        max_num_seqs=BATCH_SIZE,
+        max_num_batched_tokens=8192,
+        hf_config_override={"num_attention_heads": NUM_QUERY_HEADS},
+    )
+    vllm_config.cache_config.cache_dtype = "fp8"
+
+    spec = MLAAttentionSpec(
+        block_size=PAGE_SIZE,
+        num_kv_heads=1,
+        head_size=vllm_config.model_config.get_head_size(),
+        dtype=vllm_config.model_config.dtype,
+        cache_dtype_str="fp8",
+    )
+
+    builder_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class().get_builder_cls()
+
+    # The builder reads layer.prefill_backend from static_forward_context; a
+    # stub with the attribute is enough for metadata construction.
+    layer_name = "placeholder"
+    vllm_config.compilation_config.static_forward_context[layer_name] = (
+        types.SimpleNamespace(prefill_backend=None)
+    )
+
+    init_workspace_manager(device)
+
+    batch_spec = BatchSpec(
+        seq_lens=[CONTEXT_LEN] * BATCH_SIZE,
+        query_lens=[DECODE_QLEN] * BATCH_SIZE,
+    )
+
+    captured: dict = {}
+
+    with set_current_vllm_config(vllm_config):
+        builder = builder_cls(spec, [layer_name], vllm_config, device)
+        common_attn_metadata = create_common_attn_metadata(
+            batch_spec, PAGE_SIZE, device, arange_block_indices=True
+        )
+
+        import aiter
+
+        real_get_mla_metadata_v1 = aiter.get_mla_metadata_v1
+
+        def spy(*args, **kwargs):
+            captured["args"] = args
+            captured["kwargs"] = dict(kwargs)
+            return real_get_mla_metadata_v1(*args, **kwargs)
+
+        with patch("aiter.get_mla_metadata_v1", spy):
+            metadata = builder.build(
+                common_prefix_len=0,
+                common_attn_metadata=common_attn_metadata,
+            )
+
+    return metadata, captured
+
+
+def _compute_golden_metadata(captured: dict) -> dict[str, torch.Tensor]:
+    """Recompute the persistent metadata with explicit fp8/bf16 dtypes.
+
+    Replays ``get_mla_metadata_v1`` on the builder's exact input tensors with
+    fresh output buffers and the explicitly-correct dtypes. This reference must
+    match the builder's output when the fix is in place.
+    """
+    import aiter
+
+    args = captured["args"]
+    inputs = args[:_NUM_INPUT_ARGS]
+    # Fresh copies so the golden does not alias the builder's persistent buffers.
+    fresh_outputs = [arg.clone() for arg in args[_NUM_INPUT_ARGS:]]
+
+    golden_kwargs = dict(captured["kwargs"])
+    golden_kwargs["dtype_q"] = EXPECTED_Q_DTYPE
+    golden_kwargs["dtype_kv"] = EXPECTED_KV_DTYPE
+
+    aiter.get_mla_metadata_v1(*inputs, *fresh_outputs, **golden_kwargs)
+
+    return dict(zip(_OUTPUT_ARG_FIELDS, fresh_outputs))
+
+
+def test_persistent_decode_metadata_matches_fp8_golden():
+    """The builder's metadata must match the dtype-correct golden.
+
+    Regression guard: the fixed builder forwards fp8/bf16 dtypes so its
+    split/reduce metadata matches the golden recomputed with those explicit
+    dtypes. Dropping the dtypes (the original bug) produces a different layout
+    and fails this test.
+    """
+    metadata, captured = _build_decode_metadata()
+
+    # qlen=1 must take the persistent-metadata path for this to be meaningful.
+    assert metadata.decode is not None
+    assert metadata.decode.has_persistent_metadata
+    assert metadata.work_meta_data is not None
+
+    golden = _compute_golden_metadata(captured)
+
+    mismatched = [
+        name
+        for name in _CONTENT_METADATA_FIELDS
+        if getattr(metadata, name).shape != golden[name].shape
+        or not torch.equal(getattr(metadata, name), golden[name])
+    ]
+    assert not mismatched, (
+        "AITER MLA persistent decode metadata does not match the fp8/bf16 "
+        f"golden for fields {mismatched}; the builder must forward "
+        "dtype_q/dtype_kv to get_mla_metadata_v1."
+    )

tests/models/multimodal/generation/test_common.py

@@ -993,31 +993,6 @@ VLM_TEST_SETTINGS = {
         hf_output_post_proc=model_utils.smolvlm_trunc_hf_output,
         num_logprobs=10,
     ),
-    "tarsier": VLMTestInfo(
-        models=["omni-research/Tarsier-7b"],
-        test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
-        prompt_formatter=lambda img_prompt: f"USER: {img_prompt} ASSISTANT:",
-        max_model_len=4096,
-        max_num_seqs=2,
-        auto_cls=AutoModelForImageTextToText,
-        patch_hf_runner=model_utils.tarsier_patch_hf_runner,
-    ),
-    "tarsier2": VLMTestInfo(
-        models=["omni-research/Tarsier2-Recap-7b"],
-        test_type=(
-            VLMTestType.IMAGE,
-            VLMTestType.MULTI_IMAGE,
-            VLMTestType.VIDEO,
-        ),
-        prompt_formatter=lambda img_prompt: f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n",  # noqa: E501
-        img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>",
-        video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>",
-        max_model_len=4096,
-        max_num_seqs=2,
-        auto_cls=AutoModelForImageTextToText,
-        image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
-        marks=[pytest.mark.skip("Model initialization hangs")],
-    ),
     ### Tensor parallel / multi-gpu broadcast tests
     "chameleon-broadcast": VLMTestInfo(
         models=["facebook/chameleon-7b"],

tests/models/multimodal/generation/vlm_utils/model_utils.py

@@ -1220,18 +1220,6 @@ def qwen3_vl_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
     return hf_model
 
 
-def tarsier_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
-    from vllm.model_executor.models.tarsier import get_vision_encoder_info
-
-    vision_encoder_info = get_vision_encoder_info(hf_model.config)
-
-    hf_processor = hf_model.processor
-    if hf_processor.patch_size is None:
-        hf_processor.patch_size = vision_encoder_info.get_patch_size()
-
-    return hf_model
-
-
 def voxtral_patch_hf_runner(hf_model: "HfRunner") -> "HfRunner":
     """Patch HfRunner for Voxtral's conversation-based processor.
 

tests/models/multimodal/processing/test_moss_audio.py

@@ -17,7 +17,6 @@ from vllm.model_executor.models.moss_audio import (
     MOSS_AUDIO_PLACEHOLDER,
     MOSS_AUDIO_TOKEN,
     MOSS_AUDIO_TOKEN_ID,
-    GatedMLP,
     MossAudioConfig,
     MossAudioDummyInputsBuilder,
     MossAudioEncoder,
@@ -594,53 +593,6 @@ def test_moss_qwen3_deepstack_keys_for_pp(monkeypatch):
     assert set(forward_tensors.tensors) == set(tensors.tensors)
 
 
-@pytest.mark.parametrize("tp_size", [1, 2])
-def test_moss_audio_gated_mlp_tp_shapes_and_loading(monkeypatch, tp_size):
-    from vllm.config import VllmConfig, set_current_vllm_config
-    from vllm.config.device import DeviceConfig
-
-    _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=tp_size)
-    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
-        mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
-
-    params = dict(mlp.named_parameters())
-    assert params["gate_up_proj.weight"].shape == torch.Size([16 // tp_size, 4])
-    assert params["down_proj.weight"].shape == torch.Size([6, 8 // tp_size])
-
-    gate_weight = torch.arange(32, dtype=torch.float32).reshape(8, 4)
-    up_weight = torch.arange(100, 132, dtype=torch.float32).reshape(8, 4)
-    down_weight = torch.arange(48, dtype=torch.float32).reshape(6, 8)
-    loaded = mlp.load_weights(
-        [
-            ("gate_proj.weight", gate_weight),
-            ("up_proj.weight", up_weight),
-            ("down_proj.weight", down_weight),
-        ]
-    )
-
-    assert loaded == {"gate_up_proj.weight", "down_proj.weight"}
-    shard = 8 // tp_size
-    assert torch.equal(params["gate_up_proj.weight"][:shard], gate_weight[:shard])
-    assert torch.equal(params["gate_up_proj.weight"][shard:], up_weight[:shard])
-    assert torch.equal(params["down_proj.weight"], down_weight[:, : 8 // tp_size])
-
-    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
-        packed_mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
-    packed_params = dict(packed_mlp.named_parameters())
-    loaded = packed_mlp.load_weights(
-        [("gate_up_proj.weight", torch.cat([gate_weight, up_weight], dim=0))]
-    )
-    assert loaded == {"gate_up_proj.weight"}
-    assert torch.equal(
-        packed_params["gate_up_proj.weight"][:shard],
-        gate_weight[:shard],
-    )
-    assert torch.equal(
-        packed_params["gate_up_proj.weight"][shard:],
-        up_weight[:shard],
-    )
-
-
 def test_moss_audio_encoder_loads_realistic_attention_weight_names(monkeypatch):
     from vllm.config import VllmConfig, set_current_vllm_config
     from vllm.config.device import DeviceConfig

tests/models/registry.py

@@ -1369,22 +1369,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "fixie-ai/ultravox-v0_5-llama-3_2-1b",
         trust_remote_code=True,
     ),
-    "TarsierForConditionalGeneration": _HfExamplesInfo("omni-research/Tarsier-7b"),
-    "Tarsier2ForConditionalGeneration": _HfExamplesInfo(
-        "omni-research/Tarsier2-Recap-7b",
-        hf_overrides={
-            "architectures": ["Tarsier2ForConditionalGeneration"],
-            "model_type": "tarsier2",
-        },
-        max_transformers_version="5.3",
-        transformers_version_reason={
-            "vllm": (
-                "Qwen2VLConfig was split into Qwen2VLConfig + "
-                "Qwen2VLTextConfig in transformers v5, breaking "
-                "attribute access (num_attention_heads, hidden_size, etc.)"
-            )
-        },
-    ),
     "VoxtralForConditionalGeneration": _HfExamplesInfo(
         "mistralai/Voxtral-Mini-3B-2507",
         tokenizer_mode="mistral",

tests/renderers/test_completions.py

@@ -64,12 +64,14 @@ class DummyTokenizer:
 
     def __post_init__(self) -> None:
         self._captured_encode_kwargs: dict = {}
+        self._captured_text_len: int = 0
 
     def decode(self, tokens: list[int]):
         return str(tokens)
 
     def encode(self, text: str, **kwargs):
         self._captured_encode_kwargs = kwargs
+        self._captured_text_len = len(text)
 
         in_length = len(text)
         truncation = kwargs.get("truncation")
@@ -366,6 +368,85 @@ class TestRenderPrompt:
         assert results[0]["prompt_token_ids"] == tokens
         assert results[0]["prompt"] == "[1, 2, 3, 4]"
 
+    def test_explicit_side_tokenizer_unbounded(self):
+        renderer = _build_renderer(MockModelConfig())
+
+        prompts = renderer.render_prompts(
+            _preprocess_prompt(renderer.model_config, "x" * 500)
+        )
+        results = renderer.tokenize_prompts(
+            prompts,
+            TokenizeParams(
+                max_total_tokens=100,
+                truncate_prompt_tokens=4,
+                truncation_side="left",
+            ),
+        )
+
+        assert len(results) == 1
+        assert len(results[0]["prompt_token_ids"]) == 4
+
+        kwargs = renderer.tokenizer._captured_encode_kwargs
+        assert kwargs["truncation"] is False
+
+    def test_explicit_side_left_text(self):
+        renderer = _build_renderer(MockModelConfig())
+
+        prompts = renderer.render_prompts(
+            _preprocess_prompt(renderer.model_config, "x" * 50)
+        )
+        results = renderer.tokenize_prompts(
+            prompts,
+            TokenizeParams(
+                max_total_tokens=100,
+                truncate_prompt_tokens=5,
+                truncation_side="left",
+            ),
+        )
+
+        assert len(results) == 1
+        assert len(results[0]["prompt_token_ids"]) == 5
+        assert results[0]["prompt_token_ids"] == list(range(45, 50))
+
+    def test_explicit_side_right_text(self):
+        renderer = _build_renderer(MockModelConfig())
+
+        prompts = renderer.render_prompts(
+            _preprocess_prompt(renderer.model_config, "x" * 50)
+        )
+        results = renderer.tokenize_prompts(
+            prompts,
+            TokenizeParams(
+                max_total_tokens=100,
+                truncate_prompt_tokens=5,
+                truncation_side="right",
+            ),
+        )
+
+        assert len(results) == 1
+        assert len(results[0]["prompt_token_ids"]) == 5
+        assert results[0]["prompt_token_ids"] == list(range(5))
+
+    def test_explicit_side_text_pretokenization_guard(self):
+        renderer = _build_renderer(MockModelConfig(), max_chars_per_token=1)
+
+        prompts = renderer.render_prompts(
+            _preprocess_prompt(renderer.model_config, "x" * 500)
+        )
+        results = renderer.tokenize_prompts(
+            prompts,
+            TokenizeParams(
+                max_total_tokens=100,
+                truncate_prompt_tokens=4,
+                truncation_side="left",
+            ),
+        )
+
+        assert len(results) == 1
+        assert len(results[0]["prompt_token_ids"]) == 4
+
+        assert renderer.tokenizer._captured_text_len <= 100
+
 
 class TestRenderEmbedPrompt:
     def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes:

tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh

@@ -51,6 +51,7 @@ vllm serve $MODEL \
   --trust-remote-code \
   --enable-prefix-caching \
   --mamba-cache-mode all \
+  --attention-backend FLASHINFER \
   --kv-transfer-config "$KV_CONFIG" &
 
 # Start decode instance
@@ -68,6 +69,7 @@ vllm serve $MODEL \
   --trust-remote-code \
   --enable-prefix-caching \
   --mamba-cache-mode all \
+  --attention-backend FLASHINFER \
   --kv-transfer-config "$KV_CONFIG" &
 
 echo "Waiting for prefill instance on port $PREFILL_PORT..."

tests/v1/kv_connector/unit/test_hidden_states_connector.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU-only unit tests for ExampleHiddenStatesConnector KV-cache-group logic."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import (  # noqa: E501
+    ExampleHiddenStatesConnector,
+)
+from vllm.v1.core.kv_cache_utils import get_kv_cache_groups
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    HiddenStateCacheSpec,
+    KVCacheGroupSpec,
+    MLAAttentionSpec,
+    SlidingWindowMLASpec,
+)
+
+
+def _full(block_size: int) -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=block_size, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
+    )
+
+
+def _hidden(block_size: int) -> HiddenStateCacheSpec:
+    return HiddenStateCacheSpec(
+        block_size=block_size, num_kv_heads=6, head_size=2048, dtype=torch.bfloat16
+    )
+
+
+def _config(*specs):
+    """Minimal stand-in exposing only ``kv_cache_groups`` (all the helpers read)."""
+    return SimpleNamespace(
+        kv_cache_groups=[
+            KVCacheGroupSpec(layer_names=[f"layer.{i}"], kv_cache_spec=spec)
+            for i, spec in enumerate(specs)
+        ]
+    )
+
+
+# ---- _find_cache_kv_group_id ------------------------------------------------
+
+
+def test_find_group_id_none_config_returns_zero():
+    assert ExampleHiddenStatesConnector._find_cache_kv_group_id(None) == 0
+
+
+def test_find_group_id_single_non_hidden_group_returns_zero():
+    # Uniform (dense) model: one group, no HiddenStateCacheSpec -> group 0.
+    cfg = _config(_full(16))
+    assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 0
+
+
+def test_find_group_id_locates_hidden_group_when_not_first():
+    # Hybrid layout: the hidden-states group is not group 0.
+    cfg = _config(_full(528), _hidden(22), _full(528))
+    assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 1
+
+
+def test_find_group_id_locates_hidden_group_last():
+    cfg = _config(_full(528), _full(528), _hidden(22))
+    assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 2
+
+
+def test_find_group_id_raises_when_no_hidden_group_and_multiple_groups():
+    cfg = _config(_full(16), _full(16))
+    with pytest.raises(ValueError, match="Could not uniquely identify"):
+        ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
+
+
+def test_find_group_id_raises_when_multiple_hidden_groups():
+    cfg = _config(_hidden(22), _hidden(22))
+    with pytest.raises(ValueError, match="Could not uniquely identify"):
+        ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
+
+
+# ---- _get_cache_block_size --------------------------------------------------
+
+
+def test_get_block_size_reads_hidden_group_spec_not_global():
+    # Hidden group keeps block size 22; the global is bumped to 528 for hybrids.
+    vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=528))
+    cfg = _config(_full(528), _hidden(22))
+    block_size = ExampleHiddenStatesConnector._get_cache_block_size(
+        vllm_config, cfg, cache_kv_group_id=1
+    )
+    assert block_size == 22
+
+
+def test_get_block_size_falls_back_to_cache_config_when_no_kv_cache_config():
+    vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16))
+    block_size = ExampleHiddenStatesConnector._get_cache_block_size(
+        vllm_config, None, cache_kv_group_id=0
+    )
+    assert block_size == 16
+
+
+# ---- MLA-verifier absorption ------------------------------------------------
+
+
+def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier():
+    # HiddenStateCacheSpec subclasses MLAAttentionSpec, so an MLA + sliding-
+    # window MLA verifier absorbs it into the MLA group instead of isolating it.
+    dt = torch.bfloat16
+    spec = {
+        "layers.0.mla": MLAAttentionSpec(
+            block_size=64, num_kv_heads=1, head_size=576, dtype=dt
+        ),
+        "layers.1.swa": SlidingWindowMLASpec(
+            block_size=64, num_kv_heads=1, head_size=576, dtype=dt, sliding_window=512
+        ),
+        "cache_only_layers.61": _hidden(64),
+    }
+    vllm_config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
+        speculative_config=None,
+    )
+    groups = get_kv_cache_groups(vllm_config, spec)
+    assert not any(isinstance(g.kv_cache_spec, HiddenStateCacheSpec) for g in groups)
+    cfg = SimpleNamespace(kv_cache_groups=groups)
+    with pytest.raises(ValueError, match="MLA verifiers are unsupported"):
+        ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)

tests/v1/kv_offload/tiering/p2p/test_manager.py

@@ -13,7 +13,7 @@ from types import SimpleNamespace
 
 import numpy as np
 
-from vllm.v1.kv_offload.base import LookupResult, ReqContext
+from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext
 from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
 from vllm.v1.kv_offload.tiering.p2p import manager as manager_module
 from vllm.v1.kv_offload.tiering.p2p.manager import (
@@ -1279,7 +1279,12 @@ class TestOnScheduleEnd:
         mgr = _make_manager()
         before_sessions = dict(mgr._sessions)
         before_jobs = list(mgr._finished_jobs)
-        assert mgr.on_schedule_end() is None
+        assert (
+            mgr.on_schedule_end(
+                ScheduleEndContext(new_req_ids=[], preempted_req_ids=())
+            )
+            is None
+        )
         assert mgr._sessions == before_sessions
         assert mgr._finished_jobs == before_jobs
 

tests/v1/kv_offload/tiering/test_fs_tier.py

@@ -22,6 +22,7 @@ from vllm.v1.kv_offload.base import (
     LookupResult,
     OffloadKey,
     ReqContext,
+    ScheduleEndContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.tiering.base import JobMetadata
@@ -107,7 +108,7 @@ def lookup_and_wait(
     """Perform a full async lookup cycle and return resolved results."""
     for k in keys:
         tier.lookup(k, ctx)
-    tier.on_schedule_end()
+    tier.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=()))
     deadline = time.monotonic() + timeout
     while time.monotonic() < deadline:
         if not tier._lookup_manager._pending_results.empty():

tests/v1/kv_offload/tiering/test_obj_tier.py

@@ -21,6 +21,7 @@ from vllm.v1.kv_offload.base import (
     LookupResult,
     OffloadKey,
     ReqContext,
+    ScheduleEndContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
@@ -222,7 +223,7 @@ def lookup_and_wait(
     """Perform a full async lookup cycle and return resolved results."""
     for k in keys:
         tier.lookup(k, ctx)
-    tier.on_schedule_end()
+    tier.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=()))
     deadline = time.monotonic() + timeout
     while time.monotonic() < deadline:
         if not tier._lookup_manager._pending_results.empty():

tests/v1/kv_offload/tiering/test_tiering_offloading.py

@@ -27,6 +27,7 @@ from vllm.v1.kv_offload.base import (
     OffloadPolicy,
     ReqContext,
     RequestOffloadingContext,
+    ScheduleEndContext,
     make_offload_key,
 )
 from vllm.v1.kv_offload.tiering.base import (
@@ -233,7 +234,8 @@ class TestTieringOffloadingManager:
 
     def _simulate_on_schedule_end(self):
         """Simulate end of scheduler step: lifecycle flush + drain events."""
-        self.manager.on_schedule_end()
+        ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=())
+        self.manager.on_schedule_end(ctx)
         list(self.manager.take_events())
 
     def _start_request(self, req_context: ReqContext = _CTX):

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -309,3 +309,106 @@ def test_placeholder_draft_token_rejected():
     assert torch.equal(num_sampled, torch.ones_like(num_sampled))
     recovered = sampled[:, 0]
     assert (recovered >= 0).all() and (recovered < VOCAB_SIZE).all()
+
+
+@pytest.mark.parametrize(
+    "num_speculative_steps,temperature",
+    [
+        (1, 0.6),
+        (3, 0.6),
+        (1, 1.0),
+        (3, 1.0),
+        (5, 1.0),
+    ],
+)
+@pytest.mark.parametrize("has_draft_logits", [True, False])
+def test_block_verification_rejection_sample(
+    num_speculative_steps: int, temperature: float, has_draft_logits: bool
+):
+    """
+    Verify that block verification (Sun et al.) preserves the target
+    distribution at every accepted position, for both the full draft-logits
+    case and the one-hot (no draft logits) case. Block verification changes
+    *which* prefix is accepted, but the marginal of every output position must
+    still match the target distribution p(x).
+    """
+
+    torch.manual_seed(42)
+    device = "cuda"
+    num_trials = 10 * VOCAB_SIZE
+
+    target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
+    draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
+
+    if temperature > 0:
+        target_logits_1d /= temperature
+        draft_logits_1d /= temperature
+
+    inputs = _build_rejection_sample_inputs(
+        target_logits_1d,
+        draft_logits_1d,
+        num_speculative_steps,
+        temperature=temperature,
+        num_trials=num_trials,
+    )
+    if not has_draft_logits:
+        inputs["draft_logits"] = None
+
+    sampled, num_sampled = rejection_sample(
+        **inputs,
+        num_speculative_steps=num_speculative_steps,
+        use_block_verification=True,
+    )
+
+    target_probs = torch.softmax(target_logits_1d, dim=0)
+    for pos in range(num_speculative_steps + 1):
+        accepted_mask = num_sampled >= pos + 1
+        _assert_distribution_match(
+            sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}"
+        )
+
+
+@pytest.mark.parametrize("num_speculative_steps", [3, 5])
+def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int):
+    """
+    Block verification is designed to accept at least as long a prefix as
+    token verification in expectation. Verify the mean accepted length is no
+    worse than the standard method on the same inputs.
+    """
+
+    torch.manual_seed(0)
+    device = "cuda"
+    num_trials = 20 * VOCAB_SIZE
+    temperature = 1.0
+
+    target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
+    # A draft close to the target gives block verification room to recover
+    # prefixes that token verification would have truncated.
+    draft_logits_1d = target_logits_1d + 0.5 * torch.randn(
+        VOCAB_SIZE, device=device, dtype=torch.float32
+    )
+
+    inputs = _build_rejection_sample_inputs(
+        target_logits_1d,
+        draft_logits_1d,
+        num_speculative_steps,
+        temperature=temperature,
+        num_trials=num_trials,
+    )
+
+    _, num_sampled_standard = rejection_sample(
+        **inputs, num_speculative_steps=num_speculative_steps
+    )
+    _, num_sampled_block = rejection_sample(
+        **inputs,
+        num_speculative_steps=num_speculative_steps,
+        use_block_verification=True,
+    )
+
+    mean_standard = (num_sampled_standard - 1).float().mean().item()
+    mean_block = (num_sampled_block - 1).float().mean().item()
+    # Allow a small slack for sampling noise.
+    assert mean_block >= mean_standard - 1e-2, (
+        f"Block verification mean accepted length {mean_block:.4f} is worse "
+        f"than standard {mean_standard:.4f}."
+    )

vllm/config/speculative.py

@@ -67,7 +67,7 @@ SpeculativeMethod = Literal[
     EagleModelTypes,
     NgramGPUTypes,
 ]
-RejectionSampleMethod = Literal["standard", "synthetic"]
+RejectionSampleMethod = Literal["standard", "synthetic", "block"]
 DraftSampleMethod = Literal["greedy", "probabilistic"]
 
 
@@ -201,7 +201,9 @@ class SpeculativeConfig:
     """The rejection sampling method to use. 'standard' uses probabilistic
     rejection sampling (with or without cached draft logits, controlled by
     draft_sample_method). 'synthetic' accepts draft tokens with a decaying
-    probability calibrated to synthetic_acceptance_rate."""
+    probability calibrated to synthetic_acceptance_rate. 'block' uses block
+    verification (Sun et al.), which jointly verifies the draft tokens as a
+    block instead of one at a time."""
 
     synthetic_acceptance_rates: list[float] | None = None
     """Per-position *unconditional* acceptance rates for synthetic rejection

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

@@ -109,6 +109,49 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
         # Must be False so that drafter kv cache isn't merged with verifier's
         return False
 
+    @classmethod
+    def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int:
+        """Index of the KV cache group holding the extracted hidden states.
+
+        Located by spec type so it resolves on both scheduler and worker side.
+        """
+        if kv_cache_config is None:
+            return 0
+
+        from vllm.v1.kv_cache_interface import HiddenStateCacheSpec
+
+        groups = kv_cache_config.kv_cache_groups
+        group_ids = [
+            gid
+            for gid, group in enumerate(groups)
+            if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
+        ]
+        if len(group_ids) == 1:
+            return group_ids[0]
+        if not group_ids and len(groups) == 1:
+            return 0
+        raise ValueError(
+            "Could not uniquely identify the extract-hidden-states KV cache "
+            f"group among {len(groups)} groups; the hidden-states layer must be "
+            "isolated in its own group (MLA verifiers are unsupported)."
+        )
+
+    @staticmethod
+    def _get_cache_block_size(
+        vllm_config: "VllmConfig",
+        kv_cache_config: "KVCacheConfig | None",
+        cache_kv_group_id: int,
+    ) -> int:
+        """Block size of the hidden-states group, read from its own spec.
+
+        cache_config.block_size is bumped to a common multiple for hybrid
+        verifiers; the page-aligned hidden-states group keeps a smaller one.
+        """
+        if kv_cache_config is None:
+            return vllm_config.cache_config.block_size
+        cache_group = kv_cache_config.kv_cache_groups[cache_kv_group_id]
+        return cache_group.kv_cache_spec.block_size
+
     def __init__(
         self,
         vllm_config: "VllmConfig",
@@ -120,7 +163,12 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
             role=role,
             kv_cache_config=kv_cache_config,
         )
-        self._block_size = vllm_config.cache_config.block_size
+        # Read the hidden-states group and its block size from the group spec;
+        # cache_config.block_size is bumped (wrong) for hybrid verifiers.
+        self._cache_kv_group_id = self._find_cache_kv_group_id(kv_cache_config)
+        self._block_size = self._get_cache_block_size(
+            vllm_config, kv_cache_config, self._cache_kv_group_id
+        )
         self._storage_path = self._kv_transfer_config.get_from_extra_config(
             "shared_storage_path", "/tmp"
         )
@@ -151,13 +199,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
         # Worker-side state (set by register_kv_caches).
         self._kv_cache: torch.Tensor | None = None
 
-        # Identify which KV cache group holds the hidden-states layer.
-        self._hs_group_idx: int = 0
-        if self._kv_cache_config is not None:
-            for i, group in enumerate(self._kv_cache_config.kv_cache_groups):
-                if any("cache_only_layers" in n for n in group.layer_names):
-                    self._hs_group_idx = i
-                    break
         # Only TP rank 0 writes hidden states to disk; other TP ranks no-op.
         # Set in register_kv_caches (after distributed init).
         self._is_tp_rank_zero: bool = True
@@ -267,12 +308,14 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
         )
         self._kv_cache = kv_caches[self.cache_layers[0]]
 
-        # Find the KV cache group index for hidden states
-        if self._kv_cache_config is not None:
-            for i, group in enumerate(self._kv_cache_config.kv_cache_groups):
-                if self.cache_layers[0] in group.layer_names:
-                    self._hs_group_idx = i
-                    break
+        # Block size must match the indexed buffer, else reads hit the wrong
+        # slots. Raise (not assert) so the check survives `python -O`.
+        if self._block_size != self._kv_cache.shape[1]:
+            raise ValueError(
+                f"Hidden-states block-size mismatch: derived {self._block_size} "
+                f"but buffer block size is {self._kv_cache.shape[1]}; read slots "
+                "would be wrong (likely a hybrid block-size resolution bug)."
+            )
 
     @staticmethod
     def _write_tensors(
@@ -543,7 +586,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
         request: "Request",
         block_ids: tuple[list[int], ...],
     ) -> tuple[bool, dict[str, Any] | None]:
-        return self.request_finished(request, block_ids[self._hs_group_idx])
+        return self.request_finished(request, block_ids[self._cache_kv_group_id])
 
     @classmethod
     def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -42,6 +42,7 @@ from vllm.v1.kv_offload.base import (
     OffloadPolicy,
     ReqContext,
     RequestOffloadingContext,
+    ScheduleEndContext,
     make_offload_key,
 )
 from vllm.v1.outputs import KVConnectorOutput
@@ -1027,7 +1028,11 @@ class OffloadingConnectorScheduler:
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:
         self._update_req_states(scheduler_output)
-        self.manager.on_schedule_end()
+        schedule_end_context = ScheduleEndContext(
+            new_req_ids=[req.req_id for req in scheduler_output.scheduled_new_reqs],
+            preempted_req_ids=scheduler_output.preempted_req_ids or (),
+        )
+        self.manager.on_schedule_end(schedule_end_context)
 
         # Flush jobs for preempted requests.
         for req_id in scheduler_output.preempted_req_ids or ():

vllm/distributed/utils.py

@@ -64,6 +64,24 @@ def divide(numerator, denominator):
     return numerator // denominator
 
 
+def verify_group_size_divides_partition(
+    input_size_per_partition: int,
+    group_size: int,
+    layer_name: str | None = None,
+    extra_suggestion: str = "",
+) -> None:
+    """Validate that a TP-sharded layer holds a whole number of quant groups."""
+    if input_size_per_partition % group_size == 0:
+        return
+    location = f" for layer '{layer_name}'" if layer_name else ""
+    raise ValueError(
+        f"Weight {input_size_per_partition=}{location} is not divisible by "
+        f"{group_size=}. This happens when tensor_parallel_size splits the layer input "
+        "into shards that are not a whole number of quant groups. Consider reducing "
+        f"tensor_parallel_size{extra_suggestion}."
+    )
+
+
 def is_weak_contiguous(inp: torch.Tensor) -> bool:
     """Check that *inp* occupies a single contiguous block of memory.
 

vllm/lora/layers/fused_moe.py

@@ -39,6 +39,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
         self._ep_check()
 
         routed_experts = self.base_layer.routed_experts
+        routed_experts.lora_base_layer_prefix = "base_layer."
         assert not routed_experts.quant_method.is_monolithic, (
             "Monolithic kernels are not supported for Fused MoE LoRA."
         )

vllm/lora/utils.py

@@ -372,32 +372,26 @@ def process_packed_modules_mapping(
     model: nn.Module, force_2d_moe: bool = False
 ) -> dict[str, list[str]]:
     if is_moe_model(model):
-        if moe_packed_mapping := get_moe_expert_mapping(model):
-            # This method generates and returns a dictionary mapping packed module
-            # names to lists of their corresponding submodule names. It includes
-            # both static mappings and dynamic mappings for expert layers, where
-            # the expert indices are expanded based on the configured number
-            # of routed experts.
-            packed_modules_mapping = get_packed_modules_mapping(model)
-            # The 2D mapping is needed when the model itself is 2D, or when
-            # the engine forces the universal 2D wrapper via
-            # enable_mixed_moe_lora_format (so 3D models can also load 2D
-            # adapters through FusedMoEWithLoRA).
-            if (not model.is_3d_moe_weight) or force_2d_moe:
-                # Filter out malformed entries: non-gated MoE has empty
-                # ckpt_up_proj_name which results in weight_name containing ".."
-                # (e.g., "experts.0.." instead of "experts.0.layer_name.")
-                packed_modules_mapping["experts"] = [
-                    weight_name.rstrip(".")
-                    for _, weight_name, _, _ in moe_packed_mapping
-                    if ".." not in weight_name
-                ]
+        # This method generates and returns a dictionary mapping packed module
+        # names to lists of their corresponding submodule names. It includes
+        # both static mappings and dynamic mappings for expert layers, where
+        # the expert indices are expanded based on the configured number
+        # of routed experts.
+        packed_modules_mapping = get_packed_modules_mapping(model)
+        # The 2D mapping is needed when the model itself is 2D, or when
+        # the engine forces the universal 2D wrapper via
+        # enable_mixed_moe_lora_format (so 3D models can also load 2D
+        # adapters through FusedMoEWithLoRA).
+        if (not model.is_3d_moe_weight) or force_2d_moe:
+            # Filter out malformed entries: non-gated MoE has empty
+            # ckpt_up_proj_name which results in weight_name containing ".."
+            # (e.g., "experts.0.." instead of "experts.0.layer_name.")
+            packed_modules_mapping["experts"] = [
+                weight_name.rstrip(".")
+                for _, weight_name, _, _ in get_moe_expert_mapping(model)
+                if ".." not in weight_name
+            ]
 
-            return packed_modules_mapping
-        else:
-            raise AttributeError(
-                "To support LoRA for MoE model, "
-                "'get_expert_mapping' must be implemented"
-            )
+        return packed_modules_mapping
     else:
         return get_packed_modules_mapping(model)

vllm/model_executor/layers/fused_moe/layer.py

@@ -127,7 +127,7 @@ def FusedMoE(
     num_redundant_experts: int = 0,
     has_bias: bool = False,
     is_sequence_parallel: bool = False,
-    expert_mapping: list[tuple[str, str, int, str]] | None = None,
+    ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"),
     n_shared_experts: int | None = None,
     router_logits_dtype: torch.dtype | None = None,
     gate: torch.nn.Module | None = None,
@@ -356,7 +356,9 @@ def FusedMoE(
         moe_config,
         quant_config,
         expert_map_manager=expert_map_manager,
-        expert_mapping=expert_mapping,
+        ckpt_gate_proj_name=ckpt_names[0],
+        ckpt_down_proj_name=ckpt_names[1],
+        ckpt_up_proj_name=ckpt_names[2],
         # Extra params that are needed by quant_methods, pass along for now
         # Prefer getting these from other sources, e.g. moe_config or
         # router object

vllm/model_executor/layers/fused_moe/routed_experts.py

@@ -58,7 +58,9 @@ class RoutedExperts(PluggableLayer):
         moe_config: FusedMoEConfig,
         quant_config: QuantizationConfig | None,
         expert_map_manager: ExpertMapManager,
-        expert_mapping: list[tuple[str, str, int, str]] | None = None,
+        ckpt_gate_proj_name: str = "gate_proj",
+        ckpt_down_proj_name: str = "down_proj",
+        ckpt_up_proj_name: str = "up_proj",
         #
         # Extra params that are needed by quant_methods, pass along for now
         # Prefer getting these from other sources, e.g. moe_config or
@@ -81,7 +83,9 @@ class RoutedExperts(PluggableLayer):
         self.layer_name = layer_name
         self.moe_config = moe_config
         self.quant_config = quant_config
-        self.expert_mapping = expert_mapping
+        self.ckpt_gate_proj_name = ckpt_gate_proj_name
+        self.ckpt_down_proj_name = ckpt_down_proj_name
+        self.ckpt_up_proj_name = ckpt_up_proj_name
         self.expert_map_manager = expert_map_manager
         self.hidden_size = moe_config.hidden_dim
         self.global_num_experts = moe_config.num_experts
@@ -167,6 +171,8 @@ class RoutedExperts(PluggableLayer):
 
         self.quant_method.create_weights(layer=self, **moe_quant_params)
 
+        self.lora_base_layer_prefix = ""
+
     # TODO(bnell): Temporary hack. Get rid of this.
     def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
         self.quant_method = quant_method
@@ -863,26 +869,33 @@ class RoutedExperts(PluggableLayer):
     def load_weights(
         self, weights: Iterable[tuple[str, torch.Tensor]]
     ) -> Iterable[str]:
-        if (expert_mapping := self.expert_mapping) is None:
-            raise ValueError(
-                "`self.expert_mapping` must be provided to "
-                "load weights using `self.load_weights`."
-            )
+        expert_mapping = self.get_expert_mapping(include_fused=True)
+        unpadded_hidden = self.moe_config.hidden_dim_unpadded
         for expert_name, loaded_weight in weights:
             qual_name = f"{self.layer_name}.{expert_name}"
+            # Fused expert weights can be identified by their 3D tensors
+            is_fused = loaded_weight.dim() == 3
+            matched = False
             for param_name, weight_name, expert_id, shard_id in expert_mapping:
                 if weight_name not in qual_name:
+                    if matched and is_fused:
+                        break
                     continue
+                matched = True
                 weight_name = qual_name.replace(weight_name, param_name)
                 param_name = weight_name.removeprefix(f"{self.layer_name}.")
                 param = getattr(self, param_name)
-                # Fused expert weights can be identified by their 3D tensors
-                if loaded_weight.dim() == 3:
-                    # Repurpose expert_id as shard_idx for deconcatenating w1 and w3
+                if is_fused:
                     if shard_id in {"w1", "w3"}:
-                        shard_idx = expert_id
-                        experts_shard = loaded_weight.chunk(2, dim=1)[shard_idx]
+                        if loaded_weight.shape[-1] != unpadded_hidden:
+                            # [..., hidden, intermediate] -> [..., intermediate, hidden]
+                            loaded_weight = loaded_weight.transpose(-1, -2)
+                        # Repurpose expert_id for deconcatenating w1 and w3
+                        experts_shard = loaded_weight.chunk(2, dim=1)[expert_id]
                     else:
+                        if loaded_weight.shape[-2] != unpadded_hidden:
+                            # [..., intermediate, hidden] -> [..., hidden, intermediate]
+                            loaded_weight = loaded_weight.transpose(-1, -2)
                         experts_shard = loaded_weight
                     start = 0
                 else:
@@ -912,6 +925,27 @@ class RoutedExperts(PluggableLayer):
                         )
                         yield param_name
 
+    def get_expert_mapping(
+        self,
+        ckpt_gate_proj_name: str | None = None,
+        ckpt_down_proj_name: str | None = None,
+        ckpt_up_proj_name: str | None = None,
+        include_fused: bool = False,
+    ) -> list[tuple[str, str, int, str]]:
+        moe_config = self.moe_config
+        num_fused_shared_experts = self.expert_map_manager.num_fused_shared_experts
+        num_redundant_experts = moe_config.num_experts - moe_config.num_logical_experts
+        return self.build_expert_params_mapping(
+            ckpt_gate_proj_name or self.ckpt_gate_proj_name,
+            ckpt_down_proj_name or self.ckpt_down_proj_name,
+            ckpt_up_proj_name or self.ckpt_up_proj_name,
+            num_experts=moe_config.num_logical_experts + num_fused_shared_experts,
+            num_redundant_experts=num_redundant_experts,
+            routed_experts_prefix="",
+            lora_base_layer_prefix=self.lora_base_layer_prefix,
+            include_fused=include_fused,
+        )
+
     @staticmethod
     def make_expert_params_mapping(
         model: torch.nn.Module,
@@ -921,6 +955,36 @@ class RoutedExperts(PluggableLayer):
         num_experts: int,
         num_redundant_experts: int = 0,
         routed_experts_prefix: str = "routed_experts",
+    ) -> list[tuple[str, str, int, str]]:
+        """Build the expert mapping, detecting the LoRA `base_layer.` prefix by
+        scanning `model`'s parameters.
+
+        Legacy entry point for models that still hand-roll `load_weights`; the
+        `RoutedExperts` weight loader uses `get_expert_mapping` /
+        `build_expert_params_mapping` instead (which take the prefix directly).
+        See `build_expert_params_mapping` for the returned tuple format.
+        """
+        has_base_layer = any(".base_layer." in n for n, _ in model.named_parameters())
+        return RoutedExperts.build_expert_params_mapping(
+            ckpt_gate_proj_name,
+            ckpt_down_proj_name,
+            ckpt_up_proj_name,
+            num_experts,
+            num_redundant_experts,
+            routed_experts_prefix,
+            "base_layer." if has_base_layer else "",
+        )
+
+    @staticmethod
+    def build_expert_params_mapping(
+        ckpt_gate_proj_name: str,
+        ckpt_down_proj_name: str,
+        ckpt_up_proj_name: str,
+        num_experts: int,
+        num_redundant_experts: int = 0,
+        routed_experts_prefix: str = "routed_experts",
+        lora_base_layer_prefix: str = "",
+        include_fused: bool = False,
     ) -> list[tuple[str, str, int, str]]:
         """
         Create expert parameter mapping for weight loading with redundant experts.
@@ -929,12 +993,13 @@ class RoutedExperts(PluggableLayer):
         when loading weights with EPLB redundant experts.
 
         Args:
-            model: The model containing the MoE layer
             ckpt_gate_proj_name: Name of gate projection in checkpoint
             ckpt_down_proj_name: Name of down projection in checkpoint
             ckpt_up_proj_name: Name of up projection in checkpoint
             num_experts: Number of logical (non-redundant) experts
             num_redundant_experts: Number of redundant experts
+            lora_base_layer_prefix: Prefix to add if this layer is a LoRA base layer
+            include_fused: Prepend the fused pre-fused-checkpoint entries
 
         Returns:
             List of tuples (param_name, weight_name, expert_id, shard_id)
@@ -956,22 +1021,39 @@ class RoutedExperts(PluggableLayer):
             )
         )
 
-        base_layer = (
-            "base_layer."
-            if any(".base_layer." in name for name, _ in model.named_parameters())
-            else ""
-        )
-
         if routed_experts_prefix != "":
             routed_experts_prefix = f"{routed_experts_prefix}."
 
-        return [
+        w13 = f"experts.{routed_experts_prefix}{lora_base_layer_prefix}w13_"
+        w2 = f"experts.{routed_experts_prefix}{lora_base_layer_prefix}w2_"
+
+        fused_mapping = []
+        if include_fused:
+            gate_up = None
+            if ckpt_gate_proj_name == "gate_proj" and ckpt_up_proj_name == "up_proj":
+                gate_up = "gate_up_proj"
+            elif ckpt_gate_proj_name == "w1" and ckpt_up_proj_name == "w3":
+                gate_up = "w13"
+            else:
+                logger.warning(
+                    "Unexpected gate/up projection names: %s, %s. "
+                    "Fused gate/up mapping will be skipped.",
+                    ckpt_gate_proj_name,
+                    ckpt_up_proj_name,
+                )
+            if gate_up is not None:
+                fused_mapping = [
+                    # (param_name, weight_name, expert_id, shard_id)
+                    (f"{w13}weight", f"experts.{gate_up}", 0, "w1"),
+                    (f"{w13}weight", f"experts.{gate_up}", 1, "w3"),
+                    (f"{w2}weight", f"experts.{ckpt_down_proj_name}", 0, "w2"),
+                ]
+
+        per_expert_mapping = [
             # (param_name, weight_name, expert_id, shard_id)
             (
-                f"experts.{routed_experts_prefix}{base_layer}w13_"
-                if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name]
-                else f"experts.{routed_experts_prefix}{base_layer}w2_",
-                f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.{base_layer}",
+                w13 if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name] else w2,
+                f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.{lora_base_layer_prefix}",
                 expert_id,
                 shard_id,
             )
@@ -983,6 +1065,8 @@ class RoutedExperts(PluggableLayer):
             ]
         ]
 
+        return fused_mapping + per_expert_mapping
+
     def get_expert_weights(self) -> Iterable[torch.Tensor]:
         def _maybe_make_contiguous(
             name: str, p: torch.nn.Parameter

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -292,6 +292,11 @@ class MoERunner(MoERunnerInterface):
         # For smuggling this layer into the fused moe custom op
         register_layer_for_moe_forward_op(get_current_vllm_config(), self)
 
+    def load_weights(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[str]:
+        return self.routed_experts.load_weights(weights)
+
     def _select_forward(self) -> Callable:
         if current_platform.is_tpu() or current_platform.is_cpu():
             # TODO: Once the OOM issue for the TPU backend is resolved, we

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py

@@ -6,6 +6,7 @@ from collections.abc import Callable
 import torch
 from compressed_tensors.quantization import ActivationOrdering
 
+from vllm.distributed.utils import verify_group_size_divides_partition
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import (
     MPLinearLayerConfig,
@@ -112,7 +113,7 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
         scales_and_zp_size = input_size // group_size
 
         if partition_scales:
-            assert input_size_per_partition % group_size == 0
+            verify_group_size_divides_partition(input_size_per_partition, group_size)
             scales_and_zp_size = input_size_per_partition // group_size
 
         weight = PackedvLLMParameter(

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py

@@ -5,6 +5,7 @@ from collections.abc import Callable
 
 import torch
 
+from vllm.distributed.utils import verify_group_size_divides_partition
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import (
     MPLinearLayerConfig,
@@ -79,9 +80,8 @@ class CompressedTensorsW4A8Int(CompressedTensorsScheme):
             effective_group_size = self.group_size
 
         # Ensure group_size divides input_size_per_partition
-        assert input_size_per_partition % effective_group_size == 0, (
-            f"input_size_per_partition {input_size_per_partition}"
-            f" not divisible by group_size {effective_group_size}"
+        verify_group_size_divides_partition(
+            input_size_per_partition, effective_group_size
         )
 
         # Determine scale partitioning

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py

@@ -8,6 +8,7 @@ from fractions import Fraction
 import torch
 from compressed_tensors.quantization import ActivationOrdering
 
+from vllm.distributed.utils import verify_group_size_divides_partition
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import (
     MarlinLinearKernel,
@@ -151,7 +152,9 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
         scales_and_zp_size = input_size // group_size
 
         if partition_scales:
-            assert input_size_per_partition % group_size == 0
+            verify_group_size_divides_partition(
+                input_size_per_partition, group_size, self.layer_name
+            )
             scales_and_zp_size = input_size_per_partition // group_size
 
         packed_input_dim = math.ceil(input_size_per_partition * self.num_bits / 32)

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py

@@ -13,6 +13,7 @@ from collections.abc import Callable
 import torch
 from compressed_tensors.compressors.pack_quantized.helpers import pack_to_int32
 
+from vllm.distributed.utils import verify_group_size_divides_partition
 from vllm.model_executor.kernels.linear import (
     MPLinearLayerConfig,
     choose_mp_linear_kernel,
@@ -171,7 +172,9 @@ class CompressedTensorsWNA8O8Int(CompressedTensorsScheme):
         scales = (input_size_per_partition if partitioned else input_size) // group_size
         scale_data = torch.empty(out, scales, dtype=params_dtype)
         if partitioned:
-            assert input_size_per_partition % group_size == 0
+            verify_group_size_divides_partition(
+                input_size_per_partition, group_size, self.layer_name
+            )
             weight_scale = GroupQuantScaleParameter(
                 data=scale_data, output_dim=0, input_dim=1, weight_loader=weight_loader
             )

vllm/model_executor/layers/quantization/utils/marlin_utils.py

@@ -9,6 +9,7 @@ import torch
 
 import vllm.envs as envs
 from vllm import _custom_ops as ops
+from vllm.distributed.utils import verify_group_size_divides_partition
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import RoutedExperts
 from vllm.model_executor.layers.linear import LinearBase
@@ -193,12 +194,11 @@ def verify_marlin_supports_shape(
             "with --quantization gptq."
         )
 
-    if group_size < input_size and input_size_per_partition % group_size != 0:
-        raise ValueError(
-            f"Weight input_size_per_partition = {input_size_per_partition}"
-            f" is not divisible by group_size = {group_size}. "
-            "Consider reducing tensor_parallel_size or running "
-            "with --quantization gptq."
+    if group_size < input_size:
+        verify_group_size_divides_partition(
+            input_size_per_partition,
+            group_size,
+            extra_suggestion=" or running with --quantization gptq",
         )
 
 

vllm/model_executor/model_loader/bitsandbytes_loader.py

@@ -567,12 +567,6 @@ class BitsAndBytesModelLoader(BaseModelLoader):
 
         if is_moe_model(model):
             self.expert_params_mapping = get_moe_expert_mapping(model)
-            if not self.expert_params_mapping:
-                raise AttributeError(
-                    f"MoE Model {type(model).__name__} does not support "
-                    "BitsAndBytes quantization yet. Ensure this model has "
-                    "'get_expert_mapping' method."
-                )
         # For some models like Molmo, we need to use hf_to_vllm_mapper
         # to ensure correct loading of weights.
         if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):

vllm/model_executor/models/afmoe.py

@@ -2,8 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Inference-only AfMoE model compatible with HuggingFace weights."""
 
-import typing
-from collections.abc import Callable, Iterable
+from collections.abc import Iterable
 from itertools import islice
 
 import torch
@@ -21,7 +20,6 @@ from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
     MoERunner,
-    fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
@@ -36,10 +34,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
 from vllm.model_executor.models.interfaces import (
     EagleModelMixin,
     MixtureOfExperts,
@@ -53,7 +47,6 @@ from vllm.model_executor.models.utils import (
     PPMissingLayer,
     WeightsMapper,
     extract_layer_index,
-    is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
@@ -481,120 +474,6 @@ class AfmoeModel(nn.Module, EagleModelMixin):
             }
         )
 
-    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        # Params for weights, fp8 weight scales, fp8 activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        return fused_moe_make_expert_params_mapping(
-            self,
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.num_experts,
-        )
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        expert_params_mapping = self.get_expert_mapping()
-
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                # Skip non-stacked layers and experts (experts handled below).
-                if (weight_name not in name) or ("self_attn.gate_proj" in name):
-                    continue
-                # We have mlp.experts[0].gate_proj in the checkpoint.
-                # Since we handle the experts below in expert_params_mapping,
-                # we need to skip here BEFORE we update the name, otherwise
-                # name will be updated to mlp.experts[0].gate_up_proj, which
-                # will then be updated below in expert_params_mapping
-                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
-                if ("mlp.experts." in name) and name not in params_dict:
-                    continue
-
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                is_expert_weight = False
-                for mapping in expert_params_mapping:
-                    param_name, weight_name, expert_id, shard_id = mapping
-                    if weight_name not in name:
-                        continue
-
-                    # Anyway, this is an expert weight and should not be
-                    # attempted to load as other weights later
-                    is_expert_weight = True
-
-                    # Do not modify `name` since the loop may continue here
-                    # Instead, create a new variable
-                    name_mapped = name.replace(weight_name, param_name)
-
-                    if is_pp_missing_parameter(name_mapped, self):
-                        continue
-
-                    param = params_dict[name_mapped]
-                    # We should ask the weight loader to return success or not
-                    # here since otherwise we may skip experts with other
-                    # available replicas.
-                    weight_loader = typing.cast(
-                        Callable[..., bool], param.weight_loader
-                    )
-                    success = weight_loader(
-                        param,
-                        loaded_weight,
-                        name_mapped,
-                        shard_id=shard_id,
-                        expert_id=expert_id,
-                        return_success=True,
-                    )
-                    if success:
-                        name = name_mapped
-                        break
-                else:
-                    if is_expert_weight:
-                        # We've checked that this is an expert weight
-                        # However it's not mapped locally to this rank
-                        # So we simply skip it
-                        continue
-
-                    # Skip loading extra bias for GPTQ models.
-                    if name.endswith(".bias") and name not in params_dict:
-                        continue
-
-                    # Remapping the name of FP8 kv-scale.
-                    name = maybe_remap_kv_scale_name(name, params_dict)
-                    if name is None:
-                        continue
-
-                    if is_pp_missing_parameter(name, self):
-                        continue
-
-                    param = params_dict[name]
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-
-        return loaded_params
-
 
 class AfmoeForCausalLM(
     nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA, MixtureOfExperts
@@ -615,6 +494,17 @@ class AfmoeForCausalLM(
         orig_to_new_suffix={
             ".router.gate.weight": ".gate.weight",
         },
+        orig_to_new_stacked={
+            # weight_name: (param_name, shard_id)
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            # `.self_attn.gate_proj` is a gated-attention projection (not fused).
+            ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
+            ".mlp.up_proj": (".mlp.gate_up_proj", 1),
+            ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
+            ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
+        },
     )
 
     fall_back_to_pt_during_load = False
@@ -705,6 +595,3 @@ class AfmoeForCausalLM(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-
-    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        return self.model.get_expert_mapping()

vllm/model_executor/models/aimv2.py

@@ -20,7 +20,7 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 from vllm.transformers_utils.configs.ovis import AIMv2Config
 
 
@@ -194,6 +194,13 @@ class AIMv2Transformer(nn.Module):
 
 
 class AIMv2Model(torch.nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".fc1": (".fc13", 0),
+            ".fc3": (".fc13", 1),
+        }
+    )
+
     def __init__(
         self,
         config: AIMv2Config,
@@ -218,34 +225,13 @@ class AIMv2Model(torch.nn.Module):
         return x
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".fc13", ".fc1", 0),
-            (".fc13", ".fc3", 1),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            # post_layernorm is optional in SiglipVisionModel
-            if (
-                name.startswith("trunk.post_trunk_norm")
-                and self.trunk.post_trunk_norm is None
-            ):
-                continue
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(
+            self,
+            # post_trunk_norm is optional (absent for clip-skip backbones).
+            skip_prefixes=(
+                ["trunk.post_trunk_norm."]
+                if self.trunk.post_trunk_norm is None
+                else None
+            ),
+        )
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/apertus.py

@@ -53,10 +53,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
 from vllm.sequence import IntermediateTensors
 from vllm.v1.attention.backend import AttentionType
 
@@ -70,8 +66,8 @@ from .interfaces import (
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    WeightsMapper,
     extract_layer_index,
-    is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
@@ -404,66 +400,17 @@ class ApertusModel(nn.Module, EagleModelMixin):
             return hidden_states, aux_hidden_states
         return hidden_states
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters())
-
-        # we need to load the buffers for beta and eps (XIELU)
-        for name, buffer in self.named_buffers():
-            if name.endswith(".beta") or name.endswith(".eps"):
-                params_dict[name] = buffer
-
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
-                # Models trained using ColossalAI may include these tensors in
-                # the checkpoint. Skip them.
-                continue
-            if "scale" in name or "zero_point" in name:
-                # Remapping the name of FP8 kv-scale.
-                name = maybe_remap_kv_scale_name(name, params_dict)
-                if name is None:
-                    continue
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-
-                if is_pp_missing_parameter(name, self):
-                    continue
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-
-                if is_pp_missing_parameter(name, self):
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
-
 
 class ApertusForCausalLM(
     nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
 ):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+        }
+    )
     packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
 
     # LoRA specific attributes
@@ -548,4 +495,4 @@ class ApertusForCausalLM(
             self,
             skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
         )
-        return loader.load_weights(weights)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/bert.py

@@ -369,6 +369,14 @@ class BertModel(nn.Module, SupportsQuant):
 
     packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]}
 
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".self.query": (".self.qkv_proj", "q"),
+            ".self.key": (".self.qkv_proj", "k"),
+            ".self.value": (".self.qkv_proj", "v"),
+        }
+    )
+
     def __init__(
         self,
         *,
@@ -400,43 +408,9 @@ class BertModel(nn.Module, SupportsQuant):
 
         return self.encoder(hidden_states)
 
-    def _load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "query", "q"),
-            ("qkv_proj", "key", "k"),
-            ("qkv_proj", "value", "v"),
-        ]
-
-        loaded_stacked_params = []
-        other_weights = []
-        params_dict = dict(self.named_parameters())
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-
-                name = name.replace(weight_name, param_name)
-                if name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                loaded_stacked_params.append(name)
-                break
-            else:
-                if name in params_dict:
-                    other_weights.append((name, loaded_weight))
-
-        return other_weights, loaded_stacked_params
-
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        other_weights, loaded_stacked_params = self._load_weights(weights)
-
         loader = AutoWeightsLoader(self, skip_prefixes=["pooler."])
-        loaded_params = loader.load_weights(other_weights)
-        loaded_params.update(loaded_stacked_params)
-        return loaded_params
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 class BertPoolingModel(BertModel):
@@ -458,12 +432,8 @@ class BertPoolingModel(BertModel):
         self.pooler = BertPooler(vllm_config.model_config)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        other_weights, loaded_stacked_params = self._load_weights(weights)
-
         loader = AutoWeightsLoader(self)
-        loaded_params = loader.load_weights(other_weights)
-        loaded_params.update(loaded_stacked_params)
-        return loaded_params
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 @default_pooling_type(seq_pooling_type="CLS")

vllm/model_executor/models/blip.py

@@ -19,9 +19,9 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from .interfaces import SupportsQuant
+from .utils import AutoWeightsLoader, WeightsMapper
 
 
 def get_blip_patch_grid_length(*, image_size: int, patch_size: int) -> int:
@@ -268,6 +268,14 @@ class BlipVisionModel(nn.Module, SupportsQuant):
     main_input_name = "pixel_values"
     packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
 
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv", "q"),
+            ".k_proj": (".qkv", "k"),
+            ".v_proj": (".qkv", "v"),
+        }
+    )
+
     def __init__(
         self,
         config: BlipVisionConfig,
@@ -316,38 +324,18 @@ class BlipVisionModel(nn.Module, SupportsQuant):
         return self.post_layernorm(hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        layer_count = len(self.encoder.layers)
+        skip_prefixes: list[str] = []
+        if self.post_layernorm is None:
+            skip_prefixes.append("post_layernorm.")
+        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
 
-        for name, loaded_weight in weights:
-            # post_layernorm is not needed in BlipVisionModel
-            if name.startswith("post_layernorm") and self.post_layernorm is None:
-                continue
-
-            # omit layers when num_hidden_layers_override is set
-            if name.startswith("encoder.layers"):
-                layer_idx = int(name.split(".")[2])
-                if layer_idx >= layer_count:
+        # omit layers when num_hidden_layers_override is set
+        def _filter(ws):
+            for name, weight in ws:
+                if name.startswith("encoder.layers.") and int(
+                    name.split(".")[2]
+                ) >= len(self.encoder.layers):
                     continue
+                yield name, weight
 
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/chameleon.py

@@ -37,10 +37,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    row_parallel_weight_loader,
-)
+from vllm.model_executor.model_loader.weight_utils import row_parallel_weight_loader
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
@@ -66,7 +63,8 @@ from .interfaces import (
     SupportsQuant,
 )
 from .utils import (
-    is_pp_missing_parameter,
+    AutoWeightsLoader,
+    WeightsMapper,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
@@ -922,6 +920,16 @@ class ChameleonForConditionalGeneration(
         "gate_up_proj": ["gate_proj", "up_proj"],
     }
 
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+        },
+    )
+
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> str | None:
         if modality.startswith("image"):
@@ -1023,81 +1031,6 @@ class ChameleonForConditionalGeneration(
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-            (".gate_up_proj", ".gate_proj", 0),
-            (".gate_up_proj", ".up_proj", 1),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-
-            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
-                # Models trained using ColossalAI may include these tensors in
-                # the checkpoint. Skip them.
-                continue
-
-            # With tie_word_embeddings, we can skip lm_head.weight
-            # The weight might appear unnecessarily in the files if the model is
-            # processed with quantization, LoRA, fine-tuning, etc.
-            if self.config.tie_word_embeddings and "lm_head.weight" in name:
-                continue
-
-            use_default_weight_loading = False
-            if "vqmodel" in name:
-                if self.model.vqmodel is not None:
-                    # We only do sharding for language model and
-                    # not vqvae for now.
-                    use_default_weight_loading = True
-            else:
-                for param_name, weight_name, shard_id in stacked_params_mapping:
-                    if weight_name not in name:
-                        continue
-                    name = name.replace(weight_name, param_name)
-                    # Skip loading extra bias for GPTQ models.
-                    if name.endswith(".bias") and name not in params_dict:
-                        continue
-                    if is_pp_missing_parameter(name, self):
-                        continue
-                    param = params_dict[name]
-                    weight_loader = param.weight_loader
-                    weight_loader(param, loaded_weight, shard_id)
-                    break
-                else:
-                    # Skip loading extra bias for GPTQ models.
-                    if name.endswith(".bias") and name not in params_dict:
-                        continue
-                    # Remapping the name of FP8 kv-scale.
-                    if name.endswith("kv_scale"):
-                        remapped_kv_scale_name = name.replace(
-                            ".kv_scale", ".attn.kv_scale"
-                        )
-                        if remapped_kv_scale_name not in params_dict:
-                            logger.warning_once(
-                                "Found kv scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv-scale is not loaded.",  # noqa: E501
-                                name,
-                                remapped_kv_scale_name,
-                            )
-                            continue
-                        else:
-                            name = remapped_kv_scale_name
-                    if is_pp_missing_parameter(name, self):
-                        continue
-                    param = params_dict[name]
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight)
-            if use_default_weight_loading and name in params_dict:
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
+        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/clip.py

@@ -29,7 +29,6 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.layers.pooler import DispatchPooler
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import SupportsQuant
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
@@ -56,7 +55,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
 from .interfaces import MultiModalEmbeddings, SupportsMultiModal
 from .interfaces_base import default_pooling_type
-from .utils import AutoWeightsLoader, maybe_prefix
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 from .vision import (
     VisionEncoderInfo,
     VisionFeatureSelectStrategy,
@@ -555,6 +554,14 @@ class CLIPEncoder(nn.Module):
 
 
 class CLIPTextTransformer(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+        }
+    )
+
     def __init__(
         self,
         config: CLIPTextConfig,
@@ -605,34 +612,19 @@ class CLIPTextTransformer(nn.Module):
         return last_hidden_state
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 class CLIPVisionTransformer(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+        }
+    )
+
     def __init__(
         self,
         config: CLIPVisionConfig,
@@ -714,42 +706,21 @@ class CLIPVisionTransformer(nn.Module):
         return encoder_outputs
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        layer_count = len(self.encoder.layers)
+        skip_prefixes: list[str] = []
+        if self.post_layernorm is None:
+            skip_prefixes.append("post_layernorm.")
+        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
 
-        for name, loaded_weight in weights:
-            # post_layernorm is not needed in CLIPVisionModel
-            if name.startswith("post_layernorm") and self.post_layernorm is None:
-                continue
-
-            # omit layers when num_hidden_layers_override is set
-            if name.startswith("encoder.layers"):
-                layer_idx = int(name.split(".")[2])
-                if layer_idx >= layer_count:
+        # Drop layers beyond num_hidden_layers_override.
+        def _filter(ws):
+            for name, w in ws:
+                if name.startswith("encoder.layers.") and int(
+                    name.split(".")[2]
+                ) >= len(self.encoder.layers):
                     continue
+                yield name, w
 
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
 
 
 class CLIPVisionModel(nn.Module):

vllm/model_executor/models/cohere2_moe.py

@@ -18,7 +18,6 @@ from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
-    fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -31,8 +30,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
     row_parallel_weight_loader,
 )
 from vllm.model_executor.utils import set_weight_attrs
@@ -43,8 +40,8 @@ from .commandr import LayerNorm
 from .interfaces import SupportsPP, SupportsQuant
 from .utils import (
     AutoWeightsLoader,
+    WeightsMapper,
     extract_layer_index,
-    is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
@@ -462,88 +459,23 @@ class Cohere2MoeModel(nn.Module):
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        expert_params_mapping = fused_moe_make_expert_params_mapping(
-            self,
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.num_experts,
-        )
-
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-
-            for param_name, shard_name, shard_id in stacked_params_mapping:
-                if shard_name not in name:
-                    continue
-                if "mlp.experts" in name:
-                    continue
-                name = name.replace(shard_name, param_name)
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                for mapping in expert_params_mapping:
-                    param_name, weight_name, expert_id, shard_id = mapping
-                    if weight_name not in name:
-                        continue
-                    name = name.replace(weight_name, param_name)
-                    if is_pp_missing_parameter(name, self):
-                        continue
-                    if (
-                        name.endswith(".bias") or name.endswith("_bias")
-                    ) and name not in params_dict:
-                        continue
-                    param = params_dict[name]
-                    weight_loader = param.weight_loader
-                    weight_loader(
-                        param,
-                        loaded_weight,
-                        name,
-                        shard_id=shard_id,
-                        expert_id=expert_id,
-                    )
-                    break
-                else:
-                    if (
-                        name.endswith(".bias") or name.endswith("_bias")
-                    ) and name not in params_dict:
-                        continue
-                    if is_pp_missing_parameter(name, self):
-                        continue
-                    name = maybe_remap_kv_scale_name(name, params_dict)
-                    if name is None:
-                        continue
-                    param = params_dict[name]
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-
-        return loaded_params
-
 
 class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
     is_text_generation_model = True
 
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            # weight_name: (param_name, shard_id)
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
+            ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
+            ".mlp.up_proj": (".mlp.gate_up_proj", 1),
+            ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
+            ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
+        }
+    )
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",
@@ -599,4 +531,4 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
-        return loader.load_weights(weights)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/commandr.py

@@ -348,7 +348,11 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
             ".v_proj": (".qkv_proj", "v"),
             ".gate_proj": (".gate_up_proj", 0),
             ".up_proj": (".gate_up_proj", 1),
-        }
+        },
+        # ModelOpt NVFP4 checkpoints carry raw quantizer-module state
+        # (e.g. "*.weight_quantizer._double_scale"); drop them before loading.
+        # See #41925.
+        orig_to_new_substr={"_quantizer.": None},
     )
     packed_modules_mapping = {
         "qkv_proj": ["q_proj", "k_proj", "v_proj"],
@@ -356,9 +360,6 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
     }
     # LoRA specific attributes
     embedding_modules = {"embed_tokens": "input_embeddings"}
-    # ModelOpt NVFP4 checkpoints carry raw quantizer-module state
-    # (e.g. "*.weight_quantizer._double_scale"); drop them before loading. See #41925.
-    hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={"_quantizer.": None})
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()

vllm/model_executor/models/deepseek_eagle3.py

@@ -20,10 +20,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
 from vllm.model_executor.models.deepseek_v2 import (
     DeepseekV2ForCausalLM,
     DeepseekV2MLAAttention,
@@ -34,6 +30,7 @@ from vllm.multimodal.inputs import NestedTensors
 from .interfaces import LocalArgmaxMixin
 from .utils import (
     AutoWeightsLoader,
+    WeightsMapper,
     get_draft_quant_config,
     maybe_prefix,
     process_eagle_weight,
@@ -270,44 +267,20 @@ class DeepseekV2Eagle3Model(nn.Module):
 
         return hidden_states, aux_output
 
+    # midlayer rename + gate_up / MLA fused_qkv_a merges
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={"midlayer.": "layers.0."},
+        orig_to_new_stacked={
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+            ".q_a_proj": (".fused_qkv_a_proj", 0),
+            ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1),
+        },
+    )
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".gate_up_proj", ".gate_proj", 0),
-            (".gate_up_proj", ".up_proj", 1),
-            (".fused_qkv_a_proj", ".q_a_proj", 0),
-            (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            if "midlayer." in name:
-                name = name.replace("midlayer.", "layers.0.")
-
-            # Remapping the name FP8 kv-scale
-            if "scale" in name:
-                name = maybe_remap_kv_scale_name(name, params_dict)
-                if name is None:
-                    continue
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                if name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):

vllm/model_executor/models/dots_ocr.py

@@ -32,7 +32,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding.common import (
     ApplyRotaryEmb,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
     SupportsLoRA,
@@ -355,36 +354,6 @@ class DotsSwiGLUFFN(nn.Module):
         x, _ = self.fc2(x)
         return x
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            ("fc13", "fc1", 0),
-            ("fc13", "fc3", 1),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
-
 
 class DotsPatchEmbed(nn.Module):
     def __init__(self, config):
@@ -622,6 +591,10 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
             ".attn.qkv_proj.": ".attn.qkv.",
             ".attn.out_proj.": ".attn.proj.",
         },
+        orig_to_new_stacked={
+            ".fc1.": (".fc13.", 0),
+            ".fc3.": (".fc13.", 1),
+        },
         orig_to_new_prefix={
             "lm_head.": "language_model.lm_head.",
             "model.": "language_model.model.",

vllm/model_executor/models/ernie_mtp.py

@@ -27,7 +27,6 @@ from collections.abc import Iterable
 
 import torch
 import torch.nn as nn
-from transformers import PretrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -36,11 +35,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.sequence import IntermediateTensors
 
 from .llama import LlamaDecoderLayer
-from .utils import is_pp_missing_parameter, maybe_prefix
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 
 
 class ErnieMultiTokenPredictorLayer(nn.Module):
@@ -147,6 +145,28 @@ class ErnieMTP(nn.Module):
         super().__init__()
 
         self.config = vllm_config.model_config.hf_config
+        # MTP weights are stored under a flat `mtp_*.0.` block in the
+        # checkpoint; rewrite them into `model.layers.{spec_layer}.*`.
+        spec_layer = self.config.num_hidden_layers
+        self.hf_to_vllm_mapper = WeightsMapper(
+            orig_to_new_substr={
+                "model.mtp_emb_norm.0.": f"model.layers.{spec_layer}.mtp_emb_norm.",
+                "model.mtp_hidden_norm.0.": (
+                    f"model.layers.{spec_layer}.mtp_hidden_norm."
+                ),
+                "model.mtp_linear_proj.0.": (
+                    f"model.layers.{spec_layer}.mtp_linear_proj."
+                ),
+                "model.mtp_block.0.": f"model.layers.{spec_layer}.mtp_block.",
+            },
+            orig_to_new_stacked={
+                ".q_proj": (".qkv_proj", "q"),
+                ".k_proj": (".qkv_proj", "k"),
+                ".v_proj": (".qkv_proj", "v"),
+                ".gate_proj": (".gate_up_proj", 0),
+                ".up_proj": (".gate_up_proj", 1),
+            },
+        )
         self.model = ErnieMultiTokenPredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
@@ -185,94 +205,15 @@ class ErnieMTP(nn.Module):
         return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
+        # Checkpoint bundles the full base model; only MTP, embed_tokens and
+        # lm_head weights belong to this module.
+        def _filter(
+            weights: Iterable[tuple[str, torch.Tensor]],
+        ) -> Iterable[tuple[str, torch.Tensor]]:
+            for name, weight in weights:
+                if any(k in name for k in ("mtp", "embed_tokens", "lm_head")):
+                    yield name, weight
 
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if self.config.tie_word_embeddings and name.endswith("lm_head.weight"):
-                continue
-            if "rotary_emb.inv_freq" in name:
-                continue
-            if "mtp" in name:
-                name = self._rewrite_spec_layer_name(self.config, name)
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                # Skip non-stacked layers and experts (experts handled below).
-                if weight_name not in name:
-                    continue
-                if "mtp" not in name:
-                    continue
-                # We have mlp.experts[0].gate_proj in the checkpoint.
-                # Since we handle the experts below in expert_params_mapping,
-                # we need to skip here BEFORE we update the name, otherwise
-                # name will be updated to mlp.experts[0].gate_up_proj, which
-                # will then be updated below in expert_params_mapping
-                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
-                if ("mlp.experts." in name) and name not in params_dict:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if (
-                    name.endswith(".bias") or name.endswith("_bias")
-                ) and name not in params_dict:
-                    continue
-                # Skip layers on other devices.
-                if is_pp_missing_parameter(name, self):
-                    continue
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Skip loading extra bias for GPTQ models.
-                if (
-                    name.endswith(".bias") or name.endswith("_bias")
-                ) and name not in params_dict:
-                    continue
-                # Skip layers on other devices.
-                if is_pp_missing_parameter(name, self):
-                    continue
-
-                # According to DeepSeek-V3 Technical Report, MTP modules
-                # shares embedding layer. We only load the first weights.
-                if "mtp_" not in name and (
-                    "embed_tokens" not in name and "lm_head" not in name
-                ):
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
-
-    def _rewrite_spec_layer_name(self, config: PretrainedConfig, name: str) -> str:
-        """
-        Rewrite the weight name to match the format of the original model.
-        """
-        spec_layer_weight_names = [
-            "embed_tokens",
-            "mtp_emb_norm",
-            "mtp_hidden_norm",
-            "mtp_linear_proj",
-        ]
-        layer_idx = config.num_hidden_layers
-        for weight_name in spec_layer_weight_names:
-            if weight_name in name:
-                name = name.replace(
-                    f"model.{weight_name}.0.",
-                    f"model.layers.{layer_idx}.{weight_name}.",
-                )
-                return name
-        name = name.replace(
-            "model.mtp_block.0.", f"model.layers.{layer_idx}.mtp_block."
-        )
-        return name
+        skip_prefixes = ["lm_head"] if self.config.tie_word_embeddings else []
+        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/exaone_moe.py

@@ -15,8 +15,7 @@
 # limitations under the License.
 """Inference-only K-EXAONE-236B-A22B model compatible with HuggingFace weights."""
 
-import typing
-from collections.abc import Callable, Iterable
+from collections.abc import Iterable
 from itertools import islice
 
 import torch
@@ -32,7 +31,6 @@ from vllm.distributed import (
 )
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
-    fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import ReplicatedLinear
@@ -43,10 +41,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
 from vllm.sequence import IntermediateTensors
 
 from .exaone4 import Exaone4Attention as ExaoneMoeAttention
@@ -55,8 +49,8 @@ from .interfaces import SupportsLoRA, SupportsPP
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    WeightsMapper,
     extract_layer_index,
-    is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
@@ -325,142 +319,21 @@ class ExaoneMoeModel(nn.Module):
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
-    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        # Params for weights, fp8 weight scales, fp8 activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        return fused_moe_make_expert_params_mapping(
-            self,
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.num_experts,
-            num_redundant_experts=self.num_redundant_experts,
-        )
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-            (".gate_up_proj", ".gate_proj", 0),
-            (".gate_up_proj", ".up_proj", 1),
-        ]
-
-        # Skip loading extra parameters for GPTQ/modelopt models.
-        ignore_suffixes = (
-            ".bias",
-            "_bias",
-            ".k_scale",
-            "_k_scale",
-            ".v_scale",
-            "_v_scale",
-            ".weight_scale",
-            "_weight_scale",
-            ".input_scale",
-            "_input_scale",
-        )
-
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        expert_params_mapping = self.get_expert_mapping()
-        for name, loaded_weight in weights:
-            if name.startswith("mtp."):
-                continue
-            if "rotary_emb.inv_freq" in name:
-                continue
-            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
-                # Models trained using ColossalAI may include these tensors in
-                # the checkpoint. Skip them.
-                continue
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                if "mlp.experts" in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                is_expert_weight = False
-                for mapping in expert_params_mapping:
-                    param_name, weight_name, expert_id, shard_id = mapping
-                    if weight_name not in name:
-                        continue
-
-                    # Anyway, this is an expert weight and should not be
-                    # attempted to load as other weights later
-                    is_expert_weight = True
-
-                    # Do not modify `name` since the loop may continue here
-                    # Instead, create a new variable
-                    name_mapped = name.replace(weight_name, param_name)
-
-                    if is_pp_missing_parameter(name_mapped, self):
-                        continue
-
-                    # Skip loading extra parameters for GPTQ/modelopt models.
-                    if (
-                        name_mapped.endswith(ignore_suffixes)
-                        and name_mapped not in params_dict
-                    ):
-                        continue
-
-                    param = params_dict[name_mapped]
-                    # We should ask the weight loader to return success or not
-                    # here since otherwise we may skip experts with other
-                    # available replicas.
-                    weight_loader = typing.cast(
-                        Callable[..., bool], param.weight_loader
-                    )
-                    success = weight_loader(
-                        param,
-                        loaded_weight,
-                        name_mapped,
-                        shard_id=shard_id,
-                        expert_id=expert_id,
-                        return_success=True,
-                    )
-                    if success:
-                        name = name_mapped
-                        break
-                else:
-                    if is_expert_weight:
-                        continue
-                    # Skip loading extra bias for GPTQ models.
-                    if name.endswith(".bias") and name not in params_dict:
-                        continue
-                    # Skip loading extra parameters for GPTQ/modelopt models.
-                    if name.endswith(ignore_suffixes) and name not in params_dict:
-                        continue
-                    # Remapping the name of FP8 kv-scale.
-                    name = maybe_remap_kv_scale_name(name, params_dict)
-                    if name is None:
-                        continue
-
-                    if is_pp_missing_parameter(name, self):
-                        continue
-
-                    param = params_dict[name]
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
-
 
 class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            # weight_name: (param_name, shard_id)
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
+            ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
+            ".mlp.up_proj": (".mlp.gate_up_proj", 1),
+            ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
+            ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
+        }
+    )
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",
@@ -554,5 +427,18 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
             skip_prefixes=(
                 ["lm_head.", "mtp."] if self.config.tie_word_embeddings else ["mtp."]
             ),
+            # Skip loading extra parameters for GPTQ/modelopt models.
+            ignore_unexpected_suffixes=[
+                ".bias",
+                "_bias",
+                ".k_scale",
+                "_k_scale",
+                ".v_scale",
+                "_v_scale",
+                ".weight_scale",
+                "_weight_scale",
+                ".input_scale",
+                "_input_scale",
+            ],
         )
-        return loader.load_weights(weights)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/exaone_moe_mtp.py

@@ -18,15 +18,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.exaone_moe import ExaoneMoeDecoderLayer
 from vllm.sequence import IntermediateTensors
 
-from .utils import (
-    AutoWeightsLoader,
-    is_pp_missing_parameter,
-    maybe_prefix,
-)
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 
 logger = init_logger(__name__)
 
@@ -35,6 +30,17 @@ KVCache = tuple[torch.Tensor, torch.Tensor]
 
 @support_torch_compile
 class ExaoneMoeMultiTokenPredictor(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            # Scope to dense mlp; experts are handled separately.
+            ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
+            ".mlp.up_proj": (".mlp.gate_up_proj", 1),
+        }
+    )
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -128,55 +134,8 @@ class ExaoneMoeMultiTokenPredictor(nn.Module):
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        # Params for weights, fp8 weight scales, fp8 activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-
-                if "mlp.experts" in name:
-                    continue
-
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                # Skip layers on other devices.
-                if is_pp_missing_parameter(name, self):
-                    continue
-                if name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 @support_torch_compile

vllm/model_executor/models/extract_hidden_states.py

@@ -83,7 +83,10 @@ def basic_cache(
     kv_cache: torch.Tensor,  # shape: [num_blocks, block_size, num_heads, head_size]
     slot_mapping: torch.Tensor,  # shape: [seq_len]
 ):
+    # Padding slots are -1; redirect them to the null block (block 0, never
+    # allocated to a request) so the scatter stays branch-free and sync-free.
     block_size = kv_cache.shape[1]
+    slot_mapping = slot_mapping.clamp_min(0)
     kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache
 
 

vllm/model_executor/models/falcon_h1.py

@@ -35,10 +35,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.config import set_default_rope_theta
 
@@ -52,7 +48,7 @@ from .interfaces import (
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
-    is_pp_missing_parameter,
+    WeightsMapper,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
@@ -496,63 +492,6 @@ class FalconH1Model(nn.Module):
         hidden_states = self.final_layernorm(hidden_states)
         return hidden_states
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-
-            if "A_log" in name:
-                name = name.replace("A_log", "A")
-
-            if "mamba" in name:
-                name = name.replace("mamba", "mamba.mamba")
-
-            if "scale" in name:
-                # Remapping the name of kv-scale.
-                name = maybe_remap_kv_scale_name(name, params_dict)
-                if name is None:
-                    continue
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                # Skip layers on other devices.
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-
-        return loaded_params
-
 
 class FalconH1ForCausalLM(
     nn.Module,
@@ -567,6 +506,20 @@ class FalconH1ForCausalLM(
         "gate_up_proj": ["gate_proj", "up_proj"],
     }
 
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={
+            "A_log": "A",
+            "mamba": "mamba.mamba",
+        },
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+        },
+    )
+
     embedding_modules = {
         "embed_tokens": "input_embeddings",
         "lm_head": "output_embeddings",
@@ -694,4 +647,4 @@ class FalconH1ForCausalLM(
             self,
             skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
         )
-        return loader.load_weights(weights)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/funasr.py

@@ -31,7 +31,6 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.whisper_utils import (
     ISO639_1_SUPPORTED_LANGS,
 )
@@ -618,6 +617,14 @@ class FunASRAudioInputs(TensorSchema):
 
 
 class FunASREncoder(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".self_attn.q_proj": (".self_attn.qkv", "q"),
+            ".self_attn.k_proj": (".self_attn.qkv", "k"),
+            ".self_attn.v_proj": (".self_attn.qkv", "v"),
+        }
+    )
+
     def __init__(
         self, *, vllm_config: VllmConfig, prefix: str = "", init_in_fp32: bool = False
     ):
@@ -637,35 +644,8 @@ class FunASREncoder(nn.Module):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights with mapping from HuggingFace format."""
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("self_attn.qkv.", "self_attn.q_proj.", "q"),
-            ("self_attn.qkv.", "self_attn.k_proj.", "k"),
-            ("self_attn.qkv.", "self_attn.v_proj.", "v"),
-        ]
-        params_dict = dict(self.named_parameters(remove_duplicate=False))
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict.get(name)
-                if param is not None:
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 class FunASRModel(nn.Module):

vllm/model_executor/models/gemma3.py

@@ -291,6 +291,17 @@ class Gemma3DecoderLayer(nn.Module):
 
 @support_torch_compile
 class Gemma3Model(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            # weight_name: (param_name, shard_id)
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+        }
+    )
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         config = vllm_config.model_config.hf_config
@@ -361,18 +372,13 @@ class Gemma3Model(nn.Module):
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
 
 class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
-    hf_to_vllm_mapper = WeightsMapper(
-        orig_to_new_stacked={
-            # weight_name: (param_name, shard_id)
-            ".q_proj": (".qkv_proj", "q"),
-            ".k_proj": (".qkv_proj", "k"),
-            ".v_proj": (".qkv_proj", "v"),
-            ".gate_proj": (".gate_up_proj", 0),
-            ".up_proj": (".gate_up_proj", 1),
-        }
-    )
+    hf_to_vllm_mapper = Gemma3Model.hf_to_vllm_mapper
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",

vllm/model_executor/models/gemma4_mtp.py

@@ -44,7 +44,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.sequence import IntermediateTensors
 
 from .gemma4 import Gemma4MLP, _get_text_config
@@ -412,44 +411,6 @@ class Gemma4MultiTokenPredictor(nn.Module):
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids) * self.normalizer
 
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        params_dict = dict(self.named_parameters())
-        params_dict.update(dict(self.named_buffers()))
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-
-        return loaded_params
-
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -502,6 +463,10 @@ class Gemma4MTP(nn.Module):
             "pre_projection.": "model.pre_projection.",
             "post_projection.": "model.post_projection.",
         },
+        orig_to_new_stacked={
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+        },
     )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

vllm/model_executor/models/glm4_1v.py

@@ -76,7 +76,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.rotary_embedding.common import (
     ApplyRotaryEmb,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
@@ -618,6 +617,16 @@ class Glm4vVisionEmbeddings(nn.Module):
 
 
 class Glm4vVisionTransformer(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".attn.q.": (".attn.qkv.", "q"),
+            ".attn.k.": (".attn.qkv.", "k"),
+            ".attn.v.": (".attn.qkv.", "v"),
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+        }
+    )
+
     def __init__(
         self,
         text_config: Glm4vTextConfig,
@@ -958,33 +967,8 @@ class Glm4vVisionTransformer(nn.Module):
         return x
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("attn.qkv.", "attn.q.", "q"),
-            ("attn.qkv.", "attn.k.", "k"),
-            ("attn.qkv.", "attn.v.", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-        params_dict = dict(self.named_parameters(remove_duplicate=False))
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 class Glm4vProcessingInfo(BaseProcessingInfo):

vllm/model_executor/models/hunyuan_vision.py

@@ -48,7 +48,6 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
@@ -431,6 +430,14 @@ class HunYuanVisionPatchMerger(nn.Module):
 
 
 class HunYuanVisionTransformer(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv", "q"),
+            ".k_proj": (".qkv", "k"),
+            ".v_proj": (".qkv", "v"),
+        }
+    )
+
     def __init__(
         self,
         vision_config: HunYuanVLVisionConfig,
@@ -529,31 +536,8 @@ class HunYuanVisionTransformer(nn.Module):
         return image_embeds_list
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".qkv", ".q_proj", "q"),
-            (".qkv", ".k_proj", "k"),
-            (".qkv", ".v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters(remove_duplicate=False))
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 def _hunyuan_vl_field_config(hf_inputs: Mapping[str, torch.Tensor]):

vllm/model_executor/models/idefics2_vision_model.py

@@ -38,8 +38,8 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from .utils import AutoWeightsLoader, WeightsMapper
 from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model
 
 
@@ -352,6 +352,14 @@ class Idefics2Encoder(nn.Module):
 
 
 class Idefics2VisionTransformer(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+        }
+    )
+
     def __init__(
         self,
         config: Idefics2VisionConfig,
@@ -453,42 +461,22 @@ class Idefics2VisionTransformer(nn.Module):
         return last_hidden_state
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
+        # head is a pooling header absent from this model.
+        skip_prefixes = ["head."]
+        if not self.require_post_norm:
+            skip_prefixes.append("post_layernorm.")
+        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+
         layer_count = len(self.encoder.layers)
 
-        for name, loaded_weight in weights:
-            # skip pooling header
-            if name.startswith("head."):
-                continue
-
-            # post_layernorm is optional
-            if name.startswith("post_layernorm.") and not self.require_post_norm:
-                continue
-
-            # omit layers when num_hidden_layers_override is set
-            if name.startswith("encoder.layers."):
-                layer_idx = int(name.split(".")[2])
-                if layer_idx >= layer_count:
+        def _filter(ws: Iterable[tuple[str, torch.Tensor]]):
+            # Drop layers beyond num_hidden_layers_override.
+            for name, w in ws:
+                if (
+                    name.startswith("encoder.layers.")
+                    and int(name.split(".")[2]) >= layer_count
+                ):
                     continue
+                yield name, w
 
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name or self.use_data_parallel:
-                    continue
-                name = name.replace(weight_name, param_name)
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/isaac.py

@@ -25,9 +25,6 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-)
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
     SupportsLoRA,
@@ -652,6 +649,14 @@ class Siglip2Encoder(nn.Module):
 
 
 class Siglip2VisionTransformer(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+        }
+    )
+
     def __init__(
         self,
         config: PixelShuffleSiglip2VisionConfig,
@@ -715,31 +720,8 @@ class Siglip2VisionTransformer(nn.Module):
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 def _resolve_vision_token_id(model_config: ModelConfig, vision_token: str) -> int:

vllm/model_executor/models/jina.py

@@ -25,7 +25,7 @@ from ..layers.pooler.tokwise import (
 from .interfaces import SupportsLateInteraction
 from .interfaces_base import VllmModelForPooling
 from .qwen3 import Qwen3ForCausalLM, Qwen3Model
-from .utils import AutoWeightsLoader, maybe_prefix
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 
 logger = logging.getLogger(__name__)
 
@@ -193,6 +193,9 @@ class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling):
     """
 
     is_pooling_model = True
+    hf_to_vllm_mapper = Qwen3ForCausalLM.hf_to_vllm_mapper | WeightsMapper(
+        orig_to_new_prefix={"": "model."}
+    )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__(vllm_config=vllm_config, prefix=prefix)
@@ -254,6 +257,6 @@ class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling):
                         tensor = tensor + (lora_B @ lora_A) * scaling
                 yield name, tensor
 
-        loader = AutoWeightsLoader(self.model, ignore_unexpected_prefixes=["lm_head."])
+        loader = AutoWeightsLoader(self, ignore_unexpected_prefixes=["lm_head."])
         weights = _merge_weights(weights)
-        return loader.load_weights(weights, mapper=self.model.hf_to_vllm_mapper)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/keye.py

@@ -34,10 +34,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding.common import (
     ApplyRotaryEmb,
 )
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
@@ -76,7 +72,6 @@ from .utils import (
     AutoWeightsLoader,
     WeightsMapper,
     init_vllm_registered_model,
-    is_pp_missing_parameter,
     maybe_prefix,
 )
 from .vision import is_vit_use_data_parallel
@@ -718,6 +713,14 @@ class KeyeSiglipVisionModel(nn.Module):
     config_class = PretrainedConfig
     main_input_name = "pixel_values"
 
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+        }
+    )
+
     def __init__(
         self,
         config: PretrainedConfig,
@@ -776,53 +779,8 @@ class KeyeSiglipVisionModel(nn.Module):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters(remove_duplicate=False))
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-            if "head.attention" in name or "head.layernorm" in name:
-                continue
-            if "head.mlp" in name or "head.probe" in name:
-                continue
-            for (
-                param_name,
-                weight_name,
-                shard_id,
-            ) in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                name = maybe_remap_kv_scale_name(name, params_dict)
-                if name is None:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = getattr(
-                    param,
-                    "weight_loader",
-                    default_weight_loader,
-                )
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self, skip_prefixes=["vision_model.head."])
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 class Projector(nn.Module):

vllm/model_executor/models/kimi_audio.py

@@ -17,7 +17,6 @@ from vllm.config.multimodal import BaseDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.inputs import PromptType, TokensPrompt
 from vllm.model_executor.model_loader import DefaultModelLoader
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import (
     SupportsMultiModal,
     SupportsPP,
@@ -81,6 +80,14 @@ class KimiAudioWhisperEncoder(WhisperEncoder):
         "qkv_proj": ["q_proj", "k_proj", "v_proj"],
     }
 
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),
+            ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),
+            ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),
+        },
+    )
+
     def __init__(
         self, *, vllm_config: VllmConfig, prefix: str = "", init_in_fp32: bool = False
     ):
@@ -102,37 +109,8 @@ class KimiAudioWhisperEncoder(WhisperEncoder):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 # -----------------------------------------------------------------------------

vllm/model_executor/models/lfm2.py

@@ -32,7 +32,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.sequence import IntermediateTensors
 
 from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP, SupportsQuant
@@ -41,7 +40,6 @@ from .utils import (
     PPMissingLayer,
     WeightsMapper,
     extract_layer_index,
-    is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
@@ -297,6 +295,20 @@ class Lfm2ShortConvDecoderLayer(nn.Module):
 
 @support_torch_compile
 class Lfm2Model(nn.Module):
+    # HF uses .conv. but vLLM uses .short_conv. to avoid LoRA regex collision
+    # with the inner .conv.conv child (ShortConv has a child self.conv, so
+    # naming the container .conv too makes _match_target_modules match both).
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={".conv.": ".short_conv."},
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            ".w1": (".w13", 0),
+            ".w3": (".w13", 1),
+        },
+    )
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -375,40 +387,8 @@ class Lfm2Model(nn.Module):
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-            (".w13", ".w1", 0),
-            (".w13", ".w3", 1),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if ".conv." in name:
-                name = name.replace(".conv.", ".short_conv.", 1)
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                # Use segment-boundary matching (trailing dot) to prevent
-                # e.g. ".w1" from matching inside ".w13" in pre-fused keys.
-                if weight_name + "." not in name:
-                    continue
-                name = name.replace(weight_name + ".", param_name + ".")
-
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
 class Lfm2ForCausalLM(
@@ -427,12 +407,8 @@ class Lfm2ForCausalLM(
         "in_proj": ["in_proj"],
     }
 
-    # HF uses .conv. but vLLM uses .short_conv. to avoid LoRA regex collision
-    # with the inner .conv.conv child (ShortConv has a child self.conv, so
-    # naming the container .conv too makes _match_target_modules match both)
-    hf_to_vllm_mapper = WeightsMapper(
-        orig_to_new_substr={".conv.": ".short_conv."},
-    )
+    # Reuse the backbone mapper so LoRA/quantization see the same name mapping.
+    hf_to_vllm_mapper = Lfm2Model.hf_to_vllm_mapper
 
     # LoRA specific attributes
     embedding_modules = {

vllm/model_executor/models/lfm2_siglip2.py

@@ -23,9 +23,8 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
-from .utils import maybe_prefix
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
 from .vision import (
     is_vit_use_data_parallel,
     resolve_visual_encoder_outputs,
@@ -458,6 +457,14 @@ class Siglip2VisionTransformer(nn.Module):
 
 
 class Siglip2Model(torch.nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+        }
+    )
+
     def __init__(
         self,
         config: Siglip2VisionConfig,
@@ -501,42 +508,21 @@ class Siglip2Model(torch.nn.Module):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
+        skip_prefixes = []
+        if self.vision_model.post_layernorm is None:
+            skip_prefixes.append("vision_model.post_layernorm.")
+        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+
+        # Drop layers omitted by num_hidden_layers_override.
         layer_count = len(self.vision_model.encoder.layers)
 
-        for name, loaded_weight in weights:
-            # post_layernorm is optional in Siglip2Model
-            if (
-                name.startswith("vision_model.post_layernorm")
-                and self.vision_model.post_layernorm is None
-            ):
-                continue
-
-            # omit layers when num_hidden_layers_override is set
-            if name.startswith("vision_model.encoder.layers"):
-                layer_idx = int(name.split(".")[3])
-                if layer_idx >= layer_count:
+        def _filter(ws):
+            for n, w in ws:
+                if (
+                    n.startswith("vision_model.encoder.layers.")
+                    and int(n.split(".")[3]) >= layer_count
+                ):
                     continue
+                yield n, w
 
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/llama4_eagle.py

@@ -32,18 +32,33 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.llama4 import Llama4DecoderLayer, Llama4ForCausalLM
 from vllm.model_executor.models.utils import extract_layer_index
 
 from .interfaces import SupportsMultiModal
-from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
+from .utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    maybe_prefix,
+    process_eagle_weight,
+)
 
 logger = init_logger(__name__)
 
 
 @support_torch_compile
 class LlamaModel(nn.Module):
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={"model.": ""},
+        orig_to_new_stacked={
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
+            ".gate_proj": (".gate_up_proj", 0),
+            ".up_proj": (".gate_up_proj", 1),
+        },
+    )
+
     def __init__(
         self,
         *,
@@ -108,34 +123,8 @@ class LlamaModel(nn.Module):
         return hidden_states, hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-            (".gate_up_proj", ".gate_proj", 0),
-            (".gate_up_proj", ".up_proj", 1),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            name = name.removeprefix("model.")
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        for name in params_dict:
-            assert name in loaded_params, f"{name} is not loaded!"
-        return loaded_params
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def validate_and_update_config(
         self, start_layer_id: int, quant_config: QuantizationConfig | None = None