Overview

Low-precision training can make a correct computation look faulty because ordinary round-off obscures small corruption signals. AEGIS forms Matmul checksums from the kernel’s float32 accumulator and uses a FlashAttention backward invariant, then applies tolerances and sends suspicious deviations to cVerifier. Row versus row-and-column checksums trade coverage for copies and retained context; the methods remain operator-specific, not universal proof.

Learning outcomes

Using the float32 accumulator can separate injected faults more clearly than checking rounded bfloat16 output in the studied implementation.
Matmul equivalent paths and the FlashAttention backward invariant provide compact, operator-specific suspicion signals.
A checksum threshold crossing selects a candidate for verification; cVerifier, not the algorithmic sensor, decides whether it is an SDC.

Keywords

AEGISMatmulFlashAttentionbfloat16float32Checksum

How AEGIS Makes Low-Precision Arithmetic Observable

AEGIS’s algorithmic sensors begin with a difficult fact about modern LLM training: the computation may be wrong, while the arithmetic used to check it is noisy enough to look wrong even when it is healthy. This chapter follows the two mathematical checks AEGIS uses to work inside that constraint. The goal is to understand how mixed-precision Matmul checksums and a FlashAttention invariant produce compact, relevant suspicion signals—and why neither signal is a verdict until cVerifier confirms it.

Matmul: preserve the signal before rounding erases it

Matrix multiplication is central to Transformer workloads. If two matrices are A and B, their product is C = AB. A checksum compresses a large result into a smaller value that is easier to compare. A simple row checksum multiplies by a column vector of ones, written as 1; the result adds the entries in each row.

There are two equivalent ways to obtain that row sum. Compute the product first and then sum its rows:

C1 = (AB)1

Or sum the columns of B first and multiply A by that smaller vector:

C2 = A(B1)

In exact arithmetic, associativity says these values are equal. A checksum detector measures the difference between them, for example as ΔMatmul = ||(AB)1 − A(B1)||. If the difference exceeds an expected round-off bound, the operation becomes suspicious. The detector does not need to compare every element of a huge tensor to obtain a useful first signal.

The difficulty is that low-precision training does not perform this algebra exactly. LLM systems commonly compute with bfloat16 values, and ordinary floating-point rounding accumulates as operations proceed. The two checksum paths may therefore disagree even when the Matmul is correct. A theoretical bound designed for a higher-precision calculation can be too small for the noise produced by the actual low-precision path. If AEGIS treated every nonzero difference as corruption, normal arithmetic would generate a stream of false alarms.

The paper makes the problem visible through fault injection. It flips one bit of a bfloat16 input and compares checksum differences. When the checksum is accumulated from the final bfloat16 output, the discrepancy caused by the injected fault is masked by larger discrepancies caused by round-off. The detector is looking at the right operation, but the evidence has been degraded before it is measured.

This is an important distinction. The problem is not simply that bfloat16 is “inaccurate” and float32 is “accurate.” The question is where precision is lost relative to the evidence path. The final bfloat16 output has already been rounded. If the detector forms its checksum from that value, it inherits the truncation and may lose the small spatial or numerical change that identifies the fault. A correct detector must preserve the relevant signal long enough to compare it against the background noise.

Use the accumulator that the computation already has

AEGIS’s key design choice is to accumulate the checksum from the float32 accumulator inside the mixed-precision kernel, rather than from the final bfloat16 output. Modern accelerator kernels can pair lower-precision multiplication with higher-precision accumulation. In the paper’s setting, bfloat16 inputs and outputs are associated with a float32 accumulation path, such as the one available in NVIDIA Tensor Core execution.

The distinction can be pictured as two points at which the check might be attached. In the first arrangement, the Matmul produces a bfloat16 result, that result is rounded, and the checksum is computed afterward.

Published by DeepMox

FAQ

Why does AEGIS use a float32 accumulator for checksums?

AEGIS uses the float32 accumulator so the checksum sees higher-precision intermediate information before bfloat16 output rounding can mask a small fault signal.

How does AEGIS check FlashAttention?

AEGIS checks FlashAttention by comparing float32-accumulated column sums of dV and dO, which should match under the described backward invariant; a deviation creates a vTask rather than a verdict.

What is the cost of using row-and-column checksums?

Row-and-column checksums cost additional data movement and memory copies while extending coverage to weight-gradient Matmuls and reducing saved verification context.

Previous11 minAEGIS: How Evidence Changes Incident Actions

Next9 minAEGIS: When Recomputation Supports Bitwise Checks

AEGIS: Low-Precision Checksums and Attention Invariants

How AEGIS Makes Low-Precision Arithmetic Observable

Matmul: preserve the signal before rounding erases it

Use the accumulator that the computation already has

FAQ