hooks

Modules

fastvideo.hooks.activation_trace

Zero-overhead-when-off activation trace mode for FastVideo pipelines.

Enable by setting FASTVIDEO_TRACE_ACTIVATIONS=1. When off, this module adds zero overhead — no hooks are registered, no branches exist in the production forward path. When on, registers forward hooks on modules whose name matches FASTVIDEO_TRACE_LAYERS, computes the requested stats (FASTVIDEO_TRACE_STATS) on each output tensor, and writes JSONL records to FASTVIDEO_TRACE_OUTPUT.

Useful for parity debugging across model ports — log on both the FastVideo path and the upstream reference, diff the two JSONL files to find the first divergent layer.

Example:

FASTVIDEO_TRACE_ACTIVATIONS=1     FASTVIDEO_TRACE_LAYERS="^block\.layers\.[0-9]+$"     FASTVIDEO_TRACE_STATS="abs_mean,max,shape"     FASTVIDEO_TRACE_OUTPUT="/tmp/fv_trace.jsonl"     python examples/inference/basic/basic_magi_human.py

Classes

fastvideo.hooks.activation_trace.ActivationStatHook

ActivationStatHook(module_name: str, stats: list[tuple[str, Callable[[Tensor], Any]]], sink: JsonlSink, step_filter: set[int] | None)

Bases: ForwardHook

Forward hook that emits per-tensor stats to a JSONL sink.

Source code in fastvideo/hooks/activation_trace.py

def __init__(
    self,
    module_name: str,
    stats: list[tuple[str, Callable[[torch.Tensor], Any]]],
    sink: JsonlSink,
    step_filter: set[int] | None,
) -> None:
    self.module_name = module_name
    self.stats = stats
    self.sink = sink
    self.step_filter = step_filter

fastvideo.hooks.activation_trace.JsonlSink

JsonlSink(path: Path)

Buffered JSONL writer with thread-safe append.

Source code in fastvideo/hooks/activation_trace.py

def __init__(self, path: Path) -> None:
    self.path = path
    self.path.parent.mkdir(parents=True, exist_ok=True)
    self._fh = open(self.path, "a", buffering=1)  # noqa: SIM115
    self._lock = threading.Lock()
    logger.info("Activation trace JSONL sink: %s", self.path)

Functions:

fastvideo.hooks.activation_trace.attach_activation_trace

attach_activation_trace(model: Module | None) -> ActivationTraceManager | None

Attach activation-stat hooks to model. Returns None if trace is off.

Source code in fastvideo/hooks/activation_trace.py

def attach_activation_trace(model: nn.Module | None) -> ActivationTraceManager | None:
    """Attach activation-stat hooks to model. Returns None if trace is off."""
    if not envs.FASTVIDEO_TRACE_ACTIVATIONS or model is None:
        return None

    pattern_spec = envs.FASTVIDEO_TRACE_LAYERS
    pattern = re.compile(pattern_spec) if pattern_spec else re.compile(".*")
    stats = _resolve_stats(envs.FASTVIDEO_TRACE_STATS)
    if not stats:
        logger.warning("FASTVIDEO_TRACE_STATS yielded no valid stats; trace disabled.")
        return None

    sink = JsonlSink(_resolve_output_path(envs.FASTVIDEO_TRACE_OUTPUT))
    step_filter = _parse_step_filter(envs.FASTVIDEO_TRACE_STEPS)
    managers = []
    for name, module in model.named_modules():
        if not name or not pattern.search(name):
            continue
        manager = ModuleHookManager.get_from_or_default(module)
        manager.append_forward_hook(ActivationStatHook(name, stats, sink, step_filter))
        managers.append(manager)

    logger.info(
        "Activation trace attached to %d modules (pattern=%r, stats=%s)",
        len(managers),
        pattern_spec,
        [stat_name for stat_name, _ in stats],
    )
    return ActivationTraceManager(managers, sink)

fastvideo.hooks.activation_trace.trace_step

trace_step(step_idx: int) -> Iterator[None]

Context manager that sets the current denoise step for trace records.

Source code in fastvideo/hooks/activation_trace.py

@contextmanager
def trace_step(step_idx: int) -> Iterator[None]:
    """Context manager that sets the current denoise step for trace records."""
    prev = getattr(_TRACE_STATE, "step_idx", None)
    _TRACE_STATE.step_idx = step_idx
    try:
        yield
    finally:
        _TRACE_STATE.step_idx = prev

fastvideo.hooks.hooks

Classes

fastvideo.hooks.hooks.ForwardHook

Base class for forward hooks. Hooks are used in the way: modified_args, modified_kwargs = hook.pre_forward(module, args, kwargs) output = module.forward(modified_args, **modified_kwargs) modified_output = hook.post_forward(module, output)

Methods:

fastvideo.hooks.hooks.ForwardHook.on_attach

on_attach(module: Module)

Called once when the hook is attached to the module.

Source code in fastvideo/hooks/hooks.py

def on_attach(self, module: nn.Module):  # noqa: B027
    """Called once when the hook is attached to the module."""
    pass

fastvideo.hooks.hooks.ForwardHook.on_detach

on_detach(module: Module)

Called once when the hook is detached from the module. Note: this function is not guaranteed to be called if the module is deleted before the hook is detached.

Source code in fastvideo/hooks/hooks.py

def on_detach(self, module: nn.Module):  # noqa: B027
    """
    Called once when the hook is detached from the module.
    Note: this function is not guaranteed to be called if the module is
          deleted before the hook is detached.
    """
    pass

fastvideo.hooks.hooks.ForwardHook.post_forward

post_forward(module: Module, output: Any) -> Any

Called after the module's forward method is executed.

Source code in fastvideo/hooks/hooks.py

def post_forward(self, module: nn.Module, output: Any) -> Any:
    """Called after the module's forward method is executed."""
    return output

fastvideo.hooks.hooks.ForwardHook.pre_forward

pre_forward(module: Module, *args, **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]

Called before the module's forward method is executed.

Source code in fastvideo/hooks/hooks.py

def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]:
    """Called before the module's forward method is executed."""
    return args, kwargs

fastvideo.hooks.layerwise_offload

Classes

fastvideo.hooks.layerwise_offload.LayerwiseOffloadHook

LayerwiseOffloadHook(state: LayerwiseOffloadState)

Bases: ForwardHook

A hook that enables layerwise CPU offloading during forward pass.

Source code in fastvideo/hooks/layerwise_offload.py

def __init__(self, state: LayerwiseOffloadState) -> None:
    self.state = state

Functions: