Optimize w8a8 kernel vmem limit (#9508)

Author: kyuyeunk

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py

@@ -42,12 +42,12 @@ def matmul_kernel(
     x_abs_max_ref: jax.Array,  # (1, batch_block_size)
     out_ref: jax.Array,  # (batch_block_size, out_block_size)
     acc_scratch: jax.Array,  # (batch_block_size, out_block_size)
-    q_x_scratch: jax.Array,  # (batch_block_size, in_block_size)
+    x_q_scratch: jax.Array,  # (batch_block_size, in_block_size)
     x_scale_scratch: jax.Array,  # (batch_block_size, 1)
     *,
     quantize_activation: bool,
     save_acc: bool,
-    save_q_x: bool,
+    save_x_q: bool,
     batch_block_size: int,
     out_block_size: int,
     in_block_size: int,
@@ -66,13 +66,13 @@ def matmul_kernel(
   assert out_ref.shape == (batch_block_size,
                            out_block_size), "out_ref shape is not correct"
 
-  if save_q_x:
+  if save_x_q:
     assert quantize_activation
-    assert q_x_scratch is not None
+    assert x_q_scratch is not None
     assert x_scale_scratch is not None
     quant = (out_idx == 0)
   else:
-    assert q_x_scratch is None
+    assert x_q_scratch is None
     assert x_scale_scratch is None
     quant = quantize_activation
 
@@ -88,18 +88,18 @@ def matmul_kernel(
   def matmul_body(quant, is_first_step, is_last_step):
     if quantize_activation:
       if quant:
-        q_x_tmp, x_scale_tmp = _quantize_array(x_ref[...], x_abs_max_ref[...])
-        if save_q_x:
-          q_x_scratch[...] = q_x_tmp
+        x_q_tmp, x_scale_tmp = _quantize_array(x_ref[...], x_abs_max_ref[...])
+        if save_x_q:
+          x_q_scratch[...] = x_q_tmp
           x_scale_scratch[...] = x_scale_tmp
       else:
-        assert save_q_x
-        q_x_tmp = q_x_scratch[...]
+        assert save_x_q
+        x_q_tmp = x_q_scratch[...]
         if is_last_step:
           x_scale_tmp = x_scale_scratch[...]
 
       acc = jax.lax.dot_general(
-          q_x_tmp,
+          x_q_tmp,
           w_ref[...],
           (((1,), (1,)), ((), ())),
           preferred_element_type=jnp.int32,
@@ -130,6 +130,44 @@ def _next_multiple(x, multiple):
   return ((x + multiple - 1) // multiple) * multiple
 
 
+def _get_vmem_limit(n_bs, n_out, n_in, batch_block_size, out_block_size,
+                    in_block_size, x_bytes, w_bytes, x_q_bytes, scale_bytes,
+                    out_bytes, acc_bytes, save_acc, save_x_q):
+  # Calculate in/out VMEM size.
+  x_size = batch_block_size * in_block_size * x_bytes
+  x_abs_max_val_size = batch_block_size * scale_bytes
+  w_size = out_block_size * in_block_size * w_bytes
+  scalar_size = out_block_size * scale_bytes
+  out_size = batch_block_size * out_block_size * out_bytes
+
+  vmem_in_out = x_size + x_abs_max_val_size + w_size + scalar_size + out_size
+  vmem_in_out *= 2  # Account for compute and vreg spills.
+
+  # Account for double buffering.
+  # Double buffering is used only if there are multiple blocks per in/out.
+  vmem_in_out += x_size if (n_bs > 1 or n_in > 1) else 0
+  vmem_in_out += x_abs_max_val_size if (n_bs > 1) else 0
+  vmem_in_out += w_size if (n_out > 1 or n_in > 1) else 0
+  vmem_in_out += scalar_size if (n_out > 1) else 0
+  vmem_in_out += out_size if (n_bs > 1 or n_out > 1) else 0
+
+  # Calculate scratch VMEM size.
+  acc_size = batch_block_size * out_block_size * acc_bytes
+  x_q_scratch_size = batch_block_size * in_block_size * x_q_bytes
+  x_scale_scratch_size = batch_block_size * scale_bytes
+
+  vmem_scratch = acc_size if save_acc else 0
+  vmem_scratch += x_q_scratch_size + x_scale_scratch_size if save_x_q else 0
+  vmem_scratch *= 2  # Account for compute and vreg spills.
+
+  # Add in/out and scratch VMEM size.
+  vmem_used = vmem_in_out + vmem_scratch
+  # Specify upper limit as 96MB.
+  vmem_limit_bytes = min(vmem_used, 96 * 1024 * 1024)
+
+  return vmem_limit_bytes
+
+
 @functools.partial(
     jax.jit,
     static_argnames=[
@@ -196,32 +234,39 @@ def quantized_matmul_int8(
   assert x.shape[
       1] % in_block_size == 0, f"x.shape[1] ({x.shape[1]}) must be a multiple of block size ({in_block_size})"
 
-  acc_dtype = jnp.int32 if quantize_activation else x.dtype
-  vmem_to_be_transferred = 2 * (
-      batch_block_size * in_block_size * x.dtype.itemsize +
-      out_block_size * in_block_size * w.dtype.itemsize + out_block_size *
-      scalar.dtype.itemsize + batch_block_size * x_abs_max_val.dtype.itemsize +
-      batch_block_size * out_block_size * x.dtype.itemsize
-  ) + batch_block_size * out_block_size * jnp.dtype(acc_dtype).itemsize
-  # Within the kernel, it will use some extra VMEM for computation or vreg spills.
-  vmem_used = vmem_to_be_transferred * 2
-  vmem_limit_bytes = min(vmem_used * 2, 96 * 1024 * 1024)
-
   n_bs = padded_bs // batch_block_size
   n_out = padded_out_features // out_block_size
   n_in = padded_in_features // in_block_size
 
   save_acc = n_in > 1
   # Remove redundant input quantization logic by caching quantized input.
   # For best performance, only enable this behavior when single input block is used per batch.
-  save_q_x = quantize_activation and n_in == 1 and n_out > 1
+  save_x_q = quantize_activation and n_in == 1 and n_out > 1
+
+  acc_dtype = jnp.int32 if quantize_activation else jnp.float32
+
+  vmem_limit_bytes = _get_vmem_limit(
+      n_bs=n_bs,
+      n_out=n_out,
+      n_in=n_in,
+      batch_block_size=batch_block_size,
+      out_block_size=out_block_size,
+      in_block_size=in_block_size,
+      x_bytes=x.dtype.itemsize,
+      w_bytes=w.dtype.itemsize,
+      x_q_bytes=jnp.dtype(jnp.int8).itemsize,
+      scale_bytes=jnp.dtype(jnp.float32).itemsize,
+      out_bytes=x.dtype.itemsize,
+      acc_bytes=jnp.dtype(acc_dtype).itemsize,
+      save_acc=save_acc,
+      save_x_q=save_x_q)
 
   kernel = pl.pallas_call(
       functools.partial(
           matmul_kernel,
           quantize_activation=quantize_activation,
           save_acc=save_acc,
-          save_q_x=save_q_x,
+          save_x_q=save_x_q,
           batch_block_size=batch_block_size,
           out_block_size=out_block_size,
           in_block_size=in_block_size),
@@ -243,9 +288,9 @@ def quantized_matmul_int8(
               pltpu.VMEM((batch_block_size,
                           out_block_size), acc_dtype) if save_acc else None,
               pltpu.VMEM((batch_block_size,
-                          in_block_size), jnp.int8) if save_q_x else None,
+                          in_block_size), jnp.int8) if save_x_q else None,
               pltpu.VMEM(
-                  (batch_block_size, 1), jnp.float32) if save_q_x else None,
+                  (batch_block_size, 1), jnp.float32) if save_x_q else None,
           ],
           grid=(n_bs, n_out, n_in),
       ),