The compiler_params= argument of pl.pallas_call on GPU now uses "triton" to refer to Triton-specific parameters, instead of the repetitive "triton_params"

PiperOrigin-RevId: 623275152

Author: superbobry

jax/_src/pallas/triton/pallas_call_registration.py

@@ -145,15 +145,14 @@ def pallas_call_lowering(
     raise NotImplementedError(
         "dynamic grid bounds not supported in the Triton backend"
     )
-  triton_compiler_params = compiler_params.get("triton", compiler_params)
-  triton_params = compiler_params.get("triton_params", {})
-  num_warps = triton_compiler_params.pop("num_warps", 4)
+  triton_params = compiler_params.get("triton", compiler_params)
+  num_warps = triton_params.pop("num_warps", 4)
   if len(ctx.module_context.platforms) > 1:
     raise NotImplementedError("multi-platform lowering for Pallas kernels")
   if ctx.module_context.platforms[0] == "rocm":
-    num_stages = triton_compiler_params.pop("num_stages", 1)
+    num_stages = triton_params.pop("num_stages", 1)
   else:
-    num_stages = triton_compiler_params.pop("num_stages", 3)
+    num_stages = triton_params.pop("num_stages", 3)
 
   if debug:
     print(jaxpr)