modify triton matmul to allow the formula C=A*B+C

Signed-off-by: cliu-us <[email protected]>

Author: chichun-charlie-liu

fms_mo/custom_ext_kernels/triton_kernels.py

@@ -101,6 +101,7 @@ def matmul_kernel(
     stride_cm,
     stride_cn,
     chunk_trun_bits,
+    truncate_then_accumulate,
     # Meta-parameters
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
@@ -159,15 +160,20 @@ def matmul_kernel(
         a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
         b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
         # We accumulate along the K dimension.
-        accumulator = tl.dot(a, b, accumulator, input_precision="ieee")
+        if truncate_then_accumulate:
+            accumulator_inner = tl.dot(a, b, input_precision="ieee")
+        else:
+            accumulator_inner = tl.dot(a, b, accumulator, input_precision="ieee")
         # tl.dot() default is using TF32 approximation, not good enough for LSB truncation exp
 
         ## ------ add chunky LSB rounding/masking --------
         if chunk_trun_bits > 0:
-            accumulator = libdevice.uint_as_float(
-                (libdevice.float_as_uint(accumulator) + round_bit) & trun_mask
-            )
+            accumulator_inner = round_and_trun(accumulator_inner, round_bit, trun_mask)
         ## ---------------------------------------------------------
+        if truncate_then_accumulate:
+            accumulator += accumulator_inner
+        else:
+            accumulator = accumulator_inner
 
         # Advance the ptrs to the next K block.
         a_ptrs += BLOCK_SIZE_K * stride_ak
@@ -206,6 +212,7 @@ def imatmul_kernel(
     stride_cm,
     stride_cn,
     chunk_trun_bits,
+    truncate_then_accumulate,
     # Meta-parameters
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
@@ -244,13 +251,21 @@ def imatmul_kernel(
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
         b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
-        accumulator = tl.dot(a, b, accumulator, input_precision="ieee")
+        if truncate_then_accumulate:
+            accumulator_inner = tl.dot(a, b, input_precision="ieee")
+        else:
+            accumulator_inner = tl.dot(a, b, accumulator, input_precision="ieee")
 
         ## ------ add chunky LSB rounding/masking --------
         if chunk_trun_bits != 0:
-            accumulator = (accumulator + round_bit) >> chunk_trun_bits
-            accumulator = accumulator << chunk_trun_bits
+            accumulator_inner = (accumulator_inner + round_bit) >> chunk_trun_bits
+            accumulator_inner = accumulator_inner << chunk_trun_bits
         ## ---------------------------------------------------------
+        if truncate_then_accumulate:
+            accumulator += accumulator_inner
+        else:
+            accumulator = accumulator_inner
+
 
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
@@ -266,29 +281,163 @@ def imatmul_kernel(
     tl.store(c_ptrs, c, mask=c_mask)
 
 
+
+@triton.jit
+def matmul_kernel_DABC(
+    # Pointers to matrices
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    # Matrix dimensions
+    M,
+    N,
+    K,
+    # The stride variables represent how much to increase the ptr by when moving by 1
+    # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`
+    # by to get the element one row down (A has M rows).
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    chunk_trun_bits,
+    truncate_then_accumulate,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    ACTIVATION: tl.constexpr,
+):
+    """Kernel for computing the matmul D = A x B + C that include LSB truncation.
+    A has shape (M, K), B has shape (K, N) and C/D has shape (M, N). 
+    NOTE:
+        C should be consistent with accumulator dtype, e.g. fp8xfp8 -> fp32.
+        *D ptr is supposed to be the same as C ptr, no need to provide D as arg
+        **we can be used C to verify unintended truncation by CUDA as well. 
+    Args:
+        chunk_trun_bits (int): number of LSB to truncate/round. [0 to 23]
+    """
+    # -----------------------------------------------------------
+    # Map program ids `pid` to the block of C it should compute.
+    # This is done in a grouped ordering to promote L2 data reuse.
+    # See above `L2 Cache Optimizations` section for details.
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # ----------------------------------------------------------
+    # Create pointers for the first blocks of A and B.
+    # We will advance this pointer as we move in the K direction
+    # and accumulate
+    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
+    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
+    # See above `Pointer Arithmetic` section for details
+    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+
+    # -----------------------------------------------------------
+    # Iterate to compute a block of the C matrix.
+    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
+    # of fp32 values for higher accuracy.
+    # accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    accumulator = tl.load(c_ptrs, mask=c_mask, other=0.0)  # should have been cast to fp32 already
+    ## ------ prepare LSB rounding/truncation masks -------
+    # NOTE mask will be applied on accumulator, which is alway FP32, so we may truncate up to 23b
+    # e.g., 20b -> trun_mask = 0xFFF00000, round_bit = 0x00080000
+    #        8b -> trun_mask = 0xFFFFFF00, round_bit = 0x00000080
+    trun_mask = tl.cast((0xFFFFFFFF >> chunk_trun_bits) << chunk_trun_bits, tl.uint32)
+    round_bit = 1 << (chunk_trun_bits - 1) if chunk_trun_bits > 0 else 0
+    ## ---------------------------------------------------------
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # Load the next block of A, B, and C, generate a mask by checking the K dimension.
+        # If it is out of bounds, set it to 0.
+        # D = truncation(A*B) + C
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+        # We accumulate along the K dimension. but apply truncation on local A*B first
+        if truncate_then_accumulate:
+            accumulator_inner = tl.dot(a, b, input_precision="ieee")
+        else:
+            accumulator_inner = tl.dot(a, b, accumulator, input_precision="ieee")
+        # tl.dot() default is using TF32 approximation, not good enough for LSB truncation exp
+        # NOTE: tl.dot(a, b, c) should use one single CUDA mma instruction to handle "c = a*b+c". If
+        #       this mma instruction uses "reduced-precision" under the hood, not only a*b will
+        #       be accumulated in that precision, c most likely will be cast to that "lower"
+        #       precision first, hence, will lose some precision!
+
+        ## ------ add chunky LSB rounding/masking --------
+        if chunk_trun_bits > 0:
+            accumulator_inner = round_and_trun(accumulator_inner, round_bit, trun_mask)
+        ## ---------------------------------------------------------
+        if truncate_then_accumulate:
+            accumulator += accumulator_inner
+        else:
+            accumulator = accumulator_inner
+
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+    # You can fuse arbitrary activation functions here
+    # while the accumulator is still in FP32!
+    if ACTIVATION == "leaky_relu":
+        accumulator = leaky_relu(accumulator)
+
+    d = accumulator  # do not cast to (tl.float16) just yet
+
+    # -----------------------------------------------------------
+    # Write back the block of the output to matrix "C" with masks.
+    tl.store(c_ptrs, d, mask=c_mask)
+
+
 @triton.jit
 def leaky_relu(x):
     """Activation function that could be fused into matmul kernel"""
     return tl.where(x >= 0, x, 0.01 * x)
 
 
+@triton.jit
+def round_and_trun(x, round_bit, trun_mask):
+    """Round and truncate (usually for accumulator)."""
+    return libdevice.uint_as_float((libdevice.float_as_uint(x) + round_bit) & trun_mask)
+
+
 def tl_matmul_chunk_truncate(
     a,
     b,
     activation="",
     chunk_trun_bits=0,
     chunk_size=16,
+    truncate_then_accumulate=True,
     cast_output_to_input_dtype=None,
 ):
     """Triton matmul for HW behavior simulation. Supports float and int8.
-    a. variable chunk size (i.e., BLOCK_SIZE_K)
-    b. LSB truncation, must <23 if using float.
+    i. variable chunk size (i.e., BLOCK_SIZE_K)
+    ii. LSB truncation, must <23 if using float.
+    iii. assume D = A*B + C, where C is optional. If C exists, it will be updated inplace.
 
     Args:
         a, b: input tensors. FloatX, X in [32, 16, 8] or INT8.
         activation (str, optional): activation func to be fused, see relu example.
         chunk_trun_bits (int, optional): number of LSBs to be truncated/rounded.
         chunk_size (int, optional): BLOCK_SIZE_K, some HW has specific chunk size. must >= 16.
+        truncate_then_accumulate (bool, optional): if True, c = truncate(a*b) + c, otherwise
+                                                    c = truncate(a*b+c)
         cast_output_to_input_dtype (bool, optional): accumulator has higher prec than input, usually
                                                     FP32 or INT32. by default we cast the final
                                                     output to the same dtype as input for non-8bits.
@@ -300,6 +449,7 @@ def tl_matmul_chunk_truncate(
     use empirical way to determine BLOCK sizes, may not be optimal. But need to avoid autotune for
     real model inference. otherwise auto-tune will be triggered in every forward call.
     """
+
     # Check constraints.
     assert a.shape[1] == b.shape[0], "Incompatible dimensions"
     assert a.is_contiguous(), "Matrix A must be contiguous"
@@ -343,9 +493,18 @@ def isPowerofTwo(x):
         mm_kernel = imatmul_kernel
     else:
         acc_dtype = torch.float32
-        mm_kernel = matmul_kernel
+        mm_kernel = matmul_kernel if c == None else matmul_kernel_DABC
         assert chunk_trun_bits < 23, "FP32 accumulator only has 23 mantissa bits"
-    c = torch.zeros((M, N), device=a.device, dtype=acc_dtype)
+
+    if c == None:
+        c_org_dtype = a.dtype
+        c = torch.zeros((M, N), device=a.device, dtype=acc_dtype)
+    else:
+        # if C is in fp16, accumulate in fp32 no matter what, decide whether to cast back later
+        c_org_dtype = c.dtype
+        c = c.to(acc_dtype)
+        assert c.shape[0]==M and c.shape[1]==N, "C shape is inconsistent with A B."
+        assert acc_dtype==torch.float32, "INT truncation experiment is not yet supported."
 
     # 1D launch kernel where each block gets its own program.
     def grid(META):
@@ -386,7 +545,8 @@ def grid(META):
         c.stride(0),
         c.stride(1),
         chunk_trun_bits=chunk_trun_bits,
+        truncate_then_accumulate=truncate_then_accumulate,
         ACTIVATION=activation,
         **kernel_config,  # if using auto-tune, comment this line out.
     )
-    return c.to(a.dtype) if cast_output_to_input_dtype else c
+    return c.to(c_org_dtype) if cast_output_to_input_dtype else c