Fake keys from search_hyperplane overflow float16 -> NaN attention and degenerate generation (AdaKVPress / masked_key_indices) #250

Bug description

search_hyperplane in kvpress/attention_patch.py returns fake keys with magnitude
-1e5 * Y / ||Y||^2. With typical query norms, the resulting components routinely
exceed the float16 representable range (max 65504). When attention_patch writes them
into the key cache (key[batch_indices, head_indices, seq_indices] = k[...]) on a
model loaded in torch.float16, they are cast to ±inf. The subsequent q @ k dot
products mix +inf/-inf terms and produce NaN logits, NaN propagates through
softmax, and generation degenerates.

This affects AdaKVPress and any press that relies on module.masked_key_indices.
With meta-llama/Llama-3.2-1B-Instruct the visible symptom is the model emitting an
endless run of ! (token id 0 in the Llama-3 vocabulary), which is the classic
NaN-logits + argmax signature.

Everything works correctly with torch.bfloat16 (dynamic range ~3.4e38), which is
presumably why this does not show up in CI.

Environment

• kvpress 0.5.4 (also present in current main: no dtype clamping in attention_patch.py)
• transformers 5.2.0, torch 2.11.0+cu128
• Windows 11, RTX 3050 Laptop GPU (Ampere, sm_86), attn_implementation="sdpa"

Minimal reproduction

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from kvpress import AdaKVPress, SnapKVPress

model_id = "meta-llama/Llama-3.2-1B-Instruct"
tok = AutoTokenizer.from_pretrained(model_id)
prompt = "The pass code is 71542. " * 300 + "What is the pass code?"
ids = tok(prompt, return_tensors="pt").input_ids.cuda()

for dtype in (torch.float16, torch.bfloat16):
    model = AutoModelForCausalLM.from_pretrained(
        model_id, dtype=dtype, attn_implementation="sdpa"
    ).cuda().eval()
    press = AdaKVPress(press=SnapKVPress(compression_ratio=0.75))
    with torch.no_grad(), press(model):
        out = model.generate(ids, max_new_tokens=24, do_sample=False,
                             pad_token_id=tok.eos_token_id)
    print(dtype, repr(tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True)))

Observed (also reproduced on a real LongBench prompt of ~4000 tokens):

torch.float16  'No!!!!!!!!!!!!!!!!!!!!!!!'   # degenerate
torch.bfloat16 '<coherent answer>'           # correct

Suggested fix

Rescale the fake key per row so it stays representable in the cache dtype while
preserving the hyperplane direction (and therefore the exp(<q, k>) ~= 0 guarantee):

k = -1e5 * Y / Y.norm(dim=-1, keepdim=True) ** 2
finfo = torch.finfo(key.dtype)
scale = (0.9 * finfo.max / k.abs().amax(dim=-1, keepdim=True)).clamp(max=1.0)
k = (k * scale).to(key.dtype)

A uniform per-row rescale is preferable to component-wise clamping because clamping
distorts the direction and can weaken the <q, k> <= 0 guarantee. Alternatively (or
additionally), a warning when a masked-key press runs on a float16 model would make
the failure mode discoverable.

Happy to open a PR if the suggested fix looks reasonable.