Skip to content

MLA attention fallback - #4517

Open
jloftin-nv wants to merge 1 commit into
pytorch:mainfrom
jloftin-nv:dev-jloftin-mla-attention-fallback
Open

MLA attention fallback#4517
jloftin-nv wants to merge 1 commit into
pytorch:mainfrom
jloftin-nv:dev-jloftin-mla-attention-fallback

Conversation

@jloftin-nv

Copy link
Copy Markdown
Contributor

Description

When the SDPA validator declines (e.g. MLA K/V head-dim mismatch), the node is left as an unconvertible PyTorch op even though decompose_attention=True would work. Users must know to set that flag by hand.

After the causal-attention pass, auto-decompose SDPA nodes the validator declines so MLA models do not take one graph break per layer.

Type of change

Please delete options that are not relevant and/or add your own.

  • New feature (non-breaking change which adds functionality)

Checklist:

  • [ X] My code follows the style guidelines of this project (You can use the linters)
  • [ X] I have performed a self-review of my own code
  • [ X] I have commented my code, particularly in hard-to-understand areas and hacks
  • [ X] I have made corresponding changes to the documentation
  • [ X] I have added tests to verify my fix or my feature
  • [ X] New and existing unit tests pass locally with my changes
  • [ X] I have added the relevant labels to my PR in so that relevant reviewers are notified

@github-actions github-actions Bot added component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 18, 2026
@meta-cla meta-cla Bot added the cla signed label Aug 18, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 18, 2026 21:57
@lanluo-nvidia lanluo-nvidia added this to the v2.15.0 milestone Aug 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants