[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:
co-authored by
Jialin Ouyang
parent
cd4dd81c22
commit
f5a1434700
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user