Optimize convolution and groupnorm - #98
Merged
Merged
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. |
michaelmckinsey1
requested changes
Aug 12, 2026
Collaborator
There was a problem hiding this comment.
Status of pytest on this branch. Do you also get these two failures?
=============================================================================== short test summary info ================================================================================
FAILED tests/test_groupnorm.py::test_gpu_triton_dctensor_matches_eager_and_stays_wrapped[None] - AssertionError: Triton path was not taken (output not NDHWC)
FAILED tests/test_groupnorm.py::test_gpu_triton_dctensor_matches_eager_and_stays_wrapped[relu] - AssertionError: Triton path was not taken (output not NDHWC)
==================================================== 2 failed, 747 passed, 8 skipped, 1 xfailed, 646 warnings in 724.16s (0:12:04) =====================================================
- We should add "mi300a" to the
triton_conv3dpackage nametriton_conv3d_mi300a, so it is more clear this is architecture specific kernel implementations. - If we are expecting to keep these changes around I think it's worth having claude take another pass on cleaning up irrelevant details from the work history. These are embedded in docstrings/file descriptions and make the descriptions very convoluted. I saw there was already a pass to remove testing file paths that aren't committed. The in-line comments actually aren't too bad on this one compared to the previous PRs.
a. This includes experimental details that may likely become irrelevant soon. I think we would benefit from writing analytical descriptions of why this improves performance instead. I think claude is capable of making this change.
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.
A run gives no sign of whether the Triton rungs actually served it. Print one line per ladder on rank 0 at startup, naming Triton against everything PyTorch does -- compiled and eager together, since from outside the ladder those are the same answer. Placement is the whole of the problem. ``_triton_ok`` is a latch set when a rung first answers a call, so reporting at construction time would say "Native" about modules that have not run yet. Reporting at the end of warmup is right when there is a warmup, but ``warmup_batches <= 0`` returns before running anything, which is what the benchmark drivers configure -- so a single call site would drop the line on exactly the runs most likely to want it. Both sites call it behind a one-shot flag. The reporter is duck-typed on the latch rather than on isinstance: _rungs is imported by the modules it would otherwise have to import back, and a ladder added later is then reported without touching this code. Labels come from a _rung_label class attribute and fall back to the class name, so a ladder that does not declare one still appears. Deliberately one rank's answer and not a collective: ranks latch independently and a rung can still fall back later, so this is informational rather than a contract, and gathering it would put a barrier on a path with no other reason for one.
Reviewers asked that the comments on this branch be more concise and that specific performance numbers be removed from comments describing changes, since those numbers go stale. Comments and docstrings only: every file's AST is unchanged apart from docstrings, ruff format and check are clean, and the test runner's shell script differs only in comment lines. Gone: timings, speedups, roofline percentages, memory figures, cell tallies, dated history of what the code used to do, and pointers to the untracked scratch tree. Kept: contracts, traps, and the hardware and algorithmic constants the code depends on (buffer-op and int32 limits, tile geometry, tolerances, pinned test shapes). Where a rationale would be empty without its figure, the comment names the bench tool that regenerates it. Across the 30 files: 11.9k comment and docstring lines become 10.6k, and 703 lines carrying a measurement, date, or scratch-tree pointer become 65, each of which is a constraint rather than a measurement.
ndryden
force-pushed
the
triton-kernels
branch
from
September 9, 2026 02:48
480de78 to
2131504
Compare
ScaFFold.unet.conv3d imports the sibling package, but package discovery only listed ScaFFold*, so a wheel shipped an adapter with nothing behind it. Discovery is namespace-aware, so the test tree (and the recorded corpora it holds) is excluded by name; the suite runs from a checkout. Verified by building a wheel, installing it into an empty target and importing both packages from there.
Nothing on the runtime path reads scaffold_corpus.json or scaffold_census.json; only bench/ and the test suite do, and both run from a checkout. Move them under triton_conv3d/tests/data, point the loaders in shapes.py there, and say what each file is a record of: the corpus joins model-analysis shape dumps with MIOpen profiles of the same configurations, the census is an instrumented training run.
ScaFFold always fuses the ReLU; the default exists because the function is a stand-alone replacement for F.group_norm.
_triton_kernel_failures() resolved to () when triton.errors could not be imported, and `except ()` catches nothing, so on a Triton without that module every handler in the ladder became a bare try and the first CompilationError or OutOfResources took the rank down -- the inverse of the ladder's contract, silently. Resolve the 2.x roots (CompilationError, OutOfResources) when triton.errors is absent, and have _routing_declines refuse the rung, with one warning, when nothing at all resolved.
Every input to the Triton rung's routing verdict is a fact about one rank -- the thin-shard check in _halo_plan, the device, the layout, the shape gates, the latch -- while the halo exchange is a collective. A rank declining alone went through DistConv's exchange while its neighbour posted receives against this module's, and the mesh hung. Unreachable with the shipped shapes, where D divides evenly, but a hang. The global extent is not on a DCTensor, so the check cannot be derived locally; instead, for any call some rank may exchange on, whether to exchange through this module is the MIN of every rank's verdict, agreed once per module instance and cached. The latch, and any later change in a rank's own verdict, only choose the kernel run on the exchanged tensor, which communicates nothing. A shard that later becomes too thin for a plan on an agreed module is an error rather than a lone trip to MIOpen. Adds the first test of the real exchange: two ranks under torchrun, the sharded output checked against F.conv3d on the whole volume, then a local decline, a latch and a thin shard, each ending with both ranks on the same side of the exchange.
make_corpus.py traces the three profiled configurations with model-analysis/unet_shapes.py, the pure-stdlib shape calculator the corpus already cited as its source, and joins profile_points.json, the MIOpen profile record of the same runs reduced to its convolution kernels. Both were in the untracked scratch tree; the corpus was not regenerable from the repository. `--check` confirms the committed file is reproduced byte for byte.
conv_census.py wraps FastConv3d, the two autograd nodes and the six kernel entry points around a real `scaffold benchmark` run and writes one capture per rank; make_census.py folds the four captures under census/ into the committed file. Both were in the untracked scratch tree, where the instrument ran through a scratch training driver. Verified by running the new instrument at configuration A: its capture records the same modules, autograd nodes and kernel calls as the committed one on every field except the resolver's answers on padded backward-weight problems, which the tuning tables have moved since and the census does not read. `--check` confirms the committed census is reproduced from the captures. The census's stale `tuned_declined` field is dropped from the instrument: the decline it described no longer exists. A README in tests/data says what each file is and how to regenerate it.
michaelmckinsey1
approved these changes
Sep 16, 2026
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.