Core 加速与内存
跨步 cache、token pruning、融合 kernel、显存生命周期 wrapper、DiskMap 与 memory store。
Core 把精确实现加速与近似策略分开。融合 kernel 和 placement 调整的目标是保持模型语义；跨步 cache 和 token pruning 则明确用近似换计算量，因此会暴露 threshold、保留比例、event 和 reset 边界，让模型接入能够实际评测。
跨步 cache 示例
import torch
from worldfoundry.core.acceleration import FixedStepCache
cache = FixedStepCache(
skip_steps={1, 3},
dense_first=1,
dense_last=1,
total_steps=5,
)
calls = []
with torch.no_grad():
outputs = []
for step in range(5):
def compute(step=step):
calls.append(step)
return torch.tensor([float(step)])
outputs.append(cache.run(step, compute))
assert calls == [0, 2, 4] # 边界 step 会保持完整计算
assert [event.hit for event in cache.events] == [False, True, False, True, False]Autograd 开启时 cache 会关闭 replay。独立 denoising trajectory 之间必须调用 reset()，否则上一请求的 residual 会成为下一请求的状态。评测延迟与质量时，应把 event stream 当作实际证据，而不是假设配置的 skip step 都命中了。
Token pruning 生命周期
select_token_indices 是无状态函数；prune_tokens 返回紧凑数据和 TokenPruneState；restore_tokens 再把处理后的 token scatter 回去，用 compensation 或零填充被丢弃位置。TokenPruner 增加上一时刻 compensation：第一次调用记录 dense segment，后续调用才可以 prune；每个 prune 都必须在相同 key 下与 restore 配对。
显存 placement 状态
enable_vram_management 会把 module_map 列出的类替换成 AutoWrappedLinear 等 wrapper。每个 wrapper 可以分别设置 offload、onload、preparing 和 computation 的 dtype/device。磁盘 wrapper 通过 DiskMap 解析参数名，普通 wrapper 则在 device 之间移动或复制 tensor。
这个变换会改变 module identity，通常只应在模型构造时执行一次。存在继承重叠时，module-map 顺序会影响结果。一套 placement 配置不会自动适合新架构：需要验证峰值显存、传输重叠、输出一致性或容差，以及达到配置显存上限时的行为。
语义 Memory 与 VRAM 的区别
BaseMemory 和 MemoryStore 表示可以检索的世界模型 memory record，并不管理 GPU 分配。VRAM wrapper 管理参数放置，但不提供语义检索。把两层含义分开，才能避免把“memory”配置放到错误的系统层。
完整参考
以下为该类别的生成签名。可用本页符号索引跳转；源码链接指向各惰性导出背后的具体实现。
26 个公开符号
class AdaptiveResidualCache(threshold: float,warmup_steps: int = 1,max_consecutive_hits: int = 3,dense_last: int = 1,total_steps: int | None = None,subsample: int = 1,eps: float = 1e-06)简介
AdaptiveResidualCache — Reuse a model residual while accumulated input change stays small. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
源码 docstring
Reuse a model residual while accumulated input change stays small.
Warm up densely, estimate normalized change from a cheap signal, accumulate it across steps, and replay the last dense residual until the threshold or consecutive-hit cap is reached.
参数
thresholdfloatwarmup_stepsint默认值 1max_consecutive_hitsint默认值 3dense_lastint默认值 1total_stepsint | None默认值 Nonesubsampleint默认值 1epsfloat默认值 1e-06方法
reset — Clear residual and accumulator so a new clip cannot reuse the last one.
返回值
run(step: int,signal: torch.Tensor,compute_residual: Callable[[], torch.Tensor],total_steps: int | None = None) -> torch.Tensor源码run — Return a dense or cached residual for one denoising step.
参数
stepintsignaltorch.Tensorcompute_residualCallable[[], torch.Tensor]total_stepsint | None默认值 None返回值
class AutoTorchModule(offload_dtype: torch.dtype = None,offload_device: Union[str, torch.device] = None,onload_dtype: torch.dtype = None,onload_device: Union[str, torch.device] = None,preparing_dtype: torch.dtype = None,preparing_device: Union[str, torch.device] = None,computation_dtype: torch.dtype = None,computation_device: Union[str, torch.device] = None,vram_limit: float = None)简介
AutoTorchModule — Base state machine for modules that move between memory tiers. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
源码 docstring
Base state machine for modules that move between memory tiers.
`state is 0 while offloaded, 1 while onloaded, and 2` while kept on the computation device. Subclasses decide whether transitions copy tensors, materialize them from disk, or use temporary computation weights.
参数
offload_dtypetorch.dtype默认值 Noneoffload_deviceUnion[str, torch.device]默认值 None"disk"`.onload_dtypetorch.dtype默认值 Noneonload_deviceUnion[str, torch.device]默认值 Nonepreparing_dtypetorch.dtype默认值 Nonepreparing_deviceUnion[str, torch.device]默认值 Nonecomputation_dtypetorch.dtype默认值 Nonecomputation_deviceUnion[str, torch.device]默认值 Nonevram_limitfloat默认值 None方法
set_dtype_and_device(offload_dtype: torch.dtype = None,offload_device: Union[str, torch.device] = None,onload_dtype: torch.dtype = None,onload_device: Union[str, torch.device] = None,preparing_dtype: torch.dtype = None,preparing_device: Union[str, torch.device] = None,computation_dtype: torch.dtype = None,computation_device: Union[str, torch.device] = None,vram_limit: float = None)源码set_dtype_and_device — Update lifecycle placement, defaulting omitted stages to computation placement.
参数
offload_dtypetorch.dtype默认值 Noneoffload_deviceUnion[str, torch.device]默认值 Noneonload_dtypetorch.dtype默认值 Noneonload_deviceUnion[str, torch.device]默认值 Nonepreparing_dtypetorch.dtype默认值 Nonepreparing_deviceUnion[str, torch.device]默认值 Nonecomputation_dtypetorch.dtype默认值 Nonecomputation_deviceUnion[str, torch.device]默认值 Nonevram_limitfloat默认值 Nonecast_to — Copy one tensor to `dtype and device` without mutating the source.
参数
weightdtypedevicecheck_free_vram — Return whether current accelerator usage is below `vram_limit`.
offload — Move managed parameters to the inactive placement and set state 0.
onload — Move managed parameters to the first prefetch placement and set state 1.
keep — Keep managed parameters at computation placement and set state 2.
param_name — Return the fully qualified checkpoint key for a local parameter name.
参数
nameclass AutoWrappedLinear(module: torch.nn.Linear,offload_dtype: torch.dtype = None,offload_device: Union[str, torch.device] = None,onload_dtype: torch.dtype = None,onload_device: Union[str, torch.device] = None,preparing_dtype: torch.dtype = None,preparing_device: Union[str, torch.device] = None,computation_dtype: torch.dtype = None,computation_device: Union[str, torch.device] = None,vram_limit: float = None,name: str = '',disk_map: DiskMap = None,**kwargs)AutoWrappedLinear — Linear-layer wrapper with staged placement, disk loading, FP8, and LoRA support. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
moduletorch.nn.Lineartorch.nn.Linear` whose parameters are reused.offload_dtypetorch.dtype默认值 None"disk"`.offload_deviceUnion[str, torch.device]默认值 None"disk"`.onload_dtypetorch.dtype默认值 Noneonload_deviceUnion[str, torch.device]默认值 Nonepreparing_dtypetorch.dtype默认值 Nonepreparing_deviceUnion[str, torch.device]默认值 Nonecomputation_dtypetorch.dtype默认值 Nonecomputation_deviceUnion[str, torch.device]默认值 Nonevram_limitfloat默认值 Nonenamestr默认值 ''kwargs方法
fp8_linear(input: torch.Tensor,weight: torch.Tensor,bias: torch.Tensor = None) -> torch.Tensor源码简介
fp8_linear — Row-wise scaled FP8 matmul via `torch._scaled_mm`.
源码 docstring
Row-wise scaled FP8 matmul via `torch._scaled_mm`.
Scales are the exact values passed to the kernel so dequantization matches the quantization (CC-28). `e4m3fnuz halves the representable max; that factor is baked into fp8_max` and recovered by the kernel scales, not a post-hoc multiply.
参数
inputtorch.Tensorweighttorch.Tensorbiastorch.Tensor默认值 None返回值
简介
load_from_disk — Fetch `name.weight / name.bias` from the disk map.
源码 docstring
Fetch `name.weight / name.bias` from the disk map.
`assign=False` returns tensors without mutating this module so a temporary compute copy can leave the stored weights on meta.
参数
torch_dtypedeviceassign默认值 Trueoffload — Drop to meta on disk offload; otherwise move this Linear to offload placement.
onload — Prefetch this Linear from disk or move from offload; skip a still-disk target.
preparing — Second-stage prefetch for this Linear; materialize only when leaving disk.
简介
computation — Return weight/bias at compute placement without permanently leaving the current state.
源码 docstring
Return weight/bias at compute placement without permanently leaving the current state.
`assign=False` on the disk path so the stored Linear stays on meta.
linear_forward — Dense GEMM, or the FP8 scaled path when `computation_dtype` is float8.
参数
xweightbias简介
lora_forward — Add LoRA deltas after the base linear so merge order matches PEFT (B @ A).
源码 docstring
Add LoRA deltas after the base linear so merge order matches PEFT (B @ A).
A custom `lora_merger` stacks adapters instead of summing, for multi-adapter blends that are not a simple scale.
参数
xout简介
forward — Promote when VRAM allows, run the Linear, then apply LoRA on the output.
源码 docstring
Promote when VRAM allows, run the Linear, then apply LoRA on the output.
LoRA is applied after the base GEMM so the merge order matches PEFT (delta on the activation, not a second fused weight).
参数
xargskwargsclass AutoWrappedModule(module: torch.nn.Module,offload_dtype: torch.dtype = None,offload_device: Union[str, torch.device] = None,onload_dtype: torch.dtype = None,onload_device: Union[str, torch.device] = None,preparing_dtype: torch.dtype = None,preparing_device: Union[str, torch.device] = None,computation_dtype: torch.dtype = None,computation_device: Union[str, torch.device] = None,vram_limit: float = None,name: str = '',disk_map: DiskMap = None,**kwargs)简介
AutoWrappedModule — Wrap an arbitrary module with staged CPU/GPU/disk weight movement. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
源码 docstring
Wrap an arbitrary module with staged CPU/GPU/disk weight movement.
Forward calls opportunistically prepare weights, materialize a computation copy only when placement differs, and then delegate to the wrapped module. Attribute access falls through to the wrapped module so model code can keep using its original interface.
参数
moduletorch.nn.Moduleoffload_dtypetorch.dtype默认值 None"disk"`.offload_deviceUnion[str, torch.device]默认值 None"disk"`.onload_dtypetorch.dtype默认值 Noneonload_deviceUnion[str, torch.device]默认值 Nonepreparing_dtypetorch.dtype默认值 Nonepreparing_deviceUnion[str, torch.device]默认值 Nonecomputation_dtypetorch.dtype默认值 Nonecomputation_deviceUnion[str, torch.device]默认值 Nonevram_limitfloat默认值 Nonenamestr默认值 ''kwargs方法
简介
load_from_disk — Materialize disk-backed parameters onto *device*; buffers stay resident.
源码 docstring
Materialize disk-backed parameters onto *device*; buffers stay resident.
Non-persistent buffers cannot be reconstructed from a state dict, so they are never replaced. `copy_module=True` clones first so the stored module can remain on meta after a temporary compute load.
参数
torch_dtypedevicecopy_module默认值 Falseoffload_to_disk — Replace disk-backed parameters with meta placeholders, retaining buffers.
参数
modeltorch.nn.Module简介
offload — Return to inactive placement: meta placeholders on disk, else `.to(offload)`.
源码 docstring
Return to inactive placement: meta placeholders on disk, else `.to(offload)`.
Disk offload must not `.to(cpu)` — that would materialize the full module and defeat the mmap path.
onload — Prefetch from disk or move from offload; skip when the target is still disk.
preparing — Second-stage prefetch; materialize only when leaving disk for a real device.
cast_to — Deep-copy then move so the stored offload weights are not mutated in place.
参数
moduledtypedevice简介
computation — Return a module at compute placement without permanently leaving the current state.
源码 docstring
Return a module at compute placement without permanently leaving the current state.
Reuses `self.module when dtype/device already match. Disk + copy keeps the stored module on meta. Otherwise :meth:cast_to` clones.
forward — Promote to preparing when VRAM headroom allows, then run a compute copy.
参数
argskwargsclass AutoWrappedNonRecurseModule(module: torch.nn.Module,offload_dtype: torch.dtype = None,offload_device: Union[str, torch.device] = None,onload_dtype: torch.dtype = None,onload_device: Union[str, torch.device] = None,preparing_dtype: torch.dtype = None,preparing_device: Union[str, torch.device] = None,computation_dtype: torch.dtype = None,computation_device: Union[str, torch.device] = None,vram_limit: float = None,name: str = '',disk_map: DiskMap = None,**kwargs)AutoWrappedNonRecurseModule — Manage only a module's direct parameters while its children are wrapped separately. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
moduletorch.nn.Moduleoffload_dtypetorch.dtype默认值 Noneoffload_deviceUnion[str, torch.device]默认值 Noneonload_dtypetorch.dtype默认值 Noneonload_deviceUnion[str, torch.device]默认值 Nonepreparing_dtypetorch.dtype默认值 Nonepreparing_deviceUnion[str, torch.device]默认值 Nonecomputation_dtypetorch.dtype默认值 Nonecomputation_deviceUnion[str, torch.device]默认值 Nonevram_limitfloat默认值 Nonenamestr默认值 ''kwargs方法
load_from_disk — Materialize this module's direct parameters only; children own their own maps.
参数
torch_dtypedevicecopy_module默认值 Falseoffload_to_disk — Replace only direct parameters with meta placeholders; children stay untouched.
参数
modeltorch.nn.Module简介
cast_to — Return *module* unchanged; the architecture owns parameter casts.
源码 docstring
Return *module* unchanged; the architecture owns parameter casts.
Deep-copying would pickle child `DiskMap` readers (open safetensors handles) and fail. Callers must already be at compute placement.
参数
moduledtypedevice简介
forward — In-place disk materialize when offloaded; otherwise the recursive path.
源码 docstring
In-place disk materialize when offloaded; otherwise the recursive path.
Deep-copy of a hierarchy that contains independently disk-backed children tries to pickle their open safetensors readers.
参数
argskwargsclass BaseMemory(capacity = None,**kwargs)简介
BaseMemory — Generic multimodal memory template for VLM and generative tasks. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
源码 docstring
Generic multimodal memory template for VLM and generative tasks.
Subclasses implement the five-stage pipeline:
Command methods (mutate state): - `record(data, ...) — ingest raw interaction data. - manage()` — evict, merge, or consolidate memories.
Query methods (read state): - `select(context_query, ...) — retrieve relevant snippets. - compress(memory_items, ...) — distill selected memories. - process(refined_data, ...)` — adapt memories to model input formats.
参数
capacity默认值 Nonekwargs方法
storage — Live list of persisted record dicts (same object the store holds).
参数
self返回值
storage — Replace the backing store, preserving `capacity`.
参数
recordsIterable[Mapping[str, Any]]返回值
check_template — Return required record keys and supported content types.
参数
kwargsappend_record(content: Any,kind: str = 'other',timestamp: int | float | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any]源码append_record — Append one typed record through the bounded store.
参数
contentAnykindstr默认值 'other'timestampint | float | None默认值 NonemetadataMapping[str, Any] | None默认值 None返回值
latest_record(prefer_type: str | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any] | None源码latest_record — Return the newest matching record, or `None`.
参数
prefer_typestr | None默认值 NonemetadataMapping[str, Any] | None默认值 None返回值
record — Ingest raw interaction data and assign metadata tags.
参数
datametadata默认值 Nonekwargsselect — Retrieve memory snippets relevant to the current task context.
参数
context_querykwargscompress — Distill selected memories to reduce dimensionality or token count.
参数
_memory_itemskwargsprocess — Convert refined memories into a model-ready format such as KV cache.
参数
_refined_data_target_format默认值 'kv_cache'kwargsmanage — Maintain memory lifecycle: eviction, merging, and STM→LTM transfer.
参数
kwargsclass DiskMap(path,device,torch_dtype = None,state_dict_converter = None,buffer_size = 10 ** 9)简介
DiskMap — Lazy mapping from checkpoint parameter names to materialized tensors. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
源码 docstring
Lazy mapping from checkpoint parameter names to materialized tensors.
Safetensors files remain memory-mapped and individual weights are loaded on lookup. PyTorch binary checkpoints use an in-memory compatibility loader. The mapping may apply a state-dict converter and periodically reopen files after `buffer_size` tensor elements have been materialized.
参数
pathdevicetorch_dtype默认值 Nonestate_dict_converter默认值 Nonebuffer_size默认值 10 ** 9def enable_layerwise_cpu_offload(model: nn.Module,layer_container: str | None = None,device: torch.device | str | None = None,pin_memory: bool = True,reuse_buffers: bool = False) -> LayerwiseOffloadHandle简介
enable_layerwise_cpu_offload — Attach hooks to a ModuleList or an explicitly named Sequential stack. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：LayerwiseOffloadHandle。
源码 docstring
Attach hooks to a ModuleList or an explicitly named Sequential stack.
The helper keeps parameters on CPU and moves one layer at a time to CUDA. The next layer is prefetched on a separate CUDA stream while the current layer runs. It is intentionally opt-in and returns a disabled handle on non-CUDA systems. Keep the returned handle if you need to switch the model back to full-device residency later via `handle.disable(). Opt in with reuse_buffers=True` to pack contiguous parameters into CPU storage and two reusable CUDA slots. The per-parameter path is the default; non-contiguous weights also retain that path to preserve their layout. The reusable path is inference-only and forbids shared parameter storage across layers. Its mutation scope still supports in-place weight edits.
参数
modelnn.Modulelayer_containerstr | None默认值 Nonedevicetorch.device | str | None默认值 Nonepin_memorybool默认值 Truereuse_buffersbool默认值 Falsedef enable_vram_management(model: torch.nn.Module,module_map: dict,vram_config: dict,vram_limit = None,disk_map = None,max_num_param = None,overflow_vram_config: dict | None = None,**kwargs)enable_vram_management — Install staged VRAM management on a model and return that model. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
modeltorch.nn.Modulemodule_map`; otherwise matching descendants are mutated.module_mapdict{torch.nn.Linear: AutoWrappedLinear}`.vram_configdictvram_limit默认值 Nonedisk_map默认值 Nonemax_num_param默认值 Noneoverflow_vram_configdict | None默认值 Nonemax_num_param`.kwargs返回值
vram_management_enabled set to True`.说明
`module_map` order matters when source classes overlap. Configure complete placement keys before calling; this function changes module identity and is normally run once during model construction.
def enable_vram_management_recursively(model: torch.nn.Module,module_map: dict,vram_config: dict,vram_limit = None,name_prefix = '',disk_map = None,max_num_param = None,overflow_vram_config: dict | None = None,total_num_param = 0,**kwargs)enable_vram_management_recursively — Replace matching descendants with lifecycle-aware wrapper modules. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
modeltorch.nn.Modulemodule_mapdictvram_configdictvram_limit默认值 Nonename_prefix默认值 ''DiskMap`.disk_map默认值 Nonemax_num_param默认值 Noneoverflow_vram_config`.overflow_vram_configdict | None默认值 Nonetotal_num_param默认值 0kwargs返回值
说明
The function mutates child modules. Most callers should use `enable_vram_management so the root itself is handled correctly and the vram_management_enabled` marker is installed.
def fill_vram_config(model,vram_config)fill_vram_config — Collapse a placement policy when the root module is wrapped as one unit. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
modelvram_configclass FixedStepCache(skip_steps: Iterable[int] = (),delta_scale: float = 0.0,dense_first: int = 1,dense_last: int = 1,total_steps: int | None = None)简介
FixedStepCache — Reuse a previous denoiser output on an explicit set of steps. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
源码 docstring
Reuse a previous denoiser output on an explicit set of steps.
`delta_scale=0` replays the last dense output. A non-zero value performs first-order extrapolation from the two latest dense outputs.
参数
skip_stepsIterable[int]默认值 ()delta_scalefloat默认值 0.0dense_firstint默认值 1dense_lastint默认值 1total_stepsint | None默认值 Nonedef init_weights_on_device(device = torch.device('meta'),include_buffers: bool = False)init_weights_on_device — Temporarily redirect newly registered module weights to one device. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
device默认值 torch.device('meta')meta` device skips real allocation and initialization.include_buffersbool默认值 False返回值
说明
Global PyTorch hooks are restored in `finally`. This context is not intended to overlap across threads.
def layer_norm_scale_shift(x: torch.Tensor,scale: torch.Tensor,shift: torch.Tensor,eps: float = 1e-06,upcast: bool = False,backend: str | None = None) -> torch.Tensorlayer_norm_scale_shift — Fuse affine-free LayerNorm with AdaLN scale and shift. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：torch.Tensor。
参数
xtorch.Tensorscaletorch.Tensorshifttorch.Tensorepsfloat默认值 1e-06upcastbool默认值 Falsebackendstr | None默认值 None返回值
def layerwise_offload_mutation_scope(module: nn.Module) -> Iterator[None]layerwise_offload_mutation_scope — Temporarily materialize offloaded parameters for in-place mutations. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：Iterator[None]。
参数
modulenn.Module返回值
class LayerwiseOffloadHandle(enabled: bool,layer_count: int,reason: str = '')简介
LayerwiseOffloadHandle — Handle returned by `enable_layerwise_cpu_offload`. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
源码 docstring
Handle returned by `enable_layerwise_cpu_offload`.
Holds the registered hook handles so the offload can be undone with :meth:disable (CC-27); previously the `RemovableHandle` objects were dropped and a model could never leave offload mode without a rebuild.
属性
enabledboollayer_countintreasonstr默认值 ''方法
reset_request_window — Reset request-local counters without discarding a prefetched layer.
返回值
简介
report — Return auditable lifetime and request-local execution evidence.
源码 docstring
Return auditable lifetime and request-local execution evidence.
`enabled only proves that hooks were installed. effective is true after a forward actually used the two-layer CUDA window with asynchronous H2D copies and no emergency synchronous materialization. Pageable host tensors are surfaced explicitly because non_blocking=True` alone does not prove copy/compute overlap.
返回值
简介
disable — Remove offload hooks and restore real parameters onto each layer.
源码 docstring
Remove offload hooks and restore real parameters onto each layer.
Parameters are restored from the CPU master copies (on CPU); move the model to the desired device afterwards. Returns True when hooks were removed. Handles created for an already-enabled model (`reason="already enabled"`) do not own hooks and return False.
返回值
class MemoryStore(capacity: int | None = None,records: Iterable[Mapping[str, Any] | MemoryRecord] = ())MemoryStore — Bounded in-process memory store used by all concrete WorldFoundry memories. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
capacityint | None默认值 NonerecordsIterable[Mapping[str, Any] | MemoryRecord]默认值 ()方法
append(content: Any,kind: str = 'other',timestamp: int | float | None = None,metadata: Mapping[str, Any] | None = None,score: float | None = None) -> dict[str, Any]源码append — Append a new record; *timestamp* defaults to the current length.
参数
contentAnykindstr默认值 'other'timestampint | float | None默认值 NonemetadataMapping[str, Any] | None默认值 Nonescorefloat | None默认值 None返回值
append_record — Normalize *record* to a dict, append it, and evict if over capacity.
参数
recordMapping[str, Any] | MemoryRecord返回值
latest(prefer_type: str | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any] | None源码latest — Return the newest record matching optional type and metadata filters.
参数
prefer_typestr | None默认值 NonemetadataMapping[str, Any] | None默认值 None返回值
select — Rank stored records and return the top-*query.top_k* matches.
参数
queryMemoryQuery | None默认值 NoneoverridesAny返回值
evict — Drop oldest records until `len(self) <= capacity` (no-op if unbounded).
返回值
replace — Clear the store and append *records* in order.
参数
recordsIterable[Mapping[str, Any] | MemoryRecord]返回值
def prune_tokens(hidden_states: torch.Tensor,keep_ratio: float,method: TokenScoreMethod | str = 'feat_norm',seq_dim: int = 1,start: int = 0,end: int | None = None,previous_velocity: torch.Tensor | None = None) -> tuple[torch.Tensor, TokenPruneState]prune_tokens — Gather the selected part of a token segment before expensive blocks. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：tuple[torch.Tensor, TokenPruneState]。
参数
keep_ratiofloatmethodTokenScoreMethod | str默认值 'feat_norm'seq_dimint默认值 1startint默认值 0endint | None默认值 Noneprevious_velocitytorch.Tensor | None默认值 None返回值
def residual_gate_add(residual: torch.Tensor,update: torch.Tensor,gate: torch.Tensor,backend: str | None = None) -> torch.Tensorresidual_gate_add — Return `residual + update * gate with broadcast-aware fusion. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：torch.Tensor`。
参数
residualtorch.Tensorupdatetorch.Tensorgatetorch.Tensorbackendstr | None默认值 None返回值
def restore_tokens(processed: torch.Tensor,state: TokenPruneState,compensation: torch.Tensor | None = None) -> torch.Tensor简介
restore_tokens — Scatter processed tokens and restore dropped tokens from prior state. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：torch.Tensor。
源码 docstring
Scatter processed tokens and restore dropped tokens from prior state.
`compensation` is the previous full segment in the original tensor layout. If omitted, dropped tokens are zero-filled.
参数
返回值
def routed_swiglu_moe(hidden_states: torch.Tensor,routing_weights: torch.Tensor,selected_experts: torch.Tensor,gate_weight: torch.Tensor,up_weight: torch.Tensor,down_weight: torch.Tensor,workspace: dict[str, torch.Tensor] | Callable[[], dict[str, torch.Tensor]] | None = None,allow_triton: bool | None = None) -> torch.Tensorrouted_swiglu_moe — Dispatch packed SwiGLU MoE to Triton or the portable PyTorch reference. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：torch.Tensor。
参数
routing_weightstorch.Tensorselected_expertstorch.Tensorgate_weighttorch.Tensorup_weighttorch.Tensordown_weighttorch.Tensorworkspacedict[str, torch.Tensor] | Callable[[], dict[str, torch.Tensor]] | None默认值 Noneallow_tritonbool | None默认值 None返回值
def routed_swiglu_moe_pytorch(hidden_states: torch.Tensor,routing_weights: torch.Tensor,selected_experts: torch.Tensor,gate_weight: torch.Tensor,up_weight: torch.Tensor,down_weight: torch.Tensor) -> torch.Tensorrouted_swiglu_moe_pytorch — Evaluate a token-routed SwiGLU MoE using portable PyTorch operators. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：torch.Tensor。
参数
routing_weightstorch.Tensor[T, K]`.selected_expertstorch.Tensor[T, K]`.gate_weighttorch.Tensor[E, I, D]`.up_weighttorch.Tensor[E, I, D]`.down_weighttorch.Tensor[E, D, I]`. Only tokens routed to an expert are evaluated for that expert. This keeps the fallback substantially smaller than materializing every expert output, while retaining autograd support and working on CPU, CUDA, and other PyTorch devices.返回值
def select_token_indices(hidden_states: torch.Tensor,keep_ratio: float,method: TokenScoreMethod | str = 'feat_norm',seq_dim: int = 1,previous_velocity: torch.Tensor | None = None,random_seed: int = 42) -> torch.Tensorselect_token_indices — Return ascending indices for tokens retained by a pruning policy. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。 标注返回类型：torch.Tensor。
参数
keep_ratiofloatmethodTokenScoreMethod | str默认值 'feat_norm'seq_dimint默认值 1previous_velocitytorch.Tensor | None默认值 Nonerandom_seedint默认值 42返回值
def skip_model_initialization(device = torch.device('meta'))skip_model_initialization — Return an allocation-skipping context for constructing checkpoint-backed models. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
device默认值 torch.device('meta')class TokenPruner(keep_ratio: float,method: TokenScoreMethod | str = 'feat_norm')TokenPruner — Stateful previous-step compensation for repeated denoising calls. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
参数
keep_ratiofloatmethodTokenScoreMethod | str默认值 'feat_norm'方法
reset — Drop compensation for one branch, or every branch when `key` is omitted.
参数
keyobject | None默认值 None返回值
prune(hidden_states: torch.Tensor,key: object = 'default',seq_dim: int = 1,start: int = 0,end: int | None = None) -> tuple[torch.Tensor, TokenPruneState | None]源码prune — Prune after a dense seed call; first use only records compensation.
参数
keyobject默认值 'default'seq_dimint默认值 1startint默认值 0endint | None默认值 None返回值
restore(processed: torch.Tensor,state: TokenPruneState | None,key: object = 'default') -> torch.Tensor源码restore — Scatter a pruned step, or snapshot the dense seed as next compensation.
参数
返回值
class TokenPruneState(indices: torch.Tensor,start: int,end: int,full_length: int,seq_dim: int)TokenPruneState — Metadata needed to reconstruct a pruned token segment. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。
属性
indicestorch.Tensorstartintendintfull_lengthintseq_dimint方法
kept_length — Number of retained tokens; used to slice the compacted sequence.
参数
self