Token smearing does not correctly handle chunked prefill / chunk inference (first token of a chunk misses its cross-chunk predecessor

[Libertax-coder]

Summary

The token smearing logic in the KV-cache (kv_cache is not None) branch appears to mishandle chunk inference (a.k.a. chunked prefill, where more than one token is processed per step while a cache already holds prior context). The first token of each chunk silently skips the smear contribution from its true predecessor, which lives in the previous chunk.

Relevant code

if kv_cache is None:
    # Training / naive generate: full sequence available, use fast slice
    assert T > 1, "Training forward pass should have T > 1"
    gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
    x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
else:
    # KV cache inference: read prev embedding from cache, store current for next step
    x_pre_smear = kv_cache.prev_embedding
    kv_cache.prev_embedding = x[:, -1:, :]
    if T > 1:
        # Prefill: apply smear to positions 1+, same as training
        gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
        x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
    elif x_pre_smear is not None:
        # Decode: single token, use cached prev embedding
        gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, :, :24]))
        x = x + gate * x_pre_smear

The problem

Inside the cache branch, the code classifies the step purely by the number of incoming tokens T:

• T > 1 → treated as prefill, applies x[:, :1] (first token left untouched, no smear).
• T == 1 → treated as decode, smears using the cached x_pre_smear.

This implicitly assumes T > 1 always means "start of sequence", which is only true for the initial prefill. Under chunk inference, each step feeds a chunk of T > 1 tokens while the cache already holds earlier context. In that case the chunk's first token is not the start of the sequence — its predecessor is the last token of the previous chunk (which is exactly what x_pre_smear holds).

The T > 1 branch:

• writes kv_cache.prev_embedding = x[:, -1:, :] (so the cache is updated correctly), but
• never reads x_pre_smear, and leaves x[:, :1] unsmeared.

Result: the first token of every non-initial chunk drops the smear contribution from its cross-chunk predecessor, so chunked prefill produces different (incorrect) activations than token-by-token decode or full-sequence prefill would.

This is also inconsistent with the attention path (_sdpa_attention), which does explicitly handle the Tq != Tk chunk case via an explicit mask.

Questions

1. Is chunk inference an intended/supported path for this model, or is smearing only expected to run under training / single-shot prefill / single-token decode?
2. If chunk inference is supported, is the missing cross-chunk smear on the first token of each chunk a known limitation or an actual bug?

Suggested fix

Unify the cache branch so it no longer special-cases T. Build the full "predecessor" sequence by prepending the cached x_pre_smear to the shifted current chunk, compute gate over all positions, and smear in one shot:

x_pre_smear = kv_cache.prev_embedding
kv_cache.prev_embedding = x[:, -1:, :]

# gate over ALL tokens (note x[:, :, :24], so it aligns with every position)
gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, :, :24]))

if x_pre_smear is not None:
    # mid-sequence (chunk or decode): first token's predecessor comes from cache
    prev = torch.cat([x_pre_smear, x[:, :-1]], dim=1)
    x = x + gate * prev
else:
    # true start of sequence: first token has no predecessor
    x = torch.cat([x[:, :1], x[:, 1:] + gate[:, 1:] * x[:, :-1]], dim=1)

This makes decode (T == 1), initial prefill (T > 1, empty cache), and chunked prefill (T > 1, non-empty cache) all behave consistently, and matches the cross-chunk handling already present in the attention path.

(Note: the original T > 1 branch used x[:, :1] without gate-protecting the first token, and computed gate only over x[:, 1:, :24]; the fix above computes gate over all positions so the first token gets its own gated contribution.)

[bashir-abubakar]

Thanks for the detailed analysis and suggested fix. I’d be interested in taking this forward if you’re not already preparing a PR.

My plan would be to first add CPU regression tests comparing:

full-sequence prefill;
token-by-token cached decoding;
chunked cached prefill.

The tests would verify that the first token of a non-initial chunk uses the cached preceding embedding, while initial-prefill behaviour remains unchanged. I would then implement the unified predecessor logic you proposed.

@Libertax-coder, are you currently planning to submit this yourself, or would you be happy for me to work on it? It would also be helpful to confirm whether chunked prefill is considered a supported path before changing the production logic.