Skip to content

[Mamba] Fix checkpoint depth for prefixes that end off the radix page - #39115

Merged
ispobock merged 4 commits into
sgl-project:mainfrom
klshuster:kurt/mamba-checkpoint-depth-dcp-page
Sep 17, 2026
Merged

[Mamba] Fix checkpoint depth for prefixes that end off the radix page#39115
ispobock merged 4 commits into
sgl-project:mainfrom
klshuster:kurt/mamba-checkpoint-depth-dcp-page

Conversation

@klshuster

@klshuster klshuster commented Sep 11, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Contributor

Motivation

With DCP the radix-tree page is wider than the scheduler page (e.g. 512 vs 64), so a prefill chunk can leave prefix_indices at 64 with no node inserted. The next extend picks the Mamba checkpoint depth relative to that prefix, 64 + 512 = 576, a position no 512-token page can hold. On finish the two freed kv-row tails [512, 576) and [576, 605) share a page and the scheduler dies:

AssertionError: segment at 576 shares a page with the one ending at 576

Without the assertion the node would carry a Mamba state 64 tokens past its end. Reproduced on Kimi-K3 with --dcp-size 8 and radix cache on.

Modifications

_mamba_radix_cache_v2_req_prepare_for_extend picks the depth on the absolute checkpoint grid and donates only where an h snapshot exists ((depth - prefix_len) % cache_chunk_size == 0).

-        mask = req.extend_range.length >= checkpoint_grid
+        prefix_len = len(req.prefix_indices)
+        seq_end = prefix_len + req.extend_range.length
+        mamba_track_seqlen_aligned = (seq_end // checkpoint_grid) * checkpoint_grid
+        mask = (
+            mamba_track_seqlen_aligned > prefix_len
+            and (mamba_track_seqlen_aligned - prefix_len) % cache_chunk_size == 0
+        )

Accuracy Tests

test/registered/unit/managers/test_mamba_checkpoint_depth.py gains one table-driven case (prefix off the widened page, short extend crossing a page, extend crossing no page, prefix off the kernel chunk grid). Three of its four rows fail on main and pass with this change. Kimi-K3 dcp=8 agentic serving with the radix cache runs to completion instead of crashing.

Speed Tests and Profiling

Not affected: a few scalar ops per request at batch preparation.

Checklist

CI States

Latest PR Test (Base): 🚫 Run #35121066591
Latest PR Test (Extra): 🚫 Run #35121066143
Latest PR Test (AMD ROCm 10): ❌ Run #35121066600

@klshuster
[Mamba] Pick the radix checkpoint depth on the absolute grid when the…
07d996e 
… prefix is off the tree page

With DCP the allocator and radix-tree page widen (e.g. 64 -> 512) while
PrefillAdder still aligns chunk budgets to the scheduler page, so a prefill
chunk can end at 64. `cache_unfinished_req` inserts nothing for that chunk
(no Mamba checkpoint yet, key shorter than a page) and leaves `prefix_indices`
at 64. The next extend then measures the checkpoint depth relative to the
prefix: 64 + floor(540 / 512) * 512 = 576, a position no 512-token page can
name. On finish the Mamba component truncates the insert to 576 while the radix
key is page-floored to 512, and the two freed kv-row tails [512, 576) and
[576, 605) share a page (`_page_disjoint` assertion); without the assertion the
tree would hold a 576-token Mamba state on a 512-token node.

Pick the depth on the absolute checkpoint grid instead and donate only where
an h snapshot exists ((depth - prefix) % cache_chunk_size == 0). Extends that
cross no page boundary, or whose prefix is off the kernel chunk grid, donate
nothing.

Unit test: one table-driven case; three of its four rows fail before this change.
@ispobock ispobock self-assigned this Sep 12, 2026
@ispobock ispobock added run-ci CI: run the baseline test suite on this PR bypass-fastfail run-ci-extra CI: also run the extra suite (requires run-ci) labels Sep 12, 2026
@ispobock

ispobock commented Sep 12, 2026

Copy link
Copy Markdown
Collaborator

@codex review

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 12, 2026 •
edited
Loading

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-09-12T07:09:16.984695Z c92e8f8 Manual request
ℹ️ 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" or "@codex security review".

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

chatgpt-codex-connector Bot commented Sep 12, 2026

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Can't wait for the next one!

Reviewed commit: c92e8f8e82

ℹ️ 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".

@ispobock

ispobock commented Sep 13, 2026

Copy link
Copy Markdown
Collaborator

/rerun-failed-ci

hanming-lu
hanming-lu reviewed Sep 14, 2026
View reviewed changes
Comment thread python/sglang/srt/managers/schedule_batch.py Outdated
mamba_track_seqlen = len(req.prefix_indices) + req.extend_range.length
mamba_track_seqlen = seq_end

# mamba_track_seqlen_aligned/mamba_last_track_seqlen is actual tracked seqlen. Used to pass to

@hanming-lu hanming-lu Sep 14, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's update the comments as well.

@ispobock ispobock Sep 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @klshuster

@ispobock
ispobock merged commit 6460082 into sgl-project:main Sep 17, 2026
180 of 200 checks passed
Leoyzen added a commit to Leoyzen/sglang that referenced this pull request Sep 17, 2026
@Leoyzen
Port of upstream PR sgl-project#39565 (head sha d52e749)
914d982 
[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
@Leoyzen
Port of upstream PR sgl-project#39565 (head sha d52e749)
3298b31 
[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
@Leoyzen
Port of upstream PR sgl-project#39565 (head sha d52e749)
2cb8b6a 
[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).
@curnane-lab curnane-lab mentioned this pull request Sep 23, 2026
Closed
@yichiche yichiche mentioned this pull request Sep 23, 2026
Merged
4 of 5 tasks
@chien-an-chen chien-an-chen mentioned this pull request Sep 29, 2026
Open
4 of 5 tasks
@bingxche bingxche mentioned this pull request Oct 10, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ispobock ispobock ispobock left review comments
@hanming-lu hanming-lu hanming-lu left review comments
@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner
@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner
@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner
@xiezhq-hermann xiezhq-hermann Awaiting requested review from xiezhq-hermann xiezhq-hermann is a code owner

Assignees

@ispobock ispobock

@hanming-lu hanming-lu

Labels

bypass-fastfail run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci)

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@klshuster @ispobock @hanming-lu