[AOTI Eager] Add dynamic shapes support to AOTIPythonKernelHolder#176018
Closed
[AOTI Eager] Add dynamic shapes support to AOTIPythonKernelHolder#176018
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/176018
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit eadb030 with merge base 1342f81 ( UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
903a14d to
a1617e3
Compare
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 2, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 2, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
a1617e3 to
bc186e6
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 2, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
|
@StellarrZ has exported this pull request. If you are a Meta employee, you can view the originating Diff in D94301187. |
bc186e6 to
6898147
Compare
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
6898147 to
0ac306d
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
StellarrZ
added a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
0ac306d to
1a5288d
Compare
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
StellarrZ
added a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
1a5288d to
9e81984
Compare
9e81984 to
636f481
Compare
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
636f481 to
911158c
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
911158c to
a9c0ed8
Compare
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
a9c0ed8 to
cf1249d
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 4, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
3743410 to
c077155
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 5, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 5, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
c077155 to
9d097b4
Compare
9d097b4 to
f1aa5ca
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 5, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
f1aa5ca to
96050ea
Compare
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 6, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
96050ea to
4228b81
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 6, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 6, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
4228b81 to
675f686
Compare
675f686 to
a5cdfff
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 6, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
meta-codesync bot
pushed a commit
that referenced
this pull request
Mar 6, 2026
…76018) Summary: Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
a5cdfff to
7921dfc
Compare
StellarrZ
added a commit
that referenced
this pull request
Mar 6, 2026
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Differential Revision: D94301187
7921dfc to
a56f4aa
Compare
a56f4aa to
bb4b597
Compare
arui-meta
approved these changes
Mar 11, 2026
Summary: Populate the C++ in-memory cache after first compilation so subsequent same-shape dispatches resolve entirely in C++ without calling back into Python. Root cause: `AOTIPythonKernelHolder::cache_miss()` compiled and loaded the kernel but never added it to `aoti_kernel_cache_`. Every dispatch re-entered `produce_aoti_kernel_lib` → acquired the GIL → called Python `_compile_afg` → parsed JSON → loaded AFG from disk. Fix: after `cache_miss()` creates the kernel runner, populate `aoti_kernel_cache_` with the input metadata and runner so subsequent calls with matching shapes hit `cache_hit()` directly. Performance (`aten.bitwise_not`, shape [32,32], 100k iterations): - Before: 34,260 us/call (disk round-trip on every dispatch) - After: 21.5 us/call (in-memory cache hit, 1,593x faster) Test Plan: ``` buck run fbcode//mode/opt fbcode//mtia/host_runtime/aoti_eager:test_aoti_eager_dispatch_bin -- TestAotiEagerDispatch.test_inmemory_cache_reuse ``` Differential Revision: D93171553
…76018) Summary: Pull Request resolved: #176018 Add a `dynamic` parameter to `AOTIPythonKernelHolder` so that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. When `dynamic_=true`, the holder passes `dynamic=True` to `aoti_compile_with_persistent_cache`, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes. Changes: - `AOTIPythonKernelHolder`: added `dynamic_` member, forwarded to `produce_aoti_kernel_lib` - `AOTIKernelMetadata`: added `is_dynamic_` flag, `check()` uses `dynamic_check()` when set - `TensorMetadata::dynamic_check()`: matches by dtype/device/rank, skips exact sizes - `ParameterMetadata::dynamic_check()`: delegates to `TensorMetadata::dynamic_check()` for tensor params Test Plan: ``` buck test fbcode//caffe2/test/cpp/aoti_eager:kernel_meta_info_test ``` Reviewed By: arui-meta Differential Revision: D94301187
|
@pytorchbot merge (Initiating merge automatically since Phabricator Diff has merged) |
Collaborator
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Collaborator
Merge failedReason: 1 jobs have failed, first few of them are: inductor / inductor-cpu-test / test (cpu_inductor_timm, 1, 2, linux.2xlarge.amx) Details for Dev Infra teamRaised by workflow job |
Contributor
|
@pytorchbot merge |
Collaborator
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
|
@pytorchbot merge -i (Initiating merge automatically since Phabricator Diff has merged, merging with -i because oss signals were bypassed internally) |
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Add a
dynamicparameter toAOTIPythonKernelHolderso that the C++ dispatch path can request dynamic-shape compilation from the Python compile backend. Whendynamic_=true, the holder passesdynamic=Truetoaoti_compile_with_persistent_cache, and the in-memory cache uses rank/dtype/device matching instead of exact size/stride matching. This allows a single compiled kernel to serve multiple input shapes.Changes:
AOTIPythonKernelHolder: addeddynamic_member, forwarded toproduce_aoti_kernel_libAOTIKernelMetadata: addedis_dynamic_flag,check()usesdynamic_check()when setTensorMetadata::dynamic_check(): matches by dtype/device/rank, skips exact sizesParameterMetadata::dynamic_check(): delegates toTensorMetadata::dynamic_check()for tensor paramsTest Plan:
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo