Skip to content

Fix additive attention mask handling in the native NPU attention backend. - #14396

Open
mengchengTang wants to merge 1 commit into
huggingface:mainfrom
mengchengTang:fix/native_npu_additive_mask
Open

Fix additive attention mask handling in the native NPU attention backend.#14396
mengchengTang wants to merge 1 commit into
huggingface:mainfrom
mengchengTang:fix/native_npu_additive_mask

Conversation

@mengchengTang

@mengchengTang mengchengTang commented Aug 5, 2026

Copy link
Copy Markdown

What does this PR do?

Fixes attention mask handling in the native NPU attention backend.

PyTorch SDPA defines two mask formats:

  • Boolean masks are keep masks, where True means attend.
  • Floating-point masks are additive biases applied directly to attention scores.

Previously, _native_npu_attention passed every mask to _maybe_modify_attn_mask_npu, which casts the mask to boolean and inverts it for npu_fusion_attention. This works for boolean masks but loses the values and semantics of floating-point additive masks.

This PR makes the mask handling explicit:

  • Floating-point masks are routed through the existing native SDPA implementation, preserving their additive values.
  • Boolean masks continue using npu_fusion_attention.
  • Other mask dtypes are rejected with a clear error instead of being silently converted to boolean.

No new dependencies are introduced.

Fixes #14397

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 Aug 5, 2026
@mengchengTang mengchengTang changed the title fix the additive float mask in native_npu Fixes additive attention mask handling in the native NPU attention backend. Aug 5, 2026
@mengchengTang mengchengTang changed the title Fixes additive attention mask handling in the native NPU attention backend. Fix additive attention mask handling in the native NPU attention backend. Aug 5, 2026
@mengchengTang
mengchengTang force-pushed the fix/native_npu_additive_mask branch from e4a82e0 to e1962a8 Compare August 5, 2026 15:30
@github-actions github-actions Bot added the size/S PR with diff < 50 LOC label Aug 5, 2026
@mengchengTang
mengchengTang force-pushed the fix/native_npu_additive_mask branch from e1962a8 to c9ba560 Compare August 5, 2026 16:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fixes-issue models size/S PR with diff < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Native NPU attention incorrectly converts additive attention masks to boolean masks

1 participant