Skip to content

Use fused attention for SmolLM3#2849

Open
GargiGupta-io wants to merge 1 commit into
keras-team:masterfrom
GargiGupta-io:fix-smollm3-flash-attention
Open

Use fused attention for SmolLM3#2849
GargiGupta-io wants to merge 1 commit into
keras-team:masterfrom
GargiGupta-io:fix-smollm3-flash-attention

Conversation

@GargiGupta-io

Copy link
Copy Markdown
Contributor

Description of the change

SmolLM3's attention layer always uses the eager einsum and softmax path even
when Keras provides a fused dot-product attention operation. This bypasses the
flash-attention path available on supported JAX and PyTorch configurations.

This change follows the existing Qwen3 and Phi3 pattern: use
ops.dot_product_attention() when fused_attention_op_available() is true,
normalize the attention mask for the fused operator, and retain the existing
implementation as the fallback.

Regression tests cover masked and unmasked dispatch, the eager fallback, and
numerical equivalence between both paths.

Reference

Colab Notebook

Not applicable. This change does not add a new API or model.

Validation

  • TensorFlow: 6 passed, 2 skipped, 2 subtests passed
  • JAX: 6 passed, 2 skipped, 2 subtests passed
  • PyTorch: 6 passed, 2 skipped, 2 subtests passed
  • Ruff lint and format checks
  • The remaining macOS causal-LM test requiring tensorflow-text fails the
    same way on the unchanged master branch.

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have followed the Keras Hub Model contribution guidelines in making these changes.
  • I have followed the Keras Hub API design guidelines in making these changes.
  • I have signed the Contributor License Agreement.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request integrates fused attention support into the SmolLM3Attention layer by utilizing ops.dot_product_attention when fused_attention_op_available() is true. It also adds comprehensive unit tests in smollm3_backbone_test.py to verify the fused attention execution, its fallback path, and numerical consistency between both implementations. There are no review comments to address, and the changes conform to the repository's guidelines.

@GargiGupta-io
GargiGupta-io marked this pull request as ready for review July 22, 2026 18:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant