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

[EPD][VLM] Fix Kimi-VL 2D encoder grids (#32104)

Browse Source 
Signed-off-by: Xiaojun Zhang <zhangxiaojunhust@gmail.com>
This commit is contained in:
Xiaojun(Robin) Zhang
2026-07-29 11:46:17 +08:00
committed by GitHub
parent d6fcfe02d6
commit 1af0167493
3 changed files with 319 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/disaggregation/test_epd_disaggregation.py
+180 -1
View File
@@ -34,9 +34,10 @@ from sglang.test.vlm_utils import (
# Omni model for local testing; override via env var EPD_OMNI_MODEL
DEFAULT_OMNI_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
QWEN35_27B_MODEL = "Qwen/Qwen3.5-27B"
KIMI_VL_MODEL = "moonshotai/Kimi-VL-A3B-Instruct"
register_cuda_ci(est_time=97, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=300, stage="base-c", runner_config="4-gpu-h100")
@unittest.skipIf(
@@ -750,6 +751,184 @@ class TestEPDDisaggregationOneEncoder(MMMUMixin, PDDisaggregationServerBase):
print(f"Error killing process: {e}")
class TestEPDDisaggregationKimiVL(PDDisaggregationServerBase):
"""Regression test for Kimi-VL two-dimensional image grids in E/PD mode."""
model = KIMI_VL_MODEL
model_args = [
"--context-length=8192",
"--dtype=bfloat16",
"--mem-fraction-static=0.40",
]
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.process_encode = None
cls.encode_port = f"{int(cls.lb_port) + 307}"
cls.encode_url = f"http://{cls.base_host}:{cls.encode_port}"
cls.start_encode()
prefill_thread = threading.Thread(target=cls.start_prefill)
decode_thread = threading.Thread(target=cls.start_decode)
prefill_thread.start()
decode_thread.start()
prefill_thread.join()
decode_thread.join()
cls.wait_server_ready(cls.encode_url + "/health", process=cls.process_encode)
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_encode(cls):
encode_args = [
"--trust-remote-code",
"--encoder-only",
"--encoder-transfer-backend",
"zmq_to_scheduler",
"--tp",
"1",
"--base-gpu-id",
"0",
"--port",
cls.encode_port,
*cls.model_args,
]
cls.process_encode = popen_launch_server(
cls.model,
base_url=cls.encode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=encode_args,
)
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--language-only",
"--encoder-urls",
cls.encode_url,
"--encoder-transfer-backend",
"zmq_to_scheduler",
"--disaggregation-mode",
"prefill",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"1",
"--port",
cls.prefill_port,
*cls.model_args,
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_server(
cls.model,
base_url=cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--disaggregation-bootstrap-port",
cls.bootstrap_port,
"--tp",
"1",
"--base-gpu-id",
"2",
"--port",
cls.decode_port,
*cls.model_args,
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_server(
cls.model,
base_url=cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
@classmethod
def tearDownClass(cls):
try:
super().tearDownClass()
finally:
if cls.process_encode:
try:
kill_process_tree(cls.process_encode.pid)
except Exception as e:
print(f"Error killing encode process: {e}")
def test_multi_image_chat_completion(self):
client = openai.Client(api_key="sk-123456", base_url=f"{self.lb_url}/v1")
response = client.chat.completions.create(
model="default",
messages=[
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": IMAGE_MAN_IRONING_URL},
"modalities": "multi-images",
},
{
"type": "image_url",
"image_url": {"url": IMAGE_SGL_LOGO_URL},
"modalities": "multi-images",
},
{
"type": "text",
"text": "Describe each image separately.",
},
],
},
],
temperature=0,
max_tokens=128,
)
self.assertEqual(response.choices[0].message.role, "assistant")
text = response.choices[0].message.content
self.assertIsInstance(text, str)
self.assertGreater(len(text), 0)
text_lower = text.lower()
self.assertTrue(
any(
word in text_lower
for word in ("man", "person", "car", "vehicle", "suv", "iron")
),
f"First image was not described correctly: {text}",
)
self.assertTrue(
any(
word in text_lower
for word in ("logo", "sglang", "graphic", "stylized", "letter")
),
f"Second image was not described correctly: {text}",
)
for name, process in (
("encoder", self.process_encode),
("prefill", self.process_prefill),
("decode", self.process_decode),
("router", self.process_lb),
):
self.assertIsNone(
process.poll(),
f"{name} process exited with code {process.returncode}",
)
@unittest.skipIf(
is_in_ci(),
"Qwen3.5 EPD image/video test runs locally only",
test/registered/unit/disaggregation/test_encode_server.py
+108
View File
@@ -0,0 +1,108 @@
import pickle
import unittest
from types import SimpleNamespace
import numpy as np
import torch
from sglang.srt.disaggregation.encode_receiver import EmbeddingData
from sglang.srt.disaggregation.encode_server import MMEncoder, _get_mm_grid_dim
from sglang.srt.managers.schedule_batch import Modality
from sglang.srt.utils.common import safe_pickle_loads
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestKimiVLEPDGrid(unittest.TestCase):
@staticmethod
def _make_encoder(model_type="kimi_vl"):
encoder = MMEncoder.__new__(MMEncoder)
encoder.model_type = model_type
encoder.model_config = SimpleNamespace(
hf_config=SimpleNamespace(
vision_config=SimpleNamespace(merge_kernel_size=(2, 2))
)
)
return encoder
def test_kimi_vl_prefers_and_normalizes_hw_grid(self):
mm_inputs = {
"image_grid_hws": np.array([[40, 60]], dtype=np.int64),
"image_grid_thw": torch.tensor([[1, 20, 30]]),
"grid_thws": torch.tensor([[1, 10, 15]]),
}
grid = _get_mm_grid_dim(mm_inputs, Modality.IMAGE, "kimi_vl")
self.assertIsInstance(grid, torch.Tensor)
torch.testing.assert_close(grid, torch.tensor([[40, 60]]))
def test_kimi_k25_keeps_thw_grid_preference(self):
mm_inputs = {
"image_grid_hws": np.array([[40, 60]], dtype=np.int64),
"grid_thws": np.array([[1, 10, 15]], dtype=np.int64),
}
grid = _get_mm_grid_dim(mm_inputs, Modality.IMAGE, "kimi_k25")
torch.testing.assert_close(grid, torch.tensor([[1, 10, 15]]))
def test_kimi_vl_2d_grid_counting_and_slicing(self):
encoder = self._make_encoder()
grids = torch.tensor([[40, 60], [20, 40]])
embedding = torch.arange(800 * 2).reshape(800, 2)
self.assertEqual(
encoder.get_num_patches(grids[0], Modality.IMAGE),
2400,
)
self.assertEqual(
encoder.get_num_tokens(grids[0], Modality.IMAGE),
600,
)
slices = encoder.slice_embedding(embedding, grids, Modality.IMAGE)
self.assertEqual([item.shape for item in slices], [(600, 2), (200, 2)])
torch.testing.assert_close(slices[0], embedding[:600])
torch.testing.assert_close(slices[1], embedding[600:])
def test_kimi_3d_grid_remains_supported(self):
encoder = self._make_encoder()
grid = torch.tensor([1, 40, 60])
self.assertEqual(encoder.get_num_patches(grid, Modality.IMAGE), 2400)
self.assertEqual(encoder.get_num_tokens(grid, Modality.IMAGE), 600)
def test_kimi_k25_3d_patch_counting_is_unchanged(self):
encoder = self._make_encoder("kimi_k25")
grid = torch.tensor([2, 12, 16])
self.assertEqual(encoder.get_num_patches(grid, Modality.IMAGE), 384)
self.assertEqual(encoder.get_num_tokens(grid, Modality.IMAGE), 48)
def test_grid_metadata_is_safe_to_deserialize(self):
grid = _get_mm_grid_dim(
{"image_grid_hws": np.array([[40, 60]], dtype=np.int64)},
Modality.IMAGE,
"kimi_vl",
)
embedding_data = EmbeddingData(
req_id="test-request",
num_parts=1,
part_idx=0,
grid_dim=grid,
modality=Modality.IMAGE,
embedding=torch.zeros((600, 4)),
)
restored = safe_pickle_loads(
pickle.dumps(embedding_data.copy_without_embedding())
)
torch.testing.assert_close(restored.grid_dim, torch.tensor([[40, 60]]))
if __name__ == "__main__":
unittest.main()