Another problem is: aotriton has extended the Triton language a little bit for NPOT head dimension optimization, and consequently its kernel cannot be compiled by stock Triton.

A dedicated venv is created under the build directory regardless the value of AOTRITON_USE_EXTERNAL_TRITON, and the import triton under this venv will not work unless you have already installed one.

It works in the nixpkgs context but maybe nixpkgs is customizing something to default venvs to ineriting from site-wide packages.
Should I enable --system-site-packages when creating the venv if the triton inherit flag is set?

Another problem is: aotriton has extended the Triton language a little bit for NPOT head dimension optimization, and consequently its kernel cannot be compiled by stock Triton.

Does this apply for 0.10 or only current HEAD? Not getting any compile errors on 0.10b with upstream triton 3.3.

If a dedicated Triton is preferred, I suggest let this variable point to a .whl file and install the wheel file into the venv, which is easier to deal with and more flexible.

I don't think this would help at all with how nixpkgs' python packaging works.

Even if we have to pin a specific ROCm/triton sha it's a lot easier on our end if our aotriton build can inherit a site-packages triton that was built from that sha, otherwise we have to duplicate code needed to get triton to work for our aotriton package.

Sample configure log https://gist.github.com/LunNova/b8d22fa28ec7ccfe822c766011b2de58

Does this apply for 0.10 or only current HEAD? Not getting any compile errors on 0.10b with upstream triton 3.3.

Please double check the actual Triton version you are using to compile. In current code framework the only way to use upstream Triton is to call test_*.py under tritonsrc/. The .bit_length extension is not upstreamed yet.

It works in the nixpkgs context but maybe nixpkgs is customizing something to default venvs to ineriting from site-wide packages.

If that's true the problem is the AOTRITON_USE_EXTERNAL_TRITON does not do what is was named for. I don't have any nixos environment to validate, but the code path basically does the following thing

python -m venv build/venv
build/venv/bin/pip install -r requirements.txt
build/venv/bin/python -c "import triton; import triton.language; print('Triton found: ', triton.__version__)"

Unless nixos hacked venv pacakge, the last setup should always fail because it is intentionally avoided to add heavy weight packages (e.g., torch/triton) to requirements.txt.

In addition, to make the code above work, the standard practice is to use --system-site-packages, which means your PR only needs the current AOTRITON_VENV_INHERIT_SITE_PACKAGES (or rename it to AOTRITION_INHERIT_SYSTEM_SITE_TRITON for better self-describing)

The supposed logic of AOTRITON_USE_EXTERNAL_TRITON, from its name, should be:

python -m venv build/venv
build/venv/bin/pip install -r requirements.txt
build/venv/bin/pip install ${AOTRITON_USE_EXTERNAL_TRITON}
# Skip triton build later

Unless nixos hacked venv pacakge, the last setup should always fail because it is intentionally avoided to add heavy weight packages (e.g., torch/triton) to requirements.txt.

Yeah - it looks like nixos hacked venv creation to implicitly inherit global site-packages, hence it working earlier.

I've pushed an additional commit which adds a flag to enable inheriting site packages when creating the venv so the external triton option is usable on other distros.

To confirm, you'd be ok with a single flag AOTRITION_INHERIT_SYSTEM_SITE_TRITON which enables --system-site-packages and verifies triton is available in the venv without trying to pip install it?

To confirm, you'd be ok with a single flag AOTRITION_INHERIT_SYSTEM_SITE_TRITON which enables --system-site-packages and verifies triton is available in the venv without trying to pip install it?

Yes.
This is a more natural solution IMO.

The only concern is unlike its name it inherits everything. But the other naming scheme is also problematic: users are not informed this option expects Triton from system site and they may prefer a "clean" system site and avoid installing triton there.

Compared with the second option I think the first one's risk is lower: venv users should easily guess the "inheritance" is achieved through --system-site-packages

Updated to use AOTRITION_INHERIT_SYSTEM_SITE_TRITON

Thanks!