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). This is a silent perf change for that config; the previous path would crash, so it is unlikely anyone relied on it. The fallback still ends with hidden_states.to(input_dtype).

2. Fused attention mask: Ascend FA does not broadcast a singleton query-length dim. Expand masks shaped [B, 1|N, 1, Skv] (e.g. LTX cross-attn [B, N, 1, Skv]) to [B, 1|N, Sq, Skv], generalizing the existing [B, 1, 1, Skv] handling. The head dim is restricted to 1 or N (query.shape[2]) so unrelated head counts are not expanded.

Screenshots

Error 1 — RMSNorm / gamma is None

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

Fixes #14380

Testing

Environment

Item	Version / Configuration
OS	Linux aarch64
Hardware	Ascend NPU
Python	3.11
PyTorch	2.10.0
torch-npu	2.10.0
CANN	9.0.0
diffusers	main

Test results

Scope	Test	Before the fix	After the fix	Result
End-to-end	dg845/LTX-2.3-Diffusers with verl-omni LTX-2.3 FlowGRPO	Rollout or training stops at the RMSNorm or fused-attention mask error	Rollout and training proceed through both affected paths without error	✓ Passed
Minimal reproduction	RMSNorm(elementwise_affine=False) from #14380	npu_rms_norm receives weight=None and crashes	Falls back to the PyTorch RMSNorm path and completes successfully	✓ Passed
Minimal reproduction	Native NPU attention with mask [B, N, 1, Skv] from #14380	Ascend fused attention rejects the singleton query dimension	Mask is expanded to [B, N, Sq, Skv] and attention completes successfully	✓ Passed

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?

@yiyixuxu @sayakpaul — Ascend NPU / _native_npu attention + RMSNorm

[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.

[mengchengTang]

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.

Thanks for the careful shape table — very helpful.

Row 5: Agreed. I’ve tightened the 4-D branch to require attn_mask.shape[1] in (1, query.shape[2]), so we only expand head dim 1 (broadcast) or N, and reject unrelated head counts here instead of handing them to the kernel.

Row 3 / shape[-1] == key.shape[1]: I’ve dropped that check to avoid the behaviour change you flagged (mismatched last-dim masks falling through instead of being expanded). Expanding a wrong-Skv mask is still imperfect, but keeping the previous expand behaviour reduces surprise vs. silently moving the failure into npu_fusion_attention. Happy to revisit if maintainers prefer the stricter guard.

Row 7: Acknowledged — leaving the pre-existing 2-D vs 4-D batch-check asymmetry alone in this PR.

normalization.py: Thanks for confirming. I’ve updated the PR description to note that with elementwise_affine=False we take the pure-PyTorch path instead of the fused kernel, and that the fallback still restores dtype via hidden_states.to(input_dtype).

I’ve also verified the updated path on Ascend NPU (including the LTX-2 _native_npu case this PR targets).

[mengchengTang]

@yiyixuxu @sayakpaul @zhtmike These two Ascend NPU fixes (#14288 and #14396) are issues I ran into while adapting verl-omni — could you take a look when you get a chance? Thanks!

[zhtmike]

Hi @mengchengTang , I have replied in the issues