Integrate Triton up to [6fa4f504](https://github.com/openai/triton/co…

…mmits/6fa4f504d61e48f7cea454ce0c1f6169907d5aaa)

PiperOrigin-RevId: 681365759

jax_triton/triton_lib.py

@@ -538,7 +538,8 @@ def triton_kernel_call_lowering(
   named_args = dict(unsafe_zip(fn.arg_names, args))
 
   if isinstance(fn, autotuner.Autotuner):
-    if any(idx not in fn.key_idx for idx, _, _ in scalar_args):
+    key_idxs = [fn.arg_names.index(k) for k in fn.keys]
+    if any(idx not in key_idxs for idx, _, _ in scalar_args):
       logging.warning(
           "Auto-tuning key does not include all scalar arguments. "
           "We may perform redundant auto-tuning."