fix: v0.4.3 correctness and robustness fixes

CHANGELOG.md

@@ -4,7 +4,89 @@ All notable changes to this project will be documented here.
 Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the
 project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased]
+## [0.4.3] - 2026-06-10
+
+### Fixed
+
+- **Chunked top-k no longer crashes when `top_k` exceeds the merge width.**
+  `retrieve(Q, D, top_k=10, chunk=4)` (and the CUDA `maxsim_topk` it wraps)
+  raised "selected index k out of range": each chunk contributes at most
+  `chunk` columns, but the running merge always asked `torch.topk` for `k`.
+  The merge now clamps `k` to the merged width; the final output still
+  honors the documented `min(top_k, Nd)` contract on both paths.
+- **`MaxSimScorer.forward()` respects a caller's `torch.no_grad()`.** The
+  non-inference branch of `_score` forced `torch.enable_grad()`, so
+  `MaxSimScorer()(Q, D)` inside `torch.no_grad()` returned a tensor with
+  `requires_grad=True`. It now uses a null context and leaves the caller's
+  grad mode untouched.
+- **`maxsim_residual_varlen` handles an empty corpus (`Nd == 0`).** The
+  launcher used grid `Nq * max(Nd, 1)` while the kernel computed
+  `pid // Nd` with a constexpr `Nd=0`. The empty `[Nq, 0]` result is now
+  returned before any launch.
+- **`pack_padded` handles an empty batch (`B == 0`).** It crashed on
+  `qlen.max()` of a zero-element tensor; it now returns the trivially-empty
+  `PackedBatch` up front, so `maxsim_padded` agrees with the CPU reference
+  (`[0, C]`).
+- **PyLate legacy `mask=` is forwarded on the fallback path.** The patched
+  `colbert_scores` / `colbert_kd_scores` / `colbert_scores_pairwise` mapped
+  a legacy `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`), silently dropping the doc mask.
+- **`maxsim_backward_lowmem` raises a clean `RuntimeError` without
+  Triton/CUDA** instead of failing inside the launch — same guard as
+  `maxsim_backward_unified`, so `backward/__init__`'s import-anywhere
+  promise holds.
+- **`prune_forward` falls back to the smallest-footprint configs.** When
+  every autotune config overflowed the shared-memory budget, the fallback
+  returned the first two list entries — configs just pruned for exceeding
+  the budget. It now returns the two with the smallest estimated SMEM
+  footprint.
+- **The KD path (`maxsim` with 4-D `D`) validates device agreement** for
+  `Q` / `D` / masks like the 3-D path, instead of surfacing an internal
+  kernel error on mixed devices.
+- **FP8 fallback really works without Triton.** `maxsim_inference_fp8`'s
+  documented "transparent fallback" imported the Triton-backed `maxsim`
+  unconditionally; it now dispatches like `retrieve` (Triton on CUDA,
+  compiled reference on MPS, eager reference elsewhere), so the dequantized
+  bf16 fallback runs on CPU-only installs.
+
+### Changed
+
+- **`maxsim_varlen` autotune sweeps are amortized across batch shapes.**
+  `max_lq` / `max_ld` were constexpr autotune keys with no bucketing, so —
+  contrary to the previous docs — every distinct `(max_lq, max_ld)` pair
+  re-triggered the full autotune sweep *and* recompiled the kernel. They
+  are now exact runtime loop bounds (like the dense kernel's `Ld`), and
+  the autotune cache is keyed on power-of-two-bucketed copies (floor 16)
+  passed as key-only arguments: one sweep per bucket, zero masked
+  iterations added, argmax buffer still sized on the exact `max_lq`.
+  Results are unchanged. The PLAID kernels (`maxsim_residual` /
+  `maxsim_residual_varlen`) deliberately keep their exact `max_Ld` key: it
+  is a property of the compressed index, stable across calls, so the sweep
+  amortizes to one — bucketing the loop bound there measured 30-50%
+  steady-state throughput loss on `Ld=300` corpora for no benefit, which
+  is also why the varlen bounds stay exact.
+- **`max_seqlen_*` arguments are documented as hard loop bounds, not
+  hints**, in `maxsim_varlen`, `score_pairs_packed`, and
+  `maxsim_residual_varlen`: a too-small value silently truncated tokens and
+  returned wrong scores. Caller-supplied values are now checked against the
+  `cu_seqlens` maxima with an on-device `torch._assert_async` (no D2H sync,
+  same trade-off as `pack_padded`'s length checks).
+- **`maxsim_residual` squeezes a 2-D `Q` back to `[Nd]`** instead of
+  returning `[1, Nd]`, matching `maxsim_residual_varlen` and the `maxsim`
+  wrapper's convention. Behavior change for callers that relied on the
+  un-squeezed shape.
+- 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 also take `max_lq`
+  as a runtime arg instead of a constexpr (mirroring the `score_pairs`
+  backward kernels), so a new query-length maximum no longer recompiles
+  them.
+- Docstring corrections: `maxsim_inference_fp8` no longer mentions argmax
+  ties (the inference kernel never computes an argmax), and the `lowmem`
+  backward docs say grads are written in the *input* dtype (fp16 / bf16 /
+  fp32), not "bf16 grads".
 
 ### Removed
 

late_interaction_kernels/_autotune.py

@@ -122,28 +122,47 @@ def forward_configs():
     return base  # minimal safe shortlist for unknown GPUs
 
 
+def _smem_bytes(cfg, d_pad: int) -> int:
+    """Estimated shared-memory footprint of a config.
+
+    Triton allocates num_stages copies of Q/D in SMEM for the async-copy
+    pipeline. tl.dot accumulates S in fp32 registers (no SMEM copy), but
+    the estimate adds one bq×bd×4-byte slot as a conservative upper bound.
+    """
+    bq, bd = cfg.kwargs["BLOCK_Q"], cfg.kwargs["BLOCK_D"]
+    return (bq * d_pad + bd * d_pad) * 2 * cfg.num_stages + bq * bd * 4
+
+
 def prune_forward(configs, named_args, **kwargs):
     """Drop configs that overflow shared memory or are oversized for the problem."""
-    Lq = named_args.get("Lq", 32)
+    # The varlen kernels spell their query loop bound `max_lq` instead of `Lq`.
+    # `is None` (not `or`): a legit Lq == 0 must not fall through to 32.
+    Lq = named_args.get("Lq")
+    if Lq is None:
+        Lq = named_args.get("max_lq")
+    if Lq is None:
+        Lq = 32
     # The kernel tiles on the padded embedding dim (`d_pad = next_pow2(d)`), so
     # the SMEM estimate must use d_pad too — keying on raw d underestimates by
     # up to ~2x for non-power-of-2 d and admits configs that OOM at launch.
-    d_pad = named_args.get("d_pad") or next_pow2(named_args.get("d", 128))
+    d_pad = named_args.get("d_pad")
+    if d_pad is None:
+        d_pad = next_pow2(named_args.get("d", 128))
     gpu = detect_gpu()
     # Reserve 8 KiB for Triton scratch; the rest is ours.
     sram_budget = (_SRAM_KIB_BY_FAMILY.get(gpu, 48) - 8) * 1024
 
     keep = []
     for cfg in configs:
-        bq, bd = cfg.kwargs["BLOCK_Q"], cfg.kwargs["BLOCK_D"]
-        # Triton allocates num_stages copies of Q/D in SMEM for the async-copy
-        # pipeline. tl.dot accumulates S in fp32 registers (no SMEM copy), but
-        # the budget adds one bq×bd×4-byte slot as a conservative upper bound.
-        need = (bq * d_pad + bd * d_pad) * 2 * cfg.num_stages + bq * bd * 4
-        if need > sram_budget:
+        if _smem_bytes(cfg, d_pad) > sram_budget:
             continue
-        if bq > 2 * Lq:
+        if cfg.kwargs["BLOCK_Q"] > 2 * Lq:
             continue
         keep.append(cfg)
-    # Always return at least two configs so autotune has something to compare.
-    return keep or configs[:2]
+    if keep:
+        return keep
+    # Nothing survived the budget (huge d_pad on a small-SRAM GPU): fall back
+    # to the two smallest-footprint candidates — they were pruned too, but
+    # they are the least likely to overflow at launch. The first two list
+    # entries would be arbitrary.
+    return sorted(configs, key=lambda cfg: _smem_bytes(cfg, d_pad))[:2]

late_interaction_kernels/_utils.py

@@ -12,6 +12,44 @@ def next_pow2(x: int) -> int:
     return 1 << (x - 1).bit_length()
 
 
+def bucket_seqlen(max_len: int, floor: int = 16) -> int:
+    """Round a max sequence length up to the next power of two for use as an
+    autotune cache key.
+
+    Keying the autotune cache on an exact max length re-triggers the (5-10 s)
+    sweep for every distinct value; bucketing to {floor, 2*floor, ...} caps
+    the cache at a handful of entries. Only the *key* is bucketed — kernels
+    keep the exact value as their runtime loop bound, so no masked iterations
+    are added (bucketing the bound itself measured 30-50% throughput loss on
+    the PLAID kernels). The default floor matches the smallest block size in
+    the autotune pools.
+    """
+    if max_len <= 0:
+        return 0
+    return max(floor, next_pow2(max_len))
+
+
+def assert_max_seqlen_covers(cu_seqlens: torch.Tensor, max_seqlen: int, arg_name: str) -> None:
+    """Check a caller-supplied ``max_seqlen`` against the real ``cu_seqlens`` maxima.
+
+    ``max_seqlen_*`` arguments are hard kernel loop bounds, not hints: a
+    too-small value silently truncates tokens and returns wrong scores.
+    ``torch._assert_async`` runs the check on-device with no D2H sync (the
+    whole point of letting callers pass ``max_seqlen`` is to skip that sync),
+    so a violation surfaces as an async device-side assert rather than an
+    eager ``ValueError`` — same trade-off as the length checks in
+    :func:`late_interaction_kernels.padded.pack_padded`.
+    """
+    if cu_seqlens.numel() <= 1:
+        return
+    seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+    torch._assert_async(
+        seqlens.max().le(max_seqlen),
+        f"{arg_name}={max_seqlen} is smaller than the longest sequence in cu_seqlens; "
+        f"it is a hard loop bound and would silently truncate tokens.",
+    )
+
+
 def package_at_least(name: str, minimum: str) -> bool:
     """True when installed distribution ``name`` is >= ``minimum`` (absent → False).
 

late_interaction_kernels/autograd.py

@@ -29,6 +29,23 @@
 _WARNED_LQ_OVER_CEIL = False
 
 
+def _validate_devices(
+    Q: torch.Tensor,
+    D: torch.Tensor,
+    q_mask: torch.Tensor | None,
+    d_mask: torch.Tensor | None,
+) -> None:
+    """Fail fast on mixed devices — shared by the cross-product and KD paths."""
+    if Q.device != D.device:
+        raise ValueError(
+            f"Q and D must be on the same device; got Q.device={Q.device} vs D.device={D.device}."
+        )
+    if q_mask is not None and q_mask.device != Q.device:
+        raise ValueError(f"q_mask must be on the same device as Q; got {q_mask.device} vs {Q.device}.")
+    if d_mask is not None and d_mask.device != D.device:
+        raise ValueError(f"d_mask must be on the same device as D; got {d_mask.device} vs {D.device}.")
+
+
 def _bucket_lq(Q: torch.Tensor, q_mask: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor | None]:
     """Round Lq up to the next power of two so Triton's autotune cache reuses
     a config across batches with slightly different query lengths.
@@ -244,6 +261,7 @@ def _maxsim_kd(
         raise ValueError(
             f"d_mask must be [Nq={Nq}, K={K}, Ld={Ld}] for KD layout; got {tuple(d_mask.shape)}."
         )
+    _validate_devices(Q, D, q_mask, d_mask)
 
     # ``F.pad`` + ``reshape(Nq * K, ...)`` keeps strides contiguous so the view
     # below is a real zero-copy reinterpretation, not a hidden materialise.
@@ -367,8 +385,9 @@ def maxsim(
         backward: gradient strategy (``"auto" | "unified" | "lowmem"``).
             ``None`` (default) is treated as ``"auto"``, which routes the
             gradient-heavy shapes (KD / hard-negatives, high-contention
-            squares) to ``"lowmem"`` (bf16 grads, ~½ peak memory,
-            deterministic) and the rest to ``"unified"`` (fastest).
+            squares) to ``"lowmem"`` (grads written in the input dtype,
+            ~½ peak memory, deterministic) and the rest to ``"unified"``
+            (fastest).
             ``"lowmem"`` writes grads in the input dtype and reduces
             ``grad_D`` with half-precision matmuls, so fp32 inputs get
             fp16-precision ``grad_D``; pass ``backward="unified"`` if you
@@ -412,14 +431,7 @@ def maxsim(
             f"Q and D must share the embedding dim; got Q.shape[-1]={Q.shape[-1]} "
             f"vs D.shape[-1]={D.shape[-1]}."
         )
-    if Q.device != D.device:
-        raise ValueError(
-            f"Q and D must be on the same device; got Q.device={Q.device} vs D.device={D.device}."
-        )
-    if q_mask is not None and q_mask.device != Q.device:
-        raise ValueError(f"q_mask must be on the same device as Q; got {q_mask.device} vs {Q.device}.")
-    if d_mask is not None and d_mask.device != D.device:
-        raise ValueError(f"d_mask must be on the same device as D; got {d_mask.device} vs {D.device}.")
+    _validate_devices(Q, D, q_mask, d_mask)
 
     method = method_for_kd  # already resolved + validated above
 

late_interaction_kernels/backward/__init__.py

@@ -3,8 +3,9 @@
 :mod:`late_interaction_kernels.autograd` picks among them via ``backward=``.
 
 - :mod:`.unified` — fp32 atomic scatter; fastest single-pass path.
-- :mod:`.lowmem`  — destination-owned, bf16 grads (no fp32 buffer, no atomics);
-  lower peak memory and deterministic. Default where grad buffers dominate.
+- :mod:`.lowmem`  — destination-owned, grads written in the input dtype (no
+  fp32 buffer, no atomics); lower peak memory and deterministic. Default
+  where grad buffers dominate.
 """
 
 # Both modules guard Triton internally (@triton.jit kernels live behind a

late_interaction_kernels/backward/lowmem.py

@@ -1,15 +1,17 @@
-"""Low-memory backward: bf16 grads, no full-size fp32 buffers, no atomics.
+"""Low-memory backward: grads written in the input dtype, no full-size fp32
+buffers, no atomics.
 
 The other backends accumulate grads into full-size fp32 buffers and cast to
 the input dtype at the end, so the cast briefly holds both copies. For the 4-D
 hard-negative layout grad_D is n_neg-inflated, making that fp32 buffer + its
 transient the largest training allocation.
 
 This path accumulates in an fp32 register inside destination-owned kernels and
-stores the input dtype on the single write — same result as ``fp32 →
-.to(dtype)`` but with no fp32 buffer, no transient, and deterministic.
+stores the input dtype (fp16 / bf16 / fp32) on the single write — same result
+as ``fp32 → .to(dtype)`` but with no fp32 buffer, no transient, and
+deterministic.
 
-* grad_Q: the row-owned ``_bwd_dQ_kernel`` pointed at a bf16 buffer.
+* grad_Q: the row-owned ``_bwd_dQ_kernel`` pointed at an input-dtype buffer.
 * grad_D: one destination-owned kernel for both layouts — each
   ``(slab, doc-tile)`` reduces ``onehot^T @ (g·Q)`` from the saved argmax.
 """
@@ -209,9 +211,9 @@ def maxsim_backward_lowmem(
     *,
     kd_layout: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Backward producing bf16/fp16 grads directly (no fp32 buffers, no atomics).
+    """Backward producing grads directly in the input dtype (no fp32 buffers, no atomics).
 
-    Same gradients as the unified backend to bf16 rounding; deterministic.
+    Same gradients as the unified backend to input-dtype rounding; deterministic.
     The ``grad_D`` one-hot reduction runs at ``pick_compute_dtype`` matmul
     precision, so fp32 inputs get fp16-precision ``grad_D`` (``grad_Q`` is a
     gather and stays fp32-accurate); use ``unified`` for fp32 accumulation.
@@ -229,6 +231,11 @@ def maxsim_backward_lowmem(
 
     Returns ``(grad_Q, grad_D)`` already in ``Q`` / ``D`` dtypes.
     """
+    if not _HAS_TRITON:  # pragma: no cover
+        raise RuntimeError("maxsim_backward_lowmem requires Triton; install a CUDA-enabled Triton build.")
+    if not (Q.is_cuda and D.is_cuda):
+        raise RuntimeError("maxsim_backward_lowmem requires CUDA tensors.")
+
     Nq, Lq, d = Q.shape
     Nd_total, Ld, _ = D.shape
     Nd = Nd_total // Nq if kd_layout else Nd_total

late_interaction_kernels/fp8.py

@@ -250,11 +250,10 @@ def maxsim_inference_fp8(
         ``scores [Nq, Nd]`` fp32, squeezed for 2-D inputs.
 
     Notes:
-        - Hopper SM90+ or Blackwell required. On older GPUs (or old
-          Triton) the function transparently falls back to a dequantized
-          bf16 path with a one-time warning.
-        - Ties with argmax are not tie-broken — the first fp8-equal
-          doc token wins. Use the bf16 rerank path if ties matter.
+        Hopper SM90+ or Blackwell required for the FP8 tensor-core path.
+        Anywhere else (older GPUs, old Triton, no Triton at all — CPU /
+        MPS) the function transparently falls back to a dequantized bf16
+        path with a one-time warning.
     """
     global _WARNED_FP8_FALLBACK
 
@@ -313,11 +312,6 @@ def maxsim_inference_fp8(
                 stacklevel=2,
             )
             _WARNED_FP8_FALLBACK = True
-        # Dequantize and fall back to the regular kernel. ``Q.to(...)``
-        # detaches from autograd so the requires_grad-aware dispatch in
-        # ``maxsim`` lands on the cheap forward-only path.
-        from late_interaction_kernels.autograd import maxsim
-
         if scale_q_per_token:
             Qd = Q.to(torch.bfloat16) * scale_Q.unsqueeze(-1).to(torch.bfloat16)
         else:
@@ -326,7 +320,24 @@ def maxsim_inference_fp8(
             Dd = D.to(torch.bfloat16) * scale_D.unsqueeze(-1).to(torch.bfloat16)
         else:
             Dd = D.to(torch.bfloat16) * scale_D.to(torch.bfloat16)
-        scores = maxsim(Qd, Dd, q_mask=q_mask, d_mask=d_mask)
+        # Dequantize and dispatch by device, like ``retrieve``: the Triton
+        # bf16 kernel on CUDA, the compiled reference on MPS, the eager
+        # reference elsewhere — so the documented fallback really works
+        # without Triton. ``Q.to(...)`` detaches from autograd, so the
+        # CUDA path lands on ``maxsim``'s cheap forward-only branch.
+        if _HAS_TRITON and Qd.is_cuda:
+            from late_interaction_kernels.autograd import maxsim
+
+            scores = maxsim(Qd, Dd, q_mask=q_mask, d_mask=d_mask)
+        elif Qd.device.type == "mps":
+            from late_interaction_kernels.mps import maxsim_inference_mps
+
+            scores = maxsim_inference_mps(Qd, Dd, q_mask=q_mask, d_mask=d_mask, normalize=False)
+        else:
+            from late_interaction_kernels.reference import maxsim_reference
+
+            with torch.no_grad():
+                scores = maxsim_reference(Qd, Dd, q_mask=q_mask, d_mask=d_mask, normalize=False)
         if q_was_2d and d_was_2d:
             return scores.reshape(())
         if q_was_2d:

late_interaction_kernels/padded.py

@@ -136,6 +136,21 @@ def pack_padded(
         raise ValueError(f"doc_lengths must be [B={B}, C={C}]; got {tuple(doc_lengths.shape)}")
 
     device = queries.device
+
+    # Empty batch: ``qlen.max()`` / ``dlen.max()`` below would raise on a
+    # zero-element tensor, so build the trivially-empty PackedBatch up front
+    # (CUDA and CPU paths then agree on the [0, C] result).
+    if B == 0:
+        return PackedBatch(
+            Q_packed=queries.reshape(0, d),
+            cu_seqlens_q=torch.zeros(1, dtype=torch.int32, device=device),
+            D_packed=documents.reshape(0, d),
+            cu_seqlens_d=torch.zeros(1, dtype=torch.int32, device=device),
+            pair_q_idx=torch.zeros(0, dtype=torch.int32, device=device),
+            pair_d_idx=torch.zeros(0, dtype=torch.int32, device=device),
+            max_seqlen_q=0,
+            max_seqlen_d=0,
+        )
     qlen = query_lengths.to(device=device, dtype=torch.int32)
     dlen = doc_lengths.to(device=device, dtype=torch.int32)