CSA compressor: extend the validated ratio envelope to ratio=2

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`

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

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,17 +217,18 @@ 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
     tolerances instead of bitwise (thresholds calibrated on this suite's input
     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"),
     ],