[Pallas] Deprecate dictionary compiler_params in favor of dataclass.

PiperOrigin-RevId: 699057658

Author: justinjfu

jax/experimental/pallas/ops/gpu/layer_norm.py

@@ -94,7 +94,7 @@ def layer_norm_forward(
   ]
   method = pl.pallas_call(
       kernel,
-      compiler_params=dict(triton=dict(num_warps=num_warps)),
+      compiler_params=plgpu.TritonCompilerParams(num_warps=num_warps),
       grid=(),
       out_shape=out_shape,
       debug=False,
@@ -215,7 +215,7 @@ def layer_norm_backward(
   out_shape_dx = jax.ShapeDtypeStruct(shape=(n,), dtype=x.dtype)
   method = pl.pallas_call(
       kernel,
-      compiler_params=dict(triton=dict(num_warps=num_warps)),
+      compiler_params=plgpu.TritonCompilerParams(num_warps=num_warps),
       grid=(),
       out_shape=out_shape_dx,
       debug=False,

jax/experimental/pallas/ops/gpu/rms_norm.py

@@ -196,7 +196,7 @@ def rms_norm_backward(
   out_shape_dx = jax.ShapeDtypeStruct(shape=(n,), dtype=x.dtype)
   method = pl.pallas_call(
       kernel,
-      compiler_params=dict(triton=dict(num_warps=num_warps)),
+      compiler_params=plgpu.TritonCompilerParams(num_warps=num_warps),
       grid=(),
       out_shape=out_shape_dx,
       debug=False,

tests/pallas/tpu_pallas_pipeline_test.py

@@ -486,12 +486,11 @@ def _wait_on_prev_dma():
             + [pltpu.SemaphoreType.DMA] * 4
             + inner_allocs
         ),
-        compiler_params=dict(
-            mosaic=dict(collective_id=0,
-                        # must set scoped vmem flag *larger* than below! e.g.:
-                        # flags.FLAGS.xla_tpu_scoped_vmem_limit_kib = 131072
-                        vmem_limit_bytes=int(134217728 * 0.9)  # 0.9 * 128MiB
-                       )
+        compiler_params=pltpu.TPUCompilerParams(
+                    collective_id=0,
+                    # must set scoped vmem flag *larger* than below! e.g.:
+                    # flags.FLAGS.xla_tpu_scoped_vmem_limit_kib = 131072
+                    vmem_limit_bytes=int(134217728 * 0.9)  # 0.9 * 128MiB
         ),
     )