Add Inkling model support (#31681)

Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Yanbin Jiang <jybsuper@gmail.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Qiaolin Yu <qiaolin.yu@radixark.ai>
Co-authored-by: Zhichen Zeng <zczeng@uw.edu>
Co-authored-by: Aurick Qiao <aurick@thinkingmachines.ai>
Co-authored-by: Joseph <jk@thinkingmachines.ai>

rust/sglang-mm/tests/generate_golden.py

@@ -0,0 +1,40 @@
+import io
+import os
+import sys
+
+import numpy as np
+import torch
+from PIL import Image
+
+sys.path.insert(0, os.path.dirname(__file__))
+from bench_parity import PS, make_photo_like, ref_patchify
+
+OUT = (
+    sys.argv[1]
+    if len(sys.argv) > 1
+    else os.path.join(os.path.dirname(__file__), "..", "tests", "golden")
+)
+
+CASES = [
+    ("480x640", 480, 640, 10),
+    ("200x320", 200, 320, 11),
+    ("37x53", 37, 53, 12),
+    ("40x40", 40, 40, 13),
+]
+
+os.makedirs(OUT, exist_ok=True)
+for name, h, w, seed in CASES:
+    arr = make_photo_like(h, w, seed=seed)
+    bits = ref_patchify(arr).view(torch.uint16).numpy()
+    buf = io.BytesIO()
+    Image.fromarray(arr).save(buf, format="PNG")
+    path = os.path.join(OUT, f"golden_{name}.npz")
+    np.savez_compressed(
+        path,
+        arr=arr,
+        bits=bits,
+        png=np.frombuffer(buf.getvalue(), dtype=np.uint8),
+        patch_size=np.int64(PS),
+    )
+    print(f"  {path}: input {h}x{w}, bits {bits.shape}")
+print("GOLDEN_OK")

rust/sglang-mm/tests/test_golden.py

@@ -0,0 +1,51 @@
+import glob
+import os
+
+import numpy as np
+import pytest
+
+import sglang.srt.multimodal._core.inkling
+
+GOLDEN_DIR = os.environ.get(
+    "INKLING_MM_GOLDEN_DIR",
+    os.path.join(os.path.dirname(__file__), "..", "tests", "golden"),
+)
+GOLDENS = sorted(glob.glob(os.path.join(GOLDEN_DIR, "golden_*.npz")))
+
+
+def bf16_bits_to_f32(bits: np.ndarray) -> np.ndarray:
+    return (bits.astype(np.uint32) << 16).view(np.float32)
+
+
+@pytest.mark.parametrize("path", GOLDENS, ids=[os.path.basename(p) for p in GOLDENS])
+def test_patchify_rgb_bit_exact(path):
+    g = np.load(path)
+    got = sglang.srt.multimodal._core.inkling.patchify_rgb(
+        g["arr"], int(g["patch_size"])
+    )
+    np.testing.assert_array_equal(got, g["bits"].reshape(-1))
+
+
+@pytest.mark.parametrize("path", GOLDENS, ids=[os.path.basename(p) for p in GOLDENS])
+def test_decode_patchify_png_bit_exact(path):
+    g = np.load(path)
+    h_ref, w_ref = g["arr"].shape[:2]
+    h, w, got = sglang.srt.multimodal._core.inkling.decode_patchify(
+        g["png"].tobytes(), int(g["patch_size"])
+    )
+    assert (h, w) == (h_ref, w_ref)
+    np.testing.assert_array_equal(got, g["bits"].reshape(-1))
+
+
+def test_batch_matches_single():
+    gs = [np.load(p) for p in GOLDENS]
+    data = [g["png"].tobytes() for g in gs]
+    ps = int(gs[0]["patch_size"])
+    for (h, w, bits), g in zip(
+        sglang.srt.multimodal._core.inkling.decode_patchify_batch(data, ps), gs
+    ):
+        np.testing.assert_array_equal(bits, g["bits"].reshape(-1))
+
+
+def test_golden_fixtures_exist():
+    assert len(GOLDENS) >= 4, f"expected golden fixtures in {GOLDEN_DIR}"

rust/sglang-mm/tests/test_hash_fetch.py

@@ -0,0 +1,103 @@
+import asyncio
+import base64
+import io
+import os
+import sys
+import time
+from types import SimpleNamespace
+
+import numpy as np
+import soundfile as sf
+import torch
+from PIL import Image
+
+sys.path.insert(0, os.path.dirname(__file__))
+from bench_parity import make_photo_like
+
+from sglang.srt.managers.mm_utils import data_hash, hash_feature
+from sglang.srt.multimodal.inkling import InklingProcessor
+from sglang.srt.multimodal.processors import inkling as prc
+
+
+def png_bytes(arr):
+    buf = io.BytesIO()
+    Image.fromarray(arr).save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def wav_bytes(seconds=1.0, sr=16000):
+    t = np.linspace(0, seconds, int(sr * seconds), endpoint=False)
+    buf = io.BytesIO()
+    sf.write(
+        buf, (0.3 * np.sin(2 * np.pi * 440 * t)).astype(np.float32), sr, format="WAV"
+    )
+    return buf.getvalue()
+
+
+def make_proc():
+    proc = prc.InklingMultimodalProcessor.__new__(prc.InklingMultimodalProcessor)
+    proc.IMAGE_TOKEN_ID = 100
+    proc.AUDIO_TOKEN_ID = 101
+    proc.AUDIO_END_TOKEN_ID = 102
+    proc.inkling_processor = InklingProcessor()
+    return proc
+
+
+proc = make_proc()
+img = png_bytes(make_photo_like(200, 320, seed=7))
+aud = wav_bytes()
+
+out = proc.assemble([1, 100, 2, 101, 3], [img], [aud])
+img_item = next(i for i in out.mm_items if i.modality.name == "IMAGE")
+aud_item = next(i for i in out.mm_items if i.modality.name == "AUDIO")
+assert img_item.hash == data_hash(img), "image hash != data_hash(raw bytes)"
+assert aud_item.hash == data_hash(aud), "audio hash != data_hash(raw bytes)"
+print(f"  assemble: image hash={img_item.hash:#x} audio hash={aud_item.hash:#x} OK")
+
+h0 = img_item.hash
+img_item.set_pad_value()
+assert img_item.hash == h0 and img_item.pad_value is not None
+print(f"  set_pad_value: hash preserved, pad_value={img_item.pad_value} OK")
+
+out2 = proc.assemble([1, 100, 2], [img], [])
+assert out2.mm_items[0].hash == h0
+print("  determinism: same bytes -> same hash OK")
+
+data_url = "data:image/png;base64," + base64.b64encode(img).decode()
+req = SimpleNamespace(input_ids=[1, 100, 100, 2])
+out3 = asyncio.run(
+    proc.process_mm_data_async(
+        image_data=[data_url, data_url], audio_data=None, request_obj=req
+    )
+)
+assert all(i.hash == h0 for i in out3.mm_items), "data: URL roundtrip hash mismatch"
+print("  process_mm_data_async: concurrent resolve + hash OK")
+
+orig = prc._resolve_media_item
+prc._resolve_media_item = lambda it: (time.sleep(0.3), orig(it))[1]
+t0 = time.perf_counter()
+asyncio.run(prc._resolve_media_items([data_url] * 8))
+elapsed = time.perf_counter() - t0
+prc._resolve_media_item = orig
+assert elapsed < 1.2, f"8x 0.3s resolves took {elapsed:.2f}s; expected ~0.3s"
+print(f"  concurrency: 8 x 0.3s resolves in {elapsed:.2f}s OK")
+
+imgs_5 = [png_bytes(make_photo_like(1080, 1920, seed=s)) for s in range(5)]
+feats = [
+    torch.randn(1323, 1, 40, 40, 3, dtype=torch.bfloat16).expand(1323, 2, 40, 40, 3)
+    for _ in range(5)
+]
+t0 = time.perf_counter()
+for b in imgs_5:
+    data_hash(b)
+t_bytes = (time.perf_counter() - t0) * 1e3
+t0 = time.perf_counter()
+for f in feats:
+    hash_feature(f)
+t_feat = (time.perf_counter() - t0) * 1e3
+print(
+    f"  hash cost 5 imgs: raw bytes {t_bytes:.1f}ms vs feature tensor {t_feat:.1f}ms "
+    f"({t_feat / t_bytes:.0f}x)"
+)
+
+print("HASH_FETCH_OK")

rust/sglang-mm/tests/test_integration.py

@@ -0,0 +1,86 @@
+import io
+import os
+import sys
+
+import numpy as np
+import torch
+from PIL import Image
+
+sys.path.insert(0, os.path.dirname(__file__))
+from bench_parity import make_photo_like
+
+import sglang.srt.multimodal.inkling.image_processing as ip
+
+
+def encode(arr, fmt):
+    buf = io.BytesIO()
+    Image.fromarray(arr).save(
+        buf, format=fmt, **({"quality": 90} if fmt == "JPEG" else {})
+    )
+    return buf.getvalue()
+
+
+def run(images, use_rs: bool, rescale: bool):
+    ip._rs_module = None
+    os.environ["SGLANG_RS_MM_PREPROCESS"] = "1" if use_rs else "0"
+    kwargs = (
+        {}
+        if rescale
+        else {"rescale_image_frac": None, "rescale_image_max_upscaled_long_edge": None}
+    )
+    proc = ip.InklingImageProcessor(patch_size=40, **kwargs)
+    out = proc.preprocess(images)
+    assert (ip._rs_module is not False) == use_rs, "rust module gating mismatch"
+    return out
+
+
+def compare(tag, images, expect_exact, rescale=False):
+    ref = run(images, use_rs=False, rescale=rescale)
+    got = run(images, use_rs=True, rescale=rescale)
+    assert ref["num_patches"] == got["num_patches"], tag
+    assert ref["num_tokens"] == got["num_tokens"], tag
+    a, b = ref["vision_patches_bthwc"], got["vision_patches_bthwc"]
+    assert a.shape == b.shape and a.dtype == b.dtype, f"{tag}: {a.shape} vs {b.shape}"
+    exact = torch.equal(
+        a.contiguous().view(torch.uint16), b.contiguous().view(torch.uint16)
+    )
+    if expect_exact:
+        assert exact, f"{tag}: expected bit-exact"
+        print(f"  {tag}: bit-exact=True shape={tuple(a.shape)}")
+    else:
+        d = (a.float() - b.float()).abs()
+        print(
+            f"  {tag}: bit-exact={exact} max_abs={d.max():.6f} shape={tuple(a.shape)}"
+        )
+        assert d.max() < 0.25, f"{tag}: JPEG decoder diff too large"
+
+
+arr1 = make_photo_like(1080, 1920, seed=1)
+arr2 = make_photo_like(720, 1280, seed=2)
+arr3 = make_photo_like(480, 640, seed=3)
+
+print("=== integration: InklingImageProcessor env-gated rust path ===")
+compare("single PNG", [encode(arr1, "PNG")], expect_exact=True)
+compare("single JPEG", [encode(arr1, "JPEG")], expect_exact=False)
+compare(
+    "5x PNG batch",
+    [encode(a, "PNG") for a in [arr1, arr2, arr3, arr1, arr2]],
+    expect_exact=True,
+)
+compare(
+    "mixed JPEG/PNG batch",
+    [encode(arr1, "JPEG"), encode(arr2, "PNG")],
+    expect_exact=False,
+)
+compare("PIL input (PNG roundtrip)", [Image.fromarray(arr3)], expect_exact=True)
+compare("single PNG rescaled", [encode(arr1, "PNG")], expect_exact=True, rescale=True)
+compare(
+    "single JPEG rescaled", [encode(arr1, "JPEG")], expect_exact=False, rescale=True
+)
+compare(
+    "3x mixed rescaled",
+    [encode(arr1, "JPEG"), encode(arr2, "PNG"), encode(arr3, "PNG")],
+    expect_exact=False,
+    rescale=True,
+)
+print("INTEGRATION_OK")

rust/sglang-mm/tests/test_resize_parity.py

@@ -0,0 +1,121 @@
+import math
+import time
+from typing import Optional
+
+import numpy as np
+import pytest
+from PIL import Image
+
+import sglang.srt.multimodal._core.inkling
+
+
+def py_scaled_dims(
+    width: int,
+    height: int,
+    frac: Optional[float],
+    cap: Optional[int],
+):
+    if frac is None:
+        return width, height
+    long_edge = max(width, height)
+    if long_edge == 0:
+        return width, height
+    target = float(long_edge) * frac
+    if cap is not None:
+        target = min(target, float(max(cap, long_edge)))
+    ratio = target / float(long_edge)
+    if ratio == 1.0:
+        return width, height
+
+    def scale(value):
+        return max(1, math.floor(float(value) * ratio + 0.5))
+
+    return scale(width), scale(height)
+
+
+def pil_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray:
+    return np.array(
+        Image.fromarray(arr).resize((tw, th), resample=Image.Resampling.LANCZOS),
+        dtype=np.uint8,
+    )
+
+
+def rs_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray:
+    return sglang.srt.multimodal._core.inkling.resize_rgb(arr, tw, th).reshape(
+        th, tw, 3
+    )
+
+
+CASES = [
+    (1080, 1920, 1152, 2048),
+    (896, 896, 1792, 1792),
+    (360, 640, 720, 1280),
+    (37, 53, 74, 106),
+    (100, 100, 173, 173),
+    (1, 1, 2, 2),
+    (256, 256, 100, 100),
+    (720, 1280, 720, 1280),
+    (3, 500, 6, 1000),
+]
+
+
+@pytest.mark.parametrize(
+    "h,w,th,tw", CASES, ids=[f"{h}x{w}->{th}x{tw}" for h, w, th, tw in CASES]
+)
+def test_resize_bit_exact(h, w, th, tw):
+    rng = np.random.default_rng(h * 10000 + w)
+    arr = rng.integers(0, 256, (h, w, 3), dtype=np.uint8)
+    np.testing.assert_array_equal(rs_resize(arr, tw, th), pil_resize(arr, tw, th))
+
+
+def test_scaled_dims_sweep():
+    rng = np.random.default_rng(0)
+    sizes = [(int(a), int(b)) for a, b in rng.integers(1, 5000, (500, 2))]
+    sizes += [(2048, 1024), (2049, 100), (1024, 2048), (1, 1), (4096, 4096)]
+    for frac, cap in [(2.0, 2048), (1.5, 2048), (3.0, None), (None, None), (2.0, 1)]:
+        for w, h in sizes:
+            assert sglang.srt.multimodal._core.inkling.scaled_dims(
+                w, h, frac, cap
+            ) == py_scaled_dims(w, h, frac, cap), (
+                w,
+                h,
+                frac,
+                cap,
+            )
+
+
+def test_decode_patchify_rescaled_matches_pil_pipeline():
+    import io
+
+    import torch
+
+    rng = np.random.default_rng(7)
+    arr = rng.integers(0, 256, (1080, 1920, 3), dtype=np.uint8)
+    buf = io.BytesIO()
+    Image.fromarray(arr).save(buf, format="PNG")
+    h, w, bits = sglang.srt.multimodal._core.inkling.decode_patchify(
+        buf.getvalue(), 40, 2.0, 2048
+    )
+    assert (w, h) == py_scaled_dims(1920, 1080, 2.0, 2048)
+    ref_arr = pil_resize(arr, w, h)
+    ref_bits = sglang.srt.multimodal._core.inkling.patchify_rgb(ref_arr, 40)
+    np.testing.assert_array_equal(bits, ref_bits)
+    assert torch.from_numpy(bits).view(torch.bfloat16).shape[0] > 0
+
+
+def test_resize_bench():
+    arr = np.random.default_rng(1).integers(0, 256, (1080, 1920, 3), dtype=np.uint8)
+    tw, th = py_scaled_dims(1920, 1080, 2.0, 2048)
+    pil_resize(arr, tw, th)
+    rs_resize(arr, tw, th)
+    t0 = time.perf_counter()
+    for _ in range(10):
+        pil_resize(arr, tw, th)
+    t_pil = (time.perf_counter() - t0) / 10 * 1e3
+    t0 = time.perf_counter()
+    for _ in range(10):
+        rs_resize(arr, tw, th)
+    t_rs = (time.perf_counter() - t0) / 10 * 1e3
+    print(
+        f"\nresize 1920x1080->{tw}x{th}: PIL {t_pil:.1f}ms rust {t_rs:.1f}ms ({t_pil/t_rs:.1f}x)"
+    )