Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

triton requirement is out of date #238

Open
ywrt opened this issue Oct 29, 2023 · 1 comment
Open

triton requirement is out of date #238

ywrt opened this issue Oct 29, 2023 · 1 comment

Comments

@ywrt
Copy link

ywrt commented Oct 29, 2023

Attempting to use jax-triton nightly fails on CUDA with

 File "/home/michael/.local/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 1433, in jaxpr_subcomp
    ans = rule(rule_ctx, *rule_inputs, **eqn.params)
  File "/home/michael/.local/lib/python3.10/site-packages/jax_triton/triton_lib.py", line 383, in triton_kernel_call_lowering
    kernel, specialization = get_or_create_triton_kernel(
  File "/home/michael/.local/lib/python3.10/site-packages/jax_triton/triton_lib.py", line 252, in get_or_create_triton_kernel
    module = code_gen.ast_to_ttir(
TypeError: ast_to_ttir() got an unexpected keyword argument 'target'

This appears to be because requirements.txt lists triton_nightly-2.1.0.dev20230714011643, but a newer triton nightly is actually required.
(That version doesn't have a 'target' argument to ast_to_ttir).

I would around this by installing triton_nightly-2.1.0.dev20231014192330, but it would be great to have correct dependencies.

@sharadmv
Copy link
Collaborator

Thanks for investigating, I probably won't get around to this for a week or so. A PR is greatly appreciated though!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants