diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 73e75293f..b2c1a6da9 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -201,6 +201,14 @@ def dsa_layer_skips_topk(config: PretrainedConfig, layer_id: int) -> bool: """Return whether a DSA layer reuses the previous layer's top-k indices.""" assert is_deepseek_dsa(config) + # LongCat computes fresh top-k indices every cli_factor layers. + cli_factor = getattr(config, "cli_factor", 1) + if cli_factor is None: + cli_factor = 1 + assert cli_factor > 0, f"cli_factor must be positive, got {cli_factor}" + if cli_factor > 1: + return layer_id % cli_factor != 0 + pattern = getattr(config, "index_topk_pattern", None) if pattern is not None: return layer_id < len(pattern) and pattern[layer_id] == "S" @@ -231,11 +239,12 @@ REQUANTIZATION_METHODS = ["quark_mxfp4"] def get_num_indexer_layers(config) -> int: """Layer count for the global indexer-topk capturer's host buffer. - DSA models (V3.2) instantiate an Indexer on every transformer layer. - With index_topk_freq > 1 some layers reuse prev layer's topk; those still - get a slot (mirrored at the MLA call site). DSv4 has C4 indexers only on - layers whose compress_ratio == 4. Other architectures: set - num_indexer_layers on hf_text_config; 0 disables the capturer. + DSA models (V3.2) expose one capturer slot per transformer layer. With + index_topk_freq > 1 some layers reuse prev layer's topk; those still get a + slot mirrored at the MLA call site even if no Indexer module is built. + DSv4 has C4 indexers only on layers whose compress_ratio == 4. Other + architectures: set num_indexer_layers on hf_text_config; 0 disables the + capturer. """ if is_deepseek_dsa(config): return config.num_hidden_layers diff --git a/python/sglang/srt/mem_cache/dsa_cache_layer_split.py b/python/sglang/srt/mem_cache/dsa_cache_layer_split.py index 5915d7639..ba1cd1e52 100644 --- a/python/sglang/srt/mem_cache/dsa_cache_layer_split.py +++ b/python/sglang/srt/mem_cache/dsa_cache_layer_split.py @@ -72,7 +72,9 @@ class LayerSplitIndexKeyCache(IndexKeyCache): def _layer_num_pages(self, layer_idx: int, num_pages: int) -> int: layer_id = self.pool.start_layer + layer_idx - return num_pages if self.pool._is_layer_owned(layer_id) else 0 + if not self.pool._is_layer_owned(layer_id): + return 0 + return super()._layer_num_pages(layer_idx, num_pages) def clear(self) -> None: super().clear() @@ -150,7 +152,7 @@ class LayerSplitIndexKeyCache(IndexKeyCache): ] data_ptrs = [self.buffer[i].data_ptr() for i in owned_layer_ids] data_lens = [self.buffer[i].nbytes for i in owned_layer_ids] - item_lens = [self.buffer[i][0].nbytes for i in owned_layer_ids] + item_lens = [self._item_len(i) for i in owned_layer_ids] return data_ptrs, data_lens, item_lens def cpu_copy(self, indices): diff --git a/python/sglang/srt/mem_cache/index_key_cache.py b/python/sglang/srt/mem_cache/index_key_cache.py index 18eecec60..166d5d222 100644 --- a/python/sglang/srt/mem_cache/index_key_cache.py +++ b/python/sglang/srt/mem_cache/index_key_cache.py @@ -38,7 +38,9 @@ class IndexKeyCache: ) def _layer_num_pages(self, layer_idx: int, num_pages: int) -> int: - return num_pages + # Layers that reuse the previous layer's top-k never write index-K, so + # they get a 0-row placeholder that keeps ``buffer`` layer-aligned. + return 0 if self.pool.skip_topk_layers[layer_idx] else num_pages def clear(self) -> None: del self.buffer @@ -49,6 +51,8 @@ class IndexKeyCache: tgt_loc_flat = tgt_loc.view(-1).long() src_loc_flat = src_loc.view(-1).long() for index_k in self.buffer: + if index_k.shape[0] == 0: + continue index_k[tgt_loc_flat] = index_k[src_loc_flat] def get_local_buffer(self, layer_id: int) -> torch.Tensor: @@ -118,6 +122,8 @@ class IndexKeyCache: page_chunk_size = max(1, chunk_size // self.pool.page_size) for layer_id in range(self.pool.layer_num): index_k_cpu.append([]) + if self.buffer[layer_id].shape[0] == 0: + continue for i in range(0, len(page_indices), page_chunk_size): chunk_page_indices = page_indices[i : i + page_chunk_size] idx_cpu = self.buffer[layer_id][chunk_page_indices].to( @@ -133,17 +139,24 @@ class IndexKeyCache: chunk_size = self.pool.cpu_offloading_chunk_size page_chunk_size = max(1, chunk_size // self.pool.page_size) for layer_id in range(self.pool.layer_num): + if self.buffer[layer_id].shape[0] == 0: + continue for i in range(0, len(page_indices), page_chunk_size): chunk_page_indices = page_indices[i : i + page_chunk_size] idx_cpu = index_k_cpu[layer_id][i // page_chunk_size] assert idx_cpu.shape[0] == len(chunk_page_indices) - idx_chunk = idx_cpu.to(self.buffer[0].device, non_blocking=True) + idx_chunk = idx_cpu.to(self.buffer[layer_id].device, non_blocking=True) self.buffer[layer_id][chunk_page_indices] = idx_chunk torch.cuda.synchronize() + def _item_len(self, layer_idx: int) -> int: + # 0-row layers (skip-topk, or non-owned under CP layer split) have no item. + buf = self.buffer[layer_idx] + return 0 if buf.shape[0] == 0 else buf[0].nbytes + def state_buf_infos(self): layer_num = self.pool.layer_num data_ptrs = [self.buffer[i].data_ptr() for i in range(layer_num)] data_lens = [self.buffer[i].nbytes for i in range(layer_num)] - item_lens = [self.buffer[i][0].nbytes for i in range(layer_num)] + item_lens = [self._item_len(i) for i in range(layer_num)] return data_ptrs, data_lens, item_lens diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 18dedb76d..4d2ce9a6d 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -15,6 +15,7 @@ from sglang.srt.configs.hybrid_arch import ( ) from sglang.srt.configs.model_config import ( ModelConfig, + dsa_layer_skips_topk, get_dsa_index_head_dim, get_minimax_sparse_attention_config, get_minimax_sparse_disable_value_layer_ids, @@ -89,6 +90,17 @@ from sglang.srt.utils.common import ( logger = logging.getLogger(__name__) + +def _should_elide_dsa_index_k(*, is_draft_worker: bool) -> bool: + memory_config = get_memory() + return ( + not memory_config.enable_hisparse + and not is_draft_worker + and not memory_config.enable_hierarchical_cache + and get_disagg().disaggregation_mode == "null" + ) + + _is_hip = is_hip() @@ -1324,6 +1336,13 @@ class KVCacheConfigurator: pool_kwargs["layer_shard_size"] = dsa_cp_layer_shard_size else: PoolCls = DSATokenToKVPool + if _should_elide_dsa_index_k(is_draft_worker=self.is_draft_worker): + pool_kwargs["skip_topk_layers"] = [ + dsa_layer_skips_topk(self.model_config.hf_config, layer_id) + for layer_id in range( + self.layer_info.start_layer, self.layer_info.end_layer + ) + ] token_to_kv_pool = PoolCls( max_total_num_tokens, page_size=self.pool_page_size, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 3696ad38a..420fa25b8 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -4365,6 +4365,7 @@ class DSATokenToKVPool(MLATokenToKVPool): start_layer: Optional[int] = None, end_layer: Optional[int] = None, index_buf_size: Optional[int] = None, + skip_topk_layers: Optional[List[bool]] = None, ): override_dim = ( kv_cache_dim if kv_cache_dim != kv_lora_rank + qk_rope_head_dim else None @@ -4393,6 +4394,13 @@ class DSATokenToKVPool(MLATokenToKVPool): # num head == 1 and head dim == 128 for index_k in DSA assert index_head_dim == 128 + self.skip_topk_layers = ( + list(skip_topk_layers) + if skip_topk_layers is not None + else [False] * layer_num + ) + assert len(self.skip_topk_layers) == layer_num + if _is_hip: if aiter_can_use_preshuffle_paged_mqa(): assert ( diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 44d7f5df4..8c8a318ca 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -21,6 +21,7 @@ import torch from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.configs.model_config import ( + dsa_layer_skips_topk, get_dsa_index_head_dim, get_minimax_sparse_attention_config, get_minimax_sparse_disable_value_layer_ids, @@ -184,10 +185,33 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): and int(eagle_draft_num_layers) > 0 and int(num_layers) > 0 ): - self._cell_size = int( - self._cell_size - * (1 + int(eagle_draft_num_layers) / int(num_layers)) - ) + draft_num_layers = int(eagle_draft_num_layers) + if is_deepseek_dsa(kvc.model_config.hf_config): + target_indexer_size = self._compute_dsa_indexer_cell_size( + kvc=kvc, + num_layers=num_layers, + ) + target_kv_size = self._cell_size - target_indexer_size + from sglang.srt.layers.cp.utils import ( + get_glm_dsa_layer_split_effective_num_layers, + ) + + target_kv_num_layers = get_glm_dsa_layer_split_effective_num_layers( + kvc, num_layers + ) + draft_kv_size = int( + target_kv_size * draft_num_layers / target_kv_num_layers + ) + draft_indexer_size = self._compute_dsa_indexer_cell_size( + kvc=kvc, + num_layers=draft_num_layers, + allocate_all_layers=True, + ) + self._cell_size += draft_kv_size + draft_indexer_size + else: + self._cell_size = int( + self._cell_size * (1 + draft_num_layers / int(num_layers)) + ) # DFLASH/DSPARK: scale cell_size to account for draft model KV cache if kvc.spec_algorithm.is_dflash_family() and not kvc.is_draft_worker: @@ -254,26 +278,9 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): # Add indexer KV cache overhead for DSA models (DeepSeek V3.2) if is_deepseek_dsa(model_config.hf_config): - index_head_dim = get_dsa_index_head_dim(model_config.hf_config) - indexer_size_per_token = ( - index_head_dim - + index_head_dim // DSATokenToKVPool.quant_block_size * 4 - ) - element_size = torch._utils._element_size( - DSATokenToKVPool.index_k_with_scale_buffer_dtype - ) - indexer_ratio = 1 - if kvc.server_args.enable_hisparse: - from sglang.srt.mem_cache.sparsity import parse_hisparse_config - - indexer_ratio = parse_hisparse_config( - kvc.server_args - ).host_to_device_ratio - cell_size += int( - indexer_size_per_token - * effective_num_layers - * element_size - * indexer_ratio + cell_size += self._compute_dsa_indexer_cell_size( + kvc=kvc, + num_layers=num_layers, ) elif is_minimax_sparse(model_config.hf_config): # Mirrors MiniMaxSparseKVPool: main pool (K+V all layers) + indexer pool @@ -343,6 +350,69 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): return cell_size + def _compute_dsa_indexer_cell_size( + self, + *, + kvc: KVCacheConfigurator, + num_layers: int, + allocate_all_layers: bool = False, + ) -> int: + index_head_dim = get_dsa_index_head_dim(kvc.model_config.hf_config) + indexer_size_per_token = ( + index_head_dim + index_head_dim // DSATokenToKVPool.quant_block_size * 4 + ) + element_size = torch._utils._element_size( + DSATokenToKVPool.index_k_with_scale_buffer_dtype + ) + memory_config = get_memory() + indexer_ratio = 1 + if memory_config.enable_hisparse: + from sglang.srt.mem_cache.sparsity import parse_hisparse_config + + indexer_ratio = parse_hisparse_config(kvc.server_args).host_to_device_ratio + + from sglang.srt.mem_cache.kv_cache_configurator import ( + _should_elide_dsa_index_k, + ) + + if allocate_all_layers or not _should_elide_dsa_index_k( + is_draft_worker=kvc.is_draft_worker + ): + num_indexer_layers = num_layers + else: + active_indexer_layers = [ + layer_id + for layer_id in range( + kvc.layer_info.start_layer, kvc.layer_info.end_layer + ) + if not dsa_layer_skips_topk(kvc.model_config.hf_config, layer_id) + ] + from sglang.srt.layers.cp.utils import ( + get_glm_dsa_cp_layer_shard_info, + get_layer_shard_range, + ) + + _, shard_size = get_glm_dsa_cp_layer_shard_info(kvc) + if shard_size > 1: + active_set = set(active_indexer_layers) + max_owned = 0 + for rank in range(shard_size): + start, end = get_layer_shard_range(rank, shard_size, num_layers) + max_owned = max( + max_owned, + sum( + kvc.layer_info.start_layer + i in active_set + for i in range(start, end) + ), + ) + num_indexer_layers = max_owned + 1 + else: + num_indexer_layers = len(active_indexer_layers) + + return int( + indexer_size_per_token * num_indexer_layers * element_size * indexer_ratio + ) + def calculate_pool_sizes( self, available_bytes: int, page_size: int ) -> MemoryPoolConfig: diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index ef0077d6b..b4557f795 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -203,6 +203,9 @@ class DeepseekV2WeightLoaderMixin: with concurrent.futures.ThreadPoolExecutor() as executor: futures = [] params_dict = dict(self.named_parameters()) + indexer_present_prefixes = { + n.rsplit(".indexer.", 1)[0] for n in params_dict if ".indexer." in n + } weight_names = [] for name, loaded_weight in weights: @@ -257,6 +260,11 @@ class DeepseekV2WeightLoaderMixin: if "rotary_emb.inv_freq" in name: continue + if ".indexer." in name and ( + name.rsplit(".indexer.", 1)[0] not in indexer_present_prefixes + ): + continue + # CUDA fuses wk + weights_proj into one bf16 wk_weights_proj; the # helper returns True once it has consumed the shard. if ( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 646a943e8..4e4b48d50 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1814,27 +1814,8 @@ class DeepseekV2AttentionMLA( self.skip_topk = None self.next_skip_topk = None + self.indexer = None if self.use_dsa: - is_neox_style = not getattr(config, "indexer_rope_interleave", False) - self.indexer = Indexer( - hidden_size=hidden_size, - index_n_heads=get_dsa_index_n_heads(config), - index_head_dim=get_dsa_index_head_dim(config), - rope_head_dim=qk_rope_head_dim, - index_topk=get_dsa_index_topk(config), - q_lora_rank=q_lora_rank, - max_position_embeddings=max_position_embeddings, - rope_theta=rope_theta, - scale_fmt="ue8m0", - block_size=128, - rope_scaling=rope_scaling, - is_neox_style=is_neox_style, - prefix=add_prefix("indexer", prefix), - quant_config=quant_config, - layer_id=layer_id, - alt_stream=alt_stream, - config=config, - ) # Refer: https://arxiv.org/abs/2603.12201 for more details. # skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices. # next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices. @@ -1842,13 +1823,30 @@ class DeepseekV2AttentionMLA( self.skip_topk = True self.next_skip_topk = True else: - index_cli_factor = getattr(config, "cli_factor", 1) - if index_cli_factor > 1: - self.skip_topk = layer_id % index_cli_factor != 0 - self.next_skip_topk = (layer_id + 1) % index_cli_factor != 0 - else: - self.skip_topk = dsa_layer_skips_topk(config, layer_id) - self.next_skip_topk = dsa_layer_skips_topk(config, layer_id + 1) + self.skip_topk = dsa_layer_skips_topk(config, layer_id) + self.next_skip_topk = dsa_layer_skips_topk(config, layer_id + 1) + + if not self.skip_topk or is_nextn: + is_neox_style = not getattr(config, "indexer_rope_interleave", False) + self.indexer = Indexer( + hidden_size=hidden_size, + index_n_heads=get_dsa_index_n_heads(config), + index_head_dim=get_dsa_index_head_dim(config), + rope_head_dim=qk_rope_head_dim, + index_topk=get_dsa_index_topk(config), + q_lora_rank=q_lora_rank, + max_position_embeddings=max_position_embeddings, + rope_theta=rope_theta, + scale_fmt="ue8m0", + block_size=128, + rope_scaling=rope_scaling, + is_neox_style=is_neox_style, + prefix=add_prefix("indexer", prefix), + quant_config=quant_config, + layer_id=layer_id, + alt_stream=alt_stream, + config=config, + ) self.kv_b_proj = ColumnParallelLinear( self.kv_lora_rank, diff --git a/test/registered/unit/model_executor/test_hisparse_pool_configurator.py b/test/registered/unit/model_executor/test_hisparse_pool_configurator.py index e11cb7ac2..409ea2cdc 100644 --- a/test/registered/unit/model_executor/test_hisparse_pool_configurator.py +++ b/test/registered/unit/model_executor/test_hisparse_pool_configurator.py @@ -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 ( diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 3d254adbb..7757da848 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -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):