[AlexanderSoftCode]

Context

aether/custom_kernels/conv_kernel.py had grown to 948 lines: three CUDA/HIP templates plus their compilers, caches and launch planners. Adding another kernel there would have made it harder to follow. This PR splits it into a package with one module per kernel, then adds a fourth kernel that fuses the fp16 casts in Conv2d._backward_gpu. All timings come from the CIFAR-10 model in benchmarks/cifar10_breakdown.py (batch 256) on an RX 9070 XT (gfx1201).

There are two commits:

1. refactor(conv): split conv_kernel.py into a package, one module per kernel is mechanical. Git won't show it as a rename because the file was split four ways, so review it by comparing git show 848470c:aether/custom_kernels/conv_kernel.py with the new modules.
2. perf(conv): fuse the backward fp16 casts into one kernel.

Package layout

aether/custom_kernels/conv_kernel/ keeps the old import path. Its __init__.py re-exports everything conv.py and the tests use, so every call site and every mock.patch.object(conv_kernel, ...) is unchanged.

Module	Contents
common.py	ConvShapeMeta, ConvLaunchPlan, tile constants, CUDA/HIP vendor dicts, _COMPILE_OPTIONS
forward.py	conv_forward_wmma
dweight.py	conv_backward_dweight_wmma and its split-K sizing
dinput.py	conv_backward_dinput_wmma
cast.py	conv_backward_cast_fp16 (commit 2)

Commit 1 moves code verbatim. The only additions are imports, and the only removals are two section headers that no longer apply.

Fused backward fp16 casts

_backward_gpu built its fp16 operands with four CuPy calls per conv:

• astype(fp32), which does nothing because a cast trace showed dvalues always arrive in fp32;
• a clip to ±65504;
• astype(fp16) on dvalues;
• astype(fp16) on the saved inputs.

On the benchmark that's 5 clips and 10 fp32→fp16 copies per step.

conv_backward_cast_fp16 writes both operands in one launch:

• dvalues are clipped to the fp16 range. The clamp uses comparisons rather than fminf/fmaxf, so NaN passes through and ±inf clamp, exactly like xp.clip.
• inputs are cast unclipped, so dweight sees the same fp16 inputs the forward pass multiplied.

The layer passes both as contiguous fp32, which does nothing today, and keeps the old astype/clip chain as the fallback when the kernel fails to compile.

The four kernels now share _COMPILE_OPTIONS instead of a copy-pasted options block.

Behavior changes

• ROCm: the conv kernels no longer pass -I$ROCM_PATH/include (default /opt/rocm/include) to hiprtc.
  • hiprtc already searches its own install's include directory, which is where rocWMMA lives. All four kernels compile with ROCM_PATH and ROCM_HOME unset.
  • The -I was searched before hiprtc's defaults, so a stale ROCM_PATH, for example after a ROCm upgrade, could pull in mismatched headers.
  • If some setup's hiprtc can't find rocWMMA, the existing compile-failure path warns with [aether] and conv falls back to the vectorized path.
• Conv outputs and gradients don't change: the fused cast is bit-identical to the old astype/clip chain (fp32→fp16 round-to-nearest-even on both paths).

Testing

• python3 -m unittest discover tests: 1113 tests pass on NumPy and CuPy (54 skipped), in the ROCm devcontainer. The suite also passes with ROCM_PATH/ROCM_HOME unset.

• New tests/ops/test_conv_cast_kernel.py:

  • hand-written special values (±inf, ±70000, NaN, 0.1), including inputs of 70000 → inf, which shows inputs aren't clipped;
  • a bit-exact comparison against NumPy's clip + astype on lengths that end mid-block.

• New compile-failure case in test_conv_kernel_compile_fallback.py: a failed compile warns and returns None, and the real compiled kernel is restored for later tests.

• New test_conv_gpu_backward_parity_without_cast_kernel runs _backward_gpu's fallback branch with the kernel mocked out.

• I broke the kernel three ways and test_conv_cast_kernel failed each time:

  • fminf/fmaxf instead of comparisons;
  • clipping inputs too;
  • an off-by-one at the dvalues/inputs boundary.

• benchmarks/cifar10_breakdown.py --repeats 3:

  	main	+ refactor	+ fused cast
  model.train() ms/step	8.55	8.55 (noise)	8.31 (−2.8%)
  conv 00 backward ms	0.30	0.30	0.22 (−25%)
  conv 03 backward ms	1.28	1.28	1.21 (−5.5%)

  Every conv backward phase improved beyond run-to-run spread, and each conv's host time dropped by about 14 µs from the fewer launches.

• Not verified:

  • The CUDA build of the cast kernel (cuda_fp16.h) hasn't been compiled or run.
  • The bit-exact test assumes CuPy and NumPy both round fp32→fp16 to nearest-even, which I've confirmed only on gfx1201.
  • Dropping the -I was tested only on ROCm 7.14.

Follow-ups (out of scope)

• Loss.regularization_loss() makes about 28 no-op astype(float32) calls per step and loops over every trainable layer. That goes in a separate losses PR.
• The RTX 3050 mobile measurements for the split-K rule and the dweight knockout are still pending.