[dhkts1]

The bug

candle_nn::ops::sdpa with do_causal = true on Metal lets some queries attend to keys after
their own position when there are fewer queries than keys, which is the shape of a prefill on top
of a KV cache.

The full kernel (steel_attention in scaled_dot_product_attention.metal) lines the causal mask up
with the last keys: query i sees keys up to i + kL - qL. Each threadgroup handles a block of BQ
queries and walks the keys in blocks of BK. It applied the mask only to the key blocks with
kb >= kb_lim - ceil(BQ / BK) - int(!align_K), that is the last ceil(BQ / BK) blocks the query
block reaches, plus one when kL is not a whole number of key blocks. When kL - qL is not a
multiple of BK, the query block's diagonal starts inside a key block and crosses one block more
than that. If kL is also a whole number of key blocks, that first block is not masked, and its
queries see up to BK - 2 keys past their position.

A second, smaller problem sits beside it: the loop limit kb_lim = ceil(q_max / BK) was not
clamped to the number of key blocks. With a partial last query block, q_max runs past kL and the
kernel read one block of keys and values past the end of K and V. Those keys are masked, but a
masked score times a NaN value is still NaN, so a K/V view into a longer buffer (a preallocated
KV cache) could put NaN in the output.

We found this in fastembed-rs, reusing a shared prefix across Qwen3 reranker documents
(anush008/fastembed-rs#342). On Metal in bf16, 5 of 100 documents moved by up to 1.09 in logit
against scoring each document whole. On a random model with head dim 64 in f32 the difference was
3.0e-3. The first test there used lengths that happened to miss the bad residues and passed.

The fix

The kernel came from MLX, which already fixed both:

• Nax Refactor ml-explore/mlx#3271 masks from kb_min_causal = max(0, tid.x * BQ + qL_off) / BK, the key block
  holding the query block's first position. Every block before it holds only keys at or before
  every query in the block.
• Fix looping limit in causal attention ml-explore/mlx#1999 clamps kb_lim = min(NK, kb_lim).

This change copies both lines from MLX's current steel_attention.h. Nothing else in the kernel
changes.

Same work for shapes that were already right

When kL and kL - qL are both whole numbers of key blocks (BQ is a multiple of BK at every head
dim: 32/32, 32/16 and 8/8), kb_lim = (tid.x + 1) * BQ / BK + (kL - qL) / BK exactly, so the old
first masked block kb_lim - BQ / BK equals the new kb_min_causal: the kernel masks the same
blocks as before. When qL is a whole number of query blocks, q_max <= kL, so the clamp never
changes kb_lim. The only added work is one integer division per threadgroup. I checked this by
reading the arithmetic, not by timing: the machine was too loaded for a timing to mean much.

The tests

Two tests in candle-nn/tests/sdpa.rs, in the Metal module. They compare sdpa(.., do_causal = true, ..) with a CPU reference in f32 (matmul, the same mask aligned to the last keys, softmax,
matmul) built from the same rounded inputs:

• every kL - qL from 0 to 33;
• for each, 33 queries (a partial query block), 64 (two whole blocks) and 96 - (kL - qL) % 32
  (which makes kL a whole number of key blocks);
• head dims 64, 128 and 512, which pick key blocks of 32, 16 and 8;
• f32 and bf16 (512 is bf16 only, as the kernel allows).

The tolerance is 1e-4 in f32 and 3e-2 in bf16. With this change the largest differences are
2.0e-6 in f32 and 1.7e-2 in bf16. On main the smallest failing difference is 0.057.

sdpa_full_causal_fewer_queries_than_keys uses contiguous K and V.
sdpa_full_causal_reads_no_key_past_kv_len gives K and V as views of the first kL rows of a
buffer whose remaining 64 rows are NaN, so any read past kL shows up in the output.

• On main: the first test fails on 145 shapes, the second on 318. The 145 are exactly the shapes
  where kL is a whole number of key blocks and (kL - qL) % BK is neither 0 nor BK - 1 (31 at
  head dim 64, 29 at 128, 25 at 512). The largest difference is 4.2 (head dim 512, bf16, qL = 95,
  kL = 96). In the second test all 318 come out NaN: every failing shape has a partial last query
  block, so the kernel read the NaN rows past kL.
• With only the clamp removed, the first test passes and the second fails on 318 shapes.
• With only the old mask condition restored, both fail on the same 145 shapes.
• With this change: both pass, as do the five existing tests in sdpa.rs (7 passed).

Checks run

On an Apple M4 Pro with macOS's Metal, in this branch's own target directory:

• cargo fmt --all -- --check: clean.
• cargo clippy -p candle-metal-kernels -p candle-nn --features candle-nn/metal --lib -- -D warnings:
  clean.
• cargo clippy -p candle-nn --features metal --lib --tests -- -D warnings: clean.
• candle-metal-kernels' own tests do not build on main at present: src/tests.rs uses rand and
  rand_distr, which are not among its dependencies. This change does not touch that file.

🍄