[Bugfix] Handle GroupNorm autocast safely

[akshatvishu]

Summary

Make AITER GroupNorm follow PyTorch's FP32 autocast policy. Add dtype checks before the HIP kernel reads the input and affine parameter pointers.

Fixes #4780

Motivation

Hunyuan Image 3.0 calls GroupNorm under CUDA autocast with an FP16 input and BF16 weight and bias tensors. PyTorch registers GroupNorm as an FP32 autocast operation, but AITER calls its HIP kernel directly and does not pass through that dispatcher rule.

The AITER kernel selects its pointer type from the input dtype. Before this change, it treated the BF16 weight and bias storage as FP16 when the input was FP16. In the Hunyuan VAE, the AITER output was FP16 while the PyTorch output was FP32. The mean error was 1.0778 and the maximum error was 6.1785.

Changes

• Cast the input, weight, and bias to FP32 when CUDA autocast is active, then call the existing FP32 GroupNorm kernel with autocast disabled.
• Keep the existing FP16, BF16, and FP32 paths unchanged when autocast is disabled.
• Check that the input, output, weight, and bias have the same dtype before launching the kernel.
• Check that the scratch workspace uses FP32.
• Add the exact Hunyuan GroupNorm shape to the standalone test and include correctness and performance measurements.
• Use a BF16 tolerance for the existing BF16 benchmark.

Performance

Tests ran on one MI300X with gfx942, ROCm 7.2, and PyTorch 2.11. The autocast measurement includes the FP32 tensor conversions.

Configuration	PyTorch	AITER	Speedup	Correctness
Hunyuan autocast, shape (1, 128, 4, 1024, 1024)	15,613.13 us	2,722.52 us	5.73x	Passed with rtol=0.001, atol=0.01
Direct FP32, shape (1, 128, 4, 64, 64)	31.1 us	42.3 us	0.74x	Passed
Direct BF16, shape (1, 512, 3, 64, 64)	98.3 us	38.7 us	2.54x	Passed

The change does not alter the HIP GroupNorm arithmetic or memory access pattern. Bandwidth and roofline analysis do not apply to this correctness fix.

Testing

• Unit tests added or updated
• Performance benchmarks run
• Tested on MI250X
• Tested on MI300X

python3 op_tests/test_groupnorm.py --dtype float32 -b 32,1,128,4,64,64 --autocast-spatial-size 1024
python3 op_tests/test_groupnorm.py --dtype bfloat16 -b 32,1,512,3,64,64 --autocast-spatial-size 64
ruff check aiter/ops/groupnorm.py op_tests/test_groupnorm.py
ruff format --check aiter/ops/groupnorm.py op_tests/test_groupnorm.py

Documentation

• Docstrings updated
• User guide updated because no user guide change is needed
• Performance guide updated because no performance guide change is needed

Dependencies

• No new third party dependencies added

PyTorch compatibility question

AITER documentation currently lists PyTorch 2.0 and later as supported. The device_type argument for torch.is_autocast_enabled was added in PyTorch 2.4, while PyTorch 2.0 to 2.3 accept no arguments.

Maintainers, should GroupNorm preserve PyTorch 2.0 to 2.3 compatibility by using torch.is_autocast_enabled() without an argument, or is the current AITER baseline PyTorch 2.4 or later?

Breaking changes

A direct mixed dtype call without autocast now stops with a clear error instead of launching a kernel that interprets the affine parameter storage using the input dtype.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4779 --add-label <label>