perf(utils): opt-in chunked demean_map dispatch + corrected solver-phase memory attribution

CHANGELOG.md

@@ -19,6 +19,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   when the extension is absent or `DIFF_DIFF_BACKEND=python`; kernel-side validation errors fall
   back to the numpy engine. Affects every MAP consumer (TWFE, SunAbraham, Bacon, Wooldridge,
   DiD/MPD `absorb=`, and the survey replicate refits).
+- **Opt-in chunked dispatch for the Rust `demean_map` kernel** (internal
+  `DIFF_DIFF_DEMEAN_CHUNK_COLS` env knob; default behavior unchanged - single full-width
+  dispatch). When set to a positive integer, the wrapper feeds the kernel balanced
+  near-equal column blocks, capping the kernel's transient input/working copies on wide
+  absorbed designs for memory-constrained runs (a width near the machine's core count
+  measured best). Results are identical by construction - each column's MAP loop is fully
+  independent, so chunking changes neither demeaned values nor iteration counts (locked by
+  exact-equality tests vs single-call dispatch). Shipped opt-in rather than default-on:
+  measured end-to-end on the FE-absorption lane, the fit-level peak RSS on the widest
+  workload is dominated by the downstream solver phase, so default chunking would cost
+  ~2-7% wall-clock for only ~5-12% peak-RSS reduction (details + corrected memory
+  attribution in `docs/performance-plan.md`).
 
 ### Changed
 - **FE-absorption demeaning rewritten: factorize-once + `np.bincount` method of alternating

TODO.md

@@ -46,7 +46,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 
 | Issue | Location | Origin | Effort | Priority |
 |-------|----------|--------|--------|----------|
-| SunAbraham-family fits are now SOLVER-bound after the v3.6.x demeaning speedups: on the county-class shape (3.1k units x 60 months, ~100 interaction columns) `solve_ols` = ~60% of fit time (faer SVD 1.24s + pivoted-QR rank detection 0.70s of 3.27s), and the share grows further under the Rust demean kernel. The parked "QR+SVD redundancy is minimal vs the SVD" note predates these measurements - the pivoted QR alone is now ~20% of the whole fit. Candidates: reuse the rank-detection QR factors for the solve, or a Cholesky fast path for well-conditioned tall-skinny designs. Measured attribution in docs/performance-plan.md (FE-absorption baseline). | `linalg.py::solve_ols`, `linalg.py::_detect_rank_deficiency` | PR-C attribution | Heavy | Medium |
+| Rust-backend `solve_ols` peak memory on wide designs: the faer path holds an owned copy of the stacked design plus SVD workspace, dominating fit-level peak RSS on wide absorbed fits (firm_churn-class: ~19-21 GB rust-backend vs ~13.4 GB numpy-backend for the same fit; measured PR-D 2026-07, corrected attribution in docs/performance-plan.md - the demean kernel's transients were NOT the dominator, and opt-in chunked dispatch caps them anyway). Candidates: borrow-view instead of owned copy into faer, workspace reuse, or routing wide designs to the numpy gelsd path. Interim escape hatch: DIFF_DIFF_BACKEND=python. | `rust/src/linalg.rs`, `diff_diff/linalg.py::solve_ols` | PR-D probe | Heavy | Medium |
 | `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium |
 | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low |
 | Rust-backend HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low |
@@ -124,6 +124,7 @@ Doable in principle, but no current caller and/or explicitly out of paper scope.
 | Issue | Location | PR | Priority |
 |-------|----------|----|----------|
 | `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low |
+| SunAbraham-family fits are SOLVER-bound after the v3.6.x demeaning speedups: on the county-class shape (3.1k units x 60 months, ~100 interaction columns) `solve_ols` = ~60% of fit time (faer SVD 1.24s + pivoted-QR rank detection 0.70s of 3.27s), and the share grows further under the Rust demean kernel; the pivoted QR alone is ~20% of the whole fit (measured attribution in docs/performance-plan.md). Candidate fix is QR-reuse (one factorization for rank detection + solve). **Parked per the 2026-07 correctness-first decision**: the change perturbs every coefficient in the library at the last digits (golden/parity/bit-identity recapture) and removes the SVD's independent singular-value truncation - a second, different near-collinearity detector that is deliberate defense-in-depth for the no-silent-failures contract (the FE-span guard episode showed borderline cases are real). Re-open only on practitioner demand for wide event studies; any implementation gates on fixest/R end-to-end parity, not just internal consistency. | `linalg.py::solve_ols`, `linalg.py::_detect_rank_deficiency` | PR-C attribution | Medium |
 | dCDH parity-test SE/CI assertions only cover pure-direction scenarios; mixed-direction SE comparison is structurally apples-to-oranges (cell-count vs obs-count weighting). | `test_chaisemartin_dhaultfoeuille_parity.py` | #294 | Low |
 | `HeterogeneousAdoptionDiD` survey-design API consolidation (**scheduled: next minor bump**): drop the deprecated `survey=` / `weights=` kwargs on all 8 HAD surfaces; only `survey_design=` remains. Also fold the legacy back-end `weights=` routing into the unified `_resolve_survey_for_fit` path. DeprecationWarning has shipped; the removal is ~50 LoC gated on the semver bump. | `had.py`, `had_pretests.py` | next minor | Medium |
 | `HeterogeneousAdoptionDiD` joint cross-horizon covariance / sup-t bands: per-horizon SEs use independent sandwiches (paper-faithful pointwise CIs per Pierce-Schott Fig 2). Follow-ups (low demand): IF-based stacking for joint cross-horizon inference; analytical H×H covariance on the weighted ES path; a sup-t band on the unweighted ES path. | `had.py::_fit_event_study` | Phase 2b / 4.5 B | Low |

diff_diff/utils.py

@@ -2,6 +2,7 @@
 Utility functions for difference-in-differences estimation.
 """
 
+import os
 import warnings
 from dataclasses import dataclass, field
 from typing import Any, Dict, Iterable, List, Optional, Tuple
@@ -13,15 +14,14 @@
 # Import Rust backend if available (from _backend to avoid circular imports)
 from diff_diff._backend import (
     HAS_RUST_BACKEND,
-    _rust_demean_map,
-    _rust_project_simplex,
-    _rust_sdid_unit_weights,
-    _rust_compute_time_weights,
     _rust_compute_noise_level,
+    _rust_compute_time_weights,
+    _rust_demean_map,
     _rust_sc_weight_fw,
-    _rust_sc_weight_fw_with_convergence,
     _rust_sc_weight_fw_weighted,
     _rust_sc_weight_fw_weighted_with_convergence,
+    _rust_sc_weight_fw_with_convergence,
+    _rust_sdid_unit_weights,
 )
 from diff_diff.linalg import compute_robust_vcov as _compute_robust_vcov_linalg
 from diff_diff.linalg import solve_ols as _solve_ols_linalg
@@ -2707,6 +2707,41 @@ def _demean_map_numpy(
     return demeaned, iters
 
 
+# Opt-in column-block width for Rust demean_map kernel calls (None = single
+# dispatch, the default). The kernel holds an owned copy of its input block
+# plus the result, and chunking caps those transients at the block width while
+# leaving per-column results bit-identical: every column's MAP loop is fully
+# independent (no cross-column arithmetic), so partitioning changes neither
+# values nor iteration counts. Measured on the 2.4M x 130 firm-panel workload
+# (2026-07): dispatch-level transients shrink as designed, but the fit's peak
+# RSS is dominated by the downstream solver phase, so end-to-end peak dropped
+# only ~5-12% while wall-clock rose ~2-7% - hence OFF by default; the env
+# knob remains for memory-constrained runs (a block width near the machine's
+# core count measured best: each chunk is one full parallel wave). Peak-RSS
+# measurements need repeated runs under matched machine state - macOS memory
+# compression deflates single-run ru_maxrss readings under ambient pressure.
+_DEMEAN_MAP_CHUNK_COLS: Optional[int] = None
+
+
+def _resolve_demean_chunk_cols() -> Optional[int]:
+    """Opt-in chunk width for the Rust kernel dispatch (None = unchunked).
+
+    ``DIFF_DIFF_DEMEAN_CHUNK_COLS`` (positive integer) enables chunking;
+    invalid or non-positive values fall back silently to the module default,
+    mirroring ``DIFF_DIFF_BACKEND``'s convention. Unlike the backend switch
+    this is read PER CALL, not at import - deliberate, so benchmarks and tests
+    can A/B chunked vs unchunked dispatch within one process/build.
+    """
+    raw = os.environ.get("DIFF_DIFF_DEMEAN_CHUNK_COLS")
+    if raw is None:
+        return _DEMEAN_MAP_CHUNK_COLS
+    try:
+        value = int(raw)
+    except ValueError:
+        return _DEMEAN_MAP_CHUNK_COLS
+    return value if value > 0 else _DEMEAN_MAP_CHUNK_COLS
+
+
 def _demean_map_rust(
     x_cols: List[np.ndarray],
     codes_list: List[np.ndarray],
@@ -2715,36 +2750,61 @@ def _demean_map_rust(
     tol: float,
     max_iter: int,
 ) -> Optional[Tuple[List[np.ndarray], List[int]]]:
-    """Marshal to the Rust ``demean_map`` kernel.
+    """Marshal to the Rust ``demean_map`` kernel in column chunks.
 
     Returns None to signal "use the canonical numpy engine" (kernel absent,
     degenerate shapes, or a deliberate kernel-side validation error). Dtypes
     are coerced explicitly BEFORE the call - never rely on exception handling
     for dtype mismatches.
+
+    Variables are dispatched in blocks of ``_resolve_demean_chunk_cols()``
+    columns to bound peak memory (see ``_DEMEAN_MAP_CHUNK_COLS``); the codes
+    matrix and weights are built once and shared across chunks. Chunking is
+    exact partitioning - per-column outputs and iteration counts are identical
+    to a single-call dispatch.
     """
     if _rust_demean_map is None:
         return None
     if not x_cols or x_cols[0].shape[0] == 0:
         return None
-    x_mat = np.ascontiguousarray(np.column_stack(x_cols), dtype=np.float64)
     codes_mat = np.ascontiguousarray(np.column_stack(codes_list), dtype=np.int64)
     w = None if weights is None else np.ascontiguousarray(weights, dtype=np.float64)
-    try:
-        out, iters = _rust_demean_map(  # type: ignore[misc]
-            x_mat,
-            codes_mat,
-            [int(g) for g in n_groups_list],
-            w,
-            float(tol),
-            int(max_iter),
-        )
-    except ValueError as e:
-        if "demean_map" in str(e):
-            # deliberate kernel-side validation marker -> numpy fallback
-            return None
-        raise
-    # F-order result: per-column views are contiguous, no extra copy
-    return [out[:, j] for j in range(out.shape[1])], [int(i) for i in iters]
+    n_groups = [int(g) for g in n_groups_list]
+    k = len(x_cols)
+    chunk = _resolve_demean_chunk_cols()
+    if chunk is None:
+        chunk = k  # default: single dispatch, no partitioning
+    # Balanced partition into ceil(k / chunk) near-equal blocks. A naive
+    # fixed-stride split leaves a small remainder block (e.g. 130 columns at
+    # width 32 -> 4x32 + 2); that tail runs nearly serial across the rayon
+    # pool, and one slow-converging column in it stalls the whole dispatch.
+    n_chunks = -(-k // chunk)
+    bounds = [round(i * k / n_chunks) for i in range(n_chunks + 1)]
+    demeaned: List[np.ndarray] = []
+    iters_all: List[int] = []
+    for start, stop in zip(bounds, bounds[1:]):
+        x_mat = np.ascontiguousarray(np.column_stack(x_cols[start:stop]), dtype=np.float64)
+        try:
+            out, iters = _rust_demean_map(  # type: ignore[misc]
+                x_mat,
+                codes_mat,
+                n_groups,
+                w,
+                float(tol),
+                int(max_iter),
+            )
+        except ValueError as e:
+            if "demean_map" in str(e):
+                # deliberate kernel-side validation marker -> numpy fallback
+                # (kernel validation is chunk-invariant, so this fires on the
+                # first chunk; the whole call falls back, same contract as an
+                # unchunked dispatch)
+                return None
+            raise
+        # F-order result: per-column views are contiguous, no extra copy
+        demeaned.extend(out[:, j] for j in range(out.shape[1]))
+        iters_all.extend(int(i) for i in iters)
+    return demeaned, iters_all
 
 
 def demean_by_groups(

docs/doc-deps.yaml

@@ -904,7 +904,7 @@ sources:
       - path: docs/methodology/REGISTRY.md
         section: "Absorbed Fixed Effects with Survey Weights"
         type: methodology
-        note: "demean_by_groups / within_transform: MAP convergence contract (max_iter/tol), bincount accumulation numerics, NaN-group-key guard, snap_absorbed_regressors FE-spanned snap + exact LSMR span confirmation, rust demean_map dispatch"
+        note: "demean_by_groups / within_transform: MAP convergence contract (max_iter/tol), bincount accumulation numerics, NaN-group-key guard, snap_absorbed_regressors FE-spanned snap + exact LSMR span confirmation, rust demean_map dispatch (chunked, DIFF_DIFF_DEMEAN_CHUNK_COLS)"
       - path: docs/practitioner_getting_started.rst
         section: "Step 4: Check Whether the Result Is Trustworthy"
         type: user_guide

docs/methodology/REGISTRY.md

@@ -4351,6 +4351,11 @@ unequal selection probabilities).
   tests assert iteration-count equality plus `assert_allclose` at atol=1e-12
   (never a bit-identity claim). `DIFF_DIFF_BACKEND=python` disables the
   kernel; any kernel-side validation error falls back to the numpy engine.
+  The dispatch can partition variables into balanced column blocks
+  (`DIFF_DIFF_DEMEAN_CHUNK_COLS`, internal opt-in knob, OFF by default) to
+  bound the kernel's transient memory; per-column results and iteration
+  counts are unchanged by construction (each column's MAP loop is fully
+  independent - no cross-column arithmetic).
 - **Edge case:** NaN in an absorbed group column raises a `ValueError` naming the
   column. `pd.factorize` codes NaN keys as -1, which would otherwise silently index
   the last group's mean; the prior pandas behavior was itself silently bad

docs/performance-plan.md

@@ -76,9 +76,22 @@ submit gate), tail_stress 31.74s -> 11.37s (2.79x), geo_experiment 0.976s ->
 0.484s (2.02x, now ahead of the pyfixest yardstick's 0.623s), scanner 1.59x,
 survey 1.44x, county 1.36x, guard_small unregressed. Identity vs the
 committed after-baselines: 1e-13-1e-16 on every scenario incl. tail_stress.
-Known trade-off: the kernel holds one owned input copy plus the result
-(firm_churn peak RSS 13.4 GB -> 21.0 GB); chunking variables through the
-kernel is the noted follow-up if that bites a real workload.
+Known trade-off: the rust-backend firm_churn fit peaks at ~21 GB RSS vs
+~13.4 GB under the numpy backend. **Corrected attribution (2026-07, PR-D
+probes):** the peak is dominated by the downstream SOLVER phase (the Rust
+faer path holds an owned copy of the ~2.5 GB stacked design plus SVD
+workspace), not by the demean kernel's dispatch marshalling as this
+paragraph originally claimed. Measured evidence: an isolated width-16
+chunked dispatch cuts the kernel's transients to near-numpy footprint, yet
+the end-to-end fit still peaks ~19-21 GB (only ~5-12% below unchunked) at a
++2-7% wall-clock cost - so chunked dispatch shipped as an opt-in env knob
+(`DIFF_DIFF_DEMEAN_CHUNK_COLS`, default off) rather than default-on, and
+the solver-phase peak is tracked as its own TODO row. `DIFF_DIFF_BACKEND=
+python` remains the OOM escape hatch. Measurement note: single-run
+`ru_maxrss` on macOS is unreliable under ambient memory pressure (the
+compressor deflates resident peaks - one probe read 15.1 GB for a
+configuration that reproducibly peaks ~19-20 GB); gate memory claims on
+repeated runs under matched machine state.
 
 ### FE-absorption suite results