minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[DSA] Skip indexer KV cache for skip-topk layers (#30531)

Browse Source 
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
Co-authored-by: mmangkad <mohammad.angkad@radixark.ai>
This commit is contained in:
Mohammad Miadh Angkad
2026-08-17 02:02:23 -07:00
committed by GitHub
co-authored by Brayden Zhong mmangkad
parent 7c423cfd41
commit 8cc112d486
10 changed files with 325 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/model_executor/test_hisparse_pool_configurator.py
+18 -8
View File
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock
import torch
from sglang.srt.model_executor.pool_configurator import DefaultPoolConfigurator
from sglang.srt.runtime_context import get_parallel
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -20,31 +20,41 @@ class TestHiSparsePoolConfigurator(CustomTestCase):
enable_hisparse: bool,
host_to_device_ratio: int = 1,
) -> int:
num_layers = 2
hf_config = SimpleNamespace(
architectures=["GlmMoeDsaForCausalLM"],
index_topk=2048,
index_head_dim=128,
)
hf_config.get_text_config = lambda: hf_config
override = get_context().override_server_args(
enable_hisparse=enable_hisparse,
hisparse_config=f'{{"host_to_device_ratio": {host_to_device_ratio}}}',
enable_hierarchical_cache=False,
disaggregation_mode="null",
dsa_prefill_backend="flashmla_sparse",
dsa_decode_backend="flashmla_sparse",
)
server_args = override.install()
self.addCleanup(override.restore)
kvc = MagicMock(
use_mla_backend=True,
kv_cache_dtype=kv_cache_dtype,
is_draft_worker=False,
model_config=SimpleNamespace(
kv_lora_rank=512,
qk_rope_head_dim=64,
hf_config=hf_config,
),
server_args=SimpleNamespace(
enable_hisparse=enable_hisparse,
hisparse_config=(f'{{"host_to_device_ratio": {host_to_device_ratio}}}'),
dsa_prefill_backend="flashmla_sparse",
dsa_decode_backend="flashmla_sparse",
),
layer_info=SimpleNamespace(start_layer=0, end_layer=num_layers),
server_args=server_args,
)
with get_parallel().override(attn_tp_size=1):
configurator = object.__new__(DefaultPoolConfigurator)
return configurator._compute_cell_size(kvc, num_layers=2)
return configurator._compute_cell_size(kvc, num_layers=num_layers)
def test_mla_layout_without_hisparse(self):
for kv_cache_dtype, expected_cell_size in (
test/registered/unit/model_executor/test_pool_configurator.py
+119 -16
View File
@@ -11,8 +11,9 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.runtime_context import get_parallel, get_server_args
from sglang.srt.runtime_context import get_memory, get_parallel, get_server_args
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")
@@ -137,6 +138,7 @@ def _make_model_runner(
max_running_requests=max_running_requests,
disaggregation_decode_extra_slots=disaggregation_decode_extra_slots,
enable_hisparse=False,
enable_hierarchical_cache=False,
enable_dsa_cache_layer_split=False,
kv_cache_dtype="auto",
)
@@ -162,6 +164,13 @@ def _make_model_runner(
return mr
def _configure_dsa_model(model_runner):
hf_config = model_runner.model_config.hf_config
hf_config.architectures = ["GlmMoeDsaForCausalLM"]
hf_config.index_topk = 2048
hf_config.index_head_dim = 128
KV_SIZE = 2 # bf16
@@ -191,7 +200,7 @@ def _actual_memory_used(mr, config):
return config.max_total_num_tokens * full_pt * (nf + ns)
class TestDefaultConfigurator(unittest.TestCase):
class TestDefaultConfigurator(CustomTestCase):
"""Default (MHA): available_bytes -> tokens, memory invariant holds."""
def _run(self, available_bytes, page_size=1, **kwargs):
@@ -236,20 +245,13 @@ class TestDefaultConfigurator(unittest.TestCase):
self.assertIsNone(config.full_max_total_num_tokens)
self.assertIsNone(config.swa_max_total_num_tokens)
@patch(
"sglang.srt.model_executor.pool_configurator.get_dsa_index_head_dim",
return_value=128,
)
@patch(
"sglang.srt.model_executor.pool_configurator.is_deepseek_dsa",
return_value=True,
)
@patch(
"sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim",
side_effect=(576, 656),
)
def test_dsa_mla_cell_size_uses_backend_kv_layout(
self, mock_calculate_mla_kv_cache_dim, _mock_is_dsa, _mock_index_head_dim
self,
mock_calculate_mla_kv_cache_dim,
):
num_layers = 2
raw = _make_model_runner(
@@ -262,6 +264,8 @@ class TestDefaultConfigurator(unittest.TestCase):
num_layers=num_layers,
use_mla_backend=True,
)
_configure_dsa_model(raw)
_configure_dsa_model(packed)
with mock_cpu_env(kv_size=1):
from sglang.srt.model_executor.pool_configurator import (
@@ -277,7 +281,7 @@ class TestDefaultConfigurator(unittest.TestCase):
self.assertEqual(mock_calculate_mla_kv_cache_dim.call_count, 2)
class TestHybridSWAConfigurator(unittest.TestCase):
class TestHybridSWAConfigurator(CustomTestCase):
"""Hybrid SWA: full/swa split, ratio, memory invariant."""
def _make_swa_runner(self, full_layers=16, swa_layers=16, ratio=0.5, page_size=1):
@@ -535,7 +539,7 @@ class TestHybridSWAConfigurator(unittest.TestCase):
self.assertLessEqual(_actual_memory_used(mr, config), available)
class TestAllSWAConfigurator(unittest.TestCase):
class TestAllSWAConfigurator(CustomTestCase):
"""All-SWA (full_layers=0): special case."""
def _run(self, available_bytes, ratio=0.5, page_size=1, **kwargs):
@@ -584,7 +588,7 @@ class TestAllSWAConfigurator(unittest.TestCase):
self.assertEqual(config.swa_max_total_num_tokens, 500)
class TestEagleConfigurator(unittest.TestCase):
class TestEagleConfigurator(CustomTestCase):
"""EAGLE: draft KV cache must be accounted for so total allocation fits in budget."""
def test_eagle_does_not_exceed_budget(self):
@@ -612,8 +616,107 @@ class TestEagleConfigurator(unittest.TestCase):
used = config.max_total_num_tokens * full_pt * total_layers
self.assertLessEqual(used, available)
@patch(
"sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim",
return_value=576,
)
def test_dsa_draft_full_indexer_cost_does_not_exceed_budget(
self,
_mock_calculate_mla_kv_cache_dim,
):
"""A sharing-enabled target must not discount the draft's full index-K."""
available = 10_000_000
num_layers = 78
draft_num_layers = 1
active_indexer_layers = 21
indexer_bytes_per_token = 132
class TestFactory(unittest.TestCase):
mr = _make_model_runner(self, num_layers=num_layers, use_mla_backend=True)
_configure_dsa_model(mr)
mr.model_config.hf_config.index_topk_freq = 4
mr.model_config.hf_config.index_skip_topk_offset = 3
mr.spec_algorithm.is_eagle.return_value = True
mr.spec_algorithm.is_none.return_value = False
mr.spec_aux_config.eagle_draft_num_layers = draft_num_layers
with mock_cpu_env(kv_size=1):
from sglang.srt.model_executor.pool_configurator import (
create_memory_pool_configurator,
)
cfg = create_memory_pool_configurator(mr)
config = cfg.calculate_pool_sizes(available, page_size=1)
actual_bytes_per_token = (
576 * num_layers
+ indexer_bytes_per_token * active_indexer_layers
+ (576 + indexer_bytes_per_token) * draft_num_layers
)
self.assertEqual(cfg._cell_size, actual_bytes_per_token)
self.assertLessEqual(
config.max_total_num_tokens * actual_bytes_per_token,
available,
)
class TestDSAIndexerAllocationPolicy(CustomTestCase):
@patch(
"sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim",
return_value=576,
)
def test_resolved_hicache_override_prices_every_indexer_layer(
self,
_mock_calculate_mla_kv_cache_dim,
):
"""Post-publish HiCache overrides must keep sizing and allocation aligned."""
num_layers = 6
mr = _make_model_runner(self, num_layers=num_layers, use_mla_backend=True)
_configure_dsa_model(mr)
mr.model_config.hf_config.index_topk_freq = 4
mr.model_config.hf_config.index_skip_topk_offset = 3
with get_memory().override(enable_hierarchical_cache=True), mock_cpu_env(
kv_size=1
):
from sglang.srt.model_executor.pool_configurator import (
DefaultPoolConfigurator,
)
cfg = DefaultPoolConfigurator(mr)
self.assertEqual(cfg._cell_size, (576 + 132) * num_layers)
@patch(
"sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim",
return_value=576,
)
def test_pd_prices_every_indexer_layer(
self,
_mock_calculate_mla_kv_cache_dim,
):
"""PD must retain dense index-K metadata until transports support sparsity."""
num_layers = 6
mr = _make_model_runner(
self,
num_layers=num_layers,
use_mla_backend=True,
disaggregation_mode="prefill",
)
_configure_dsa_model(mr)
mr.model_config.hf_config.index_topk_freq = 4
mr.model_config.hf_config.index_skip_topk_offset = 3
with mock_cpu_env(kv_size=1):
from sglang.srt.model_executor.pool_configurator import (
DefaultPoolConfigurator,
)
cfg = DefaultPoolConfigurator(mr)
self.assertEqual(cfg._cell_size, (576 + 132) * num_layers)
class TestFactory(CustomTestCase):
def test_default_for_non_swa(self):
mr = _make_model_runner(self, is_hybrid_swa=False)
with mock_cpu_env():
@@ -671,7 +774,7 @@ class TestFactory(unittest.TestCase):
self.assertNotIsInstance(_cfg(None), SWAChunkCapPoolConfigurator)
class TestDflashDraftKvBudget(unittest.TestCase):
class TestDflashDraftKvBudget(CustomTestCase):
"""DFLASH draft KV pool as a flat bytes/token term on the target's budget."""
def test_bytes_per_token_from_draft_geometry(self):