Skip to content

[Feat] sgl-router: Add opt-in decode cache affinity - #44516

Open
rudrakshkarpe wants to merge 1 commit into
sgl-project:mainfrom
rudrakshkarpe:feat/decode-cache-routing
Open

rudrakshkarpe wants to merge 1 commit into
sgl-project:mainfrom
rudrakshkarpe:feat/decode-cache-routing

Conversation

@rudrakshkarpe

@rudrakshkarpe rudrakshkarpe commented Oct 15, 2026 •

Copy link
Copy Markdown

Motivation

Decode workers can reuse cached prefixes, but bucket-first routing currently rejects cache_aware on decode groups. A repeated request can therefore go to a cold decode worker even when an eligible worker already holds the prefix.

This implements the decode-selection part of the decode-aware pairing item in #43699, related to #43098 and #25760.

Modifications

  • Allow decode groups to opt into cache_aware with prefer or balanced affinity. The default remains power-of-two.
  • Rank admitted workers by their longest confirmed GPU-or-host prefix, then GPU-resident depth and decode load. Disk/external copies and pending placements do not count.
  • Keep bucket and version-group boundaries, admission limits, and load-based fallback. Balanced mode uses running-request load.
  • Limit cache affinity to single-DP-rank workers using the local radix tree. Multi-rank workers can still be selected by the load fallback; remote-index affinity needs a tier-aware API first.

The README includes the bucket configuration. Prefill selection and engine transfer behavior are unchanged.

Accuracy Tests

Added policy regressions for tier changes, eviction, admission, stale load, partial-page thresholds, cache-salt/LoRA isolation, and missing cache signals. An HTTP PD test checks selection within the prefill's version group and preserves the shared bootstrap room.

Passed locally on macOS arm64:

cargo +1.92.0 test --locked --release --workspace -- --skip parity_matrix
cargo +1.92.0 clippy --locked --workspace --all-targets -- -D warnings
cargo +nightly fmt --all -- --check
# From the repository root:
pre-commit run --all-files

The tokenizer parity tests require external fixtures and were excluded. No model accuracy or GPU serving run was performed.

Speed Tests and Profiling

Added a paired CPU lookup benchmark on the same tree, with half the owners on GPU and half on host memory. Apple M4 Pro, Rust 1.92.0 release build, 20 samples, 1s warmup, 2s measurement:

Shared prefix Existing any-tier lookup New resident lookup
8 workers, 64 blocks 8.11 us 13.83 us
32 workers, 512 blocks 232.26 us 360.90 us
cargo +1.92.0 bench --locked --bench tree_lookup -- resident_prefix --sample-size 20 --warm-up-time 1 --measurement-time 2

These measure tree lookup cost, excluding token hashing and dispatch. The additional walk distinguishes GPU residency from host-backed prefixes. Reduced KV transfer and TTFT still need a real PD serving benchmark; this PR makes no claim about those gains.

Checklist

  • Format code with pre-commit.
  • Add regression tests.
  • Update documentation.
  • Provide GPU accuracy and serving performance results.
  • Follow the SGLang code style guidance.

@hzh0425 @sherlockwu, could you review the tier-selection behavior and single-rank scope for the decode-aware pairing work?


CI States

Latest PR Test (Base): ✅ Run #38765013308
Latest PR Test (Extra): ❌ Run #38765012969
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

Signed-off-by: rudrakshkarpe <rudraksh.karpe@gmail.com>
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Oct 15, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant