Fix eagle radix cache

[ispobock]

Motivation

python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct --speculative-algo EAGLE3 \
    --speculative-draft-model-path lmsys/sglang-EAGLE3-LLaMA3.1-Instruct-8B  \
    --speculative-num-steps 2 --speculative-eagle-topk 1 --speculative-num-draft-tokens 3 \
    --dtype float16
    
cd benchmark/mtbench
python3 bench_sglang_eagle.py --parallel 1 --num-questions 10

main:

w/o radix cache:
#questions: 10, Throughput: 269.86 token/s, Acceptance length: 2.39

w/ radix cache:
#questions: 10, Throughput: 255.74 token/s, Acceptance length: 2.26

this PR:

w/o radix cache:
#questions: 10, Throughput: 269.45 token/s, Acceptance length: 2.39

w/ radix cache:
#questions: 10, Throughput: 270.85 token/s, Acceptance length: 2.39

compatibility test:

• page size 16
• page size 16 + chunked prefill 64
• page size 2 + chunked prefill 64
• page size 1 + chunked prefill 64
• HiCache (this fix is not adapted for HiCache, but doesn't break current HiCache)
• multi prompts in one batch (chunked prefill corner case, ref: Fix spec filter batch when target extend  #10991)

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[JustinTong0323]

Awesome!

[ispobock]

The page_size>1 still have some issue, will fix it later.

[ispobock]

@xiezhq-hermann @merrymercy This PR is ready for review.

[xiezhq-hermann]

hiradix (and other trees like swa) override the insert function, would that be a problem since eagle worker shared the same tree?

[ispobock]

Yes, in current design, we need to adapt this change to other trees like swa and hiradix if they override these functions. This PR just makes the main radix tree ready. HiCache and swa need extra work and test to make them ready.

[xiezhq-hermann]

while hiradix does not, swa tree override this implementation as well

[ispobock]

The swa cache inherits from BaseRadixCache, so it seems all the changes should be implemented again on it. HiCache is from RadixCache, we just need to do some adaptation on it with less override. But for HiCache, the main thing I'm concerning is that the chunked prefill size is a little changed. If the chunked prefill size is 64, actually only 63 bigram keys are inserted to the tree. Maybe it's not efficient for cache offloading with block.

[xiezhq-hermann]

My understanding is what we are doing is primarily to resolve conflict with eagle workers since it shares the same radix tree but has its own pool, but not to have hicache support for eagle workers, i.e., eagle workers to fetch kv caches from host memory, which seems unnecessary and potentially complicated. Is it correct?

[ispobock]

I agree that the kv cache for eagle worker is unnecessary to store into host memory since it's only one layer. If we use HiCache only for target model, can we still share the kv indices between target and draft pool?

[xiezhq-hermann]

yes I think it should be fine just wanted to confirm that we are aligned on this

[ispobock]

cc: @merrymercy