Skip Pallas and Mosaic GPU tests that don't fit on RTX 6000 PRO

Author: andportnoy

jax/_src/test_util.py

@@ -468,6 +468,47 @@ def is_device_tpu(version: int | None = None, variant: str = "") -> bool:
     return "v6 lite" in device_kind
   return expected_version in device_kind
 
+def device_kind_match(device_patterns: str | Sequence[str]):
+  if isinstance(device_patterns, list):
+    device_patterns = tuple(device_patterns)
+  elif not isinstance(device_patterns, tuple):
+    device_patterns = (device_patterns,)  # type: ignore
+
+  current_device = jax.devices()[0].device_kind
+  for device_pattern in device_patterns:
+    if device_pattern in current_device:
+      return device_pattern
+  return None
+
+def skip_if_device_kind_matches(device_patterns: str | Sequence[str]):
+  """A decorator for test methods to skip the test when run on specified devices."""
+  matching_pattern = device_kind_match(device_patterns)
+
+  reason = f"Device kind matches \"{matching_pattern}\""
+  def skip(test_method):
+    return unittest.skipIf(matching_pattern is not None, reason)(test_method)
+  return skip
+
+def skip_if_exceeds_shared_memory_on(device_patterns: str | Sequence[str]):
+  """Skip if a kernel exceeds shared memory on a specified device."""
+  def skip(test_method):
+    @functools.wraps(test_method)
+    def test_method_wrapper(self, *args, **kwargs):
+      try:
+        return test_method(self, *args, **kwargs)
+      except Exception as e:
+        mosaic_gpu_pattern = "kernel exceeds available shared memory"
+        triton_pattern = "Shared memory size limit exceeded"
+        shared_memory_exceeded = (mosaic_gpu_pattern in str(e) or
+                                  triton_pattern in str(e))
+        matching_device = device_kind_match(device_patterns)
+        if shared_memory_exceeded and matching_device is not None:
+          reason = f"Kernel exceeds shared memory on {matching_device}"
+          self.skipTest(reason)
+        raise
+    return test_method_wrapper
+  return skip
+
 def is_cuda_compute_capability_at_least(capability: str) -> bool:
   if not is_device_cuda():
     return False

tests/mosaic/gpu_test.py

@@ -664,6 +664,7 @@ def kernel(ctx, inp, out, smem):
           fa.WGMMA_LAYOUT_UPCAST_4X,
       ),
   )
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_optimized_conversion(self, jax_dtype_from_to, layout):
     jax_dtype_from, jax_dtype_to = jax_dtype_from_to
     mlir_dtype_from = utils.dtype_to_ir_type(jax_dtype_from)
@@ -2752,6 +2753,7 @@ def kernel(ctx, src, dst, _):
     np.testing.assert_allclose(result, np.full((128, 32), 3.14, np.float32))
 
   @parameterized.product(in_shape=((128, 128), (128, 64), (64, 128)))
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_strided_load_store(self, in_shape):
     def kernel(ctx, *args):
       gmem_input, gmem_output, (smem_input, smem_output) = args
@@ -3070,6 +3072,7 @@ def kernel(ctx, dst, _):
       row_tiling=[8, 64],
   )
   @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
+  @jtu.skip_if_device_kind_matches("RTX PRO 6000 Blackwell")
   def test_copy_tiled(self, dtype, swizzle, num_col_tiles, row_tiling):
     mlir_dtype = utils.dtype_to_ir_type(dtype)
     bw = bytewidth(mlir_dtype)
@@ -3154,6 +3157,7 @@ def kernel(ctx, in_, out, smems):
           (fa.TCGEN05_LAYOUT, fa.TCGEN05_TRANSPOSED_LAYOUT),
       ],
   )
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_transpose_tiled(self, dtype, swizzle, layouts):
     mlir_dtype = utils.dtype_to_ir_type(dtype)
     bw = bytewidth(mlir_dtype)
@@ -3198,6 +3202,7 @@ def kernel(ctx, in_, out, smems):
       (fa.WGMMA_LAYOUT_UPCAST_4X, fa.WGMMA_LAYOUT, jnp.int4, jnp.int4, 2),
   )
   @jtu.thread_unsafe_test()  # Modifies ``os.environ``.
+  @jtu.skip_if_device_kind_matches("RTX PRO 6000 Blackwell")
   def test_upcast_to_wgmma(
       self, start_layout, end_layout, in_dtype, cast_dtype, shfl_per_reg
   ):

tests/pallas/mosaic_gpu_test.py

@@ -226,6 +226,7 @@ def kernel(x_ref, y_ref, o_ref):
     np.testing.assert_array_equal(kernel(x, y), x + y[0])
 
   @parameterized.product(shape=[(128,), (128, 128)])
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_reduce_sum(self, shape):
     @functools.partial(
         self.pallas_call, out_shape=jax.ShapeDtypeStruct(shape, jnp.float32)
@@ -719,6 +720,7 @@ def kernel(x_ref, o_ref, barrier_ref):
     x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
     np.testing.assert_array_equal(f(x), x)
 
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_scoped_copy_with_transforms(self):
     self.skip_if_wg_semantics()
 
@@ -742,6 +744,7 @@ def body(tmp_ref):
     x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
     np.testing.assert_array_equal(f(x), x * 2)
 
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_scoped_copy_with_user_transforms(self):
     def kernel(x_ref, o_ref, barrier_ref):
       def body(tmp_ref):
@@ -762,6 +765,7 @@ def body(tmp_ref):
     x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
     np.testing.assert_array_equal(f(x), x * 2)
 
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_copy_with_transforms_and_indexing(self):
     self.skip_if_wg_semantics()
 
@@ -811,6 +815,7 @@ def kernel(x_ref, o_ref):
     x = jnp.arange(2 * 128, dtype=jnp.float32).reshape(2, 128)
     np.testing.assert_array_equal(kernel(x), x)
 
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_indexing_before_transpose(self):
     self.skip_if_wg_semantics()
 
@@ -2021,6 +2026,7 @@ def kernel(x_ref, y_ref, smem_ref, smem_out_ref, barrier_ref):
     np.testing.assert_array_equal(kernel(x), jnp.broadcast_to(x[:, None], (128, 128)))
 
   @parameterized.named_parameters((l.name.lower(), l) for l in plgpu.Layout)
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_copy_layout(self, layout):
     self.skip_if_wg_semantics()
     if layout in {

tests/pallas/ops_test.py

@@ -560,6 +560,7 @@ def kernel(*refs):
       for name, func, strategy in UNARY_FUNCTIONS
   )
   @hp.given(hps.data())
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_unary_primitives(self, name, func, shape_dtype_strategy, data):
     if name in ["abs", "log1p", "pow2", "reciprocal", "relu", "sin", "sqrt"]:
       self.skip_if_mosaic_gpu()
@@ -1897,6 +1898,7 @@ def f(x_ref, o_ref):
       trans_x=[False, True],
       trans_y=[False, True],
   )
+  @jtu.skip_if_exceeds_shared_memory_on("RTX PRO 6000 Blackwell")
   def test_dot(self, lhs_and_rhs_shape, dtype, trans_x, trans_y):
     self.skip_if_mosaic_gpu()