Repository navigation
Declare hybrid Mamba indexer host pool - #40915
Merged
Merged
Conversation
ispobock
requested review from
Ying1123,
alphabetc1,
hanming-lu,
hnyls2002,
huangtingwei9988,
hzh0425,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
September 23, 2026 13:04
This was referenced Sep 23, 2026
ispobock
force-pushed
the
hicache-decl-b
branch
from
September 23, 2026 13:25
1270033 to
4bcb754
Compare
ispobock
force-pushed
the
hicache-decl-c
branch
from
September 23, 2026 13:25
1ad08dd to
3d9a643
Compare
| is_anchor=True, | ||
| packed_draft_device_pools=mtp_draft_device_pools, | ||
| ), | ||
| entries = _build_declared_entries(configs, root_host_pool=kv_host_pool) + [ |
Collaborator
There was a problem hiding this comment.
Question on the INDEXER entry for kpool models: it copies index rows by KV page, but with index_kpool=4 the writer stores the compressed rows for all four 64-token pages of a 256-token group in the group's first page (token_page_row = pool_page_group * POOL_SIZE,
). mamba_checkpoint_grid is lcm(64, 64) = 64 here, so a parent node can end at 192. With write_through it is backed up once while page 0 holds rows 0-47; a child then writes rows for 192-255 into slots 48-63 of that same page, and those slots are never re-backed-up. A host restore past 192 would read stale index rows.Is something keeping radix/checkpoint boundaries 256-aligned for kpool models that I missed? If not, this probably needs the page-local layout from #40134 (or a 256 tree page as in #38212) before or alongside this PR; the shared-page layout was also raised in #39156.
2 tasks done
ispobock
force-pushed
the
hicache-decl-b
branch
from
October 1, 2026 08:55
4bcb754 to
5730ea4
Compare
ispobock
force-pushed
the
hicache-decl-c
branch
from
October 1, 2026 13:33
3d9a643 to
ed6e707
Compare
ispobock
requested review from
ByronHsu,
Duyi-Wang,
Fridge003,
HaiShaw,
Qiaolin-Yu,
ShangmingCai and
hebiao064
as code owners
October 1, 2026 13:33
ispobock
requested review from
BBuf,
DarkSharpness,
HydraQYH,
celve,
sogalin and
yuan-luo
as code owners
October 1, 2026 13:33
github-actions
Bot
added
documentation
ispobock
added
run-ci
This was referenced Oct 4, 2026
# Conflicts: # python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
|
Thank you for getting this merged. Looking forward to the next rebuild for GLM 5.3 Flash. |
Open
2 tasks done
ormandj
added a commit
to ormandj/sglang
that referenced
this pull request
Oct 10, 2026
…s separate Rebuilt on the declared host pools from sgl-project#40913 and sgl-project#40915, which now restore the target and draft DSA index buffers. What main does not cover is kept: - Hybrid Mamba stacks scale the explicit KV host budget for declared dependent pools (the DSA indexer), as they already do for QSA indexes. - DSAIndexerPoolHost.destroy unregisters its index buffer, which the inherited cleanup of kv_buffer misses. - An MTP draft whose MLA KV rows differ from the target's (for example a separate draft KV dtype) uses the separate draft sidecar instead of failing the packed-row check; several mismatched runners are rejected because the sidecar keeps one.
This was referenced Oct 13, 2026
3 of 12 tasks
ormandj
added a commit
to ormandj/sglang
that referenced
this pull request
Oct 13, 2026
…s separate Rebuilt on the declared host pools from sgl-project#40913 and sgl-project#40915, which now restore the target and draft DSA index buffers. What main does not cover is kept: - Hybrid Mamba stacks scale the explicit KV host budget for declared dependent pools (the DSA indexer), as they already do for QSA indexes. - An MTP draft whose MLA KV rows differ from the target's (for example a separate draft KV dtype) uses the separate draft sidecar instead of failing the packed-row check; several mismatched runners are rejected because the sidecar keeps one. Separate draft recurrent state still requires matching rows, because only packed MTP backs it up.
5 tasks done
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
3 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
Part 3 of 5 of #40608. Hybrid Mamba plus DSA models (GLM-5.3-Flash) restore KV from host without the indexer on main, so a host hit runs attention on stale block selection. Host restore KL against warm is 0.06 to 0.71 on main. build_hybrid_mamba_stack now assembles the full attention pool's declarations next to the Mamba host pool, and _MambaStrategy raises when a declared pool is missing. HybridLinearKVPool.host_pool_decls() declares only the buffers of its full attention sub pool. Mamba state stays with its owner and the existing Mamba host path.
The DSA indexer is already declared by its pool since #40913, but each strategy assembled its own host pools and the Mamba path never added it. With declarations the pool owner lists its host buffers once, and every strategy holding that pool builds them from the same list instead of repeating a per-strategy branch. A strategy that drops a declared pool now fails at startup instead of restoring stale state. On main the same check only logs
declares host pools ['indexer'] that _MambaStrategy did not assemble, which is how this gap showed up.Stack: #40913 (declarations and DSA target), #40914 (separate draft), #40915 (hybrid Mamba), #40916 (QSA), #40917 (plain KV and cleanup). Each PR is based on the previous one.
Accuracy
GLM-5.3-Flash tp8, page 64, hicache size 20 GB, extra_buffer strategy: stack KV + INDEXER + MAMBA, indexer host 1.76 GB, host restore KL against warm 1.9e-3, 1.6e-3, 3.5e-3 at the finished sequence length. Main on the same runs: 0.06 to 0.71.
CI States
Latest PR Test (Base): 🚫 Run #37302418283
Latest PR Test (Extra): 🚫 Run #37302417380
Latest PR Test (AMD ROCm 10): ❌ Run #37302418055