triton
triton copied to clipboard
triton-lang
Allow setting kwargs to special value during autotuning
Some of our kernels have side-effects, such as incrementing an extrernal counter. This counter is consumed in "real-time" by a concurrent kernel.
It breaks our code if this counter is increased more than once per (logical) kernel launch. We'd thus like a way to disable this increment while autotuning, and only do it for the "final" actual kernel run.
We cannot use the restore feature because this wouldn't prevent the concurrent kernel from seeing the update. Instead, we'd love a way for our kernel to "know" if it's still autotuning, so that we can just disable this increment. Hence in this PR I propose a way to set the kwargs of a kernel to some particular value during autotune.
Why not using triton.Config(kwargs={'BLOCK_SIZE': 32, 'skip_increment': True})?
@Jokeren if I set that value in the config it will be set both during benchmarking and during the "real" launch, right? What I want is to only set this value during benchmarking, but disable it for real launches.
Could you take a look into the do_not_specialize option? In this case, if this is not a constexpr, we could skip specialization so that skip_increment=True and skip_increment=False generate the same kernel
I don't see how that would help: do_not_specialize only affects the triton.jit decorator, but my issue isn't about compiling the kernel multiple times, it's about running the kernel multiple times during autotuning.
I'll give another (toy) example of where this feature would be helpful, to try to be clearer. Imagine I have a kernel that prints one line each time it runs. During the first time I launch the kernel it would actually print many hundreds of lines, because the autotuning step launches this kernel over and over in order to benchmark it. I want it to print exactly one line each time, I would need a way to disable this print during benchmarking and only turn the print back on for the "final" actual run. Does this make sense?
OK, I see it now. Thanks for the explanation. Can you add some comments to overrides?
