From fc954b7e0824de51f896b044355d0ba3ffe7c9c5 Mon Sep 17 00:00:00 2001 From: abing Date: Wed, 16 Sep 2026 00:26:02 -0700 Subject: [PATCH] add test (cherry picked from commit 094abd5e90af77507acd7c91513f69b5b29b685b) --- python/sglang/test/dsv41_cp_test_utils.py | 79 +++++ .../disaggregation/test_dsv41_dspark_pd.py | 286 ++++++++++++++++ .../test_fake_transfer_chunk_progress.py | 46 +++ .../models/test_deepseek_v41_multimodal_cp.py | 316 ++++++++++++++++++ .../unit/models/test_deepseek_v41_text_cp.py | 219 ++++++++++++ .../test_deepseek_v41_text_only_pd.py | 73 ++++ 6 files changed, 1019 insertions(+) create mode 100644 python/sglang/test/dsv41_cp_test_utils.py create mode 100644 test/registered/unit/disaggregation/test_dsv41_dspark_pd.py create mode 100644 test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py create mode 100644 test/registered/unit/models/test_deepseek_v41_multimodal_cp.py create mode 100644 test/registered/unit/models/test_deepseek_v41_text_cp.py create mode 100644 test/registered/unit/server_args/test_deepseek_v41_text_only_pd.py diff --git a/python/sglang/test/dsv41_cp_test_utils.py b/python/sglang/test/dsv41_cp_test_utils.py new file mode 100644 index 000000000..5c6430078 --- /dev/null +++ b/python/sglang/test/dsv41_cp_test_utils.py @@ -0,0 +1,79 @@ +"""Small CPU tensors; production CP slicing/gather, mocked collective transport.""" + +from contextlib import ExitStack, contextmanager, nullcontext +from types import SimpleNamespace as NS +from unittest.mock import patch + +import torch + +from sglang.srt.layers.cp.interleave import InterleaveCPStrategy +from sglang.srt.layers.cp.padding import pad_logical_token_to_physical +from sglang.srt.model_executor.forward_batch_info import ForwardMode + +CP = "sglang.srt.layers.cp" + + +@contextmanager +def cp_context(size, rank, lengths=(3, 6), prefix_lengths=(7, 13)): + """Keep real interleave indexing/padding; replace only runtime context.""" + strategy = InterleaveCPStrategy(size) + parallel = NS(attn_cp_size=size, attn_cp_rank=rank, attn_cp_group=None) + batch = NS( + forward_mode=ForwardMode.EXTEND, + input_ids=torch.arange(1, sum(lengths) + 1), + positions=torch.cat( + [ + torch.arange(prefix, prefix + length) + for prefix, length in zip(prefix_lengths, lengths) + ] + ), + extend_seq_lens_cpu=list(lengths), + extend_prefix_lens_cpu=list(prefix_lengths), + mm_inputs=None, + spec_info=None, + ) + batch.attn_cp_metadata = strategy.build_metadata( + sum(lengths), [p + n for p, n in zip(prefix_lengths, lengths)], list(lengths) + ) + with ExitStack() as stack: + for module in ("base", "utils", "padding", "interleave"): + stack.enter_context( + patch(CP + "." + module + ".get_parallel", return_value=parallel) + ) + stack.enter_context(patch(CP + ".utils.get_cp_strategy", return_value=strategy)) + stack.enter_context( + patch(CP + ".padding.get_cp_padding_align_size", return_value=size) + ) + stack.enter_context( + patch( + CP + ".utils.get_moe_a2a_backend", return_value=NS(is_none=lambda: True) + ) + ) + pad_logical_token_to_physical(batch.attn_cp_metadata) + yield strategy, batch + + +@contextmanager +def simulated_collective(strategy, batch, global_tensor): + """Inject peer buffers into all-gather; retain production unpadding/reordering.""" + physical = max(batch.attn_cp_metadata.per_rank_actual_token) + buffers = [] + for rank in range(strategy.cp_size): + buf = global_tensor.new_zeros((physical, *global_tensor.shape[1:])) + local = global_tensor[rank :: strategy.cp_size] + buf[: len(local)] = local + buffers.append(buf) + + def gather(output, local): + torch.testing.assert_close(local, buffers[strategy.cp_rank], rtol=0, atol=0) + output.copy_(torch.cat(buffers)) + + with ( + patch( + CP + ".interleave.use_symmetric_memory", + side_effect=lambda *a, **k: nullcontext(), + ), + patch(CP + ".interleave.is_allocation_symmetric", return_value=False), + patch(CP + ".interleave.attn_cp_all_gather_into_tensor", side_effect=gather), + ): + yield diff --git a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py new file mode 100644 index 000000000..69e58c505 --- /dev/null +++ b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py @@ -0,0 +1,286 @@ +import asyncio +import copy +import json +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock, patch + +from sglang.srt.disaggregation.base.conn import StateType +from sglang.srt.disaggregation.common.conn import ( + CommonKVBootstrapServer, + CommonKVManager, +) +from sglang.srt.disaggregation.decode import DecodePreallocQueue +from sglang.srt.disaggregation.utils import get_dsv41_spec_layout +from sglang.srt.mem_cache.common import retraction_backup +from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool +from sglang.srt.runtime_context import get_context +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +def make_layout(): + args = SimpleNamespace( + mla_compression_ratios=[0, 2, 1], + kv_layer_ids=[1, 2], + kv_item_lens=[512, 1024], + state_types=[StateType.SWA, StateType.DSV4_REQUEST_STATE, StateType.SWA], + state_item_lens=[[512], [32768], [512]], + ) + with get_context().override_server_args( + speculative_algorithm="DSPARK", speculative_num_draft_tokens=6 + ): + return get_dsv41_spec_layout(args) + + +class TestDSV41DSparkPD(CustomTestCase): + def test_bootstrap_validates_before_caching(self): + layout = make_layout() + cases = [("matching", layout, layout, 4, True), ("legacy", None, None, 2, True)] + for key, value in ( + ("num_draft_tokens", 5), + ("kv_layer_ids", [2, 1]), + ("kv_item_lens", [256, 1024]), + ("state_types", ["swa", "c128_state"]), + ("state_item_lens", [[512], [8192], [512]]), + ): + different = copy.deepcopy(layout) + different[key] = value + cases.append((key, layout, different, 4, False)) + cases += [ + ("prefill_only", None, layout, 4, False), + ("decode_only_or_old_prefill", layout, None, 4, False), + ("tp_mismatch", layout, layout, 2, False), + ] + for name, local, peer, tp_size, supported in cases: + with self.subTest(name=name): + manager = object.__new__(CommonKVManager) + manager.prefill_info_table = {} + manager.kv_args = SimpleNamespace(page_size=256) + manager.kv_cache_dtype_str = "fp8_e4m3" + manager.dsv41_spec_layout = local + manager.attn_tp_size = 4 + manager.dcp_size = 1 + manager._resolve_rank_mapping = Mock() + response = Mock(status_code=200) + response.json.return_value = dict( + attn_tp_size=tp_size, + attn_cp_size=1, + dp_size=1, + pp_size=1, + page_size=256, + kv_cache_dtype="fp8_e4m3", + follow_bootstrap_room=True, + dsv41_spec_layout=peer, + ) + with patch( + "sglang.srt.disaggregation.common.conn.requests.get", + return_value=response, + ) as fetch: + if supported: + self.assertTrue( + manager.try_ensure_parallel_info("prefill:8998") + ) + self.assertTrue( + manager.try_ensure_parallel_info("prefill:8998") + ) + fetch.assert_called_once() + else: + with self.assertRaisesRegex( + RuntimeError, "DeepSeek-V4.1 DSpark PD" + ): + manager.try_ensure_parallel_info("prefill:8998") + self.assertFalse(manager.prefill_info_table) + manager._resolve_rank_mapping.assert_not_called() + + def test_python_bootstrap_preserves_layout_and_rejects_mixed_ranks(self): + with patch.object(CommonKVBootstrapServer, "run"): + server = CommonKVBootstrapServer("127.0.0.1", 8998) + layout = make_layout() + payload = dict( + attn_tp_size=1, + attn_tp_rank=0, + attn_cp_size=1, + attn_cp_rank=0, + attn_dp_size=1, + attn_dp_rank=0, + pp_size=1, + pp_rank=0, + system_dp_size=1, + system_dp_rank=0, + rank_ip="127.0.0.1", + rank_port=1234, + page_size=256, + kv_cache_dtype="fp8_e4m3", + dsv41_spec_layout=layout, + ) + request = Mock(json=AsyncMock(return_value=payload)) + self.assertEqual(asyncio.run(server._handle_route_put(request)).status, 200) + query = Mock( + query={ + key: "-1" + for key in ( + "prefill_dp_rank", + "prefill_cp_rank", + "target_tp_rank", + "target_pp_rank", + ) + } + ) + response = asyncio.run(server._handle_route_get(query)) + self.assertEqual(json.loads(response.text)["dsv41_spec_layout"], layout) + payload["dsv41_spec_layout"] = None + self.assertEqual(asyncio.run(server._handle_route_put(request)).status, 400) + self.assertEqual(server._registered_count, 1) + self.assertEqual(server.dsv41_spec_layout, layout) + + def test_retraction_recomputes_from_prefill_and_replays_boundary_token(self): + pool = object.__new__(DeepSeekV4TokenToKVPool) + pool.compression_ratios = [0, 2, 1] + pool.device = "cuda" + allocator = Mock(get_kvcache=Mock(return_value=pool)) + for algorithm in (None, "DSPARK"): + with ( + self.subTest(algorithm=algorithm), + get_context().override_server_args(speculative_algorithm=algorithm), + patch("torch.get_device_module") as device_module, + ): + req = SimpleNamespace( + output_ids=[7, 8], + bootstrap_host="prefill", + time_stats=Mock(), + offload_kv_cache=Mock(), + ) + request_pool = Mock() + self.assertTrue( + retraction_backup( + req, Mock(), request_pool, allocator, "cpu_tensor" + ) + ) + queue = SimpleNamespace( + token_to_kv_pool_allocator=allocator, + _check_if_req_exceed_kv_capacity=Mock(return_value=False), + _create_receiver_and_enqueue=Mock(), + _resolve_prefill_dp_rank=Mock(return_value=0), + retracted_queue=[], + pending_reqs=[], + ) + DecodePreallocQueue.add(queue, req, is_retracted=True) + if algorithm == "DSPARK": + req.offload_kv_cache.assert_not_called() + device_module.return_value.synchronize.assert_called_once_with( + "cuda" + ) + self.assertEqual(req.output_ids, [7]) + self.assertEqual(req.pd_rebootstrap_forced_output_id, 8) + self.assertTrue(req.pd_rebootstrap_in_progress) + queue._create_receiver_and_enqueue.assert_called_once_with( + req, is_rebootstrap=True + ) + self.assertFalse(queue.retracted_queue) + else: + req.offload_kv_cache.assert_called_once_with( + request_pool, allocator + ) + device_module.assert_not_called() + self.assertEqual(req.output_ids, [7, 8]) + self.assertEqual(queue.retracted_queue, [req]) + + +class TestDSV41CPPDHandshake(CustomTestCase): + def make(self, rank=0, hybrid=True): + m = object.__new__(CommonKVManager) + m.prefill_info_table = {} + m.kv_args = SimpleNamespace(page_size=256, engine_rank=rank) + m.kv_cache_dtype_str = "fp8_e4m3" + m.dsv41_spec_layout = {"kv_item_lens": [512], "state_item_lens": [[32768]]} + m.attn_tp_size = 4 + m.attn_cp_size = 1 + m.attn_cp_rank = 0 + m.dcp_size = 1 + m.is_mla_backend = False + m.is_hybrid_mla_backend = hybrid + m.enable_all_cp_ranks_for_transfer = True + m.pp_size = 1 + m.pp_rank = 0 + return m + + def fetch(self, m, tp, cp, layout=None): + response = Mock(status_code=200) + response.json.return_value = dict( + attn_tp_size=tp, + attn_cp_size=cp, + dp_size=1, + pp_size=1, + page_size=256, + kv_cache_dtype="fp8_e4m3", + follow_bootstrap_room=True, + dsv41_spec_layout=layout or m.dsv41_spec_layout, + ) + with patch( + "sglang.srt.disaggregation.common.conn.requests.get", return_value=response + ): + return m.try_ensure_parallel_info("prefill:8761") + + def test_cp4_maps_all_shards_to_each_decode_rank(self): + for rank in range(4): + m = self.make(rank) + self.assertTrue(self.fetch(m, 1, 4)) + info = m.prefill_info_table["prefill:8761"] + self.assertEqual(info.target_tp_ranks, [0]) + self.assertEqual(info.target_cp_ranks, [0, 1, 2, 3]) + self.assertEqual(info.required_prefill_response_num, 4) + self.assertEqual(info.required_dst_info_num, 4) + + def test_dsv4_pool_is_classified_as_mla(self): + from sglang.srt.disaggregation.utils import is_mla_backend + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + + pool = object.__new__(DeepSeekV4TokenToKVPool) + self.assertTrue(is_mla_backend(pool)) + m = self.make(hybrid=False) + m.is_mla_backend = is_mla_backend(pool) + self.assertTrue(self.fetch(m, 1, 4)) + self.assertEqual( + m.prefill_info_table["prefill:8761"].required_prefill_response_num, 4 + ) + + def test_cp2_tp2_maps_corresponding_tp_and_both_cp_ranks(self): + for rank in range(4): + m = self.make(rank) + self.assertTrue(self.fetch(m, 2, 2)) + info = m.prefill_info_table["prefill:8761"] + self.assertEqual(info.target_tp_ranks, [rank // 2]) + self.assertEqual(info.target_cp_ranks, [0, 1]) + self.assertEqual(info.required_prefill_response_num, 2) + + def test_plain_tp4_unchanged(self): + m = self.make(3) + self.assertTrue(self.fetch(m, 4, 1)) + info = m.prefill_info_table["prefill:8761"] + self.assertEqual(info.target_tp_ranks, [3]) + self.assertEqual(info.target_cp_ranks, [0]) + + def test_unequal_model_tp_rejected(self): + for tp, cp in [(2, 1), (1, 2), (1, 8)]: + m = self.make() + with self.assertRaisesRegex(RuntimeError, "same TP size"): + self.fetch(m, tp, cp) + self.assertFalse(m.prefill_info_table) + + def test_nonhybrid_cp_mismatch_rejected(self): + m = self.make(hybrid=False) + with self.assertRaisesRegex(RuntimeError, "same TP size"): + self.fetch(m, 1, 4) + + def test_layout_mismatch_still_rejected(self): + m = self.make() + with self.assertRaisesRegex(RuntimeError, "layout mismatch"): + self.fetch(m, 1, 4, {"kv_item_lens": [1024], "state_item_lens": [[32768]]}) + self.assertFalse(m.prefill_info_table) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py b/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py new file mode 100644 index 000000000..840d77a5a --- /dev/null +++ b/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py @@ -0,0 +1,46 @@ +"""Chunk continuation when fake PD transfer skips shared radix insertion.""" + +import unittest +from types import SimpleNamespace as NS +from unittest.mock import Mock + +import torch + +from sglang.srt.mem_cache.common import maybe_cache_unfinished_req +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestFakeTransferChunkProgress(CustomTestCase): + def test_chunk_progress_without_shared_insert(self): + slots = torch.arange(16385, dtype=torch.int32).reshape(1, -1) + cache = NS( + req_to_token_pool=NS(req_to_token=slots), cache_unfinished_req=Mock() + ) + req = NS( + skip_radix_cache_insert=True, + kv=NS(req_pool_idx=0, cache_protected_len=0), + get_fill_ids=lambda: range(16384), + ) + maybe_cache_unfinished_req(req, cache, chunked=True) + self.assertEqual(16385 - len(req.prefix_indices), 1) + self.assertEqual(req.kv.cache_protected_len, 0) + cache.cache_unfinished_req.assert_not_called() + self.assertEqual(req.prefix_indices.dtype, torch.int64) + slots[0, 0] = -1 + self.assertEqual(req.prefix_indices[0].item(), 0) + req.get_fill_ids = lambda: range(16385) + maybe_cache_unfinished_req(req, cache, chunked=True) + self.assertEqual(len(req.prefix_indices), 16385) + + def test_real_transfer_preserves_cache_path(self): + req = NS(skip_radix_cache_insert=False) + cache = NS(cache_unfinished_req=Mock()) + maybe_cache_unfinished_req(req, cache, chunked=True) + cache.cache_unfinished_req.assert_called_once_with(req, chunked=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py b/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py new file mode 100644 index 000000000..d0a52b632 --- /dev/null +++ b/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py @@ -0,0 +1,316 @@ +"""V4.1 image/text CP input contracts; vision and model compute are mocked.""" + +import unittest +from contextlib import contextmanager +from types import SimpleNamespace as NS +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.managers.schedule_batch import MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.model_executor.runner.eager_runner import EagerRunner +from sglang.srt.models.deepseek_v4 import MM_PAD_SHIFT_VALUE, DeepseekV4ForCausalLM +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.dsv41_cp_test_utils import cp_context, simulated_collective +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +MODEL = "sglang.srt.models.deepseek_v4" +RUNNER = "sglang.srt.model_executor.runner.eager_runner" +IMAGE_ID = 129264 + + +class TestDSV41MultimodalCP(CustomTestCase): + def test_image_spans_cross_ranks_before_shard_and_gather(self): + # Two image spans with distinct cache hashes; first request is text-only. + original = torch.tensor( + [ + 7, + 8, + 9, + MM_PAD_SHIFT_VALUE + 11, + MM_PAD_SHIFT_VALUE + 11, + 10, + MM_PAD_SHIFT_VALUE + 23, + MM_PAD_SHIFT_VALUE + 23, + 12, + ] + ) + normalized = torch.tensor( + [7, 8, 9, IMAGE_ID, IMAGE_ID, 10, IMAGE_ID, IMAGE_ID, 12] + ) + full = torch.arange(36, dtype=torch.float32).reshape(9, 4) + # Distinct image features expose using text embeddings or wrong row order. + full[3:5] += 1000 + full[6:8] += 2000 + for size in (2, 4): + for rank in range(size): + with ( + self.subTest(size=size, rank=rank), + cp_context(size, rank) as (strategy, batch), + ): + batch.input_ids = original.clone() + batch.mm_inputs = [None, MultimodalInputs(mm_items=[])] + model = NS( + vision=object(), + config=NS(image_token_id=IMAGE_ID), + get_input_embeddings=Mock( + side_effect=AssertionError( + "Raw image hashes entered text embeddings" + ) + ), + _prepare_mm_embeddings=Mock(return_value=full), + capture_aux_hidden_states=False, + pp_group=NS(is_last_rank=True), + lm_head=object(), + logits_processor=Mock(return_value="ok"), + ) + model.prepare_language_model_inputs = lambda ids, fb, emb: ( + DeepseekV4ForCausalLM.prepare_language_model_inputs( + model, ids, fb, emb + ) + ) + + def body(ids, positions, fb, input_embeds): + model._prepare_mm_embeddings.assert_called_once_with( + batch.input_ids, batch + ) + n = len(normalized[rank::size]) + torch.testing.assert_close(ids[:n], normalized[rank::size]) + torch.testing.assert_close(input_embeds[:n], full[rank::size]) + torch.testing.assert_close( + positions[:n], batch.positions[rank::size] + ) + self.assertFalse( + (fb.input_ids_global >= MM_PAD_SHIFT_VALUE).any().item() + ) + return input_embeds + + model.model = body + with ( + simulated_collective(strategy, batch, full), + patch(RUNNER + ".torch.cuda.current_stream", return_value=None), + ): + result = EagerRunner._execute_extend_cp( + NS(model_runner=NS(model=model)), batch, {} + ) + self.assertEqual(result, "ok") + torch.testing.assert_close( + model.logits_processor.call_args.args[0], normalized + ) + torch.testing.assert_close( + model.logits_processor.call_args.args[1], full + ) + torch.testing.assert_close(batch.input_ids, original) + self.assertFalse(hasattr(batch, "input_ids_global")) + + def test_chunk_prefix_metadata_and_scheduler_hashes_survive_embedder(self): + for prefixes, lengths in (([0, 0], [3, 6]), ([16384, 127], [3, 6])): + with self.subTest(prefixes=prefixes): + ids = torch.tensor([7, 8, 9] + [MM_PAD_SHIFT_VALUE + 17] * 6) + original = ids.clone() + image = MultimodalInputs(mm_items=[]) + batch = NS( + mm_inputs=[None, image], + extend_prefix_lens_cpu=prefixes, + extend_seq_lens_cpu=lengths, + ) + full = torch.arange(27, dtype=torch.float32).reshape(9, 3) + embedding = object() + model = NS(get_input_embeddings=lambda: embedding) + + def embed(**kwargs): + self.assertEqual(kwargs["extend_prefix_lens"], prefixes) + self.assertEqual(kwargs["extend_seq_lens"], lengths) + self.assertIs(kwargs["mm_inputs_list"][1], image) + self.assertEqual(kwargs["mm_inputs_list"][0].mm_items, []) + self.assertIs(kwargs["input_embedding"], embedding) + self.assertNotEqual(kwargs["input_ids"].data_ptr(), ids.data_ptr()) + kwargs["input_ids"].zero_() + return full, {} + + with patch(MODEL + ".embed_mm_inputs", side_effect=embed) as mocked: + result = DeepseekV4ForCausalLM._prepare_mm_embeddings( + model, ids, batch + ) + mocked.assert_called_once() + self.assertIs(result, full) + self.assertIs(batch.mm_input_embeds, full) + torch.testing.assert_close(ids, original) + + def test_vision_enabled_text_batch_skips_image_encoder(self): + for mm_inputs in (None, [None, None], []): + with self.subTest(mm_inputs=mm_inputs): + ids = torch.tensor([4, 5, 6]) + model = NS( + vision=object(), + config=NS(image_token_id=IMAGE_ID), + _prepare_mm_embeddings=Mock(), + ) + batch = NS(forward_mode=ForwardMode.EXTEND, mm_inputs=mm_inputs) + result, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + model, ids, batch + ) + torch.testing.assert_close(result, ids) + self.assertIsNone(embeds) + model._prepare_mm_embeddings.assert_not_called() + + def test_decode_idle_and_verify_preserve_vocab_ids(self): + for mode in (ForwardMode.DECODE, ForwardMode.IDLE, ForwardMode.TARGET_VERIFY): + with self.subTest(mode=mode): + ids = torch.tensor([4, IMAGE_ID, 6]) + model = NS( + vision=object(), + config=NS(image_token_id=IMAGE_ID), + _prepare_mm_embeddings=Mock(), + ) + batch = NS(forward_mode=mode, mm_inputs=None) + result, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + model, ids, batch + ) + self.assertIs(result, ids) + self.assertIsNone(embeds) + model._prepare_mm_embeddings.assert_not_called() + + def test_image_embedding_failure_does_not_mutate_scheduler_ids(self): + ids = torch.tensor([7, MM_PAD_SHIFT_VALUE + 12, 8]) + original = ids.clone() + batch = NS( + mm_inputs=[MultimodalInputs(mm_items=[])], + extend_prefix_lens_cpu=[0], + extend_seq_lens_cpu=[3], + ) + model = NS(get_input_embeddings=lambda: object()) + + def fail(**kwargs): + kwargs["input_ids"].zero_() + raise RuntimeError("vision failure") + + with patch(MODEL + ".embed_mm_inputs", side_effect=fail): + with self.assertRaisesRegex(RuntimeError, "vision failure"): + DeepseekV4ForCausalLM._prepare_mm_embeddings(model, ids, batch) + torch.testing.assert_close(ids, original) + self.assertFalse(hasattr(batch, "mm_input_embeds")) + + +class TestDSV41MultimodalInputs(CustomTestCase): + def setUp(self): + self.ids = torch.tensor( + [7, MM_PAD_SHIFT_VALUE + 12, MM_PAD_SHIFT_VALUE + 12, 9, 10] + ) + self.original = self.ids.clone() + self.embeds = torch.arange(15, dtype=torch.float32).reshape(5, 3) + self.model = NS( + vision=object(), + config=NS(image_token_id=129264), + _prepare_mm_embeddings=Mock(return_value=self.embeds), + ) + self.batch = NS( + input_ids=self.ids, + forward_mode=ForwardMode.EXTEND, + mm_inputs=[MultimodalInputs(mm_items=[])], + ) + + def test_prepare_global_embeddings_and_normalized_ids(self): + ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch + ) + self.assertEqual(ids.tolist(), [7, 129264, 129264, 9, 10]) + self.assertIs(embeds, self.embeds) + self.model._prepare_mm_embeddings.assert_called_once_with(self.ids, self.batch) + self.assertTrue(torch.equal(self.ids, self.original)) + + def test_reject_preembedded_images(self): + with self.assertRaisesRegex(ValueError, "Cannot combine"): + DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch, self.embeds + ) + + def test_text_only_model_keeps_existing_embeddings(self): + self.model.vision = None + ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch, self.embeds + ) + self.assertIs(ids, self.ids) + self.assertIs(embeds, self.embeds) + self.model._prepare_mm_embeddings.assert_not_called() + + def test_text_subclass_without_vision_module(self): + del self.model.vision + ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch, self.embeds + ) + self.assertIs(ids, self.ids) + self.assertIs(embeds, self.embeds) + + def test_decode_keeps_vocabulary_ids(self): + self.batch.forward_mode = ForwardMode.DECODE + ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch + ) + self.assertIs(ids, self.ids) + self.assertIsNone(embeds) + self.model._prepare_mm_embeddings.assert_not_called() + + def test_embedding_does_not_mutate_scheduler_hashes(self): + self.batch.extend_prefix_lens_cpu = [0] + self.batch.extend_seq_lens_cpu = [5] + self.model.get_input_embeddings = lambda: None + + def embed(**kwargs): + kwargs["input_ids"].zero_() + return (self.embeds, {}) + + with patch("sglang.srt.models.deepseek_v4.embed_mm_inputs", side_effect=embed): + result = DeepseekV4ForCausalLM._prepare_mm_embeddings( + self.model, self.ids, self.batch + ) + self.assertTrue(torch.equal(self.ids, self.original)) + self.assertIs(result, self.batch.mm_input_embeds) + + def test_cp_runner_prepares_before_sharding_and_uses_model_ids_for_logits(self): + normalized = torch.tensor([7, 129264, 129264, 9, 10]) + self.batch.positions = torch.arange(5) + self.model.prepare_language_model_inputs = lambda ids, batch, emb: ( + DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, ids, batch, emb + ) + ) + self.model.get_input_embeddings = Mock( + side_effect=AssertionError("Raw hashes must not enter text embedding") + ) + self.model.model = Mock(return_value=self.embeds[1::4]) + self.model.capture_aux_hidden_states = False + self.model.pp_group = NS(is_last_rank=True) + self.model.lm_head = object() + self.model.logits_processor = Mock(return_value="ok") + + @contextmanager + def shard(embeds, positions, batch, ids): + self.assertIs(embeds, self.embeds) + self.assertTrue(torch.equal(ids, normalized)) + yield (embeds[1::4], positions[1::4], ids[1::4]) + + runner = NS(model_runner=NS(model=self.model)) + module = "sglang.srt.model_executor.runner.eager_runner" + with ( + patch(module + ".cp_shard_model_inputs", side_effect=shard), + patch(module + ".cp_gather_after_forward", return_value=self.embeds), + patch(module + ".torch.cuda.current_stream", return_value=None), + ): + result = EagerRunner._execute_extend_cp(runner, self.batch, {}) + self.assertEqual(result, "ok") + args, kwargs = self.model.model.call_args + self.assertEqual(args[0].tolist(), [129264]) + self.assertTrue(torch.equal(kwargs["input_embeds"], self.embeds[1::4])) + self.assertTrue( + torch.equal(self.model.logits_processor.call_args.args[0], normalized) + ) + self.assertTrue(torch.equal(self.batch.input_ids, self.original)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_deepseek_v41_text_cp.py b/test/registered/unit/models/test_deepseek_v41_text_cp.py new file mode 100644 index 000000000..7689b118c --- /dev/null +++ b/test/registered/unit/models/test_deepseek_v41_text_cp.py @@ -0,0 +1,219 @@ +"""Pure-language V4.1 CP input, padding and DSpark state regressions.""" + +import unittest +from types import SimpleNamespace as NS +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.layers.cp.utils import ( + cp_gather_after_forward, + cp_shard_model_inputs, + is_cp_active, +) +from sglang.srt.model_executor.runner.eager_runner import EagerRunner +from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.dsv41_cp_test_utils import cp_context, simulated_collective +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +RUNNER = "sglang.srt.model_executor.runner.eager_runner" + + +class TestDSV41TextCP(CustomTestCase): + def test_interleave_roundtrip_mixed_lengths_prefix_and_padding(self): + for size in (2, 4): + for length in (4, 5, 9, 127, 128, 129): + for rank in range(size): + with ( + self.subTest(size=size, length=length, rank=rank), + cp_context(size, rank, (1, length - 1), (0, 16384)) as ( + strategy, + batch, + ), + ): + embeddings = torch.arange( + length * 3, dtype=torch.float32 + ).reshape(length, 3) + original_ids = batch.input_ids.clone() + with cp_shard_model_inputs( + embeddings, batch.positions, batch, batch.input_ids + ) as (local, positions, ids): + count = len(embeddings[rank::size]) + torch.testing.assert_close( + local[:count], embeddings[rank::size] + ) + torch.testing.assert_close( + positions[:count], batch.positions[rank::size] + ) + torch.testing.assert_close( + ids[:count], batch.input_ids[rank::size] + ) + self.assertEqual( + torch.count_nonzero(local[count:]).item(), 0 + ) + self.assertEqual(torch.count_nonzero(ids[count:]).item(), 0) + with simulated_collective(strategy, batch, embeddings): + restored = cp_gather_after_forward(local, batch) + torch.testing.assert_close( + restored, embeddings, rtol=0, atol=0 + ) + torch.testing.assert_close(batch.input_ids, original_ids) + self.assertFalse(hasattr(batch, "input_ids_global")) + + def test_speculative_state_and_global_ids_restored_on_exception(self): + for size in (2, 4): + for rank in range(size): + for had_global in (False, True): + with ( + self.subTest(size=size, rank=rank, had_global=had_global), + cp_context(size, rank) as (_, batch), + ): + full = torch.arange(36, dtype=torch.float32).reshape(9, 4) + batch.spec_info = NS(hidden_states=full) + previous = object() + if had_global: + batch.input_ids_global = previous + with self.assertRaisesRegex(RuntimeError, "injected"): + with cp_shard_model_inputs( + full, batch.positions, batch, batch.input_ids + ): + n = len(full[rank::size]) + torch.testing.assert_close( + batch.spec_info.hidden_states[:n], full[rank::size] + ) + # Global MoE IDs are in rank-major order, with padding. + physical = sum( + batch.attn_cp_metadata.per_rank_actual_token + ) + padded = batch.input_ids.new_zeros(physical) + padded[:9] = batch.input_ids + expected = torch.cat( + [padded[r::size] for r in range(size)] + ) + torch.testing.assert_close( + batch.input_ids_global, expected + ) + raise RuntimeError("injected") + self.assertIs(batch.spec_info.hidden_states, full) + if had_global: + self.assertIs(batch.input_ids_global, previous) + else: + self.assertFalse(hasattr(batch, "input_ids_global")) + + def test_short_prompt_falls_back_from_cp(self): + with cp_context(4, 0, (1, 2), (0, 0)) as (_, batch): + self.assertFalse(is_cp_active(batch)) + + def test_runner_text_embedding_and_preembedded_paths(self): + for preembedded in (False, True): + for rank in range(4): + with ( + self.subTest(preembedded=preembedded, rank=rank), + cp_context(4, rank) as (strategy, batch), + ): + full = torch.arange(27, dtype=torch.float32).reshape(9, 3) + embedding = Mock(return_value=full) + model = NS( + vision=None, + _prepare_mm_embeddings=Mock( + side_effect=AssertionError("Text must not invoke vision") + ), + get_input_embeddings=Mock(return_value=embedding), + pp_group=NS(is_last_rank=True), + lm_head=object(), + capture_aux_hidden_states=False, + logits_processor=Mock(return_value="ok"), + ) + model.prepare_language_model_inputs = lambda ids, fb, emb: ( + DeepseekV4ForCausalLM.prepare_language_model_inputs( + model, ids, fb, emb + ) + ) + + def body(ids, positions, fb, input_embeds): + n = len(full[rank::4]) + torch.testing.assert_close(ids[:n], batch.input_ids[rank::4]) + torch.testing.assert_close( + positions[:n], batch.positions[rank::4] + ) + torch.testing.assert_close(input_embeds[:n], full[rank::4]) + return input_embeds + + model.model = body + with ( + simulated_collective(strategy, batch, full), + patch(RUNNER + ".torch.cuda.current_stream", return_value=None), + ): + result = EagerRunner._execute_extend_cp( + NS(model_runner=NS(model=model)), + batch, + {"input_embeds": full} if preembedded else {}, + ) + self.assertEqual(result, "ok") + if preembedded: + model.get_input_embeddings.assert_not_called() + else: + embedding.assert_called_once_with(batch.input_ids) + model._prepare_mm_embeddings.assert_not_called() + args = model.logits_processor.call_args.args + torch.testing.assert_close(args[0], batch.input_ids) + torch.testing.assert_close(args[1], full) + + def test_dspark_aux_tensor_and_list_gathered_without_pre_norm_override(self): + for as_list in (False, True): + with self.subTest(as_list=as_list), cp_context(4, 2) as (strategy, batch): + full = torch.arange(27, dtype=torch.float32).reshape(9, 3) + local = strategy.shard_hidden_states(full, batch) + aux = [local.clone(), local.clone()] if as_list else local.clone() + model = NS( + get_input_embeddings=lambda: lambda ids: full, + model=Mock(return_value=((local, local.clone()), aux)), + capture_aux_hidden_states=True, + pp_group=NS(is_last_rank=True), + lm_head=object(), + logits_processor=Mock(return_value="ok"), + ) + with ( + simulated_collective(strategy, batch, full), + patch(RUNNER + ".torch.cuda.current_stream", return_value=None), + ): + EagerRunner._execute_extend_cp( + NS(model_runner=NS(model=model)), batch, {} + ) + args, kwargs = model.logits_processor.call_args + torch.testing.assert_close(args[1], full) + for tensor in args[4] if as_list else [args[4]]: + torch.testing.assert_close(tensor, full) + self.assertNotIn("hidden_states_before_norm", kwargs) + + def test_target_hidden_states_before_norm_preserved_without_dspark_aux(self): + with cp_context(4, 1) as (strategy, batch): + full = torch.arange(27, dtype=torch.float32).reshape(9, 3) + local = strategy.shard_hidden_states(full, batch) + model = NS( + get_input_embeddings=lambda: lambda ids: full, + model=Mock(return_value=(local, local.clone())), + capture_aux_hidden_states=False, + pp_group=NS(is_last_rank=True), + lm_head=object(), + logits_processor=Mock(return_value="ok"), + ) + with ( + simulated_collective(strategy, batch, full), + patch(RUNNER + ".torch.cuda.current_stream", return_value=None), + ): + EagerRunner._execute_extend_cp( + NS(model_runner=NS(model=model)), batch, {} + ) + torch.testing.assert_close( + model.logits_processor.call_args.kwargs["hidden_states_before_norm"], + full, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_deepseek_v41_text_only_pd.py b/test/registered/unit/server_args/test_deepseek_v41_text_only_pd.py new file mode 100644 index 000000000..e07a4a6c5 --- /dev/null +++ b/test/registered/unit/server_args/test_deepseek_v41_text_only_pd.py @@ -0,0 +1,73 @@ +"""V4.1 language-model-only PD configuration validation.""" + +import unittest +from types import SimpleNamespace as NS +from unittest.mock import patch + +from sglang.srt.arg_groups import model_hook +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestDSV41TextOnlyPDPolicy(CustomTestCase): + def validate( + self, + model_type="deepseek_v41", + mode="prefill", + arch="DeepseekV4ForCausalLM", + **flags, + ): + cfg = NS( + language_model_only=True, + encoder_only=False, + language_only=False, + enable_prefix_mm_cache=False, + enable_broadcast_mm_inputs_process=False, + mm_enable_dp_encoder=False, + disaggregation_mode=mode, + ) + for name, value in flags.items(): + setattr(cfg, name, value) + model = NS(hf_config=NS(model_type=model_type, architectures=[arch])) + args = NS( + LANGUAGE_MODEL_ONLY_ARCHITECTURES=ServerArgs.LANGUAGE_MODEL_ONLY_ARCHITECTURES + ) + with ( + patch.object(model_hook, "resolving_view", return_value=cfg), + patch.object(model_hook, "model_config_of", return_value=model), + ): + model_hook.handle_language_model_only(args) + + def test_v41_modes(self): + for mode in ("null", "prefill", "decode"): + with self.subTest(mode=mode): + self.validate(mode=mode) + + def test_other_models_still_reject_pd(self): + with self.assertRaisesRegex(ValueError, "incompatible"): + self.validate(model_type="cosmos3", arch="Cosmos3ForConditionalGeneration") + + def test_encoder_options_still_rejected(self): + for flag in ( + "encoder_only", + "language_only", + "enable_prefix_mm_cache", + "enable_broadcast_mm_inputs_process", + "mm_enable_dp_encoder", + ): + with ( + self.subTest(flag=flag), + self.assertRaisesRegex(ValueError, "cannot be combined"), + ): + self.validate(**{flag: True}) + + def test_unknown_arch_rejected(self): + with self.assertRaisesRegex(ValueError, "does not support"): + self.validate(arch="UnknownArchitecture") + + +if __name__ == "__main__": + unittest.main()