[HiCache] Document transfer arguments (#40239)

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
@@ -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,