Skip to content
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.

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

perf: accelerate JIT compilation speed #618

Merged
merged 4 commits into from
Nov 20, 2024
Merged

perf: accelerate JIT compilation speed #618

merged 4 commits into from
Nov 20, 2024

Conversation

yzh119
Copy link
Collaborator

@yzh119 yzh119 commented Nov 20, 2024

Current JIT compilation is slow because we rely on a huge header <torch/extension.h> which is too heavy for our use case.

This PR refactors the codebase to only include necessary headers for pybind, and moves most of torch runtime API calls from C++ to python.

The compilation time was reduced from 48 seconds to 18 seconds for lightweight operators such as norm.

@yzh119
Copy link
Collaborator Author

yzh119 commented Nov 20, 2024

Using nanobind or ctypes/cython for ffi could make the compilation speed even faster, but requires some fundamental changes to the codebase, I prefer sticking to pybind at this moment.

upd

bugfix
@yzh119 yzh119 merged commit eaf73fd into main Nov 20, 2024
yzh119 added a commit that referenced this pull request Nov 20, 2024
Some of the commits for fixing MLA are missing in #618, this PR add them
back.
yzh119 added a commit that referenced this pull request Nov 20, 2024
#618 breaks aot compilation because of function signature mismatch, this
PR fixes the issue.
@zhyncs zhyncs deleted the faster-jit branch November 23, 2024 06:43
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