Evict swa kv cache during decoding

[ispobock]

Motivation

• Before: SWARadixCache only inserts new tokens into the radix tree when (chunked) prefill finishes or a whole request finishes, so out of window tokens can’t be evicted during decoding.
• Issue: Long generations waste memory and limit swa’s benefit.
• Now: Only cache the last sliding_window_size tokens for finished requests (for multi-turn prefix reuse). Out of window tokens can be evicted early without hurting cache hit rate.

Modifications

• Unify SWA eviction for SWARadixCache and SWAChunkCache. Since these tokens are not inserted into RadixCache during decoding, move the eviction logic from the cache implementation into the schedule batch.
• Update cache_finished_req to consider the evicted tokens.

Benchmarking and Profiling

python3 -m sglang.launch_server --model-path XiaomiMiMo/MiMo-V2-Flash --tp-size 4 --trust-remote-code --mem-fraction-static 0.7 --model-loader-extra-config '{"enable_multithread_load": "true","num_threads": 64}' --attention-backend fa3 --page-size 64 --swa-full-tokens-ratio 0.2

python3 benchmark/hicache/bench_multiturn.py --model-path XiaomiMiMo/MiMo-V2-Flash --disable-random-sample \
--output-length 1 --request-length 300 \
--num-clients 10 --num-rounds 10 --max-parallel 10 --request-rate 16 \
--ready-queue-policy random --disable-auto-run --enable-round-barrier --output-length 3000

main:

[2026-01-19 02:58:09 TP0] Decode batch, #running-req: 10, #full token: 330688, full token usage: 0.37, #swa token: 48384, swa token usage: 0.27, cuda graph: True, gen throughput (token/s): 836.43, #queue-req: 0, 
[2026-01-19 02:58:10 TP0] Decode batch, #running-req: 10, #full token: 331072, full token usage: 0.37, #swa token: 48768, swa token usage: 0.27, cuda graph: True, gen throughput (token/s): 836.62, #queue-req: 0, 
[2026-01-19 02:58:10 TP0] Decode batch, #running-req: 10, #full token: 331392, full token usage: 0.37, #swa token: 49088, swa token usage: 0.27, cuda graph: True, gen throughput (token/s): 837.37, #queue-req: 0, 
[2026-01-19 02:58:11 TP0] Decode batch, #running-req: 10, #full token: 331840, full token usage: 0.37, #swa token: 49536, swa token usage: 0.28, cuda graph: True, gen throughput (token/s): 836.79, #queue-req: 0, 


Performance metrics summary:
  Total requests: 100 at 16.0 requests per second
  Average Prompt Length: 15242.46 tokens
  Average Output Length: 2999.55 tokens
  Average TTFT: 0.49
  P90 TTFT: 1.03
  P99 TTFT: 2.85
  Median TTFT: 0.27
  Max TTFT: 2.85
  Average latency: 36.79
  P90 latency: 38.20
  P99 latency: 38.67
  Median latency: 36.69
  Max latency: 38.67
  Input token throughput: 4104.34 tokens per second
  Output token throughput: 807.69 tokens per second
  Request Throughput: 0.27 requests per second
  Cache Hit Rate: 0.874651
Per-round metrics:
  Round 0: Average TTFT = 0.33s, Cache Hit Rate = 0.000000 (10 requests)
  Round 1: Average TTFT = 0.19s, Cache Hit Rate = 0.827510 (10 requests)
  Round 2: Average TTFT = 0.23s, Cache Hit Rate = 0.922246 (10 requests)
  Round 3: Average TTFT = 0.35s, Cache Hit Rate = 0.914384 (10 requests)
  Round 4: Average TTFT = 0.21s, Cache Hit Rate = 0.954352 (10 requests)
  Round 5: Average TTFT = 0.26s, Cache Hit Rate = 0.963636 (10 requests)
  Round 6: Average TTFT = 0.46s, Cache Hit Rate = 0.856698 (10 requests)
  Round 7: Average TTFT = 0.69s, Cache Hit Rate = 0.874354 (10 requests)
  Round 8: Average TTFT = 0.29s, Cache Hit Rate = 0.987110 (10 requests)
  Round 9: Average TTFT = 1.87s, Cache Hit Rate = 0.691597 (10 requests)

this PR:

[2026-01-19 02:49:00 TP0] Decode batch, #running-req: 10, #full token: 330368, full token usage: 0.37, #swa token: 5248, swa token usage: 0.03, cuda graph: True, gen throughput (token/s): 814.85, #queue-req: 0, 
[2026-01-19 02:49:00 TP0] Decode batch, #running-req: 10, #full token: 330624, full token usage: 0.37, #swa token: 5376, swa token usage: 0.03, cuda graph: True, gen throughput (token/s): 813.99, #queue-req: 0, 
[2026-01-19 02:49:01 TP0] Decode batch, #running-req: 10, #full token: 331072, full token usage: 0.37, #swa token: 5824, swa token usage: 0.03, cuda graph: True, gen throughput (token/s): 814.25, #queue-req: 0, 
[2026-01-19 02:49:01 TP0] Decode batch, #running-req: 10, #full token: 331584, full token usage: 0.37, #swa token: 6336, swa token usage: 0.04, cuda graph: True, gen throughput (token/s): 813.87, #queue-req: 0, 
[2026-01-19 02:49:02 TP0] Decode batch, #running-req: 10, #full token: 331904, full token usage: 0.37, #swa token: 5376, swa token usage: 0.03, cuda graph: True, gen throughput (token/s): 815.36, #queue-req: 0,

All requests completed
Performance metrics summary:
  Total requests: 100 at 16.0 requests per second
  Average Prompt Length: 15253.59 tokens
  Average Output Length: 2999.55 tokens
  Average TTFT: 0.25
  P90 TTFT: 0.40
  P99 TTFT: 0.75
  Median TTFT: 0.23
  Max TTFT: 0.75
  Average latency: 37.21
  P90 latency: 38.12
  P99 latency: 38.31
  Median latency: 37.41
  Max latency: 38.31
  Input token throughput: 4057.60 tokens per second
  Output token throughput: 797.91 tokens per second
  Request Throughput: 0.27 requests per second
  Cache Hit Rate: 0.965354
Per-round metrics:
  Round 0: Average TTFT = 0.16s, Cache Hit Rate = 0.000000 (10 requests)
  Round 1: Average TTFT = 0.16s, Cache Hit Rate = 0.818212 (10 requests)
  Round 2: Average TTFT = 0.22s, Cache Hit Rate = 0.861090 (10 requests)
  Round 3: Average TTFT = 0.29s, Cache Hit Rate = 0.966613 (10 requests)
  Round 4: Average TTFT = 0.23s, Cache Hit Rate = 0.974975 (10 requests)
  Round 5: Average TTFT = 0.28s, Cache Hit Rate = 0.979822 (10 requests)
  Round 6: Average TTFT = 0.26s, Cache Hit Rate = 0.967138 (10 requests)
  Round 7: Average TTFT = 0.21s, Cache Hit Rate = 0.971974 (10 requests)
  Round 8: Average TTFT = 0.38s, Cache Hit Rate = 0.976123 (10 requests)
  Round 9: Average TTFT = 0.28s, Cache Hit Rate = 0.988668 (10 requests)

This PR can keep much lower swa token usage during decoding and have better cache hit rate in muti-turn benchmark.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @ispobock, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the memory management strategy for Sliding Window Attention (SWA) KV caches. By centralizing and refining the eviction logic within the scheduling batch, it ensures that only the most relevant tokens within the sliding window are retained, thereby significantly reducing memory consumption during long token generation processes. This change enhances the efficiency and scalability of SWA-enabled models by preventing unnecessary memory accumulation.

Highlights

• Improved SWA KV Cache Eviction: The pull request introduces early eviction of out-of-window tokens from the Sliding Window Attention (SWA) KV cache during the decoding phase, addressing previous memory waste issues during long generation sequences.
• Unified Eviction Logic: The SWA eviction logic has been unified for both SWARadixCache and SWAChunkCache and moved from the individual cache implementations to the schedule_batch manager, centralizing control.
• Enhanced Radix Cache Insertion: The SWARadixCache's insertion mechanism has been updated to consider swa_evicted_seqlen, allowing for more granular handling of tokens that fall outside the sliding window during insertion or re-insertion of finished requests.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the SWA KV cache eviction logic by centralizing it in ScheduleBatch, which is a good improvement for code clarity and consistency. The changes allow for evicting out-of-window tokens during the decoding phase, which should improve memory efficiency for long generations. The new logic in SWARadixCache to handle partially evicted sequences is complex but seems correct. I have a few suggestions to improve maintainability and robustness.

[hanming-lu]

I see. you are actually evicting them instead of unlocking

[ispobock]

Yes, the generated tokens are not in the tree until the sequence finished. So we cannot unlock it.

[hanming-lu]

add some comments for each branch please

[hanming-lu]

Mainly paid attention to swa radix cache logic and evict frequency. Looks good. Added some comments, mainly on naming and add some comments please.

Also, please share how you tested the change, and add test coverage for the long decode. I am not aware of any test covering the swa radix cache? Maybe we can remove some self.disable_hybrid_swa_memory = True in the server_args? e.g. GptOssForCausalLM

[yinghai]

neat

[ispobock]

/tag-and-rerun-ci

[ispobock]

swa radix cache kl test passed https://github.com/sgl-project/sglang/actions/runs/21115701448/job/60738198124?pr=17220#step:5:25851

[sitabulaixizawaluduo]

Is this optimization intended for non-PD Disaggregation scenarios? I see that there are already corresponding optimizations for PD disaggregation scenarios.

[yizhang2077]

Is this optimization intended for non-PD Disaggregation scenarios? I see that there are already corresponding optimizations for PD disaggregation scenarios.

I think it is mainly optimized for non-PD scenarios. This optimization is mainly focused on long-decode cases with radix-cache. While in PD settings decode node's radix cache is closed and the swa optimization in Chunk cache has been implemented