[Unified Tree] Port SWA Branching-Point Caching to the Rust TreeCore - #37584
Conversation
81d36f5 to
58a9d2f
Compare
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 58a9d2fd8a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
58a9d2f to
be0fbe9
Compare
be0fbe9 to
d6d7aea
Compare
d6d7aea to
5bf1f81
Compare
|
/rerun-test registered/unit/mem_cache/test_rust_tree_core_integration.py |
|
Results for 🚀 |
|
/rerun-test -c |
|
Results for 🚀 🚀 |
…hing # Conflicts: # test/registered/unit/mem_cache/test_rust_tree_core_integration.py
…hing # Conflicts: # rust/sglang-radix-tree/src/tests/components/swa.rs
…ment lock protocol
ispobock
left a comment
ispobock
left a comment
There was a problem hiding this comment.
next todo is to align the swa checkpoint with the mamba checkpoint
3 participants
Motivation
#34565 added SWA branching-point caching entirely in Python. The Rust TreeCore
(
SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND=rust) is a parallel implementation of thesame interface, not a binding over the Python one, so it did not inherit any of it:
rust/sglang-radix-tree/src/components/swa.rsdid a single-nodeBackupHostwith no dirty window and no
nodes_to_loadSwaComponenthad noneeds_incremental_backupoverride, so it used the traitdefault
falseMatchResult.swa_branching_seqlen,InsertParams.swa_branching_seqlenandInsertResult.swa_branch_inserteddid not exist anywhere in the crateThe adapter builds its result objects from an explicit field list, so the missing
fields did not raise; they defaulted to
None/False. On the Rust backend #34565 wastherefore a silent no-op: no crash, no log, no branching-point caching. The e2e
below shows exactly that on
main(armmain_rust) and that this PR closes the gap.Changes
Contract fields, threaded core -> bindings -> adapter the way
mamba_branching_seqlenalready is:
MatchResult.swa_branching_seqlenInsertParams.swa_branching_seqlen(carried across the resumable insert's walk state)InsertResult.swa_branch_insertedcomponents/swa.rs:collect_unbacked_swa_nodes_in_window_: the device-only nodes within one slidingwindow, stopping at a node an in-flight backup already owns, so two acks never claim
the same node
needs_incremental_backupoverrides the trait default with!unbacked.is_empty()finalize_match_result_in_tree_corecomputes the page-aligned branching pointcommit_insert_component_datareports whether the insert reached that boundaryBackupHostbuild emits the whole window ancestors-first withnodes_to_load; commitscatters the acked host span back across those nodes by device length
Write-back: #34565 dropped
not self.is_write_backfrom_should_backup_after_insertand restricted the write-back incremental check to SWA (Full and Mamba defer to eviction,
SWA cannot because out-of-window frees never reach eviction). The Rust core still carried
the old
!self.is_write_backgate, so under write-back a reconstructed SWA window wasnever published and its device slots could be freed with the host tier still tombstoned.
should_backup_after_insert_/needs_incremental_component_backup_now mirror thePython core; two native tests pin both halves (SWA backs up, Mamba still defers).
Buffer mode: #34565's SWA component reads
self.cache.host_memory_mode, which the Rust component cannot; it has no cache handle.The flag moves onto the tree core (
set_host_memory_buffer_only()plus a read-onlyis_host_memory_buffer_onlyon both backends) and both components, plus the Pythonprefetch path, now read the same source. Without this,
buffer_mode/pipeline.pysizesits storage keys off a windowed transfer and writes the wrong number of pages.
Tests. #34565 gated three shared-suite tests behind
_skip_swa_branching_on_rust.The gate is gone; the tests now drive eviction and the write-through pending id through
the backend-neutral test inspector (
evict_component,get_write_through_pending_id)instead of
resolve_node_handle()/node_by_id, so the same assertions run on bothcores. One test in
test_rust_tree_core_integration.pypins theswa_branch_inserted/swa_branching_seqlenfield threading through the PyO3 binding; the window and scatterbehavior is covered by the native tests and the shared suite.
Validation
All on one H200 devbox (
lmsysorg/sglang:dev), this branch rebased onmainat 756d0e0.cargo testtest_rust_tree_core_integration.pytest_unified_radix_cache_unittest.py, Python coretest_rust_unified_radix_cache_unittest.py, Rust coreThe skip delta of four is exactly the three un-gated tests on their matching fixtures
(
Test_FULL_SWA_ps1_sw4x3,Test_FULL_SWA_ps4_sw4x1), identical to what the Pythonentry executes.
pre-commit(ruff, clippy-D warnings, rustfmt) passes.E2E: the branching fingerprint on DSV4-Flash
DeepSeek-V4-Flash-0731, 4xH200, DSPARK,
--chunked-prefill-size 16384,--swa-full-tokens-ratio 0.1,SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS=True;same box, same command, four arms = {
main756d0e0, this PR} x{
SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND=python,rust}. Each arm's log confirms thetree core it ran (
Tree Core: UnifiedTreeCore/RustUnifiedTreeCore).Probe (
/generatewithinput_ids,max_new_tokens=1, readmeta_info.cached_tokens):A seeds a 24576-token shared prefix followed by a 16384-token question (40960 tokens = 3
prefill chunks, so the prefix end sits in an interior chunk and its SWA gets freed);
B, C, D reuse the prefix with fresh questions. Two groups per arm, results identical
across groups.
main, pythonmain, rustB = 16384 then C/D = 24576 is the branching fingerprint: B hits Full for the whole prefix,
gets a branching point, and its insert re-publishes the SWA window; C and D then hit the
full 24576. On
mainthe Rust core never gets past 16384 (the silent no-op from theMotivation section). With this PR the Rust core matches the Python core byte for byte,
and the Python core has the same behavior for this Full+SWA configuration.
Shared-prefix bench (
generated-shared-prefix, 8 groups x 8 prompts,system 24576 / question 16384 / output 128, request rate 2, concurrency 8; two passes
per arm, warm pass reported, cold pass is JIT-confounded):
main, rust(
main, python ran first on the box and its warm pass still carried first-arm JIT cost,46.7% / 0.64 req/s; its hit rate matches the PR arms, which is the number that matters.)
Zero
CUDA out of memory/illegal memory/ scheduler exceptions in any arm.Hybrid SWA/Mamba checkpoint correctness
Fixes #38815. In hybrid caches,
SWAComponent.prepare_for_caching_reqleaves the insertion length to Mamba after preserving the SWA eviction boundary. This keeps each checkpoint attached to its actual prefix depth. Full+SWA caches without Mamba retain branching-point caching.An ad hoc E2E experiment ran real Inkling requests with Triton, page size 1, and the Rust TreeCore. Checkpoints at 64/128, a branch at 96, and prefill to 192 previously attached state192 to key96. A subsequent probe reused 96 tokens and could change greedy output. The guarded insertion stores key192/state192 and the probe correctly resumes at checkpoint64. Across three input corpora, warm output token IDs and logprobs now match cold-cache results exactly.
Validation for the guard: the real-model experiment passed on H200 and failed all three subcases with the guard removed (
96 != 64); both Python and Rust shared-suite-k swa_branchruns pass (76 discovered, 5 executed, 71 skipped per backend); pre-commit passed for the guard and experiment. The experiment is retained as reproduction evidence in #38815; this PR does not add it to permanent CI or change the Inkling CI time estimate. The broader historical suite and DSV4 results above predate this guard.CI States
Latest PR Test (Base): ✅ Run #34507885445
Latest PR Test (Extra): ❌ Run #34507885297
Latest PR Test (AMD ROCm 10): ❌ Run #34507885442