Fix Ascend NPU RMSNorm and fused-attention mask shapes

[mengchengTang]

What does this PR do?

Fixes two Ascend NPU incompatibilities hit when running LTX-2 with
attn_backend=_native_npu:

1. RMSNorm with elementwise_affine=False leaves weight=None.
   torch_npu.npu_rms_norm requires a gamma tensor, so we only use the
   fused kernel when weight is not None, and otherwise fall back to the
   existing PyTorch path (same as CUDA).

2. Fused attention mask: Ascend FA does not broadcast a singleton
   query-length dim. Expand masks shaped [B, N, 1, Skv] (e.g. LTX
   cross-attn) to [B, N, Sq, Skv], generalizing the existing
   [B, 1, 1, Skv] handling.

Screenshots

Error 1 — RMSNorm / gamma is None

Error 2 — FA mask shape [1, 32, 1, 1024]

Fixes # (issue)

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[github-actions]

Hi @mengchengTang, 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.

[ErenAta16]

Ran both conditions over a spread of mask shapes. The [B, N, 1, Skv] case the PR is after does get picked up, and the added shape[-1] == key.shape[1] check is a genuine tightening — but two rows behave in ways the new comment doesn't cover.

query is [B, Sq, N, D] so query.shape[1] == Sq; B=2, N=4, Sq=7, Skv=11:

mask shape	before	after	result shape (after)
[B, 1, 1, Skv]	expand	expand	(2, 1, 7, 11)
[B, N, 1, Skv] — the target	skip	expand	(2, 4, 7, 11)
[B, 1, 1, Skv+3]	expand	skip	(2, 1, 1, 14)
[B, N, 1, Skv+3]	skip	skip	(2, 4, 1, 14)
[B, 99, 1, Skv] — nonsense head count	skip	expand	(2, 99, 7, 11)
[B, 1, Sq, Skv] — already correct	skip	skip	unchanged
[B+5, 1, 1, Skv] — wrong batch	expand	expand	(7, 1, 7, 11)

Row 3 is a behaviour change worth stating: a 4-D mask whose last dim doesn't match key.shape[1] used to be expanded anyway and now falls through to npu_fusion_attention as-is. That's more correct — silently expanding a mismatched mask was the worse option — but it moves where the failure surfaces from this function to the kernel.

Row 5 is the one I'd tighten. expand(-1, -1, Sq, -1) accepts any shape[1], so dropping shape[1] == 1 without replacing it means a mask with a head count that matches nothing gets reshaped and handed to the kernel. shape[1] in (1, query.shape[2]) would keep the intent ("1 or N") and reject the rest here rather than downstream.

Row 7 is pre-existing on both sides — the 2-D branch checks attn_mask.shape[0] == query.shape[0] and the 4-D branch never has. Not this PR's problem, just noting the asymmetry since the 4-D branch is being touched.

On the normalization.py half: the fix looks right, and self.bias can only be non-None when elementwise_affine=True (it's created inside that branch), so the if self.bias is not None left in the NPU path is consistent with self.weight is not None now gating it. Worth saying out loud in the PR body that NPU users with elementwise_affine=False now take the pure-PyTorch path rather than the fused kernel — that's the correct trade, but it's a silent perf change for anyone who was previously… well, crashing, so probably nobody. The else branch ends with hidden_states.to(input_dtype), so the dtype contract holds on the fallback.

Shape table produced with plain torch tensors on CPU (_maybe_modify_attn_mask_npu's reshaping is device-independent). I have no Ascend hardware, so I can't confirm what npu_fusion_attention does with the row-3 and row-5 masks it now receives — that's the part worth a check from someone with an NPU.