fovi.arch.knn_pool_cuda

Native CUDA (CuPy NVRTC) kernels for KNNPoolingLayer.

The baseline layer (fovi/arch/knn.py) materializes a [B, C, K, Nout] NaN-padded gather (194 MiB at alexp0 B=512 fp16) and reduces it with NaN-aware torch ops; the max path additionally performs a boolean-mask in-place fill whose nonzero triggers a host sync every call, and its backward runs a max-scatter + index_put backward + fp16 scatter_add_ chain (~3x the forward). This module replaces both directions with one kernel each:

• Forward (knn_pool_forward): fused gather + reduce. One CTA stages RS consecutive (b, c) rows of x (rows are contiguous in global memory, so the staging copy is a single flat coalesced range) plus the persistent [K, Nout] pad-token index table into shared memory, then reduces every (row, n) in registers. The B*C*K*Nout intermediate is never materialized and no NaN sentinel is needed (pad entries are detected by index == Nin). max emits a uint8 argmax slot per (b, c, n) for the backward, using the strict > / ascending-k update whose tie rule (first occurrence of the maximum) matches torch.max on CUDA (verified empirically; the baseline backward routes ties identically). avg accumulates valid neighbors in fp32 and divides by the precomputed valid count with an IEEE division, matching nanmean’s sum-then-divide bit-for-bit (0/0 = NaN reproduces the empty-slice NaN for all-pad neighborhoods).

• Backward (knn_pool_backward): deterministic, atomic-free reverse-CSR gather. For each input node m the persistent CSR (built once per layer) lists the packed (k, n) pairs referencing it; every dx[b, c, m] is accumulated in fp32 registers and written exactly once (so a NaN-prefill canary works for gradients too, unlike accumulate-into-zeros designs). max adds gy[b, c, n] where the stored argmax equals k; avg adds gy[b, c, n] / count[n]. This removes the baseline’s fp16 scatter-add gradient noise entirely. Pooling is parameter-free and channel-wise, so there is no weight gradient and no Cin*K contraction anywhere.

Host-path doctrine from the conv kernels applies (sub-ms calls are host-bound): kernel arguments are raw data_ptr() values, launches go through the conv module’s cached ExternalStream path, and all derived index structures are built once per layer and cached. Unlike the conv Function, the pooling backward needs neither x nor any re-layout of grad_y — only the uint8 argmax map is saved for training (max mode), an activation-memory win over the baseline’s saved gather intermediates.

NVRTC note: torch must be imported before CuPy (see fovi.arch.knn_cuda); importing that module first guarantees the ordering and provides the shared launch machinery.

class fovi.arch.knn_pool_cuda.PoolMeta(nin: int, nout: int, k: int, indices: Tensor, indices_i64: Tensor, count: Tensor, inv_count: Tensor, rev_rowptr: Tensor, rev_kn: Tensor, device: device)

Bases: object

Derived, device-resident index tables for one pooling layer.

nin: int

nout: int

k: int

indices: Tensor

indices_i64: Tensor

count: Tensor

inv_count: Tensor

rev_rowptr: Tensor

rev_kn: Tensor

device: device

__init__(nin: int, nout: int, k: int, indices: Tensor, indices_i64: Tensor, count: Tensor, inv_count: Tensor, rev_rowptr: Tensor, rev_kn: Tensor, device: device) → None

fovi.arch.knn_pool_cuda.pool_meta_from_indices(indices, nin, device=None)

Build PoolMeta from a [K, Nout] pad-token index table.

indices follows the layer convention: entries >= nin (the layer uses exactly nin) are padding. The reverse CSR drops padding entries, packs j = (k << 24) | n and orders entries by (input node, then flattened k*Nout + n) via a stable sort, so the backward’s fp32 register accumulation is deterministic.

fovi.arch.knn_pool_cuda.ensure_pool_metadata(layer, device)

Build (or fetch cached) PoolMeta for a KNNPoolingLayer.

Cached on layer._knn_pool_cuda_meta (never enters state_dict); the index table is fixed at layer construction, so the cache is keyed by device only.

fovi.arch.knn_pool_cuda.pool_forward(meta, x, mode, rs=None, need_aux=True)

Fused KNN pooling forward.

Parameters:

• meta – PoolMeta for the layer’s index table.

• x – [B, C, Nin] CUDA tensor, float16 / bfloat16 / float32.

• mode – "max" or "avg".

• rs – optional rows-per-CTA override (default: heuristic).

• need_aux – emit the argmax map (max mode only; skip at inference to save traffic).

Returns:

y is [B, C, Nout] in x.dtype; aux is the uint8 argmax map (max mode with need_aux) or None.

Return type:

(y, aux)

fovi.arch.knn_pool_cuda.pool_backward(meta, grad_y, aux, mode, rs=None)

Deterministic reverse-CSR KNN pooling backward.

Parameters:

• meta – PoolMeta (must be the forward’s meta).

• grad_y – [B, C, Nout] CUDA tensor, float16 / bfloat16 / float32.

• aux – the forward’s uint8 argmax map (required for max; ignored for avg).

• mode – "max" or "avg".

• rs – optional rows-per-CTA override.

Returns:

[B, C, Nin] gradient in grad_y.dtype (fp32 register accumulation, one final rounding — strictly less gradient noise than the baseline’s fp16 scatter_add_ chain).

class fovi.arch.knn_pool_cuda.KNNPoolFunction(*args, **kwargs)

Bases: Function

Autograd wrapper: fused forward + deterministic rev-CSR backward.

Pooling is parameter-free and not autocast-eligible (the baseline runs in the input dtype under AMP, consuming the preceding conv’s half-precision output), so no dtype cast happens here: compute dtype == x.dtype. Only the uint8 argmax map is saved for max mode — neither x nor any gather intermediate is retained.

static forward(ctx, x, meta, mode, need_aux)

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_pool_cuda.pool_function(x, indices, mode)

Benchmark-harness-compatible entry point: (x, [K, Nout] indices, mode) -> y with full autograd support (meta cached by index-table identity).

fovi.arch.knn_pool_cuda.optimized_pool_forward(layer, x)

Layer-level dispatch mirroring fovi.arch.knn_optimization.optimized_forward().

Returns the optimized result, or None when the caller should fall back to the baseline (unsupported mode/dtype/device, shape outside the shared-memory envelope). Reads only layer.mode, layer.knn_indices_pad_token and layer.knn_pad_token_val (in_coords.shape[0] fallback); caches metadata on layer._knn_pool_cuda_meta.

fovi.arch.knn_pool_cuda.clear_pool_cache()

Drop compiled-module references, canary state, and the raw-indices meta cache (per-layer metadata lives on the layer; CuPy’s on-disk compile cache is unaffected).