[BUILD] Support building out of tree plugins

[nhat-nguyen]

This PR adds support for building out of tree plugins with triton.
External plugins can provide an environment variable before invoking python setup.py to let triton know where to find the backends.

• TRITON_PLUGIN_DIRS: semicolon-separated list of directories containing the plugins. The last components of the paths will be used as plugin names.

For instance, if we have triton_shared under /home/user/github/triton_shared, to build triton_shared with triton, we can do:

export TRITON_PLUGIN_DIRS="/home/user/github/triton_shared"
python3 setup.py build

This also assumes that the external backends have moved over to the new external backend architecture; they have to expose a backend folder under their root directory.

With this change, all external backends will still be copied to triton/runtime and be used like:

import triton
from triton.backends.triton_shared.driver import CPUDriver

triton.runtime.driver.set_active(CPUDriver())

[manbearian]

@nhat-nguyen Thanks for doing this work!

The changes look good; i can confirm this should work for the OSS triton-shared project as well as what Microsoft needs to support Maia 100.

[ptillet]

seems like there is a lot of copy pasted code between copy_backends and copy_external_backends

[etiotto]

This PR provides a way for external backends to be built and included into Triton so it has value. Users will need to clone the external backend and so they will actually need to be aware that the backend exists somewhere in a different github repository.

I personally would prefer an external backend to be a submodule. This has the advantage of not requiring users to explicitly clone the third party code, and makes it immediately clear that the backend exists. If the submodule is a complete fork of the OpenAI code (plus vendor changes of course) there would be no maintenance costs in the OpenAI project and no extra efforts for the end users of the project.

[ptillet]

I think this looks overengineered now... I am not convinced that setup.py needs +60lines to discovered external backends...

[nhat-nguyen]

I think this looks overengineered now... I am not convinced that setup.py needs +60lines to discovered external backends...

Thanks! I just simplified this the best I could and kept almost the same code structure as before. Let me know what you think!

[nhat-nguyen]

@ptillet When you have some cycles, would you mind taking a look at this again? Thank you!

[nhat-nguyen]

@ptillet I played around with this idea a bit more and would like to hear what you think.

I feel automatically using the last component of the plugin path isn't good a good dev experience since it requires the directory to also be a valid python identifier (for when we actually copy the backend dirs to triton/runtime during the build). So backends containing dashes - have to be renamed during clone to use correctly.

Using a list like name1,path1;name2,path2 is also error-prone and makes the cmake / python code more complicated because we have to parse strings; paths can't also contain commas otherwise everything will be messed up.

What are your thoughts on requiring the plugin backend to contain a file describing the name to use? Like in the case of triton-shared, we will put something under triton-shared/backend a file called name.conf with the python module name to be used? Thank you!

[aaronsm]

I think this could be done similar to LLVM_TARGETS_TO_BUILD.

$ TRITON_CODEGEN_BACKENDS="amd;nvidia;triton-shared" \
  pip3 install -e python

Here are example changes -- https://github.com/aaronsm/triton/commit/17c6f6f08d8f835d30c9cf0294f02f57c0b679e1

[nhat-nguyen]

I think this could be done similar to LLVM_TARGETS_TO_BUILD.

$ TRITON_CODEGEN_BACKENDS="amd;nvidia;triton-shared" \
  pip3 install -e python

Here are example changes -- aaronsm@dbecc92

The issue though is both triton-shared and intel xpu are no longer in-tree. Only nvidia and amd are.

During the last community meeting we had a discussion about this but didn’t reach a definitive conclusion.

Phil mentioned that third-party plugins will now have to include triton as a submodule. With this requirement, we will need a way to tell triton where to build the plugins through an env variable.

We can still definitely reuse the old approach like you suggested, but that means 1) we will now have a fork of triton that includes the source code of the backend and 2) it doesn’t quite match the structure that Phil suggested: having triton as a submodule.