feat(kvstore): support mamba l2 cache transfers

[XucSh]

Summary

This PR adds L2 cache support for Mamba cache alongside the existing KV cache L2 path.

The core change is to generalize cache movement into CacheTransferUnit, which carries the cache kind (KV or Mamba) plus source and destination slots/pages. WriteBackOperation and LoadBackOperation can now carry mixed KV/Mamba transfer units without adding Mamba-specific scheduler branches everywhere.

Key behavior:

• Mamba cache can be written back from GPU to host memory.
• Mamba cache can be loaded back from host memory to GPU.
• Mamba host memory supports eviction when host Mamba slots are exhausted.
• Mamba loadback completion is synchronized through stream fence semantics, not scheduler loadback-done events.
• Prefix reuse now requires all cache types for a token prefix to exist. KV-only cache is not considered reusable when the corresponding Mamba state is missing.
• Mamba state depth is tracked explicitly so page-aligned prefix states and non-exact retraction recovery states are not confused.

Important design points:

• CacheTransferUnit avoids duplicating writeback/loadback logic for KV and Mamba.
• TreeNode tracks device/host Mamba slots and their token depth.
• FindLastMambaNode(..., require_exact_depth=true) is used for prefix cache reuse.
• Retraction recovery may use non-exact live Mamba state, but normal prefix reuse may not.
• Host-side Mamba eviction detaches only the host Mamba state; future prefix reuse still requires KV and Mamba to both be present.

Test Plan

CUDA_VISIBLE_DEVICES=4,5,6,7 \ ts serve Qwen/Qwen3.5-122B-A10B \ --tp 4 \ --max-num-seqs 16 \ --max-total-tokens 160000 \ --max-model-len 80000 \ --chunked-prefill-size 128 \ --max-prefill-tokens 128 \ --block-size 64 \ --max-mamba-cache-size 64 \ --mamba-track-interval 64 \ --kvstore-ratio 1.0 \ --kvstore-io-backend kernel \ --disable-prefill-graph \ --enforce-eager \ --disable-overlap-schedule \ --enable-cache-report \ --host 127.0.0.1 \ --port 8000

evalscope eval \ --model Qwen/Qwen3.5-122B-A10B \ --api-url http://127.0.0.1:8000/v1 \ --api-key EMPTY_TOKEN \ --datasets aime25 \ --eval-batch-size 16 \ --generation-config '{"max_tokens":100000}'

mamba loadback/writeback can be seen. The result:

[tuanzhangCS]

Currently, we only call write_back when a request finishes or is retracted, and we only write back the last Mamba slot of that request. Should we also drain Mamba slots during the prefill stage to improve the cache hit rate?

[tuanzhangCS]

Will the mamba_host_slot associated with a retract request be protected from eviction?

[SimonCqk]

QQ, in KVTransferBackend, there are draft_pool write/load operations, will mamba be possible with a draft pool ?

[tuanzhangCS]

The current draft version of Qwen3.5 does not include Mamba layers; however, I cannot confirm whether a Mamba pool will be introduced in the future or other models.

[SimonCqk]

mamba_loadback_diff is a vector and for now its length is alway 1, but overwrite mamba_cow_src_index in for-loop seems weird, should it be lifted as a vector too?

[SimonCqk]

should the memory of mamba host pool be pinned?

[SimonCqk]

LGTM overall, since #146 was just merged, there are conflicts in some files, resolve them to merge 🚀🙌🏻