[torch.library.custom_op] Add _register_nvfuser_translator
#2481
Conversation
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
github-actions
bot
added
the
documentation
crcrpar
force-pushed
the
custom_op/override-fwd-def
branch
3 times, most recently
from
c772160 to
e84169d
Compare
crcrpar
force-pushed
the
custom_op/override-fwd-def
branch
from
86ab3da to
6de2392
Compare
crcrpar
commented
Sep 18, 2025
|
the failure seems caused by the "polluted" translation_map & torch->symbol map |
crcrpar
force-pushed
the
custom_op/override-fwd-def
branch
from
6de2392 to
d6a6b29
Compare
crcrpar
force-pushed
the
custom_op/override-fwd-def
branch
from
d6a6b29 to
c3e7eb3
Compare
crcrpar
requested review from
KaelanDt,
lantiga,
mruberry and
t-vi
as code owners
crcrpar
force-pushed
the
custom_op/override-fwd-def
branch
from
c3e7eb3 to
72ba46c
Compare
crcrpar
changed the title
thunder.executors.custom_op_ex._override_custom_op_forwardtorch.library.custom_op] Add _register_nvfuser_translator
crcrpar
force-pushed
the
custom_op/override-fwd-def
branch
from
d4a59a4 to
6ce7b86
Compare
which allows overriding forward of `torch.library.custom_op`'s forward Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
…ranslator` Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
crcrpar
force-pushed
the
custom_op/override-fwd-def
branch
from
3e0fa38 to
0f88252
Compare
mattteochen
approved these changes
Oct 14, 2025
| ): | ||
| if any(sub_bsym.sym.id == _symbol.id for sub_bsym in bsym.subsymbols): | ||
| nvfuser_def_for_custom_op_found = True | ||
| assert nvfuser_def_for_custom_op_found |
There was a problem hiding this comment.
Maybe we can assert the number of custom op we expect to find in fuser definiton and backward trace
Comment on lines
+531
to
+534
| # Register the custom_op of `mul` with :func:`thunder.torch._register_custom_op` | ||
| _symbol = _register_custom_op(mul) | ||
| # Register custom nvfuser definition for the already registered custom_op of mul | ||
| _register_nvfuser_translator(_symbol, mul_translator) |
There was a problem hiding this comment.
Is this the final API usage proposal (symbol creation and translation registration)?
| from thunder.executors.torchex import _always_executable | ||
|
|
||
| register_supported(symbol, translator_for_nvfuser, checker or _always_executable) | ||
| register_supported(symbol.id, translator_for_nvfuser, checker or _always_executable) |
There was a problem hiding this comment.
Can you clarify why both calls are needed?
There was a problem hiding this comment.
for most cases the first one should be enough. I don't quite remember the exact reason why I have these two here.
mattteochen
approved these changes
Oct 14, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
What does this PR do?
custom_op_exto the default executor list whenthunder.torch.custom_op._register_custom_opis usedregister_executor(custom_op_ex)tothunder/executors/custom_op_ex.py_register_nvfuser_translatorthat lets us register nvfuser translator rule for an already registeredtorch.library.custom_opclose #2605
cc @Borda @lantiga