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

[HiCache] Rework the buffer-mode storage prefetch pipeline and retry bookkeeping (#39283)

Browse Source
This commit is contained in:
Zhiqiang Xie
2026-09-15 10:48:49 -07:00
committed by GitHub
parent 03ea13a545
commit 7f5dd19256
43 changed files with 3062 additions and 869 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-radix-tree/src/components/full.rs
+1
View File
@@ -391,6 +391,7 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
_host_indices: Option<Tensor>,
_token_ids: Option<&[i64]>,
_prefetch_tokens: usize,
_staging_tokens: usize,
_last_hash: Option<&str>,
) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
Ok(match phase {
rust/sglang-radix-tree/src/components/mamba.rs
+6 -3
View File
@@ -494,6 +494,7 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
host_indices: Option<Tensor>,
_token_ids: Option<&[i64]>,
_prefetch_tokens: usize,
staging_tokens: usize,
_last_hash: Option<&str>,
) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
Ok(match phase {
@@ -560,11 +561,13 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
}])
}
CacheTransferPhase::Prefetch => {
let host_indices =
host_indices.expect("Mamba PREFETCH build requires host indices");
if staging_tokens == 0 {
return Ok(None);
}
// Staging is allocated once the hit is known; the placeholder
// key carries the single trailing page this pool loads.
Some(vec![PoolTransfer {
name: PoolName::Mamba,
host_indices: Some(host_indices),
keys: Some(vec!["__placeholder__".to_string()]),
hit_policy: PoolHitPolicy::TrailingPages,
..Default::default()
rust/sglang-radix-tree/src/components/mod.rs
+2
View File
@@ -372,6 +372,7 @@ pub trait TreeComponent<K: ChildKeyType> {
host_indices: Option<Tensor>,
token_ids: Option<&[i64]>,
prefetch_tokens: usize,
staging_tokens: usize,
last_hash: Option<&str>,
) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
// Python reference — base.py::TreeComponent.build_hicache_transfers:
@@ -384,6 +385,7 @@ pub trait TreeComponent<K: ChildKeyType> {
// host_indices: Optional[torch.Tensor] = None,
// token_ids: Optional[Sequence[int]] = None,
// prefetch_tokens: int = 0,
// staging_tokens: int = 0,
// last_hash: Optional[str] = None,
// ) -> Optional[list[PoolTransfer]]:
// """Build transfer descriptors for this component in the given phase.
rust/sglang-radix-tree/src/components/swa.rs
+10 -5
View File
@@ -925,9 +925,10 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
node_id: NodeIdx_,
phase: CacheTransferPhase,
_mamba_pool_idx: Option<Tensor>,
host_indices: Option<Tensor>,
_host_indices: Option<Tensor>,
_token_ids: Option<&[i64]>,
_prefetch_tokens: usize,
staging_tokens: usize,
_last_hash: Option<&str>,
) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
// unified_kv keeps SWA as a device-only ring.
@@ -1024,12 +1025,16 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
}])
}
CacheTransferPhase::Prefetch => {
let host_indices = host_indices.expect("SWA PREFETCH build requires host indices");
let sw_pages = host_indices.numel() / tree_core.page_size;
// Staging is allocated once the hit is known; the placeholders
// carry the planned page count and the trailing hashes fill in
// at commit.
let num_pages = staging_tokens / tree_core.page_size;
if num_pages == 0 {
return Ok(None);
}
Some(vec![PoolTransfer {
name: PoolName::Swa,
host_indices: Some(host_indices),
keys: Some(vec!["__placeholder__".to_string(); sw_pages]),
keys: Some(vec!["__placeholder__".to_string(); num_pages]),
hit_policy: PoolHitPolicy::TrailingPages,
..Default::default()
}])
rust/sglang-radix-tree/src/python_bindings.rs
+49 -1
View File
@@ -999,6 +999,19 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
MatchResultBinding::from_match_result(py, result)
}
/// Read-only FULL-device match, independent of auxiliary components.
fn match_full_device_prefix(
&self,
py: Python<'_>,
params: &MatchParamsBinding,
) -> (usize, NodeId, usize) {
let key = K::key_from(Cow::Borrowed(&params.key));
let key = key.as_ref();
let namespace =
KeyNamespaceRef::new(params.extra_key.as_deref(), params.cache_salt.as_deref());
py.allow_threads(|| self.core().match_full_device_prefix(key, namespace))
}
/// The empty match result anchored at the root.
fn empty_match_result(&self, py: Python<'_>) -> PyResult<MatchResultBinding> {
let result = py.allow_threads(|| self.core().empty_match_result());
@@ -1113,6 +1126,18 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
Ok(IncLockRefResultBinding::from_result(result))
}
/// Pin only the FULL device values on a node's root path.
fn inc_full_pin(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> {
py.allow_threads(|| self.core().inc_full_pin(node_id))
.map_err(node_access_error)
}
/// Release a FULL-only root-path pin.
fn dec_full_pin(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> {
py.allow_threads(|| self.core().dec_full_pin(node_id))
.map_err(node_access_error)
}
/// Decrease the reference count on a node's component locks.
fn dec_lock_ref(
&self,
@@ -1486,6 +1511,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
host_indices: Option<PyTensor>,
token_ids: Option<Vec<i64>>,
prefetch_tokens: usize,
staging_tokens: usize,
last_hash: Option<String>,
) -> PyResult<Option<Vec<Py<PyAny>>>> {
let component_type = parse_component_type(component_type)?;
@@ -1500,6 +1526,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
host_indices,
token_ids.as_deref(),
prefetch_tokens,
staging_tokens,
last_hash.as_deref(),
)
})
@@ -2423,6 +2450,15 @@ macro_rules! tree_core_binding {
self.inner.match_prefix(py, params)
}
/// Read-only FULL-device match, independent of auxiliary components.
fn match_full_device_prefix(
&self,
py: Python<'_>,
params: &MatchParamsBinding,
) -> (usize, NodeId, usize) {
self.inner.match_full_device_prefix(py, params)
}
/// The empty match result anchored at the root.
fn empty_match_result(&self, py: Python<'_>) -> PyResult<MatchResultBinding> {
self.inner.empty_match_result(py)
@@ -2472,6 +2508,16 @@ macro_rules! tree_core_binding {
self.inner.inc_lock_ref(py, node_id, skip_lock_components)
}
/// Pin only the FULL device values on a node's root path.
fn inc_full_pin(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> {
self.inner.inc_full_pin(py, node_id)
}
/// Release a FULL-only root-path pin.
fn dec_full_pin(&self, py: Python<'_>, node_id: NodeId) -> PyResult<()> {
self.inner.dec_full_pin(py, node_id)
}
/// Decrease the reference count on a node's component locks. The
/// receipt is required: a release must replay its acquire's evidence.
#[pyo3(signature = (node_id, params, skip_swa = false))]
@@ -2718,7 +2764,7 @@ macro_rules! tree_core_binding {
/// Route a build_hicache_transfers call to the component for the given type.
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (component_type, node_id, phase, host_indices = None, token_ids = None, prefetch_tokens = 0, last_hash = None))]
#[pyo3(signature = (component_type, node_id, phase, host_indices = None, token_ids = None, prefetch_tokens = 0, staging_tokens = 0, last_hash = None))]
fn build_hicache_transfers(
&self,
py: Python<'_>,
@@ -2728,6 +2774,7 @@ macro_rules! tree_core_binding {
host_indices: Option<PyTensor>,
token_ids: Option<Vec<i64>>,
prefetch_tokens: usize,
staging_tokens: usize,
last_hash: Option<String>,
) -> PyResult<Option<Vec<Py<PyAny>>>> {
self.inner.build_hicache_transfers(
@@ -2738,6 +2785,7 @@ macro_rules! tree_core_binding {
host_indices,
token_ids,
prefetch_tokens,
staging_tokens,
last_hash,
)
}
rust/sglang-radix-tree/src/tests/components/full.rs
+4 -1
View File
@@ -2830,7 +2830,7 @@ fn build_hicache_transfers_returns_none_for_non_load_back_phases() {
.build_hicache_transfers(
&tc, a, phase, /* mamba_pool_idx = */ None, /* host_indices = */ None,
/* token_ids = */ None, /* prefetch_tokens = */ 0,
/* last_hash = */ None,
/* staging_tokens = */ 0, /* last_hash = */ None,
)
.unwrap();
assert!(transfers.is_none());
@@ -2850,6 +2850,7 @@ fn load_back_build_collects_the_evicted_suffix_ancestors_first() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
@@ -2884,6 +2885,7 @@ fn load_back_build_returns_an_empty_cpu_transfer_for_a_device_backed_node() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
@@ -2917,6 +2919,7 @@ fn load_back_build_panics_on_an_evicted_unbacked_node() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
);
}
rust/sglang-radix-tree/src/tests/components/mamba.rs
+19 -6
View File
@@ -1071,6 +1071,7 @@ fn backup_host_build_carries_the_device_slot() {
None,
None,
0,
0,
None,
)
.expect("live test node")
@@ -1102,6 +1103,7 @@ fn backup_host_build_carries_the_device_slot() {
None,
None,
0,
0,
None,
)
.expect("live test node")
@@ -1122,6 +1124,7 @@ fn load_back_build_restores_the_host_only_node() {
None,
None,
0,
0,
None,
)
.expect("live test node")
@@ -1151,6 +1154,7 @@ fn load_back_build_skips_device_backed_and_bare_nodes() {
None,
None,
0,
0,
None,
)
.expect("live test node")
@@ -1173,6 +1177,7 @@ fn load_back_build_adds_the_per_request_cow_transfer() {
None,
None,
0,
0,
None,
)
.unwrap()
@@ -1388,6 +1393,7 @@ fn backup_storage_build_keys_the_trailing_hash() {
None,
None,
0,
0,
None,
)
.expect("live test node")
@@ -1403,6 +1409,7 @@ fn backup_storage_build_keys_the_trailing_hash() {
None,
None,
0,
0,
None,
)
.expect("live test node")
@@ -1417,6 +1424,7 @@ fn backup_storage_build_keys_the_trailing_hash() {
None,
None,
0,
0,
None,
)
.expect("live test node")
@@ -1434,23 +1442,28 @@ fn backup_storage_build_keys_the_trailing_hash() {
}
#[test]
fn prefetch_build_wraps_the_host_buffer_with_a_placeholder_key() {
fn prefetch_build_carries_a_placeholder_key_for_the_planned_slot() {
let tc = mamba_core(/* page_size = */ 1);
let transfers = tc
.build_hicache_transfers(
let root_id = tc.arena.node(tc.arena.root()).id;
let build = |staging_tokens: usize| {
tc.build_hicache_transfers(
MAMBA,
tc.arena.node(tc.arena.root()).id,
root_id,
CacheTransferPhase::Prefetch,
Some(Tensor::from_slice(&[30i64])),
None,
None,
0,
staging_tokens,
None,
)
.expect("live test node")
.unwrap();
};
let transfers = build(1).unwrap();
assert_eq!(transfers.len(), 1);
assert_eq!(transfers[0].keys, Some(vec!["__placeholder__".to_string()]));
assert_eq!(transfers[0].hit_policy, PoolHitPolicy::TrailingPages);
assert!(transfers[0].host_indices.is_none());
assert!(build(0).is_none());
}
#[test]
rust/sglang-radix-tree/src/tests/components/swa.rs
+33 -22
View File
@@ -3595,6 +3595,7 @@ fn backup_storage_transfers_carry_trailing_page_keys() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
@@ -3625,6 +3626,7 @@ fn backup_storage_is_none_without_host_value_or_hashes() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
};
@@ -3646,7 +3648,7 @@ fn build_transfers_are_gated_off_until_the_swa_host_pool_is_wired() {
.build_hicache_transfers(
&tc, a, phase, /* mamba_pool_idx = */ None, /* host_indices = */ None,
/* token_ids = */ None, /* prefetch_tokens = */ 0,
/* last_hash = */ None,
/* staging_tokens = */ 0, /* last_hash = */ None,
)
.unwrap();
assert!(transfers.is_none());
@@ -3662,6 +3664,7 @@ fn build_transfers_are_gated_off_until_the_swa_host_pool_is_wired() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap();
@@ -3684,6 +3687,7 @@ fn backup_host_build_wraps_the_device_value_as_int64() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
@@ -3715,6 +3719,7 @@ fn backup_host_build_returns_none_for_a_tombstone() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap();
@@ -3765,6 +3770,7 @@ fn load_back_build_collects_host_only_nodes_within_the_window() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
@@ -3805,6 +3811,7 @@ fn load_back_build_stops_at_the_window_boundary() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
@@ -3838,6 +3845,7 @@ fn load_back_build_returns_none_when_the_window_is_on_device() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap();
@@ -3857,6 +3865,7 @@ fn load_back_build_rejects_a_bare_window_node() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
),
Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { node_id })
@@ -3880,6 +3889,7 @@ fn fallible_load_back_boundaries_reject_a_bare_window_node() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
),
Err(TreeCoreRuntimeError::SwaLoadBackMissingValue { node_id: missing })
@@ -4112,21 +4122,26 @@ fn commit_hicache_transfers_routes_to_the_component() {
}
#[test]
fn prefetch_build_wraps_the_host_buffer_with_placeholder_keys() {
fn prefetch_build_sizes_the_placeholder_keys_from_the_staging_tokens() {
let tc = swa_core(/* window = */ 4, /* page_size = */ 1);
let transfers = swa_component(4)
.build_hicache_transfers(
&tc,
tc.arena.root(),
CacheTransferPhase::Prefetch,
/* mamba_pool_idx = */ None,
/* host_indices = */ Some(Tensor::from_slice(&[30i64, 31])),
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
.unwrap();
let build = |staging_tokens: usize| {
swa_component(4)
.build_hicache_transfers(
&tc,
tc.arena.root(),
CacheTransferPhase::Prefetch,
/* mamba_pool_idx = */ None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
staging_tokens,
/* last_hash = */ None,
)
.unwrap()
};
// Staging is allocated once the hit is known: the build carries only the
// planned page count, never a host buffer.
let transfers = build(2).unwrap();
assert_eq!(transfers.len(), 1);
assert_eq!(
transfers[0].keys,
@@ -4136,13 +4151,8 @@ fn prefetch_build_wraps_the_host_buffer_with_placeholder_keys() {
])
);
assert_eq!(transfers[0].hit_policy, PoolHitPolicy::TrailingPages);
assert!(
transfers[0]
.host_indices
.as_ref()
.unwrap()
.equal(&Tensor::from_slice(&[30i64, 31]))
);
assert!(transfers[0].host_indices.is_none());
assert!(build(0).is_none());
}
#[test]
@@ -5192,6 +5202,7 @@ fn backup_transfers(
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
rust/sglang-radix-tree/src/tests/unified_tree_core.rs
+20
View File
@@ -1295,6 +1295,25 @@ fn match_prefix_splits_on_a_partial_match() {
);
}
#[test]
fn match_full_device_prefix_is_read_only_and_accounts_the_pinned_node() {
let mut tc = core();
let (a, _b) = matched_chain(&mut tc);
let (matched_len, node_id, pinned_len) =
tc.match_full_device_prefix(&vec![1, 9], KeyNamespaceRef::new(None, None));
assert_eq!(matched_len, 1);
assert_eq!(node_id, tc.arena.node(a).id);
assert_eq!(pinned_len, 2);
assert_eq!(tc.arena.node(a).key, vec![1, 2]);
tc.inc_full_pin(node_id).unwrap();
assert_eq!(tc.arena.node(a).device_lock_ref(FULL), 1);
tc.dec_full_pin(node_id).unwrap();
assert_eq!(tc.arena.node(a).device_lock_ref(FULL), 0);
}
#[test]
fn match_prefix_stops_at_a_dead_node() {
// An evicted, unbackuped child ends the traversal before it.
@@ -4346,6 +4365,7 @@ fn fallible_node_boundaries_reject_stale_handles() {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
),
Err(TreeCoreRuntimeError::NodeAccess(NodeAccessError { node_id }))
rust/sglang-radix-tree/src/unified_tree_core.rs
+72
View File
@@ -850,6 +850,34 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
Ok(result)
}
/// Pin only the FULL device values on a node's root path.
pub fn inc_full_pin(&mut self, node_id: NodeId) -> Result<(), NodeAccessError> {
let node_idx = self.arena.resolve(node_id)?;
let full = self.component_by_type_(FULL);
full.acquire_component_lock(
self,
node_idx,
IncLockRefResult::default(),
/* lock_host = */ false,
);
self.update_evictable_leaf_sets_(node_idx);
Ok(())
}
/// Release a FULL-only root-path pin.
pub fn dec_full_pin(&mut self, node_id: NodeId) -> Result<(), NodeAccessError> {
let node_idx = self.arena.resolve(node_id)?;
let full = self.component_by_type_(FULL);
full.release_component_lock(
self,
node_idx,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
self.update_evictable_leaf_sets_(node_idx);
Ok(())
}
/// A receipt releases only the node its acquire returned; a mispaired
/// node would silently release (or steal) another holder's segment.
fn assert_receipt_anchor_(&self, node_idx: NodeIdx_, params: &DecLockRefParams) {
@@ -1035,6 +1063,44 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
)
}
/// Read-only FULL-device match, independent of auxiliary components.
/// Returns the request match and the complete root-path length pinned by
/// the deepest node; they differ when the key ends inside that node.
pub fn match_full_device_prefix(
&self,
key: &K,
namespace: KeyNamespaceRef<'_>,
) -> (usize, NodeId, usize) {
let aligned_key_len = key.atom_len() / self.page_size * self.page_size;
let mut node_id = self.arena.root();
let mut offset = 0;
let mut pinned_len = 0;
while offset < aligned_key_len {
let Some(child_id) = self.arena.child_on_page_in_namespace(
node_id,
namespace,
key.page_at(offset, self.page_size),
) else {
break;
};
let child = self.arena.node(child_id);
if !child.has_device_value(FULL) {
break;
}
let prefix_len = key.match_len(offset, &child.key, self.page_size);
if prefix_len == 0 {
break;
}
offset += prefix_len;
pinned_len += child.device_value_len(FULL);
node_id = child_id;
if prefix_len < child.key.atom_len() {
break;
}
}
(offset, self.arena.node(node_id).id, pinned_len)
}
/// Walk the tree for `key`; returns matched value chunks, the best match,
/// the best device-resident match, its device value length, and any split action.
pub fn match_prefix_helper_(
@@ -3178,6 +3244,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap();
@@ -3217,6 +3284,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap();
@@ -3244,6 +3312,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
host_indices: Option<Tensor>,
token_ids: Option<&[i64]>,
prefetch_tokens: usize,
staging_tokens: usize,
last_hash: Option<&str>,
) -> Result<Option<Vec<PoolTransfer>>, TreeCoreRuntimeError> {
let node_id = self.arena.resolve(node_id)?;
@@ -3256,6 +3325,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
host_indices,
token_ids,
prefetch_tokens,
staging_tokens,
last_hash,
)
}
@@ -3281,6 +3351,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)?
.unwrap();
@@ -3299,6 +3370,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* staging_tokens = */ 0,
/* last_hash = */ None,
)?;
if let Some(transfers) = transfers