[dhkts1]

Every row Qwen3Reranker::score builds starts with the same tokens: the system prompt, then <Instruct>: ..., <Query>: ... and <Document>:. Until now each row ran them through all 28 layers again. On the model card's example, 69 of a row's 85 tokens are the same for every document.

What changes

• score tokenizes the shared head once. It runs the head's first multiple of 32 tokens once as [1, P] and keeps each layer's k and v after RoPE. Then it runs each length group of tails (the rest of the head, the document and the suffix) against them, causal, at positions P onward.
• A row shares only when its token ids start with those P ids. When a document's first characters merge with the head's last token, the merge falls in the part that goes with the tail. A row cut inside the shared ids is scored whole, as before.
• Qwen3Attention::forward and Qwen3DecoderLayer::forward take optional earlier k and v and can return their own. They hold k and v as [B, S, Nkv, D], the layout the rope kernel and v_proj already write. That way joining the shared part to a tail is one plain copy of each. My first version joined in [B, Nkv, S, D]. It copied a transposed view and a broadcast before each join, and long documents, which mostly run one row per length group, came out about 9% slower than whole rows.
• Qwen3Model::forward_causal_after runs rows after shared k and v. On Metal it uses the fused kernel's causal mask. Elsewhere it uses an additive [1, 1, T, P + T] mask shifted by P.

Why a multiple of 32

When candle 0.11's fused Metal attention (call_sdpa_full) gets fewer queries than keys, it lines its causal mask up with the last keys. It masks only the last ceil((BQ + BK - 1) / BK) key blocks, plus one when the key count is not a whole number of blocks. If keys - queries is not a multiple of the key block, a query block's diagonal crosses one block more than that. Then some positions see keys after their own. My first measurement found it: on Metal bf16, 5 of 100 documents of about 500 tokens moved by up to 1.09 in logit against whole rows. BK is 8, 16 or 32, so a shared length that is a multiple of 32 keeps the diagonal inside the masked blocks. This is worth an issue upstream in candle.

Tests

• reranker_scores_with_the_shared_head_as_with_whole_rows (CPU, f32) compares score with whole rows within 1e-5 over documents of 0 to 40 words. It runs with a head that is not a multiple of 32, a 64-token head where one document leaves the shared ids and is scored whole, and a max length that cuts every row inside the shared ids. I broke three things one at a time: removing the fallback, the RoPE offset, or the mask offset each turns it red.
• reranker_scores_with_the_shared_head_as_with_whole_rows_on_metal covers every tail length modulo 32 at head dims 64 and 128, in f32 (1e-4) and bf16 (2e-2). It failed on the version without the rounding (3.0e-3 at head dim 64, f32) and fails again with the block set to 1.
• qwen3_reranker_06b_matches_python (CPU f32 against the model card's Python, 1e-4, plus the mid-relevance logits within 1e-3) and qwen3_reranker_06b_metal_bf16_near_python pass with sharing on.

Measured

score with Qwen3-Reranker-0.6B, Metal bf16, query "What is the capital of China?". The base is the reranker as #339 merged it. Base and new runs alternated in three rounds on a loaded machine (load average 200 to 270).

documents	base (ms)	shared head (ms)
100 card-length (one sentence)	1369, 1370, 1313	568, 528, 490
100 of about 500 tokens	7723, 7869, 7724	6754, 6892, 6793

Card-length documents run about 2.6 times faster, and long ones about 12% faster. In every round, all 200 scores match the base at the 8 decimals the bench prints, and so do the logits computed from them. The long-document token count is an estimate from the text's length; I did not count it.

🍄

[greptile-apps]

Greptile reviewed this. Push back if it's wrong - it has no feelings.

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge, with the previous finding fixed and no new blocking issues found.

_{Reviews (2) · Last reviewed commit: "test: compare the reranker tests' score ..." · Reviewed by Greptile}

[github-actions]

🎉 This PR is included in version 7.5.0 🎉

The release is available on:

• GitHub release
• v7.5.0

Your semantic-release bot 📦🚀

[dhkts1]

Opened huggingface/candle#4062 for the Metal causal-mask issue described above. Once a candle release has it, the multiple-of-32 rounding here can go. 🍄