debug-quantized-kernel-accuracy

Skill: Debug Quantized Kernel Accuracy

Purpose

Guide the agent through a systematic process for diagnosing and isolating accuracy degradation in a quantized (INT8, FP8, or low-bit) kernel, from measuring the error to identifying the specific computational step responsible.

Use this when

• A quantized kernel produces outputs that differ from the fp32 reference by more than the expected quantization error bound.
• A model using quantized kernels shows accuracy degradation that exceeds what is expected for the chosen quantization scheme.
• A quantization refactor introduced a regression and the specific step that broke is not obvious.
• Debugging a quantized kernel that works correctly on some input shapes or batch sizes but fails on others.

Do not use this when

• The error is within the expected quantization error bound (approximately 0.5 * scale per element for well-calibrated INT8) and the downstream task accuracy loss is acceptable.
• The issue is clearly a non-accuracy bug (segfault, wrong shape, miscompilation) — fix the structural bug first.
• The degradation is due to model-level quantization sensitivity (certain layers or operators being inherently sensitive to quantization), which requires a quantization-aware training or mixed-precision approach rather than kernel debugging.

Inputs the agent should gather first

• The exact mathematical specification of what the quantized kernel is supposed to compute, written in terms of the original unquantized operation.
• The quantization scheme: per-tensor, per-channel, or per-token; symmetric or asymmetric; INT8, INT4, or FP8; signed or unsigned range.
• The scale computation method: offline calibration, per-batch dynamic quantization, or per-token dynamic quantization.
• The accumulation dtype: INT32, FP32, FP16, or FP8.
• The dequantization epilogue: where is scale applied, in what order, and what is the output dtype.
• A fp32 reference output for the same inputs (required for comparison).
• Whether the error is consistent across runs (deterministic) or varies (stochastic — possible race condition or non-deterministic reduction).

Required reasoning process

1. Measure the error precisely. Do not characterize accuracy failures as "wrong" without quantifying them:

   • Compute max absolute error (MAE_max), mean absolute error (MAE_mean), and relative error (RE = |output - reference| / (|reference| + epsilon)).
   • Compute per-element SNR if the output represents a signal: 20 * log10(rms(reference) / rms(output - reference)).
   • Check whether the error is distributed uniformly across the output tensor or concentrated in specific rows, columns, or batch elements. Localized error strongly suggests a scale indexing bug.
   • Check whether the error scales with the magnitude of the input values. If error is proportional to input magnitude, the scale factor is wrong. If error is constant, the zero-point is wrong.

2. Isolate the quantization step responsible. The quantization pipeline is:

   fp32 input → quantize → int8 input
   int8 GEMM (or other op) with int32 accumulation
   int32 accumulation → dequantize → fp32/fp16 output
   

   Test each boundary:

   • Quant-dequant roundtrip: quantize the fp32 input to INT8, then immediately dequantize back to fp32. Compare to the original. Max error should be approximately 0.5 * scale per element. If it exceeds this, the scale computation or the quantization formula is wrong.
   • INT8 kernel with exact INT8 inputs: if you can construct inputs where the exact INT8 representation is known (e.g., all values are multiples of the scale), run the INT8 kernel and verify that INT32 accumulation produces the expected result before dequantization.
   • Dequantization epilogue only: take known-correct INT32 accumulation results, apply the dequantization epilogue, and verify the fp32 output. If this step introduces error, the scale application formula is wrong.

3. Check scale computation for outliers. A single large-magnitude input value can dominate the scale computation for a per-tensor or per-row quantization scheme, causing all other values to be mapped to a narrow range of INT8 values (near zero). This is called the outlier problem:

   • Compute the histogram of INT8 values for the weight tensor and the activation tensor. If most values cluster near 0 with only a few values near ±127, outliers are the cause of accuracy loss.
   • Check per-tensor vs per-channel granularity: if per-tensor quantization produces a scale dominated by one outlier channel, switching to per-channel quantization may resolve the accuracy issue.
   • Examine the distribution of absolute values in the activation tensor. If a significant fraction are near the max representable value in fp16, fp16 overflow may be affecting the scale computation itself.

4. Verify zero-point handling for asymmetric quantization. Zero-point bugs produce systematic bias in the output:

   • The correct dequantization is x_fp32 = (q_int8 - zero_point) * scale.
   • If zero_point is added instead of subtracted, every output element has a systematic bias of 2 * zero_point * scale.
   • If zero_point is a signed quantity but stored as unsigned (or vice versa), the bias direction is wrong.
   • For symmetric quantization: zero-point must be exactly 0. If a non-zero zero-point leaks in from misconfiguration, it introduces a constant offset.

5. Verify accumulation dtype is INT32. INT8 accumulation overflows for dot products of length > ~10 with typical INT8 values. Symptoms: outputs are wildly wrong and the error grows with K (the reduction dimension). Verify by reading the kernel source and confirming the accumulator variable is int32_t, not int8_t or int16_t.

6. Check scale application order for per-channel / per-token quantization. The scale must be applied along the correct dimension:

   • For per-channel weight quantization with per-tensor activation quantization: output[m][n] = int32_acc[m][n] * act_scale * weight_scale[n].
   • If weight_scale[n] is applied along the M dimension instead of N, every output row m is multiplied by weight_scale[m] instead of weight_scale[n]. The error pattern will show column-wise or row-wise distortion.
   • For per-token activation quantization: output[m][n] = int32_acc[m][n] * act_scale[m] * weight_scale[n]. Verify that act_scale[m] loads the scale for token m, not token m+1 or another token.

7. Check INT32 accumulation overflow. For very long dot products with INT8 inputs, INT32 can overflow:

   • Maximum accumulation value = 127 * 127 * K. This exceeds INT32 max (~2.1B) when K > ~131,000.
   • For standard LLM linear layer sizes (K up to ~16384), overflow is not a concern for symmetric INT8. For larger K, check.
   • If INT32 overflow is detected, the fix is to split the K accumulation into segments and dequantize partially, or to accumulate in INT64 for the final segment.

8. Check the dequantization epilogue sequence. The order of operations in the epilogue matters:

   • Correct order: INT32 accumulation → cast to fp32 → multiply by scale → add bias (in fp32) → cast to output dtype.
   • Incorrect order: INT32 accumulation → multiply by scale (in INT32, which may overflow for large scales) → cast to fp32.
   • Incorrect order: INT32 accumulation → cast to fp32 → add bias (using wrong bias scale) → multiply by scale.
   • Write out the epilogue operations in order and verify each step against the mathematical definition.

9. Check for shape-dependent failures. If the kernel is correct on some shapes but wrong on others, the likely causes are:

   • Partial K tile (K not a multiple of dp4a width = 4): the tail elements are either skipped or computed incorrectly.
   • Per-token scale with variable sequence length: if the scale tensor is indexed incorrectly when sequence length changes, scales from the wrong position are applied.
   • Batch size 1 vs larger batch: if batch dimension handling has an off-by-one, batch size 1 may work correctly but larger batches fail.

Kernel design rules

• Never accumulate in INT8. Always accumulate in INT32. This is non-negotiable.
• The dequantization step must be a separate, identifiable section of the kernel. Do not interleave scale application with the INT8 accumulation loop.
• Scale tensors must be loaded with explicit indexing that is reviewed against the mathematical formula. Do not use flat pointer arithmetic without verifying the index maps to the correct scale.
• For debugging: add an assert or printf in the kernel for the first output element to verify the INT32 accumulation value before dequantization. Compare this to a Python reference that computes int32_t(a_int8) @ int32_t(b_int8) (integer GEMM).
• Zero-point values must be stored and loaded in the same signed/unsigned convention throughout the pipeline. Mixing signed and unsigned zero-points is a correctness bug.

Correctness requirements

• Quant-dequant roundtrip error must be bounded by 0.5 * scale per element. If it exceeds this, the scale computation or quantization formula is wrong and must be fixed before debugging the kernel.
• After confirming the roundtrip is correct, the INT8 kernel output must match the fp32 reference to within N * (0.5 * scale) absolute error, where N is the number of terms in the accumulation. In practice, well-implemented INT8 GEMM matches fp32 GEMM to within 1–3 ULPs at the output scale.
• Per-element error must be randomly distributed across the output tensor, not structured (not concentrated in specific rows, columns, or blocks). Structured error implies a scale indexing bug, not just quantization noise.
• The kernel must produce identical results across multiple runs on the same inputs (deterministic). Non-deterministic INT8 accumulation is a sign of a race condition or unsynchronized reduction.

Performance requirements

• Debugging steps (quant-dequant roundtrip tests, INT32 accumulation

Content truncated.