minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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
python/sglang/srt/configs/model_config.py
+14 -5
View File
@@ -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
python/sglang/srt/mem_cache/dsa_cache_layer_split.py
+4 -2
View File
@@ -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):
python/sglang/srt/mem_cache/index_key_cache.py
+16 -3
View File
@@ -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
python/sglang/srt/mem_cache/kv_cache_configurator.py
+19
View File
@@ -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,
python/sglang/srt/mem_cache/memory_pool.py
+8
View File
@@ -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 (
python/sglang/srt/model_executor/pool_configurator.py
+94 -24
View File
@@ -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:
python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
+8
View File
@@ -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 (
python/sglang/srt/models/deepseek_v2.py
+25 -27
View File
@@ -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,
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):