Skip to content

Include extra specialization in specialization_key#1883

Merged
jansel merged 4 commits intopytorch:mainfrom
hinriksnaer:include-extra-specialization
Apr 3, 2026
Merged

Include extra specialization in specialization_key#1883
jansel merged 4 commits intopytorch:mainfrom
hinriksnaer:include-extra-specialization

Conversation

@hinriksnaer
Copy link
Copy Markdown
Collaborator

Addresses #1209

specialization_key() now returns the full specialization key, including constraints discovered during compilation. There might be a more elegant solution to this but would likely involve more risk and potential for introducing bugs. Hoping that this fully addresses #1209.

Changes

  • Extracted _base_specialization_key() from specialization_key() for internal use
  • specialization_key() appends compilation-discovered extras (e.g. from hl.specialize()) when available
  • Updated bind(), local_cache.py, and aot_cache.py to use _base_specialization_key() internally
  • Added regression test in test_specialize.py

Note

Before the first compilation, extras are not yet known, so the key may still be incomplete. I am not sure if this would cause any existing autotuning cache results to miss.

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 30, 2026
Copy link
Copy Markdown
Collaborator

@fulvius31 fulvius31 left a comment

Choose a reason for hiding this comment

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

LGTM but base_search.py should use _base_specialization_key() so that FROM_BEST_AVAILABLE matches across different hl.specialize values and seeds the autotuner with prior good configs instead of falling back to defaults.

It's not the scope of this PR, I'll follow up with another PR when this get merged.

I am not sure if this would cause any existing autotuning cache results to miss.

Yes but I don't think is relevant.

@hinriksnaer hinriksnaer requested a review from jansel March 31, 2026 14:57
@hinriksnaer hinriksnaer requested a review from jansel April 1, 2026 16:55
@jansel jansel merged commit 43e6d83 into pytorch:main Apr 3, 2026
18 checks passed
fulvius31 added a commit to fulvius31/helion that referenced this pull request Apr 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants