Skip to content

[Torch] Support dynamic head_dim in attention conversion#23680

Closed
keshavvinayak01 wants to merge 1 commit intoiree-org:mainfrom
keshavvinayak01:users/keshavvinayak01/sdpa-dynamic-head-dim
Closed

[Torch] Support dynamic head_dim in attention conversion#23680
keshavvinayak01 wants to merge 1 commit intoiree-org:mainfrom
keshavvinayak01:users/keshavvinayak01/sdpa-dynamic-head-dim

Conversation

@keshavvinayak01
Copy link
Copy Markdown
Contributor

When head_dim is dynamic, scale is computed at runtime via index_cast + sitofp + rsqrt instead of being folded to a constant.

Static head_dim behavior is unchanged.

Remove the NYI error for dynamic head_dim and compute scale at runtime
via math.rsqrt when the dimension is not statically known. Add lit tests
for dynamic head_dim cases.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>
@keshavvinayak01
Copy link
Copy Markdown
Contributor Author

Closing, duplicate of #23636

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