non-dynamo case passed.

Author: vanbasten23

test/test_pallas.py

@@ -939,8 +939,9 @@ def quantized_matmul_wrapper(x, w_int, scalar):
 
     self.assertEqual(actual.shape, expected.shape)
     self.assertEqual(actual.dtype, expected.dtype)
-    torch.testing.assert_close(
-        actual, expected, atol=1.5)
+    self.assertTrue(
+        torch.allclose(
+            actual, expected, atol=atol))
 
 
   @parameterized.product(
@@ -967,19 +968,19 @@ def test_quantized_matmul_with_dynamo(
     ...
 
   # @parameterized.product(
-  #     dtype=[torch.bfloat16, torch.float32],
-  #     bs=[128, 256],
-  #     n_input_features=[128, 256],
-  #     n_output_features=[128, 256],
+  #     dtype=[torch.bfloat16],
+  #     bs=[128],
+  #     n_input_features=[128],
+  #     n_output_features=[128],
   #     quantize_activation=[True],
   #     # block_sizes=[(None, None, None), (128, 128, 128)],
-  #     block_sizes=[(128, 128, 128)],
+  #     kernel_block_sizes=[(128, 128, 128)],
   # )
   @parameterized.product(
-      dtype=[torch.bfloat16],
-      bs=[128],
-      n_input_features=[128],
-      n_output_features=[128],
+      dtype=[torch.bfloat16, torch.float32],
+      bs=[128, 256],
+      n_input_features=[128, 256],
+      n_output_features=[128, 256],
       quantize_activation=[True],
       # block_sizes=[(None, None, None), (128, 128, 128)],
       kernel_block_sizes=[(128, 128, 128)],

torch_xla/experimental/custom_kernel.py

@@ -1068,7 +1068,6 @@ def quantized_matmul(
   scalar: torch.Tensor,
   zero_point: torch.Tensor | None = None,
   block_size: torch.Tensor | None = None,
-  int4_weight: bool = False,
   quantize_activation: bool = False,
   batch_block_size: int | None = None,
   out_block_size: int | None = None,
@@ -1078,7 +1077,7 @@ def quantized_matmul(
   from torch_xla.experimental.pallas_kernels.quantized_matmul_kernel import quantized_matmul
   return xb.call_jax(
     quantized_matmul, 
-    (x, w, scalar, zero_point, block_size, int4_weight, quantize_activation), 
+    (x, w, scalar, zero_point, block_size, quantize_activation), 
     {"batch_block_size": batch_block_size, "out_block_size": out_block_size, "in_block_size": in_block_size, "vmem_limit_bytes": vmem_limit_bytes}
   )