[Rust Frontend] Reduce copy in auxiliary frame resolution (#54148)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
Aug 29, 2026 at 02:01 UTC
11012d2a35090a3c624d0aeeec0f407026fc4b12
4 files changed
+69
-74
rust/src/engine-core-client/src/protocol/logprobs.rs
+3
-8
@@ -8,6 +8,7 @@ mod wire;
8
9
use std::ops::{Deref, DerefMut};
10
11
+use bytes::Bytes;
12
use enum_as_inner::EnumAsInner;
13
use serde::{Deserialize, Deserializer, Serialize};
14
@@ -162,10 +163,7 @@ impl Serialize for MaybeWireLogprobs {
163
impl MaybeWireLogprobs {
164
/// Resolve the wire representation into decoded logprobs by looking up aux
165
/// frames and decoding raw views as needed.
165
- pub(super) fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
166
- where
167
- Frame: AsRef<[u8]>,
168
- {
166
+ pub(super) fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result<Self> {
167
match self {
168
Self::Direct(value) => Ok(Self::Direct(value)),
169
Self::Wire(value) => value.resolve(frames, field_prefix).map(Self::Direct),
@@ -229,10 +227,7 @@ impl WireLogprobs {
227
/// Resolve the wire-format logprobs into semantic [`Logprobs`] records by
228
/// looking up aux frames, decoding raw views, and grouping each row
229
/// into one [`PositionLogprobs`].
232
- fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Logprobs>
233
- where
234
- Frame: AsRef<[u8]>,
235
- {
230
+ fn resolve(self, frames: &[Bytes], field_prefix: &str) -> Result<Logprobs> {
231
if let Some(indices) = self.cu_num_generated_tokens {
232
bail_ext_value_decode!(
233
"{field_prefix}.cu_num_generated_tokens: \
rust/src/engine-core-client/src/protocol/logprobs/array.rs
+24
-54
@@ -8,7 +8,7 @@ use bytes::Bytes;
8
use itertools::Itertools as _;
9
10
use crate::error::{Error, Result, ext_value_decode};
11
-use crate::protocol::tensor::{ShapeExt as _, WireArrayData, WireNdArray};
11
+use crate::protocol::tensor::{ShapeExt as _, WireNdArray};
12
13
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14
pub(super) enum ScalarType {
@@ -31,14 +31,11 @@ pub(super) struct DecodedArray2<T> {
31
pub data: Vec<T>,
32
}
33
34
-pub(super) fn decode_array2_u32<Frame>(
34
+pub(super) fn decode_array2_u32(
35
value: WireNdArray,
36
field: &str,
37
- frames: &[Frame],
38
-) -> Result<DecodedArray2<u32>>
39
-where
40
- Frame: AsRef<[u8]>,
41
-{
37
+ frames: &[Bytes],
38
+) -> Result<DecodedArray2<u32>> {
39
let (shape, bytes, scalar, endianness) =
40
decode_array_metadata(value, field, frames, &[ScalarType::I32, ScalarType::I64])?;
41
if shape.len() != 2 {
@@ -66,14 +63,11 @@ where
63
})
64
}
65
69
-pub(super) fn decode_array1_u32<Frame>(
66
+pub(super) fn decode_array1_u32(
67
value: WireNdArray,
68
field: &str,
72
- frames: &[Frame],
73
-) -> Result<Vec<u32>>
74
-where
75
- Frame: AsRef<[u8]>,
76
-{
69
+ frames: &[Bytes],
70
+) -> Result<Vec<u32>> {
71
let (shape, bytes, scalar, endianness) =
72
decode_array_metadata(value, field, frames, &[ScalarType::I32, ScalarType::I64])?;
73
if shape.len() != 1 {
@@ -97,14 +91,11 @@ where
91
Ok(data)
92
}
93
100
-pub(super) fn decode_array2_f32<Frame>(
94
+pub(super) fn decode_array2_f32(
95
value: WireNdArray,
96
field: &str,
103
- frames: &[Frame],
104
-) -> Result<DecodedArray2<f32>>
105
-where
106
- Frame: AsRef<[u8]>,
107
-{
97
+ frames: &[Bytes],
98
+) -> Result<DecodedArray2<f32>> {
99
let (shape, bytes, _, endianness) =
100
decode_array_metadata(value, field, frames, &[ScalarType::F32])?;
101
if shape.len() != 2 {
@@ -122,25 +113,29 @@ where
113
})
114
}
115
125
-pub(super) fn decode_array_metadata<Frame>(
116
+pub(super) fn decode_array_metadata(
117
value: WireNdArray,
118
field: &str,
128
- frames: &[Frame],
119
+ frames: &[Bytes],
120
expected_scalars: &[ScalarType],
130
-) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)>
131
-where
132
- Frame: AsRef<[u8]>,
133
-{
134
- let WireNdArray { dtype, shape, data } = value;
135
- let (scalar, endianness) = parse_dtype(&dtype, field)?;
121
+) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)> {
122
+ let mut value = value;
123
+ let (scalar, endianness) = parse_dtype(&value.dtype, field)?;
124
if !expected_scalars.contains(&scalar) {
125
return Err(decode_error(
126
field,
139
- &format!("expected dtype in {:?}, got {}", expected_scalars, dtype),
127
+ &format!(
128
+ "expected dtype in {:?}, got {}",
129
+ expected_scalars, value.dtype
130
+ ),
131
));
132
}
133
143
- let bytes = resolve_array_bytes(data, field, frames)?;
134
+ value
135
+ .resolve_aux_frame(frames)
136
+ .map_err(|message| decode_error(field, &message))?;
137
+ let WireNdArray { shape, data, .. } = value;
138
+ let bytes = data.into_raw_view().expect("auxiliary frame reference was resolved above");
139
validate_byte_length(shape.as_slice(), bytes.len(), field, scalar)?;
140
Ok((shape, bytes, scalar, endianness))
141
}
@@ -168,31 +163,6 @@ pub(super) fn parse_dtype(dtype: &str, field: &str) -> Result<(ScalarType, Endia
163
Ok((scalar, endianness))
164
}
165
171
-pub(super) fn resolve_array_bytes<Frame>(
172
- value: WireArrayData,
173
- field: &str,
174
- frames: &[Frame],
175
-) -> Result<Bytes>
176
-where
177
- Frame: AsRef<[u8]>,
178
-{
179
- match value {
180
- WireArrayData::RawView(bytes) => Ok(bytes),
181
- WireArrayData::AuxIndex(index) => {
182
- let frame = frames.get(index).ok_or_else(|| {
183
- decode_error(
184
- field,
185
- &format!(
186
- "aux frame index {index} out of range for {} frames",
187
- frames.len()
188
- ),
189
- )
190
- })?;
191
- Ok(Bytes::copy_from_slice(frame.as_ref()))
192
- }
193
- }
194
-}
195
-
166
pub(super) fn validate_byte_length(
167
shape: &[usize],
168
byte_len: usize,
rust/src/engine-core-client/src/protocol/output.rs
+4
-12
@@ -3,6 +3,7 @@
3
4
use std::collections::BTreeSet;
5
6
+use bytes::Bytes;
7
use enum_as_inner::EnumAsInner;
8
use serde::{Deserialize, Deserializer, Serialize, Serializer};
9
use serde_default::DefaultFromSerde;
@@ -140,10 +141,7 @@ impl EngineCoreOutput {
141
142
/// Resolve all wire-format fields in-place by looking up aux frames and
143
/// decoding raw-view payloads as needed.
143
- fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
144
- where
145
- Frame: AsRef<[u8]>,
146
- {
144
+ fn resolve_in_place(&mut self, frames: &[Bytes]) -> Result<()> {
145
self.new_logprobs = (self.new_logprobs.take())
146
.map(|value| value.resolve(frames, "new_logprobs"))
147
.transpose()?;
@@ -246,10 +244,7 @@ impl From<DpControlOutput> for EngineCoreOutputs {
244
impl EngineCoreOutputs {
245
/// Resolve all wire-format fields in-place by looking up aux frames and
246
/// decoding raw-view payloads as needed.
249
- fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
250
- where
251
- Frame: AsRef<[u8]>,
252
- {
247
+ fn resolve_in_place(&mut self, frames: &[Bytes]) -> Result<()> {
248
if let Self::RequestBatch(batch) = self {
249
for output in &mut batch.outputs {
250
output.resolve_in_place(frames)?;
@@ -365,10 +360,7 @@ impl<'de> Deserialize<'de> for EngineCoreOutputs {
360
361
/// Decode one ordinary or multipart engine-core output message into the strong
362
/// typed public protocol shape.
368
-pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
369
-where
370
- Frame: AsRef<[u8]>,
371
-{
363
+pub fn decode_engine_core_outputs(frames: &[Bytes]) -> Result<EngineCoreOutputs> {
364
let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
365
366
let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
rust/src/engine-core-client/src/protocol/tensor.rs
+38
@@ -166,6 +166,11 @@ impl WireNdArray {
166
pub(crate) fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
167
self.data.extract_aux_frame(aux_frames, threshold);
168
}
169
+
170
+ /// Resolve an auxiliary-frame reference into owned raw bytes.
171
+ pub(crate) fn resolve_aux_frame(&mut self, frames: &[Bytes]) -> Result<(), String> {
172
+ self.data.resolve_aux_frame(frames)
173
+ }
174
}
175
176
/// Validate that the number of elements implied by the shape matches the length
@@ -207,6 +212,22 @@ pub enum WireArrayData {
212
}
213
214
impl WireArrayData {
215
+ /// Resolve an auxiliary-frame reference into owned raw bytes.
216
+ pub(crate) fn resolve_aux_frame(&mut self, frames: &[Bytes]) -> Result<(), String> {
217
+ let Self::AuxIndex(index) = self else {
218
+ return Ok(());
219
+ };
220
+ let index = *index;
221
+ let frame = frames.get(index).ok_or_else(|| {
222
+ format!(
223
+ "auxiliary frame index {index} is out of range for {} frames",
224
+ frames.len()
225
+ )
226
+ })?;
227
+ *self = Self::RawView(frame.clone());
228
+ Ok(())
229
+ }
230
+
231
/// Replace a sufficiently large raw view with its one-based auxiliary-frame index.
232
fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
233
let Self::RawView(bytes) = self else {
@@ -344,4 +365,21 @@ mod tests {
365
let err = WireNdArray::from_f32(vec![2, 2], vec![1.0, 2.0]).unwrap_err();
366
assert!(err.contains("does not match shape"));
367
}
368
+
369
+ #[test]
370
+ fn resolve_aux_frame_shares_bytes_using_one_based_message_index() {
371
+ let mut tensor = WireNdArray {
372
+ dtype: "float32".to_string(),
373
+ shape: vec![2],
374
+ data: WireArrayData::AuxIndex(1),
375
+ };
376
+ let bytes = Bytes::from_static(b"payload");
377
+ let bytes_ptr = bytes.as_ptr();
378
+
379
+ tensor.resolve_aux_frame(&[Bytes::new(), bytes]).unwrap();
380
+
381
+ let resolved = tensor.data.as_raw_view().unwrap();
382
+ assert_eq!(resolved.as_ptr(), bytes_ptr);
383
+ assert_eq!(resolved.as_ref(), b"payload");
384
+ }
385
}