Support DCP for Kimi Linear model (#32612)
Co-authored-by: Julien Lin <jullin@nvidia.com> Co-authored-by: kpham-sgl <khoa.pham@radixark.ai>
This commit is contained in:
Baizhou Zhang
co-authored by
Julien Lin
kpham-sgl
Julien Lin
kpham-sgl
parent
c4fc241fd3
commit
ef6c07008b
@@ -12,12 +12,19 @@ Usage:
|
||||
python test_dcp_layout_unit.py
|
||||
"""
|
||||
|
||||
import math
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.dcp.layout import get_dcp_lens
|
||||
from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
|
||||
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
|
||||
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")
|
||||
|
||||
@@ -36,7 +43,7 @@ def _legacy_inplace_formula(length: int, n: int, rank: int) -> int:
|
||||
return (length - rank - 1) // n + 1
|
||||
|
||||
|
||||
class TestGetDcpLens(unittest.TestCase):
|
||||
class TestGetDcpLens(CustomTestCase):
|
||||
def test_start_none_matches_owner_count(self):
|
||||
for n in DCP_SIZES:
|
||||
for rank in range(n):
|
||||
@@ -86,6 +93,156 @@ class TestGetDcpLens(unittest.TestCase):
|
||||
lens = torch.tensor(LENS, dtype=torch.int32)
|
||||
self.assertTrue(torch.equal(get_dcp_lens(lens, 1, 0), lens))
|
||||
|
||||
def test_paged_allocator_exposes_dcp_virtual_capacity(self):
|
||||
real_kv_size = 1024
|
||||
dcp_size = 4
|
||||
physical_page_size = 64
|
||||
allocator = PagedTokenToKVPoolAllocator(
|
||||
size=real_kv_size * dcp_size,
|
||||
page_size=physical_page_size * dcp_size,
|
||||
dtype=torch.bfloat16,
|
||||
device="cpu",
|
||||
kvcache=object(),
|
||||
need_sort=False,
|
||||
)
|
||||
|
||||
allocations = [allocator.alloc(physical_page_size * dcp_size) for _ in range(4)]
|
||||
self.assertTrue(all(indices is not None for indices in allocations))
|
||||
virtual_indices = torch.cat(allocations)
|
||||
|
||||
self.assertEqual(allocator.size, real_kv_size * dcp_size)
|
||||
self.assertEqual(allocator.page_size, physical_page_size * dcp_size)
|
||||
self.assertEqual(allocator.num_pages, real_kv_size // physical_page_size)
|
||||
self.assertEqual(
|
||||
len(torch.unique(virtual_indices // dcp_size)),
|
||||
len(virtual_indices) // dcp_size,
|
||||
)
|
||||
self.assertLess(
|
||||
int((virtual_indices // dcp_size).max()),
|
||||
real_kv_size + physical_page_size,
|
||||
)
|
||||
|
||||
def test_configurator_scales_only_the_virtual_dcp_allocator(self):
|
||||
physical_kv_size = 1024
|
||||
physical_page_size = 64
|
||||
physical_kv_cache = SimpleNamespace(
|
||||
size=physical_kv_size,
|
||||
page_size=physical_page_size,
|
||||
)
|
||||
sizes = SimpleNamespace(
|
||||
max_total_num_tokens=physical_kv_size,
|
||||
full_max_total_num_tokens=None,
|
||||
swa_max_total_num_tokens=None,
|
||||
)
|
||||
allocators = {}
|
||||
|
||||
for dcp_size in (1, 4):
|
||||
configurator = SimpleNamespace(
|
||||
server_args=SimpleNamespace(
|
||||
disaggregation_mode="null",
|
||||
enable_hisparse=False,
|
||||
page_size=physical_page_size,
|
||||
dcp_size=dcp_size,
|
||||
),
|
||||
hybrid_gdn_config=None,
|
||||
is_hybrid_swa=False,
|
||||
kv_cache_dtype=torch.bfloat16,
|
||||
device="cpu",
|
||||
is_draft_worker=False,
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.mem_cache.kv_cache_configurator.current_platform.is_out_of_tree",
|
||||
return_value=False,
|
||||
):
|
||||
allocators[dcp_size] = (
|
||||
KVCacheConfigurator._build_token_to_kv_pool_allocator(
|
||||
configurator,
|
||||
sizes=sizes,
|
||||
token_to_kv_pool=physical_kv_cache,
|
||||
is_dsv4_model=False,
|
||||
req_to_token_pool=object(),
|
||||
token_to_kv_pool_allocator=None,
|
||||
)
|
||||
)
|
||||
|
||||
dcp1_allocator = allocators[1]
|
||||
dcp4_allocator = allocators[4]
|
||||
self.assertIs(dcp1_allocator.get_kvcache(), physical_kv_cache)
|
||||
self.assertIs(dcp4_allocator.get_kvcache(), physical_kv_cache)
|
||||
self.assertEqual(dcp1_allocator.size, 1024)
|
||||
self.assertEqual(dcp1_allocator.page_size, 64)
|
||||
self.assertEqual(dcp1_allocator.num_pages, 16)
|
||||
self.assertEqual(dcp4_allocator.size, 4096)
|
||||
self.assertEqual(dcp4_allocator.page_size, 256)
|
||||
self.assertEqual(dcp4_allocator.num_pages, 16)
|
||||
|
||||
def test_live_cell_and_page_ownership_formulas(self):
|
||||
dcp_size = 4
|
||||
physical_page_size = 64
|
||||
ragged_lengths = (0, 1, 2, 3, 4, 63, 64, 65, 255, 256, 257, 515)
|
||||
|
||||
per_rank_counts = []
|
||||
for rank in range(dcp_size):
|
||||
expected_counts = [
|
||||
length // dcp_size + int(rank < length % dcp_size)
|
||||
for length in ragged_lengths
|
||||
]
|
||||
actual_counts = [
|
||||
_owner_count(length, dcp_size, rank, 0) for length in ragged_lengths
|
||||
]
|
||||
self.assertEqual(actual_counts, expected_counts)
|
||||
per_rank_counts.append(sum(actual_counts))
|
||||
|
||||
allocated_pages = [
|
||||
math.ceil(length / (physical_page_size * dcp_size))
|
||||
for length in ragged_lengths
|
||||
]
|
||||
active_pages = [
|
||||
math.ceil(count / physical_page_size) for count in actual_counts
|
||||
]
|
||||
self.assertTrue(
|
||||
all(
|
||||
active <= allocated
|
||||
for active, allocated in zip(active_pages, allocated_pages)
|
||||
)
|
||||
)
|
||||
self.assertTrue(
|
||||
all(
|
||||
allocated - active <= 1
|
||||
for active, allocated in zip(active_pages, allocated_pages)
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(sum(per_rank_counts), sum(ragged_lengths))
|
||||
|
||||
aligned_lengths = (256, 512, 768, 1024)
|
||||
full_replica_cells = sum(aligned_lengths)
|
||||
full_replica_pages = sum(
|
||||
length // physical_page_size for length in aligned_lengths
|
||||
)
|
||||
for rank in range(dcp_size):
|
||||
local_cells = sum(
|
||||
_owner_count(length, dcp_size, rank, 0) for length in aligned_lengths
|
||||
)
|
||||
local_pages = sum(
|
||||
math.ceil(_owner_count(length, dcp_size, rank, 0) / physical_page_size)
|
||||
for length in aligned_lengths
|
||||
)
|
||||
self.assertEqual(local_cells * dcp_size, full_replica_cells)
|
||||
self.assertEqual(local_pages * dcp_size, full_replica_pages)
|
||||
|
||||
def test_hybrid_pool_reports_the_backing_attention_shape(self):
|
||||
pool = object.__new__(HybridLinearKVPool)
|
||||
pool.start_layer = 0
|
||||
pool.layer_transfer_counter = None
|
||||
pool.full_attention_layer_id_mapping = {3: 0, 7: 1}
|
||||
pool.full_kv_pool = MagicMock()
|
||||
expected = (torch.Size([1024, 1, 576]), torch.Size([1024, 1, 576]))
|
||||
pool.full_kv_pool.get_kv_buffer_shape.return_value = expected
|
||||
|
||||
self.assertEqual(pool.get_kv_buffer_shape(), expected)
|
||||
pool.full_kv_pool.get_kv_buffer_shape.assert_called_once_with()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Four-Blackwell acceptance coverage for Kimi Linear TokenSpeed MLA DCP.
|
||||
|
||||
The captured-shape and eager-shape requests deliberately straddle
|
||||
``--cuda-graph-max-bs-decode=64``. This guards both the regular CUDA graph
|
||||
decode path and the full-capacity eager DCP LSE scratch-buffer path.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-b200")
|
||||
|
||||
KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
|
||||
|
||||
|
||||
def _has_four_blackwell_gpus() -> bool:
|
||||
if not torch.cuda.is_available() or torch.cuda.device_count() < 4:
|
||||
return False
|
||||
return all(
|
||||
torch.cuda.get_device_capability(device_index) >= (10, 0)
|
||||
for device_index in range(4)
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_has_four_blackwell_gpus(),
|
||||
"TokenSpeed MLA DCP acceptance requires four Blackwell GPUs",
|
||||
)
|
||||
class TestKimiLinearDCP4(GSM8KMixin, CustomTestCase):
|
||||
model = KIMI_LINEAR_MODEL
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
gsm8k_score_threshold = 0.90
|
||||
gsm8k_num_examples = 200
|
||||
# Keep accuracy evaluation within the captured decode batch sizes so its
|
||||
# score is batch-invariant. The separate smoke test still exercises the
|
||||
# eager path with batch size 65.
|
||||
gsm8k_num_threads = 4
|
||||
gsm8k_num_shots = 5
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.process = None
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
|
||||
other_args=[
|
||||
"--tp-size",
|
||||
"4",
|
||||
"--dcp-size",
|
||||
"4",
|
||||
"--attention-backend",
|
||||
"tokenspeed_mla",
|
||||
"--kv-cache-dtype",
|
||||
"fp8_e4m3",
|
||||
"--dcp-comm-backend",
|
||||
"a2a",
|
||||
"--dcp-replicate-q-proj",
|
||||
"--trust-remote-code",
|
||||
"--random-seed",
|
||||
"0",
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--cuda-graph-max-bs-decode",
|
||||
"64",
|
||||
"--cuda-graph-backend-prefill",
|
||||
"disabled",
|
||||
"--mem-fraction-static",
|
||||
"0.80",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if hasattr(cls, "process") and cls.process:
|
||||
kill_process_tree(cls.process.pid, wait_timeout=60)
|
||||
|
||||
def _assert_batch_completes(self, batch_size: int):
|
||||
prompts = [
|
||||
f"Reply with one short word for request {index}: the sky is"
|
||||
for index in range(batch_size)
|
||||
]
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompts,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 8,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
},
|
||||
timeout=180,
|
||||
)
|
||||
response.raise_for_status()
|
||||
outputs = response.json()
|
||||
self.assertIsInstance(outputs, list)
|
||||
self.assertEqual(len(outputs), batch_size)
|
||||
for output in outputs:
|
||||
self.assertTrue(output["text"].strip())
|
||||
self.assertGreater(output["meta_info"]["completion_tokens"], 0)
|
||||
|
||||
def test_decode_cuda_graph_and_eager_batch(self):
|
||||
# Batch two replays a captured shape; batch 65 is above the configured
|
||||
# regular CUDA graph maximum and therefore exercises eager decode.
|
||||
self._assert_batch_completes(2)
|
||||
self._assert_batch_completes(2)
|
||||
self._assert_batch_completes(65)
|
||||
|
||||
def test_physical_capacity_sanity(self):
|
||||
response = requests.get(self.base_url + "/server_info", timeout=30)
|
||||
response.raise_for_status()
|
||||
self.assertGreater(response.json()["max_total_num_tokens"], 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user