minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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, node_id: NodeId,
phase: CacheTransferPhase, 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, host_indices: Optional[torch.Tensor] = None,
token_ids: Optional[Sequence[int]] = None, token_ids: Optional[Sequence[int]] = None,
prefetch_tokens: int = 0, 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, MAMBA,
tc.arena.node(a).id, tc.arena.node(a).id,
CacheTransferPhase::BackupHost, CacheTransferPhase::BackupHost,
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
0, /* staging_tokens = */ 0,
None, /* last_hash = */ None,
) )
.expect("live test node") .expect("live test node")
.unwrap(); .unwrap();
@@ -1100,11 +1100,11 @@ fn backup_host_build_carries_the_device_slot() {
MAMBA, MAMBA,
tc.arena.node(b).id, tc.arena.node(b).id,
CacheTransferPhase::BackupHost, CacheTransferPhase::BackupHost,
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
0, /* staging_tokens = */ 0,
None, /* last_hash = */ None,
) )
.expect("live test node") .expect("live test node")
.is_none() .is_none()
@@ -1121,11 +1121,11 @@ fn load_back_build_restores_the_host_only_node() {
MAMBA, MAMBA,
tc.arena.node(a).id, tc.arena.node(a).id,
CacheTransferPhase::LoadBack, CacheTransferPhase::LoadBack,
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
0, /* staging_tokens = */ 0,
None, /* last_hash = */ None,
) )
.expect("live test node") .expect("live test node")
.unwrap(); .unwrap();
@@ -1151,11 +1151,11 @@ fn load_back_build_skips_device_backed_and_bare_nodes() {
MAMBA, MAMBA,
tc.arena.node(node).id, tc.arena.node(node).id,
CacheTransferPhase::LoadBack, CacheTransferPhase::LoadBack,
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
0, /* staging_tokens = */ 0,
None, /* last_hash = */ None,
) )
.expect("live test node") .expect("live test node")
.is_none() .is_none()
@@ -1174,11 +1174,11 @@ fn load_back_build_adds_the_per_request_cow_transfer() {
a, a,
CacheTransferPhase::LoadBack, CacheTransferPhase::LoadBack,
/* mamba_pool_idx = */ Some(Tensor::from_slice(&[3i64]).squeeze()), /* mamba_pool_idx = */ Some(Tensor::from_slice(&[3i64]).squeeze()),
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
0, /* staging_tokens = */ 0,
None, /* last_hash = */ None,
) )
.unwrap() .unwrap()
.unwrap(); .unwrap();
@@ -1390,11 +1390,11 @@ fn backup_storage_build_keys_the_trailing_hash() {
MAMBA, MAMBA,
tc.arena.node(a).id, tc.arena.node(a).id,
CacheTransferPhase::BackupStorage, CacheTransferPhase::BackupStorage,
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
0, /* staging_tokens = */ 0,
None, /* last_hash = */ None,
) )
.expect("live test node") .expect("live test node")
.is_none() .is_none()
@@ -1406,11 +1406,11 @@ fn backup_storage_build_keys_the_trailing_hash() {
MAMBA, MAMBA,
tc.arena.node(a).id, tc.arena.node(a).id,
CacheTransferPhase::BackupStorage, CacheTransferPhase::BackupStorage,
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
0, /* staging_tokens = */ 0,
None, /* last_hash = */ None,
) )
.expect("live test node") .expect("live test node")
.is_none() .is_none()
@@ -1421,11 +1421,11 @@ fn backup_storage_build_keys_the_trailing_hash() {
MAMBA, MAMBA,
tc.arena.node(a).id, tc.arena.node(a).id,
CacheTransferPhase::BackupStorage, CacheTransferPhase::BackupStorage,
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
0, /* staging_tokens = */ 0,
None, /* last_hash = */ None,
) )
.expect("live test node") .expect("live test node")
.unwrap(); .unwrap();
@@ -1450,11 +1450,11 @@ fn prefetch_build_carries_a_placeholder_key_for_the_planned_slot() {
MAMBA, MAMBA,
root_id, root_id,
CacheTransferPhase::Prefetch, CacheTransferPhase::Prefetch,
None, /* host_indices = */ None,
None, /* token_ids = */ None,
0, /* prefetch_tokens = */ 0,
staging_tokens, /* staging_tokens = */ staging_tokens,
None, /* last_hash = */ None,
) )
.expect("live test node") .expect("live test node")
}; };