Skip to content

Support request-level cache for hybrid sliding-window attention - #1545

Open
shihaobai wants to merge 15 commits into
mainfrom
bsh/hybrid-sliding-window-cache
Open

Support request-level cache for hybrid sliding-window attention#1545
shihaobai wants to merge 15 commits into
mainfrom
bsh/hybrid-sliding-window-cache

Conversation

@shihaobai

@shihaobai shihaobai commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator

Summary

Gemma4 keeps full-attention KV token-granular and stores sliding-window KV in a request-owned pool. Prefix-cache hits restore window checkpoints through main's shared hybrid-attention interfaces.

  • Reserve W slots per request and reuse them during single-token decode. Larger prefill chunks borrow temporary slots and release obsolete slots after all layers finish.
  • Support KV-shared layers and CPU-pinned window checkpoints for big/small-page GPU prefix-cache hits. CPU-cache pages pack full KV together with window state.
  • Bind sliding prefill/decode state to its own token table. Text-only sliding prefill uses the common Triton backend; image tokens retain Gemma's bidirectional mask. Full attention stays on Triton for 512-dimensional heads.
  • Co-locate SlidingWindowCacheConfig and SlidingWindowStateCacheManager. PD, MTP, and quantized KV remain unsupported for Gemma4.

Validation

  • 95 focused tests passed on H200: config/CPU-page layout, checkpoint copying, CPU-cache transfer, sliding prefill/decode, request-window wraparound/release/restoration, and full/sliding attention numerical comparisons including head_dim=512.
  • Repository pre-commit Black and Flake8 passed.
  • Full-model inference and CUDA Graph replay were not rerun for this update. Earlier PR revisions passed E4B TP2 and 31B TP4 prefix-cache comparisons; the E4B graph check required communication optimizations to be disabled due to a FlashInfer lazy-init capture issue.

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