Remove the HiMambaRadixTree that is no longer in use (#33468)
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -34,7 +34,6 @@ if TYPE_CHECKING:
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
||||||
from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache
|
|
||||||
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
|
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
|
||||||
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
|
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
|
||||||
from sglang.srt.server_args import ServerArgs
|
from sglang.srt.server_args import ServerArgs
|
||||||
@@ -1883,58 +1882,3 @@ def attach_hybrid_dsa_pool_to_hiradix_cache(
|
|||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("attach_hybrid_dsa_pool_to_hiradix_cache failed")
|
logger.exception("attach_hybrid_dsa_pool_to_hiradix_cache failed")
|
||||||
raise
|
raise
|
||||||
|
|
||||||
|
|
||||||
def attach_hybrid_pool_to_mamba_cache(
|
|
||||||
mamba_cache: HiMambaRadixCache,
|
|
||||||
params: CacheInitParams,
|
|
||||||
server_args: ServerArgs,
|
|
||||||
*,
|
|
||||||
extra_config: dict,
|
|
||||||
prefetch_threshold: int,
|
|
||||||
load_cache_event,
|
|
||||||
enable_storage_metrics: bool = False,
|
|
||||||
) -> None:
|
|
||||||
"""Attach HostPoolGroup (KV + Mamba) + HybridCacheController for HiMambaRadixCache.
|
|
||||||
|
|
||||||
This entrypoint is currently intended only for HiMambaRadixCache.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
hybrid_kv = mamba_cache.hybrid_kv_cache
|
|
||||||
kvcache = mamba_cache.kvcache
|
|
||||||
full_layer_mapping = dict(hybrid_kv.full_attention_layer_id_mapping)
|
|
||||||
mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
|
|
||||||
host_pool_group, cache_controller = build_hybrid_mamba_stack(
|
|
||||||
params=params,
|
|
||||||
server_args=server_args,
|
|
||||||
kv_pool=kvcache,
|
|
||||||
mamba_pool=params.req_to_token_pool.mamba_pool,
|
|
||||||
full_layer_mapping=full_layer_mapping,
|
|
||||||
mamba_layer_mapping=mamba_layer_mapping,
|
|
||||||
load_cache_event=load_cache_event,
|
|
||||||
storage_backend=server_args.hicache_storage_backend,
|
|
||||||
use_mla=hybrid_kv.use_mla,
|
|
||||||
host_mamba_evict_fn=mamba_cache.evict_mamba_host,
|
|
||||||
device_mamba_evict_fn=mamba_cache.evict_mamba,
|
|
||||||
prefetch_threshold=prefetch_threshold,
|
|
||||||
model_name=server_args.served_model_name,
|
|
||||||
storage_backend_extra_config=extra_config,
|
|
||||||
enable_storage_metrics=enable_storage_metrics,
|
|
||||||
)
|
|
||||||
mamba_cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV)
|
|
||||||
mamba_cache.mamba_pool_host = host_pool_group.get_pool(PoolName.MAMBA)
|
|
||||||
mamba_cache.transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping)
|
|
||||||
mamba_cache.host_pool_group = host_pool_group
|
|
||||||
mamba_cache.cache_controller = cache_controller
|
|
||||||
params.req_to_token_pool.register_layer_transfer_counter(
|
|
||||||
cache_controller.layer_done_counter
|
|
||||||
)
|
|
||||||
hybrid_kv.register_layer_transfer_counter(cache_controller.layer_done_counter)
|
|
||||||
logger.info(
|
|
||||||
"Attached hybrid Mamba pool stack to HiMambaRadixCache: pools=KV + MAMBA, "
|
|
||||||
"transfer_layer_num=%s",
|
|
||||||
mamba_cache.transfer_layer_num,
|
|
||||||
)
|
|
||||||
except Exception:
|
|
||||||
logger.exception("attach_hybrid_pool_to_mamba_cache failed")
|
|
||||||
raise
|
|
||||||
|
|||||||
@@ -36,7 +36,7 @@ The scale axis (reduces over d_v) matches the per-k-channel decay diag(alpha), s
|
|||||||
the large state entries keep ~bf16 precision and the error concentrates on small
|
the large state entries keep ~bf16 precision and the error concentrates on small
|
||||||
entries that barely affect the readout. Storing cached states int8 gives ~2x the
|
entries that barely affect the readout. Storing cached states int8 gives ~2x the
|
||||||
cached-prefix capacity at fixed memory, and composes with host-offload
|
cached-prefix capacity at fixed memory, and composes with host-offload
|
||||||
(HiMambaRadixCache) which it also halves.
|
(--enable-hierarchical-cache) which it also halves.
|
||||||
|
|
||||||
This is strategy-agnostic: whether the active slot to be cached was produced by
|
This is strategy-agnostic: whether the active slot to be cached was produced by
|
||||||
the ``no_buffer`` donate (copy_from) or the ``extra_buffer`` ping-pong track
|
the ``no_buffer`` donate (copy_from) or the ``extra_buffer`` ping-pong track
|
||||||
|
|||||||
@@ -242,7 +242,7 @@ export SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR=/tmp/sglang_nixl_e2e_storage
|
|||||||
|
|
||||||
Expected behavior for this validated setup:
|
Expected behavior for this validated setup:
|
||||||
|
|
||||||
- the server starts with `Attached hybrid Mamba pool stack to HiMambaRadixCache: pools=KV + MAMBA`
|
- the server starts with `Attached hybrid pool stack to UnifiedRadixCache: pools=KV + MAMBA`
|
||||||
- NIXL logs show `Backend POSIX was instantiated`
|
- NIXL logs show `Backend POSIX was instantiated`
|
||||||
- the server logs `HiCacheNixl: registered hybrid host pool mamba zero_copy=...`
|
- the server logs `HiCacheNixl: registered hybrid host pool mamba zero_copy=...`
|
||||||
- the storage directory contains KV files plus Mamba sidecar files such as `..._0_2_mamba_temporal` and `..._0_2_mamba_conv_0`
|
- the storage directory contains KV files plus Mamba sidecar files such as `..._0_2_mamba_temporal` and `..._0_2_mamba_conv_0`
|
||||||
|
|||||||
@@ -6009,7 +6009,7 @@ class ServerArgs:
|
|||||||
|
|
||||||
def _handle_int8_mamba_checkpoint(self):
|
def _handle_int8_mamba_checkpoint(self):
|
||||||
# The int8 mamba checkpoint pool is only wired into the built-in
|
# The int8 mamba checkpoint pool is only wired into the built-in
|
||||||
# MambaRadixCache. The host-offload variant (HiMambaRadixCache, enabled by
|
# MambaRadixCache. The host-offload path (enabled by
|
||||||
# --enable-hierarchical-cache) and custom radix-cache backends are NOT
|
# --enable-hierarchical-cache) and custom radix-cache backends are NOT
|
||||||
# int8-aware: they would read int8 checkpoint slots as bf16 active slots
|
# int8-aware: they would read int8 checkpoint slots as bf16 active slots
|
||||||
# (wrong pool / out-of-range). Reject the combination up front rather than
|
# (wrong pool / out-of-range). Reject the combination up front rather than
|
||||||
@@ -6020,7 +6020,7 @@ class ServerArgs:
|
|||||||
raise ValueError(
|
raise ValueError(
|
||||||
"--enable-int8-mamba-checkpoint is not supported together with "
|
"--enable-int8-mamba-checkpoint is not supported together with "
|
||||||
"--enable-hierarchical-cache: the host-offload path "
|
"--enable-hierarchical-cache: the host-offload path "
|
||||||
"(HiMambaRadixCache) is not int8-aware. Disable one of them."
|
"is not int8-aware. Disable one of them."
|
||||||
)
|
)
|
||||||
if self.radix_cache_backend is not None:
|
if self.radix_cache_backend is not None:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
|
|||||||
@@ -16,8 +16,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
||||||
from sglang.srt.mem_cache.common import available_and_evictable_str
|
from sglang.srt.mem_cache.common import available_and_evictable_str
|
||||||
from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache
|
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
|
||||||
from sglang.srt.mem_cache.mamba_radix_cache import LRUList, MambaRadixCache, TreeNode
|
|
||||||
from sglang.srt.mem_cache.memory_pool import (
|
from sglang.srt.mem_cache.memory_pool import (
|
||||||
HybridLinearKVPool,
|
HybridLinearKVPool,
|
||||||
HybridReqToTokenPool,
|
HybridReqToTokenPool,
|
||||||
@@ -628,53 +627,6 @@ class TestMamba(unittest.TestCase):
|
|||||||
|
|
||||||
return tree, allocator, req_to_token_pool, make_dummy_req
|
return tree, allocator, req_to_token_pool, make_dummy_req
|
||||||
|
|
||||||
def test_hi_mamba_tombstone_cleanup_respects_host_ref(self):
|
|
||||||
tree = object.__new__(HiMambaRadixCache)
|
|
||||||
root = TreeNode()
|
|
||||||
parent = TreeNode()
|
|
||||||
deleted = TreeNode()
|
|
||||||
|
|
||||||
root.key = RadixKey(array("q", []))
|
|
||||||
parent.key = RadixKey(array("q", [1]))
|
|
||||||
deleted.key = RadixKey(array("q", [2]))
|
|
||||||
parent.parent = root
|
|
||||||
deleted.parent = parent
|
|
||||||
parent.value = torch.tensor([1], dtype=torch.int64)
|
|
||||||
parent.protect_host()
|
|
||||||
root.children[parent.key.child_key(1)] = parent
|
|
||||||
|
|
||||||
class RecordingCacheController:
|
|
||||||
def __init__(self):
|
|
||||||
self.device_evictions = []
|
|
||||||
self.host_evictions = []
|
|
||||||
|
|
||||||
def evict_device(self, value):
|
|
||||||
self.device_evictions.append(value)
|
|
||||||
|
|
||||||
def evict_host(self, value):
|
|
||||||
self.host_evictions.append(value)
|
|
||||||
|
|
||||||
tree.root_node = root
|
|
||||||
tree.page_size = 1
|
|
||||||
tree.full_lru_list = LRUList(mamba=False)
|
|
||||||
tree.full_lru_list.insert_mru(parent)
|
|
||||||
tree.cache_controller = RecordingCacheController()
|
|
||||||
tree.full_evictable_size_ = len(parent.value)
|
|
||||||
tree.evictable_full_device_leaves = {parent}
|
|
||||||
tree.evictable_full_host_leaves = set()
|
|
||||||
|
|
||||||
result_node, full_evicted, mamba_evicted = (
|
|
||||||
tree._iteratively_delete_tombstone_leaf(deleted)
|
|
||||||
)
|
|
||||||
|
|
||||||
self.assertIs(result_node, deleted)
|
|
||||||
self.assertEqual(full_evicted, 0)
|
|
||||||
self.assertEqual(mamba_evicted, 0)
|
|
||||||
self.assertIs(root.children[parent.key.child_key(1)], parent)
|
|
||||||
self.assertTrue(tree.full_lru_list.in_list(parent))
|
|
||||||
self.assertEqual(tree.cache_controller.device_evictions, [])
|
|
||||||
self.assertEqual(tree.cache_controller.host_evictions, [])
|
|
||||||
|
|
||||||
def test_mamba_pool_cpu_offload(self):
|
def test_mamba_pool_cpu_offload(self):
|
||||||
"""MambaPool.get_cpu_copy / load_cpu_copy round-trips conv and temporal state."""
|
"""MambaPool.get_cpu_copy / load_cpu_copy round-trips conv and temporal state."""
|
||||||
_, _, req_to_token_pool, _ = self._setup_tree_and_allocator()
|
_, _, req_to_token_pool, _ = self._setup_tree_and_allocator()
|
||||||
|
|||||||
Reference in New Issue
Block a user