Skip to content

feat: add DSA cache and PP support#134

Open
zhjc1124 wants to merge 1 commit intotaco-project:feat/layerwise_rebasefrom
zhjc1124:feat/layerwise_rebase
Open

feat: add DSA cache and PP support#134
zhjc1124 wants to merge 1 commit intotaco-project:feat/layerwise_rebasefrom
zhjc1124:feat/layerwise_rebase

Conversation

@zhjc1124
Copy link
Copy Markdown
Contributor

@zhjc1124 zhjc1124 commented Apr 2, 2026

Summary

This PR adds DSA (Dynamic Sparse Attention) cache support and Pipeline Parallelism (PP) support to FlexKV.

Changes

New Features

DSA/NSA Indexer Cache Support

  • Added dataclass in to hold indexer-specific cache configuration (e.g., , , ) for DSA/NSA sparse attention models
  • Extended to manage separate indexer storage handles () for CPU, SSD, and REMOTE devices
  • Extended to accept optional indexer GPU blocks

Pipeline Parallelism (PP) Support

  • Added parameter to so each PP rank only manages its own layers instead of the full model layer count
  • Fixed resolution to use total heads (not per-rank heads) for correct KV layout

@YconquestY YconquestY self-requested a review April 2, 2026 08:57
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