Kl Consistency Test

KL Consistency Tests

What the test is for

kl_test_utils scores the same token twice -- once as a prefill input logprob, once as a decode output logprob -- and compares. The two paths run different kernels over different shapes, so agreement is a statement about state, not about answer quality: it catches a radix-cache prefix that does not reproduce a fresh prefill, a stale conv/mamba checkpoint, a SWA pool that evicted something it still needed.

gsm8k passing says nothing about this. Accuracy is insensitive to a handful of corrupted tokens; the KL check is not.

Two independent conditions produce a zero

Reaching bit-identity needs both, and they fail for unrelated reasons. Knowing which one a nonzero belongs to is most of the debugging.

1. Every operator on the path is batch-invariant. A token's result must not depend on how many tokens share its forward. Note this is a property across the two paths, not a property of each: a kernel can be perfectly reproducible at M=1 and again at M=N while disagreeing between them, which is exactly what a tile-size switch or a message-size-dependent reduction does.

2. The two paths compute the same function. Decode's context and state at a position must equal what a fresh prefill computes there -- the same KV set, the same sliding window, the same conv/mamba state, a restored cache prefix that reproduces a recomputed one. This is logic, not arithmetic, and it survives any amount of numerical hygiene.

The conditions are independent, and one measurement separates them: with (1) satisfied, match and decode_cache_hit read exactly 0 while prefill_cache_hit stays nonzero when a prefix restore is wrong. Same server, same prompts -- float noise cannot pick a code path, so a helper-specific divergence is (2).

Order the work accordingly. Settle (1) first: until it holds, its noise is orders of magnitude above anything (2) produces and hides it completely.

The three helpers differ in what touches the cache

KLDivergenceMixin runs the last two. Pick deliberately -- they are not interchangeable, and only the cache-hit pair exercises prefix reuse.

A divergence confined to one helper is diagnostic. match clean but prefill_cache_hit dirty means the restore path is wrong, not the arithmetic -- float noise does not pick a code path.

Run it the way CI runs it

KLDivergenceMixin defaults: max_samples=32, max_new_tokens=512. Do not characterize with fewer.

avg_kl_div is the k3 estimator, exp(logr) - 1 - logr, applied to the sampled token's logprob. It is exponentially sensitive to the tail, so the mean is carried by a handful of tokens. At 4 samples the same config measured 0.049 to 0.158 -- a 3x spread that invalidates any A/B comparison drawn from it.

When characterizing rather than gating, report tail statistics -- the fraction of tokens past a threshold, and the max -- rather than the mean.

Generate past the sliding window if the model has one, so decode carries the window through the handover from prompt tokens to generated ones.

Condition 1: determinism is not batch-invariance

This distinction decides whether a threshold means anything.

• Deterministic: same input, same shape, same result on every run.
• Batch-invariant: a token's result does not depend on how many other tokens share its batch.

The KL check compares a prefill of thousands of tokens against decode steps of one, so it measures the second. --enable-deterministic-inference buys both -- it swaps the aten kernels for fixed-reduction versions and pins the NCCL algorithm and channel count -- but only for kernels it covers. Custom kernels that never reach an aten op are outside batch_invariant_ops and stay shape-dependent.

The consequence: a nonzero KL under deterministic inference that appears in every helper alike means some kernel on the path is still batch-dependent. Localize it (below) rather than widening the threshold.

Background, and the source of the fixed-reduction approach the aten overrides take: Defeating nondeterminism in LLM inference.

How much batch-invariance an operator needs

For a token-wise operator -- GEMM, norm, activation, the router's linear -- a token's output depends only on that token's row, so pinning the reduction order is the whole requirement. Once its result is independent of how many rows share the launch, it is done.

Two kinds need more than that, and they are where the remaining nonzero usually lives:

• Operators that reduce across tokens -- attention over a KV range, and any collective. Fixing the arithmetic order is not enough if the extent still varies: an all-reduce whose tree shape follows the message size, or an attention split whose block boundary follows the query count, gives a token a different reduction depending on its batch. Pin the shape, not just the order.
• Operators that carry state across calls -- conv windows, SSM checkpoints. These are batch-invariant per call and still diverge, because what they store is reused by a later request. That is condition 2, and no amount of reduction-order work reaches it.

So "make everything batch-invariant" closes condition 1 for the token-wise majority, and the residual after that is concentrated in these two classes.

MoE amplifies this to a degree dense models do not

Top-k routing is a discrete decision over near-tied scores. A 1e-8 difference in gate weights flips which experts a token is routed to, the outputs diverge completely, and 42 layers compound it. Measured on one MoE checkpoint: a gate GEMM that switched tiling between M=8 and M=16 produced a 1.6e-5 logits difference, which became 20-37 nat on individual high-confidence tokens and a KL of 0.177.

A dense model of comparable size shows the same root cause as ~1e-4. So a KL in the hundredths is not evidence of a worse bug on a MoE model -- it is the same class of numerical difference, amplified. Do not calibrate a MoE threshold by analogy to a dense one.

Condition 2: the two paths must compute the same function

Once condition 1 holds, whatever remains is a state bug, and the helper it appears in names the path. A restore that does not reproduce a recomputed prefix shows up in prefill_cache_hit alone; the other two stay at exactly 0.

What the signature looks like, and how to read it:

• Which sequences. Divergence concentrated in a couple of requests out of a batch, with the rest bit-identical, is a condition triggered by those requests -- not a systematic offset. Compare their prompt lengths, cached_tokens, and page and checkpoint-interval remainders against the ones that pass.
• Where in the generation. Contiguous from the first generated token means the state was already wrong when generation began, so the fault is in the prefix restore rather than in decode. Divergence starting mid-generation points instead at something that happens during decode -- a window handover, a checkpoint rotation.
• Whether it is a race. Re-run under different configurations that should not matter (page size, TP degree, buffer strategy). Bit-identical numbers across them mean a deterministic logic fault, which is far cheaper to chase than a race.

Generate past the sliding window if the model has one: the handover from prompt tokens to generated ones inside the window is where eviction and checkpoint rotation actually run.

Choosing a threshold

Once every kernel on the path is batch-invariant, prefill and decode agree bit for bit and the honest assertion is a stray-ulp floor, not a tolerance:

KL_DIV_THRESHOLD = 1e-9   # measured 0; anything a state bug produces is orders above

A loose threshold tolerates float noise and small logic errors alike, which is how a state-reuse bug hides. Prefer running the KL case on its own deterministic server and asserting near-zero, and keep the accuracy case on the production numerics -- one server cannot serve both.

Thresholds are per (model, tp). A value calibrated at tp=1 does not transfer: tp=1 has no all-reduce, so it never exercises the source that dominates at tp>1.