fovi.arch.knn_autograd

Training-capable compact backends for KNNConvLayer.

The baseline layer computes

y[b, o, n] = bias[o] + sum_{c,k} x[b, c, knn[k, n]] * weight[o, c*V + rf_index[n, k]]

through a dense one-hot einsum followed by F.linear. The backends here use the compact contraction width P = Cin*K instead of Q = Cin*V and provide a custom autograd path so training no longer falls back to the baseline.

Structure:

• TrainingMeta bundles the derived index tables every backend consumes, including a reverse-CSR structure mapping each input node to the (k, n) pairs that reference it (needed for deterministic grad-input kernels).

• KNNConvFunction is the single torch.autograd.Function; per-backend compute is looked up in OPS_REGISTRY so kernel backends (Warp, CUDA) only need to register an ops object with forward / grad_input / grad_weight.

• torch_scatter replaces the dense one-hot einsum with scatter_add_ and keeps native autograd; torch_compact is the compact bmm formulation with the custom backward.

Gradient accumulation is always fp32; scattered writes go through index_add_ on fp32 buffers (never fp16 atomics). Under autocast, inputs are cast to the autocast dtype inside the Function, grad_input is returned in x.dtype and grad_weight / grad_bias in the parameter dtype (fp32 master weights under AMP).

fovi.arch.knn_autograd.dataclass(cls=None, /, *, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False)

Add dunder methods based on the fields defined in the class.

Examines PEP 526 __annotations__ to determine fields.

If init is true, an __init__() method is added to the class. If repr is true, a __repr__() method is added. If order is true, rich comparison dunder methods are added. If unsafe_hash is true, a __hash__() method is added. If frozen is true, fields may not be assigned to after instance creation. If match_args is true, the __match_args__ tuple is added. If kw_only is true, then by default all fields are keyword-only. If slots is true, a new class with a __slots__ attribute is returned.

class fovi.arch.knn_autograd.TrainingMeta(cin: int, cout: int, nin: int, nout: int, k: int, v: int, p: int, p64: int, q: int, input_linear: Tensor, weight_linear: Tensor, input_linear_flat: Tensor, weight_linear_flat: Tensor, rev_rowptr: Tensor, rev_col: Tensor, device: device)

Bases: object

Derived, device-resident index tables shared by all training backends.

cin: int

cout: int

nin: int

nout: int

k: int

v: int

p: int

p64: int

q: int

input_linear: Tensor

weight_linear: Tensor

input_linear_flat: Tensor

weight_linear_flat: Tensor

rev_rowptr: Tensor

rev_col: Tensor

device: device

__init__(cin: int, cout: int, nin: int, nout: int, k: int, v: int, p: int, p64: int, q: int, input_linear: Tensor, weight_linear: Tensor, input_linear_flat: Tensor, weight_linear_flat: Tensor, rev_rowptr: Tensor, rev_col: Tensor, device: device) → None

fovi.arch.knn_autograd.ensure_training_metadata(layer, device: device) → TrainingMeta

Build (or fetch cached) TrainingMeta for layer on device.

Cached on layer._knn_training_metadata; released by fovi.arch.knn_optimization.clear_cache(). Never enters state_dict.

fovi.arch.knn_autograd._nout_chunk(meta: TrainingMeta, batch: int, itemsize: int) → int

Largest Nout slice keeping per-chunk staging (W_eff + A + dA) under budget.

fovi.arch.knn_autograd._pad_flat_input(meta: TrainingMeta, x: Tensor) → Tensor

[B, Cin, Nin] -> [B, Cin*Nin + 1] with a trailing zero pad column.

fovi.arch.knn_autograd._effective_weight_t(meta: TrainingMeta, weight: Tensor) → Tensor

weight [Cout, Q] -> transposed compact operand source [Q, Cout], contiguous.

class fovi.arch.knn_autograd.CompactTorchOps

Bases: object

Pure-torch compact ops: gather -> bmm forward, bmm -> fp32 index_add_ backward.

Serves as the correctness oracle and the structural skeleton for kernel backends. All three entry points take the compute-dtype tensors prepared by KNNConvFunction (casting/AMP policy handled there, not here).

name = 'torch_compact'

static forward(meta: TrainingMeta, x: Tensor, weight: Tensor, bias) → Tensor

static grad_input(meta: TrainingMeta, grad_y: Tensor, weight: Tensor) → Tensor

static grad_weight(meta: TrainingMeta, grad_y: Tensor, x: Tensor) → Tensor

fovi.arch.knn_autograd.register_ops(ops) → None

Register a kernel backend’s ops object (must expose name/forward/grad_input/grad_weight).

class fovi.arch.knn_autograd.KNNConvFunction(*args, **kwargs)

Bases: Function

static forward(ctx, x, weight, bias, meta, ops_name)

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

• See Combined or separate forward() and setup_context() for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.

• Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

• See Extending torch.autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

static backward(ctx, grad_y)

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

fovi.arch.knn_autograd.compact_forward(layer, x: Tensor, ops_name: str = 'torch_compact') → Tensor

Layer-level entry point for registry-backed backends (training and inference).

fovi.arch.knn_autograd.scatter_forward(layer, x: Tensor) → Tensor

Baseline-equivalent forward with scatter_add_ instead of the one-hot einsum.

Exact-arithmetic-identical to the baseline (padding neighbors gather the zero pad node and scatter zeros; rf-bin collisions sum, exactly as the einsum does), fully differentiable through native autograd, and valid for every dtype/device.