AttributeError: '_NestedLoopCheckVisitor' object has no attribute '_inst'

[Serpientw]

As I try out mamba3 today, I encounter a weird bug which didn't occur yesterday:

Traceback (most recent call last):
File "/checkpoint/binary/train_package/scripts/test_mamba3_env.py", line 16, in <module>
y = model(x)
File "/root/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/root/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
return forward_call(*args, **kwargs)
File "/checkpoint/binary/train_package/mamba_ssm/modules/mamba3.py", line 187, in forward
y = mamba3_mimo_combined(
File "/checkpoint/binary/train_package/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py", line 290, in mamba3_mimo
return _Mamba3Function.apply(
File "/root/.local/lib/python3.10/site-packages/torch/autograd/function.py", line 581, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
File "/checkpoint/binary/train_package/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py", line 80, in forward
Out, Final_SSM_State, Final_K = mamba_mimo_forward(
File "/checkpoint/binary/train_package/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py", line 439, in mamba_mimo_forward
kernel = mamba_mimo_fwd(B, S, H, G, N, P, R,
File "/root/.local/lib/python3.10/site-packages/tilelang/jit/__init__.py", line 440, in __call__
kernel = self.compile(*args, **kwargs)
File "/root/.local/lib/python3.10/site-packages/tilelang/jit/__init__.py", line 375, in compile
kernel_result = compile(
File "/root/.local/lib/python3.10/site-packages/tilelang/jit/__init__.py", line 98, in compile
return cached(
File "/root/.local/lib/python3.10/site-packages/tilelang/cache/__init__.py", line 74, in cached
return _dispatch_map[execution_backend].cached(
File "/root/.local/lib/python3.10/site-packages/tilelang/cache/kernel_cache.py", line 264, in cached
kernel = JITKernel(
File "/root/.local/lib/python3.10/site-packages/tilelang/jit/kernel.py", line 137, in __init__
adapter = self._compile_and_create_adapter(func, out_idx)
File "/root/.local/lib/python3.10/site-packages/tilelang/jit/kernel.py", line 241, in _compile_and_create_adapter
artifact = tilelang.lower(
File "/root/.local/lib/python3.10/site-packages/tilelang/engine/lower.py", line 246, in lower
PreLowerSemanticCheck(mod)
File "/root/.local/lib/python3.10/site-packages/tilelang/engine/phase.py", line 132, in PreLowerSemanticCheck
tilelang.analysis.NestedLoopChecker()(mod)
File "/root/.local/lib/python3.10/site-packages/tilelang/3rdparty/tvm/python/tvm/ir/transform.py", line 167, in __call__
return _ffi_transform_api.RunPass(self, mod)
File "python/tvm_ffi/cython/function.pxi", line 929, in tvm_ffi.core.Function.__call__
File "<unknown>", line 0, in tvm::transform::Pass::operator()(tvm::IRModule) const
File "<unknown>", line 0, in tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
File "<unknown>", line 0, in tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
File "python/tvm_ffi/cython/function.pxi", line 1083, in tvm_ffi.core.tvm_ffi_callback
File "/root/.local/lib/python3.10/site-packages/tilelang/analysis/nested_loop_checker.py", line 116, in pass_fn
_NestedLoopCheckVisitor().visit_stmt(func.body)
File "/root/.local/lib/python3.10/site-packages/tilelang/3rdparty/tvm/python/tvm/runtime/support.py", line 159, in __init__
self._inst = cls(*args, **kwargs)
File "/root/.local/lib/python3.10/site-packages/tilelang/3rdparty/tvm/python/tvm/runtime/support.py", line 197, in __setattr__
super(TVMDerivedObject, self).__setattr__(name, value)
AttributeError: '_NestedLoopCheckVisitor' object has no attribute '_inst'

As it turns out, mamba3 relies on tilelang==0.1.8, which relies on "apache-tvm-ffi~=0.1.0,>=0.1.2", however, apache-tvm-ffi releases its latest version 0.1.10 13 hours ago as I write this, which causes the above problem, I'm not sure whether this is a bug in tvm-ffi or simply a compatibility issue, but for now I can only avoid this bug by fixing the version of apache-tvm-ffi to 0.1.9

[KylekG]

I also encountered this issue and was able to get past it by adding the following line into the pyproject.toml dependencies:

"apache-tvm-ffi<=0.1.9",

Removing the line again before rebuilding my container causes the following traceback:

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/workspace/main.py", line 16, in <module>
    y = model(x)
        ^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1790, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mamba_ssm/modules/mamba3.py", line 187, in forward
    y = mamba3_mimo_combined(
        ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py", line 290, in mamba3_mimo
    return _Mamba3Function.apply(
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 596, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py", line 80, in forward
    Out, Final_SSM_State, Final_K = mamba_mimo_forward(
                                    ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py", line 439, in mamba_mimo_forward
    kernel = mamba_mimo_fwd(B, S, H, G, N, P, R, 
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tilelang/jit/__init__.py", line 440, in __call__
    kernel = self.compile(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tilelang/jit/__init__.py", line 375, in compile
    kernel_result = compile(
                    ^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tilelang/jit/__init__.py", line 98, in compile
    return cached(
           ^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tilelang/cache/__init__.py", line 74, in cached
    return _dispatch_map[execution_backend].cached(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tilelang/cache/kernel_cache.py", line 264, in cached
    kernel = JITKernel(
             ^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tilelang/jit/kernel.py", line 137, in __init__
    adapter = self._compile_and_create_adapter(func, out_idx)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tilelang/jit/kernel.py", line 241, in _compile_and_create_adapter
    artifact = tilelang.lower(
               ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tilelang/engine/lower.py", line 246, in lower
    PreLowerSemanticCheck(mod)
  File "/usr/local/lib/python3.12/dist-packages/tilelang/engine/phase.py", line 132, in PreLowerSemanticCheck
    tilelang.analysis.NestedLoopChecker()(mod)
  File "/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/tvm/python/tvm/ir/transform.py", line 167, in __call__
    return _ffi_transform_api.RunPass(self, mod)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/tvm_ffi/cython/function.pxi", line 929, in tvm_ffi.core.Function.__call__
  File "<unknown>", line 0, in tvm::transform::Pass::operator()(tvm::IRModule) const
  File "<unknown>", line 0, in tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  File "<unknown>", line 0, in tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  File "python/tvm_ffi/cython/function.pxi", line 1083, in tvm_ffi.core.tvm_ffi_callback
  File "/usr/local/lib/python3.12/dist-packages/tilelang/analysis/nested_loop_checker.py", line 116, in pass_fn
    _NestedLoopCheckVisitor().visit_stmt(func.body)

  File "/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/tvm/python/tvm/runtime/support.py", line 159, in __init__
    self._inst = cls(*args, **kwargs)

  File "/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/tvm/python/tvm/runtime/support.py", line 197, in __setattr__
    super(TVMDerivedObject, self).__setattr__(name, value)

AttributeError: '_NestedLoopCheckVisitor' object has no attribute '_inst'

[skaae]

I solved it by compiling tilelang from source. In uv you can do something like:

dependencies = [
    "tilelang",
    "mamba-ssm",
]

[tool.uv.sources]
mamba_ssm = { git = "https://github.com/state-spaces/mamba", rev = "7438488222dc44eb9146e4f39c0764a8f651ede6" }
tilelang = { git = "https://github.com/tile-ai/tilelang", rev = "853e805b2ae8d1d06cee215367a1bdf775a83058" }

[tridao]

cc @berlinchen7

[berlinchen7]

The _NestedLoopCheckVisitor object in Tilelang appears to be broken with the new apache-tvm-ffi version (0.1.10). I filed an issue there: tile-ai/tilelang#2040

In the meanwhile, I will add apache-tvm-ffi<=0.1.9 to the dependency.