[AlexanderSoftCode]

Context

PR #10 left the bigger conv kernel rewrite waiting on an RTX 3050 knockout. That knockout showed the tensor cores busy only ~20 % of the time; the rest went to scalar 2-byte tile staging and half-empty 64-wide tiles on 32-channel layers. The PR follows the plan from that knockout, one commit per step: d920273 (16-byte staging), d89b37e (per-shape block tile), a665463 (dinput weight tile), 0646072 (fp32 forward input). fdbc9ac is a test-only fix to a dW/dB tolerance. No commit is mechanical.

RTX 3050, benchmarks/cifar10_breakdown.py (synthetic data, batch 256):

	ms/step
main	42.69
+ d920273	33.90	−20.6 %
+ d89b37e	31.28	−7.7 %
+ a665463	30.04	−4.0 %
+ 0646072	29.78	−0.9 % (within noise)

Total −30 %. Conv 03 (the largest layer) forward 3.27 → ~1.1 ms, backward 6.55 → 2.3 ms.

16-byte staging (d920273)

Every element of every A/B tile was its own 2-byte global load, LDS store and div/mod index decomposition, 16 k-deep per iteration. In the conv 11 knockout the gather alone was 36–48 % of each kernel. When the channel count is a multiple of 8, an 8-aligned k (or n) chunk stays inside one channel run, so a lane now moves 8 channels as one uint4: forward A and B, dweight A and B, and dinput A. In forward, each A lane's pixel (s, h, w) is decomposed once before the K loop. In dweight, the bias is summed in 8 per-lane registers and flushed with 8 atomicAdds. An #if on C % 8 picks the path per shape, so C_in = 3 keeps the scalar loop. The shared tiles are alignas(32), and _forward_gpu copies an input that is non-contiguous or not 16-byte aligned.

render_*_source(shape_meta, vectorize=False) renders the scalar staging, so tests can compare the two paths directly.

Non-contiguous forward inputs

On main, _forward_gpu passed inputs.astype(float16, copy=False) straight to the kernel. astype keeps the source's memory order, so a transposed or otherwise non-contiguous view reached a kernel that assumes C-contiguous NHWC, and the output was silently wrong. The new test gets 2048 / 2048 mismatched outputs on main. xp.ascontiguousarray fixes it, and the same step adds the 16-byte alignment check.

Per-shape block tile (d89b37e)

The block tile was 64 × 64 for every shape. On a 32-channel GEMM N axis (C_out for forward/dweight, C_in for dinput), half its warps had no columns, and all 512 threads still staged the half-zero tile. choose_block_tile_n() now returns 32 for N ≤ 32 channels. dweight's M axis (K_TOTAL) takes a 32-row tile up to 32 filter indices (conv 00, K_TOTAL = 27: 0.45 → 0.30 ms). The tiles depend only on the shape, so the kernel cache keys don't change. The block size follows the tile, and choose_dweight_split_k already scales its slot count by it.

Conv 03 with the 32-wide tile: forward 1.21 → 0.74 ms, dweight 1.29 → 0.77 ms, dinput 2.36 → 1.44 ms (interleaved A/B in one process).

dinput weight tile (a665463)

After d920273, dinput's remaining cost was its scalar weight gather. The weights are contiguous along c_out, which is the GEMM k axis in dinput. So the tile is now staged transposed (Bs[n][k]), 8 c_out per lane, and read back with a col-major B fragment (new frag_layout_b_col; dweight's A operand uses the same trick). A and B vectorize under the same C_OUT % 8 condition. dinput kernel time roughly halved again (conv 03: 1.44 → 0.78 ms).

fp32 forward input (0646072)

conv_forward_wmma takes half or float input (a template type, now part of the forward cache key, see forward_cache_key). float is converted as it's gathered, rounding exactly as astype(float16) does. The plan expected this to win on every conv, but that was measured with the scalar prototype. With 16-byte staging, each element is read about F_H × F_W times, so doubling its bytes costs more than one cast pass at copy bandwidth: convs 03–16 take 3.38 ms with fp32 reads vs 3.27 ms with the cast. choose_forward_input_dtype() therefore reads fp32 directly only when the gather is scalar (C_in % 8 != 0; conv 00: 0.53 → 0.47 ms). Other shapes keep the cast.

Behavior changes

• Non-contiguous inputs to a GPU Conv2d forward now give correct output instead of silently wrong output.
• Forward rounds fp64 inputs to fp32 before fp16, matching what backward's cast already did. In rare cases the forward output can differ from main's by the double rounding.
• dW/dB use a different split-K slice count on 32-wide tiles, so their rounding changes (~1e-5 relative, same order as the existing run-to-run atomic noise). Forward and dX are bit-identical to main.
• None to seeding or serialization.

Testing

• python -m unittest discover -s tests -t .: 1140 tests, 1069 pass on NumPy and CuPy (69 skipped). The 2 errors are test_save_replaces_symlink_instead_of_following_it on both backends, a Windows symlink issue that also fails on main.
• test_conv_gpu_forward_view_inputs_match_contiguous: the transposed case fails on main (2048 / 2048 mismatched). The misaligned fp16 view case fails on this branch without the alignment guard (cudaErrorMisalignedAddress).
• New parity tests (they guard the new paths rather than catch old bugs, so they also pass on main where they apply):
  • vectorized vs scalar staging: forward and dX bit-exact, dW norm-relative (tests/ops/test_conv_vector_staging.py)
  • 32-wide vs 64-wide tile at C = 8/24/32 and K_TOTAL = 27 (tests/ops/test_conv_tile_choice.py)
  • fp32 and fp64 forward inputs vs fp16, on both gathers (test_conv_gpu_forward_input_dtypes_round_like_astype)
• test_conv_kernel_compile_fallback reads the forward cache key from forward_cache_key(), since that key now includes the input dtype.
• fdbc9ac: test_backward_without_input_grad_keeps_parameter_grads compares dW/dB against the tensor's scale (_assert_gemm_close, 1e-5) instead of a per-element rtol, which split-K atomic ordering could trip.
• GPU: everything above ran on an RTX 3050 6GB Laptop (sm_86) with CuPy 14.2. Laptop clocks drift ±10–20 % between runs, so kernel comparisons were interleaved A/B in one process.

Follow-ups (out of scope)

• HIP / gfx1201 is untested. The build of alignas, uint4, __floats2half2_rn and the rocWMMA col-major B fragment, and whether each step also wins on the RX 9070 XT, need a run .

• Next kernel lever. Knockout after this PR (RTX 3050, share of base time):

  conv 11	base	MMA	gather	barriers
  forward	0.90 ms	27 %	23 %	26 %
  dweight	0.93 ms	36 %	10 %	10 %
  dinput	1.19 ms	30 %	27 %	19 %