Skip to content

Conversation

@danieldk
Copy link

@danieldk danieldk commented Nov 28, 2025

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

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