Update kernels for conv_2d, batch_norm, pool_2d, concat, reshape, split, batch_matmul, transpose and upsample

[elliottslaughter]

Generated via Claude to update the remaining kernels used in YOLOv10.

My instructions to Claude were to explicitly follow the style of the element binary PR #1667 and to update the style, API design, and tests, while mostly leaving the implementations alone IF they were correct. Claude ended up finding several correctness bugs in the kernels, summarized in Claude's report.

Marked as a draft because it stacks on top of #1667.

Claude prompt:

Contents of human-written prompt provided to Claude

I'm working through an audit of the kernels (lib/kernels) in this repository. Initial work has been done in the branches element-binary and softmax. Now I want to go through and audit all of the remaining kernels used in the YOLOv10 model (lib/models/src/models/yolov10). (We can ignore any kernels not used in that model for now.)

Follow the style of the existing element-binary and softmax branches. Changes include both stylistic changes as well as API changes (e.g., to follow a consistent argument order and use newer APIs).

The kernels themselves can be left alone IF THEY ARE CORRECT. Their APIs may need updating but the implementations themselves can be left alone. Of course, if the kernels are NOT correct, they should be fixed.

If a kernel is not yet implemented for CPUs, leave it unimplemented. While I did implement element binary kernels, the focus right now is just on updating style, API design, correctness and testing of existing kernels.

All of the kernels should be tested. Existing tests may be inadequate, missing, or in the wrong location. In particular, I found that several pre-existing tests only test that ANY non-zero value is produced, while seeding inputs with random values. Instead we should follow the style of the new element binary tests: we should seed inputs with small, meaningful values, and we should check the actual outputs that are produced. If a tensor is written before it is read, we should initialize it with random values to ensure that none are used in computing the final result. Finally, we should cross-check the actual values against those computed by an equivalent PyTorch program.

So far I've been able to use bitwise identical equality checks, and those have worked. However, it is acceptable to use an approximate equality check if the floating point values do not match exactly as long as the epsilon is reasonable (e.g., 1e-20 or better). There is an accessors_within_epsilon function that can be used for this purpose.

Try to keep all changes within lib/kernels and prefer to avoid touching files already modified in the various outstanding branches (merged into mocha-2026 but not into master) so that the patches developed here are commutative. As an exception, it's ok to rely on element-binary and build on top of functions used in that branch specifically. Note that, for the purposes of testing, update-nix is mandatory for running on this machine, so we'll base work off of mocha-2026 to begin with. I'll take care of rebasing later.

Test should be run from the Nix shell:

NIXPKGS_ALLOW_UNFREE=1 nix develop .#gpu --accept-flake-config --impure
proj test

Note that proj test does a full build before testing so no manual build steps are required.

Some existing tests are currently failing and can be ignored:

• local-execution:LocalBackend e2e Training
• realm-execution:RealmBackend e2e Training Replicate Op (GPU Model Parallelism)
• realm-execution:RealmBackend e2e Training (GPU Model Parallelism)
• realm-execution:RealmBackend e2e Training Gradient Reduction Op (GPU Model Parallelism)

A PyTorch environment is available in /home/eslaught/flexflow/test-pytorch and can be used via uv:

cd /home/eslaught/flexflow/test-pytorch
uv run python

Local changes are ok. Do not modify any existing branches. New branches should be created with the prefix claude/. Do not push.

Claude's summary of the results:

Contents of Claude-generated summary of the results

Branch claude/kernel-audit-yolov10, based on mocha-2026. Full suite: 891 passed / 4 failed, and the 4 are exactly the pre-existing failures you listed. Kernel tests alone: 91 passed / 0 failed.

Kernels audited

All nine YOLOv10 kernels not already covered by element-binary/softmax: conv_2d, batch_norm, pool_2d, concat, reshape, split, batch_matmul, transpose, upsample.

Each moved out of FlexFlow::Kernels::<Op> into FlexFlow with <op>_{gpu_,cpu_,}{init,forward,backward,cleanup}_kernel names, accessors instead of raw pointers, attrs instead of exploded scalars, the (stream, handle, per_device_state, attrs, …) order, backward as (output, output_grad, input, input_grad), and PerDeviceFFHandle dropped from the state structs.

Correctness bugs found and fixed

• transpose — three separate bugs: forward did out += out*beta + in on an uninitialized output; backward applied the forward permutation instead of its inverse (invisible to the old 2-D self-inverse test) and doubled the existing gradient; strides were stride[i] = stride[i-1]*dim[i] instead of dim[i-1] (invisible with 10×10). The CPU dispatch also called cpu_forward_kernel from backward_kernel.
• upsample — CPU backward had input_grad.at(output_coord) += output_grad.at(input_coord), i.e. coordinates swapped.
• batch_norm — attrs.eps and attrs.momentum ignored (hardcoded CUDNN_BN_MIN_EPSILON / 1.0); runningVar init to 0 not 1; a runningMean parameter written before being read; init created and destroyed a Legion stream. The task impl read nonexistent SCALE/BIAS slots (op-attrs declares GAMMA/BETA) and passed an uninitialized float *.
• conv_2d — used cudnnFind*AlgorithmEx, which writes into the input, output and filter-grad buffers it's handed during init; switched to the cudnnGet*Algorithm_v7 heuristics so init only needs shapes.
• element_unary — SILU was unimplemented on GPU (assert(false)) despite YOLOv10 using it. Added it via the scalar path with beta defaulting to 1, matching the CPU kernel.
• reshape — forward used a synchronous cudaMemcpy that ignored the stream.
• create_4d_accessor_w_with_contents — wrote directly into the target allocator, so it couldn't be used with a CUDA allocator (the 1/2/3-D versions stage through CPU memory).
• batch_matmul — the .cu included upsample_kernels_gpu.h, and tensor_accessor_batch_matmul was declared with a different signature than its definition, so it could never have been linked against.
• split task impl only ever handled one output tensor; rewritten against the variadic output slots.

Silently-ignored fused activations on conv_2d/batch_norm/pool_2d now assert instead of producing wrong results. None are used by YOLOv10 and a correct backward needs scratch space, so I did not implement them — but note inception_v3 uses batch_norm with Activation::RELU, which was previously applied in backward but not in forward, so it was already wrong.

Newly implemented

upsample and batch_matmul had no GPU kernel at all, which meant YOLOv10 could not run on GPU; reshape, concat, split and transpose had no CPU kernel.

• upsample (GPU) — nearest-neighbour, NCHW. Forward maps each output element back to the input element it samples; backward iterates over input elements and sums the scale_factor² block of output gradients that read from each one, which keeps the accumulation deterministic and avoids atomics.
• batch_matmul (GPU) — cublasSgemmStridedBatched. Our tensors are row-major and cuBLAS is column-major, so the forward pass computes outputᵀ = input_rhsᵀ · input_lhsᵀ, which needs no transposes and no data movement because a row-major (r, c) matrix is bit-for-bit a column-major (c, r) matrix. Its kernels now also take a handle (for the cuBLAS handle) in the standard position.
• reshape (CPU) — forward is a copy; backward views input_grad with output_grad's shape (both contiguous with the same element count, so their elements line up) and accumulates elementwise.
• concat / split (CPU) — mirror images, walking each variadic tensor's coordinate set and offsetting the axis coordinate by a running total.
• transpose (CPU) — relabels coordinates with the same permutation permute_tensor_dims applies to the dims. Both directions iterate the input side, so unlike the GPU kernel backward needs the same permutation rather than its inverse.

All four CPU kernels are datatype-generic via DataTypeDispatch1 and validate their shapes up front.

Tests

Replaced the "random input, check anything non-zero" tests for batch_norm, concat, pool_2d, reshape, split, transpose; added first tests for conv_2d, for batch_matmul on GPU, and for upsample, reshape, concat, split and transpose on both CPU and GPU; fixed the element_unary GPU tests, which were calling the CPU kernels. All moved to lib/kernels/test/src/kernels/. Every value is cross-checked against an equivalent PyTorch program, bitwise exact except cuDNN's batch-norm data gradient and CUDA's expf (1e-6). Where an op has both a CPU and a GPU kernel the two tests share the same expected values, so they double as a differential check. Upsample is covered at scale factors 2 and 3, and transpose at a non-self-inverse 3-D permutation, so the cases that hid the original bugs are exercised.

Still missing on CPU

softmax (~35–40 lines), pool_2d (~50), batch_norm (~50) and conv_2d (~75) are still NOT_IMPLEMENTED on CPU. All are writable by hand with no new dependencies, but none fit in the budget the four above did — pool_2d's MAX backward needs to rescan each window for the argmax, batch_norm's backward formula is long and easy to get subtly wrong, and conv_2d needs a 7-deep forward loop plus three separate backward loops.

Two things for you to decide

1. Backward accumulate vs. overwrite is inconsistent across ops. Most kernels accumulate, but softmax and batch_matmul (both CPU and GPU) overwrite, and element_unary overwrites on CPU while mostly accumulating on GPU. I preserved each op's existing behaviour and made every kernel I wrote agree with its own counterpart on the other device rather than pick a side. Given fix-bwd-multi-use is adding explicit gradient-reduction ops, which way this should settle is your call.
2. No *_cleanup_kernel is called anywhere outside lib/kernels, so cuDNN descriptors and batch_norm's running-stats allocation leak. Pre-existing; out of scope here since the fix lives in task-spec/execution.

---

This change is