Skip to content

Fix Ascend NPU RMSNorm and fused-attention mask shapes#14288

Open
mengchengTang wants to merge 1 commit into
huggingface:mainfrom
mengchengTang:fix/npu-rmsnorm-and-attn-mask-main
Open

Fix Ascend NPU RMSNorm and fused-attention mask shapes#14288
mengchengTang wants to merge 1 commit into
huggingface:mainfrom
mengchengTang:fix/npu-rmsnorm-and-attn-mask-main

Conversation

@mengchengTang

@mengchengTang mengchengTang commented Jul 25, 2026

Copy link
Copy Markdown

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

image

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

image

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 github-actions Bot added size/S PR with diff < 50 LOC models and removed size/S PR with diff < 50 LOC labels Jul 25, 2026
@mengchengTang
mengchengTang force-pushed the fix/npu-rmsnorm-and-attn-mask-main branch from d77a8f5 to 609e5b1 Compare July 25, 2026 06:28
@github-actions github-actions Bot added the size/S PR with diff < 50 LOC label Jul 25, 2026
@mengchengTang
mengchengTang force-pushed the fix/npu-rmsnorm-and-attn-mask-main branch from 609e5b1 to 04c4ed7 Compare July 25, 2026 06:43
@github-actions

Copy link
Copy Markdown
Contributor

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

models size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant