RMSNorm crashes on NPU when elementwise_affine=False (weight=None): npu_rms_norm requires a real gamma tensor

[xucqX]

Describe the bug

diffusers.models.normalization.RMSNorm crashes on Ascend NPU when constructed
with elementwise_affine=False (i.e. self.weight is None).

The NPU branch of RMSNorm.forward passes self.weight straight into the
fused CANN op torch_npu.npu_rms_norm, but that op's schema declares gamma
as a required Tensor (not Optional), so a None weight is invalid and the
forward errors out. The non-NPU else branch handles weight=None correctly.

Reproduction

import torch
from diffusers.models.normalization import RMSNorm

# LTX-2 style block norm: no learnable affine => weight is None
norm = RMSNorm(dim=4096, eps=1e-6, elementwise_affine=False).npu().to(torch.bfloat16)
x = torch.randn(1, 256, 4096, device="npu", dtype=torch.bfloat16)
out = norm(x)   # raises on NPU

(Note: this needs an actual Ascend NPU device; on a non-NPU build it fails earlier
with a backend-availability error rather than the gamma=None error.)

Logs

On Ascend NPU the call to `torch_npu.npu_rms_norm(hidden_states, None, eps)` violates
the op schema and the forward aborts. The op's schema (from `torch_npu`) is:


npu::npu_rms_norm(Tensor input, Tensor gamma, float epsilon=1e-6) -> (Tensor, Tensor)


`gamma` is a required `Tensor`, while `RMSNorm` can legitimately have
`self.weight = None` (when `elementwise_affine=False`, i.e. gamma == 1, no affine).

System Info

• 🤗 Diffusers version: 0.40.0.dev0
• Platform: Linux-5.15.0-119-generic-x86_64-with-glibc2.39
• Running on Google Colab?: No
• Python version: 3.11.15
• PyTorch version (GPU?): 2.9.0+cpu (False)
• Huggingface_hub version: 1.27.0
• Transformers version: 5.3.0
• Accelerate version: 1.10.1
• PEFT version: 0.18.0
• bitsandbytes version: 0.49.2
• optimum-quanto version: 0.2.7
• Safetensors version: 0.8.0
• xFormers version: not installed
• Accelerator: Ascend NPU (torch_npu 2.9.0.dev20260207)
• Using GPU in script?: Yes (Ascend NPU via torch_npu, is_torch_npu_available()==True)
• Using distributed or parallel set-up in script?: Yes (accelerate multi-NPU)

Who can help?

No response