Optimize convolution and groupnorm - #98
Open
ndryden wants to merge 7 commits into
Open
Conversation
michaelmckinsey1
requested changes
Aug 5, 2026
michaelmckinsey1
left a comment
michaelmckinsey1
left a comment
Collaborator
There was a problem hiding this comment.
Can you add
export SCAFFOLD_GROUPNORM_TRITON=1
export SCAFFOLD_CONV_TRITON=1
to scripts/scaffold-tuolumne-torchpypi.job so we make sure we are running with this?
Collaborator
Author
|
@michaelmckinsey1 They should be enabled by default whenever safe. |
ScaFFold's GroupNorm was the only operator forcing a layout change: ATen's kernel launches one workgroup per (batch, group) row -- 8 of them at the benchmark's defaults -- so on a 228-CU MI300A it ran at a small fraction of achievable bandwidth, and it was the sole reason channels-last broke. This is an NDHWC-native kernel: Welford statistics, fp32/bf16/fp16 with torch's autocast contract, an int64 tile-base path for tensors past 2^31 elements, and a fused ReLU that is bit-exact against the unfused form. Registered as a custom op with fake kernels and register_autograd, so it traces under torch.compile(fullgraph=True) and composes with DistConv's DCTensor. Determinism is a contract here, not an accident: no float atomics, and the grid, split, tile and reduction order are pure functions of the shape. Verified run-to-run and across interpreters. Nothing imports this yet -- the wiring is the next commit.
FastGroupNorm becomes a three-rung ladder -- the native Triton kernel, a torch.compile'd functional, and stock eager -- with the routing decision made per call and the rejections tested in order: an explicit opt-out, a rung that has already failed in this process, an active torch.func transform, a non-CUDA tensor, a tensor subclass other than DistConv's DCTensor, a GPU the launch tables were not tuned on, and anything the kernel's own is_supported rejects. The ladder primitives live in _rungs.py, which the convolution ladder will share. The hardware guard is a preference rather than a correctness condition -- the kernel is correct anywhere Triton lowers it, and what is unknown elsewhere is only its speed -- so an explicit opt-in overrides it and nothing else on that list can be overridden at all. DoubleConv now asks the GroupNorm for its ReLU: the Triton kernel folds the activation into its forward store, removing a streaming pass worth 38% of the forward at the shapes that dominate the step. The nn.ReLU slots stay occupied by nn.Identity so nn.Sequential does not renumber its children and existing checkpoints keep loading.
A self-contained NDHWC implicit-GEMM convolution for MI300A / gfx942, built to replace MIOpen in this benchmark and to be upstreamable to DistConv on its own terms. It imports nothing from ScaFFold. Four kernels serve seven operator-directions. Backward-data is not a kernel at all -- it is the forward contraction on a permuted weight -- and the transposed operator's backward directions are its own forward kernel and backward-weight with the operands swapped. That reuse is the main structural result. The corpus is recorded from real ScaFFold calls, and it records the *shape form*: the logical convolution, the halo'd unpadded one DistConv hands a backend, and the padded one this package's adapter actually issues. Those are different problems, and conflating them has been the most expensive mistake in this work, so ConvProblem names which is which and the benchmark driver takes --form. Backward-weight is deterministic by default: split-K with a reduction tree whose split count, tile and order are pure functions of the shape. Correctness is checked against fp64 references under a three-tier tolerance policy, with a bitwise corpus for the exactly-representable cases. Also here: the benchmark harness, which times kernels through CUDA-graph replay with 95% intervals and an online iteration count, and which can run without the MIOpen control -- 98% of a two-arm capture's wall clock was MIOpen's find, not measurement.
FastConv3d and FastConvTranspose3d mirror FastGroupNorm: drop-in nn.Conv3d and nn.ConvTranspose3d, same parameters under the same names, no buffers, a rung ladder sharing _rungs.py, and MIOpen underneath everything the kernel declines. The transposed operator gets a factory of its own rather than a flag, because it is a different operator -- the weight's channel axes are the other way round and a different set of kernels sits behind it. The adapter performs the halo exchange itself, above autograd, rather than leaving it to the one DistConv does below. That has a consequence for the shape the kernel sees, and it is the thing to know when reading any number from this work: only the split axis is halo'd, so padding=1 survives on the other two and every k=3 convolution in the network is padded at every configuration -- unsharded there is nothing to halo at all. The benchmark corpus calls that the adapter form, and it is what production issues.
Every comment and docstring in the committed tree stood on its own path into the untracked work/ scratch directory: 70 references across 20 files, all of them dangling for anyone who clones this branch. The measurement each one supported is kept and stated as a result -- the number, the direction of the effect, and the "this was tried and it loses" warnings that stop a closed question being re-litigated -- while the path, the capture filename, the section number of an unshipped document and the blow-by-blow methodology go. Three pointers were dead even with work/ present and are simply gone: the transposed benchmark driver the conv_bench docstring narrated a refactor away from, a tuned table's named source capture, and the review commit SHA in the GroupNorm wiring tests. The corpus JSON keeps model-analysis/unet_shapes.py as its provenance; nothing parses either file's "source" field.
The style workflow runs `ruff format --diff .` and `ruff check .`, and both were failing: 22 files would be reformatted and there were 19 lint errors, all of the latter in the Triton package. Most of this is whitespace. Three fixes are not mechanical: `gemm_probe` built two closures over `a` and `b` in functions whose `finally` deletes both names, so each was correct only because the harness happens to call it before the cleanup runs. They now bind the tensors as default arguments, which captures at definition time and does not depend on call order. This is also what ruff was reporting as F821. `baseline._callable` assigned two lambdas to names, now plain functions. Six imports were unused and are gone. `ScaFFold/viz/standard_viz.py` is not part of this branch's work -- it fails the format check on round2-fixes too, and is reformatted here only so the check can pass. Both suites re-run afterwards, since several tests read their subject's source: triton_conv3d 1425 passed / 16 skipped, ScaFFold 742 passed / 8 skipped / 1 xfailed. Reformatting changes Triton's JIT cache key, so the first run after this recompiles every kernel and takes ~5x longer.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds Triton kernels for GroupNorm (to support channels-last) and convolution (to address some pathological cases, enable determinism, and generally speed things up). I expect to eventually shift a bunch of this code out of ScaFFold and to DistConv upstream, but these should unblock us.
Performance on Tuolumne:
The Triton code should be gated to MI300A in SPX mode, since it is not optimized or tuned for any other arch. Disable Triton GroupNorm with
SCAFFOLD_GROUPNORM_TRITON=0and convolution withSCAFFOLD_CONV_TRITON=0.Tagging @tbennun as a reviewer for the Triton code.
Code by Claude.