[Prer1n1]

Description of the change

When attention uses SDPA, it is traced as a single fused op, so qmodel_prep (TorchScript path) finds no bmm/matmul to attach QBmm to, and BMM quantization (nbits_bmm1/nbits_bmm2) is silently skipped.

This PR logs a warning in that case, as the Dynamo path already does (find_and_prep_bmm_gm), and points users to attn_implementation="eager". It only warns when no QBmm was attached and an aten::scaled_dot_product_attention node exists in the traced graph, so eager models are unaffected.

It does not add SDPA support itself; that is left for the issue.

Related issues or PRs

Refs #208

How to verify the PR

Run qmodel_prep(model, inputs, qcfg, use_dynamo=False) on a model that uses SDPA with nbits_bmm1/nbits_bmm2 set; the new warning appears. With attn_implementation="eager" it does not.

Was the PR tested

Added two tests in tests/models/test_qmodelprep.py using a small randomly initialised BERT (no download): sdpa expects the warning and no QBmm; eager expects no warning and QBmm attached. The sdpa test fails without the fix. Also ran the rest of tests/models/test_qmodelprep.py except the ResNet/ViT/BERT tests that need model downloads (513 passed).

• I have added >=1 unit test(s) for every new method I have added (if that coverage is difficult, please briefly explain the reason)
• I have ensured all unit tests pass

Checklist for passing CI/CD:

• All commits are signed showing "Signed-off-by: Name <email@domain.com>" with git commit -signoff or equivalent
• PR title and commit messages adhere to Conventional Commits
• Contribution is formatted with tox -e fix
• Contribution passes linting with tox -e lint
• Contribution passes spellcheck with tox -e spellcheck
• Contribution passes all unit tests with tox -e unit

Formatting and linting were checked by running ruff 0.5.0, isort 5.11.5 and pylint directly on the changed files (all clean), not via tox. Spellcheck and the full unit suite were not run locally; relying on CI for those.

[Prer1n1]

Hi! Gentle nudge: could a maintainer approve the workflows when you get a chance? Happy to address anything CI or review flags. Thanks!

[chichun-charlie-liu]

@Prer1n1 Thanks for the PR. Adding a log message looks fine to me. I'm just curious which of your use cases is still using torchscript? We have been planning to retire the entire torchscript path since PyTorch itself no long supports torchscript since 2.10. It probably won't be long before fms-mo's minimum torch version catch up...
Thanks.

[Prer1n1]

Thanks for taking a look! Honest answer: I don't have a production use case. I picked this up from #208 as my first contribution to fms-mo, and it was reported against the TorchScript path (use_dynamo=False), which is why I added the warning there to mirror the existing one in the Dynamo path. If you're planning to retire TorchScript soon, I'm happy to close this PR, or to put the effort into something on the Dynamo path instead. Do you have a preference? Is there a Dynamo-side issue you'd like help with?

[chichun-charlie-liu]

Thanks for the honest answer, and no worries about not having a production use case. #208 asked for exactly this warning, and the change is small and correct, so I'm happy to move forward with it even though TorchScript is on its way out.

I have one small request. Could you add a comment next to the warning saying that its text is matched by the unit tests, and naming them? Something like:

# NOTE: the substring "Likely SDPA is enabled" is asserted by these tests in
# tests/models/test_qmodelprep.py, keep them in sync if the message changes:
#   - test_torchscript_sdpa_warns_no_bmm
#   - test_torchscript_eager_no_sdpa_warning

Otherwise someone rewording the message would break the tests... Once that's in, I'll approve the workflows so CI can run, and we can go from there.

[Prer1n1]

Done, added the note in 75fb2ab. Thanks for the review!

[chichun-charlie-liu]

lgtm

[Prer1n1]

Update after the re-run: Test passes. The other jobs fail in CI setup before reaching my changes. Lint: tox -e lint/ruff error with "could not find python interpreter matching any of the specs python3.12" (lint.yml sets up 3.11, tox.ini pins 3.12). Build: hynek/build-and-inspect-python-package fails installing its own tools because its environment is on Python 3.15.2 and pydantic-core==2.41.1 can't build with PyO3 0.26 (max supported 3.14). Both look like runner/toolchain changes rather than anything in this PR. Locally ruff, isort and pylint pass. Happy to open a separate PR for the lint.yml Python version if useful.

[chichun-charlie-liu]

Hi! Gentle nudge: could a maintainer approve the workflows when you get a chance? Happy to address anything CI or review flags. Thanks!

@andrea-fasoli and I have updated the main, CI should be fixed now. Please click "Update branch" on your fork (rebase) and see if it works now. Thx.

[Prer1n1]

Rebased onto the updated main. Let's see how CI looks now. Thanks for the quick fix!