From f9950b6db735100333b6d8637635d9bac44a800b Mon Sep 17 00:00:00 2001
From: Tony Wu <28306721+tonywu71@users.noreply.github.com>
Date: Tue, 9 Jun 2026 21:35:53 +0200
Subject: [PATCH 1/4] fix: v0.4.3 correctness and robustness fixes

- chunked top-k: clamp k to the merge width (retrieve + maxsim_topk)
- MaxSimScorer.forward: nullcontext instead of forced enable_grad
- maxsim_residual_varlen: early-return on Nd == 0 (grid used max(Nd, 1))
- pack_padded: handle B == 0 up front
- pylate_compat: forward legacy mask= on the fallback path
- maxsim_varlen: bucket max_lq/max_ld constexpr autotune keys; same for
  the plaid residual kernels' max_Ld
- max_seqlen_*: documented as hard loop bounds + on-device assert
- maxsim_residual: squeeze 2-D Q back to [Nd]
- fp8: drop stale argmax-ties note; fallback works without Triton
- lowmem: Triton/CUDA guard + input-dtype (not bf16) wording
- plaid/varlen kernels: drop dead B / Lq / Ld params
- prune_forward: smallest-footprint fallback instead of configs[:2]
- KD path: same device validation as the 3-D path
---
 CHANGELOG.md                                  | 74 ++++++++++++++++-
 late_interaction_kernels/_autotune.py         | 32 ++++---
 late_interaction_kernels/_utils.py            | 37 +++++++++
 late_interaction_kernels/autograd.py          | 32 ++++---
 late_interaction_kernels/backward/__init__.py |  5 +-
 late_interaction_kernels/backward/lowmem.py   | 19 +++--
 late_interaction_kernels/fp8.py               | 33 +++++---
 late_interaction_kernels/padded.py            | 15 ++++
 late_interaction_kernels/plaid.py             | 40 ++++++---
 late_interaction_kernels/pylate_compat.py     | 13 +--
 late_interaction_kernels/retrieve.py          | 18 +++-
 late_interaction_kernels/score_pairs.py       | 12 ++-
 late_interaction_kernels/topk.py              |  7 +-
 late_interaction_kernels/varlen.py            | 31 +++++--
 tests/test_autotune_kwargs.py                 | 16 ++++
 tests/test_backward_lowmem.py                 | 18 +++-
 tests/test_fp8.py                             | 24 ++++++
 tests/test_padded.py                          | 42 ++++++++++
 tests/test_plaid.py                           | 34 ++++++++
 tests/test_pylate_compat_fallback.py          | 83 +++++++++++++++++++
 tests/test_retrieve_cpu.py                    | 43 ++++++++++
 tests/test_topk.py                            | 29 +++++++
 tests/test_utils.py                           | 26 +++++-
 tests/test_varlen.py                          | 51 ++++++++++++
 24 files changed, 659 insertions(+), 75 deletions(-)
 create mode 100644 tests/test_pylate_compat_fallback.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 338bd2a..1a1e561 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,79 @@ 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 (unreleased)
+
+### 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` now buckets `max_lq` / `max_ld` to the next power of
+  two** (floor 16) before they reach the kernel. Both 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. The kernel masks on the actual `cu_seqlens` bounds, so results are
+  unchanged; the argmax buffer is sized on the bucketed `max_lq` (up to 2×
+  larger, rows `-1`-padded). The same bucketing now applies to the `max_Ld`
+  loop bound of `maxsim_residual` / `maxsim_residual_varlen`.
+- **`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) and the unused `Lq` / `Ld` placeholder
+  params from the varlen forward kernel.
+- 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
 
diff --git a/late_interaction_kernels/_autotune.py b/late_interaction_kernels/_autotune.py
index 784371a..4a25432 100644
--- a/late_interaction_kernels/_autotune.py
+++ b/late_interaction_kernels/_autotune.py
@@ -122,9 +122,21 @@ 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`.
+    Lq = named_args.get("Lq") or named_args.get("max_lq") or 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.
@@ -135,15 +147,15 @@ def prune_forward(configs, named_args, **kwargs):
 
     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]
diff --git a/late_interaction_kernels/_utils.py b/late_interaction_kernels/_utils.py
index cc015f9..4e85d1a 100644
--- a/late_interaction_kernels/_utils.py
+++ b/late_interaction_kernels/_utils.py
@@ -12,6 +12,43 @@ def next_pow2(x: int) -> int:
     return 1 << (x - 1).bit_length()
 
 
+def bucket_seqlen(max_len: int, floor: int = 16) -> int:
+    """Round a kernel's sequence-loop bound up to the next power of two.
+
+    Several kernels take a max-length loop bound that is both a ``constexpr``
+    and an autotune key, so every distinct value re-triggers the (5-10 s)
+    autotune sweep. Bucketing to {floor, 2*floor, ...} caps the cache at a
+    handful of entries. Safe because every such kernel masks its loads and
+    reductions on the *actual* per-sequence length — a larger loop bound only
+    adds fully-masked iterations. 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).
 
diff --git a/late_interaction_kernels/autograd.py b/late_interaction_kernels/autograd.py
index 7207d9d..034d2fc 100644
--- a/late_interaction_kernels/autograd.py
+++ b/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
 
diff --git a/late_interaction_kernels/backward/__init__.py b/late_interaction_kernels/backward/__init__.py
index 77a8696..d2c41dd 100644
--- a/late_interaction_kernels/backward/__init__.py
+++ b/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
diff --git a/late_interaction_kernels/backward/lowmem.py b/late_interaction_kernels/backward/lowmem.py
index 1efda2c..0c16427 100644
--- a/late_interaction_kernels/backward/lowmem.py
+++ b/late_interaction_kernels/backward/lowmem.py
@@ -1,4 +1,5 @@
-"""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
@@ -6,10 +7,11 @@
 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
diff --git a/late_interaction_kernels/fp8.py b/late_interaction_kernels/fp8.py
index c92cb9d..f1d25cb 100644
--- a/late_interaction_kernels/fp8.py
+++ b/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:
diff --git a/late_interaction_kernels/padded.py b/late_interaction_kernels/padded.py
index 7860751..87c9638 100644
--- a/late_interaction_kernels/padded.py
+++ b/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)
 
diff --git a/late_interaction_kernels/plaid.py b/late_interaction_kernels/plaid.py
index e6f12bc..a5ef780 100644
--- a/late_interaction_kernels/plaid.py
+++ b/late_interaction_kernels/plaid.py
@@ -17,7 +17,7 @@
 import triton.language as tl
 
 from late_interaction_kernels._autotune import autotune_kwargs, forward_configs, prune_forward
-from late_interaction_kernels._utils import next_pow2
+from late_interaction_kernels._utils import assert_max_seqlen_covers, bucket_seqlen, next_pow2
 
 # -----------------------------------------------------------------------------
 # C1. plaid_approx_score
@@ -30,7 +30,6 @@ def _plaid_approx_score_kernel(
     codes_ptr,  # [B, max_Ld] int64 centroid codes (padded)
     doc_len_ptr,  # [B] int64 real lengths
     out_ptr,  # [B] fp32 scores
-    B: tl.constexpr,
     n_centroids: tl.constexpr,
     Lq: tl.constexpr,
     max_Ld: tl.constexpr,
@@ -119,7 +118,6 @@ def plaid_approx_score(
         codes,
         doc_lengths,
         out,
-        B,
         n_centroids,
         Lq,
         max_Ld,
@@ -449,6 +447,9 @@ def _maxsim_residual_forward(
         raise ValueError(
             f"residuals last dim {packed_dim} != expected {expected_pd} for d={d}, nbits={nbits}"
         )
+    # max_Ld is constexpr + autotune key; bucket it so distinct padded widths
+    # share the cache entry (the kernel masks on the real per-doc length).
+    max_Ld = bucket_seqlen(max_Ld)
 
     d_pad = next_pow2(d)
     codes_per_byte = 8 // nbits
@@ -639,22 +640,26 @@ def maxsim_residual(
             the kernel (standard PLAID convention).
 
     Returns:
-        scores: ``[Nq, Nd]`` fp32.
+        scores: ``[Nq, Nd]`` fp32. ``[Nd]`` if Q was 2-D.
 
     Autograd flows into ``Q`` (only). ``codes`` / ``residuals`` are integer;
     ``centroids`` / ``bucket_weights`` are treated as frozen.
     The argmax save and fused backward kernel run only when
     ``Q.requires_grad`` is True.
     """
-    if Q.dim() == 2:
+    q_squeeze = Q.dim() == 2
+    if q_squeeze:
         Q = Q.unsqueeze(0)
     if Q.requires_grad:
-        return _MaxSimResidualFn.apply(
+        scores = _MaxSimResidualFn.apply(
             Q, codes, residuals, doc_lengths, centroids, bucket_weights, nbits, normalize
         )
-    scores, _, _ = _maxsim_residual_forward(
-        Q, codes, residuals, doc_lengths, centroids, bucket_weights, nbits, normalize, False
-    )
+    else:
+        scores, _, _ = _maxsim_residual_forward(
+            Q, codes, residuals, doc_lengths, centroids, bucket_weights, nbits, normalize, False
+        )
+    if q_squeeze:
+        return scores.squeeze(0)
     return scores
 
 
@@ -843,7 +848,9 @@ def maxsim_residual_varlen(
         centroids: ``[n_centroids, d]`` fp16/bf16/fp32.
         bucket_weights: ``[2**nbits]`` fp32.
         nbits: one of ``{2, 4, 8}``.
-        max_seqlen_d: optional; inferred from ``cu_seqlens_d`` otherwise.
+        max_seqlen_d: hard doc-loop bound (not a hint); must be >= the
+            longest doc or tokens are dropped — checked by an on-device
+            assert. Inferred from ``cu_seqlens_d`` (one D2H sync) if omitted.
         normalize: L2-normalize Q and the reconstructed embedding inside
             the kernel (standard PLAID convention).
 
@@ -887,7 +894,11 @@ def maxsim_residual_varlen(
         starts = cu_seqlens_d[:-1]
         ends = cu_seqlens_d[1:]
         max_seqlen_d = int((ends - starts).max().item()) if Nd > 0 else 0
-    max_seqlen_d = max(int(max_seqlen_d), 1)
+    else:
+        assert_max_seqlen_covers(cu_seqlens_d, int(max_seqlen_d), "max_seqlen_d")
+    # max_seqlen_d is constexpr + autotune key; bucket it so distinct batch
+    # maxima share the cache entry (the kernel masks on the real doc length).
+    max_seqlen_d = bucket_seqlen(max(int(max_seqlen_d), 1))
 
     d_pad = next_pow2(d)
     codes_per_byte = 8 // nbits
@@ -895,9 +906,14 @@ def maxsim_residual_varlen(
     tl_dtype = tl.float16 if compute_dtype == torch.float16 else tl.bfloat16
 
     out = torch.empty(Nq, Nd, device=Q.device, dtype=torch.float32)
+    # No docs -> nothing to launch. The kernel's `pid // Nd` indexing would
+    # divide by the constexpr Nd=0; the empty [Nq, 0] result needs no kernel.
+    if Nd == 0:
+        return out.squeeze(0) if q_squeeze else out
+
     argmax = torch.empty(1, device=Q.device, dtype=torch.int32)  # unused placeholder
 
-    grid = (Nq * max(Nd, 1),)
+    grid = (Nq * Nd,)
     _maxsim_residual_varlen_kernel[grid](
         Q,
         codes_flat,
diff --git a/late_interaction_kernels/pylate_compat.py b/late_interaction_kernels/pylate_compat.py
index 3da4931..d1aea8d 100644
--- a/late_interaction_kernels/pylate_compat.py
+++ b/late_interaction_kernels/pylate_compat.py
@@ -107,11 +107,14 @@ def _resolve_masks(queries_mask, documents_mask, legacy_mask):
     * legacy (< 1.1):    ``mask=`` (document mask only)
     * positional legacy: 3rd positional == document mask
 
-    Whichever arrived, hand back ``(q_mask, d_mask)`` as bool tensors.
+    Whichever arrived, hand back ``(q_mask, d_mask, documents_mask)`` — the
+    first two as bool tensors for the fused path, the last as the resolved
+    raw document mask so the fallback path forwards a legacy ``mask=``
+    instead of dropping it.
     """
     if documents_mask is None and legacy_mask is not None:
         documents_mask = legacy_mask
-    return _mask_as_bool(queries_mask), _mask_as_bool(documents_mask)
+    return _mask_as_bool(queries_mask), _mask_as_bool(documents_mask), documents_mask
 
 
 def patched_colbert_scores(
@@ -127,7 +130,7 @@ def patched_colbert_scores(
 
     Q = convert_to_tensor(queries_embeddings)
     D = convert_to_tensor(documents_embeddings)
-    q_mask, d_mask = _resolve_masks(queries_mask, documents_mask, mask)
+    q_mask, d_mask, documents_mask = _resolve_masks(queries_mask, documents_mask, mask)
 
     path = _device_path(Q, D)
     if path is None:
@@ -153,7 +156,7 @@ def patched_colbert_kd_scores(
 
     Q = convert_to_tensor(queries_embeddings)
     D = convert_to_tensor(documents_embeddings)
-    q_mask, d_mask = _resolve_masks(queries_mask, documents_mask, mask)
+    q_mask, d_mask, documents_mask = _resolve_masks(queries_mask, documents_mask, mask)
 
     path = _device_path(Q, D)
     if path is None:
@@ -188,7 +191,7 @@ def patched_colbert_scores_pairwise(
 
     Q = convert_to_tensor(queries_embeddings)
     D = convert_to_tensor(documents_embeddings)
-    q_mask, d_mask = _resolve_masks(queries_mask, documents_mask, mask)
+    q_mask, d_mask, documents_mask = _resolve_masks(queries_mask, documents_mask, mask)
 
     path = _device_path(Q, D)
     # CUDA gets the fused K=1 KD launch; everything else defers to PyLate's
diff --git a/late_interaction_kernels/retrieve.py b/late_interaction_kernels/retrieve.py
index b857e64..c32f6a0 100644
--- a/late_interaction_kernels/retrieve.py
+++ b/late_interaction_kernels/retrieve.py
@@ -11,6 +11,7 @@
 unit-testable on macOS / Windows before renting a CUDA box.
 """
 
+import contextlib
 from typing import Literal
 
 import torch
@@ -56,8 +57,10 @@ def _score(
         # `inference=True` is a stronger contract than auto-dispatch:
         # `score()` must not participate in autograd even when inputs have
         # ``requires_grad=True``, so force ``no_grad()``. Training
-        # (`inference=False`) routes through autograd as usual.
-        with torch.no_grad() if inference else torch.enable_grad():
+        # (`inference=False`) must NOT force ``enable_grad()`` — that would
+        # override a caller's ``torch.no_grad()`` — so it leaves the caller's
+        # grad mode untouched.
+        with torch.no_grad() if inference else contextlib.nullcontext():
             return maxsim(Q, D, q_mask=q_mask, d_mask=d_mask, normalize=normalize, backward=backward)
 
     if Q.device.type == "mps":
@@ -69,7 +72,7 @@ def _score(
 
     from late_interaction_kernels.reference import maxsim_reference
 
-    with torch.no_grad() if inference else torch.enable_grad():
+    with torch.no_grad() if inference else contextlib.nullcontext():
         return maxsim_reference(Q, D, q_mask=q_mask, d_mask=d_mask, normalize=normalize)
 
 
@@ -291,7 +294,14 @@ def _score_chunk(Q_, D_, qm_, dm_):
             else:
                 merged_scores = torch.cat([topk_s, chunk_scores], dim=-1)
                 merged_indices = torch.cat([topk_i, chunk_indices], dim=-1)
-                topk_s, gather_idx = torch.topk(merged_scores, k, dim=-1, largest=largest, sorted=sorted)
+                # Early chunks can contribute fewer than k columns total
+                # (e.g. top_k=10, chunk=4), so clamp k at the merge; the
+                # final width still reaches min(top_k, Nd) once all chunks
+                # are merged.
+                k_merge = min(k, merged_scores.shape[-1])
+                topk_s, gather_idx = torch.topk(
+                    merged_scores, k_merge, dim=-1, largest=largest, sorted=sorted
+                )
                 topk_i = torch.gather(merged_indices, -1, gather_idx)
 
     if q_was_2d:
diff --git a/late_interaction_kernels/score_pairs.py b/late_interaction_kernels/score_pairs.py
index 8e5a7af..6b6f2bd 100644
--- a/late_interaction_kernels/score_pairs.py
+++ b/late_interaction_kernels/score_pairs.py
@@ -19,7 +19,7 @@
 import triton.language as tl
 
 from late_interaction_kernels._autotune import autotune_kwargs, forward_configs, prune_forward
-from late_interaction_kernels._utils import next_pow2, pick_compute_dtype
+from late_interaction_kernels._utils import assert_max_seqlen_covers, next_pow2, pick_compute_dtype
 
 
 @triton.autotune(
@@ -287,8 +287,12 @@ def _scatter_forward(
     Nd = cu_seqlens_d.numel() - 1
     if max_seqlen_q is None:
         max_seqlen_q = int((cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item()) if Nq else 0
+    else:
+        assert_max_seqlen_covers(cu_seqlens_q, int(max_seqlen_q), "max_seqlen_q")
     if max_seqlen_d is None:
         max_seqlen_d = int((cu_seqlens_d[1:] - cu_seqlens_d[:-1]).max().item()) if Nd else 0
+    else:
+        assert_max_seqlen_covers(cu_seqlens_d, int(max_seqlen_d), "max_seqlen_d")
 
     d = Q_packed.shape[1]
     d_pad = next_pow2(d)
@@ -433,8 +437,10 @@ def score_pairs_packed(
         cu_seqlens_d: ``[Nd + 1]`` int32 cumulative offsets into ``D_packed``.
         pair_q_idx: ``[num_pairs]`` int32 query indices.
         pair_d_idx: ``[num_pairs]`` int32 doc indices.
-        max_seqlen_q / max_seqlen_d: hints; computed from ``cu_seqlens`` if
-            omitted (one D2H sync per call).
+        max_seqlen_q / max_seqlen_d: hard kernel loop bounds, NOT hints — a
+            value smaller than the longest sequence would silently drop
+            tokens, so it is rejected by an on-device assert (no D2H sync).
+            Computed from ``cu_seqlens`` (one D2H sync) if omitted.
 
     Returns:
         scores: ``[num_pairs]`` fp32. ``scores[k]`` is the MaxSim of
diff --git a/late_interaction_kernels/topk.py b/late_interaction_kernels/topk.py
index 64ecf45..0de5938 100644
--- a/late_interaction_kernels/topk.py
+++ b/late_interaction_kernels/topk.py
@@ -71,9 +71,14 @@ def maxsim_topk(
             else:
                 merged_scores = torch.cat([topk_scores, chunk_scores], dim=-1)
                 merged_indices = torch.cat([topk_idx, chunk_indices], dim=-1)
+                # Early chunks can contribute fewer than k columns total
+                # (e.g. k=10, chunk_size=4), so clamp k at the merge; the
+                # final width still reaches min(k, Nd) once all chunks are
+                # merged.
+                k_merge = min(k, merged_scores.shape[-1])
                 topk_scores, gather_idx = torch.topk(
                     merged_scores,
-                    k,
+                    k_merge,
                     dim=-1,
                     largest=largest,
                     sorted=sorted,
diff --git a/late_interaction_kernels/varlen.py b/late_interaction_kernels/varlen.py
index d3f7a61..cb3518b 100644
--- a/late_interaction_kernels/varlen.py
+++ b/late_interaction_kernels/varlen.py
@@ -12,7 +12,12 @@
 import triton.language as tl
 
 from late_interaction_kernels._autotune import autotune_kwargs, forward_configs, prune_forward
-from late_interaction_kernels._utils import next_pow2, pick_compute_dtype
+from late_interaction_kernels._utils import (
+    assert_max_seqlen_covers,
+    bucket_seqlen,
+    next_pow2,
+    pick_compute_dtype,
+)
 
 
 @triton.autotune(
@@ -44,8 +49,6 @@ def _varlen_fwd_kernel(
     stride_am_n,
     stride_am_d,
     stride_am_l,
-    Lq: tl.constexpr,
-    Ld,  # unused; runtime arg kept for call-site compat (constexpr would recompile per Ld).
     BLOCK_Q: tl.constexpr,
     BLOCK_D: tl.constexpr,
     COMPUTE_DTYPE: tl.constexpr,
@@ -292,8 +295,20 @@ def _varlen_forward(
 
     if max_seqlen_q is None:
         max_seqlen_q = int((cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item()) if Nq else 0
+    else:
+        assert_max_seqlen_covers(cu_seqlens_q, int(max_seqlen_q), "max_seqlen_q")
     if max_seqlen_d is None:
         max_seqlen_d = int((cu_seqlens_d[1:] - cu_seqlens_d[:-1]).max().item()) if Nd else 0
+    else:
+        assert_max_seqlen_covers(cu_seqlens_d, int(max_seqlen_d), "max_seqlen_d")
+
+    # max_lq / max_ld are constexpr loop bounds AND autotune keys, so without
+    # bucketing every distinct (max_lq, max_ld) pair re-triggers the full
+    # autotune sweep. The kernel masks on the actual cu_seqlens bounds, so a
+    # larger loop bound only adds fully-masked iterations. The argmax buffer
+    # is sized on the bucketed value (its rows are -1-padded; backward skips).
+    max_seqlen_q = bucket_seqlen(int(max_seqlen_q))
+    max_seqlen_d = bucket_seqlen(int(max_seqlen_d))
 
     d_pad = next_pow2(d)
     compute_dtype = pick_compute_dtype(Q_packed, D_packed)
@@ -333,8 +348,6 @@ def _varlen_forward(
         am_strides[0],
         am_strides[1],
         am_strides[2],
-        max_seqlen_q,
-        max_seqlen_d,  # Lq, Ld placeholders
         COMPUTE_DTYPE=tl_dtype,
         SAVE_ARGMAX=save_argmax,
     )
@@ -441,8 +454,12 @@ def maxsim_varlen(
         D_packed: ``[sum(Ld_j), d]``.
         cu_seqlens_q: ``[Nq + 1]`` int32 cumulative offsets.
         cu_seqlens_d: ``[Nd + 1]`` int32 cumulative offsets.
-        max_seqlen_q, max_seqlen_d: tile-count hints; inferred from
-            ``cu_seqlens`` if omitted.
+        max_seqlen_q, max_seqlen_d: hard kernel loop bounds, NOT hints — a
+            value smaller than the longest sequence would silently drop
+            tokens, so it is rejected by an on-device assert (no D2H sync).
+            Inferred from ``cu_seqlens`` (one D2H sync) if omitted. Both are
+            rounded up to the next power of two internally so the autotune
+            cache is reused across batches with different maxima.
 
     Returns:
         scores: ``[Nq, Nd]`` fp32.
diff --git a/tests/test_autotune_kwargs.py b/tests/test_autotune_kwargs.py
index 39aaf23..1fdd167 100644
--- a/tests/test_autotune_kwargs.py
+++ b/tests/test_autotune_kwargs.py
@@ -29,6 +29,22 @@ def test_autotune_kwargs_matches_triton_signature():
     )
 
 
+def test_prune_forward_fallback_returns_smallest_footprint():
+    """When every config overflows the SMEM budget, the fallback must hand
+    back the smallest-footprint candidates, not the first two list entries."""
+    from late_interaction_kernels._autotune import _smem_bytes, forward_configs, prune_forward
+
+    configs = forward_configs()
+    # d_pad large enough that no config fits any family's budget.
+    named_args = {"Lq": 32, "d_pad": 1 << 20}
+    kept = prune_forward(configs, named_args)
+
+    assert len(kept) == 2
+    d_pad: int = named_args["d_pad"]
+    smallest = sorted(_smem_bytes(cfg, d_pad) for cfg in configs)[:2]
+    assert sorted(_smem_bytes(cfg, d_pad) for cfg in kept) == smallest
+
+
 def test_autotune_kwargs_safe_to_unpack_into_decorator():
     """Whatever ``autotune_kwargs()`` returns must compose with our other
     autotune args — i.e. the call site
diff --git a/tests/test_backward_lowmem.py b/tests/test_backward_lowmem.py
index 4cc7165..500bce1 100644
--- a/tests/test_backward_lowmem.py
+++ b/tests/test_backward_lowmem.py
@@ -1,8 +1,8 @@
-"""Parity + memory for the low-memory (bf16, atomic-free) backward.
+"""Parity + memory for the low-memory (input-dtype, atomic-free) backward.
 
 The lowmem path produces grads directly in the input dtype via fp32-register
 accumulation in destination-owned kernels (one-hot matmul for both layouts).
-It must match the unified (saved-argmax, fp32-then-cast) backend to bf16
+It must match the unified (saved-argmax, fp32-then-cast) backend to input-dtype
 rounding, and use strictly less training memory on the n_neg-inflated KD case.
 """
 
@@ -10,6 +10,20 @@
 import torch
 
 
+def test_lowmem_rejects_non_cuda_tensors():
+    """Same clean RuntimeError contract as ``maxsim_backward_unified`` —
+    CPU tensors must not reach the kernel launch. Runs everywhere: without
+    Triton the guard raises on the missing dependency instead."""
+    from late_interaction_kernels.backward import maxsim_backward_lowmem
+
+    Q = torch.randn(1, 4, 8)
+    D = torch.randn(2, 4, 8)
+    grad_s = torch.randn(1, 2)
+    argmax = torch.zeros(2, 4, dtype=torch.int32)
+    with pytest.raises(RuntimeError, match="maxsim_backward_lowmem requires"):
+        maxsim_backward_lowmem(grad_s, Q, D, argmax, None)
+
+
 @pytest.mark.cuda
 @pytest.mark.parametrize(
     "shape",
diff --git a/tests/test_fp8.py b/tests/test_fp8.py
index 7a0a3a1..56064da 100644
--- a/tests/test_fp8.py
+++ b/tests/test_fp8.py
@@ -28,6 +28,30 @@ def _make(Nq, Nd, Lq, Ld, d, dtype=torch.bfloat16, device="cuda"):
 # --------------------------------------------------------------------------- #
 
 
+@pytest.mark.skipif(fp8_dtype is None, reason="torch has no FP8 dtype")
+def test_fp8_fallback_works_on_cpu():
+    """The documented "transparent fallback" must run without Triton/CUDA:
+    on CPU it dequantizes to bf16 and routes through the eager reference."""
+    import warnings
+
+    from late_interaction_kernels.fp8 import maxsim_inference_fp8, quantize_fp8_per_tensor
+    from late_interaction_kernels.reference import maxsim_reference
+
+    torch.manual_seed(0)
+    Q, D = _make(2, 4, 8, 16, 32, dtype=torch.float32, device="cpu")
+    Qq, sq = quantize_fp8_per_tensor(Q)
+    Dq, sd = quantize_fp8_per_tensor(D)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", RuntimeWarning)
+        scores = maxsim_inference_fp8(Qq, Dq, scale_Q=sq, scale_D=sd)
+
+    ref = maxsim_reference(Q, D)
+    assert scores.shape == (2, 4)
+    err = (scores - ref).abs().max().item() / max(1e-6, ref.abs().max().item())
+    assert err < 0.05  # fp8 quantization + bf16 dequant noise
+
+
 @pytest.mark.skipif(fp8_dtype is None, reason="torch has no FP8 dtype")
 def test_quantize_roundtrip_per_tensor():
     from late_interaction_kernels.fp8 import (
diff --git a/tests/test_padded.py b/tests/test_padded.py
index 59ed80d..8abae75 100644
--- a/tests/test_padded.py
+++ b/tests/test_padded.py
@@ -148,6 +148,48 @@ def test_pack_padded_validate_flag_catches_zero_length() -> None:
         pack_padded(queries, documents, qlen_bad, dlen, validate=True)
 
 
+def test_pack_padded_empty_batch() -> None:
+    """B == 0 must produce an empty PackedBatch instead of crashing on
+    ``qlen.max()`` of a zero-element tensor."""
+    B, C, Lq, Ld, d = 0, 3, 8, 12, 16
+    queries = torch.randn(B, Lq, d)
+    documents = torch.randn(B, C, Ld, d)
+    qlen = torch.zeros(B, dtype=torch.int32)
+    dlen = torch.zeros(B, C, dtype=torch.int32)
+
+    batch = pack_padded(queries, documents, qlen, dlen)
+    assert batch.Q_packed.shape == (0, d)
+    assert batch.D_packed.shape == (0, d)
+    assert batch.cu_seqlens_q.tolist() == [0]
+    assert batch.cu_seqlens_d.tolist() == [0]
+    assert batch.pair_q_idx.numel() == 0
+    assert batch.pair_d_idx.numel() == 0
+    assert batch.max_seqlen_q == 0
+    assert batch.max_seqlen_d == 0
+
+
+def test_maxsim_padded_empty_batch_cpu() -> None:
+    B, C = 0, 4
+    queries = torch.randn(B, 8, 16)
+    documents = torch.randn(B, C, 12, 16)
+    qlen = torch.zeros(B, dtype=torch.int32)
+    dlen = torch.zeros(B, C, dtype=torch.int32)
+    scores = maxsim_padded(queries, documents, qlen, dlen)
+    assert scores.shape == (B, C)
+
+
+@pytest.mark.cuda
+def test_maxsim_padded_empty_batch_cuda() -> None:
+    """The CUDA path must agree with CPU on the empty batch (no crash)."""
+    B, C = 0, 4
+    queries = torch.randn(B, 8, 16, device="cuda")
+    documents = torch.randn(B, C, 12, 16, device="cuda")
+    qlen = torch.zeros(B, dtype=torch.int32, device="cuda")
+    dlen = torch.zeros(B, C, dtype=torch.int32, device="cuda")
+    scores = maxsim_padded(queries, documents, qlen, dlen)
+    assert scores.shape == (B, C)
+
+
 # ---------------------------------------------------------------------------
 # Always-on async assert (subprocess-isolated)
 #
diff --git a/tests/test_plaid.py b/tests/test_plaid.py
index 0b265f5..65269fe 100644
--- a/tests/test_plaid.py
+++ b/tests/test_plaid.py
@@ -367,6 +367,40 @@ def test_residual_varlen_handles_empty_docs():
     assert scores[0, 2].item() == 0.0
 
 
+def test_residual_2d_query_squeezes():
+    """A 2-D Q must return [Nd], matching ``maxsim_residual_varlen`` and the
+    ``maxsim`` wrapper's squeeze convention."""
+    from late_interaction_kernels.plaid import maxsim_residual
+
+    idx = _make_quant_index(Nd=3, max_Ld=16, d=128, n_centroids=32, nbits=4)
+    Q3 = torch.randn(1, 32, 128, device="cuda", dtype=torch.bfloat16)
+
+    args = (idx["codes"], idx["residuals"], idx["doc_lengths"], idx["centroids"], idx["bucket_weights"])
+    out3 = maxsim_residual(Q3, *args, nbits=4, normalize=True)
+    out2 = maxsim_residual(Q3[0], *args, nbits=4, normalize=True)
+    assert out2.shape == (3,)
+    torch.testing.assert_close(out2, out3[0])
+
+
+def test_residual_varlen_empty_corpus():
+    """Nd == 0 must return an empty [Nq, 0] result instead of launching the
+    kernel with a constexpr Nd=0 divisor."""
+    from late_interaction_kernels.plaid import maxsim_residual_varlen
+
+    centroids = torch.randn(32, 128, device="cuda", dtype=torch.float32)
+    buckets = torch.linspace(-0.1, 0.1, 16, device="cuda", dtype=torch.float32)
+    codes = torch.zeros(0, device="cuda", dtype=torch.int64)
+    res = torch.zeros(0, 64, device="cuda", dtype=torch.uint8)  # nbits=4, d=128
+    cu = torch.zeros(1, device="cuda", dtype=torch.int32)  # Nd = 0
+    Q = torch.randn(2, 32, 128, device="cuda", dtype=torch.bfloat16)
+
+    scores = maxsim_residual_varlen(Q, codes, res, cu, centroids, buckets, nbits=4, normalize=True)
+    assert scores.shape == (2, 0)
+
+    scores_2d = maxsim_residual_varlen(Q[0], codes, res, cu, centroids, buckets, nbits=4, normalize=True)
+    assert scores_2d.shape == (0,)
+
+
 def test_residual_matches_dense_maxsim():
     """With nbits=8 and identity bucket weights, residual scoring should
     recover exact dense MaxSim up to rounding."""
diff --git a/tests/test_pylate_compat_fallback.py b/tests/test_pylate_compat_fallback.py
new file mode 100644
index 0000000..41fef6d
--- /dev/null
+++ b/tests/test_pylate_compat_fallback.py
@@ -0,0 +1,83 @@
+"""CPU-safe regression tests for the patched scorers' fallback path.
+
+On CPU ``_device_path`` returns ``None``, so the patched wrappers defer to
+PyLate's original implementation. A legacy ``mask=`` (single document mask)
+must be forwarded as ``documents_mask`` on that path — it used to be
+silently dropped. We inject a minimal fake ``pylate`` (just
+``convert_to_tensor``) and stub ``_ORIGINAL`` directly; no GPU, no real
+PyLate required.
+"""
+
+import sys
+import types
+
+import pytest
+import torch
+
+
+@pytest.fixture
+def fake_pylate_tensor(monkeypatch):
+    """Provide ``pylate.utils.tensor.convert_to_tensor`` for the wrappers."""
+    tensor_mod = types.ModuleType("pylate.utils.tensor")
+    tensor_mod.convert_to_tensor = lambda x: x
+    utils_mod = types.ModuleType("pylate.utils")
+    utils_mod.tensor = tensor_mod
+    pylate_mod = types.ModuleType("pylate")
+    pylate_mod.utils = utils_mod
+    monkeypatch.setitem(sys.modules, "pylate", pylate_mod)
+    monkeypatch.setitem(sys.modules, "pylate.utils", utils_mod)
+    monkeypatch.setitem(sys.modules, "pylate.utils.tensor", tensor_mod)
+
+
+def _vanilla_colbert_scores(q, d, queries_mask=None, documents_mask=None):
+    """Minimal PyLate-style scorer: doc mask multiplied in before the max."""
+    sim = torch.einsum("ald,bsd->abls", q, d)
+    if documents_mask is not None:
+        sim = sim * documents_mask[None, :, None, :].to(sim.dtype)
+    scores = sim.max(dim=-1).values
+    if queries_mask is not None:
+        scores = scores * queries_mask[:, None, :].to(scores.dtype)
+    return scores.sum(dim=-1)
+
+
+def test_fallback_forwards_legacy_mask(fake_pylate_tensor, monkeypatch):
+    from late_interaction_kernels import pylate_compat
+
+    received: dict[str, torch.Tensor | None] = {}
+
+    def recording_original(q, d, queries_mask=None, documents_mask=None):
+        received["documents_mask"] = documents_mask
+        return _vanilla_colbert_scores(q, d, queries_mask=queries_mask, documents_mask=documents_mask)
+
+    monkeypatch.setattr(pylate_compat, "_ORIGINAL", {"colbert_scores": recording_original})
+
+    torch.manual_seed(0)
+    Q = torch.nn.functional.normalize(torch.randn(2, 6, 16), dim=-1)
+    D = torch.nn.functional.normalize(torch.randn(3, 8, 16), dim=-1)
+    d_mask = torch.ones(3, 8)
+    d_mask[:, -3:] = 0
+
+    out = pylate_compat.patched_colbert_scores(Q, D, mask=d_mask)
+
+    assert received["documents_mask"] is d_mask
+    expected = _vanilla_colbert_scores(Q, D, documents_mask=d_mask)
+    torch.testing.assert_close(out, expected)
+
+
+def test_kd_fallback_forwards_legacy_mask(fake_pylate_tensor, monkeypatch):
+    from late_interaction_kernels import pylate_compat
+
+    received: dict[str, torch.Tensor | None] = {}
+
+    def recording_original(q, d, queries_mask=None, documents_mask=None):
+        received["documents_mask"] = documents_mask
+        return torch.zeros(q.shape[0], d.shape[1])
+
+    monkeypatch.setattr(pylate_compat, "_ORIGINAL", {"colbert_kd_scores": recording_original})
+
+    Q = torch.randn(2, 6, 16)
+    D = torch.randn(2, 3, 8, 16)
+    d_mask = torch.ones(2, 3, 8)
+
+    pylate_compat.patched_colbert_kd_scores(Q, D, mask=d_mask)
+    assert received["documents_mask"] is d_mask
diff --git a/tests/test_retrieve_cpu.py b/tests/test_retrieve_cpu.py
index 051deb4..02d1396 100644
--- a/tests/test_retrieve_cpu.py
+++ b/tests/test_retrieve_cpu.py
@@ -72,6 +72,19 @@ def test_scorer_score_is_no_grad():
     assert not out.requires_grad
 
 
+def test_scorer_forward_respects_callers_no_grad():
+    """Regression: ``forward()`` used to force ``enable_grad()``, overriding a
+    caller's ``torch.no_grad()`` and returning ``requires_grad=True``."""
+    from late_interaction_kernels import MaxSimScorer
+
+    Q = torch.randn(2, 8, 16, requires_grad=True)
+    D = torch.randn(3, 12, 16, requires_grad=True)
+    scorer = MaxSimScorer(normalize=True)
+    with torch.no_grad():
+        out = scorer(Q, D)
+    assert not out.requires_grad
+
+
 def test_scorer_composes_inside_nn_module():
     from late_interaction_kernels import MaxSimScorer
 
@@ -187,6 +200,36 @@ def test_retrieve_chunked_equals_unchunked():
     assert torch.equal(full_i, chunked_i)
 
 
+def test_retrieve_chunked_top_k_larger_than_chunk():
+    """Regression: each chunk contributes at most ``chunk`` columns, so the
+    running merge starts narrower than ``top_k`` and used to crash with
+    "selected index k out of range"."""
+    from late_interaction_kernels import retrieve
+
+    torch.manual_seed(0)
+    Q = torch.randn(2, 8, 16, dtype=torch.float32)
+    D = torch.randn(20, 6, 16, dtype=torch.float32)
+    full_s, full_i = retrieve(Q, D, top_k=10, normalize=True)
+    chunked_s, chunked_i = retrieve(Q, D, top_k=10, normalize=True, chunk=4)
+    assert chunked_s.shape == (2, 10)
+    assert torch.allclose(full_s, chunked_s, atol=1e-5)
+    assert torch.equal(full_i, chunked_i)
+
+
+def test_retrieve_chunked_top_k_larger_than_nd():
+    """``top_k > Nd`` with chunking still clamps the output width to Nd."""
+    from late_interaction_kernels import retrieve
+
+    torch.manual_seed(0)
+    Q = torch.randn(2, 8, 16, dtype=torch.float32)
+    D = torch.randn(6, 6, 16, dtype=torch.float32)
+    full_s, full_i = retrieve(Q, D, top_k=10, normalize=True)
+    chunked_s, chunked_i = retrieve(Q, D, top_k=10, normalize=True, chunk=4)
+    assert chunked_s.shape == (2, 6)
+    assert torch.allclose(full_s, chunked_s, atol=1e-5)
+    assert torch.equal(full_i, chunked_i)
+
+
 def test_retrieve_clamps_k_to_nd():
     from late_interaction_kernels import retrieve
 
diff --git a/tests/test_topk.py b/tests/test_topk.py
index 0550b27..ae6dde4 100644
--- a/tests/test_topk.py
+++ b/tests/test_topk.py
@@ -42,6 +42,35 @@ def test_topk_chunked_matches_unchunked(chunk):
     assert torch.equal(i_full, i_chunk)
 
 
+def test_topk_chunked_k_larger_than_chunk():
+    """Regression: chunk_size < k means the running merge starts narrower
+    than k; the merge used to call torch.topk with an out-of-range k."""
+    from late_interaction_kernels.topk import maxsim_topk
+
+    Q = torch.randn(2, 16, 64, device="cuda", dtype=torch.bfloat16)
+    D = torch.randn(20, 24, 64, device="cuda", dtype=torch.bfloat16)
+
+    s_full, i_full = maxsim_topk(Q, D, 10)
+    s_chunk, i_chunk = maxsim_topk(Q, D, 10, chunk_size=4)
+    assert s_chunk.shape == (2, 10)
+    assert torch.allclose(s_full, s_chunk)
+    assert torch.equal(i_full, i_chunk)
+
+
+def test_topk_chunked_k_larger_than_nd():
+    """k > Nd with chunking still clamps the output width to Nd."""
+    from late_interaction_kernels.topk import maxsim_topk
+
+    Q = torch.randn(2, 16, 64, device="cuda", dtype=torch.bfloat16)
+    D = torch.randn(6, 24, 64, device="cuda", dtype=torch.bfloat16)
+
+    s_full, i_full = maxsim_topk(Q, D, 10)
+    s_chunk, i_chunk = maxsim_topk(Q, D, 10, chunk_size=4)
+    assert s_chunk.shape == (2, 6)
+    assert torch.allclose(s_full, s_chunk)
+    assert torch.equal(i_full, i_chunk)
+
+
 def test_topk_with_masks():
     from late_interaction_kernels import maxsim
     from late_interaction_kernels.topk import maxsim_topk
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 20d570b..000bdae 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -9,8 +9,9 @@
 import importlib.metadata as metadata
 
 import pytest
+import torch
 
-from late_interaction_kernels._utils import package_at_least
+from late_interaction_kernels._utils import assert_max_seqlen_covers, bucket_seqlen, package_at_least
 
 
 @pytest.mark.parametrize(
@@ -38,3 +39,26 @@ def _raise(name: str):
 
     monkeypatch.setattr(metadata, "version", _raise)
     assert package_at_least("definitely-not-installed", "1.0.0") is False
+
+
+@pytest.mark.parametrize(
+    ("max_len", "expected"),
+    [(0, 0), (1, 16), (5, 16), (16, 16), (17, 32), (1030, 2048)],
+)
+def test_bucket_seqlen(max_len: int, expected: int):
+    assert bucket_seqlen(max_len) == expected
+
+
+def test_assert_max_seqlen_covers_accepts_exact_and_generous():
+    cu = torch.tensor([0, 4, 12], dtype=torch.int32)
+    assert_max_seqlen_covers(cu, 8, "max_seqlen_q")  # longest is 8
+    assert_max_seqlen_covers(cu, 64, "max_seqlen_q")
+    assert_max_seqlen_covers(torch.zeros(1, dtype=torch.int32), 0, "max_seqlen_q")  # no sequences
+
+
+def test_assert_max_seqlen_covers_rejects_too_small():
+    """On CPU ``torch._assert_async`` evaluates eagerly; on CUDA the same
+    violation surfaces as a device-side assert (see test_varlen.py)."""
+    cu = torch.tensor([0, 4, 12], dtype=torch.int32)
+    with pytest.raises(RuntimeError, match="max_seqlen_q=4 is smaller"):
+        assert_max_seqlen_covers(cu, 4, "max_seqlen_q")
diff --git a/tests/test_varlen.py b/tests/test_varlen.py
index 8ab1484..9b8fe71 100644
--- a/tests/test_varlen.py
+++ b/tests/test_varlen.py
@@ -65,6 +65,57 @@ def test_varlen_matches_padded_path():
     assert (a - b).abs().max().item() < 1e-3
 
 
+def test_varlen_explicit_max_seqlen_matches_inferred():
+    """Caller-supplied max_seqlen_* (non-power-of-2) must give the same
+    scores as the inferred path — exercises the loop-bound bucketing."""
+    from late_interaction_kernels import maxsim_varlen
+
+    q_lens = [5, 9]
+    d_lens = [13, 30, 7]
+    Qp, cu_q = _build_varlen(q_lens, 128)
+    Dp, cu_d = _build_varlen(d_lens, 128)
+
+    inferred = maxsim_varlen(Qp, Dp, cu_q, cu_d)
+    explicit = maxsim_varlen(Qp, Dp, cu_q, cu_d, max_seqlen_q=max(q_lens), max_seqlen_d=max(d_lens))
+    torch.testing.assert_close(inferred, explicit)
+
+    # A generous bound is also fine — the kernel masks on cu_seqlens.
+    generous = maxsim_varlen(Qp, Dp, cu_q, cu_d, max_seqlen_q=64, max_seqlen_d=64)
+    torch.testing.assert_close(inferred, generous)
+
+
+# Too-small max_seqlen_* triggers ``torch._assert_async``, which aborts the
+# CUDA context for the whole process — run in a subprocess like the
+# pack_padded async-assert tests.
+_VARLEN_ASSERT_SCRIPT = """
+import torch
+from late_interaction_kernels import maxsim_varlen
+
+Qp = torch.randn(8, 64, device="cuda", dtype=torch.float16)
+Dp = torch.randn(24, 64, device="cuda", dtype=torch.float16)
+cu_q = torch.tensor([0, 4, 8], device="cuda", dtype=torch.int32)
+cu_d = torch.tensor([0, 12, 24], device="cuda", dtype=torch.int32)
+try:
+    maxsim_varlen(Qp, Dp, cu_q, cu_d, max_seqlen_d=4)  # real max is 12
+    torch.cuda.synchronize()
+except RuntimeError:
+    raise SystemExit(7)
+raise SystemExit(0)
+"""
+
+
+def test_varlen_too_small_max_seqlen_asserts():
+    """max_seqlen_* is a hard loop bound: a too-small value used to silently
+    truncate tokens; now the on-device assert rejects it."""
+    import subprocess
+    import sys
+
+    result = subprocess.run([sys.executable, "-c", _VARLEN_ASSERT_SCRIPT], capture_output=True)
+    assert result.returncode == 7, (
+        f"expected exit 7 (RuntimeError caught); got {result.returncode}\n{result.stderr.decode()}"
+    )
+
+
 # -----------------------------------------------------------------------------
 # Backward parity: check grad_Q and grad_D match the padded autograd path by
 # building two equivalent inputs — one varlen-packed, one padded with masks —

From 526f1af81412a47b53f113b17a6b3235a51d8975 Mon Sep 17 00:00:00 2001
From: Tony Wu <28306721+tonywu71@users.noreply.github.com>
Date: Tue, 9 Jun 2026 22:03:05 +0200
Subject: [PATCH 2/4] perf(plaid): keep exact max_Ld, no power-of-two bucketing
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

H100 measurements on the fastplaid e2e bench showed the bucketed loop
bound costing 30-50% steady-state throughput on Ld=300 corpora
(300 buckets to 512; the masked decompress+dot iterations are not
free), while max_Ld is a property of the compressed index — stable
across calls — so the per-value autotune sweep already amortizes to
one. maxsim_varlen keeps its bucketing: training batches have
genuinely varying maxima and the sweep cost there is the documented
pain point.
---
 CHANGELOG.md                      | 11 ++++++++---
 late_interaction_kernels/plaid.py | 15 ++++++++-------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1a1e561..c174200 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -57,9 +57,14 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   autotune keys with no bucketing, so — contrary to the previous docs —
   every distinct `(max_lq, max_ld)` pair re-triggered the full autotune
   sweep. The kernel masks on the actual `cu_seqlens` bounds, so results are
-  unchanged; the argmax buffer is sized on the bucketed `max_lq` (up to 2×
-  larger, rows `-1`-padded). The same bucketing now applies to the `max_Ld`
-  loop bound of `maxsim_residual` / `maxsim_residual_varlen`.
+  unchanged; the trade-off (mirroring the dense path's `Lq` bucketing) is
+  up to one power-of-two of fully-masked tile iterations per program, and
+  the argmax buffer is sized on the bucketed `max_lq` (up to 2× larger,
+  rows `-1`-padded). The PLAID kernels (`maxsim_residual` /
+  `maxsim_residual_varlen`) deliberately keep their exact `max_Ld`: it is a
+  property of the compressed index, stable across calls, so the sweep
+  amortizes to one — bucketing there measured 30-50% steady-state
+  throughput loss on `Ld=300` corpora for no benefit.
 - **`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
diff --git a/late_interaction_kernels/plaid.py b/late_interaction_kernels/plaid.py
index a5ef780..091f1eb 100644
--- a/late_interaction_kernels/plaid.py
+++ b/late_interaction_kernels/plaid.py
@@ -17,7 +17,7 @@
 import triton.language as tl
 
 from late_interaction_kernels._autotune import autotune_kwargs, forward_configs, prune_forward
-from late_interaction_kernels._utils import assert_max_seqlen_covers, bucket_seqlen, next_pow2
+from late_interaction_kernels._utils import assert_max_seqlen_covers, next_pow2
 
 # -----------------------------------------------------------------------------
 # C1. plaid_approx_score
@@ -447,9 +447,11 @@ def _maxsim_residual_forward(
         raise ValueError(
             f"residuals last dim {packed_dim} != expected {expected_pd} for d={d}, nbits={nbits}"
         )
-    # max_Ld is constexpr + autotune key; bucket it so distinct padded widths
-    # share the cache entry (the kernel masks on the real per-doc length).
-    max_Ld = bucket_seqlen(max_Ld)
+    # max_Ld stays exact (no power-of-two bucketing, unlike maxsim_varlen):
+    # it is a property of the compressed index, stable across calls, so the
+    # per-value autotune sweep amortizes to one — while a bucketed loop bound
+    # costs masked decompress+dot iterations on every call (measured -30-50%
+    # throughput on Ld=300 corpora, where 300 would bucket to 512).
 
     d_pad = next_pow2(d)
     codes_per_byte = 8 // nbits
@@ -896,9 +898,8 @@ def maxsim_residual_varlen(
         max_seqlen_d = int((ends - starts).max().item()) if Nd > 0 else 0
     else:
         assert_max_seqlen_covers(cu_seqlens_d, int(max_seqlen_d), "max_seqlen_d")
-    # max_seqlen_d is constexpr + autotune key; bucket it so distinct batch
-    # maxima share the cache entry (the kernel masks on the real doc length).
-    max_seqlen_d = bucket_seqlen(max(int(max_seqlen_d), 1))
+    # Exact, not bucketed — see the matching note in _maxsim_residual_forward.
+    max_seqlen_d = max(int(max_seqlen_d), 1)
 
     d_pad = next_pow2(d)
     codes_per_byte = 8 // nbits

From dccd1276adcb63325b8f9557f75392c52b10aa43 Mon Sep 17 00:00:00 2001
From: Tony Wu <28306721+tonywu71@users.noreply.github.com>
Date: Tue, 9 Jun 2026 22:33:07 +0200
Subject: [PATCH 3/4] docs: use the Keep-a-Changelog Unreleased heading

AGENTS.md mandates entries land under '## [Unreleased]'; the version
number gets assigned at release time.
---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c174200..5d43ba1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,7 @@ 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).
 
-## 0.4.3 (unreleased)
+## [Unreleased]
 
 ### Fixed
 

From 1fc735d39b13afa9c61bbca205e5b8f5a2e93987 Mon Sep 17 00:00:00 2001
From: Tony Wu <28306721+tonywu71@users.noreply.github.com>
Date: Wed, 10 Jun 2026 09:45:49 +0200
Subject: [PATCH 4/4] perf: keep exact varlen loop bounds, bucket only the
 autotune key

- Make max_lq/max_ld plain runtime args of the varlen forward kernel and
  key its autotune cache on bucketed copies, so sweeps stay capped with
  zero masked loop iterations (same scheme as the dense kernel's Ld)
- Take max_lq as a runtime arg in the varlen backward kernels and drop
  the dead Nq constexpr, avoiding recompiles on new query maxima
- Replace falsy-or fallbacks in prune_forward with explicit None checks
  so a legit Lq == 0 is not treated as missing
- Pin the CHANGELOG section to v0.4.3 (2026-06-10)
---
 CHANGELOG.md                          | 37 +++++++++++++++------------
 late_interaction_kernels/_autotune.py | 11 ++++++--
 late_interaction_kernels/_utils.py    | 19 +++++++-------
 late_interaction_kernels/varlen.py    | 37 +++++++++++++++------------
 tests/test_autotune_kwargs.py         | 23 +++++++++++++++++
 tests/test_varlen.py                  |  3 ++-
 6 files changed, 85 insertions(+), 45 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5d43ba1..0504bbe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,7 @@ 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
 
@@ -52,19 +52,20 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ### Changed
 
-- **`maxsim_varlen` now buckets `max_lq` / `max_ld` to the next power of
-  two** (floor 16) before they reach the kernel. Both 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. The kernel masks on the actual `cu_seqlens` bounds, so results are
-  unchanged; the trade-off (mirroring the dense path's `Lq` bucketing) is
-  up to one power-of-two of fully-masked tile iterations per program, and
-  the argmax buffer is sized on the bucketed `max_lq` (up to 2× larger,
-  rows `-1`-padded). The PLAID kernels (`maxsim_residual` /
-  `maxsim_residual_varlen`) deliberately keep their exact `max_Ld`: it is a
-  property of the compressed index, stable across calls, so the sweep
-  amortizes to one — bucketing there measured 30-50% steady-state
-  throughput loss on `Ld=300` corpora for no benefit.
+- **`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
@@ -76,8 +77,12 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   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) and the unused `Lq` / `Ld` placeholder
-  params from the varlen forward kernel.
+  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 /
diff --git a/late_interaction_kernels/_autotune.py b/late_interaction_kernels/_autotune.py
index 4a25432..0c9017f 100644
--- a/late_interaction_kernels/_autotune.py
+++ b/late_interaction_kernels/_autotune.py
@@ -136,11 +136,18 @@ def _smem_bytes(cfg, d_pad: int) -> int:
 def prune_forward(configs, named_args, **kwargs):
     """Drop configs that overflow shared memory or are oversized for the problem."""
     # The varlen kernels spell their query loop bound `max_lq` instead of `Lq`.
-    Lq = named_args.get("Lq") or named_args.get("max_lq") or 32
+    # `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
diff --git a/late_interaction_kernels/_utils.py b/late_interaction_kernels/_utils.py
index 4e85d1a..48a2da0 100644
--- a/late_interaction_kernels/_utils.py
+++ b/late_interaction_kernels/_utils.py
@@ -13,15 +13,16 @@ def next_pow2(x: int) -> int:
 
 
 def bucket_seqlen(max_len: int, floor: int = 16) -> int:
-    """Round a kernel's sequence-loop bound up to the next power of two.
-
-    Several kernels take a max-length loop bound that is both a ``constexpr``
-    and an autotune key, so every distinct value re-triggers the (5-10 s)
-    autotune sweep. Bucketing to {floor, 2*floor, ...} caps the cache at a
-    handful of entries. Safe because every such kernel masks its loads and
-    reductions on the *actual* per-sequence length — a larger loop bound only
-    adds fully-masked iterations. The default floor matches the smallest
-    block size in the autotune pools.
+    """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
diff --git a/late_interaction_kernels/varlen.py b/late_interaction_kernels/varlen.py
index cb3518b..f7b71de 100644
--- a/late_interaction_kernels/varlen.py
+++ b/late_interaction_kernels/varlen.py
@@ -22,7 +22,7 @@
 
 @triton.autotune(
     configs=forward_configs(),
-    key=["max_lq", "max_ld", "d_pad"],
+    key=["lq_key", "ld_key", "d_pad"],
     prune_configs_by={"early_config_prune": prune_forward},
     **autotune_kwargs(),
 )
@@ -36,8 +36,10 @@ def _varlen_fwd_kernel(
     argmax_ptr,
     Nq: tl.constexpr,
     Nd: tl.constexpr,
-    max_lq: tl.constexpr,
-    max_ld: tl.constexpr,
+    max_lq,
+    max_ld,
+    lq_key,  # bucketed max_lq; unused, autotune cache key only
+    ld_key,  # bucketed max_ld; unused, autotune cache key only
     d: tl.constexpr,
     d_pad: tl.constexpr,
     stride_q_t,
@@ -143,9 +145,8 @@ def _varlen_bwd_dQ_kernel(
     argmax_ptr,  # [Nq, Nd, max_lq]
     grad_s_ptr,  # [Nq, Nd] fp32
     grad_Q_ptr,  # [sum_Lq, d] fp32
-    Nq: tl.constexpr,
     Nd: tl.constexpr,
-    max_lq: tl.constexpr,
+    max_lq,
     d: tl.constexpr,
     d_pad: tl.constexpr,
     stride_d_t,
@@ -211,7 +212,7 @@ def _varlen_bwd_dD_kernel(
     grad_s_ptr,
     grad_D_ptr,  # [sum_Ld, d] fp32
     Nd: tl.constexpr,
-    max_lq: tl.constexpr,
+    max_lq,
     d: tl.constexpr,
     d_pad: tl.constexpr,
     stride_q_t,
@@ -302,13 +303,14 @@ def _varlen_forward(
     else:
         assert_max_seqlen_covers(cu_seqlens_d, int(max_seqlen_d), "max_seqlen_d")
 
-    # max_lq / max_ld are constexpr loop bounds AND autotune keys, so without
-    # bucketing every distinct (max_lq, max_ld) pair re-triggers the full
-    # autotune sweep. The kernel masks on the actual cu_seqlens bounds, so a
-    # larger loop bound only adds fully-masked iterations. The argmax buffer
-    # is sized on the bucketed value (its rows are -1-padded; backward skips).
-    max_seqlen_q = bucket_seqlen(int(max_seqlen_q))
-    max_seqlen_d = bucket_seqlen(int(max_seqlen_d))
+    # max_lq / max_ld stay exact runtime loop bounds (zero masked iterations,
+    # like the dense kernel's Ld); only the autotune cache is keyed on the
+    # bucketed copies below, so distinct (max_lq, max_ld) pairs share one
+    # sweep per power-of-two bucket instead of each re-triggering it.
+    max_seqlen_q = int(max_seqlen_q)
+    max_seqlen_d = int(max_seqlen_d)
+    lq_key = bucket_seqlen(max_seqlen_q)
+    ld_key = bucket_seqlen(max_seqlen_d)
 
     d_pad = next_pow2(d)
     compute_dtype = pick_compute_dtype(Q_packed, D_packed)
@@ -337,6 +339,8 @@ def _varlen_forward(
         Nd,
         max_seqlen_q,
         max_seqlen_d,
+        lq_key,
+        ld_key,
         d,
         d_pad,
         Q_packed.stride(0),
@@ -387,7 +391,6 @@ def backward(ctx, grad_scores):
                 argmax,
                 grad_scores,
                 grad_Q,
-                Nq,
                 Nd,
                 max_q,
                 d,
@@ -457,9 +460,9 @@ def maxsim_varlen(
         max_seqlen_q, max_seqlen_d: hard kernel loop bounds, NOT hints — a
             value smaller than the longest sequence would silently drop
             tokens, so it is rejected by an on-device assert (no D2H sync).
-            Inferred from ``cu_seqlens`` (one D2H sync) if omitted. Both are
-            rounded up to the next power of two internally so the autotune
-            cache is reused across batches with different maxima.
+            Inferred from ``cu_seqlens`` (one D2H sync) if omitted. The
+            autotune cache is keyed on power-of-two buckets of both, so it
+            is reused across batches with different maxima.
 
     Returns:
         scores: ``[Nq, Nd]`` fp32.
diff --git a/tests/test_autotune_kwargs.py b/tests/test_autotune_kwargs.py
index 1fdd167..121ed57 100644
--- a/tests/test_autotune_kwargs.py
+++ b/tests/test_autotune_kwargs.py
@@ -45,6 +45,29 @@ def test_prune_forward_fallback_returns_smallest_footprint():
     assert sorted(_smem_bytes(cfg, d_pad) for cfg in kept) == smallest
 
 
+def test_prune_forward_reads_varlen_max_lq_spelling():
+    """The varlen kernels spell their query loop bound ``max_lq``; the
+    oversized-BLOCK_Q rule must apply to it like to the dense ``Lq``."""
+    from late_interaction_kernels._autotune import forward_configs, prune_forward
+
+    kept = prune_forward(forward_configs(), {"max_lq": 16, "d_pad": 64})
+    assert kept
+    assert all(cfg.kwargs["BLOCK_Q"] <= 32 for cfg in kept)
+
+
+def test_prune_forward_zero_lq_is_not_treated_as_missing():
+    """A legit ``max_lq == 0`` must not fall through to the default 32: every
+    BLOCK_Q is oversized for it, so the smallest-footprint fallback applies."""
+    from late_interaction_kernels._autotune import _smem_bytes, forward_configs, prune_forward
+
+    configs = forward_configs()
+    kept = prune_forward(configs, {"max_lq": 0, "d_pad": 128})
+
+    assert len(kept) == 2
+    smallest = sorted(_smem_bytes(cfg, 128) for cfg in configs)[:2]
+    assert sorted(_smem_bytes(cfg, 128) for cfg in kept) == smallest
+
+
 def test_autotune_kwargs_safe_to_unpack_into_decorator():
     """Whatever ``autotune_kwargs()`` returns must compose with our other
     autotune args — i.e. the call site
diff --git a/tests/test_varlen.py b/tests/test_varlen.py
index 9b8fe71..5da18b3 100644
--- a/tests/test_varlen.py
+++ b/tests/test_varlen.py
@@ -67,7 +67,8 @@ def test_varlen_matches_padded_path():
 
 def test_varlen_explicit_max_seqlen_matches_inferred():
     """Caller-supplied max_seqlen_* (non-power-of-2) must give the same
-    scores as the inferred path — exercises the loop-bound bucketing."""
+    scores as the inferred path — exercises the explicit loop-bound path
+    and its bucketed autotune key."""
     from late_interaction_kernels import maxsim_varlen
 
     q_lens = [5, 9]