Optimize convolution and groupnorm by ndryden · Pull Request #98 · LBANN/ScaFFold

ndryden · 2026-08-05T18:23:23Z

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:

┌────────────────────┬────────────────────────┬────────────────────┬────────┐
│       config       │         Triton         │       MIOpen       │ ratio  │                                                       
├────────────────────┼────────────────────────┼────────────────────┼────────┤
│ A — scale 7, 1 GPU │ 74.19 ms ±0.57 (5×22)  │ 94.38 ms (stored)  │ 1.273x │                                                       
├────────────────────┼────────────────────────┼────────────────────┼────────┤
│ B — scale 8, 1 GPU │ 457.7 ms ±0.41 (5×22)  │ does not run       │ —      │                                                       
├────────────────────┼────────────────────────┼────────────────────┼────────┤
│ C — scale 8, 2 GPU │ 280.8 ms ±2.58 (8×20)  │ 53,972 ms (stored) │ 192.6x │                                                       
├────────────────────┼────────────────────────┼────────────────────┼────────┤
│ D — scale 8, 4 GPU │ 180.2 ms ±1.33 (10×20) │ 222.8 ms (stored)  │ 1.231x │                                                      
└────────────────────┴────────────────────────┴────────────────────┴────────┘

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=0 and convolution with SCAFFOLD_CONV_TRITON=0.

Tagging @tbennun as a reviewer for the Triton code.

Code by Claude.

michaelmckinsey1

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?

ndryden · 2026-08-05T22:26:04Z

@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.

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.

michaelmckinsey1

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_conv3d package name triton_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.

michaelmckinsey1 · 2026-08-12T20:11:13Z

+    silently produced one at 1 shard and ``None`` at 2, the four sites would go
+    back to MIOpen on every multi-GPU run and nothing would say so.
+    """
+    x = torch.empty(1, 8, 8, 8, 8)


guard needed for _dc (we do this elsewhere)

Suggested change

x = torch.empty(1, 8, 8, 8, 8)

import torch.distributed as dist

distconv = pytest.importorskip("distconv")

x = torch.empty(1, 8, 8, 8, 8)

michaelmckinsey1 · 2026-08-12T20:35:30Z

+    could have failed inside one -- and the predicate would have declined the
+    rung long before.
+    """
+    global _TRITON_KERNEL_FAILURES


_triton_kernel_failures() returns () when the import fails, and except () catches nothing — so every try: ... except _triton_kernel_failures() in this module becomes a bare try, and the first CompilationError/OutOfResources propagates out of forward and takes the rank down. That's the inverse of this ladder's contract, and it's silent.

Claude is saying drop try/except's and let the ImportError fail instead of silently passing. I am a proponent of the hard failure with this much code.

michaelmckinsey1 · 2026-08-12T20:39:01Z

+            return None
+        if halo == 0:  # k == 1: no neighbour voxel is ever read
+            continue
+        if int(x.shape[dim]) < 2 * halo:


"
if int(x.shape[dim]) < 2 * halo: return None is a local decision gating a collective. With a ragged split axis a thin shard declines and routes to MIOpen while its neighbour posts irecv against it and blocks. Not reachable with the shipped shapes (D divides evenly), but the failure is a hang, and strategy.num_shards plus the global extent are both in hand here — derive it globally, or assert that every rank agrees. Same class of thing at :933 and :954, where the platform verdict and the layout test also decide per rank whether this rank exchanges.
" - opus

michaelmckinsey1 · 2026-08-12T20:44:16Z

For triton_conv3d sibling package, pyproject.toml must be updated for the install

[tool.setuptools.packages.find] include = ["ScaFFold*", "triton_conv3d*"] [tool.setuptools.package-data] ScaFFold = [ "package_data/weights_ins145.csv", "configs/*", ] triton_conv3d = [ "scaffold_corpus.json", "scaffold_census.json", ]

michaelmckinsey1 · 2026-08-12T21:37:09Z

@@ -0,0 +1,4871 @@
+{
+ "source": "model-analysis/unet_shapes.py",


The generator for this and triton_conv3d/scaffold_census.json are not committed?

michaelmckinsey1 · 2026-08-12T22:10:54Z

Would it be more prudent to name this triton_group_norm_MI300A.py for the files that are tuned for MI300A only?

michaelmckinsey1 · 2026-08-12T22:21:05Z



-def _group_norm(num_groups, num_channels):
+def _group_norm(num_groups, num_channels, activation=None):


Is this not always activation="relu"?

michaelmckinsey1 · 2026-08-12T22:21:59Z

I think we rename this package to triton_conv3d_mi300a, since this is so specific.

michaelmckinsey1 · 2026-08-12T22:28:43Z

+command.
+
+The transposed benchmarks used to live in a separate driver
+(``work/triton-conv/bin/m5_convT_bench.py``, deleted at this commit), on the


Is it possible to have fable take a pass on cleaning up irrelevant information, like a file being deleted at a certain commit. It just adds bloat.

michaelmckinsey1 · 2026-08-12T22:29:26Z

+kernel compiles a different ``PADDED`` body either way.  (``bwd_weight_config``
+also used to change its answer on it; that clause went on 2026-08-05, and the
+forms are still three different measurements without it.)  ``--form`` chooses


more detail related to a specific date that is irrelevant

ndryden requested review from PatrickRMiles, michaelmckinsey1 and tbennun August 5, 2026 18:23

michaelmckinsey1 requested changes Aug 5, 2026

View reviewed changes

Base automatically changed from round2-fixes to main August 5, 2026 22:27

ndryden added 6 commits August 5, 2026 15:27

Add helper script to run tests.

77cf50d

ndryden force-pushed the triton-kernels branch from cbd7baf to 77cf50d Compare August 5, 2026 22:27

ndryden added 2 commits August 5, 2026 15:52

michaelmckinsey1 mentioned this pull request Aug 6, 2026

ROCm/7.2 #102

Merged

4 tasks

michaelmckinsey1 requested changes Aug 12, 2026

View reviewed changes

Navigation Menu

Uh oh!

Optimize convolution and groupnorm - #98

Optimize convolution and groupnorm#98
ndryden wants to merge 8 commits into
mainfrom
triton-kernels

ndryden commented Aug 5, 2026

Uh oh!

michaelmckinsey1 left a comment

Uh oh!

ndryden commented Aug 5, 2026

Uh oh!

michaelmckinsey1 left a comment •

edited

Loading

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

michaelmckinsey1 Aug 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -0,0 +1,4871 @@
		{
		"source": "model-analysis/unet_shapes.py",



		def _group_norm(num_groups, num_channels):
		def _group_norm(num_groups, num_channels, activation=None):

Uh oh!

Conversation

ndryden commented Aug 5, 2026

Uh oh!

michaelmckinsey1 left a comment

Choose a reason for hiding this comment

Uh oh!

ndryden commented Aug 5, 2026

Uh oh!

michaelmckinsey1 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

michaelmckinsey1 left a comment •

edited

Loading