xnn.common.models.fast

Fast paths: optimized kernels behind the same parameters as the reference code.

A module with a fast path keeps exactly the parameters, buffers and state_dict of its reference implementation and adds a second way of evaluating them (for example a fused GPU kernel). Checkpoints, the model hub and the trainer never see the difference, and the reference stays the canonical implementation: TorchScript export, CPU runs and anything the fast backend does not support use it.

Every such module subclasses FastPathModule. A model owning them chooses, once per evaluation, whether they run fast:

use_fast=False

always the reference code.

use_fast=True

the fast path wherever it is available (the reference otherwise).

use_fast="auto" (the default)

the fast path when it is available and measured to be faster for this input, i.e. on a GPU and above a size threshold per GPU model (see AutoPolicy).

The fast kernels compute the same function with a different summation order, so results agree with the reference to rounding: about 1e-6 relative in float32 and 1e-14 in float64. A one-time FastPathPrecisionWarning says so when a fast path first runs in float32.

Functions

deactivate(model)	Switch every fast path inside model off (before TorchScript export).
fast_modules(model)	Every FastPathModule inside model.
iter_fast_capable(model, device, dtype)	The fast-path modules of model that can run on device in dtype.
resolve_use_fast(use_fast)	Validate a use_fast value (True, False or "auto").
select(modules, use_fast, policy, device, ...)	Set FastPathModule.fast_active on modules for one evaluation.
set_use_fast(model, use_fast)	Set use_fast on every model inside model that has fast paths.
warn_float32_once()	Issue the float32 precision warning, once per process.

Classes

AutoPolicy([min_edges, min_edges_float64, ...])	When use_fast="auto" takes the fast path.
FastPathModule()	Mixin of a module with a reference and a fast implementation.

Exceptions

FastPathPrecisionWarning

A fast path ran in float32: results differ from the reference at ~1e-6.