iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Rust Frontend] Reduce copy in auxiliary frame resolution (#54148)

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-08-28 09:01:17 -07:00
committed by GitHub
parent 94a54f581e
commit 11012d2a35
4 changed files with 69 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/engine-core-client/src/protocol/logprobs.rs
+3 -8
View File
@@ -8,6 +8,7 @@ mod wire;
use std::ops::{Deref, DerefMut};
use bytes::Bytes;
use enum_as_inner::EnumAsInner;
use serde::{Deserialize, Deserializer, Serialize};
@@ -162,10 +163,7 @@ impl Serialize for MaybeWireLogprobs {
impl MaybeWireLogprobs {
/// Resolve the wire representation into decoded logprobs by looking up aux
/// frames and decoding raw views as needed.
pub(super) fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
where
Frame: AsRef<[u8]>,
{
pub(super) fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result<Self> {
match self {
Self::Direct(value) => Ok(Self::Direct(value)),
Self::Wire(value) => value.resolve(frames, field_prefix).map(Self::Direct),
@@ -229,10 +227,7 @@ impl WireLogprobs {
/// Resolve the wire-format logprobs into semantic [`Logprobs`] records by
/// looking up aux frames, decoding raw views, and grouping each row
/// into one [`PositionLogprobs`].
fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Logprobs>
where
Frame: AsRef<[u8]>,
{
fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result<Logprobs> {
if let Some(indices) = self.cu_num_generated_tokens {
bail_ext_value_decode!(
"{field_prefix}.cu_num_generated_tokens: \
rust/src/engine-core-client/src/protocol/logprobs/array.rs
+24 -54
View File
@@ -8,7 +8,7 @@ use bytes::Bytes;
use itertools::Itertools as _;
use crate::error::{Error, Result, ext_value_decode};
use crate::protocol::tensor::{ShapeExt as _, WireArrayData, WireNdArray};
use crate::protocol::tensor::{ShapeExt as _, WireNdArray};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ScalarType {
@@ -31,14 +31,11 @@ pub(super) struct DecodedArray2<T> {
pub data: Vec<T>,
}
pub(super) fn decode_array2_u32<Frame>(
pub(super) fn decode_array2_u32(
value: WireNdArray,
field: &str,
frames: &[Frame],
) -> Result<DecodedArray2<u32>>
where
Frame: AsRef<[u8]>,
{
frames: &[Bytes],
) -> Result<DecodedArray2<u32>> {
let (shape, bytes, scalar, endianness) =
decode_array_metadata(value, field, frames, &[ScalarType::I32, ScalarType::I64])?;
if shape.len() != 2 {
@@ -66,14 +63,11 @@ where
})
}
pub(super) fn decode_array1_u32<Frame>(
pub(super) fn decode_array1_u32(
value: WireNdArray,
field: &str,
frames: &[Frame],
) -> Result<Vec<u32>>
where
Frame: AsRef<[u8]>,
{
frames: &[Bytes],
) -> Result<Vec<u32>> {
let (shape, bytes, scalar, endianness) =
decode_array_metadata(value, field, frames, &[ScalarType::I32, ScalarType::I64])?;
if shape.len() != 1 {
@@ -97,14 +91,11 @@ where
Ok(data)
}
pub(super) fn decode_array2_f32<Frame>(
pub(super) fn decode_array2_f32(
value: WireNdArray,
field: &str,
frames: &[Frame],
) -> Result<DecodedArray2<f32>>
where
Frame: AsRef<[u8]>,
{
frames: &[Bytes],
) -> Result<DecodedArray2<f32>> {
let (shape, bytes, _, endianness) =
decode_array_metadata(value, field, frames, &[ScalarType::F32])?;
if shape.len() != 2 {
@@ -122,25 +113,29 @@ where
})
}
pub(super) fn decode_array_metadata<Frame>(
pub(super) fn decode_array_metadata(
value: WireNdArray,
field: &str,
frames: &[Frame],
frames: &[Bytes],
expected_scalars: &[ScalarType],
) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)>
where
Frame: AsRef<[u8]>,
{
let WireNdArray { dtype, shape, data } = value;
let (scalar, endianness) = parse_dtype(&dtype, field)?;
) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)> {
let mut value = value;
let (scalar, endianness) = parse_dtype(&value.dtype, field)?;
if !expected_scalars.contains(&scalar) {
return Err(decode_error(
field,
&format!("expected dtype in {:?}, got {}", expected_scalars, dtype),
&format!(
"expected dtype in {:?}, got {}",
expected_scalars, value.dtype
),
));
}
let bytes = resolve_array_bytes(data, field, frames)?;
value
.resolve_aux_frame(frames)
.map_err(|message| decode_error(field, &message))?;
let WireNdArray { shape, data, .. } = value;
let bytes = data.into_raw_view().expect("auxiliary frame reference was resolved above");
validate_byte_length(shape.as_slice(), bytes.len(), field, scalar)?;
Ok((shape, bytes, scalar, endianness))
}
@@ -168,31 +163,6 @@ pub(super) fn parse_dtype(dtype: &str, field: &str) -> Result<(ScalarType, Endia
Ok((scalar, endianness))
}
pub(super) fn resolve_array_bytes<Frame>(
value: WireArrayData,
field: &str,
frames: &[Frame],
) -> Result<Bytes>
where
Frame: AsRef<[u8]>,
{
match value {
WireArrayData::RawView(bytes) => Ok(bytes),
WireArrayData::AuxIndex(index) => {
let frame = frames.get(index).ok_or_else(|| {
decode_error(
field,
&format!(
"aux frame index {index} out of range for {} frames",
frames.len()
),
)
})?;
Ok(Bytes::copy_from_slice(frame.as_ref()))
}
}
}
pub(super) fn validate_byte_length(
shape: &[usize],
byte_len: usize,
rust/src/engine-core-client/src/protocol/output.rs
+4 -12
View File
@@ -3,6 +3,7 @@
use std::collections::BTreeSet;
use bytes::Bytes;
use enum_as_inner::EnumAsInner;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_default::DefaultFromSerde;
@@ -140,10 +141,7 @@ impl EngineCoreOutput {
/// Resolve all wire-format fields in-place by looking up aux frames and
/// decoding raw-view payloads as needed.
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
where
Frame: AsRef<[u8]>,
{
fn resolve_in_place(&mut self, frames: &[Bytes]) -> Result<()> {
self.new_logprobs = (self.new_logprobs.take())
.map(|value| value.resolve(frames, "new_logprobs"))
.transpose()?;
@@ -246,10 +244,7 @@ impl From<DpControlOutput> for EngineCoreOutputs {
impl EngineCoreOutputs {
/// Resolve all wire-format fields in-place by looking up aux frames and
/// decoding raw-view payloads as needed.
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
where
Frame: AsRef<[u8]>,
{
fn resolve_in_place(&mut self, frames: &[Bytes]) -> Result<()> {
if let Self::RequestBatch(batch) = self {
for output in &mut batch.outputs {
output.resolve_in_place(frames)?;
@@ -365,10 +360,7 @@ impl<'de> Deserialize<'de> for EngineCoreOutputs {
/// Decode one ordinary or multipart engine-core output message into the strong
/// typed public protocol shape.
pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
where
Frame: AsRef<[u8]>,
{
pub fn decode_engine_core_outputs(frames: &[Bytes]) -> Result<EngineCoreOutputs> {
let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
rust/src/engine-core-client/src/protocol/tensor.rs
+38
View File
@@ -166,6 +166,11 @@ impl WireNdArray {
pub(crate) fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
self.data.extract_aux_frame(aux_frames, threshold);
}
/// Resolve an auxiliary-frame reference into owned raw bytes.
pub(crate) fn resolve_aux_frame(&mut self, frames: &[Bytes]) -> Result<(), String> {
self.data.resolve_aux_frame(frames)
}
}
/// Validate that the number of elements implied by the shape matches the length
@@ -207,6 +212,22 @@ pub enum WireArrayData {
}
impl WireArrayData {
/// Resolve an auxiliary-frame reference into owned raw bytes.
pub(crate) fn resolve_aux_frame(&mut self, frames: &[Bytes]) -> Result<(), String> {
let Self::AuxIndex(index) = self else {
return Ok(());
};
let index = *index;
let frame = frames.get(index).ok_or_else(|| {
format!(
"auxiliary frame index {index} is out of range for {} frames",
frames.len()
)
})?;
*self = Self::RawView(frame.clone());
Ok(())
}
/// 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 {
@@ -344,4 +365,21 @@ mod tests {
let err = WireNdArray::from_f32(vec![2, 2], vec![1.0, 2.0]).unwrap_err();
assert!(err.contains("does not match shape"));
}
#[test]
fn resolve_aux_frame_shares_bytes_using_one_based_message_index() {
let mut tensor = WireNdArray {
dtype: "float32".to_string(),
shape: vec![2],
data: WireArrayData::AuxIndex(1),
};
let bytes = Bytes::from_static(b"payload");
let bytes_ptr = bytes.as_ptr();
tensor.resolve_aux_frame(&[Bytes::new(), bytes]).unwrap();
let resolved = tensor.data.as_raw_view().unwrap();
assert_eq!(resolved.as_ptr(), bytes_ptr);
assert_eq!(resolved.as_ref(), b"payload");
}
}