Skip to content

[mem_cache] Run mamba models on UnifiedRadixCache when the radix cache is disabled - #42354

Merged
hnyls2002 merged 7 commits into
mainfrom
lsyin/mamba-disable-unified
Oct 3, 2026
Merged

hnyls2002 merged 7 commits into
mainfrom
lsyin/mamba-disable-unified

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Oct 3, 2026 •

Copy link
Copy Markdown
Collaborator
  • With --disable-radix-cache, mamba models (GDN, Mamba2, KDA, lightning, incl. MiniCPM SALA and the PD decode side) now use UnifiedRadixCache's disabled mode instead of ChunkCache; the mamba component owns mamba-state release in both modes.
  • MambaComponent requires page_size == 1 only when the tree caches states; a disabled tree caches none. This also fixes --disable-radix-cache --enable-streaming-session on mamba models, which crashed at startup since [Session] Run streaming sessions only on UnifiedRadixCache; reject unverified tree caches #42295.
  • Gate the mamba COW headroom (allocation.py) and exact_chunk_fill on the tree actually caching mamba states.
  • MLX: drop the pool-owned auxiliary-state release path; MlxAuxiliaryStateComponent releases in both modes.
  • Reject mamba models on tree caches that do not manage mamba state (e.g. --enable-flexkv, third-party backends) as not verified, and drop the release_kv_cache fallback that freed mamba state for them.
  • Behavior change: beam search now rejects mamba models with --disable-radix-cache as well (it already rejected them with the radix cache on). Previously ChunkCache reported supports_mamba() == False and slipped past the check, but beam forking copies only the KV page mapping (beam_search/fork.py) and never the mamba state, so those beams ran on the wrong recurrent state.

Verification (Qwen3.5-4B, --disable-radix-cache, strict busy mem check + KV page invariants + tree sanity check on): GSM8K 400q default / retract / no-overlap 0.8775 / 0.8875 / 0.8825 (main with ChunkCache: 0.875-0.8925), concurrent abort storms, streaming-session kit (default, retract) and abort-leak repro all pass with no leak. Output throughput with checks off matches main within noise (avg 2692 vs 2627 tok/s over 3 alternating pairs).


CI States

Latest PR Test (Base): 🚫 Run #37115199321
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.
Latest PR Test (AMD ROCm 10): ❌ Run #37115199488

@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/rerun-test test_disaggregation_hybrid_attention.py test_qwen35_gdn_multi_item_scoring.py test_nvidia_nemotron_3_super_bf16_mtp.py test_streaming_session.py test_mamba2_extra_buffer_kl.py

@github-actions

github-actions Bot commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

Results for /rerun-test test_disaggregation_hybrid_attention.py test_qwen35_gdn_multi_item_scoring.py test_nvidia_nemotron_3_super_bf16_mtp.py test_streaming_session.py test_mamba2_extra_buffer_kl.py:

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_hybrid_attention.py
cd test/ && python3 registered/e2e/models/test_nvidia_nemotron_3_super_bf16_mtp.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/e2e/gdn/test_qwen35_gdn_multi_item_scoring.py
cd test/ && python3 registered/sessions/test_streaming_session.py
cd test/ && python3 registered/radix_cache/test_mamba2_extra_buffer_kl.py

@hnyls2002 hnyls2002 added the run-ci CI: run the baseline test suite on this PR label Oct 3, 2026
@hnyls2002 hnyls2002 added the parallel-stages CI: stages dispatch together instead of waiting on each other label Oct 3, 2026
@hnyls2002 hnyls2002 added the bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) label Oct 3, 2026
@hnyls2002
hnyls2002 merged commit df53c97 into main Oct 3, 2026
216 of 279 checks passed
@hnyls2002
hnyls2002 deleted the lsyin/mamba-disable-unified branch October 3, 2026 21:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

apple-silicon bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) parallel-stages CI: stages dispatch together instead of waiting on each other run-ci CI: run the baseline test suite on this PR unified-radix-cache

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants