Skip to content

[Bug] PIC addition aborts a batch when segment-state allocation cannot evict #4

Description

@morluto

Observed behavior

A PIC addition request can satisfy the token and KV-cache budgets but still fail admission when the required Mamba segment-state slots cannot be allocated. If all possible eviction candidates are locked, the allocator raises RuntimeError and aborts the entire batch.

The affected request could instead be processed through normal full-prompt recomputation.

PIC candidate
     |
     v
state slots available? -- yes --> admit reduced PIC work
     |
     no
     v
evict unlocked entries
     |
     v
enough now? ---------- yes --> admit reduced PIC work
     |
     no
     v
price full recompute against token/KV budgets
     |
     +-- fits --> admit full prompt
     |
     +-- no fit --> return NO_TOKEN

Expected behavior

PIC admission should account for required segment-state slots before it commits the request. When eviction cannot provide enough slots, the scheduler should reconsider that request using its full prompt:

  • Admit full recomputation if the normal token and KV budgets permit it.
  • Return NO_TOKEN if neither PIC admission nor full recomputation fits.
  • Do not fail unrelated requests in the batch.

Scope

The reproduced path is non-RoPE PIC addition mode. Transition-mode behavior does not need to change to address it.

Activity

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions