DynamicCache(offloading=True) returns wrong logits after first cache reuse on hybrid Gated-DeltaNet + attention models (e.g. Qwen3.8-27B)

[com-junkawasaki]

System Info

• transformers version: 5.17.0
• Platform: Linux (Modal container), single NVIDIA RTX PRO 6000 Blackwell (96GB)
• Python: 3.12
• torch: 2.14.0, CUDA
• Model: Qwen/Qwen3.8-27B (hybrid architecture — config.layer_types reports 48/64 layers linear_attention (Gated DeltaNet-style recurrent layers) and 16/64 layers full_attention)
• dtype: bfloat16, attn_implementation="sdpa"

Who can help?

Not sure who owns hybrid-architecture (Gated DeltaNet / linear-attention) cache layers specifically — feel free to redirect.

Reproduction

DynamicCache(config=model.config, offloading=True) returns numerically wrong logits as soon as the cache is reused in a second forward call, on this hybrid Gated-DeltaNet + attention model. The first forward call (prefill) is bit-identical to offloading=False; the very next call that reuses past_key_values already diverges far beyond bf16 rounding noise, and the divergence compounds every subsequent step under greedy decoding (different argmax tokens from step 1 onward).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

model_id = "Qwen/Qwen3.8-27B"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, dtype=torch.bfloat16, attn_implementation="sdpa", device_map="cuda"
).eval()

g = torch.Generator(device="cuda").manual_seed(999)
ids = torch.randint(1000, 50000, (2, 4096), device="cuda", generator=g)

def run(offloading: bool, n_steps: int = 4):
    past = DynamicCache(config=model.config, offloading=offloading) if offloading else None
    with torch.no_grad():
        out = model(input_ids=ids, past_key_values=past, use_cache=True, logits_to_keep=1)
    logits_hist = [out.logits.clone()]
    past = out.past_key_values
    next_ids = out.logits[:, -1:].argmax(-1)
    for _ in range(n_steps - 1):
        with torch.no_grad():
            out = model(input_ids=next_ids, past_key_values=past, use_cache=True)
        logits_hist.append(out.logits.clone())
        past = out.past_key_values
        next_ids = out.logits[:, -1:].argmax(-1)
    return logits_hist

base = run(False)
off = run(True)
for step, (lb, lo) in enumerate(zip(base, off)):
    diff = (lb.float() - lo.float()).abs()
    print(step, "max_abs_diff", float(diff.max()), "argmax_match", bool((lb.argmax(-1) == lo.argmax(-1)).all()))

Output:

0 max_abs_diff 0.0 argmax_match True
1 max_abs_diff 2.1171875 argmax_match False
2 max_abs_diff 7.625 argmax_match False
3 max_abs_diff 17.4375 argmax_match False

Expected behavior

offloading=True is documented as changing only where cache tensors live ("Whether to perform offloading of the layers to cpu, to save GPU memory") — it should not change the numerical result. Step 0 confirms this holds for a single forward call (max abs diff 0.0, exact match). But the second call, which is the first call that actually has to move an offloaded layer's state back for use, already produces a max abs diff of 2.12 on a normal-range logit scale (base top-1 logit ≈ 6.2) — far larger than bf16 rounding (~1e-2 scale at this magnitude), and it is exactly reproducible (not run-to-run noise; re-running gives the same numbers).

This was found while investigating DynamicCache offloading as a KV-memory-reduction technique for a hybrid Gated-DeltaNet + attention model. Memory and prefill/decode-speed numbers with offloading=True were otherwise plausible (lower peak GPU memory, ~4-6x slower decode, both directionally expected for CPU offload) — only the correctness of the linear_attention layers' recurrent/conv state under offload looks broken. I have not traced this further than "the bug appears exactly at the first post-prefill forward call that reuses the cache"; I have not isolated which specific tensor (recurrent state vs conv state vs something else) fails to round-trip through CPU offload for this layer type, since I don't have bandwidth to go deeper right now — flagging with a minimal, exact repro in case it's useful, and happy to help narrow it down further if that's helpful.

Logits (and therefore generated tokens) with offloading=True should match offloading=False up to ordinary bf16 nondeterminism, at every step, not just the first.

[Rocketknight1]

cc @Cyrilvallez for cache