from __future__ import annotations import logging import math from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Optional import msgspec import torch from sglang.srt.configs.hybrid_arch import hybrid_gdn_config, mambaish_config from sglang.srt.configs.model_config import ( ModelConfig, get_dsa_index_head_dim, get_minimax_sparse_attention_config, get_minimax_sparse_disable_value_layer_ids, get_minimax_sparse_layer_ids, is_deepseek_dsa, is_deepseek_v4, is_minimax_sparse, ) from sglang.srt.distributed.parallel_state import get_world_group from sglang.srt.environ import envs from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( get_kv_cache_quant_method, resolve_kv_cache_quant, ) from sglang.srt.mem_cache.allocation_sizing import get_req_to_token_extra_context_len from sglang.srt.mem_cache.allocator import ( BaseTokenToKVPoolAllocator, PagedTokenToKVPoolAllocator, TokenToKVPoolAllocator, ) from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, HiSparseTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.allocator.swa import ( PureSWATokenToKVPoolAllocator, SWATokenToKVPoolAllocator, ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool from sglang.srt.mem_cache.memory_pool import ( DSATokenToKVPool, HybridLinearKVPool, HybridReqToTokenPool, KVCache, MHATokenToKVPool, MHATokenToKVPoolFP4, MHATokenToKVPoolMXFP8, MiniMaxSparseKVPool, MLATokenToKVPool, MLATokenToKVPoolFP4, NoOpMHATokenToKVPool, PageMajorMHATokenToKVPool, ReqToTokenPool, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils.common import ( get_available_gpu_memory, get_device_memory_capacity, is_float4_e2m1fn_x2, is_hip, is_npu, ) logger = logging.getLogger(__name__) _is_hip = is_hip() def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]: dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower() if dtype_name in ("float32", "fp32"): return torch.float32, torch.float32 if dtype_name in ("bfloat16", "bf16"): return torch.bfloat16, torch.bfloat16 raise ValueError( "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE=" f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16." ) _is_npu = is_npu() def _should_enable_lazy_compaction() -> bool: """Lazy compaction default — ON unless `SGLANG_DISABLE_LAZY_COMPACTION=1` (escape hatch for A/B / rollback). Centralized here so both unified-memory-pool factory call sites stay in sync. """ return not envs.SGLANG_DISABLE_LAZY_COMPACTION.get() # the ratio of mamba cache pool size to max_running_requests MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3 MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2 MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP_LAZY = 1 MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1 if TYPE_CHECKING: from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.mem_cache.unified_memory_pool import ( UnifiedKVPool, UnifiedPoolBundle, ) from sglang.srt.model_executor.model_runner_components.layer_setup import ( ModelLayerInfo, ) from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import ( SpecAuxHiddenStateConfig, ) from sglang.srt.model_executor.pool_configurator import ( MemoryPoolConfig, ) class KVCacheConfigResult(msgspec.Struct, frozen=True, kw_only=True): max_total_num_tokens: int max_running_requests: int full_max_total_num_tokens: Optional[int] swa_max_total_num_tokens: Optional[int] req_to_token_pool: ReqToTokenPool token_to_kv_pool: KVCache token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator memory_pool_config: MemoryPoolConfig unified_memory_pool: Optional[UnifiedKVPool] = None class _InitializedPools(msgspec.Struct, frozen=True, kw_only=True): req_to_token_pool: ReqToTokenPool token_to_kv_pool: KVCache token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator unified_memory_pool: Optional[UnifiedKVPool] = None class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True): max_total_num_tokens: int max_running_requests: int full_max_total_num_tokens: Optional[int] swa_max_total_num_tokens: Optional[int] c4_max_total_num_tokens: int c128_max_total_num_tokens: int c4_state_pool_size: int c128_state_pool_size: int c4_state_dtype: Optional[torch.dtype] c128_state_dtype: Optional[torch.dtype] @dataclass(slots=True, kw_only=True) class KVCacheConfigurator: device: str gpu_id: int ps: ParallelState pp_group: Any model: Any model_config: ModelConfig server_args: ServerArgs kv_cache_dtype: torch.dtype model_dtype: torch.dtype page_size: int sliding_window_size: Optional[int] spec_algorithm: SpeculativeAlgorithm is_draft_worker: bool post_capture_kv_active: bool spec_aux_config: SpecAuxHiddenStateConfig is_hybrid_swa: bool is_hybrid_swa_compress: bool use_mla_backend: bool layer_info: ModelLayerInfo forward_stream: Any req_to_token_pool: Optional[ReqToTokenPool] token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] memory_pool_config: Optional[MemoryPoolConfig] draft_model_idx: Optional[int] = None mambaish_config: Optional[Any] = field(init=False) hybrid_gdn_config: Optional[Any] = field(init=False) is_inkling_mtp_draft: bool = field(init=False) draft_swa_full_capacity: bool = field(init=False) def __post_init__(self) -> None: self.mambaish_config = mambaish_config(self.model_config) self.hybrid_gdn_config = hybrid_gdn_config(self.model_config) # Each multi-layer EAGLE MTP head owns one transformer block at # layer_id=draft_model_idx; heads at a banded 's' depth route that layer # into the SWA ring sub-pool (draft_swa_full_capacity) so the SWA # store/read path activates for this depth, exactly like a trunk local # layer. self.is_inkling_mtp_draft = ( self.is_draft_worker and self.draft_model_idx is not None and self.model_config.hf_config.architectures[0] == "InklingForConditionalGenerationMTP" ) self.draft_swa_full_capacity = self.is_inkling_mtp_draft and ( self.draft_model_idx in set(self.model_config.hf_text_config.mtp_local_layer_ids) ) def _build_fp4_quant_method(self, *, num_layers: int): if not is_float4_e2m1fn_x2(self.kv_cache_dtype): return None quant_name = resolve_kv_cache_quant(self.server_args.kv_cache_dtype) if quant_name is None: return None quant_method = get_kv_cache_quant_method( quant_name, num_layers=num_layers, device=self.device, ) quant_method.load_scales_from_model(self.model) return quant_method def configure(self, *, pre_model_load_memory: int) -> KVCacheConfigResult: """Apply a resolved MemoryPoolConfig and initialize pools.""" if not self.spec_algorithm.is_none() and self.is_draft_worker: assert ( self.memory_pool_config is not None ), "Draft worker requires memory_pool_config" config = self.memory_pool_config else: config = self._resolve_memory_pool_config(pre_model_load_memory) sizes = self._derive_pool_sizes(config=config) pools = self._init_pools( sizes=sizes, req_to_token_pool=self.req_to_token_pool, token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, ) logger.info( f"Memory pool end. " f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB" ) return KVCacheConfigResult( max_total_num_tokens=sizes.max_total_num_tokens, max_running_requests=sizes.max_running_requests, full_max_total_num_tokens=sizes.full_max_total_num_tokens, swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, req_to_token_pool=pools.req_to_token_pool, token_to_kv_pool=pools.token_to_kv_pool, token_to_kv_pool_allocator=pools.token_to_kv_pool_allocator, memory_pool_config=config, unified_memory_pool=pools.unified_memory_pool, ) def _derive_pool_sizes(self, *, config: MemoryPoolConfig) -> _PoolSizes: max_total_num_tokens = config.max_total_num_tokens max_running_requests = config.max_running_requests full_max_total_num_tokens = None swa_max_total_num_tokens = None if self.is_hybrid_swa: full_max_total_num_tokens = config.full_max_total_num_tokens swa_max_total_num_tokens = config.swa_max_total_num_tokens # DSV4 compressed-attention pool sizes. Draft worker reuses target's # full/swa sizes but does NOT own c4/c128/state pools (those live on # the target rank only); zero them out regardless of what config holds. if self.is_draft_worker: c4_max_total_num_tokens = 0 c128_max_total_num_tokens = 0 c4_state_pool_size = 0 c128_state_pool_size = 0 else: c4_max_total_num_tokens = config.c4_max_total_num_tokens c128_max_total_num_tokens = config.c128_max_total_num_tokens c4_state_pool_size = config.c4_state_pool_size c128_state_pool_size = config.c128_state_pool_size # Draft worker does not own the compression-state pools, but keep the # dtype attributes initialized so _init_pools can share one code path. c4_state_dtype: Optional[torch.dtype] = None c128_state_dtype: Optional[torch.dtype] = None if is_deepseek_v4(self.model_config.hf_config): c4_state_dtype, c128_state_dtype = _get_dsv4_compress_state_dtypes() return _PoolSizes( max_total_num_tokens=max_total_num_tokens, max_running_requests=max_running_requests, full_max_total_num_tokens=full_max_total_num_tokens, swa_max_total_num_tokens=swa_max_total_num_tokens, c4_max_total_num_tokens=c4_max_total_num_tokens, c128_max_total_num_tokens=c128_max_total_num_tokens, c4_state_pool_size=c4_state_pool_size, c128_state_pool_size=c128_state_pool_size, c4_state_dtype=c4_state_dtype, c128_state_dtype=c128_state_dtype, ) def _init_pools( self, *, sizes: _PoolSizes, req_to_token_pool: Optional[ReqToTokenPool], token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator], ) -> _InitializedPools: """Initialize the memory pools.""" token_to_kv_pool = None # Unified-pool fast path: build req_to_token + token_to_kv pool + allocator # from one byte buffer, then return. Gated to the target worker # (req_to_token_pool is None); supports hybrid Mamba and hybrid SWA (not DSV4). if ( self.server_args.enable_unified_memory and self.server_args.disaggregation_mode == "null" and req_to_token_pool is None ): if self.mambaish_config is not None: bundle = self._init_unified_mamba_pools( max_num_reqs=sizes.max_running_requests, max_total_num_tokens=sizes.max_total_num_tokens, ) elif self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config): bundle = self._init_unified_swa_pools( max_num_reqs=sizes.max_running_requests, full_max_total_num_tokens=sizes.full_max_total_num_tokens, swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, ) else: # Fail loud, not silently fall through to the normal pools (which would # leave the flag a no-op). The feature replaces the HYBRID pools only. raise ValueError( "--enable-unified-memory only supports hybrid Mamba and " "hybrid sliding-window-attention models (DeepSeek-V4 excluded); " f"the current model ({self.model_config.hf_config.architectures}) " "is neither, so the unified memory pool cannot be built. Drop " "--enable-unified-memory for this model." ) return _InitializedPools( req_to_token_pool=bundle.req_to_token_pool, token_to_kv_pool=bundle.token_to_kv_pool, token_to_kv_pool_allocator=bundle.token_to_kv_pool_allocator, unified_memory_pool=bundle.unified_memory_pool, ) # Initialize req_to_token_pool if req_to_token_pool is None: req_to_token_pool = self._build_req_to_token_pool( max_num_reqs=sizes.max_running_requests ) else: # Draft worker shares req_to_token_pool with the target worker. assert self.is_draft_worker # Each multi-layer EAGLE MTP head owns one transformer block at # layer_id=draft_model_idx and needs its own sconv/mamba cache while # sharing the target's request-to-token mapping. if self.is_inkling_mtp_draft and isinstance( req_to_token_pool, HybridReqToTokenPool ): # speculative_num_draft_tokens=None: draft heads never run # TARGET_VERIFY, so their pools skip the per-step intermediate # (SpeculativeState) buffers only the target pool consumes. req_to_token_pool = req_to_token_pool.clone_with_new_mamba( mamba_size=self.server_args.max_mamba_cache_size, mamba_spec_state_size=sizes.max_running_requests, cache_params=self.mambaish_config.mamba2_cache_params, device=self.device, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), draft_model_idx=self.draft_model_idx, speculative_eagle_topk=self.server_args.speculative_eagle_topk, ) # Initialize token_to_kv_pool is_dsa_model = is_deepseek_dsa(self.model_config.hf_config) is_dsv4_model = is_deepseek_v4(self.model_config.hf_config) self._validate_prefill_only_disable_kv_cache_pool_family( is_dsa_model, is_dsv4_model, current_platform ) token_to_kv_pool = self._build_token_to_kv_pool( sizes=sizes, is_dsa_model=is_dsa_model, is_dsv4_model=is_dsv4_model, req_to_token_pool=req_to_token_pool, ) token_to_kv_pool_allocator = self._build_token_to_kv_pool_allocator( sizes=sizes, token_to_kv_pool=token_to_kv_pool, is_dsv4_model=is_dsv4_model, req_to_token_pool=req_to_token_pool, token_to_kv_pool_allocator=token_to_kv_pool_allocator, ) # Defensive check: the explicit validation above should reject known # unsupported pool families before allocation. Keep this guard here so # future pool-selection refactors fail at boot instead of on first use. if ( self.server_args.prefill_only_disable_kv_cache and not self.is_draft_worker and not isinstance(token_to_kv_pool, NoOpMHATokenToKVPool) ): raise RuntimeError( "--prefill-only-disable-kv-cache expected NoOpMHATokenToKVPool but the " f"runtime pool is {type(token_to_kv_pool).__name__}. This pool " "family is not yet supported by --prefill-only-disable-kv-cache. " "Supported configurations today: plain MHA models on CUDA with the FA " "(fa3/fa4) prefill backend, --is-embedding, --chunked-prefill-size=-1, " "--disable-radix-cache, no context-parallel attention, no HiSparse, " "and --kv-cache-dtype not in {nvfp4, fp4_mx_block16}." ) return _InitializedPools( req_to_token_pool=req_to_token_pool, token_to_kv_pool=token_to_kv_pool, token_to_kv_pool_allocator=token_to_kv_pool_allocator, ) def _init_unified_mamba_pools( self, *, max_num_reqs: int, max_total_num_tokens: int ) -> UnifiedPoolBundle: """Build the shared-KV-pool stack for a hybrid-Mamba model: one byte buffer split between the full-attn MHA KV pool and the per-request Mamba state pool, with virtual slot ids above the allocator.""" from sglang.srt.mem_cache.unified_memory_pool import init_unified_mamba_pools config = self.mambaish_config assert config is not None assert ( not self.use_mla_backend ), "unified memory pool does not support MLA-hybrid-Mamba yet" # The full sub-pool is page-aware (via `MultiEndedAllocator(page_size=...)`); # the mamba sub-pool stays page=1. assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}" # Mirror the non-shared path's extra_max_context_len computation. extra_max_context_len = 4 if self.server_args.speculative_num_draft_tokens is not None: extra_max_context_len += self.server_args.speculative_num_draft_tokens mamba_layer_ids = [ i for i in config.mamba2_cache_params.layers if self.layer_info.start_layer <= i < self.layer_info.end_layer ] full_attention_layer_ids = [ i for i in config.full_attention_layer_ids if self.layer_info.start_layer <= i < self.layer_info.end_layer ] bundle = init_unified_mamba_pools( device=self.device, kv_cache_dtype=self.kv_cache_dtype, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, page_size=self.page_size, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, is_draft_worker=self.is_draft_worker, use_mla_backend=self.use_mla_backend, mamba_layer_ids=mamba_layer_ids, full_attention_layer_ids=full_attention_layer_ids, mamba2_cache_params=config.mamba2_cache_params, model_context_len=self.model_config.context_len, extra_max_context_len=extra_max_context_len, max_total_num_tokens=max_total_num_tokens, max_mamba_cache_size=self.server_args.max_mamba_cache_size, max_num_reqs=max_num_reqs, enable_memory_saver=self.server_args.enable_memory_saver, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, disable_overlap_schedule=self.server_args.disable_overlap_schedule, need_sort=self.server_args.disaggregation_mode in ("decode", "prefill"), mamba_full_memory_ratio=self.server_args.mamba_full_memory_ratio, # Overlap mode: the allocator's `free` drops a wait_stream(forward_stream) # barrier so eager compaction serializes after the in-flight forward's # v2p/KV reads. Near-no-op in normal mode. forward_stream=self.forward_stream, # Lazy compaction: default ON, env-var escape hatch for rollback / A/B. lazy_compaction=_should_enable_lazy_compaction(), ) return bundle def _init_unified_swa_pools( self, *, max_num_reqs: int, full_max_total_num_tokens: Optional[int], swa_max_total_num_tokens: Optional[int], ) -> UnifiedPoolBundle: """Build the unified-pool stack for a hybrid-SWA model (Triton): one byte buffer split between the full-attention and SWA KV pools.""" from sglang.srt.mem_cache.unified_memory_pool import ( UnifiedPoolBundle, init_unified_swa_pools, ) assert self.is_hybrid_swa, "_init_unified_swa_pools called on a non-SWA model" # Both sub-pools are page-aware; the SWA composite runs alloc_extend_kernel # once in virtual space and binds the new pages on both sub-allocators. assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}" assert ( not self.use_mla_backend ), "unified memory pool does not support MLA-SWA hybrid yet" # Mirror the non-shared path's extra_max_context_len computation. extra_max_context_len = 4 if self.server_args.speculative_num_draft_tokens is not None: extra_max_context_len += self.server_args.speculative_num_draft_tokens req_to_token_pool = ReqToTokenPool( size=max_num_reqs, max_context_len=self.model_config.context_len + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, ) head_num = self.model_config.get_num_kv_heads(get_parallel().attn_tp_size) head_dim = self.model_config.head_dim if self.is_hybrid_swa_compress: # Asymmetric head dims between full and SWA (NPU compress path): # pull SWA-specific dims from the hf text config. v_head_dim = self.model_config.hf_text_config.v_head_dim swa_head_num = max( 1, self.model_config.hf_text_config.swa_num_key_value_heads // get_parallel().attn_tp_size, ) swa_head_dim = self.model_config.hf_text_config.swa_head_dim swa_v_head_dim = self.model_config.hf_text_config.swa_v_head_dim else: v_head_dim = head_dim swa_head_num = head_num swa_head_dim = head_dim swa_v_head_dim = head_dim # Filter layer ids to this worker's [start_layer, end_layer) range. swa_attention_layer_ids = [ i for i in self.model_config.swa_attention_layer_ids if self.layer_info.start_layer <= i < self.layer_info.end_layer ] full_attention_layer_ids = [ i for i in self.model_config.full_attention_layer_ids if self.layer_info.start_layer <= i < self.layer_info.end_layer ] bundle = init_unified_swa_pools( device=self.device, kv_cache_dtype=self.kv_cache_dtype, head_num=head_num, head_dim=head_dim, v_head_dim=v_head_dim, swa_head_num=swa_head_num, swa_head_dim=swa_head_dim, swa_v_head_dim=swa_v_head_dim, page_size=self.page_size, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, swa_attention_layer_ids=swa_attention_layer_ids, full_attention_layer_ids=full_attention_layer_ids, full_max_total_num_tokens=full_max_total_num_tokens, swa_max_total_num_tokens=swa_max_total_num_tokens, enable_memory_saver=self.server_args.enable_memory_saver, need_sort=self.server_args.disaggregation_mode in ("decode", "prefill"), # Overlap mode: same wait_stream(forward_stream) rationale as # `_init_unified_mamba_pools`. forward_stream=self.forward_stream, # Lazy compaction: default ON, with env var escape hatch for rollback / A/B. lazy_compaction=_should_enable_lazy_compaction(), ) return UnifiedPoolBundle( unified_memory_pool=bundle.unified_memory_pool, token_to_kv_pool=bundle.token_to_kv_pool, token_to_kv_pool_allocator=bundle.token_to_kv_pool_allocator, req_to_token_pool=req_to_token_pool, ) def _validate_prefill_only_disable_kv_cache_pool_family( self, is_dsa_model: bool, is_dsv4_model: bool, current_platform, ): if not self.server_args.prefill_only_disable_kv_cache or self.is_draft_worker: return unsupported_pool_family = None if is_dsv4_model: unsupported_pool_family = "DeepSeekV4TokenToKVPool" elif current_platform.is_out_of_tree() and not self.mambaish_config: unsupported_pool_family = "out-of-tree platform KV pool" elif ( self.server_args.attention_backend == "ascend" and not self.mambaish_config ): unsupported_pool_family = "NPU/Ascend KV pool" elif self.use_mla_backend and is_dsa_model: unsupported_pool_family = "DSA/MLA KV pool" elif self.use_mla_backend and not self.mambaish_config: unsupported_pool_family = "MLA KV pool" elif self.is_hybrid_swa: unsupported_pool_family = "SWA KV pool" elif self.mambaish_config: unsupported_pool_family = "hybrid linear/Mamba KV pool" elif is_float4_e2m1fn_x2(self.kv_cache_dtype): unsupported_pool_family = "FP4 MHA KV pool" if unsupported_pool_family is not None: raise RuntimeError( "--prefill-only-disable-kv-cache is not supported for " f"{unsupported_pool_family}. Supported configurations today: plain MHA " "models on CUDA with the FA (fa3/fa4) prefill backend, --is-embedding, " "--chunked-prefill-size=-1, --disable-radix-cache, no context-parallel " "attention, no HiSparse, and --kv-cache-dtype not in {nvfp4, fp4_mx_block16}." ) def _build_req_to_token_pool(self, *, max_num_reqs: int) -> ReqToTokenPool: extra_max_context_len = get_req_to_token_extra_context_len(self.server_args) if self.server_args.disaggregation_mode == "decode": # Extra slots for pre-allocated requests pre_alloc_size = self.server_args.disaggregation_decode_extra_slots if self.mambaish_config: req_to_token_pool = self._build_hybrid_mamba_decode_req_pool( max_num_reqs=max_num_reqs, extra_max_context_len=extra_max_context_len, pre_alloc_size=pre_alloc_size, ) else: req_to_token_pool = self._build_decode_req_pool( max_num_reqs=max_num_reqs, extra_max_context_len=extra_max_context_len, pre_alloc_size=pre_alloc_size, ) elif self.mambaish_config: req_to_token_pool = self._build_hybrid_req_pool( max_num_reqs=max_num_reqs, extra_max_context_len=extra_max_context_len, ) else: req_to_token_pool = self._build_default_req_pool( max_num_reqs=max_num_reqs, extra_max_context_len=extra_max_context_len, ) return req_to_token_pool def _build_hybrid_mamba_decode_req_pool( self, *, max_num_reqs: int, extra_max_context_len: int, pre_alloc_size: int, ) -> ReqToTokenPool: from sglang.srt.disaggregation.decode import ( HybridMambaDecodeReqToTokenPool, ) req_to_token_pool = HybridMambaDecodeReqToTokenPool( size=max_num_reqs, max_context_len=self.model_config.context_len + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, cache_params=self.mambaish_config.mamba2_cache_params, mamba_layer_ids=( [ i for i in self.mambaish_config.mamba2_cache_params.layers if self.layer_info.start_layer <= i < self.layer_info.end_layer ] ), speculative_num_draft_tokens=self.server_args.max_speculative_num_draft_tokens, speculative_eagle_topk=self.server_args.speculative_eagle_topk, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), pre_alloc_size=pre_alloc_size, enable_overlap_schedule=not self.server_args.disable_overlap_schedule, mamba_size=self.server_args.max_mamba_cache_size, start_layer=self.layer_info.start_layer, ) return req_to_token_pool def _build_decode_req_pool( self, *, max_num_reqs: int, extra_max_context_len: int, pre_alloc_size: int, ) -> ReqToTokenPool: from sglang.srt.disaggregation.decode import DecodeReqToTokenPool req_to_token_pool = DecodeReqToTokenPool( size=max_num_reqs, max_context_len=self.model_config.context_len + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, pre_alloc_size=pre_alloc_size, ) return req_to_token_pool def _build_hybrid_req_pool( self, *, max_num_reqs: int, extra_max_context_len: int, ) -> ReqToTokenPool: req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=self.server_args.max_mamba_cache_size, mamba_spec_state_size=max_num_reqs, max_context_len=self.model_config.context_len + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, cache_params=self.mambaish_config.mamba2_cache_params, mamba_layer_ids=( [ i for i in self.mambaish_config.mamba2_cache_params.layers if self.layer_info.start_layer <= i < self.layer_info.end_layer ] ), enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), enable_mamba_extra_buffer_lazy=self.server_args.enable_mamba_extra_buffer_lazy(), speculative_num_draft_tokens=self.server_args.max_speculative_num_draft_tokens, speculative_eagle_topk=self.server_args.speculative_eagle_topk, enable_overlap_schedule=not self.server_args.disable_overlap_schedule, start_layer=self.layer_info.start_layer, enable_linear_replayssm=self.server_args.enable_linear_replayssm, linear_replayssm_cache_len=self.server_args.linear_replayssm_cache_len, mamba_envelope_layout=self.server_args.enable_page_major_kv_layout, # ReplaySSM spec-verify is GDN-only: activate the pool machinery # (rings + cursors + the intermediate_ssm gate) only for GDN-hybrid # models, so any other mamba-ish model (Mamba2/Nemotron, lightning, # ...) run with the flag set stays byte-identical to flag-off. enable_gdn_replayssm_spec=( self.server_args.enable_gdn_replayssm_spec and self.hybrid_gdn_config is not None ), ) return req_to_token_pool def _build_default_req_pool( self, *, max_num_reqs: int, extra_max_context_len: int, ) -> ReqToTokenPool: # DSV4 on NPU needs an extended ReqToTokenPool holding per-req # swa/c4/c128/c{4,128}_state tables; others stay on the stock one. req_to_token_pool_cls = ReqToTokenPool if _is_npu and is_deepseek_v4(self.model_config.hf_config): from sglang.srt.hardware_backend.npu.dsv4.dsv4_req_to_token_pool import ( DSV4NPUReqToTokenPool, ) req_to_token_pool_cls = DSV4NPUReqToTokenPool req_to_token_pool = req_to_token_pool_cls( size=max_num_reqs, max_context_len=self.model_config.context_len + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, ) return req_to_token_pool def _build_token_to_kv_pool( self, *, sizes: _PoolSizes, is_dsa_model: bool, is_dsv4_model: bool, req_to_token_pool: ReqToTokenPool, ) -> KVCache: # Page-granularity envelope layout for the MHA-shaped (full / SWA) pools, # selected by swapping in the PageMajorMHATokenToKVPool subclass. The # default keeps upstream's per-layer layout. The Mamba state pool is routed # separately via `mamba_envelope_layout` on the req-to-token pool above. enable_page_major = self.server_args.enable_page_major_kv_layout mha_pool_class = ( PageMajorMHATokenToKVPool if enable_page_major else MHATokenToKVPool ) if is_dsv4_model: token_to_kv_pool = self._build_dsv4_kv_pool( max_running_requests=sizes.max_running_requests, swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, c4_max_total_num_tokens=sizes.c4_max_total_num_tokens, c128_max_total_num_tokens=sizes.c128_max_total_num_tokens, c4_state_pool_size=sizes.c4_state_pool_size, c128_state_pool_size=sizes.c128_state_pool_size, c4_state_dtype=sizes.c4_state_dtype, c128_state_dtype=sizes.c128_state_dtype, req_to_token_pool=req_to_token_pool, ) elif current_platform.is_out_of_tree() and not self.mambaish_config: if self.use_mla_backend and is_dsa_model: token_to_kv_pool = self._build_oot_dsa_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, ) elif self.use_mla_backend: token_to_kv_pool = self._build_oot_mla_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, is_dsa_model=is_dsa_model, ) else: token_to_kv_pool = self._build_oot_mha_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, ) elif ( self.server_args.attention_backend == "ascend" and not self.mambaish_config ): if self.is_hybrid_swa: token_to_kv_pool = self._build_ascend_swa_kv_pool( full_max_total_num_tokens=sizes.full_max_total_num_tokens, swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, ) elif self.use_mla_backend: token_to_kv_pool = self._build_ascend_mla_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, is_dsa_model=is_dsa_model, ) else: token_to_kv_pool = self._build_ascend_mha_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, ) elif self.use_mla_backend and is_dsa_model: token_to_kv_pool = self._build_dsa_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, ) elif self.use_mla_backend and not self.mambaish_config: assert not is_dsa_model if is_float4_e2m1fn_x2(self.kv_cache_dtype): token_to_kv_pool = self._build_mla_fp4_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, ) else: token_to_kv_pool = self._build_mla_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, ) else: if self.is_hybrid_swa: token_to_kv_pool = self._build_hybrid_swa_kv_pool( full_max_total_num_tokens=sizes.full_max_total_num_tokens, swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, mha_pool_class=mha_pool_class, ) elif is_minimax_sparse(self.model_config.hf_config): token_to_kv_pool = self._build_minimax_sparse_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, ) elif self.mambaish_config: token_to_kv_pool = self._build_hybrid_linear_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, req_to_token_pool=req_to_token_pool, mha_pool_class=mha_pool_class, ) else: quant_method = None if is_float4_e2m1fn_x2(self.kv_cache_dtype): assert ( not enable_page_major ), "page-major KV layout is not supported with fp4 KV cache" quant_method = self._build_fp4_quant_method( num_layers=self.layer_info.num_effective_layers ) token_to_kv_pool = self._build_mha_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, mha_pool_class=mha_pool_class, quant_method=quant_method, ) return token_to_kv_pool def _build_dsv4_kv_pool( self, *, max_running_requests: int, swa_max_total_num_tokens: Optional[int], c4_max_total_num_tokens: int, c128_max_total_num_tokens: int, c4_state_pool_size: int, c128_state_pool_size: int, c4_state_dtype: Optional[torch.dtype], c128_state_dtype: Optional[torch.dtype], req_to_token_pool: ReqToTokenPool, ) -> KVCache: swa_page_size = self.server_args.page_size if not _is_npu: assert swa_page_size == 256, "In paged swa mode, page_size must be 256." if self.is_draft_worker: from sglang.srt.models.agnes_nextn import ( COMPRESS_RATIO_NEXTN_LAYER, ) compression_ratios = [ COMPRESS_RATIO_NEXTN_LAYER ] * self.layer_info.num_effective_layers else: compression_ratios = self.model_config.compress_ratios # NPU + DSV4 → paged-state subclass: the fused compressor kernel # needs cache_mode=1 (paged); Atlas A3 rejects cache_mode=2 (ring), # so the CUDA ring-buffer state path can't be shared. CUDA keeps # DeepSeekV4TokenToKVPool unchanged; NPU recomputes state sizes below. if _is_npu: from sglang.srt.hardware_backend.npu.dsv4.dsv4_memory_pool import ( DSV4NPUTokenToKVPool, npu_state_pool_size, ) pool_cls = DSV4NPUTokenToKVPool # Recompute state pool sizes for the NPU paged formula (CUDA's # ring sizes are dropped here). Tail-only allocation keeps the # per-req-budget formula sufficient at any prefill length: long # prompts allocate only ``tail+128`` (c4) / ``tail`` (c128) # slots (tail = seq_len % 128), and decode is drained by # sliding eviction in ``ScheduleBatch._evict_swa``. c4_state_pool_size = npu_state_pool_size( ratio=4, page_size=self.server_args.page_size, max_num_reqs=max_running_requests, ) c128_state_pool_size = npu_state_pool_size( ratio=128, page_size=self.server_args.page_size, max_num_reqs=max_running_requests, ) else: pool_cls = DeepSeekV4TokenToKVPool c4_state_pool_size = c4_state_pool_size c128_state_pool_size = c128_state_pool_size token_to_kv_pool = pool_cls( max_num_reqs=max_running_requests, # SWA ring is indexed by req_pool_idx; PD decode inflates req_to_token # past max_running_requests (pre-alloc), so size to the real capacity. num_req_slots=req_to_token_pool.req_to_token.shape[0], swa_size=swa_max_total_num_tokens, c4_size=c4_max_total_num_tokens, c128_size=c128_max_total_num_tokens, c4_state_pool_size=c4_state_pool_size, c128_state_pool_size=c128_state_pool_size, page_size=self.server_args.page_size, swa_page_size=swa_page_size, sliding_window=self.model_config.window_size, dtype=self.kv_cache_dtype, c4_state_dtype=c4_state_dtype, c128_state_dtype=c128_state_dtype, qk_nope_head_dim=self.model_config.qk_nope_head_dim, qk_rope_head_dim=self.model_config.qk_rope_head_dim, indexer_head_dim=self.model_config.index_head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, compression_ratios=compression_ratios, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, enable_hisparse=self.server_args.enable_hisparse, online_mtp_max_draft_tokens=( self.server_args.max_speculative_num_draft_tokens or 0 ), ) return token_to_kv_pool def _build_oot_dsa_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: PoolCls = current_platform.get_dsa_kv_pool_cls() token_to_kv_pool = PoolCls( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, kv_cache_dim=calculate_mla_kv_cache_dim( model_config=self.model_config, kv_cache_dtype=self.kv_cache_dtype, server_args=self.server_args, ), enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), ) return token_to_kv_pool def _build_oot_mla_kv_pool( self, *, max_total_num_tokens: int, is_dsa_model: bool ) -> KVCache: PoolCls = current_platform.get_mla_kv_pool_cls() token_to_kv_pool = PoolCls( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, index_head_dim=(self.model_config.index_head_dim if is_dsa_model else None), layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, ) return token_to_kv_pool def _build_oot_mha_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: PoolCls = current_platform.get_mha_kv_pool_cls() token_to_kv_pool = PoolCls( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, ) return token_to_kv_pool def _build_ascend_swa_kv_pool( self, *, full_max_total_num_tokens: Optional[int], swa_max_total_num_tokens: Optional[int], ) -> KVCache: from sglang.srt.hardware_backend.npu.memory_pool_npu import ( NPUMHATokenToKVPool, ) kwargs = {} if self.is_hybrid_swa_compress: kwargs = { "swa_head_num": max( 1, self.model_config.hf_text_config.swa_num_key_value_heads // get_parallel().attn_tp_size, ), "swa_head_dim": self.model_config.swa_head_dim, "swa_v_head_dim": self.model_config.swa_v_head_dim, "v_head_dim": self.model_config.v_head_dim, } token_to_kv_pool = SWAKVPool( size=full_max_total_num_tokens, size_swa=swa_max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, post_capture_active=self.post_capture_kv_active, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, swa_attention_layer_ids=self.model_config.swa_attention_layer_ids, full_attention_layer_ids=self.model_config.full_attention_layer_ids, device=self.device, token_to_kv_pool_class=NPUMHATokenToKVPool, **kwargs, ) return token_to_kv_pool def _build_ascend_mla_kv_pool( self, *, max_total_num_tokens: int, is_dsa_model: bool ) -> KVCache: from sglang.srt.hardware_backend.npu.memory_pool_npu import ( NPUMLATokenToKVPool, ) token_to_kv_pool = NPUMLATokenToKVPool( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, index_head_dim=(self.model_config.index_head_dim if is_dsa_model else None), layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, ) return token_to_kv_pool def _build_ascend_mha_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: from sglang.srt.hardware_backend.npu.memory_pool_npu import ( NPUMHATokenToKVPool, ) token_to_kv_pool = NPUMHATokenToKVPool( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, ) return token_to_kv_pool def _build_dsa_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: from sglang.srt.layers.cp.utils import get_glm_dsa_cp_layer_shard_info ( dsa_cp_layer_shard_rank, dsa_cp_layer_shard_size, ) = get_glm_dsa_cp_layer_shard_info(self) pool_kwargs = {} if self.server_args.enable_hisparse: PoolCls = HiSparseDSATokenToKVPool from sglang.srt.mem_cache.sparsity import parse_hisparse_config pool_kwargs["host_to_device_ratio"] = parse_hisparse_config( self.server_args ).host_to_device_ratio elif dsa_cp_layer_shard_rank is not None: # DSA cache layer split: shard KV/indexer layers across CP ranks. from sglang.srt.mem_cache.dsa_cache_layer_split import ( LayerSplitDSATokenToKVPool, ) PoolCls = LayerSplitDSATokenToKVPool pool_kwargs["layer_shard_rank"] = dsa_cp_layer_shard_rank pool_kwargs["layer_shard_size"] = dsa_cp_layer_shard_size else: PoolCls = DSATokenToKVPool token_to_kv_pool = PoolCls( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, kv_cache_dim=calculate_mla_kv_cache_dim( model_config=self.model_config, kv_cache_dtype=self.kv_cache_dtype, server_args=self.server_args, ), enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), **pool_kwargs, ) return token_to_kv_pool def _build_mla_fp4_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: token_to_kv_pool = MLATokenToKVPoolFP4( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, ) return token_to_kv_pool def _build_mla_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: token_to_kv_pool = MLATokenToKVPool( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, ) return token_to_kv_pool def _build_hybrid_swa_kv_pool( self, *, full_max_total_num_tokens: Optional[int], swa_max_total_num_tokens: Optional[int], mha_pool_class: type, ) -> KVCache: kwargs = {} if self.is_hybrid_swa_compress: kwargs = { "swa_head_num": max( 1, self.model_config.hf_text_config.swa_num_key_value_heads // get_parallel().attn_tp_size, ), "swa_head_dim": self.model_config.swa_head_dim, "swa_v_head_dim": self.model_config.swa_v_head_dim, "v_head_dim": self.model_config.v_head_dim, } swa_pool_class = ( MHATokenToKVPoolMXFP8 if self.server_args.kv_cache_dtype == "mxfp8" else mha_pool_class ) swa_attention_layer_ids = self.model_config.swa_attention_layer_ids full_attention_layer_ids = self.model_config.full_attention_layer_ids if self.is_inkling_mtp_draft: if self.draft_swa_full_capacity: # Banded 's' depth: route the draft's single layer into the SWA # ring sub-pool so use_sliding_window_kv_pool activates the SWA # store/read path for this depth, exactly like a trunk local # layer. swa_attention_layer_ids = [self.draft_model_idx] full_attention_layer_ids = [] else: swa_attention_layer_ids = [] full_attention_layer_ids = [self.draft_model_idx] # Size the banded draft's SWA ring to FULL draft capacity (not the # trunk-window-derived swa_max): with the identity full->swa mapping # registered in _build_token_to_kv_pool_allocator, every logical slot # the shared target allocator hands out (up to full_max) must be # addressable in the ring, whatever the head-vs-trunk window # relationship. size_swa = ( full_max_total_num_tokens if self.draft_swa_full_capacity else swa_max_total_num_tokens ) token_to_kv_pool = SWAKVPool( size=full_max_total_num_tokens, size_swa=size_swa, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, post_capture_active=self.post_capture_kv_active, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, swa_attention_layer_ids=swa_attention_layer_ids, full_attention_layer_ids=full_attention_layer_ids, device=self.device, enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None), token_to_kv_pool_class=swa_pool_class, **kwargs, ) return token_to_kv_pool def _build_minimax_sparse_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: _hf_config = self.model_config.hf_config sparse_cfg = get_minimax_sparse_attention_config(_hf_config) dense_layer_ids, sparse_layer_ids = get_minimax_sparse_layer_ids(sparse_cfg) disable_value_sparse_layer_ids = get_minimax_sparse_disable_value_layer_ids( sparse_cfg ) token_to_kv_pool = MiniMaxSparseKVPool( size=max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, index_dtype=self.model_dtype, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, idx_head_dim=sparse_cfg["sparse_index_dim"], dense_layer_ids=dense_layer_ids, sparse_layer_ids=sparse_layer_ids, disable_value_sparse_layer_ids=disable_value_sparse_layer_ids, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, ) return token_to_kv_pool def _build_hybrid_linear_kv_pool( self, *, max_total_num_tokens: int, req_to_token_pool: ReqToTokenPool, mha_pool_class: type, ) -> KVCache: extra_args = {} if self.use_mla_backend: extra_args = { "kv_lora_rank": self.model_config.kv_lora_rank, "qk_rope_head_dim": self.model_config.qk_rope_head_dim, } full_attention_layer_ids = ( [0] if self.is_draft_worker else [ i for i in self.mambaish_config.full_attention_layer_ids if self.layer_info.start_layer <= i < self.layer_info.end_layer ] ) quant_method = self._build_fp4_quant_method( num_layers=len(full_attention_layer_ids) ) # MXFP8 KV cache needs the block-scaled pool (data + UE8M0 scale # buffers) for the full-attention layers, same as the SWA branch. full_pool_class = ( MHATokenToKVPoolMXFP8 if self.server_args.kv_cache_dtype == "mxfp8" and not self.use_mla_backend else mha_pool_class ) token_to_kv_pool = HybridLinearKVPool( page_size=self.server_args.page_size, size=max_total_num_tokens, dtype=self.kv_cache_dtype, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, # if draft worker, we only need 1 attention layer's kv pool full_attention_layer_ids=full_attention_layer_ids, device=self.device, mamba_pool=req_to_token_pool.mamba_pool, enable_memory_saver=self.server_args.enable_memory_saver, enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None), use_mla=self.use_mla_backend, start_layer=self.layer_info.start_layer, full_kv_pool_class=full_pool_class, quant_method=quant_method, post_capture_active=self.post_capture_kv_active and quant_method is None, **extra_args, ) return token_to_kv_pool def _build_mha_fp4_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: token_to_kv_pool = MHATokenToKVPoolFP4( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, v_head_dim=self.model_config.v_head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, enable_alt_stream=not self.server_args.enable_pdmux, enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None), ) return token_to_kv_pool def _build_mha_kv_pool( self, *, max_total_num_tokens: int, mha_pool_class: type, quant_method=None ) -> KVCache: if self.server_args.kv_cache_dtype == "mxfp8": pool_cls = MHATokenToKVPoolMXFP8 else: pool_cls = ( NoOpMHATokenToKVPool if self.server_args.prefill_only_disable_kv_cache else mha_pool_class ) pool_kwargs = {} if quant_method is not None: pool_kwargs["quant_method"] = quant_method else: pool_kwargs["post_capture_active"] = self.post_capture_kv_active token_to_kv_pool = pool_cls( max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_dim=self.model_config.head_dim, v_head_dim=self.model_config.v_head_dim, layer_num=self.layer_info.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, enable_alt_stream=not self.server_args.enable_pdmux, enable_kv_cache_copy=(self.server_args.speculative_algorithm is not None), **pool_kwargs, ) return token_to_kv_pool def _build_token_to_kv_pool_allocator( self, *, sizes: _PoolSizes, token_to_kv_pool: KVCache, is_dsv4_model: bool, req_to_token_pool: ReqToTokenPool, token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator], ) -> BaseTokenToKVPoolAllocator: # Initialize token_to_kv_pool_allocator need_sort = self.server_args.disaggregation_mode in ("decode", "prefill") if token_to_kv_pool_allocator is None: if current_platform.is_out_of_tree(): AllocatorCls = current_platform.get_paged_allocator_cls() token_to_kv_pool_allocator = AllocatorCls( sizes.max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, ) elif _is_npu and ( self.server_args.attention_backend == "ascend" or is_dsv4_model or self.hybrid_gdn_config is not None ): if self.is_hybrid_swa: # DSV4 on NPU: SWA allocator subclass that also drives the # c4/c128 allocators, producing a DSV4OutCacheLoc per alloc. if is_dsv4_model: from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import ( DSV4NPUTokenToKVPoolAllocator, ) swa_allocator_cls = DSV4NPUTokenToKVPoolAllocator else: swa_allocator_cls = SWATokenToKVPoolAllocator token_to_kv_pool_allocator = swa_allocator_cls( sizes.full_max_total_num_tokens, sizes.swa_max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, ) else: from sglang.srt.hardware_backend.npu.allocator_npu import ( NPUPagedTokenToKVPoolAllocator, ) token_to_kv_pool_allocator = NPUPagedTokenToKVPoolAllocator( sizes.max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, ) else: if self.is_hybrid_swa and sizes.full_max_total_num_tokens == 0: token_to_kv_pool_allocator = PureSWATokenToKVPoolAllocator( sizes.swa_max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, ) elif self.is_hybrid_swa: token_to_kv_pool_allocator = SWATokenToKVPoolAllocator( sizes.full_max_total_num_tokens, sizes.swa_max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, ) else: if self.server_args.enable_hisparse: from sglang.srt.mem_cache.sparsity import ( parse_hisparse_config, ) hisparse_cfg = parse_hisparse_config(self.server_args) token_to_kv_pool_allocator = HiSparseTokenToKVPoolAllocator( sizes.max_total_num_tokens, page_size=self.server_args.page_size, dtype=self.kv_cache_dtype, device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, host_to_device_ratio=hisparse_cfg.host_to_device_ratio, ) elif ( self.server_args.page_size == 1 and self.server_args.dcp_size == 1 ): token_to_kv_pool_allocator = TokenToKVPoolAllocator( sizes.max_total_num_tokens, dtype=self.kv_cache_dtype, device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, ) else: token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator( sizes.max_total_num_tokens * self.server_args.dcp_size, page_size=self.server_args.page_size * self.server_args.dcp_size, dtype=self.kv_cache_dtype, device=self.device, kvcache=token_to_kv_pool, need_sort=need_sort, ) if self.server_args.enable_hisparse and is_dsv4_model: assert self.is_hybrid_swa, "DeepSeek V4 HiSparse requires SWA mode." token_to_kv_pool_allocator = DeepSeekV4HiSparseTokenToKVPoolAllocator( token_to_kv_pool_allocator ) # DSV4-NPU: wire allocator back-ref into req_to_token_pool so its # free(req) can release c4/c128 pool pages alongside the slot. if hasattr(req_to_token_pool, "register_dsv4_allocator"): req_to_token_pool.register_dsv4_allocator(token_to_kv_pool_allocator) else: assert self.is_draft_worker if self.is_hybrid_swa: if self.draft_swa_full_capacity: # Banded depth: the SWA ring is full draft capacity, so use # an IDENTITY full->swa mapping — store and read locs both # equal out_cache_loc, and a slot is never evicted before # the request frees it. The window itself is enforced by the # FA sliding-window kernel, not by the ring. Layout mirrors # SWATokenToKVPoolAllocator's mapping (size + page_size # entries + trailing -1 sentinel so a -1 last_loc maps # to -1). n = sizes.full_max_total_num_tokens + self.page_size identity_mapping = torch.arange( n + 1, dtype=torch.int64, device=self.device ) identity_mapping[-1] = -1 token_to_kv_pool.register_mapping(identity_mapping) else: swa_allocator = getattr( token_to_kv_pool_allocator, "logical_attn_allocator", token_to_kv_pool_allocator, ) assert isinstance(swa_allocator, SWATokenToKVPoolAllocator) token_to_kv_pool.register_mapping( swa_allocator.full_to_swa_index_mapping ) return token_to_kv_pool_allocator def _profile_available_bytes(self, pre_model_load_memory: int) -> int: # KV pool budget = currently-free GPU memory minus the non-static runtime # slack (pre_model_load_memory * (1 - mem_fraction_static)). Whatever is # already resident (model weights, etc.) is thus charged against it. available_gpu_memory = get_available_gpu_memory( self.device, self.gpu_id, distributed=get_world_group().world_size > 1, cpu_group=get_world_group().cpu_group, ) slack_gb = pre_model_load_memory * (1 - self.server_args.mem_fraction_static) if self.mambaish_config is not None and self.post_capture_kv_active: # Mamba state is a fixed pre-capture allocation, so it can't ride the ~0 post-capture slack. slack_gb = max( slack_gb, self.server_args.mamba_pre_capture_reserve_mb( get_device_memory_capacity(self.device) ) / 1024, ) rest_memory = available_gpu_memory - slack_gb if self.mambaish_config is not None: rest_memory = self._handle_max_mamba_cache(rest_memory) # Loaded weights (target + draft) can exceed the static budget if rest_memory <= 0: minimum_mem_fraction_static = ( 1 - available_gpu_memory / pre_model_load_memory ) suggested_mem_fraction_static = ( math.ceil(minimum_mem_fraction_static * 1000) / 1000 ) raise ValueError( f"Loaded weights leave no GPU memory for the KV cache under " f"--mem-fraction-static={self.server_args.mem_fraction_static}. " f"Raise --mem-fraction-static above " f"{suggested_mem_fraction_static:.3f} " f"(minimum viable = 1 - available/pre = " f"{minimum_mem_fraction_static:.4f}). If using speculative " f"decoding, draft weights are now counted." ) return int(rest_memory * (1 << 30)) # return in bytes def _calculate_mamba_ratio(self) -> int: if self.server_args.disable_radix_cache: return 1 additional_ratio = 0 if self.server_args.enable_mamba_extra_buffer(): # ping-pong buffer size is 2 when overlap schedule is on, 1 otherwise. # Lazy mode saves 1 slot (2 → 1) for overlap; non-overlap already uses 1. if not self.server_args.disable_overlap_schedule: if self.server_args.enable_mamba_extra_buffer_lazy(): additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP_LAZY else: additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP else: assert ( not self.server_args.enable_mamba_extra_buffer_lazy() ), "Lazy extra buffer requires overlap schedule (--disable-overlap-schedule is incompatible)" additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP return MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio def _apply_token_constraints(self, token_capacity: int) -> int: """Apply external constraints to token capacity: user cap, PP sync. Page alignment is handled by the configurator, not here. If constraints change the value, the configurator re-runs and re-aligns. """ user_limit = self.server_args.max_total_tokens # Apply user-specified upper bound if user_limit is not None: if user_limit > token_capacity: logging.warning( f"max_total_tokens={user_limit} is larger than the profiled value " f"{token_capacity}. Use the profiled value instead." ) token_capacity = min(token_capacity, user_limit) # Sync across PP ranks (each may have different layer counts) if self.server_args.pp_size > 1: tensor = torch.tensor(token_capacity, dtype=torch.int64) torch.distributed.all_reduce( tensor, op=torch.distributed.ReduceOp.MIN, group=get_world_group().cpu_group, ) token_capacity = tensor.item() return token_capacity def resolve_max_num_reqs(self, token_capacity: int) -> int: """Compute max concurrent requests (per dp worker) from the finalized token capacity.""" # Estimate pool size (used as upper bound when user specifies max_running_requests) estimated = int(token_capacity / self.model_config.context_len * 512) estimated = max(min(estimated, 4096), 2048) max_num_reqs = self.server_args.max_running_requests if max_num_reqs is not None: requested_per_worker = max_num_reqs // self.ps.attn_dp_size max_num_reqs = min(requested_per_worker, token_capacity // 2) else: requested_per_worker = None max_num_reqs = min(estimated, token_capacity // 2) if self.mambaish_config is not None: ratio = self._calculate_mamba_ratio() max_num_reqs = min( max_num_reqs, self.server_args.max_mamba_cache_size // ratio ) if max_num_reqs <= 0: raise RuntimeError( f"Hybrid (mamba/linear-attention) state cache is too small to serve " f"any requests. max_mamba_cache_size={self.server_args.max_mamba_cache_size}, " f"mamba_ratio={ratio}, resulting max_num_reqs={max_num_reqs}. " f"Try: (1) reduce --max-running-requests, " f"(2) increase --mem-fraction-static, or " f"(3) use GPUs with more memory." ) if requested_per_worker is not None and max_num_reqs < requested_per_worker: logger.warning( "max_running_requests was reduced from the requested %d to %d " "(per dp worker) due to the available KV cache capacity.", requested_per_worker, max_num_reqs, ) return max_num_reqs def _resolve_memory_pool_config( self, pre_model_load_memory: int ) -> MemoryPoolConfig: """Profile GPU memory and resolve all pool parameters into a config.""" from sglang.srt.model_executor.pool_configurator import ( create_memory_pool_configurator, ) available_bytes = self._profile_available_bytes(pre_model_load_memory) config = self.config_from_budget(available_bytes) config.max_running_requests = self.resolve_max_num_reqs( config.max_total_num_tokens ) configurator = create_memory_pool_configurator(self) config = configurator.finalize_with_max_running_requests(config) config.mem_fraction_static = self.server_args.mem_fraction_static return config def config_from_budget( self, budget_bytes: int, *, cap_tokens: Optional[int] = None ) -> MemoryPoolConfig: """Turn a KV byte budget into a pool config via the configurator, re-applying the external token constraints (user cap, page alignment, PP sync) and the optional ``cap_tokens`` clamp.""" # Local import avoids a pool_configurator import cycle. from sglang.srt.model_executor.pool_configurator import ( create_memory_pool_configurator, ) configurator = create_memory_pool_configurator(self) config = configurator.calculate_pool_sizes( budget_bytes, self.server_args.page_size ) max_tokens = self._apply_token_constraints(config.max_total_num_tokens) if cap_tokens is not None: max_tokens = min(max_tokens, cap_tokens) if max_tokens != config.max_total_num_tokens: config = configurator.calculate_pool_sizes_from_max_tokens( max_tokens, self.server_args.page_size ) return config def _handle_max_mamba_cache(self, total_rest_memory): config = self.mambaish_config server_args = self.server_args assert config is not None has_spec_dec = not self.spec_algorithm.is_none() if has_spec_dec: assert server_args.speculative_num_draft_tokens is not None assert server_args.max_running_requests is not None if server_args.max_mamba_cache_size is not None: # Use explicitly set max_mamba_cache_size server_args.override( "mamba_pool.per_dp_shard", max_mamba_cache_size=server_args.max_mamba_cache_size // self.ps.attn_dp_size, ) # Reserve intermediate memory based on capped max_num_reqs if has_spec_dec: ratio = self._calculate_mamba_ratio() capped_reqs = min( server_args.max_running_requests // self.ps.attn_dp_size, server_args.max_mamba_cache_size // ratio, ) intermediate_size = ( config.mamba2_cache_params.mamba_cache_per_req * capped_reqs * server_args.speculative_num_draft_tokens ) total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30)) elif ( server_args.disable_radix_cache and server_args.max_running_requests is not None ): # Use explicitly set max_running_requests when radix cache is disabled server_args.override( "mamba_pool.from_max_running_requests", max_mamba_cache_size=server_args.max_running_requests // self.ps.attn_dp_size, ) # Reserve intermediate memory based on capped max_num_reqs if has_spec_dec: intermediate_size = ( config.mamba2_cache_params.mamba_cache_per_req * server_args.max_mamba_cache_size * server_args.speculative_num_draft_tokens ) total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30)) else: # Use ratio-based calculation to auto-fit available memory assert config.mamba2_cache_params.mamba_cache_per_req > 0 per_req = config.mamba2_cache_params.mamba_cache_per_req # Solve jointly for max_mamba_cache_size accounting for intermediate memory. # The mamba budget (from the ratio split) must cover both: # 1. main mamba state: max_mamba_cache_size * per_req # 2. intermediate states: (max_mamba_cache_size / ratio) * D * per_req # So: max_mamba_cache_size * per_req * (1 + D/ratio) = mamba_budget_bytes mamba_budget = ( total_rest_memory * server_args.mamba_full_memory_ratio / (1 + server_args.mamba_full_memory_ratio) ) mamba_budget_bytes = mamba_budget * (1 << 30) if has_spec_dec: ratio = self._calculate_mamba_ratio() D = server_args.speculative_num_draft_tokens # Joint solve: main_state + intermediate = mamba_budget server_args.override( "mamba_pool.memory_budget_spec", max_mamba_cache_size=int( mamba_budget_bytes // (per_req * (1 + D / ratio)) ), ) # Intermediate memory is included in mamba_budget, subtract it # so the return value only has main_state subtracted from total capped_reqs = min( server_args.max_running_requests // self.ps.attn_dp_size, server_args.max_mamba_cache_size // ratio, ) intermediate_size = per_req * capped_reqs * D total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30)) else: server_args.override( "mamba_pool.memory_budget", max_mamba_cache_size=int(mamba_budget_bytes // per_req), ) # Validate: max_mamba_cache_size must be positive after memory allocation. # A non-positive value means GPU memory is insufficient for the requested # configuration. Fail fast with actionable advice instead of silently # producing garbled output at runtime. if server_args.max_mamba_cache_size <= 0: raise RuntimeError( f"Not enough GPU memory for hybrid (mamba/linear-attention) state cache. " f"Computed max_mamba_cache_size={server_args.max_mamba_cache_size} " f"(total_rest_memory={total_rest_memory:.2f} GB, " f"mamba_cache_per_req={config.mamba2_cache_params.mamba_cache_per_req / (1 << 20):.2f} MB). " f"Try: (1) reduce --max-running-requests, " f"(2) increase --mem-fraction-static, " f"(3) reduce --speculative-num-draft-tokens, or " f"(4) use GPUs with more memory." ) mamba_state_memory = ( server_args.max_mamba_cache_size * config.mamba2_cache_params.mamba_cache_per_req / (1 << 30) ) return total_rest_memory - mamba_state_memory def calculate_mla_kv_cache_dim( *, model_config: ModelConfig, kv_cache_dtype: torch.dtype, server_args: ServerArgs, ) -> int: is_dsa_model = is_deepseek_dsa(model_config.hf_config) kv_cache_dtype = kv_cache_dtype kv_lora_rank = model_config.kv_lora_rank qk_rope_head_dim = model_config.qk_rope_head_dim kv_cache_dim = kv_lora_rank + qk_rope_head_dim # default mla kv cache dim # For non-DSA models, MLA kv cache dim is simply kv_lora_rank + qk_rope_head_dim if not is_dsa_model: return kv_cache_dim # TRTLLM backend does not override kv_cache_dim for MLA kv cache # Assuming dsa prefill and decode backends are the same when using trtllm MLA backend, # since it is not compatible for trtllm and other mla attn backend due to the different # kv cache layout. if ( server_args.dsa_prefill_backend == "trtllm" or server_args.dsa_decode_backend == "trtllm" ): return kv_cache_dim # On HIP, TileLang and AITER DSA kernels consume the raw MLA KV layout: # nope(512 fp8) + rope(64 fp8), without extra per-block scales. if _is_hip and ( server_args.dsa_prefill_backend in ("tilelang", "aiter") or server_args.dsa_decode_backend in ("tilelang", "aiter") ): return kv_cache_dim quant_block_size = DSATokenToKVPool.quant_block_size rope_storage_dtype = DSATokenToKVPool.rope_storage_dtype # Calculate override_kv_cache_dim for FP8 storage in backends that use scaled KV layout # (excluding TRTLLM and HIP raw-layout kernels). # kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage # Note: rope dimension is stored in original dtype (bf16), not quantized to fp8 if kv_cache_dtype == torch.float8_e4m3fn: assert ( kv_lora_rank % quant_block_size == 0 ), f"kv_lora_rank {kv_lora_rank} must be multiple of quant_block_size {quant_block_size}" return ( kv_lora_rank + kv_lora_rank // quant_block_size * 4 + qk_rope_head_dim * rope_storage_dtype.itemsize ) return kv_cache_dim