minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

create rust workspace (#32014)

Browse Source
This commit is contained in:
Rain Jiang
2026-07-23 12:02:41 -07:00
committed by GitHub
parent d0b9689805
commit 7fe82dd02e
25 changed files with 550 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-mm/.gitignore
-1
View File
@@ -1 +0,0 @@
/target
rust/sglang-mm/Cargo.toml
+21 -13
View File
@@ -1,23 +1,31 @@
[package]
name = "sglang-mm"
version = "0.1.0"
edition = "2024"
description = "Rust-accelerated multimodal preprocessing for SGLang"
license = "Apache-2.0"
version.workspace = true
edition.workspace = true
license.workspace = true
# Consumed by python/setup.py: registers this crate as a PyO3 extension module
# of the main sglang wheel at the given import path. debug = false keeps
# editable installs on release builds (image preprocessing is perf-sensitive).
[package.metadata.sglang]
python-module = "sglang.srt.multimodal._core"
debug = false
[lib]
name = "_core"
# Unique per-crate artifact name (the shared workspace target/ dir cannot hold
# two `lib_core` cdylibs). The importable Python module is still `_core`: the
# name comes from the `#[pymodule]` entry point and the setuptools-rust
# `target` in python/pyproject.toml, which renames the built artifact.
name = "sglang_mm_core"
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.23", features = ["extension-module"] }
numpy = "0.23"
rayon = "1.10"
pyo3 = { workspace = true }
base64 = "0.22"
blake3 = "1"
half = "2.4"
image = { version = "0.25", default-features = false, features = ["jpeg", "png"] }
blake3 = "1"
base64 = "0.22"
[profile.release]
lto = true
codegen-units = 1
numpy = "0.29"
rayon = "1.10"
rust/sglang-mm/pyproject.toml
+2 -2
View File
@@ -4,9 +4,9 @@ build-backend = "maturin"
[project]
name = "sglang-mm"
version = "0.1.0"
dynamic = ["version"]
description = "Rust-accelerated multimodal preprocessing for SGLang"
requires-python = ">=3.9"
requires-python = ">=3.10"
[tool.maturin]
module-name = "_core"
rust/sglang-mm/src/common/mod.rs
+9 -11
View File
@@ -7,7 +7,6 @@ use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
pub fn pool() -> &'static rayon::ThreadPool {
static POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
POOL.get_or_init(|| {
@@ -72,10 +71,9 @@ pub fn resize_rgb<'py>(
.as_slice()
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
.to_vec();
let out = py.allow_threads(move || {
pool().install(|| resize::resize_lanczos_rgb(&data, h, w, out_h, out_w))
});
Ok(out.into_pyarray_bound(py))
let out =
py.detach(move || pool().install(|| resize::resize_lanczos_rgb(&data, h, w, out_h, out_w)));
Ok(out.into_pyarray(py))
}
#[pyfunction]
@@ -95,14 +93,14 @@ pub fn image_decode_rgb<'py>(
data: Vec<u8>,
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u8>>)> {
let (rgb, h, w) = py
.allow_threads(move || decode_rgb(&data))
.detach(move || decode_rgb(&data))
.map_err(PyValueError::new_err)?;
Ok((h, w, rgb.into_pyarray_bound(py)))
Ok((h, w, rgb.into_pyarray(py)))
}
#[pyfunction]
pub fn data_hash(py: Python<'_>, data: Vec<u8>) -> u64 {
py.allow_threads(move || {
py.detach(move || {
let digest = blake3::hash(&data);
u64::from_be_bytes(digest.as_bytes()[..8].try_into().unwrap())
})
@@ -115,17 +113,17 @@ pub fn base64_decode<'py>(
) -> PyResult<Bound<'py, pyo3::types::PyBytes>> {
use base64::Engine;
let decoded = py
.allow_threads(|| {
.detach(|| {
base64::engine::general_purpose::STANDARD
.decode(encoded)
.map_err(|e| format!("base64 decode error: {e}"))
})
.map_err(PyValueError::new_err)?;
Ok(pyo3::types::PyBytes::new_bound(py, &decoded))
Ok(pyo3::types::PyBytes::new(py, &decoded))
}
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
let m = PyModule::new_bound(parent.py(), "common")?;
let m = PyModule::new(parent.py(), "common")?;
m.add_function(wrap_pyfunction!(resize_rgb, &m)?)?;
m.add_function(wrap_pyfunction!(scaled_dims, &m)?)?;
m.add_function(wrap_pyfunction!(image_decode_rgb, &m)?)?;
rust/sglang-mm/src/common/resize.rs
+22 -35
View File
@@ -46,14 +46,14 @@ fn precompute_coeffs(in_size: usize, out_size: usize) -> Coeffs {
let count = (xmax - xmin) as usize;
let k = &mut kkf[xx * ksize..(xx + 1) * ksize];
let mut ww = 0.0f64;
for x in 0..count {
for (x, kv) in k[..count].iter_mut().enumerate() {
let w = lanczos((x as f64 + xmin as f64 - center + 0.5) * ss);
k[x] = w;
*kv = w;
ww += w;
}
if ww != 0.0 {
for x in 0..count {
k[x] /= ww;
for kv in k[..count].iter_mut() {
*kv /= ww;
}
}
bounds[xx] = (xmin as usize, count);
@@ -111,35 +111,27 @@ fn resample_horizontal(src: &[u8], h: usize, w: usize, out_w: usize, c: &Coeffs)
fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec<u8> {
let mut out = vec![0u8; out_h * w * 3];
out.par_chunks_mut(w * 3)
.enumerate()
.for_each(|(yy, row)| {
let (ymin, count) = c.bounds[yy];
let k = &c.kk[yy * c.ksize..yy * c.ksize + count];
for x in 0..w {
let mut s = [1i32 << (PRECISION_BITS - 1); 3];
for (y, &coef) in k.iter().enumerate() {
let p = ((ymin + y) * w + x) * 3;
s[0] += src[p] as i32 * coef;
s[1] += src[p + 1] as i32 * coef;
s[2] += src[p + 2] as i32 * coef;
}
let o = x * 3;
row[o] = clip8(s[0]);
row[o + 1] = clip8(s[1]);
row[o + 2] = clip8(s[2]);
out.par_chunks_mut(w * 3).enumerate().for_each(|(yy, row)| {
let (ymin, count) = c.bounds[yy];
let k = &c.kk[yy * c.ksize..yy * c.ksize + count];
for x in 0..w {
let mut s = [1i32 << (PRECISION_BITS - 1); 3];
for (y, &coef) in k.iter().enumerate() {
let p = ((ymin + y) * w + x) * 3;
s[0] += src[p] as i32 * coef;
s[1] += src[p + 1] as i32 * coef;
s[2] += src[p + 2] as i32 * coef;
}
});
let o = x * 3;
row[o] = clip8(s[0]);
row[o + 1] = clip8(s[1]);
row[o + 2] = clip8(s[2]);
}
});
out
}
pub fn resize_lanczos_rgb(
src: &[u8],
h: usize,
w: usize,
out_h: usize,
out_w: usize,
) -> Vec<u8> {
pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec<u8> {
let need_h = out_w != w;
let need_v = out_h != h;
if need_h && need_v {
@@ -158,12 +150,7 @@ pub fn resize_lanczos_rgb(
}
}
pub fn scaled_dims(
w: usize,
h: usize,
frac: Option<f64>,
cap: Option<i64>,
) -> (usize, usize) {
pub fn scaled_dims(w: usize, h: usize, frac: Option<f64>, cap: Option<i64>) -> (usize, usize) {
let Some(frac) = frac else {
return (w, h);
};
rust/sglang-mm/src/common/transforms.rs
+7 -3
View File
@@ -2,6 +2,10 @@
//!
//! Model-specific processors compose these to build their preprocessing
//! pipelines. All functions operate on flat RGB byte arrays (HWC layout).
//!
//! Not every primitive is wired into a compiled-in processor yet; they are
//! kept available for upcoming model integrations.
#![allow(dead_code)]
/// Normalize u8 RGB pixels to f32 in a single pass: `(pixel/255 - mean) / std`.
///
@@ -37,8 +41,8 @@ pub fn pad_to_grid(
grid_w: usize,
pad_value: &[f32],
) -> (Vec<f32>, usize, usize) {
let new_h = ((h + grid_h - 1) / grid_h) * grid_h;
let new_w = ((w + grid_w - 1) / grid_w) * grid_w;
let new_h = h.div_ceil(grid_h) * grid_h;
let new_w = w.div_ceil(grid_w) * grid_w;
let mut out = vec![0.0f32; new_h * new_w * channels];
// Fill with pad value
for i in 0..new_h * new_w {
@@ -89,5 +93,5 @@ pub fn extract_patches_hwc(
/// Compute the patch grid dimensions for a given image size and patch size.
#[inline]
pub fn patch_grid(h: usize, w: usize, patch_h: usize, patch_w: usize) -> (usize, usize) {
((h + patch_h - 1) / patch_h, (w + patch_w - 1) / patch_w)
(h.div_ceil(patch_h), w.div_ceil(patch_w))
}
rust/sglang-mm/src/inkling/mod.rs
+52 -40
View File
@@ -8,6 +8,15 @@ use rayon::prelude::*;
use crate::common;
/// `(height, width, patches_as_u16_bits)` for one decoded image.
type Patches = (usize, usize, Vec<u16>);
/// [`Patches`] plus the image content hash.
type HashedPatches = (usize, usize, Vec<u16>, u64);
/// [`Patches`] with the patch data as a numpy array bound to `'py`.
type PyPatches<'py> = (usize, usize, Bound<'py, PyArray1<u16>>);
/// [`PyPatches`] plus the image content hash.
type PyHashedPatches<'py> = (usize, usize, Bound<'py, PyArray1<u16>>, u64);
const MEAN: [f32; 3] = [
0.48145466f64 as f32,
0.4578275f64 as f32,
@@ -40,7 +49,7 @@ fn luts() -> &'static [[u16; 256]; 3] {
#[inline]
pub fn grid(h: usize, w: usize, ps: usize) -> (usize, usize) {
((h + ps - 1) / ps, w / ps + 1)
(h.div_ceil(ps), w / ps + 1)
}
fn patchify_into(arr: &[u8], h: usize, w: usize, ps: usize, out: &mut [u16]) {
@@ -95,7 +104,9 @@ fn patchify_alloc(arr: &[u8], h: usize, w: usize, ps: usize) -> Vec<u16> {
fn check_ps(ps: usize) -> PyResult<()> {
if ps == 0 {
return Err(PyValueError::new_err("patch_size must be greater than zero"));
return Err(PyValueError::new_err(
"patch_size must be greater than zero",
));
}
Ok(())
}
@@ -118,8 +129,8 @@ fn patchify_rgb<'py>(
.as_slice()
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
.to_vec();
let out = py.allow_threads(move || patchify_alloc(&data, h, w, patch_size));
Ok(out.into_pyarray_bound(py))
let out = py.detach(move || patchify_alloc(&data, h, w, patch_size));
Ok(out.into_pyarray(py))
}
#[pyfunction]
@@ -133,14 +144,14 @@ fn decode_patchify<'py>(
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u16>>)> {
check_ps(patch_size)?;
let (h, w, out) = py
.allow_threads(move || {
.detach(move || {
common::pool().install(|| {
let (rgb, h, w) = common::decode_rescale(&data, rescale_frac, rescale_cap)?;
Ok::<_, String>((h, w, patchify_alloc(&rgb, h, w, patch_size)))
})
})
.map_err(PyValueError::new_err)?;
Ok((h, w, out.into_pyarray_bound(py)))
Ok((h, w, out.into_pyarray(py)))
}
#[pyfunction]
@@ -151,25 +162,24 @@ fn decode_patchify_batch<'py>(
patch_size: usize,
rescale_frac: Option<f64>,
rescale_cap: Option<i64>,
) -> PyResult<Vec<(usize, usize, Bound<'py, PyArray1<u16>>)>> {
) -> PyResult<Vec<PyPatches<'py>>> {
check_ps(patch_size)?;
let results: Vec<Result<(usize, usize, Vec<u16>), String>> =
py.allow_threads(move || {
common::pool().install(|| {
datas
.par_iter()
.map(|data| {
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size)))
})
.collect()
})
});
let results: Vec<Result<Patches, String>> = py.detach(move || {
common::pool().install(|| {
datas
.par_iter()
.map(|data| {
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size)))
})
.collect()
})
});
results
.into_iter()
.map(|r| {
let (h, w, v) = r.map_err(PyValueError::new_err)?;
Ok((h, w, v.into_pyarray_bound(py)))
Ok((h, w, v.into_pyarray(py)))
})
.collect()
}
@@ -182,26 +192,25 @@ fn preprocess_images<'py>(
patch_size: usize,
rescale_frac: Option<f64>,
rescale_cap: Option<i64>,
) -> PyResult<Vec<(usize, usize, Bound<'py, PyArray1<u16>>, u64)>> {
) -> PyResult<Vec<PyHashedPatches<'py>>> {
check_ps(patch_size)?;
let results: Vec<Result<(usize, usize, Vec<u16>, u64), String>> =
py.allow_threads(move || {
common::pool().install(|| {
datas
.par_iter()
.map(|data| {
let hash = common::sha256_u64(data);
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash))
})
.collect()
})
});
let results: Vec<Result<HashedPatches, String>> = py.detach(move || {
common::pool().install(|| {
datas
.par_iter()
.map(|data| {
let hash = common::sha256_u64(data);
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash))
})
.collect()
})
});
results
.into_iter()
.map(|r| {
let (h, w, v, hash) = r.map_err(PyValueError::new_err)?;
Ok((h, w, v.into_pyarray_bound(py), hash))
Ok((h, w, v.into_pyarray(py), hash))
})
.collect()
}
@@ -237,7 +246,6 @@ impl crate::registry::ImageProcessorSpec for InklingProcessor {
}
}
#[pyfunction]
#[pyo3(signature = (arr, raw_bytes, patch_size, rescale_frac=None, rescale_cap=None))]
fn rescale_patchify_hash<'py>(
@@ -261,22 +269,26 @@ fn rescale_patchify_hash<'py>(
.as_slice()
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
.to_vec();
let (oh, ow, out) = py.allow_threads(move || {
let (oh, ow, out) = py.detach(move || {
common::pool().install(|| {
let (tw, th) = common::resize::scaled_dims(w, h, rescale_frac, rescale_cap);
let (rgb, h, w) = if (tw, th) != (w, h) {
(common::resize::resize_lanczos_rgb(&rgb, h, w, th, tw), th, tw)
(
common::resize::resize_lanczos_rgb(&rgb, h, w, th, tw),
th,
tw,
)
} else {
(rgb, h, w)
};
(h, w, patchify_alloc(&rgb, h, w, patch_size))
})
});
Ok((oh, ow, out.into_pyarray_bound(py), hash))
Ok((oh, ow, out.into_pyarray(py), hash))
}
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
let m = PyModule::new_bound(parent.py(), "inkling")?;
let m = PyModule::new(parent.py(), "inkling")?;
m.add_function(wrap_pyfunction!(patchify_rgb, &m)?)?;
m.add_function(wrap_pyfunction!(decode_patchify, &m)?)?;
m.add_function(wrap_pyfunction!(decode_patchify_batch, &m)?)?;
rust/sglang-mm/src/registry.rs
+13 -6
View File
@@ -3,8 +3,8 @@
//! Each model implements `ImageProcessorSpec` and registers itself. The Python
//! layer looks up a processor by model name at init time.
use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;
/// `(height, width, patches_as_u16_bits, content_hash)` for one image.
pub type PreprocessedImage = (usize, usize, Vec<u16>, u64);
/// Trait that each model's image processor must implement.
pub trait ImageProcessorSpec: Send + Sync {
@@ -12,15 +12,13 @@ pub trait ImageProcessorSpec: Send + Sync {
fn name(&self) -> &'static str;
/// Process a batch of raw image bytes: decode + preprocess + hash.
///
/// Returns `(height, width, patches_as_u16_bits, content_hash)` per image.
fn preprocess_batch(
&self,
datas: &[Vec<u8>],
patch_size: usize,
rescale_frac: Option<f64>,
rescale_cap: Option<i64>,
) -> Result<Vec<(usize, usize, Vec<u16>, u64)>, String>;
) -> Result<Vec<PreprocessedImage>, String>;
}
/// Global registry of available processors.
@@ -28,6 +26,12 @@ pub struct ProcessorRegistry {
specs: Vec<Box<dyn ImageProcessorSpec>>,
}
impl Default for ProcessorRegistry {
fn default() -> Self {
Self::new()
}
}
impl ProcessorRegistry {
pub fn new() -> Self {
Self { specs: Vec::new() }
@@ -38,7 +42,10 @@ impl ProcessorRegistry {
}
pub fn lookup(&self, name: &str) -> Option<&dyn ImageProcessorSpec> {
self.specs.iter().find(|s| s.name() == name).map(|s| s.as_ref())
self.specs
.iter()
.find(|s| s.name() == name)
.map(|s| s.as_ref())
}
pub fn list_names(&self) -> Vec<&'static str> {