f8fca72635
Prefill CP previously replicated NSA/MLA persistent KV on every CP rank, so CP8 consumed eight copies of KV memory while exposing only one rank of logical cache capacity. This change splits logical KV locs from per-rank physical storage, shards MLA latent KV and NSA index K/scale by deterministic page ownership, and keeps existing NSA attention kernels working through a full-view runtime materialization layer. Mooncake PD transfer now sends each prefill CP rank's owned physical pages with explicit logical page positions so non-CP decode can reconstruct full-layout KV. The implementation is guarded by an explicit server flag and startup checks, and the design documentation records the implemented scope, debug environment, and Phase 3 boundary. Constraint: Phase 2 must preserve existing NSA attention/index kernels via runtime full-view materialization Constraint: Decode side remains non-CP and receives full KV through Mooncake Rejected: Shard-aware NSA attention in this change | belongs to Phase 3 because it requires distributed topk/softmax/output contracts Rejected: Request-contiguous CP ownership | unstable under chunked prefill and tied to attention split mode Confidence: medium Scope-risk: broad Directive: Do not enable round-robin CP shared KV without wiring runtime materialization/PD transfer contracts for that split mode Directive: Keep SGLANG_DEBUG_CP_SHARED_KV disabled for perf measurements; it intentionally enables CUDA-syncing diagnostics Tested: Remote py_compile for shared-KV touched Python files in g0034 container Tested: Remote pytest selected cp_shared/shared_kv/nsa suite: 37 passed, 34 deselected Not-tested: Full GLM5 multi-node throughput/regression run after final doc update Not-tested: Phase 3 shard-aware runtime, round-robin CP mode, and non-Mooncake PD backends
450 lines
17 KiB
Python
450 lines
17 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from sglang.srt.server_args import PortArgs, ServerArgs, prepare_server_args
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
|
|
CustomTestCase,
|
|
)
|
|
|
|
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
|
|
|
# Mock get_device() so all tests run on CPU-only CI runners
|
|
_mock_device = patch("sglang.srt.server_args.get_device", return_value="cuda")
|
|
_mock_device.start()
|
|
|
|
|
|
class TestPrepareServerArgs(CustomTestCase):
|
|
def test_prepare_server_args(self):
|
|
server_args = prepare_server_args(
|
|
[
|
|
"--model-path",
|
|
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
|
|
"--json-model-override-args",
|
|
'{"rope_scaling": {"factor": 2.0, "rope_type": "linear"}}',
|
|
]
|
|
)
|
|
self.assertEqual(server_args.model_path, DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN)
|
|
self.assertEqual(
|
|
json.loads(server_args.json_model_override_args),
|
|
{"rope_scaling": {"factor": 2.0, "rope_type": "linear"}},
|
|
)
|
|
|
|
|
|
def test_enable_nsa_prefill_cp_shared_kv_parser_flag():
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
ServerArgs.add_cli_args(parser)
|
|
raw_args = parser.parse_args(
|
|
[
|
|
"--model-path",
|
|
"dummy",
|
|
"--enable-nsa-prefill-context-parallel",
|
|
"--enable-nsa-prefill-cp-shared-kv",
|
|
"--nsa-prefill-cp-mode",
|
|
"in-seq-split",
|
|
]
|
|
)
|
|
args = ServerArgs.from_cli_args(raw_args)
|
|
assert args.enable_nsa_prefill_cp_shared_kv is True
|
|
|
|
|
|
class TestLoadBalanceMethod(unittest.TestCase):
|
|
def test_non_pd_defaults_to_round_robin(self):
|
|
server_args = ServerArgs(model_path="dummy", disaggregation_mode="null")
|
|
self.assertEqual(server_args.load_balance_method, "round_robin")
|
|
|
|
def test_pd_prefill_defaults_to_follow_bootstrap_room(self):
|
|
server_args = ServerArgs(model_path="dummy", disaggregation_mode="prefill")
|
|
self.assertEqual(server_args.load_balance_method, "follow_bootstrap_room")
|
|
|
|
def test_pd_decode_defaults_to_round_robin(self):
|
|
server_args = ServerArgs(model_path="dummy", disaggregation_mode="decode")
|
|
self.assertEqual(server_args.load_balance_method, "round_robin")
|
|
|
|
|
|
class TestPortArgs(unittest.TestCase):
|
|
@patch("sglang.srt.server_args.get_free_port")
|
|
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
|
|
def test_init_new_with_nccl_port_none(self, mock_temp_file, mock_get_free_port):
|
|
"""Test that get_free_port() is called when nccl_port is None"""
|
|
mock_temp_file.return_value.name = "temp_file"
|
|
mock_get_free_port.return_value = 45678 # Mock ephemeral port
|
|
|
|
# Use MagicMock here to verify get_free_port is called
|
|
server_args = MagicMock()
|
|
server_args.nccl_port = None
|
|
server_args.enable_dp_attention = False
|
|
server_args.tokenizer_worker_num = 1
|
|
|
|
port_args = PortArgs.init_new(server_args)
|
|
|
|
# Verify get_free_port was called
|
|
mock_get_free_port.assert_called_once()
|
|
|
|
# Verify the returned port is used
|
|
self.assertEqual(port_args.nccl_port, 45678)
|
|
|
|
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
|
|
def test_init_new_standard_case(self, mock_temp_file):
|
|
mock_temp_file.return_value.name = "temp_file"
|
|
|
|
server_args = ServerArgs(model_path="dummy")
|
|
server_args.port = 30000
|
|
server_args.nccl_port = None
|
|
server_args.enable_dp_attention = False
|
|
|
|
port_args = PortArgs.init_new(server_args)
|
|
|
|
self.assertTrue(port_args.tokenizer_ipc_name.startswith("ipc://"))
|
|
self.assertTrue(port_args.scheduler_input_ipc_name.startswith("ipc://"))
|
|
self.assertTrue(port_args.detokenizer_ipc_name.startswith("ipc://"))
|
|
self.assertIsInstance(port_args.nccl_port, int)
|
|
|
|
def test_init_new_with_single_node_dp_attention(self):
|
|
|
|
server_args = ServerArgs(model_path="dummy")
|
|
server_args.port = 30000
|
|
server_args.nccl_port = None
|
|
server_args.enable_dp_attention = True
|
|
server_args.nnodes = 1
|
|
server_args.dist_init_addr = None
|
|
|
|
port_args = PortArgs.init_new(server_args)
|
|
|
|
self.assertTrue(port_args.tokenizer_ipc_name.startswith("tcp://127.0.0.1:"))
|
|
self.assertTrue(
|
|
port_args.scheduler_input_ipc_name.startswith("tcp://127.0.0.1:")
|
|
)
|
|
self.assertTrue(port_args.detokenizer_ipc_name.startswith("tcp://127.0.0.1:"))
|
|
self.assertIsInstance(port_args.nccl_port, int)
|
|
|
|
def test_init_new_with_dp_rank(self):
|
|
server_args = ServerArgs(model_path="dummy")
|
|
server_args.port = 30000
|
|
server_args.nccl_port = None
|
|
server_args.enable_dp_attention = True
|
|
server_args.nnodes = 1
|
|
server_args.dist_init_addr = "192.168.1.1:25000"
|
|
|
|
worker_ports = [25006, 25007, 25008, 25009]
|
|
port_args = PortArgs.init_new(server_args, dp_rank=2, worker_ports=worker_ports)
|
|
|
|
self.assertTrue(port_args.scheduler_input_ipc_name.endswith(":25008"))
|
|
|
|
self.assertTrue(port_args.tokenizer_ipc_name.startswith("tcp://192.168.1.1:"))
|
|
self.assertTrue(port_args.detokenizer_ipc_name.startswith("tcp://192.168.1.1:"))
|
|
self.assertIsInstance(port_args.nccl_port, int)
|
|
|
|
def test_init_new_with_ipv4_address(self):
|
|
server_args = ServerArgs(model_path="dummy")
|
|
server_args.port = 30000
|
|
server_args.nccl_port = None
|
|
|
|
server_args.enable_dp_attention = True
|
|
server_args.nnodes = 2
|
|
server_args.dist_init_addr = "192.168.1.1:25000"
|
|
|
|
port_args = PortArgs.init_new(server_args)
|
|
|
|
self.assertTrue(port_args.tokenizer_ipc_name.startswith("tcp://192.168.1.1:"))
|
|
self.assertTrue(
|
|
port_args.scheduler_input_ipc_name.startswith("tcp://192.168.1.1:")
|
|
)
|
|
self.assertTrue(port_args.detokenizer_ipc_name.startswith("tcp://192.168.1.1:"))
|
|
self.assertIsInstance(port_args.nccl_port, int)
|
|
|
|
def test_init_new_with_malformed_ipv4_address(self):
|
|
server_args = ServerArgs(model_path="dummy")
|
|
server_args.port = 30000
|
|
server_args.nccl_port = None
|
|
|
|
server_args.enable_dp_attention = True
|
|
server_args.nnodes = 2
|
|
server_args.dist_init_addr = "192.168.1.1"
|
|
|
|
with self.assertRaises(ValueError) as context:
|
|
PortArgs.init_new(server_args)
|
|
|
|
self.assertIn("Missing port", str(context.exception))
|
|
|
|
def test_init_new_with_malformed_ipv4_address_invalid_port(self):
|
|
server_args = ServerArgs(model_path="dummy")
|
|
server_args.port = 30000
|
|
server_args.nccl_port = None
|
|
|
|
server_args.enable_dp_attention = True
|
|
server_args.nnodes = 2
|
|
server_args.dist_init_addr = "192.168.1.1:abc"
|
|
|
|
with self.assertRaises(ValueError):
|
|
PortArgs.init_new(server_args)
|
|
|
|
|
|
class TestSSLArgs(unittest.TestCase):
|
|
def test_default_ssl_fields_are_none(self):
|
|
server_args = ServerArgs(model_path="dummy")
|
|
self.assertIsNone(server_args.ssl_keyfile)
|
|
self.assertIsNone(server_args.ssl_certfile)
|
|
self.assertIsNone(server_args.ssl_ca_certs)
|
|
self.assertIsNone(server_args.ssl_keyfile_password)
|
|
|
|
def test_ssl_keyfile_without_certfile_raises(self):
|
|
with self.assertRaises(ValueError) as context:
|
|
ServerArgs(model_path="dummy", ssl_keyfile="key.pem")
|
|
self.assertIn("--ssl-certfile", str(context.exception))
|
|
|
|
def test_ssl_certfile_without_keyfile_raises(self):
|
|
with self.assertRaises(ValueError) as context:
|
|
ServerArgs(model_path="dummy", ssl_certfile="cert.pem")
|
|
self.assertIn("--ssl-keyfile", str(context.exception))
|
|
|
|
@patch("os.path.isfile", return_value=True)
|
|
def test_ssl_both_keyfile_and_certfile_accepted(self, _mock_isfile):
|
|
server_args = ServerArgs(
|
|
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
|
)
|
|
self.assertEqual(server_args.ssl_keyfile, "key.pem")
|
|
self.assertEqual(server_args.ssl_certfile, "cert.pem")
|
|
|
|
def test_url_returns_http_without_ssl(self):
|
|
server_args = ServerArgs(model_path="dummy")
|
|
self.assertTrue(server_args.url().startswith("http://"))
|
|
|
|
def test_url_rewrites_all_interfaces_to_loopback(self):
|
|
server_args = ServerArgs(model_path="dummy", host="0.0.0.0")
|
|
self.assertEqual(server_args.url(), "http://127.0.0.1:30000")
|
|
|
|
def test_url_rewrites_empty_host_to_loopback(self):
|
|
server_args = ServerArgs(model_path="dummy", host="")
|
|
self.assertEqual(server_args.url(), "http://127.0.0.1:30000")
|
|
|
|
@patch("os.path.isfile", return_value=True)
|
|
def test_url_returns_https_with_ssl(self, _mock_isfile):
|
|
server_args = ServerArgs(
|
|
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
|
)
|
|
self.assertTrue(server_args.url().startswith("https://"))
|
|
|
|
@patch("os.path.isfile", return_value=True)
|
|
def test_ssl_cli_args_parsed(self, _mock_isfile):
|
|
server_args = prepare_server_args(
|
|
[
|
|
"--model-path",
|
|
"dummy",
|
|
"--ssl-keyfile",
|
|
"key.pem",
|
|
"--ssl-certfile",
|
|
"cert.pem",
|
|
"--ssl-ca-certs",
|
|
"ca.pem",
|
|
"--ssl-keyfile-password",
|
|
"secret",
|
|
]
|
|
)
|
|
self.assertEqual(server_args.ssl_keyfile, "key.pem")
|
|
self.assertEqual(server_args.ssl_certfile, "cert.pem")
|
|
self.assertEqual(server_args.ssl_ca_certs, "ca.pem")
|
|
self.assertEqual(server_args.ssl_keyfile_password, "secret")
|
|
|
|
def test_ssl_verify_without_ssl(self):
|
|
server_args = ServerArgs(model_path="dummy")
|
|
self.assertIs(server_args.ssl_verify(), True)
|
|
|
|
@patch("os.path.isfile", return_value=True)
|
|
def test_ssl_verify_with_ssl_no_ca(self, _mock_isfile):
|
|
server_args = ServerArgs(
|
|
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
|
)
|
|
self.assertIs(server_args.ssl_verify(), False)
|
|
|
|
@patch("os.path.isfile", return_value=True)
|
|
def test_ssl_verify_with_ssl_and_ca(self, _mock_isfile):
|
|
server_args = ServerArgs(
|
|
model_path="dummy",
|
|
ssl_keyfile="key.pem",
|
|
ssl_certfile="cert.pem",
|
|
ssl_ca_certs="ca.pem",
|
|
)
|
|
self.assertEqual(server_args.ssl_verify(), "ca.pem")
|
|
|
|
def test_ssl_ca_certs_without_certfile_raises(self):
|
|
with self.assertRaises(ValueError) as context:
|
|
ServerArgs(model_path="dummy", ssl_ca_certs="ca.pem")
|
|
self.assertIn("--ssl-ca-certs", str(context.exception))
|
|
|
|
def test_ssl_keyfile_password_without_certfile_raises(self):
|
|
with self.assertRaises(ValueError) as context:
|
|
ServerArgs(model_path="dummy", ssl_keyfile_password="secret")
|
|
self.assertIn("--ssl-keyfile-password", str(context.exception))
|
|
|
|
def test_ssl_keyfile_not_found_raises(self):
|
|
with self.assertRaises(ValueError) as context:
|
|
ServerArgs(
|
|
model_path="dummy",
|
|
ssl_keyfile="/nonexistent/key.pem",
|
|
ssl_certfile="/nonexistent/cert.pem",
|
|
)
|
|
self.assertIn("not found", str(context.exception))
|
|
|
|
def test_ssl_certfile_not_found_raises(self):
|
|
with tempfile.NamedTemporaryFile(suffix=".pem") as keyfile:
|
|
with self.assertRaises(ValueError) as context:
|
|
ServerArgs(
|
|
model_path="dummy",
|
|
ssl_keyfile=keyfile.name,
|
|
ssl_certfile="/nonexistent/cert.pem",
|
|
)
|
|
self.assertIn("SSL certificate file not found", str(context.exception))
|
|
|
|
def test_ssl_ca_certs_not_found_raises(self):
|
|
with tempfile.NamedTemporaryFile(suffix=".pem") as keyfile:
|
|
with tempfile.NamedTemporaryFile(suffix=".pem") as certfile:
|
|
with self.assertRaises(ValueError) as context:
|
|
ServerArgs(
|
|
model_path="dummy",
|
|
ssl_keyfile=keyfile.name,
|
|
ssl_certfile=certfile.name,
|
|
ssl_ca_certs="/nonexistent/ca.pem",
|
|
)
|
|
self.assertIn(
|
|
"SSL CA certificates file not found", str(context.exception)
|
|
)
|
|
|
|
def test_enable_ssl_refresh_default_false(self):
|
|
server_args = ServerArgs(model_path="dummy")
|
|
self.assertFalse(server_args.enable_ssl_refresh)
|
|
|
|
def test_enable_ssl_refresh_without_ssl_raises(self):
|
|
with self.assertRaises(ValueError) as context:
|
|
ServerArgs(model_path="dummy", enable_ssl_refresh=True)
|
|
self.assertIn("--enable-ssl-refresh", str(context.exception))
|
|
self.assertIn("--ssl-certfile", str(context.exception))
|
|
|
|
@patch("os.path.isfile", return_value=True)
|
|
def test_enable_ssl_refresh_with_ssl_accepted(self, _mock_isfile):
|
|
server_args = ServerArgs(
|
|
model_path="dummy",
|
|
ssl_keyfile="key.pem",
|
|
ssl_certfile="cert.pem",
|
|
enable_ssl_refresh=True,
|
|
)
|
|
self.assertTrue(server_args.enable_ssl_refresh)
|
|
|
|
@patch("os.path.isfile", return_value=True)
|
|
def test_enable_ssl_refresh_cli_flag(self, _mock_isfile):
|
|
server_args = prepare_server_args(
|
|
[
|
|
"--model-path",
|
|
"dummy",
|
|
"--ssl-keyfile",
|
|
"key.pem",
|
|
"--ssl-certfile",
|
|
"cert.pem",
|
|
"--enable-ssl-refresh",
|
|
]
|
|
)
|
|
self.assertTrue(server_args.enable_ssl_refresh)
|
|
|
|
|
|
class TestHiCacheArgs(unittest.TestCase):
|
|
def _make_args(self, **overrides) -> ServerArgs:
|
|
args = ServerArgs(model_path="dummy")
|
|
for key, value in overrides.items():
|
|
setattr(args, key, value)
|
|
return args
|
|
|
|
def _assert_hicache_fields(
|
|
self,
|
|
args: ServerArgs,
|
|
*,
|
|
expected_io_backend: str,
|
|
expected_mem_layout: str,
|
|
expected_decode_backend: str | None = None,
|
|
):
|
|
self.assertEqual(args.hicache_io_backend, expected_io_backend)
|
|
self.assertEqual(args.hicache_mem_layout, expected_mem_layout)
|
|
if expected_decode_backend is not None:
|
|
self.assertEqual(args.decode_attention_backend, expected_decode_backend)
|
|
|
|
def test_hicache_io_backend_and_mem_layout_compatibility(self):
|
|
cases = [
|
|
{
|
|
"name": "kernel_with_page_first_direct",
|
|
"overrides": {
|
|
"enable_hierarchical_cache": True,
|
|
"hicache_io_backend": "kernel",
|
|
"hicache_mem_layout": "page_first_direct",
|
|
},
|
|
"expected_io_backend": "direct",
|
|
"expected_mem_layout": "page_first_direct",
|
|
},
|
|
{
|
|
"name": "direct_with_page_first",
|
|
"overrides": {
|
|
"enable_hierarchical_cache": True,
|
|
"hicache_io_backend": "direct",
|
|
"hicache_mem_layout": "page_first",
|
|
},
|
|
"expected_io_backend": "direct",
|
|
"expected_mem_layout": "page_first_direct",
|
|
},
|
|
{
|
|
"name": "mooncake_with_layer_first",
|
|
"overrides": {
|
|
"enable_hierarchical_cache": True,
|
|
"hicache_storage_backend": "mooncake",
|
|
"hicache_io_backend": "direct",
|
|
"hicache_mem_layout": "layer_first",
|
|
},
|
|
"expected_io_backend": "direct",
|
|
"expected_mem_layout": "page_first_direct",
|
|
},
|
|
{
|
|
"name": "fa3_kernel_with_explicit_decode_backend",
|
|
"overrides": {
|
|
"enable_hierarchical_cache": True,
|
|
"hicache_io_backend": "kernel",
|
|
"hicache_mem_layout": "page_first",
|
|
"attention_backend": "triton",
|
|
"decode_attention_backend": "fa3",
|
|
},
|
|
"expected_io_backend": "direct",
|
|
"expected_mem_layout": "page_first_direct",
|
|
},
|
|
]
|
|
|
|
for case in cases:
|
|
with self.subTest(case=case["name"]):
|
|
args = self._make_args(**case["overrides"])
|
|
args._handle_hicache()
|
|
self._assert_hicache_fields(
|
|
args,
|
|
expected_io_backend=case["expected_io_backend"],
|
|
expected_mem_layout=case["expected_mem_layout"],
|
|
)
|
|
|
|
@patch.object(ServerArgs, "use_mla_backend", return_value=False)
|
|
@patch("sglang.srt.server_args.is_flashinfer_available", return_value=False)
|
|
def test_decode_attention_backend_with_implicit_fa3(
|
|
self, _mock_flashinfer, _mock_use_mla_backend
|
|
):
|
|
args = self._make_args(
|
|
enable_hierarchical_cache=True,
|
|
hicache_io_backend="kernel",
|
|
attention_backend="fa3",
|
|
decode_attention_backend=None,
|
|
)
|
|
|
|
args._handle_hicache()
|
|
|
|
self.assertEqual(args.decode_attention_backend, "triton")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|