[tracking] Minimal prefix-sharing kv cache

[renxida]

Block Trie Attention Implementation Plan

Also known as Radix Attention / Block Trie attention

Project Goal

Implement a KV cache that:

• Preserves old KV values with their associated tokens
• Matches them against new incoming sequences to avoid redundant computation
• Will eventually support:
  • Cache page sharing and KV computation of common prefixes between inference requests in the same batch
  • Large page copying upon partial matches

Implementation Tasks

Phase 1: Preparation

• Enable decode requests to access token list (Give InferenceExecRequest full token context #582)
• Implement TrieAttentionPagedCache
  • Complete prerequisite PRs:
    • Replace AttnPagedCache with BasePagedAttentionCache #565
    • Implement PageAllocation as a handle into a PagedAttentionCache, allowing publishing and releasing an allocation via handle rather than cache #608
  • Merge PR TriePagedAttentionCache - 2 #628: TriePagedAttentionCache
    • Unit test implementation
      • Verify cache sharing through multiple sequential cache allocations with shared cache lines
      • Verify proper eviction functionality
      • Verify reference counting prevents premature eviction
    • Pass unit test suite
    • Pass integration test suite
    • Create new PR from draft implementation (from PR Draft Blocktrie PR #600)

Phase 2: Janky implementation -> minimum working implementation

Note: At this stage, the class should work for LLM integration tests but may fail sglang integration tests. Stephen's involvement will be crucial for hardening against concurrency and optimizing benefits.

• Concurrent access testing
• Multi-request batch testing with shared prefixes
  • Verify error-free execution of requests
  • Same-batch common-prefix related tests
    • Demonstrate that same-batch inputs with common prefixes do not crash system / produce bad tokens.
    • Demonstrate that same-batch inputs with common prefixes are not able to share cache pages and have no performance advantage over same-batch inputs without common prefixes. Do so by implementing a test that expects improved performance and xfailing it
• Performance comparison testing against BaseAttentionCache. The tests should expect substantial performance improvement vs Base. Xfail if we don't find the improvement and figure out why.

Phase 3: Benchmarking / polishing

Reference Implementations

• sglang
• lmdeploy/pytorch/paging/block_trie.py (InternLM/lmdeploy)

Priority Note

Focus on achieving minimum viable solution to facilitate Stephen's involvement in the project.