Skip to content

refactor: generalize hybrid attention checkpoint management - #1559

Merged
shihaobai merged 4 commits into
mainfrom
bsh/refactor-hybrid-cache
Sep 10, 2026
Merged

refactor: generalize hybrid attention checkpoint management#1559
shihaobai merged 4 commits into
mainfrom
bsh/refactor-hybrid-cache

Conversation

@shihaobai

@shihaobai shihaobai commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

Summary

Extract the existing linear/full checkpoint flow into reusable hybrid interfaces, without including the Gemma/sliding-window implementation.

  • Add HybridAttentionReqManager for request-state initialization, big/small-page snapshots, restoration and model-specific MTP state updates. Conv/SSM access stays inside the linear implementation.
  • Move checkpoint storage to common/state_cache_manager/: shared slot allocator and layer storage, with the linear config and buffers together in linear_att.py.
  • Generalize request fields, paged radix-cache types, multi-level cache routing and PD state-page dispatch across the repository.
  • Preserve checkpoint creation timing, size-first CPU layouts, page matching/eviction boundaries, CPU/disk packing and existing --linear_att_* options. Model-specific PD packing/restoration and payload layouts remain unchanged. Rename the PD page kind from linear_att_state to att_state and document its boundary-state semantics; P and D must be updated together because the kind is part of the task identity.

@shihaobai
shihaobai merged commit 8d3a424 into main Sep 10, 2026
1 check passed
@shihaobai
shihaobai deleted the bsh/refactor-hybrid-cache branch September 10, 2026 11:14
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.

1 participant