fovi.arch.knn_gather_gemm

Gather+GEMM ops for the degenerate K=1 / V=1 KNN convolution class.

fovi-resnet18’s downsample convolutions (res18_ds2/ds3/ds4) have exactly one neighbor per output node (K=1) and a single reference-grid cell (V=1), so the general operator

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

degenerates to a pure indexed gather followed by ONE shared dense GEMM:

y[b, o, n] = bias[o] + sum_c x[b, c, idx[n]] * weight[o, c]

with no per-node weight indexing at all (weight is [Cout, Cin]). Forward and backward are each an index_select/index_add_ plus a single (batched) matmul, which works for every dtype (fp32 / fp16 / bf16) and is within launch-overhead distance of the dense Conv2d floor.

Registered as "gather_gemm" in the fovi.arch.knn_autograd ops registry. Routing predicate for the dispatcher: meta.k == 1 and meta.v == 1 (layer-level: layer._k == 1 and layer.local_rf.shape[2] == 1).

Padding semantics match the baseline: padding neighbors (input_linear entries equal to Cin*Nin) contribute exact zeros forward and receive no grad_input scatter.

fovi.arch.knn_gather_gemm._node_index(meta)

Derive the [Nout] gather index and (optional) valid mask from meta.input_linear.

For K=1 the c=0 column of input_linear is 0*Nin + idx[n] for real neighbors and Cin*Nin (out of range) for padding neighbors. Returns (index, valid) with valid is None when there is no padding (the common case for the res18 downsamples).

fovi.arch.knn_gather_gemm._scatter_matrix(meta, dtype)

Cached one-hot [Nout, Nin] scatter operand: S[n, idx[n]] = 1 for real neighbors.

grad_input then becomes a single GEMM d_selected @ S — collisions (several output nodes sharing one input node) sum inside cuBLAS’s fp32 accumulators, padding rows are all-zero, and the fp32 zeros/index_add/cast round-trips disappear. Dense S is only sensible because this class is tiny (Nout*Nin <= a few 1e4 for the res18 downsamples).

class fovi.arch.knn_gather_gemm.GatherGemmOps

Bases: object

Ops-registry backend for the K=1/V=1 class: index_select + one dense GEMM.

name = 'gather_gemm'

static forward(meta, x, weight, bias)

static grad_input(meta, grad_y, weight)

static grad_weight(meta, grad_y, x)

fovi.arch.knn_gather_gemm._register_ops(ops) → None

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