[Inductor] Cache generated user defined triton kernels on tensor dtype and non tensor parameters

[oulgen]

Stack from ghstack (oldest at bottom):

• [AOTI] Support non auto-tuned triton kernels in aoti #113090
• [Inductor] Allow None values to be passed in as arguments to triton kernels #113056
• [Inductor] Improve reinplace_scatters pass #112801
• Move all triton related testing utils into shared file #113008
• -> [Inductor] Cache generated user defined triton kernels on tensor dtype and non tensor parameters #112752

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/112752

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (3 Unrelated Failures)

As of commit 6cf317d with merge base 7715b47 ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_distributed, 1, 1, linux.g5.12xlarge.nvidia.gpu) (gh)
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
• pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 1, 3, linux.8xlarge.nvidia.gpu) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[oulgen]

The cache key needs to be dtype of all tensors and values of all non tensors since our triton autotune implementation assumes every input is implicitly in "keys" list.

[aakhundov]

Can we run into a (rather weird, but possible) edge case that the user defines two functionally different Triton kernels in different scopes with the same name and things added to the cache_key below?

[oulgen]

This caching is for the same fx graph so I dont think it is possible for user to unambiguously define two kernels with same name.

[jansel]

I don't think the kernel name is enough for a cache key.

from file1 import my_cool_triton_kernel as k1
from file2 import my_cool_triton_kernel as k2

or

import file1
import file2
...
file1.call_my_cool_triton_kernel(...)
file2.call_my_cool_triton_kernel(...)

[aakhundov]

... since our triton autotune implementation assumes every input is implicitly in "keys" list.

Does our autotune implementation not assume that the key doesn't exist? If the key is taken into account, how is it used? Thanks!

[oulgen]

Does our autotune implementation not assume that the key doesn't exist? If the key is taken into account, how is it used? Thanks!

our autotune implementation assumes that all non tensor arguments to kernel are part of the key (because we specialize on all non tensor values -- this is how dynamo works)

[jansel]

I don't think the kernel name is enough for a cache key.

from file1 import my_cool_triton_kernel as k1
from file2 import my_cool_triton_kernel as k2

or

import file1
import file2
...
file1.call_my_cool_triton_kernel(...)
file2.call_my_cool_triton_kernel(...)

[oulgen]

Swapped over to using code id, added a test case where kernel.__name__ is the same but we do not cache since code id does not match

[zou3519]

Is it possible that kernel.fn ends up being GC'ed in Python and we end up with an id that points to nothing? We've run into some situations in the past where something similar has caused issues

[oulgen]

@zou3519 This is not a part of python I am familiar with. What would you recommend using in place of id?
In general, why would kernel.fn get GCed when we hold a reference to it?

[zou3519]

I read the code a bit more. We should be fine here. As long as we hold a reference to kernel.fn then it won't get GC'ed (and it looks like we keep one in the kernel_side_table).

In some other places in Dynamo (the allowed_functions dict) we use the id of a nested function as a key to a dictionary but don't store a reference to it, resulting in the object getting GC'ed