Precision Tuning¶

A practical guide to keeping a pypto-lib kernel numerically faithful to its torch reference — and to diagnosing the divergence when it is not. It pairs with debugging.md (the mechanics of golden_data replay and tensor dump) and performance-tuning.md (when a precision fix and a perf fix pull in opposite directions).

The recurring lesson: most "the kernel is wrong" mismatches are not logic bugs. They are an avoidable precision loss — a wrong rounding mode, a dtype that silently widened or narrowed at a boundary, an extra cast hop, or a test that was measuring the wrong thing on near-zero output. Work through the checklist below before assuming the algorithm is broken.

1. Pick the right `pl.cast` rounding mode¶

pl.cast(x, target_type, mode=...) takes a rounding mode that changes the last-bit result. The accepted names (pypto.ir.utils.CAST_MODE_NAMES):

mode	int	meaning
`none`	0	reinterpret / no explicit rounding (Acc→Vec move, identity-width)
`rint`	1	round to nearest, ties to even (RNE)
`round`	2	round to nearest, ties away from zero (the default)
`floor`	3	toward −∞
`ceil`	4	toward +∞
`trunc`	5	toward zero
`odd`	6	round to nearest, ties to odd

The default is mode="round" (ties-away). torch's .to(dtype) uses RNE (ties-to-even). So if you want to match a torch reference for a float-narrowing cast, you must pass mode="rint" explicitly — the default diverges from torch on exact ties.

For fp32 → bf16 the two modes differ only on exact ties. How often those ties occur and their whole-tensor impact are workload-dependent, so measure them rather than relying on a fixed error estimate. The golden and kernel should still use the same rule.

Recommended modes vs. torch-CPU¶

conversion	recommended `mode`	why
`fp32 → bf16`	`rint`	matches torch `.to(torch.bfloat16)` (RNE)
`fp32 → fp16`	`rint`	matches torch `.to(torch.float16)` (RNE)
`fp32 → int8` (quant)	`rint`	round-to-nearest-even on the scaled value, matching a torch `round().to(int8)` quantizer
`int32 → fp16/fp32` (de-quant)	`rint`	matches torch when rounding is needed; every integer is exactly representable only through magnitude 2^11 in fp16 and 2^24 in fp32
`acc → fp32` (cube/Acc move)	`none`	width-preserving move, no rounding involved
float → `int` for indices / lane math	`trunc` or `floor`	deterministic floor/truncation, never ties-away (which would jump an index)

When the golden deliberately reorders an algebraically-equivalent op for the kernel's benefit (e.g. folding a scale), align the golden to the kernel's order rather than de-optimizing the kernel — except when the golden is the ground-truth numeric (an RNE cast), in which case fix the kernel. See docs/debug-and-tune/debugging.md §2 for the replay loop.

2. Make the kernel and golden implementations identical¶

The golden is the device-emulating reference, so it must compute the same thing the same way as the kernel — identical op order and identical dtype at every step. Floating-point arithmetic is non-associative and lossy at each narrowing, so any divergence shows up as a spurious mismatch that looks like a kernel bug but is really a golden bug.

Op order. An algebraically-equivalent reorder — folding a scale, changing accumulation order, reassociating adds, casting at a different point — changes the last bits. Keep the two in lockstep; when the kernel reorders an op for perf, mirror that order in the golden (see §1's note) rather than letting them drift.
Dtype at each step. The golden must narrow and accumulate at the same dtype the kernel uses at each stage. A golden that stays fp32 where the kernel drops to bf16 (or quantizes to int8) under-reports the kernel's true error and hides the real bottleneck.

3. Align every input and output dtype¶

The single most common silent error: a kernel parameter, a golden tensor, and the real weight disagree on dtype. When injecting a real weight, cast it to the kernel's declared spec dtype precisely to avoid this:

if v.dtype != spec.dtype:
    v = v.to(spec.dtype)                 # force the real weight to match

Checklist:

The TensorSpec.dtype of every input must equal the dtype the kernel's pl.load / matmul actually expects. A bf16 param fed an fp32 buffer (or vice-versa) either mis-reads bytes or silently widens.
Cube matmul on a2a3 honors the stored dtype of each operand — fp32×fp32, fp16×fp16, bf16×bf16, int8 are all real paths. A weight stored FP32 runs an fp32 matmul; one stored BF16 runs bf16. Do not assume "everything is bf16": match the matmul precision to the weight's stored dtype, and make the golden do the same.
Output dtype must match between kernel and golden, or the comparator casts one side and hides a real narrowing. Compare at the dtype the kernel emits.

4. Keep intermediates wide; never cast in two hops¶

Activations and intermediate tensors should stay in the highest precision that is free — almost always fp32 — and only narrow at the boundary that genuinely requires it (a bf16 matmul input, an int8 quantized buffer, the final output).

Two rules:

Prefer fp32 for intermediates. Accumulate, normalize, apply RoPE, and do residual adds in fp32. Narrow to bf16/int8 only at the op that consumes the narrow type. The dsv4 layers carry acc/res_row/y_row in fp32 and cast to bf16 only on the store.
Never cast through an intermediate dtype. A direct fp32 → int8 is strictly better than fp32 → bf16 → int8: the second hop throws away mantissa bits before quantization, so the int8 result is rounded off a value that is already wrong. Quantize from the widest source you have. Likewise fp32 → int8 for quant, not fp32 → fp16 → int8.

# BAD: loses fp32 mantissa before quantizing
x_bf16 = pl.cast(x_fp32, pl.BF16, mode="rint")
x_i8   = pl.cast(pl.mul(x_bf16, inv_scale), pl.INT8, mode="rint")

# GOOD: quantize straight from fp32
x_i8   = pl.cast(pl.mul(x_fp32, inv_scale), pl.INT8, mode="rint")

The same applies to the int8→float dequant path: go straight from the integer to fp32, not via fp16.

5. Choose the quantization scheme deliberately¶

When a quantized stage is the precision bottleneck, the scheme is a knob, not a constant. Things to vary and measure:

Granularity — per-tensor vs. per-channel / per-token scales. A single per-tensor scale loses badly when one channel has a much larger dynamic range; per-channel (or per-token activation) scales recover it. The dsv4 weights ship per-channel weight_scale; match that granularity in the kernel and golden.
Symmetric vs. asymmetric — symmetric (offset 0) is cheaper and is what the dsv4 W8A8 checkpoints use (weight_offset == 0). Only reach for a zero-point if the data is genuinely one-sided.
Scale source — compute the activation scale from the actual dynamic range of the activation, not a static constant. Real-weight activations can have a substantially different range from random fixtures, so measure the resulting quantization error on representative data.
Rounding — RNE (mode="rint") on the scaled value to match a torch round() quantizer (see §1).

Measure each variant with the error-distribution report (§6) rather than a single pass/fail — the schemes differ in where the error lands, not just whether it passes.

6. Sweep threshold levels to see the error shape¶

A single tolerance hides the shape of the error. Use the error_distribution comparator from the golden harness to see the precision distribution instead:

from golden import error_distribution

run_jit(
    ...,
    compare_fn={
        "x_next": error_distribution(),    # measure, never fails
    },
)

It always passes (it is a measurement, not a gate) and prints, for the named output: whole-tensor rel-L2 and cosine (the trustworthy verdict for quantized / low-magnitude tensors), a frac>thd table showing what fraction of points exceed each threshold (so you can pick the tolerance the output actually needs), and percentiles of the relative diff, absolute diff, and golden magnitude (the magnitude row tells you whether a large relative diff is just an output pressed low near zero). Custom levels: error_distribution(diff_thds=(1e-3, 1e-2), quantiles=(0.5, 0.99, 1.0)).

Once the shape is known, pick a real gate (ratio_allclose or ratio_reldiff) at the threshold and bad-point budget the distribution justifies.

7. Localize the offending code with args dump + error-distribution¶

error_distribution (§6) tells you how much and what shape the output is off, but not which stage. Combine it with args dump to pin the exact op — to go from "the whole kernel is off" to "the q-proj dequant is the source" (full dump procedure in debugging.md §2 and §5):

Pin the inputs with golden_data=<dir> so every re-run sees identical tensors.
Tag the suspect intermediates with pl.dump_tag(t) and run with enable_dump_args=1 (partial dump — keep it to a few tags to avoid the full-dump AICPU timeouts).
For each dumped intermediate, run error_distribution against its torch reference and walk them in dependency order. The first stage whose distribution blows up is the culprit; everything downstream just inherits its error.

This narrows a norm → matmul → dequant → activation chain to the single op where the error first appears, which is where one of §1–§5 was violated.

8. Test with real weights and matched data distribution¶

Random fixtures hide precision bugs that only real weights expose, because the error is data-dependent:

Dynamic range drives quant error. Real activations can cover a very different range from torch.randn fixtures. A standalone module can look accurate on random data and still be the dominant error term in the full model, so report measurements with the exact workload and checkpoint.
Seed your fixtures. Bare torch.randn is unseeded and flaky run-to-run; the strict cache tolerances (max_error_ratio=0.0) need a seeded normal so the distribution is reproducible. The prefill fixtures use a seeded normal for exactly this.
Inject real weights by extracting the converted checkpoint per layer, casting each tensor to its spec dtype (§3), and overriding the input via TensorSpec.init_value. Run with real weights before trusting any precision number.
Measure by rel-L2 / cosine, not per-element reldiff, on real-weight runs — near-zero entries make per-element relative diff meaningless (§6).