Skip to content

[fix(externalize)] Skip fully-specialised dims in submodule re-export#7

Draft
gokulkrishna98 wants to merge 1 commit into
apple:mainfrom
gokulkrishna98:dev/gokul/fix_sdpa_externalization
Draft

[fix(externalize)] Skip fully-specialised dims in submodule re-export#7
gokulkrishna98 wants to merge 1 commit into
apple:mainfrom
gokulkrishna98:dev/gokul/fix_sdpa_externalization

Conversation

@gokulkrishna98

Copy link
Copy Markdown
Contributor

When externalising an SDPA submodule, ``_dynamic_shapes_from_node``
requested a Dim for every SymInt in the upstream FakeTensor's shape
— including SymInts whose ``.node.expr`` had been specialised to a
literal int by the parent program. Those re-exported as an unbounded
``Dim(min=1)`` and ``torch.export`` rejected the submodule with
``L['key'].size()[2] <= IntInfinity()`` whenever a model used a static
query length and a dynamic KV-context length (the prefill / decode
shape used by hybrid linear-attention models).

Skip SymInts whose expr is already a number — they are static dims and
should not appear in ``dynamic_shapes`` at all.

Signed-off-by: gokulkrishna98 <gokulkrishna98@users.noreply.github.com>
@gokulkrishna98 gokulkrishna98 force-pushed the dev/gokul/fix_sdpa_externalization branch from 7ef0b8e to ac0e0de Compare June 11, 2026 18:42
@gokulkrishna98 gokulkrishna98 self-assigned this Jun 16, 2026
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