fovi.arch.knn_cuda

Native CUDA (CuPy NVRTC) kernels for compact KNN convolution.

This module implements the custom-CUDA forward path for KNNConvLayer:

y[b, o, n] = bias[o] + sum_{p} A_n[b, p] * W_n[p, o]

where per output node n the operand tiles are gathered through two persistent int32 tables (input_linear/weight_linear, shape [Nout, P64] with P64 = ceil(Cin*K/64)*64; see fovi.arch.knn_optimization). Out-of-range input_linear entries equal Cin*Nin and must contribute zero.

Key design points versus an earlier WMMA prototype:

• Both gathered operands are re-laid-out host-side so the contiguous axis is the non-indexed one: x is transposed to xt[Cin*Nin + 1, Bpad] (last row zeros = padding row) and the weight to wt[Cin*V, Cpad]. Each random index then selects a row whose copy is a clean run of 16-byte chunks, so staging is fully vectorized and coalesced.

• Shared-memory staging is double-buffered with cp.async (__pipeline_memcpy_async).

• One CTA processes RB consecutive batch sub-tiles against a single staged weight tile (multi-b-tile weight reuse; the killer was re-streaming W_n for every 16-row batch tile).

• Output is written to a [B, Nout, Cout] buffer (coalesced along Cout) and returned as a [B, Cout, Nout] transposed view, matching the baseline layer’s return convention.

• Everything is a C++ template over the scalar type (__half/__nv_bfloat16) and tile geometry, instantiated on demand through name_expressions. Accumulation is always fp32.

Host-path note: kernel arguments are passed as raw data_ptr() values (np.uint64 packs byte-identically to an ndarray pointer in CuPy’s launch ABI), skipping per-call DLPack export/stream negotiation — the small-shape training cells are host-bound. This also sidesteps CuPy’s lack of a bfloat16 dtype entirely.

NVRTC note: CuPy resolves libnvrtc.so.12 from whatever is already loaded in the process. import torch (done below, before cupy) loads torch’s bundled NVRTC (12.8 for this repo’s stack), which is required to emit native sm_120 (Blackwell) cubins. Importing cupy first in a fresh process can silently bind an older NVRTC and fail with CUDA_ERROR_NO_BINARY_FOR_GPU.

fovi.arch.knn_cuda.forward(x, weight, bias, input_linear, weight_linear, config=None)

Compute the compact KNN convolution forward pass with the native CUDA kernel.

Parameters:

• x – [B, Cin, Nin] CUDA tensor, float16 or bfloat16 (fp32 inputs are cast).

• weight – [Cout, Cin * V] tensor on the same device (cast to x’s compute dtype).

• bias – [Cout] tensor or None; accumulated in fp32.

• input_linear – [Nout, P64] int32 table; entries == Cin * Nin select the zero row.

• weight_linear – [Nout, P64] int32 table into the transposed weight.

• config – optional KernelConfig override.

Returns:

[B, Cout, Nout] tensor (transposed view of a contiguous [B, Nout, Cout] buffer, matching the baseline layer’s return convention).

fovi.arch.knn_cuda.grad_input(grad_y, weight, input_linear, weight_linear, cin, nin, config=None)

Input gradient: dx[b, c, m] = sum over (n, p) with input_linear[n, p] == c*Nin + m of (grad_y[:, :, n] @ W_n^T)[b, p], accumulated in fp32 via atomics.

Parameters:

• grad_y – [B, Cout, Nout] CUDA tensor, float16 or bfloat16, contiguous.

• weight – [Cout, Cin * V] tensor (cast to grad_y’s dtype for the gather).

• weight_linear (input_linear /) – [Nout, P64] int32 tables.

• cin – input geometry (iw = cin * nin addresses the pad-absorbing column).

• nin – input geometry (iw = cin * nin addresses the pad-absorbing column).

Returns:

[B, Cin, Nin] float32 gradient.

fovi.arch.knn_cuda.grad_weight(grad_y, x, input_linear, weight_linear, q, config=None, ksplit=None)

Weight gradient: dW[o, weight_linear[n, p]] += (A_n^T @ grad_y[:, :, n])[p, o], accumulated in fp32 via atomics into the transposed [Q, Cout] buffer.

Padding p entries gather the zero xt row and therefore add exact zeros to row 0, matching CompactTorchOps.grad_weight semantics.

Parameters:

• grad_y – [B, Cout, Nout] CUDA tensor, float16 or bfloat16, contiguous.

• x – [B, Cin, Nin] tensor (cast to grad_y’s dtype for the gather).

• weight_linear (input_linear /) – [Nout, P64] int32 tables.

• q – Cin * V (the weight’s flattened second dimension).

• ksplit – split-K-over-batch factor (None = heuristic, see default_grad_weight_ksplit(); 1 = unsplit behavior). Encoded in gridDim.y = p_tiles * ksplit — no separate kernel instantiation.

Returns:

[Cout, Q] float32 gradient.

fovi.arch.knn_cuda.backward_combined(grad_y, x, weight, input_linear, weight_linear, cin, nin, q)

Both gradients from ONE host-side staging build.

Semantically identical to grad_input(...) + grad_weight(...) — the same two kernels are launched with the same default configs — but the host path runs once: one pad computation, one gt re-layout (no cache round-trips), one xt/wt fetch, one stream lookup shared by both launches. The separate entries stay for needs_input_grad edge cases (the autograd Function routes there when only one gradient is required).

Returns:

(dx [B, Cin, Nin] fp32, dW [Cout, Q] fp32).

class fovi.arch.knn_cuda.CudaOps

Bases: object

Registry ops backed by the native CUDA kernels (forward + fused-atomic backward).

fp32 compute (no AMP) delegates to the torch compact oracle: the WMMA kernels are fp16/bf16 only, and dense cuBLAS serves fp32 well. Gradients are computed fully fused (fp16/bf16 operand staging -> fp32 WMMA accumulators -> fp32 global atomics) with no reduced-precision intermediates, so pre-scaled AMP gradients are strictly more overflow-robust here than in the baseline autograd.

name = 'cuda'

static forward(meta, x, weight, bias)

static grad_input(meta, grad_y, weight)

static grad_weight(meta, grad_y, x)

static backward_combined(meta, grad_y, x, weight)

Both grads from one staging build.

KNNConvFunction calls this automatically when BOTH input and weight grads are needed; the split grad_input/grad_weight entries above continue to serve the needs_input_grad edge cases. Same kernels, same configs, same math — the saving is host-path only (one prep instead of two op dispatches).

class fovi.arch.knn_cuda.KernelConfig(bm, bn, bk, rb, wm, wn, ksteps, async_copy, vec)

Bases: tuple

async_copy

Alias for field number 7

bk

Alias for field number 2

bm

Alias for field number 0

bn

Alias for field number 1

ksteps

Alias for field number 6

rb

Alias for field number 3

vec

Alias for field number 8

wm

Alias for field number 4

wn

Alias for field number 5

fovi.arch.knn_cuda.default_config(batch, cout, p64, dtype=torch.float16, nout=None)

Heuristic geometry from the tuning sweep.

Large batches want ~128 batch rows per CTA (weight-stream reuse saturates there) with a 128-wide cout tile; small batches want the minimal 16-row tile with synchronous staging (cp.async overhead does not amortize with a single k-pipeline in flight per CTA).

A small-Nout occupancy fallback (shrinking bm/bn to raise CTA count) was REFUTED by measurement: even at Nout=16, B=512 (128-CTA grid on 188 SMs) the big tiles win — per-CTA staging efficiency beats residency. nout is accepted for signature stability but no longer alters the choice.

fovi.arch.knn_cuda.default_grad_input_config(batch, cout)

dx kernel geometry: batch-tiled like the forward (rows/CTA law), p-tile fixed at 64.

fovi.arch.knn_cuda.default_grad_weight_config(batch, cout, p64)

dW kernel geometry: p-tile 128 when P64 allows, cout-tile 128 for wide layers.

fovi.arch.knn_cuda.default_grad_weight_ksplit(cfg, bpad, cpad, p64, nout, device_index)

Split-K-over-batch factor for the dW kernel.

At small Nout the (o-tile, p-tile, node) grid cannot fill the GPU (the dW smem footprint caps residency at ~1 CTA/SM for the 128x128 geometry), so idle SMs are bought with extra contraction slices; each slice adds one [BM, BN] fp32 atomic pass, which is cheap for the high-Cin mid-tier shapes (tens of MiB) but NOT for wide-Q alexnet shapes — hence the grid-based gate rather than an unconditional split. Splitting below one BK batch step is impossible (ksteps caps the factor).

fovi.arch.knn_cuda.clear_kernel_cache()

Drop compiled-module references and derived-tensor caches (CuPy’s on-disk cache is unaffected).