stash

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

.buildkite/intel_jobs/lora_intel.yaml

@@ -81,7 +81,7 @@ steps:
       'cd tests &&
       export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
       set -o pipefail &&
-      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-32000-32-4-4]" --deselect="tests/lora/test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-32000-32-4-4]"'
+      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]"'
 
 - label: LoRA Punica FP8/XPU Ops
   timeout_in_minutes: 45

examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py

@@ -327,9 +327,6 @@ async def handle_request(api: str, request: Request):
         session, decode_response = await decode_request_task
         stream_generator = stream_decode_response(session, decode_response, request_id)
         response = await make_response(stream_generator)
-        response.headers["Content-Type"] = decode_response.headers.get(
-            "Content-Type", "application/json"
-        )
         return response
     except Exception as e:
         logger.exception("An error occurred while handling the request: %s", e)

rust/Cargo.lock

@@ -5099,6 +5099,7 @@ dependencies = [
  "uuid",
  "vllm-engine-core-client",
  "vllm-llm",
+ "vllm-model-files",
  "vllm-parser",
  "vllm-text",
  "vllm-tokenizer",
@@ -5236,6 +5237,20 @@ dependencies = [
  "zeromq",
 ]
 
+[[package]]
+name = "vllm-model-files"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "hf-hub",
+ "serde",
+ "serde_json",
+ "tempfile",
+ "thiserror 2.0.18",
+ "thiserror-ext",
+ "tokio",
+]
+
 [[package]]
 name = "vllm-parser"
 version = "0.1.0"
@@ -5324,7 +5339,6 @@ dependencies = [
  "enum-as-inner",
  "expect-test",
  "futures",
- "hf-hub",
  "itertools 0.14.0",
  "reqwest",
  "serde",
@@ -5339,6 +5353,7 @@ dependencies = [
  "trait-set",
  "vllm-engine-core-client",
  "vllm-llm",
+ "vllm-model-files",
  "vllm-tokenizer",
 ]
 

rust/Cargo.toml

@@ -6,6 +6,7 @@ members = [
     "src/llm",
     "src/managed-engine",
     "src/metrics",
+    "src/model-files",
     "src/mock-engine",
     "src/parser",
     "src/parser/python",
@@ -126,6 +127,7 @@ vllm-chat = { path = "src/chat" }
 vllm-engine-core-client = { path = "src/engine-core-client" }
 vllm-llm = { path = "src/llm" }
 vllm-managed-engine = { path = "src/managed-engine" }
+vllm-model-files = { path = "src/model-files" }
 vllm-metrics = { path = "src/metrics" }
 vllm-parser = { path = "src/parser" }
 vllm-server = { path = "src/server" }

rust/src/chat/Cargo.toml

@@ -32,6 +32,7 @@ trait-set.workspace = true
 uuid.workspace = true
 vllm-engine-core-client.workspace = true
 vllm-llm.workspace = true
+vllm-model-files.workspace = true
 vllm-parser.workspace = true
 vllm-text.workspace = true
 vllm-tokenizer.workspace = true

rust/src/chat/src/error.rs

@@ -1,5 +1,6 @@
 use thiserror::Error;
 use thiserror_ext::Macro;
+use vllm_model_files::Error as ModelFilesError;
 
 type BoxedError = Box<dyn std::error::Error + Send + Sync>;
 
@@ -69,6 +70,8 @@ pub enum Error {
     #[error(transparent)]
     Text(#[from] vllm_text::Error),
     #[error(transparent)]
+    ModelFiles(#[from] ModelFilesError),
+    #[error(transparent)]
     Tokenizer(#[from] vllm_tokenizer::TokenizerError),
 }
 

rust/src/chat/src/lib.rs

@@ -50,8 +50,7 @@ mod request;
 mod stream;
 
 use vllm_engine_core_client::EngineCoreClient;
-use vllm_engine_core_client::protocol::dtype::ModelDtype;
-use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
+use vllm_engine_core_client::protocol::{ModelDtype, ReasoningParserKwargs};
 use vllm_llm::Llm;
 use vllm_text::{Prompt, TextLlm, TextRequest};
 

rust/src/chat/src/multimodal.rs

@@ -22,7 +22,7 @@ use llm_multimodal::{
     TrackedMedia,
 };
 use tracing::warn;
-use vllm_engine_core_client::protocol::dtype::ModelDtype;
+use vllm_engine_core_client::protocol::ModelDtype;
 use vllm_engine_core_client::protocol::multimodal::{
     MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem,
     MmSharedField, MmSlice, PlaceholderRange, SliceSpec,

rust/src/chat/src/multimodal/tensor.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use half::{bf16, f16};
 use llm_multimodal::{ModelSpecificValue, PreprocessedImages};
-use vllm_engine_core_client::protocol::dtype::ModelDtype;
+use vllm_engine_core_client::protocol::ModelDtype;
 use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
 use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
 

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

@@ -1,9 +1,7 @@
 //! Applies xgrammar structural-tag constraints for strict tool calling.
 
 use thiserror_ext::AsReport;
-use vllm_engine_core_client::protocol::structured_outputs::{
-    StructuredOutputBackend, StructuredOutputsParams,
-};
+use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
 use vllm_parser::tool::StructuralTagModel;
 use xgrammar_structural_tag::{
     FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
@@ -78,9 +76,7 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option<StructuralTagTool
 #[cfg(test)]
 mod tests {
     use serde_json::{Value, json};
-    use vllm_engine_core_client::protocol::structured_outputs::{
-        StructuredOutputBackend, StructuredOutputsParams,
-    };
+    use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
     use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser};
 
     use super::*;

rust/src/chat/src/renderer/hf/mod.rs

@@ -6,7 +6,7 @@ use thiserror_ext::AsReport as _;
 use tracing::{info, trace, warn};
 use vllm_text::Prompt;
 use vllm_text::backend::hf::{
-    HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config,
+    ResolvedModelFiles, SpecialTokens, TokenizerConfig, load_tokenizer_config,
 };
 
 use self::format::{
@@ -42,7 +42,7 @@ pub struct HfChatRenderer {
     default_template: Option<CompiledChatTemplate>,
     default_template_kwargs: HashMap<String, JsonValue>,
     content_format: ContentFormatOption,
-    special_tokens: Option<HfSpecialTokens>,
+    special_tokens: Option<SpecialTokens>,
     multimodal: Option<MultimodalRenderInfo>,
 }
 
@@ -67,7 +67,7 @@ impl HfChatRenderer {
         })
     }
 
-    pub fn with_special_tokens(mut self, special_tokens: Option<HfSpecialTokens>) -> Self {
+    pub fn with_special_tokens(mut self, special_tokens: Option<SpecialTokens>) -> Self {
         self.special_tokens = special_tokens;
         self
     }
@@ -83,7 +83,7 @@ impl HfChatRenderer {
         options: LoadModelBackendsOptions,
         multimodal: Option<MultimodalRenderInfo>,
     ) -> Result<Self> {
-        let HfTokenizerConfig {
+        let TokenizerConfig {
             special_tokens,
             chat_template,
             ..
@@ -451,7 +451,7 @@ mod tests {
     use expect_test::expect;
     use serde_json::Value;
     use vllm_text::Prompt;
-    use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
+    use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens};
 
     use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo};
     use crate::request::{
@@ -675,7 +675,7 @@ mod tests {
     #[test]
     fn chat_template_injects_special_tokens_into_context() {
         let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
-        let special_tokens = HfSpecialTokens {
+        let special_tokens = SpecialTokens {
             bos_token: Some(NamedSpecialToken::Text("<bos>".to_string())),
             ..Default::default()
         };

rust/src/chat/src/renderer/hf/template.rs

@@ -2,7 +2,7 @@
 //!
 //! This module is inlined from SMG's tokenizer crate with local adaptations:
 //! - thinking-related detection/state is removed
-//! - special tokens are wired to `vllm_text::backends::hf::HfSpecialTokens`
+//! - special tokens are wired to `vllm_text::backends::hf::SpecialTokens`
 
 use std::collections::HashMap;
 use std::fs;
@@ -11,7 +11,7 @@ use std::path::Path;
 use minijinja::Environment;
 use serde::{Deserialize, Serialize};
 use serde_json::{self};
-use vllm_text::backend::hf::HfSpecialTokens;
+use vllm_text::backend::hf::SpecialTokens;
 
 use super::error::TemplateError;
 use super::format::{
@@ -46,7 +46,7 @@ pub(super) struct TemplateContext<'a> {
     pub(super) tools: Option<&'a [TemplateTool]>,
     pub(super) documents: Option<&'a [serde_json::Value]>,
     #[serde(flatten)]
-    pub(super) special_tokens: Option<&'a HfSpecialTokens>,
+    pub(super) special_tokens: Option<&'a SpecialTokens>,
     #[serde(flatten)]
     pub(super) template_kwargs: Option<&'a HashMap<String, serde_json::Value>>,
 }
@@ -133,7 +133,7 @@ mod tests {
     use std::fs;
 
     use tempfile::TempDir;
-    use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
+    use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens};
 
     use super::*;
 
@@ -170,7 +170,7 @@ mod tests {
             CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
                 .unwrap();
 
-        let special_tokens = HfSpecialTokens {
+        let special_tokens = SpecialTokens {
             bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
             eos_token: Some(NamedSpecialToken::Text("</s>".to_string())),
             ..Default::default()
@@ -205,7 +205,7 @@ mod tests {
             CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
                 .unwrap();
 
-        let special_tokens = HfSpecialTokens {
+        let special_tokens = SpecialTokens {
             bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
             eos_token: None,
             ..Default::default()

rust/src/chat/tests/chat.rs

@@ -15,10 +15,9 @@ use vllm_chat::{
 use vllm_engine_core_client::protocol::logprobs::{
     Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
 };
-use vllm_engine_core_client::protocol::output::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, StopReason,
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
 };
-use vllm_engine_core_client::protocol::request::EngineCoreRequest;
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
 use vllm_llm::Llm;

rust/src/engine-core-client/examples/external_engine_logprobs.rs

@@ -5,9 +5,9 @@ use clap::Parser;
 use futures::StreamExt as _;
 use tokio::time::timeout;
 use tracing_subscriber::EnvFilter;
-use vllm_engine_core_client::protocol::output::EngineCoreFinishReason;
-use vllm_engine_core_client::protocol::request::EngineCoreRequest;
-use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
+};
 use vllm_engine_core_client::{
     EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode,
 };

rust/src/engine-core-client/src/client.rs

@@ -11,11 +11,10 @@ use tracing::{debug, info, trace};
 use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop};
 use crate::coordinator::CoordinatorHandle;
 use crate::error::{Error, Result};
-use crate::protocol::dtype::ModelDtype;
 use crate::protocol::handshake::EngineCoreReadyResponse;
 use crate::protocol::lora::LoraRequest;
-use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
 use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
+use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
 use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};
 use crate::transport::{self, ConnectedEngine};
 

rust/src/engine-core-client/src/client/imp.rs

@@ -16,11 +16,12 @@ use crate::client::stream::EngineCoreStreamOutput;
 use crate::client::{AbortCause, AbortRequest};
 use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
 use crate::metrics::{LoraInfoExporter, record_scheduler_stats};
-use crate::protocol::encode_msgpack;
-use crate::protocol::output::{ClassifiedEngineCoreOutputs, EngineCoreOutput, EngineCoreOutputs};
-use crate::protocol::request::EngineCoreRequestType;
 use crate::protocol::stats::SchedulerStats;
 use crate::protocol::utility::UtilityOutput;
+use crate::protocol::{
+    ClassifiedEngineCoreOutputs, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequestType,
+    encode_msgpack,
+};
 use crate::transport::{ConnectedEngine, EngineId};
 use crate::{Error, Result, transport};
 

rust/src/engine-core-client/src/client/state.rs

@@ -7,9 +7,9 @@ use tracing::trace;
 use crate::EngineId;
 use crate::client::stream::EngineCoreStreamOutput;
 use crate::error::{Error, Result};
-use crate::protocol::output::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput};
 use crate::protocol::stats::SchedulerStats;
 use crate::protocol::utility::UtilityOutput;
+use crate::protocol::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput};
 use crate::transport::ConnectedEngine;
 
 pub type OutputSender = mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>;
@@ -452,7 +452,7 @@ mod tests {
         EngineLoadSnapshot, EngineRoutingState, RequestRegistry, UtilityRegistry,
     };
     use crate::mock_engine::default_ready_response;
-    use crate::protocol::output::{
+    use crate::protocol::{
         EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
     };
     use crate::transport::ConnectedEngine;

rust/src/engine-core-client/src/client/stream.rs

@@ -10,7 +10,7 @@ use tracing::{debug, error, warn};
 
 use crate::client::AbortRequest;
 use crate::client::state::OutputReceiver;
-use crate::protocol::output::{EngineCoreFinishReason, EngineCoreOutput};
+use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput};
 use crate::{AbortCause, Error, Result};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

rust/src/engine-core-client/src/coordinator/inproc.rs

@@ -10,9 +10,10 @@ use zeromq::{XPubSocket, ZmqMessage};
 use crate::client::imp::ClientInner;
 use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState};
 use crate::error::{Error, Result, bail_unexpected_coordinator_output};
-use crate::protocol::encode_msgpack;
-use crate::protocol::output::{ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs};
-use crate::protocol::request::EngineCoreRequestType;
+use crate::protocol::{
+    ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs, EngineCoreRequestType,
+    encode_msgpack,
+};
 
 /// Coordinator-to-engine `START_DP_WAVE` control payload encoded on the
 /// engine-facing coordinator socket.

rust/src/engine-core-client/src/mock_engine.rs

@@ -8,9 +8,8 @@ use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
 
 use crate::EngineId;
 use crate::error::{Error, Result, bail_unexpected_handshake_message};
-use crate::protocol::dtype::ModelDtype;
 use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
-use crate::protocol::{decode_msgpack, encode_msgpack};
+use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack};
 
 /// Default model length advertised by reusable mock engine helpers.
 pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024;

rust/src/engine-core-client/src/protocol/classified_outputs.rs

@@ -1,164 +1,10 @@
 use std::collections::BTreeSet;
 
 use enum_as_inner::EnumAsInner;
-use serde::{Deserialize, Serialize};
-use serde_default::DefaultFromSerde;
-use serde_repr::{Deserialize_repr, Serialize_repr};
-use serde_tuple::{Deserialize_tuple, Serialize_tuple};
 
 use super::utility::UtilityOutput;
-use crate::error::{Error, Result, ext_value_decode};
-use crate::protocol::logprobs::MaybeWireLogprobs;
-use crate::protocol::stats::{PrefillStats, SchedulerStats};
-use crate::protocol::{OpaqueValue, decode_msgpack};
-
-/// The stop reason associated with a finished output.
-///
-/// Python models this as the union-typed `stop_reason: int | str | None`
-/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum.
-///
-/// Original Python field:
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L155>
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
-#[serde(untagged)]
-pub enum StopReason {
-    TokenId(u32),
-    Text(String),
-}
-
-/// Reason a request finished: stop, length, abort, error, or repetition.
-///
-/// This mirrors the Python enum and uses integer encoding for compact wire
-/// representation.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L41-L63>
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
-#[repr(u8)]
-pub enum EngineCoreFinishReason {
-    /// A stop string was emitted.
-    Stop = 0,
-    /// `max_tokens` or `max_model_len` was reached.
-    Length = 1,
-    /// The request was aborted by the client.
-    Abort = 2,
-    /// A retryable request-level internal error occurred.
-    Error = 3,
-    /// A repetitive token pattern was detected.
-    Repetition = 4,
-}
-
-/// Event types emitted by engine-core for one request.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L113-L118>
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
-#[repr(u8)]
-pub enum EngineCoreEventType {
-    Queued = 1,
-    Scheduled = 2,
-    Preempted = 3,
-}
-
-/// A timestamped engine-core event associated with one request.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L121-L130>
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct EngineCoreEvent {
-    pub r#type: EngineCoreEventType,
-    pub timestamp: f64,
-}
-
-/// Engine-core output for a single request.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/engine/__init__.py#L154-L184>
-#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
-pub struct EngineCoreOutput {
-    pub request_id: String,
-    pub new_token_ids: Vec<u32>,
-    /// Decoded sample logprobs for the newly generated positions in this
-    /// output.
-    #[serde(default)]
-    pub new_logprobs: Option<MaybeWireLogprobs>,
-    /// Decoded prompt logprobs for the scored prompt positions emitted in this
-    /// output.
-    #[serde(default)]
-    pub new_prompt_logprobs_tensors: Option<MaybeWireLogprobs>,
-    #[serde(default)]
-    pub pooling_output: Option<OpaqueValue>,
-    #[serde(default)]
-    pub finish_reason: Option<EngineCoreFinishReason>,
-    #[serde(default)]
-    pub stop_reason: Option<StopReason>,
-    #[serde(default)]
-    pub events: Option<Vec<EngineCoreEvent>>,
-    #[serde(default)]
-    pub kv_transfer_params: Option<serde_json::Value>,
-    #[serde(default)]
-    pub trace_headers: Option<OpaqueValue>,
-    /// Breakdown of the scheduled prefill computation, set on the first output
-    /// of a newly scheduled prefill and elided for subsequent decode outputs.
-    #[serde(default)]
-    pub prefill_stats: Option<PrefillStats>,
-    #[serde(default)]
-    pub routed_experts: Option<OpaqueValue>,
-    /// Number of NaNs seen in logits. Values above zero indicate corruption.
-    #[serde(default)]
-    pub num_nans_in_logits: u32,
-}
-
-impl EngineCoreOutput {
-    /// Returns whether this output is terminal for the request.
-    pub fn finished(&self) -> bool {
-        self.finish_reason.is_some()
-    }
-
-    /// Resolve all wire-format fields in-place by looking up aux frames and
-    /// decoding raw-view payloads as needed.
-    fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
-    where
-        Frame: AsRef<[u8]>,
-    {
-        self.new_logprobs = (self.new_logprobs.take())
-            .map(|value| value.resolve(frames, "new_logprobs"))
-            .transpose()?;
-        self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
-            .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
-            .transpose()?;
-        Ok(())
-    }
-}
-
-/// Batch of engine-core outputs returned to a frontend client.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
-#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
-pub struct EngineCoreOutputs {
-    #[serde(default)]
-    pub engine_index: u32,
-    /// Outputs grouped for this client in the current engine tick.
-    #[serde(default)]
-    pub outputs: Vec<EngineCoreOutput>,
-    #[serde(default)]
-    pub scheduler_stats: Option<Box<SchedulerStats>>,
-    #[serde(default)]
-    pub timestamp: f64,
-    #[serde(default)]
-    pub utility_output: Option<UtilityOutput>,
-    #[serde(default)]
-    pub finished_requests: Option<BTreeSet<String>>,
-    /// In DP mode, signals that the current wave finished and engines are
-    /// paused.
-    #[serde(default)]
-    pub wave_complete: Option<u32>,
-    /// In DP mode, signals that a request arrived for an old wave and the next
-    /// wave needs to start in other engines.
-    #[serde(default)]
-    pub start_wave: Option<u32>,
-}
+use super::{EngineCoreOutput, EngineCoreOutputs};
+use crate::protocol::stats::SchedulerStats;
 
 /// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -203,18 +49,6 @@ pub enum ClassifiedEngineCoreOutputs {
 }
 
 impl EngineCoreOutputs {
-    /// Resolve all wire-format fields in-place by looking up aux frames and
-    /// decoding raw-view payloads as needed.
-    fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
-    where
-        Frame: AsRef<[u8]>,
-    {
-        for output in &mut self.outputs {
-            output.resolve_in_place(frames)?;
-        }
-        Ok(())
-    }
-
     /// Classify the raw wire message into a more semantic Rust enum.
     pub fn classify(self) -> ClassifiedEngineCoreOutputs {
         let has_request_payload = !self.outputs.is_empty()
@@ -258,62 +92,12 @@ impl EngineCoreOutputs {
     }
 }
 
-/// Decode one ordinary or multipart engine-core output message into the strong
-/// typed public protocol shape.
-pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
-where
-    Frame: AsRef<[u8]>,
-{
-    let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
-
-    let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
-    outputs.resolve_in_place(frames)?;
-    Ok(outputs)
-}
-
 #[cfg(test)]
 mod tests {
     use std::collections::BTreeSet;
 
     use super::*;
-    use crate::protocol::output::EngineCoreOutput;
-    use crate::protocol::{decode_msgpack, encode_msgpack};
-
-    #[test]
-    fn engine_core_outputs_roundtrip_finished_fields() {
-        let outputs = EngineCoreOutputs {
-            outputs: vec![EngineCoreOutput {
-                request_id: "req-1".to_string(),
-                new_token_ids: vec![42],
-                new_logprobs: None,
-                new_prompt_logprobs_tensors: None,
-                pooling_output: None,
-                finish_reason: Some(EngineCoreFinishReason::Length),
-                stop_reason: Some(StopReason::Text("stop".to_string())),
-                events: None,
-                kv_transfer_params: None,
-                trace_headers: None,
-                prefill_stats: None,
-                routed_experts: None,
-                num_nans_in_logits: 0,
-            }],
-            finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
-            ..Default::default()
-        };
-
-        let encoded = encode_msgpack(&outputs).unwrap();
-        let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap();
-
-        assert_eq!(decoded.outputs.len(), 1);
-        assert_eq!(
-            decoded.outputs[0].finish_reason,
-            Some(EngineCoreFinishReason::Length)
-        );
-        assert_eq!(
-            decoded.finished_requests,
-            Some(BTreeSet::from(["req-1".to_string()]))
-        );
-    }
+    use crate::protocol::EngineCoreOutput;
 
     #[test]
     fn engine_core_outputs_classify_request_batch() {

rust/src/engine-core-client/src/protocol/handshake.rs

@@ -2,8 +2,7 @@ use std::collections::BTreeMap;
 
 use serde::{Deserialize, Serialize};
 
-use crate::protocol::OpaqueValue;
-use crate::protocol::dtype::ModelDtype;
+use crate::protocol::{ModelDtype, OpaqueValue};
 
 /// Decoded engine startup-handshake payload sent on the handshake socket.
 ///

rust/src/engine-core-client/src/protocol/logprobs.rs

@@ -9,7 +9,8 @@ use enum_as_inner::EnumAsInner;
 use serde::{Deserialize, Deserializer, Serialize};
 
 use self::wire::*;
-use crate::error::{Error, Result, bail_ext_value_decode};
+use super::{EngineCoreOutput, EngineCoreOutputs, decode_msgpack};
+use crate::error::{Error, Result, bail_ext_value_decode, ext_value_decode};
 use crate::protocol::tensor::{WireArrayData, WireNdArray};
 
 /// One token candidate and its logprob metadata for a single sequence position.
@@ -159,7 +160,7 @@ impl Serialize for MaybeWireLogprobs {
 impl MaybeWireLogprobs {
     /// Resolve the wire representation into decoded logprobs by looking up aux
     /// frames and decoding raw views as needed.
-    pub(super) fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
+    fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
     where
         Frame: AsRef<[u8]>,
     {
@@ -170,6 +171,37 @@ impl MaybeWireLogprobs {
     }
 }
 
+impl EngineCoreOutputs {
+    /// Resolve all wire-format fields in-place by looking up aux frames and
+    /// decoding raw-view payloads as needed.
+    fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
+    where
+        Frame: AsRef<[u8]>,
+    {
+        for output in &mut self.outputs {
+            output.resolve_in_place(frames)?;
+        }
+        Ok(())
+    }
+}
+
+impl EngineCoreOutput {
+    /// Resolve all wire-format fields in-place by looking up aux frames and
+    /// decoding raw-view payloads as needed.
+    fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
+    where
+        Frame: AsRef<[u8]>,
+    {
+        self.new_logprobs = (self.new_logprobs.take())
+            .map(|value| value.resolve(frames, "new_logprobs"))
+            .transpose()?;
+        self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
+            .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
+            .transpose()?;
+        Ok(())
+    }
+}
+
 impl WireLogprobs {
     /// Convert semantic per-position logprobs into the Python wire tuple shape.
     ///
@@ -283,3 +315,16 @@ impl WireLogprobs {
         Ok(Logprobs { positions })
     }
 }
+
+/// Decode one ordinary or multipart engine-core output message into the strong
+/// typed public protocol shape.
+pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
+where
+    Frame: AsRef<[u8]>,
+{
+    let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
+
+    let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
+    outputs.resolve_in_place(frames)?;
+    Ok(outputs)
+}

rust/src/engine-core-client/src/protocol/logprobs/tests.rs

@@ -3,8 +3,8 @@ use std::collections::BTreeSet;
 use bytes::Bytes;
 use rmpv::Value;
 
-use super::{Logprobs, PositionLogprobs, TokenLogprob};
-use crate::protocol::output::{EngineCoreFinishReason, decode_engine_core_outputs};
+use super::{Logprobs, PositionLogprobs, TokenLogprob, decode_engine_core_outputs};
+use crate::protocol::EngineCoreFinishReason;
 
 fn encode_value(value: &Value) -> Vec<u8> {
     let mut out = Vec::new();

rust/src/engine-core-client/src/protocol/mod.rs

@@ -1,11 +1,28 @@
 use std::any::type_name;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::io::Cursor;
 
+use bytes::Bytes;
 use rmpv::Value;
 use serde::{Deserialize, Serialize};
+use serde_default::DefaultFromSerde;
+use serde_repr::{Deserialize_repr, Serialize_repr};
+use serde_tuple::{Deserialize_tuple, Serialize_tuple};
 use thiserror_ext::AsReport;
 
 use crate::error::{Error, Result};
+use crate::protocol::logprobs::MaybeWireLogprobs;
+use crate::protocol::multimodal::MmFeatures;
+use crate::protocol::stats::{PrefillStats, SchedulerStats};
+use crate::protocol::utility::UtilityOutput;
+
+// TODO: This module currently mixes reusable frontend-facing semantic types
+// (for example `FinishReason`, `StopReason`, `RequestOutputKind`, and future
+// cleaned-up frontend sampling types) with engine-core-specific wire DTOs and
+// handshake/control messages. While the Rust frontend is still evolving
+// quickly, keep them co-located here for iteration speed. Once the higher-level
+// API boundary stabilizes, move the truly reusable semantic types into a
+// lower-level common crate and keep the engine transport/wire messages here.
 
 /// Dynamic msgpack value used for schema positions that are preserved but not
 /// yet strongly typed in the early-stage Rust client.
@@ -19,18 +36,499 @@ fn is_false(v: &bool) -> bool {
     !v
 }
 
+fn default_top_p() -> f32 {
+    1.0
+}
+
+fn default_repetition_penalty() -> f32 {
+    1.0
+}
+
+fn default_temperature() -> f32 {
+    1.0
+}
+
+fn default_max_tokens() -> u32 {
+    16
+}
+
+mod classified_outputs;
 pub mod dtype;
 pub mod handshake;
 pub mod logprobs;
 pub mod lora;
 pub mod multimodal;
-pub mod output;
-pub mod request;
-pub mod sampling;
 pub mod stats;
-pub mod structured_outputs;
 pub mod tensor;
 pub mod utility;
+pub use classified_outputs::{
+    ClassifiedEngineCoreOutputs, DpControlMessage, RequestBatchOutputs, UtilityCallOutput,
+};
+pub use dtype::ModelDtype;
+pub use logprobs::decode_engine_core_outputs;
+
+/// Request types are encoded as single-byte protocol constants so they can be
+/// sent over the ZMQ socket without an extra encoding step.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L217-L228>
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[repr(u8)]
+pub enum EngineCoreRequestType {
+    Add = 0,
+    Abort = 1,
+    StartDpWave = 2,
+    Utility = 3,
+}
+
+impl EngineCoreRequestType {
+    /// Decode the single-byte request type frame used on the engine input
+    /// socket. Returns `None` for unrecognized values.
+    pub fn from_frame(frame: &[u8]) -> Option<Self> {
+        let [value] = frame else {
+            return None;
+        };
+
+        match value {
+            0 => Some(Self::Add),
+            1 => Some(Self::Abort),
+            2 => Some(Self::StartDpWave),
+            3 => Some(Self::Utility),
+            _ => None,
+        }
+    }
+
+    /// Encode the request type as the single-byte frame used on the engine
+    /// input socket.
+    pub fn to_frame(self) -> Bytes {
+        Bytes::from_static(match self {
+            Self::Add => b"\x00",
+            Self::Abort => b"\x01",
+            Self::StartDpWave => b"\x02",
+            Self::Utility => b"\x03",
+        })
+    }
+}
+
+/// Reason a request finished: stop, length, abort, error, or repetition.
+///
+/// This mirrors the Python enum and uses integer encoding for compact wire
+/// representation.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L41-L63>
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
+#[repr(u8)]
+pub enum EngineCoreFinishReason {
+    /// A stop string was emitted.
+    Stop = 0,
+    /// `max_tokens` or `max_model_len` was reached.
+    Length = 1,
+    /// The request was aborted by the client.
+    Abort = 2,
+    /// A retryable request-level internal error occurred.
+    Error = 3,
+    /// A repetitive token pattern was detected.
+    Repetition = 4,
+}
+
+/// Event types emitted by engine-core for one request.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L113-L118>
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
+#[repr(u8)]
+pub enum EngineCoreEventType {
+    Queued = 1,
+    Scheduled = 2,
+    Preempted = 3,
+}
+
+/// A timestamped engine-core event associated with one request.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L121-L130>
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct EngineCoreEvent {
+    pub r#type: EngineCoreEventType,
+    pub timestamp: f64,
+}
+
+/// Controls how intermediate outputs are returned to the frontend.
+///
+/// `Cumulative = 0` is intentionally not supported in Rust frontend.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L146-L152>
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize_repr, Deserialize_repr)]
+#[repr(u8)]
+pub enum RequestOutputKind {
+    /// Return only token deltas in each update.
+    #[default]
+    Delta = 1,
+    /// Suppress intermediate updates and return only the final output.
+    FinalOnly = 2,
+}
+
+/// Structured-output backend selected for EngineCore grammar compilation.
+///
+/// Python vLLM stores this in `StructuredOutputsParams._backend` after request
+/// validation. The Rust frontend currently always lowers structured-output
+/// requests to guidance, while ignoring any user-supplied `_backend` value.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub enum StructuredOutputBackend {
+    Xgrammar,
+    #[default]
+    Guidance,
+    Outlines,
+    LmFormatEnforcer,
+}
+
+/// The stop reason associated with a finished output.
+///
+/// Python models this as the union-typed `stop_reason: int | str | None`
+/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum.
+///
+/// Original Python field:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L155>
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum StopReason {
+    TokenId(u32),
+    Text(String),
+}
+
+/// Parameters for configuring structured outputs (guided decoding).
+///
+/// Exactly one constraint field (`json`, `regex`, `choice`, `grammar`,
+/// `json_object`, or `structural_tag`) should be set. The engine-core
+/// backend selects the appropriate grammar compiler based on which field
+/// is present.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
+#[serde(default)]
+pub struct StructuredOutputsParams {
+    /// JSON schema (as a dict/object or JSON string) constraining the output.
+    pub json: Option<serde_json::Value>,
+    /// Regular expression the output must match.
+    pub regex: Option<String>,
+    /// List of allowed output strings (the model must produce one of these).
+    pub choice: Option<Vec<String>>,
+    /// Context-free grammar (in EBNF-like notation) the output must conform to.
+    pub grammar: Option<String>,
+    /// When `true`, output must be valid JSON (free-form, no schema).
+    pub json_object: Option<bool>,
+    /// Disable any additional whitespace in guided JSON output.
+    #[serde(skip_serializing_if = "crate::protocol::is_false")]
+    pub disable_any_whitespace: bool,
+    /// Disable `additionalProperties` in JSON schema output.
+    #[serde(skip_serializing_if = "crate::protocol::is_false")]
+    pub disable_additional_properties: bool,
+    /// Custom whitespace pattern for guided JSON output.
+    pub whitespace_pattern: Option<String>,
+    /// Structural tag configuration (JSON-encoded string).
+    pub structural_tag: Option<String>,
+    /// Structured-output backend, mirroring Python's internal `_backend`.
+    ///
+    /// User-supplied values are ignored during deserialization. This matches
+    /// Python's request boundary, where `_backend` is set by validation rather
+    /// than accepted as a request-level backend selector.
+    #[serde(
+        default,
+        rename = "_backend",
+        deserialize_with = "serde_with::rust::deserialize_ignore_any"
+    )]
+    pub backend: StructuredOutputBackend,
+}
+
+/// Engine-core-facing sampling parameters for text generation.
+///
+/// This is the normalized southbound subset used by the Rust frontend when it
+/// talks to Python engine-core over the wire. User-facing request semantics
+/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are
+/// intentionally handled by higher layers before values reach this DTO.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
+// Python's SamplingParams is `omit_defaults=True`, so msgpack drops
+// default-valued keys; default the whole struct. Per-field fns cover the
+// non-zero defaults.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)]
+#[serde(default)]
+pub struct EngineCoreSamplingParams {
+    /// Controls randomness. Lower values are more deterministic; zero means
+    /// greedy sampling.
+    #[serde(default = "default_temperature")]
+    pub temperature: f32,
+    /// Cumulative probability threshold for nucleus sampling.
+    #[serde(default = "default_top_p")]
+    pub top_p: f32,
+    /// Maximum number of top tokens to consider. `0` means all tokens.
+    pub top_k: u32,
+    /// Random seed used by the sampler when present.
+    pub seed: Option<i64>,
+    /// Maximum number of tokens to generate per output sequence.
+    #[serde(default = "default_max_tokens")]
+    pub max_tokens: u32,
+    /// Minimum number of tokens to generate before EOS or stop-token handling.
+    pub min_tokens: u32,
+    /// Maximum number of reasoning ("thinking") tokens to emit before the
+    /// reasoning section is force-closed. `None` means unlimited; the
+    /// user-facing `-1` sentinel is normalized to `None` by the frontend before
+    /// reaching this DTO, so only non-negative values are sent. Enforced
+    /// engine-side (and only when a reasoning parser is configured).
+    pub thinking_token_budget: Option<u64>,
+    /// Number of log probabilities to return per generated token.
+    ///
+    /// `None` disables sample logprobs. `-1` requests the full vocabulary.
+    pub logprobs: Option<i32>,
+    /// Number of log probabilities to return per prompt token.
+    ///
+    /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
+    pub prompt_logprobs: Option<i32>,
+    /// Minimum probability threshold for token sampling.
+    pub min_p: f32,
+    /// Frequency penalty applied by the sampler.
+    pub frequency_penalty: f32,
+    /// Presence penalty applied by the sampler.
+    pub presence_penalty: f32,
+    /// Repetition penalty applied by the sampler.
+    #[serde(default = "default_repetition_penalty")]
+    pub repetition_penalty: f32,
+    /// Token IDs that stop generation.
+    pub stop_token_ids: Vec<u32>,
+    /// Primary EOS token ID used by engine-core's dedicated EOS stop path.
+    ///
+    /// This mirrors Python's internal `_eos_token_id` field and is derived by
+    /// the frontend from tokenizer/model metadata rather than supplied directly
+    /// by end users.
+    #[serde(rename = "_eos_token_id")]
+    pub eos_token_id: Option<u32>,
+    /// Complete stop-token set used by engine-core for `min_tokens` masking.
+    ///
+    /// This mirrors Python's internal `_all_stop_token_ids` field and should
+    /// contain explicit `stop_token_ids` plus any frontend-derived EOS token
+    /// IDs.
+    #[serde(rename = "_all_stop_token_ids")]
+    pub all_stop_token_ids: BTreeSet<u32>,
+    /// Logit biases to apply during sampling.
+    /// Keys are token IDs
+    pub logit_bias: Option<HashMap<u32, f32>>,
+    /// Restrict output to these token IDs only.
+    pub allowed_token_ids: Option<Vec<u32>>,
+    /// Tokenized bad words to avoid during generation.
+    #[serde(rename = "_bad_words_token_ids")]
+    pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
+    /// Parameters for configuring structured outputs (guided decoding).
+    pub structured_outputs: Option<StructuredOutputsParams>,
+    /// Specific token IDs for which log probabilities should be returned at
+    /// each position.
+    ///
+    /// When set, the engine returns logprobs for exactly these tokens in
+    /// addition to the sampled/scored token. Mutually exclusive with the
+    /// `logprobs` count field in practice.
+    pub logprob_token_ids: Option<Vec<u32>>,
+    /// If `Some(true)`, the request will not attempt to read from the prefix
+    /// cache; newly computed blocks may still populate the cache. `None`
+    /// defers to engine-core defaults.
+    pub skip_reading_prefix_cache: Option<bool>,
+    /// Additional request parameters for custom extensions (from `vllm_xargs`).
+    pub extra_args: Option<HashMap<String, serde_json::Value>>,
+}
+
+impl EngineCoreSamplingParams {
+    /// Constructs a default sampling params for testing purposes only.
+    pub fn for_test() -> Self {
+        Self {
+            temperature: 1.0,
+            top_p: 1.0,
+            top_k: 0,
+            seed: None,
+            max_tokens: 65536,
+            min_tokens: 0,
+            thinking_token_budget: None,
+            logprobs: None,
+            prompt_logprobs: None,
+            min_p: 0.0,
+            frequency_penalty: 0.0,
+            presence_penalty: 0.0,
+            repetition_penalty: 1.0,
+            stop_token_ids: Vec::new(),
+            eos_token_id: None,
+            all_stop_token_ids: BTreeSet::new(),
+            logit_bias: None,
+            allowed_token_ids: None,
+            bad_words_token_ids: None,
+            structured_outputs: None,
+            logprob_token_ids: None,
+            skip_reading_prefix_cache: None,
+            extra_args: None,
+        }
+    }
+}
+
+/// Extra kwargs consumed by engine-side reasoning parsers.
+///
+/// Original Python construction point:
+/// <https://github.com/vllm-project/vllm/blob/cec2ec11760f9f3beabd4c90451936078bf91533/vllm/entrypoints/openai/chat_completion/serving.py#L367-L369>
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct ReasoningParserKwargs {
+    /// Effective kwargs visible to the chat template for this request.
+    pub chat_template_kwargs: HashMap<String, serde_json::Value>,
+}
+
+/// Engine-core add-request payload sent from frontend to engine.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/3f5bd482f5c1a5dbdffbbf68d624e20bb7032013/vllm/v1/engine/__init__.py#L80-L129>
+#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
+pub struct EngineCoreRequest {
+    pub request_id: String,
+    pub prompt_token_ids: Option<Vec<u32>>,
+    /// Multimodal features attached to the request.
+    pub mm_features: Option<MmFeatures>,
+    pub sampling_params: Option<EngineCoreSamplingParams>,
+    /// Pooling parameters are preserved in the schema but not yet strongly
+    /// typed.
+    pub pooling_params: Option<OpaqueValue>,
+    pub arrival_time: f64,
+    #[serde(default)]
+    pub lora_request: Option<lora::LoraRequest>,
+    #[serde(default)]
+    pub cache_salt: Option<String>,
+    #[serde(default)]
+    pub data_parallel_rank: Option<u32>,
+    /// Unsupported in the first-stage Rust client because Python uses a custom
+    /// tensor/aux-frame encoding path for this field.
+    #[serde(default)]
+    pub prompt_embeds: Option<OpaqueValue>,
+    /// Per-position mask for mixed-mode inputs (e.g. chat completion with
+    /// `prompt_embeds` content parts). `Some(true)` means real token id;
+    /// `Some(false)` means the position uses a pre-computed entry from
+    /// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
+    #[serde(default)]
+    pub prompt_is_token_ids: Option<Vec<bool>>,
+    /// Index of the client, used to ensure outputs are sent back to the same
+    /// client when scaling out the frontend.
+    #[serde(default)]
+    pub client_index: u32,
+    /// In DP mode, indicates which wave this request is expected to belong to.
+    #[serde(default)]
+    pub current_wave: u32,
+    #[serde(default)]
+    pub priority: i32,
+    #[serde(default)]
+    pub trace_headers: Option<BTreeMap<String, String>>,
+    #[serde(default)]
+    pub resumable: bool,
+    /// Original user-provided request ID, used for output reporting and aborts.
+    #[serde(default)]
+    pub external_req_id: Option<String>,
+    #[serde(default)]
+    pub reasoning_ended: Option<bool>,
+    /// Reasoning-parser kwargs forwarded from the frontend to the
+    /// structured-output backend.
+    #[serde(default)]
+    pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,
+    /// If `true`, the request should be added to the scheduler's waiting queue
+    /// and immediately aborted, so connector-side cleanup runs via the
+    /// standard `request_finished` hook.
+    #[serde(default)]
+    pub abort_immediately: bool,
+}
+
+impl EngineCoreRequest {
+    /// Validate fields intentionally not supported in the first-stage client.
+    pub fn validate(&self) -> Result<()> {
+        if self.prompt_embeds.is_some() {
+            return Err(Error::UnsupportedField {
+                context: "EngineCoreRequest",
+                field: "prompt_embeds",
+            });
+        }
+        Ok(())
+    }
+}
+
+/// Engine-core output for a single request.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/engine/__init__.py#L154-L184>
+#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
+pub struct EngineCoreOutput {
+    pub request_id: String,
+    pub new_token_ids: Vec<u32>,
+    /// Decoded sample logprobs for the newly generated positions in this
+    /// output.
+    #[serde(default)]
+    pub new_logprobs: Option<MaybeWireLogprobs>,
+    /// Decoded prompt logprobs for the scored prompt positions emitted in this
+    /// output.
+    #[serde(default)]
+    pub new_prompt_logprobs_tensors: Option<MaybeWireLogprobs>,
+    #[serde(default)]
+    pub pooling_output: Option<OpaqueValue>,
+    #[serde(default)]
+    pub finish_reason: Option<EngineCoreFinishReason>,
+    #[serde(default)]
+    pub stop_reason: Option<StopReason>,
+    #[serde(default)]
+    pub events: Option<Vec<EngineCoreEvent>>,
+    #[serde(default)]
+    pub kv_transfer_params: Option<serde_json::Value>,
+    #[serde(default)]
+    pub trace_headers: Option<OpaqueValue>,
+    /// Breakdown of the scheduled prefill computation, set on the first output
+    /// of a newly scheduled prefill and elided for subsequent decode outputs.
+    #[serde(default)]
+    pub prefill_stats: Option<PrefillStats>,
+    #[serde(default)]
+    pub routed_experts: Option<OpaqueValue>,
+    /// Number of NaNs seen in logits. Values above zero indicate corruption.
+    #[serde(default)]
+    pub num_nans_in_logits: u32,
+}
+
+impl EngineCoreOutput {
+    /// Returns whether this output is terminal for the request.
+    pub fn finished(&self) -> bool {
+        self.finish_reason.is_some()
+    }
+}
+
+/// Batch of engine-core outputs returned to a frontend client.
+///
+/// Original Python definition:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
+#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
+pub struct EngineCoreOutputs {
+    #[serde(default)]
+    pub engine_index: u32,
+    /// Outputs grouped for this client in the current engine tick.
+    #[serde(default)]
+    pub outputs: Vec<EngineCoreOutput>,
+    #[serde(default)]
+    pub scheduler_stats: Option<Box<SchedulerStats>>,
+    #[serde(default)]
+    pub timestamp: f64,
+    #[serde(default)]
+    pub utility_output: Option<UtilityOutput>,
+    #[serde(default)]
+    pub finished_requests: Option<BTreeSet<String>>,
+    /// In DP mode, signals that the current wave finished and engines are
+    /// paused.
+    #[serde(default)]
+    pub wave_complete: Option<u32>,
+    /// In DP mode, signals that a request arrived for an old wave and the next
+    /// wave needs to start in other engines.
+    #[serde(default)]
+    pub start_wave: Option<u32>,
+}
 
 /// Encode a Rust value into msgpack using the protocol crate's serde model.
 pub fn encode_msgpack<T>(value: &T) -> Result<Vec<u8>>
@@ -66,17 +564,81 @@ where
     })
 }
 
-/// Decode a msgpack payload into a dynamic value for diagnostics and tests.
 pub fn decode_value(bytes: &[u8]) -> Result<Value> {
     Ok(rmpv::decode::read_value(&mut Cursor::new(bytes))?)
 }
 
 #[cfg(test)]
 mod tests {
-    use std::collections::BTreeMap;
+    use std::collections::BTreeSet;
 
     use super::*;
 
+    #[test]
+    fn engine_core_request_serializes_as_full_array() {
+        let request = EngineCoreRequest {
+            request_id: "req-1".to_string(),
+            prompt_token_ids: Some(vec![1, 2, 3]),
+            sampling_params: Some(EngineCoreSamplingParams {
+                max_tokens: 8,
+                ..EngineCoreSamplingParams::for_test()
+            }),
+            arrival_time: 1234.5,
+            client_index: 7,
+            ..EngineCoreRequest::default()
+        };
+
+        let encoded = encode_msgpack(&request).unwrap();
+        let value = decode_value(&encoded).unwrap();
+        let array = match value {
+            Value::Array(array) => array,
+            other => panic!("expected array, got {other:?}"),
+        };
+
+        assert_eq!(array.len(), 20);
+        assert_eq!(array[0], Value::from("req-1"));
+        assert_eq!(array[2], Value::Nil);
+        assert_eq!(array[4], Value::Nil);
+        assert_eq!(array[10], Value::Nil);
+        assert_eq!(array[11], Value::from(7));
+    }
+
+    #[test]
+    fn engine_core_outputs_roundtrip_finished_fields() {
+        let outputs = EngineCoreOutputs {
+            outputs: vec![EngineCoreOutput {
+                request_id: "req-1".to_string(),
+                new_token_ids: vec![42],
+                new_logprobs: None,
+                new_prompt_logprobs_tensors: None,
+                pooling_output: None,
+                finish_reason: Some(EngineCoreFinishReason::Length),
+                stop_reason: Some(StopReason::Text("stop".to_string())),
+                events: None,
+                kv_transfer_params: None,
+                trace_headers: None,
+                prefill_stats: None,
+                routed_experts: None,
+                num_nans_in_logits: 0,
+            }],
+            finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
+            ..Default::default()
+        };
+
+        let encoded = encode_msgpack(&outputs).unwrap();
+        let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap();
+
+        assert_eq!(decoded.outputs.len(), 1);
+        assert_eq!(
+            decoded.outputs[0].finish_reason,
+            Some(EngineCoreFinishReason::Length)
+        );
+        assert_eq!(
+            decoded.finished_requests,
+            Some(BTreeSet::from(["req-1".to_string()]))
+        );
+    }
+
     #[test]
     fn decode_msgpack_includes_type_name_and_value_fallback() {
         let error = decode_msgpack::<u64>(
@@ -86,4 +648,72 @@ mod tests {
 
         expect_test::expect![[r#"messagepack decode failed for u64: wrong msgpack marker FixMap(1); value fallback: {"status": "READY"}"#]].assert_eq(&error.to_report_string());
     }
+
+    #[test]
+    fn structured_outputs_backend_ignores_deserialized_value() {
+        let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({
+            "json_object": true,
+            "_backend": "xgrammar",
+        }))
+        .unwrap();
+
+        assert_eq!(params.backend, StructuredOutputBackend::Guidance);
+
+        let value = serde_json::to_value(params).unwrap();
+        assert_eq!(value["_backend"], "guidance");
+    }
+
+    /// A real `sampling_params` is a sparse `omit_defaults` map; absent fields
+    /// must fall back to defaults. `python_compat` can't catch this since Rust
+    /// encodes full maps (see `engine_core_request_serializes_as_full_array`).
+    #[test]
+    fn decodes_sampling_params_with_omitted_defaults() {
+        let sampling_params = Value::Map(vec![
+            (
+                Value::from("stop_token_ids"),
+                Value::Array(vec![Value::from(151643u32)]),
+            ),
+            (Value::from("skip_reading_prefix_cache"), Value::from(false)),
+        ]);
+        let request = Value::Array(vec![
+            Value::from("req-omit-defaults"),
+            Value::Array(vec![
+                Value::from(1u32),
+                Value::from(2u32),
+                Value::from(3u32),
+            ]),
+            Value::Nil,
+            sampling_params,
+            Value::Nil,
+            Value::from(1.0f64),
+        ]);
+
+        let mut bytes = Vec::new();
+        rmpv::encode::write_value(&mut bytes, &request).unwrap();
+
+        let decoded: EngineCoreRequest = decode_msgpack(&bytes)
+            .expect("a real omit_defaults request must decode (regression: missing field)");
+
+        assert_eq!(decoded.request_id, "req-omit-defaults");
+        let sampling = decoded.sampling_params.expect("sampling params present");
+
+        assert_eq!(sampling.stop_token_ids, vec![151643]);
+        assert_eq!(sampling.skip_reading_prefix_cache, Some(false));
+
+        // Omitted fields -> Python defaults.
+        assert_eq!(sampling.temperature, 1.0);
+        assert_eq!(sampling.top_p, 1.0);
+        assert_eq!(sampling.top_k, 0);
+        assert_eq!(sampling.seed, None);
+        assert_eq!(sampling.max_tokens, 16);
+        assert_eq!(sampling.min_tokens, 0);
+        assert_eq!(sampling.min_p, 0.0);
+        assert_eq!(sampling.frequency_penalty, 0.0);
+        assert_eq!(sampling.presence_penalty, 0.0);
+        assert_eq!(sampling.repetition_penalty, 1.0);
+        assert_eq!(sampling.logprobs, None);
+        assert_eq!(sampling.prompt_logprobs, None);
+        assert_eq!(sampling.eos_token_id, None);
+        assert!(sampling.all_stop_token_ids.is_empty());
+    }
 }

rust/src/engine-core-client/src/protocol/request.rs

@@ -1,175 +0,0 @@
-use std::collections::{BTreeMap, HashMap};
-
-use bytes::Bytes;
-use serde::{Deserialize, Serialize};
-use serde_default::DefaultFromSerde;
-use serde_tuple::{Deserialize_tuple, Serialize_tuple};
-
-use crate::protocol::multimodal::MmFeatures;
-use crate::protocol::sampling::EngineCoreSamplingParams;
-use crate::protocol::{OpaqueValue, lora};
-use crate::{Error, Result};
-
-/// Request types are encoded as single-byte protocol constants so they can be
-/// sent over the ZMQ socket without an extra encoding step.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L217-L228>
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-#[repr(u8)]
-pub enum EngineCoreRequestType {
-    Add = 0,
-    Abort = 1,
-    StartDpWave = 2,
-    Utility = 3,
-}
-
-impl EngineCoreRequestType {
-    /// Decode the single-byte request type frame used on the engine input
-    /// socket. Returns `None` for unrecognized values.
-    pub fn from_frame(frame: &[u8]) -> Option<Self> {
-        let [value] = frame else {
-            return None;
-        };
-
-        match value {
-            0 => Some(Self::Add),
-            1 => Some(Self::Abort),
-            2 => Some(Self::StartDpWave),
-            3 => Some(Self::Utility),
-            _ => None,
-        }
-    }
-
-    /// Encode the request type as the single-byte frame used on the engine
-    /// input socket.
-    pub fn to_frame(self) -> Bytes {
-        Bytes::from_static(match self {
-            Self::Add => b"\x00",
-            Self::Abort => b"\x01",
-            Self::StartDpWave => b"\x02",
-            Self::Utility => b"\x03",
-        })
-    }
-}
-
-/// Extra kwargs consumed by engine-side reasoning parsers.
-///
-/// Original Python construction point:
-/// <https://github.com/vllm-project/vllm/blob/cec2ec11760f9f3beabd4c90451936078bf91533/vllm/entrypoints/openai/chat_completion/serving.py#L367-L369>
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct ReasoningParserKwargs {
-    /// Effective kwargs visible to the chat template for this request.
-    pub chat_template_kwargs: HashMap<String, serde_json::Value>,
-}
-
-/// Engine-core add-request payload sent from frontend to engine.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/3f5bd482f5c1a5dbdffbbf68d624e20bb7032013/vllm/v1/engine/__init__.py#L80-L129>
-#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
-pub struct EngineCoreRequest {
-    pub request_id: String,
-    pub prompt_token_ids: Option<Vec<u32>>,
-    /// Multimodal features attached to the request.
-    pub mm_features: Option<MmFeatures>,
-    pub sampling_params: Option<EngineCoreSamplingParams>,
-    /// Pooling parameters are preserved in the schema but not yet strongly
-    /// typed.
-    pub pooling_params: Option<OpaqueValue>,
-    pub arrival_time: f64,
-    #[serde(default)]
-    pub lora_request: Option<lora::LoraRequest>,
-    #[serde(default)]
-    pub cache_salt: Option<String>,
-    #[serde(default)]
-    pub data_parallel_rank: Option<u32>,
-    /// Unsupported in the first-stage Rust client because Python uses a custom
-    /// tensor/aux-frame encoding path for this field.
-    #[serde(default)]
-    pub prompt_embeds: Option<OpaqueValue>,
-    /// Per-position mask for mixed-mode inputs (e.g. chat completion with
-    /// `prompt_embeds` content parts). `Some(true)` means real token id;
-    /// `Some(false)` means the position uses a pre-computed entry from
-    /// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
-    #[serde(default)]
-    pub prompt_is_token_ids: Option<Vec<bool>>,
-    /// Index of the client, used to ensure outputs are sent back to the same
-    /// client when scaling out the frontend.
-    #[serde(default)]
-    pub client_index: u32,
-    /// In DP mode, indicates which wave this request is expected to belong to.
-    #[serde(default)]
-    pub current_wave: u32,
-    #[serde(default)]
-    pub priority: i32,
-    #[serde(default)]
-    pub trace_headers: Option<BTreeMap<String, String>>,
-    #[serde(default)]
-    pub resumable: bool,
-    /// Original user-provided request ID, used for output reporting and aborts.
-    #[serde(default)]
-    pub external_req_id: Option<String>,
-    #[serde(default)]
-    pub reasoning_ended: Option<bool>,
-    /// Reasoning-parser kwargs forwarded from the frontend to the
-    /// structured-output backend.
-    #[serde(default)]
-    pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,
-    /// If `true`, the request should be added to the scheduler's waiting queue
-    /// and immediately aborted, so connector-side cleanup runs via the
-    /// standard `request_finished` hook.
-    #[serde(default)]
-    pub abort_immediately: bool,
-}
-
-impl EngineCoreRequest {
-    /// Validate fields intentionally not supported in the first-stage client.
-    pub fn validate(&self) -> Result<()> {
-        if self.prompt_embeds.is_some() {
-            return Err(Error::UnsupportedField {
-                context: "EngineCoreRequest",
-                field: "prompt_embeds",
-            });
-        }
-        Ok(())
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use rmpv::Value;
-
-    use super::*;
-    use crate::protocol::sampling::EngineCoreSamplingParams;
-    use crate::protocol::{decode_value, encode_msgpack};
-
-    #[test]
-    fn engine_core_request_serializes_as_full_array() {
-        let request = EngineCoreRequest {
-            request_id: "req-1".to_string(),
-            prompt_token_ids: Some(vec![1, 2, 3]),
-            sampling_params: Some(EngineCoreSamplingParams {
-                max_tokens: 8,
-                ..EngineCoreSamplingParams::for_test()
-            }),
-            arrival_time: 1234.5,
-            client_index: 7,
-            ..EngineCoreRequest::default()
-        };
-
-        let encoded = encode_msgpack(&request).unwrap();
-        let value = decode_value(&encoded).unwrap();
-        let array = match value {
-            Value::Array(array) => array,
-            other => panic!("expected array, got {other:?}"),
-        };
-
-        assert_eq!(array.len(), 20);
-        assert_eq!(array[0], Value::from("req-1"));
-        assert_eq!(array[2], Value::Nil);
-        assert_eq!(array[4], Value::Nil);
-        assert_eq!(array[10], Value::Nil);
-        assert_eq!(array[11], Value::from(7));
-    }
-}

rust/src/engine-core-client/src/protocol/sampling.rs

@@ -1,211 +0,0 @@
-use std::collections::{BTreeSet, HashMap};
-
-use serde::{Deserialize, Serialize};
-use serde_default::DefaultFromSerde;
-
-use crate::protocol::structured_outputs::StructuredOutputsParams;
-
-fn default_top_p() -> f32 {
-    1.0
-}
-
-fn default_repetition_penalty() -> f32 {
-    1.0
-}
-
-fn default_temperature() -> f32 {
-    1.0
-}
-
-fn default_max_tokens() -> u32 {
-    16
-}
-
-/// Engine-core-facing sampling parameters for text generation.
-///
-/// This is the normalized southbound subset used by the Rust frontend when it
-/// talks to Python engine-core over the wire. User-facing request semantics
-/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are
-/// intentionally handled by higher layers before values reach this DTO.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
-// Python's SamplingParams is `omit_defaults=True`, so msgpack drops
-// default-valued keys; default the whole struct. Per-field fns cover the
-// non-zero defaults.
-#[serde_with::skip_serializing_none]
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)]
-#[serde(default)]
-pub struct EngineCoreSamplingParams {
-    /// Controls randomness. Lower values are more deterministic; zero means
-    /// greedy sampling.
-    #[serde(default = "default_temperature")]
-    pub temperature: f32,
-    /// Cumulative probability threshold for nucleus sampling.
-    #[serde(default = "default_top_p")]
-    pub top_p: f32,
-    /// Maximum number of top tokens to consider. `0` means all tokens.
-    pub top_k: u32,
-    /// Random seed used by the sampler when present.
-    pub seed: Option<i64>,
-    /// Maximum number of tokens to generate per output sequence.
-    #[serde(default = "default_max_tokens")]
-    pub max_tokens: u32,
-    /// Minimum number of tokens to generate before EOS or stop-token handling.
-    pub min_tokens: u32,
-    /// Maximum number of reasoning ("thinking") tokens to emit before the
-    /// reasoning section is force-closed. `None` means unlimited; the
-    /// user-facing `-1` sentinel is normalized to `None` by the frontend before
-    /// reaching this DTO, so only non-negative values are sent. Enforced
-    /// engine-side (and only when a reasoning parser is configured).
-    pub thinking_token_budget: Option<u64>,
-    /// Number of log probabilities to return per generated token.
-    ///
-    /// `None` disables sample logprobs. `-1` requests the full vocabulary.
-    pub logprobs: Option<i32>,
-    /// Number of log probabilities to return per prompt token.
-    ///
-    /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
-    pub prompt_logprobs: Option<i32>,
-    /// Minimum probability threshold for token sampling.
-    pub min_p: f32,
-    /// Frequency penalty applied by the sampler.
-    pub frequency_penalty: f32,
-    /// Presence penalty applied by the sampler.
-    pub presence_penalty: f32,
-    /// Repetition penalty applied by the sampler.
-    #[serde(default = "default_repetition_penalty")]
-    pub repetition_penalty: f32,
-    /// Token IDs that stop generation.
-    pub stop_token_ids: Vec<u32>,
-    /// Primary EOS token ID used by engine-core's dedicated EOS stop path.
-    ///
-    /// This mirrors Python's internal `_eos_token_id` field and is derived by
-    /// the frontend from tokenizer/model metadata rather than supplied directly
-    /// by end users.
-    #[serde(rename = "_eos_token_id")]
-    pub eos_token_id: Option<u32>,
-    /// Complete stop-token set used by engine-core for `min_tokens` masking.
-    ///
-    /// This mirrors Python's internal `_all_stop_token_ids` field and should
-    /// contain explicit `stop_token_ids` plus any frontend-derived EOS token
-    /// IDs.
-    #[serde(rename = "_all_stop_token_ids")]
-    pub all_stop_token_ids: BTreeSet<u32>,
-    /// Logit biases to apply during sampling.
-    /// Keys are token IDs
-    pub logit_bias: Option<HashMap<u32, f32>>,
-    /// Restrict output to these token IDs only.
-    pub allowed_token_ids: Option<Vec<u32>>,
-    /// Tokenized bad words to avoid during generation.
-    #[serde(rename = "_bad_words_token_ids")]
-    pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
-    /// Parameters for configuring structured outputs (guided decoding).
-    pub structured_outputs: Option<StructuredOutputsParams>,
-    /// Specific token IDs for which log probabilities should be returned at
-    /// each position.
-    ///
-    /// When set, the engine returns logprobs for exactly these tokens in
-    /// addition to the sampled/scored token. Mutually exclusive with the
-    /// `logprobs` count field in practice.
-    pub logprob_token_ids: Option<Vec<u32>>,
-    /// If `Some(true)`, the request will not attempt to read from the prefix
-    /// cache; newly computed blocks may still populate the cache. `None`
-    /// defers to engine-core defaults.
-    pub skip_reading_prefix_cache: Option<bool>,
-    /// Additional request parameters for custom extensions (from `vllm_xargs`).
-    pub extra_args: Option<HashMap<String, serde_json::Value>>,
-}
-
-impl EngineCoreSamplingParams {
-    /// Constructs a default sampling params for testing purposes only.
-    pub fn for_test() -> Self {
-        Self {
-            temperature: 1.0,
-            top_p: 1.0,
-            top_k: 0,
-            seed: None,
-            max_tokens: 65536,
-            min_tokens: 0,
-            thinking_token_budget: None,
-            logprobs: None,
-            prompt_logprobs: None,
-            min_p: 0.0,
-            frequency_penalty: 0.0,
-            presence_penalty: 0.0,
-            repetition_penalty: 1.0,
-            stop_token_ids: Vec::new(),
-            eos_token_id: None,
-            all_stop_token_ids: BTreeSet::new(),
-            logit_bias: None,
-            allowed_token_ids: None,
-            bad_words_token_ids: None,
-            structured_outputs: None,
-            logprob_token_ids: None,
-            skip_reading_prefix_cache: None,
-            extra_args: None,
-        }
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use rmpv::Value;
-
-    use crate::protocol::decode_msgpack;
-    use crate::protocol::request::EngineCoreRequest;
-
-    /// A real `sampling_params` is a sparse `omit_defaults` map; absent fields
-    /// must fall back to defaults. `python_compat` can't catch this since Rust
-    /// encodes full maps (see `engine_core_request_serializes_as_full_array`).
-    #[test]
-    fn decodes_sampling_params_with_omitted_defaults() {
-        let sampling_params = Value::Map(vec![
-            (
-                Value::from("stop_token_ids"),
-                Value::Array(vec![Value::from(151643u32)]),
-            ),
-            (Value::from("skip_reading_prefix_cache"), Value::from(false)),
-        ]);
-        let request = Value::Array(vec![
-            Value::from("req-omit-defaults"),
-            Value::Array(vec![
-                Value::from(1u32),
-                Value::from(2u32),
-                Value::from(3u32),
-            ]),
-            Value::Nil,
-            sampling_params,
-            Value::Nil,
-            Value::from(1.0f64),
-        ]);
-
-        let mut bytes = Vec::new();
-        rmpv::encode::write_value(&mut bytes, &request).unwrap();
-
-        let decoded: EngineCoreRequest = decode_msgpack(&bytes)
-            .expect("a real omit_defaults request must decode (regression: missing field)");
-
-        assert_eq!(decoded.request_id, "req-omit-defaults");
-        let sampling = decoded.sampling_params.expect("sampling params present");
-
-        assert_eq!(sampling.stop_token_ids, vec![151643]);
-        assert_eq!(sampling.skip_reading_prefix_cache, Some(false));
-
-        // Omitted fields -> Python defaults.
-        assert_eq!(sampling.temperature, 1.0);
-        assert_eq!(sampling.top_p, 1.0);
-        assert_eq!(sampling.top_k, 0);
-        assert_eq!(sampling.seed, None);
-        assert_eq!(sampling.max_tokens, 16);
-        assert_eq!(sampling.min_tokens, 0);
-        assert_eq!(sampling.min_p, 0.0);
-        assert_eq!(sampling.frequency_penalty, 0.0);
-        assert_eq!(sampling.presence_penalty, 0.0);
-        assert_eq!(sampling.repetition_penalty, 1.0);
-        assert_eq!(sampling.logprobs, None);
-        assert_eq!(sampling.prompt_logprobs, None);
-        assert_eq!(sampling.eos_token_id, None);
-        assert!(sampling.all_stop_token_ids.is_empty());
-    }
-}

rust/src/engine-core-client/src/protocol/structured_outputs.rs

@@ -1,81 +0,0 @@
-use serde::{Deserialize, Serialize};
-
-/// Structured-output backend selected for EngineCore grammar compilation.
-///
-/// Python vLLM stores this in `StructuredOutputsParams._backend` after request
-/// validation. The Rust frontend currently always lowers structured-output
-/// requests to guidance, while ignoring any user-supplied `_backend` value.
-#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
-#[serde(rename_all = "kebab-case")]
-pub enum StructuredOutputBackend {
-    Xgrammar,
-    #[default]
-    Guidance,
-    Outlines,
-    LmFormatEnforcer,
-}
-
-/// Parameters for configuring structured outputs (guided decoding).
-///
-/// Exactly one constraint field (`json`, `regex`, `choice`, `grammar`,
-/// `json_object`, or `structural_tag`) should be set. The engine-core
-/// backend selects the appropriate grammar compiler based on which field
-/// is present.
-///
-/// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
-#[serde_with::skip_serializing_none]
-#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
-#[serde(default)]
-pub struct StructuredOutputsParams {
-    /// JSON schema (as a dict/object or JSON string) constraining the output.
-    pub json: Option<serde_json::Value>,
-    /// Regular expression the output must match.
-    pub regex: Option<String>,
-    /// List of allowed output strings (the model must produce one of these).
-    pub choice: Option<Vec<String>>,
-    /// Context-free grammar (in EBNF-like notation) the output must conform to.
-    pub grammar: Option<String>,
-    /// When `true`, output must be valid JSON (free-form, no schema).
-    pub json_object: Option<bool>,
-    /// Disable any additional whitespace in guided JSON output.
-    #[serde(skip_serializing_if = "crate::protocol::is_false")]
-    pub disable_any_whitespace: bool,
-    /// Disable `additionalProperties` in JSON schema output.
-    #[serde(skip_serializing_if = "crate::protocol::is_false")]
-    pub disable_additional_properties: bool,
-    /// Custom whitespace pattern for guided JSON output.
-    pub whitespace_pattern: Option<String>,
-    /// Structural tag configuration (JSON-encoded string).
-    pub structural_tag: Option<String>,
-    /// Structured-output backend, mirroring Python's internal `_backend`.
-    ///
-    /// User-supplied values are ignored during deserialization. This matches
-    /// Python's request boundary, where `_backend` is set by validation rather
-    /// than accepted as a request-level backend selector.
-    #[serde(
-        default,
-        rename = "_backend",
-        deserialize_with = "serde_with::rust::deserialize_ignore_any"
-    )]
-    pub backend: StructuredOutputBackend,
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn structured_outputs_backend_ignores_deserialized_value() {
-        let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({
-            "json_object": true,
-            "_backend": "xgrammar",
-        }))
-        .unwrap();
-
-        assert_eq!(params.backend, StructuredOutputBackend::Guidance);
-
-        let value = serde_json::to_value(params).unwrap();
-        assert_eq!(value["_backend"], "guidance");
-    }
-}

rust/src/engine-core-client/src/tests/client.rs

@@ -22,14 +22,13 @@ use crate::protocol::multimodal::{
     MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange,
     SliceSpec,
 };
-use crate::protocol::output::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, decode_engine_core_outputs,
-};
-use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
-use crate::protocol::sampling::EngineCoreSamplingParams;
 use crate::protocol::stats::SchedulerStats;
 use crate::protocol::tensor::WireTensor;
 use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
+use crate::protocol::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
+    EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs,
+};
 use crate::test_utils::{
     IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
     setup_mock_engine_with_init, spawn_mock_engine_task,

rust/src/engine-core-client/src/transport.rs

@@ -18,8 +18,9 @@ use crate::error::{Error, Result, bail_unexpected_handshake_message};
 use crate::protocol::handshake::{
     EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage,
 };
-use crate::protocol::output::{EngineCoreOutputs, decode_engine_core_outputs};
-use crate::protocol::{decode_msgpack, encode_msgpack};
+use crate::protocol::{
+    EngineCoreOutputs, decode_engine_core_outputs, decode_msgpack, encode_msgpack,
+};
 
 /// Dedicated single-frame sentinel emitted by Python `EngineCoreProc` when the
 /// engine dies.

rust/src/llm/examples/external_engine_smoke.rs

@@ -5,7 +5,7 @@ use clap::Parser;
 use futures::StreamExt as _;
 use tokio::time::timeout;
 use tracing_subscriber::EnvFilter;
-use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
 use vllm_llm::{FinishReason, GenerateOutputStream, GenerateRequest, Llm};
 

rust/src/llm/src/output.rs

@@ -8,7 +8,7 @@ use futures::stream::FusedStream;
 use futures::{Stream, StreamExt as _, pin_mut};
 use serde::{Deserialize, Serialize};
 use vllm_engine_core_client::protocol::logprobs::Logprobs;
-use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason};
+use vllm_engine_core_client::protocol::{EngineCoreFinishReason, StopReason};
 use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
 
 use crate::error::Result;

rust/src/llm/src/request.rs

@@ -4,8 +4,9 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use uuid::Uuid;
 use vllm_engine_core_client::protocol::lora::LoraRequest;
 use vllm_engine_core_client::protocol::multimodal::MmFeatures;
-use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningParserKwargs};
-use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::{
+    EngineCoreRequest, EngineCoreSamplingParams, ReasoningParserKwargs,
+};
 
 use crate::error::{Error, Result};
 
@@ -133,8 +134,7 @@ fn current_unix_timestamp_secs() -> f64 {
 mod tests {
     use std::collections::BTreeMap;
 
-    use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
-    use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+    use vllm_engine_core_client::protocol::{EngineCoreSamplingParams, ReasoningParserKwargs};
 
     use super::GenerateRequest;
     use crate::error::Error;

rust/src/llm/src/request_metrics.rs

@@ -1,9 +1,7 @@
 use std::time::{SystemTime, UNIX_EPOCH};
 
-use vllm_engine_core_client::protocol::output::{
-    EngineCoreEvent, EngineCoreEventType, EngineCoreOutput,
-};
 use vllm_engine_core_client::protocol::stats::PrefillStats;
+use vllm_engine_core_client::protocol::{EngineCoreEvent, EngineCoreEventType, EngineCoreOutput};
 use vllm_metrics::{
     EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics,
 };
@@ -330,8 +328,8 @@ pub(crate) fn current_unix_timestamp_secs() -> f64 {
 
 #[cfg(test)]
 mod tests {
-    use vllm_engine_core_client::protocol::output::{EngineCoreEvent, EngineCoreEventType};
     use vllm_engine_core_client::protocol::stats::PrefillStats;
+    use vllm_engine_core_client::protocol::{EngineCoreEvent, EngineCoreEventType};
 
     use super::{RequestMetricsTracker, diff_or_zero};
 
@@ -343,7 +341,7 @@ mod tests {
             2,
             10.0,
             100.2,
-            &vllm_engine_core_client::protocol::output::EngineCoreOutput {
+            &vllm_engine_core_client::protocol::EngineCoreOutput {
                 request_id: "req-1".to_string(),
                 new_token_ids: vec![1],
                 finish_reason: None,
@@ -371,7 +369,7 @@ mod tests {
             2,
             11.5,
             100.4,
-            &vllm_engine_core_client::protocol::output::EngineCoreOutput {
+            &vllm_engine_core_client::protocol::EngineCoreOutput {
                 request_id: "req-1".to_string(),
                 new_token_ids: vec![2, 3],
                 finish_reason: None,

rust/src/llm/tests/generate.rs

@@ -9,13 +9,11 @@ use uuid::Uuid;
 use vllm_engine_core_client::protocol::logprobs::{
     Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
 };
-use vllm_engine_core_client::protocol::output::{
-    EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
-    EngineCoreOutputs,
-};
-use vllm_engine_core_client::protocol::request::EngineCoreRequest;
-use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
 use vllm_engine_core_client::protocol::stats::PrefillStats;
+use vllm_engine_core_client::protocol::{
+    EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
+    EngineCoreOutputs, EngineCoreRequest, EngineCoreSamplingParams,
+};
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
 use vllm_llm::{

rust/src/mock-engine/src/engine.rs

@@ -11,13 +11,12 @@ use tokio::sync::mpsc;
 use tokio::task::yield_now;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, info, warn};
-use vllm_engine_core_client::protocol::output::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
-};
-use vllm_engine_core_client::protocol::request::EngineCoreRequest;
 use vllm_engine_core_client::protocol::utility::{
     EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope,
 };
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
+};
 
 use super::Opt;
 

rust/src/mock-engine/src/io.rs

@@ -4,9 +4,10 @@ use tokio::sync::mpsc;
 use tokio_util::sync::CancellationToken;
 use tracing::warn;
 use vllm_engine_core_client::mock_engine::MockEngineDataSockets;
-use vllm_engine_core_client::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
 use vllm_engine_core_client::protocol::utility::EngineCoreUtilityRequest;
-use vllm_engine_core_client::protocol::{decode_msgpack, encode_msgpack};
+use vllm_engine_core_client::protocol::{
+    EngineCoreRequest, EngineCoreRequestType, decode_msgpack, encode_msgpack,
+};
 use zeromq::{DealerSocket, PushSocket, SocketRecv as _, SocketSend as _, ZmqMessage};
 
 use crate::engine::{EngineInput, EngineOutput};

rust/src/mock-engine/src/tests.rs

@@ -5,9 +5,9 @@ use anyhow::Result;
 use futures::StreamExt as _;
 use tokio::time::timeout;
 use tokio_util::sync::CancellationToken;
-use vllm_engine_core_client::protocol::output::EngineCoreFinishReason;
-use vllm_engine_core_client::protocol::request::EngineCoreRequest;
-use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
+};
 use vllm_engine_core_client::test_utils::IpcNamespace;
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
 

rust/src/model-files/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "vllm-model-files"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+hf-hub.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+thiserror-ext.workspace = true
+
+[dev-dependencies]
+tempfile.workspace = true
+tokio.workspace = true
+
+[lints]
+workspace = true

rust/src/model-files/src/config.rs

@@ -0,0 +1,21 @@
+use std::path::Path;
+
+use serde::Deserialize;
+
+use crate::error::Result;
+use crate::json::read_json_file;
+
+/// Minimal subset of `tokenizer_config.json` needed by tokenizer selection.
+#[derive(Debug, Default, Deserialize)]
+#[serde(default)]
+pub(crate) struct TokenizerConfig {
+    /// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
+    /// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
+    /// containing "Tiktoken" which can be used as a hint for backend
+    /// selection.
+    pub tokenizer_class: Option<String>,
+}
+
+pub(crate) fn load_tokenizer_config(path: Option<&Path>) -> Result<TokenizerConfig> {
+    read_json_file(path)
+}

rust/src/model-files/src/error.rs

@@ -0,0 +1,15 @@
+use thiserror::Error as ThisError;
+
+/// Error returned while resolving or reading model files.
+#[derive(Debug, ThisError)]
+#[error("model file error: {0}")]
+pub struct Error(String);
+
+impl Error {
+    pub(crate) fn new(message: impl Into<String>) -> Self {
+        Self(message.into())
+    }
+}
+
+/// Result type used by model-file discovery helpers.
+pub type Result<T> = std::result::Result<T, Error>;

rust/src/model-files/src/json.rs

@@ -0,0 +1,31 @@
+use std::fs;
+use std::path::Path;
+
+use serde::Deserialize;
+use thiserror_ext::AsReport as _;
+
+use crate::error::{Error, Result};
+
+/// Read an optional JSON file into `T`, returning `T::default()` when absent.
+pub fn read_json_file<T>(path: Option<&Path>) -> Result<T>
+where
+    T: for<'de> Deserialize<'de> + Default,
+{
+    let Some(path) = path else {
+        return Ok(T::default());
+    };
+    let content = fs::read_to_string(path).map_err(|error| {
+        Error::new(format!(
+            "failed to read {}: {}",
+            path.display(),
+            error.as_report()
+        ))
+    })?;
+    serde_json::from_str(&content).map_err(|error| {
+        Error::new(format!(
+            "failed to parse {}: {}",
+            path.display(),
+            error.as_report()
+        ))
+    })
+}

rust/src/model-files/src/lib.rs

@@ -0,0 +1,10 @@
+//! Hugging Face model file discovery shared by Rust frontend crates.
+
+mod config;
+mod error;
+mod json;
+mod model_files;
+
+pub use error::{Error, Result};
+pub use json::read_json_file;
+pub use model_files::{ResolvedModelFiles, TokenizerSource};

rust/src/model-files/src/model_files.rs

@@ -4,7 +4,7 @@ use hf_hub::Cache;
 use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo};
 use thiserror_ext::AsReport as _;
 
-use super::config::{HfTokenizerConfig, load_tokenizer_config};
+use crate::config::{TokenizerConfig, load_tokenizer_config};
 use crate::error::{Error, Result};
 
 const HF_TOKEN_ENV: &str = "HF_TOKEN";
@@ -26,6 +26,28 @@ pub enum TokenizerSource {
 }
 
 impl TokenizerSource {
+    /// Select a tokenizer source from a tokenizer file path.
+    pub fn from_path(path: impl Into<PathBuf>) -> Result<Self> {
+        let path = path.into();
+        let file_name = path.file_name().and_then(|name| name.to_str()).ok_or_else(|| {
+            Error::new(format!(
+                "tokenizer path has no file name: {}",
+                path.display()
+            ))
+        })?;
+
+        match file_name {
+            "tekken.json" => Ok(Self::Tekken(path)),
+            "tokenizer.json" => Ok(Self::HuggingFace(path)),
+            _ if is_tiktoken_file(&path) => Ok(Self::Tiktoken(path)),
+            _ => Err(Error::new(format!(
+                "unsupported tokenizer file '{}'",
+                path.display()
+            ))),
+        }
+    }
+
+    /// Return the local filesystem path for this tokenizer source.
     pub fn path(&self) -> &Path {
         match self {
             Self::HuggingFace(path) | Self::Tiktoken(path) | Self::Tekken(path) => path,
@@ -38,10 +60,15 @@ impl TokenizerSource {
 pub struct ResolvedModelFiles {
     /// The selected tokenizer source for this model.
     pub tokenizer: TokenizerSource,
+    /// Path to `tokenizer_config.json` when present.
     pub tokenizer_config_path: Option<PathBuf>,
+    /// Path to `generation_config.json` when present.
     pub generation_config_path: Option<PathBuf>,
+    /// Path to `preprocessor_config.json` when present.
     pub preprocessor_config_path: Option<PathBuf>,
+    /// Path to a discovered chat template file when present.
     pub chat_template_path: Option<PathBuf>,
+    /// Path to `config.json` when present.
     pub config_path: Option<PathBuf>,
 }
 
@@ -76,10 +103,10 @@ fn resolve_local_model_files(model_dir: &Path) -> Result<ResolvedModelFiles> {
 }
 
 async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles> {
-    let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?;
+    let api = build_api().map_err(|error| Error::new(format!("{}", error.as_report())))?;
     let repo = api.model(model_id.to_string());
     let info = repo.info().await.map_err(|error| {
-        Error::Tokenizer(format!(
+        Error::new(format!(
             "failed to fetch model '{model_id}': {}",
             error.as_report()
         ))
@@ -138,9 +165,10 @@ fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFile
         None => return Ok(None),
     };
 
-    let model_dir = tokenizer.path().parent().ok_or_else(|| {
-        Error::Tokenizer("resolved tokenizer file has no parent directory".to_string())
-    })?;
+    let model_dir = tokenizer
+        .path()
+        .parent()
+        .ok_or_else(|| Error::new("resolved tokenizer file has no parent directory"))?;
     let generation_config_path = cache_repo.get("generation_config.json");
     let preprocessor_config_path = cache_repo.get("preprocessor_config.json");
     let chat_template_path = discover_chat_template_in_dir(model_dir);
@@ -162,113 +190,88 @@ async fn resolve_remote_tokenizer_source(
     siblings: &std::collections::BTreeSet<&str>,
     tokenizer_class: Option<&str>,
 ) -> Result<TokenizerSource> {
-    if let Some(tekken_path) = download_if_present(repo, model_id, siblings, "tekken.json").await? {
-        return Ok(TokenizerSource::Tekken(tekken_path));
-    }
-
-    let tokenizer_path = if siblings.contains("tokenizer.json") {
+    let tokenizer_path = if siblings.contains("tekken.json") {
+        download_known_file(repo, model_id, "tekken.json").await?
+    } else if siblings.contains("tokenizer.json") {
         download_known_file(repo, model_id, "tokenizer.json").await?
     } else if let Some(tiktoken_name) = find_tiktoken_sibling(siblings) {
         download_known_file(repo, model_id, tiktoken_name).await?
     } else {
-        return Err(Error::Tokenizer(format!(
+        return Err(Error::new(format!(
             "model '{model_id}' does not expose a supported tokenizer file \
              (tokenizer.json, tiktoken.model, or *.tiktoken) on Hugging Face"
         )));
     };
 
-    Ok(resolve_tokenizer_source(
-        tokenizer_path,
-        tokenizer_class,
-        None,
-    ))
+    resolve_tokenizer_source(tokenizer_path, tokenizer_class)
 }
 
 fn resolve_cached_tokenizer_source(
     cache_repo: &hf_hub::CacheRepo,
-    tokenizer_config: &HfTokenizerConfig,
+    tokenizer_config: &TokenizerConfig,
 ) -> Result<Option<TokenizerSource>> {
-    let tekken_path = cache_repo.get("tekken.json");
-
-    if let Some(tekken_path) = tekken_path {
-        return Ok(Some(TokenizerSource::Tekken(tekken_path)));
-    }
-
-    let Some(tokenizer_path) = cache_repo.get("tokenizer.json").or_else(|| {
-        // tiktoken.model is the most common name, try it first.
-        cache_repo.get("tiktoken.model").or_else(|| {
-            // Scan for any *.tiktoken file in the cache snapshot directory.
-            let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
-            discover_tiktoken_in_dir(&snapshot_dir)
+    let Some(tokenizer_path) = cache_repo
+        .get("tekken.json")
+        .or_else(|| cache_repo.get("tokenizer.json"))
+        .or_else(|| {
+            // tiktoken.model is the most common name, try it first.
+            cache_repo.get("tiktoken.model").or_else(|| {
+                // Scan for any *.tiktoken file in the cache snapshot directory.
+                let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
+                discover_tiktoken_in_dir(&snapshot_dir)
+            })
         })
-    }) else {
+    else {
         return Ok(None);
     };
 
     Ok(Some(resolve_tokenizer_source(
         tokenizer_path,
         tokenizer_config.tokenizer_class.as_deref(),
-        None,
-    )))
+    )?))
 }
 
 fn resolve_local_tokenizer_source(
     model_dir: &Path,
-    tokenizer_config: &HfTokenizerConfig,
+    tokenizer_config: &TokenizerConfig,
 ) -> Result<TokenizerSource> {
-    let tekken_path = local_file_if_exists(model_dir, "tekken.json");
-    if let Some(tekken_path) = tekken_path {
-        return Ok(TokenizerSource::Tekken(tekken_path));
-    }
-
-    let tokenizer_path = local_file_if_exists(model_dir, "tokenizer.json")
+    let tokenizer_path = local_file_if_exists(model_dir, "tekken.json")
+        .or_else(|| local_file_if_exists(model_dir, "tokenizer.json"))
         .or_else(|| local_file_if_exists(model_dir, "tiktoken.model"))
         .or_else(|| discover_tiktoken_in_dir(model_dir))
         .ok_or_else(|| {
-            Error::Tokenizer(format!(
+            Error::new(format!(
                 "local model directory '{}' does not contain a supported tokenizer file \
                  (tokenizer.json, tiktoken.model, or *.tiktoken)",
                 model_dir.display()
             ))
         })?;
 
-    Ok(resolve_tokenizer_source(
-        tokenizer_path,
-        tokenizer_config.tokenizer_class.as_deref(),
-        None,
-    ))
+    resolve_tokenizer_source(tokenizer_path, tokenizer_config.tokenizer_class.as_deref())
 }
 
 /// Choose the tokenizer.
 ///
 /// Selection order:
-/// 1. `tekken.json` — Mistral native tokenizer (preferred over HF `tokenizer.json` because the HF
-///    version has a known regex bug for Mistral models).
-/// 2. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
-/// 3. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
+/// 1. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
+/// 2. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
 ///    insensitive) trigger tiktoken loading from a sibling BPE file.
-/// 4. Default — `tokenizer.json` in HuggingFace format.
+/// 3. Default — `tokenizer.json` in HuggingFace format.
 fn resolve_tokenizer_source(
     tokenizer_path: PathBuf,
     tokenizer_class: Option<&str>,
-    tekken_path: Option<PathBuf>,
-) -> TokenizerSource {
-    if let Some(tekken_path) = tekken_path {
-        return TokenizerSource::Tekken(tekken_path);
-    }
+) -> Result<TokenizerSource> {
+    let tokenizer = TokenizerSource::from_path(tokenizer_path)?;
 
-    if is_tiktoken_file(&tokenizer_path) {
-        return TokenizerSource::Tiktoken(tokenizer_path);
-    }
-
-    if tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
-        && let Some(dir) = tokenizer_path.parent()
+    if let TokenizerSource::HuggingFace(path) = &tokenizer
+        && tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
+        && let Some(dir) = path.parent()
         && let Some(tiktoken_path) = discover_tiktoken_in_dir(dir)
     {
-        return TokenizerSource::Tiktoken(tiktoken_path);
+        return Ok(TokenizerSource::Tiktoken(tiktoken_path));
     }
 
-    TokenizerSource::HuggingFace(tokenizer_path)
+    Ok(tokenizer)
 }
 
 /// Download `filename` only if it exists in `siblings`.
@@ -286,7 +289,7 @@ async fn download_if_present(
 
 async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) -> Result<PathBuf> {
     repo.get(filename).await.map_err(|error| {
-        Error::Tokenizer(format!(
+        Error::new(format!(
             "failed to download '{filename}' for model '{model_id}': {}",
             error.as_report()
         ))
@@ -317,7 +320,7 @@ fn find_tiktoken_sibling<'a>(siblings: &std::collections::BTreeSet<&'a str>) ->
 }
 
 /// Discover a tiktoken model file in a local directory.
-pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
+fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
     let tiktoken_model = dir.join("tiktoken.model");
     if tiktoken_model.exists() {
         return Some(tiktoken_model);
@@ -337,7 +340,7 @@ pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf>
 }
 
 /// Returns `true` if `path` points to a tiktoken-format file (by name).
-pub(super) fn is_tiktoken_file(path: &std::path::Path) -> bool {
+fn is_tiktoken_file(path: &std::path::Path) -> bool {
     path.file_name()
         .and_then(|n| n.to_str())
         .is_some_and(|name| name == "tiktoken.model" || name.ends_with(".tiktoken"))
@@ -368,7 +371,6 @@ mod tests {
     use std::fs;
 
     use tempfile::tempdir;
-    use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
 
     use super::{ResolvedModelFiles, TokenizerSource};
 
@@ -399,61 +401,4 @@ mod tests {
             Some(dir.path().join("tokenizer_config.json"))
         );
     }
-
-    #[tokio::test]
-    #[ignore = "too slow for CI and requires network access to Hugging Face"]
-    async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
-        let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
-            .await
-            .expect("resolve real Kimi K2.5 model files");
-
-        let tokenizer_path = match &files.tokenizer {
-            TokenizerSource::Tiktoken(path) => path.clone(),
-            other => panic!("expected tiktoken tokenizer source, got {other:?}"),
-        };
-
-        for backend in [
-            TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
-            TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
-        ] {
-            let think_id = backend.token_to_id("<think>").expect("resolve <think>");
-            let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
-            let tool_section_id = backend
-                .token_to_id("<|tool_calls_section_begin|>")
-                .expect("resolve tool call section marker");
-            let contraction_heavy_text =
-                "I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
-            let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
-
-            assert_eq!(
-                (think_id, end_think_id, tool_section_id),
-                (163606, 163607, 163595)
-            );
-            assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
-            assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
-            assert_eq!(
-                backend.decode(&[tool_section_id], true).unwrap(),
-                "<|tool_calls_section_begin|>"
-            );
-
-            // This demonstrates that we're using Kimi's custom BPE pattern.
-            // With CL100K this will be 23 tokens instead.
-            assert_eq!(
-                contraction_heavy_ids,
-                vec![
-                    17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
-                    1573, 56229, 13922, 13,
-                ]
-            );
-            assert_eq!(contraction_heavy_ids.len(), 17);
-            assert_eq!(
-                backend.decode(&contraction_heavy_ids, false).unwrap(),
-                contraction_heavy_text
-            );
-
-            // Special-looking text that is not actually registered should fail gracefully.
-            assert_eq!(backend.token_to_id("◁think▷"), None);
-            assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
-        }
-    }
 }

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

@@ -5,13 +5,14 @@
 
 mod adapter;
 
-pub(super) use adapter::UnifiedToolParserAdapter;
 use futures::FutureExt as _;
 use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool};
 use tool_parser::traits::ToolParser as ExternalToolParser;
 use vllm_parser::tool::test_utils::collect_stream;
 use vllm_parser::tool::{Tool, ToolParser};
 
+pub(super) use adapter::UnifiedToolParserAdapter;
+
 pub(super) fn openai_tools(tools: &[Tool]) -> Vec<OpenAiTool> {
     tools
         .iter()

rust/src/parser/python/src/lib.rs

@@ -231,9 +231,8 @@ fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
 
 #[cfg(test)]
 mod tests {
-    use serde_json::json;
-
     use super::*;
+    use serde_json::json;
 
     fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
         Python::initialize();

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

@@ -49,8 +49,10 @@ mod tests {
     use std::sync::Arc;
 
     use super::SeedOssReasoningParser;
-    use crate::reasoning::ReasoningParser;
-    use crate::reasoning::tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer};
+    use crate::reasoning::{
+        ReasoningParser,
+        tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer},
+    };
 
     #[test]
     fn without_prompt_markers_expects_start_token() {

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

@@ -127,8 +127,10 @@ mod tests {
     use std::sync::Arc;
 
     use super::Step3p5ReasoningParser;
-    use crate::reasoning::ReasoningParser;
-    use crate::reasoning::tests::{THINK_START_ID, fake_tokenizer};
+    use crate::reasoning::{
+        ReasoningParser,
+        tests::{THINK_START_ID, fake_tokenizer},
+    };
 
     #[test]
     fn picks_up_prompt_start_boundary() {

rust/src/parser/src/tool/mod.rs

@@ -14,6 +14,8 @@ mod parameters;
 mod qwen_coder;
 #[cfg(any(test, feature = "test-util"))]
 pub mod test_utils;
+use crate::utils;
+
 use std::collections::{BTreeMap, btree_map};
 
 pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser};
@@ -33,8 +35,6 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 pub use xgrammar_structural_tag::Model as StructuralTagModel;
 
-use crate::utils;
-
 /// One function-style tool made available to the model.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Tool {

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

@@ -2,10 +2,11 @@
 
 use vllm_tokenizer::DynTokenizer;
 
-use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
 use crate::reasoning::ReasoningParser;
 use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput};
 
+use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
+
 /// Unified parser that composes existing reasoning and tool parsers.
 pub struct CombinedParser {
     reasoning: Option<Box<dyn ReasoningParser>>,

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

@@ -1,5 +1,4 @@
 use serde_json::{Map, Number, Value};
-use vllm_tokenizer::DynTokenizer;
 use winnow::ascii::multispace0 as ws0;
 use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated};
 use winnow::error::{ContextError, ErrMode, ModalResult};
@@ -7,6 +6,8 @@ use winnow::prelude::*;
 use winnow::stream::{Partial, Stream};
 use winnow::token::{literal, take_till, take_until};
 
+use vllm_tokenizer::DynTokenizer;
+
 use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
 use crate::reasoning::last_reasoning_boundary;
 use crate::tool::{Tool, ToolCallDelta};

rust/src/parser/src/unified/mod.rs

@@ -3,12 +3,13 @@
 mod combined;
 mod gemma4;
 
-pub use combined::CombinedParser;
-pub use gemma4::Gemma4UnifiedParser;
 use thiserror::Error;
 use thiserror_ext::Macro;
 use vllm_tokenizer::DynTokenizer;
 
+pub use combined::CombinedParser;
+pub use gemma4::Gemma4UnifiedParser;
+
 use crate::reasoning::ReasoningError;
 use crate::tool::{
     StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,

rust/src/server/src/error.rs

@@ -1,7 +1,8 @@
 use axum::Json;
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
-use thiserror_ext::{AsReport as _, Construct, Macro};
+use thiserror_ext::AsReport as _;
+use thiserror_ext::{Construct, Macro};
 
 use crate::routes::openai::utils::types::{ErrorDetail, ErrorResponse};
 

rust/src/server/src/grpc/convert.rs

@@ -3,8 +3,7 @@
 
 use tonic::Status;
 use uuid::Uuid;
-use vllm_engine_core_client::protocol::output::StopReason;
-use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;
+use vllm_engine_core_client::protocol::{StopReason, StructuredOutputsParams};
 use vllm_text::{
     DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams,
     TextDecodeOptions, TextRequest,
@@ -503,7 +502,7 @@ impl ResponseOpts {
 
 #[cfg(test)]
 mod tests {
-    use vllm_engine_core_client::protocol::output::StopReason;
+    use vllm_engine_core_client::protocol::StopReason;
     use vllm_text::{FinishReason, Finished, Prompt};
 
     use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason};

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

@@ -18,10 +18,9 @@ use vllm_chat::{
     ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
     DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
 };
-use vllm_engine_core_client::protocol::output::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
 };
-use vllm_engine_core_client::protocol::request::EngineCoreRequest;
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
 use vllm_llm::Llm;

rust/src/server/src/listener.rs

@@ -152,8 +152,8 @@ impl axum::serve::Listener for Listener {
 
 /// Allow the unified listener to be adaptable to `tls_listener`.
 impl AsyncAccept for Listener {
-    type Address = ListenerAddr;
     type Connection = ListenerIo;
+    type Address = ListenerAddr;
     type Error = std::io::Error;
 
     fn poll_accept(

rust/src/server/src/middleware/offload.rs

@@ -57,9 +57,9 @@ where
     S::Error: Send + 'static,
     B: Send + 'static,
 {
+    type Response = S::Response;
     type Error = S::Error;
     type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
-    type Response = S::Response;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         self.inner.poll_ready(cx)

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

@@ -18,10 +18,9 @@ use vllm_chat::{
     ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
     DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
 };
-use vllm_engine_core_client::protocol::output::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
 };
-use vllm_engine_core_client::protocol::request::EngineCoreRequest;
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
 use vllm_llm::Llm;

rust/src/server/src/routes/openai/chat_completions.rs

@@ -21,7 +21,7 @@ use vllm_chat::{
     AssistantBlockKind, AssistantMessageExt as _, ChatEvent, ChatEventStream, ChatEventStreamTrait,
     CollectedAssistantMessage, FinishReason,
 };
-use vllm_engine_core_client::protocol::output::StopReason;
+use vllm_engine_core_client::protocol::StopReason;
 
 use self::convert::{ResponseOptions, prepare_chat_request};
 use crate::config::ApiServerOptions;
@@ -825,7 +825,7 @@ mod tests {
     use vllm_chat::{
         AssistantBlockKind, AssistantContentBlock, AssistantToolCall, ChatEvent, FinishReason,
     };
-    use vllm_engine_core_client::protocol::output::StopReason;
+    use vllm_engine_core_client::protocol::StopReason;
     use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
 
     use super::{

rust/src/server/src/routes/openai/utils/structured_outputs.rs

@@ -1,6 +1,6 @@
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;
+use vllm_engine_core_client::protocol::StructuredOutputsParams;
 
 use crate::error::ApiError;
 

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

@@ -24,15 +24,14 @@ use vllm_chat::{
     NewChatOutputProcessorOptions,
 };
 use vllm_engine_core_client::mock_engine::default_ready_response;
-use vllm_engine_core_client::protocol::decode_value;
 use vllm_engine_core_client::protocol::logprobs::{
     Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
 };
-use vllm_engine_core_client::protocol::output::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, StopReason,
-};
-use vllm_engine_core_client::protocol::request::EngineCoreRequest;
 use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
+    decode_value,
+};
 use vllm_engine_core_client::test_utils::{
     IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready,
 };

rust/src/server/src/state.rs

@@ -1,5 +1,6 @@
+use std::sync::Arc;
+use std::sync::OnceLock;
 use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::{Arc, OnceLock};
 
 use serde_json::Value;
 use sha2::{Digest, Sha256};

rust/src/text/Cargo.toml

@@ -10,7 +10,6 @@ asynk-strim-attr.workspace = true
 easy-ext.workspace = true
 enum-as-inner.workspace = true
 futures.workspace = true
-hf-hub.workspace = true
 itertools.workspace = true
 reqwest.workspace = true
 serde.workspace = true
@@ -22,6 +21,7 @@ tracing.workspace = true
 trait-set.workspace = true
 vllm-engine-core-client.workspace = true
 vllm-llm.workspace = true
+vllm-model-files.workspace = true
 vllm-tokenizer.workspace = true
 
 [dev-dependencies]

rust/src/text/src/backend/hf/config.rs

@@ -1,24 +1,18 @@
 use std::collections::BTreeSet;
-use std::fs;
 use std::path::Path;
 
 use serde::{Deserialize, Serialize};
-use thiserror_ext::AsReport as _;
+use vllm_model_files::read_json_file;
 
 use crate::error::{Error, Result};
 
 /// Minimal subset of `tokenizer_config.json` needed by chat/EOS handling.
 #[derive(Debug, Default, Deserialize)]
 #[serde(default)]
-pub struct HfTokenizerConfig {
+pub struct TokenizerConfig {
     #[serde(flatten)]
-    pub special_tokens: HfSpecialTokens,
+    pub special_tokens: SpecialTokens,
     pub chat_template: Option<String>,
-    /// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
-    /// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
-    /// containing "Tiktoken" which can be used as a hint for backend
-    /// selection.
-    pub tokenizer_class: Option<String>,
 }
 
 /// Hugging Face named special tokens may be serialized as a string or an
@@ -61,14 +55,14 @@ impl NamedSpecialToken {
 #[serde_with::skip_serializing_none]
 #[derive(Debug, Clone, Default, Deserialize, Serialize)]
 #[serde(default)]
-pub struct HfSpecialTokens {
+pub struct SpecialTokens {
     pub bos_token: Option<NamedSpecialToken>,
     pub eos_token: Option<NamedSpecialToken>,
     pub unk_token: Option<NamedSpecialToken>,
     pub pad_token: Option<NamedSpecialToken>,
 }
 
-impl HfSpecialTokens {
+impl SpecialTokens {
     /// Returns true if we don't discover any special tokens in the config.
     pub fn is_empty(&self) -> bool {
         self.bos_token.is_none()
@@ -234,42 +228,19 @@ impl ModelConfig {
     }
 }
 
-/// Load the tokenizer-side EOS metadata if a config file is present.
-pub fn load_tokenizer_config(path: Option<&Path>) -> Result<HfTokenizerConfig> {
-    read_json_file(path)
-}
-
 /// Load the generation-side EOS metadata if a config file is present.
 pub(super) fn load_generation_config(path: Option<&Path>) -> Result<GenerationConfig> {
-    read_json_file(path)
+    Ok(read_json_file(path)?)
+}
+
+/// Load the tokenizer-side EOS metadata if a config file is present.
+pub fn load_tokenizer_config(path: Option<&Path>) -> Result<TokenizerConfig> {
+    Ok(read_json_file(path)?)
 }
 
 /// Load the model-side config (`config.json`) if present.
 pub fn load_model_config(path: Option<&Path>) -> Result<ModelConfig> {
-    read_json_file(path)
-}
-
-fn read_json_file<T>(path: Option<&Path>) -> Result<T>
-where
-    T: for<'de> Deserialize<'de> + Default,
-{
-    let Some(path) = path else {
-        return Ok(T::default());
-    };
-    let content = fs::read_to_string(path).map_err(|error| {
-        Error::Tokenizer(format!(
-            "failed to read {}: {}",
-            path.display(),
-            error.as_report()
-        ))
-    })?;
-    serde_json::from_str(&content).map_err(|error| {
-        Error::Tokenizer(format!(
-            "failed to parse {}: {}",
-            path.display(),
-            error.as_report()
-        ))
-    })
+    Ok(read_json_file(path)?)
 }
 
 #[cfg(test)]

rust/src/text/src/backend/hf/mod.rs

@@ -1,5 +1,4 @@
 mod config;
-mod model_files;
 
 use std::collections::BTreeSet;
 use std::sync::Arc;
@@ -9,12 +8,12 @@ use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, Tiktok
 
 use self::config::{GenerationConfig, load_generation_config};
 pub use self::config::{
-    HfSpecialTokens, HfTokenizerConfig, ModelConfig, NamedSpecialToken, load_model_config,
+    ModelConfig, NamedSpecialToken, SpecialTokens, TokenizerConfig, load_model_config,
     load_tokenizer_config,
 };
-pub use self::model_files::{ResolvedModelFiles, TokenizerSource};
 use crate::backend::{SamplingHints, TextBackend};
 use crate::error::Result;
+pub use vllm_model_files::{ResolvedModelFiles, TokenizerSource};
 
 fn load_tokenizer(tokenizer: &TokenizerSource) -> Result<DynTokenizer> {
     match tokenizer {
@@ -125,3 +124,67 @@ impl TextBackend for HfTextBackend {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
+
+    use super::{ResolvedModelFiles, TokenizerSource};
+
+    #[tokio::test]
+    #[ignore = "too slow for CI and requires network access to Hugging Face"]
+    async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
+        let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
+            .await
+            .expect("resolve real Kimi K2.5 model files");
+
+        let tokenizer_path = match &files.tokenizer {
+            TokenizerSource::Tiktoken(path) => path.clone(),
+            other => panic!("expected tiktoken tokenizer source, got {other:?}"),
+        };
+
+        for backend in [
+            TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
+            TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
+        ] {
+            let think_id = backend.token_to_id("<think>").expect("resolve <think>");
+            let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
+            let tool_section_id = backend
+                .token_to_id("<|tool_calls_section_begin|>")
+                .expect("resolve tool call section marker");
+            let contraction_heavy_text =
+                "I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
+            let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
+
+            assert_eq!(
+                (think_id, end_think_id, tool_section_id),
+                (163606, 163607, 163595)
+            );
+            assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
+            assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
+            assert_eq!(
+                backend.decode(&[tool_section_id], true).unwrap(),
+                "<|tool_calls_section_begin|>"
+            );
+
+            // This demonstrates that we're using Kimi's custom BPE pattern.
+            // With CL100K this will be 23 tokens instead.
+            assert_eq!(
+                contraction_heavy_ids,
+                vec![
+                    17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
+                    1573, 56229, 13922, 13,
+                ]
+            );
+            assert_eq!(contraction_heavy_ids.len(), 17);
+            assert_eq!(
+                backend.decode(&contraction_heavy_ids, false).unwrap(),
+                contraction_heavy_text
+            );
+
+            // Special-looking text that is not actually registered should fail gracefully.
+            assert_eq!(backend.token_to_id("◁think▷"), None);
+            assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
+        }
+    }
+}

rust/src/text/src/error.rs

@@ -1,6 +1,7 @@
 use thiserror::Error;
 use vllm_engine_core_client::Error as EngineCoreError;
 use vllm_llm::Error as LlmError;
+use vllm_model_files::Error as ModelFilesError;
 
 pub use crate::lower::logprobs::LogprobsError;
 pub use crate::lower::token_ids::TokenIdsError;
@@ -20,6 +21,8 @@ pub enum Error {
     Logprobs(#[from] LogprobsError),
     #[error(transparent)]
     TokenIds(#[from] TokenIdsError),
+    #[error(transparent)]
+    ModelFiles(#[from] ModelFilesError),
     #[error(
         "`min_tokens` must be less than or equal to `max_tokens`, \
          got min_tokens={min_tokens}, max_tokens={max_tokens}"

rust/src/text/src/lower.rs

@@ -3,15 +3,15 @@ use std::collections::BTreeSet;
 pub(crate) mod logprobs;
 pub(crate) mod token_ids;
 
-use logprobs::validate_logprobs;
-use token_ids::{validate_prompt_token_ids, validate_vocab_range};
-use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
 use vllm_llm::GenerateRequest;
 use vllm_tokenizer::Tokenizer;
 
 use crate::backend::{SamplingHints, SamplingLimits};
 use crate::error::{Error, Result};
 use crate::request::{SamplingParams, TextRequest};
+use logprobs::validate_logprobs;
+use token_ids::{validate_prompt_token_ids, validate_vocab_range};
 
 /// One text request after it has been lowered into the raw generate boundary.
 #[derive(Debug)]

rust/src/text/src/lower/logprobs.rs

@@ -3,9 +3,8 @@
 //! `-1` is expanded only for bounds checks. The original request values are
 //! passed through to engine-core.
 
-use thiserror::Error;
-
 use crate::backend::SamplingLimits;
+use thiserror::Error;
 
 #[derive(Debug, Error)]
 pub enum LogprobsError {

rust/src/text/src/lower/token_ids.rs

@@ -1,7 +1,7 @@
 use std::result::Result;
 
 use thiserror::Error;
-use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
 
 use crate::SamplingLimits;
 

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

@@ -5,7 +5,7 @@ use futures::{Stream, StreamExt};
 use serde::{Deserialize, Serialize};
 use tracing::{Level, debug, trace};
 use vllm_engine_core_client::AbortCause;
-use vllm_engine_core_client::protocol::output::StopReason;
+use vllm_engine_core_client::protocol::StopReason;
 use vllm_llm::{FinishReason, GenerateOutput, TokenUsage};
 use vllm_tokenizer::{DynTokenizer, IncrementalDecoder};
 

rust/src/text/src/request.rs

@@ -5,8 +5,7 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use vllm_engine_core_client::protocol::lora::LoraRequest;
 use vllm_engine_core_client::protocol::multimodal::MmFeatures;
-use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
-use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;
+use vllm_engine_core_client::protocol::{ReasoningParserKwargs, StructuredOutputsParams};
 
 use crate::error::{Error, Result};
 use crate::output::TextDecodeOptions;

rust/src/tokenizer/src/hf/added_tokens.rs

@@ -1,12 +1,11 @@
-use std::fs;
-use std::path::Path;
-
 use serde::{Deserialize, Serialize};
 use thiserror_ext::AsReport as _;
 use tracing::warn;
 
 use crate::Result;
 
+use std::{fs, path::Path};
+
 /// Minimal `tokenizer.json` projection used to patch `added_tokens` while
 /// preserving the rest of the tokenizer definition verbatim.
 #[derive(Debug, Deserialize, Serialize)]

tests/benchmarks/test_audio_dataset.py

@@ -34,16 +34,6 @@ class _Tokenizer:
         return _TokenizedPrompt(prompt)
 
 
-class CohereAsrTokenizer(_Tokenizer):
-    def __init__(self, name_or_path: str = "/models/cohere-transcribe") -> None:
-        super().__init__(name_or_path)
-
-
-class _CohereNameOnlyTokenizer(_Tokenizer):
-    def __init__(self) -> None:
-        super().__init__("cohere/some-local-checkpoint")
-
-
 def _write_wav(path: Path, duration_s: float = 0.1, sample_rate: int = 16_000) -> None:
     num_samples = int(duration_s * sample_rate)
     sf.write(path, np.zeros(num_samples, dtype=np.float32), sample_rate)
@@ -208,64 +198,3 @@ def test_async_request_openai_audio_handles_decoded_audio_arrays(
     assert session.uploaded_bytes is not None
     assert output.success is True
     assert output.generated_text == "hello"
-
-
-_COHERE_ASR_PROMPT = (
-    "<|startofcontext|><|startoftranscript|>"
-    "<|emo:undefined|><|en|><|en|><|pnc|><|noitn|>"
-    "<|notimestamp|><|nodiarize|>"
-)
-
-
-def _make_asr_dataset(tmp_path: Path) -> datasets_module.ASRDataset:
-    audio_path = tmp_path / "sample.wav"
-    _write_wav(audio_path, duration_s=0.1)
-    dataset = object.__new__(datasets_module.ASRDataset)
-    dataset.data = [
-        {
-            "audio": {"path": str(audio_path), "bytes": None},
-            "text": "hello world",
-        }
-    ]
-    return dataset
-
-
-def test_asr_dataset_cohere_class_name_gets_decoder_prompt(tmp_path: Path) -> None:
-    dataset = _make_asr_dataset(tmp_path)
-    samples = dataset.sample(
-        tokenizer=CohereAsrTokenizer(),
-        num_requests=1,
-        output_len=32,
-        asr_min_audio_len_sec=0.0,
-        asr_max_audio_len_sec=1.0,
-    )
-    assert len(samples) == 1
-    assert samples[0].prompt == _COHERE_ASR_PROMPT
-
-
-def test_asr_dataset_cohere_name_or_path_fallback_gets_decoder_prompt(
-    tmp_path: Path,
-) -> None:
-    dataset = _make_asr_dataset(tmp_path)
-    samples = dataset.sample(
-        tokenizer=_CohereNameOnlyTokenizer(),
-        num_requests=1,
-        output_len=32,
-        asr_min_audio_len_sec=0.0,
-        asr_max_audio_len_sec=1.0,
-    )
-    assert len(samples) == 1
-    assert samples[0].prompt == _COHERE_ASR_PROMPT
-
-
-def test_asr_dataset_unknown_tokenizer_gets_empty_prompt(tmp_path: Path) -> None:
-    dataset = _make_asr_dataset(tmp_path)
-    samples = dataset.sample(
-        tokenizer=_Tokenizer(name_or_path="some-other/asr-model"),
-        num_requests=1,
-        output_len=32,
-        asr_min_audio_len_sec=0.0,
-        asr_max_audio_len_sec=1.0,
-    )
-    assert len(samples) == 1
-    assert samples[0].prompt == ""

tests/models/registry.py

@@ -1407,19 +1407,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         max_num_seqs=32,
         min_transformers_version="4.56.3",  # Required for Qwen3Next
     ),
-    # [DSpark]
-    "DSparkDraftModel": _HfExamplesInfo(
-        "deepseek-ai/DeepSeek-V4-Pro-DSpark",
-        speculative_model="deepseek-ai/DeepSeek-V4-Pro-DSpark",  # draft in mtp.*
-        is_available_online=False,
-        use_original_num_layers=True,  # DSpark has >1 draft block
-    ),
-    "Qwen3DSparkModel": _HfExamplesInfo(
-        "Qwen/Qwen3-8B",
-        speculative_model="deepseek-ai/dspark_qwen3_8b_block7",
-        is_available_online=False,
-        use_original_num_layers=True,  # DSpark backbone requires all layers
-    ),
     # [Eagle]
     "EagleCohereForCausalLM": _HfExamplesInfo(
         "/host/engines/cohere-moe",

tests/models/test_registry.py

@@ -48,10 +48,6 @@ def test_registry_imports(model_arch):
                 "(see #41376)"
             )
 
-    # DSpark draft model is NVIDIA-only; class is stubbed to None on ROCm/XPU.
-    if model_arch == "DSparkDraftModel" and not current_platform.is_cuda():
-        pytest.skip("DSparkDraftModel is only supported on CUDA")
-
     # Ensure all model classes can be imported successfully
     model_cls = ModelRegistry._try_load_model_cls(model_arch)
     assert model_cls is not None

tests/utils_/test_extensible_tensor.py

@@ -1,141 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-from vllm.utils.extensible_tensor import ExtensibleTensor
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
-
-
-def test_extensible_tensor_grows_without_moving() -> None:
-    buffer = ExtensibleTensor(4096, device="cuda")
-    try:
-        base_ptr = buffer.base_ptr
-        first_view = buffer.resize_(1024)
-        assert first_view.data_ptr() == base_ptr
-        first_view.fill_(7)
-
-        second_view = buffer.resize_(2048)
-        assert second_view.data_ptr() == base_ptr
-        assert torch.equal(second_view[:1024], torch.full_like(second_view[:1024], 7))
-
-        second_view[1024:].fill_(3)
-        assert torch.equal(buffer.tensor, second_view)
-
-        full_view = buffer.full_view()
-        assert full_view.data_ptr() == base_ptr
-        assert full_view.numel() == 4096
-    finally:
-        buffer.free()
-
-
-def test_extensible_tensor_rejects_shrink_and_overflow() -> None:
-    buffer = ExtensibleTensor(1024, device="cuda")
-    try:
-        buffer.resize_(512)
-        with pytest.raises(ValueError, match="grow-only"):
-            buffer.resize_(256)
-        with pytest.raises(ValueError, match="exceeds the segment capacity"):
-            buffer.resize_(1025)
-    finally:
-        buffer.free()
-
-
-def test_segments_grow_in_lockstep_and_zero_new() -> None:
-    """Each segment's committed prefix grows in lockstep.
-
-    Data written to a segment's committed prefix survives a grow; the newly
-    committed range of each segment is zeroed with `zero_new=True` while old
-    bytes are preserved.
-    """
-    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
-    try:
-        assert et.num_segments == 2
-        assert et.segment_capacity_bytes == 4096
-
-        et.resize_per_segment_(256, zero_new=True)
-        assert et.bytes_per_segment == 256
-        assert et.num_bytes == 512
-        fv = et.full_view()
-        assert fv.shape == (8192,)
-        # Committed prefixes start zeroed.
-        assert torch.count_nonzero(fv[:256]) == 0
-        assert torch.count_nonzero(fv[4096 : 4096 + 256]) == 0
-
-        pattern_a = torch.arange(256, device="cuda", dtype=torch.uint8)
-        pattern_b = 255 - pattern_a
-        fv[:256].copy_(pattern_a)
-        fv[4096 : 4096 + 256].copy_(pattern_b)
-
-        et.resize_per_segment_(1024, zero_new=True)
-        fv2 = et.full_view()
-        assert fv2.data_ptr() == fv.data_ptr()
-        # Old bytes of both segments preserved; freshly committed ranges zeroed.
-        assert torch.equal(fv2[:256], pattern_a)
-        assert torch.equal(fv2[4096 : 4096 + 256], pattern_b)
-        assert torch.count_nonzero(fv2[256:1024]) == 0
-        assert torch.count_nonzero(fv2[4096 + 256 : 4096 + 1024]) == 0
-    finally:
-        et.free()
-
-
-def test_segments_at_granularity_scale() -> None:
-    """Segments spanning multiple mapping granules commit correctly.
-
-    Uses a segment capacity that is not a multiple of the allocation
-    granularity, so a granule straddles the segment boundary and is shared by
-    the first commit of one segment and a later commit of the other -- it must
-    be mapped exactly once.
-    """
-    probe = ExtensibleTensor(max_num_bytes=1, device="cuda")
-    granularity = probe.capacity_bytes
-    probe.free()
-    # Two segments of 1.5 granules each; the middle granule straddles the
-    # boundary.
-    max_num_bytes = 3 * granularity
-    et = ExtensibleTensor(max_num_bytes=max_num_bytes, device="cuda", num_segments=2)
-    try:
-        seg = et.segment_capacity_bytes
-        assert seg == max_num_bytes // 2
-
-        step = granularity // 2
-        et.resize_per_segment_(step, zero_new=True)
-        fv = et.full_view()
-        fv[:step].fill_(1)
-        fv[seg : seg + step].fill_(2)
-
-        # Grow to the full segment capacity: previously mapped granules
-        # (including the boundary-straddling one) are reused, new ones are
-        # committed and zeroed.
-        et.resize_per_segment_(seg, zero_new=True)
-        fv2 = et.full_view()
-        assert torch.all(fv2[:step] == 1)
-        assert torch.all(fv2[seg : seg + step] == 2)
-        assert torch.count_nonzero(fv2[step:seg]) == 0
-        assert torch.count_nonzero(fv2[seg + step :]) == 0
-    finally:
-        et.free()
-
-
-def test_multi_segment_invalid_usage_raises() -> None:
-    """Prefix-view APIs and invalid segment configs raise for multi-segment
-    buffers."""
-    with pytest.raises(ValueError):
-        ExtensibleTensor(max_num_bytes=100, device="cuda", num_segments=3)
-
-    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
-    try:
-        with pytest.raises(ValueError):
-            _ = et.tensor
-        with pytest.raises(ValueError):
-            et.resize_(256)
-
-        et.resize_per_segment_(256)
-        with pytest.raises(ValueError):
-            et.resize_per_segment_(128)  # shrink
-        with pytest.raises(ValueError):
-            et.resize_per_segment_(et.segment_capacity_bytes + 1)  # over capacity
-    finally:
-        et.free()

tests/v1/attention/test_dspark_noncausal_sparse_mla.py

@@ -1,529 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Correctness tests for DSpark non-causal sliding-window MLA via sparse indices.
-
-DSpark drafts a block of N tokens whose attention is NON-CAUSAL within the block:
-every block token attends to the sliding window of context AND to all block
-tokens (including ones at later positions than itself).
-
-We can implement this using the existing sparse-MLA pathway by expanding the window size
-to include the rest of the block tokens: instead of setting topk indices to the 127
-previous tokens, we expand it to the next power of 2 (256) and include up to
-swa_size + block_size - 1 topk indices, so that each query attends to the rest. The
-remaining slots are filled with padding.
-
-The sparse-MLA decode kernels (FlashMLA on SM90/SM100, FlashInfer TRTLLM on
-SM100/SM120) are index-driven: each query attends over exactly the slots in its
-index list, with no causal mask (see ``flash_mla_with_kvcache(..., indices=...)``
-and ``_forward_decode``'s "attend only by generated indices"). The existing
-``test_sparse_mla_backends`` suite already validates arbitrary index lists, but
-only ones whose entries are <= the query's own position. This test suite specifically
-ensures correctness of the non-causal attention case.
-
-This reuses the harness/helpers of ``test_sparse_mla_backends.py`` (same model
-shapes, fp8_ds_mla round-trip, mock indexer, MockSparseMLAAttentionLayer); only
-the index construction differs.
-"""
-
-import math
-from types import MethodType, SimpleNamespace
-
-import pytest
-import torch
-
-from tests.v1.attention.test_mla_backends import (
-    BatchSpec,
-    MockSparseMLAAttentionLayer,
-    create_and_prepopulate_kv_cache,
-)
-from tests.v1.attention.test_sparse_mla_backends import (
-    _quantize_dequantize_fp8_ds_mla,
-)
-from tests.v1.attention.utils import (
-    create_common_attn_metadata,
-    create_standard_kv_cache_spec,
-    create_vllm_config,
-)
-from vllm.config import set_current_vllm_config
-from vllm.model_executor.layers.linear import ColumnParallelLinear
-from vllm.platforms import current_platform
-
-if not current_platform.is_cuda():
-    pytest.skip(
-        "DSpark non-causal sparse MLA tests currently only support CUDA.",
-        allow_module_level=True,
-    )
-
-from vllm.utils.math_utils import cdiv
-from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
-    FlashInferMLASparseTRTLLMBackend,
-)
-from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend
-from vllm.v1.attention.ops import flashmla
-
-DEVICE_TYPE = current_platform.device_type
-
-# (window, block_size, topk_width). topk_width must be a multiple of the kernel's
-# B_TOPK (= padded query-head count, 64 or 128); we use 128-multiples to cover
-# both. The "wide" case needs window + block > 128 -> width must grow past 128.
-_DSPARK_CONFIGS = {
-    "small_block": (8, 4, 128),
-    "full_window_block": (128, 5, 256),
-}
-
-
-def _build_dspark_noncausal_indices(
-    seq_lens: list[int],
-    query_lens: list[int],
-    window: int,
-    topk_width: int,
-    device: torch.device,
-) -> torch.Tensor:
-    """Per-token sparse indices for the DSpark non-causal block.
-
-    For a request with context length ``ctx`` and a query block of ``q_len``
-    tokens (block positions ``ctx .. ctx+q_len-1``), EVERY block query attends to
-    the same set: the trailing ``window`` context positions plus all block
-    positions, i.e. the contiguous range ``[max(ctx-window,0) .. ctx+q_len-1]``.
-    This is non-causal: an early block query's list contains later block tokens
-    (future-pointing). The list is padded to ``topk_width`` with ``-1``.
-    """
-    total_query_tokens = sum(query_lens)
-    sparse_indices = torch.full(
-        (total_query_tokens, topk_width), -1, dtype=torch.int32, device=device
-    )
-    gt = 0
-    for s_len, q_len in zip(seq_lens, query_lens):
-        ctx_len = s_len - q_len
-        lo = max(ctx_len - window, 0)
-        hi = ctx_len + q_len  # exclusive: window context + the full block
-        idx_list = torch.arange(lo, hi, dtype=torch.int32, device=device)
-        n = idx_list.numel()
-        assert n <= topk_width, (
-            f"index list ({n}) exceeds aligned topk width ({topk_width})"
-        )
-        for _ in range(q_len):
-            sparse_indices[gt, :n] = idx_list
-            gt += 1
-    return sparse_indices
-
-
-def _run_sparse_backend_vs_sdpa(
-    backend_cls,
-    seq_lens: list[int],
-    query_lens: list[int],
-    sparse_indices: torch.Tensor,
-    kv_cache_dtype: str,
-    block_size: int,
-    num_heads: int,
-    device: torch.device,
-    force_future_dominance: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Run a sparse-MLA backend with the given per-token indices and compute a
-    dense per-token SDPA reference over the SAME indices.
-
-    Mirrors ``test_sparse_mla_backends.test_sparse_backend_decode_correctness``
-    but with externally-supplied (non-causal) ``sparse_indices``.
-
-    ``num_heads`` selects the kernel's B_TOPK (= padded q-head count): 128 -> 128,
-    64 -> 64. The aligned widths (128/256) are multiples of both, so num_heads=64
-    exercises the head64 decode path that the SM100 alignment assert guards.
-
-    ``force_future_dominance`` scales the LAST block token's latent KV so it
-    dominates the softmax for every query that attends to it. With random data the
-    few future block tokens carry negligible attention mass (especially with a wide
-    window), so causal and non-causal outputs coincide; this knob makes the
-    future-token contribution provably large for the differentiation test. It is
-    OFF for the correctness test (which needs sensitivity to all tokens).
-
-    Returns (backend_output, noncausal_reference, causal_reference). The causal
-    reference restricts each query to indices <= its own absolute position.
-    """
-    batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
-    topk_tokens = sparse_indices.shape[1]
-    dtype = torch.bfloat16
-    use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
-
-    kv_lora_rank = 512
-    qk_nope_head_dim = 128
-    qk_rope_head_dim = 64
-    v_head_dim = 128
-    head_size = kv_lora_rank + qk_rope_head_dim
-
-    max_seqlen = max(seq_lens)
-    total_cache_tokens = sum(seq_lens)
-
-    vllm_config = create_vllm_config(
-        model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
-        tensor_parallel_size=1,
-        max_model_len=max_seqlen,
-        num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
-        block_size=block_size,
-        hf_config_override={
-            "index_topk": topk_tokens,
-            "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
-        },
-    )
-    model_config = vllm_config.model_config
-    model_config.hf_text_config = SimpleNamespace(
-        q_lora_rank=None,
-        kv_lora_rank=kv_lora_rank,
-        qk_nope_head_dim=qk_nope_head_dim,
-        qk_rope_head_dim=qk_rope_head_dim,
-        v_head_dim=v_head_dim,
-        model_type="deepseek_v2",
-    )
-    model_config.dtype = dtype
-    model_config.get_num_attention_heads = MethodType(
-        lambda self, parallel_config: num_heads, model_config
-    )
-    model_config.get_num_kv_heads = MethodType(
-        lambda self, parallel_config: 1, model_config
-    )
-    model_config.get_head_size = MethodType(lambda self: head_size, model_config)
-    model_config.get_sliding_window = MethodType(lambda self: None, model_config)
-
-    kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
-
-    torch.manual_seed(0)
-    scale = 1.0 / math.sqrt(head_size)
-
-    # Shared MLA projection weights, used by both reference and backend.
-    W_UK = torch.rand(
-        kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device
-    )
-    W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
-
-    all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
-    kv_c_contexts, k_pe_contexts = [], []
-    reference_outputs = []
-    # Causal counterpart of the reference: same index lists, but each query is
-    # restricted to indices <= its own absolute position (drops future-pointing
-    # block tokens). Used to prove the non-causal result is genuinely different.
-    causal_reference_outputs = []
-
-    kv_cache_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
-    global_token_idx = 0
-
-    for s_len, q_len in zip(seq_lens, query_lens):
-        ctx_len = s_len - q_len
-
-        q_c = torch.rand(
-            q_len,
-            num_heads,
-            qk_nope_head_dim + qk_rope_head_dim,
-            dtype=dtype,
-            device=device,
-        )
-        kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
-        k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
-
-        if force_future_dominance:
-            # Scale the last block token's latent KV so its key/value dominate the
-            # softmax for any query attending to it. 4x in the latent dot makes its
-            # pre-softmax score exceed the others by a wide margin, so non-causal
-            # queries (which include it) diverge sharply from causal ones (which,
-            # for all but the last query, exclude it). Applied before quantization
-            # so cache and reference stay consistent.
-            kv_c_full[s_len - 1] = kv_c_full[s_len - 1] * 4.0 + 2.0
-
-        if use_fp8_ds_mla_quantization:
-            is_sm100 = torch.cuda.get_device_capability()[0] >= 10
-            kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla(
-                kv_c_full,
-                k_pe_full.squeeze(1),
-                block_size=block_size,
-                scale=kv_cache_scale,
-                simulate_sm100_e8m0_scales=is_sm100,
-            )
-            k_pe_full = k_pe_squeezed.unsqueeze(1)
-
-        q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
-        ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)
-        q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
-
-        k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
-        v_mqa = kv_c_full
-
-        # Per-token sparse SDPA reference over the supplied (non-causal) indices.
-        def _sparse_sdpa(idx_tensor, q_tok, k_mqa=k_mqa, v_mqa=v_mqa):
-            k_sparse = k_mqa[idx_tensor].unsqueeze(1).expand(-1, num_heads, -1)
-            v_sparse = v_mqa[idx_tensor].unsqueeze(1).expand(-1, num_heads, -1)
-            out = torch.nn.functional.scaled_dot_product_attention(
-                q_tok.unsqueeze(0).transpose(1, 2),
-                k_sparse.unsqueeze(0).transpose(1, 2),
-                v_sparse.unsqueeze(0).transpose(1, 2),
-                scale=scale,
-            )
-            out = out.transpose(1, 2).squeeze(0)
-            out = torch.einsum("qnl,lnv->qnv", out, W_UV)
-            return out.flatten(start_dim=-2)
-
-        for q_idx in range(q_len):
-            tok_sparse_idx = sparse_indices[global_token_idx]
-            valid_indices = tok_sparse_idx[tok_sparse_idx >= 0].long()
-
-            q_tok = q_mqa[q_idx : q_idx + 1]
-            reference_outputs.append(_sparse_sdpa(valid_indices, q_tok))
-
-            # Causal: drop indices pointing past this query's own position.
-            abs_pos = ctx_len + q_idx
-            causal_indices = valid_indices[valid_indices <= abs_pos]
-            causal_reference_outputs.append(_sparse_sdpa(causal_indices, q_tok))
-            global_token_idx += 1
-
-        all_q_vllm.append(q_c)
-        all_kv_c_vllm.append(kv_c_full[ctx_len:])
-        all_k_pe_vllm.append(k_pe_full[ctx_len:])
-        kv_c_contexts.append(kv_c_full[: ctx_len + 1])
-        k_pe_contexts.append(k_pe_full[: ctx_len + 1])
-
-    query_vllm = torch.cat(all_q_vllm, dim=0)
-    kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
-    k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
-    sdpa_reference = torch.cat(reference_outputs, dim=0)
-    causal_reference = torch.cat(causal_reference_outputs, dim=0)
-
-    vllm_config.cache_config.cache_dtype = kv_cache_dtype
-    vllm_config.model_config.hf_config.index_topk = topk_tokens
-
-    common_attn_metadata = create_common_attn_metadata(
-        batch_spec, block_size, device, arange_block_indices=True
-    )
-    kv_cache = create_and_prepopulate_kv_cache(
-        kv_c_contexts=kv_c_contexts,
-        k_pe_contexts=k_pe_contexts,
-        block_size=block_size,
-        head_size=head_size,
-        dtype=dtype,
-        device=device,
-        num_blocks=vllm_config.cache_config.num_gpu_blocks,
-        common_attn_metadata=common_attn_metadata,
-        randomize_blocks=False,
-        kv_cache_dtype=kv_cache_dtype,
-        scale=kv_cache_scale,
-    )
-
-    builder = backend_cls.get_builder_cls()(
-        kv_cache_spec, ["placeholder"], vllm_config, device
-    )
-    metadata = builder.build(
-        common_prefix_len=0, common_attn_metadata=common_attn_metadata
-    )
-
-    mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)
-
-    kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1).view(
-        kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
-    )
-    mock_kv_b_proj = ColumnParallelLinear(
-        input_size=kv_lora_rank,
-        output_size=num_heads * (qk_nope_head_dim + v_head_dim),
-        bias=False,
-    ).to(device=device, dtype=dtype)
-    mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
-
-    with set_current_vllm_config(vllm_config):
-        impl = backend_cls.get_impl_cls()(
-            num_heads=num_heads,
-            head_size=head_size,
-            scale=scale,
-            num_kv_heads=1,
-            alibi_slopes=None,
-            sliding_window=None,
-            kv_cache_dtype=vllm_config.cache_config.cache_dtype,
-            logits_soft_cap=None,
-            attn_type="decoder",
-            kv_sharing_target_layer_name=None,
-            q_lora_rank=None,
-            kv_lora_rank=kv_lora_rank,
-            qk_nope_head_dim=qk_nope_head_dim,
-            qk_rope_head_dim=qk_rope_head_dim,
-            qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
-            v_head_dim=v_head_dim,
-            kv_b_proj=mock_kv_b_proj,
-            indexer=mock_indexer,
-        )
-        impl.process_weights_after_loading(dtype)
-        mock_layer = MockSparseMLAAttentionLayer(
-            impl=impl,
-            num_heads=num_heads,
-            qk_nope_head_dim=qk_nope_head_dim,
-            qk_rope_head_dim=qk_rope_head_dim,
-            v_head_dim=v_head_dim,
-            kv_lora_rank=kv_lora_rank,
-            device=device,
-            W_UK=W_UK,
-            W_UV=W_UV,
-            q_scale=1.0,
-            k_scale=1.0,
-        )
-
-    out_buffer = torch.empty(
-        metadata.num_actual_tokens, num_heads * v_head_dim, dtype=dtype, device=device
-    )
-    with torch.inference_mode():
-        backend_output = mock_layer.forward_impl(
-            query_vllm, kv_c_vllm, k_pe_vllm, kv_cache, metadata, out_buffer
-        )
-    return backend_output, sdpa_reference, causal_reference
-
-
-def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: int):
-    if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
-        pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
-    if (
-        backend_cls is FlashMLASparseBackend
-        and kv_cache_dtype.startswith("fp8")
-        and kv_cache_dtype != "fp8_ds_mla"
-    ):
-        pytest.skip("FlashMLA Sparse fp8 only supports fp8_ds_mla kv-cache dtype")
-    if block_size not in backend_cls.get_supported_kernel_block_sizes():
-        pytest.skip(
-            f"{backend_cls.get_name()} does not support block_size={block_size}"
-        )
-    if backend_cls is FlashMLASparseBackend:
-        ok, reason = flashmla.is_flashmla_sparse_supported()
-        if not ok:
-            pytest.skip(reason)
-    elif backend_cls is FlashInferMLASparseTRTLLMBackend:
-        cap = current_platform.get_device_capability()
-        if cap is None or not backend_cls.supports_compute_capability(cap):
-            pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability")
-
-
-@pytest.mark.parametrize(
-    "backend_cls",
-    [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
-    ids=["FlashMLA", "FlashInferTRTLLM"],
-)
-@pytest.mark.parametrize("config_name", list(_DSPARK_CONFIGS.keys()))
-# Per backend, the skip logic routes fp8 to the supported flavor: FlashMLA tests
-# auto + fp8_ds_mla (and skips per-tensor "fp8", which it aliases to ds_mla);
-# FlashInfer TRTLLM tests auto + per-tensor "fp8" (and skips fp8_ds_mla, which it
-# does not implement). So both backends get a bf16 case and an fp8 case.
-@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_ds_mla", "fp8"])
-@pytest.mark.parametrize("block_size", [64])
-# h_q=128 -> B_TOPK=128; h_q=64 -> B_TOPK=64 (covers the head64 decode path the
-# SM100 alignment assert specifically guards). Aligned widths (128/256) satisfy both.
-@pytest.mark.parametrize("num_heads", [128, 64], ids=["h128", "h64"])
-def test_dspark_noncausal_sparse_mla_matches_sdpa(
-    default_vllm_config,
-    dist_init,
-    workspace_init,
-    backend_cls,
-    config_name,
-    kv_cache_dtype,
-    block_size,
-    num_heads,
-):
-    """Non-causal (window ∪ block, future-pointing) per-token indices must match
-    a dense SDPA reference over the same indices, for both sparse-MLA backends."""
-    _skip_if_backend_unavailable(backend_cls, kv_cache_dtype, block_size)
-
-    window, block, topk_width = _DSPARK_CONFIGS[config_name]
-    device = torch.device(DEVICE_TYPE)
-
-    # Decode-style batch: each request has `block` query tokens and enough
-    # context for a full sliding window.
-    seq_lens = [window + block + 123, window + block + 50]
-    query_lens = [block, block]
-
-    sparse_indices = _build_dspark_noncausal_indices(
-        seq_lens, query_lens, window, topk_width, device
-    )
-
-    # Sanity: the construction must actually be non-causal (an early block query
-    # must reference a later block position than itself).
-    ctx0 = seq_lens[0] - query_lens[0]
-    first_query_valid = sparse_indices[0][sparse_indices[0] >= 0]
-    assert int(first_query_valid.max()) >= ctx0 + query_lens[0] - 1, (
-        "expected the first block query to attend to a future block token"
-    )
-
-    backend_output, sdpa_reference, _ = _run_sparse_backend_vs_sdpa(
-        backend_cls,
-        seq_lens,
-        query_lens,
-        sparse_indices,
-        kv_cache_dtype,
-        block_size,
-        num_heads,
-        device,
-    )
-
-    assert backend_output.shape == sdpa_reference.shape
-    assert backend_output.dtype == sdpa_reference.dtype
-    assert torch.isfinite(backend_output).all()
-    if kv_cache_dtype.startswith("fp8"):
-        rtol, atol = 0.065, 0.05
-    else:
-        rtol, atol = 0.01, 0.01
-    torch.testing.assert_close(backend_output, sdpa_reference, rtol=rtol, atol=atol)
-
-
-@pytest.mark.parametrize(
-    "backend_cls",
-    [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
-    ids=["FlashMLA", "FlashInferTRTLLM"],
-)
-@pytest.mark.parametrize("config_name", list(_DSPARK_CONFIGS.keys()))
-@pytest.mark.parametrize("block_size", [64])
-def test_dspark_noncausal_differs_from_causal(
-    default_vllm_config,
-    dist_init,
-    workspace_init,
-    backend_cls,
-    config_name,
-    block_size,
-):
-    """Differentiation guard: prove the backend genuinely attends to the
-    future-pointing indices (not silently applying a causal mask, and not merely
-    coinciding with a causal result because future tokens carry little weight).
-
-    With random data the few future block tokens are a negligible fraction of the
-    attended set (especially with a wide window), so causal and non-causal outputs
-    are numerically indistinguishable -- that is correct physics, not a backend
-    bug. To make the check meaningful we use ``force_future_dominance`` so the last
-    block token dominates the softmax: the backend must then match the non-causal
-    reference and diverge sharply from the causal one. bf16 (``auto``) suffices;
-    the property is dtype-independent and fp8 correctness is covered above.
-    """
-    _skip_if_backend_unavailable(backend_cls, "auto", block_size)
-
-    window, block, topk_width = _DSPARK_CONFIGS[config_name]
-    device = torch.device(DEVICE_TYPE)
-    seq_lens = [window + block + 123, window + block + 50]
-    query_lens = [block, block]
-
-    sparse_indices = _build_dspark_noncausal_indices(
-        seq_lens, query_lens, window, topk_width, device
-    )
-
-    backend_output, sdpa_reference, causal_reference = _run_sparse_backend_vs_sdpa(
-        backend_cls,
-        seq_lens,
-        query_lens,
-        sparse_indices,
-        "auto",
-        block_size,
-        128,
-        device,
-        force_future_dominance=True,
-    )
-
-    # The two references must be clearly distinguishable for the check to mean
-    # anything (dominance guarantees this).
-    ref_gap = (sdpa_reference - causal_reference).abs().max().item()
-    assert ref_gap > 0.1, (
-        f"non-causal and causal references are too close (gap={ref_gap}); "
-        "force_future_dominance did not create a separable scenario"
-    )
-
-    # Backend must track the NON-causal reference, not the causal one.
-    torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)
-    causal_err = (backend_output - causal_reference).abs().max().item()
-    assert causal_err > 0.1, (
-        f"non-causal backend output matches the causal reference "
-        f"(max abs diff={causal_err}); future-pointing indices are not attended to"
-    )

tests/v1/core/test_encoder_cache_manager.py

@@ -149,30 +149,6 @@ def test_has_cache_restores_from_freeable():
     assert manager.num_freeable_slots == 6
 
 
-def test_make_profiling_reservation():
-    assert (
-        EncoderCacheManager.make_profiling_reservation(
-            cache_size=0,
-            embed_size=8,
-            dtype=torch.float16,
-            device="cpu",
-        )
-        is None
-    )
-
-    reservation = EncoderCacheManager.make_profiling_reservation(
-        cache_size=7,
-        embed_size=8,
-        dtype=torch.float16,
-        device="cpu",
-    )
-
-    assert reservation is not None
-    assert reservation.shape == (7, 8)
-    assert reservation.dtype == torch.float16
-    assert reservation.device.type == "cpu"
-
-
 def test_get_freed_mm_hashes_clears_freed_list():
     manager = EncoderCacheManager(cache_size=10)
     req1 = MockRequest("reqA", ["a"], [5])

tests/v1/e2e/spec_decode/test_spec_decode.py

@@ -1412,67 +1412,6 @@ def test_dflash_acceptance_rates(
     cleanup_dist_env_and_memory()
 
 
-@pytest.fixture
-def dspark_config():
-    target_model = "Qwen/Qwen3-4B-FP8"
-    draft_model = "deepseek-ai/dspark_qwen3_4b_block7"
-
-    return dict(
-        model=target_model,
-        trust_remote_code=True,
-        speculative_config={
-            "method": "dspark",
-            "model": draft_model,
-            "num_speculative_tokens": 7,
-            "attention_backend": "FLASH_ATTN",
-            "draft_sample_method": "probabilistic",
-        },
-        max_model_len=4096,
-        disable_log_stats=False,
-    )
-
-
-@single_gpu_only
-@large_gpu_mark(min_gb=24)
-def test_dspark_correctness_and_acceptance_rate(dspark_config):
-    """
-    E2E test for DSpark speculative decoding: acceptance rate/length
-    regression coverage plus GSM8K correctness, at temperature=1.0 to
-    exercise the probabilistic draft-sampling/rejection-sampling path
-    (not just greedy).
-
-    Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft
-    model. Reference: measured over 12 runs of the full GSM8K set at
-    temperature=1.0 (prefix caching disabled to avoid cross-run reuse):
-      accuracy:        min=0.782 max=0.814 mean=0.801
-      acceptance_rate: min=0.418 max=0.434 mean=0.428
-      acceptance_len:  min=3.928 max=4.037 mean=3.994
-    Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
-    """
-    spec_llm = LLM(**dspark_config)
-
-    results = evaluate_gsm8k_offline(spec_llm, temperature=1.0)
-    gsm8k_accuracy = results["accuracy"]
-
-    metrics = spec_llm.get_metrics()
-    acceptance_rate = compute_acceptance_rate(metrics)
-    acceptance_len = compute_acceptance_len(metrics)
-
-    print(
-        f"DSpark acceptance_rate={acceptance_rate:.2f}, "
-        f"acceptance_len={acceptance_len:.2f}, "
-        f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
-    )
-
-    assert acceptance_rate >= 0.428 * 0.9
-    assert acceptance_len >= 3.994 * 0.9
-    assert gsm8k_accuracy >= 0.801 * 0.9
-
-    del spec_llm
-    torch.accelerator.empty_cache()
-    cleanup_dist_env_and_memory()
-
-
 @single_gpu_only
 def test_synthetic_acceptance_rate():
     """Verify that synthetic rejection sampling produces an acceptance

tests/v1/engine/test_engine_args.py

@@ -49,18 +49,6 @@ def test_prefix_caching_from_cli():
         args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"])
 
 
-def test_extensible_kv_cache_from_cli():
-    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
-
-    args = parser.parse_args([])
-    engine_args = EngineArgs.from_cli_args(args=args)
-    assert not engine_args.enable_extensible_kv_cache
-
-    args = parser.parse_args(["--enable-extensible-kv-cache"])
-    engine_args = EngineArgs.from_cli_args(args=args)
-    assert engine_args.enable_extensible_kv_cache
-
-
 @pytest.mark.skipif(_xxhash is None, reason="xxhash not installed")
 def test_prefix_caching_xxhash_from_cli():
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())

tests/v1/tracing/test_tracing.py

@@ -25,10 +25,11 @@ def test_traces(
 ):
     with monkeypatch.context() as m:
         m.setenv(OTEL_EXPORTER_OTLP_TRACES_INSECURE, "true")
-        # The fake OTLP server starts gRPC worker threads before the engine
-        # core is launched. gRPC's C-core is not fork-safe and can segfault
-        # if forked.
-        m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+        if current_platform.is_rocm():
+            # The fake OTLP server starts gRPC worker threads before the engine
+            # core is launched. On ROCm CI, forking while those threads are
+            # active can segfault in gRPC during engine startup or teardown.
+            m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
 
         sampling_params = SamplingParams(
             temperature=0.01,

tests/v1/worker/test_extensible_kv_cache.py

@@ -1,369 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""GPU integration tests for the extensible KV cache allocation paths.
-
-Drives `GPUModelRunner._allocate_kv_cache_tensors` / `_reshape_kv_cache_tensors`
-/ `extend_kv_cache` directly with fake attention backends, covering the buffer
-layouts the extensible flow supports: block-major (one committed prefix),
-K/V-split (one prefix per half), Mamba (block-major per layer), and hybrid
-attention + Mamba (attention re-strided to block-major). Buffer sizes exceed
-the CUDA VMM allocation granularity so touching a block that the commit logic
-missed would fault instead of silently passing.
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm.v1.attention.backend import AttentionBackend
-from vllm.v1.kv_cache_interface import (
-    FullAttentionSpec,
-    KVCacheConfig,
-    KVCacheGroupSpec,
-    KVCacheTensor,
-    MambaSpec,
-)
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-from vllm.v1.worker.utils import AttentionGroup
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
-
-BLOCK_SIZE = 16
-NUM_BLOCKS = 256
-
-
-class _SplitKVBackend(AttentionBackend):
-    """Fake backend with a K/V-split layout, like FlashAttention."""
-
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (2, num_blocks, block_size, num_kv_heads, head_size)
-
-
-class _BlockMajorBackend(AttentionBackend):
-    """Fake backend with a num-blocks-first layout, like FlashInfer."""
-
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, 2, block_size, num_kv_heads, head_size)
-
-
-class _StrideOrderBackend(AttentionBackend):
-    """Fake backend whose stride order makes a kv-first shape block-major."""
-
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (2, num_blocks, block_size, num_kv_heads, head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (1, 0, 2, 3, 4)
-
-
-def _full_attention_spec() -> FullAttentionSpec:
-    # page_size_bytes = 2 (K+V) * 16 * 8 * 128 * 2 bytes = 64 KiB; 256 blocks
-    # = 16 MiB, several VMM granules per buffer.
-    return FullAttentionSpec(
-        block_size=BLOCK_SIZE,
-        num_kv_heads=8,
-        head_size=128,
-        dtype=torch.bfloat16,
-    )
-
-
-def _mamba_spec() -> MambaSpec:
-    # page_size_bytes = (8*128 + 16*64) * 4 bytes = 8 KiB per block per layer.
-    return MambaSpec(
-        block_size=BLOCK_SIZE,
-        shapes=((8, 128), (16, 64)),
-        dtypes=(torch.float32, torch.float32),
-    )
-
-
-def _make_runner(kv_cache_config: KVCacheConfig, attn_groups) -> GPUModelRunner:
-    runner = object.__new__(GPUModelRunner)
-    runner.device = torch.device("cuda:0")
-    runner.kv_cache_config = kv_cache_config
-    runner.attn_groups = attn_groups
-    runner.runner_only_attn_layers = set()
-    runner.cache_config = SimpleNamespace(cache_dtype="auto")
-    return runner
-
-
-def _attention_config(spec: FullAttentionSpec, backend) -> tuple[KVCacheConfig, list]:
-    kv_cache_config = KVCacheConfig(
-        num_blocks=NUM_BLOCKS,
-        kv_cache_tensors=[
-            KVCacheTensor(size=NUM_BLOCKS * spec.page_size_bytes, shared_by=["layer.0"])
-        ],
-        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)],
-    )
-    attn_groups = [
-        [
-            AttentionGroup(
-                backend=backend,
-                layer_names=["layer.0"],
-                kv_cache_spec=spec,
-                kv_cache_group_id=0,
-            )
-        ]
-    ]
-    return kv_cache_config, attn_groups
-
-
-def _free_buffers(runner: GPUModelRunner) -> None:
-    for buffer, _ in getattr(runner, "_extensible_kv_cache_buffers", []):
-        buffer.free()
-
-
-def test_kv_cache_num_segments_by_layer() -> None:
-    """Segment counts follow the physical layout of each layer's backend."""
-    spec = _full_attention_spec()
-    for backend, expected in (
-        (_SplitKVBackend, 2),
-        (_BlockMajorBackend, 1),
-        # kv-first logical shape but block-major physical order -> 1 segment.
-        (_StrideOrderBackend, 1),
-    ):
-        kv_cache_config, attn_groups = _attention_config(spec, backend)
-        runner = _make_runner(kv_cache_config, attn_groups)
-        assert runner._kv_cache_num_segments_by_layer() == {"layer.0": expected}
-
-
-def test_extensible_split_layout_grows_both_halves() -> None:
-    """A K/V-split layer keeps its natural layout and both halves grow in
-    lockstep."""
-    spec = _full_attention_spec()
-    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
-    runner = _make_runner(kv_cache_config, attn_groups)
-    try:
-        raw_tensors = runner._allocate_kv_cache_tensors(
-            kv_cache_config, extensible=True
-        )
-        kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
-        kv_cache = kv_caches["layer.0"]
-        assert kv_cache.shape == (2, NUM_BLOCKS, BLOCK_SIZE, 8, 128)
-        [(buffer, bytes_per_block_per_segment)] = runner._extensible_kv_cache_buffers
-        assert buffer.num_segments == 2
-        assert bytes_per_block_per_segment == spec.page_size_bytes // 2
-
-        # Only block 0 is committed -- in each half.
-        kv_cache[0, 0].fill_(1)  # K, block 0
-        kv_cache[1, 0].fill_(2)  # V, block 0
-        torch.cuda.synchronize()
-
-        runner.extend_kv_cache(NUM_BLOCKS)
-        # Old data survives the grow; new blocks are usable in both halves and
-        # zeroed.
-        assert torch.all(kv_cache[0, 0] == 1)
-        assert torch.all(kv_cache[1, 0] == 2)
-        kv_cache[0, NUM_BLOCKS - 1].fill_(3)
-        kv_cache[1, NUM_BLOCKS - 1].fill_(4)
-        torch.cuda.synchronize()
-        assert torch.all(kv_cache[0, NUM_BLOCKS - 1] == 3)
-        assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 4)
-        assert torch.count_nonzero(kv_cache[:, 1 : NUM_BLOCKS - 1]) == 0
-    finally:
-        _free_buffers(runner)
-
-
-def test_extensible_block_major_layout() -> None:
-    """A layer whose physical layout is block-major uses a single segment."""
-    spec = _full_attention_spec()
-    kv_cache_config, attn_groups = _attention_config(spec, _BlockMajorBackend)
-    runner = _make_runner(kv_cache_config, attn_groups)
-    try:
-        raw_tensors = runner._allocate_kv_cache_tensors(
-            kv_cache_config, extensible=True
-        )
-        kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
-        kv_cache = kv_caches["layer.0"]
-        assert kv_cache.shape == (NUM_BLOCKS, 2, BLOCK_SIZE, 8, 128)
-        [(buffer, bytes_per_block_per_segment)] = runner._extensible_kv_cache_buffers
-        assert buffer.num_segments == 1
-        assert bytes_per_block_per_segment == spec.page_size_bytes
-
-        kv_cache[0].fill_(1)
-        runner.extend_kv_cache(NUM_BLOCKS)
-        kv_cache[NUM_BLOCKS - 1].fill_(2)
-        torch.cuda.synchronize()
-        assert torch.all(kv_cache[0] == 1)
-        assert torch.all(kv_cache[NUM_BLOCKS - 1] == 2)
-        assert torch.count_nonzero(kv_cache[1 : NUM_BLOCKS - 1]) == 0
-    finally:
-        _free_buffers(runner)
-
-
-def test_legacy_split_layout_commits_everything() -> None:
-    """Without `extensible`, the full buffer is committed up front."""
-    spec = _full_attention_spec()
-    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
-    runner = _make_runner(kv_cache_config, attn_groups)
-    raw_tensors = runner._allocate_kv_cache_tensors(kv_cache_config, extensible=False)
-    kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
-    kv_cache = kv_caches["layer.0"]
-    kv_cache[0, NUM_BLOCKS - 1].fill_(1)
-    kv_cache[1, NUM_BLOCKS - 1].fill_(2)
-    torch.cuda.synchronize()
-    assert torch.all(kv_cache[0, NUM_BLOCKS - 1] == 1)
-    assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 2)
-    with pytest.raises(RuntimeError, match="extensible"):
-        runner.extend_kv_cache(NUM_BLOCKS)
-
-
-def test_extensible_mamba_grows_per_layer() -> None:
-    """Mamba per-layer buffers are block-major and grow with the KV cache."""
-    spec = _mamba_spec()
-    num_blocks = 512
-    layer_names = ["mamba.0", "mamba.1"]
-    kv_cache_config = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=[
-            KVCacheTensor(size=num_blocks * spec.page_size_bytes, shared_by=[name])
-            for name in layer_names
-        ],
-        kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)],
-    )
-    attn_groups = [
-        [
-            AttentionGroup(
-                backend=_BlockMajorBackend,
-                layer_names=layer_names,
-                kv_cache_spec=spec,
-                kv_cache_group_id=0,
-            )
-        ]
-    ]
-    runner = _make_runner(kv_cache_config, attn_groups)
-    try:
-        raw_tensors = runner._allocate_kv_cache_tensors(
-            kv_cache_config, extensible=True
-        )
-        kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
-        assert set(kv_caches) == set(layer_names)
-        assert len(runner._extensible_kv_cache_buffers) == len(layer_names)
-        for buffer, bytes_per_block_per_segment in runner._extensible_kv_cache_buffers:
-            assert buffer.num_segments == 1
-            assert bytes_per_block_per_segment == spec.page_size_bytes
-
-        # Write block 0 of every state of every layer (the committed
-        # prefixes), then grow.
-        for name in layer_names:
-            for state_tensor in kv_caches[name]:
-                state_tensor[0].fill_(1)
-        torch.cuda.synchronize()
-        runner.extend_kv_cache(num_blocks)
-        for name in layer_names:
-            for state_tensor in kv_caches[name]:
-                state_tensor[num_blocks - 1].fill_(2)
-        torch.cuda.synchronize()
-        for name in layer_names:
-            for state_tensor in kv_caches[name]:
-                assert torch.all(state_tensor[0] == 1)
-                assert torch.all(state_tensor[num_blocks - 1] == 2)
-                assert torch.count_nonzero(state_tensor[1 : num_blocks - 1]) == 0
-    finally:
-        _free_buffers(runner)
-
-
-def test_extensible_hybrid_attention_mamba() -> None:
-    """In hybrid models the attention cache is re-strided to block-major, so
-    its buffer must use a single segment."""
-    attn_spec = _full_attention_spec()
-    mamba_spec = _mamba_spec()
-    kv_cache_config = KVCacheConfig(
-        num_blocks=NUM_BLOCKS,
-        kv_cache_tensors=[
-            KVCacheTensor(
-                size=NUM_BLOCKS * attn_spec.page_size_bytes, shared_by=["attn.0"]
-            ),
-            KVCacheTensor(
-                size=NUM_BLOCKS * mamba_spec.page_size_bytes, shared_by=["mamba.0"]
-            ),
-        ],
-        kv_cache_groups=[
-            KVCacheGroupSpec(layer_names=["attn.0"], kv_cache_spec=attn_spec),
-            KVCacheGroupSpec(layer_names=["mamba.0"], kv_cache_spec=mamba_spec),
-        ],
-    )
-    attn_groups = [
-        [
-            AttentionGroup(
-                backend=_SplitKVBackend,
-                layer_names=["attn.0"],
-                kv_cache_spec=attn_spec,
-                kv_cache_group_id=0,
-            )
-        ],
-        [
-            AttentionGroup(
-                backend=_BlockMajorBackend,
-                layer_names=["mamba.0"],
-                kv_cache_spec=mamba_spec,
-                kv_cache_group_id=1,
-            )
-        ],
-    ]
-    runner = _make_runner(kv_cache_config, attn_groups)
-    try:
-        # The K/V-split attention layer is forced to one segment by the hybrid
-        # block-major re-stride.
-        assert runner._kv_cache_num_segments_by_layer() == {"attn.0": 1, "mamba.0": 1}
-
-        raw_tensors = runner._allocate_kv_cache_tensors(
-            kv_cache_config, extensible=True
-        )
-        kv_caches = runner._reshape_kv_cache_tensors(
-            raw_tensors, [BLOCK_SIZE, BLOCK_SIZE]
-        )
-        attn_cache = kv_caches["attn.0"]
-        # `_update_hybrid_attention_mamba_layout` re-strides to interleave K/V
-        # per block: block b spans one contiguous page.
-        hidden_size = attn_cache.shape[2:].numel()
-        assert attn_cache.stride()[:2] == (hidden_size, 2 * hidden_size)
-
-        attn_cache[0, 0].fill_(1)  # K, block 0
-        attn_cache[1, 0].fill_(2)  # V, block 0
-        for state_tensor in kv_caches["mamba.0"]:
-            state_tensor[0].fill_(3)
-        torch.cuda.synchronize()
-
-        runner.extend_kv_cache(NUM_BLOCKS)
-        attn_cache[0, NUM_BLOCKS - 1].fill_(4)
-        attn_cache[1, NUM_BLOCKS - 1].fill_(5)
-        for state_tensor in kv_caches["mamba.0"]:
-            state_tensor[NUM_BLOCKS - 1].fill_(6)
-        torch.cuda.synchronize()
-        assert torch.all(attn_cache[0, 0] == 1)
-        assert torch.all(attn_cache[1, 0] == 2)
-        assert torch.all(attn_cache[0, NUM_BLOCKS - 1] == 4)
-        assert torch.all(attn_cache[1, NUM_BLOCKS - 1] == 5)
-        assert torch.count_nonzero(attn_cache[:, 1 : NUM_BLOCKS - 1]) == 0
-        for state_tensor in kv_caches["mamba.0"]:
-            assert torch.all(state_tensor[0] == 3)
-            assert torch.all(state_tensor[NUM_BLOCKS - 1] == 6)
-            assert torch.count_nonzero(state_tensor[1 : NUM_BLOCKS - 1]) == 0
-    finally:
-        _free_buffers(runner)

tests/v1/worker/test_gpu_input_batch.py

@@ -435,40 +435,6 @@ def test_pooling_prompt_lens_not_aliased(device: str):
     )
 
 
-def test_placeholder_spec_token_ids_written_verbatim():
-    input_batch = InputBatch(
-        max_num_reqs=1,
-        max_model_len=8,
-        max_num_batched_tokens=8,
-        device=torch.device("cpu"),
-        vocab_size=VOCAB_SIZE,
-        block_sizes=[16],
-        kernel_block_sizes=[16],
-    )
-    req = CachedRequestState(
-        req_id="req",
-        prompt_token_ids=[10, 11],
-        mm_features=[],
-        sampling_params=SamplingParams(),
-        block_ids=([],),
-        generator=None,
-        num_computed_tokens=3,
-        output_token_ids=[12],
-    )
-    input_batch.add_request(req)
-
-    input_batch.update_req_spec_token_ids(
-        req,
-        {"req": [13, -1, -1]},
-    )
-
-    # Placeholders (-1) are kept verbatim in both the spec_token_ids list and
-    # the token buffer; they are clamped to 0 only at the embedding boundary
-    # (GPUModelRunner._preprocess).
-    assert input_batch.spec_token_ids[0] == [13, -1, -1]
-    assert input_batch.token_ids_cpu[0, 3:6].tolist() == [13, -1, -1]
-
-
 @pytest.mark.parametrize(
     ("pooling_params", "expect_device_prompt_token_ids", "expect_cpu_prompt_token_ids"),
     [

tests/v1/worker/test_gpu_model_runner.py

@@ -862,28 +862,6 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler():
     assert torch.equal(passed_draft_probs, expected_draft_probs)
 
 
-def test_invalid_draft_suffixes_remain_rejected_in_metadata():
-    runner = object.__new__(GPUModelRunner)
-    runner.device = torch.device("cpu")
-    runner.arange_np = np.arange(64, dtype=np.int64)
-    runner._arange_scratch = np.empty(64, dtype=np.int64)
-    # Placeholder (-1) drafts are kept in input_ids (clamped to 0 only at the
-    # embedding boundary). For num_draft_tokens=[2, 1, 2] the draft positions
-    # are [1, 2, 4, 6, 7], so the gather carries the -1s straight into the
-    # rejection-sampling metadata.
-    runner.input_ids = SimpleNamespace(
-        gpu=torch.tensor([99, 10, -1, 99, 12, 99, 13, -1], dtype=torch.int32),
-    )
-
-    metadata = GPUModelRunner._calc_spec_decode_metadata(
-        runner,
-        np.array([2, 1, 2], dtype=np.int32),
-        np.array([3, 5, 8], dtype=np.int32),
-    )
-
-    assert metadata.draft_token_ids.tolist() == [10, -1, 12, 13, -1]
-
-
 def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config):
     torch.set_default_dtype(torch.float16)
     layer_0 = "model.layers.0.self_attn.attn"

vllm/benchmarks/datasets/datasets.py

@@ -2452,7 +2452,6 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
                 num_requests=args.num_prompts,
                 tokenizer=tokenizer,
                 output_len=args.speed_bench_output_len,
-                skip_chat_template=args.skip_chat_template,
                 chat_template_kwargs=getattr(args, "chat_template_kwargs", None),
                 enable_multimodal_chat=args.enable_multimodal_chat,
                 request_id_prefix=args.request_id_prefix,
@@ -4156,21 +4155,8 @@ class ASRDataset(HuggingFaceDataset):
         **kwargs,
     ) -> list[SampleRequest]:
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
-        name_or_path = getattr(tokenizer, "name_or_path", "")
-        tok_class = type(tokenizer).__name__
-        if "openai" in name_or_path:
+        if "openai" in getattr(tokenizer, "name_or_path", ""):
             prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
-        elif tok_class == "CohereAsrTokenizer" or "cohere" in name_or_path.lower():
-            # CohereAsrTokenizer does not inject a decoder start token, so the
-            # decoder prompt must supply the full control-token sequence.
-            # Token order: context boundary, transcript start, emotion (default
-            # undefined), language (en), transcription directive (en), punctuation
-            # enabled, no ITN, no timestamp, no diarization.
-            prompt = (
-                "<|startofcontext|><|startoftranscript|>"
-                "<|emo:undefined|><|en|><|en|><|pnc|><|noitn|>"
-                "<|notimestamp|><|nodiarize|>"
-            )
         else:
             prompt = ""
         prompt_len = len(tokenizer(prompt).input_ids)

vllm/config/cache.py

@@ -174,15 +174,6 @@ class CacheConfig:
     gpu_memory_utilization. Note that kv_cache_memory_bytes
     (when not-None) ignores gpu_memory_utilization"""
 
-    enable_extensible_kv_cache: bool = False
-    """Use CUDA virtual memory to reserve the KV cache address range before
-    CUDA graph capture and commit the final size after capture.
-
-    This makes automatic KV sizing account for the actual CUDA graph pool.
-    Supported for all V1 CUDA attention backends (block-major and K/V-split
-    KV cache layouts) and for Mamba / linear-attention models.
-    """
-
     kv_offloading_size: float | None = None
     """Size of the KV cache offloading buffer in GiB. When TP > 1, this is
     the total buffer size summed across all TP ranks. By default, this is set
@@ -226,8 +217,6 @@ class CacheConfig:
             "kv_cache_max_concurrency",
             # WIP feature toggle not impacting compiled graph shape
             "kv_sharing_fast_prefill",
-            # Runtime memory allocation strategy, not graph structure.
-            "enable_extensible_kv_cache",
         }
 
         from vllm.config.utils import get_hash_factors, hash_factors

vllm/config/speculative.py

@@ -54,7 +54,6 @@ MTPModelTypes = Literal[
 ]
 NgramGPUTypes = Literal["ngram_gpu"]
 DFlashModelTypes = Literal["dflash"]
-DSparkModelTypes = Literal["dspark"]
 EagleModelTypes = Literal[
     "eagle", "eagle3", "extract_hidden_states", MTPModelTypes, DFlashModelTypes
 ]
@@ -67,7 +66,6 @@ SpeculativeMethod = Literal[
     "custom_class",
     EagleModelTypes,
     NgramGPUTypes,
-    DSparkModelTypes,
 ]
 RejectionSampleMethod = Literal["standard", "synthetic", "block"]
 DraftSampleMethod = Literal["greedy", "probabilistic"]
@@ -293,7 +291,6 @@ class SpeculativeConfig:
             "eagle3",
             "extract_hidden_states",
             "dflash",
-            "dspark",
         )
         factors.append(uses_aux_hidden_states)
 
@@ -611,13 +608,6 @@ class SpeculativeConfig:
                 # --quantization fp8 with a bf16 checkpoint.
                 if not self.quantization:
                     self.quantization = self.target_model_config.quantization
-            elif self.method == "dspark":
-                # DeepSeek DSpark can ship the weights inside the target checkpoint
-                if self.target_model_config is None:
-                    raise ValueError("target_model_config must be present for dspark")
-                self.model = self.target_model_config.model
-                if not self.quantization:
-                    self.quantization = self.target_model_config.quantization
             elif self.method in ("ngram", "[ngram]"):
                 self.model = "ngram"
             elif self.method == "ngram_gpu":
@@ -765,24 +755,18 @@ class SpeculativeConfig:
                         draft_hf.truncated_vocab_size = target_vocab
 
                 # Automatically detect the method
-                if self.method in ("eagle", "eagle3", "dflash", "dspark"):
+                if self.method in ("eagle", "eagle3", "dflash"):
                     pass
                 # examples:
                 # yuhuili/EAGLE-LLaMA3-Instruct-8B
                 # yuhuili/EAGLE3-LLaMA3.1-Instruct-8B
                 # AngelSlim/Qwen3-8B_eagle3
-                # deepseek-ai/dspark_qwen3_8b_block7
                 elif "eagle-" in self.draft_model_config.model.lower():
                     self.method = "eagle"
                 elif "eagle3" in self.draft_model_config.model.lower():
                     self.method = "eagle3"
                 elif "dflash" in self.draft_model_config.model.lower():
                     self.method = "dflash"
-                elif (
-                    "dspark" in self.draft_model_config.model.lower()
-                    or "Qwen3DSparkModel" in self.draft_model_config.architectures
-                ):
-                    self.method = "dspark"
                 elif self.draft_model_config.hf_config.model_type == "medusa":
                     self.method = "medusa"
                 elif self.draft_model_config.hf_config.model_type == "mlp_speculator":
@@ -829,18 +813,7 @@ class SpeculativeConfig:
                         self.draft_model_config.hf_config = eagle_config
                         self.update_arch_()
 
-                if self.method == "dspark" and (
-                    "Qwen3DSparkModel" not in self.draft_model_config.architectures
-                ):
-                    # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config
-                    # and its weights ship in the target checkpoint.
-                    self.draft_model_config.hf_config.model_type = "deepseek_v4"
-                    self.draft_model_config.hf_config.architectures = [
-                        "DSparkDraftModel"
-                    ]
-                    self.update_arch_()
-
-                if self.method in ("dflash", "dspark"):
+                if self.method == "dflash":
                     self.parallel_drafting = True
 
                 if self.num_speculative_tokens is not None and hasattr(
@@ -1156,17 +1129,11 @@ class SpeculativeConfig:
         )
 
     def use_eagle(self) -> bool:
-        # NOTE: This method is usually a stand-in for "speculative decoding using
-        # target model hidden states"
-        # TODO(ben): Refactor this so the naming is clearer
-        return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark")
+        return self.method in ("eagle", "eagle3", "mtp", "dflash")
 
     def use_dflash(self) -> bool:
         return self.method == "dflash"
 
-    def use_dspark(self) -> bool:
-        return self.method == "dspark"
-
     def uses_dynamic_speculative_decoding(self) -> bool:
         return self.num_speculative_tokens_per_batch_size is not None
 

vllm/config/vllm.py

@@ -524,16 +524,6 @@ class VllmConfig:
         if use_v2_model_runner is not None:
             return use_v2_model_runner
 
-        # DSpark is implemented only by the V2 GPU model runner, and DeepSeek-V4
-        # is not otherwise a default-V2 architecture, so force V2 for it. If V2
-        # is unsupported for the rest of the config, _validate_v2_model_runner
-        # raises rather than silently falling back to V1 (which can't run dspark).
-        if (
-            self.speculative_config is not None
-            and self.speculative_config.method == "dspark"
-        ):
-            return True
-
         if self.model_config is not None and self.model_config.is_diffusion:
             return True
 
@@ -968,11 +958,10 @@ class VllmConfig:
                     self.speculative_config.method not in get_args(EagleModelTypes)
                     and self.speculative_config.method not in get_args(NgramGPUTypes)
                     and self.speculative_config.method != "draft_model"
-                    and self.speculative_config.method != "dspark"
                 ):
                     raise ValueError(
                         "Currently, async scheduling is only supported "
-                        "with EAGLE/MTP/Draft Model/NGram GPU/DSpark kind of "
+                        "with EAGLE/MTP/Draft Model/NGram GPU kind of "
                         "speculative decoding"
                     )
                 if self.speculative_config.disable_padded_drafter_batch:
@@ -1000,7 +989,6 @@ class VllmConfig:
                 self.speculative_config is not None
                 and self.speculative_config.method not in get_args(EagleModelTypes)
                 and self.speculative_config.method not in get_args(NgramGPUTypes)
-                and self.speculative_config.method != "dspark"
             ):
                 logger.warning_once(
                     "Async scheduling not supported with %s-based "
@@ -2050,24 +2038,17 @@ class VllmConfig:
             # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet
             if speculative_config.method in ("ngram", "ngram_gpu"):
                 unsupported.append("ngram/ngram_gpu speculative decoding")
-            elif speculative_config.method not in (
-                "eagle",
-                "eagle3",
-                "mtp",
-                "dflash",
-                "dspark",
-            ):
+            elif speculative_config.method not in ("eagle", "eagle3", "mtp", "dflash"):
                 unsupported.append(f"speculative method '{speculative_config.method}'")
 
             if speculative_config.uses_dynamic_speculative_decoding():
                 unsupported.append("dynamic speculative decoding")
 
-            # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle).
-            # DFlash and DSpark use parallel drafting natively in V2 via their
-            # own speculators.
+            # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle)
+            # DFlash uses parallel drafting natively in V2 via DFlashSpeculator.
             if (
                 speculative_config.parallel_drafting
-                and speculative_config.method not in ("dflash", "dspark")
+                and speculative_config.method != "dflash"
             ):
                 unsupported.append("parallel drafting for EAGLE speculative decoding")
 
@@ -2112,9 +2093,6 @@ class VllmConfig:
             # Will be added by https://github.com/vllm-project/vllm/pull/35045
             unsupported.append("KV sharing fast prefill")
 
-        if self.cache_config.enable_extensible_kv_cache:
-            unsupported.append("extensible KV cache")
-
         if self.ec_transfer_config is not None:
             # Will be added by https://github.com/vllm-project/vllm/pull/38390
             unsupported.append("EC transfer")

vllm/engine/arg_utils.py

@@ -522,7 +522,6 @@ class EngineArgs:
     offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params")
     gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
     kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
-    enable_extensible_kv_cache: bool = CacheConfig.enable_extensible_kv_cache
     max_num_batched_tokens: int | None = None
     max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills
     max_long_partial_prefills: int = SchedulerConfig.max_long_partial_prefills
@@ -1153,10 +1152,6 @@ class EngineArgs:
         cache_group.add_argument(
             "--kv-cache-memory-bytes", **cache_kwargs["kv_cache_memory_bytes"]
         )
-        cache_group.add_argument(
-            "--enable-extensible-kv-cache",
-            **cache_kwargs["enable_extensible_kv_cache"],
-        )
         cache_group.add_argument("--kv-cache-dtype", **cache_kwargs["cache_dtype"])
         cache_group.add_argument(
             "--num-gpu-blocks-override", **cache_kwargs["num_gpu_blocks_override"]
@@ -1874,7 +1869,6 @@ class EngineArgs:
             block_size=self.block_size,  # type: ignore[arg-type]
             gpu_memory_utilization=self.gpu_memory_utilization,
             kv_cache_memory_bytes=self.kv_cache_memory_bytes,
-            enable_extensible_kv_cache=self.enable_extensible_kv_cache,
             cache_dtype=resolved_cache_dtype,  # type: ignore[arg-type]
             is_attention_free=model_config.is_attention_free,
             num_gpu_blocks_override=self.num_gpu_blocks_override,

vllm/entrypoints/llm.py

@@ -119,11 +119,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
             compared with using gpu_memory_utilization. Note that
             kv_cache_memory_bytes (when not-None) ignores
             gpu_memory_utilization
-        enable_extensible_kv_cache: Use CUDA virtual memory to reserve the KV
-            cache address range before CUDA graph capture and commit the final
-            cache size after capture. Supported by V1 CUDA workers for all
-            attention backends (block-major and K/V-split KV cache layouts)
-            and for Mamba / linear-attention models.
         cpu_offload_gb: The size (GiB) of CPU memory to use for offloading
             the model weights. This virtually increases the GPU memory space
             you can use to hold the model weights, at the cost of CPU-GPU data
@@ -216,7 +211,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
         profiler_config: dict[str, Any] | ProfilerConfig | None = None,
         attention_config: dict[str, Any] | AttentionConfig | None = None,
         kv_cache_memory_bytes: int | None = None,
-        enable_extensible_kv_cache: bool = False,
         compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
         quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None,
         logits_processors: list[str | type[LogitsProcessor]] | None = None,
@@ -326,7 +320,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
             seed=seed,
             gpu_memory_utilization=gpu_memory_utilization,
             kv_cache_memory_bytes=kv_cache_memory_bytes,
-            enable_extensible_kv_cache=enable_extensible_kv_cache,
             cpu_offload_gb=cpu_offload_gb,
             offload_group_size=offload_group_size,
             offload_num_in_group=offload_num_in_group,

vllm/lora/ops/xpu_ops/lora_ops.py

@@ -3,27 +3,29 @@
 
 import torch
 
-from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
 
 
-def _bgmv_shrink_impl(
+def bgmv_shrink(
     inputs: torch.Tensor,
     lora_a_weights: torch.Tensor,
     output_tensor: torch.Tensor,
     lora_indices_tensor: torch.Tensor,
-    scaling: float,
+    scaling: float = 1.0,
 ) -> None:
     torch.ops._xpu_C.bgmv_shrink(
         output_tensor, inputs, lora_a_weights, lora_indices_tensor, scaling
     )
 
 
-def _bgmv_expand_impl(
+def bgmv_expand(
     inputs: torch.Tensor,
     lora_b_weights: torch.Tensor,
     output_tensor: torch.Tensor,
     lora_indices_tensor: torch.Tensor,
-    add_inputs: bool,
+    add_inputs: bool = True,
 ) -> None:
     weight_out_dim = lora_b_weights.size(-2)
     output_dim = output_tensor.size(1)
@@ -63,14 +65,14 @@ def _bgmv_expand_impl(
         )
 
 
-def _bgmv_expand_slice_impl(
+def bgmv_expand_slice(
     inputs: torch.Tensor,
     lora_b_weights: torch.Tensor,
     output_tensor: torch.Tensor,
     lora_indices_tensor: torch.Tensor,
     slice_offset: int,
     slice_size: int,
-    add_inputs: bool,
+    add_inputs: bool = True,
 ) -> None:
     assert slice_size == lora_b_weights.size(-2)
     assert slice_offset + slice_size <= output_tensor.size(1)
@@ -83,101 +85,3 @@ def _bgmv_expand_slice_impl(
         slice_size,
         add_inputs,
     )
-
-
-def _bgmv_shrink_fake(
-    inputs: torch.Tensor,
-    lora_a_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    scaling: float,
-) -> None:
-    return None
-
-
-def _bgmv_expand_fake(
-    inputs: torch.Tensor,
-    lora_b_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    add_inputs: bool,
-) -> None:
-    return None
-
-
-def _bgmv_expand_slice_fake(
-    inputs: torch.Tensor,
-    lora_b_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    slice_offset: int,
-    slice_size: int,
-    add_inputs: bool,
-) -> None:
-    return None
-
-
-direct_register_custom_op(
-    op_name="xpu_bgmv_shrink",
-    op_func=_bgmv_shrink_impl,
-    mutates_args=["output_tensor"],
-    fake_impl=_bgmv_shrink_fake,
-)
-
-direct_register_custom_op(
-    op_name="xpu_bgmv_expand",
-    op_func=_bgmv_expand_impl,
-    mutates_args=["output_tensor"],
-    fake_impl=_bgmv_expand_fake,
-)
-
-direct_register_custom_op(
-    op_name="xpu_bgmv_expand_slice",
-    op_func=_bgmv_expand_slice_impl,
-    mutates_args=["output_tensor"],
-    fake_impl=_bgmv_expand_slice_fake,
-)
-
-
-def bgmv_shrink(
-    inputs: torch.Tensor,
-    lora_a_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    scaling: float = 1.0,
-) -> None:
-    torch.ops.vllm.xpu_bgmv_shrink(
-        inputs, lora_a_weights, output_tensor, lora_indices_tensor, scaling
-    )
-
-
-def bgmv_expand(
-    inputs: torch.Tensor,
-    lora_b_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    add_inputs: bool = True,
-) -> None:
-    torch.ops.vllm.xpu_bgmv_expand(
-        inputs, lora_b_weights, output_tensor, lora_indices_tensor, add_inputs
-    )
-
-
-def bgmv_expand_slice(
-    inputs: torch.Tensor,
-    lora_b_weights: torch.Tensor,
-    output_tensor: torch.Tensor,
-    lora_indices_tensor: torch.Tensor,
-    slice_offset: int,
-    slice_size: int,
-    add_inputs: bool = True,
-) -> None:
-    torch.ops.vllm.xpu_bgmv_expand_slice(
-        inputs,
-        lora_b_weights,
-        output_tensor,
-        lora_indices_tensor,
-        slice_offset,
-        slice_size,
-        add_inputs,
-    )

vllm/lora/punica_wrapper/punica_xpu.py

@@ -62,10 +62,6 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             captured_lora_counts=captured_lora_counts,
         )
 
-        # When speculative decoding is enabled, max_num_samples is
-        # max_batches * (num_speculative_decoding_tokens + 1).
-        # This line can be optimized by replacing max_num_batched_tokens
-        # to  max_batches * (num_speculative_decoding_tokens + 1).
         self.prompt_mapping_meta = LoRAKernelMeta.make(
             self.max_loras,
             max_num_batched_tokens,
@@ -110,14 +106,6 @@ class PunicaWrapperXPU(PunicaWrapperBase):
         add_inputs: bool,
     ):
         token_lora_indices = self._get_token_lora_indices(x)
-        # After tensor-parallel all-gather (non-fully-sharded LoRA), x may
-        # have been gathered along the rank dim so x.size(1) == max_lora_rank
-        # * tp_size, while lora_b only uses max_lora_rank elements. The XPU
-        # C++ kernel requires inputs.size(1) == lora_b.size(-1), so truncate
-        # to the actual rank. x[:, :rank] is non-contiguous, hence the copy.
-        rank = w_t_all.size(-1)
-        if x.size(1) != rank:
-            x = x[:, :rank].contiguous()
         bgmv_expand_slice(
             x, w_t_all, y, token_lora_indices, y_offset, y_slice_size, add_inputs
         )
@@ -191,7 +179,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
                 add_inputs=add_inputs,
             )
             offset_start += output_slices[slice_idx]
-        y = y.view_as(y_org)
+        y.view_as(y_org)
 
     def add_lora_embedding(
         self,
@@ -239,6 +227,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
                     @ lora_b_stacked[indices[i], layer_idx, :, :]
                     * scale
                     ).squeeze(0)
+
         Args:
             y (torch.Tensor): Output tensor. Will be changed in-place.
             x (torch.Tensor): Input tensor
@@ -251,17 +240,13 @@ class PunicaWrapperXPU(PunicaWrapperBase):
 
         assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
 
-        assert buffer is None, (
-            "To minimize overhead, the buffer should be created by "
-            ".add_lora_linear() instead of being passed in."
-        )
-        r = lora_b_stacked[0].size(-1)
-        buffer = torch.zeros(  # type: ignore
-            (len(output_slices), x.size(0), r),
-            dtype=x.dtype,
-            device=x.device,
-        )
-        add_inputs = kwargs.pop("add_inputs", True)
+        if buffer is None:
+            r = lora_b_stacked[0].size(-1)
+            buffer = torch.zeros(  # type: ignore
+                (len(output_slices), x.size(0), r),
+                dtype=x.dtype,
+                device=x.device,
+            )
         self.add_shrink(
             buffer,  # type: ignore
             x,
@@ -274,7 +259,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             buffer,  # type: ignore
             lora_b_stacked,
             output_slices,
-            add_inputs=add_inputs,
+            add_inputs=True,
             **kwargs,
         )
 
@@ -315,16 +300,12 @@ class PunicaWrapperXPU(PunicaWrapperBase):
         y = y.view(-1, y.shape[-1])
         x = x.view(-1, x.shape[-1])
         r = lora_b_stacked.size(-1)
-
-        assert buffer is None, (
-            "To minimize overhead, the buffer should be created by "
-            ".add_lora_linear() instead of being passed in."
-        )
-        buffer = torch.zeros((x.size(0), r), dtype=x.dtype, device=x.device)
+        if buffer is None:
+            buffer = torch.zeros((x.size(0), r), dtype=x.dtype, device=x.device)
         sampler_indices = torch.narrow(self._sampler_indices, 0, 0, x.size(0))
         bgmv_shrink(x, lora_a_stacked, buffer, sampler_indices, scale)
         bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True)
-        y = y.view_as(y_org)
+        return y.view_as(y_org)
 
     def moe_lora_align_block_size(
         self,
@@ -337,62 +318,33 @@ class PunicaWrapperXPU(PunicaWrapperBase):
         expert_map: torch.Tensor | None = None,
         pad_sorted_ids: bool = False,
         naive_block_assignment: bool = False,
-        token_lora_mapping: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Aligns tokens and experts into block-sized chunks for LoRA-based
         mixture-of-experts (MoE) execution.
-
-        When `token_lora_mapping` is provided, it overrides the global mapping
-        read from `self.token_mapping_meta`. This is how EP+LoRA injects the
-        per-rank-local token→LoRA map after all-to-all dispatch.
         """
-        (
-            token_lora_mapping_meta,
-            _,
-            _,
-            _,
-            lora_ids,
-            _,
-            _,
-        ) = self.token_mapping_meta.meta_args(
-            num_tokens, self.lora_config.specialize_active_lora
-        )
-        if token_lora_mapping is None:
-            token_lora_mapping = token_lora_mapping_meta
-        # Under EP the caller passes local_num_experts but topk_ids carries
-        # GLOBAL expert indices. The CUDA kernel uses num_experts to size
-        # its bucketing table; with EP we must size by global_num_experts
-        # so global topk_ids don't overflow. expert_map inside the kernel
-        # then translates global→local so the output expert_ids are local
-        # (mirrors the non-LoRA moe_align_block_size behavior).
-        kernel_num_experts = (
-            expert_map.numel() if expert_map is not None else num_experts
+        (token_lora_mapping, _, _, _, lora_ids, _, _) = (
+            self.token_mapping_meta.meta_args(
+                num_tokens, self.lora_config.specialize_active_lora
+            )
         )
         if naive_block_assignment:
             expert_ids = topk_ids.reshape(-1)
             sorted_ids = None
             num_tokens_post_pad = None
         else:
-            max_num_tokens_padded = topk_ids.numel() + kernel_num_experts * (
-                block_size - 1
-            )
+            max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
             if pad_sorted_ids:
                 max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
-            if topk_ids.numel() < kernel_num_experts:
-                max_num_tokens_padded = topk_ids.numel() * block_size
             sorted_ids = torch.empty(
                 (max_loras * max_num_tokens_padded,),
                 dtype=torch.int32,
                 device=topk_ids.device,
             )
             max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
-            # Expert ids are initialized to -1 so unused (lora, expert)
-            # slots don't drive the LoRA Triton kernel into the wrong bucket.
-            # The kernel overwrites only active slots.
-            expert_ids = torch.full(
+            # Expert ids must be set default to -1 to prevent a blank block
+            expert_ids = torch.empty(
                 (max_loras * max_num_m_blocks,),
-                -1,
                 dtype=torch.int32,
                 device=topk_ids.device,
             )
@@ -403,7 +355,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             ops.moe_lora_align_block_size(
                 topk_ids,
                 token_lora_mapping,
-                kernel_num_experts,
+                num_experts,
                 block_size,
                 max_loras,
                 max_num_tokens_padded,
@@ -413,10 +365,11 @@ class PunicaWrapperXPU(PunicaWrapperBase):
                 num_tokens_post_pad,
                 adapter_enabled,
                 lora_ids,
-                expert_map,
             )
+            if expert_map is not None:
+                expert_ids = expert_map[expert_ids]
 
-        return token_lora_mapping, sorted_ids, expert_ids, num_tokens_post_pad
+        return None, sorted_ids, expert_ids, num_tokens_post_pad
 
     def add_lora_fused_moe(
         self,
@@ -572,8 +525,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
 
         SPARSITY_FACTOR = 8
         naive_block_assignment = (
-            not fully_sharded
-            and expert_map is None
+            expert_map is None
             and num_tokens * top_k * SPARSITY_FACTOR <= local_num_experts * max_loras
         )
 
@@ -591,7 +543,6 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             adapter_enabled,
             expert_map,
             naive_block_assignment=naive_block_assignment,
-            token_lora_mapping=token_lora_mapping,
         )
 
         _sorted = sorted_token_ids_lora
@@ -616,7 +567,6 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             adapter_enabled,
             fully_sharded=fully_sharded,
             token_lora_mapping=token_lora_mapping,
-            add_inputs=add_inputs,
         )
 
         return (
@@ -730,5 +680,4 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             fully_sharded=fully_sharded,
             offset=offset,
             token_lora_mapping=token_lora_mapping,
-            add_inputs=add_inputs,
         )

vllm/model_executor/models/granitemoeshared.py

@@ -214,11 +214,11 @@ class GraniteMoeSharedModel(nn.Module):
                 for e in range(p.size(0)):
                     w1_name = n.replace(
                         ".block_sparse_moe.input_linear.weight",
-                        f".block_sparse_moe.experts.{e}.w1.weight",
+                        f".block_sparse_moe.experts.routed_experts.{e}.w1.weight",
                     )
                     w3_name = n.replace(
                         ".block_sparse_moe.input_linear.weight",
-                        f".block_sparse_moe.experts.{e}.w3.weight",
+                        f".block_sparse_moe.experts.routed_experts.{e}.w3.weight",
                     )
                     w1_param, w3_param = p[e].chunk(2, dim=0)
                     assert w1_name not in new_weights
@@ -229,7 +229,7 @@ class GraniteMoeSharedModel(nn.Module):
                 for e in range(p.size(0)):
                     w2_name = n.replace(
                         ".block_sparse_moe.output_linear.weight",
-                        f".block_sparse_moe.experts.{e}.w2.weight",
+                        f".block_sparse_moe.experts.routed_experts.{e}.w2.weight",
                     )
                     w2_param = p[e]
                     assert w2_name not in new_weights

vllm/model_executor/models/qwen3_dflash.py

@@ -472,7 +472,7 @@ class DFlashQwen3Model(nn.Module):
         self,
         context_states: torch.Tensor,
         context_positions: torch.Tensor,
-        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
+        context_slot_mapping: torch.Tensor | None = None,
     ) -> None:
         """Precompute K/V for context states write them into each layer's KV cache.
 
@@ -551,13 +551,7 @@ class DFlashQwen3Model(nn.Module):
 
         # --- Per-layer cache insert ---
         all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
-        per_layer = isinstance(context_slot_mapping, (list, tuple))
         for i in range(L):
-            slot_mapping = (
-                context_slot_mapping[i] if per_layer else context_slot_mapping
-            )
-            if slot_mapping is None:
-                continue  # dummy run: skip cache ops
             attn = self._attn_layers[i]
             kv_cache = attn.kv_cache
             attn.impl.do_kv_cache_update(
@@ -565,7 +559,7 @@ class DFlashQwen3Model(nn.Module):
                 all_k_final[i],
                 all_v[i],
                 kv_cache,
-                slot_mapping,
+                context_slot_mapping,
             )
 
     def forward(
@@ -707,7 +701,7 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
         self,
         context_states: torch.Tensor,
         context_positions: torch.Tensor,
-        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
+        context_slot_mapping: torch.Tensor | None = None,
     ) -> None:
         """Precompute projected + RoPE'd K/V and write to cache."""
         self.model.precompute_and_store_context_kv(

vllm/model_executor/models/qwen3_dspark.py

@@ -1,153 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Qwen3 DSpark draft model for semi-autoregressive drafting.
-
-DSpark drafts a whole block in one parallel pass (DFlash-style: context-KV
-precompute + a non-causal query-block forward) and then injects intra-block
-dependency with a lightweight sequential Markov head.
-
-The parallel backbone is a standard Qwen3 decoder stack reused from the
-DFlash Qwen3 draft (see qwen3_dflash.py). DSpark adds:
-  * ``markov_head``: low-rank V x r / r x V transition bias added to the base
-    logits, sampled left-to-right by the speculator (the sequential stage).
-
-DSparkMarkovHead is shared with the DSV4-style DSpark model.
-"""
-
-from collections.abc import Iterable
-
-import torch
-import torch.nn as nn
-
-from vllm.config import VllmConfig
-from vllm.logger import init_logger
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead,
-    VocabParallelEmbedding,
-)
-
-from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
-from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
-
-logger = init_logger(__name__)
-
-
-class DSparkMarkovHead(nn.Module):
-    """Sequential transition-bias head (low-rank V x r, r x V).
-
-    ``markov_w1[token]`` is an r-dim embedding of the previously sampled token;
-    ``markov_w2`` projects it back to a vocab-size bias added to the base logits.
-    """
-
-    def __init__(self, vocab_size: int, markov_rank: int, prefix: str) -> None:
-        super().__init__()
-        # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
-        self.markov_w1 = VocabParallelEmbedding(
-            vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
-        )
-        self.markov_w2 = ParallelLMHead(
-            vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
-        )
-
-    def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
-        """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
-        return self.markov_w1(token_ids)
-
-    def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor:
-        """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
-        return logits_processor(self.markov_w2, markov_embed)
-
-
-class Qwen3DSparkModel(DFlashQwen3Model):
-    """DFlash Qwen3 backbone + DSpark Markov head."""
-
-    def __init__(
-        self,
-        *,
-        vllm_config: VllmConfig,
-        start_layer_id: int = 0,
-        prefix: str = "",
-    ) -> None:
-        super().__init__(
-            vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix
-        )
-        config = self.config
-        self.markov_head = DSparkMarkovHead(
-            config.vocab_size,
-            config.markov_rank,
-            prefix=maybe_prefix(prefix, "markov_head"),
-        )
-
-
-class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        nn.Module.__init__(self)
-        self.draft_model_config = vllm_config.speculative_config.draft_model_config
-        self.config = self.draft_model_config.hf_config
-        if getattr(self.config, "draft_vocab_size", None) is None:
-            self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
-        target_layer_num = vllm_config.model_config.get_num_layers(
-            vllm_config.parallel_config
-        )
-        self.model = Qwen3DSparkModel(
-            vllm_config=vllm_config,
-            prefix=maybe_prefix(prefix, "model"),
-            start_layer_id=target_layer_num,
-        )
-
-        logit_scale = getattr(self.config, "logit_scale", 1.0)
-        self.lm_head = ParallelLMHead(
-            self.config.draft_vocab_size,
-            self.config.hidden_size,
-            prefix=maybe_prefix(prefix, "lm_head"),
-        )
-        self.logits_processor = LogitsProcessor(
-            self.config.draft_vocab_size, scale=logit_scale
-        )
-        target_vocab_size = vllm_config.model_config.get_vocab_size()
-        if self.config.draft_vocab_size != target_vocab_size:
-            self.draft_id_to_target_id = nn.Parameter(
-                torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
-                requires_grad=False,
-            )
-        else:
-            self.draft_id_to_target_id = None
-
-    def get_draft_kv_cache_layer_names(self) -> list[str]:
-        return [layer.self_attn.attn.layer_name for layer in self.model.layers]
-
-    def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
-        return self.model.markov_head.embed(token_ids)
-
-    def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
-        return self.model.markov_head.bias(markov_embed, self.logits_processor)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        model_weights = {}
-        includes_embed_tokens = False
-        includes_lm_head = False
-        for name, loaded_weight in weights:
-            if "lm_head" not in name:
-                name = "model." + name
-            if "embed_tokens" in name:
-                includes_embed_tokens = True
-            if "lm_head" in name:
-                includes_lm_head = True
-            model_weights[name] = loaded_weight
-            # Sets has_own_embed_tokens / has_own_lm_head so load_dspark_model
-            # knows whether to keep these or alias the target's.
-            process_eagle_weight(self, name)
-
-        # mask_embedding is an unused placeholder param; DSpark masks via the vocab row.
-        # confidence_head is not wired into inference yet; skip its weights.
-        # embed_tokens / lm_head are optional; when omitted they are shared from
-        # the target by load_dspark_model, so skip the unloaded params here.
-        skip_substrs = ["mask_embedding", "confidence_head"]
-        if not includes_embed_tokens:
-            skip_substrs.append("embed_tokens")
-        if not includes_lm_head:
-            skip_substrs.append("lm_head")
-        loader = AutoWeightsLoader(self, skip_substrs=skip_substrs)
-        loader.load_weights(model_weights.items())
-        self.model._build_fused_kv_buffers()

vllm/model_executor/models/registry.py

@@ -585,8 +585,6 @@ _SPECULATIVE_DECODING_MODELS = {
     "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
     "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
     "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
-    "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
-    "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
     "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

vllm/models/deepseek_v4/__init__.py

@@ -17,23 +17,14 @@ from .quant_config import DeepseekV4FP8Config
 if current_platform.is_rocm():
     from .amd.model import DeepseekV4ForCausalLM
     from .amd.mtp import DeepSeekV4MTP
-
-    # DSpark is NVIDIA-only for now.
-    DSparkDeepseekV4ForCausalLM = None  # type: ignore[assignment]
 elif current_platform.is_xpu():
     from .xpu.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
     from .xpu.mtp import DeepSeekV4MTP  # type: ignore[assignment]
-
-    DSparkDeepseekV4ForCausalLM = None  # type: ignore[assignment]
 else:
-    from .nvidia.dspark import (  # type: ignore[assignment]
-        DSparkDeepseekV4ForCausalLM,
-    )
     from .nvidia.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
     from .nvidia.mtp import DeepSeekV4MTP  # type: ignore[assignment]
 
 __all__ = [
-    "DSparkDeepseekV4ForCausalLM",
     "DeepSeekV4MTP",
     "DeepseekV4FP8Config",
     "DeepseekV4ForCausalLM",

vllm/models/deepseek_v4/nvidia/dspark.py

@@ -1,477 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""DSpark draft model for DeepSeek-V4 (semi-autoregressive speculative decoding).
-
-See: qwen3_dspark.py for base architecture. This one is specialized to the DSV4 DSpark,
-which reuses the target model's architecture similarly to MTP.
-
-To implement non-causal attention, we leverage the sparse attention implementation to
-include the future query tokens in the top-k indices for each query token.
-"""
-
-from collections.abc import Iterable
-
-import regex as re
-import torch
-import torch.nn as nn
-
-from vllm.config import VllmConfig, get_current_vllm_config
-from vllm.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
-from vllm.logger import init_logger
-from vllm.model_executor.kernels.mhc.tilelang import (
-    hc_head_fused_kernel_tilelang,
-    mhc_post_tilelang,
-)
-from vllm.model_executor.layers.fused_moe import (
-    fused_moe_make_expert_params_mapping,
-)
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-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.qwen3_dspark import (
-    DSparkMarkovHead,
-)
-from vllm.model_executor.models.utils import maybe_prefix
-
-from .model import (
-    DeepseekV4DecoderLayer,
-    make_deepseek_v4_expert_params_mapping,
-)
-
-logger = init_logger(__name__)
-
-# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
-# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
-_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
-
-
-class DSparkDeepseekV4Model(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        super().__init__()
-        assert vllm_config.speculative_config is not None
-        config = vllm_config.speculative_config.draft_model_config.hf_config
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.hc_mult = config.hc_mult
-        self.hc_eps = config.hc_eps
-        self.rms_norm_eps = config.rms_norm_eps
-        self.num_hidden_layers = config.num_hidden_layers
-        self.target_layer_ids = tuple(config.dspark_target_layer_ids)
-
-        self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3
-
-        # Shared with the target (aliased by the speculator's loading utility).
-        self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size,
-            config.hidden_size,
-            prefix=maybe_prefix(prefix, "embed_tokens"),
-        )
-
-        self.main_proj = ReplicatedLinear(
-            config.hidden_size * len(self.target_layer_ids),
-            config.hidden_size,
-            bias=False,
-            return_bias=False,
-            quant_config=vllm_config.quant_config,
-            prefix=maybe_prefix(prefix, "main_proj"),
-        )
-        self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-        current_vllm_config = get_current_vllm_config()
-        self.layers = nn.ModuleList(
-            [
-                DeepseekV4DecoderLayer(
-                    current_vllm_config,
-                    prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
-                )
-                for i in range(self.num_dspark_layers)
-            ]
-        )
-
-        # Heads: final norm + hc_head, and the Markov head
-        # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        hc_dim = self.hc_mult * config.hidden_size
-        self.hc_head_fn = nn.Parameter(
-            torch.empty(self.hc_mult, hc_dim, dtype=torch.float32),
-            requires_grad=False,
-        )
-        self.hc_head_base = nn.Parameter(
-            torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False
-        )
-        self.hc_head_scale = nn.Parameter(
-            torch.empty(1, dtype=torch.float32), requires_grad=False
-        )
-        self.markov_head = DSparkMarkovHead(
-            config.vocab_size,
-            config.dspark_markov_rank,
-            prefix=maybe_prefix(prefix, "markov_head"),
-        )
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
-
-    def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
-        """main_x = main_norm(main_proj(concat of target aux hidden states)).
-
-        ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)].
-        """
-        return self.main_norm(self.main_proj(aux_hidden_states))
-
-    @torch.inference_mode()
-    def precompute_and_store_context_kv(
-        self,
-        main_x: torch.Tensor,
-        context_positions: torch.Tensor,
-        context_slot_mappings: list[torch.Tensor | None] | None = None,
-    ) -> None:
-        """Insert the sliding-window context KV for every draft layer.
-
-        Mirrors the reference DSparkAttention: each layer derives its context KV
-        from the SAME projected target hidden ``main_x``, via that layer's own
-        ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the
-        layer's context slots.
-
-        ``context_slot_mappings`` is a per-layer list (each entry is the context
-        slot mapping for that layer's kv-cache group, since the hybrid manager may
-        place draft layers in different groups). ``None`` (or a ``None`` entry)
-        runs the projection to reserve workspace but writes nothing (profiling).
-        """
-        for i, layer in enumerate(self.layers):
-            slot_mapping = (
-                None if context_slot_mappings is None else context_slot_mappings[i]
-            )
-            attn = layer.attn
-            # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
-            # (q_lora part discarded), then RoPE/quant/insert via the fused op.
-            qr_kv, _ = attn.fused_wqa_wkv(main_x)
-            kv = qr_kv[..., attn.q_lora_rank :]
-            kv = attn.kv_norm(kv)
-            if slot_mapping is None:
-                continue
-            _insert_context_kv(attn, kv, context_positions, slot_mapping)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        inputs_embeds: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_input_ids(input_ids)
-        # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
-        hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
-
-        residual = post_mix = res_mix = None
-        for layer in self.layers:
-            hidden_states, residual, post_mix, res_mix = layer(
-                hidden_states,
-                positions,
-                input_ids,
-                post_mix,
-                res_mix,
-                residual,
-            )
-        hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
-        # hc_head reduces the hc copies; return the PRE-norm head hidden
-        hidden_states = hc_head_fused_kernel_tilelang(
-            hidden_states,
-            self.hc_head_fn,
-            self.hc_head_scale,
-            self.hc_head_base,
-            self.rms_norm_eps,
-            self.hc_eps,
-        )
-        return hidden_states
-
-
-def _insert_context_kv(
-    attn: nn.Module,
-    kv: torch.Tensor,
-    positions: torch.Tensor,
-    slot_mapping: torch.Tensor,
-) -> None:
-    """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
-
-    Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
-    query and discard it, since context tokens have no query). Mirrors
-    ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
-    """
-    swa_cache = attn.swa_cache_layer.kv_cache
-    block_size = attn.swa_cache_layer.block_size
-    cos_sin_cache = attn.rotary_emb.cos_sin_cache
-    cache_dtype = swa_cache.dtype
-    n_ctx = kv.shape[0]
-    dummy_q = torch.zeros(
-        (n_ctx, attn.n_local_heads, attn.head_dim),
-        dtype=kv.dtype,
-        device=kv.device,
-    )
-    if cache_dtype == torch.uint8:
-        # fp8_ds_mla UE8M0 paged layout
-        swa_2d = swa_cache.view(swa_cache.shape[0], -1)
-        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-            dummy_q,
-            kv,
-            swa_2d,
-            slot_mapping,
-            positions,
-            cos_sin_cache,
-            attn.padded_heads,
-            attn.eps,
-            block_size,
-        )
-    elif cache_dtype == torch.bfloat16:
-        swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
-        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
-            dummy_q,
-            kv,
-            swa_3d,
-            slot_mapping,
-            positions,
-            cos_sin_cache,
-            attn.eps,
-            block_size,
-        )
-    else:  # per-tensor fp8 (torch.float8_e4m3fn)
-        # TODO(ben): double-check if this is being dispatched correctly for FI backend
-        swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
-        dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
-        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
-            dummy_q,
-            kv,
-            dummy_q_fp8,
-            swa_3d,
-            slot_mapping,
-            positions,
-            cos_sin_cache,
-            attn._flashinfer_fp8_kv_scale,
-            attn._flashinfer_fp8_q_scale_inv,
-            attn.eps,
-            block_size,
-        )
-
-
-class DSparkDeepseekV4ForCausalLM(nn.Module):
-    # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so
-    # load_dspark_model always aliases the target's.
-    has_own_embed_tokens = False
-    has_own_lm_head = False
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        super().__init__()
-        assert vllm_config.speculative_config is not None
-        self.draft_model_config = vllm_config.speculative_config.draft_model_config
-        self.config = self.draft_model_config.hf_config
-        self.model = DSparkDeepseekV4Model(
-            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
-        )
-        # Shared with the target (aliased by the speculator's load utility).
-        self.lm_head = ParallelLMHead(
-            self.config.vocab_size,
-            self.config.hidden_size,
-            prefix=maybe_prefix(prefix, "lm_head"),
-        )
-        self.logits_processor = LogitsProcessor(self.config.vocab_size)
-
-    # --- Hooks used by the speculator -------------------------------------
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.model.embed_input_ids(input_ids)
-
-    def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
-        return self.model.combine_hidden_states(aux_hidden_states)
-
-    def get_draft_kv_cache_layer_names(self) -> list[str]:
-        # DSV4 MLA path: each draft layer's sliding-window cache is a separate
-        # layer, named by its prefix.
-        return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers]
-
-    def precompute_and_store_context_kv(
-        self,
-        context_states: torch.Tensor,
-        context_positions: torch.Tensor,
-        context_slot_mappings: list[torch.Tensor | None] | None = None,
-    ) -> None:
-        self.model.precompute_and_store_context_kv(
-            context_states, context_positions, context_slot_mappings
-        )
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        inputs_embeds: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        # Returns the pre-norm hc_head hidden ([T, hidden_size]).
-        return self.model(input_ids, positions, inputs_embeds)
-
-    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """Base logits U_k = lm_head(norm(head_hidden))."""
-        return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
-
-    def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
-        return self.model.markov_head.embed(token_ids)
-
-    def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
-        return self.model.markov_head.bias(markov_embed, self.logits_processor)
-
-    # --- Weight loading ----------------------------------------------------
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.
-
-        Non-mtp weights (embed/head/main layers) belong to the target model and
-        are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
-        """
-        first_layer = self.model.layers[0]
-        use_mega_moe = first_layer.ffn.use_mega_moe
-        if use_mega_moe:
-            expert_mapping = make_deepseek_v4_expert_params_mapping(
-                self.config.n_routed_experts
-            )
-        else:
-            expert_mapping = fused_moe_make_expert_params_mapping(
-                self,
-                ckpt_gate_proj_name="w1",
-                ckpt_down_proj_name="w2",
-                ckpt_up_proj_name="w3",
-                num_experts=self.config.n_routed_experts,
-            )
-        expert_scale_suffix = (
-            ".weight_scale"
-            if getattr(self.config, "expert_dtype", "fp4") == "fp4"
-            else ".weight_scale_inv"
-        )
-
-        # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params.
-        stacked_params_mapping = [
-            ("gate_up_proj", "w1", 0),
-            ("gate_up_proj", "w3", 1),
-            ("attn.fused_wqa_wkv", "attn.wq_a", 0),
-            ("attn.fused_wqa_wkv", "attn.wkv", 1),
-        ]
-
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-
-        tp_size = get_tensor_model_parallel_world_size()
-        tp_rank = get_tensor_model_parallel_rank()
-        n_local_head = self.config.num_attention_heads // tp_size
-        head_start = n_local_head * tp_rank
-        head_end = n_local_head * (tp_rank + 1)
-
-        for name, loaded_weight in weights:
-            mapped = self._remap_dspark_name(name)
-            if mapped is None:
-                continue
-            name = mapped
-
-            # ``.scale`` -> per-method scale suffix.
-            if name.endswith(".scale"):
-                suffix = (
-                    expert_scale_suffix
-                    if _EXPERT_SCALE_RE.search(name)
-                    else ".weight_scale_inv"
-                )
-                name = name.removesuffix(".scale") + suffix
-
-            # E8M0 expert scales: keep raw exponent bytes.
-            if ".experts." in name:
-                if (
-                    "weight_scale" in name
-                    and loaded_weight.dtype == torch.float8_e8m0fnu
-                ):
-                    loaded_weight = loaded_weight.view(torch.uint8)
-                for param_name, weight_name, expert_id, shard_id in expert_mapping:
-                    if weight_name not in name:
-                        continue
-                    name_mapped = name.replace(weight_name, param_name)
-                    param = params_dict[name_mapped]
-                    success = param.weight_loader(
-                        param,
-                        loaded_weight,
-                        name_mapped,
-                        shard_id=shard_id,
-                        expert_id=expert_id,
-                        return_success=True,
-                    )
-                    if success:
-                        loaded_params.add(name_mapped)
-                        break
-                continue
-
-            # Stacked rules only apply to decoder-layer weights. Head-stack params
-            # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.
-            # "markov_w1" would collide with the "w1" shard rule.
-            is_layer_param = name.startswith("model.layers.")
-            for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
-                if not is_layer_param or weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                param = params_dict[name]
-                param.weight_loader(param, loaded_weight, stacked_shard_id)
-                loaded_params.add(name)
-                break
-            else:
-                if "attn_sink" in name:
-                    narrow = loaded_weight[head_start:head_end]
-                    params_dict[name][: narrow.shape[0]].copy_(narrow)
-                    loaded_params.add(name)
-                    continue
-                if ".shared_experts.w2" in name:
-                    name = name.replace(
-                        ".shared_experts.w2", ".shared_experts.down_proj"
-                    )
-                if name.endswith(".ffn.gate.bias"):
-                    name = name.replace(
-                        ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
-                    )
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-                loaded_params.add(name)
-
-        self._finalize_moe()
-        logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
-        return loaded_params
-
-    def _finalize_moe(self) -> None:
-        for layer in self.model.layers:
-            layer.ffn.finalize_mega_moe_weights()
-
-    def _remap_dspark_name(self, name: str) -> str | None:
-        """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
-
-        Returns None for non-mtp weights (owned by the target model).
-        """
-        m = re.match(r"mtp\.(\d+)\.(.*)", name)
-        if m is None:
-            return None
-        stage = int(m.group(1))
-        rest = m.group(2)
-        # The confidence head is not wired into inference yet; drop its weights.
-        if rest.startswith("confidence_head."):
-            return None
-        # Head-stack params live at model level (mtp.last), context combiner at
-        # model level (mtp.0); everything else is a per-layer decoder block.
-        head_prefixes = (
-            "norm.",
-            "hc_head_fn",
-            "hc_head_base",
-            "hc_head_scale",
-            "markov_head.",
-        )
-        if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
-            head_prefixes
-        ):
-            return f"model.{rest}"
-        return f"model.layers.{stage}.{rest}"

vllm/models/deepseek_v4/nvidia/model.py

@@ -48,12 +48,7 @@ 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 (
-    EagleModelMixin,
-    MixtureOfExperts,
-    SupportsEagle3,
-    SupportsPP,
-)
+from vllm.model_executor.models.interfaces import MixtureOfExperts, SupportsPP
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
@@ -938,7 +933,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         return x, residual, post_mix, res_mix
 
 
-class DeepseekV4Model(nn.Module, EagleModelMixin):
+class DeepseekV4Model(nn.Module):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -1079,12 +1074,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             input_ids = input_ids.to(torch.int64)
 
         residual, post_mix, res_mix = None, None, None
-        aux_hidden_states: list[torch.Tensor] = []
-        final_aux_recon: torch.Tensor | None = None  # avoid duplicate mhc_post call
-        for idx, layer in enumerate(
-            islice(self.layers, self.start_layer, self.end_layer),
-            start=self.start_layer,
-        ):
+        for layer in islice(self.layers, self.start_layer, self.end_layer):
             hidden_states, residual, post_mix, res_mix = layer(
                 hidden_states,
                 positions,
@@ -1093,21 +1083,10 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                 res_mix,
                 residual,
             )
-            if idx + 1 in self.aux_hidden_state_layers:
-                # Reconstruct the aux hidden state for draft models
-                aux_recon = mhc_post_tilelang(
-                    hidden_states, residual, post_mix, res_mix
-                )
-                aux_hidden_states.append(aux_recon.mean(dim=1))
-                final_aux_recon = aux_recon
         if layer is not None:
-            # Reuse if the last layer was captured as an aux hidden state
-            if self.end_layer in self.aux_hidden_state_layers:
-                hidden_states = final_aux_recon
-            else:
-                hidden_states = mhc_post_tilelang(
-                    hidden_states, residual, post_mix, res_mix
-                )
+            hidden_states = mhc_post_tilelang(
+                hidden_states, residual, post_mix, res_mix
+            )
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})
@@ -1125,8 +1104,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             self.hc_eps,
         )
         hidden_states = self.norm(hidden_states)
-        if len(aux_hidden_states) > 0:
-            return hidden_states, aux_hidden_states
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -1353,9 +1330,7 @@ class DeepseekV4MixtureOfExperts(MixtureOfExperts):
             moe.experts.update_expert_map()
 
 
-class DeepseekV4ForCausalLM(
-    nn.Module, SupportsPP, SupportsEagle3, DeepseekV4MixtureOfExperts
-):
+class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts):
     model_cls = DeepseekV4Model
 
     # Default mapper assumes the original FP4-expert checkpoint layout.

vllm/utils/extensible_tensor.py

@@ -1,520 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Growable CUDA byte buffers backed by CUDA virtual memory management."""
-
-from __future__ import annotations
-
-import ctypes
-from contextlib import suppress
-
-import torch
-
-_CUDA_SUCCESS = 0
-_CU_MEM_ALLOCATION_TYPE_PINNED = 1
-_CU_MEM_LOCATION_TYPE_DEVICE = 1
-_CU_MEM_ALLOC_GRANULARITY_MINIMUM = 0
-_CU_MEM_ACCESS_FLAGS_PROT_READWRITE = 3
-_CU_MEM_ALLOCATION_COMP_NONE = 0
-
-
-class _CUmemLocation(ctypes.Structure):
-    _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)]
-
-
-class _CUmemAllocFlags(ctypes.Structure):
-    _fields_ = [
-        ("compressionType", ctypes.c_ubyte),
-        ("gpuDirectRDMACapable", ctypes.c_ubyte),
-        ("usage", ctypes.c_ushort),
-        ("reserved", ctypes.c_ubyte * 4),
-    ]
-
-
-class _CUmemAllocationProp(ctypes.Structure):
-    _fields_ = [
-        ("type", ctypes.c_int),
-        ("requestedHandleTypes", ctypes.c_int),
-        ("location", _CUmemLocation),
-        ("win32HandleMetaData", ctypes.c_void_p),
-        ("allocFlags", _CUmemAllocFlags),
-    ]
-
-
-class _CUmemAccessDesc(ctypes.Structure):
-    _fields_ = [("location", _CUmemLocation), ("flags", ctypes.c_int)]
-
-
-_CUdeviceptr = ctypes.c_ulonglong
-_CUmemHandle = ctypes.c_ulonglong
-_CUcontext = ctypes.c_void_p
-
-_libcuda: ctypes.CDLL | None = None
-
-
-def _find_loaded_library(lib_name: str) -> str | None:
-    try:
-        with open("/proc/self/maps") as f:
-            for line in f:
-                if lib_name not in line:
-                    continue
-                start = line.index("/")
-                return line[start:].strip()
-    except (OSError, ValueError):
-        return None
-    return None
-
-
-def _load_libcuda() -> ctypes.CDLL:
-    for name in ("libcuda.so.1", "libcuda.so"):
-        try:
-            return ctypes.CDLL(name)
-        except OSError:
-            continue
-    if path := _find_loaded_library("libcuda"):
-        return ctypes.CDLL(path)
-    raise RuntimeError(
-        "Could not load libcuda. The CUDA driver library is required for "
-        "ExtensibleTensor."
-    )
-
-
-def _configure_signatures(lib: ctypes.CDLL) -> None:
-    pointer = ctypes.POINTER
-    lib.cuGetErrorString.argtypes = [ctypes.c_int, pointer(ctypes.c_char_p)]
-    lib.cuCtxGetCurrent.argtypes = [pointer(_CUcontext)]
-    lib.cuDevicePrimaryCtxRetain.argtypes = [pointer(_CUcontext), ctypes.c_int]
-    lib.cuCtxSetCurrent.argtypes = [_CUcontext]
-    lib.cuMemGetAllocationGranularity.argtypes = [
-        pointer(ctypes.c_size_t),
-        pointer(_CUmemAllocationProp),
-        ctypes.c_int,
-    ]
-    lib.cuMemAddressReserve.argtypes = [
-        pointer(_CUdeviceptr),
-        ctypes.c_size_t,
-        ctypes.c_size_t,
-        _CUdeviceptr,
-        ctypes.c_ulonglong,
-    ]
-    lib.cuMemCreate.argtypes = [
-        pointer(_CUmemHandle),
-        ctypes.c_size_t,
-        pointer(_CUmemAllocationProp),
-        ctypes.c_ulonglong,
-    ]
-    lib.cuMemMap.argtypes = [
-        _CUdeviceptr,
-        ctypes.c_size_t,
-        ctypes.c_size_t,
-        _CUmemHandle,
-        ctypes.c_ulonglong,
-    ]
-    lib.cuMemSetAccess.argtypes = [
-        _CUdeviceptr,
-        ctypes.c_size_t,
-        pointer(_CUmemAccessDesc),
-        ctypes.c_size_t,
-    ]
-    lib.cuMemUnmap.argtypes = [_CUdeviceptr, ctypes.c_size_t]
-    lib.cuMemRelease.argtypes = [_CUmemHandle]
-    lib.cuMemAddressFree.argtypes = [_CUdeviceptr, ctypes.c_size_t]
-
-    for fn in (
-        lib.cuGetErrorString,
-        lib.cuCtxGetCurrent,
-        lib.cuDevicePrimaryCtxRetain,
-        lib.cuCtxSetCurrent,
-        lib.cuMemGetAllocationGranularity,
-        lib.cuMemAddressReserve,
-        lib.cuMemCreate,
-        lib.cuMemMap,
-        lib.cuMemSetAccess,
-        lib.cuMemUnmap,
-        lib.cuMemRelease,
-        lib.cuMemAddressFree,
-    ):
-        fn.restype = ctypes.c_int
-
-
-def _cuda() -> ctypes.CDLL:
-    global _libcuda
-    if _libcuda is None:
-        lib = _load_libcuda()
-        _configure_signatures(lib)
-        _libcuda = lib
-    return _libcuda
-
-
-def _check(result: int) -> None:
-    if result == _CUDA_SUCCESS:
-        return
-    msg = ctypes.c_char_p()
-    _cuda().cuGetErrorString(result, ctypes.byref(msg))
-    detail = msg.value.decode() if msg.value else "unknown error"
-    raise RuntimeError(f"CUDA driver error {result}: {detail}")
-
-
-def _ensure_context(device_index: int) -> None:
-    pctx = _CUcontext()
-    _check(_cuda().cuCtxGetCurrent(ctypes.byref(pctx)))
-    if pctx.value:
-        return
-    _check(_cuda().cuDevicePrimaryCtxRetain(ctypes.byref(pctx), device_index))
-    _check(_cuda().cuCtxSetCurrent(pctx))
-
-
-def _make_alloc_prop(device_index: int) -> _CUmemAllocationProp:
-    prop = _CUmemAllocationProp()
-    prop.type = _CU_MEM_ALLOCATION_TYPE_PINNED
-    prop.location.type = _CU_MEM_LOCATION_TYPE_DEVICE
-    prop.location.id = device_index
-    prop.allocFlags.compressionType = _CU_MEM_ALLOCATION_COMP_NONE
-    return prop
-
-
-def _round_up(value: int, multiple: int) -> int:
-    return ((value + multiple - 1) // multiple) * multiple
-
-
-class _VirtualBuffer:
-    """Own one device VA reservation and the physical chunks mapped into it.
-
-    Physical memory is committed incrementally, at granularity-sized granules,
-    via `ensure_committed_range`; granules already mapped by an earlier
-    (possibly overlapping) range are skipped, so ranges may abut or overlap
-    freely.
-    """
-
-    def __init__(self, max_bytes: int, device_index: int) -> None:
-        _ensure_context(device_index)
-        self.device_index = device_index
-
-        prop = _make_alloc_prop(device_index)
-        granularity = ctypes.c_size_t()
-        _check(
-            _cuda().cuMemGetAllocationGranularity(
-                ctypes.byref(granularity),
-                ctypes.byref(prop),
-                _CU_MEM_ALLOC_GRANULARITY_MINIMUM,
-            )
-        )
-        self.granularity: int = granularity.value
-        self.reserved_size: int = _round_up(max(max_bytes, 1), self.granularity)
-
-        dptr = _CUdeviceptr()
-        _check(
-            _cuda().cuMemAddressReserve(ctypes.byref(dptr), self.reserved_size, 0, 0, 0)
-        )
-        self.base_ptr: int = dptr.value
-
-        # Granule indices (VA offset // granularity) that have physical
-        # memory mapped.
-        self._mapped_granules: set[int] = set()
-        # Each entry is (handle, va_offset, size) for one mapped physical chunk.
-        self._handles: list[tuple[int, int, int]] = []
-        self._freed: bool = False
-
-    @property
-    def committed_bytes(self) -> int:
-        """Total physically mapped bytes (a multiple of the granularity)."""
-        return len(self._mapped_granules) * self.granularity
-
-    def ensure_committed(self, nbytes: int) -> None:
-        """Map physical pages so that at least the first `nbytes` are backed."""
-        self.ensure_committed_range(0, nbytes)
-
-    def ensure_committed_range(self, start: int, end: int) -> None:
-        """Map physical pages so that the byte range `[start, end)` is backed.
-
-        The range is widened outward to granule boundaries; granules mapped by
-        earlier calls are skipped, so a granule shared by two requested ranges
-        is mapped once.
-        """
-        if not 0 <= start <= end:
-            raise ValueError(f"Invalid range [{start}, {end}).")
-        if end > self.reserved_size:
-            raise ValueError(
-                f"Requested range end {end} exceeds reserved capacity "
-                f"{self.reserved_size}."
-            )
-        if start == end:
-            return
-        first = start // self.granularity
-        last = (end + self.granularity - 1) // self.granularity  # exclusive
-        run_start: int | None = None
-        for g in range(first, last + 1):
-            unmapped = g < last and g not in self._mapped_granules
-            if unmapped and run_start is None:
-                run_start = g
-            elif not unmapped and run_start is not None:
-                self._map_chunk_at(
-                    run_start * self.granularity, (g - run_start) * self.granularity
-                )
-                self._mapped_granules.update(range(run_start, g))
-                run_start = None
-
-    def _map_chunk_at(self, offset: int, size: int) -> None:
-        """Create one physical chunk of `size` bytes and map it at `offset`."""
-        _ensure_context(self.device_index)
-        prop = _make_alloc_prop(self.device_index)
-
-        handle = _CUmemHandle()
-        _check(_cuda().cuMemCreate(ctypes.byref(handle), size, ctypes.byref(prop), 0))
-
-        addr = self.base_ptr + offset
-        try:
-            _check(_cuda().cuMemMap(addr, size, 0, handle, 0))
-        except RuntimeError:
-            _cuda().cuMemRelease(handle)
-            raise
-
-        desc = _CUmemAccessDesc()
-        desc.location.type = _CU_MEM_LOCATION_TYPE_DEVICE
-        desc.location.id = self.device_index
-        desc.flags = _CU_MEM_ACCESS_FLAGS_PROT_READWRITE
-        _check(_cuda().cuMemSetAccess(addr, size, ctypes.byref(desc), 1))
-
-        self._handles.append((handle.value, offset, size))
-
-    def free(self) -> None:
-        if self._freed:
-            return
-        self._freed = True
-        _ensure_context(self.device_index)
-        if self._handles:
-            torch.cuda.synchronize(self.device_index)
-        for handle, offset, size in self._handles:
-            _check(_cuda().cuMemUnmap(self.base_ptr + offset, size))
-            _check(_cuda().cuMemRelease(handle))
-        if self.base_ptr:
-            _check(_cuda().cuMemAddressFree(self.base_ptr, self.reserved_size))
-        self._handles = []
-        self._mapped_granules = set()
-        self.base_ptr = 0
-
-    def __del__(self) -> None:
-        with suppress(Exception):
-            self.free()
-
-
-_K_DL_CUDA = 2
-_K_DL_UINT = 1
-_UINT8_BITS = 8
-
-
-class _DLDevice(ctypes.Structure):
-    _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int)]
-
-
-class _DLDataType(ctypes.Structure):
-    _fields_ = [
-        ("code", ctypes.c_uint8),
-        ("bits", ctypes.c_uint8),
-        ("lanes", ctypes.c_uint16),
-    ]
-
-
-class _DLTensor(ctypes.Structure):
-    _fields_ = [
-        ("data", ctypes.c_void_p),
-        ("device", _DLDevice),
-        ("ndim", ctypes.c_int),
-        ("dtype", _DLDataType),
-        ("shape", ctypes.POINTER(ctypes.c_int64)),
-        ("strides", ctypes.POINTER(ctypes.c_int64)),
-        ("byte_offset", ctypes.c_uint64),
-    ]
-
-
-class _DLManagedTensor(ctypes.Structure):
-    pass
-
-
-_DLDeleter = ctypes.CFUNCTYPE(None, ctypes.POINTER(_DLManagedTensor))
-_DLManagedTensor._fields_ = [
-    ("dl_tensor", _DLTensor),
-    ("manager_ctx", ctypes.c_void_p),
-    ("deleter", _DLDeleter),
-]
-
-_KEEPALIVE: dict[int, tuple[object, object, object]] = {}
-_PyCapsule_New = ctypes.pythonapi.PyCapsule_New
-_PyCapsule_New.restype = ctypes.py_object
-_PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
-
-
-def _uint8_tensor_from_ptr(ptr: int, num_bytes: int, device_index: int) -> torch.Tensor:
-    shape_arr = (ctypes.c_int64 * 1)(num_bytes)
-
-    managed = _DLManagedTensor()
-    managed.dl_tensor.data = ctypes.c_void_p(ptr)
-    managed.dl_tensor.device = _DLDevice(_K_DL_CUDA, device_index)
-    managed.dl_tensor.ndim = 1
-    managed.dl_tensor.dtype = _DLDataType(_K_DL_UINT, _UINT8_BITS, 1)
-    managed.dl_tensor.shape = ctypes.cast(shape_arr, ctypes.POINTER(ctypes.c_int64))
-    managed.dl_tensor.strides = None
-    managed.dl_tensor.byte_offset = 0
-    managed.manager_ctx = None
-
-    key = ctypes.addressof(managed)
-
-    def _deleter(_managed_ptr: object) -> None:
-        _KEEPALIVE.pop(key, None)
-
-    deleter = _DLDeleter(_deleter)
-    managed.deleter = deleter
-    _KEEPALIVE[key] = (managed, shape_arr, deleter)
-
-    capsule = _PyCapsule_New(ctypes.addressof(managed), b"dltensor", None)
-    return torch.from_dlpack(capsule)
-
-
-class ExtensibleTensor:
-    """A 1-D CUDA byte buffer that can grow without moving its base pointer.
-
-    With `num_segments > 1` the reservation is divided into that many equal
-    segments that grow in lockstep via `resize_per_segment_`: the committed
-    bytes form a prefix of each segment (segment `i` spans
-    `[i * segment_capacity_bytes, (i + 1) * segment_capacity_bytes)` of
-    `full_view()`). This backs layouts whose block dimension is not outermost,
-    e.g. a K/V-split KV cache (`num_segments=2`). `resize_` / `tensor` /
-    `append` assume a single contiguous prefix and are only valid when
-    `num_segments == 1`.
-    """
-
-    def __init__(
-        self,
-        max_num_bytes: int,
-        device: torch.device | str | int | None = None,
-        num_segments: int = 1,
-    ) -> None:
-        if max_num_bytes < 0:
-            raise ValueError("max_num_bytes must be non-negative.")
-        if num_segments < 1:
-            raise ValueError(f"num_segments must be positive, got {num_segments}.")
-        if max_num_bytes % num_segments != 0:
-            raise ValueError(
-                f"max_num_bytes ({max_num_bytes}) must be divisible by "
-                f"num_segments ({num_segments})."
-            )
-
-        if device is None:
-            device = torch.cuda.current_device()
-        dev = device if isinstance(device, torch.device) else torch.device(device)
-        if dev.type != "cuda":
-            raise ValueError(f"ExtensibleTensor requires a cuda device, got {dev}.")
-        self._device_index: int = (
-            dev.index if dev.index is not None else torch.cuda.current_device()
-        )
-
-        torch.cuda.init()
-
-        self._max_num_bytes: int = max_num_bytes
-        self._num_segments: int = num_segments
-        self._segment_capacity_bytes: int = max_num_bytes // num_segments
-        self._buffer: _VirtualBuffer = _VirtualBuffer(max_num_bytes, self._device_index)
-        self._bytes_per_segment: int = 0
-
-    @property
-    def tensor(self) -> torch.Tensor:
-        """Return a uint8 tensor view of the currently committed prefix."""
-        if self._num_segments != 1:
-            raise ValueError(
-                "tensor (a single committed prefix) is only valid for "
-                "num_segments=1; use full_view() and index segments explicitly."
-            )
-        return _uint8_tensor_from_ptr(
-            self._buffer.base_ptr, self._bytes_per_segment, self._device_index
-        )
-
-    def full_view(self) -> torch.Tensor:
-        """Return a uint8 tensor view spanning the requested maximum size."""
-        return _uint8_tensor_from_ptr(
-            self._buffer.base_ptr, self._max_num_bytes, self._device_index
-        )
-
-    def resize_(self, num_bytes: int) -> torch.Tensor:
-        """Grow the buffer to `num_bytes` and return the committed-prefix view."""
-        if self._num_segments != 1:
-            raise ValueError(
-                "resize_ (a single committed prefix) is only valid for "
-                "num_segments=1; use resize_per_segment_."
-            )
-        self.resize_per_segment_(num_bytes)
-        return self.tensor
-
-    def resize_per_segment_(
-        self, bytes_per_segment: int, zero_new: bool = False
-    ) -> None:
-        """Grow every segment's committed prefix to `bytes_per_segment` bytes.
-
-        Existing bytes are preserved and the base pointer is unchanged. With
-        `zero_new=True` the newly committed byte range of each segment is
-        zeroed (bytes committed earlier are left intact). Raises if
-        `bytes_per_segment` is smaller than the current per-segment size
-        (shrink is unsupported) or larger than `segment_capacity_bytes`.
-        """
-        old = self._bytes_per_segment
-        if bytes_per_segment < old:
-            raise ValueError(
-                f"ExtensibleTensor is grow-only: cannot resize from {old} "
-                f"to {bytes_per_segment} bytes per segment."
-            )
-        if bytes_per_segment > self._segment_capacity_bytes:
-            raise ValueError(
-                f"Requested {bytes_per_segment} bytes per segment exceeds the "
-                f"segment capacity {self._segment_capacity_bytes}."
-            )
-        if bytes_per_segment == old:
-            return
-        for i in range(self._num_segments):
-            start = i * self._segment_capacity_bytes
-            self._buffer.ensure_committed_range(start + old, start + bytes_per_segment)
-        self._bytes_per_segment = bytes_per_segment
-        if zero_new:
-            full = self.full_view()
-            for i in range(self._num_segments):
-                start = i * self._segment_capacity_bytes
-                full[start + old : start + bytes_per_segment].zero_()
-
-    def append(self, num_bytes: int) -> torch.Tensor:
-        """Grow by `num_bytes` additional bytes and return the new view."""
-        if num_bytes < 0:
-            raise ValueError("num_bytes to append must be non-negative.")
-        return self.resize_(self._bytes_per_segment + num_bytes)
-
-    @property
-    def num_bytes(self) -> int:
-        """Current committed size in bytes, summed over all segments."""
-        return self._bytes_per_segment * self._num_segments
-
-    @property
-    def bytes_per_segment(self) -> int:
-        """Current committed prefix size of each segment in bytes."""
-        return self._bytes_per_segment
-
-    @property
-    def num_segments(self) -> int:
-        """Number of equal segments the reservation is divided into."""
-        return self._num_segments
-
-    @property
-    def segment_capacity_bytes(self) -> int:
-        """Maximum size of each segment (`max_num_bytes / num_segments`)."""
-        return self._segment_capacity_bytes
-
-    @property
-    def capacity_bytes(self) -> int:
-        return self._buffer.reserved_size
-
-    @property
-    def base_ptr(self) -> int:
-        return self._buffer.base_ptr
-
-    @property
-    def device(self) -> torch.device:
-        return torch.device("cuda", self._device_index)
-
-    def free(self) -> None:
-        self._buffer.free()
-        self._bytes_per_segment = 0