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)Overview
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_stepsintdefault 1max_consecutive_hitsintdefault 3dense_lastintdefault 1total_stepsint | Nonedefault Nonesubsampleintdefault 1epsfloatdefault 1e-06Methods
Clear residual and accumulator so a new clip cannot reuse the last one.
Returns
run(step: int,signal: torch.Tensor,compute_residual: Callable[[], torch.Tensor],total_steps: int | None = None) -> torch.TensorSourceReturn a dense or cached residual for one denoising step.
Parameters
stepintsignaltorch.Tensorcompute_residualCallable[[], torch.Tensor]total_stepsint | Nonedefault NoneReturns
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)Overview
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.dtypedefault Noneoffload_deviceUnion[str, torch.device]default None"disk"`.onload_dtypetorch.dtypedefault Noneonload_deviceUnion[str, torch.device]default Nonepreparing_dtypetorch.dtypedefault Nonepreparing_deviceUnion[str, torch.device]default Nonecomputation_dtypetorch.dtypedefault Nonecomputation_deviceUnion[str, torch.device]default Nonevram_limitfloatdefault NoneMethods
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)SourceUpdate lifecycle placement, defaulting omitted stages to computation placement.
Parameters
offload_dtypetorch.dtypedefault Noneoffload_deviceUnion[str, torch.device]default Noneonload_dtypetorch.dtypedefault Noneonload_deviceUnion[str, torch.device]default Nonepreparing_dtypetorch.dtypedefault Nonepreparing_deviceUnion[str, torch.device]default Nonecomputation_dtypetorch.dtypedefault Nonecomputation_deviceUnion[str, torch.device]default Nonevram_limitfloatdefault NoneCopy one tensor to `dtype and device` without mutating the source.
Parameters
weightdtypedeviceReturn whether current accelerator usage is below `vram_limit`.
Move managed parameters to the inactive placement and set state 0.
Move managed parameters to the first prefetch placement and set state 1.
Keep managed parameters at computation placement and set state 2.
Return the fully qualified checkpoint key for a local parameter name.
Parameters
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)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.Lineartorch.nn.Linear` whose parameters are reused.offload_dtypetorch.dtypedefault None"disk"`.offload_deviceUnion[str, torch.device]default None"disk"`.onload_dtypetorch.dtypedefault Noneonload_deviceUnion[str, torch.device]default Nonepreparing_dtypetorch.dtypedefault Nonepreparing_deviceUnion[str, torch.device]default Nonecomputation_dtypetorch.dtypedefault Nonecomputation_deviceUnion[str, torch.device]default Nonevram_limitfloatdefault Nonenamestrdefault ''kwargsMethods
fp8_linear(input: torch.Tensor,weight: torch.Tensor,bias: torch.Tensor = None) -> torch.TensorSourceOverview
Row-wise scaled FP8 matmul via `torch._scaled_mm`.
Source 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.
Parameters
inputtorch.Tensorweighttorch.Tensorbiastorch.Tensordefault NoneReturns
Overview
Fetch `name.weight / name.bias` from the disk map.
Source 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.
Parameters
torch_dtypedeviceassigndefault TrueDrop to meta on disk offload; otherwise move this Linear to offload placement.
Prefetch this Linear from disk or move from offload; skip a still-disk target.
Second-stage prefetch for this Linear; materialize only when leaving disk.
Overview
Return weight/bias at compute placement without permanently leaving the current state.
Source 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.
Dense GEMM, or the FP8 scaled path when `computation_dtype` is float8.
Parameters
xweightbiasOverview
Add LoRA deltas after the base linear so merge order matches PEFT (B @ A).
Source 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.
Parameters
xoutOverview
Promote when VRAM allows, run the Linear, then apply LoRA on the output.
Source 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).
Parameters
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)Overview
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.Moduleoffload_dtypetorch.dtypedefault None"disk"`.offload_deviceUnion[str, torch.device]default None"disk"`.onload_dtypetorch.dtypedefault Noneonload_deviceUnion[str, torch.device]default Nonepreparing_dtypetorch.dtypedefault Nonepreparing_deviceUnion[str, torch.device]default Nonecomputation_dtypetorch.dtypedefault Nonecomputation_deviceUnion[str, torch.device]default Nonevram_limitfloatdefault Nonenamestrdefault ''kwargsMethods
Overview
Materialize disk-backed parameters onto *device*; buffers stay resident.
Source 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.
Parameters
torch_dtypedevicecopy_moduledefault FalseReplace disk-backed parameters with meta placeholders, retaining buffers.
Parameters
modeltorch.nn.ModuleOverview
Return to inactive placement: meta placeholders on disk, else `.to(offload)`.
Source 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.
Prefetch from disk or move from offload; skip when the target is still disk.
Second-stage prefetch; materialize only when leaving disk for a real device.
Deep-copy then move so the stored offload weights are not mutated in place.
Parameters
moduledtypedeviceOverview
Return a module at compute placement without permanently leaving the current state.
Source 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.
Promote to preparing when VRAM headroom allows, then run a compute copy.
Parameters
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)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.dtypedefault Noneoffload_deviceUnion[str, torch.device]default Noneonload_dtypetorch.dtypedefault Noneonload_deviceUnion[str, torch.device]default Nonepreparing_dtypetorch.dtypedefault Nonepreparing_deviceUnion[str, torch.device]default Nonecomputation_dtypetorch.dtypedefault Nonecomputation_deviceUnion[str, torch.device]default Nonevram_limitfloatdefault Nonenamestrdefault ''kwargsMethods
Materialize this module's direct parameters only; children own their own maps.
Parameters
torch_dtypedevicecopy_moduledefault FalseReplace only direct parameters with meta placeholders; children stay untouched.
Parameters
modeltorch.nn.ModuleOverview
Return *module* unchanged; the architecture owns parameter casts.
Source 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.
Parameters
moduledtypedeviceOverview
In-place disk materialize when offloaded; otherwise the recursive path.
Source 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.
Parameters
argskwargsBaseMemory
class BaseMemory(capacity = None,**kwargs)Overview
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
capacitydefault NonekwargsMethods
Live list of persisted record dicts (same object the store holds).
Parameters
selfReturns
Replace the backing store, preserving `capacity`.
Parameters
recordsIterable[Mapping[str, Any]]Returns
Return required record keys and supported content types.
Parameters
kwargsappend_record(content: Any,kind: str = 'other',timestamp: int | float | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any]SourceAppend one typed record through the bounded store.
Parameters
contentAnykindstrdefault 'other'timestampint | float | Nonedefault NonemetadataMapping[str, Any] | Nonedefault NoneReturns
latest_record(prefer_type: str | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any] | NoneSourceReturn the newest matching record, or `None`.
Parameters
prefer_typestr | Nonedefault NonemetadataMapping[str, Any] | Nonedefault NoneReturns
Ingest raw interaction data and assign metadata tags.
Parameters
datametadatadefault NonekwargsRetrieve memory snippets relevant to the current task context.
Parameters
context_querykwargsDistill selected memories to reduce dimensionality or token count.
Parameters
_memory_itemskwargsConvert refined memories into a model-ready format such as KV cache.
Parameters
_refined_data_target_formatdefault 'kv_cache'kwargsMaintain memory lifecycle: eviction, merging, and STM→LTM transfer.
Parameters
kwargsclass DiskMap(path,device,torch_dtype = None,state_dict_converter = None,buffer_size = 10 ** 9)Overview
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
pathdevicetorch_dtypedefault Nonestate_dict_converterdefault Nonebuffer_sizedefault 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) -> LayerwiseOffloadHandleOverview
Attach hooks to a ModuleList or an explicitly named Sequential stack. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels). Annotated return type: LayerwiseOffloadHandle.
Source 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.
Parameters
modelnn.Modulelayer_containerstr | Nonedefault Nonedevicetorch.device | str | Nonedefault Nonepin_memorybooldefault Truereuse_buffersbooldefault FalseReturns
enable_vram_management
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)Install staged VRAM management on a model and return that model. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
modeltorch.nn.Modulemodule_map`; otherwise matching descendants are mutated.module_mapdict{torch.nn.Linear: AutoWrappedLinear}`.vram_configdictvram_limitdefault Nonedisk_mapdefault Nonemax_num_paramdefault Noneoverflow_vram_configdict | Nonedefault Nonemax_num_param`.kwargsReturns
vram_management_enabled set to True`.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)Replace matching descendants with lifecycle-aware wrapper modules. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
modeltorch.nn.Modulemodule_mapdictvram_configdictvram_limitdefault Nonename_prefixdefault ''DiskMap`.disk_mapdefault Nonemax_num_paramdefault Noneoverflow_vram_config`.overflow_vram_configdict | Nonedefault Nonetotal_num_paramdefault 0kwargsReturns
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
def fill_vram_config(model,vram_config)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_configFixedStepCache
class FixedStepCache(skip_steps: Iterable[int] = (),delta_scale: float = 0.0,dense_first: int = 1,dense_last: int = 1,total_steps: int | None = None)Overview
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_scalefloatdefault 0.0dense_firstintdefault 1dense_lastintdefault 1total_stepsint | Nonedefault NoneMethods
Drop stored outputs so the next `run` is a seed, not a stale hit.
Returns
Compute or replay one step output.
Parameters
stepintcomputeCallable[[], T]total_stepsint | Nonedefault NoneReturns
init_weights_on_device
def init_weights_on_device(device = torch.device('meta'),include_buffers: bool = False)Temporarily redirect newly registered module weights to one device. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
devicedefault torch.device('meta')meta` device skips real allocation and initialization.include_buffersbooldefault FalseReturns
Notes
Global PyTorch hooks are restored in `finally`. This context is not intended to overlap across threads.
layer_norm_scale_shift
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.TensorFuse 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.Tensorepsfloatdefault 1e-06upcastbooldefault Falsebackendstr | Nonedefault NoneReturns
def layerwise_offload_mutation_scope(module: nn.Module) -> Iterator[None]Overview
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.ModuleReturns
class LayerwiseOffloadHandle(enabled: bool,layer_count: int,reason: str = '')Overview
Handle returned by `enable_layerwise_cpu_offload`. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Source 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.
Attributes
enabledboollayer_countintreasonstrdefault ''Methods
Reset request-local counters without discarding a prefetched layer.
Returns
Overview
Return auditable lifetime and request-local execution evidence.
Source 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.
Returns
Overview
Remove offload hooks and restore real parameters onto each layer.
Source 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.
Returns
MemoryStore
class MemoryStore(capacity: int | None = None,records: Iterable[Mapping[str, Any] | MemoryRecord] = ())Bounded in-process memory store used by all concrete WorldFoundry memories. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
capacityint | Nonedefault NonerecordsIterable[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]SourceAppend a new record; *timestamp* defaults to the current length.
Parameters
contentAnykindstrdefault 'other'timestampint | float | Nonedefault NonemetadataMapping[str, Any] | Nonedefault Nonescorefloat | Nonedefault NoneReturns
Normalize *record* to a dict, append it, and evict if over capacity.
Parameters
recordMapping[str, Any] | MemoryRecordReturns
latest(prefer_type: str | None = None,metadata: Mapping[str, Any] | None = None) -> dict[str, Any] | NoneSourceReturn the newest record matching optional type and metadata filters.
Parameters
prefer_typestr | Nonedefault NonemetadataMapping[str, Any] | Nonedefault NoneReturns
Rank stored records and return the top-*query.top_k* matches.
Parameters
queryMemoryQuery | Nonedefault NoneoverridesAnyReturns
Drop oldest records until `len(self) <= capacity` (no-op if unbounded).
Returns
Clear the store and append *records* in order.
Parameters
recordsIterable[Mapping[str, Any] | MemoryRecord]Returns
prune_tokens
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]Overview
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 | strdefault 'feat_norm'seq_dimintdefault 1startintdefault 0endint | Nonedefault Noneprevious_velocitytorch.Tensor | Nonedefault NoneReturns
residual_gate_add
def residual_gate_add(residual: torch.Tensor,update: torch.Tensor,gate: torch.Tensor,backend: str | None = None) -> torch.TensorReturn `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.Tensorbackendstr | Nonedefault NoneReturns
restore_tokens
def restore_tokens(processed: torch.Tensor,state: TokenPruneState,compensation: torch.Tensor | None = None) -> torch.TensorOverview
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
Returns
routed_swiglu_moe
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.TensorOverview
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]] | Nonedefault Noneallow_tritonbool | Nonedefault NoneReturns
routed_swiglu_moe_pytorch
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.TensorOverview
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[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.Returns
select_token_indices
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.TensorReturn 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 | strdefault 'feat_norm'seq_dimintdefault 1previous_velocitytorch.Tensor | Nonedefault Nonerandom_seedintdefault 42Returns
skip_model_initialization
def skip_model_initialization(device = torch.device('meta'))Return an allocation-skipping context for constructing checkpoint-backed models. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
devicedefault torch.device('meta')TokenPruner
class TokenPruner(keep_ratio: float,method: TokenScoreMethod | str = 'feat_norm')Stateful previous-step compensation for repeated denoising calls. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Parameters
keep_ratiofloatmethodTokenScoreMethod | strdefault 'feat_norm'Methods
Drop compensation for one branch, or every branch when `key` is omitted.
Parameters
keyobject | Nonedefault NoneReturns
prune(hidden_states: torch.Tensor,key: object = 'default',seq_dim: int = 1,start: int = 0,end: int | None = None) -> tuple[torch.Tensor, TokenPruneState | None]SourcePrune after a dense seed call; first use only records compensation.
Parameters
keyobjectdefault 'default'seq_dimintdefault 1startintdefault 0endint | Nonedefault NoneReturns
restore(processed: torch.Tensor,state: TokenPruneState | None,key: object = 'default') -> torch.TensorSourceScatter a pruned step, or snapshot the dense seed as next compensation.
Parameters
Returns
class TokenPruneState(indices: torch.Tensor,start: int,end: int,full_length: int,seq_dim: int)Metadata needed to reconstruct a pruned token segment. Belongs to Core acceleration and memory (caches, offload, VRAM, kernels).
Attributes
indicestorch.Tensorstartintendintfull_lengthintseq_dimintMethods
Number of retained tokens; used to slice the compacted sequence.
Parameters
self