Core 加速与内存

跨步 cache、token pruning、融合 kernel、显存生命周期 wrapper、DiskMap 与 memory store。

最后更新: 2026年7月22日

本页内容

跨步 cache 示例 Token pruning 生命周期显存 placement 状态语义 Memory 与 VRAM 的区别完整参考

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 都命中了。

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 个公开符号

`AdaptiveResidualCache`

cls

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)

worldfoundry.core.acceleration.AdaptiveResidualCachefrom worldfoundry.core.acceleration import AdaptiveResidualCache

源码

简介

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.

参数

thresholdfloat
warmup_stepsint: 默认值: 1
max_consecutive_hitsint: 默认值: 3
dense_lastint: 默认值: 1
total_stepsint | None: 默认值: None
subsampleint: 默认值: 1
epsfloat: 默认值: 1e-06

方法

methreset() -> None源码

简介

该类型上的公开 method。

返回值: None

meth

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.

参数

stepint
signaltorch.Tensor
compute_residualCallable[[], torch.Tensor]
total_stepsint | None: 默认值: None

返回值: torch.Tensor

`AutoTorchModule`

cls

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)

worldfoundry.core.AutoTorchModulefrom worldfoundry.core import AutoTorchModule

源码

简介

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: Storage dtype while the module is inactive.默认值: None
offload_deviceUnion[str, torch.device]: Storage device while inactive; subclasses may also accept the sentinel `"disk"`.默认值: None
onload_dtypetorch.dtype: Dtype after the first prefetch transition.默认值: None
onload_deviceUnion[str, torch.device]: Device used for prefetched weights.默认值: None
preparing_dtypetorch.dtype: Dtype used by the optional second prefetch stage.默认值: None
preparing_deviceUnion[str, torch.device]: Device used by the preparing stage.默认值: None
computation_dtypetorch.dtype: Dtype used for the actual forward operation.默认值: None
computation_deviceUnion[str, torch.device]: Device used for the actual forward operation.默认值: None
vram_limitfloat: Optional used-memory limit in GiB. Below the limit, wrappers may keep prepared weights resident.默认值: None

方法

meth

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: 默认值: None
offload_deviceUnion[str, torch.device]: 默认值: None
onload_dtypetorch.dtype: 默认值: None
onload_deviceUnion[str, torch.device]: 默认值: None
preparing_dtypetorch.dtype: 默认值: None
preparing_deviceUnion[str, torch.device]: 默认值: None
computation_dtypetorch.dtype: 默认值: None
computation_deviceUnion[str, torch.device]: 默认值: None
vram_limitfloat: 默认值: None

methcast_to(weight,dtype,device)源码

简介

cast_to — Copy one tensor to `dtype and device` without mutating the source.

参数

weight
dtype
device

methcheck_free_vram()源码

简介

check_free_vram — Return whether current accelerator usage is below `vram_limit`.

methoffload()源码

简介

offload — Move managed parameters to the inactive placement and set state 0.

methonload()源码

简介

onload — Move managed parameters to the first prefetch placement and set state 1.

methkeep()源码

简介

keep — Keep managed parameters at computation placement and set state 2.

methparam_name(name)源码

简介

param_name — Return the fully qualified checkpoint key for a local parameter name.

参数

name

`AutoWrappedLinear`

cls

class 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)

worldfoundry.core.AutoWrappedLinearfrom worldfoundry.core import AutoWrappedLinear

源码

简介

AutoWrappedLinear — Linear-layer wrapper with staged placement, disk loading, FP8, and LoRA support. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

moduletorch.nn.Linear: Source `torch.nn.Linear` whose parameters are reused.
offload_dtypetorch.dtype: Inactive storage dtype or `"disk"`.默认值: None
offload_deviceUnion[str, torch.device]: Inactive storage device or `"disk"`.默认值: None
onload_dtypetorch.dtype: First prefetch dtype.默认值: None
onload_deviceUnion[str, torch.device]: First prefetch device.默认值: None
preparing_dtypetorch.dtype: Second prefetch dtype.默认值: None
preparing_deviceUnion[str, torch.device]: Second prefetch device.默认值: None
computation_dtypetorch.dtype: Matmul dtype; supported FP8 dtypes use scaled MM.默认值: None
computation_deviceUnion[str, torch.device]: Matmul device.默认值: None
vram_limitfloat: Optional used-memory limit in GiB.默认值: None
namestr: Qualified checkpoint prefix for weight and bias.默认值: ''
disk_mapDiskMap: Required lazy tensor map when disk offload is selected.默认值: None
kwargs: Reserved for module-map compatibility.

方法

methfp8_linear(input: torch.Tensor,weight: torch.Tensor,bias: torch.Tensor = None) -> torch.Tensor源码

简介

该类型上的公开 method。

参数

inputtorch.Tensor
weighttorch.Tensor
biastorch.Tensor: 默认值: None

返回值: torch.Tensor

methload_from_disk(torch_dtype,device,assign = True)源码

简介

该类型上的公开 method。

参数

torch_dtype
device
assign: 默认值: True

methoffload()源码

简介

该类型上的公开 method。

methonload()源码

简介

该类型上的公开 method。

methpreparing()源码

简介

该类型上的公开 method。

该类型上的公开 method。

参数

x
args
kwargs

`AutoWrappedModule`

cls

class 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)

worldfoundry.core.AutoWrappedModulefrom worldfoundry.core import AutoWrappedModule

源码

简介

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.Module: Original PyTorch module.
offload_dtypetorch.dtype: Inactive storage dtype or `"disk"`.默认值: None
offload_deviceUnion[str, torch.device]: Inactive storage device or `"disk"`.默认值: None
onload_dtypetorch.dtype: First-stage prefetch dtype.默认值: None
onload_deviceUnion[str, torch.device]: First-stage prefetch device.默认值: None
preparing_dtypetorch.dtype: Second-stage prefetch dtype.默认值: None
preparing_deviceUnion[str, torch.device]: Second-stage prefetch device.默认值: None
computation_dtypetorch.dtype: Forward-pass dtype.默认值: None
computation_deviceUnion[str, torch.device]: Forward-pass device.默认值: None
vram_limitfloat: Optional used-memory limit in GiB.默认值: None
namestr: Qualified module name used to resolve disk-map keys.默认值: ''
disk_mapDiskMap: Lazy checkpoint tensor mapping required for disk offload.默认值: None
kwargs: Reserved for wrapper-compatible module maps.

方法

methload_from_disk(torch_dtype,device,copy_module = False)源码

简介

该类型上的公开 method。

参数

torch_dtype
device
copy_module: 默认值: False

methoffload_to_disk(model: torch.nn.Module)源码

简介

该类型上的公开 method。

参数

modeltorch.nn.Module

methoffload()源码

简介

该类型上的公开 method。

methonload()源码

简介

该类型上的公开 method。

methpreparing()源码

简介

该类型上的公开 method。

methcast_to(module,dtype,device)源码

简介

该类型上的公开 method。

参数

module
dtype
device

methcomputation()源码

简介

该类型上的公开 method。

methforward(args, **kwargs)源码

简介

该类型上的公开 method。

参数

args
kwargs

`AutoWrappedNonRecurseModule`

cls

class 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)

worldfoundry.core.AutoWrappedNonRecurseModulefrom worldfoundry.core import AutoWrappedNonRecurseModule

源码

简介

AutoWrappedNonRecurseModule — Manage only a module's direct parameters while its children are wrapped separately. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

moduletorch.nn.Module
offload_dtypetorch.dtype: 默认值: None
offload_deviceUnion[str, torch.device]: 默认值: None
onload_dtypetorch.dtype: 默认值: None
onload_deviceUnion[str, torch.device]: 默认值: None
preparing_dtypetorch.dtype: 默认值: None
preparing_deviceUnion[str, torch.device]: 默认值: None
computation_dtypetorch.dtype: 默认值: None
computation_deviceUnion[str, torch.device]: 默认值: None
vram_limitfloat: 默认值: None
namestr: 默认值: ''
disk_mapDiskMap: 默认值: None
kwargs

module
dtype
device

`BaseMemory`

cls

class BaseMemory(capacity = None, **kwargs)

worldfoundry.core.memory.BaseMemoryfrom worldfoundry.core.memory import BaseMemory

源码

简介

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: 默认值: None
kwargs

方法

propstorage -> list[dict[str, Any]]源码

简介

该类型上的公开 property。

参数

self

返回值: list[dict[str, Any]]

methstorage(records: Iterable[Mapping[str, Any]]) -> None源码

简介

该类型上的公开 method。

参数

recordsIterable[Mapping[str, Any]]

返回值: None

methcheck_template(**kwargs)源码

简介

check_template — Return required record keys and supported content types.

参数

kwargs

meth

append_record(content: Any,kind: str = 'other',timestamp: int | float | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any]

源码

简介

该类型上的公开 method。

参数

contentAny
kindstr: 默认值: 'other'
timestampint | float | None: 默认值: None
metadataMapping[str, Any] | None: 默认值: None

返回值: dict[str, Any]

methlatest_record(prefer_type: str | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any] | None源码

简介

该类型上的公开 method。

参数

prefer_typestr | None: 默认值: None
metadataMapping[str, Any] | None: 默认值: None

返回值: dict[str, Any] | None

methreset_records() -> None源码

简介

该类型上的公开 method。

kwargs

`DiskMap`

cls

class DiskMap(path,device,torch_dtype = None,state_dict_converter = None,buffer_size = 10 ** 9)

worldfoundry.core.DiskMapfrom worldfoundry.core import DiskMap

源码

简介

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.

参数

path: Checkpoint path or list of paths.
device: Device on which fetched tensors are materialized.
torch_dtype: Optional dtype conversion applied on lookup.默认值: None
state_dict_converter: Optional callable that remaps public model keys to keys stored in the checkpoint.默认值: None
buffer_size: Number of fetched tensor elements after which file handles are refreshed.默认值: 10 ** 9

方法

methflush_files()源码

简介

flush_files — Open or refresh checkpoint readers and reset the materialized-element count.

methfetch_rename_dict(state_dict_converter)源码

简介

fetch_rename_dict — Build the optional model-key to checkpoint-key mapping.

参数

state_dict_converter

`enable_layerwise_cpu_offload`

func

def enable_layerwise_cpu_offload(model: nn.Module,layer_container: str | None = None,device: torch.device | str | None = None,pin_memory: bool = True) -> LayerwiseOffloadHandle

worldfoundry.core.enable_layerwise_cpu_offloadfrom worldfoundry.core import enable_layerwise_cpu_offload

源码

简介

enable_layerwise_cpu_offload — Attach layerwise CPU offload hooks to the first or named `ModuleList. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。标注返回类型：LayerwiseOffloadHandle`。

源码 docstring

Attach layerwise CPU offload hooks to the first or named `ModuleList`.

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.

参数

modelnn.Module
layer_containerstr | None: 默认值: None
devicetorch.device | str | None: 默认值: None
pin_memorybool: 默认值: True

返回值: LayerwiseOffloadHandle

`enable_vram_management`

func

def 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)

worldfoundry.core.vram.enable_vram_managementfrom worldfoundry.core.vram import enable_vram_management

源码

简介

enable_vram_management — Install staged VRAM management on a model and return that model. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

modeltorch.nn.Module: Model to wrap. It may be replaced when its own class appears in `module_map`; otherwise matching descendants are mutated.
module_mapdict: Mapping from original module classes to wrapper classes, for example `{torch.nn.Linear: AutoWrappedLinear}`.
vram_configdict: Dictionary containing offload, onload, preparing, and computation dtype/device pairs.
vram_limit: Optional used-memory limit in GiB.默认值: None
disk_map: Lazy checkpoint mapping used by disk-backed wrappers.默认值: None
max_num_param: Optional parameter budget for the primary policy.默认值: None
overflow_vram_configdict | None: Placement used after `max_num_param`.默认值: None
kwargs: Extra wrapper constructor arguments.

说明

`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.

`enable_vram_management_recursively`

func

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)

worldfoundry.core.vram.enable_vram_management_recursivelyfrom worldfoundry.core.vram import enable_vram_management_recursively

源码

简介

enable_vram_management_recursively — Replace matching descendants with lifecycle-aware wrapper modules. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

modeltorch.nn.Module: Root module mutated in place.
module_mapdict: Mapping from source module classes to wrapper classes.
vram_configdict: Default offload/onload/preparing/computation placements.
vram_limit: Optional used-memory limit in GiB forwarded to wrappers.默认值: None
name_prefix: Prefix used when resolving checkpoint keys through a `DiskMap`.默认值: ''
disk_map: Optional lazy checkpoint mapping for disk-backed wrappers.默认值: None
max_num_param: Parameter budget that switches later modules to `overflow_vram_config`.默认值: None
overflow_vram_configdict | None: Placement policy used after the parameter budget.默认值: None
total_num_param: Running parameter count for recursive calls.默认值: 0
kwargs: Extra wrapper constructor arguments.

说明

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.

`fill_vram_config`

func

def fill_vram_config(model, vram_config)

worldfoundry.core.vram.fill_vram_configfrom worldfoundry.core.vram import fill_vram_config

源码

简介

fill_vram_config — Collapse a placement policy when the root module is wrapped as one unit. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

model
vram_config

`FixedStepCache`

cls

class FixedStepCache(skip_steps: Iterable[int] = (),delta_scale: float = 0.0,dense_first: int = 1,dense_last: int = 1,total_steps: int | None = None)

worldfoundry.core.acceleration.FixedStepCachefrom worldfoundry.core.acceleration import FixedStepCache

源码

简介

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.0
dense_firstint: 默认值: 1
dense_lastint: 默认值: 1
total_stepsint | None: 默认值: None

方法

methreset() -> None源码

简介

该类型上的公开 method。

返回值: None

methrun(step: int,compute: Callable[[], T],total_steps: int | None = None) -> T源码

简介

run — Compute or replay one step output.

参数

stepint
computeCallable[[], T]
total_stepsint | None: 默认值: None

返回值: T

`init_weights_on_device`

func

def init_weights_on_device(device = torch.device('meta'), include_buffers: bool = False)

worldfoundry.core.init_weights_on_devicefrom worldfoundry.core import init_weights_on_device

源码

简介

init_weights_on_device — Temporarily redirect newly registered module weights to one device. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

device: Destination for parameters created inside the context. The default `meta` device skips real allocation and initialization.默认值: torch.device('meta')
include_buffersbool: Also redirect registered buffers and common tensor constructors. Leave disabled unless module construction allocates large persistent buffers.默认值: False

说明

Global PyTorch hooks are restored in `finally`. This context is not intended to overlap across threads.

`layer_norm_scale_shift`

func

def layer_norm_scale_shift(x: torch.Tensor,scale: torch.Tensor,shift: torch.Tensor,eps: float = 1e-06,upcast: bool = False) -> torch.Tensor

worldfoundry.core.kernels.layer_norm_scale_shiftfrom worldfoundry.core.kernels import layer_norm_scale_shift

源码

简介

layer_norm_scale_shift — Fuse affine-free LayerNorm with AdaLN scale and shift. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。标注返回类型：torch.Tensor。

参数

xtorch.Tensor
scaletorch.Tensor
shifttorch.Tensor
epsfloat: 默认值: 1e-06
upcastbool: 默认值: False

返回值: torch.Tensor

`layerwise_offload_mutation_scope`

func

def layerwise_offload_mutation_scope(module: nn.Module) -> Iterator[None]

worldfoundry.core.layerwise_offload_mutation_scopefrom worldfoundry.core import layerwise_offload_mutation_scope

源码

简介

layerwise_offload_mutation_scope — Temporarily materialize offloaded parameters for in-place mutations. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。标注返回类型：Iterator[None]。

参数

modulenn.Module

返回值: Iterator[None]

`LayerwiseOffloadHandle`

cls

class LayerwiseOffloadHandle(enabled: bool,layer_count: int,reason: str = '')

worldfoundry.core.LayerwiseOffloadHandlefrom worldfoundry.core import LayerwiseOffloadHandle

源码

简介

LayerwiseOffloadHandle — Handle returned by `enable_layerwise_cpu_offload`. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

属性

enabledbool
layer_countint
reasonstr: 默认值: ''

`MemoryStore`

cls

class MemoryStore(capacity: int | None = None,records: Iterable[Mapping[str, Any] | MemoryRecord] = ())

worldfoundry.core.memory.MemoryStorefrom worldfoundry.core.memory import MemoryStore

源码

简介

MemoryStore — Bounded in-process memory store used by all concrete WorldFoundry memories. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

capacityint | None: 默认值: None
recordsIterable[Mapping[str, Any] | MemoryRecord]: 默认值: ()

方法

meth

append(content: Any,kind: str = 'other',timestamp: int | float | None = None,metadata: Mapping[str, Any] | None = None,score: float | None = None) -> dict[str, Any]

源码

简介

该类型上的公开 method。

参数

contentAny
kindstr: 默认值: 'other'
timestampint | float | None: 默认值: None
metadataMapping[str, Any] | None: 默认值: None
scorefloat | None: 默认值: None

返回值: dict[str, Any]

methappend_record(record: Mapping[str, Any] | MemoryRecord) -> dict[str, Any]源码

简介

该类型上的公开 method。

参数

recordMapping[str, Any] | MemoryRecord

返回值: dict[str, Any]

methlatest(prefer_type: str | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any] | None源码

简介

该类型上的公开 method。

参数

prefer_typestr | None: 默认值: None
metadataMapping[str, Any] | None: 默认值: None

返回值: dict[str, Any] | None

methselect(query: MemoryQuery | None = None, **overrides: Any) -> MemorySelection源码

简介

select — Rank stored records and return the top-*query.top_k* matches.

参数

queryMemoryQuery | None: 默认值: None
overridesAny

返回值: MemorySelection

methevict() -> None源码

简介

该类型上的公开 method。

返回值: None

methreset() -> None源码

简介

该类型上的公开 method。

返回值: None

methreplace(records: Iterable[Mapping[str, Any] | MemoryRecord]) -> None源码

简介

该类型上的公开 method。

参数

recordsIterable[Mapping[str, Any] | MemoryRecord]

返回值: None

`prune_tokens`

func

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]

worldfoundry.core.acceleration.prune_tokensfrom worldfoundry.core.acceleration import prune_tokens

源码

简介

prune_tokens — Gather the selected part of a token segment before expensive blocks. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。标注返回类型：tuple[torch.Tensor, TokenPruneState]。

参数

hidden_statestorch.Tensor
keep_ratiofloat
methodTokenScoreMethod | str: 默认值: 'feat_norm'
seq_dimint: 默认值: 1
startint: 默认值: 0
endint | None: 默认值: None
previous_velocitytorch.Tensor | None: 默认值: None

返回值: tuple[torch.Tensor, TokenPruneState]

`residual_gate_add`

func

def residual_gate_add(residual: torch.Tensor,update: torch.Tensor,gate: torch.Tensor) -> torch.Tensor

worldfoundry.core.kernels.residual_gate_addfrom worldfoundry.core.kernels import residual_gate_add

源码

简介

residual_gate_add — Return `residual + update * gate with broadcast-aware fusion. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。标注返回类型：torch.Tensor`。

参数

residualtorch.Tensor
updatetorch.Tensor
gatetorch.Tensor

返回值: torch.Tensor

`restore_tokens`

func

def restore_tokens(processed: torch.Tensor,state: TokenPruneState,compensation: torch.Tensor | None = None) -> torch.Tensor

worldfoundry.core.acceleration.restore_tokensfrom worldfoundry.core.acceleration import restore_tokens

源码

简介

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.

参数

processedtorch.Tensor
stateTokenPruneState
compensationtorch.Tensor | None: 默认值: None

返回值: torch.Tensor

`routed_swiglu_moe`

func

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.Tensor

worldfoundry.core.routed_swiglu_moefrom worldfoundry.core import routed_swiglu_moe

源码

简介

routed_swiglu_moe — Dispatch packed SwiGLU MoE to Triton or the portable PyTorch reference. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。标注返回类型：torch.Tensor。

参数

hidden_statestorch.Tensor
routing_weightstorch.Tensor
selected_expertstorch.Tensor
gate_weighttorch.Tensor
up_weighttorch.Tensor
down_weighttorch.Tensor
workspacedict[str, torch.Tensor] | Callable[[], dict[str, torch.Tensor]] | None: 默认值: None
allow_tritonbool | None: 默认值: None

返回值: torch.Tensor

`routed_swiglu_moe_pytorch`

func

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.Tensor

worldfoundry.core.routed_swiglu_moe_pytorchfrom worldfoundry.core import routed_swiglu_moe_pytorch

源码

简介

routed_swiglu_moe_pytorch — Evaluate a token-routed SwiGLU MoE using portable PyTorch operators. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。标注返回类型：torch.Tensor。

参数

hidden_statestorch.Tensor: Flattened token activations with shape `[T, D]`.
routing_weightstorch.Tensor: Per-route weights with shape `[T, K]`.
selected_expertstorch.Tensor: Expert indices with shape `[T, K]`.
gate_weighttorch.Tensor: Packed gate weights with shape `[E, I, D]`.
up_weighttorch.Tensor: Packed up-projection weights with shape `[E, I, D]`.
down_weighttorch.Tensor: Packed down-projection weights with shape `[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.

返回值: torch.Tensor

`select_token_indices`

func

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.Tensor

worldfoundry.core.acceleration.select_token_indicesfrom worldfoundry.core.acceleration import select_token_indices

源码

简介

select_token_indices — Return ascending indices for tokens retained by a pruning policy. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。标注返回类型：torch.Tensor。

参数

hidden_statestorch.Tensor
keep_ratiofloat
methodTokenScoreMethod | str: 默认值: 'feat_norm'
seq_dimint: 默认值: 1
previous_velocitytorch.Tensor | None: 默认值: None
random_seedint: 默认值: 42

返回值: torch.Tensor

`skip_model_initialization`

func

def skip_model_initialization(device = torch.device('meta'))

worldfoundry.core.skip_model_initializationfrom worldfoundry.core import skip_model_initialization

源码

简介

skip_model_initialization — Return an allocation-skipping context for constructing checkpoint-backed models. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

device: 默认值: torch.device('meta')

`TokenPruner`

cls

class TokenPruner(keep_ratio: float, , method: TokenScoreMethod | str = 'feat_norm')

worldfoundry.core.acceleration.TokenPrunerfrom worldfoundry.core.acceleration import TokenPruner

源码

简介

TokenPruner — Stateful previous-step compensation for repeated denoising calls. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

参数

keep_ratiofloat
methodTokenScoreMethod | str: 默认值: 'feat_norm'

方法

methreset(key: object | None = None) -> None源码

简介

该类型上的公开 method。

参数

keyobject | None: 默认值: None

返回值: None

meth

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.

参数

hidden_statestorch.Tensor
keyobject: 默认值: 'default'
seq_dimint: 默认值: 1
startint: 默认值: 0
endint | None: 默认值: None

返回值: tuple[torch.Tensor, TokenPruneState | None]

methrestore(processed: torch.Tensor,state: TokenPruneState | None,key: object = 'default') -> torch.Tensor源码

简介

该类型上的公开 method。

参数

processedtorch.Tensor
stateTokenPruneState | None
keyobject: 默认值: 'default'

返回值: torch.Tensor

`TokenPruneState`

cls

class TokenPruneState(indices: torch.Tensor,start: int,end: int,full_length: int,seq_dim: int)

worldfoundry.core.acceleration.TokenPruneStatefrom worldfoundry.core.acceleration import TokenPruneState

源码

简介

TokenPruneState — Metadata needed to reconstruct a pruned token segment. 属于 Core 加速与内存（cache、offload、VRAM、kernels）。

属性

indicestorch.Tensor
startint
endint
full_lengthint
seq_dimint

方法

propkept_length -> int源码

简介

该类型上的公开 property。

参数

self

返回值: int