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

[Perf] Speed up the Kimi-K2.5 vision path and match PIL bicubic in the GPU resize (#33349)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Liangsheng Yin
2026-08-03 21:57:49 -07:00
committed by GitHub
co-authored by Mick
parent 7ba393dd15
commit afc868517b
11 changed files with 761 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/vlm/verify_k25_equivalence.py
+145
View File
@@ -0,0 +1,145 @@
"""Prove the two GPU-only rewrites in the K2.5 port are equivalent to main.
1. normalize_and_patchify(scale/bias) == pad -> /255 -> (x-mean)*inv_std -> patchify
2. apply_fused_qk_complex_rope_inplace == the torch complex reference
"""
import torch
import torch.nn.functional as F
from sglang.kernels.ops.attention.vision_rope import (
apply_fused_qk_complex_rope_inplace,
prepare_fused_qk_complex_rope_inplace,
)
from sglang.kernels.ops.mm.process import normalize_and_patchify
MEAN = [0.5, 0.5, 0.5]
STD = [0.5, 0.5, 0.5]
ASYM_MEAN = [0.481, 0.457, 0.408]
ASYM_STD = [0.268, 0.261, 0.275]
def reference_preprocess(batch_u8, mean, std, patch_size, padded_h, padded_w):
"""Exactly what main does, in main's order."""
image_mean = torch.tensor(mean, device="cuda", dtype=torch.float32).view(1, 3, 1, 1)
image_std_inv = (1.0 / torch.tensor(std, device="cuda", dtype=torch.float32)).view(
1, 3, 1, 1
)
x = batch_u8.float()
pad_h = padded_h - x.shape[-2]
pad_w = padded_w - x.shape[-1]
if pad_h > 0 or pad_w > 0:
x = F.pad(x, (0, pad_w, 0, pad_h), value=0.0)
x = x / 255.0
x = (x - image_mean) * image_std_inv
B, C, H, W = x.shape
gh, gw = H // patch_size, W // patch_size
x = x.view(B, C, gh, patch_size, gw, patch_size)
return x.permute(0, 2, 4, 1, 3, 5).reshape(B, -1, C, patch_size, patch_size)
def check_patchify():
print("== normalize_and_patchify vs main's pad/normalize/patchify ==")
torch.manual_seed(0)
cases = [
# (H, W, padded_h, padded_w, patch, mean, std, label)
(32, 24, 32, 24, 8, MEAN, STD, "no padding, symmetric norm"),
(30, 22, 32, 24, 8, MEAN, STD, "padded, symmetric norm"),
(30, 22, 32, 24, 8, ASYM_MEAN, ASYM_STD, "padded, per-channel norm"),
(64, 64, 64, 64, 16, ASYM_MEAN, ASYM_STD, "large patch"),
]
ok = True
for h, w, ph, pw, patch, mean, std, label in cases:
raw = torch.randint(0, 256, (3, 3, h, w), dtype=torch.uint8, device="cuda")
ref = reference_preprocess(raw, mean, std, patch, ph, pw)
scale = torch.tensor(
[1.0 / (255.0 * s) for s in std], device="cuda", dtype=torch.float32
).view(1, 3, 1, 1)
bias = torch.tensor(
[-m / s for m, s in zip(mean, std)], device="cuda", dtype=torch.float32
).view(1, 3, 1, 1)
got = normalize_and_patchify(raw.float(), scale, bias, patch, ph, pw)
max_abs = (got - ref).abs().max().item()
# The padded rows must carry -mean/std, not zero.
pad_ok = True
if ph > h or pw > w:
pad_ok = torch.allclose(
got.flatten()[(got - ref).abs().argmax()],
ref.flatten()[(got - ref).abs().argmax()],
atol=1e-5,
)
good = max_abs < 1e-5 and pad_ok
ok &= good
print(f" {'PASS' if good else 'FAIL'} {label:32s} max|d|={max_abs:.3e}")
return ok
def check_padded_value_is_not_zero():
"""The old pipeline padded in raw space, so pad cells become -mean/std."""
print("== padded cells carry -mean/std, not 0 ==")
raw = torch.full((1, 3, 8, 8), 128, dtype=torch.uint8, device="cuda")
scale = torch.tensor(
[1.0 / (255.0 * s) for s in ASYM_STD], device="cuda", dtype=torch.float32
).view(1, 3, 1, 1)
bias = torch.tensor(
[-m / s for m, s in zip(ASYM_MEAN, ASYM_STD)],
device="cuda",
dtype=torch.float32,
).view(1, 3, 1, 1)
got = normalize_and_patchify(raw.float(), scale, bias, 8, 16, 16)
# patch index 1 is the (row 0, col 1) patch -- entirely padding.
pad_patch = got[0, 1]
expected = bias.view(3, 1, 1).expand(3, 8, 8)
good = torch.allclose(pad_patch, expected, atol=1e-6)
print(
f" {'PASS' if good else 'FAIL'} pad cell = {pad_patch[0, 0, 0].item():.6f}, "
f"expected -mean/std = {expected[0, 0, 0].item():.6f}"
)
return good
def reference_rope(xq, xk, freqs_cis):
freqs_cis = freqs_cis.unsqueeze(-2)
xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2))
xk_ = torch.view_as_complex(xk.float().view(*xk.shape[:-1], -1, 2))
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2)
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2)
return xq_out.type_as(xq), xk_out.type_as(xk)
def check_rope():
print("== fused vision RoPE vs the torch complex reference ==")
torch.manual_seed(0)
ok = True
for dtype, tol in ((torch.bfloat16, 8e-3), (torch.float16, 2e-3)):
for tokens, heads, head_dim in ((1024, 16, 72), (4096, 8, 128), (37, 4, 64)):
xq = torch.randn(tokens, heads, head_dim, device="cuda", dtype=dtype)
xk = torch.randn(tokens, heads, head_dim, device="cuda", dtype=dtype)
angle = torch.randn(tokens, head_dim // 2, device="cuda")
freqs_cis = torch.polar(torch.ones_like(angle), angle)
ref_q, ref_k = reference_rope(xq, xk, freqs_cis)
prepared = prepare_fused_qk_complex_rope_inplace(freqs_cis)
got_q, got_k = apply_fused_qk_complex_rope_inplace(
xq.clone(), xk.clone(), prepared
)
dq = (got_q.float() - ref_q.float()).abs().max().item()
dk = (got_k.float() - ref_k.float()).abs().max().item()
good = dq < tol and dk < tol
ok &= good
print(
f" {'PASS' if good else 'FAIL'} {str(dtype):16s} "
f"t={tokens:5d} h={heads:2d} d={head_dim:3d} "
f"max|dq|={dq:.2e} max|dk|={dk:.2e}"
)
return ok
if __name__ == "__main__":
results = [check_patchify(), check_padded_value_is_not_zero(), check_rope()]
print()
print("ALL PASS" if all(results) else "SOME CHECKS FAILED")
raise SystemExit(0 if all(results) else 1)
test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py
+29
View File
@@ -20,6 +20,8 @@ from sglang.kernels.ops.attention.set_mla_kv_concat_q import (
from sglang.kernels.ops.attention.utils import concat_mla_absorb_q_general
from sglang.kernels.ops.attention.vision_rope import (
apply_fused_qk_complex_rope,
apply_fused_qk_complex_rope_inplace,
prepare_fused_qk_complex_rope_inplace,
)
from sglang.kernels.ops.elementwise import add3
from sglang.kernels.ops.gemm.tiny_gemm import (
@@ -427,6 +429,33 @@ class TestKimiK3PrerequisiteOps(CustomTestCase):
torch.testing.assert_close(actual_q, reference(q), rtol=0, atol=atol)
torch.testing.assert_close(actual_k, reference(k), rtol=0, atol=atol)
def test_vision_rope_inplace(self):
# VisionAttention hands the applier contiguous q/k, which is what the
# in-place kernel requires; mirror that rather than qkv.unbind views.
for dtype in (torch.bfloat16, torch.float16):
torch.manual_seed(4)
q = torch.randn(480, 12, 128, device="cuda", dtype=dtype)
k = torch.randn(480, 12, 128, device="cuda", dtype=dtype)
angles = torch.randn(480, 64, device="cuda")
freqs = torch.polar(torch.ones_like(angles), angles)
freqs_expanded = freqs.unsqueeze(-2)
def reference(x):
value = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))
return torch.view_as_real(value * freqs_expanded).flatten(-2).type_as(x)
expected_q, expected_k = reference(q), reference(k)
prepared = prepare_fused_qk_complex_rope_inplace(freqs)
actual_q, actual_k = apply_fused_qk_complex_rope_inplace(q, k, prepared)
atol = 2 * torch.finfo(dtype).eps
torch.testing.assert_close(actual_q, expected_q, rtol=0, atol=atol)
torch.testing.assert_close(actual_k, expected_k, rtol=0, atol=atol)
def test_vision_rope_inplace_rejects_non_complex_frequencies(self):
with self.assertRaises(ValueError):
prepare_fused_qk_complex_rope_inplace(torch.randn(8, 64, device="cuda"))
def test_normalize_and_patchify(self):
torch.manual_seed(5)
image = torch.randn(2, 3, 17, 19, device="cuda")
test/registered/unit/models/test_kimi_k25.py
+242 -7
View File
@@ -3,10 +3,12 @@
from types import SimpleNamespace
from unittest.mock import Mock, patch
import numpy as np
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from sglang.srt.managers.schedule_batch import (
Modality,
@@ -14,9 +16,19 @@ from sglang.srt.managers.schedule_batch import (
MultimodalInputs,
MultimodalProcessorOutput,
)
from sglang.srt.models.kimi_k25 import KimiK25ForConditionalGeneration
from sglang.srt.models.kimi_k25 import (
KimiK25ForConditionalGeneration,
mm_projection_auto,
)
from sglang.srt.models.kimi_vl_moonvit import tpool_patch_merger
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin
from sglang.srt.multimodal.processors.kimi_k25 import (
KimiGPUProcessorWrapper,
_ensure_chw_rgb,
_expand_image_token_ids,
_resize_bicubic_if_needed,
_resize_images_by_source_shape,
)
from sglang.srt.runtime_context import get_parallel
@@ -63,14 +75,12 @@ def _image_item(feature, grid_thw):
def test_kimi_gpu_preprocess_batches_only_source_compatible_images():
torch.manual_seed(0)
indexed_images = [
(0, torch.randn(3, 32, 24)),
(1, torch.randn(3, 32, 24)),
(2, torch.randn(3, 28, 20)),
(0, torch.randint(0, 256, (3, 32, 24), dtype=torch.uint8)),
(1, torch.randint(0, 256, (3, 32, 24), dtype=torch.uint8)),
(2, torch.randint(0, 256, (3, 28, 20), dtype=torch.uint8)),
]
expected = [
F.interpolate(
image.unsqueeze(0), size=(16, 12), mode="bicubic", align_corners=False
)
_resize_bicubic_if_needed(image.unsqueeze(0), 16, 12)
for _, image in indexed_images
]
real_interpolate = F.interpolate
@@ -92,6 +102,209 @@ def test_kimi_gpu_preprocess_batches_only_source_compatible_images():
torch.testing.assert_close(result, reference)
def test_kimi_resize_tracks_the_checkpoint_processors_pil_bicubic():
# Plain F.interpolate skips PIL's implicit antialiasing on downscale and
# drifts far outside 8-bit rounding; photo-like content, not pure noise.
rng = np.random.default_rng(0)
yy, xx = np.mgrid[0:512, 0:512].astype(np.float32)
plane = np.clip(
128
+ 90 * np.sin(xx / 40) * np.cos(yy / 55)
+ 40 * ((xx // 37 + yy // 41) % 2)
+ rng.normal(0, 6, (512, 512)),
0,
255,
)
array = np.stack([plane, np.roll(plane, 7, 0), np.roll(plane, 13, 1)], -1).astype(
np.uint8
)
pil = torch.from_numpy(
np.asarray(Image.fromarray(array).resize((252, 252), Image.BICUBIC)).astype(
np.float32
)
).permute(2, 0, 1)
source = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
resized = _resize_bicubic_if_needed(source, 252, 252)
assert resized.shape == (1, 3, 252, 252)
torch.testing.assert_close(resized, resized.round())
assert resized.min() >= 0.0 and resized.max() <= 255.0
# Within a couple of 8-bit levels of PIL; the non-antialiased resize is off
# by an order of magnitude more, which is the regression this guards.
assert (resized[0] - pil).abs().max() <= 4.0
naive = F.interpolate(
source.float(), size=(252, 252), mode="bicubic", align_corners=False
)
assert (naive[0] - pil).abs().max() > 20.0
def test_kimi_resize_is_a_dtype_only_cast_when_already_at_target():
image = torch.randint(0, 256, (1, 3, 16, 12), dtype=torch.uint8)
resized = _resize_bicubic_if_needed(image, 16, 12)
assert resized.dtype == torch.float32
torch.testing.assert_close(resized, image.float())
def test_kimi_expands_one_placeholder_per_image_from_existing_ids():
# 7 is the placeholder; the two images claim 3 and 2 tokens.
input_ids = [1, 7, 2, 7, 3]
expanded = _expand_image_token_ids(
input_ids, image_token_id=7, image_token_counts=[3, 2]
)
assert expanded.tolist() == [[1, 7, 7, 7, 2, 7, 7, 3]]
def test_kimi_expansion_rejects_a_placeholder_count_mismatch():
with pytest.raises(ValueError, match="placeholder"):
_expand_image_token_ids([1, 7, 2], image_token_id=7, image_token_counts=[3, 2])
def test_kimi_expansion_matches_the_base_retokenize_avoidance_rebuild():
# preserve_processor_input_ids skips the base rebuild, which is only safe
# while both produce the same sequence. Reference is the original loop.
def reference(original_ids, counts, placeholder):
rebuilt, next_image = [], 0
for token_id in original_ids:
if token_id == placeholder:
rebuilt.extend([placeholder] * counts[next_image])
next_image += 1
else:
rebuilt.append(token_id)
return rebuilt
rng = np.random.default_rng(0)
for n_images in (1, 3, 8):
# Placeholder 7 is below the random range, so only the inserted
# positions count as placeholders.
ids = rng.integers(100, 5000, 400).tolist()
for slot in range(n_images):
ids.insert(slot * 37 + 5, 7)
counts = rng.integers(1, 400, n_images).tolist()
expected = reference(ids, counts, 7)
assert BaseMultimodalProcessor._expand_input_ids(ids, counts, 7) == expected
wrapper = _expand_image_token_ids(
ids, image_token_id=7, image_token_counts=counts
)
assert wrapper.flatten().tolist() == expected
def test_kimi_cpu_fallback_keeps_the_request_tokens():
# preserve_processor_input_ids disables the base rebuild on every path.
hf_processor = Mock()
hf_processor.media_processor.media_tokens_calculator = Mock(return_value=3)
hf_processor.return_value = {"input_ids": torch.tensor([[99, 99, 99]])}
wrapper = KimiGPUProcessorWrapper.__new__(KimiGPUProcessorWrapper)
wrapper._hf_processor = hf_processor
wrapper._image_token = "<|media_pad|>"
wrapper._image_token_id = 7
out = wrapper._cpu_call(
"a<|media_pad|>b", ["img"], original_input_ids=[1, 7, 2], medias=None
)
# Not the [99, 99, 99] the HF processor returned.
assert out["input_ids"].flatten().tolist() == [1, 7, 7, 7, 2]
def test_kimi_cpu_fallback_falls_back_to_the_hf_tokens_without_request_ids():
hf_processor = Mock()
hf_processor.media_processor.media_tokens_calculator = Mock(return_value=3)
hf_processor.return_value = {"input_ids": torch.tensor([[99, 99, 99]])}
wrapper = KimiGPUProcessorWrapper.__new__(KimiGPUProcessorWrapper)
wrapper._hf_processor = hf_processor
wrapper._image_token = "<|media_pad|>"
wrapper._image_token_id = 7
out = wrapper._cpu_call("a<|media_pad|>b", ["img"], medias=None)
assert out["input_ids"].flatten().tolist() == [99, 99, 99]
def test_kimi_refuses_already_normalized_float_pixels():
with pytest.raises(ValueError, match="uint8"):
_ensure_chw_rgb(torch.rand(3, 8, 8))
def test_kimi_placeholder_count_only_reads_real_token_ids():
count = KimiGridMMDataMixin.count_image_placeholders
assert count([1, 7, 2, 7], 7) == 2
assert count(torch.tensor([[1, 7, 2]]), 7) == 1
assert count([1, 2, 3], 7) == 0
# A prompt string carries no token IDs, so the caller must not take the
# tokenized fast path.
assert count("<|media_pad|>", 7) is None
def test_kimi_single_frame_pool_matches_the_temporal_mean():
torch.manual_seed(0)
x = torch.randn(1 * 4 * 4, 8)
grid_thws = torch.tensor([[1, 4, 4]])
(merged,) = tpool_patch_merger(x, grid_thws)
# t == 1 skips the mean; it must stay bit-identical to averaging one frame.
reference = (
x.view(1, 2, 2, 2, 2, 8).permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0)
)
assert torch.equal(merged, reference.view(4, 4, 8))
def test_kimi_multi_frame_pool_still_averages_across_frames():
torch.manual_seed(0)
x = torch.randn(3 * 4 * 4, 8)
grid_thws = torch.tensor([[3, 4, 4]])
(merged,) = tpool_patch_merger(x, grid_thws)
reference = (
x.view(3, 2, 2, 2, 2, 8).permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0)
)
assert merged.shape == (4, 4, 8)
torch.testing.assert_close(merged, reference.view(4, 4, 8))
class _IdentityProjector(nn.Module):
"""Stands in for K2VLMultiModalProjector, which is never None in production."""
def __init__(self):
super().__init__()
self.seen = None
def forward(self, x):
self.seen = x
return x
def test_kimi_projection_returns_one_flattened_feature_tensor():
torch.manual_seed(0)
per_image = [torch.randn(4, 2, 8), torch.randn(6, 2, 8)]
packed = mm_projection_auto(_IdentityProjector(), per_image)
assert packed.shape == (20, 8)
torch.testing.assert_close(packed, torch.cat(per_image, dim=0).reshape(-1, 8))
def test_kimi_projection_does_not_copy_a_single_image():
single = torch.randn(4, 2, 8)
projector = _IdentityProjector()
packed = mm_projection_auto(projector, [single])
# The projector must receive the tensor itself, not a one-element cat of it.
assert projector.seen.data_ptr() == single.data_ptr()
assert packed.data_ptr() == single.data_ptr()
def test_dp_helper_supports_moonvit3d_packed_embeddings_on_tp1():
tower = _MoonViT3dTower()
pixel_values = torch.randn(4, 2)
@@ -248,6 +461,28 @@ def test_kimi_k25_encoder_dp_selects_packed_moonvit_contract():
assert callable(run_dp.call_args.kwargs["load_local_pixel_values"])
def test_kimi_non_dp_keeps_grid_thws_on_the_host():
model = KimiK25ForConditionalGeneration.__new__(KimiK25ForConditionalGeneration)
nn.Module.__init__(model)
model.use_data_parallel = False
model.vision_tower = _MoonViT3dTower()
# Not the host, so a stray .to(tower.device) shows up without a GPU.
model.vision_tower.device = torch.device("meta")
model.mm_projector = _IdentityProjector()
items = [_image_item(torch.randn(4, 2), [[1, 2, 2]])]
with get_parallel().override(
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
), patch(
"sglang.srt.models.kimi_k25.get_server_args",
return_value=SimpleNamespace(tp_size=1),
):
model.get_image_feature(items)
# A device copy would cost one sync per .tolist() inside MoonViT3d.
assert model.vision_tower.grid_thws.device.type == "cpu"
def test_kimi_lazy_ipc_feature_skips_scheduler_reconstruction():
proxy = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy)
proxy.reconstruct_on_target_device = Mock()