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

Fix disagg PP MTP for GLM-5.2 (#39378)

Browse Source 
Co-authored-by: Julien Lin <jullin@nvidia.com>
Co-authored-by: YAMY1234 <74099316+YAMY1234@users.noreply.github.com>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
This commit is contained in:
Po-Han Huang (NVIDIA)
2026-09-19 13:55:07 -07:00
committed by GitHub
co-authored by Julien Lin YAMY1234 Yangmin Li
parent 8139a1740e
commit 3a64faa1f2
12 changed files with 627 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/disaggregation/test_pp_mla_kv_transfer.py
+258
View File
@@ -0,0 +1,258 @@
"""Unit tests for NIXL KV transfer from a prefill with pp_size > 1 to a decode
peer that is not pipelined, on plain-MLA models (MLATokenToKVPool)."""
import unittest
from types import SimpleNamespace
from sglang.srt.disaggregation.common.conn import CommonKVManager
from sglang.srt.disaggregation.nixl.conn import NixlKVManager
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
NUM_LAYERS = 78
ITEM_LEN = 576
def _pp_layer_split(total: int, pp: int) -> list:
"""(start_layer, num_layers) per PP rank, remainder to the last ranks."""
base, rem = divmod(total, pp)
out, start = [], 0
for rank in range(pp):
num = base + (1 if rank >= pp - rem else 0)
out.append((start, num))
start += num
return out
class _StubKVManager:
"""Carries the real span primitive, so the cases exercise the shared mapping
rather than a reimplementation of it."""
_mla_kv_entry_span_with_pp = CommonKVManager._mla_kv_entry_span_with_pp
get_mla_kv_ptrs_with_pp = CommonKVManager.get_mla_kv_ptrs_with_pp
def __init__(
self,
*,
pp_size: int,
prefill_start_layer: int,
n_src: int,
is_mla_backend: bool = True,
is_hybrid_mla_backend: bool = False,
kv_layer_ids: list = (),
mla_compression_ratios=None,
item_len: int = ITEM_LEN,
):
self.pp_size = pp_size
self.is_mla_backend = is_mla_backend
self.is_hybrid_mla_backend = is_hybrid_mla_backend
self.kv_args = SimpleNamespace(
prefill_start_layer=prefill_start_layer,
kv_layer_ids=list(kv_layer_ids),
kv_item_lens=[item_len] * n_src,
mla_compression_ratios=mla_compression_ratios,
)
def _peer(*, n_dst: int, dst_kv_layer_ids: list = (), item_len: int = ITEM_LEN):
return SimpleNamespace(
dst_kv_layer_ids=list(dst_kv_layer_ids),
dst_kv_item_lens=[item_len] * n_dst,
)
def _resolve(manager, peer, n_src, n_dst):
return NixlKVManager._build_transfer_dst_indices(
manager, peer_info=peer, n_src=n_src, n_dst=n_dst
)
class TestPpPrefillToUnpipelinedDecode(CustomTestCase):
"""Bug regression: a pp>1 prefill paired with a pp=1 decode never transferred
KV on NIXL, and failed after /health was green so the request hung rather
than the launch failing. Each stage must write into its own layer range of
the peer's pool."""
def test_stages_tile_the_decode_pool_exactly_once(self):
covered = []
for start, num in _pp_layer_split(NUM_LAYERS, 4):
with self.subTest(start_layer=start):
indices = _resolve(
_StubKVManager(pp_size=4, prefill_start_layer=start, n_src=num),
_peer(n_dst=NUM_LAYERS),
num,
NUM_LAYERS,
)
self.assertEqual(indices, list(range(start, start + num)))
covered += indices
self.assertEqual(sorted(covered), list(range(NUM_LAYERS)))
self.assertEqual(len(covered), len(set(covered)))
def test_decode_side_draft_regions_are_never_targeted(self):
"""Decode-only speculative decoding appends draft buffers after the
target's, so the span must stay inside the target layer range."""
start, num = _pp_layer_split(NUM_LAYERS, 4)[3]
indices = _resolve(
_StubKVManager(pp_size=4, prefill_start_layer=start, n_src=num),
_peer(n_dst=NUM_LAYERS + 1),
num,
NUM_LAYERS + 1,
)
self.assertEqual(indices, list(range(start, start + num)))
self.assertNotIn(NUM_LAYERS, indices)
def test_cell_geometry_disagreement_raises(self):
"""A per-layer item_len mismatch means the peers did not build the same
KV geometry, which must fail loudly instead of corrupting the pool."""
start, num = _pp_layer_split(NUM_LAYERS, 4)[3]
with self.assertRaisesRegex(RuntimeError, "geometry differs"):
_resolve(
_StubKVManager(pp_size=4, prefill_start_layer=start, n_src=num),
_peer(n_dst=NUM_LAYERS, item_len=ITEM_LEN + 80),
num,
NUM_LAYERS,
)
class TestExistingPairingIsUnchanged(CustomTestCase):
"""All layouts outside heterogeneous plain MLA retain existing pairing."""
def test_unpipelined_prefill_uses_positional_pairing(self):
for n_dst in (NUM_LAYERS, NUM_LAYERS + 1):
with self.subTest(n_dst=n_dst):
self.assertEqual(
_resolve(
_StubKVManager(
pp_size=1, prefill_start_layer=0, n_src=NUM_LAYERS
),
_peer(n_dst=n_dst),
NUM_LAYERS,
n_dst,
),
list(range(NUM_LAYERS)),
)
def test_matched_pp_uses_identity_pairing(self):
for start, num in _pp_layer_split(NUM_LAYERS, 4):
with self.subTest(start_layer=start):
self.assertEqual(
_resolve(
_StubKVManager(pp_size=4, prefill_start_layer=start, n_src=num),
_peer(n_dst=num),
num,
num,
),
list(range(num)),
)
def test_one_sided_layer_ids_are_rejected(self):
start, num = _pp_layer_split(NUM_LAYERS, 4)[2]
with self.assertRaisesRegex(RuntimeError, "both PD peers or neither"):
_resolve(
_StubKVManager(
pp_size=4,
prefill_start_layer=start,
n_src=num,
kv_layer_ids=range(start, start + num),
),
_peer(n_dst=NUM_LAYERS),
num,
NUM_LAYERS,
)
with self.assertRaisesRegex(RuntimeError, "both PD peers or neither"):
_resolve(
_StubKVManager(pp_size=4, prefill_start_layer=start, n_src=num),
_peer(n_dst=NUM_LAYERS, dst_kv_layer_ids=range(NUM_LAYERS)),
num,
NUM_LAYERS,
)
def test_heterogeneous_non_plain_mla_requires_layer_ids(self):
start, num = _pp_layer_split(NUM_LAYERS, 4)[2]
for label, kwargs, n_dst in (
("mha", {"is_mla_backend": False}, 2 * NUM_LAYERS),
("hybrid_mla", {"is_hybrid_mla_backend": True}, NUM_LAYERS),
(
"compressed_mla",
{"mla_compression_ratios": [4] * NUM_LAYERS},
2 * NUM_LAYERS,
),
):
with self.subTest(pool=label):
with self.assertRaisesRegex(
RuntimeError, "PP-heterogeneous transfer requires layer ids"
):
_resolve(
_StubKVManager(
pp_size=4,
prefill_start_layer=start,
n_src=num,
**kwargs,
),
_peer(n_dst=n_dst),
num,
n_dst,
)
class TestMatchedPpWithDraftRegions(CustomTestCase):
"""Derived property: bootstrap admits a peer at our pp or at 1, so under
matched pp a decode-side draft buffer makes n_src != n_dst and reaches the
span. Stage 0 degenerates to the identity and stays correct; every stage
above it must be rejected by the bound rather than writing at an offset the
peer does not have."""
def test_rank0_identity_span_stays_correct(self):
start, num = _pp_layer_split(NUM_LAYERS, 4)[0]
self.assertEqual(start, 0)
self.assertEqual(
_resolve(
_StubKVManager(pp_size=4, prefill_start_layer=start, n_src=num),
_peer(n_dst=num + 1),
num,
num + 1,
),
list(range(num)),
)
def test_later_ranks_are_rejected(self):
for start, num in _pp_layer_split(NUM_LAYERS, 4)[1:]:
with self.subTest(start_layer=start):
with self.assertRaisesRegex(
RuntimeError, "PP-heterogeneous transfer requires layer ids"
):
_resolve(
_StubKVManager(pp_size=4, prefill_start_layer=start, n_src=num),
_peer(n_dst=num + 1),
num,
num + 1,
)
class TestPointerAndIndexViewsAgree(CustomTestCase):
"""Derived property: get_mla_kv_ptrs_with_pp (pointer view) and
_build_transfer_dst_indices (index view) read the same span, so a change to
one must not silently diverge from the other."""
def test_views_select_the_same_destination_entries(self):
dst_ptrs = [9000 + i for i in range(NUM_LAYERS)]
for start, num in _pp_layer_split(NUM_LAYERS, 4):
with self.subTest(start_layer=start):
manager = _StubKVManager(
pp_size=4, prefill_start_layer=start, n_src=num
)
src_ptrs = [1000 + i for i in range(num)]
_, sliced_dst, count = manager.get_mla_kv_ptrs_with_pp(
src_ptrs, dst_ptrs
)
indices = _resolve(manager, _peer(n_dst=NUM_LAYERS), num, NUM_LAYERS)
self.assertEqual(count, num)
self.assertEqual(sliced_dst, [dst_ptrs[j] for j in indices])
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_auxiliary_output.py
+45
View File
@@ -25,6 +25,7 @@ from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import publish, reset_context
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.eagle_info import EagleDraftInput
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import published_topology
@@ -689,6 +690,50 @@ def test_pipeline_parallel_auxiliary_output_round_trip():
receiver.future_map.stash.assert_called_once()
@pytest.mark.parametrize("dsa_topk_indices", [None, torch.tensor([[2, 5, 7]])])
def test_pipeline_parallel_dsa_seed_round_trip(dsa_topk_indices):
draft_input = EagleDraftInput(
topk_p=torch.tensor([[0.8, 0.2]]),
topk_index=torch.tensor([[11, 13]]),
hidden_states=torch.tensor([[1.0, 2.0]]),
dsa_topk_indices=dsa_topk_indices,
)
result = GenerationBatchResult(
logits_output=None,
next_token_ids=torch.tensor([17]),
next_draft_input=draft_input,
)
batch = SimpleNamespace(
return_logprob=False,
req_pool_indices=torch.tensor([3]),
input_ids=torch.tensor([5]),
spec_info=None,
)
tensors = Scheduler._pp_prepare_tensor_dict(
object.__new__(Scheduler), result, batch
)
if dsa_topk_indices is None:
assert "draft_dsa_topk_indices" not in tensors
else:
assert torch.equal(tensors["draft_dsa_topk_indices"], dsa_topk_indices)
receiver = object.__new__(Scheduler)
receiver.pp_group = SimpleNamespace(is_first_rank=False)
receiver.future_map = SimpleNamespace(stash=Mock())
Scheduler._pp_prep_batch_result(
receiver,
batch,
PPBatchMetadata(can_run_cuda_graph=True),
PPProxyTensors(tensors),
)
if dsa_topk_indices is None:
assert batch.spec_info.dsa_topk_indices is None
else:
assert torch.equal(batch.spec_info.dsa_topk_indices, dsa_topk_indices)
def test_pipeline_parallel_auxiliary_output_stays_packed_before_first_rank():
device_output = DeviceOutput(torch.tensor([1.0]))
result = GenerationBatchResult(
test/registered/unit/models/test_deepseek_pp_mtp_embedding.py
+58
View File
@@ -0,0 +1,58 @@
"""Pipeline-stage embedding ownership for DeepSeek/GLM NextN."""
import unittest
from types import SimpleNamespace
import torch
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.models.deepseek_v2 import pp_stage_needs_embedding
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")
def _pp_group(*, rank: int, size: int) -> SimpleNamespace:
return SimpleNamespace(is_first_rank=rank == 0, is_last_rank=rank == size - 1)
class TestDeepseekPPStageEmbedding(CustomTestCase):
def test_without_speculative_decoding_only_the_first_stage_embeds(self):
for rank in range(4):
with self.subTest(rank=rank):
self.assertEqual(
pp_stage_needs_embedding(_pp_group(rank=rank, size=4), None),
rank == 0,
)
def test_speculative_decoding_adds_the_last_stage(self):
wants = [
pp_stage_needs_embedding(_pp_group(rank=rank, size=4), "EAGLE")
for rank in range(4)
]
self.assertEqual(wants, [True, False, False, True])
def test_middle_stages_never_embed(self):
self.assertFalse(pp_stage_needs_embedding(_pp_group(rank=2, size=4), "EAGLE"))
def test_without_pipeline_parallelism_the_single_stage_embeds_either_way(self):
for spec in (None, "EAGLE"):
with self.subTest(spec=spec):
self.assertTrue(
pp_stage_needs_embedding(_pp_group(rank=0, size=1), spec)
)
def test_pp_missing_layer_registers_no_parameters(self):
placeholder = PPMissingLayer()
self.assertEqual(dict(placeholder.named_parameters()), {})
def test_a_real_embedding_registers_weight_under_its_prefix(self):
model = torch.nn.Module()
model.embed_tokens = torch.nn.Embedding(8, 4)
params = dict(model.named_parameters())
self.assertIn("embed_tokens.weight", params)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_server_args.py
+116 -1
View File
@@ -13,7 +13,12 @@ import msgspec
import msgspec.structs
import sglang.srt.server_args as server_args_module
from sglang.srt.arg_groups import parallel_hook, pd_disaggregation_hook, serving_hook
from sglang.srt.arg_groups import (
parallel_hook,
pd_disaggregation_hook,
serving_hook,
validation_hook,
)
from sglang.srt.arg_groups.attention_hook import (
handle_attention_backend_compatibility,
handle_deterministic_inference,
@@ -68,6 +73,7 @@ from sglang.srt.arg_groups.serving_hook import (
)
from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
from sglang.srt.arg_groups.validation_hook import (
check_pipeline_parallel_compat,
check_two_batch_overlap,
)
from sglang.srt.entrypoints.sidecar import (
@@ -2364,6 +2370,115 @@ class TestCudaGraphConfigDataclassAccess(CustomTestCase):
self.assertEqual(config.compiler, "eager")
class TestPipelineParallelCompat(CustomTestCase):
"""Features supported with `pipeline-parallel-size > 1`."""
_SUPPORTED_ARCH = "GlmMoeDsaForCausalLM"
@staticmethod
def _cfg(**overrides):
cfg = dict(
disable_overlap_schedule=True,
speculative_algorithm=None,
enable_multi_layer_eagle=False,
disaggregation_mode="prefill",
min_free_slots_delay=None,
)
cfg.update(overrides)
return SimpleNamespace(**cfg)
def test_overlap_schedule_must_be_off(self):
with self.assertRaisesRegex(AssertionError, "overlap schedule"):
check_pipeline_parallel_compat(self._cfg(disable_overlap_schedule=False))
def test_no_speculative_decoding_is_fine(self):
check_pipeline_parallel_compat(self._cfg())
def test_eagle_is_allowed_on_prefill(self):
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="EAGLE"),
model_architecture=self._SUPPORTED_ARCH,
)
def test_eagle_is_rejected_outside_prefill(self):
for mode in ("decode", "null"):
with self.subTest(disaggregation_mode=mode):
with self.assertRaisesRegex(AssertionError, "prefill nodes"):
check_pipeline_parallel_compat(
self._cfg(
speculative_algorithm="EAGLE", disaggregation_mode=mode
),
model_architecture=self._SUPPORTED_ARCH,
)
def test_eagle_is_rejected_for_unsupported_model(self):
with self.assertRaisesRegex(AssertionError, "DeepSeek/GLM/Qwen3.5 models"):
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="EAGLE"),
model_architecture="LlamaForCausalLM",
)
def test_supported_architectures(self):
for architecture in (
"DeepseekV2ForCausalLM",
"DeepseekV3ForCausalLM",
"DeepseekV32ForCausalLM",
"GlmMoeDsaForCausalLM",
"Qwen3_5ForCausalLM",
"Qwen3_5MoeForCausalLM",
"Qwen3_5ForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
):
with self.subTest(architecture=architecture):
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="EAGLE"),
model_architecture=architecture,
)
def test_pp_spec_env_gate_allows_aggregate_and_rejects_pd(self):
cfg = self._cfg(
speculative_algorithm="EAGLE",
disaggregation_mode="null",
speculative_adaptive=False,
enable_dp_attention=False,
)
with patch.object(
validation_hook.envs.SGLANG_ENABLE_PP_SPEC, "get", return_value=True
):
check_pipeline_parallel_compat(cfg, model_architecture="LlamaForCausalLM")
with self.assertRaisesRegex(AssertionError, "SGLANG_ENABLE_PP_SPEC"):
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="EAGLE"),
model_architecture=self._SUPPORTED_ARCH,
)
def test_nextn_resolves_to_eagle_and_is_allowed(self):
"""`--speculative-algorithm NEXTN` has collapsed to EAGLE by the time the
validation hook runs, so the check only ever sees the resolved name."""
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="eagle"),
model_architecture=self._SUPPORTED_ARCH,
)
def test_non_eagle_speculative_algorithms_are_rejected(self):
with self.assertRaisesRegex(AssertionError, "only supports EAGLE"):
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="EAGLE3"),
model_architecture=self._SUPPORTED_ARCH,
)
def test_multi_layer_eagle_is_rejected(self):
with self.assertRaisesRegex(AssertionError, "only supports EAGLE"):
check_pipeline_parallel_compat(
self._cfg(speculative_algorithm="EAGLE", enable_multi_layer_eagle=True),
model_architecture=self._SUPPORTED_ARCH,
)
def test_min_free_slots_delay_is_rejected(self):
with self.assertRaisesRegex(AssertionError, "min-free-slots-delay"):
check_pipeline_parallel_compat(self._cfg(min_free_slots_delay=4))
class TestCudaGraphPrefillMaxContextResolution(CustomTestCase):
@staticmethod
def _make_args(max_context_size, model_context_len=4096, page_size=64):