[Kernel] Update tile tuning key for future mixed precision support (#715)

Author: kyuyeunk

tests/kernels/quantized_matmul_kernel_test.py

@@ -64,9 +64,7 @@ def _test_quantized_matmul(
         n_input_features: int,
         n_output_features: int,
         quantize_activation: bool,
-        batch_block_size=None,
-        out_block_size=None,
-        in_block_size=None,
+        tuned_value=None,
         atol=0.5,
         rtol=0.5,
     ):
@@ -88,14 +86,13 @@ def _test_quantized_matmul(
         w_scale = jnp.squeeze(w_scale)
         assert w_scale.shape == (n_output_features, )
 
+        x_q_dtype = w_q.dtype if quantize_activation else dtype
         output = quantized_matmul_kernel(
             x,
             w_q,
             w_scale,
-            quantize_activation=quantize_activation,
-            batch_block_size=batch_block_size,
-            out_block_size=out_block_size,
-            in_block_size=in_block_size,
+            x_q_dtype=x_q_dtype,
+            tuned_value=tuned_value,
         )
         expected = reference_quantized_matmul(
             x, w_q, w_scale, quantize_activation=quantize_activation)
@@ -130,9 +127,7 @@ def test_quantized_matmul_various_input_shapes(
             n_input_features,
             n_output_features,
             quantize_activation=quantize_activation,
-            batch_block_size=128,
-            out_block_size=128,
-            in_block_size=128,
+            tuned_value=None,
         )
 
     @parameterized.product(
@@ -159,9 +154,7 @@ def test_quantized_matmul_unaligned_input_shapes(
             n_input_features,
             n_output_features,
             quantize_activation=quantize_activation,
-            batch_block_size=128,
-            out_block_size=128,
-            in_block_size=128,
+            tuned_value=None,
         )
 
     @parameterized.parameters(
@@ -190,9 +183,7 @@ def test_quantized_matmul_use_tuned_block_sizes(
             n_input_features,
             n_output_features,
             quantize_activation=quantize_activation,
-            batch_block_size=None,
-            out_block_size=None,
-            in_block_size=None,
+            tuned_value=None,
         )