fix: v0.4.3 correctness and robustness fixes

[tonywu71]

Context

A post-v0.4.2 review of the library surfaced five verified bugs (two reproducible crashes among them) plus ten API-contract, docstring, and hardening issues. This PR fixes all fifteen, with regression tests for everything CPU-testable and GPU-gated tests for the rest.

The fixes were validated end-to-end on an H100: the full GPU test suite passes, and the affected benchmarks were re-run to compare against the published tables. That benchmark pass caught a real 30–50% PLAID regression from the initial max_Ld bucketing change (a 300→512 loop bound took lik_full from 3.78 to 5.41 ms at 10k×300), reverted for the plaid kernels in 298ffac. Full numbers are in the benchmark validation comment.

Note

Behavior changes, all listed in the CHANGELOG under [0.4.3]:

• maxsim_varlen keys its autotune cache on power-of-two buckets of max_lq/max_ld; the kernel loop bounds stay exact (perf domain only; results unchanged, sweep count capped, zero masked iterations added).
• maxsim_residual squeezes a 2-D Q back to [Nd] instead of [1, Nd], matching its siblings.
• The FP8 fallback dispatches by device and no longer requires Triton.
• Caller-supplied max_seqlen_* values are validated with an on-device torch._assert_async (kept async to preserve the sync-free hot path these arguments exist for).

None of these affect the native LIK backends in PyLate (pylate#222) or colpali-engine (colpali#412): both only call maxsim / maxsim_pairs / maxsim_mps, whose signatures and return shapes are unchanged (pinned by tests/test_public_api.py); the behavior changes above are all in functions neither imports.

Benchmark tables in the docs are untouched: after the plaid revert, every re-measured number matched the published tables within run noise, so changing them would only churn noise.

Validation

• CPU suite: 279 passed, 318 skipped (CUDA-gated); ruff check, ruff format --check, ty check clean
• Full GPU suite on the H100 cluster: exit 0 with -x (only colpali_engine/MPS skips)
• Benchmarks re-run on 1×H100 via SkyPilot (forward, backward_method, decompress_maxsim, fastplaid_e2e, lowmem backward); cluster torn down afterwards

Changes

Click here to expand

Verified bugs

• Chunked top-k crash when top_k exceeds the merge width — retrieve(Q, D, top_k=10, chunk=4) raised "selected index k out of range" on both the CPU path (retrieve.py) and the CUDA path (topk.py): each chunk contributes at most chunk columns but the merge always asked torch.topk for k. The merge now clamps k to the merged width; the final output still honors min(top_k, Nd).
• MaxSimScorer.forward() overrode a caller's torch.no_grad() — _score's non-inference branch forced torch.enable_grad(), so scoring inside torch.no_grad() returned requires_grad=True tensors. Now contextlib.nullcontext().
• maxsim_residual_varlen broken for Nd == 0 — the grid was Nq * max(Nd, 1) while the kernel computes pid // Nd with constexpr Nd = 0. The empty [Nq, 0] result is returned before launching.
• pack_padded crashed on an empty batch (B == 0, qlen.max() on a zero-element tensor) — early-exit returning the trivially-empty PackedBatch; CUDA and CPU maxsim_padded now agree on [0, C].
• PyLate legacy mask= silently dropped on the fallback path — the patched wrappers mapped mask= into the fused path but called the original function with a still-None documents_mask when deferring to PyLate (CPU / sub-Ampere / LIK_DISABLE=1). The resolved mask is now forwarded in all three wrappers.

API contracts and docstrings

• maxsim_varlen autotune sweeps capped by a bucketed cache key — max_lq/max_ld were constexpr autotune keys with no bucketing, so every distinct pair re-triggered the full sweep (and a recompile), contrary to the docs. They are now exact runtime loop bounds — zero masked iterations, exact-sized argmax buffer — and the autotune cache is keyed on power-of-two-bucketed copies (floor 16) passed as key-only dummy args, the same trick the dense kernel uses to keep Ld out of its key. An earlier draft of this fix bucketed the loop bounds themselves; the H100 benchmark pass measured that idea at 30–50% steady-state loss on the plaid kernels (298ffac), and since nothing in benchmarks/ exercises maxsim_varlen, the same regression there would have gone unmeasured — hence exact bounds everywhere, bucketing only in the cache key. Plaid keeps exact max_Ld in its key, which is index-stable per corpus, so its sweep already amortizes to one.
• max_seqlen_* documented as hard loop bounds, not hints (maxsim_varlen, score_pairs_packed, maxsim_residual_varlen) — a too-small value silently truncated tokens. Values are now validated against the cu_seqlens maxima on-device.
• maxsim_residual 2-D Q contract — returns [Nd], matching maxsim_residual_varlen and maxsim.
• Stale FP8 docstring — removed the argmax-ties note (the inference kernel never computes an argmax).
• lowmem "bf16 grads" wording — grads are written in the input dtype (fp16/fp32 included); module, function, and maxsim docstrings corrected.

Hardening and dead code

• Removed the dead B constexpr from _plaid_approx_score_kernel (it forced a recompile per batch size), the unused Lq/Ld placeholder params from the varlen forward kernel, and the dead Nq constexpr from _varlen_bwd_dQ_kernel; the varlen backward kernels take max_lq as a runtime arg (mirroring the score_pairs backward kernels), so a new query-length maximum no longer recompiles them. prune_forward reads max_lq so varlen pruning is unchanged, with explicit is None fallbacks so a legit Lq == 0 is not treated as missing.
• maxsim_backward_lowmem gets the same clean Triton/CUDA RuntimeError guard as maxsim_backward_unified.
• prune_forward falls back to the two smallest-SMEM configs when nothing survives the budget (was: the first two arbitrary entries, which had just been pruned for overflowing).
• The KD path (maxsim with 4-D D) validates Q/D/mask device agreement like the 3-D path (shared _validate_devices helper).
• The FP8 fallback dispatches by device like retrieve (Triton bf16 kernel on CUDA, compiled reference on MPS, eager reference elsewhere), with a CPU regression test — chosen over a docstring-only fix.

Tests

• New CPU regression tests: chunked top-k clamp, no_grad scorer, empty-batch pack_padded, legacy-mask fallback parity, FP8 CPU fallback, bucket_seqlen / assert_max_seqlen_covers, lowmem guard, prune fallback, prune max_lq spelling + Lq == 0 handling.
• New GPU-gated tests: chunked top-k on CUDA, residual 2-D squeeze, empty varlen corpus, explicit-max-seqlen parity, async-assert subprocess isolation.

Next steps

GPU CI needs the run-gpu-tests label (triage+) or a workflow_dispatch of the GPU CI workflow — the new GPU-gated tests have only run on the SkyPilot H100, not in CI. Note the varlen exact-bounds change (1fc735d, post-review) and the rebase onto main landed after that H100 run, so the GPU suite must be re-run before merge.

[tonywu71]

H100 benchmark validation (SkyPilot, 1×H100 80GB, torch 2.8.0+cu128)

Reran the benches these fixes could plausibly move, via scripts/sky_run_all_benchmarks.yaml (RUN_ONLY="forward backward_method decompress_maxsim fastplaid_e2e") plus bench_backward_lowmem.py and the full GPU test suite on the same cluster.

Outcome: no doc tables needed updating — but the run caught one real regression, fixed in 298ffac.

• PLAID bucketing regression caught and reverted (298ffac). With max_Ld bucketed to the next power of two, the fastplaid e2e bench regressed 30–50% on non-power-of-2 corpora (lik_full 3.78→5.41 ms on 10k×300, 9.25→13.90 ms on 25k×300; ×512 corpora unchanged) — the masked decompress+dot iterations from a 300→512 loop bound are not free. Since max_Ld is a property of the compressed index (stable across calls), the per-value autotune sweep amortizes to one and bucketing buys nothing there. The PLAID kernels now keep their exact max_Ld; maxsim_varlen keeps its bucketing (training batches have genuinely varying maxima, and that sweep cost is the documented pain point — same trade as the dense path's Lq bucketing).
• After the revert, all PLAID numbers match the documented tables within noise: full/partial speedups 15.9×/18.8×, 12.2×/26.4×, 13.8×/30.9×, 22.3×/49.5×, 7.9×/42.9× vs documented 15.9/18.3, 12.5/27.3, 14.2/32.0, 22.9/51.0, 8.3/45.1. Decompress-vs-transliteration: 2.98–4.04× (documented 3.0–4.0×).
• Forward headline table confirmed: substantial shapes match exactly (corpus-10k 0.247 ms, train-batch-128 0.227 ms, lateon-big 0.260 ms vs documented 0.247/0.226/0.260); launch-bound ~0.1 ms shapes show only ±10 µs timer noise. README ranges unchanged.
• Backward (unified/lowmem): peak-memory column reproduces exactly (96/52, 141/72, 1086/543, 4329/2165, 179/107 MB); step times within session noise, same auto-routing winners.
• Full GPU test suite passed on H100 against the final branch code (pytest -x, exit 0; only skips: colpali_engine not installed, 3 MPS tests). This covers all the new GPU-gated regression tests (chunked top-k, residual 2-D squeeze, empty varlen corpus, explicit max_seqlen parity, async-assert subprocess). GPU CI via the run-gpu-tests label would still be a good belt-and-suspenders on A10G.

[h-aurelien-lac]

Solid PR, the fixes and tests all check out. One real concern before merge though:

The varlen max_ld bucketing is the same pattern that cost PLAID 30-50% (which you caught and reverted) — but nothing in benchmarks/ actually exercises maxsim_varlen, so the re-run couldn't have caught a regression there. For retrieval over a fixed packed corpus, max_ld is just as stable as PLAID's max_Ld, and 300→512 means ~70% more dot iterations per program.

Since both loops in the kernel are plain range() (not static_range), couldn't max_lq/max_ld just become exact runtime args, with bucketed values passed as separate dummy args for the autotune key? That'd amortize the sweep with zero masked iterations — same trick the dense kernel already uses for Ld.

Two smaller things:

• named_args.get("Lq") or named_args.get("max_lq") or 32 — a legit Lq == 0 silently falls through, explicit is None checks would be safer.
• Given the behavior changes (residual squeeze, new device assert), this feels more like 0.5.0 than a patch.

[tonywu71]

Solid PR, the fixes and tests all check out. One real concern before merge though:

The varlen max_ld bucketing is the same pattern that cost PLAID 30-50% (which you caught and reverted) — but nothing in benchmarks/ actually exercises maxsim_varlen, so the re-run couldn't have caught a regression there. For retrieval over a fixed packed corpus, max_ld is just as stable as PLAID's max_Ld, and 300→512 means ~70% more dot iterations per program.

Since both loops in the kernel are plain range() (not static_range), couldn't max_lq/max_ld just become exact runtime args, with bucketed values passed as separate dummy args for the autotune key? That'd amortize the sweep with zero masked iterations — same trick the dense kernel already uses for Ld.

Two smaller things:

• named_args.get("Lq") or named_args.get("max_lq") or 32 — a legit Lq == 0 silently falls through, explicit is None checks would be safer.
• Given the behavior changes (residual squeeze, new device assert), this feels more like 0.5.0 than a patch.

Thanks for the review! There are no API breaking changes wrt to signatures and types. device assert is acceptable imo, and the newly introduced residual squeeze only affects the PLAID entrypoint which is not used by any of our package users (pylate and colpali-engine).

As agreed, I'll keep this release as a minor.