[RUNTIME] Restored do_not_specialize (#374)

2021-11-12 15:06:55 -08:00
parent e66bf76354
commit 01cc3d4503
3 changed files with 35 additions and 7 deletions
--- a/python/triton/code_gen.py
+++ b/python/triton/code_gen.py
@@ -667,7 +667,7 @@ class Kernel:
        stream = ((bits & 0xFFFFFFFFFFFF) ^ mask) - mask
        # stream = torch.cuda.current_stream(device).cuda_stream
        # make key for cache
-        return _triton.runtime.launch(wargs, self.fn.cache_key + cc, self.fn.arg_names, device, stream,
+        return _triton.runtime.launch(wargs, self.fn.do_not_specialize, self.fn.cache_key + cc, self.fn.arg_names, device, stream,
                                      self.fn.bin_cache, num_warps, num_stages, self.add_to_cache, grid)