FP32LayerNorm makes two full-size float32 copies on every call #15221

Is your feature request related to a problem? Please describe.

FP32LayerNorm.forward computes

F.layer_norm(inputs.float(), self.normalized_shape, self.weight.float(), self.bias.float(), self.eps).to(origin_dtype)

For bf16 / fp16 inputs this makes a full-size float32 copy of the input and another of the output on every call. The CUDA layer norm kernel already computes in float32 for half inputs, so with half-precision weights the copies only add memory traffic: on an RTX 5090 the op takes 668 us for [1, 32760, 1536] bf16, of which the layer norm itself is ~131 us. In HunyuanDiT-v1.2 the copies are about 5.6% of the GPU time of a pipeline call.

Describe the solution you'd like.

On CUDA, when the weight / bias dtype matches the input, call F.layer_norm on the input directly. Output is bit-identical (fp16 / bf16, with and without affine), and a full pipeline call gets faster in eager mode: AuraFlow v0.3 +6.6-6.9%, HunyuanDiT-v1.2 +4.0-9.5% (RTX 5090 / Radeon AI PRO R9700). I have a PR ready and will link it here.

Describe alternatives you've considered.

Under torch.compile Inductor already fuses the casts, so compiled models are unchanged either way; only eager benefits.

Additional context.

The CPU layer norm kernel gives slightly different bits for half inputs than the float32 path, so CPU (and float32 weights with half inputs) keep the current code.