Skip to content

Fix bool attention mask bias conversion#4354

Open
fallintoplace wants to merge 1 commit into
pytorch:mainfrom
fallintoplace:fix-bool-attn-mask-select
Open

Fix bool attention mask bias conversion#4354
fallintoplace wants to merge 1 commit into
pytorch:mainfrom
fallintoplace:fix-bool-attn-mask-select

Conversation

@fallintoplace

Copy link
Copy Markdown

Summary

  • Convert bool attention masks to additive bias with ISelectLayer instead of arithmetic involving -inf.
  • Return 0.0 for valid mask entries and -inf for masked entries.
  • Add a deterministic bool SDPA mask regression that checks the engine output has no NaN values.

Testing

  • git diff --check
  • TORCH_PATH="$PWD/.venv/lib/python3.12/site-packages/torch" bazelisk test //tests/core/conversion/converters:test_scaled_dot_product_attention --test_filter=ATenScaledDotProductAttnMaskBoolDoesNotProduceNaN was attempted locally. Bazel stayed in analysis without reaching compilation or test execution after about 96 seconds, so the run was interrupted.
@meta-cla meta-cla Bot added the cla signed label Jun 20, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: converters Issues re: Specific op converters labels Jun 20, 2026
@github-actions github-actions Bot requested a review from zewenli98 June 20, 2026 22:16
@zewenli98

Copy link
Copy Markdown
Collaborator

@fallintoplace Thanks for the contribution. Torchscript is being removed soon (like pytorch 2.14). Please use Dynamo instead.

@zewenli98 zewenli98 closed this Jun 23, 2026
@zewenli98 zewenli98 reopened this Jun 30, 2026

@zewenli98 zewenli98 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@fallintoplace Thanks for the contribution. Torchscript is being removed soon (like pytorch 2.14). Please use Dynamo instead.

@fallintoplace Pytorch decided not to remove Torchscript support in 2.14. This PR looks good to me. Can you rebase the PR to the main?

@fallintoplace

Copy link
Copy Markdown
Author

@zewenli98 Sure, let me push in a few minutes.

@fallintoplace fallintoplace force-pushed the fix-bool-attn-mask-select branch from b4a67e8 to 19e2a2b Compare June 30, 2026 22:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: tests Issues re: Tests

2 participants