fix: npu_rms_norm does not support gamma=None, use all-ones tensor in…

[xucqX]

fixes: #14590

What does this PR do?

Fixes a crash in RMSNorm.forward on Ascend NPU when the layer is built with
elementwise_affine=False (i.e. self.weight is None).

Context

RMSNorm supports elementwise_affine=False, which sets self.weight = None
(no learnable gamma, mathematically gamma == 1). The non-NPU else branch
handles this fine, but the NPU branch explicitly calls the fused CANN op

torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.eps)

and passes self.weight straight through. npu_rms_norm does not accept
gamma=None (it requires a real tensor), so any model that uses weightless
RMSNorm blocks crashes on Ascend the moment forward hits one of them.

Concrete in-repo trigger: LTX2VideoTransformerBlock
(src/diffusers/models/transformers/transformer_ltx2.py) constructs its block
norms (norm1/norm2/norm3, the audio variants, and the a2v/v2a cross-attn
norms) with elementwise_affine=False, so LTX-2 is unusable on NPU without
this fix.

The fix

When self.weight is None, substitute an all-ones tensor of the same shape,
dtype and device:

weight = self.weight
if weight is None:
    weight = torch.ones(self.dim, device=hidden_states.device, dtype=hidden_states.dtype)
if weight.dtype in [torch.float16, torch.bfloat16]:
    hidden_states = hidden_states.to(weight.dtype)
hidden_states = torch_npu.npu_rms_norm(hidden_states, weight, epsilon=self.eps)[0]

Multiplying by an all-ones tensor is identical to "no affine" (gamma=1), so
this preserves the original semantics while giving the CANN op a valid tensor.
The CPU/GPU else branch and bias handling are unchanged.

No new dependencies.

Self-review notes (AI-assisted)

• Math equivalence: ones ≡ gamma=1 ≡ the "no affine" semantics the
  else branch already implements for the weight=None case. No behaviour
  change on CPU/GPU.
• Alternative considered: fall back to the else (decomposed) branch on
  NPU when weight is None. Rejected — the explicit NPU branch exists
  precisely to get the fused CANN kernel in eager mode; the ones-shim keeps
  the fast path rather than degrading to a multi-launch decomposition.
• dtype/device correctness: the substitute weight is created with
  hidden_states.device / hidden_states.dtype, so the subsequent
  .to(weight.dtype) is a no-op for the None case and does not silently
  promote/demote precision.
• Tests not added: the NPU path only runs when is_torch_npu_available()
  is true, which is not the case in HF CI. Mocking the guard would not
  exercise the real CANN op and would be fragile. Happy to add a
  device-agnostic test (assert the weight=None path matches a reference
  RMSNorm with gamma=1 on CPU) if reviewers want one.
• Intentionally did not touch the else branch or bias handling.

Before submitting

• Did you use an AI agent (Claude Code) to help with this PR?
  • Read the Coding with AI agents guide.
  • Ran self-review on the diff (notes above).
  • Shared final self-review notes above.
• Read the contributor guideline.
• Read the philosophy doc (change is small and localized, bugfix only).
• Discussed via issue/forum: none yet — opening this PR as the discussion. Can file a tracking issue if preferred.
• Documentation: no doc changes (internal forward path, no public API change).
• Tests: not added, see self-review notes.

Who can review?

NPU doesn't have a dedicated owner in the list; this is a model/normalization
change, so tagging:

@yiyixuxu @dg845

[github-actions]

Hi @xucqX, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[xucqX]

fixes: #14590

[xucqX]

Follow-up to commit 5d489e0: while resolving the merge conflict with upstream/main, the upstream side of the conflict added and self.weight is not None to the NPU branch guard. That guard silently reverted this PR's fix — weight=None (i.e. elementwise_affine=False) was pushed back onto the pure-PyTorch else path, losing the fused CANN kernel that this PR exists to provide.

Commit 46c7553 restores the intended behaviour: on NPU, when self.weight is None, we substitute an all-ones tensor so npu_rms_norm still gets a valid gamma and stays on the fast path. The upstream else branch is otherwise unchanged.