Repository navigation
[RFC] GDN2 intermediate precision for strict BF16 context-parallel accuracy #1550
Description
Summary
The strict GDN2 CP2 tests fail on correlated BF16 inputs even when no sequence crosses a rank boundary. The error also occurs in the ordinary single-GPU chunk operator. This is separate from the local-log decay correction in #1404.
The remaining error needs a precision decision for stored chunk intermediates. Please agree on the accuracy target and which intermediate tensors may retain FP32 before an implementation changes the existing precision contract.
Reproduction and evidence
Run on two CUDA GPUs with strict comparisons:
FLA_CI_ENV=0 FLA_CACHE_RESULTS=0 pytest -s -v tests/context_parallel/test_cp_gdn2.py::test_cp2_boundary_alignedThe test uses BF16 operands, T=10240, H=12, K=V=128, sequence lengths [5120,5120], fused bounded gates with lower_bound=-5, and asymmetric Q/K inputs. Neither sequence crosses a rank boundary.
Earlier paired H200 runs at upstream 329b534245d3d7bb27cc03a99e87bdc51fb28d06 and candidate 3d00b00 reproduced the same approximately 1.5% output error. The following diagnosis used 7b2abb6122af469472209d7751c01cd94febba5a on B300, PyTorch 2.13.0+cu130, Triton 3.7.1, Python 3.12.13, driver 580.173.02.
| Diagnostic | Relative RMS output error |
|---|---|
| Original CP2 boundary-aligned test | 1.5099% |
| Same test, reference Q/K normalized in FP32 then cast back | 1.3515% |
| Ordinary chunk, fused safe gates, FP32-normalized reference | 1.3509% |
| Ordinary chunk, precomputed safe gates | 1.3509% |
| Ordinary chunk, raw-log gates | 1.3465% |
| Fused recurrent operator, same inputs/reference | 0.4297% |
| Chunk with IEEE triangular-solve products | 1.3043% |
| Chunk with IEEE solve and default FP32 dot products | 1.3050% |
The single-GPU comparison uses the first 5120 tokens from the test's identical global input generation. IEEE overrides were diagnostic only. No production precision, input distribution, assertion, or tolerance was changed.
The CP reference currently normalizes BF16 Q/K in BF16. The kernel and ordinary GDN2 operator tests normalize in FP32 before casting back. Correcting that mismatch helps but does not resolve the failure.
Intermediate rounding diagnosis
For the first 64 tokens of the same input, FP64 chunk algebra agrees with the FP32 recurrent reference to relative RMS 7.05e-8. Rounding individual intermediate tensors to BF16 produces:
| Rounded intermediate | Relative RMS error |
|---|---|
| Triangular inverse only | 0.3430% |
| Causal QK score matrix only | 0.7154% |
| Write product only | 0.1657% |
WY value tensor u only |
0.8027% |
Inverse, score matrix, write product, u, and output |
1.1573% |
The production chunk result for those first 64 tokens has 1.1580% error. This supports intermediate rounding as the dominant issue; it is not evidence of a CP communication defect. The quantization experiment is algebraic diagnosis, not a proposed replacement kernel or a performance measurement.
Proposed scope
- Correct reference normalization to match the established operator contract.
- Preserve the existing 0.008 threshold and input distribution.
- Evaluate a targeted precision policy for score/WY intermediates, including backward and shared state/output kernels, while preserving public output dtypes.
- Validate dense/packed inputs, CP2/4/8 where supported, retained/recomputed states, final states, gradients, and graph replay.
- Measure latency and peak memory on the affected backends before selecting a production policy.
No production precision change is included in this RFC. #1455 discusses fused triangular-solve precision on other hardware, but changing that solve alone did not resolve this case.
Activity