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/Cargo.toml
+33
View File
@@ -0,0 +1,33 @@
[workspace]
resolver = "3"
members = ["sglang-grpc", "sglang-mm"]
[workspace.package]
version = "0.1.0"
edition = "2024"
license = "Apache-2.0"
# Single source of truth for dependency versions across workspace members.
# Members pull these in with `dep = { workspace = true }` (plus extra features
# where needed).
[workspace.dependencies]
async-stream = "0.3"
pyo3 = { version = "0.29.0", features = ["extension-module"] }
serde_json = "1"
tokio = { version = "1", features = ["full"] }
tokio-stream = { version = "0.1", features = ["net"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
uuid = { version = "1", features = ["v4"] }
# Profiles only take effect at the workspace root; per-member profile sections
# are ignored by cargo. Per-package tweaks live in the override tables below.
[profile.release]
lto = true
strip = true
opt-level = 3
codegen-units = 1
[profile.dev]
opt-level = 0
debug = 1
rust/sglang-grpc/rust-toolchain.toml → rust/rust-toolchain.toml
+1
View File
@@ -1,3 +1,4 @@
[toolchain]
channel = "1.90"
profile = "minimal"
components = ["clippy", "rustfmt"]
rust/sglang-grpc/Cargo.toml
+26 -22
View File
@@ -1,38 +1,42 @@
[package]
name = "sglang-grpc"
version = "0.1.0"
edition = "2024"
description = "In-process Rust gRPC server 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.
[package.metadata.sglang]
python-module = "sglang.srt.grpc._core"
[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_grpc_core"
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.23", features = ["extension-module"] }
tokio = { version = "1", features = ["full"] }
tonic = { version = "0.12", features = ["gzip", "transport"] }
async-stream = { workspace = true }
pyo3 = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true }
tokio-stream = { workspace = true }
tracing = { workspace = true }
tracing-subscriber = { workspace = true }
uuid = { workspace = true }
prost = "0.13"
uuid = { version = "1", features = ["v4"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
serde_json = "1"
tokenizers = { version = "0.21", default-features = false, features = ["onig"] }
tokio-stream = { version = "0.1", features = ["net"] }
async-stream = "0.3"
tonic = { version = "0.12", features = ["gzip", "transport"] }
[build-dependencies]
tonic-build = "0.12"
# Fallback protoc for machines without a system install (see build.rs); keeps
# `cargo clippy/check/build` and lint CI working without apt/brew protobuf.
protoc-bin-vendored = "3"
[features]
default = ["pyo3/extension-module"]
[profile.release]
opt-level = 2
lto = "thin"
strip = true
[profile.dev]
opt-level = 0
debug = 1
rust/sglang-grpc/build.rs
+12
View File
@@ -1,4 +1,16 @@
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Prefer an explicitly configured protoc; otherwise use the vendored
// binary so builds (including `cargo clippy` and the pre-commit hook on
// machines/CI runners without protobuf installed) are self-contained.
// protoc_bin_path() errs on platforms the vendored crate doesn't cover;
// those fall back to prost-build's own lookup of `protoc` on PATH.
if std::env::var_os("PROTOC").is_none()
&& let Ok(vendored) = protoc_bin_vendored::protoc_bin_path()
{
// SAFETY: build scripts are single-threaded at this point.
unsafe { std::env::set_var("PROTOC", vendored) };
}
let proto_path = "../../proto/sglang/runtime/v1/sglang.proto";
tonic_build::configure()
rust/sglang-grpc/pyproject.toml
+17
View File
@@ -0,0 +1,17 @@
# Standalone dev builds only (`maturin build` / `maturin develop`); the
# extension ships to users inside the main sglang wheel via setuptools-rust
# (python/pyproject.toml, target sglang.srt.grpc._core). Requires a repo
# checkout (build.rs reads ../../proto), so this is not publishable as an
# sdist. protoc is not required: build.rs falls back to a vendored binary.
[build-system]
requires = ["maturin>=1.5,<2"]
build-backend = "maturin"
[project]
name = "sglang-grpc"
dynamic = ["version"]
description = "In-process Rust gRPC server for SGLang"
requires-python = ">=3.10"
[tool.maturin]
module-name = "_core"
rust/sglang-grpc/src/bridge.rs
+31 -28
View File
@@ -40,7 +40,7 @@ type BridgeStateRef = Arc<Mutex<BridgeState>>;
struct BridgeState {
channels: HashMap<String, Sender<ResponseChunk>>,
pending_sends: HashSet<String>,
ready_callbacks: HashMap<String, PyObject>,
ready_callbacks: HashMap<String, Py<PyAny>>,
ready_signals: HashSet<String>,
terminal_errors: HashMap<String, TerminalError>,
}
@@ -66,7 +66,10 @@ impl TerminalError {
}
}
#[pyclass(eq, eq_int)]
// skip_from_py_object: this enum is only returned to Python, never received
// from it, so it opts out of pyo3's (deprecated-by-default) FromPyObject
// derive for Clone pyclasses.
#[pyclass(eq, eq_int, skip_from_py_object)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChunkSendStatus {
Ready,
@@ -83,7 +86,7 @@ fn lock_or_recover<'a, T>(mutex: &'a Mutex<T>, name: &'static str) -> MutexGuard
/// Holds a reference to the Python RuntimeHandle and manages per-request channels.
pub struct PyBridge {
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
state: BridgeStateRef,
rust_tokenizer: Option<RustTokenizer>,
context_len: i32,
@@ -93,7 +96,7 @@ pub struct PyBridge {
impl PyBridge {
pub fn new(
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
rust_tokenizer: Option<RustTokenizer>,
context_len: i32,
response_channel_capacity: usize,
@@ -144,7 +147,7 @@ impl PyBridge {
Ok(receiver)
}
fn make_chunk_callback(&self, py: Python<'_>, rid: String) -> PyResult<PyObject> {
fn make_chunk_callback(&self, py: Python<'_>, rid: String) -> PyResult<Py<PyAny>> {
let callback = ChunkCallback {
rid,
state: self.state.clone(),
@@ -155,7 +158,7 @@ impl PyBridge {
Ok(py_callback.into_any())
}
fn make_json_callback(&self, py: Python<'_>, rid: String) -> PyResult<PyObject> {
fn make_json_callback(&self, py: Python<'_>, rid: String) -> PyResult<Py<PyAny>> {
let callback = JsonChunkCallback {
rid,
state: self.state.clone(),
@@ -183,7 +186,7 @@ impl PyBridge {
let receiver = self.create_channel(rid)?;
let rid_owned = rid.to_string();
let result = Python::with_gil(|py| -> PyResult<()> {
let result = Python::attach(|py| -> PyResult<()> {
let py_req_dict = json_map_to_pydict(py, &req_dict)?;
let callback = self.make_chunk_callback(py, rid_owned)?;
@@ -253,7 +256,7 @@ impl PyBridge {
return Ok(());
}
Python::with_gil(|py| {
Python::attach(|py| {
self.runtime_handle
.call_method1(py, "abort", (rid, abort_all))?;
Ok(())
@@ -265,21 +268,21 @@ impl PyBridge {
// ------------------------------------------------------------------
pub fn get_model_info(&self) -> PyResult<String> {
Python::with_gil(|py| {
Python::attach(|py| {
let result = self.runtime_handle.call_method0(py, "get_model_info")?;
result.extract::<String>(py)
})
}
pub fn get_server_info(&self) -> PyResult<String> {
Python::with_gil(|py| {
Python::attach(|py| {
let result = self.runtime_handle.call_method0(py, "get_server_info")?;
result.extract::<String>(py)
})
}
pub fn health_check(&self) -> PyResult<bool> {
Python::with_gil(|py| {
Python::attach(|py| {
let result = self.runtime_handle.call_method0(py, "health_check")?;
result.extract::<bool>(py)
})
@@ -287,7 +290,7 @@ impl PyBridge {
/// Tokenize via Python (fallback when Rust tokenizer unavailable).
pub fn tokenize_py(&self, text: &str, add_special_tokens: bool) -> PyResult<String> {
Python::with_gil(|py| {
Python::attach(|py| {
let result =
self.runtime_handle
.call_method1(py, "tokenize", (text, add_special_tokens))?;
@@ -297,7 +300,7 @@ impl PyBridge {
/// Detokenize via Python (fallback when Rust tokenizer unavailable).
pub fn detokenize_py(&self, tokens: Vec<i32>) -> PyResult<String> {
Python::with_gil(|py| {
Python::attach(|py| {
let result = self
.runtime_handle
.call_method1(py, "detokenize", (tokens,))?;
@@ -306,7 +309,7 @@ impl PyBridge {
}
pub fn list_models(&self) -> PyResult<String> {
Python::with_gil(|py| {
Python::attach(|py| {
let result = self.runtime_handle.call_method0(py, "list_models")?;
result.extract::<String>(py)
})
@@ -314,13 +317,13 @@ impl PyBridge {
fn submit_json<F>(&self, rid: &str, call: F) -> PyResult<Receiver<ResponseChunk>>
where
F: for<'py> FnOnce(Python<'py>, &PyObject, PyObject) -> PyResult<()>,
F: for<'py> FnOnce(Python<'py>, &Py<PyAny>, Py<PyAny>) -> PyResult<()>,
{
// Closure args are: current Python token, RuntimeHandle, and the JSON chunk callback.
let receiver = self.create_channel(rid)?;
let rid_owned = rid.to_string();
let result = Python::with_gil(|py| -> PyResult<()> {
let result = Python::attach(|py| -> PyResult<()> {
let callback = self.make_json_callback(py, rid_owned)?;
call(py, &self.runtime_handle, callback)
});
@@ -450,7 +453,7 @@ fn close_channel_with_error(
py: Python<'_>,
rid: &str,
state: &BridgeStateRef,
runtime_handle: &PyObject,
runtime_handle: &Py<PyAny>,
error: TerminalError,
) {
let mut state = lock_or_recover(state.as_ref(), "state");
@@ -478,7 +481,7 @@ fn register_pending_send(rid: &str, state: &BridgeStateRef) -> bool {
state.pending_sends.insert(rid.to_string())
}
fn mark_send_ready(py: Python<'_>, rid: &str, state: &BridgeStateRef) -> Option<PyObject> {
fn mark_send_ready(py: Python<'_>, rid: &str, state: &BridgeStateRef) -> Option<Py<PyAny>> {
let mut state = lock_or_recover(state.as_ref(), "state");
state.pending_sends.remove(rid);
if let Some(callback) = state.ready_callbacks.get(rid) {
@@ -489,7 +492,7 @@ fn mark_send_ready(py: Python<'_>, rid: &str, state: &BridgeStateRef) -> Option<
}
}
fn notify_ready(py: Python<'_>, rid: &str, callback: PyObject) {
fn notify_ready(py: Python<'_>, rid: &str, callback: Py<PyAny>) {
if let Err(err) = callback.call0(py) {
tracing::warn!(rid, "gRPC on_ready callback failed: {}", err);
}
@@ -499,7 +502,7 @@ fn set_on_ready_for_rid(
py: Python<'_>,
rid: &str,
state: &BridgeStateRef,
on_ready: PyObject,
on_ready: Py<PyAny>,
) -> PyResult<()> {
let should_notify = {
let mut state = lock_or_recover(state.as_ref(), "state");
@@ -525,7 +528,7 @@ fn try_send_chunk(
py: Python<'_>,
rid: &str,
state: &BridgeStateRef,
runtime_handle: &PyObject,
runtime_handle: &Py<PyAny>,
tokio_handle: &Handle,
sender: &Sender<ResponseChunk>,
msg: ResponseChunk,
@@ -569,14 +572,14 @@ fn try_send_chunk(
return;
}
Python::with_gil(|py| {
Python::attach(|py| {
if let Some(callback) = mark_send_ready(py, &rid_owned, &state) {
notify_ready(py, &rid_owned, callback);
}
});
}
Err(_) => {
Python::with_gil(|py| {
Python::attach(|py| {
close_channel_with_error(
py,
&rid_owned,
@@ -611,7 +614,7 @@ fn try_send_chunk(
struct ChunkCallback {
rid: String,
state: BridgeStateRef,
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
tokio_handle: Handle,
}
@@ -619,7 +622,7 @@ struct ChunkCallback {
impl ChunkCallback {
/// Register before producing chunks. If a parked chunk drained before registration,
/// Rust fires `on_ready` immediately so late registration cannot miss the edge.
fn set_on_ready(&self, py: Python<'_>, on_ready: PyObject) -> PyResult<()> {
fn set_on_ready(&self, py: Python<'_>, on_ready: Py<PyAny>) -> PyResult<()> {
set_on_ready_for_rid(py, &self.rid, &self.state, on_ready)
}
@@ -699,7 +702,7 @@ impl ChunkCallback {
struct JsonChunkCallback {
rid: String,
state: BridgeStateRef,
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
tokio_handle: Handle,
}
@@ -707,7 +710,7 @@ struct JsonChunkCallback {
impl JsonChunkCallback {
/// Register before producing chunks. If a parked chunk drained before registration,
/// Rust fires `on_ready` immediately so late registration cannot miss the edge.
fn set_on_ready(&self, py: Python<'_>, on_ready: PyObject) -> PyResult<()> {
fn set_on_ready(&self, py: Python<'_>, on_ready: Py<PyAny>) -> PyResult<()> {
set_on_ready_for_rid(py, &self.rid, &self.state, on_ready)
}
@@ -785,7 +788,7 @@ impl JsonChunkCallback {
fn extract_meta_info(chunk: &Bound<'_, PyDict>) -> HashMap<String, String> {
let mut meta = HashMap::new();
if let Ok(Some(meta_obj)) = chunk.get_item("meta_info")
&& let Ok(meta_dict) = meta_obj.downcast::<PyDict>()
&& let Ok(meta_dict) = meta_obj.cast::<PyDict>()
{
for (k, v) in meta_dict.iter() {
// The proto schema is map<string, string>; encode each Python value as JSON
rust/sglang-grpc/src/lib.rs
+7 -7
View File
@@ -53,10 +53,10 @@ struct TokenizerInfo {
/// fall back to Python tokenization.
fn try_get_attr(
py: Python<'_>,
obj: &PyObject,
obj: &Py<PyAny>,
attr: &'static str,
context: &'static str,
) -> Option<PyObject> {
) -> Option<Py<PyAny>> {
obj.getattr(py, attr).map(Some).unwrap_or_else(|err| {
tracing::debug!("{}.{} is unavailable: {}", context, attr, err);
None
@@ -65,7 +65,7 @@ fn try_get_attr(
fn try_get_attr_str(
py: Python<'_>,
obj: &PyObject,
obj: &Py<PyAny>,
attr: &'static str,
context: &'static str,
) -> Option<String> {
@@ -79,7 +79,7 @@ fn try_get_attr_str(
fn try_get_attr_i32(
py: Python<'_>,
obj: &PyObject,
obj: &Py<PyAny>,
attr: &'static str,
context: &'static str,
) -> Option<i32> {
@@ -91,8 +91,8 @@ fn try_get_attr_i32(
})
}
fn extract_tokenizer_info(runtime_handle: &PyObject) -> PyResult<TokenizerInfo> {
Python::with_gil(|py| {
fn extract_tokenizer_info(runtime_handle: &Py<PyAny>) -> PyResult<TokenizerInfo> {
Python::attach(|py| {
let tm = runtime_handle
.getattr(py, "tokenizer_manager")
.map_err(|err| {
@@ -152,7 +152,7 @@ fn extract_tokenizer_info(runtime_handle: &PyObject) -> PyResult<TokenizerInfo>
fn start_server(
host: String,
port: u16,
runtime_handle: PyObject,
runtime_handle: Py<PyAny>,
worker_threads: usize,
response_channel_capacity: usize,
response_timeout_secs: u64,
rust/sglang-grpc/src/server.rs
+4 -3
View File
@@ -64,7 +64,7 @@ fn resolve_max_message_size() -> usize {
/// Everything else (typically `PyRuntimeError`, but also Python tracebacks
/// from inside the tokenizer manager) maps to `INTERNAL`.
fn pyerr_to_status(err: PyErr, context: &str) -> Status {
let is_client_error = Python::with_gil(|py| {
let is_client_error = Python::attach(|py| {
err.is_instance_of::<PyValueError>(py) || err.is_instance_of::<PyTypeError>(py)
});
let msg = format!("{}: {}", context, err);
@@ -125,9 +125,10 @@ impl Drop for RequestAbortGuard {
fn spawn_abort(bridge: Arc<PyBridge>, rid: String) {
match tokio::runtime::Handle::try_current() {
Ok(handle) => {
let _ = handle.spawn_blocking(move || {
// Fire-and-forget: dropping the JoinHandle detaches the task.
drop(handle.spawn_blocking(move || {
let _ = bridge.abort(&rid, false);
});
}));
}
Err(_) => {
tracing::warn!(
rust/sglang-grpc/src/utils/py_utils.rs
+2 -2
View File
@@ -2,7 +2,7 @@ use pyo3::prelude::*;
use pyo3::types::{PyAny, PyDict, PyList};
use std::collections::HashMap;
fn json_value_to_py<'py>(py: Python<'py>, v: &serde_json::Value) -> PyResult<PyObject> {
fn json_value_to_py<'py>(py: Python<'py>, v: &serde_json::Value) -> PyResult<Py<PyAny>> {
match v {
serde_json::Value::Null => Ok(py.None()),
serde_json::Value::Bool(b) => Ok(b.into_pyobject(py)?.to_owned().into_any().unbind()),
@@ -17,7 +17,7 @@ fn json_value_to_py<'py>(py: Python<'py>, v: &serde_json::Value) -> PyResult<PyO
}
serde_json::Value::String(s) => Ok(s.into_pyobject(py)?.into_any().unbind()),
serde_json::Value::Array(arr) => {
let items: Vec<PyObject> = arr
let items: Vec<Py<PyAny>> = arr
.iter()
.map(|item| json_value_to_py(py, item))
.collect::<PyResult<_>>()?;
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> {