Fix processing of relative imports in AST preprocessing #2821
+4
−1
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
Suppose that we have a file add.py in a package:
import cutlass.cute as cute
from . import hello
@cute.jit
def add(a, b):
return a + b
and then in a file in the same directory we have import this function
and compile it:
import cutlass.cute as cute
from .add import add
jit_executor = cute.compile(add, 1, 2)
The compilation will fail with:
jit_executor = cute.compile(add, 1, 2)
^^^^^^^^^^^^^^^^^^^^^^^
File "[...]/nvidia_cutlass_dsl/python_packages/cutlass/base_dsl/ast_preprocessor.py", line 378, in exec_imports
raise ImportError(
ImportError: Failed to import mypkg.None: No module named 'mypkg.None'
I debugged `ast_processor.py` and found that the issue is with the
relative import in the first file (`from . import hello`), where the
package is set correctly, but the module name is `None`. This seems
correct, because the function is exposed at the top-level. This small
patch handles this case.
I have tested this both with a toy example and with flash-attention 4,
where I have replaced all the absolute imports to relative imports.
Author
|
@hwu36 @Junkai-Wu Any chance this (or another fix for this issue) could be merged soon? It is blocking us to support CuTe DSL kernels through https://huggingface.co/docs/kernels/ |
|
This PR has been labeled |
Author
|
This is still relevant. |
Collaborator
|
@fengxie @brandon-yujie-sun to help review. |
Collaborator
|
Thanks for proposing the change. This has been adopted and fixed in |
Author
|
Awesome, thanks! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 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.
Suppose that we have a file
add.pyin a package:and then in a file in the same directory we have import this function
and compile it:
The compilation will fail with:
I debugged
ast_processor.pyand found that the issue is with the relative import in the first file (from . import hello), where the package is set correctly, but the module name isNone. This seems correct, because the function is exposed at the top-level. This small patch handles this case.I have tested this both with a toy example and with flash-attention 4, where I have replaced all the absolute imports to relative imports.
I could not find a good existing place to put a test (if necessary).