[UnifiedTree]: Fix UnifiedRadixCache device match semantics with HiCache - #25277
Merged
ispobock merged 5 commits intoMay 15, 2026
Merged
Conversation
hzh0425
requested review from
Ying1123,
hanming-lu,
hnyls2002,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
May 14, 2026 11:39
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the UnifiedRadixCache to support HiCache by distinguishing between host-backed and device-resident matches. Key changes include updating component match validators with a match_device_only flag, tracking separate device anchors during prefix matching, removing the unused readonly match helper, and refactoring init_load_back to gather prefix indices after a load operation. Review feedback suggests optimizing the _all_valid helper with a generator expression, correcting a return type hint, and removing dead code in the init_load_back logic.
hzh0425
force-pushed
the
hybrid_tree/fix-last-device-node-define
branch
from
May 14, 2026 11:55
c825ed0 to
327fcb3
Compare
Collaborator
Author
Collaborator
Author
|
/rerun-test test/registered/radix_cache/test_unified_radix_hicache_kl.py |
Contributor
|
🚀 |
ispobock
approved these changes
May 15, 2026
Collaborator
Author
Almost all CI jobs passed in the last run.
|
Fridge003
pushed a commit
that referenced
this pull request
May 15, 2026
SpencerGarnets
added a commit
to ai-blaise/optimization-playground
that referenced
this pull request
May 16, 2026
Merging upstream sgl-project/sglang main (d523ae1, 52 commits) into ai-blaise/optimization-playground main, preserving the full custom stack (IndexCache + TurboQuant 2.5-bit + SMC-SD, NVFP4 IndexCache+HISA 4:1, HIGGS dense 2-bit MLA KV, HF config dispatcher with indexcache-hisa, LayerSplit, NCCLX collective, Warp Decode, TokenSpeed/SMC/indexer server_args defaults, W4A4 MoE WarpDecode). Conflict resolutions (3 files): 1. python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py COMPOSE - kept custom override_kv_cache_dim=compressed_mla_host_dim(kv) semantics; adopted upstream function rename (build_shared_anchor_stack -> build_anchor_sidecar_stack from PR sgl-project#25277 UnifiedRadixCache device match semantics). The renamed function still accepts override_kv_cache_dim (verified at hybrid_pool_assembler.py:583), so this is backwards compatible. The full call site at L946 already feeds the override into the kwarg. 2. sgl-kernel/cmake/flashmla.cmake TAKE UPSTREAM - upstream PR sgl-project#24130 (fix sm90 compile flashmla failed) renamed the FlashMLA-internal SM100 gate from SGL_KERNEL_ENABLE_SM100A to FLASHMLA_ENABLE_SM100 and re-introduced the SM100 source list under the new macro. HEADs older SGL_KERNEL_ENABLE_SM100A source-list block was a duplicate of upstreams new FLASHMLA_ENABLE_SM100 block; dropped the duplicate to avoid CMake emitting the same source files twice. The outer SGL_KERNEL_ENABLE_SM100A cmake option in CMakeLists.txt is preserved as it gates additional sgl-kernel targets beyond FlashMLA. 3. sgl-kernel/csrc/flashmla_extension.cc TAKE UPSTREAM - upstream PR sgl-project#24130 switched the FlashMLA-internal #ifdef gate at L94 from SGL_KERNEL_ENABLE_SM100A to FLASHMLA_ENABLE_SM100. flashmla.cmake (post-resolution 2) emits target_compile_definitions PRIVATE FLASHMLA_ENABLE_SM100 at L173, so the new macro is the one actually defined at FlashMLA extension compile time. The unrelated #ifndef SGL_KERNEL_ENABLE_SM100A at L21 (different code path) is left untouched. Verification gates: GATE 1 git diff --check PASS GATE 2 no conflict markers PASS GATE 3 python3.12 -m compileall srt/jit_kernel/ci PASS GATE 4 cargo check sgl-model-gateway -j1 PASS GATE 5 test_quantization_config_dispatch (33 tests) PASS 33/33 GATE 6 test_nsa_layersplit (26 tests) PASS 26/26 GATE 7 test_higgs_dense_2bit_kv unit BLOCKED (GPU saturation: 8x H200 fully occupied by sglang::scheduler_TP0-7 from another job; CUDA error: devices busy or unavailable. File compiles cleanly, imports succeed, not merge-induced.) GATE 8 test_higgs_dense_2bit_kv_integration BLOCKED (same) GATE 9 test_torchcomms_ncclx CPU subset PASS 4/4 (remaining 8 are CUDA-only, expected skip on busy GPU) GATE 10 focused custom-stack CPU subset PASS GPU-dependent tests (HiCache transfer, pool host integration) BLOCKED on same GPU saturation - reproduced identical cudaErrorDevicesUnavailable across all GPU-touching tests. Custom-stack file integrity check (all 10 sentinel files present): layersplit.py, higgs_dense_2bit_kv.cuh, higgs_dense_2bit.py, torchcomms_ncclx.py, nvfp4_indexer_quant.cuh, nvfp4_indexer.py, hisa_tilelang_kernels/hisa.py, indexer_policy.py, indexer_quantization.py, turboquant_dense_kv.cuh. Custom server_args defaults preserved (verified by grep): tokenspeed_mla, nsa_indexer_mode, nsa_indexcache_freq, enable_nsa_nvfp4_hisa, enable_turboquant_dense_kv_cache, enable_higgs_dense_2bit_kv_cache, enable_hisparse, torchcomms_ncclx_*, enable_return_indexer_topk, enable_nsa_prefill_context_parallel. Upstream additions accepted into shared infra (no custom regression): - Multi-detokenizer worker support (PR sgl-project#24944, server_args.py). - SGLANG_OPT_FP8_WO_A_GEMM auto-disable on non-Blackwell (PR sgl-project#25181, server_args.py). - SpeculativeAlgorithm.supports_target_verify_for_draft + create_future_map extension points (PR sgl-project#24999, spec_info.py). - Frozen-KV MTP bonus_tokens=None fix (PR sgl-project#25204, frozen_kv_mtp_utils.py - uses existing draft_input.input_ids fallback; spec-decoding naming rule still satisfied). - HiCache UnifiedRadixCache device match semantics (PR sgl-project#25277, hybrid_pool_assembler function rename). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This was referenced May 19, 2026
5 tasks
Chronostasys
pushed a commit
to MindLab-Research/sglang
that referenced
this pull request
Aug 24, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This PR clarifies UnifiedRadixCache match semantics when HiCache is enabled:
best_match_nodefrom the device-resident match anchor used aslast_device_nodelast_device_nodeas the deepest node whose active components are all device-residentinit_load_backtests to validate prefix index delta collection, aux-only load-back, and fallback behaviorModifications
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test: Run #25925574846⚠️ Not enabled — add
Latest PR Test (Extra):
run-ci-extralabel to opt in.