Core acceleration and memory
Cross-step caches, token pruning, fused kernels, VRAM lifecycle wrappers, disk maps, and memory stores.
On this page
Core separates exact implementation acceleration from approximation policy. Fused kernels and placement changes aim to preserve model semantics. Cross-step caches and token pruning intentionally trade computation for an approximation and therefore expose thresholds, retained ratios, events, and reset boundaries that integrations can evaluate.
Cross-step cache example
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] # boundary steps remain dense
assert [event.hit for event in cache.events] == [False, True, False, True, False]Caches disable replay while autograd is enabled. Call reset() between independent denoising trajectories; otherwise prior residuals become state for the next request. Treat the event stream as evidence during latency/quality evaluation rather than assuming requested skip steps were all used.
Token pruning lifecycle
select_token_indices is stateless. prune_tokens returns compact data plus TokenPruneState, and restore_tokens scatters processed tokens back, filling dropped positions from compensation or zeros. TokenPruner adds previous-step compensation: its first call records a dense segment, later calls can prune, and every prune must be paired with restore under the same key.
VRAM placement state
enable_vram_management replaces classes listed in module_map with wrappers such as AutoWrappedLinear. Each wrapper can use separate offload, onload, preparing, and computation dtype/device pairs. Disk-backed wrappers resolve parameter names through DiskMap; ordinary wrappers move or copy tensors between devices.
This transformation mutates module identity and should normally run once during model construction. The module-map order matters for overlapping classes. A placement configuration is not automatically safe for a new architecture: validate peak memory, transfer overlap, output equality/tolerance, and behavior when the configured VRAM limit is reached.
Memory records versus VRAM
The BaseMemory and MemoryStore APIs represent retrievable world-model memory records. They do not manage GPU allocation. VRAM wrappers manage parameter placement but do not provide semantic retrieval. Keeping those meanings separate prevents a “memory” setting from being applied at the wrong layer.
Complete reference
The blocks below are the generated signatures for this category. Use the on-page symbol index to jump; source links open the defining implementation behind each lazy export.
26 public symbols
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 AdaptiveResidualCacheOverview
Reuse a model residual while accumulated input change stays small. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Source 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.
Parameters
thresholdfloatwarmup_stepsint- default:
1 max_consecutive_hitsint- default:
3 dense_lastint- default:
1 total_stepsint | None- default:
None subsampleint- default:
1 epsfloat- default:
1e-06
Methods
run(step: int,signal: torch.Tensor,compute_residual: Callable[[], torch.Tensor],total_steps: int | None = None) -> torch.TensorsourceOverview
Return a dense or cached residual for one denoising step.
Parameters
stepintsignaltorch.Tensorcompute_residualCallable[[], torch.Tensor]total_stepsint | None- default:
None
Returns: torch.Tensor
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 AutoTorchModuleOverview
Base state machine for modules that move between memory tiers. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Source 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.
Parameters
offload_dtypetorch.dtype- Storage dtype while the module is inactive.default:
None offload_deviceUnion[str, torch.device]- Storage device while inactive; subclasses may also accept the sentinel `
"disk"`.default:None onload_dtypetorch.dtype- Dtype after the first prefetch transition.default:
None onload_deviceUnion[str, torch.device]- Device used for prefetched weights.default:
None preparing_dtypetorch.dtype- Dtype used by the optional second prefetch stage.default:
None preparing_deviceUnion[str, torch.device]- Device used by the preparing stage.default:
None computation_dtypetorch.dtype- Dtype used for the actual forward operation.default:
None computation_deviceUnion[str, torch.device]- Device used for the actual forward operation.default:
None vram_limitfloat- Optional used-memory limit in GiB. Below the limit, wrappers may keep prepared weights resident.default:
None
Methods
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)sourceOverview
Update lifecycle placement, defaulting omitted stages to computation placement.
Parameters
offload_dtypetorch.dtype- default:
None offload_deviceUnion[str, torch.device]- default:
None onload_dtypetorch.dtype- default:
None onload_deviceUnion[str, torch.device]- default:
None preparing_dtypetorch.dtype- default:
None preparing_deviceUnion[str, torch.device]- default:
None computation_dtypetorch.dtype- default:
None computation_deviceUnion[str, torch.device]- default:
None vram_limitfloat- default:
None
Overview
Copy one tensor to `dtype and device` without mutating the source.
Parameters
weightdtypedevice
Overview
Return whether current accelerator usage is below `vram_limit`.
Overview
Move managed parameters to the inactive placement and set state 0.
Overview
Move managed parameters to the first prefetch placement and set state 1.
Overview
Keep managed parameters at computation placement and set state 2.
Overview
Return the fully qualified checkpoint key for a local parameter name.
Parameters
name
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 AutoWrappedLinearOverview
Linear-layer wrapper with staged placement, disk loading, FP8, and LoRA support. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
moduletorch.nn.Linear- Source `
torch.nn.Linear` whose parameters are reused. offload_dtypetorch.dtype- Inactive storage dtype or `
"disk"`.default:None offload_deviceUnion[str, torch.device]- Inactive storage device or `
"disk"`.default:None onload_dtypetorch.dtype- First prefetch dtype.default:
None onload_deviceUnion[str, torch.device]- First prefetch device.default:
None preparing_dtypetorch.dtype- Second prefetch dtype.default:
None preparing_deviceUnion[str, torch.device]- Second prefetch device.default:
None computation_dtypetorch.dtype- Matmul dtype; supported FP8 dtypes use scaled MM.default:
None computation_deviceUnion[str, torch.device]- Matmul device.default:
None vram_limitfloat- Optional used-memory limit in GiB.default:
None namestr- Qualified checkpoint prefix for weight and bias.default:
'' disk_mapDiskMap- Required lazy tensor map when disk offload is selected.default:
None kwargs- Reserved for module-map compatibility.
Methods
fp8_linear(input: torch.Tensor,weight: torch.Tensor,bias: torch.Tensor = None) -> torch.TensorsourceOverview
Public method on this type.
Parameters
inputtorch.Tensorweighttorch.Tensorbiastorch.Tensor- default:
None
Returns: torch.Tensor
Overview
Public method on this type.
Parameters
torch_dtypedeviceassign- default:
True
Overview
Public method on this type.
Overview
Public method on this type.
Overview
Public method on this type.
Overview
Public method on this type.
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 AutoWrappedModuleOverview
Wrap an arbitrary module with staged CPU/GPU/disk weight movement. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Source 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.
Parameters
moduletorch.nn.Module- Original PyTorch module.
offload_dtypetorch.dtype- Inactive storage dtype or `
"disk"`.default:None offload_deviceUnion[str, torch.device]- Inactive storage device or `
"disk"`.default:None onload_dtypetorch.dtype- First-stage prefetch dtype.default:
None onload_deviceUnion[str, torch.device]- First-stage prefetch device.default:
None preparing_dtypetorch.dtype- Second-stage prefetch dtype.default:
None preparing_deviceUnion[str, torch.device]- Second-stage prefetch device.default:
None computation_dtypetorch.dtype- Forward-pass dtype.default:
None computation_deviceUnion[str, torch.device]- Forward-pass device.default:
None vram_limitfloat- Optional used-memory limit in GiB.default:
None namestr- Qualified module name used to resolve disk-map keys.default:
'' disk_mapDiskMap- Lazy checkpoint tensor mapping required for disk offload.default:
None kwargs- Reserved for wrapper-compatible module maps.
Methods
Overview
Public method on this type.
Parameters
torch_dtypedevicecopy_module- default:
False
Overview
Public method on this type.
Parameters
modeltorch.nn.Module
Overview
Public method on this type.
Overview
Public method on this type.
Overview
Public method on this type.
Overview
Public method on this type.
Parameters
moduledtypedevice
Overview
Public method on this type.
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 AutoWrappedNonRecurseModuleOverview
Manage only a module's direct parameters while its children are wrapped separately. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
moduletorch.nn.Moduleoffload_dtypetorch.dtype- default:
None offload_deviceUnion[str, torch.device]- default:
None onload_dtypetorch.dtype- default:
None onload_deviceUnion[str, torch.device]- default:
None preparing_dtypetorch.dtype- default:
None preparing_deviceUnion[str, torch.device]- default:
None computation_dtypetorch.dtype- default:
None computation_deviceUnion[str, torch.device]- default:
None vram_limitfloat- default:
None namestr- default:
'' disk_mapDiskMap- default:
None kwargs
Methods
Overview
Public method on this type.
Parameters
torch_dtypedevicecopy_module- default:
False
Overview
Public method on this type.
Parameters
modeltorch.nn.Module
Overview
Public method on this type.
Parameters
moduledtypedevice
BaseMemory
clsclass BaseMemory(capacity = None, **kwargs)worldfoundry.core.memory.BaseMemoryfrom worldfoundry.core.memory import BaseMemoryOverview
Generic multimodal memory template for VLM and generative tasks. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Source 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.
Parameters
capacity- default:
None kwargs
Methods
Overview
Public property on this type.
Parameters
self
Returns: list[dict[str, Any]]
Overview
Public method on this type.
Parameters
recordsIterable[Mapping[str, Any]]
Returns: None
Overview
Return required record keys and supported content types.
Parameters
kwargs
append_record(content: Any,kind: str = 'other',timestamp: int | float | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any]sourceOverview
Public method on this type.
Parameters
contentAnykindstr- default:
'other' timestampint | float | None- default:
None metadataMapping[str, Any] | None- default:
None
Returns: dict[str, Any]
latest_record(prefer_type: str | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any] | NonesourceOverview
Public method on this type.
Parameters
prefer_typestr | None- default:
None metadataMapping[str, Any] | None- default:
None
Returns: dict[str, Any] | None
Overview
Ingest raw interaction data and assign metadata tags.
Parameters
datametadata- default:
None kwargs
Overview
Retrieve memory snippets relevant to the current task context.
Parameters
context_querykwargs
Overview
Distill selected memories to reduce dimensionality or token count.
Parameters
_memory_itemskwargs
Overview
Convert refined memories into a model-ready format such as KV cache.
Parameters
_refined_data_target_format- default:
'kv_cache' kwargs
Overview
Maintain memory lifecycle: eviction, merging, and STM→LTM transfer.
Parameters
kwargs
DiskMap
clsclass DiskMap(path,device,torch_dtype = None,state_dict_converter = None,buffer_size = 10 ** 9)worldfoundry.core.DiskMapfrom worldfoundry.core import DiskMapOverview
Lazy mapping from checkpoint parameter names to materialized tensors. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Source 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.
Parameters
path- Checkpoint path or list of paths.
device- Device on which fetched tensors are materialized.
torch_dtype- Optional dtype conversion applied on lookup.default:
None state_dict_converter- Optional callable that remaps public model keys to keys stored in the checkpoint.default:
None buffer_size- Number of fetched tensor elements after which file handles are refreshed.default:
10 ** 9
def enable_layerwise_cpu_offload(model: nn.Module,layer_container: str | None = None,device: torch.device | str | None = None,pin_memory: bool = True) -> LayerwiseOffloadHandleworldfoundry.core.enable_layerwise_cpu_offloadfrom worldfoundry.core import enable_layerwise_cpu_offloadOverview
Attach layerwise CPU offload hooks to the first or named `ModuleList. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: LayerwiseOffloadHandle`.
Source 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.
Parameters
modelnn.Modulelayer_containerstr | None- default:
None devicetorch.device | str | None- default:
None pin_memorybool- default:
True
Returns: LayerwiseOffloadHandle
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_managementOverview
Install staged VRAM management on a model and return that model. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
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.default:
None disk_map- Lazy checkpoint mapping used by disk-backed wrappers.default:
None max_num_param- Optional parameter budget for the primary policy.default:
None overflow_vram_configdict | None- Placement used after `
max_num_param`.default:None kwargs- Extra wrapper constructor arguments.
Notes
`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)worldfoundry.core.vram.enable_vram_management_recursivelyfrom worldfoundry.core.vram import enable_vram_management_recursivelyOverview
Replace matching descendants with lifecycle-aware wrapper modules. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
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.default:
None name_prefix- Prefix used when resolving checkpoint keys through a `
DiskMap`.default:'' disk_map- Optional lazy checkpoint mapping for disk-backed wrappers.default:
None max_num_param- Parameter budget that switches later modules to `
overflow_vram_config`.default:None overflow_vram_configdict | None- Placement policy used after the parameter budget.default:
None total_num_param- Running parameter count for recursive calls.default:
0 kwargs- Extra wrapper constructor arguments.
Notes
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
funcdef fill_vram_config(model, vram_config)worldfoundry.core.vram.fill_vram_configfrom worldfoundry.core.vram import fill_vram_configOverview
Collapse a placement policy when the root module is wrapped as one unit. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
modelvram_config
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 FixedStepCacheOverview
Reuse a previous denoiser output on an explicit set of steps. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Source 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.
Parameters
skip_stepsIterable[int]- default:
() delta_scalefloat- default:
0.0 dense_firstint- default:
1 dense_lastint- default:
1 total_stepsint | None- default:
None
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_deviceOverview
Temporarily redirect newly registered module weights to one device. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
device- Destination for parameters created inside the context. The default `
meta` device skips real allocation and initialization.default:torch.device('meta') include_buffersbool- Also redirect registered buffers and common tensor constructors. Leave disabled unless module construction allocates large persistent buffers.default:
False
Notes
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) -> torch.Tensorworldfoundry.core.kernels.layer_norm_scale_shiftfrom worldfoundry.core.kernels import layer_norm_scale_shiftOverview
Fuse affine-free LayerNorm with AdaLN scale and shift. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: torch.Tensor.
Parameters
xtorch.Tensorscaletorch.Tensorshifttorch.Tensorepsfloat- default:
1e-06 upcastbool- default:
False
Returns: torch.Tensor
def layerwise_offload_mutation_scope(module: nn.Module) -> Iterator[None]worldfoundry.core.layerwise_offload_mutation_scopefrom worldfoundry.core import layerwise_offload_mutation_scopeOverview
Temporarily materialize offloaded parameters for in-place mutations. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: Iterator[None].
Parameters
modulenn.Module
Returns: Iterator[None]
class LayerwiseOffloadHandle(enabled: bool,layer_count: int,reason: str = '')worldfoundry.core.LayerwiseOffloadHandlefrom worldfoundry.core import LayerwiseOffloadHandleOverview
Handle returned by `enable_layerwise_cpu_offload`. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Attributes
enabledboollayer_countintreasonstr- default:
''
MemoryStore
clsclass MemoryStore(capacity: int | None = None,records: Iterable[Mapping[str, Any] | MemoryRecord] = ())worldfoundry.core.memory.MemoryStorefrom worldfoundry.core.memory import MemoryStoreOverview
Bounded in-process memory store used by all concrete WorldFoundry memories. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
capacityint | None- default:
None recordsIterable[Mapping[str, Any] | MemoryRecord]- default:
()
Methods
append(content: Any,kind: str = 'other',timestamp: int | float | None = None,metadata: Mapping[str, Any] | None = None,score: float | None = None) -> dict[str, Any]sourceOverview
Public method on this type.
Parameters
contentAnykindstr- default:
'other' timestampint | float | None- default:
None metadataMapping[str, Any] | None- default:
None scorefloat | None- default:
None
Returns: dict[str, Any]
Overview
Public method on this type.
Parameters
recordMapping[str, Any] | MemoryRecord
Returns: dict[str, Any]
latest(prefer_type: str | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any] | NonesourceOverview
Public method on this type.
Parameters
prefer_typestr | None- default:
None metadataMapping[str, Any] | None- default:
None
Returns: dict[str, Any] | None
Overview
Rank stored records and return the top-*query.top_k* matches.
Parameters
queryMemoryQuery | None- default:
None overridesAny
Returns: MemorySelection
Overview
Public method on this type.
Parameters
recordsIterable[Mapping[str, Any] | MemoryRecord]
Returns: None
prune_tokens
funcdef 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_tokensOverview
Gather the selected part of a token segment before expensive blocks. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: tuple[torch.Tensor, TokenPruneState].
Parameters
keep_ratiofloatmethodTokenScoreMethod | str- default:
'feat_norm' seq_dimint- default:
1 startint- default:
0 endint | None- default:
None previous_velocitytorch.Tensor | None- default:
None
Returns: tuple[torch.Tensor, TokenPruneState]
def residual_gate_add(residual: torch.Tensor,update: torch.Tensor,gate: torch.Tensor) -> torch.Tensorworldfoundry.core.kernels.residual_gate_addfrom worldfoundry.core.kernels import residual_gate_addOverview
Return `residual + update * gate with broadcast-aware fusion. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: torch.Tensor`.
Parameters
residualtorch.Tensorupdatetorch.Tensorgatetorch.Tensor
Returns: torch.Tensor
restore_tokens
funcdef restore_tokens(processed: torch.Tensor,state: TokenPruneState,compensation: torch.Tensor | None = None) -> torch.Tensorworldfoundry.core.acceleration.restore_tokensfrom worldfoundry.core.acceleration import restore_tokensOverview
Scatter processed tokens and restore dropped tokens from prior state. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: torch.Tensor.
Source 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.
Parameters
processedtorch.TensorstateTokenPruneStatecompensationtorch.Tensor | None- default:
None
Returns: torch.Tensor
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.Tensorworldfoundry.core.routed_swiglu_moefrom worldfoundry.core import routed_swiglu_moeOverview
Dispatch packed SwiGLU MoE to Triton or the portable PyTorch reference. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: torch.Tensor.
Parameters
routing_weightstorch.Tensorselected_expertstorch.Tensorgate_weighttorch.Tensorup_weighttorch.Tensordown_weighttorch.Tensorworkspacedict[str, torch.Tensor] | Callable[[], dict[str, torch.Tensor]] | None- default:
None allow_tritonbool | None- default:
None
Returns: torch.Tensor
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.Tensorworldfoundry.core.routed_swiglu_moe_pytorchfrom worldfoundry.core import routed_swiglu_moe_pytorchOverview
Evaluate a token-routed SwiGLU MoE using portable PyTorch operators. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: torch.Tensor.
Parameters
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.
Returns: torch.Tensor
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.Tensorworldfoundry.core.acceleration.select_token_indicesfrom worldfoundry.core.acceleration import select_token_indicesOverview
Return ascending indices for tokens retained by a pruning policy. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: torch.Tensor.
Parameters
keep_ratiofloatmethodTokenScoreMethod | str- default:
'feat_norm' seq_dimint- default:
1 previous_velocitytorch.Tensor | None- default:
None random_seedint- default:
42
Returns: torch.Tensor
def skip_model_initialization(device = torch.device('meta'))worldfoundry.core.skip_model_initializationfrom worldfoundry.core import skip_model_initializationOverview
Return an allocation-skipping context for constructing checkpoint-backed models. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
device- default:
torch.device('meta')
TokenPruner
clsclass TokenPruner(keep_ratio: float, , method: TokenScoreMethod | str = 'feat_norm')worldfoundry.core.acceleration.TokenPrunerfrom worldfoundry.core.acceleration import TokenPrunerOverview
Stateful previous-step compensation for repeated denoising calls. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
keep_ratiofloatmethodTokenScoreMethod | str- default:
'feat_norm'
Methods
Overview
Public method on this type.
Parameters
keyobject | None- default:
None
Returns: 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]sourceOverview
Prune after a dense seed call; first use only records compensation.
Parameters
keyobject- default:
'default' seq_dimint- default:
1 startint- default:
0 endint | None- default:
None
Returns: tuple[torch.Tensor, TokenPruneState | None]
restore(processed: torch.Tensor,state: TokenPruneState | None,key: object = 'default') -> torch.TensorsourceOverview
Public method on this type.
Parameters
processedtorch.TensorstateTokenPruneState | Nonekeyobject- default:
'default'
Returns: torch.Tensor
class TokenPruneState(indices: torch.Tensor,start: int,end: int,full_length: int,seq_dim: int)worldfoundry.core.acceleration.TokenPruneStatefrom worldfoundry.core.acceleration import TokenPruneStateOverview
Metadata needed to reconstruct a pruned token segment. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Attributes
indicestorch.Tensorstartintendintfull_lengthintseq_dimint