feat: support Kimi Linear PD disaggregation with DCP (#32837)

Co-authored-by: Yangmin Li <yangminl@nvidia.com>
This commit is contained in:
Khoa Pham
2026-07-31 02:14:09 -07:00
committed by GitHub
co-authored by Yangmin Li
parent 33c27d8e7f
commit 2573190b93
13 changed files with 1084 additions and 48 deletions
@@ -8,6 +8,7 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import sglang.srt.server_args as server_args_module
from sglang.srt.arg_groups import pd_disaggregation_hook
from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
from sglang.srt.entrypoints.sidecar import (
SGLANG_GRPC_ENDPOINT_ENV,
@@ -219,6 +220,56 @@ class TestLoadBalanceMethod(unittest.TestCase):
server_args = self._load_balance_args(disaggregation_mode="decode")
self.assertEqual(server_args.load_balance_method, "round_robin")
def test_pd_prefill_dcp_warns_about_performance(self):
server_args = ServerArgs(
model_path="dummy",
disaggregation_mode="prefill",
dcp_size=4,
)
with self.assertLogs(pd_disaggregation_hook.logger, level="WARNING") as logs:
server_args._handle_pd_disaggregation()
self.assertIn("without improving prefill performance", "\n".join(logs.output))
def test_pd_decode_dcp_forces_chunk_cache(self):
server_args = self._load_balance_args(
disaggregation_mode="decode",
disaggregation_transfer_backend="mooncake",
dcp_size=4,
)
self.assertTrue(server_args.disable_radix_cache)
def test_pd_decode_dcp_rejects_unsupported_transfer_backend(self):
server_args = ServerArgs(
model_path="dummy",
disaggregation_mode="decode",
disaggregation_transfer_backend="fake",
dcp_size=4,
)
with self.assertRaisesRegex(ValueError, "mooncake or nixl"):
server_args._handle_pd_disaggregation()
def test_pd_decode_dcp_rejects_radix_cache(self):
server_args = ServerArgs(
model_path="dummy",
disaggregation_mode="decode",
disaggregation_transfer_backend="nixl",
disaggregation_decode_enable_radix_cache=True,
dcp_size=4,
)
with self.assertRaisesRegex(ValueError, "currently requires chunk cache"):
server_args._handle_pd_disaggregation()
def test_pd_decode_dcp_rejects_hierarchical_cache(self):
server_args = ServerArgs(
model_path="dummy",
disaggregation_mode="decode",
disaggregation_transfer_backend="nixl",
enable_hierarchical_cache=True,
dcp_size=4,
)
with self.assertRaisesRegex(ValueError, "--enable-hierarchical-cache"):
server_args._handle_pd_disaggregation()
def test_pd_decode_radix_cache_rejects_hisparse(self):
server_args = ServerArgs(
model_path="dummy",