chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
47cb9b9d116d1c5c112e894110ec94de80f2db6b
sglang/test/registered/unit/server_args/test_server_args.py
laoyao0822 3a43727216 Bound CP prefill batching by estimated temp memory 
CP shared-KV bs>1 batching was only bounded by request count, extend tokens, and cached tokens. That left temporary GPU buffers such as MLA/index materialization, remap metadata, logits windows, and transfer descriptors implicit, and raw extend-token limits could exceed the active chunked-prefill budget.\n\nThis adds an explicit max-buffer-size admission gate with a CPU-only stream-aware estimator, wires it through PrefillAdder/Scheduler, performs a startup CUDA smoke allocation when configured, and reports the estimate in the scheduler admission benchmark. When chunked prefill is active, the effective CP extend-token limit is capped by the current chunk budget so the CP path does not advertise unreachable batch capacity or lift max-prefill-tokens too far.\n\nConstraint: Admission estimation must stay CPU-only on the scheduler hot path; CUDA allocation is limited to startup smoke checking.\nConstraint: Single oversized requests must still be allowed to run alone to avoid scheduler deadlock.\nRejected: Rely only on --max-prefill-tokens | it does not reliably bound the first oversized request and does not model cache-hit/load-back pressure.\nRejected: Let CP extend limit exceed chunked-prefill size | it creates an unreachable effective capacity and misleading budget lift.\nConfidence: medium\nScope-risk: moderate\nDirective: If bs>1 L1 prefetch is enabled later, update CPSharedKVPrefillBufferEstimatorContext.bs_gt1_l1_prefetch_enabled and include the live prefetch dense buffers in overlap windows.\nTested: local py_compile for touched files\nTested: local PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py (4 passed)\nTested: remote cjy-glm5-new targeted pytest for new server_args, PrefillAdder, estimator, and benchmark cases (10 passed)\nTested: remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_prefill_scheduler_admission_bench.py (29 passed before chunk cap, then test_prefill_adder.py 21 passed after chunk cap)\nNot-tested: full server_args suite because existing TestPrepareServerArgs tries to reach HuggingFace and fails under container DNS/network\nNot-tested: GLM5 ETE smoke with --cp-shared-kv-prefill-max-buffer-size
2026-06-11 01:33:28 +08:00

752 lines
27 KiB
Python
Raw Blame History

import json
import tempfile
import unittest
from unittest.mock import MagicMock, patch
import pytest
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
def test_cp_shared_kv_prefill_bs_gt1_parser_limits():
import argparse
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
raw_args = parser.parse_args(
[
"--model-path",
"dummy",
"--enable-cp-shared-kv-prefill-bs-gt1",
"--cp-shared-kv-prefill-max-batch-requests",
"4",
"--cp-shared-kv-prefill-max-total-extend-tokens",
"8192",
]
)
args = ServerArgs.from_cli_args(raw_args)
assert args.enable_cp_shared_kv_prefill_bs_gt1 is True
assert args.cp_shared_kv_prefill_max_batch_requests == 4
assert args.cp_shared_kv_prefill_max_total_extend_tokens == 8192
def test_cp_shared_kv_prefill_max_buffer_size_defaults_to_gb():
import argparse
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
raw_args = parser.parse_args(
[
"--model-path",
"dummy",
"--enable-cp-shared-kv-prefill-bs-gt1",
"--cp-shared-kv-prefill-max-buffer-size",
"8",
]
)
args = ServerArgs.from_cli_args(raw_args)
assert args.cp_shared_kv_prefill_max_buffer_size == 8_000_000_000
def test_cp_shared_kv_prefill_max_buffer_size_accepts_iec_suffix():
import argparse
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
raw_args = parser.parse_args(
[
"--model-path",
"dummy",
"--enable-cp-shared-kv-prefill-bs-gt1",
"--cp-shared-kv-prefill-max-buffer-size",
"8Gi",
]
)
args = ServerArgs.from_cli_args(raw_args)
assert args.cp_shared_kv_prefill_max_buffer_size == 8 * 2**30
def test_cp_shared_kv_prefill_max_buffer_size_rejects_non_positive_value():
import argparse
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
raw_args = parser.parse_args(
[
"--model-path",
"dummy",
"--enable-cp-shared-kv-prefill-bs-gt1",
"--cp-shared-kv-prefill-max-buffer-size",
"0",
]
)
with pytest.raises(ValueError, match="cp_shared_kv_prefill_max_buffer_size"):
ServerArgs.from_cli_args(raw_args)
def test_hicache_mem_layout_parser_accepts_layer_page_first():
import argparse
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
raw_args = parser.parse_args(
[
"--model-path",
"dummy",
"--hicache-mem-layout",
"layer_page_first",
]
)
args = ServerArgs.from_cli_args(raw_args)
assert args.hicache_mem_layout == "layer_page_first"
def test_nsa_index_topk_cli_overrides_model_config_json():
import argparse
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
raw_args = parser.parse_args(
[
"--model-path",
"dummy",
"--json-model-override-args",
'{"rope_scaling": {"factor": 2.0}, "index_topk_freq": 2}',
"--nsa-index-topk-freq",
"4",
"--nsa-index-skip-topk-offset",
"1",
]
)
args = ServerArgs.from_cli_args(raw_args)
override_args = json.loads(args.json_model_override_args)
assert args.nsa_index_topk_freq == 4
assert args.nsa_index_skip_topk_offset == 1
assert override_args["rope_scaling"] == {"factor": 2.0}
assert override_args["index_topk_freq"] == 4
assert override_args["index_skip_topk_offset"] == 1
def test_nsa_index_skip_topk_offset_zero_means_unset():
import argparse
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
raw_args = parser.parse_args(
[
"--model-path",
"dummy",
"--nsa-index-topk-freq",
"4",
"--nsa-index-skip-topk-offset",
"0",
]
)
args = ServerArgs.from_cli_args(raw_args)
override_args = json.loads(args.json_model_override_args)
assert args.nsa_index_skip_topk_offset == 0
assert override_args == {"index_topk_freq": 4}
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(CustomTestCase):
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 _make_cp_hicache_args(self, **overrides) -> ServerArgs:
args = self._make_args(
enable_hierarchical_cache=True,
enable_nsa_prefill_context_parallel=True,
enable_nsa_prefill_cp_shared_kv=True,
nsa_prefill_cp_mode="in-seq-split",
disaggregation_mode="prefill",
page_size=64,
enable_hisparse=False,
)
for key, value in overrides.items():
setattr(args, key, value)
return args
def _normalize_and_validate_cp_hicache_args(self, **overrides) -> ServerArgs:
args = self._make_cp_hicache_args(**overrides)
args._handle_hicache()
args._handle_cp_hicache_layout_validation()
return args
def test_cp_shared_kv_rejects_hicache_storage_backend(self):
with (
self.assertRaisesRegex(
AssertionError,
"enable_nsa_prefill_cp_shared_kv.*hicache_storage_backend",
),
patch.object(
ServerArgs,
"get_model_config",
side_effect=AssertionError("get_model_config should not be called"),
),
):
ServerArgs(
model_path="dummy",
enable_hierarchical_cache=True,
enable_nsa_prefill_context_parallel=True,
enable_nsa_prefill_cp_shared_kv=True,
nsa_prefill_cp_mode="in-seq-split",
disaggregation_mode=None,
page_size=64,
enable_hisparse=False,
hicache_storage_backend="mooncake",
)
def test_cp_shared_kv_prefill_rejects_nixl_transfer_backend(self):
with (
self.assertRaisesRegex(
AssertionError,
"enable_nsa_prefill_cp_shared_kv.*disaggregation_transfer_backend.*mooncake",
),
patch.object(
ServerArgs,
"get_model_config",
side_effect=AssertionError("get_model_config should not be called"),
),
):
ServerArgs(
model_path="dummy",
enable_nsa_prefill_context_parallel=True,
enable_nsa_prefill_cp_shared_kv=True,
nsa_prefill_cp_mode="in-seq-split",
disaggregation_mode="prefill",
disaggregation_transfer_backend="nixl",
page_size=64,
enable_hisparse=False,
)
def test_cp_shared_kv_prefill_batch_limits_must_be_positive(self):
with self.assertRaisesRegex(
ValueError, "cp_shared_kv_prefill_max_batch_requests.*positive"
):
ServerArgs(
model_path="dummy",
cp_shared_kv_prefill_max_batch_requests=0,
)
with self.assertRaisesRegex(
ValueError, "cp_shared_kv_prefill_max_total_extend_tokens.*positive"
):
ServerArgs(
model_path="dummy",
cp_shared_kv_prefill_max_total_extend_tokens=0,
)
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")
def test_cp_hicache_accepts_supported_backend_layout_pairs(self):
cases = [
("kernel", "layer_first"),
("kernel", "page_first"),
("direct", "layer_first"),
("direct", "page_first_direct"),
("direct", "layer_page_first"),
]
for io_backend, mem_layout in cases:
with self.subTest(io_backend=io_backend, mem_layout=mem_layout):
args = self._normalize_and_validate_cp_hicache_args(
hicache_io_backend=io_backend,
hicache_mem_layout=mem_layout,
)
self.assertEqual(args.hicache_io_backend, io_backend)
self.assertEqual(args.hicache_mem_layout, mem_layout)
def test_cp_hicache_normalizes_supported_alias_pairs(self):
args = self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel",
hicache_mem_layout="page_first_direct",
)
self.assertEqual(args.hicache_io_backend, "direct")
self.assertEqual(args.hicache_mem_layout, "page_first_direct")
args = self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="direct",
hicache_mem_layout="page_first",
)
self.assertEqual(args.hicache_io_backend, "direct")
self.assertEqual(args.hicache_mem_layout, "page_first_direct")
def test_cp_hicache_rejects_page_head_layout(self):
with self.assertRaisesRegex(ValueError, "CP shared KV HiCache.*page_head"):
self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel",
hicache_mem_layout="page_head",
)
def test_cp_hicache_rejects_cuda_kv_split_layout(self):
with self.assertRaisesRegex(
ValueError, "CP shared KV HiCache.*page_first_kv_split"
):
self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel",
hicache_mem_layout="page_first_kv_split",
)
def test_cp_hicache_rejects_kernel_layer_page_first_layout(self):
with self.assertRaisesRegex(
ValueError, "CP shared KV HiCache.*kernel.*layer_page_first"
):
self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel",
hicache_mem_layout="layer_page_first",
)
def test_hicache_storage_rejects_layer_page_first_layout(self):
args = self._make_args(
enable_hierarchical_cache=True,
hicache_storage_backend="mooncake",
hicache_io_backend="direct",
hicache_mem_layout="layer_page_first",
)
with self.assertRaisesRegex(ValueError, "layer_page_first.*storage"):
args._handle_hicache()
def test_cp_hicache_rejects_kernel_ascend_backend(self):
with self.assertRaisesRegex(ValueError, "CP shared KV HiCache.*kernel_ascend"):
self._normalize_and_validate_cp_hicache_args(
hicache_io_backend="kernel_ascend",
hicache_mem_layout="page_first_kv_split",
)
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink