Repository navigation
Gate mamba extra-buffer predicates on uses_mamba_radix_cache - #37474
Merged
ispobock merged 7 commits intoSep 17, 2026
Merged
Conversation
An explicitly passed --mamba-radix-cache-strategy=extra_buffer steered non-mamba models into the mamba scheduler tracking paths, crashing on the first prefill batch (req.mamba_ping_pong_track_buffer is None). Require the architecture-derived uses_mamba_radix_cache leaf in the extra-buffer predicates so the strategy stays inert for models without mamba state, and resolve the leaf in the Inkling overrides (which pin extra_buffer outside the linear-attn registry). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012yzAauSvQdQD835r5feCFw
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: 170d14ba09
ℹ️ 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".
| # The extra-buffer predicates also require the arch-derived leaf, which the | ||
| # generic resolution likewise never sets for Inkling. | ||
| if not cfg.disable_radix_cache: | ||
| overrides["uses_mamba_radix_cache"] = True |
There was a problem hiding this comment.
With radix caching enabled (the normal Inkling configuration), this marker makes the later unconditional handle_mamba_radix_cache() call treat Inkling as a generic mamba-cache architecture and invoke validate_mamba_extra_buffer(). That validator asserts supports_mamba_cache_extra_buffer(view, model_arch), but neither Inkling architecture is in _MAMBA_EXTRA_BUFFER_ARCHS, so default Inkling startup now fails during argument resolution with extra_buffer is not supported for InklingForConditionalGeneration. Either register Inkling as supported by that validator or avoid routing this model-specific marker through the generic validation path.
Useful? React with 👍 / 👎.
With uses_mamba_radix_cache now resolved for Inkling, the unconditional mamba radix-cache validation slot runs _validate_mamba_extra_buffer for these archs; supports_mamba_cache_extra_buffer must accept them or default Inkling startup fails during argument resolution. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012yzAauSvQdQD835r5feCFw
ispobock
added
run-ci
Collaborator
|
@sshleifer could you fix the conflict? |
Collaborator
|
Since |
Collaborator
|
alternative to #39430 |
Instead of gating the extra-buffer predicates on uses_mamba_radix_cache, handle_mamba_radix_cache raises when an explicit extra_buffer strategy reaches a model without mamba state. Predicates and test fixtures return to main. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Collaborator
|
/rerun-failed-ci |
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 17, 2026
[Unified Cache][9/N] add opt-in MLA load deduplication for Mooncake Linker. Squashes 5 upstream commits (cc4431b, 455f5bb, 310b76e, 8acfd72, d52e749 — three of them merge-main commits whose only net effect on this PR's files is what lands below). New opt-in flag --enable-linker-mla-dedup: with the Mooncake external linker, rank 0 loads replicated MLA KV from L3 and broadcasts each layer to the other TP ranks instead of every rank reading the same objects. Files: - python/sglang/srt/arg_groups/fields/memory.py: add enable_linker_mla_dedup. - python/sglang/srt/arg_groups/hicache_hook.py: require the Mooncake linker. - python/sglang/srt/mem_cache/unified_cache/linker_mla_dedup.py (new): LinkerMLADedupBroadcaster, adapting hybrid device-pool geometry to the existing MLAHostDedupBroadcaster primitive. - python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_direct_linker.py: LayerWiseLoadCounter on_layer_ready hook; broadcaster construction; per-batch logical-plan all_gather gate; rank-0-only load queueing; forward-stream broadcasts; broadcast-completion event tracking in num_completed_loads; reset/close teardown. - test/registered/unit/server_args/test_server_args.py: flag validation test. Manual merge notes: - Our fork's mooncake_direct_linker.py carries fork-local behavior (fail-soft revalidate_load, retrying session start, SGLANG_LINKER_DEBUG_KEY probes, idx-checksum probe, rank-key handling). All of it is preserved untouched; the upstream hunks apply around it (imports, LayerWiseLoadCounter, __init__ broadcaster wiring, num_completed_loads, start_layer_wise_loading, load_thread_func, reset/close). - The upstream branch was based on main before sgl-project#39835/sgl-project#39115/sgl-project#37474; those are unrelated to this PR's files (verified by diffing the PR range against the merge base), so the squash is a faithful net port. - Depends on MLAHostDedupBroadcaster and attn_tp_cache_group/tp_cache_group in CacheInitParams, both already present in our base (upstream parent == ours).
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 17, 2026
[Unified Cache][9/N] add opt-in MLA load deduplication for Mooncake Linker. Squashes 5 upstream commits (cc4431b, 455f5bb, 310b76e, 8acfd72, d52e749 — three of them merge-main commits whose only net effect on this PR's files is what lands below). New opt-in flag --enable-linker-mla-dedup: with the Mooncake external linker, rank 0 loads replicated MLA KV from L3 and broadcasts each layer to the other TP ranks instead of every rank reading the same objects. Files: - python/sglang/srt/arg_groups/fields/memory.py: add enable_linker_mla_dedup. - python/sglang/srt/arg_groups/hicache_hook.py: require the Mooncake linker. - python/sglang/srt/mem_cache/unified_cache/linker_mla_dedup.py (new): LinkerMLADedupBroadcaster, adapting hybrid device-pool geometry to the existing MLAHostDedupBroadcaster primitive. - python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_direct_linker.py: LayerWiseLoadCounter on_layer_ready hook; broadcaster construction; per-batch logical-plan all_gather gate; rank-0-only load queueing; forward-stream broadcasts; broadcast-completion event tracking in num_completed_loads; reset/close teardown. - test/registered/unit/server_args/test_server_args.py: flag validation test. Manual merge notes: - Our fork's mooncake_direct_linker.py carries fork-local behavior (fail-soft revalidate_load, retrying session start, SGLANG_LINKER_DEBUG_KEY probes, idx-checksum probe, rank-key handling). All of it is preserved untouched; the upstream hunks apply around it (imports, LayerWiseLoadCounter, __init__ broadcaster wiring, num_completed_loads, start_layer_wise_loading, load_thread_func, reset/close). - The upstream branch was based on main before sgl-project#39835/sgl-project#39115/sgl-project#37474; those are unrelated to this PR's files (verified by diffing the PR range against the merge base), so the squash is a faithful net port. - Depends on MLAHostDedupBroadcaster and attn_tp_cache_group/tp_cache_group in CacheInitParams, both already present in our base (upstream parent == ours).
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 18, 2026
[Unified Cache][9/N] add opt-in MLA load deduplication for Mooncake Linker. Squashes 5 upstream commits (cc4431b, 455f5bb, 310b76e, 8acfd72, d52e749 — three of them merge-main commits whose only net effect on this PR's files is what lands below). New opt-in flag --enable-linker-mla-dedup: with the Mooncake external linker, rank 0 loads replicated MLA KV from L3 and broadcasts each layer to the other TP ranks instead of every rank reading the same objects. Files: - python/sglang/srt/arg_groups/fields/memory.py: add enable_linker_mla_dedup. - python/sglang/srt/arg_groups/hicache_hook.py: require the Mooncake linker. - python/sglang/srt/mem_cache/unified_cache/linker_mla_dedup.py (new): LinkerMLADedupBroadcaster, adapting hybrid device-pool geometry to the existing MLAHostDedupBroadcaster primitive. - python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_direct_linker.py: LayerWiseLoadCounter on_layer_ready hook; broadcaster construction; per-batch logical-plan all_gather gate; rank-0-only load queueing; forward-stream broadcasts; broadcast-completion event tracking in num_completed_loads; reset/close teardown. - test/registered/unit/server_args/test_server_args.py: flag validation test. Manual merge notes: - Our fork's mooncake_direct_linker.py carries fork-local behavior (fail-soft revalidate_load, retrying session start, SGLANG_LINKER_DEBUG_KEY probes, idx-checksum probe, rank-key handling). All of it is preserved untouched; the upstream hunks apply around it (imports, LayerWiseLoadCounter, __init__ broadcaster wiring, num_completed_loads, start_layer_wise_loading, load_thread_func, reset/close). - The upstream branch was based on main before sgl-project#39835/sgl-project#39115/sgl-project#37474; those are unrelated to this PR's files (verified by diffing the PR range against the merge base), so the squash is a faithful net port. - Depends on MLAHostDedupBroadcaster and attn_tp_cache_group/tp_cache_group in CacheInitParams, both already present in our base (upstream parent == ours).
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
An explicitly passed
--mamba-radix-cache-strategy extra_bufferon a model without mamba state steers the scheduler into the mamba-only tracking paths and crashes on the first prefill batch:autois already arch-gated, so only an explicit misconfiguration reaches this state. We hit it when a shared test launcher started passing the strategy for every model in its zoo.Modifications
handle_mamba_radix_cache(the unconditional call slot) fails fast when the resolved view has nouses_mamba_radix_cachebut the extra-buffer predicate is armed:--mamba-radix-cache-strategy extra_buffer needs mamba state, got <arch>.The predicates themselves are unchanged. With the radix cache disabled the flag is genuinely inert everywhere, so that combination is left alone.extra_bufferthrough its model override rather than the linear-attn registry, so the override now also declaresuses_mamba_radix_cache=True; model overrides are collected before the unconditional slot runs, so an explicitextra_bufferon Inkling does not trip the check. With the leaf declared, the validation slot runs for Inkling, so_MAMBA_EXTRA_BUFFER_ARCHSaccepts the Inkling archs.Tests
test/registered/unit/test_model_overrides.py:LlamaForCausalLMwith explicitextra_bufferraises with the message above; the same with--disable-radix-cacheresolves without error. Existing mamba fixtures and the runtime-context parity test are untouched from main.-Robot
CI States
Latest PR Test (Base): ❌ Run #35190739916
Latest PR Test (Extra): ❌ Run #35190740005
Latest PR Test (AMD ROCm 10): ⏳ Run #35190739931