fovi.arch.knn_optimization

Optimized backends for KNNConvLayer (inference and training).

All paths exploit the fact that local_rf is one-hot: each KNN neighbor selects exactly one reference-grid weight. Inference-only backends (torch_cached/warp_cached/warp_memory) cache derived operands keyed on the weight version. Training-capable backends (torch_scatter/torch_compact, plus kernel backends registered in fovi.arch.knn_autograd) provide autograd-correct forward/backward and are selected when gradients are enabled.

fovi.arch.knn_optimization._is_gather_gemm_layer(layer) → bool

K=1/V=1 layers (resnet-style downsample convs) degenerate to gather + one dense GEMM.

fovi.arch.knn_optimization._native_cuda_supported(device: device) → bool

Return whether the fused convolution kernel supports this CUDA device.

fovi.arch.knn_optimization.clear_cache(layer) → None

Release cached effective weights and derived index metadata (inference + training).

fovi.arch.knn_optimization.select_backend(layer, x: Tensor) → str

fovi.arch.knn_optimization.optimized_forward(layer, x: Tensor) → Tensor | None

Return an optimized result, or None when the baseline should run.