Skip to content

Enable direct bindings in Thunder #2502

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
t-vi merged 6 commits into from
Oct 6, 2025
Merged

Enable direct bindings in Thunder #2502

t-vi merged 6 commits into from
Oct 6, 2025
+18 −55

Conversation

@rdspring1
Copy link
Collaborator

@rdspring1 rdspring1 commented Sep 4, 2025 •
edited by kevinstephano
Loading

This PR enables the direct bindings for nvfuser executor if DIRECT_BINDINGS_SUPPORTED_VERSION >= 0.2.34.

  • Legacy bindings will remain in the nvfuser pip wheel under import nvfuser, but it will be deprecated and no longer receive new operations.
  • dockers/ubuntu-cuda/Dockerfile, scripts/validate_build.py, thunder/tests/conftest.py, thunder/executors/nvfuserex.py contain import nvfuser for version checks and FusionCache.reset(). These files will need to be updated when legacy bindings are removed.

What is direct bindings?

  • TL;DR: It is direct pybind11 mapping from nvfuser CPP API to python.
  • The python API is the same as legacy bindings, so Thunder is minimally affected.
  • Direct bindings does not have a static Fusion Definition cache like legacy bindings.

Why direct bindings?

  • The legacy bindings utilized an intermediate representation from nvFuser's underlying FusionIR where only the definition ops were exposed in this intermediate representation to allow for graph caching. By moving to direct bindings, we are allowing for the possibility of adding more bindings to expose the scheduling of operations from python to allow users to more directly interact with nvFuser internals from python. One example where exposing scheduling operations is useful is multi-gpu support where the mesh dimensions of a tensor are annotated via a scheduling operation.
  • It is difficult to expose multi-gpu scheduling primitives with legacy bindings.
  • The static Fusion Definition cache can be a memory hog and cumbersome to update.

@rdspring1 rdspring1 force-pushed the nvf_direct branch 2 times, most recently from 48aa575 to 27d9531 Compare September 4, 2025 22:49
@rdspring1 rdspring1 mentioned this pull request Sep 4, 2025
Merged
rdspring1 added a commit to NVIDIA/Fuser that referenced this pull request Sep 5, 2025
@rdspring1
Fixes to direct bindings from Thunder CI (#5119)
ed11f2b 
This PR is a set of fixes for issues found from
Lightning-AI/lightning-thunder#2502.
* Increase `max_length` default in `FusionDefinition` to `9999`
* Rename `num_dims` function to `ndim` property in TensorView
* Use size of logical domain without reductions as number of dimensions
instead of `TensorView::nDims()`
rdspring1
rdspring1 commented Sep 6, 2025
View reviewed changes
thunder/executors/nvfuserex_impl.py
# TODO Review splititng very large fusions or removing the max length restriction completely
# See "Very large nvFuser fusions hit max_length"
fd = FusionDefinition(max_length=MAX_LENGTH)
fd = FusionDefinition()
Copy link
Collaborator Author

@rdspring1 rdspring1 Sep 6, 2025

Choose a reason for hiding this comment

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

nvfuser 0.2.32 bumped the default max length for legacy bindings to 9999 while the direct bindings does not have a fusion cache.

rdspring1
rdspring1 commented Sep 6, 2025
View reviewed changes
thunder/executors/nvfuserex_impl.py
else:
self.set_fuel(FUEL_LEVEL.UNLIMITED)

env_var_save_serde = os.getenv("ENABLE_NVFUSER_SERIALIZATION", None)
Copy link
Collaborator Author

@rdspring1 rdspring1 Sep 6, 2025

Choose a reason for hiding this comment

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

Direct bindings does not support serialization. Removing it for simplicity.

Copy link
Collaborator

@kshitij12345 kshitij12345 Sep 20, 2025

Choose a reason for hiding this comment

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

Just curious to know if there is a plan to support this in future.

Copy link
Collaborator Author

@rdspring1 rdspring1 Sep 20, 2025

Choose a reason for hiding this comment

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

Serialization would need to be added on top of [RFC] Create LRU cache for direct bindings. #4893. It isn't on the roadmap but it has been accounted for.

@rdspring1 rdspring1 force-pushed the nvf_direct branch 5 times, most recently from 204100a to 296167e Compare September 6, 2025 23:26
@rdspring1 rdspring1 changed the title [WIP] Enable direct bindings in Thunder Enable direct bindings in Thunder Sep 7, 2025
@rdspring1 rdspring1 marked this pull request as ready for review September 7, 2025 00:33
@rdspring1 rdspring1 requested review from kshitij12345 and removed request for kshitij12345 September 7, 2025 00:33
@rdspring1 rdspring1 mentioned this pull request Sep 16, 2025
Merged
kshitij12345
kshitij12345 approved these changes Sep 20, 2025
View reviewed changes
Copy link
Collaborator

@kshitij12345 kshitij12345 left a comment

Choose a reason for hiding this comment

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

LGTM, thanks @rdspring1

thunder/executors/nvfuserex_impl.py
else:
self.set_fuel(FUEL_LEVEL.UNLIMITED)

env_var_save_serde = os.getenv("ENABLE_NVFUSER_SERIALIZATION", None)
Copy link
Collaborator

@kshitij12345 kshitij12345 Sep 20, 2025

Choose a reason for hiding this comment

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

Just curious to know if there is a plan to support this in future.

@rdspring1 rdspring1 requested a review from KaelanDt as a code owner October 4, 2025 17:27
@rdspring1
Copy link
Collaborator Author

rdspring1 commented Oct 6, 2025

@t-vi This PR is ready for your review. Also, what are the 10 required tests that haven't run yet?

@t-vi
Copy link
Collaborator

t-vi commented Oct 6, 2025

@rdspring1 unfortunately, the CI is not up currently. I'm running the tests manually now and will merge when they pass.

t-vi
t-vi approved these changes Oct 6, 2025
View reviewed changes
Copy link
Collaborator

@t-vi t-vi left a comment

Choose a reason for hiding this comment

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

Thank you @rdspring1 @crcrpar @kshitij12345

@t-vi t-vi enabled auto-merge (squash) October 6, 2025 15:47
@t-vi t-vi disabled auto-merge October 6, 2025 15:49
@t-vi t-vi merged commit 8b542cf into Lightning-AI:main Oct 6, 2025
39 of 51 checks passed
This was referenced Oct 7, 2025
Merged
Merged
@riccardofelluga riccardofelluga mentioned this pull request Oct 29, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@crcrpar crcrpar crcrpar approved these changes

@kshitij12345 kshitij12345 kshitij12345 approved these changes

@t-vi t-vi t-vi approved these changes

@mruberry mruberry Awaiting requested review from mruberry mruberry is a code owner

@lantiga lantiga Awaiting requested review from lantiga lantiga is a code owner

@KaelanDt KaelanDt Awaiting requested review from KaelanDt KaelanDt is a code owner

Assignees

No one assigned

Labels

nvfuser

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@rdspring1 @t-vi @crcrpar @kshitij12345 @IvanYashchuk