Unable to Install Sparse Attention -- Unable to Import Triton

When I try to install DeepSpeed Sparse Attention using DS_BUILD_SPARSE_ATTN=1 ./install.sh, I get the error Warning: Unable to import triton, sparse attention will not be accessible, and the final output shows [WARNING] sparse attention is NOT installed. Everything else successfully installs.

The installer does install Triton to my conda environment, but python is unable to import it. If I run python -c "import triton", the following exception is thrown:

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/triton/__init__.py", line 1, in <module>
    from .kernel import *
  File "/home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/triton/kernel.py", line 1, in <module>
    import triton._C.libtriton as libtriton
ImportError: /home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/triton/_C/libtriton.so: undefined symbol: _ZN3c105ErrorC1ENS_14SourceLocationESs

The command ldd -r libtriton.so shows quite a few undefined symbols. Here is the full output:

	linux-vdso.so.1 (0x00007ffe819bf000)
	libLLVM-9.so.1 => /usr/lib/llvm-9/lib/libLLVM-9.so.1 (0x00007fbd3e028000)
	libc10_cuda.so => /home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/torch/lib/libc10_cuda.so (0x00007fbd3ddf0000)
	libtorch_cpu.so => /home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so (0x00007fbd36b5a000)
	libstdc++.so.6 => /usr/lib/x86_64-linux-gnu/libstdc++.so.6 (0x00007fbd367d1000)
	libgcc_s.so.1 => /lib/x86_64-linux-gnu/libgcc_s.so.1 (0x00007fbd365b9000)
	libc.so.6 => /lib/x86_64-linux-gnu/libc.so.6 (0x00007fbd361c8000)
	/lib64/ld-linux-x86-64.so.2 (0x00007fbd42684000)
	libffi.so.6 => /usr/lib/x86_64-linux-gnu/libffi.so.6 (0x00007fbd35fc0000)
	libedit.so.2 => /usr/lib/x86_64-linux-gnu/libedit.so.2 (0x00007fbd35d89000)
	libz.so.1 => /lib/x86_64-linux-gnu/libz.so.1 (0x00007fbd35b6c000)
	librt.so.1 => /lib/x86_64-linux-gnu/librt.so.1 (0x00007fbd35964000)
	libdl.so.2 => /lib/x86_64-linux-gnu/libdl.so.2 (0x00007fbd35760000)
	libtinfo.so.5 => /lib/x86_64-linux-gnu/libtinfo.so.5 (0x00007fbd35536000)
	libpthread.so.0 => /lib/x86_64-linux-gnu/libpthread.so.0 (0x00007fbd35317000)
	libm.so.6 => /lib/x86_64-linux-gnu/libm.so.6 (0x00007fbd34f79000)
	libcudart.so.10.2 => /usr/local/cuda/lib64/libcudart.so.10.2 (0x00007fbd34cfb000)
	libc10.so => /home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/torch/lib/libc10.so (0x00007fbd34a85000)
	libgomp.so.1 => /usr/lib/x86_64-linux-gnu/libgomp.so.1 (0x00007fbd34856000)
	libmkl_intel_lp64.so => /home/nathaniel/Programs/anaconda3/envs/nlp/lib/libmkl_intel_lp64.so (0x00007fbd33b45000)
	libmkl_gnu_thread.so => /home/nathaniel/Programs/anaconda3/envs/nlp/lib/libmkl_gnu_thread.so (0x00007fbd31e1a000)
	libmkl_core.so => /home/nathaniel/Programs/anaconda3/envs/nlp/lib/libmkl_core.so (0x00007fbd2d854000)
undefined symbol: PyInstanceMethod_Type	(./libtriton.so)
undefined symbol: PyExc_RuntimeError	(./libtriton.so)
undefined symbol: PyExc_MemoryError	(./libtriton.so)
undefined symbol: PyBaseObject_Type	(./libtriton.so)
undefined symbol: PyModule_Type	(./libtriton.so)
undefined symbol: PyExc_TypeError	(./libtriton.so)
undefined symbol: PyCFunction_Type	(./libtriton.so)
undefined symbol: PyMethod_Type	(./libtriton.so)
undefined symbol: PyExc_ImportError	(./libtriton.so)
undefined symbol: PyExc_ValueError	(./libtriton.so)
undefined symbol: _Py_FalseStruct	(./libtriton.so)
undefined symbol: PyDict_Type	(./libtriton.so)
undefined symbol: PyExc_BufferError	(./libtriton.so)
undefined symbol: PyType_Type	(./libtriton.so)
undefined symbol: PyProperty_Type	(./libtriton.so)
undefined symbol: _Py_NoneStruct	(./libtriton.so)
undefined symbol: _Py_NotImplementedStruct	(./libtriton.so)
undefined symbol: PyExc_SystemError	(./libtriton.so)
undefined symbol: PyCapsule_Type	(./libtriton.so)
undefined symbol: _Py_TrueStruct	(./libtriton.so)
undefined symbol: PyFloat_Type	(./libtriton.so)
undefined symbol: PyExc_IndexError	(./libtriton.so)
undefined symbol: PyDict_DelItemString	(./libtriton.so)
undefined symbol: _Py_Dealloc	(./libtriton.so)
undefined symbol: PyLong_FromLong	(./libtriton.so)
undefined symbol: PyModule_Create2	(./libtriton.so)
undefined symbol: PyLong_FromSsize_t	(./libtriton.so)
undefined symbol: PyObject_SetAttrString	(./libtriton.so)
undefined symbol: PyType_IsSubtype	(./libtriton.so)
undefined symbol: PyErr_Fetch	(./libtriton.so)
undefined symbol: PyException_SetTraceback	(./libtriton.so)
undefined symbol: PyDict_Copy	(./libtriton.so)
undefined symbol: _PyThreadState_UncheckedGet	(./libtriton.so)
undefined symbol: PyDict_Size	(./libtriton.so)
undefined symbol: PyUnicode_DecodeUTF8	(./libtriton.so)
undefined symbol: PyUnicode_FromString	(./libtriton.so)
undefined symbol: PyNumber_Check	(./libtriton.so)
undefined symbol: PyThread_tss_create	(./libtriton.so)
undefined symbol: PyObject_HasAttrString	(./libtriton.so)
undefined symbol: PyObject_GetBuffer	(./libtriton.so)
undefined symbol: PyBytes_AsString	(./libtriton.so)
undefined symbol: PyCapsule_GetContext	(./libtriton.so)
undefined symbol: _PyType_Lookup	(./libtriton.so)
undefined symbol: PyThreadState_Clear	(./libtriton.so)
undefined symbol: PyErr_Clear	(./libtriton.so)
undefined symbol: PyCallable_Check	(./libtriton.so)
undefined symbol: PyNumber_Xor	(./libtriton.so)
undefined symbol: PyCapsule_New	(./libtriton.so)
undefined symbol: PyMem_Calloc	(./libtriton.so)
undefined symbol: PyErr_ExceptionMatches	(./libtriton.so)
undefined symbol: PyBuffer_Release	(./libtriton.so)
undefined symbol: PyBytes_Size	(./libtriton.so)
undefined symbol: PyNumber_Or	(./libtriton.so)
undefined symbol: PySequence_Size	(./libtriton.so)
undefined symbol: PyLong_AsUnsignedLong	(./libtriton.so)
undefined symbol: PyEval_InitThreads	(./libtriton.so)
undefined symbol: PyNumber_Long	(./libtriton.so)
undefined symbol: PyType_Ready	(./libtriton.so)
undefined symbol: PyTuple_GetItem	(./libtriton.so)
undefined symbol: _PyObject_GetDictPtr	(./libtriton.so)
undefined symbol: PyObject_Malloc	(./libtriton.so)
undefined symbol: PyTuple_SetItem	(./libtriton.so)
undefined symbol: PyThread_tss_get	(./libtriton.so)
undefined symbol: PyErr_Restore	(./libtriton.so)
undefined symbol: PyThread_tss_alloc	(./libtriton.so)
undefined symbol: PyLong_FromSize_t	(./libtriton.so)
undefined symbol: PyList_Size	(./libtriton.so)
undefined symbol: PyGILState_GetThisThreadState	(./libtriton.so)
undefined symbol: PyBytes_AsStringAndSize	(./libtriton.so)
undefined symbol: PyUnicode_AsEncodedString	(./libtriton.so)
undefined symbol: PyTuple_New	(./libtriton.so)
undefined symbol: PyObject_CallFunctionObjArgs	(./libtriton.so)
undefined symbol: PyErr_Occurred	(./libtriton.so)
undefined symbol: PyThreadState_DeleteCurrent	(./libtriton.so)
undefined symbol: _ZN3c105ErrorC1ENS_14SourceLocationESs	(./libtriton.so)
undefined symbol: PySequence_Tuple	(./libtriton.so)
undefined symbol: PyWeakref_NewRef	(./libtriton.so)
undefined symbol: PyUnicode_AsUTF8String	(./libtriton.so)
undefined symbol: PyEval_GetBuiltins	(./libtriton.so)
undefined symbol: PyUnicode_FromFormat	(./libtriton.so)
undefined symbol: PyNumber_And	(./libtriton.so)
undefined symbol: PyCapsule_GetName	(./libtriton.so)
undefined symbol: PyEval_AcquireThread	(./libtriton.so)
undefined symbol: PyErr_SetString	(./libtriton.so)
undefined symbol: PyObject_CallObject	(./libtriton.so)
undefined symbol: PyThread_tss_set	(./libtriton.so)
undefined symbol: PyList_GetItem	(./libtriton.so)
undefined symbol: PyObject_IsInstance	(./libtriton.so)
undefined symbol: PyList_Append	(./libtriton.so)
undefined symbol: PyMem_Free	(./libtriton.so)
undefined symbol: PyThreadState_Get	(./libtriton.so)
undefined symbol: PyThreadState_New	(./libtriton.so)
undefined symbol: PyObject_Repr	(./libtriton.so)
undefined symbol: PyFrame_GetLineNumber	(./libtriton.so)
undefined symbol: PyErr_NormalizeException	(./libtriton.so)
undefined symbol: PyObject_ClearWeakRefs	(./libtriton.so)
undefined symbol: PyDict_Next	(./libtriton.so)
undefined symbol: PyCapsule_GetPointer	(./libtriton.so)
undefined symbol: PyInstanceMethod_New	(./libtriton.so)
undefined symbol: PyDict_New	(./libtriton.so)
undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKSs	(./libtriton.so)
undefined symbol: PyObject_RichCompareBool	(./libtriton.so)
undefined symbol: PyModule_AddObject	(./libtriton.so)
undefined symbol: PySequence_GetItem	(./libtriton.so)
undefined symbol: PyObject_GetItem	(./libtriton.so)
undefined symbol: PyObject_SetItem	(./libtriton.so)
undefined symbol: PyDict_Contains	(./libtriton.so)
undefined symbol: PyDict_GetItemString	(./libtriton.so)
undefined symbol: _ZN3c106detail12infer_schema20make_function_schemaEOSsS2_NS_8ArrayRefINS1_11ArgumentDefEEES5_	(./libtriton.so)
undefined symbol: PyLong_AsLong	(./libtriton.so)
undefined symbol: PyEval_SaveThread	(./libtriton.so)
undefined symbol: PyCapsule_SetContext	(./libtriton.so)
undefined symbol: Py_GetVersion	(./libtriton.so)
undefined symbol: PySequence_Check	(./libtriton.so)
undefined symbol: PyObject_GetAttrString	(./libtriton.so)
undefined symbol: PyList_New	(./libtriton.so)
undefined symbol: PyObject_Str	(./libtriton.so)
undefined symbol: PyErr_Format	(./libtriton.so)
undefined symbol: PyObject_SetAttr	(./libtriton.so)
undefined symbol: PyCFunction_NewEx	(./libtriton.so)
undefined symbol: PyTuple_Size	(./libtriton.so)

I’ve tried installing Triton through pip as well as building from source, and I have the same problem in either case. I have llvm-9-dev installed which is the only dependency listed on Triton’s documentation. There was already an open issue on Triton’s github which looked similar to mine – https://github.com/ptillet/triton/issues/42, however the author has not replied yet. Do you have any idea why this might be happening? Let me know if you need any more information. Thank you.

EDIT: If I set -D_GLIBCXX_USE_CXX11_ABI=1 in CMakeLists.txt before compiling, I no longer get an import error. However, now running python einsum.py gives the following error:

Traceback (most recent call last):
  File "einsum.py", line 195, in <module>
    triton.ops.einsum(expr, ta, tb, tc, arrays = arrays, bench = True)
AttributeError: module 'triton' has no attribute 'ops'

And if I try python -c import triton.ops, I get this error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/triton/python/triton/ops/__init__.py", line 1, in <module>
    from .einsum import _einsum, einsum
  File "/home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/triton/python/triton/ops/einsum.py", line 19, in <module>
    class _einsum(torch.autograd.Function):
  File "/home/nathaniel/Programs/anaconda3/envs/nlp/lib/python3.8/site-packages/triton/python/triton/ops/einsum.py", line 642, in _einsum
    registry = triton.utils.id_dict()
AttributeError: module 'triton' has no attribute 'utils'

[anti4m]

Alright, I’ll try that now and let you know. Thanks.

[arashashari]

@anti4m: for installing Sparse Attention, you have the following dependencies:

• llvm: sudo apt-get install llvm-9-dev
• cmake: sudo apt-get install cmake
• triton>=0.2.2: this will be handled by the installer script, if there was any issue, please install it through pip Then you need to set the install flag: DS_BUILD_SPARSE_ATTN=1 (for sparse attention) Based on the information you provided above, seems like you have all dependencies respected. However, you still have issue with importing Triton from Python. Here is one suggestion: with this commit, we confirmed installing Sparse Attention successfully. Could you please change your Triton source code accordingly and check again? (@jeffra maybe able to comment on this as well) Meanwhile, I will look more into this and will contact Triton owner for his insight on it.