cutlass icon indicating copy to clipboard operation
cutlass copied to clipboard

Fix processing of relative imports in AST preprocessing

Open danieldk opened this issue 1 month ago • 0 comments

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).

danieldk avatar Nov 28 '25 13:11 danieldk