Towards a Reliable Kernel Correctness Check in Matrix Multiplication

In the previous blog , we discussed defenses against hacks in automatic kernel generation. In that discussion, we emphasized the importance of a robust correctness-check function as it is a prerequisite for many defense strategies. Without it, any defensive approach would be riddled with holes .

One of the key challenges we encountered in developing fast matrix multiplication (matmul) kernels in the project CUDA-L2 is the lack of correctness check. When we first embarked on this project, it was natural to expect that a mature , universal and reliable system or pipeline for kernel correctness checking already exists. After all, thousands, and possibly tens of thousands, of kernels are used everyday by millions of developers trillion of times. If someone tells me that there is no reliable method to check whether a kernel is correct, I think he must be kidding. So the question becomes: do we actually have a mature, universal, and reliable system for verifying kernel correctness?

Unfortunately, the answer is no, especially when the kernel operates at 16 bit precision such as BF16 or FP16, or at even lower precision.

Let’s first look at how kernel correctness check is commonly done now.

Checking Numerical Differences for Correctness

The most straightforward way to is to compare the numerical differences between custom and reference kernels. This strategy is widely applied and used in Flashinfer-Bench and KernelBench .

How to Choose the Reference Kernel

The reference kernel is relatively easy to pick. The safest choice is the FP32 CPU implementation due to its well-established correctness and numerical stability.

Numerical Differences

The most convenient and thus widely-used function for checking Numerical Differences between two vectors is torch.allclose , which returns True if all elements satisfy the condition

where atol (absolute tolerance, default 1e-8) and rtol (relative tolerance, default 1e-5) control the tolerance of difference.

Threshold values are set for atol and rtol to decide whether a custom kernel’s output is numerically acceptable (here, I use atol/rtol and difference threshold interchangeably). If the output difference is within the threshold, the kernel is considered correct; otherwise, it is considered incorrect.

The Failure of the Associative Law

One key reason numerical-difference-based correctness checks described above are flawed is that associativity does not hold in floating-point arithmetic:

Floating-point Representation

Here, I plan to explain GPU floating-point representation in a bit detail, as it is a prerequisite for building the proposed correctness-checking function for matmul in CUDA-L2.

Here I use an FP16 example for illustration. FP16 consists of 1 sign bit, 5 exponent bits, and 10 fraction bits. The idea will be described also applies to BF16, with the only difference being that BF16 uses a different bit allocation (1 sign bit, 8 exponent bits, and 7 fraction bits). For normalized FP16 values, the numerical value is :

For example, if fraction=25, then 1.fraction=1.25. Let’s see how the 10 fraction bits works:

Suppose the fraction bits are \(b_1 b_2 \cdots b_{10}\), each \(b_i\in\{1,\ldots,10\}\) is binary \(0\) or \(1\).

The fraction value is given as follows:

This makes 1.fraction equal to the following:

So the value is:

This is actually how the number 2.5 is represented in FP16.

The Failure of the Associative Law

Let’s use the following example to illustrate why associativity does not hold. We will get back to this example in the correctness-checking function section.

FP16 can precisely represent \(2048\):

Can we precisely represent 2049 ?

Actually, no. This is because the 10 fraction bits give us 1024 (\(2^{10}\)) different patterns, which will be multiplied with (\(2^{11}\)) = 2048 in the exponent if we want the final value lies in range [2048, 4096).

So we can precisely represent

But we cannot precisely represent 2049, and 2049 will be rounded to its nearest representable number, which is 2048. Similarly, all numbers in the range [2048, 2049) will be rounded to 2048.

Since 2049 is exactly halfway between 2048 and 2050, the IEEE 754 "round to nearest, ties to even" rule applies here: 2048 is chosen because its fraction bits are all zeros (even), whereas 2050 has \(b_{10}=1\) (odd).

We are FINALLY done with the floating-point background. We can now clearly see that:

Associativity does not hold!

Why Non-Associativity Makes Correctness Checking Fail

How does non-associativity relate to kernel correctness checking?

On a GPU, many threads work on different parts of the matrix at the same time in parallel. This means the additions don't always happen in the same order. Even if two GPU kernels are both correct, their outputs might be different because they accumulate values in different orders.

It means if a custom kernel produces different outputs from a reference/golden kernel, you cannot immediately conclude that the custom kernel is incorrect, because the difference can come from the non-associativity of floating-point arithmetic.

Of course, the difference should not be too large. But how large is “too large” before it indicates incorrectness; if the difference falls closely to this “too large” threshold, will the kernel be treated as correct or not ?

Another question is that does the same tolerance apply across different kernels, different inputs, or even different precision formats? Clearly not. Numerical differences are inevitably larger for a (10000×10000) × (10000×10000) matrix multiplication than for a (2×2) × (2×2) one, and larger in FP8 than in FP16.

Let’s run a simple test: compare the maximum difference between FP16 and FP32 matmul outputs for (10000×10000) × (10000×10000) and (2×2) × (2×2) .

import torch


def get_precision_diff(M, K, N):
   """
   Compare fp16 vs fp32 matrix multiplication precision.
  
   Computes: A (M x K) @ B (K x N) = C (M x N)
  
   Returns max absolute difference between fp16 and fp32 results.
   """
   # Create random matrices in fp32
   A_fp32 = torch.randn(M, K, device='cuda', dtype=torch.float32)
   B_fp32 = torch.randn(K, N, device='cuda', dtype=torch.float32)
  
   # FP32 matmul (ground truth)
   C_fp32 = torch.matmul(A_fp32, B_fp32)
  
   # FP16 matmul (convert inputs, compute, convert back for comparison)
   A_fp16 = A_fp32.half()
   B_fp16 = B_fp32.half()
   C_fp16 = torch.matmul(A_fp16, B_fp16).float()  # Convert result back to fp32 for comparison
  
   # Compute max absolute difference
   max_diff = (C_fp32 - C_fp16).abs().max().item()
  
   return max_diff


if __name__ == "__main__":
   torch.cuda.synchronize()
  
   # Test 1: 10000 x 1000 x 1000
   print("=" * 50)
   print("Matrix multiplication: (10000 x 1000) @ (1000 x 1000)")
   print("=" * 50)
   max_diff_large = get_precision_diff(10000, 1000, 1000)
   print(f"Max absolute difference (fp16 vs fp32): {max_diff_large:.3g}")
  
   print()
  
   # Test 2: 2 x 2 x 2
   print("=" * 50)
   print("Matrix multiplication: (2 x 2) @ (2 x 2)")
   print("=" * 50)
   max_diff_small = get_precision_diff(2, 2, 2)
   print(f"Max absolute difference (fp16 vs fp32): {max_diff_small:.3g}")

Here is the output when the code is executed:

==================================================
Matrix multiplication: (10000 x 1000) @ (1000 x 1000)
==================================================
Max absolute difference (fp16 vs fp32): 0.0947

==================================================
Matrix multiplication: (2 x 2) @ (2 x 2)
==================================================
Max absolute difference (fp16 vs fp32): 0.000424

This leads to a fundamental question: how do we decide what level of difference is acceptable? Or does a mature, universal, and reliable criterion for kernel correctness checking simply not exist?

Unfortunately, the answer is no.

Solution: Exact Match with binary Inputs

In this section, we will describe the correctness checking function used for matrix multiplication in CUDA-L2 , which does not rely on numerical difference–based strategies.

For matrix multiplication, we multiply matrix \(A \in \mathbb{R}^{M\times K}\) with matrix \(B \in \mathbb{R}^{K\times N}\) to obtain matrix \(C \in \mathbb{R}^{M\times N}\). Each element of \(C\) is given by:

One direct solution for correctness check is to restrict the computation to a situation where floating-point associativity holds so that precise computation can be guaranteed. In that case, exact match between outputs must be achieved to ensure the matmul kernel correctness.

Then the question becomes how we can design this situation. Back to the associativity example we talked about, FP16 can precisely represent all integers in the range [0, 2048]. Therefore, precise computation can be guaranteed if the following condition is satisfied:

Condition 1: not only the final result in the equation above lies within this range, but all intermediate values in the sum sequence must also remain below \(2048\).

For example, if \(A \in \mathbb{R}^{3\times 3}\) and \(B \in \mathbb{R}^{3\times 3}\), we have

To satisfy the condition above, the following must hold:

Why must all intermediate values in the sum sequence also remain below 2048 ?

This is because once an intermediate value exceeds 2048, exactness is immediately lost. As a result, the exactness of the final value cannot be guaranteed.

Condition 2: the partial sums are monotonically non-decreasing.

As can be seen, Condition 2 implies Condition 1.

Condition 3: all elements in \(A\) and \(B\) are non-negative.

Clearly, Condition 3 implies Condition 2: since each product \(A_{ik}B_{kj}\ge 0\), the partial sums are monotonically non-decreasing.

In matmul, for large matrices when computing dot products involve a large number of additions, we want each product \(a_{ik} \times b_{kj}\) to remain as small as possible, since the accumulated sum can easily exceed the 2048 threshold. Therefore, we can simply restrict \(A\) and \(B\) to be binary.

However, for very large matrices, even binary-valued \(A\) and \(B\) can cause the accumulation to exceed the 2048 threshold. In this case, we further control the sampling distribution of the binary values by assigning a higher probability to zero, thereby reducing the expected magnitude of the accumulated sum.

The full picture of the correctness checking function is given as follows:

We first compute the reference output \(C^{ref}\) using FP32 on CPU, which provides exact integer results. We then compute the output \(C^{custom}\) using the custom kernel. For each position \((i, j)\) where \(c^{ref}_{ij} \le 2048\), we require \(c^{custom}_{ij} = c^{ref}_{ij}\) exactly, and we ignore positions \(c^{ref}_{ij} > 2048\).

We repeat this process multiple times with different random inputs; if any single iteration fails, the kernel is considered incorrect.

The following is the pseudo code:

procedure CHECK_KERNEL_CORRECTNESS(custom_kernel, num_trials):

  for trial = 1 to num_trials do

    # 1. Sample binary inputs with biased distribution
    A <- SAMPLE_BINARY_MATRIX(p_zero)
    B <- SAMPLE_BINARY_MATRIX(p_zero)

    # 2. Compute reference result (exact integers)
    C_ref <- MATMUL_FP32_CPU(A, B)

    # 3. Compute custom kernel result
    C_custom <- MATMUL_CUSTOM_KERNEL(A, B)

    # 4. Verify correctness under threshold
    for each index (i, j) do
      if C_ref[i, j] < 2048 then
        if C_custom[i, j] != C_ref[i, j] then
          return INCORRECT
        end if
      end if
    end for

  end for

  return CORRECT

end procedure

While passing the check does not strictly guarantee correctness, a custom kernel that consistently passes across many iterations, where each with a reference output containing a significant fraction of values in the range [0,2048], is very unlikely to be incorrect.

What if we want to use fractional values

One might consider using fractional values such as \(\tfrac{1}{2},\ \tfrac{1}{4},\ \tfrac{1}{8}\), but this does not fundamentally change the analysis; the key factor is the smallest representable increment at a given magnitude.

For example, consider the case where elements are sampled from \(\{0,\tfrac{1}{2},1\}\).

The individual products lie in \(\{0,\tfrac{1}{4},\tfrac{1}{2},1\}\), and the accumulated sum can take values such as \(\tfrac{3}{4}\).

In FP16, values spaced by 0.25 are exactly representable only up to magnitude 512. In particular, exact representation is guaranteed for

But not \(512.25,\ 512.75\), etc.

How well does the Exact Match strategy generalize

For all linear kernels, the extract match strategy we used in matmul for CUDA-L2 can be readily adapted.

But for kernels involving non-linear operations, unfortunately, this approach no longer applies ☹️. Non-linear operations can easily break the exactness. For example, for functions such as log and exp , the outputs are typically irrational numbers and therefore can never be represented exactly, regardless of mantissa length.

Alleviating this issue and building more robust, though not perfect, kernel correctness checking functions is a long-term goal of our work. We sincerely welcome feedback and discussion.

Thank you for your time and attention.