diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 5a7080ba6..fdad40e34 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1513,6 +1513,27 @@ class ShmPointerMMData: self.shm.close() def __getstate__(self): + if not hasattr(self, "shm") or self.shm is None: + tensor = getattr(self, "cpu_tensor", None) + if tensor is None: + tensor = getattr(self, "tensor", None) + if tensor is None: + raise RuntimeError( + "ShmPointerMMData cannot recreate shared memory without tensor" + ) + + cpu_tensor = tensor.cpu().contiguous() + self.shape = cpu_tensor.shape + self.dtype = cpu_tensor.dtype + + nbytes = cpu_tensor.numel() * cpu_tensor.element_size() + self.shm = shared_memory.SharedMemory(create=True, size=nbytes) + try: + shm_view = np.ndarray((nbytes,), dtype=np.uint8, buffer=self.shm.buf) + shm_view[:] = cpu_tensor.view(torch.uint8).numpy().flatten() + finally: + self.shm.close() + return { "shm_name": self.shm.name, "shape": self.shape, @@ -1521,12 +1542,15 @@ class ShmPointerMMData: def __setstate__(self, state): self.shm_name = state["shm_name"] + self.shape = state["shape"] + self.dtype = state["dtype"] + self.shm = None shm_handle = shared_memory.SharedMemory(name=self.shm_name) try: self.tensor = ( - torch.frombuffer(shm_handle.buf, dtype=state["dtype"]) - .reshape(state["shape"]) + torch.frombuffer(shm_handle.buf, dtype=self.dtype) + .reshape(self.shape) .clone() ) finally: diff --git a/test/registered/distributed/test_dp_attention.py b/test/registered/distributed/test_dp_attention.py index e36f7a2e3..455685e46 100644 --- a/test/registered/distributed/test_dp_attention.py +++ b/test/registered/distributed/test_dp_attention.py @@ -3,6 +3,7 @@ from types import SimpleNamespace import requests +from sglang.lang.chat_template import get_chat_template_by_model_path from sglang.srt.environ import envs from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci @@ -13,6 +14,7 @@ from sglang.test.kits.radix_cache_server_kit import run_radix_attention_test from sglang.test.kits.regex_constrained_kit import TestRegexConstrainedMixin from sglang.test.run_eval import run_eval from sglang.test.test_utils import ( + DEFAULT_IMAGE_URL, DEFAULT_MLA_MODEL_NAME_FOR_TEST, DEFAULT_MODEL_NAME_FOR_TEST_MLA, DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN, @@ -182,5 +184,52 @@ class TestDPAttentionDP2TP2DeepseekV3MTP( self.assertGreater(avg_spec_accept_length, 2.5) +class TestDPAttentionDP2TP2VLM(CustomTestCase): + @classmethod + def setUpClass(cls): + # TODO(FlamingoPg): Use Kimi-VL-A3B-Instruct temporarily + # cauz Qwen3-VL use mrope which has bug in DP attention mode + cls.model = "moonshotai/Kimi-VL-A3B-Instruct" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.image_url = DEFAULT_IMAGE_URL + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp", + "2", + "--enable-dp-attention", + "--dp", + "2", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_vlm_generate(self): + chat_template = get_chat_template_by_model_path(self.model) + prompt = f"{chat_template.image_token}What is in this image?" + response = requests.post( + self.base_url + "/generate", + json={ + "text": prompt, + "image_data": [self.image_url], + "sampling_params": { + "temperature": 0, + "max_new_tokens": 16, + }, + }, + ) + response.raise_for_status() + response_json = response.json() + print(response_json) + self.assertIn("output_ids", response_json) + self.assertGreater(len(response_json["output_ids"]), 0) + + if __name__ == "__main__": unittest.main()