diff --git a/docs/fe-oss-apis/csa.md b/docs/fe-oss-apis/csa.md index 99ab62ed6..6a08bb529 100644 --- a/docs/fe-oss-apis/csa.md +++ b/docs/fe-oss-apis/csa.md @@ -48,12 +48,13 @@ gradients on such padding rows. All arithmetic is fp32 with one final bf16 rounding; `mul.rn.f32` / `fma.rn.f32` are pinned in PTX so results do not depend on compiler FMA contraction. The numerics contract is **per ratio family** (the two families intentionally differ — do not -assume the ratio=4 guarantees at ratio=128): +assume the ratio {2, 4} guarantees at ratio=128): -- **`ratio == 4`** (production, unchanged): against an fp32-intermediate eager - reference (same op order, fp32 throughout), `dKV`/`dScore` are **bit-identical** - and the forward matches within one bf16 rounding step on a tiny fraction of - elements. +- **`ratio in {2, 4}`** (generic kernels; the bitwise contract — ratio=4 as originally + validated, ratio=2 added on the same terms): against an + fp32-intermediate eager reference (same op order, fp32 throughout), `dKV`/`dScore` + are **bit-identical** and the forward matches within one bf16 rounding step on a + tiny fraction of elements. - **`ratio == 128`** (deterministic tolerance contract): the kernels are **deterministic and faithful to the fp32-intermediate eager reference** (the same eager region computed with fp32 intermediates and one final bf16 rounding — @@ -132,8 +133,9 @@ that need a fully deterministic backward must use an eager implementation. - Compute-capability major **>= 10** (SM100 and newer; the kernels use no architecture-specific features beyond the SM100 baseline) -- `ratio == 4`, `coff in {1, 2}` (`coff == 2` is the production CSA/HCA configuration, - `coff == 1` the own-block window form) — served by the generic kernels, which are +- `ratio in {2, 4}`, `coff in {1, 2}` (`coff == 2` is the production CSA/HCA configuration, + `coff == 1` the own-block window form; `ratio == 2` is the configuration used in + production training) — served by the generic kernels, which are generic over `(ratio, head_dim, coff in {1, 2})` but keep the whole pooling window in registers (register-bound beyond `ratio = 32`) - `ratio == 128`, `coff in {1, 2}`, `head_dim in {128, 512}` — served by dedicated @@ -338,7 +340,7 @@ never-consumed slots and fp32-atomic `dAPE`, exactly as at ratio=4. The backward defaults to the same fast exp outside the d=128 small-pack (vec=1) buckets. **The ratio=128 numerics contract is the deterministic tolerance contract described in -[Numerics](#numerics), NOT the ratio=4 bitwise-`dKV`/`dScore` contract.** The +[Numerics](#numerics), NOT the generic-family (ratio {2, 4}) bitwise-`dKV`/`dScore` contract.** The reduction orders (forward chunk merge, backward `den`/`S` partial merge), the backward's hoisted reciprocal, and the fast-exp buckets all differ from the eager op order by design, each adopted on a measured same-GPU win and gated on tolerance + @@ -415,7 +417,7 @@ committed scripts `benchmark/csa/gate_csa_compressor_r128.py` and ``` The tests validate numerics against an fp32-intermediate eager reference (bitwise -`dKV`/`dScore` at ratio=4; the deterministic tolerance contract at ratio=128, with +`dKV`/`dScore` at ratio {2, 4}; the deterministic tolerance contract at ratio=128, with fp64-oracle parity on finite-intermediate inputs), the upstream eager numerics, plus ragged packs, static-capacity padding, kernel-side zero-writes into uninitialized gradient buffers (NaN-canary with exact-zero assertions on every never-consumed slot class, and the `total_comp == 0` diff --git a/python/cudnn/csa/compressor/api.py b/python/cudnn/csa/compressor/api.py index 710eda509..9358294df 100644 --- a/python/cudnn/csa/compressor/api.py +++ b/python/cudnn/csa/compressor/api.py @@ -22,7 +22,7 @@ ``out``, FP32 ``ape``, int32 ``cu_seqlens``/``cu_seqlens_comp``, int32 flat offsets (``total_tokens * coff * head_dim < 2**31``), and per ratio: -- ``ratio == 4``, ``coff in {1, 2}`` (``coff == 2`` is the production CSA/HCA +- ``ratio in {2, 4}``, ``coff in {1, 2}`` (``coff == 2`` is the production CSA/HCA configuration, ``coff == 1`` the own-block window form) — served by the generic kernels in ``compressor_sm100.py`` (whole window in registers; optimal at small ratios, register-bound beyond ``ratio = 32``); @@ -35,7 +35,7 @@ Numerics contract (see the kernel modules and docs/fe-oss-apis/csa.md for details): fp32 arithmetic with one final bf16 rounding, ``mul.rn``/``fma.rn`` pinned in PTX. Forward, ``dKV`` and ``dScore`` are bitwise run-to-run deterministic in BOTH families. -At ``ratio == 4`` dKV/dScore are additionally bit-identical to the fp32-intermediate +At ``ratio in {2, 4}`` dKV/dScore are additionally bit-identical to the fp32-intermediate eager autograd; at ``ratio == 128`` the contract is faithfulness to that fp32-intermediate eager reference: out/dKV/dScore match it within final-bf16 rounding at the gate tolerances (differing elements <= max(1, 0.1%), max_abs <= 1.6e-2, @@ -198,10 +198,10 @@ def check_support(self) -> bool: APIs; there is no soft fallback path inside this API. """ self._logger.debug("Entering check_support") - if self.ratio == 4: + if self.ratio in (2, 4): self._value_error_if( self.coff not in (1, 2), - f"CSA compressor at ratio=4 is validated for coff in {{1, 2}} (coff=2 is the production CSA/HCA form), got coff={self.coff}", + f"CSA compressor at ratio={self.ratio} is validated for coff in {{1, 2}} (coff=2 is the production CSA/HCA form), got coff={self.coff}", ) elif self.ratio == 128: self._value_error_if( @@ -211,7 +211,7 @@ def check_support(self) -> bool: else: self._value_error_if( True, - f"CSA compressor is validated for ratio in {{4, 128}} only, got ratio={self.ratio}, coff={self.coff}", + f"CSA compressor is validated for ratio in {{2, 4, 128}} only, got ratio={self.ratio}, coff={self.coff}", ) self._value_error_if( self.kv_desc.ndim != 2, @@ -599,11 +599,11 @@ def csa_compressor_forward_wrapper( cu_seqlens_comp: ``(B + 1,)`` int32 cumulative compressed-block counts, ``cu_seqlens_comp[b + 1] - cu_seqlens_comp[b] == seqlen_b // ratio``. ratio: compression ratio (tokens per output block); validated envelope: - {4, 128} (the wrappers route to the matching kernel family by ratio). + {2, 4, 128} (the wrappers route to the matching kernel family by ratio). head_dim: output feature dimension; inferred from ``kv`` width when omitted. coff: 1 for the own-block window form (window = ``ratio`` tokens, no overlap) or 2 for the overlapping-window form (window = ``2 * ratio``); validated - envelope: {1, 2} at both ratios (ratio=128 additionally requires head_dim + envelope: {1, 2} at all ratios (ratio=128 additionally requires head_dim in {128, 512}). total_comp: output row count. Defaults to ``cu_seqlens_comp[-1]`` (synchronizes); pass it explicitly (e.g. a static CUDA-graph capacity, which must be diff --git a/test/python/fe_api/csa/test_CSA_compressor.py b/test/python/fe_api/csa/test_CSA_compressor.py index 382783f98..91821657e 100644 --- a/test/python/fe_api/csa/test_CSA_compressor.py +++ b/test/python/fe_api/csa/test_CSA_compressor.py @@ -7,7 +7,7 @@ measurements and numerics in https://github.com/NVIDIA/Megatron-LM/issues/5968). Covers: - numerics of the fused region vs an fp32-intermediate eager reference, per ratio - family: at ``ratio == 4`` ``dKV``/``dScore`` are bit-identical and the forward is + family: at ``ratio in {2, 4}`` ``dKV``/``dScore`` are bit-identical and the forward is within one bf16 rounding step on a tiny fraction of elements; at ``ratio == 128`` the contract is faithfulness to the fp32 eager reference: all three match it within tolerance thresholds (differing elements <= max(1, 0.1%), max_abs <= @@ -37,7 +37,7 @@ bitwise against a direct call), and the loud error when the first call for a configuration would JIT under capture; - ``check_support`` boundaries (validated envelope: compute-capability major >= 10; - ratio 4 with coff + ratio {2, 4} with coff {1, 2}, ratio 128 with coff {1, 2} x head_dim {128, 512}; BF16 kv/score, FP32 ape, int32 cu_seqlens and int32 flat-offset bounds). @@ -217,7 +217,8 @@ def _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go): def _assert_grads_vs_fp32(gkv, gs, ref_kv, ref_s, ratio): """dKV/dScore vs the fp32-intermediate eager reference, per the ratio's contract. - ratio=4: bit-identical (the production bitwise contract, unchanged). + ratio {2, 4} (generic kernels): bit-identical (the production bitwise contract, + unchanged). ratio=128: deterministic tolerance contract — the fused backward reorders the den/S reductions (fixed chunk merge) and hoists 1/den, and some forward buckets use the ex2.approx fast exp, so dKV/dScore match eager within the forward-style @@ -225,9 +226,9 @@ def _assert_grads_vs_fp32(gkv, gs, ref_kv, ref_s, ratio): distribution; they stay bitwise run-to-run, and the fp64-oracle parity assertion below keeps the accuracy honest on the tested finite-intermediate inputs). """ - if ratio == 4: - assert torch.equal(gkv, ref_kv), "dKV must be bit-identical to the fp32 reference at ratio=4" - assert torch.equal(gs, ref_s), "dScore must be bit-identical to the fp32 reference at ratio=4" + if ratio in (2, 4): + assert torch.equal(gkv, ref_kv), f"dKV must be bit-identical to the fp32 reference at ratio={ratio}" + assert torch.equal(gs, ref_s), f"dScore must be bit-identical to the fp32 reference at ratio={ratio}" return for name, fused_t, ref_t in (("dKV", gkv, ref_kv), ("dScore", gs, ref_s)): diff = (fused_t.float() - ref_t.float()).abs() @@ -254,6 +255,22 @@ def _assert_grads_vs_fp32(gkv, gs, ref_kv, ref_s, ratio): pytest.param([3, 515, 1024, 129], 128, 4, 1, id="short-seg-d128-r4-coff1"), pytest.param([260], 65, 4, 1, id="b1-d65-odd-r4-coff1"), pytest.param([64, 0, 253, 3], 128, 4, 1, id="empty-seg-d128-r4-coff1"), + # ratio=2 (generic kernels; the compression ratio used in production training): + # the same shape classes as the ratio=4 rows above, both window forms. The + # short-seg pack leads with a 1-token segment so the seqlen < ratio class is + # still exercised at ratio=2. + pytest.param([2048], 128, 2, 2, id="b1-d128-r2"), + pytest.param([1023, 2048, 509], 128, 2, 2, id="ragged3-d128-r2"), + pytest.param([2048], 512, 2, 2, id="b1-d512-r2"), + pytest.param([1, 515, 1024, 129], 128, 2, 2, id="short-seg-d128-r2"), + pytest.param([260], 65, 2, 2, id="b1-d65-odd-r2"), + pytest.param([64, 0, 253, 3], 128, 2, 2, id="empty-seg-d128-r2"), + pytest.param([2048], 128, 2, 1, id="b1-d128-r2-coff1"), + pytest.param([1023, 2048, 509], 128, 2, 1, id="ragged3-d128-r2-coff1"), + pytest.param([2048], 512, 2, 1, id="b1-d512-r2-coff1"), + pytest.param([1, 515, 1024, 129], 128, 2, 1, id="short-seg-d128-r2-coff1"), + pytest.param([260], 65, 2, 1, id="b1-d65-odd-r2-coff1"), + pytest.param([64, 0, 253, 3], 128, 2, 1, id="empty-seg-d128-r2-coff1"), # ratio=128 (dedicated r128 kernels; coff {1, 2} x head_dim {128, 512}). The edge # pack covers zero-block segments (127, 3), a literal empty segment, an # exactly-one-block segment (128), a 1-token tail (129) and other tails. @@ -275,7 +292,7 @@ def _assert_grads_vs_fp32(gkv, gs, ref_kv, ref_s, ratio): @pytest.mark.L0 @pytest.mark.parametrize("lens,d,ratio,coff", _SHAPES) def test_numerics_vs_references(lens, d, ratio, coff): - """Fused fwd+bwd vs fp32-eager (bitwise dKV/dScore at ratio=4, tolerance at + """Fused fwd+bwd vs fp32-eager (bitwise dKV/dScore at ratio {2, 4}, tolerance at ratio=128), upstream eager, and fp64 oracle.""" _require_sm100() kv, score, ape, cu, cuc, total_comp, go = _make_inputs(lens, d, ratio, coff) @@ -286,7 +303,7 @@ def test_numerics_vs_references(lens, d, ratio, coff): r_fp64 = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="fp64") # vs fp32-intermediate eager reference (the fused kernels' numerics contract): - # dKV / dScore bit-identical at ratio=4 / within the forward-style tolerances at + # dKV / dScore bit-identical at ratio {2, 4} / within the forward-style tolerances at # ratio=128 (see _assert_grads_vs_fp32); forward within one bf16 rounding step on # a tiny fraction of elements; dAPE within fp32 atomics reorder noise. _assert_grads_vs_fp32(r_fused[1], r_fused[2], r_fp32[1], r_fp32[2], ratio) @@ -304,8 +321,8 @@ def test_numerics_vs_references(lens, d, ratio, coff): # vs the fp64 oracle. ratio=128: the contract's fp64-parity gate — per tensor the # fused output must be at least as close to the oracle as the FP32-INTERMEDIATE # eager reference, the comparator the contract names (comparing against the - # bf16-weight upstream path instead would be materially looser). ratio=4 keeps its - # historical check against the upstream numerics it replaced (its contract pins + # bf16-weight upstream path instead would be materially looser). ratio {2, 4} keeps + # its historical check against the upstream numerics it replaced (its contract pins # dKV/dScore bitwise-to-fp32-eager above and has no fp64-parity clause). eager_ref = r_fp32 if ratio == 128 else r_up for i in range(4): @@ -315,12 +332,13 @@ def test_numerics_vs_references(lens, d, ratio, coff): @pytest.mark.L0 +@pytest.mark.parametrize("ratio", [2, 4]) @pytest.mark.parametrize("coff", [1, 2]) -def test_replay_determinism(coff): +def test_replay_determinism(ratio, coff): """Forward, dKV and dScore replay bitwise identically run to run (dAPE is exempt).""" _require_sm100() - kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1023, 2048, 509], 128, 4, coff) - runs = [_run_fused(kv, score, ape, cu, cuc, total_comp, 4, 128, coff, go) for _ in range(3)] + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1023, 2048, 509], 128, ratio, coff) + runs = [_run_fused(kv, score, ape, cu, cuc, total_comp, ratio, 128, coff, go) for _ in range(3)] for other in runs[1:]: assert torch.equal(runs[0][0], other[0]) assert torch.equal(runs[0][1], other[1]) @@ -485,6 +503,16 @@ def test_r128_envelope_execution(coff, d, nb_rows): ([1023, 2048, 509], 128, 4, 2, 8), ([3, 515, 1024, 129], 128, 4, 1, 8), ([1023, 2048, 509], 128, 4, 1, 8), + # ratio=2: same classes; the leading 1-token segment keeps the padding-row gather + # spanning a segment boundary (seqlen < ratio at ratio=2). The [1, 1] pack isolates + # the padding-only path (nb_valid == 0 with a positive static capacity): every + # output row is a padding row gathering from token 0. + ([1, 515, 1024, 129], 128, 2, 2, 8), + ([1023, 2048, 509], 128, 2, 2, 8), + ([1, 1], 128, 2, 2, 8), + ([1, 515, 1024, 129], 128, 2, 1, 8), + ([1023, 2048, 509], 128, 2, 1, 8), + ([1, 1], 128, 2, 1, 8), ] @@ -620,6 +648,27 @@ def _never_consumed_mask(lens, total_tokens, d, ratio, coff): pytest.param([1023, 2048, 509], 128, 4, 1, 8, 0, id="ragged3-d128-padded-coff1"), pytest.param([1023, 2048, 509], 128, 4, 1, 0, 37, id="ragged3-d128-tokpad-coff1"), pytest.param([3, 515, 1024, 129], 128, 4, 1, 8, 21, id="short-seg-d128-padded-tokpad-coff1"), + # ratio=2 (generic kernels): the same never-consumed slot classes as the ratio=4 + # rows above; the tiny-segment packs shrink to lengths 0-1 so the zero-block + # (seqlen < ratio) class is still hit at ratio=2. + pytest.param([2048], 128, 2, 2, 0, 0, id="b1-d128-r2"), + pytest.param([1023, 2048, 509], 128, 2, 2, 0, 0, id="ragged3-d128-r2"), + pytest.param([1, 515, 1024, 129], 128, 2, 2, 0, 0, id="short-seg-d128-r2"), + pytest.param([1, 2, 3], 128, 2, 2, 0, 0, id="all-tiny-d128-r2"), + pytest.param([1023, 2048, 509], 512, 2, 2, 0, 0, id="ragged3-d512-r2"), + pytest.param([1, 515, 1024, 129], 128, 2, 2, 8, 0, id="short-seg-d128-r2-padded"), + pytest.param([1023, 2048, 509], 128, 2, 2, 8, 0, id="ragged3-d128-r2-padded"), + pytest.param([1023, 2048, 509], 128, 2, 2, 0, 37, id="ragged3-d128-r2-tokpad"), + pytest.param([1, 515, 1024, 129], 128, 2, 2, 8, 21, id="short-seg-d128-r2-padded-tokpad"), + pytest.param([2048], 128, 2, 1, 0, 0, id="b1-d128-r2-coff1"), + pytest.param([1023, 2048, 509], 128, 2, 1, 0, 0, id="ragged3-d128-r2-coff1"), + pytest.param([1, 515, 1024, 129], 128, 2, 1, 0, 0, id="short-seg-d128-r2-coff1"), + pytest.param([1, 2, 3], 128, 2, 1, 0, 0, id="all-tiny-d128-r2-coff1"), + pytest.param([1023, 2048, 509], 512, 2, 1, 0, 0, id="ragged3-d512-r2-coff1"), + pytest.param([1, 515, 1024, 129], 128, 2, 1, 8, 0, id="short-seg-d128-r2-padded-coff1"), + pytest.param([1023, 2048, 509], 128, 2, 1, 8, 0, id="ragged3-d128-r2-padded-coff1"), + pytest.param([1023, 2048, 509], 128, 2, 1, 0, 37, id="ragged3-d128-r2-tokpad-coff1"), + pytest.param([1, 515, 1024, 129], 128, 2, 1, 8, 21, id="short-seg-d128-r2-padded-tokpad-coff1"), # ratio=128: the zero classes are up to 127 tokens each (tails, zero-block # segments) plus the coff=2 last-block first-half (128 rows). pytest.param([1023, 2048, 509], 128, 128, 1, 0, 0, id="ragged3-d128-r128c1"), @@ -685,7 +734,7 @@ def run(poison): assert torch.allclose(gape_nan, gape_ref, rtol=0, atol=1e-3) # And the zero-slot pattern matches autograd: never-consumed slots are exact zeros, - # exactly as the fp32 eager reference computes them (bitwise at ratio=4, tolerance + # exactly as the fp32 eager reference computes them (bitwise at ratio {2, 4}, tolerance # at ratio=128 — the zero slots themselves are exact in both). (The fused backward # ignores incoming gradients on static-capacity padding rows by design, so the # eager reference runs with those rows zeroed.) @@ -1040,15 +1089,16 @@ def _meta_samples( @pytest.mark.L0 +@pytest.mark.parametrize("ratio", [2, 4]) @pytest.mark.parametrize("coff", [1, 2]) -def test_check_support_accepts_envelope(coff): +def test_check_support_accepts_envelope(ratio, coff): """Metadata-only samples inside the validated envelope pass check_support.""" _require_sm100() compressor = _import_compressor() for cls in (compressor.CSACompressorForward, compressor.CSACompressorBackward): - api = cls(**_meta_samples(coff=coff), ratio=4, coff=coff) + api = cls(**_meta_samples(ratio=ratio, coff=coff), ratio=ratio, coff=coff) assert api.check_support() is True - assert api.head_dim == 128 and api.total_tokens == 512 and api.total_comp == 128 + assert api.head_dim == 128 and api.total_tokens == 512 and api.total_comp == 512 // ratio @pytest.mark.L0 @@ -1089,9 +1139,11 @@ def test_check_support_rejects_compute_capability_major_below_10(monkeypatch): "kwargs,ctor,match", [ (dict(), dict(ratio=128, coff=3), "ratio=128 supports coff"), - (dict(), dict(ratio=8, coff=2), "ratio in \\{4, 128\\}"), + (dict(), dict(ratio=8, coff=2), "ratio in \\{2, 4, 128\\}"), (dict(coff=3), dict(ratio=4, coff=3), "coff in"), (dict(), dict(ratio=4, coff=0), "coff in"), + (dict(ratio=2, coff=3), dict(ratio=2, coff=3), "coff in"), + (dict(ratio=2), dict(ratio=2, coff=0), "coff in"), (dict(kv_dtype=torch.float16), dict(), "kv"), (dict(ape_dtype=torch.bfloat16), dict(), "ape"), (dict(cu_dtype=torch.int64), dict(), "cu_seqlens"), @@ -1101,6 +1153,7 @@ def test_check_support_rejects_compute_capability_major_below_10(monkeypatch): (dict(total=2**25, d=128), dict(), "int32 flat offsets"), (dict(total=4, d=8388482, total_comp=1), dict(), "head_dim"), (dict(total=2, total_comp=1), dict(), "requires at least ratio"), + (dict(total=1, total_comp=1, ratio=2), dict(ratio=2), "requires at least ratio"), (dict(kv_stride=(512, 2)), dict(), "contiguous"), (dict(total=1024, d=96, ratio=128, coff=1), dict(ratio=128, coff=1), "ratio=128 is validated for head_dim"), ],