Skip to content

refactor: generalize hybrid attention checkpoint management - #1559

Open
shihaobai wants to merge 2 commits into
mainfrom
bsh/refactor-hybrid-cache
Open

refactor: generalize hybrid attention checkpoint management#1559
shihaobai wants to merge 2 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. PD I/O structures, page_kind literals and model-specific PD packing/restoration remain identical to main.

No changes to attention computation, Gemma/sliding runtime allocation or the small-page policy in this PR.

Validation

Revalidated the review cleanup at 5f1c96a3 on 2026-09-10 with the same cases and main baseline.

Executed in bsh_dsv4 with Qwen3.5-27B, BF16, TP=1 on H200 GPUs, against pristine main 1eb4810c.

  • Normal deployment and 1P1D with NCCL transport: 19 requests per mode, 38 requests / 608 generated token IDs, all exactly identical before and after the refactor. Normal and PD outputs also agree for all cases.
  • Covered cold inputs, 1024-token big-page hits, 2816-token small-page hits, concurrent requests and GPU eviction followed by CPU restoration. Both normal and PD runs confirmed GPU hit = 0, CPU hit = 2816 on the restoration case.
  • Decode CUDA Graph enabled; chunked prefill and CPU cache enabled (CPU cache on P only in PD mode).
  • Existing regression tests: 158 passed. The unrelated test_radix_cache.py::test_case10 fails on pristine main as well because its tree has no mem_manager; reproduced separately and excluded from the final selected run.
  • pre-commit (Black and Flake8) passed inside the container.

The validation launcher starts P before D to avoid main's existing concurrent NCCL control-port selection race. No unrelated startup fix is included.

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