[Core] Enable extensible KV cache for all attention backends and Mamba

Port of internal D110967544. The extensible KV cache flow (reserve KV
virtual address space up front, capture CUDA graphs first, then size and
commit the KV cache from post-capture free memory) previously required a
block-major attention backend and rejected Mamba models. This enables it
for every backend layout and for Mamba / linear attention:

- ExtensibleTensor gains num_segments: the reservation is divided into
  equal segments that grow in lockstep, with committed bytes forming a
  prefix of each segment. Physical pages are mapped at
  allocation-granularity granules and deduped across overlapping ranges,
  so a granule straddling a segment boundary is mapped exactly once.
  resize_per_segment_(bytes, zero_new=True) zeroes only the newly
  committed logical range of each segment.
- Each KV cache buffer keeps its layers' physical layout and is committed
  as one prefix per layout segment. The segment count is derived from the
  backend's get_kv_cache_shape / get_kv_cache_block_dim / stride order:
  K/V-split layouts (e.g. FlashAttention) get one prefix per half,
  block-major layouts (e.g. FlashInfer, MLA) a single prefix. Mamba state
  pages are block-major per layer, and hybrid-model attention caches are
  re-strided to block-major, so both use a single segment.
- Removed the supports_extensible_kv_cache gate plumbing from EngineCore,
  Executor, Worker, WorkerBase and GPUModelRunner; a CUDA platform check
  remains in EngineCore.
- enable_extensible_kv_cache is reported as unsupported by the V2 model
  runner so V2-default models fall back to the V1 runner (which implements
  the flow); also fixed initialize_kv_cache being called with the
  extensible kwarg on runners that do not accept it, which broke every
  default V2-runner boot on this branch.

Tested on H100:
- tests/utils_/test_extensible_tensor.py (5 passed, incl. new segmented
  lockstep-grow/zero, granule-dedup and invalid-usage tests)
- tests/v1/worker/test_extensible_kv_cache.py (new, 6 passed: segment
  derivation, split grows both halves, block-major, legacy full commit,
  Mamba per-layer growth, hybrid attention+Mamba)
- E2E Qwen3-0.6B greedy with VLLM_ATTENTION_BACKEND=FLASH_ATTN (a K/V-split
  backend the old gate rejected): extensible generations byte-identical to
  the legacy path; log shows reserve then "Extended KV cache to 34663
  blocks". V2->V1 auto-fallback path verified as well.

.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]"'
+      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]"'
 
 - label: LoRA Punica FP8/XPU Ops
   timeout_in_minutes: 45

examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py

@@ -327,6 +327,9 @@ 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/src/chat/src/lib.rs

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

rust/src/chat/tests/chat.rs

@@ -15,9 +15,10 @@ use vllm_chat::{
 use vllm_engine_core_client::protocol::logprobs::{
     Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
 };
-use vllm_engine_core_client::protocol::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
+use vllm_engine_core_client::protocol::output::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, 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::{
-    EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
-};
+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::{
     EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode,
 };

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

@@ -11,10 +11,11 @@ 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,12 +16,11 @@ 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::{
+    use crate::protocol::output::{
         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::{EngineCoreFinishReason, EngineCoreOutput};
+use crate::protocol::output::{EngineCoreFinishReason, EngineCoreOutput};
 use crate::{AbortCause, Error, Result};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

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

@@ -10,10 +10,9 @@ 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::{
-    ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs, EngineCoreRequestType,
-    encode_msgpack,
-};
+use crate::protocol::encode_msgpack;
+use crate::protocol::output::{ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs};
+use crate::protocol::request::EngineCoreRequestType;
 
 /// 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,8 +8,9 @@ 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::{ModelDtype, decode_msgpack, encode_msgpack};
+use crate::protocol::{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/handshake.rs

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

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

@@ -9,8 +9,7 @@ use enum_as_inner::EnumAsInner;
 use serde::{Deserialize, Deserializer, Serialize};
 
 use self::wire::*;
-use super::{EngineCoreOutput, EngineCoreOutputs, decode_msgpack};
-use crate::error::{Error, Result, bail_ext_value_decode, ext_value_decode};
+use crate::error::{Error, Result, bail_ext_value_decode};
 use crate::protocol::tensor::{WireArrayData, WireNdArray};
 
 /// One token candidate and its logprob metadata for a single sequence position.
@@ -160,7 +159,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.
-    fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
+    pub(super) fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
     where
         Frame: AsRef<[u8]>,
     {
@@ -171,37 +170,6 @@ 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.
     ///
@@ -315,16 +283,3 @@ 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, decode_engine_core_outputs};
-use crate::protocol::EngineCoreFinishReason;
+use super::{Logprobs, PositionLogprobs, TokenLogprob};
+use crate::protocol::output::{EngineCoreFinishReason, decode_engine_core_outputs};
 
 fn encode_value(value: &Value) -> Vec<u8> {
     let mut out = Vec::new();

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

@@ -1,28 +1,11 @@
 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.
@@ -36,499 +19,18 @@ 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>>
@@ -564,81 +66,17 @@ 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::BTreeSet;
+    use std::collections::BTreeMap;
 
     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>(
@@ -648,72 +86,4 @@ 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/output.rs

@@ -1,10 +1,164 @@
 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 super::{EngineCoreOutput, EngineCoreOutputs};
-use crate::protocol::stats::SchedulerStats;
+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>,
+}
 
 /// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -49,6 +203,18 @@ 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()
@@ -92,12 +258,62 @@ 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::EngineCoreOutput;
+    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()]))
+        );
+    }
 
     #[test]
     fn engine_core_outputs_classify_request_batch() {

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

@@ -0,0 +1,175 @@
+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

@@ -0,0 +1,211 @@
+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

@@ -0,0 +1,81 @@
+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,13 +22,14 @@ 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,9 +18,8 @@ use crate::error::{Error, Result, bail_unexpected_handshake_message};
 use crate::protocol::handshake::{
     EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage,
 };
-use crate::protocol::{
-    EngineCoreOutputs, decode_engine_core_outputs, decode_msgpack, encode_msgpack,
-};
+use crate::protocol::output::{EngineCoreOutputs, decode_engine_core_outputs};
+use crate::protocol::{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::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::sampling::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::{EngineCoreFinishReason, StopReason};
+use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason};
 use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
 
 use crate::error::Result;

rust/src/llm/src/request.rs

@@ -4,9 +4,8 @@ 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::{
-    EngineCoreRequest, EngineCoreSamplingParams, ReasoningParserKwargs,
-};
+use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningParserKwargs};
+use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
 
 use crate::error::{Error, Result};
 
@@ -134,7 +133,8 @@ fn current_unix_timestamp_secs() -> f64 {
 mod tests {
     use std::collections::BTreeMap;
 
-    use vllm_engine_core_client::protocol::{EngineCoreSamplingParams, ReasoningParserKwargs};
+    use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
+    use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
 
     use super::GenerateRequest;
     use crate::error::Error;

rust/src/llm/src/request_metrics.rs

@@ -1,7 +1,9 @@
 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,
 };
@@ -328,8 +330,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};
 
@@ -341,7 +343,7 @@ mod tests {
             2,
             10.0,
             100.2,
-            &vllm_engine_core_client::protocol::EngineCoreOutput {
+            &vllm_engine_core_client::protocol::output::EngineCoreOutput {
                 request_id: "req-1".to_string(),
                 new_token_ids: vec![1],
                 finish_reason: None,
@@ -369,7 +371,7 @@ mod tests {
             2,
             11.5,
             100.4,
-            &vllm_engine_core_client::protocol::EngineCoreOutput {
+            &vllm_engine_core_client::protocol::output::EngineCoreOutput {
                 request_id: "req-1".to_string(),
                 new_token_ids: vec![2, 3],
                 finish_reason: None,

rust/src/llm/tests/generate.rs

@@ -9,11 +9,13 @@ use uuid::Uuid;
 use vllm_engine_core_client::protocol::logprobs::{
     Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
 };
-use vllm_engine_core_client::protocol::stats::PrefillStats;
-use vllm_engine_core_client::protocol::{
+use vllm_engine_core_client::protocol::output::{
     EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
-    EngineCoreOutputs, EngineCoreRequest, EngineCoreSamplingParams,
+    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::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,12 +11,13 @@ 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,10 +4,9 @@ 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::{
-    EngineCoreRequest, EngineCoreRequestType, decode_msgpack, encode_msgpack,
-};
+use vllm_engine_core_client::protocol::{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::{
-    EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
-};
+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::test_utils::IpcNamespace;
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
 

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

@@ -5,14 +5,13 @@
 
 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,10 @@ fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use serde_json::json;
 
+    use super::*;
+
     fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
         Python::initialize();
         Python::attach(f)

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

@@ -49,10 +49,8 @@ mod tests {
     use std::sync::Arc;
 
     use super::SeedOssReasoningParser;
-    use crate::reasoning::{
-        ReasoningParser,
-        tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer},
-    };
+    use crate::reasoning::ReasoningParser;
+    use crate::reasoning::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,10 +127,8 @@ mod tests {
     use std::sync::Arc;
 
     use super::Step3p5ReasoningParser;
-    use crate::reasoning::{
-        ReasoningParser,
-        tests::{THINK_START_ID, fake_tokenizer},
-    };
+    use crate::reasoning::ReasoningParser;
+    use crate::reasoning::tests::{THINK_START_ID, fake_tokenizer};
 
     #[test]
     fn picks_up_prompt_start_boundary() {

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

@@ -14,8 +14,6 @@ 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};
@@ -35,6 +33,8 @@ 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,11 +2,10 @@
 
 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,4 +1,5 @@
 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};
@@ -6,8 +7,6 @@ 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,13 +3,12 @@
 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,8 +1,7 @@
 use axum::Json;
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
-use thiserror_ext::AsReport as _;
-use thiserror_ext::{Construct, Macro};
+use thiserror_ext::{AsReport as _, Construct, Macro};
 
 use crate::routes::openai::utils::types::{ErrorDetail, ErrorResponse};
 

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

@@ -3,7 +3,8 @@
 
 use tonic::Status;
 use uuid::Uuid;
-use vllm_engine_core_client::protocol::{StopReason, StructuredOutputsParams};
+use vllm_engine_core_client::protocol::output::StopReason;
+use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;
 use vllm_text::{
     DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams,
     TextDecodeOptions, TextRequest,
@@ -502,7 +503,7 @@ impl ResponseOpts {
 
 #[cfg(test)]
 mod tests {
-    use vllm_engine_core_client::protocol::StopReason;
+    use vllm_engine_core_client::protocol::output::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,9 +18,10 @@ use vllm_chat::{
     ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
     DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
 };
-use vllm_engine_core_client::protocol::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
+use vllm_engine_core_client::protocol::output::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
 };
+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 Connection = ListenerIo;
     type Address = ListenerAddr;
+    type Connection = ListenerIo;
     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,9 +18,10 @@ use vllm_chat::{
     ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
     DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
 };
-use vllm_engine_core_client::protocol::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
+use vllm_engine_core_client::protocol::output::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
 };
+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::StopReason;
+use vllm_engine_core_client::protocol::output::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::StopReason;
+    use vllm_engine_core_client::protocol::output::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::StructuredOutputsParams;
+use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;
 
 use crate::error::ApiError;
 

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

@@ -24,14 +24,15 @@ 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::utility::{UtilityOutput, UtilityResultEnvelope};
-use vllm_engine_core_client::protocol::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
-    decode_value,
+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::test_utils::{
     IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready,
 };

rust/src/server/src/state.rs

@@ -1,6 +1,5 @@
-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/src/lower.rs

@@ -3,15 +3,15 @@ use std::collections::BTreeSet;
 pub(crate) mod logprobs;
 pub(crate) mod token_ids;
 
-use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
+use logprobs::validate_logprobs;
+use token_ids::{validate_prompt_token_ids, validate_vocab_range};
+use vllm_engine_core_client::protocol::sampling::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,10 @@
 //! `-1` is expanded only for bounds checks. The original request values are
 //! passed through to engine-core.
 
-use crate::backend::SamplingLimits;
 use thiserror::Error;
 
+use crate::backend::SamplingLimits;
+
 #[derive(Debug, Error)]
 pub enum LogprobsError {
     #[error("{parameter} must be non-negative or -1, got {value}")]

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::EngineCoreSamplingParams;
+use vllm_engine_core_client::protocol::sampling::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::StopReason;
+use vllm_engine_core_client::protocol::output::StopReason;
 use vllm_llm::{FinishReason, GenerateOutput, TokenUsage};
 use vllm_tokenizer::{DynTokenizer, IncrementalDecoder};
 

rust/src/text/src/request.rs

@@ -5,7 +5,8 @@ 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::{ReasoningParserKwargs, StructuredOutputsParams};
+use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
+use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;
 
 use crate::error::{Error, Result};
 use crate::output::TextDecodeOptions;

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

@@ -1,11 +1,12 @@
+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,6 +34,16 @@ 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)
@@ -198,3 +208,64 @@ 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,6 +1407,19 @@ _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,6 +48,10 @@ 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

@@ -0,0 +1,141 @@
+# 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

@@ -0,0 +1,529 @@
+# 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,6 +149,30 @@ 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,6 +1412,67 @@ 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,6 +49,18 @@ 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,11 +25,10 @@ def test_traces(
 ):
     with monkeypatch.context() as m:
         m.setenv(OTEL_EXPORTER_OTLP_TRACES_INSECURE, "true")
-        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")
+        # 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")
 
         sampling_params = SamplingParams(
             temperature=0.01,

tests/v1/worker/test_extensible_kv_cache.py

@@ -0,0 +1,369 @@
+# 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,6 +435,40 @@ 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,6 +862,28 @@ 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,6 +2452,7 @@ 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,
@@ -4155,8 +4156,21 @@ class ASRDataset(HuggingFaceDataset):
         **kwargs,
     ) -> list[SampleRequest]:
         output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
-        if "openai" in getattr(tokenizer, "name_or_path", ""):
+        name_or_path = getattr(tokenizer, "name_or_path", "")
+        tok_class = type(tokenizer).__name__
+        if "openai" in 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,6 +174,15 @@ 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
@@ -217,6 +226,8 @@ 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,6 +54,7 @@ MTPModelTypes = Literal[
 ]
 NgramGPUTypes = Literal["ngram_gpu"]
 DFlashModelTypes = Literal["dflash"]
+DSparkModelTypes = Literal["dspark"]
 EagleModelTypes = Literal[
     "eagle", "eagle3", "extract_hidden_states", MTPModelTypes, DFlashModelTypes
 ]
@@ -66,6 +67,7 @@ SpeculativeMethod = Literal[
     "custom_class",
     EagleModelTypes,
     NgramGPUTypes,
+    DSparkModelTypes,
 ]
 RejectionSampleMethod = Literal["standard", "synthetic", "block"]
 DraftSampleMethod = Literal["greedy", "probabilistic"]
@@ -291,6 +293,7 @@ class SpeculativeConfig:
             "eagle3",
             "extract_hidden_states",
             "dflash",
+            "dspark",
         )
         factors.append(uses_aux_hidden_states)
 
@@ -608,6 +611,13 @@ 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":
@@ -755,18 +765,24 @@ class SpeculativeConfig:
                         draft_hf.truncated_vocab_size = target_vocab
 
                 # Automatically detect the method
-                if self.method in ("eagle", "eagle3", "dflash"):
+                if self.method in ("eagle", "eagle3", "dflash", "dspark"):
                     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":
@@ -813,7 +829,18 @@ class SpeculativeConfig:
                         self.draft_model_config.hf_config = eagle_config
                         self.update_arch_()
 
-                if self.method == "dflash":
+                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"):
                     self.parallel_drafting = True
 
                 if self.num_speculative_tokens is not None and hasattr(
@@ -1129,11 +1156,17 @@ class SpeculativeConfig:
         )
 
     def use_eagle(self) -> bool:
-        return self.method in ("eagle", "eagle3", "mtp", "dflash")
+        # 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")
 
     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,6 +524,16 @@ 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
 
@@ -958,10 +968,11 @@ 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 kind of "
+                        "with EAGLE/MTP/Draft Model/NGram GPU/DSpark kind of "
                         "speculative decoding"
                     )
                 if self.speculative_config.disable_padded_drafter_batch:
@@ -989,6 +1000,7 @@ 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 "
@@ -2038,17 +2050,24 @@ 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"):
+            elif speculative_config.method not in (
+                "eagle",
+                "eagle3",
+                "mtp",
+                "dflash",
+                "dspark",
+            ):
                 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 uses parallel drafting natively in V2 via DFlashSpeculator.
+            # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle).
+            # DFlash and DSpark use parallel drafting natively in V2 via their
+            # own speculators.
             if (
                 speculative_config.parallel_drafting
-                and speculative_config.method != "dflash"
+                and speculative_config.method not in ("dflash", "dspark")
             ):
                 unsupported.append("parallel drafting for EAGLE speculative decoding")
 
@@ -2093,6 +2112,9 @@ 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,6 +522,7 @@ 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
@@ -1152,6 +1153,10 @@ 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"]
@@ -1869,6 +1874,7 @@ 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,6 +119,11 @@ 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
@@ -211,6 +216,7 @@ 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,
@@ -320,6 +326,7 @@ 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,29 +3,27 @@
 
 import torch
 
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
+from vllm.utils.torch_utils import direct_register_custom_op
 
 
-def bgmv_shrink(
+def _bgmv_shrink_impl(
     inputs: torch.Tensor,
     lora_a_weights: torch.Tensor,
     output_tensor: torch.Tensor,
     lora_indices_tensor: torch.Tensor,
-    scaling: float = 1.0,
+    scaling: float,
 ) -> None:
     torch.ops._xpu_C.bgmv_shrink(
         output_tensor, inputs, lora_a_weights, lora_indices_tensor, scaling
     )
 
 
-def bgmv_expand(
+def _bgmv_expand_impl(
     inputs: torch.Tensor,
     lora_b_weights: torch.Tensor,
     output_tensor: torch.Tensor,
     lora_indices_tensor: torch.Tensor,
-    add_inputs: bool = True,
+    add_inputs: bool,
 ) -> None:
     weight_out_dim = lora_b_weights.size(-2)
     output_dim = output_tensor.size(1)
@@ -65,14 +63,14 @@ def bgmv_expand(
         )
 
 
-def bgmv_expand_slice(
+def _bgmv_expand_slice_impl(
     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,
+    add_inputs: bool,
 ) -> None:
     assert slice_size == lora_b_weights.size(-2)
     assert slice_offset + slice_size <= output_tensor.size(1)
@@ -85,3 +83,101 @@ def bgmv_expand_slice(
         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,6 +62,10 @@ 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,
@@ -106,6 +110,14 @@ 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
         )
@@ -179,7 +191,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
                 add_inputs=add_inputs,
             )
             offset_start += output_slices[slice_idx]
-        y.view_as(y_org)
+        y = y.view_as(y_org)
 
     def add_lora_embedding(
         self,
@@ -227,7 +239,6 @@ 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
@@ -240,13 +251,17 @@ class PunicaWrapperXPU(PunicaWrapperBase):
 
         assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
 
-        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,
-            )
+        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)
         self.add_shrink(
             buffer,  # type: ignore
             x,
@@ -259,7 +274,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             buffer,  # type: ignore
             lora_b_stacked,
             output_slices,
-            add_inputs=True,
+            add_inputs=add_inputs,
             **kwargs,
         )
 
@@ -300,12 +315,16 @@ class PunicaWrapperXPU(PunicaWrapperBase):
         y = y.view(-1, y.shape[-1])
         x = x.view(-1, x.shape[-1])
         r = lora_b_stacked.size(-1)
-        if buffer is None:
-            buffer = torch.zeros((x.size(0), r), dtype=x.dtype, device=x.device)
+
+        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)
         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)
-        return y.view_as(y_org)
+        y = y.view_as(y_org)
 
     def moe_lora_align_block_size(
         self,
@@ -318,33 +337,62 @@ 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, _, _, _, lora_ids, _, _) = (
-            self.token_mapping_meta.meta_args(
-                num_tokens, self.lora_config.specialize_active_lora
-            )
+        (
+            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
         )
         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() + num_experts * (block_size - 1)
+            max_num_tokens_padded = topk_ids.numel() + kernel_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 must be set default to -1 to prevent a blank block
-            expert_ids = torch.empty(
+            # 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(
                 (max_loras * max_num_m_blocks,),
+                -1,
                 dtype=torch.int32,
                 device=topk_ids.device,
             )
@@ -355,7 +403,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             ops.moe_lora_align_block_size(
                 topk_ids,
                 token_lora_mapping,
-                num_experts,
+                kernel_num_experts,
                 block_size,
                 max_loras,
                 max_num_tokens_padded,
@@ -365,11 +413,10 @@ 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 None, sorted_ids, expert_ids, num_tokens_post_pad
+        return token_lora_mapping, sorted_ids, expert_ids, num_tokens_post_pad
 
     def add_lora_fused_moe(
         self,
@@ -525,7 +572,8 @@ class PunicaWrapperXPU(PunicaWrapperBase):
 
         SPARSITY_FACTOR = 8
         naive_block_assignment = (
-            expert_map is None
+            not fully_sharded
+            and expert_map is None
             and num_tokens * top_k * SPARSITY_FACTOR <= local_num_experts * max_loras
         )
 
@@ -543,6 +591,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             adapter_enabled,
             expert_map,
             naive_block_assignment=naive_block_assignment,
+            token_lora_mapping=token_lora_mapping,
         )
 
         _sorted = sorted_token_ids_lora
@@ -567,6 +616,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
             adapter_enabled,
             fully_sharded=fully_sharded,
             token_lora_mapping=token_lora_mapping,
+            add_inputs=add_inputs,
         )
 
         return (
@@ -680,4 +730,5 @@ 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.routed_experts.{e}.w1.weight",
+                        f".block_sparse_moe.experts.{e}.w1.weight",
                     )
                     w3_name = n.replace(
                         ".block_sparse_moe.input_linear.weight",
-                        f".block_sparse_moe.experts.routed_experts.{e}.w3.weight",
+                        f".block_sparse_moe.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.routed_experts.{e}.w2.weight",
+                        f".block_sparse_moe.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 | None = None,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
     ) -> None:
         """Precompute K/V for context states write them into each layer's KV cache.
 
@@ -551,7 +551,13 @@ 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(
@@ -559,7 +565,7 @@ class DFlashQwen3Model(nn.Module):
                 all_k_final[i],
                 all_v[i],
                 kv_cache,
-                context_slot_mapping,
+                slot_mapping,
             )
 
     def forward(
@@ -701,7 +707,7 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
         self,
         context_states: torch.Tensor,
         context_positions: torch.Tensor,
-        context_slot_mapping: torch.Tensor | None = None,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | 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

@@ -0,0 +1,153 @@
+# 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,6 +585,8 @@ _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,14 +17,23 @@ 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

@@ -0,0 +1,477 @@
+# 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,7 +48,12 @@ 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 MixtureOfExperts, SupportsPP
+from vllm.model_executor.models.interfaces import (
+    EagleModelMixin,
+    MixtureOfExperts,
+    SupportsEagle3,
+    SupportsPP,
+)
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
@@ -933,7 +938,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         return x, residual, post_mix, res_mix
 
 
-class DeepseekV4Model(nn.Module):
+class DeepseekV4Model(nn.Module, EagleModelMixin):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -1074,7 +1079,12 @@ class DeepseekV4Model(nn.Module):
             input_ids = input_ids.to(torch.int64)
 
         residual, post_mix, res_mix = None, None, None
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
+        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,
+        ):
             hidden_states, residual, post_mix, res_mix = layer(
                 hidden_states,
                 positions,
@@ -1083,10 +1093,21 @@ class DeepseekV4Model(nn.Module):
                 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:
-            hidden_states = mhc_post_tilelang(
-                hidden_states, residual, post_mix, res_mix
-            )
+            # 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
+                )
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})
@@ -1104,6 +1125,8 @@ class DeepseekV4Model(nn.Module):
             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]:
@@ -1330,7 +1353,9 @@ class DeepseekV4MixtureOfExperts(MixtureOfExperts):
             moe.experts.update_expert_map()
 
 
-class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts):
+class DeepseekV4ForCausalLM(
+    nn.Module, SupportsPP, SupportsEagle3, DeepseekV4MixtureOfExperts
+):
     model_cls = DeepseekV4Model
 
     # Default mapper assumes the original FP4-expert checkpoint layout.

vllm/utils/extensible_tensor.py

@@ -0,0 +1,520 @@
+# 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

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -306,16 +306,19 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         )
 
         # Handle MTP: adjust decode_threshold like the indexer does
+        spec_config = self.vllm_config.speculative_config
         self.num_speculative_tokens = (
-            self.vllm_config.speculative_config.num_speculative_tokens
-            if self.vllm_config.speculative_config
-            else 0
+            spec_config.num_speculative_tokens if spec_config else 0
+        )
+        # Decode can have query_len up to
+        #   1 + (2 if parallel drafting else 1) * num_speculative_tokens.
+        # This MUST match the flashmla_sparse / indexer threshold so that
+        # all backends agree on the decode/prefill split.
+        spec_mult = (
+            2 if (spec_config is not None and spec_config.parallel_drafting) else 1
         )
-        # With MTP, decode can have query_len up to 1 + num_speculative_tokens.
-        # Must match the threshold used by the indexer and flashmla_sparse so
-        # that all backends agree on the decode/prefill split.
         self.decode_threshold = (
-            self.reorder_batch_threshold + self.num_speculative_tokens
+            self.reorder_batch_threshold + spec_mult * self.num_speculative_tokens
         )
 
         hf_config = self.vllm_config.model_config.hf_config
@@ -368,6 +371,20 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
             device=self.device,
         )
 
+        # DSpark draft: the block is non-causal (every query attends to the
+        # trailing window of context PLUS all query tokens, including future ones),
+        # so its per-token index list is wider than `window_size`. The kernel pads
+        # the q-head count to B_TOPK (64/128), which requires the index width to be
+        # a multiple of 128.
+        self.is_dspark = spec_config is not None and spec_config.use_dspark()
+        self.noncausal_index_width = (
+            cdiv(self.window_size + self.num_speculative_tokens, 128) * 128
+            if self.is_dspark
+            else 0
+        )
+        self.decode_swa_indices_noncausal: torch.Tensor | None = None
+        self._max_tokens = max_tokens
+
     def build(
         self,
         common_prefix_len: int,
@@ -407,23 +424,56 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
         is_valid_token.copy_(slot_mapping >= 0)
 
+        non_causal = not common_attn_metadata.causal
+        decode_swa_indices = self.decode_swa_indices
         if num_decode_tokens > 0:
             self.decode_swa_lens[num_decode_tokens:] = 0
-            _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
-                self.decode_swa_indices,
-                self.decode_swa_indices.stride(0),
-                self.decode_swa_lens,
-                self.window_size,
-                query_start_loc,
-                seq_lens,
-                token_to_req_indices,
-                is_valid_token,
-                block_table,
-                block_table.stride(0),
-                self.block_size,
-                token_offset=0,
-                TRITON_BLOCK_SIZE=1024,
-            )
+            if non_causal:
+                assert self.is_dspark, (
+                    "Non-causal DeepseekV4 SWA is only supported for the DSpark "
+                    "speculation mode, but causal=False was set without DSpark."
+                )
+                if self.decode_swa_indices_noncausal is None:
+                    self.decode_swa_indices_noncausal = torch.zeros(
+                        self._max_tokens,
+                        1,
+                        self.noncausal_index_width,
+                        dtype=torch.int32,
+                        device=self.device,
+                    )
+                decode_swa_indices = self.decode_swa_indices_noncausal
+                _compute_dspark_noncausal_swa_indices_kernel[(num_decode_tokens,)](
+                    decode_swa_indices,
+                    decode_swa_indices.stride(0),
+                    self.decode_swa_lens,
+                    self.window_size,
+                    self.noncausal_index_width,
+                    query_start_loc,
+                    seq_lens,
+                    token_to_req_indices,
+                    is_valid_token,
+                    block_table,
+                    block_table.stride(0),
+                    self.block_size,
+                    token_offset=0,
+                    TRITON_BLOCK_SIZE=1024,
+                )
+            else:
+                _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
+                    decode_swa_indices,
+                    decode_swa_indices.stride(0),
+                    self.decode_swa_lens,
+                    self.window_size,
+                    query_start_loc,
+                    seq_lens,
+                    token_to_req_indices,
+                    is_valid_token,
+                    block_table,
+                    block_table.stride(0),
+                    self.block_size,
+                    token_offset=0,
+                    TRITON_BLOCK_SIZE=1024,
+                )
 
         # Prefill SWA indices live in paged coordinates. `token_offset` lets
         # the kernel read is_valid_token / token_to_req_indices at absolute
@@ -471,7 +521,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
             slot_mapping=slot_mapping,
             is_valid_token=is_valid_token,
             token_to_req_indices=token_to_req_indices,
-            decode_swa_indices=self.decode_swa_indices[:num_decode_tokens],
+            decode_swa_indices=decode_swa_indices[:num_decode_tokens],
             decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
             prefill_swa_indices=(
                 self.prefill_swa_indices[:num_prefill_tokens]
@@ -665,3 +715,69 @@ def _compute_swa_indices_and_lens_kernel(
             slot_ids,
             mask=offset < window_size,
         )
+
+
+# TODO(ben): unify this kernel to reduce duplication
+@triton.jit(do_not_specialize=["token_offset"])
+def _compute_dspark_noncausal_swa_indices_kernel(
+    swa_indices_ptr,
+    swa_indices_stride,
+    swa_lens_ptr,
+    window_size,
+    index_width,
+    query_start_loc_ptr,
+    seq_lens_ptr,
+    token_to_req_indices_ptr,
+    is_valid_token_ptr,
+    block_table_ptr,
+    block_table_stride,
+    block_size,
+    token_offset,
+    TRITON_BLOCK_SIZE: tl.constexpr,
+):
+    """Non-causal per-token indices for the DSpark draft block.
+
+    Here, we populate the topk indices with the trailing window of context tokens,
+    plus all query tokens (including future ones).
+    """
+    pid = tl.program_id(0)
+    token_idx = pid + token_offset
+    is_valid = tl.load(is_valid_token_ptr + token_idx)
+    if not is_valid:
+        tl.store(swa_lens_ptr + pid, 0)
+        return
+
+    req_idx = tl.load(token_to_req_indices_ptr + token_idx)
+
+    query_start = tl.load(query_start_loc_ptr + req_idx)
+    query_end = tl.load(query_start_loc_ptr + req_idx + 1)
+    query_len = query_end - query_start
+
+    seq_len = tl.load(seq_lens_ptr + req_idx)
+    prefix_len = seq_len - query_len
+
+    # Block-anchored window (shared by every token in the block) + full block.
+    start_pos = tl.maximum(prefix_len - window_size, 0)
+    end_pos = seq_len
+
+    swa_len = end_pos - start_pos
+    tl.store(swa_lens_ptr + pid, swa_len)
+
+    for i in range(0, index_width, TRITON_BLOCK_SIZE):
+        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+
+        pos_offset = start_pos + offset
+        block_indices = pos_offset // block_size
+        block_numbers = tl.load(
+            block_table_ptr + req_idx * block_table_stride + block_indices,
+            mask=pos_offset < end_pos,
+        )
+        block_offsets = pos_offset % block_size
+        slot_ids = block_numbers * block_size + block_offsets
+
+        slot_ids = tl.where(offset < swa_len, slot_ids, -1)
+        tl.store(
+            swa_indices_ptr + pid * swa_indices_stride + offset,
+            slot_ids,
+            mask=offset < index_width,
+        )

vllm/v1/attention/ops/chunked_prefill_paged_decode.py

@@ -156,6 +156,11 @@ def kernel_paged_attention_2d(
         # Supports non-contiguous mapping
         # from logical blocks to physical blocks
         abs_token_idx = start_n + offs_n
+        # Slots >= seq_len are unwritten KV cache and may hold NaN/garbage
+        # (e.g. the tail of the last partial block). They are score-masked
+        # below, but 0 * NaN = NaN would still poison the output, so exclude
+        # them from the K/V loads too.
+        kv_load_mask = abs_token_idx < seq_len
         l_block_idx = abs_token_idx // PHYSICAL_BLOCK_SIZE
         # Vectorized loading of physical block IDs
         p_block_idx = tl.load(block_tables_ptr + block_table_offset + l_block_idx)
@@ -181,7 +186,7 @@ def kernel_paged_attention_2d(
         # K : (HEAD_SIZE, BLOCK_SIZE)
         K_load = tl.load(
             key_cache_ptr + k_offset,
-            mask=dim_mask[:, None],
+            mask=dim_mask[:, None] & kv_load_mask[None, :],
             other=0.0,
             eviction_policy="evict_last",
         )
@@ -194,7 +199,7 @@ def kernel_paged_attention_2d(
         # V : (BLOCK_SIZE, HEAD_SIZE)
         V_load = tl.load(
             value_cache_ptr + v_offset,
-            mask=dim_mask[None, :],
+            mask=dim_mask[None, :] & kv_load_mask[:, None],
             other=0.0,
             eviction_policy="evict_last",
         )

vllm/v1/attention/ops/triton_attention_helpers.py

@@ -183,10 +183,12 @@ def compute_tile_loop_bounds(
         + 1
     )
     if USE_MM_PREFIX or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL):
-        # Non-causal or mixed batches need the full sequence range.
-        # Per-element masking in compute_kv_seq_mask handles the
-        # actual causal/non-causal boundary per sequence.
-        max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len)
+        # Read the full sequence but never past seq_len: the causal-style
+        # formula above can overshoot for non-causal sequences, and slots
+        # >= seq_len are unwritten KV (last-block tail) that may hold NaN
+        # (0 * NaN poisons the output). Per-element masking in
+        # compute_kv_seq_mask handles the causal/non-causal boundary.
+        max_seq_prefix_len = seq_len
     else:
         max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
 

vllm/v1/core/encoder_cache_manager.py

@@ -5,6 +5,8 @@ from collections import OrderedDict
 from collections.abc import Mapping
 from typing import TYPE_CHECKING
 
+import torch
+
 from vllm.logger import init_logger
 from vllm.v1.request import Request
 
@@ -78,6 +80,17 @@ class EncoderCacheManager:
         self.freeable: OrderedDict[str, int] = OrderedDict()
         self.freed: list[str] = []
 
+    @staticmethod
+    def make_profiling_reservation(
+        cache_size: int,
+        embed_size: int,
+        dtype: torch.dtype,
+        device: torch.device | str,
+    ) -> torch.Tensor | None:
+        if cache_size <= 0:
+            return None
+        return torch.empty((cache_size, embed_size), dtype=dtype, device=device)
+
     def reset(self) -> None:
         """Reset the encoder cache to its initial state.
 

vllm/v1/core/sched/scheduler.py

@@ -247,6 +247,11 @@ class Scheduler(SchedulerInterface):
                 # decoding instead of standard next-token sampling, so it has a query
                 # for the last sampled token plus queries for each draft token.
                 self.num_lookahead_tokens = self.num_spec_tokens + 1
+            if speculative_config.use_dspark():
+                # DSpark drafts a block of num_spec_tokens query tokens in which the
+                # anchor itself is the first prediction position (no separate bonus
+                # query), so it needs exactly num_spec_tokens lookahead slots.
+                self.num_lookahead_tokens = self.num_spec_tokens
 
         # Create the KV cache manager.
         if hash_block_size is None:

vllm/v1/engine/core.py

@@ -89,6 +89,7 @@ from vllm.version import __version__ as VLLM_VERSION
 logger = init_logger(__name__)
 
 HANDSHAKE_TIMEOUT_MINS = 5
+_CUDAGRAPH_MEMORY_BUFFER_BYTES = 150 * (1 << 20)
 
 _R = TypeVar("_R")  # Return type for collective_rpc
 
@@ -288,37 +289,73 @@ class EngineCore:
 
         assert len(kv_cache_specs) == len(available_gpu_memory)
 
-        # Track max_model_len before KV cache config to detect auto-fit changes
-        max_model_len_before = vllm_config.model_config.max_model_len
+        use_extensible_kv_cache = (
+            has_kv_cache and vllm_config.cache_config.enable_extensible_kv_cache
+        )
+        if use_extensible_kv_cache:
+            from vllm.platforms import current_platform
 
+            if vllm_config.cache_config.kv_cache_memory_bytes is not None:
+                raise ValueError(
+                    "enable_extensible_kv_cache=True is not supported with "
+                    "kv_cache_memory_bytes. The extensible path requires "
+                    "automatic KV cache sizing."
+                )
+            if not current_platform.is_cuda():
+                raise ValueError(
+                    "enable_extensible_kv_cache=True is only supported on CUDA."
+                )
+
+        # Track max_model_len before KV cache config to detect auto-fit changes
+        # made by get_kv_cache_configs().
+        max_model_len_before = vllm_config.model_config.max_model_len
         kv_cache_configs = get_kv_cache_configs(
             vllm_config, kv_cache_specs, available_gpu_memory
         )
+        scheduler_kv_cache_config = self._apply_kv_cache_config(
+            vllm_config,
+            kv_cache_configs,
+            max_model_len_before,
+        )
 
-        # If auto-fit reduced max_model_len, sync the new value to workers.
-        # This is needed because workers were spawned before memory profiling
-        # and have the original (larger) max_model_len cached.
-        max_model_len_after = vllm_config.model_config.max_model_len
-        if max_model_len_after != max_model_len_before:
-            self.collective_rpc("update_max_model_len", args=(max_model_len_after,))
-
-        scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
-        vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
-        kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
-        if kv_cache_groups:
-            vllm_config.cache_config.block_size = min(
-                g.kv_cache_spec.block_size for g in kv_cache_groups
+        # Initialize KV cache and warm up execution. With extensible KV cache,
+        # this reserves the upper-bound address range, commits one block, and
+        # captures CUDA graphs before committing the post-capture KV size.
+        compilation_times = self.model_executor.initialize_from_config(
+            kv_cache_configs,
+            extensible=use_extensible_kv_cache,
+        )
+        if use_extensible_kv_cache:
+            if len(compilation_times) != len(available_gpu_memory):
+                raise RuntimeError(
+                    "Expected one CompilationTimes result per worker when "
+                    "initializing extensible KV cache, but got "
+                    f"{len(compilation_times)} results for "
+                    f"{len(available_gpu_memory)} workers."
+                )
+            final_available_gpu_memory = [
+                max(
+                    available_memory
+                    - times.cuda_graph
+                    - _CUDAGRAPH_MEMORY_BUFFER_BYTES,
+                    0,
+                )
+                for available_memory, times in zip(
+                    available_gpu_memory, compilation_times, strict=True
+                )
+            ]
+            max_model_len_before = vllm_config.model_config.max_model_len
+            kv_cache_configs = get_kv_cache_configs(
+                vllm_config,
+                kv_cache_specs,
+                final_available_gpu_memory,
             )
-            num_tokens, max_concurrency = get_kv_cache_capacity(
-                vllm_config, scheduler_kv_cache_config
+            scheduler_kv_cache_config = self._apply_kv_cache_config(
+                vllm_config,
+                kv_cache_configs,
+                max_model_len_before,
             )
-            vllm_config.cache_config.kv_cache_size_tokens = num_tokens
-            vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency
-
-        vllm_config.validate_block_size()
-
-        # Initialize kv cache and warmup the execution
-        self.model_executor.initialize_from_config(kv_cache_configs)
+            self.model_executor.extend_kv_cache(scheduler_kv_cache_config.num_blocks)
 
         elapsed = time.time() - start
         compile_time = vllm_config.compilation_config.compilation_time
@@ -347,6 +384,35 @@ class EngineCore:
             )
         return scheduler_kv_cache_config
 
+    def _apply_kv_cache_config(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_configs: list[KVCacheConfig],
+        max_model_len_before: int,
+    ) -> KVCacheConfig:
+        # If auto-fit reduced max_model_len, sync the new value to workers.
+        # This is needed because workers were spawned before memory profiling
+        # and have the original (larger) max_model_len cached.
+        max_model_len_after = vllm_config.model_config.max_model_len
+        if max_model_len_after != max_model_len_before:
+            self.collective_rpc("update_max_model_len", args=(max_model_len_after,))
+
+        scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
+        vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
+        kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
+        if kv_cache_groups:
+            vllm_config.cache_config.block_size = min(
+                g.kv_cache_spec.block_size for g in kv_cache_groups
+            )
+            num_tokens, max_concurrency = get_kv_cache_capacity(
+                vllm_config, scheduler_kv_cache_config
+            )
+            vllm_config.cache_config.kv_cache_size_tokens = num_tokens
+            vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency
+
+        vllm_config.validate_block_size()
+        return scheduler_kv_cache_config
+
     def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
         return self.model_executor.supported_tasks
 

vllm/v1/executor/abstract.py

@@ -115,12 +115,20 @@ class Executor(ABC):
     def _init_executor(self) -> None:
         raise NotImplementedError
 
-    def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
+    def initialize_from_config(
+        self,
+        kv_cache_configs: list[KVCacheConfig],
+        extensible: bool = False,
+    ) -> list[CompilationTimes]:
         """
         Initialize the KV caches and begin the model execution loop of the
         underlying workers.
         """
-        self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
+        self.collective_rpc(
+            "initialize_from_config",
+            args=(kv_cache_configs,),
+            kwargs={"extensible": extensible} if extensible else None,
+        )
         compilation_times: list[CompilationTimes] = self.collective_rpc(
             "compile_or_warm_up_model"
         )
@@ -135,6 +143,7 @@ class Executor(ABC):
             self.vllm_config.compilation_config.encoder_compilation_time = max(
                 t.encoder for t in compilation_times
             )
+        return compilation_times
 
     def register_failure_callback(self, callback: FailureCallback):  # noqa: B027
         """
@@ -149,6 +158,9 @@ class Executor(ABC):
     def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
         return self.collective_rpc("get_kv_cache_spec")
 
+    def extend_kv_cache(self, num_blocks: int) -> None:
+        self.collective_rpc("extend_kv_cache", args=(num_blocks,))
+
     @overload
     def collective_rpc(
         self,

vllm/v1/worker/gpu/model_runner.py

@@ -193,7 +193,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             if self.is_last_pp_rank:
                 self.speculator = init_speculator(self.vllm_config, self.device)
 
-            if self.speculative_config.method in ("eagle3", "dflash"):
+            if self.speculative_config.method in ("eagle3", "dflash", "dspark"):
                 # Drafting may require auxiliary hidden states from target model outputs
                 self.use_aux_hidden_state_outputs = True
                 if self.use_pp:

vllm/v1/worker/gpu/sample/gumbel.py

@@ -223,6 +223,11 @@ def gumbel_sample(
     output_processed_logits_col: torch.Tensor | None = None,
     use_fp64: bool = False,
 ) -> torch.Tensor:
+    # Enforce contiguity on non-strided input tensors
+    expanded_idx_mapping = expanded_idx_mapping.contiguous()
+    pos = pos.contiguous()
+    if output_processed_logits_col is not None:
+        output_processed_logits_col = output_processed_logits_col.contiguous()
     num_tokens, vocab_size = logits.shape
     BLOCK_SIZE = 1024
     num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)

vllm/v1/worker/gpu/spec_decode/__init__.py

@@ -14,6 +14,12 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):
         )
 
         return DFlashSpeculator(vllm_config, device)
+    elif speculative_config.method == "dspark":
+        from vllm.v1.worker.gpu.spec_decode.dspark.speculator import (
+            DSparkSpeculator,
+        )
+
+        return DSparkSpeculator(vllm_config, device)
     elif speculative_config.use_gemma4_mtp():
         from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import (
             Gemma4Speculator,

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -29,6 +29,8 @@ logger = init_logger(__name__)
 
 
 class DFlashSpeculator(DraftModelSpeculator):
+    _speculator_name = "DFlash"  # For logging, so we can share methods with subclasses
+
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
         super().__init__(vllm_config, device)
 
@@ -48,15 +50,17 @@ class DFlashSpeculator(DraftModelSpeculator):
 
         self.dflash_causal = get_dflash_causal(self.draft_model_config)
 
-        # Buffers for context K/V precomputation. Populated by prepare_dflash_inputs,
-        # and processed by the model's precompute_and_store_context_kv method.
-        # NOT captured by CUDA graphs.
+        # Whether the anchor query position is itself a prediction. DFlash default uses
+        # the anchor as the bonus token (only mask tokens predict); DSpark samples from
+        # the anchor and the N-1 mask token positions. See _prepare_dflash_inputs_kernel
+        self.sample_from_anchor = False
+
+        # Context positions for the K/V precompute. Populated by
+        # prepare_dflash_inputs, and processed by the model's
+        # precompute_and_store_context_kv method. NOT captured by CUDA graphs.
         self.context_positions = torch.zeros(
             self.max_num_tokens, dtype=torch.int64, device=device
         )
-        self.context_slot_mapping = torch.zeros(
-            self.max_num_tokens, dtype=torch.int64, device=device
-        )
 
         # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens).
         max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps
@@ -94,7 +98,7 @@ class DFlashSpeculator(DraftModelSpeculator):
         )
 
     def capture(self, attn_states: dict | None = None) -> None:
-        logger.info("Capturing model for DFlash speculator...")
+        logger.info("Capturing model for %s speculator...", self._speculator_name)
         # Reset sampling indices to zero to prevent stale values from prior
         # dummy runs from being baked into the captured graph.
         self.sample_indices.zero_()
@@ -108,7 +112,7 @@ class DFlashSpeculator(DraftModelSpeculator):
             self.attn_groups,
             self.kv_cache_config,
             self.max_model_len,
-            progress_bar_desc="Capturing dflash CUDA graphs",
+            progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs",
         )
 
     def load_draft_model(
@@ -126,18 +130,39 @@ class DFlashSpeculator(DraftModelSpeculator):
     ) -> None:
         super().set_attn(model_state, kv_cache_config, block_tables)
 
-        # DFlash precomputes context K/V with a single block_size; mixing
-        # kv-cache groups would silently corrupt the cache for the non-matching group.
-        draft_groups = [gid for gid, g in enumerate(self.attn_groups) if g]
-        assert len(draft_groups) == 1, (
-            "DFlash currently requires all draft attention layers to share "
-            "a single kv-cache group."
-        )
-        self.draft_kv_cache_group_id = draft_groups[0]
+        self.draft_kv_cache_group_ids = [
+            gid for gid, g in enumerate(self.attn_groups) if g
+        ]
+        assert self.draft_kv_cache_group_ids, "No draft attention groups found."
+        self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0]
         self.draft_block_size = self.block_tables.block_sizes[
             self.draft_kv_cache_group_id
         ]
 
+        # Per-group context slot buffers for the precompute (one row per group).
+        self._context_slot_mappings = torch.zeros(
+            len(self.draft_kv_cache_group_ids),
+            self.max_num_tokens,
+            dtype=torch.int64,
+            device=self.device,
+        )
+
+        # Map each draft decoder layer to the index (within draft_kv_cache_group_ids)
+        # of the kv-cache group its cache belongs to. Models that share a single group
+        # leave this as None and share one context slot mapping.
+        self._layer_group_idx: list[int] | None = None
+        if hasattr(self.model, "get_draft_kv_cache_layer_names"):
+            name_to_gid = {
+                ln: gid
+                for gid, group in enumerate(kv_cache_config.kv_cache_groups)
+                for ln in group.layer_names
+            }
+            gid_to_idx = {gid: i for i, gid in enumerate(self.draft_kv_cache_group_ids)}
+            self._layer_group_idx = [
+                gid_to_idx[name_to_gid[name]]
+                for name in self.model.get_draft_kv_cache_layer_names()
+            ]
+
     @torch.inference_mode()
     def _run_model(
         self,
@@ -183,9 +208,11 @@ class DFlashSpeculator(DraftModelSpeculator):
 
         num_sample = num_reqs * self.num_speculative_steps
         sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]]
+        # sample_pos is the predicted token's position Q; verification keys
+        # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2.
         draft_tokens = self.sample_draft(
             sample_hidden_states,
-            self.sample_pos[:num_sample],
+            self.sample_pos[:num_sample] - 2,
             self.sample_idx_mapping[:num_sample],
             self.temperature,
             self.seeds,
@@ -294,40 +321,49 @@ class DFlashSpeculator(DraftModelSpeculator):
         # The query slot mapping is written into the shared BlockTables slot_mappings.
         # That buffer's address is what the captured CUDA graph reads from at replay.
         assert self.draft_kv_cache_group_id >= 0
-        query_slot_mapping = self.block_tables.slot_mappings[
-            self.draft_kv_cache_group_id
-        ]
-        prepare_dflash_inputs(
-            self.input_buffers,
-            query_slot_mapping,
-            self.context_positions,
-            self.context_slot_mapping,
-            self.sample_indices,
-            self.sample_pos,
-            self.sample_idx_mapping,
-            input_batch,
-            num_sampled,
-            num_rejected,
-            last_sampled,
-            next_prefill_tokens,
-            self.block_tables.input_block_tables[self.draft_kv_cache_group_id],
-            self.draft_block_size,
-            self.parallel_drafting_token_id,
-            self.num_query_per_req,
-            self.num_speculative_steps,
-            self.max_num_reqs,
-            self.max_num_tokens,
-        )
+        # Support multiple draft KV cache groups by preparing inputs once for each
+        for i, gid in enumerate(self.draft_kv_cache_group_ids):
+            prepare_dflash_inputs(
+                self.input_buffers,
+                self.block_tables.slot_mappings[gid],
+                self.context_positions,
+                self._context_slot_mappings[i],
+                self.sample_indices,
+                self.sample_pos,
+                self.sample_idx_mapping,
+                input_batch,
+                num_sampled,
+                num_rejected,
+                last_sampled,
+                next_prefill_tokens,
+                self.block_tables.input_block_tables[gid],
+                self.block_tables.block_sizes[gid],
+                self.parallel_drafting_token_id,
+                self.num_query_per_req,
+                self.num_speculative_steps,
+                self.max_num_reqs,
+                self.max_num_tokens,
+                self.max_model_len,
+                self.sample_from_anchor,
+            )
 
         # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph
         # because the context shape varies per step. During dummy runs the block tables
         # are placeholders, so we skip the cache write to avoid clobbering real entries.
+        # Each layer uses the context slots of its own kv-cache group.
+        if dummy_run:
+            context_slots: torch.Tensor | list[torch.Tensor | None] | None = None
+        elif self._layer_group_idx is not None:
+            context_slots = [
+                self._context_slot_mappings[gidx][:num_target_tokens]
+                for gidx in self._layer_group_idx
+            ]
+        else:
+            context_slots = self._context_slot_mappings[0][:num_target_tokens]
         self.model.precompute_and_store_context_kv(
             self.hidden_states[:num_target_tokens],
             self.context_positions[:num_target_tokens],
-            context_slot_mapping=(
-                None if dummy_run else self.context_slot_mapping[:num_target_tokens]
-            ),
+            context_slots,
         )
 
         # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs
@@ -408,6 +444,8 @@ def _prepare_dflash_inputs_kernel(
     num_speculative_steps,
     max_num_reqs,
     max_num_tokens,
+    max_model_len,
+    SAMPLE_FROM_ANCHOR: tl.constexpr,
     PAD_SLOT_ID: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
@@ -468,14 +506,21 @@ def _prepare_dflash_inputs_kernel(
     q_slot = q_block_id * block_size + (query_pos % block_size)
 
     tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
-    tl.store(out_query_positions_ptr + query_idx, query_pos, mask=is_query)
+    clamped_query_pos = tl.minimum(query_pos, max_model_len - 1)
+    tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query)
     tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
 
-    # --- Sample indices / positions / idx_mapping (mask tokens only) ---
-    is_sample = is_query & (query_off > 0)
-    sample_idx = req_idx * num_speculative_steps + (query_off - 1)
+    # --- Sample indices / positions / idx_mapping ---
+    # When SAMPLE_FROM_ANCHOR (DSpark), so we sample at EVERY query position
+    # and each position k predicts the NEXT token (sampled position = query_pos + 1).
+    # Otherwise (DFlash default) the anchor is the bonus token and only the mask tokens
+    # at offsets > 0 are sampled from, each AT its own position.
+    sample_off = 0 if SAMPLE_FROM_ANCHOR else 1
+    is_sample = is_query & (query_off >= sample_off)
+    sample_idx = req_idx * num_speculative_steps + (query_off - sample_off)
+    sample_pos = query_pos + 1 if SAMPLE_FROM_ANCHOR else query_pos
     tl.store(out_sample_indices_ptr + sample_idx, query_idx, mask=is_sample)
-    tl.store(out_sample_pos_ptr + sample_idx, query_pos, mask=is_sample)
+    tl.store(out_sample_pos_ptr + sample_idx, sample_pos, mask=is_sample)
     tl.store(out_sample_idx_mapping_ptr + sample_idx, req_state_idx, mask=is_sample)
 
     if block_idx == 0:
@@ -542,6 +587,8 @@ def prepare_dflash_inputs(
     num_speculative_steps: int,
     max_num_reqs: int,
     max_num_tokens: int,
+    max_model_len: int,
+    sample_from_anchor: bool = False,
 ) -> None:
     num_reqs = input_batch.num_reqs
     assert num_reqs > 0
@@ -577,6 +624,8 @@ def prepare_dflash_inputs(
         num_speculative_steps,
         max_num_reqs,
         max_num_tokens,
+        max_model_len,
+        SAMPLE_FROM_ANCHOR=sample_from_anchor,
         PAD_SLOT_ID=PAD_SLOT_ID,
         BLOCK_SIZE=BLOCK_SIZE,
     )

vllm/v1/worker/gpu/spec_decode/dspark/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/v1/worker/gpu/spec_decode/dspark/speculator.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DSpark speculator: semi-autoregressive parallel drafting.
+
+DSpark drafts a block of ``num_speculative_tokens`` tokens in one parallel pass
+(reusing the DFlash machinery: context-KV precompute + a query-block forward),
+then injects intra-block dependency with a lightweight sequential Markov head.
+
+Differences from DFlash:
+  * Anchor-as-first-prediction: each request emits exactly ``N =
+    num_speculative_tokens`` query tokens (anchor + N-1 noise), NOT ``1 + N``.
+    Every query position is a prediction (the anchor predicts the first draft
+    token), so we sample at all N positions and ``sample_pos = query_pos + 1``
+    (standard next-token), whereas DFlash's masks sit AT the predicted position.
+    This is the ``sample_from_anchor`` path in the shared prepare-inputs kernel.
+  * Sequential Markov sampling: instead of DFlash's single parallel sample, we
+    sample left-to-right, adding a prefix-dependent Markov bias derived from the
+    previously sampled token at each step.
+
+CUDA graphs (FULL, mirroring DFlash) cover the whole draft step: the parallel
+backbone forward AND the sequential Markov sampling.
+"""
+
+from typing import Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model
+
+
+class DSparkSpeculator(DFlashSpeculator):
+    _speculator_name = "DSpark"
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        super().__init__(vllm_config, device)
+
+        # Anchor-first: N query tokens per request (anchor + N-1 noise), not 1+N.
+        self.num_query_per_req = self.num_speculative_steps
+
+        # DSpark consumes mean-pooled target aux hidden states at the target
+        # layers, combined to hidden_size via main_proj. Store that combined
+        # main_x (hidden_size wide). DSpark does not use the same pre-allocated buffer
+        # that DeepSeek-V4's MTP uses.
+        draft_hidden = self.draft_model_config.get_hidden_size()
+        self.hidden_states = torch.zeros(
+            self.max_num_tokens, draft_hidden, dtype=self.dtype, device=device
+        )
+
+        self.dflash_causal = False
+
+        # The anchor query position is itself a prediction (see module docstring).
+        self.sample_from_anchor = True
+
+        self._step_cols = torch.arange(
+            self.num_speculative_steps, dtype=torch.int32, device=device
+        )
+
+        self._anchor_idx = (
+            torch.arange(self.max_num_reqs, dtype=torch.int64, device=device)
+            * self.num_query_per_req
+        )
+
+    def load_draft_model(
+        self,
+        target_model: torch.nn.Module,
+        target_attn_layer_names: set[str],
+    ) -> torch.nn.Module:
+        return load_dspark_model(target_model, self.vllm_config)
+
+    def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
+        # Sequential Markov sampling over the backbone's output hidden states.
+        n_spec = self.num_speculative_steps
+        num_sample = num_reqs * n_spec
+        # Per-(req, position) head hidden, ordered (req, step).
+        sample_hidden = head_hidden[self.sample_indices[:num_sample]]
+        base_logits = self.model.compute_logits(sample_hidden)
+        vocab_size = base_logits.shape[-1]
+        base_logits = base_logits.view(num_reqs, n_spec, vocab_size)
+
+        idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)
+        sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec)
+
+        # Anchor (bonus) token per request = the input id at query offset 0,
+        # read via the precomputed persistent index (fixed buffer for capture).
+        prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]]
+
+        for i in range(n_spec):
+            # Sequential stage: Markov bias from the previously sampled token.
+            markov_embed = self.model.markov_embed(prev)
+            bias = self.model.markov_bias(markov_embed)
+            logits_i = base_logits[:, i] + bias
+            if self.draft_logits is not None:
+                # sample_pos is the predicted token's position Q; the target
+                # verifies it with the predecessor's Gumbel key (Q-1). Pass Q-1.
+                draft_i = gumbel_sample(
+                    logits_i,
+                    idx_map[:, i],
+                    self.temperature,
+                    self.seeds,
+                    sample_pos[:, i] - 1,
+                    apply_temperature=True,
+                    output_processed_logits=self.draft_logits,
+                    output_processed_logits_col=self._step_cols[i],
+                    use_fp64=self.use_fp64_gumbel,
+                )
+            else:
+                draft_i = logits_i.argmax(dim=-1)
+            self.draft_tokens[:num_reqs, i] = draft_i
+            prev = draft_i
+
+    def _generate_draft(
+        self,
+        num_reqs: int,
+        num_tokens_padded: int,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        num_tokens_across_dp: torch.Tensor | None,
+        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+    ) -> None:
+        # Full draft step (captured under CUDA graph): parallel backbone forward
+        # then sequential Markov sampling over its hidden state outputs.
+        head_hidden = self._run_model(
+            num_tokens_padded,
+            attn_metadata,
+            slot_mappings,
+            num_tokens_across_dp,
+            cudagraph_runtime_mode,
+        )
+        self._sample_sequential(num_reqs, head_hidden)

vllm/v1/worker/gpu/spec_decode/dspark/utils.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch.nn as nn
+
+from vllm.config import VllmConfig, replace
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.model_executor.model_loader import get_model
+from vllm.v1.worker.gpu.spec_decode.eagle.utils import _should_share
+
+
+def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
+    speculative_config = vllm_config.speculative_config
+    assert speculative_config is not None
+    draft_model_config = speculative_config.draft_model_config
+
+    from vllm.compilation.backends import set_model_tag
+
+    # DSpark uses non-causal attention.
+    causal = False
+    draft_vllm_config = replace(
+        vllm_config,
+        attention_config=replace(
+            vllm_config.attention_config,
+            use_non_causal=not causal,
+            backend=speculative_config.attention_backend,
+        ),
+    )
+
+    with set_model_tag("dspark_head"):
+        draft_model = get_model(
+            vllm_config=draft_vllm_config, model_config=draft_model_config
+        )
+
+    if get_pp_group().world_size != 1:
+        raise NotImplementedError("DSpark does not support pipeline parallelism.")
+
+    target_language_model = (
+        target_model.get_language_model()
+        if hasattr(target_model, "get_language_model")
+        else target_model
+    )
+    target_inner = target_language_model.model
+    draft_inner = draft_model.model
+
+    target_embed = getattr(target_inner, "embed_tokens", None)
+    draft_embed = getattr(draft_inner, "embed_tokens", None)
+    if target_embed is not None and _should_share(
+        draft_model, "has_own_embed_tokens", draft_embed, target_embed
+    ):
+        if draft_embed is not None:
+            del draft_inner.embed_tokens
+        draft_inner.embed_tokens = target_embed
+
+    target_lm_head = getattr(target_model, "lm_head", None)
+    draft_lm_head = getattr(draft_model, "lm_head", None)
+    if target_lm_head is not None and _should_share(
+        draft_model, "has_own_lm_head", draft_lm_head, target_lm_head
+    ):
+        if draft_lm_head is not None:
+            del draft_model.lm_head
+        draft_model.lm_head = target_lm_head
+
+    return draft_model

vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py

@@ -44,6 +44,15 @@ def get_eagle3_aux_layers_from_config(
         if dflash_config and isinstance(dflash_config, dict):
             # Add 1 to convert DFlash's aux layer id semantics
             layer_ids = [i + 1 for i in (dflash_config.get("target_layer_ids") or [])]
+    if not layer_ids:
+        dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None)
+        if dspark_layer_ids:
+            layer_ids = [i + 1 for i in dspark_layer_ids]
+    if not layer_ids:
+        # Dense DSpark (e.g. Qwen3) also uses different aux layer semantics.
+        target_layer_ids = getattr(hf_config, "target_layer_ids", None)
+        if target_layer_ids:
+            layer_ids = [i + 1 for i in target_layer_ids]
     if layer_ids and isinstance(layer_ids, (list, tuple)):
         return tuple(layer_ids)
     return None

vllm/v1/worker/gpu/spec_decode/utils.py

@@ -54,17 +54,23 @@ class DraftTokensHandler:
 def get_parallel_drafting_token_id(hf_config) -> int:
     """Resolve the mask token id used for parallel drafting slots.
 
-    Checks (in order): `dflash_config.mask_token_id`, `pard_token`,
-    `ptd_token_id`. Raises ValueError if none are present.
+    Checks (in order): `dflash_config.mask_token_id`, top-level `mask_token_id`,
+    `dspark_noise_token_id`, `pard_token`, `ptd_token_id`. Raises ValueError if
+    none are present.
     """
     dflash_config = getattr(hf_config, "dflash_config", None) or {}
     if "mask_token_id" in dflash_config:
         return int(dflash_config["mask_token_id"])
+    if getattr(hf_config, "mask_token_id", None) is not None:
+        return int(hf_config.mask_token_id)
+    if hasattr(hf_config, "dspark_noise_token_id"):
+        return int(hf_config.dspark_noise_token_id)
     if hasattr(hf_config, "pard_token"):
         return int(hf_config.pard_token)
     if hasattr(hf_config, "ptd_token_id"):
         return int(hf_config.ptd_token_id)
     raise ValueError(
         "Model config must specify `dflash_config.mask_token_id`,"
-        " `pard_token`, or `ptd_token_id` for parallel drafting."
+        " `mask_token_id`, `dspark_noise_token_id`, `pard_token`, or"
+        " `ptd_token_id` for parallel drafting."
     )