minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[HiCache] Document transfer arguments (#40239)

Browse Source 
Clarify the legacy host_indices argument and label Mamba test arguments,
including the current staging_tokens parameter. Executable code is unchanged.

Co-authored-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
This commit is contained in:
Lianmin Zheng
2026-09-18 14:08:19 -07:00
committed by GitHub
co-authored by Jialin Ouyang
parent cd4dd81c22
commit f5a1434700
2 changed files with 48 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py
+3
View File
@@ -510,6 +510,9 @@ class UnifiedTreeCoreInterface(ABC):
node_id: NodeId,
phase: CacheTransferPhase,
*,
# TODO(Jialin): Remove the legacy `host_indices` argument from the
# interface and bindings. Prefetch uses `staging_tokens`;
# `PoolTransfer.host_indices` is populated after the storage hit.
host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0,
rust/sglang-radix-tree/src/tests/components/mamba.rs
+45 -45
View File
@@ -1068,11 +1068,11 @@ fn backup_host_build_carries_the_device_slot() {
MAMBA,
tc.arena.node(a).id,
CacheTransferPhase::BackupHost,
None,
None,
0,
0,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.expect("live test node")
.unwrap();
@@ -1100,11 +1100,11 @@ fn backup_host_build_carries_the_device_slot() {
MAMBA,
tc.arena.node(b).id,
CacheTransferPhase::BackupHost,
None,
None,
0,
0,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.expect("live test node")
.is_none()
@@ -1121,11 +1121,11 @@ fn load_back_build_restores_the_host_only_node() {
MAMBA,
tc.arena.node(a).id,
CacheTransferPhase::LoadBack,
None,
None,
0,
0,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.expect("live test node")
.unwrap();
@@ -1151,11 +1151,11 @@ fn load_back_build_skips_device_backed_and_bare_nodes() {
MAMBA,
tc.arena.node(node).id,
CacheTransferPhase::LoadBack,
None,
None,
0,
0,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.expect("live test node")
.is_none()
@@ -1174,11 +1174,11 @@ fn load_back_build_adds_the_per_request_cow_transfer() {
a,
CacheTransferPhase::LoadBack,
/* mamba_pool_idx = */ Some(Tensor::from_slice(&[3i64]).squeeze()),
None,
None,
0,
0,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
.unwrap();
@@ -1390,11 +1390,11 @@ fn backup_storage_build_keys_the_trailing_hash() {
MAMBA,
tc.arena.node(a).id,
CacheTransferPhase::BackupStorage,
None,
None,
0,
0,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.expect("live test node")
.is_none()
@@ -1406,11 +1406,11 @@ fn backup_storage_build_keys_the_trailing_hash() {
MAMBA,
tc.arena.node(a).id,
CacheTransferPhase::BackupStorage,
None,
None,
0,
0,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.expect("live test node")
.is_none()
@@ -1421,11 +1421,11 @@ fn backup_storage_build_keys_the_trailing_hash() {
MAMBA,
tc.arena.node(a).id,
CacheTransferPhase::BackupStorage,
None,
None,
0,
0,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.expect("live test node")
.unwrap();
@@ -1450,11 +1450,11 @@ fn prefetch_build_carries_a_placeholder_key_for_the_planned_slot() {
MAMBA,
root_id,
CacheTransferPhase::Prefetch,
None,
None,
0,
staging_tokens,
None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ staging_tokens,
/* last_hash = */ None,
)
.expect("live test node")
};