feat: support Kimi Linear PD disaggregation with DCP (#32837)
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
This commit is contained in:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user