Skip to content

[Fix] Compose PIC addition state causally across multiple misses - #7

Open
morluto wants to merge 3 commits into
redai-studio:mainfrom
morluto:agent/pic-multimiss
Open

morluto wants to merge 3 commits into
redai-studio:mainfrom
morluto:agent/pic-multimiss

Conversation

@morluto

@morluto morluto commented Jul 28, 2026

Copy link
Copy Markdown

Motivation

When one request contains multiple cache misses, each missed segment must begin
from the accumulated state of everything preceding it, including cached
segments between misses. The previous GDN and KDA addition paths seeded only
the first miss and ignored prompt order, so later misses lost causal state and
earlier misses could receive state from future cached segments.

The request state also needs to retain cached segments after the final miss.

Modifications

  • compute reusable zero-start states for each missed segment
  • merge cached and missed segments in prompt order and snapshot the causal
    prefix before every miss
  • rerun misses with those prefix states to produce history-aware outputs
  • preserve zero-start miss states in their cache slots
  • add cached suffix state after the final miss to the request's post-kernel
    recurrent state
  • share the composition helper between GDN and KDA
  • allocate the prefix and suffix workspaces only for addition batches that
    contain a multi-miss request
  • preserve the one-pass path when every request has at most one miss

Transition matrices are not introduced into addition mode.

Accuracy Tests

The deterministic regression uses:

next_state = 0.5 * state + token_value

For miss(S=1) → hit(S=10) → miss(S=2) → hit(S=20), it verifies:

miss outputs:        [1.0, 7.5]
final request state: 27.5

The tests exercise the GDN and KDA addition orchestration while replacing only
the hardware-boundary kernels with deterministic CPU implementations. No
full-model or GPU numerical accuracy claim is made.

Speed Tests and Profiling

A batch containing a multi-miss request performs two recurrent passes over its
miss segments: one to obtain reusable zero-start states and one to produce
causally seeded outputs. Batches where every request has at most one miss retain
the existing one-pass behavior.

The additional workspaces are allocated only for multi-miss addition batches
and reused by later calls. No end-to-end performance claim is made.

Test Plan

test/registered/unit/layers/attention/linear/test_pic_addition.py — 5 passed.

Coverage includes:

  • GDN and KDA interleaved miss/hit composition
  • cached suffix state after the final miss
  • mixed single-miss and multi-miss request shapes
  • zero-start miss-state persistence
  • single-miss one-pass preservation
  • addition-only workspace allocation

Closes #2.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] PIC addition loses causal state across multiple misses

1 participant