[Pallas] Fix pallas_call lowering mutating compiler params during Triton lowering.

Addresses: #25714
PiperOrigin-RevId: 712566709

Author: justinjfu

jax/_src/pallas/triton/pallas_call_registration.py

@@ -61,14 +61,14 @@ def pallas_call_lowering(
         "scalar prefetch not implemented in the Triton backend"
     )
   triton_params = compiler_params.get("triton", compiler_params)
-  num_warps = triton_params.pop("num_warps", 4)
+  num_warps = triton_params.get("num_warps", 4)
   num_warps = 4 if num_warps is None else num_warps
   [lowering_platform] = ctx.platforms or ctx.module_context.platforms
   if lowering_platform == "rocm":
-    num_stages = triton_params.pop("num_stages", 1)
+    num_stages = triton_params.get("num_stages", 1)
     num_stages = 1 if num_stages is None else num_stages
   else:
-    num_stages = triton_params.pop("num_stages", 3)
+    num_stages = triton_params.get("num_stages", 3)
     num_stages = 3 if num_stages is None else num_stages
 
   if debug:

tests/pallas/pallas_test.py

@@ -30,19 +30,23 @@
 from jax._src import api_util
 from jax._src import checkify
 from jax._src import config
+from jax._src import core as jax_core
 from jax._src import dtypes
 from jax._src import test_util as jtu
 from jax._src.lax.control_flow.for_loop import for_loop
 from jax._src.pallas import core as pallas_core
+from jax._src.pallas import pallas_call
 from jax._src.pallas.pallas_call import _trace_kernel_to_jaxpr
 from jax.experimental import pallas as pl
 import jax.numpy as jnp
 import numpy as np
 
 if sys.platform != "win32":
   from jax.experimental.pallas import tpu as pltpu
+  from jax.experimental.pallas import triton as plgpu
 else:
   pltpu = None
+  plgpu = None
 
 
 # TODO(sharadmv): Update signatures of pallas_call to correct inputs/outputs.
@@ -2361,5 +2365,47 @@ class PallasCallNamedGridInterpretTest(PallasCallNamedGridTest):
   INTERPRET = True
 
 
+def _find_pallas_call_in_jaxpr(
+    jaxpr: jax_core.Jaxpr) -> jax_core.JaxprEqn | None:
+  for eqn in jaxpr.eqns:
+    call_eqn = None
+    if eqn.primitive == pallas_call.pallas_call_p:
+      call_eqn = eqn
+    elif 'jaxpr' in eqn.params:
+      call_eqn = _find_pallas_call_in_jaxpr(eqn.params['jaxpr'])
+    if call_eqn is not None:
+      return call_eqn
+  return None
+
+
+class PallasCompilerParamsTest(PallasBaseTest):
+  def test_triton_params_consistent_across_double_jit(self):
+    # Test for https://github.com/jax-ml/jax/issues/25714
+    if not jtu.test_device_matches(["gpu"]):
+      self.skipTest("Triton backend only works on GPU.")
+    params = plgpu.TritonCompilerParams(num_warps=8)
+
+    @jax.jit
+    @functools.partial(
+        self.pallas_call, out_shape=jax.ShapeDtypeStruct((), jnp.float32),
+        compiler_params=params)
+    def copy_kernel(x_ref, o_ref):
+      o_ref[...] = x_ref[...]
+
+    @functools.partial(jax.jit, static_argnames=["z"])
+    def plus_z(x, z):
+      return copy_kernel(x+z)
+
+    x = 0.
+    extracted_params = _find_pallas_call_in_jaxpr(
+        plus_z.trace(x, 1).jaxpr).params["compiler_params"]
+    self.assertEqual(plus_z(0., 1.), 1.)
+    self.assertEqual(extracted_params["triton"]["num_warps"], 8)
+    extracted_params = _find_pallas_call_in_jaxpr(
+        plus_z.trace(x, 2).jaxpr).params["compiler_params"]
+    self.assertEqual(plus_z(0., 2.), 2.)
+    self.assertEqual(extracted_params["triton"]["num_warps"], 8)
+
+
 if __name__ == "__main__":
   absltest.main()