Shorten precision analysis to the requested English text

Translate the selected six-part explanation and append measured workload errors, including the KDA individual-state outlier. Keep reproduction commands in README.

Files changed (2) hide show

PRECISION_ANALYSIS.md +145 -283
README.md +3 -2

PRECISION_ANALYSIS.md CHANGED Viewed

@@ -1,283 +1,145 @@
-# 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).
-- Source hashes and release details: [GDN](gdn/README.md#files-and-provenance), [KDA](kda/README.md#files-and-provenance).
-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`.

+# Precision Analysis
+The input/output formats and acceptance thresholds are unchanged. Internally,
+FP32 state/residual retention has been added, while some lower-precision
+approximations are used in projection, reduction, and scale computation. I have
+rerun the original baseline and current implementation on B300 with exactly
+the same inputs.
+First, distinguish three things:
+- **Mathematical reference:** mainly uses FP32 arithmetic, then outputs FP16
+  low-rank factors and an INT8 residual.
+- **Original TileLang baseline:** already uses mixed precision, including
+  FP16/BF16 multiplication, FP32 accumulation, and FP16 intermediate storage.
+- **Current CuTe implementation:** still uses mixed precision, but changes
+  which intermediates retain FP32 and where approximations are allowed.
+The external data types are completely unchanged:
+| Data | Original baseline and current implementation |
+|---|---|
+| Residual `codes` | INT8, range [-127, 127] |
+| Low-rank factors `u`, `q` | FP16 |
+| Update buffers `kbuf`, `ubuf` | BF16 |
+| Scales `s_k`, `s_v` | FP32 |
+| `w`, `p`, `pcum`, `q0` | FP32 |
+| KDA gate history `gbuf` | FP16 |
+The numerical changes occur at the following stages.
+1. **GDN state and residual: reduce FP16 rounding of matrix intermediates.**
+   The original baseline first computes the quantized residual plus the window
+   updates to obtain `B`, then stores `B` in FP16. The old rank-4 term participates
+   separately through FP32 algebraic corrections. After the final residual is
+   computed, it is stored in FP16 again before scales and codes are generated.
+   The current implementation retains the full rebuilt state `S` in FP32 TMEM.
+   It uses a scaled FP16 copy to extract the low-rank component, but computes
+   $$
+   E=S-q_{\rm new}^{T}u_{\rm new}
+   $$
+   using the retained FP32 state. The residual also remains in FP32 until
+   quantization.
+   This is the main precision improvement: when subtracting two large values
+   produces a small residual, rounding the state to FP16 beforehand can make
+   the rounding error large relative to that residual.
+2. **KDA residual: FP32 reconstruction already existed; the current kernel
+   removes the subsequent FP16 round trip.**
+   The original KDA baseline enables `KEXACT=True` by default: it decodes the
+   checkpoint again and recomputes the residual in FP32, avoiding direct use
+   of the earlier FP16 state. However, it then scales the residual by a power
+   of two per block and stores it in FP16.
+   The current implementation retains the FP32 state directly, subtracts the
+   new low-rank term from it, and keeps the residual in FP32 TMEM. The change
+   therefore removes the FP16 store and reload after residual computation;
+   it is not a change from an entirely FP16 pipeline to an entirely FP32 one.
+3. **Product compensation: retain BF16 high/low splitting and reorganize the
+   arithmetic.**
+   The original baseline already uses BF16 high/low splitting. For an FP32
+   value, the idea is
+   $$
+   x\approx x_{\rm hi}+x_{\rm lo},\qquad
+   Ax\approx Ax_{\rm hi}+Ax_{\rm lo}.
+   $$
+   The current GDN implementation folds the weight into the key operand and
+   then splits it into BF16 high and low parts; the baseline mainly folds
+   the weight into the value operand. These are mathematically equivalent,
+   but the rounding locations and accumulation order change.
+   The current implementation also uses TF32 high/low compensation for
+   weighted or decayed old low-rank terms. This retains more information
+   than a single TF32 conversion, but still does not establish bitwise
+   equivalence to IEEE FP32 computation.
+4. **Starting matrix and dynamic range: add scaling and low-part protection.**
+   The original baseline converts `q0` directly to FP16. The current
+   implementation first scales each column by a power of two and retains
+   FP16 high and low parts. KDA also conditions the starting basis through
+   orthogonalization.
+   The current implementation also chooses a per-head power-of-two scale
+   during state reconstruction, keeping the FP16 projection copy within a
+   suitable numerical range. The KDA baseline already had per-block scaling
+   at the residual stage; the current kernel moves this protection earlier
+   into the state stage.
+   Scaling helps avoid overflow and underflow, but FP16 rounding remains.
+5. **Orthogonalization: reduce computation according to numerical conditioning,
+   while retaining corrections for difficult inputs.**
+   The original baseline defaults to `GRAM16=False`, so the Gram matrices
+   in both QR passes are computed in FP32.
+   The current implementation skips the second orthogonalization pass when
+   conditioning is good and adds corrections when conditioning is poor. GDN
+   also uses state-low-part compensation and reorthogonalization for
+   ill-conditioned cases. These choices change numerical results and are
+   numerical tradeoffs.
+   Also, `f32x2` in the code means computing two FP32 values simultaneously;
+   it does not increase the precision of either value.
+6. **Scales and reductions: these do introduce lower-precision approximations.**
+   - **GDN:** part of the residual column summation uses packed FP16 addition
+     before returning to FP32. The baseline reads the FP16 residual and
+     accumulates those sums in FP32.
+   - **KDA:** the column-scale reciprocal is rounded to FP16, then used to
+     generate codes and the corresponding FP32 output scale. A scale stored
+     in FP32 therefore need not have been computed entirely at FP32 precision.
+   - Current row-scale normalization mainly uses approximate reciprocals;
+     the baseline uses a correctly rounded reciprocal. GDN retains a more
+     accurate fallback for extreme scales.
+   - The final INT8 nearest-rounding rule and value range remain unchanged.
+These paths can be compared directly in the original
+[GDN baseline](gdn/baseline.py), [KDA baseline](kda/baseline.py), and current
+[GDN](build/torch-cuda/_gdn.py) and [KDA](build/torch-cuda/_kda.py) implementations.
+On the six given B300 workloads (batches 16, 32, 64, 128, 256, and 512), tested
+with structured, second-window, and random inputs, both implementations pass
+all 18 original cases per kernel. The worst reconstruction-error ratio relative
+to the FP32 reference improves from **1.00324455 to 1.00041398 for GDN** and from
+**1.00296773 to 1.00105573 for KDA**, with both errors measured against the same
+FP64 rebuilt state. The current kernels' maximum individual-state
+`MAE / mean(abs(S64))` is **0.576760% for GDN** and **0.612224% for KDA**.
+An additional check of 96,768 individual states per kernel finds no GDN
+violations and one KDA violation: a batch-128 random state has an error ratio
+of **1.01256057**, exceeding the individual limit of 1.01. The original judge
+averages over the batch dimension and still passes that case. Full measurements
+are in the [GDN](gdn/precision_comparison.json) and
+[KDA](kda/precision_comparison.json) reports.

README.md CHANGED Viewed

@@ -65,8 +65,9 @@ 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

 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 are linked in the analysis document.
+Reproduce with `python compare_precision.py gdn` and
+`python compare_precision.py kda`.
 ## Reproduce