fovi.arch.knn_warp

Optional Warp kernels for compact KNN convolution inference and training.

Forward kernels plus adjoint kernels. Warp autodiff is never used (enable_backward=False everywhere); gradients are explicit adjoint kernels registered with the fovi.arch.knn_autograd ops registry as warp_train.

fovi.arch.knn_warp._make_uncached_r2_kernel(tile_m: int, tile_n: int)

Uncached kernel with a 2-way batch-tile weight-reuse loop.

Each CTA gathers every indexed weight tile once and applies it to two consecutive batch tiles, halving weight-stream traffic relative to one-accumulator kernels at equal tile_m.

fovi.arch.knn_warp._make_uncached_r4_kernel(tile_m: int, tile_n: int)

Uncached kernel with a 4-way batch-tile weight-reuse loop.

fovi.arch.knn_warp._canary_launch_ok(kernel, tile_m, tile_n, r, block_dim, device)

Return True when the kernel actually launches on device.

Kernel launches that exceed CTA resource limits fail with a CUDA “invalid argument” that Warp reports on stderr without raising, leaving the (uninitialized) output untouched. A NaN-prefilled single-node canary makes that failure mode detectable.

fovi.arch.knn_warp._get_uncached_batch_config(name, device=None)

Build (and canary-validate per device) a large-batch config; None when unlaunchable.

fovi.arch.knn_warp._select_uncached_batch_config(batch, out_channels, out_nodes, p_padded)

Heuristic for the large-batch (B >= 64) uncached kernel.

Fit to the Ada fp16 sweeps. With the coalesced weight operand the wide m64n128r2 tile wins or ties every measured shape at B >= 128 whenever Cout does not waste most of the 128-wide tile; narrower output tiles only pay off for small Cout. Below 128 rows the 2-way reuse kernels waste half their compute on out-of-bounds subtiles, so a single 64-row accumulator is the safe choice. out_nodes / p_padded are kept for future tuning (they decided earlier fits).

fovi.arch.knn_warp._pad_bias_to_tiles(bias, cout, tile_n)

Zero-pad bias so 1D tail tile loads never read past the buffer (tile-overread doctrine).

fovi.arch.knn_warp.run_cached(x, effective_weight, bias, input_linear, cout=None)

Cached-effective-weight fp16 forward.

cout is the true output-channel count; it defaults to effective_weight.shape[2] for backward compatibility. Pass it explicitly when the effective-weight cache is padded to a tile multiple (recommended: pad to a multiple of 64) so the weight tile loads never read past the buffer; unpadded caches with non-multiple Cout are padded here per call, which costs a full cache copy — pad the cache instead.

fovi.arch.knn_warp._make_grad_input_stage2_kernel(tile_p: int, tile_b: int, fp32_stage: bool)

v2 gi staging: one CTA per (n, b_tile), the P64 loop lives inside the CTA.

Amortizes per-CTA fixed costs over P64/tile_p accumulator rounds (the earlier v1 layout ran P64/64 x b_tiles CTAs of only Cout/64 matmul steps each).

fovi.arch.knn_warp._make_grad_weight_stage3_kernel(tile_p: int, tile_n: int, fp32_stage: bool)

v3 gw staging: v1’s 3D (n, p_tile, o_tile) grid with the v2 transposed-x operand.

For small-Nout shapes the v2 (n, o_tile) grid under-occupies the GPU; this keeps per-CTA work small but parallel while still using contiguous xt gathers and no tile_transpose.

fovi.arch.knn_warp._make_grad_weight_stage4_kernel(tile_p: int, tile_n: int, fp32_stage: bool)

v4 gw staging: v3’s 3D grid, but each CTA computes TWO adjacent o tiles.

Each gathered xt tile is reused for both output-column tiles, halving the x-gather stream (the forward’s reuse principle applied to the weight-gradient’s batch contraction). Requires g_nbo padded to a multiple of 2*tile_n columns.

fovi.arch.knn_warp._make_grad_weight_stage2_kernel(tile_p: int, tile_n: int, fp32_stage: bool)

v2 gw staging: one CTA per (n, o_tile), P64 loop in-CTA, transposed x operand.

Takes xt [Cin*Nin+1, Bpad] (batch contiguous, cross-pollinated from the CUDA track): the indexed gather then reads contiguous 128-byte rows AND directly yields the A^T tile, eliminating wp.tile_transpose entirely.

fovi.arch.knn_warp._canary_grad_stage_ok(kernel, tile_p, tile_wide, block_dim, version, weight_operand, device)

NaN-canary a grad staging kernel (same silent-launch trap as the forward).

fovi.arch.knn_warp._get_grad_kernel(kind, name, device)

kind in {‘gi’, ‘gw’, ‘csr’}; returns a per-device-validated kernel config or None.

fovi.arch.knn_warp._pad_flat_input_fp16(x, meta)

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

fovi.arch.knn_warp.grad_input(meta, grad_y, weight, config=None)

dx [B, Cin, Nin] fp32 for the compact operator; deterministic reverse-CSR reduction.

grad_y [B, Cout, Nout] and weight [Cout, Q] must be fp16 (the autograd Function guarantees contiguity and compute dtype).

fovi.arch.knn_warp.grad_weight(meta, grad_y, x, config=None)

dW [Cout, Q] fp32; per-node dWeff staged on GPU, finished with fp32 index_add_.

Mirrors CompactTorchOps.grad_weight semantics exactly: input pad entries gather the zero pad column (exact-zero contributions) and weight pad entries scatter zeros into row 0 of the [Q, Cout] accumulator.

class fovi.arch.knn_warp.WarpCompactOps

Bases: object

Forward-only ops object matching the fovi.arch.knn_autograd registry interface.

name = 'warp_compact'

static forward(meta, x, weight, bias)

static grad_input(meta, grad_y, weight)

static grad_weight(meta, grad_y, x)

class fovi.arch.knn_warp.WarpTrainOps

Bases: object

Train-capable Warp ops: forward + adjoint kernels (fp16 compute, fp32 grads).

name = 'warp_train'

static forward(meta, x, weight, bias)

static grad_input(meta, grad_y, weight)

static grad_weight(meta, grad_y, x)

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

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

fovi.arch.knn_warp.run_uncached(x, weight, bias, input_linear, weight_linear, config=None)

Uncached (per-call weight gather) fp16 forward.

config optionally forces a large-batch kernel by name (see _UNCACHED_BATCH_CONFIG_SPECS); by default a kernel is selected from the batch size.