cutlass
cutlass copied to clipboard
Fix processing of relative imports in AST preprocessing
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.
I could not find a good existing place to put a test (if necessary).