[ROCm][CI] Fix Kimi K3 KDA on ROCm (#50262 )

Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
[Spec Decode][Perf] Replicate DSpark Markov head across TP ranks (#49731 )
2026-07-29 15:47:50 +00:00 · 2026-07-29 11:43:51 -04:00 · 2026-07-29 06:41:54 -07:00 · 2026-07-29 06:30:45 -07:00 · 2026-07-29 13:16:13 +00:00 · 2026-07-29 12:21:09 +00:00
47 changed files with 1996 additions and 376 deletions
@@ -1576,10 +1576,14 @@ steps:
  - vllm/third_party/flash_linear_attention/ops/kda.py
  - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
  - vllm/third_party/flash_linear_attention/ops/l2norm.py
-  - tests/kernels/test_kda.py
+  - vllm/models/kimi_k3/nvidia/kda.py
+  - vllm/models/kimi_k3/nvidia/kda_metadata.py
+  - vllm/models/kimi_k3/nvidia/ops/third_party/kda/
+  - tests/models/kimi_k3/test_kda.py
+  - tests/models/kimi_k3/test_kda_metadata.py
  - vllm/platforms/rocm.py
  commands:
-  - pytest -v -s kernels/test_kda.py
+  - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py

 - label: Kernels Mamba Test # TBD
  timeout_in_minutes: 180
@@ -222,7 +222,9 @@ mod tests {
        assert_eq!(tensor.shape, vec![expected.len()]);
        assert_eq!(
            tensor.data,
-            WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect())
+            WireArrayData::RawView(
+                expected.iter().map(|value| u8::from(*value)).collect::<Vec<_>>().into(),
+            )
        );
    }

@@ -2,28 +2,58 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use std::collections::HashMap;
+use std::mem::size_of;

 use half::{bf16, f16};
 use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs};
 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};
+use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireArrayData, WireTensor};

 use crate::error::{Error, Result, bail_multimodal, multimodal};

+/// Element type retained alongside an encoded tensor during multimodal lowering.
+#[derive(Debug, Clone, Copy)]
+pub(super) enum TensorKind {
+    /// 32-bit floating point.
+    F32,
+    /// IEEE 16-bit floating point.
+    F16,
+    /// Brain floating point.
+    Bf16,
+    /// Signed 64-bit integer.
+    I64,
+    /// Unsigned 32-bit integer.
+    U32,
+}
+
+impl TensorKind {
+    const fn element_size(self) -> usize {
+        match self {
+            Self::F32 => size_of::<f32>(),
+            Self::F16 => size_of::<f16>(),
+            Self::Bf16 => size_of::<bf16>(),
+            Self::I64 => size_of::<i64>(),
+            Self::U32 => size_of::<u32>(),
+        }
+    }
+
+    const fn wire_dtype(self) -> &'static str {
+        match self {
+            Self::F32 => "float32",
+            Self::F16 => "float16",
+            Self::Bf16 => "bfloat16",
+            Self::I64 => "int64",
+            Self::U32 => "uint32",
+        }
+    }
+}
+
 /// Representation for multimodal kwarg values for transformation.
 #[derive(Debug)]
 pub(super) enum KwargValue {
-    /// Float tensor with row-major flat data and shape.
-    F32Tensor { data: Vec<f32>, shape: Vec<usize> },
-    /// Float16 tensor with row-major flat data and shape.
-    F16Tensor { data: Vec<f16>, shape: Vec<usize> },
-    /// BFloat16 tensor with row-major flat data and shape.
-    Bf16Tensor { data: Vec<bf16>, shape: Vec<usize> },
-    /// Signed integer tensor with row-major flat data and shape.
-    I64Tensor { data: Vec<i64>, shape: Vec<usize> },
-    /// Unsigned integer tensor with row-major flat data and shape.
-    U32Tensor { data: Vec<u32>, shape: Vec<usize> },
+    /// Tensor with row-major flat data and shape.
+    Tensor { kind: TensorKind, wire: WireTensor },
    /// Non-tensor kwarg value that is shared or copied as-is.
    Passthrough(ProtocolKwargValue),
 }
@@ -67,8 +97,14 @@ impl KwargValue {
            ModelSpecificValue::Tensor { data, shape } => {
                Self::from_f32_tensor(data, shape, float_dtype)?
            }
-            ModelSpecificValue::IntTensor { data, shape } => Self::I64Tensor { data, shape },
-            ModelSpecificValue::UintTensor { data, shape } => Self::U32Tensor { data, shape },
+            ModelSpecificValue::IntTensor { data, shape } => {
+                let wire = WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?;
+                Self::tensor(TensorKind::I64, wire)
+            }
+            ModelSpecificValue::UintTensor { data, shape } => {
+                let wire = WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?;
+                Self::tensor(TensorKind::U32, wire)
+            }
            ModelSpecificValue::Int(value) => Self::Passthrough(Int(value)),
            ModelSpecificValue::Float(value) => Self::Passthrough(Float(value)),
            ModelSpecificValue::IntVec(values) => {
@@ -93,17 +129,23 @@ impl KwargValue {
    /// Convert a float tensor to the target float dtype if needed, keeping the
    /// same shape.
    fn from_f32_tensor(data: Vec<f32>, shape: Vec<usize>, float_dtype: ModelDtype) -> Result<Self> {
-        match float_dtype {
-            ModelDtype::Float16 => Ok(Self::F16Tensor {
-                data: data.into_iter().map(f16::from_f32).collect(),
-                shape,
-            }),
-            ModelDtype::BFloat16 => Ok(Self::Bf16Tensor {
-                data: data.into_iter().map(bf16::from_f32).collect(),
-                shape,
-            }),
-            ModelDtype::Float32 => Ok(Self::F32Tensor { data, shape }),
-        }
+        let (kind, wire) = match float_dtype {
+            ModelDtype::Float16 => (
+                TensorKind::F16,
+                WireTensor::from_f16(shape, data.into_iter().map(f16::from_f32).collect()),
+            ),
+            ModelDtype::BFloat16 => (
+                TensorKind::Bf16,
+                WireTensor::from_bf16(shape, data.into_iter().map(bf16::from_f32).collect()),
+            ),
+            ModelDtype::Float32 => (TensorKind::F32, WireTensor::from_f32(shape, data)),
+        };
+        wire.map(|wire| Self::tensor(kind, wire)).map_err(Error::Multimodal)
+    }
+
+    fn tensor(kind: TensorKind, wire: WireTensor) -> Self {
+        debug_assert_eq!(wire.dtype, kind.wire_dtype());
+        Self::Tensor { kind, wire }
    }
 }

@@ -111,15 +153,11 @@ impl TryFrom<&KwargValue> for ProtocolKwargValue {
    type Error = Error;

    fn try_from(value: &KwargValue) -> Result<Self> {
-        let tensor = match value {
-            KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
-            KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
-            KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
-            KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
-            KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
+        let wire = match value {
+            KwargValue::Tensor { wire, .. } => wire.clone(),
            KwargValue::Passthrough(value) => return Ok(value.clone()),
        };
-        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
+        Ok(ProtocolKwargValue::Tensor(wire))
    }
 }

@@ -127,11 +165,7 @@ impl KwargValue {
    /// First-axis length for tensor values; `None` for passthrough kwargs.
    pub(super) fn first_dim(&self) -> Option<usize> {
        match self {
-            Self::F32Tensor { shape, .. }
-            | Self::F16Tensor { shape, .. }
-            | Self::Bf16Tensor { shape, .. }
-            | Self::I64Tensor { shape, .. }
-            | Self::U32Tensor { shape, .. } => shape.first().copied(),
+            Self::Tensor { wire, .. } => wire.shape.first().copied(),
            Self::Passthrough(_) => None,
        }
    }
@@ -161,30 +195,13 @@ impl KwargValue {
        end: usize,
        drop_axis: bool,
    ) -> Result<ProtocolKwargValue> {
-        let tensor = match self {
-            Self::F32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_f32(shape, data)
-            }
-            Self::F16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_f16(shape, data)
-            }
-            Self::Bf16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_bf16(shape, data)
-            }
-            Self::I64Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_i64(shape, data)
-            }
-            Self::U32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_u32(shape, data)
+        let wire = match self {
+            Self::Tensor { kind, wire } => {
+                slice_first_axis_range(wire, kind.element_size(), start, end, drop_axis)
            }
            Self::Passthrough(value) => return Ok(value.clone()),
        };
-        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
+        wire.map(ProtocolKwargValue::Tensor)
    }
 }

@@ -208,44 +225,53 @@ pub(super) fn flat_range_for_index(
 /// Read a tensor value as per-image sizes for flat slicing.
 fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
    match tensor {
-        KwargValue::I64Tensor { data, .. } => data
-            .iter()
+        KwargValue::Tensor {
+            kind: TensorKind::I64,
+            wire,
+        } => raw_tensor_bytes(wire, size_of::<i64>())?
+            .chunks_exact(size_of::<i64>())
+            .map(|bytes| i64::from_ne_bytes(bytes.try_into().expect("exact int64 chunk")))
            .map(|value| {
-                usize::try_from(*value)
+                usize::try_from(value)
                    .map_err(|_| multimodal!("negative flat tensor size `{value}`"))
            })
            .collect(),
-        KwargValue::U32Tensor { data, .. } => {
-            Ok(data.iter().map(|value| *value as usize).collect())
-        }
+        KwargValue::Tensor {
+            kind: TensorKind::U32,
+            wire,
+        } => Ok(raw_tensor_bytes(wire, size_of::<u32>())?
+            .chunks_exact(size_of::<u32>())
+            .map(|bytes| u32::from_ne_bytes(bytes.try_into().expect("exact uint32 chunk")) as usize)
+            .collect()),
        _ => Err(multimodal!("flat tensor sizes must be int64 or uint32")),
    }
 }

 /// Slice a flat row-major tensor along its first axis.
-fn slice_first_axis_range<'a, T>(
-    shape: &[usize],
-    data: &'a [T],
+fn slice_first_axis_range(
+    tensor: &WireTensor,
+    element_size: usize,
    start: usize,
    end: usize,
    drop_axis: bool,
-) -> Result<(Vec<usize>, &'a [T])> {
+) -> Result<WireTensor> {
+    let shape = tensor.shape.as_slice();
+    raw_tensor_bytes(tensor, element_size)?;
    let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
    if start > end || end > first_dim {
        bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
    }
-    let expected_len = shape
-        .checked_numel()
-        .ok_or_else(|| multimodal!("tensor shape {shape:?} has too many elements"))?;
-    if expected_len != data.len() {
-        bail_multimodal!(
-            "tensor shape {shape:?} expects {expected_len} elements, got {}",
-            data.len()
-        );
-    }
-    let stride = shape[1..].iter().product::<usize>();
-    let data_start = start * stride;
-    let data_end = end * stride;
+    let stride = shape[1..]
+        .iter()
+        .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))
+        .and_then(|stride| stride.checked_mul(element_size))
+        .ok_or_else(|| multimodal!("tensor shape {shape:?} byte stride overflowed usize"))?;
+    let data_start = start
+        .checked_mul(stride)
+        .ok_or_else(|| multimodal!("tensor slice start byte offset overflowed usize"))?;
+    let data_end = end
+        .checked_mul(stride)
+        .ok_or_else(|| multimodal!("tensor slice end byte offset overflowed usize"))?;
    let out_shape = if drop_axis {
        shape[1..].to_vec()
    } else {
@@ -253,7 +279,38 @@ fn slice_first_axis_range<'a, T>(
        shape[0] = end - start;
        shape
    };
-    Ok((out_shape, &data[data_start..data_end]))
+    let WireArrayData::RawView(data) = &tensor.data else {
+        return Err(multimodal!("cannot slice an aux tensor buffer"));
+    };
+    Ok(WireTensor::from_raw_bytes(
+        tensor.dtype.clone(),
+        out_shape,
+        data.slice(data_start..data_end),
+    ))
+}
+
+fn raw_tensor_bytes(tensor: &WireTensor, element_size: usize) -> Result<&[u8]> {
+    let WireArrayData::RawView(data) = &tensor.data else {
+        return Err(multimodal!("expected an inline tensor buffer"));
+    };
+    let expected_bytes = tensor
+        .shape
+        .checked_numel()
+        .and_then(|numel| numel.checked_mul(element_size))
+        .ok_or_else(|| {
+            multimodal!(
+                "tensor shape {:?} byte length overflowed usize",
+                tensor.shape
+            )
+        })?;
+    if expected_bytes != data.len() {
+        bail_multimodal!(
+            "tensor shape {:?} expects {expected_bytes} bytes, got {}",
+            tensor.shape,
+            data.len()
+        );
+    }
+    Ok(data)
 }

 #[cfg(test)]
@@ -262,28 +319,32 @@ mod tests {

    #[test]
    fn batched_wire_value_at_drops_first_axis() {
-        let value = KwargValue::F32Tensor {
-            data: vec![1.0, 2.0, 3.0, 4.0],
-            shape: vec![2, 2],
-        };
+        let data = vec![1.0_f32, 2.0, 3.0, 4.0];
+        let expected_ptr = data.as_ptr().cast::<u8>().wrapping_add(2 * size_of::<f32>());
+        let value = KwargValue::tensor(
+            TensorKind::F32,
+            WireTensor::from_f32(vec![2, 2], data).unwrap(),
+        );

        let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
            panic!("expected tensor");
        };

        assert_eq!(tensor.shape, vec![2]);
+        let raw_view = tensor.data.into_raw_view().unwrap();
+        assert_eq!(raw_view.as_ptr(), expected_ptr);
        assert_eq!(
-            tensor.data.into_raw_view().unwrap(),
+            raw_view,
            [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
        );
    }

    #[test]
    fn flat_wire_value_range_keeps_first_axis() {
-        let value = KwargValue::U32Tensor {
-            data: (0..10).collect(),
-            shape: vec![5, 2],
-        };
+        let value = KwargValue::tensor(
+            TensorKind::U32,
+            WireTensor::from_u32(vec![5, 2], (0..10_u32).collect()).unwrap(),
+        );

        let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
            panic!("expected tensor");
@@ -298,10 +359,10 @@ mod tests {

    #[test]
    fn flat_range_for_index_uses_size_tensor() {
-        let sizes = KwargValue::I64Tensor {
-            data: vec![2, 3, 4],
-            shape: vec![3],
-        };
+        let sizes = KwargValue::tensor(
+            TensorKind::I64,
+            WireTensor::from_i64(vec![3], vec![2_i64, 3, 4]).unwrap(),
+        );

        assert_eq!(
            flat_range_for_index(&sizes, "image_grid_thw", 1).unwrap(),
@@ -311,10 +372,11 @@ mod tests {

    #[test]
    fn slice_first_axis_range_errors_on_shape_data_mismatch() {
-        let error = slice_first_axis_range(&[2, 2], &[1.0_f32, 2.0, 3.0], 0, 1, true).unwrap_err();
+        let tensor = WireTensor::from_raw("float32", vec![2, 2], vec![0; 3 * size_of::<f32>()]);
+        let error = slice_first_axis_range(&tensor, size_of::<f32>(), 0, 1, true).unwrap_err();

        assert!(
-            matches!(error, Error::Multimodal(message) if message.contains("expects 4 elements"))
+            matches!(error, Error::Multimodal(message) if message.contains("expects 16 bytes"))
        );
    }

@@ -519,7 +519,7 @@ impl EngineCoreClient {
                "registered request to engine"
            );

-            self.inner.send_to_engine(&engine_id, EngineCoreRequestType::Add, &req).await?;
+            self.inner.send_request_to_engine(&engine_id, req).await?;
            Ok(())
        }
        .await;
@@ -6,6 +6,7 @@ use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};

 use arc_swap::ArcSwapOption;
+use bytes::Bytes;
 use parking_lot::Mutex;
 use thiserror_ext::AsReport as _;
 use tokio::runtime::Handle;
@@ -21,17 +22,30 @@ use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_outpu
 use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
 use crate::protocol::encode_msgpack;
 use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
-use crate::protocol::request::EngineCoreRequestType;
+use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
 use crate::protocol::stats::SchedulerStats;
 use crate::protocol::utility::UtilityOutput;
 use crate::transport::{ConnectedEngine, EngineId};
 use crate::{Error, Result, transport};

+const MSGPACK_ZERO_COPY_THRESHOLD_ENV: &str = "VLLM_MSGPACK_ZERO_COPY_THRESHOLD";
+const DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD: usize = 256;
+
+fn msgpack_zero_copy_threshold() -> usize {
+    std::env::var(MSGPACK_ZERO_COPY_THRESHOLD_ENV)
+        .ok()
+        .and_then(|value| value.parse().ok())
+        .unwrap_or(DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD)
+}
+
 pub(crate) struct ClientInner {
    input_send: RouterSendHalf,
    /// The runtime handle used for sending messages to the engine.
    handle: Handle,
    model_name: String,
+    /// Per-tensor byte threshold loaded from env variable
+    /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created.
+    msgpack_zero_copy_threshold: usize,
    scheduler_stats_recorder: SchedulerStatsRecorder,
    request_reg: Mutex<RequestRegistry>,
    utility_reg: Mutex<UtilityRegistry>,
@@ -54,6 +68,7 @@ impl ClientInner {
            input_send,
            handle,
            model_name,
+            msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(),
            scheduler_stats_recorder,
            request_reg: Mutex::new(RequestRegistry::new(engines)),
            utility_reg: Mutex::new(UtilityRegistry::default()),
@@ -229,9 +244,29 @@ impl ClientInner {
    where
        T: serde::Serialize + std::fmt::Debug,
    {
-        // TODO: for `EngineCoreRequest`, split outbound tensor raw views into aux
-        // frames instead of always producing a single msgpack frame.
-        let payload = encode_msgpack(payload)?;
+        let payload = Bytes::from(encode_msgpack(payload)?);
+        self.send_encoded_to_engine(engine_id, request_type, payload, Vec::new()).await
+    }
+
+    /// Send an add request, moving large tensor buffers into auxiliary frames.
+    pub async fn send_request_to_engine(
+        &self,
+        engine_id: &EngineId,
+        mut payload: EngineCoreRequest,
+    ) -> Result<()> {
+        let aux_frames = payload.extract_aux_frames(self.msgpack_zero_copy_threshold);
+        let payload = Bytes::from(encode_msgpack(&payload)?);
+        self.send_encoded_to_engine(engine_id, EngineCoreRequestType::Add, payload, aux_frames)
+            .await
+    }
+
+    async fn send_encoded_to_engine(
+        &self,
+        engine_id: &EngineId,
+        request_type: EngineCoreRequestType,
+        payload: Bytes,
+        aux_frames: Vec<Bytes>,
+    ) -> Result<()> {
        let mut input_send = self.input_send.clone();
        let engine_id = engine_id.clone();

@@ -242,6 +277,7 @@ impl ClientInner {
                    &engine_id,
                    request_type.to_frame(),
                    payload,
+                    aux_frames,
                )
                .await
            })
@@ -209,17 +209,17 @@ impl WireLogprobs {
            logprob_token_ids: WireNdArray {
                dtype: "<i8".to_string(),
                shape: vec![rows, cols],
-                data: WireArrayData::RawView(token_ids),
+                data: WireArrayData::RawView(token_ids.into()),
            },
            logprobs: WireNdArray {
                dtype: "<f4".to_string(),
                shape: vec![rows, cols],
-                data: WireArrayData::RawView(logprobs),
+                data: WireArrayData::RawView(logprobs.into()),
            },
            token_ranks: WireNdArray {
                dtype: "<i8".to_string(),
                shape: vec![rows],
-                data: WireArrayData::RawView(token_ranks),
+                data: WireArrayData::RawView(token_ranks.into()),
            },
            cu_num_generated_tokens: None,
        })
@@ -4,6 +4,7 @@
 use std::io::Cursor;

 use byteorder::{BigEndian, LittleEndian, NativeEndian, ReadBytesExt};
+use bytes::Bytes;
 use itertools::Itertools as _;

 use crate::error::{Error, Result, ext_value_decode};
@@ -126,7 +127,7 @@ pub(super) fn decode_array_metadata<Frame>(
    field: &str,
    frames: &[Frame],
    expected_scalars: &[ScalarType],
-) -> Result<(Vec<usize>, Vec<u8>, ScalarType, Endianness)>
+) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)>
 where
    Frame: AsRef<[u8]>,
 {
@@ -171,7 +172,7 @@ pub(super) fn resolve_array_bytes<Frame>(
    value: WireArrayData,
    field: &str,
    frames: &[Frame],
-) -> Result<Vec<u8>>
+) -> Result<Bytes>
 where
    Frame: AsRef<[u8]>,
 {
@@ -187,7 +188,7 @@ where
                    ),
                )
            })?;
-            Ok(frame.as_ref().to_vec())
+            Ok(Bytes::copy_from_slice(frame.as_ref()))
        }
    }
 }
@@ -3,6 +3,7 @@

 use std::collections::BTreeMap;

+use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use serde_tuple::{Deserialize_tuple, Serialize_tuple};

@@ -105,6 +106,44 @@ pub enum MmKwargValue {
    List(Vec<MmKwargValue>),
 }

+impl MmFeatureSpec {
+    /// Extract large tensor buffers from this feature in serialized field order.
+    pub(crate) fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        if let Some(data) = &mut self.data {
+            for elem in data.values_mut() {
+                elem.extract_aux_frames(aux_frames, threshold);
+            }
+        }
+        if let Some(is_embed) = &mut self.mm_position.is_embed {
+            is_embed.extract_aux_frame(aux_frames, threshold);
+        }
+    }
+}
+
+impl MmFieldElem {
+    /// Extract large tensor buffers from this field element.
+    fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        if let Some(data) = &mut self.data {
+            data.extract_aux_frames(aux_frames, threshold);
+        }
+    }
+}
+
+impl MmKwargValue {
+    /// Recursively extract large tensor buffers from this nested value.
+    fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        match self {
+            Self::Tensor(tensor) => tensor.extract_aux_frame(aux_frames, threshold),
+            Self::List(values) => {
+                for value in values {
+                    value.extract_aux_frames(aux_frames, threshold);
+                }
+            }
+            Self::Int(_) | Self::Float(_) => {}
+        }
+    }
+}
+
 /// Defines how to interpret tensor data belonging to a keyword argument for
 /// `MultiModalKwargsItems`, and vice versa.
 ///
@@ -137,6 +137,17 @@ impl EngineCoreRequest {
        }
        Ok(())
    }
+
+    /// Extract large request tensors into ordered auxiliary frames.
+    pub(crate) fn extract_aux_frames(&mut self, threshold: usize) -> Vec<Bytes> {
+        let mut aux_frames = Vec::new();
+        if let Some(features) = &mut self.mm_features {
+            for feature in features {
+                feature.extract_aux_frames(&mut aux_frames, threshold);
+            }
+        }
+        aux_frames
+    }
 }

 #[cfg(test)]
@@ -144,9 +155,15 @@ mod tests {
    use rmpv::Value;

    use super::*;
+    use crate::protocol::multimodal::{
+        MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, PlaceholderRange,
+    };
    use crate::protocol::sampling::EngineCoreSamplingParams;
+    use crate::protocol::tensor::{WireArrayData, WireTensor};
    use crate::protocol::{decode_value, encode_msgpack};

+    const AUX_FRAME_THRESHOLD: usize = 256;
+
    #[test]
    fn engine_core_request_serializes_as_full_array() {
        let request = EngineCoreRequest {
@@ -175,4 +192,84 @@ mod tests {
        assert_eq!(array[10], Value::Nil);
        assert_eq!(array[11], Value::from(7));
    }
+
+    #[test]
+    fn engine_core_request_extracts_large_nested_tensors_in_wire_order() {
+        let inline = vec![1_u8; AUX_FRAME_THRESHOLD - 1];
+        let first_aux = vec![2_u8; AUX_FRAME_THRESHOLD];
+        let second_aux = vec![3_u8; AUX_FRAME_THRESHOLD + 1];
+        let first_aux_ptr = first_aux.as_ptr();
+        let second_aux_ptr = second_aux.as_ptr();
+        let mut request = EngineCoreRequest {
+            mm_features: Some(vec![MmFeatureSpec {
+                data: Some(BTreeMap::from([
+                    (
+                        "inline".to_string(),
+                        MmFieldElem {
+                            data: Some(MmKwargValue::Tensor(WireTensor::from_raw(
+                                "uint8",
+                                vec![inline.len()],
+                                inline,
+                            ))),
+                            field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+                        },
+                    ),
+                    (
+                        "nested".to_string(),
+                        MmFieldElem {
+                            data: Some(MmKwargValue::List(vec![
+                                MmKwargValue::Int(7),
+                                MmKwargValue::Tensor(WireTensor::from_raw(
+                                    "uint8",
+                                    vec![first_aux.len()],
+                                    first_aux,
+                                )),
+                            ])),
+                            field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+                        },
+                    ),
+                ])),
+                modality: "image".to_string(),
+                identifier: "id".to_string(),
+                mm_position: PlaceholderRange {
+                    offset: 0,
+                    length: second_aux.len(),
+                    is_embed: Some(WireTensor::from_raw(
+                        "bool",
+                        vec![second_aux.len()],
+                        second_aux,
+                    )),
+                },
+                mm_hash: None,
+            }]),
+            ..EngineCoreRequest::default()
+        };
+
+        let aux_frames = request.extract_aux_frames(AUX_FRAME_THRESHOLD);
+
+        assert_eq!(aux_frames.len(), 2);
+        assert_eq!(aux_frames[0].as_ptr(), first_aux_ptr);
+        assert_eq!(aux_frames[1].as_ptr(), second_aux_ptr);
+        let feature = &request.mm_features.as_ref().unwrap()[0];
+        let MmKwargValue::Tensor(inline) =
+            feature.data.as_ref().unwrap()["inline"].data.as_ref().unwrap()
+        else {
+            panic!("expected inline tensor");
+        };
+        assert!(matches!(inline.data, WireArrayData::RawView(_)));
+        let MmKwargValue::List(nested) =
+            feature.data.as_ref().unwrap()["nested"].data.as_ref().unwrap()
+        else {
+            panic!("expected nested tensor list");
+        };
+        let MmKwargValue::Tensor(nested_tensor) = &nested[1] else {
+            panic!("expected nested tensor");
+        };
+        assert_eq!(nested_tensor.data, WireArrayData::AuxIndex(1));
+        assert_eq!(
+            feature.mm_position.is_embed.as_ref().unwrap().data,
+            WireArrayData::AuxIndex(2)
+        );
+        assert!(request.extract_aux_frames(AUX_FRAME_THRESHOLD).is_empty());
+    }
 }
@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-use bytemuck::allocation::pod_collect_to_vec;
+use bytemuck::{Pod, cast_slice};
+use bytes::Bytes;
 use enum_as_inner::EnumAsInner;
 use half::{bf16, f16};
 use rmpv::Value;
@@ -20,6 +21,21 @@ struct MsgpackExtRef<'a>((i8, ByteSlice<'a>));

 struct ByteSlice<'a>(&'a [u8]);

+struct PodVec<T: Pod>(Vec<T>);
+
+impl<T: Pod> AsRef<[u8]> for PodVec<T> {
+    fn as_ref(&self) -> &[u8] {
+        cast_slice(&self.0)
+    }
+}
+
+fn bytes_from_pod_vec<T>(data: Vec<T>) -> Bytes
+where
+    T: Pod + Send + 'static,
+{
+    Bytes::from_owner(PodVec(data))
+}
+
 impl Serialize for ByteSlice<'_> {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
@@ -55,58 +71,63 @@ pub struct WireNdArray {

 impl WireNdArray {
    /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "float32".to_string(),
+        Ok(Self::from_raw_bytes(
+            "float32",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "float16".to_string(),
+        Ok(Self::from_raw_bytes(
+            "float16",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "bfloat16".to_string(),
+        Ok(Self::from_raw_bytes(
+            "bfloat16",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "int64".to_string(),
+        Ok(Self::from_raw_bytes(
+            "int64",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "uint32".to_string(),
+        Ok(Self::from_raw_bytes(
+            "uint32",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build a bool tensor/ndarray backed by raw-view bytes.
@@ -118,7 +139,9 @@ impl WireNdArray {
        Ok(Self {
            dtype: "bool".to_string(),
            shape,
-            data: WireArrayData::RawView(data.iter().map(|value| u8::from(*value)).collect()),
+            data: WireArrayData::RawView(Bytes::from(
+                data.into_iter().map(u8::from).collect::<Vec<_>>(),
+            )),
        })
    }

@@ -127,12 +150,22 @@ impl WireNdArray {
    /// Use this as an escape hatch when the caller already owns bytes that
    /// match the requested `dtype` and `shape`.
    pub fn from_raw(dtype: impl Into<String>, shape: Vec<usize>, data: Vec<u8>) -> Self {
+        Self::from_raw_bytes(dtype, shape, Bytes::from(data))
+    }
+
+    /// Build a tensor/ndarray from an owned immutable raw-view buffer.
+    pub fn from_raw_bytes(dtype: impl Into<String>, shape: Vec<usize>, data: Bytes) -> Self {
        Self {
            dtype: dtype.into(),
            shape,
            data: WireArrayData::RawView(data),
        }
    }
+
+    /// Move a sufficiently large inline buffer into the ordered auxiliary-frame list.
+    pub(crate) fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        self.data.extract_aux_frame(aux_frames, threshold);
+    }
 }

 /// Validate that the number of elements implied by the shape matches the length
@@ -170,7 +203,25 @@ pub enum WireArrayData {
    /// stored.
    AuxIndex(usize),
    /// The raw bytes of this array/tensor.
-    RawView(Vec<u8>),
+    RawView(Bytes),
+}
+
+impl WireArrayData {
+    /// Replace a sufficiently large raw view with its one-based auxiliary-frame index.
+    fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        let Self::RawView(bytes) = self else {
+            return;
+        };
+        if bytes.len() < threshold {
+            return;
+        }
+
+        let index = aux_frames.len() + 1;
+        let bytes = std::mem::replace(self, Self::AuxIndex(index))
+            .into_raw_view()
+            .expect("raw view was matched above");
+        aux_frames.push(bytes);
+    }
 }

 impl<'de> Deserialize<'de> for WireArrayData {
@@ -180,7 +231,9 @@ impl<'de> Deserialize<'de> for WireArrayData {
    {
        let value = Value::deserialize(deserializer)?;
        match value {
-            Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => Ok(Self::RawView(bytes)),
+            Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => {
+                Ok(Self::RawView(Bytes::from(bytes)))
+            }
            Value::Ext(tag, _) => Err(serde::de::Error::custom(format!(
                "unsupported extension type code {tag}"
            ))),
@@ -201,9 +254,6 @@ impl Serialize for WireArrayData {
    where
        S: Serializer,
    {
-        // TODO: outbound request serialization currently only supports inline
-        // raw-view bytes. Emitting aux frames needs transport-level plumbing;
-        // serializing `AuxIndex` here only preserves an already-built reference.
        match self {
            Self::AuxIndex(index) => serializer.serialize_u64(*index as u64),
            Self::RawView(bytes) => {
@@ -221,7 +271,7 @@ mod tests {
    fn raw_view_serializes_as_msgpack_ext() {
        let bytes = vec![1, 2, 3, 4];
        let encoded =
-            rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode");
+            rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone().into())).expect("encode");
        let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()))
            .expect("encode expected");

@@ -234,11 +284,15 @@ mod tests {

    #[test]
    fn constructors_build_raw_view_tensors() {
-        let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap();
+        let f32_data = vec![1.0, 2.5];
+        let f32_data_ptr = f32_data.as_ptr().cast::<u8>();
+        let f32_tensor = WireNdArray::from_f32(vec![2], f32_data).unwrap();
        assert_eq!(f32_tensor.dtype, "float32");
        assert_eq!(f32_tensor.shape, vec![2]);
+        let f32_raw_view = f32_tensor.data.into_raw_view().expect("raw view");
+        assert_eq!(f32_raw_view.as_ptr(), f32_data_ptr);
        assert_eq!(
-            f32_tensor.data.into_raw_view().expect("raw view"),
+            f32_raw_view,
            [1.0_f32, 2.5].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
        );

@@ -258,15 +312,15 @@ mod tests {
        let i64_tensor = WireNdArray::from_i64(vec![1], vec![-7]).unwrap();
        assert_eq!(i64_tensor.dtype, "int64");
        assert_eq!(
-            i64_tensor.data.into_raw_view().expect("raw view"),
-            (-7_i64).to_ne_bytes()
+            i64_tensor.data.into_raw_view().expect("raw view").as_ref(),
+            (-7_i64).to_ne_bytes().as_ref()
        );

        let u32_tensor = WireNdArray::from_u32(vec![1], vec![42]).unwrap();
        assert_eq!(u32_tensor.dtype, "uint32");
        assert_eq!(
-            u32_tensor.data.into_raw_view().expect("raw view"),
-            42_u32.to_ne_bytes()
+            u32_tensor.data.into_raw_view().expect("raw view").as_ref(),
+            42_u32.to_ne_bytes().as_ref()
        );

        let bool_tensor = WireNdArray::from_bool(vec![2], vec![false, true]).unwrap();
@@ -32,7 +32,7 @@ use crate::protocol::output::{
 use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
 use crate::protocol::sampling::EngineCoreSamplingParams;
 use crate::protocol::stats::SchedulerStats;
-use crate::protocol::tensor::WireTensor;
+use crate::protocol::tensor::{WireArrayData, WireTensor};
 use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
 use crate::test_utils::{
    IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
@@ -1728,6 +1728,86 @@ async fn client_decodes_multipart_logprob_outputs() {
    client.shutdown().await.unwrap();
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn client_sends_large_multimodal_tensor_as_aux_frame() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+    let engine_id = b"engine-multimodal-aux".to_vec();
+    let tensor_data = (0..64).map(|value| value as f32).collect::<Vec<_>>();
+    let expected_bytes =
+        tensor_data.iter().flat_map(|value| value.to_ne_bytes()).collect::<Vec<_>>();
+    let mut request = sample_multimodal_request();
+    request.mm_features.as_mut().unwrap()[0]
+        .data
+        .as_mut()
+        .unwrap()
+        .get_mut("pixel_values")
+        .unwrap()
+        .data = Some(MmKwargValue::Tensor(
+        WireTensor::from_f32(vec![64], tensor_data).unwrap(),
+    ));
+
+    let (shutdown_tx, engine_task) = spawn_mock_engine_task(
+        handshake_address.clone(),
+        engine_id.clone(),
+        move |dealer, push| {
+            Box::pin(async move {
+                let add = recv_engine_message(dealer).await;
+                assert_eq!(add.len(), 3);
+                assert_eq!(add[0].as_ref(), &[0x00]);
+                let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
+                let MmKwargValue::Tensor(tensor) =
+                    request.mm_features.as_ref().unwrap()[0].data.as_ref().unwrap()["pixel_values"]
+                        .data
+                        .as_ref()
+                        .unwrap()
+                else {
+                    panic!("expected tensor");
+                };
+                assert_eq!(tensor.data, WireArrayData::AuxIndex(1));
+                assert_eq!(add[2].as_ref(), expected_bytes);
+
+                send_outputs(
+                    push,
+                    RequestBatchOutputs {
+                        outputs: vec![request_output(
+                            "req-mm",
+                            vec![],
+                            Some(EngineCoreFinishReason::Length),
+                        )],
+                        finished_requests: Some(BTreeSet::from(["req-mm".to_string()])),
+                        ..Default::default()
+                    }
+                    .into(),
+                )
+                .await;
+            })
+        },
+    );
+
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+
+    let outputs = client.call(request).await.unwrap().collect::<Vec<_>>().await;
+    assert_eq!(outputs.len(), 1);
+    assert!(outputs[0].is_ok());
+
+    let _ = shutdown_tx.send(());
+    engine_task.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
    init_tracing();
@@ -512,14 +512,14 @@ pub async fn send_message(
    input_send: &mut RouterSendHalf,
    engine_id: &EngineId,
    request_type: Bytes,
-    payload: Vec<u8>,
+    payload: Bytes,
+    aux_frames: Vec<Bytes>,
 ) -> Result<()> {
-    let message = ZmqMessage::try_from(vec![
-        engine_id.to_frame(),
-        request_type,
-        Bytes::from(payload),
-    ])
-    .expect("router messages must contain identity and payload");
+    let mut frames = Vec::with_capacity(3 + aux_frames.len());
+    frames.extend([engine_id.to_frame(), request_type, payload]);
+    frames.extend(aux_frames);
+    let message =
+        ZmqMessage::try_from(frames).expect("router messages must contain identity and payload");

    trace!(
        ?engine_id,
@@ -2,9 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Equivalence test for ``precopy_mamba_align_fused_kernel``.

-The V2 "align" pre-copy must migrate mamba state across block boundaries with
-byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` /
-``get_temporal_copy_spec``):
+The fused "align" pre-copy must migrate mamba state across block boundaries
+with byte-identical semantics to the scalar V1 copy specs
+(``get_conv_copy_spec`` / ``get_temporal_copy_spec``):

 * conv state (SD layout, conv_width > 0): shift the sliding window by
  ``token_bias`` tokens -- ``state[bt[src_col], token_bias:]`` ->
@@ -14,20 +14,27 @@ byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` /
  ``state[bt[dst_col]]``.

 The kernel must also no-op when ``src_col < 0`` (fresh request) or
-``src_col == dst_col`` (no boundary crossed).
+``src_col == dst_col`` (no boundary crossed). V2 callers pass an explicit
+``idx_mapping``; V1 align preprocessing launches in batch order with
+``idx_mapping=None``.
 """

 from __future__ import annotations

+from types import SimpleNamespace
+
+import numpy as np
 import torch

+from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils
 from vllm.platforms import current_platform
+from vllm.v1.worker import mamba_utils as worker_mamba_utils
 from vllm.v1.worker.mamba_utils import precopy_mamba_align_fused_kernel

 try:
    import pytest

-    pytestmark = pytest.mark.skipif(
+    _cuda_required = pytest.mark.skipif(
        not current_platform.is_cuda(),
        reason="precopy_mamba_align_fused_kernel needs CUDA/Triton",
    )
@@ -35,6 +42,9 @@ try:
 except ModuleNotFoundError:  # allow running directly as ``python <thisfile>``
    pytest = None

+    def _cuda_required(fn):
+        return fn
+
    def _parametrize(_name, _values):
        def _deco(fn):
            return fn
@@ -49,13 +59,20 @@ SSM_SHAPE = (4, 16, 16)
 MAX_COLS = 8


-def _build_state(num_blocks, device):
+def _build_state(num_blocks, device, conv_state_dim_first):
    """Per-layer (conv SD [nb, width, dim] bf16, ssm [nb, *shape] fp32) pools."""
    convs, ssms = [], []
    for _ in range(NUM_LAYERS):
+        conv_shape = (
+            (num_blocks, CONV_DIM, CONV_WIDTH)
+            if conv_state_dim_first
+            else (num_blocks, CONV_WIDTH, CONV_DIM)
+        )
        convs.append(
            torch.randn(
-                num_blocks, CONV_WIDTH, CONV_DIM, dtype=torch.bfloat16, device=device
+                *conv_shape,
+                dtype=torch.bfloat16,
+                device=device,
            )
        )
        ssms.append(
@@ -64,7 +81,7 @@ def _build_state(num_blocks, device):
    return convs, ssms


-def _build_meta(convs, ssms, device):
+def _build_meta(convs, ssms, device, conv_state_dim_first):
    """Flattened per-(layer, state-type) metadata, ordered conv, ssm per layer."""
    n = NUM_LAYERS * 2
    base = torch.zeros(n, dtype=torch.int64, device=device)
@@ -82,8 +99,14 @@ def _build_meta(convs, ssms, device):
        base[i] = conv.data_ptr()
        blk_stride[i] = conv.stride(0) * conv.element_size()
        elem[i] = conv.element_size()
-        width[i] = conv.size(1)
-        inner[i] = conv.stride(1)
+        if conv_state_dim_first:
+            width[i] = conv.size(2)
+            inner[i] = 1
+            drc[i] = conv.size(1)
+            drs[i] = conv.stride(1) * conv.element_size()
+        else:
+            width[i] = conv.size(1)
+            inner[i] = conv.stride(1)
        i += 1
        # ssm (temporal): width = 0, inner = elems per block
        base[i] = ssm.data_ptr()
@@ -95,7 +118,7 @@ def _build_meta(convs, ssms, device):
    return base, blk_stride, elem, inner, width, group, drc, drs


-def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs):
+def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs, conv_dim_first):
    """Apply the V1 copy semantics on clones, reading from the pre-copy state."""
    conv_pre = [c.clone() for c in convs]
    ssm_pre = [s.clone() for s in ssms]
@@ -108,14 +131,24 @@ def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs):
        sblk, dblk = int(bt[r, sc]), int(bt[r, dc])
        tblk = int(bt[r, sc + tb])  # temporal src column shifted by bias
        for layer in range(NUM_LAYERS):
-            conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:]
+            if conv_dim_first:
+                conv_ref[layer][dblk, :, : CONV_WIDTH - tb] = conv_pre[layer][
+                    sblk, :, tb:
+                ]
+            else:
+                conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:]
            ssm_ref[layer][dblk] = ssm_pre[layer][tblk]
    return conv_ref, ssm_ref


+@_parametrize("conv_state_dim_first", [False, True])
@_parametrize("num_reqs", [1, 4, 16])
@_parametrize("token_bias", [0, 1, 2])
-def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
+@_parametrize("has_idx_mapping", [True, False])
+@_cuda_required
+def test_precopy_matches_v1_copy_specs(
+    num_reqs, token_bias, has_idx_mapping, conv_state_dim_first
+):
    device = torch.device("cuda")
    torch.manual_seed(0)
    # Distinct physical block per (req, col) so copies never alias.
@@ -136,13 +169,20 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
    if num_reqs >= 2:
        dst_col[1] = 1  # src_col == dst_col -> no copy

-    convs, ssms = _build_state(num_blocks, device)
+    convs, ssms = _build_state(num_blocks, device, conv_state_dim_first)
    conv_ref, ssm_ref = _reference(
-        convs, ssms, bt.cpu(), src_col.cpu(), dst_col.cpu(), bias.cpu(), num_reqs
+        convs,
+        ssms,
+        bt.cpu(),
+        src_col.cpu(),
+        dst_col.cpu(),
+        bias.cpu(),
+        num_reqs,
+        conv_state_dim_first,
    )

    base, blk_stride, elem, inner, width, group, drc, drs = _build_meta(
-        convs, ssms, device
+        convs, ssms, device, conv_state_dim_first
    )
    bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device)
    idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
@@ -161,10 +201,11 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
        group,
        drc,
        drs,
-        idx_mapping,
+        idx_mapping if has_idx_mapping else None,
        num_reqs,
        COPY_BLOCK_SIZE=1024,
-        CONV_STATE_DIM_FIRST=False,
+        CONV_STATE_DIM_FIRST=conv_state_dim_first,
+        HAS_IDX_MAPPING=has_idx_mapping,
    )
    torch.accelerator.synchronize()

@@ -173,8 +214,204 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
        torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0)


+def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch):
+    monkeypatch.setattr(
+        layer_mamba_utils,
+        "is_conv_state_dim_first",
+        lambda: True,
+    )
+    state = torch.empty((2, CONV_DIM, CONV_WIDTH), dtype=torch.bfloat16)
+
+    with pytest.raises(AssertionError, match="num_accepted_tokens > 1"):
+        layer_mamba_utils.get_conv_copy_spec(
+            state=state,
+            block_ids=[0, 1],
+            cur_block_idx=0,
+            num_accepted_tokens=3,
+        )
+
+
+class _FakeCpuGpuBuffer:
+    def __init__(self, n):
+        self.np = np.zeros(n, dtype=np.int32)
+        self.gpu = object()
+        self.copy_sizes = []
+
+    def copy_to_gpu(self, n=None):
+        self.copy_sizes.append(n)
+        return self.gpu
+
+
+class _FakePrecopyContext:
+    def __init__(self, n):
+        self.is_initialized = True
+        self.mamba_group_ids = [0]
+        self.mamba_state_idx_buf = _FakeCpuGpuBuffer(n)
+        self.precopy_src_col_buf = _FakeCpuGpuBuffer(n)
+        self.precopy_token_bias_buf = _FakeCpuGpuBuffer(n)
+        self.calls = []
+
+    def initialize_from_forward_context(self, *args, **kwargs):
+        raise AssertionError("test context is pre-initialized")
+
+    def run_fused_precopy(
+        self,
+        *,
+        num_reqs,
+        state_idx_gpu,
+        src_col_gpu,
+        token_bias_gpu,
+        idx_mapping,
+    ):
+        self.calls.append(
+            {
+                "num_reqs": num_reqs,
+                "state_idx": self.mamba_state_idx_buf.np[:num_reqs].copy(),
+                "src_col": self.precopy_src_col_buf.np[:num_reqs].copy(),
+                "token_bias": self.precopy_token_bias_buf.np[:num_reqs].copy(),
+                "idx_mapping": idx_mapping,
+            }
+        )
+
+
+def _make_preprocess_case(token_bias):
+    req_ids = ["fresh", "same", "cross_a", "cross_b"]
+    scheduler_output = SimpleNamespace(
+        finished_req_ids=set(),
+        preempted_req_ids=set(),
+        scheduled_cached_reqs=SimpleNamespace(resumed_req_ids=set()),
+        num_scheduled_tokens={
+            "fresh": 1,
+            "same": 1,
+            "cross_a": 1,
+            "cross_b": 2,
+        },
+    )
+    input_batch = SimpleNamespace(
+        req_ids=req_ids,
+        num_accepted_tokens_cpu=np.array(
+            [token_bias + 1, token_bias + 1, token_bias + 1, 2],
+            dtype=np.int32,
+        ),
+    )
+    requests = {
+        "fresh": SimpleNamespace(req_id="fresh", num_computed_tokens=0),
+        "same": SimpleNamespace(req_id="same", num_computed_tokens=5),
+        "cross_a": SimpleNamespace(req_id="cross_a", num_computed_tokens=8),
+        "cross_b": SimpleNamespace(req_id="cross_b", num_computed_tokens=7),
+    }
+    mamba_state_idx = {"same": 1, "cross_a": 0, "cross_b": 1}
+    return scheduler_output, input_batch, requests, mamba_state_idx
+
+
+@_parametrize("token_bias", [1, 2])
+def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bias):
+    block_size = 4
+    mamba_spec = SimpleNamespace(block_size=block_size, num_speculative_blocks=1)
+    cache_config = SimpleNamespace(enable_prefix_caching=True)
+    kv_cache_config = SimpleNamespace()
+    scalar_copy_calls = []
+
+    def fake_collect(
+        copy_bufs,
+        kv_cache_config,
+        mamba_state_copy_funcs,
+        mamba_group_ids,
+        src_block_idx,
+        dest_block_idx,
+        accept_token_bias,
+        req_state,
+        forward_context,
+    ):
+        scalar_copy_calls.append(
+            (
+                req_state.req_id,
+                src_block_idx,
+                dest_block_idx,
+                accept_token_bias,
+            )
+        )
+
+    monkeypatch.setattr(worker_mamba_utils, "collect_mamba_copy_meta", fake_collect)
+    monkeypatch.setattr(
+        worker_mamba_utils, "do_mamba_copy_block", lambda copy_bufs: None
+    )
+
+    scalar_case = _make_preprocess_case(token_bias)
+    fused_case = _make_preprocess_case(token_bias)
+
+    scalar_copy_bufs = SimpleNamespace(
+        mamba_group_ids=[0],
+        mamba_spec=mamba_spec,
+        offset=0,
+    )
+    worker_mamba_utils.preprocess_mamba(
+        scheduler_output=scalar_case[0],
+        kv_cache_config=kv_cache_config,
+        cache_config=cache_config,
+        mamba_state_idx=scalar_case[3],
+        input_batch=scalar_case[1],
+        requests=scalar_case[2],
+        forward_context={},
+        mamba_state_copy_funcs=(),
+        copy_bufs=scalar_copy_bufs,
+    )
+
+    ctx = _FakePrecopyContext(len(fused_case[1].req_ids))
+    fused_copy_bufs = SimpleNamespace(
+        mamba_group_ids=[0],
+        mamba_spec=mamba_spec,
+        offset=0,
+    )
+    worker_mamba_utils.preprocess_mamba(
+        scheduler_output=fused_case[0],
+        kv_cache_config=kv_cache_config,
+        cache_config=cache_config,
+        mamba_state_idx=fused_case[3],
+        input_batch=fused_case[1],
+        requests=fused_case[2],
+        forward_context={},
+        mamba_state_copy_funcs=(),
+        copy_bufs=fused_copy_bufs,
+        align_ctx=ctx,
+    )
+
+    assert fused_case[3] == scalar_case[3]
+    np.testing.assert_array_equal(
+        fused_case[1].num_accepted_tokens_cpu,
+        scalar_case[1].num_accepted_tokens_cpu,
+    )
+    assert scalar_copy_calls == [
+        ("cross_a", 0, 2, token_bias),
+        ("cross_b", 1, 2, 1),
+    ]
+    assert len(ctx.calls) == 1
+    call = ctx.calls[0]
+    assert call["num_reqs"] == len(fused_case[1].req_ids)
+    assert call["idx_mapping"] is None
+    np.testing.assert_array_equal(call["state_idx"], np.array([0, 1, 2, 2]))
+    np.testing.assert_array_equal(call["src_col"], np.array([-1, -1, 0, 1]))
+    np.testing.assert_array_equal(call["token_bias"], np.array([0, 0, token_bias, 1]))
+    fused_copy_calls = [
+        (req_id, int(src), int(dst), int(bias))
+        for req_id, src, dst, bias in zip(
+            fused_case[1].req_ids,
+            call["src_col"],
+            call["state_idx"],
+            call["token_bias"],
+        )
+        if int(src) != -1 and int(src) != int(dst)
+    ]
+    assert fused_copy_calls == scalar_copy_calls
+
+
 if __name__ == "__main__":
    for nr in (1, 4, 16):
        for tb in (0, 1, 2):
-            test_precopy_matches_v1_copy_specs(nr, tb)
-            print(f"OK num_reqs={nr} token_bias={tb}")
+            for mapping in (True, False):
+                for dim_first in (False, True):
+                    test_precopy_matches_v1_copy_specs(nr, tb, mapping, dim_first)
+                    print(
+                        f"OK num_reqs={nr} token_bias={tb} "
+                        f"has_idx_mapping={mapping} conv_dim_first={dim_first}"
+                    )
@@ -15,7 +15,9 @@ from vllm.config import PoolerConfig
    ["Qwen/Qwen3-Embedding-0.6B"],
 )
@torch.inference_mode
-def test_embed_models(hf_runner, vllm_runner, model: str):
+def test_embed_models(hf_runner, vllm_runner, monkeypatch, model: str):
+    # Keep token_embed on MRV1 when sequence pooling becomes MRV2 by default.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
    chunk_size = 10
    n_prompt_tokens = [55, 56, 57]
    token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
@@ -104,3 +106,44 @@ def test_last_pool_score_chunked_prefill_matches_unchunked(vllm_runner, model: s
    assert chunked == pytest.approx(unchunked, abs=5e-2), (
        f"chunked score {chunked} diverged from unchunked {unchunked}"
    )
+
+
+@torch.inference_mode
+def test_sequence_embed_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "Qwen/Qwen3-Embedding-0.6B"
+    chunk_size = 10
+    token_prompts = [[1024 + i for i in range(n)] for n in (25, 27)]
+    prompts = [TokensPrompt(prompt_token_ids=t) for t in token_prompts]
+
+    with hf_runner(model, auto_cls=AutoModel) as hf_model:
+        hf_outputs = []
+        for token_prompt in token_prompts:
+            inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])})
+            output = hf_model.model(inputs["input_ids"])
+            embedding = torch.nn.functional.normalize(
+                output.last_hidden_state.float()[0, -1], dim=0
+            )
+            hf_outputs.append(embedding.cpu().tolist())
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        pooler_config=PoolerConfig(task="embed"),
+        max_model_len=64,
+        max_num_batched_tokens=chunk_size,
+        max_num_seqs=2,
+        gpu_memory_utilization=0.25,
+        enforce_eager=True,
+        enable_chunked_prefill=True,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = vllm_model.embed(prompts)
+
+    check_embeddings_close(
+        embeddings_0_lst=hf_outputs,
+        embeddings_1_lst=vllm_outputs,
+        name_0="hf",
+        name_1="vllm",
+        tol=1e-2,
+    )
@@ -49,3 +49,60 @@ def test_models(
            vllm_output,
            rtol=2e-3 if dtype == "float" else 1e-2,
        )
+
+
+@pytest.mark.core_model
+def test_bert_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "cross-encoder/ms-marco-TinyBERT-L-2-v2"
+    score_inputs = (
+        "What is the capital of France?",
+        [
+            "Paris.",
+            "Paris is the capital and largest city of France.",
+            "William Shakespeare wrote Hamlet in the early seventeenth century.",
+        ],
+    )
+    prompt_batches = [
+        ["short input"],
+        [
+            "short input",
+            "a longer input that exercises mixed sequence lengths",
+        ],
+    ]
+
+    with hf_runner(
+        model, dtype="half", auto_cls=AutoModelForSequenceClassification
+    ) as hf_model:
+        # HfRunner uses problem_type to preserve the model's
+        # sbert_ce_default_activation_function=Identity raw logits.
+        hf_model.config.problem_type = "regression"
+        hf_outputs = [hf_model.classify(prompts) for prompts in prompt_batches]
+
+    text_1, text_2 = score_inputs
+    text_pairs = [[text_1, document] for document in text_2]
+    with hf_runner(model, dtype="half", is_cross_encoder=True) as hf_model:
+        hf_scores = hf_model.predict(text_pairs).tolist()
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype="half",
+        max_model_len=64,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = [vllm_model.classify(prompts) for prompts in prompt_batches]
+        vllm_scores = vllm_model.score(*score_inputs)
+
+    for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs):
+        hf_tensor = torch.tensor(hf_batch)
+        vllm_tensor = torch.tensor(vllm_batch)
+        assert vllm_tensor.shape == hf_tensor.shape
+        assert torch.allclose(vllm_tensor, hf_tensor, rtol=1e-2, atol=1e-4)
+
+    assert torch.allclose(
+        torch.tensor(vllm_scores),
+        torch.tensor(hf_scores),
+        rtol=1e-2,
+        atol=1e-4,
+    )
@@ -5,6 +5,7 @@ import pytest
 import torch
 from transformers import AutoModel

+from vllm import PoolingParams
 from vllm.config import PoolerConfig

 from ...utils import check_embeddings_close
@@ -132,6 +133,7 @@ def test_encoder_only_model_runner_v2_attention(
            task="embed", seq_pooling_type="LAST", use_activation=True
        ),
    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
        vllm_outputs = vllm_model.embed(prompts)

    check_embeddings_close(
@@ -141,3 +143,77 @@ def test_encoder_only_model_runner_v2_attention(
        name_1="vllm",
        tol=1e-2,
    )
+
+
+@pytest.mark.core_model
+def test_encoder_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "sentence-transformers/all-MiniLM-L6-v2"
+    prompt_batches = [
+        ["short input"],
+        [
+            "short input",
+            "a longer input that exercises mixed sequence lengths",
+        ],
+    ]
+
+    with hf_runner(model, is_sentence_transformer=True) as hf_model:
+        hf_outputs = [hf_model.encode(prompts) for prompts in prompt_batches]
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        max_model_len=64,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = [vllm_model.embed(prompts) for prompts in prompt_batches]
+
+    for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs):
+        check_embeddings_close(
+            embeddings_0_lst=hf_batch,
+            embeddings_1_lst=vllm_batch,
+            name_0="hf",
+            name_1="vllm",
+            tol=1e-2,
+        )
+
+
+@pytest.mark.core_model
+def test_matryoshka_dimensions_model_runner_v2(
+    hf_runner, vllm_runner, monkeypatch
+) -> None:
+    model = "Snowflake/snowflake-arctic-embed-m-v1.5"
+    prompts = ["short input", "a longer input for a different output width"]
+    dimensions = [None, 256]
+
+    with hf_runner(model, is_sentence_transformer=True) as hf_model:
+        hf_outputs = hf_model.encode(prompts)
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        max_model_len=64,
+        gpu_memory_utilization=0.25,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = vllm_model.embed(
+            prompts,
+            pooling_params=[PoolingParams(dimensions=d) for d in dimensions],
+        )
+
+    expected_outputs = []
+    for output, dimension in zip(hf_outputs, dimensions):
+        output = torch.as_tensor(output)
+        if dimension is not None:
+            output = torch.nn.functional.normalize(output[:dimension], dim=0)
+        expected_outputs.append(output.tolist())
+
+    assert [len(output) for output in vllm_outputs] == [768, 256]
+    check_embeddings_close(
+        embeddings_0_lst=expected_outputs,
+        embeddings_1_lst=vllm_outputs,
+        name_0="hf",
+        name_1="vllm",
+        tol=1e-2,
+    )
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+from unittest.mock import MagicMock, patch
+
+import numpy as np
 import pytest
 import torch
 import torch.nn as nn
@@ -9,8 +12,13 @@ from vllm.model_executor.models.bert import (
    BertMLMHead,
    SPLADESparsePooler,
 )
+from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
+from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
+from vllm.v1.worker.gpu.states import RequestState

 # ---------------------------------------------------------------------
 # Functional test: SPLADE formula correctness (no HF download needed)
@@ -91,3 +99,70 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):
        rtol=1e-4,
        atol=1e-4,
    )
+
+
+def test_pooling_runner_gathers_required_token_ids() -> None:
+    runner = PoolingRunner.__new__(PoolingRunner)
+    pooling_params = PoolingParams(task="embed", requires_token_ids=True)
+    runner.pooling_params = {1: pooling_params, 3: pooling_params}
+    runner.pooling_states = {1: PoolingStates(), 3: PoolingStates()}
+    runner.prompt_token_ids = {
+        1: torch.tensor([101, 102]),
+        3: torch.tensor([101, 11, 102]),
+    }
+
+    input_batch = MagicMock(spec=InputBatch)
+    input_batch.idx_mapping_np = np.array([3, 1], dtype=np.int32)
+    input_batch.num_reqs = 2
+    req_states = MagicMock(spec=RequestState)
+    req_states.prompt_len = MagicMock(np=np.array([0, 2, 0, 3], dtype=np.int32))
+    metadata = runner._get_pooling_metadata(
+        input_batch, req_states, torch.device(current_platform.device_type)
+    )
+
+    expected = torch.tensor([[101, 11, 102], [101, 102, 0]])
+    assert metadata.prompt_token_ids_cpu is not None
+    assert metadata.prompt_token_ids is not None
+    assert metadata.prompt_token_ids_cpu.is_pinned() == PIN_MEMORY
+    torch.testing.assert_close(
+        metadata.prompt_lens, torch.tensor([3, 2], dtype=torch.int32)
+    )
+    torch.testing.assert_close(metadata.prompt_token_ids_cpu, expected)
+    torch.testing.assert_close(metadata.prompt_token_ids.cpu(), expected)
+
+
+def test_pooling_runner_stores_only_required_token_ids() -> None:
+    runner = PoolingRunner.__new__(PoolingRunner)
+    runner.model = MagicMock()
+    runner.supported_tasks = frozenset({"embed"})
+    runner.pooling_params = {}
+    runner.pooling_states = {}
+    runner.prompt_token_ids = {}
+
+    runner.add_request(1, PoolingParams(task="embed"), [101, 102])
+    runner.add_request(
+        2,
+        PoolingParams(task="embed", requires_token_ids=True),
+        [101, 11, 102],
+    )
+
+    assert 1 not in runner.prompt_token_ids
+    torch.testing.assert_close(runner.prompt_token_ids[2], torch.tensor([101, 11, 102]))
+
+
+def test_pooling_runner_rejects_unsupported_selected_task() -> None:
+    model = MagicMock()
+    model.pooler.get_supported_tasks.return_value = {
+        "embed",
+        "embed&token_classify",
+        "token_classify",
+    }
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_seqs = 2
+    vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify"
+
+    with (
+        patch.object(PoolingRunner, "get_supported_tasks", return_value=["embed"]),
+        pytest.raises(ValueError, match="selects 'embed&token_classify'"),
+    ):
+        PoolingRunner(model, vllm_config)
@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest

+from vllm.exceptions import VLLMValidationError
+
 MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
 max_model_len = 128

@@ -60,7 +62,7 @@ def test_bigger_truncation_size(
    truncate_prompt_tokens = max_model_len + 1

    with (
-        pytest.raises(ValueError),
+        pytest.raises(VLLMValidationError),
        vllm_runner(
            model_name, runner="pooling", max_model_len=max_model_len
        ) as vllm_model,
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import uuid
+from io import BytesIO
 from pathlib import Path

 import numpy as np
@@ -10,6 +11,7 @@ from PIL import Image, ImageDraw

 from vllm.multimodal.hasher import MultiModalHasher
 from vllm.multimodal.media.base import MediaWithBytes
+from vllm.multimodal.media.image import ImageMediaIO
 from vllm.multimodal.parse import MultiModalDataParser

 pytestmark = pytest.mark.cpu_test
@@ -134,3 +136,37 @@ def test_hash_image_exif_id():
    assert hasher.hash_kwargs(image=image1) == hasher.hash_kwargs(image=id.bytes)
    # second image has non-UUID in ImageID, so it should hash to the image data
    assert hasher.hash_kwargs(image=image2) == hasher.hash_kwargs(image=image2a)
+
+
+def _rgba_png_bytes() -> bytes:
+    image = Image.new("RGBA", (8, 8), (255, 0, 0, 128))
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def test_hash_collision_media_io_config():
+    data = _rgba_png_bytes()
+    white = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data)
+    black = ImageMediaIO(rgba_background_color=(0, 0, 0)).load_bytes(data)
+    white2 = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data)
+    keep = ImageMediaIO(image_mode=None).load_bytes(data)
+
+    hasher = MultiModalHasher
+    assert hasher.hash_kwargs(image=white) != hasher.hash_kwargs(image=black)
+    assert hasher.hash_kwargs(image=white) != hasher.hash_kwargs(image=keep)
+    assert hasher.hash_kwargs(image=white) == hasher.hash_kwargs(image=white2)
+
+
+def test_hash_media_io_noop_config_preserves_hash():
+    image = Image.new("RGB", (8, 8), (0, 128, 255))
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    data = buf.getvalue()
+
+    loaded = ImageMediaIO().load_bytes(data)
+    assert loaded.io_config is None
+
+    plain = MediaWithBytes(loaded.media, data)
+    hasher = MultiModalHasher
+    assert hasher.hash_kwargs(image=loaded) == hasher.hash_kwargs(image=plain)
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections import deque
+from collections import defaultdict, deque
+from collections.abc import Callable
 from unittest.mock import Mock

 import pytest
@@ -329,46 +330,311 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
    assert not scheduler.running


-def test_no_placeholder_underflow_on_discarded_spec_frame():
-    num_spec = 5
+class PipelinedEngine:
+    """Drive a real AsyncScheduler like EngineCore.step_with_batch_queue:
+    schedule until the batch queue is full, then process the oldest step's
+    output. Async PP runs pp_size+1 concurrent batches, so up to pp_size
+    steps are in flight at each schedule() call -- the window in which
+    preemption must handle output that has not yet returned. (Single-GPU e2e
+    tests can never create this window: at PP=1, exactly one step is in
+    flight and it is processed before a preempted request can resume.)
+
+    The model runner is emulated with the V2 runner's own bookkeeping, from
+    only what the scheduler serializes to it: slots flushed on
+    preempted_req_ids, resumed requests re-added from the NewRequestData
+    snapshot, sampling when a step reaches the end of the runner's own view
+    of the sequence. This makes preemption races observable: a stale token
+    delivered after a resume is scheduled extends the scheduler's sequence
+    but not the runner's.
+
+    Every sample emits a globally unique token tagged with its sampled
+    position, so tests can assert exact delivery.
+    """
+
+    def __init__(
+        self,
+        scheduler: AsyncScheduler,
+        queue_size: int,
+        accept_drafts: Callable[[int, str, int], int] | None = None,
+    ):
+        self.scheduler = scheduler
+        self.queue_size = queue_size
+        self.accept_drafts = accept_drafts
+        # In-flight steps: (scheduler_output, new_reqs snapshot) in FIFO order.
+        self.queue: deque[tuple[SchedulerOutput, list[tuple[str, int, int]]]] = deque()
+        # Runner-side request state: req_id -> [seq_len, num_computed] as the
+        # runner sees them (its own sampled tokens, not the scheduler's).
+        self.runner_view: dict[str, list[int]] = {}
+        # All tokens the fake runner ever sampled, per request, in order.
+        self.emitted: dict[str, list[int]] = defaultdict(list)
+        # Sequence position each (globally unique) token was sampled for.
+        self.emitted_position: dict[int, int] = {}
+        self.step_idx = 0
+        self._next_token = 1000
+
+    def _schedule(self) -> bool:
+        scheduler_output = self.scheduler.schedule()
+        self.step_idx += 1
+        # Snapshot what NewRequestData serializes at schedule time (both new
+        # and resumed requests for the V2 runner).
+        new_reqs = [
+            (r.req_id, len(r.prefill_token_ids), r.num_computed_tokens)
+            for r in scheduler_output.scheduled_new_reqs
+        ]
+        # Enqueue empty steps too (the engine executes them), so the runner
+        # still observes their preempted/finished request ids in step order.
+        self.queue.appendleft((scheduler_output, new_reqs))
+        return True
+
+    def _process_oldest_step(self) -> None:
+        scheduler_output, new_reqs = self.queue.pop()
+        # Worker-side state updates, in step order: flush preempted/finished
+        # slots, then (re-)add new/resumed requests.
+        for req_id in scheduler_output.preempted_req_ids or ():
+            self.runner_view.pop(req_id, None)
+        for req_id in scheduler_output.finished_req_ids or ():
+            self.runner_view.pop(req_id, None)
+        for req_id, seq_len, num_computed in new_reqs:
+            self.runner_view[req_id] = [seq_len, num_computed]
+
+        req_ids = list(scheduler_output.num_scheduled_tokens.keys())
+        sampled_token_ids: list[list[int]] = []
+        for req_id in req_ids:
+            num_scheduled = scheduler_output.num_scheduled_tokens[req_id]
+            view = self.runner_view.get(req_id)
+            if view is None:
+                # Slot already flushed (request finished/aborted mid-flight).
+                sampled_token_ids.append([])
+                continue
+            seq_len, num_computed = view
+            end = num_computed + num_scheduled
+            if end < seq_len:
+                # Partial prefill by the runner's own bookkeeping: no sample.
+                view[1] = end
+                sampled_token_ids.append([])
+                continue
+            drafts = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())
+            num_accepted = (
+                min(self.accept_drafts(self.step_idx, req_id, len(drafts)), len(drafts))
+                if drafts and self.accept_drafts
+                else 0
+            )
+            num_rejected = len(drafts) - num_accepted
+            tokens = list(range(self._next_token, self._next_token + 1 + num_accepted))
+            self._next_token += 1 + num_accepted
+            self.emitted[req_id].extend(tokens)
+            sampled_token_ids.append(tokens)
+            # Rejected drafts roll back computed; the sampled tokens extend
+            # the runner's sequence.
+            view[1] = end - num_rejected
+            view[0] = view[1] + 1
+            for offset, token in enumerate(tokens):
+                self.emitted_position[token] = view[0] - len(tokens) + offset
+        model_runner_output = ModelRunnerOutput(
+            req_ids=req_ids,
+            req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},
+            sampled_token_ids=sampled_token_ids,
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        )
+        self.scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    def run(
+        self,
+        max_steps: int = 2000,
+        before_step: Callable[[int, "PipelinedEngine"], None] | None = None,
+    ) -> None:
+        for i in range(max_steps):
+            if not self.scheduler.has_requests() and not self.queue:
+                return
+            if before_step is not None:
+                before_step(i, self)
+            scheduled = (
+                self.scheduler.has_requests()
+                and len(self.queue) < self.queue_size
+                and self._schedule()
+            )
+            if scheduled and len(self.queue) < self.queue_size:
+                # Queue not yet full: the engine returns without blocking.
+                continue
+            if self.queue:
+                self._process_oldest_step()
+        raise AssertionError("engine loop did not converge")
+
+
+def _create_async_pp_scheduler(
+    num_spec: int, pp_size: int = 3, num_blocks: int = 5
+) -> AsyncScheduler:
    scheduler = create_scheduler(
        async_scheduling=True,
-        num_speculative_tokens=num_spec,
-        speculative_method="ngram_gpu",
+        num_speculative_tokens=num_spec or None,
+        speculative_method="ngram_gpu" if num_spec else None,
+        use_v2_model_runner=True,
+        num_blocks=num_blocks,
+        block_size=16,
+        max_num_batched_tokens=512,
    )
-    req = create_requests(num_requests=1, max_tokens=20)[0]
-    req.num_computed_tokens = req.num_tokens
-    scheduler.requests[req.request_id] = req
-    scheduler.running.append(req)
-    req.status = RequestStatus.RUNNING
+    # Emulate PP at the scheduler level; constructing with
+    # pipeline_parallel_size>1 requires that many visible GPUs. Drive with
+    # queue_size=pp_size+1 (V2 async PP runs pp_size+1 concurrent batches).
+    scheduler.pp_size = pp_size
+    scheduler.use_pp = pp_size > 1
+    return scheduler

-    req.num_output_placeholders = 1
-    req.async_tokens_to_discard = num_spec
-    computed_before = req.num_computed_tokens

-    scheduler_output = SchedulerOutput(
-        scheduled_new_reqs=[],
-        scheduled_cached_reqs=CachedRequestData.make_empty(),
-        num_scheduled_tokens={req.request_id: num_spec + 1},
-        total_num_scheduled_tokens=num_spec + 1,
-        scheduled_encoder_inputs={},
-        scheduled_spec_decode_tokens={req.request_id: [10] * num_spec},
-        num_common_prefix_blocks=[],
-        finished_req_ids=set(),
-        free_encoder_mm_hashes=[],
+def _assert_ordered_subset(delivered: list[int], emitted: list[int]) -> None:
+    """Delivered tokens must be an order-preserving subset of the emitted
+    tokens with no duplicates (tokens are globally unique)."""
+    it = iter(emitted)
+    for token in delivered:
+        assert token in it, f"token {token} delivered out of order or twice"
+
+
+def _assert_positions_consistent(req, engine: PipelinedEngine) -> None:
+    """The i-th delivered output token must be one the runner sampled for
+    exactly sequence position prompt_len + i: catches a preempted request's
+    stale output landing on a position the resumed request resampled (or
+    vice versa), which token-stream equality alone cannot see."""
+    for i, token in enumerate(req.output_token_ids):
+        expected = req.num_prompt_tokens + i
+        actual = engine.emitted_position[token]
+        assert actual == expected, (
+            f"output {i} of {req.request_id}: token sampled for position "
+            f"{actual}, delivered as position {expected}"
+        )
+
+
+@pytest.mark.parametrize("num_spec", [0, 3])
+def test_kv_pressure_preemption_with_inflight_output(num_spec: int):
+    """KV-pressure preemption of requests with in-flight async output.
+
+    PP=3 + async scheduling (batch queue of 4), a block pool small enough
+    that decodes contend and preempt mid-flight, and staggered arrivals so
+    the batch queue actually pipelines. A preempted request's in-flight steps
+    still return: their tokens must be delivered exactly once, their stale
+    spec-rejection counts must not corrupt the rolled-back counters, and the
+    resume must not resample a position that output later delivers.
+
+    Regression for the num_output_placeholders underflow EngineCore crash:
+    with the fix reverted, the num_spec=3 variant fails with exactly
+    ``assert request.num_output_placeholders >= 0`` when a stale spec output
+    returns after the preempted request was resumed and sampled.
+    """
+    max_tokens = 24
+    scheduler = _create_async_pp_scheduler(num_spec)
+    requests = create_requests(
+        num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True
    )
-    model_runner_output = ModelRunnerOutput(
-        req_ids=[req.request_id],
-        req_id_to_index={req.request_id: 0},
-        sampled_token_ids=[[999]],
-        logprobs=None,
-        prompt_logprobs_dict={},
-        pooler_output=[],
+    pending = list(requests)
+    for _ in range(2):
+        scheduler.add_request(pending.pop(0))
+
+    # Observe that the scenario under test actually occurs.
+    preempts_with_inflight_output = 0
+    orig_preempt = scheduler._preempt_request
+
+    def counting_preempt(request, timestamp, **kwargs):
+        nonlocal preempts_with_inflight_output
+        if request.num_in_flight_tokens > 0:
+            preempts_with_inflight_output += 1
+        return orig_preempt(request, timestamp, **kwargs)
+
+    scheduler._preempt_request = counting_preempt
+
+    def add_requests(step: int, engine: PipelinedEngine):
+        if pending:
+            scheduler.add_request(pending.pop(0))
+
+    engine = PipelinedEngine(
+        scheduler,
+        queue_size=4,
+        # Deterministically vary spec acceptance so stale outputs carry
+        # nonzero rejection counts.
+        accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1),
    )
+    engine.run(before_step=add_requests)

-    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert preempts_with_inflight_output > 0, (
+        "test did not exercise preemption with in-flight output"
+    )
+    for req in requests:
+        assert req.is_finished()
+        assert req.num_output_tokens == max_tokens
+        # Lossless: delivered tokens are exactly the sampled tokens, in order
+        # (the excluded tail was emitted after the request finished).
+        emitted = engine.emitted[req.request_id]
+        assert list(req.output_token_ids) == emitted[:max_tokens]
+        _assert_positions_consistent(req, engine)

-    assert req.num_output_placeholders == 1
-    assert req.num_computed_tokens == computed_before
-    assert req.async_tokens_to_discard == num_spec - 1
-    assert req.status == RequestStatus.RUNNING
+
+@pytest.mark.parametrize("pp_size", [1, 3])
+def test_reset_prefix_cache_with_inflight_output_under_kv_pressure(pp_size: int):
+    """reset_prefix_cache(reset_running_requests=True) resumes requests in
+    the same step it preempts them, so in-flight output must be dropped (the
+    resume resamples those positions).
+
+    pp_size=1: regression for the frame-based discard this fix replaces,
+    which with spec decode drained one *token* count per output frame and
+    over-discarded, corrupting the fresh frames after the resume.
+    pp_size=3: back-to-back resets, so the second re-preempts requests whose
+    dropped stale share is still in flight -- it must be recorded once (not
+    accumulated) and stay dropped.
+    """
+    max_tokens = 24
+    scheduler = _create_async_pp_scheduler(num_spec=3, pp_size=pp_size)
+    requests = create_requests(
+        num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True
+    )
+    pending = list(requests)
+    for _ in range(2):
+        scheduler.add_request(pending.pop(0))
+
+    # Observe re-preemptions with an undrained stale share (the
+    # double-count hazard).
+    repreempts_with_stale = 0
+    orig_preempt = scheduler._preempt_request
+
+    def counting_preempt(request, timestamp, **kwargs):
+        nonlocal repreempts_with_stale
+        if getattr(request, "num_stale_output_tokens", 0) > 0:
+            repreempts_with_stale += 1
+        return orig_preempt(request, timestamp, **kwargs)
+
+    scheduler._preempt_request = counting_preempt
+
+    resets = 0
+    reset_steps = {6, 14} if pp_size == 1 else {6, 7, 18, 19}
+
+    def before_step(step: int, engine: PipelinedEngine):
+        nonlocal resets
+        if pending:
+            scheduler.add_request(pending.pop(0))
+        if step in reset_steps and (engine.queue or scheduler.running):
+            scheduler.reset_prefix_cache(reset_running_requests=True)
+            resets += 1
+
+    engine = PipelinedEngine(
+        scheduler,
+        queue_size=pp_size + 1,
+        accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1),
+    )
+    engine.run(before_step=before_step)
+
+    assert resets > 0, "test did not exercise reset_prefix_cache"
+    if pp_size > 1:
+        # The re-preempt-while-stale-pending window needs pipeline depth.
+        assert repreempts_with_stale > 0, (
+            "test did not exercise re-preemption with an undrained stale share"
+        )
+    for req in requests:
+        assert req.is_finished()
+        assert req.num_output_tokens == max_tokens
+        # Dropped tokens are never delivered; order must be preserved with
+        # no duplicates.
+        _assert_ordered_subset(
+            list(req.output_token_ids), engine.emitted[req.request_id]
+        )
+        _assert_positions_consistent(req, engine)
+        # All stale shares fully drained by the end.
+        assert getattr(req, "num_stale_output_tokens", 0) == 0
@@ -1107,7 +1107,7 @@ def test_preemption_re_records_prefix_cache_query():
    request = create_requests(num_requests=1)[0]
    scheduler.add_request(request)

-    scheduler.schedule()
+    scheduler_output = scheduler.schedule()
    stats = scheduler.kv_cache_manager.prefix_cache_stats
    assert stats is not None
    assert (stats.requests, stats.preempted_requests) == (1, 0)
@@ -1116,6 +1116,21 @@ def test_preemption_re_records_prefix_cache_query():
    scheduler._preempt_request(request, 0.0)
    assert request.status == RequestStatus.PREEMPTED

+    scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[1000]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        ),
+    )
+    assert request.num_stale_output_tokens == 0
+    stats = scheduler.kv_cache_manager.prefix_cache_stats
+    assert stats is not None
+
    scheduler.schedule()
    assert request.status == RequestStatus.RUNNING
    assert stats.preempted_requests == 1
@@ -338,6 +338,30 @@ def get_fake_process_mamba_fn(
            assert copy_info[0][-1] == expected_temporal_src
            assert copy_info[1][-1] == expected_temporal_dest

+    def check_fused_copy_info(
+        action: tuple[int, int],
+        align_ctx: mamba_utils.MambaSpecDecodeGPUContext,
+    ):
+        # Align + spec-decode on a hybrid model routes the pre-copy through the
+        # fused kernel (preprocess_mamba -> run_fused_precopy) instead of
+        # do_mamba_copy_block, so copy_info is never populated. Verify from the
+        # fused buffers (req-0 scope, mirroring check_copy_info):
+        #   - the copy DECISION: src_col is -1 iff no pre-copy is scheduled;
+        #   - the DESTINATION column: state_idx == action[1] (curr_state_idx;
+        #     maps directly through block_ids, exactly as check_copy_info's dst).
+        # The source column is NOT asserted here: on the scalar path the source
+        # address is produced by the per-state copy func with an accept-token
+        # bias offset (collect_mamba_copy_meta), so prev_state_idx does not map
+        # to action[0] by plain equality. Source block-level exactness (incl.
+        # the accept-bias) is covered by test_precopy_mamba_align.py.
+        src_col = int(align_ctx.precopy_src_col_buf.np[0])
+        state_idx = int(align_ctx.mamba_state_idx_buf.np[0])
+        if action == (-1, -1):
+            assert src_col == -1
+        else:
+            assert src_col != -1
+            assert state_idx == action[1]
+
    def fake_preprocess_mamba_fn(
        scheduler_output: SchedulerOutput,
        kv_cache_config: KVCacheConfig,
@@ -348,6 +372,7 @@ def get_fake_process_mamba_fn(
        forward_context: dict[str, Any],
        mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
        copy_bufs: mamba_utils.MambaCopyBuffers,
+        align_ctx: mamba_utils.MambaSpecDecodeGPUContext | None = None,
    ):
        nonlocal copy_info
        copy_info = None
@@ -361,14 +386,21 @@ def get_fake_process_mamba_fn(
            forward_context,
            mamba_state_copy_funcs,
            copy_bufs,
+            align_ctx,
        )
        if cur_step_action is not None:
-            check_copy_info(
-                cur_step_action.preprocess_copy_idx,
-                kv_cache_config,
-                forward_context,
-                input_batch,
-            )
+            if align_ctx is not None:
+                check_fused_copy_info(
+                    cur_step_action.preprocess_copy_idx,
+                    align_ctx,
+                )
+            else:
+                check_copy_info(
+                    cur_step_action.preprocess_copy_idx,
+                    kv_cache_config,
+                    forward_context,
+                    input_batch,
+                )
        return ret

    def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers):
@@ -14,6 +14,7 @@ import torch
 from vllm import LLM, SamplingParams
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
    UnquantizedEmbeddingMethod,
 )

@@ -28,6 +29,7 @@ class _FakeLmHead:
        self.weight = weight
        self.quant_method = object() if quantized else UnquantizedEmbeddingMethod()
        self.shard_indices = shard_indices
+        self.tp_size = 1


 def _build_processor(vocab_size: int) -> LogitsProcessor:
@@ -135,11 +137,59 @@ def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):
        lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None)


+def test_replicated_lm_head_skips_tp_communication_and_preserves_processing(
+    default_vllm_config,
+):
+    from unittest import mock
+
+    vocab_size, hidden_size = 12, 8
+    soft_cap, scale = 2.0, 0.5
+    lp = LogitsProcessor(
+        vocab_size,
+        soft_cap=soft_cap,
+        scale=scale,
+    )
+    lp.head_dtype = torch.float32
+
+    hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
+    weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
+    world_size_getter = (
+        "vllm.model_executor.layers.vocab_parallel_embedding."
+        "get_tensor_model_parallel_world_size"
+    )
+    with mock.patch(world_size_getter, return_value=2):
+        lm_head = ParallelLMHead(
+            vocab_size,
+            hidden_size,
+            params_dtype=torch.bfloat16,
+            disable_tp=True,
+        )
+    lm_head.weight_loader(lm_head.weight, weight)
+    assert lm_head.tp_size == 1
+
+    with mock.patch.object(lp, "_gather_logits") as gather_mock:
+        logits = lp(lm_head, hidden_states)
+
+    gather_mock.assert_not_called()
+
+    expected = torch.nn.functional.linear(hidden_states.float(), weight.float())
+    expected = torch.tanh(expected / soft_cap) * soft_cap * scale
+    torch.testing.assert_close(logits, expected)
+
+    all_gather_path = (
+        "vllm.model_executor.layers.logits_processor.tensor_model_parallel_all_gather"
+    )
+    with mock.patch(all_gather_path) as all_gather:
+        top = lp.get_top_tokens(lm_head, hidden_states)
+
+    all_gather.assert_not_called()
+    assert torch.equal(top, expected.argmax(dim=-1))
+
+
 def test_get_top_tokens_honors_head_dtype(default_vllm_config):
    # The spec-decode local-argmax path (get_top_tokens) must run the lm_head
    # in head_dtype too, not just _get_logits.
    import types
-    from unittest import mock

    vocab_size, hidden_size = 64, 16
    lp = _build_processor(vocab_size)
@@ -154,13 +204,7 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config):
        ),
    )

-    with mock.patch(
-        "vllm.model_executor.layers.logits_processor."
-        "get_tensor_model_parallel_world_size",
-        return_value=1,
-    ):
-        top = lp.get_top_tokens(lm_head, hidden_states, None)
-
+    top = lp.get_top_tokens(lm_head, hidden_states, None)
    expected = torch.nn.functional.linear(hidden_states.float(), weight.float()).argmax(
        dim=-1
    )
@@ -41,6 +41,7 @@ def mock_model_runner_with_req_states():
    runner.sampler = None
    runner.prompt_logprobs_worker = None
    runner.is_last_pp_rank = False
+    runner.pooling_runner = None

    # Mock staged writes — they use Triton kernels that require GPU
    runner.req_states.apply_staged_writes = Mock()
@@ -644,7 +644,7 @@ class SpeechToTextBaseServing(GenerateBaseServing):
                        TranscriptionResponseVerbose(
                            text=text,
                            language=request.language,
-                            duration=str(duration_s),
+                            duration=duration_s,
                            segments=total_segments,
                        ),
                    )
@@ -658,7 +658,7 @@ class SpeechToTextBaseServing(GenerateBaseServing):
                        TranslationResponseVerbose(
                            text=text,
                            language=request.language,
-                            duration=str(duration_s),
+                            duration=duration_s,
                            segments=total_segments,
                        ),
                    )
@@ -384,7 +384,7 @@ class TranscriptionSegment(OpenAIBaseModel):


 class TranscriptionResponseVerbose(OpenAIBaseModel):
-    duration: str
+    duration: float
    """The duration of the input audio."""

    language: str
@@ -357,7 +357,7 @@ class TranslationSegment(OpenAIBaseModel):


 class TranslationResponseVerbose(OpenAIBaseModel):
-    duration: str
+    duration: float
    """The duration of the input audio."""

    language: str
@@ -30,9 +30,6 @@ def rms_norm(
    x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
 ) -> Tensor:
    assert variance_size is None
-    if weight is None:
-        # Kernel requires weight tensor, pass ones
-        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
    output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
    torch.ops._C.rms_norm(output, x, weight, epsilon)
    return output
@@ -58,8 +55,5 @@ def fused_add_rms_norm(
    variance_size: int | None = None,
 ) -> tuple[Tensor, Tensor]:
    assert variance_size is None
-    if weight is None:
-        # Kernel requires weight tensor, pass ones
-        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
    torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
    return x, x_residual
@@ -7,7 +7,6 @@ import torch.nn.functional as F

 from vllm.config import get_current_vllm_config
 from vllm.distributed import (
-    get_tensor_model_parallel_world_size,
    tensor_model_parallel_all_gather,
    tensor_model_parallel_gather,
 )
@@ -145,7 +144,8 @@ class LogitsProcessor(PluggableLayer):
        logits = self._apply_head(lm_head, hidden_states, embedding_bias)

        # Gather logits for TP
-        logits = self._gather_logits(logits)
+        if lm_head.tp_size > 1:
+            logits = self._gather_logits(logits)

        # Remove paddings in vocab (if any).
        if logits is not None:
@@ -169,7 +169,7 @@ class LogitsProcessor(PluggableLayer):
                "The local argmax reduction optimization is not supported for "
                "non-positive logit scaling factors."
            )
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = lm_head.tp_size

        logits = self._apply_head(lm_head, hidden_states, embedding_bias)
        if self.soft_cap is not None:
@@ -232,6 +232,7 @@ class VocabParallelEmbedding(PluggableLayer):
        padding_size: padding size for the vocabulary.
        quant_config: quant config for the layer
        prefix: full name of the layer in the state dict
+        disable_tp: If true, tensor parallelism will be disabled for this layer.
    """  # noqa: E501

    # --8<-- [end:vocab_parallel_embedding]
@@ -245,12 +246,19 @@ class VocabParallelEmbedding(PluggableLayer):
        padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
+        *,
+        disable_tp: bool = False,
    ):
        super().__init__()

        # Keep the input dimensions.
-        tp_rank = get_tensor_model_parallel_rank()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.disable_tp = disable_tp
+        if disable_tp:
+            tp_rank, self.tp_size = 0, 1
+        else:
+            tp_rank = get_tensor_model_parallel_rank()
+            self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = tp_rank
        self.num_embeddings = num_embeddings
        self.padding_size = padding_size
        self.org_vocab_size = org_num_embeddings or num_embeddings
@@ -323,6 +331,13 @@ class VocabParallelEmbedding(PluggableLayer):
            params_dtype=params_dtype,
            weight_loader=self.weight_loader,
        )
+        self.update_param_tp_status()
+
+    def update_param_tp_status(self):
+        for param in self.parameters():
+            if isinstance(param, BasevLLMParameter):
+                param.tp_rank = self.tp_rank
+                param.tp_size = self.tp_size

    @classmethod
    def _get_indices(
@@ -487,9 +502,9 @@ class VocabParallelEmbedding(PluggableLayer):
        # Mask the output embedding.
        if self.tp_size > 1:
            output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
-        # Reduce across all the model parallel GPUs.
-        output = tensor_model_parallel_all_reduce(output_parallel)
-        return output
+            # Reduce across all the model parallel GPUs.
+            return tensor_model_parallel_all_reduce(output_parallel)
+        return output_parallel

    def extra_repr(self) -> str:
        s = f"num_embeddings={self.num_embeddings_per_partition}"
@@ -516,6 +531,7 @@ class ParallelLMHead(VocabParallelEmbedding):
        params_dtype: type of the parameters.
        org_num_embeddings: original vocabulary size (without LoRA).
        padding_size: padding size for the vocabulary.
+        disable_tp: If true, tensor parallelism will be disabled for this layer.
    """

    # --8<-- [end:parallel_lm_head]
@@ -530,6 +546,8 @@ class ParallelLMHead(VocabParallelEmbedding):
        padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
+        *,
+        disable_tp: bool = False,
    ):
        super().__init__(
            num_embeddings,
@@ -539,6 +557,7 @@ class ParallelLMHead(VocabParallelEmbedding):
            padding_size,
            quant_config,
            prefix,
+            disable_tp=disable_tp,
        )
        self.quant_config = quant_config
        if bias:
@@ -24,7 +24,6 @@ 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
@@ -40,6 +39,10 @@ class DSparkMarkovHead(nn.Module):
    ``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias
    (``draft_vocab_size``) added to the base draft logits. The two sizes
    coincide for full-vocab drafts.
+
+    Both weights are replicated because the head runs sequentially for every
+    draft position. Sharding them would add an all-reduce and a full-vocab
+    gather to each position.
    """

    def __init__(
@@ -50,19 +53,24 @@ class DSparkMarkovHead(nn.Module):
        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_w1 = nn.Embedding(vocab_size, markov_rank)
        self.markov_w2 = ParallelLMHead(
-            draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
+            draft_vocab_size,
+            markov_rank,
+            bias=False,
+            prefix=maybe_prefix(prefix, "markov_w2"),
+            disable_tp=True,
        )

    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:
+    def bias(
+        self,
+        markov_embed: torch.Tensor,
+        logits_processor: LogitsProcessor,
+    ) -> torch.Tensor:
        """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
        return logits_processor(self.markov_w2, markov_embed)

@@ -164,6 +164,8 @@ def is_flashkda_supported(
    dtype: torch.dtype,
    lower_bound: float | None,
 ) -> bool:
+    if not current_platform.is_cuda():
+        return False
    capability = current_platform.get_device_capability()
    return (
        capability is not None
@@ -81,6 +81,11 @@ class MultiModalHasher:
            ):
                return (exif[Image.ExifTags.Base.ImageID].bytes,)

+            if obj.io_config:
+                return cls.iter_item_to_bytes(
+                    "image",
+                    {"io_config": obj.io_config, "data": obj.original_bytes},
+                )
            return cls.iter_item_to_bytes("image", obj.original_bytes)

        if isinstance(obj, MediaWithBytes) and isinstance(obj.media, np.ndarray):
@@ -29,6 +29,9 @@ class MediaWithBytes(Generic[_T]):

    media: _T
    original_bytes: bytes = field(repr=False)
+    io_config: dict[str, Any] | None = None
+    """Decode settings that altered the media relative to `original_bytes`
+    (e.g. `image_mode` conversion), so they participate in cache hashing."""

    def __array__(self, *args, **kwargs) -> np.ndarray:
        """Allow np.array(obj) to return np.array(obj.media)."""
@@ -86,10 +86,17 @@ class ImageMediaIO(MediaIO[Image.Image]):
                )
            image = normalize_image(image)
            image.load()
-            image = self._convert_image_mode(image)
+            converted = self._convert_image_mode(image)
        except (OSError, Image.UnidentifiedImageError) as e:
            raise ValueError(f"Failed to load image: {e}") from e
-        return MediaWithBytes(image, data)
+
+        io_config = None
+        if converted is not image:
+            io_config = {
+                "image_mode": self.image_mode,
+                "rgba_background_color": self.rgba_background_color,
+            }
+        return MediaWithBytes(converted, data, io_config)

    def load_base64(self, media_type: str, data: str) -> MediaWithBytes[Image.Image]:
        return self.load_bytes(pybase64.b64decode(data, validate=True))
@@ -65,7 +65,7 @@ class ProcessorInputs:
                        **{modality: item},
                        **hf_processor_mm_kwargs,
                    )
-                    for item in data_items
+                    for item in data_items.get_all_items_for_hash()
                ]

        return mm_hashes
@@ -49,23 +49,18 @@ class AsyncScheduler(Scheduler):
                request.next_decode_eligible_step = self.current_step + self.pp_size

    def _update_request_with_output(
-        self, request: Request, new_token_ids: list[int]
+        self, request: Request, new_token_ids: list[int], is_stale: bool = False
    ) -> tuple[list[int], bool]:
-        if request.async_tokens_to_discard > 0:
-            # The request was force-preempted in reset_prefix_cache; drop one
-            # stale in-flight async output frame per call until the counter
-            # is drained.
-            request.async_tokens_to_discard -= 1
-            return [], False
-
        status_before_update = request.status
        new_token_ids, stopped = super()._update_request_with_output(
            request, new_token_ids
        )

-        # Update the number of output placeholders.
-        request.num_output_placeholders -= len(new_token_ids)
-        assert request.num_output_placeholders >= 0
+        # Placeholders were zeroed at preemption; a stale delivery must not
+        # decrement them (it would underflow).
+        if not is_stale:
+            request.num_output_placeholders -= len(new_token_ids)
+            assert request.num_output_placeholders >= 0

        # Cache the new tokens. Preempted requests should be skipped.
        if status_before_update == RequestStatus.RUNNING:
@@ -694,6 +694,17 @@ class Scheduler(SchedulerInterface):
                    step_skipped_waiting.prepend_request(request)
                    continue

+                if (
+                    request.num_stale_output_tokens > 0
+                    and not request.drop_stale_output
+                ):
+                    # Deliverable stale output still in flight: resuming now
+                    # could resample a position that output later delivers.
+                    # It drains within the pipeline depth.
+                    request_queue.pop_request()
+                    step_skipped_waiting.prepend_request(request)
+                    continue
+
                # Check that adding the request still respects the max_loras
                # constraint.
                if (
@@ -1244,11 +1255,17 @@ class Scheduler(SchedulerInterface):

        return new_block_ids_to_zero or None

-    def _preempt_request(self, request: Request, timestamp: float) -> None:
+    def _preempt_request(
+        self, request: Request, timestamp: float, drop_stale_output: bool = False
+    ) -> None:
        """Preempt a request and put it back to the waiting queue.

        NOTE: The request should be popped from the running queue outside of this
        method.
+
+        drop_stale_output: drop (rather than deliver) any in-flight output; used
+        by reset_prefix_cache, whose same-step resume would otherwise deliver
+        tokens out of order.
        """
        assert request.status == RequestStatus.RUNNING, (
            "Only running requests can be preempted"
@@ -1260,6 +1277,18 @@ class Scheduler(SchedulerInterface):
        request.num_computed_tokens = 0
        if request.spec_token_ids:
            request.spec_token_ids = []
+        # Async scheduling: mark all in-flight output as stale. Its tokens are
+        # still delivered on return (dropping them would perturb spec-decode
+        # acceptance) but must not mutate the reset counters; each step drains
+        # its share in update_from_output. num_in_flight_tokens already
+        # includes any undrained stale share, so assign rather than accumulate.
+        # An undrained drop-mode share stays dropped: its positions have
+        # already been resampled.
+        request.drop_stale_output = drop_stale_output or (
+            request.drop_stale_output and request.num_stale_output_tokens > 0
+        )
+        request.num_stale_output_tokens = request.num_in_flight_tokens
+        request.num_output_placeholders = 0
        request.num_preemptions += 1
        if self.log_stats:
            request.record_event(EngineCoreEventType.PREEMPTED, timestamp)
@@ -1687,8 +1716,14 @@ class Scheduler(SchedulerInterface):
        for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
            assert num_tokens_scheduled > 0
            request = self.requests.get(req_id)
+            output_is_stale = False
            if request is not None:
                request.num_in_flight_tokens -= num_tokens_scheduled
+                # Drain any stale share (see _preempt_request) in lockstep.
+                if request.num_stale_output_tokens > 0:
+                    output_is_stale = True
+                    request.num_stale_output_tokens -= num_tokens_scheduled
+                    assert request.num_stale_output_tokens >= 0
            if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
                # skip failed or rescheduled requests from KV load failure
                continue
@@ -1702,6 +1737,10 @@ class Scheduler(SchedulerInterface):
                # In this case, we use is_finished() to check.
                continue

+            # Drop-mode stale output (same-step resume) is discarded entirely.
+            if output_is_stale and request.drop_stale_output:
+                continue
+
            req_index = model_runner_output.req_id_to_index[req_id]
            generated_token_ids = (
                sampled_token_ids[req_index] if sampled_token_ids else []
@@ -1710,28 +1749,22 @@ class Scheduler(SchedulerInterface):
            scheduled_spec_token_ids = (
                scheduler_output.scheduled_spec_decode_tokens.get(req_id)
            )
-            # Skip a stale frame still pending discard (async_tokens_to_discard
-            # > 0): its pre-reset rejection count would underflow the counters.
-            if (
-                scheduled_spec_token_ids
-                and (generated_token_ids or self.num_sampled_tokens_per_step == 0)
-                and request.async_tokens_to_discard == 0
+            if scheduled_spec_token_ids and (
+                generated_token_ids or self.num_sampled_tokens_per_step == 0
            ):
                num_draft_tokens = len(scheduled_spec_token_ids)
                num_sampled = self.num_sampled_tokens_per_step
                num_accepted = max(len(generated_token_ids) - num_sampled, 0)
                num_rejected = num_draft_tokens - num_accepted
-                # num_computed_tokens represents the number of tokens
-                # processed in the current step, considering scheduled
-                # tokens and rejections. If some tokens are rejected,
-                # num_computed_tokens is decreased by the number of rejected
-                # tokens.
-                if request.num_computed_tokens > 0:
-                    request.num_computed_tokens -= num_rejected
-                # If async scheduling, num_output_placeholders also includes
-                # the scheduled spec tokens count and so is similarly adjusted.
-                if request.num_output_placeholders > 0:
-                    request.num_output_placeholders -= num_rejected
+                # Rejections roll back num_computed_tokens (and, under async
+                # scheduling, num_output_placeholders, which covers the spec
+                # tokens). A stale rejection count predates the preemption
+                # rollback and must not apply.
+                if not output_is_stale:
+                    if request.num_computed_tokens > 0:
+                        request.num_computed_tokens -= num_rejected
+                    if request.num_output_placeholders > 0:
+                        request.num_output_placeholders -= num_rejected
                spec_decoding_stats = self.make_spec_decoding_stats(
                    spec_decoding_stats,
                    num_draft_tokens=num_draft_tokens,
@@ -1757,7 +1790,7 @@ class Scheduler(SchedulerInterface):
            # Check for stop and update request status.
            if new_token_ids:
                new_token_ids, stopped = self._update_request_with_output(
-                    request, new_token_ids
+                    request, new_token_ids, is_stale=output_is_stale
                )
            elif request.pooling_params and pooler_output is not None:
                # Pooling stops as soon as there is output.
@@ -1899,6 +1932,7 @@ class Scheduler(SchedulerInterface):
        if stopped_preempted_reqs:
            # This is a rare case and unlikely to impact performance.
            self.waiting.remove_requests(stopped_preempted_reqs)
+            self.skipped_waiting.remove_requests(stopped_preempted_reqs)

        error_req_ids = set(self.grammar_compile_error_reqs)
        self.grammar_compile_error_reqs.clear()
@@ -2041,8 +2075,9 @@ class Scheduler(SchedulerInterface):
        return False

    def _update_request_with_output(
-        self, request: Request, new_token_ids: list[int]
+        self, request: Request, new_token_ids: list[int], is_stale: bool = False
    ) -> tuple[list[int], bool]:
+        # is_stale is only used by the AsyncScheduler override.
        # Append generated tokens and check for stop. Note that if
        # a request is still being prefilled, we expect the model runner
        # to return empty token ids for the request.
@@ -2388,15 +2423,7 @@ class Scheduler(SchedulerInterface):
            # running queue in FIFO order.
            while self.running:
                request = self.running.pop()
-                self._preempt_request(request, timestamp)
-                # For async scheduling, any output frames already in flight at
-                # preemption time are now stale and must be discarded when they
-                # return. num_output_placeholders is exactly that count: 0 if
-                # the engine has drained (e.g. pause_generation(keep) waited
-                # for idle), 1 for vanilla async mid-step, or 1 + spec/PP frames
-                # otherwise.
-                request.async_tokens_to_discard = request.num_output_placeholders
-                request.num_output_placeholders = 0
+                self._preempt_request(request, timestamp, drop_stale_output=True)

            # Clear scheduled request ids cache. Since we are forcing preemption
            # + resumption in the same step, we must act as if these requests were
@@ -149,7 +149,12 @@ class Request:

        # Used in async scheduling.
        self.num_output_placeholders = 0
-        self.async_tokens_to_discard = 0
+        # Tokens of output in flight when the request was preempted: delivered
+        # on return, but must not mutate the reset counters.
+        self.num_stale_output_tokens = 0
+        # Drop the stale output instead, for same-step preempt + resume
+        # (reset_prefix_cache).
+        self.drop_stale_output = False

        # Tokens of steps whose output is not yet processed (async scheduling
        # and PP run ahead of the GPU); `num_computed_tokens` counts them
@@ -6,7 +6,12 @@ import numpy as np
 import torch

 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
-from vllm.v1.outputs import AsyncModelRunnerOutput, LogprobsTensors, ModelRunnerOutput
+from vllm.v1.outputs import (
+    AsyncModelRunnerOutput,
+    LogprobsTensors,
+    ModelRunnerOutput,
+    PoolerOutput,
+)
 from vllm.v1.worker.gpu.sample.output import SamplerOutput


@@ -89,34 +94,42 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput):
    def __init__(
        self,
        model_runner_output: ModelRunnerOutput,
-        pooler_output: torch.Tensor,
-        is_valid: torch.Tensor | None,
+        pooler_output: PoolerOutput,
+        finished_mask: list[bool],
        main_stream: torch.cuda.Stream,
        copy_stream: torch.cuda.Stream,
    ):
        self.model_runner_output = model_runner_output
        self.pooler_output = pooler_output
-        self.is_valid = is_valid
        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
        self.copy_event = torch.cuda.Event(blocking=True)

        with stream(copy_stream, main_stream):
            copy_stream.wait_stream(main_stream)
-            self.pooler_output_cpu = self.pooler_output.to("cpu", non_blocking=True)
-            if self.is_valid is not None:
-                self.is_valid_cpu = self.is_valid.to("cpu", non_blocking=True)
+            if isinstance(self.pooler_output, torch.Tensor) and all(finished_mask):
+                self.pooler_output_cpu: PoolerOutput = self.pooler_output.to(
+                    "cpu", non_blocking=True
+                )
            else:
-                self.is_valid_cpu = None
+                outputs = (
+                    self.pooler_output.unbind()
+                    if isinstance(self.pooler_output, torch.Tensor)
+                    else self.pooler_output
+                )
+                self.pooler_output_cpu = [
+                    None
+                    if output is None or not is_finished
+                    else output.to("cpu", non_blocking=True)
+                    for output, is_finished in zip(outputs, finished_mask, strict=True)
+                ]
            self.copy_event.record(copy_stream)

    def get_output(self) -> ModelRunnerOutput:
-        pooler_output = list(self.pooler_output_cpu.unbind(dim=0))
+        if isinstance(self.pooler_output_cpu, torch.Tensor):
+            pooler_output = list(self.pooler_output_cpu.unbind(dim=0))
+        else:
+            pooler_output = self.pooler_output_cpu
        self.copy_event.synchronize()
-        if self.is_valid_cpu is not None:
-            is_valid_cpu = self.is_valid_cpu.tolist()
-            for i, is_valid in enumerate(is_valid_cpu):
-                if not is_valid:
-                    pooler_output[i] = None
        self.model_runner_output.pooler_output = pooler_output
        return self.model_runner_output

@@ -366,7 +366,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
            )

        if self.is_pooling_model and self.is_last_pp_rank:
-            self.pooling_runner = PoolingRunner(self.model)
+            self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
        eplb_models_added |= self.eplb.maybe_register_model(
            self.model,
            self.model_config,
@@ -787,6 +787,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
        req_idx = self.req_states.remove_request(req_id)
        if req_idx is None:
            return False
+        if self.pooling_runner is not None:
+            self.pooling_runner.remove_request(req_idx)
        if self.pp_handler is not None:
            self.pp_handler.on_req_idx_freed(req_idx)
        if self.encoder_cache is not None:
@@ -839,6 +841,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
            )
            req_index = self.req_states.req_id_to_index[req_id]

+            if self.pooling_runner is not None:
+                assert new_req_data.pooling_params is not None
+                self.pooling_runner.add_request(
+                    req_index,
+                    new_req_data.pooling_params,
+                    new_req_data.prompt_token_ids,
+                )
+
            if self.encoder_cache is not None:
                self.encoder_cache.add_request(req_id, new_req_data.mm_features)

@@ -1591,7 +1601,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
            return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)

        assert self.pooling_runner is not None
-        pooler_output, is_valid = self.pooling_runner.pool(
+        pooler_output, finished_mask = self.pooling_runner.pool(
            hidden_states, input_batch, self.req_states
        )

@@ -1604,7 +1614,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
        async_output = AsyncPoolingOutput(
            model_runner_output=model_runner_output,
            pooler_output=pooler_output,
-            is_valid=is_valid,
+            finished_mask=finished_mask,
            main_stream=self.main_stream,
            copy_stream=self.output_copy_stream,
        )
@@ -8,12 +8,13 @@ import torch.nn as nn
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.config.compilation import CUDAGraphMode
 from vllm.model_executor.layers.attention import Attention
-from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.attention.backend import (
    AttentionCGSupport,
    AttentionType,
    CommonAttentionMetadata,
 )
+from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.kv_cache_interface import (
    AttentionSpec,
    EncoderOnlyAttentionSpec,
@@ -22,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+from vllm.v1.worker.gpu.states import RequestState
 from vllm.v1.worker.utils import AttentionGroup


@@ -95,6 +97,52 @@ class EncoderOnlyModelState(DefaultModelState):
        self._dummy_slot_mapping = torch.zeros(
            self.max_num_tokens, dtype=torch.int64, device=device
        )
+        self.token_type_ids: dict[str, torch.Tensor] = {}
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        super().add_request(req_index, new_req_data)
+        pooling_params = new_req_data.pooling_params
+        if pooling_params is None or pooling_params.extra_kwargs is None:
+            return
+
+        token_type_start = pooling_params.extra_kwargs.get("compressed_token_type_ids")
+        if token_type_start is not None:
+            assert new_req_data.prompt_token_ids is not None
+            self.token_type_ids[new_req_data.req_id] = (
+                torch.arange(len(new_req_data.prompt_token_ids), dtype=torch.int32)
+                >= token_type_start
+            ).to(torch.int32)
+
+    def remove_request(self, req_id: str) -> None:
+        super().remove_request(req_id)
+        self.token_type_ids.pop(req_id, None)
+
+    def prepare_inputs(
+        self, input_batch: InputBatch, req_states: RequestState
+    ) -> dict[str, torch.Tensor | None]:
+        model_inputs = super().prepare_inputs(input_batch, req_states)
+        if not self.token_type_ids:
+            return model_inputs
+
+        token_type_ids_cpu = torch.zeros(
+            input_batch.num_tokens_after_padding,
+            dtype=torch.int32,
+            pin_memory=PIN_MEMORY,
+        )
+        offset = 0
+        for i, req_id in enumerate(input_batch.req_ids):
+            num_tokens = int(input_batch.num_scheduled_tokens[i])
+            request_token_type_ids = self.token_type_ids.get(req_id)
+            if request_token_type_ids is not None:
+                start = int(input_batch.num_computed_tokens_np[i])
+                token_type_ids_cpu[offset : offset + num_tokens].copy_(
+                    request_token_type_ids[start : start + num_tokens]
+                )
+            offset += num_tokens
+        model_inputs["token_type_ids"] = token_type_ids_cpu.to(
+            self.device, non_blocking=True
+        )
+        return model_inputs

    def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]:
        # Encoder groups are built here rather than in init_attn_backend, so
@@ -9,10 +9,6 @@ import torch.nn as nn

 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
-from vllm.model_executor.layers.mamba.mamba_utils import (
-    get_conv_copy_spec,
-    is_conv_state_dim_first,
-)
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
@@ -133,14 +129,11 @@ class MambaHybridModelState(DefaultModelState):
    ) -> MambaSpecDecodeGPUContext:
        if self._mamba_ctx is None:
            copy_funcs = self.model.get_mamba_state_copy_func()
-            # The fused copy kernels shift conv windows assuming the SD layout;
-            # the DS layout cannot express a >0 spec-decode shift as a single
-            # contiguous copy (mirrors get_conv_copy_spec's NotImplementedError).
-            if get_conv_copy_spec in copy_funcs and is_conv_state_dim_first():
-                assert self.vllm_config.speculative_config is None, (
-                    "DS conv state layout does not support mamba align state "
-                    "copies with speculative decoding"
-                )
+            # Both SD and DS conv layouts support a >0 spec-decode shift: the
+            # fused pre-copy kernel (``_copy_mamba_state_block``) applies the
+            # ``token_bias = num_accepted - 1`` window shift per conv layout
+            # (SD: contiguous slice; DS: per-dim-row strided slice), matching
+            # the V1 ``get_conv_copy_spec`` semantics.
            self._mamba_ctx = MambaSpecDecodeGPUContext.create(
                max_num_reqs=self.max_num_reqs,
                kv_cache_config=kv_cache_config,
@@ -2,45 +2,190 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import cast

+import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F

+from vllm.config import VllmConfig
 from vllm.model_executor.models import VllmModelForPooling, is_pooling_model
+from vllm.pooling_params import PoolingParams
 from vllm.tasks import PoolingTask
+from vllm.utils.torch_utils import PIN_MEMORY
+from vllm.v1.outputs import PoolerOutput
+from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.states import RequestState

+_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset({"embed", "classify"})
+

-# NOTE(woosuk): Currently, this class only supports the "LAST" pooling task
-# on decoder-only models. How to support other pooling tasks and models
-# is to be determined.
 class PoolingRunner:
-    def __init__(self, model: nn.Module):
+    def __init__(self, model: nn.Module, vllm_config: VllmConfig):
        self.model = cast(VllmModelForPooling, model)
+        self.model_config = vllm_config.model_config
+        self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
+        model_tasks = tuple(sorted(self.model.pooler.get_supported_tasks()))
+        selected_task = self.model_config.get_pooling_task(model_tasks)
+        if selected_task not in _SUPPORTED_TASKS:
+            raise ValueError(
+                "Model Runner V2 supports only sequence-level pooling tasks "
+                f"{sorted(_SUPPORTED_TASKS)}, but this model selects "
+                f"{selected_task!r} from {list(model_tasks)}. Set an explicitly "
+                "supported task or VLLM_USE_V2_MODEL_RUNNER=0."
+            )
+        self.supported_tasks = frozenset(self.get_supported_tasks(model))
+        if not self.supported_tasks:
+            raise ValueError(
+                "Model Runner V2 supports only sequence-level pooling tasks "
+                f"{sorted(_SUPPORTED_TASKS)}, but this model supports "
+                f"{list(model_tasks)}. "
+                "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."
+            )
+        self.pooling_params: dict[int, PoolingParams] = {}
+        self.pooling_states: dict[int, PoolingStates] = {}
+        self.prompt_token_ids: dict[int, torch.Tensor] = {}

    @staticmethod
    def get_supported_tasks(model: nn.Module) -> list[PoolingTask]:
        if not is_pooling_model(model):
            return []
-        assert "embed" in model.pooler.get_supported_tasks()
-        return ["embed"]
+        return sorted(model.pooler.get_supported_tasks() & _SUPPORTED_TASKS)
+
+    def add_request(
+        self,
+        req_index: int,
+        pooling_params: PoolingParams,
+        prompt_token_ids: list[int],
+    ) -> None:
+        task = pooling_params.task
+        if task not in self.supported_tasks:
+            raise ValueError(
+                f"Unsupported task: {task!r}. "
+                f"Supported tasks: {sorted(self.supported_tasks)}"
+            )
+        self.model.pooler.get_pooling_updates(task).apply(pooling_params)
+        self.pooling_params[req_index] = pooling_params
+        self.pooling_states[req_index] = PoolingStates()
+        if pooling_params.requires_token_ids:
+            self.prompt_token_ids[req_index] = torch.tensor(
+                prompt_token_ids, dtype=torch.int64
+            )
+
+    def remove_request(self, req_index: int) -> None:
+        self.pooling_params.pop(req_index, None)
+        if state := self.pooling_states.pop(req_index, None):
+            state.clean()
+        self.prompt_token_ids.pop(req_index, None)
+
+    def _get_pooling_metadata(
+        self,
+        input_batch: InputBatch,
+        req_states: RequestState,
+        device: torch.device,
+    ) -> PoolingMetadata:
+        req_indices = input_batch.idx_mapping_np.tolist()
+        pooling_params = [self.pooling_params[i] for i in req_indices]
+        pooling_states = [self.pooling_states[i] for i in req_indices]
+        prompt_lens = torch.from_numpy(
+            req_states.prompt_len.np[input_batch.idx_mapping_np].copy()
+        )
+
+        prompt_token_ids_cpu = None
+        prompt_token_ids = None
+        if any(params.requires_token_ids for params in pooling_params):
+            max_prompt_len = int(prompt_lens.max())
+            prompt_token_ids_cpu = torch.zeros(
+                (input_batch.num_reqs, max_prompt_len),
+                dtype=torch.int64,
+                pin_memory=PIN_MEMORY,
+            )
+            for i, (req_index, params) in enumerate(zip(req_indices, pooling_params)):
+                if not params.requires_token_ids:
+                    continue
+                token_ids = self.prompt_token_ids[req_index]
+                prompt_token_ids_cpu[i, : token_ids.numel()] = token_ids
+            prompt_token_ids = prompt_token_ids_cpu.to(device, non_blocking=True)
+
+        return PoolingMetadata(
+            prompt_lens=prompt_lens,
+            prompt_token_ids=prompt_token_ids,
+            prompt_token_ids_cpu=prompt_token_ids_cpu,
+            pooling_params=pooling_params,
+            pooling_states=pooling_states,
+        )

    def pool(
        self,
        hidden_states: torch.Tensor,
        input_batch: InputBatch,
        req_states: RequestState,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        # TODO(woosuk): Support different types of pooling tasks.
-        last_hidden_states = hidden_states[input_batch.logits_indices]
-        # TODO(woosuk): Make normalization optional.
-        last_hidden_states = F.normalize(last_hidden_states, p=2, dim=-1)
+    ) -> tuple[PoolerOutput, list[bool]]:
+        hidden_states = hidden_states[: input_batch.num_tokens]
+        pooling_metadata = self._get_pooling_metadata(
+            input_batch, req_states, hidden_states.device
+        )
+        num_reqs = input_batch.num_reqs
+        # Pooling has no speculative tokens, so this CPU upper bound is exact.
+        seq_lens_cpu = input_batch.seq_lens_cpu_upper_bound[:num_reqs]
+        pooling_metadata.build_pooling_cursor(
+            input_batch.num_scheduled_tokens,
+            seq_lens_cpu,
+            device=hidden_states.device,
+            query_start_loc_gpu=input_batch.query_start_loc[: num_reqs + 1],
+        )
+        pooler_output = self.model.pooler(hidden_states, pooling_metadata)

-        prompt_len = req_states.prompt_len.gpu[input_batch.idx_mapping]
-        is_valid = input_batch.seq_lens == prompt_len
-        return last_hidden_states, is_valid
+        finished_mask = pooling_metadata.get_pooling_cursor().is_finished().tolist()
+        return pooler_output, finished_mask
+
+    def _dummy_pooler_run_task(
+        self, hidden_states: torch.Tensor, task: PoolingTask
+    ) -> PoolerOutput:
+        num_tokens = hidden_states.shape[0]
+        num_reqs = min(num_tokens, self.max_num_reqs)
+        base_tokens = num_tokens // num_reqs
+        num_extra = num_tokens % num_reqs
+        num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32)
+        if num_extra > 0:
+            num_scheduled_tokens[-num_extra:] += 1
+        prompt_lens = torch.from_numpy(num_scheduled_tokens)
+
+        pooling_params = PoolingParams(task=task)
+        pooling_params.verify(self.model_config)
+        self.model.pooler.get_pooling_updates(task).apply(pooling_params)
+        prompt_token_ids = None
+        if pooling_params.requires_token_ids:
+            prompt_token_ids = torch.zeros(
+                (num_reqs, int(prompt_lens.max())),
+                dtype=torch.int64,
+                device=hidden_states.device,
+            )
+        pooling_metadata = PoolingMetadata(
+            prompt_lens=prompt_lens,
+            prompt_token_ids=prompt_token_ids,
+            prompt_token_ids_cpu=None
+            if prompt_token_ids is None
+            else prompt_token_ids.cpu(),
+            pooling_params=[pooling_params] * num_reqs,
+            pooling_states=[PoolingStates() for _ in range(num_reqs)],
+        )
+        pooling_metadata.build_pooling_cursor(
+            num_scheduled_tokens,
+            seq_lens_cpu=prompt_lens,
+            device=hidden_states.device,
+        )
+        try:
+            return self.model.pooler(hidden_states, pooling_metadata)
+        except RuntimeError as e:
+            if "out of memory" not in str(e):
+                raise
+            raise RuntimeError(
+                "CUDA out of memory occurred when warming up pooler "
+                f"({task=}) with {num_reqs} dummy requests. Please try "
+                "lowering `max_num_seqs` or `gpu_memory_utilization` when "
+                "initializing the engine."
+            ) from e

    def dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
-        F.normalize(hidden_states, p=2, dim=-1)
-        return
+        for task in sorted(self.supported_tasks):
+            self._dummy_pooler_run_task(hidden_states, task)
@@ -230,7 +230,11 @@ def warmup_kernels(
    # SamplingParams exercising all sampling features.
    if model_runner.is_pooling_model:
        sampling_params = None
-        pooling_params = PoolingParams()
+        pooling_task = model_runner.model_config.get_pooling_task(
+            model_runner.get_supported_tasks()
+        )
+        pooling_params = PoolingParams(task=pooling_task)
+        pooling_params.verify(model_runner.model_config)
    else:
        sampling_params = SamplingParams.for_sampler_warmup()
        pooling_params = None
@@ -4335,6 +4335,7 @@ class GPUModelRunner(
                    self.compilation_config.static_forward_context,
                    self.model.get_mamba_state_copy_func(),
                    mamba_bufs.preprocess,
+                    align_ctx=mamba_bufs.postprocess_align,
                )
                # preprocess_mamba resets num_accepted_tokens_cpu to 1
                # for requests whose state was copied to a new block.
@@ -3,7 +3,7 @@
 import dataclasses
 import itertools
 from collections.abc import Callable
-from typing import Any
+from typing import Any, NamedTuple

 import torch

@@ -348,8 +348,9 @@ def precopy_mamba_align_fused_kernel(
    num_reqs,
    COPY_BLOCK_SIZE: tl.constexpr,
    CONV_STATE_DIM_FIRST: tl.constexpr,
+    HAS_IDX_MAPPING: tl.constexpr = True,
 ):
-    """Pre-copy mamba "align" state across block boundaries on the V2 runner.
+    """Pre-copy mamba "align" state across block boundaries.

    Before the forward pass, copy each request's last SSM/conv state from its
    previous block column into the new window block column, so the kernels read
@@ -359,16 +360,20 @@ def precopy_mamba_align_fused_kernel(
    copy specs), but driven by the GPU-resident src columns so it needs no
    CPU-GPU sync (async-scheduling safe).

-    Grid: (num_reqs, num_layers * num_state_types); block tables are indexed by
-    batch row, per-request state by req_idx via idx_mapping (V2 layout).
+    Grid: (num_reqs, num_layers * num_state_types). V2 passes a batch-to-state
+    idx_mapping; V1 already stores the staged arrays in batch order and uses
+    HAS_IDX_MAPPING=False.
    """
    batch_idx = tl.program_id(0)
    state_idx = tl.program_id(1)
    if batch_idx >= num_reqs:
        return
-    req_idx = tl.load(idx_mapping_ptr + batch_idx)
-    if req_idx < 0:
-        return
+    if HAS_IDX_MAPPING:
+        req_idx = tl.load(idx_mapping_ptr + batch_idx)
+        if req_idx < 0:
+            return
+    else:
+        req_idx = batch_idx

    src_col = tl.load(src_col_ptr + req_idx)
    dst_col = tl.load(mamba_state_idx_ptr + req_idx)
@@ -527,6 +532,8 @@ class MambaSpecDecodeGPUContext:
    num_scheduled_tokens_buf: CpuGpuBuffer | None = None
    num_computed_tokens_buf: CpuGpuBuffer | None = None
    num_draft_tokens_buf: CpuGpuBuffer | None = None
+    precopy_src_col_buf: CpuGpuBuffer | None = None
+    precopy_token_bias_buf: CpuGpuBuffer | None = None

    # Flag to track if metadata has been populated
    is_initialized: bool = False
@@ -590,6 +597,8 @@ class MambaSpecDecodeGPUContext:
            num_scheduled_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            num_computed_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            num_draft_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
+            precopy_src_col_buf=make_buffer(max_num_reqs, dtype=torch.int32),
+            precopy_token_bias_buf=make_buffer(max_num_reqs, dtype=torch.int32),
            is_initialized=False,
        )

@@ -797,17 +806,18 @@ class MambaSpecDecodeGPUContext:
        state_idx_gpu: torch.Tensor,
        src_col_gpu: torch.Tensor,
        token_bias_gpu: torch.Tensor,
-        idx_mapping: torch.Tensor,
+        idx_mapping: torch.Tensor | None,
    ) -> None:
        """Pre-copy each request's previous running block into its new window
-        block before the forward pass (V2 align boundary migration).
+        block before the forward pass (align boundary migration).

        Args:
            num_reqs: Number of active requests (batch order).
            state_idx_gpu: [max_reqs] post-advance dst block column per req slot.
            src_col_gpu: [max_reqs] pre-advance src block column (-1 = fresh).
            token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1).
-            idx_mapping: [num_reqs] batch_idx -> req_state_idx (-1 to skip).
+            idx_mapping: optional [num_reqs] batch_idx -> req_state_idx.
+                None means V1 batch order already equals request state order.
        """
        if num_reqs == 0 or not self.is_initialized:
            return
@@ -831,6 +841,7 @@ class MambaSpecDecodeGPUContext:
            num_reqs,
            COPY_BLOCK_SIZE=1024,
            CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
+            HAS_IDX_MAPPING=idx_mapping is not None,
        )

    def run_fused_postprocess_align(
@@ -989,6 +1000,36 @@ def cleanup_mamba_state_idx(
        mamba_state_idx.pop(req_id, None)


+class _FusedPrecopy(NamedTuple):
+    """Resolved fused align pre-copy resources (all non-None once resolved)."""
+
+    ctx: "MambaSpecDecodeGPUContext"
+    state_idx: CpuGpuBuffer
+    src_col: CpuGpuBuffer
+    token_bias: CpuGpuBuffer
+
+
+def _resolve_fused_precopy(
+    align_ctx: "MambaSpecDecodeGPUContext | None",
+) -> _FusedPrecopy | None:
+    """Bundle the fused-path buffers, or None for the scalar path.
+
+    Returning one non-None bundle lets callers narrow all four members with a
+    single ``is not None`` check instead of re-asserting each buffer per use.
+    """
+    if align_ctx is None:
+        return None
+    assert align_ctx.mamba_state_idx_buf is not None
+    assert align_ctx.precopy_src_col_buf is not None
+    assert align_ctx.precopy_token_bias_buf is not None
+    return _FusedPrecopy(
+        align_ctx,
+        align_ctx.mamba_state_idx_buf,
+        align_ctx.precopy_src_col_buf,
+        align_ctx.precopy_token_bias_buf,
+    )
+
+
 def preprocess_mamba(
    scheduler_output: SchedulerOutput,
    kv_cache_config: KVCacheConfig,
@@ -999,11 +1040,13 @@ def preprocess_mamba(
    forward_context: dict[str, Any],
    mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
    copy_bufs: MambaCopyBuffers,
+    align_ctx: MambaSpecDecodeGPUContext | None = None,
 ):
    """
    Copy the mamba state of previous step to the last
    (1 + num_speculative_blocks) block.
    """
+    fused = _resolve_fused_precopy(align_ctx)
    mamba_group_ids = copy_bufs.mamba_group_ids
    mamba_spec = copy_bufs.mamba_spec
    num_speculative_blocks = mamba_spec.num_speculative_blocks
@@ -1013,20 +1056,37 @@ def preprocess_mamba(
    cleanup_mamba_state_idx(scheduler_output, mamba_state_idx)

    copy_bufs.offset = 0
+    num_reqs = len(input_batch.req_ids)
+
+    if fused is not None:
+        if num_reqs == 0:
+            return
+        if not fused.ctx.is_initialized:
+            fused.ctx.initialize_from_forward_context(
+                kv_cache_config,
+                forward_context,
+                mamba_state_copy_funcs,
+                [
+                    input_batch.block_table[gid].get_device_tensor(num_reqs)
+                    for gid in fused.ctx.mamba_group_ids
+                ],
+            )
+
+        fused.src_col.np[:num_reqs] = -1
+        fused.token_bias.np[:num_reqs] = 0
+
    for i, req_id in enumerate(input_batch.req_ids):
        req_state = requests[req_id]
        prev_state_idx = mamba_state_idx.get(req_id)
        if prev_state_idx is None:
-            # new / resumed request, no previous state
-            # if num_computed_tokens is 0, prev_state_idx will be -1
+            # New / resumed request; num_computed_tokens == 0 gives -1.
            prev_state_idx = (req_state.num_computed_tokens - 1) // block_size

        num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
-        num_blocks: int = (
+        num_blocks = (
            cdiv(req_state.num_computed_tokens + num_scheduled_tokens, block_size)
            + num_speculative_blocks
        )
-
        # We always save the current running state at the last
        # (1 + num_speculative_blocks) block.
        # A corner case worth mention here: assume we have block_size = 4 and
@@ -1039,20 +1099,42 @@ def preprocess_mamba(
        # And use block 1 to save the running state.
        curr_state_idx = num_blocks - 1 - num_speculative_blocks
        mamba_state_idx[req_id] = curr_state_idx
+        if fused is not None:
+            fused.state_idx.np[i] = curr_state_idx
+
        if prev_state_idx != -1 and prev_state_idx != curr_state_idx:
-            collect_mamba_copy_meta(
-                copy_bufs,
-                kv_cache_config,
-                mamba_state_copy_funcs,
-                mamba_group_ids,
-                prev_state_idx,
-                curr_state_idx,
-                input_batch.num_accepted_tokens_cpu[i] - 1,
-                req_state,
-                forward_context,
-            )
+            accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1
+            if fused is not None:
+                assert accept_token_bias >= 0
+                fused.src_col.np[i] = prev_state_idx
+                fused.token_bias.np[i] = accept_token_bias
+            else:
+                collect_mamba_copy_meta(
+                    copy_bufs,
+                    kv_cache_config,
+                    mamba_state_copy_funcs,
+                    mamba_group_ids,
+                    prev_state_idx,
+                    curr_state_idx,
+                    accept_token_bias,
+                    req_state,
+                    forward_context,
+                )
            input_batch.num_accepted_tokens_cpu[i] = 1
-    do_mamba_copy_block(copy_bufs)
+
+    if fused is not None:
+        fused.state_idx.copy_to_gpu(num_reqs)
+        fused.src_col.copy_to_gpu(num_reqs)
+        fused.token_bias.copy_to_gpu(num_reqs)
+        fused.ctx.run_fused_precopy(
+            num_reqs=num_reqs,
+            state_idx_gpu=fused.state_idx.gpu,
+            src_col_gpu=fused.src_col.gpu,
+            token_bias_gpu=fused.token_bias.gpu,
+            idx_mapping=None,
+        )
+    else:
+        do_mamba_copy_block(copy_bufs)


 def postprocess_mamba_all(
Author	SHA1	Message	Date
stefankoncarevicandGitHub	381b691620	[ROCm][CI] Fix Kimi K3 KDA on ROCm (#50262 ) Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>	2026-07-29 15:47:50 +00:00
Michael Goin GitHub OpenAI Codex	d6247d7173	[Spec Decode][Perf] Replicate DSpark Markov head across TP ranks (#49731 ) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-07-29 11:43:51 -04:00
Nick Hill GitHub Chris Eastwood Yifan Qiao	a0c092ee72	[BugFix] Fix `num_output_placeholders` preemption underflow (#48245 ) Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev> Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Chris Eastwood <chris.eastwood@pwn4g3.dev> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>	2026-07-29 06:41:54 -07:00
Taneem IbrahimandGitHub	43eaefba5a	[ModelRunner V2] Enable sequence pooling for embedding and classification models (#48791 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>	2026-07-29 06:30:45 -07:00
christianandGitHub	e0cfa52d22	[Bugfix][Frontend] Return transcription and translation verbose as float (#49073 ) Signed-off-by: Lucas Christian <lucaschgf7@gmail.com>	2026-07-29 13:16:13 +00:00
Reid GitHub Bugen Zhao	242c591d5a	[Rust Frontend] Send multimodal tensors in auxiliary frames (#49341 ) Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-29 12:21:09 +00:00
stefankoncarevicandGitHub	625871b52c	[CI][Test] Fix pooling truncation test after VLLMError hierarchy change (#50241 ) Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>	2026-07-29 12:15:56 +00:00
Yintong Lu GitHub Yongqi Wang Kunshang Ji	72297d859b	[XPU] Route weightless RMSNorm to _C dispatch (#47121 ) Signed-off-by: Yintong Lu <yintong.lu@intel.com> Co-authored-by: Yongqi Wang <yongqi.wang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-07-29 18:59:58 +08:00
sungsoo ha GitHub Claude Opus 4.8 Jiangyun Zhu	f51193b9ae	[Kernel][Mamba] Fused-kernel support for align-mode DS-conv state migration with num_accepted_tokens > 1 (#49291 ) Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>	2026-07-29 18:54:52 +08:00
Guan-Ming ChiuandGitHub	aeaa50a71c	[Bugfix][Multimodal] Include media IO config in MM cache hash (#49975 ) Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>	2026-07-29 10:48:40 +00:00