[Qwen3-next] support mamba radix cache for overlap scheduler by hanming-lu · Pull Request #14792 · sgl-project/sglang

hanming-lu · 2025-12-10T06:27:49Z

Motivation

Currently, qwen3-next doesn't support 1) overlap scheduler; 2) page_size > 1; 3) branching point caching

Modifications

Support overlap scheduler for qwen3-next
Support page_size > 1 for qwen3-next
Support branching point caching for qwen3-next
tested for (ps = 1, ps > 1) x (non-spec dec, sd topk1, sd topk>1). All work except for ps > 1 + sd topk > 1, which is not supported on main yet
enable with --mamba-scheduler-strategy extra_buffer
Better memory allocation for ssm spec dec - instead of coupling spec dec intermediate state size with total ssm states, couple it with max running requests.

Accuracy Tests

Added mamba radix cache KL tests for prefill and decode
cover both --mamba-scheduler-strategy extra_buffer and no_buffer

Benchmarking and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.
Work with maintainers to merge your PR. See the PR Merge Process

gemini-code-assist · 2025-12-10T06:27:52Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

hanming-lu · 2025-12-11T01:30:40Z

/tag-and-rerun-ci

merrymercy · 2025-12-13T03:55:43Z

+        self.enable_mamba_radix_cache_v2 = False
        self._init_mamba_pool(
-            size + pre_alloc_size, cache_params, device, speculative_num_draft_tokens
+            size + pre_alloc_size,
+            size + pre_alloc_size,
+            cache_params,
+            device,
+            self.enable_mamba_radix_cache_v2,
+            speculative_num_draft_tokens,


@ShangmingCai Once this is merged. Can you help support PD + this PR?

merrymercy · 2025-12-14T12:24:10Z

 | `--max-mamba-cache-size` | The maximum size of the mamba cache. | `None` | Type: int |
 | `--mamba-ssm-dtype` | The data type of the SSM states in mamba cache. | `float32` | `float32`, `bfloat16` |
 | `--mamba-full-memory-ratio` | The ratio of mamba state memory to full kv cache memory. | `0.2` | Type: float |
+| `--mamba-scheduler-strategy` | The strategy to use for mamba scheduler. `auto` currently defaults to `no_buffer`. 1. `no_buffer` does not support overlap scheduler due to not allocating extra mamba state buffers. Branching point caching support is feasible but not implemented. 2. `extra_buffer` supports overlap schedule by allocating extra mamba state buffers to track mamba state for caching (increases mamba state usage per running req by 2 for non-speculative decoding, 1 for speculative decoding). 2a. `extra_buffer` is strictly better for non-KV-cache-bound cases; for KV-cache-bound cases, the tradeoff depends on whether enabling overlap outweighs reduced max running requests. 2b. mamba caching at radix cache branching point is strictly better than non-branch but requires kernel support (currently only FLA backend), currently only extra_buffer supports branching. | `auto` | `auto`, `no_buffer`, `extra_buffer` |


What is the meaning of (increases mamba state usage per running req by 2 for non-speculative decoding)? Can we use descriptions such as increase by 1.2x or 1.5x?

ShangmingCai · 2025-12-14T12:44:56Z

@ShangmingCai Once this is merged. Can you help support PD + this PR?

@merrymercy Sure. I will support PD + this PR ASAP.

…ject#14792)

Ageliss · 2026-01-20T05:22:07Z

Is there any performance details before and after the overlap scheduler support of qwen3-next?

…ject#14792)

mingyuedang · 2026-06-22T19:34:56Z

求指教，非投机解码开启extra_buffer给每个请求分配2个slot作用是啥

Hanming Lu added 7 commits December 8, 2025 20:36

wip

2f74618

wip

311f582

radix cache file done

7c8fa41

memory_pool done

0fd76d2

forward batch info and cuda graph runner done

11c08a6

model runner done

7f6b177

kernel done; hybrid linear backend done; eagle done; output process done

554f655

github-actions Bot added the npu label Dec 10, 2025

Hanming Lu added 3 commits December 10, 2025 06:32

schedule batch done; can start testing

42f35b3

TestQwen3Next passing

c63ff7e

add branching test; need to fix branching

cdd6123

yizhang2077 reviewed Dec 10, 2025

View reviewed changes

Comment thread python/sglang/srt/managers/scheduler_output_processor_mixin.py Outdated

Hanming Lu and others added 4 commits December 10, 2025 21:15

qwen3-next fully working

5083e92

address comment

58c53c5

Merge branch 'main' into hanming/mamba-cache-p1

6e8764c

ready for ci; qwen3-next, kimi linear, and nemotron tests passed

70e42af

hanming-lu changed the title ~~[Qwen3-next] Prefix cache for qwen3-next~~ [Qwen3-next] radix cache v2 for qwen3-next Dec 11, 2025

Merge branch 'main' into hanming/mamba-cache-p1

4f87145

hanming-lu marked this pull request as ready for review December 11, 2025 01:30

hanming-lu requested review from Fridge003, Ying1123, hebiao064, hnyls2002, ispobock, merrymercy, xiezhq-hermann and zhyncs as code owners December 11, 2025 01:30

github-actions Bot added the run-ci label Dec 11, 2025

Hanming Lu added 2 commits December 13, 2025 23:51

documentation

0483a37

minor

92f89ba

github-actions Bot added the documentation Improvements or additions to documentation label Dec 13, 2025

Hanming Lu added 4 commits December 13, 2025 23:53

minor

e9db49f

minor

457fdc9

cleanup

d885731

address comments

acca13b

hanming-lu changed the title ~~[Qwen3-next] radix cache v2 for qwen3-next~~ [Qwen3-next] support mamba radix cache for overlap scheduler Dec 14, 2025

merrymercy approved these changes Dec 14, 2025

View reviewed changes

Hanming Lu and others added 2 commits December 14, 2025 19:13

comment

d540b8c

Merge branch 'main' into hanming/mamba-cache-p1

606f906

merrymercy merged commit e61dabf into main Dec 15, 2025
136 of 146 checks passed

merrymercy deleted the hanming/mamba-cache-p1 branch December 15, 2025 02:54

ShangmingCai mentioned this pull request Dec 15, 2025

[Qwen3-next] Add PD disaggregation support for mamba with extra_buffer #15180

Merged

6 tasks

roikoren755 mentioned this pull request Dec 15, 2025

Add missing assertion in NemotronH path #15193

Merged

6 tasks

ShangmingCai mentioned this pull request Dec 16, 2025

[CI] Improve flaky 4 GPU test success rate #15234

Merged

6 tasks

elvischenv mentioned this pull request Dec 17, 2025

[Hotfix] Fix required enable_mamba_track argument for Flashinfer autotune path #15314

Merged

6 tasks

lw9527 mentioned this pull request Jan 6, 2026

[Bug] [SwaPool]Qwen380Bnext prefill Memory region not registered. #16439

Closed

5 tasks

This was referenced Jan 7, 2026

Fix qwen3_80_next radix error pd #16656

Closed

fix spec qwen3 pd error #16708

Merged

alisonshao mentioned this pull request Jan 8, 2026

Triton kernel causal_conv1d_update bug with padded batches #16714

Closed

YChange01 pushed a commit to YChange01/sglang that referenced this pull request Jan 13, 2026

[Qwen3-next] support mamba radix cache for overlap scheduler (sgl-pro…

d838c75

…ject#14792)

popsiclexu pushed a commit to popsiclexu/sglang that referenced this pull request Mar 25, 2026

[Qwen3-next] support mamba radix cache for overlap scheduler (sgl-pro…

772259f

…ject#14792)

0826joyce pushed a commit to 0826joyce/sglang-perf-opt that referenced this pull request May 19, 2026

[Qwen3-next] support mamba radix cache for overlap scheduler (sgl-pro…

cd0c89c

…ject#14792)

ntgiang71096 mentioned this pull request Jun 3, 2026

[ROCm] Support radix-cache with mamba-extra-buffer for Qwen3.5 #27141

Closed

5 tasks

adityakamat24 mentioned this pull request Jul 21, 2026

[Bug] NemotronH --mamba-scheduler-strategy extra_buffer accuracy drop on AIME26 (Nemotron-3-Super-120B) #31833

Closed

5 tasks

adityakamat24 mentioned this pull request Aug 4, 2026

[Mamba] Fix wrong SSM state cached by the mamba radix cache under extra_buffer #33450

Open

4 tasks

Navigation Menu

Uh oh!

[Qwen3-next] support mamba radix cache for overlap scheduler - #14792

[Qwen3-next] support mamba radix cache for overlap scheduler#14792
merrymercy merged 50 commits into
mainfrom
hanming/mamba-cache-p1

hanming-lu commented Dec 10, 2025 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Dec 10, 2025

Uh oh!

Uh oh!

hanming-lu commented Dec 11, 2025

Uh oh!

merrymercy Dec 13, 2025

Uh oh!

merrymercy Dec 14, 2025

Uh oh!

ShangmingCai commented Dec 14, 2025

Uh oh!

Uh oh!

Ageliss commented Jan 20, 2026

Uh oh!

mingyuedang commented Jun 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Uh oh!

Conversation

hanming-lu commented Dec 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Uh oh!

gemini-code-assist Bot commented Dec 10, 2025

Uh oh!

Uh oh!

hanming-lu commented Dec 11, 2025

Uh oh!

merrymercy Dec 13, 2025

Choose a reason for hiding this comment

Uh oh!

merrymercy Dec 14, 2025

Choose a reason for hiding this comment

Uh oh!

ShangmingCai commented Dec 14, 2025

Uh oh!

Uh oh!

Ageliss commented Jan 20, 2026

Uh oh!

mingyuedang commented Jun 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

hanming-lu commented Dec 10, 2025 •

edited

Loading