Add standalone English precision analysis

Explain baseline and current arithmetic, matched B300 measurements, and the individual-state KDA outlier; link the document from README.

Files changed (2) hide show

PRECISION_ANALYSIS.md +284 -0
README.md +10 -27

PRECISION_ANALYSIS.md ADDED Viewed

	@@ -0,0 +1,284 @@

+# Precision Analysis: TileLang Baselines and CuTe DSL Kernels
+This document compares the original TileLang baselines with the released GDN
+and KDA CuTe DSL implementations. It describes unchanged formats, arithmetic
+changes, precision improvements, deliberate approximations, and measured limits.
+The paired measurements were collected on NVIDIA B300 on 2026-10-26.
+The output format and the original benchmark thresholds are unchanged. The
+current implementations retain the rebuilt state and residual in FP32 Tensor
+Memory (TMEM), while using reduced-precision operands and selected approximations
+elsewhere. FP32 storage or accumulation does not imply FP32 precision at every
+operation, or bitwise equivalence to the mathematical reference.
+## 1. Three different meanings of precision
+The following are separate implementations and measurement roles:
+1. **Mathematical reference:** `definition.json` computes the flush mainly with
+   FP32 PyTorch operations. It still rounds the new rank-4 factors to FP16 and
+   quantizes the residual to INT8. Its decoded output already has format error.
+2. **Original TileLang baseline:** already uses mixed precision, including
+   FP16/BF16 tensor-core operands, FP32 accumulation, FP16 intermediate storage,
+   and FP32 corrections. It is not an all-FP32 implementation.
+3. **Current CuTe DSL implementation:** also uses mixed precision. It changes
+   which intermediates retain FP32, how products are compensated, and where
+   lower-precision approximations are accepted.
+The judge rebuilds a separate target state in FP64 from the same supplied
+inputs. These inputs already contain INT8, FP16, and BF16 values. This target is
+not an unquantized model state from before those input conversions. FP64 is used
+for validation, not by the production flush arithmetic.
+## 2. Input and output formats are unchanged
+The fixed geometry is 32 heads, a 128 by 128 state, a 16-entry update window,
+and rank 4.
+| Data | Original baseline | Current implementation |
+|---|---|---|
+| Residual `codes` | INT8, values in [-127, 127] | Same |
+| Checkpoint scales `s_k`, `s_v` | FP32 | Same |
+| Low-rank factors `u`, `q` | FP16 | Same |
+| Update buffers `kbuf`, `ubuf` | BF16 | Same |
+| Weights `w`, GDN scalar gate `p` | FP32 | Same |
+| KDA gate history `gbuf` | FP16 | Same |
+| KDA cumulative log-gate `pcum` | FP32 | Same |
+| Supplied starting basis `q0` | FP32 | Same |
+The checkpoint continues to decode as
+$$
+S_{\mathrm{old}}[v,k]
+=\sum_{r=1}^{4}q[r,v]u[r,k]
++C[v,k]\,s_k[k]s_v[v]/127.
+$$
+The rank, window length, supplied starting basis, and one round of subspace
+iteration are retained. The new factors are rounded to FP16 before forming the
+final residual, so that the residual compensates for the factors actually stored.
+INT8 codes retain nearest rounding and the same value range.
+`Pool.raw` is an FP32 backing allocation with typed views into its storage.
+That backing dtype does not make its INT8 codes or FP16 factors FP32 data.
+## 3. State and residual storage
+### GDN baseline
+The baseline first forms a partial state `B` containing the dequantized old
+residual and buffered updates, then stores `B` in FP16 shared memory. The old
+rank-4 term is carried separately through FP32 algebraic corrections rather than
+being irreversibly folded into that FP16 matrix. The final residual is computed
+with FP32 accumulation, then stored in FP16 before scale calculation and coding.
+This distinction matters: the baseline does not simply cast the entire state
+to FP16 and discard the old low-rank correction.
+### KDA baseline
+The default configuration has `KEXACT=True`. For the final residual it performs
+another dequantization/rebuild from the staged inputs, accumulates the terms in
+FP32, and applies the per-channel checkpoint decay in FP32. It preserves the
+unscaled old factors for this residual calculation.
+The residual is then scaled by a power of two per block and stored in FP16.
+Scaling protects small residuals from FP16 underflow, but does not remove FP16
+significand rounding. The projection stage still uses an FP16 partial state and
+an FP16 copy of the decayed old factors.
+### Current implementations
+Both kernels retain the full rebuilt state in FP32 TMEM. They also create a
+scaled FP16 shared-memory copy for the projection products. After producing the
+new FP16 factors, they subtract the low-rank product from the retained FP32
+state and keep the residual in FP32 through the quantization passes.
+This removes an important FP16 storage-rounding step from the final residual
+path. It also lets KDA avoid the baseline's second full residual rebuild.
+However, the retained FP32 state was itself formed using mixed-precision
+products, so it is not necessarily the same bit pattern as a sequence of FP32
+reference operations.
+Residual accuracy is sensitive to cancellation: when `S` and the new low-rank
+product are large but their difference is small, a rounding error measured
+relative to `S` can be large relative to the residual.
+## 4. Product precision and high/low compensation
+The original baselines already use BF16 high/low splitting for weighted or
+decayed updates. Conceptually, a value is represented as
+$$
+x\approx x_{\mathrm{hi}}+x_{\mathrm{lo}},\qquad
+Ax\approx Ax_{\mathrm{hi}}+Ax_{\mathrm{lo}}.
+$$
+The current GDN implementation folds the update weight into the key operand,
+whereas the baseline primarily folds it into the value operand. The current GDN
+high BF16 part is truncated and the low part compensates for the discarded
+bits. These are mathematically equivalent organizations with different
+floating-point rounding and accumulation orders.
+The current kernels also use TF32 high/low compensation for scaled or decayed
+old low-rank terms. This retains more information than a single TF32 conversion,
+but is not a guarantee of exact FP32 multiplication or reduction.
+The projection products primarily use FP16 operands with FP32 accumulation.
+The fast projection path can still lose information in the FP16 state copy;
+retaining a separate FP32 state improves the residual path without making the
+computed low-rank subspace identical to the reference.
+`f32x2` means two independent FP32 operations executed in a packed instruction.
+It does not mean FP64, double-double arithmetic, or additional significant bits
+for either value.
+## 5. Starting basis and dynamic range
+The baseline converts `q0` directly to FP16 for the projection. The current GDN
+kernel first rescales its columns by powers of two and represents the basis
+with FP16 high and low parts. KDA also conditions the starting basis with
+Cholesky orthogonalization before storing its high/low representation.
+The current kernels choose a per-head power-of-two scale for state construction
+and projection. The scale is undone when the factors and checkpoint scales are
+written. For normal representable values, multiplication by a power of two
+changes the exponent without introducing the same rounding as an arbitrary
+scale factor. Range limits and subsequent FP16 conversion still matter.
+KDA already used block scaling in the baseline's residual-storage stage. The
+current implementation extends range management to the full state path while
+retaining its FP32 main copy. It does not make FP16 and BF16 interchangeable.
+## 6. Orthogonalization changes
+The baseline defaults to `GRAM16=False`: both CholeskyQR passes compute their
+Gram matrices in FP32. A comment describing an optional FP16 second Gram pass
+does not describe this default path.
+The current kernels can skip the second orthogonalization pass when the first
+pass meets a conditioning test. KDA performs the second pass when needed. GDN
+also has a more conservative path with state-low-part compensation and repeated
+Gram-Schmidt projections for difficult cases; its second Cholesky pass can
+suppress numerically dependent directions.
+These choices alter numerical results and are not merely scheduling changes.
+Two QR passes also do not mean two rounds of the outer subspace-iteration
+algorithm; these are different operations.
+## 7. Deliberate lower-precision approximations
+Some current operations trade local precision for throughput:
+- **GDN column sums:** parts of the absolute-residual column reduction use
+  FP16 packed addition before returning to FP32. The baseline reads an FP16
+  residual but accumulates those column sums in FP32.
+- **KDA column-scale reciprocal:** the reciprocal is rounded to FP16. That same
+  rounded value is used for coding and to derive the corresponding FP32 output
+  scale. A scale stored in FP32 can therefore still reflect a lower-precision
+  internal computation.
+- **Row normalization:** the current kernels primarily use approximate
+  reciprocals, whereas the baseline uses a correctly rounded reciprocal here.
+  GDN retains a more accurate fallback for extreme row scales.
+- **Conditional QR:** skipping the second pass changes the amount of numerical
+  correction performed on well-conditioned inputs.
+For example, ignoring zero-value guards and clamps, KDA's scaled residual
+`E_prime` yields a column reciprocal approximately of the form
+$$
+r_h=\operatorname{FP16}\left(
+  \operatorname{rsqrt}(\operatorname{mean}_v|E_{\mathrm{prime}}|)
+\right).
+$$
+Coding uses `r_h`, and the stored column scale is derived from its reciprocal
+with the state scale undone. This keeps the encoded values and saved scales
+consistent, while moving the quantization grid relative to a fully FP32 scale
+calculation. Individual codes and scales can consequently differ from the
+reference even when decoded-state error remains close to it.
+## 8. What the accuracy metrics measure
+For the original benchmark, each head index has an error averaged over batch
+items and matrix elements:
+$$
+e_h=\operatorname{mean}_{b,v,k}
+\left|\operatorname{decode}(\mathrm{checkpoint})-S_{64}\right|.
+$$
+On the relative-error branch, the reported ratio is
+$$
+\rho_h=e_h^{\mathrm{implementation}}/e_h^{\mathrm{FP32\ reference}}.
+$$
+The original requirements are unchanged:
+- If reference error is nonzero and at least `1e-4 * mean(abs(S64))`, candidate
+  error must be at most `1.01 * reference_error`.
+- Otherwise, the bound is `1e-3 * mean(abs(S64))`.
+- Outputs must be finite and codes must remain in [-127, 127].
+The reference output is itself quantized. A ratio of 1.0004 means approximately
+0.04% more reconstruction error than that reference, not a state error of 0.04%.
+This criterion does not require elementwise-identical codes, factors, or scales.
+## 9. Matched B300 measurements
+Both implementations received identical tensors. The comparison used six batch
+sizes (16, 32, 64, 128, 256, 512) and three input families: structured edge cases,
+a second window starting from the reference checkpoint, and random in-domain
+inputs. Each implementation passed all 18 original benchmark cases per kernel.
+| Metric | GDN baseline | GDN current | KDA baseline | KDA current |
+|---|---:|---:|---:|---:|
+| Worst original-benchmark reference-relative error ratio | 1.00324455 | 1.00041398 | 1.00296773 | 1.00105573 |
+| Maximum individual-state MAE / mean(abs(S64)) | 1.052615% | 0.576760% | 0.613289% | 0.612224% |
+| Original cases passed | 18/18 | 18/18 | 18/18 | 18/18 |
+| Failures in the additional individual-state check | 30/96,768 | 0/96,768 | 44/96,768 | 1/96,768 |
+The individual-state check applies the same threshold formulas separately to
+each `(batch item, head)`, including evaluating the relative/absolute branch
+separately for that state. It does not average over batch items and is stricter
+than the original benchmark. The failure counts in the last row are not failures
+of the original aggregated scoring procedure.
+### Current KDA boundary case
+One individual state in the batch-128 random case has a reference-relative MAE
+ratio of **1.01256057**, exceeding the individual relative limit of 1.01. Its
+original aggregated benchmark case passes. Therefore, the data supports better
+measured original-benchmark ratios, but not a claim that every individual state
+satisfies the same bound.
+No stage-by-stage numerical ablation was performed. The isolated failure cannot
+be assigned to the FP16 reciprocal, projection, conditional QR, or another
+specific approximation from these measurements alone.
+These are finite test results. They do not establish bounds for every possible
+input, long sequences of self-generated checkpoints, or full-model perplexity.
+In particular, the second-window test starts from a reference checkpoint rather
+than a long chain of each implementation's own outputs.
+## 10. Evidence and reproduction
+- Original mathematical definitions: [GDN](gdn/definition.json), [KDA](kda/definition.json).
+- Unmodified TileLang baselines: [GDN](gdn/baseline.py), [KDA](kda/baseline.py).
+- Released CuTe sources: [GDN](build/torch-cuda/_gdn.py), [KDA](build/torch-cuda/_kda.py).
+- Matched measurements: [GDN](gdn/precision_comparison.json), [KDA](kda/precision_comparison.json).
+- Earlier deployment validation: [GDN](gdn/validation_results.json), [KDA](kda/validation_results.json).
+- Source hashes and provenance: [GDN](gdn/provenance.json), [KDA](kda/provenance.json).
+Using the pinned [dependencies](requirements.txt), run on an idle B300:
+```bash
+python compare_precision.py gdn
+python compare_precision.py kda
+```
+The [comparison script](compare_precision.py) creates the inputs once for both
+implementations, rebuilds the FP64 target, invokes the original judge, and
+records the extra individual-state check separately. It writes results to each
+component's `precision_comparison.json`.

README.md CHANGED Viewed

@@ -57,33 +57,16 @@ and provenance. Dependency pins and Apache-2.0 licensing are shared at the root.
 ## Precision
-Output formats remain INT8 codes, FP32 scales, and FP16 rank-4 factors;
-update buffers remain BF16. Both baseline and current implementations use
-mixed-precision arithmetic. The current kernels keep the state and residual in
-FP32 TMEM, while projections use scaled FP16 operands and selected operations
-use local approximations. FP32 storage does not imply FP32 precision at every step.
-On 2026-10-26, both implementations were run on identical B300 inputs: six batch
-sizes and three input families. Each passed all 18 original benchmark cases.
-Ratios below compare decoded-state MAE to the FP32 reference's MAE, both measured
-against the same FP64 rebuilt state; the normal acceptance limit is 1.01.
-| Kernel | Baseline worst original-judge ratio | Current worst original-judge ratio | Baseline / current failures in extra individual-state check |
-|---|---:|---:|---:|
-| GDN | 1.00324455 | 1.00041398 | 30 / 0 |
-| KDA | 1.00296773 | 1.00105573 | 44 / 1 |
-The extra check independently tests 96,768 (batch item, head) states per kernel.
-The original judge averages over the batch dimension, so this extra check is
-stricter. The current KDA outlier occurs in batch 128 random inputs, with an
-individual reference-relative MAE ratio of 1.01256057. This exceeds 1.01 for that
-individual state even though the original aggregated benchmark passes. These
-results do not establish a universal bound on every possible input.
-Full records are in [gdn/precision_comparison.json](gdn/precision_comparison.json)
-and [kda/precision_comparison.json](kda/precision_comparison.json). Reproduce on
-an idle B300 with `python compare_precision.py gdn` and
-`python compare_precision.py kda` using the pinned dependencies.
 ## Reproduce

 ## Precision
+See the standalone [precision analysis](PRECISION_ANALYSIS.md) for the original
+and current arithmetic paths, deliberate approximations, matched measurements,
+and validation limits. Input/output formats and original scoring thresholds
+are unchanged; both implementations pass all 18 original comparison cases.
+The additional individual-state check found one current KDA outlier in the
+batch-128 random case (error ratio 1.01256057 against the individual limit 1.01).
+The original benchmark averages the batch dimension; the extra check treats
+individual states separately. Full records and reproduction commands are in
+the analysis document.
 ## Reproduce