Support ragged patch grids in attentional_pool

[amazloumi]

Summary

• attentional_pool now pools a patch grid the pool window does not evenly divide, instead of rejecting it. The grid is padded to ceil(grid/window) * window and each partial edge window pools only its real patches — the padded patches are masked out of that window's K/V and out of the mean-query. Output length is ceil(grid/window)², so any window pools any grid (e.g. a 3×3 window on a 14×14 SigLIP2 grid → 5×5 = 25 tokens, previously refused).
• Both pooling adapters share a _pad_grid_to_windows helper for the pad + validity mask, and avgpool now masks its window sum through the same valid tensor rather than relying on the pad value being zero.
• A divisible grid builds no mask and takes the same unmasked path as before, so no existing config's token count, shapes, or loss change.
• BREAKING: pooled_token_count loses its keyword-only require_divisible parameter and DIVISIBLE_ONLY_POOL_TYPES is removed. The parameter existed only to raise on ragged grids, which is the behaviour being removed, so there is no shim — out-of-tree callers passing require_divisible= must drop the argument.

Testing

• uv run pytest tests/unit/ -v --timeout=120 — 1757 passed, 3 skipped (skips are pre-existing gates: reserved arches, mlflow absent, RUN_HF_TESTS)
• uv run pytest examples/vlm/tests -q — 33 passed; uv run pytest examples/vlm/eval/tests/unit -q — 104 passed
• uv run ruff check kempnerforge/ tests/ scripts/ passes; uv run ruff format --check kempnerforge/ tests/ scripts/ — 166 files already formatted
• uv run pyright kempnerforge/ — 0 errors, 0 warnings
• tests/unit/test_adapter.py — 118 passed. Corner cases: grid % window == 1 (7×7 @ 3 → the corner window has exactly one real patch), window > grid (2×2 @ 3 → 1 token), window == grid (4×4 @ 4 → 1 token, divisible path), non-perfect-square token counts still raise.
• Ragged semantics cross-checked against an independent brute-force reference that attends over each window's real patches with no padded tensor at all, over 7 (grid, window) pairs incl. divisible ones — outputs agree to atol=1e-5. Separately, padding with 1e4 instead of 0 leaves both adapters' output unchanged, so padding is masked rather than merely zero.
• Mutation-tested detection power: dropping the SDPA attn_mask fails 11 tests; replacing the ragged denominator with w*w fails the avgpool edge-window test; reverting avgpool's numerator masking fails only the 4 pad-value tests (i.e. that change is bit-exact under zero padding).
• Token-count invariance across all 13 shipped examples/vlm/configs/*.toml, comparing this branch's pooled_token_count against main's (extracted verbatim from git show origin/main:…/adapter.py): 10 configs use mlp_2layer (identity, unaffected), 3 leave vision_encoder.num_tokens=0 (probed at build), 0 use attentional_pool. Because that set is weak on its own, also swept the whole space — grid 1..64 × window 1..16 = 1024 pairs: 211 divisible pairs unchanged, 0 changed; 813 ragged pairs previously rejected are now allowed; 0 ragged counts changed. The checker refuses (exit 2) on an empty/incomplete config enumeration, a config that fails to load, and a truncated/empty/helper-less main source — all 6 refusals verified against genuinely bad input.
• 4-GPU H200 (kempner_eng), ragged run — the capability: examples/vlm/train.py with attentional_pool, pool_window=3, 14×14 grid → 25 tokens/image, 20 steps forward+backward under FSDP2 (dp_shard=4), loss 10.81 → 8.63, step_20 committed. The same config on origin/main exits 1 with ValueError: this pooling connector requires the patch grid (14x14) be divisible by the pool window (3).
• Defaults unchanged, divisible run: same config at pool_window=2 (14×14 → 49 tokens) on this branch and on origin/main, 4×H200 each from an empty checkpoint dir — num_image_tokens=49 and 29,831,680 params on both, and the per-step loss is identical at every one of the 20 steps (10.8125 → 8.3750, max |Δ| = 0.000000). The comparison tool refuses on a missing/empty log, a log with no step lines, a log short of max_steps, a log ending in a traceback, and asymmetric step ids (7 guards verified).
• 4-GPU H200 torchrun -m pytest: tests/distributed/test_checkpoint.py 8 passed, test_resilience.py 5 passed, test_vlm_fsdp.py 6 passed — the adapter is FSDP2-wrapped on the VLM path.
• SIGTERM preemption drill on the ragged config (4×H200): SIGTERM at step 63 → Shutdown requested at step 63 — saving emergency checkpoint → step_63 committed → Graceful shutdown complete. Auto-resume logs Resumed from step 63, 290,304 tokens seen and Resumed DataLoader: epoch=0, skip_batches=63, then continues from step 64 with loss in the pre-interrupt range (3.72 → 4.19).
• 2-node run of the ragged config via scripts/slurm/multinode.sh: 2 nodes x 2 H200 (holygpu8a[10201,10601]), cross-node NCCL/gloo over ib0, srun one task per GPU, world_size=4, num_image_tokens=25, Applied VLM FSDP2: dp_mesh=('dp_shard',) — 20 steps, step_20 committed, exit 0. Its per-step loss is identical to the single-node 4-GPU ragged run at every step (max |delta| = 0.000000), so the node layout does not perturb the ragged path.
• Ragged attentional_pool also exercised through all four VLM arches (joint_decoder, mot, moma, cross_attention) forward+backward — the pooled count feeds the residual/positional layout, and all four give finite grads at 25 pooled tokens.
• torch.compile(fullgraph=True) over both pooling adapters at 5 (grid, window) pairs incl. 3 ragged: no graph break, outputs match eager to atol=1e-5, backward finite.
• No new or changed checkpoint key: attentional_pool's 10 state-dict keys and avgpool's 2 are byte-identical across pool_window 2/3/5/7, and window-2 weights load unchanged into a window-3 module. Activation checkpointing wraps only TransformerBlock/Attention, never the adapter, so there is no recompute interaction.
• Design plan re-verified against kf-codebase-design-new.html after implementation: qualifies for core as a general model knob / capability improvement to an existing core model block (Part 3 lists "Ragged pooling → To core"), entering through the existing adapter registry key and adapter.pool_window field — no new registry category, no new config field, no core→examples/ dependency, and it helps every pooling VLM rather than one experiment.
• If training loop / parallelism / optimizers changed: n/a — no change to the training loop, parallelism, optimizers, collectives, or checkpoint keys.

Closes #135

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
kempnerforge/config/adapter.py	100.00% <100.00%> (ø)
kempnerforge/model/adapter.py	97.43% <100.00%> (+0.15%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Naeemkh]

I checked the two places this could go wrong and both hold up: the validity mask and the patch data are reshaped in the same window order, and since the number of windows is a ceiling division, every window keeps at least one real patch, so a fully masked window, which would make attention return NaN, is unreachable rather than merely untested. The breaking removal is self-contained; every reference to the dropped parameter lives in a file this PR already changes, so nothing is orphaned.