diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 260cdef32..b2856d0d7 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -379,6 +379,7 @@ class MambaPool: self.size = size self.device = device + self.debug_memory_pool = envs.SGLANG_DEBUG_MEMORY_POOL.get() self.enable_linear_replayssm = enable_linear_replayssm self.linear_replayssm_cache_len = linear_replayssm_cache_len @@ -700,7 +701,7 @@ class MambaPool: caps the donate to the last flush boundary. The dst cursor is reset to 0 (the copied checkpoint has no pending ring entries). """ - if self.replayssm_write_pos is not None and envs.SGLANG_DEBUG_MEMORY_POOL.get(): + if self.replayssm_write_pos is not None and self.debug_memory_pool: # Debug-only (syncs): catch any copy of an active, un-flushed slot. src_wp = self.replayssm_write_pos[src_indices] assert bool((src_wp == 0).all().item()), ( @@ -1381,6 +1382,7 @@ class MHATokenToKVPool(KVCache): # X = 16 / dtype_bytes — AITER-only (ignored elsewhere, no consumer kernel). # HND and vectorized_5d are mutually exclusive; HND takes precedence. self.use_hnd = envs.SGLANG_USE_HND_KVCACHE.get() + self.use_native_move_kv_cache = envs.SGLANG_NATIVE_MOVE_KV_CACHE.get() if kv_cache_layout is not None: # Explicit physical-layout selector wins over the platform default. # This is a label only; layouts that change buffer identity (e.g. the @@ -1981,7 +1983,7 @@ class MHATokenToKVPool(KVCache): # Physical move strategy. Override for layouts that change buffer identity # (e.g. PageMajorMHATokenToKVPool always uses the native move). The 3-D # per-layer buffers here ignore page_size in move_kv_cache_native. - if envs.SGLANG_NATIVE_MOVE_KV_CACHE.get(): + if self.use_native_move_kv_cache: move_kv_cache_native(self.k_buffer, self.v_buffer, tgt_loc, src_loc) return @@ -3558,6 +3560,9 @@ class MiniMaxSparseKVPool(KVCache): self.page_size = page_size self.dtype = dtype self.device = device + self.use_minimax_fused_kv_index_store = ( + envs.SGLANG_OPT_USE_MINIMAX_FUSED_KV_INDEX_STORE.get() + ) local_dense_layer_ids = [ lid for lid in dense_layer_ids if start_layer <= lid < end_layer @@ -3770,7 +3775,7 @@ class MiniMaxSparseKVPool(KVCache): head byte size shared by main and index caches.""" main = self.main_pool return ( - envs.SGLANG_OPT_USE_MINIMAX_FUSED_KV_INDEX_STORE.get() + self.use_minimax_fused_kv_index_store and _is_cuda # No dtype conversion / fp8 scaling on either side (the fused kernel # is a raw byte copy, it does not quantize). diff --git a/python/sglang/srt/mem_cache/storage/eic/eic_storage.py b/python/sglang/srt/mem_cache/storage/eic/eic_storage.py index adb1e4675..f9f4f2be2 100644 --- a/python/sglang/srt/mem_cache/storage/eic/eic_storage.py +++ b/python/sglang/srt/mem_cache/storage/eic/eic_storage.py @@ -267,6 +267,7 @@ class EICStorage(HiCacheStorage): self.world_size = hicache_config.tp_size self.page_size = self.memory_pool_host.page_size self.use_zero_copy = self.memory_pool_host.layout == "page_first" + self.mha_zero_copy = self.use_zero_copy and not self.is_mla_model if not self.use_zero_copy: self.kv_cache_shape = self.memory_pool_host.get_data_page( 0, flat=True @@ -410,7 +411,7 @@ class EICStorage(HiCacheStorage): ) -> int: if len(keys) == 0: return 0 - if self.use_zero_copy and not self.is_mla_model: + if self.mha_zero_copy: keys = self._get_mha_zero_copy_keys(keys) exist_mask = self._batch_exists_impl(keys) prefix_success = 0 @@ -419,7 +420,7 @@ class EICStorage(HiCacheStorage): prefix_success += 1 else: break - if not self.is_mla_model and self.use_zero_copy: + if self.mha_zero_copy: prefix_success = prefix_success // 2 return prefix_success @@ -713,7 +714,7 @@ class EICStorage(HiCacheStorage): ] ) - if self.use_zero_copy and not self.is_mla_model: + if self.mha_zero_copy: keys = self._get_mha_zero_copy_keys(keys) values = self._get_mha_zero_copy_values(values) @@ -760,7 +761,7 @@ class EICStorage(HiCacheStorage): for i in range(page_num) ] - if self.use_zero_copy and not self.is_mla_model: + if self.mha_zero_copy: keys = self._get_mha_zero_copy_keys(keys) values = self._get_mha_zero_copy_values(values) diff --git a/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py b/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py index abf07c680..84ab181eb 100644 --- a/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py +++ b/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py @@ -520,7 +520,7 @@ class HiCacheHF3FS(HiCacheStorage): self, keys: List[str], extra_info: Optional[HiCacheStorageExtraInfo] = None ) -> int: factor = 1 - if self.is_zero_copy and not self.is_mla_model: + if self.mha_zero_copy: keys = self._get_mha_zero_copy_keys(keys) factor = 2 @@ -571,6 +571,7 @@ class HiCacheHF3FS(HiCacheStorage): "page_first", "page_first_direct", ] + self.mha_zero_copy = self.is_zero_copy and not self.is_mla_model logger.info(f"{self.is_zero_copy=}, layout={self.mem_pool_host.layout}") @@ -646,7 +647,7 @@ class HiCacheHF3FS(HiCacheStorage): ] ) - if self.is_zero_copy and not self.is_mla_model: + if self.mha_zero_copy: keys = self._get_mha_zero_copy_keys(keys) values = self._get_mha_zero_copy_values(values) @@ -899,7 +900,7 @@ class HiCacheHF3FS(HiCacheStorage): for i in range(page_num) ] - if self.is_zero_copy and not self.is_mla_model: + if self.mha_zero_copy: keys = self._get_mha_zero_copy_keys(keys) values = self._get_mha_zero_copy_values(values) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index ef195d9f4..e12841165 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -542,8 +542,9 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): self.mla_suffix = "" self.storage_config = storage_config + self.should_split_heads = storage_config.should_split_heads self.split_factor = 0 - if self.storage_config.should_split_heads: + if self.should_split_heads: self.split_factor = ( self.storage_config.tp_lcm_size // self.storage_config.tp_size ) @@ -952,7 +953,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): if self.is_mla_backend: return self._get_mla_buffer_meta(keys, host_indices) else: - if self.storage_config.should_split_heads: + if self.should_split_heads: return self._get_mha_split_heads_buffer_meta(keys, host_indices) else: return self._get_mha_buffer_meta(keys, host_indices) @@ -972,7 +973,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): key_multiplier = 1 else: key_multiplier = 2 - if self.storage_config.should_split_heads: + if self.should_split_heads: key_multiplier *= self.split_factor result_groups = [ @@ -1209,7 +1210,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): key_multiplier = 1 else: query_keys = [] - if self.storage_config.should_split_heads: + if self.should_split_heads: for key in keys: for suffix in self.mha_suffix: query_keys.append(f"{key}_{suffix}_k")