Add TRITON_IGNORE_LIBTRITON_HASH for hash stability between Python versions

[woct0rdho]

This is for reducing the disk usage of the JIT-compiled CUDA binaries (such as add_kernel.cubin). For example, when running a large AI model such as Wan with torch.compile autotune, the cached binaries can take 500 MB of disk space, and I would like to avoid recompiling them when switching the Python version.

Conceptually, the CUDA binaries generated by Triton should not depend on the Python version (they don't involve any C-Python binding), and I've tested that their binary contents are actually the same between Python versions. However, in the current implementation, the cache folder hash depends on libtriton, which involves the C-Python binding. Also, the binary content of libtriton can depend on some debug infomation such as the tmp path where it's compiled.

I propose to add the environment variable TRITON_IGNORE_LIBTRITON_HASH, which can opt in to ignore this dependency. Then the cache can be reused between Python versions, and the user is responsible for clearing the cache if they somehow updated the functionality of libtriton without changing the rest Python code in the package.

New contributor declaration

• I am not making a trivial change, such as fixing a typo in a comment.

• I have written a PR description following these
  rules.

• I have run pre-commit run --from-ref origin/main --to-ref HEAD.

• Select one of the following.

  • I have added tests.
    • /test for lit tests
    • /unittest for C++ tests
    • /python/test for end-to-end tests
  • This PR does not need a test because: The test requires two Python versions, which is out of the current scope of pytest. I've tested locally.

• Select one of the following.

  • I have not added any lit tests.
  • The lit tests I have added follow these best practices,
    including the "tests should be minimal" section. (Usually running Python code
    and using the instructions it generates is not minimal.)

[ThomasRaoux]

When triton changes we should definitely invalidate the cache. I don't understand the description, this is not about updating python, if you use a different triton build you can't expect the cache to be compatible.

[woct0rdho]

Yes we should definitely invalidate the cache when Triton's functionality changes. By 'functionality' I mostly mean how we compile the Triton language -> TTIR -> PTX -> cubin.

Conceptually, this functionality should be the same in all Python versions (3.9 to 3.13). This is why I say the compiled cubin should not depend on the Python version.

However, in the current implementation, the cache hash depends on the binary content of libtriton, which not only includes the above functionality of Triton, but also includes the C-Python binding (and other things like debug information). When I change the Python version, even if the compiled cubin is exactly the same, the cache hash still changes. That's the problem.

In this PR I provide an opt-in way to ignore this dependency. Maybe there is a better way to define the 'functionality' of Triton in the cache hash, such as the Triton package version or the git hash. What do you think?

[ThomasRaoux]

Yes we should definitely invalidate the cache when Triton's functionality changes. By 'functionality' I mostly mean how we compile the Triton language -> TTIR -> PTX -> cubin.

Conceptually, this functionality should be the same in all Python versions (3.9 to 3.13). This is why I say the compiled cubin should not depend on the Python version.

However, in the current implementation, the cache hash depends on the binary content of libtriton, which not only includes the above functionality of Triton, but also includes the C-Python binding (and other things like debug information). When I change the Python version, even if the compiled cubin is exactly the same, the cache hash still changes. That's the problem.

I don't follow how changing your python version changes libtriton.so, are you saying that you rebuild libtriton with different python headers?

In this PR I provide an opt-in way to ignore this dependency.

But this also ignores a lot of real dependencies, I don't see how this can ever be correct.

Maybe there is a better way to define the 'functionality' of Triton in the cache hash, such as the Triton package version or the git hash. What do you think?

I don't think we want the build to depend on git.

I do think that any rebuild of triton backend should invalidate the cache and there is no way to know what could have changed in the binary. Are you rebuilding triton with different python headers that often?

[woct0rdho]

libtriton.so does not use the Python stable ABI, so we need to build a different libtriton.so for each Python version. That's discussed in #5879

The problem is not that I rebuild Triton. Usually (when I release my Windows wheels) I build the same version of Triton for all Python versions (3.9 to 3.13), and users may switch between them.

If we use the git hash to define the 'functionality' of Triton, we can get this value and add it to __init__.py or libtriton.so when building the package, then the package does not depend on git when the user uses it. Currently setup.py already depends on git to add this value in the metadata for pip (such as triton 3.3.0+git917437a4), but not in the package itself (__init__.py or libtriton.so) yet. See the new commit.

[woct0rdho]

Now I get the git hash in the package version using importlib. It's a Python stdlib and does not introduce any new package as dependency.

The package version is like triton 3.3.0+git917437a4, which is set in setup.py when building the package (either normal installation/editable installation/wheel). The package does not depend on git after it's built (when the user uses it).

We can assume that the git hash changes whenever the functionality of libtriton changes, and it does not change when the Python version changes.

Also, I rebased because of #6467 . What do you think?

[ThomasRaoux]

Now I get the git hash in the package version using importlib. It's a Python stdlib and does not introduce any new package as dependency.

The package version is like triton 3.3.0+git917437a4, which is set in setup.py when building the package (either normal installation/editable installation/wheel). The package does not depend on git after it's built (when the user uses it).

We can assume that the git hash changes whenever the functionality of libtriton changes, and it does not change when the Python version changes.

Also, I rebased because of #6467 . What do you think?

it still feels like a footgun to me as any changes not committed would not be considered by the caching, right?

[danzimm]

@ThomasRaoux what do you think about allowing triton_key hook / config? I know on our side of the world we have a patch to pre-compute the hash key and freeze it, so selfishly we could replace the patch with that knob. On @woct0rdho 's side they could implement the hook with the impl in this PR.

[woct0rdho]

it still feels like a footgun to me as any changes not committed would not be considered by the caching, right?

Indeed, currently the git hash in setup.py does not consider the uncommitted changes. If we really want we can add a few git commands in setup.py to compute a hash for the uncommitted changes, such as https://stackoverflow.com/a/48213033/29980299 . This is also what setuptools-scm or versioningit may do.

[Jokeren]

@ThomasRaoux what do you think about allowing triton_key hook / config? I know on our side of the world we have a patch to pre-compute the hash key and freeze it, so selfishly we could replace the patch with that knob. On @woct0rdho 's side they could implement the hook with the impl in this PR.

The idea seems fine to me, but the description of the PR has to be updated.

[woct0rdho]

@danzimm I agree that a hook/config to directly provide triton_key also works in my case, and it's simpler than all the git crafts. After all, this config is opt-in, and the user should know what they're doing. In my case it's meant to be used in the stable production environment rather than the development.