(cherry picked from commit 094abd5e90af77507acd7c91513f69b5b29b685b)
This commit is contained in:
abing
2026-09-20 22:07:24 +08:00
committed by minke.yu
parent c2059c4fb2
commit fc954b7e08
6 changed files with 1019 additions and 0 deletions
+79
View File
@@ -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
@@ -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()
@@ -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()
@@ -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()
@@ -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()
@@ -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()