Fix GroupNorm NHWC one-pass backward execution

[Aidyn-A]

Disclaimer: This PR was vibe-coded with Codex GPT v5.5 model.

Description

Fix GroupNorm NHWC one-pass backward execution on GPUs where the number of activation tiles exceeds the maximum number of cooperatively resident blocks (like Thor).

For example, the failing configuration requires 128 blocks per activation slice, while the target GPU supports only 120 resident blocks. This previously triggered:

Assertion `blocks_per_slice <= max_blocks_per_grid' failed.

Solution

The cooperative grid is capped at the device’s residency limit. When additional activation tiles remain, resident blocks process them using a grid-stride loop.

Overflow tiles participate in the same reductions and are completed within the same kernel launch, so the backward implementation remains strictly one-pass. There is no fallback to the
two-pass implementation.

Configurations whose grids already fit retain their existing processing path.

Performance impact

For affected configurations, some blocks process an additional activation tile and reload that tile’s input after the global reduction. In the failing case 128 → 120 case, this applies to 8 of
128 tiles, adding approximately 6.25% to the x/dy read traffic and introducing minor work imbalance.

Unaffected configurations only incur inexpensive loop-condition checks.

Validation

Tested locally on Thor machine.

cc @crcrpar

[crcrpar]

Excuse me for the delay, I think this targets data center grade devices so a guard on the device prop could be better

[Aidyn-A]

Whaaat did the bot just do? 🫨

[crcrpar]

Apparently most of the changes are identical to #2022. I'm speculating merging this per-commit hook would let them go vanish

[crcrpar]

I think there may be one SM80 gap here: SM80 is included in the capped-grid setup and in the later overflow dx store path, but not in the first extra-tile accumulation guard.

That first loop contributes mean_1, mean_2, dgamma, and dbeta. On an A100, overflow tiles are therefore omitted from the parameter-gradient reductions and dx uses incomplete means.

I reproduced this on an A100 PCIe (SM80, 108 SMs) with two-pass forward and forced one-pass backward at N=1, C=128, G=16, H=1024, W=2048. Against an fp64 PyTorch reference, dx was within tolerance, while dgamma and dbeta differed by 376.16 and 339.40 respectively.

Could we include SM80 in the first accumulation guard as well, to match the later store path?

#if __CUDA_ARCH__ / 100 == 8 || __CUDA_ARCH__ / 100 == 11 || __CUDA_ARCH__ / 100 == 12

For regression coverage, the large GroupNorm case currently follows the two-pass path, so it does not execute this one-pass backward loop. Here is the focused reproducer I used (after building with APEX_GROUP_NORM=1 python -m pip install -v --no-build-isolation -e .):

import torch
import torch.nn.functional as F
import group_norm_cuda

torch.manual_seed(0)
n, c, g, h, w = 1, 128, 16, 1024, 2048
eps = 1e-5
x = torch.randn(n, c, h, w, device="cuda", dtype=torch.float16)
x = x.to(memory_format=torch.channels_last)
weight = torch.randn(c, device="cuda", dtype=torch.float32)
bias = torch.randn(c, device="cuda", dtype=torch.float32)
dy = torch.randn_like(x)

# Match the large-HW dispatch while isolating the changed backward kernel.
_, sums = group_norm_cuda.forward(x, g, weight, bias, eps, 2, False)
dx, dweight, dbias = group_norm_cuda.backward(
    dy, sums, x, g, weight, bias, eps, 1, False
)

x_ref = x.double().detach().requires_grad_()
weight_ref = weight.double().detach().requires_grad_()
bias_ref = bias.double().detach().requires_grad_()
F.group_norm(x_ref, g, weight_ref, bias_ref, eps).backward(dy.double())

torch.testing.assert_close(dx, x_ref.grad.to(dx.dtype), atol=1e-2, rtol=1e-2)
torch.testing.assert_close(dweight, weight_ref.grad.to(dweight.dtype), atol=1e-2, rtol=1e-2)
torch.testing.assert_close(dbias, bias_ref.grad.to(dbias.dtype), atol=1e-2, rtol=1e-2)

The latter two assertions fail before the guard change and should pass after it.

[Aidyn-A]

Oh, indeed I missed that. Thanks for noticing!