GPU pallas_call loses compiler params during second call when double jit-wrapped

[hanzhi713]

Description

If you add a print(triton_params) after this line

jax/jax/_src/pallas/triton/pallas_call_registration.py

Line 63 in 57b2154

triton_params = compiler_params.get("triton", compiler_params)

and run the following reproducer, you'll get

{'num_warps': 8}
{}

This causes some performance problems in production as kernels can't get the right compiler params.

Reproducer

from functools import partial

import jax
from jax.experimental import pallas as pl
import jax.numpy as jnp


def add_vectors_kernel(x_ref, y_ref, o_ref):
  x, y = x_ref[...], y_ref[...]
  o_ref[...] = x + y
  

@partial(jax.jit, static_argnames="z")
def dummy(x, y, z):
    x += z
    return add_vectors(x, y)
  
  
@jax.jit
def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
  return pl.pallas_call(
      add_vectors_kernel,
      out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
      compiler_params=dict(triton=dict(num_warps=8))
  )(x, y)
  
a = jnp.arange(8)
b = jnp.arange(8)
dummy(a, b, 1)
dummy(a, b, 2)

System info (python version, jaxlib version, accelerator, etc.)

jax                                      0.4.39.dev20250102
jax-cuda12-pjrt                          0.4.39.dev20250102
jax-cuda12-plugin                        0.4.39.dev20250102
jaxlib                                   0.4.39.dev20250102

[justinjfu]

Looks like a silly bug where the params are mutated because we pop 'num_warps' instead of get. #25735 should address.