[mm] refactor mm code for rust tokenizer manager (#34660)

Co-authored-by: Rain Jiang <rain-jiang@outlook.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

rust/sglang-mm/src/qwen_vl/mod.rs

@@ -344,31 +344,36 @@ pub fn mrope_image_only(
 
 /// The qwen scheduler-drain shape, extracted from the generic driver
 /// [`Output`](crate::driver::Output). Shared by `sglang-server`'s MM worker
-/// and the parity binding so the mapping can't drift; replaced by a generic
-/// named-tensor handoff once a second family needs a different shape.
-pub struct QwenDrain {
+/// and the parity binding so the mapping can't drift. TODO(mm-families):
+/// replace with a generic named-tensor handoff once a second family needs a
+/// different shape.
+pub struct QwenPackedOutput {
     pub input_ids: Vec<i32>,
-    /// All items' `pixel_values`, concatenated in prompt order.
+    /// All items' `pixel_values`, concatenated in prompt order; flattened
+    /// `[Σ t·h·w, 3·temporal_patch_size·patch_size²]`.
     pub features: Vec<f32>,
+    /// Per item `[t, h, w]` patch grid.
     pub grids: Vec<[u32; 3]>,
     pub hashes: Vec<u64>,
+    /// Per item inclusive token range in `input_ids`.
     pub offsets: Vec<(u32, u32)>,
+    /// Flattened row-major `[3, input_len]` M-RoPE positions.
     pub mrope: Vec<i64>,
     pub mrope_delta: i64,
 }
 
-pub fn pack_drain(output: crate::driver::Output) -> Result<QwenDrain, String> {
+pub fn pack_output(output: crate::driver::Output) -> Result<QwenPackedOutput, String> {
     use crate::pipeline::PositionOutput;
 
     let PositionOutput::MRope { positions, delta } = output.positions else {
-        return Err("qwen_vl drain: expected M-RoPE positions".into());
+        return Err("qwen_vl pack: expected M-RoPE positions".into());
     };
     let mut features = Vec::new();
     let mut grids = Vec::with_capacity(output.items.len());
     let mut hashes = Vec::with_capacity(output.items.len());
     for item in output.items {
         let TensorData::F32(pixel_values) = item.feature.data else {
-            return Err("qwen_vl drain: expected f32 feature".into());
+            return Err("qwen_vl pack: expected f32 feature".into());
         };
         features.extend(pixel_values);
         let grid = item
@@ -378,11 +383,11 @@ pub fn pack_drain(output: crate::driver::Output) -> Result<QwenDrain, String> {
                 ("image_grid_thw", TensorData::I64(v)) => Some(v),
                 _ => None,
             })
-            .ok_or("qwen_vl drain: missing image_grid_thw")?;
+            .ok_or("qwen_vl pack: missing image_grid_thw")?;
         grids.push([grid[0] as u32, grid[1] as u32, grid[2] as u32]);
         hashes.push(item.hash);
     }
-    Ok(QwenDrain {
+    Ok(QwenPackedOutput {
         input_ids: output.input_ids,
         features,
         grids,
@@ -504,23 +509,27 @@ mod python {
             input_ids,
             images,
         };
-        let drain = py
+        let packed = py
             .detach(move || {
                 let family = crate::registry::pipeline_from_spec(&spec_json)?;
                 let output = crate::driver::process(family.as_ref(), input, |_| {
                     Err("native parity API requires input_ids".into())
                 })?;
-                pack_drain(output)
+                pack_output(output)
             })
             .map_err(PyValueError::new_err)?;
         Ok((
-            drain.input_ids,
-            drain.features.into_pyarray(py),
-            drain.grids.into_iter().map(|[t, h, w]| (t, h, w)).collect(),
-            drain.hashes,
-            drain.offsets,
-            drain.mrope.into_pyarray(py),
-            drain.mrope_delta,
+            packed.input_ids,
+            packed.features.into_pyarray(py),
+            packed
+                .grids
+                .into_iter()
+                .map(|[t, h, w]| (t, h, w))
+                .collect(),
+            packed.hashes,
+            packed.offsets,
+            packed.mrope.into_pyarray(py),
+            packed.mrope_delta,
         ))
     }