[Bug] GroupNorm does not follow FP32 autocast policy

[akshatvishu]

Problem

AITER GroupNorm does not follow PyTorch's FP32 autocast policy. Hunyuan Image 3.0 calls GroupNorm under CUDA autocast with an FP16 input and BF16 weight and bias tensors.

The AITER HIP kernel selects its pointer type from the input dtype. It therefore reads the BF16 weight and bias storage as FP16 when the input is FP16. AITER returns an FP16 output, while PyTorch returns an FP32 output.

In the Hunyuan VAE input with shape (1, 128, 4, 1024, 1024), the mean error was 1.0778 and the maximum error was 6.1785.

Source references

• The PyTorch AMP documentation lists group_norm among the CUDA operations that autocast to FP32.
• The PyTorch source includes group_norm in AT_FORALL_FP32 and registers that list with the FP32 policy for AutocastCUDA.
• The AITER Python wrapper calls the HIP binding directly, so the call does not pass through PyTorch's autocast dispatcher.
• The AITER HIP launcher selects its template type only from the input dtype. The selected type is then used for the input, weight, bias, and output pointers.

Reproduction

import torch

from aiter.ops.groupnorm import GroupNorm

norm = GroupNorm(
    32,
    128,
    eps=1e-6,
    affine=True,
    device="cuda",
    dtype=torch.bfloat16,
)
x = torch.randn(
    (1, 128, 4, 64, 64),
    device="cuda",
    dtype=torch.float16,
)

with torch.inference_mode(), torch.amp.autocast("cuda", dtype=torch.float16):
    expected = norm(x, use_torch=True)
    actual = norm(x)

print(expected.dtype)
print(actual.dtype)
print((actual.float() - expected).abs().mean())
print((actual.float() - expected).abs().max())

Expected behavior

AITER should cast the input, weight, and bias to FP32 while autocast is active. The AITER output should be FP32 and should match PyTorch GroupNorm.

Environment

• GPU: AMD MI300X with gfx942
• ROCm: 7.2
• PyTorch: 2.11
• AITER: 0.1.16.post3
• Consumer: vLLM Omni with Hunyuan Image 3.0