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

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.)