Repository navigation
perf(sglang): reduce admission checkpoint copies - #321
Conversation
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. Important Review skippedReview was skipped as selected files did not have any reviewable changes. ⚙️ Run configurationConfiguration used: Repository: ai-dynamo/aisimulate/.coderabbit.yaml Review profile: ASSERTIVE Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: ai-dynamo/aisimulate/.coderabbit.yaml Review profile: ASSERTIVE Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🔗 Linked repositories identifiedCodeRabbit considers these linked repositories for cross-repo context during reviews:
Included review availability: Your plan provides up to 12 included reviews per hour; 9 remain after this review. 📜 Recent review details⏰ Context from checks skipped due to timeout. (21)
🧰 Additional context used📓 Path-based instructions (3)Review scheduler state transitions and event ordering, including cancel, preempt, drain, retry, duplicate, empty, and terminal paths.⚙️ CodeRabbit configuration file Files:
Read REVIEW.md before commenting.⚙️ CodeRabbit configuration file Files:
Before making any change under: `python/aisimulate/src/aiconfigurator/generator/**` MUST read: `python/aisimulate/.claude/rules/generator-development.md` Before making any change under `python/aisimulate/collector/**` MUST read: `python/ais...📄 CodeRabbit inference engine (AGENTS.md) Files:
🔀 Multi-repo context ai-dynamo/dynamo, ai-dynamo/aiconfiguratorLinked repositories findingsai-dynamo/dynamo
ai-dynamo/aiconfigurator
🔇 Additional comments (5)
📝 SummaryRisk: Medium. Human review should focus on:
Behavior and contracts: Admission now starts without copying manager state. It checkpoints before mutation, shares unchanged radix-edge arrays, and records lease page counts rather than copying page IDs. Partial eviction retains only the necessary prefix. Evidence supplied: The change adds rollback tests for failed validation and partial-prefix matching. The author reports 591 passing engine tests, successful formatting and diff checks, and exact per-request and simulated-summary parity. The author reports current-head CPU median differences of -1.42%, -0.55%, and -0.30%; all paired intervals include zero, so these results do not establish a clear speedup. Separate historical measurements against the exact Evidence missing and readiness: Test and benchmark results have not been independently verified here. Current review findings and severity counts were not supplied. Technical review status and merge readiness remain undetermined. Bot review is not approval. WalkthroughRadix cache edges now share key and value vectors through ChangesSGLang admission rollback
Estimated code review effort: 3 (Moderate) | ~25 minutes Merge Risk: ⚪ Minimal · up to This change makes admission checkpoints cheaper by sharing radix-cache edge storage and taking the snapshot only when state may change. Rollback still restores the cache, leases, queue and events after a failed admission. No concrete correctness or data-integrity defect was established at the current head. The remaining notes are performance and robustness follow-ups, so the change looks ready to merge. 🚥 Pre-merge checks | ✅ 8✅ Passed checks (8 passed)
Comment |
jthomson04
left a comment
jthomson04
left a comment
There was a problem hiding this comment.
No correctness bug found. Inside the admission window, only allocate_for_request_lease_capped and extend_allocation change the cache, and both call checkpoint_admission() before their first change. The late snapshot therefore restores the same state the old up-front one did. All 10 admission_validation_rollback tests pass on this head.
Inline comments cover one latent rollback risk, a test gap, a small API point and some remaining copy costs. Two items point at lines outside the diff:
RadixRequestLease::admission_checkpoint(sglang_backend.rs:48) still cloneslease.pagesfor every waiting request on every pass. During admission, pages only grow:extend_allocationappends, and allocation fills an empty, inactive lease. Storingpages.len()and truncating on restore would be enough, aspage_hashes_lenalready does. Otherwise a long chunked continuation at the head of the queue (for example 100k pages at page_size 1) is copied on every pass, including the "nothing fits" passes this PR targets.- Design: after this change the cost per admitting pass is still proportional to the cache size, not to the number of changes admission makes. Taking the snapshot late and making edges copy-on-write lower the cost, but under load most passes admit something. Two ways to go further:
- An undo journal that records the few changes admission makes (split node ids, lock increments, evicted leaves/pages, allocated pages, event-id and
page_to_block_hashwrites) and replays them backwards on rollback. - Checking the planned
(tokens_computed, prefix_tokens)before changing anything.get_new_batch_prefillalready gets prefix lengths from the read-onlyprefix_match_hashes_len.
- An undo journal that records the few changes admission makes (split node ids, lock increments, evicted leaves/pages, allocated pages, event-id and
|
Implemented the local changes in 7b1f8bb. All 591 engine tests pass; moving the snapshot after prefix matching makes the new rollback test fail as intended. The lease-copy suggestion from the review body is also addressed: checkpoints save the page count, and rollback truncates appended pages. The existing continuation failure test covers restoration of existing page IDs, capacity, queue order and events. The undo journal and validation-before-mutation proposals are deferred. Both would change transaction planning beyond the local copy reductions. Validation alone also does not cover later timing prediction or duration-conversion failures. We retain the rollback boundary and add the small debug checks rather than redesigning it. Three alternating timing pairs per workload (18 runs) against the previous PR head show CPU median changes of -1.42% for AgentX with 8 sessions, -0.55% for 32 sessions, and -0.30% for the 32-chain decode control. All intervals include zero, so this does not establish a clear end-to-end speedup. No repeatable slowdown was observed. Every workload has exact per-request and summary parity. The full table, revisions and method are in the PR description. The individual inline replies describe the fixes and the remaining full-cache snapshot limit. Review threads are left open for confirmation. |
Harrilee
left a comment
Harrilee
left a comment
There was a problem hiding this comment.
Reviewed commit 7b1f8bb. Approve; no actionable correctness issues found.
The lazy checkpoint precedes prefix matching (including split/touch/lock), and all allocation/extension mutations inside admission are covered. Length-only lease rollback is valid because existing page IDs are not rewritten before commit. Shared radix edges preserve snapshot contents through splitting and LRU/Belady eviction; failure/retry tests exercise caller-visible rollback and event suppression.
Independent validation: cargo +1.96.1 test -p aisimulate-core --locked --lib engine:: — 591 passed, 0 failed; Rust formatting and diff checks passed. Full CI run 35889924285 succeeded on this exact SHA.
I did not rerun AgentX benchmarks. The historical 2.82x result is a single capture-enabled comparison, as disclosed; the current-head matched campaign measures additional review changes and does not establish a further speedup. Approval does not turn that historical figure into a new independently reproduced result.
7b1f8bb to
3d32556
Compare
|
/ok to test 3d32556 |
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
3d32556 to
294e3a1
Compare
|
/ok to test 294e3a1 |
2 participants
Summary
The broader SGLang admission rollback checks from #158 copy resident radix edges and waiting leases before timing validation. Large AgentX traces spend substantial host CPU time copying unchanged cache data.
Delay the cache snapshot until the first mutation and share unchanged radix edge arrays. Lease checkpoints save page counts; shared-edge splits copy each half once, and partial eviction copies only the retained prefix. Unique edges retain their existing mutation path. When prefix caching is disabled, rejected allocations do not need a cache snapshot.
Review map
sglang_backend.rs, then shared-edge mutations inradix_cache.rs.Validation
Tested commit:
7b1f8bb04e365122add1adaee9c601f9f10bcaa7, based on main2e01b14a543889f96639a4b0bb4c1d17523b545b.cargo +1.96.1 test -p aisimulate-core --locked --lib engine::: 591 passed.cargo +1.96.1 fmt --all -- --checkandgit diff --check: pass.7b1f8bb04e365122add1adaee9c601f9f10bcaa7. Hosted checks and external reviews must be read at this same head; local results do not imply hosted CI success.Additional review changes: matched benchmarks
Compared
fff607b35f468928ca89f29daf2b6b9c9c0c5304with7b1f8bb04e365122add1adaee9c601f9f10bcaa7: three alternating pairs per workload, 18 timed runs total. Negative CPU change means less host CPU time.All three workloads have exact per-request and simulated-summary parity, excluding generated UUIDs and three host-timing fields. All 18 timed runs completed the expected request/token counts and matched simulated summaries. No runs were excluded or replaced; no overlapping builds, profiles or simulation jobs were detected.
The differences are small and all intervals include zero: this campaign does not establish a clear end-to-end speedup. No repeatable slowdown was observed.
Point estimates are ratios of CPU medians. Intervals are nominal paired Student-t intervals on log ratios (paired geometric means); three pairs have limited precision. These results measure the combined review changes, not the isolated benefit of any one optimization. Qualification runs are separate from timing estimates.
Original #158 regression evidence
Historical measurements on #158's exact merge (
8772127ba718239cecf26d80f59d49c67bb9b058), with and without the original lazy/shared-edge patch:The 8-session row compares five original and three fixed timing runs. The 32-session row compares single capture-enabled runs and has lower precision. These older measurements are separate from the current-head comparison above. Session count is the trace subset, not fixed request concurrency. A separate three-pair comparison against pre-#158 retained +3.46% CPU for 8 sessions, +5.42% for 32 sessions and +0.71% for the decode control.
Modeling and measurement provenance
Host-cost measurements use the public
aisimulate predict --stack engineCLI: Qwen/Qwen3.8-2.4T-A95B-FP8, GB300, SGLang 0.5.17, AIC op-level timing, prefix caching and trace arrival times. Each new comparison uses the same inputs, dependency locks and release flags: Python 3.12.12, Rust 1.96.1, x86-64-v3, thin LTO, one codegen unit and frame pointers. Intel Core Ultra 9 285K; replay on CPU 4, observer on CPU 0. Separate native extensions and their hashes were verified before running.These results measure simulator CPU/wall time, not predicted GPU or serving performance. Related: #158.