Optimize w8a8 quantized matmul kernel (#9412)

vanbasten23 · web-flow · commit 4101ea567597 · 2025-07-01T09:42:01.000-07:00
diff --git a/test/test_pallas.py b/test/test_pallas.py
@@ -1,6 +1,7 @@
 import logging
 import sys
 import unittest
+from unittest.mock import patch
 from absl.testing import parameterized
 
 import torch
@@ -877,6 +878,11 @@ def test_ragged_paged_attention_wrapper_without_dynamo(
         use_dynamo=False,
     )
 
+  def _compute_rel_error(self, x, q_x):
+    return torch.mean(torch.sqrt(torch.mean(torch.square(q_x - x),
+                                            axis=1))) / torch.sqrt(
+                                                torch.mean(torch.square(x)))
+
   def _test_quantized_matmul_int8(
       self,
       dtype,
@@ -885,10 +891,6 @@ def _test_quantized_matmul_int8(
       n_output_features,
       quantize_activation,
       use_dynamo,
-      batch_block_size=None,
-      out_block_size=None,
-      in_block_size=None,
-      atol=1.5,
       n_bits=8,
   ):
     x = torch.randn((bs, n_input_features), dtype=dtype)
@@ -918,17 +920,9 @@ def _test_quantized_matmul_int8(
     scalar_xla = scalar.to('xla')
     if use_dynamo:
 
-      def quantized_matmul_int8_wrapper(x, w_int, scalar, quantize_activation,
-                                        batch_block_size, out_block_size,
-                                        in_block_size):
+      def quantized_matmul_int8_wrapper(x, w_int, scalar, quantize_activation):
         return torch.ops.xla.quantized_matmul_int8(
-            x,
-            w_int,
-            scalar,
-            quantize_activation=quantize_activation,
-            batch_block_size=batch_block_size,
-            out_block_size=out_block_size,
-            in_block_size=in_block_size)
+            x, w_int, scalar, quantize_activation=quantize_activation)
 
       quantized_matmul_int8 = torch.compile(
           quantized_matmul_int8_wrapper, backend="openxla")
@@ -941,46 +935,90 @@ def quantized_matmul_int8_wrapper(x, w_int, scalar, quantize_activation,
         w_int_xla,
         scalar_xla,
         quantize_activation=quantize_activation,
-        batch_block_size=batch_block_size,
-        out_block_size=out_block_size,
-        in_block_size=in_block_size).cpu()
+    ).cpu()
+
+    # print(f'Output max diff: {torch.max(torch.abs(expected - actual))}')
+    # print(f'Output mean diff: {torch.mean(torch.abs(expected - actual))}')
+    rel_error = self._compute_rel_error(expected, actual)
 
     self.assertEqual(actual.shape, expected.shape)
     self.assertEqual(actual.dtype, expected.dtype)
-    self.assertTrue(torch.allclose(actual, expected, atol=atol))
+    self.assertTrue(rel_error < 3e-2)
+
+  @parameterized.product(
+      dtype=[torch.bfloat16
+            ],  # not testing float32 because we haven't tuned for float32 case.
+      quantize_activation=[True],
+      use_dynamo=[True, False],
+  )
+  @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 5,
+                   "This test only works on TPUv5+.")
+  @patch(
+      'torch_xla.experimental.pallas_kernels.quantized_matmul_kernel.get_tpu_version'
+  )
+  def test_quantized_matmul_int8_wrapper_key_exists_in_table(
+      self,
+      get_tpu_version,
+      dtype,
+      quantize_activation,
+      use_dynamo,
+  ):
+    from torch_xla.experimental.pallas_kernels.quantized_matmul_kernel import TUNED_BLOCK_SIZES
+    num_cases_to_test = 2
+    if len(TUNED_BLOCK_SIZES) < num_cases_to_test:
+      self.fail(
+          "Not enough tuned block sizes for quantized matmul int8 test. But we should have {num_cases_to_test} block sizes to test."
+      )
+    input_shapes = []
+    for key in TUNED_BLOCK_SIZES.keys():
+      if len(input_shapes) >= num_cases_to_test:
+        break
+      _, batch_size, n_output_features, n_input_features, *_ = key
+      input_shapes.append((batch_size, n_output_features, n_input_features))
+    tpu_version_to_use = 6
+    get_tpu_version.return_value = tpu_version_to_use
+    for batch_size, n_output_features, n_input_features in input_shapes:
+      self._test_quantized_matmul_int8(
+          dtype,
+          batch_size,
+          n_input_features,
+          n_output_features,
+          quantize_activation,
+          use_dynamo=use_dynamo,
+      )
 
   @parameterized.product(
       dtype=[torch.bfloat16, torch.float32],
       bs=[256, 512],
       n_input_features=[256, 512],
       n_output_features=[256, 512],
       quantize_activation=[True],
-      kernel_block_sizes=[(None, None, None), (256, 256, 256)],
       use_dynamo=[True, False],
   )
   @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 5,
                    "This test only works on TPUv5+.")
-  def test_quantized_matmul_int8_wrapper(
+  @patch(
+      'torch_xla.experimental.pallas_kernels.quantized_matmul_kernel.get_tuned_block_sizes'
+  )
+  def test_quantized_matmul_int8_wrapper_key_not_exists_in_table(
       self,
+      get_tuned_block_sizes,
       dtype,
       bs,
       n_input_features,
       n_output_features,
       quantize_activation,
-      kernel_block_sizes,
       use_dynamo,
   ):
-    batch_block_size, out_block_size, in_block_size = kernel_block_sizes
+    get_tuned_block_sizes.return_value = (None, None, None)
     self._test_quantized_matmul_int8(
         dtype,
         bs,
         n_input_features,
         n_output_features,
         quantize_activation,
         use_dynamo=use_dynamo,
-        batch_block_size=batch_block_size,
-        out_block_size=out_block_size,
-        in_block_size=in_block_size)
+    )
 
   @unittest.skipIf(xr.device_type() != 'TPU' or tpu.version() < 4,
                    "This test only works on TPUv4+.")
diff --git a/test/test_quantized_matmul_pallas_kernel.py b/test/test_quantized_matmul_pallas_kernel.py
@@ -128,7 +128,8 @@ def test_quantized_matmul_retrieve_block_sizes(self, get_tpu_version):
         break
     expected_block_sizes = TUNED_BLOCK_SIZES[key0]
     _, bs, n_output_features, n_input_features, activation_dtype, quantize_activation = key0
-    actual_block_sizes = get_tuned_block_sizes(bs, n_output_features,
+    actual_block_sizes = get_tuned_block_sizes(TUNED_BLOCK_SIZES, bs,
+                                               n_output_features,
                                                n_input_features,
                                                activation_dtype,
                                                quantize_activation)
@@ -145,12 +146,17 @@ def test_quantized_matmul_use_tuned_block_sizes(self, dtype, bs,
                                                   n_input_features,
                                                   n_output_features,
                                                   quantize_activation):
-    self._test_quantized_matmul(
-        dtype,
-        bs,
-        n_input_features,
-        n_output_features,
-        quantize_activation=quantize_activation)
+    with self.assertRaises(AssertionError):
+      self._test_quantized_matmul(
+          dtype,
+          bs,
+          n_input_features,
+          n_output_features,
+          quantize_activation=quantize_activation,
+          batch_block_size=None,
+          out_block_size=None,
+          in_block_size=None,
+      )
 
 
 if __name__ == "__main__":
diff --git a/torch_xla/experimental/custom_kernel.py b/torch_xla/experimental/custom_kernel.py
@@ -1075,17 +1075,32 @@ def quantized_matmul_int8(
     in_block_size: int | None = None,
     vmem_limit_bytes: int | None = 64 * 1024 * 1024,
 ) -> torch.Tensor:
-  from torch_xla.experimental.pallas_kernels.quantized_matmul_kernel import quantized_matmul_int8
-  return xb.call_jax(
-      quantized_matmul_int8, (x, w, scalar), {
-          "zero_point": zero_point,
-          "quant_block_size": quant_block_size,
-          "quantize_activation": 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
-      })
+  from torch_xla.experimental.pallas_kernels.quantized_matmul_kernel import (
+      quantized_matmul_int8,
+      get_tuned_block_sizes,
+      TUNED_BLOCK_SIZES,
+  )
+  bs, n_in_features = x.shape
+  n_out_features, _ = w.shape
+  jax_dtype = convert_torch_dtype_to_jax(x.dtype)
+  import jax.numpy as jnp
+  batch_block_size, out_block_size, in_block_size = get_tuned_block_sizes(
+      TUNED_BLOCK_SIZES, bs, n_out_features, n_in_features,
+      jnp.dtype(jax_dtype).name, quantize_activation)
+  if batch_block_size is not None and out_block_size is not None and in_block_size is not None:
+    return xb.call_jax(
+        quantized_matmul_int8, (x, w, scalar), {
+            "zero_point": zero_point,
+            "quant_block_size": quant_block_size,
+            "quantize_activation": 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
+        })
+  from torch_xla.experimental.xla_quantized_matmul import quantized_matmul_xla
+  return quantized_matmul_xla(
+      x, w, scalar, quantize_activation=quantize_activation)
 
 
 def _multi_queries_paged_attention_nonkernel(
diff --git a/torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py b/torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py
