SGLang Qwen3/kimi-k3 Hybrid Linear Mamba KV cache管理

Ref

Kimi K3 Tech Blog: Open Frontier Intelligence

当 Prefix Cache 遇见 KDA：Mooncake 如何 Day-0 支持 Kimi K3

Hybrid Models Meet SGLang: More than Full Attention -- PyTorch

Qwen3.5 GDN(Gated Delta Networks) 原理与代码分析

LLMs-from-scratch/ch04/08_deltanet/README.md at main · rasbt/LLMs-from-scratch · GitHub

谈谈Kimi K3 的KDA(1): KDA 如何与 Gated MLA 以及 AttnRes 协同

Mamba模型简介

mamba-ish 可以理解为：

"具有 Mamba 类似状态管理特征的模型"。

采用Linear和dense layer 3:1混合的模型，当前有Qwen 3.6和kimi k3等。

Qwen 3.6是GQA + Linear，而kimi-k3是DeepSeek MLA + Linear

sglang内部采用hybrid mamba模式来管理混合full和linear attention的kv cache.

• Dense/full-attention 层：KV cache 按 token 追加，显存随序列长度增长。
• Mamba/Gated-DeltaNet 层：每层、每个请求只有一份 conv state + recurrent/SSM state。每处理一个 token 都会递推更新，但属于原地覆盖，不会为每个 token 保存一份。

虽然 GDN/KDA 从算法命名上不叫 Mamba，但它们的推理状态具有相同的工程性质：

• 每个请求需要一个持久状态；
• 每生成一个 token 更新一次；
• 状态不是按 token 索引的 KV cache；
• prefix cache、复制、回滚和 speculative commit 都需要特殊处理。

因此它们也被纳入 mamba-ish。

核心特征

普通 Transformer attention 保存每个历史 token 的 KV：

token 1 → K₁, V₁
token 2 → K₂, V₂
...
token T → Kₜ, Vₜ

缓存随着序列长度增长：

KV cache memory ∝ sequence length

Mamba、SSM (State Space Model)和部分线性注意力模型则维护递归状态：

state_t = update(state_{t-1}, input_t)

处理完一个 token 后，历史信息被压缩进固定形状的状态中，不需要为每个历史 token 保存一份 KV：

state memory ≈ 固定大小 / request

SGLang 把需要这种"每请求一个递归状态 slot"管理方式的模型统称为 mamba-ish。

状态通常包含什么

SGLang 中一个 Mamba-style cache通常包含两部分：

conv state
+
temporal/SSM state

Temporal/SSM state

保存长期递归状态。

Mamba2 中可能是标准 SSM state；GDN/KDA 中通常是类似下面的线性注意力状态矩阵：

[head_num, v_head_dim * k_head_dim]

原始的linear attention:

常规的MHA Attention是保存每个token对应的这个vt和kt作为kv cache，而linear attention是保存vt kt矩阵乘的累加状态st作为kv cache。

Gated Delta Networks(GDN)

Kimi Delta Attention (KDA)改进后的计算逻辑：

这里vt, kt是常规的kv cache，而St则是linear attention需要保存的状态。

Conv state

Short Conv ：对 QKV 投影后的向量做 depthwise causal convolution（kernel_size=4），这在 DeltaNet、GDN 等线性注意力模型中几乎是标配。其作用可以从 TTT（Test-Time Training）视角理解：线性注意力的隐状态更新等价于用 (K, V) 对在线训练一个小模型（3.2 节），但当 K 和 V 来自同一个输入（"键值同源"），在线训练的目标退化为自预测，训练信号不足。Short Conv 通过对 K 做局部混合，使其变成类似 n-gram 的特征，从而将训练目标从"预测自己"转化为"预测下一个 token"，提供更有效的学习信号。实验表明，给 K 加 Short Conv 的收益最大，给 Q/V 加也有辅助效果（详见苏神的为什么线性注意力要加Short Conv？）。

源码级别的理解，参考

transformers/src/transformers/models/qwen3_5/modeling_qwen3_5.py at v5.12.1 · huggingface/transformers · GitHub

卷积的细节如上图，当前token的q k v cache在进行token之间的mixing，通常conv kernel size为4，也就是当前token只有当前token+之前的4-1个token进行计算。最终用mixing后的token q k v向量进行GDN计算。

对于某一个 channel c，如果后续代码只保留前 L 个输出，那么第 t 个位置可以理解为：

那么为什么除了上面的ssm state还需要存这个conv state:

因为当前的q k v向量需要与之前的conv kernel size -1个token进行一次mixing，因此需要把之前的conv kernel size -1个token q k v向量存下来。这里有点像SWA模型，需要存最近的N个token kv cache，不过这里的q也需要存储。

因此这部分的state cache大小为：

layer * (conv kernel size - 1) * head_num * (q_head_dim + k_head_dim + v_head_dim) * size_of_dtype
=
layer * (conv kernel size - 1) * head_num * (2 * k_head_dim + v_head_dim) * size_of_dtype

(conv kernel size - 1) 目前通常为3，如果k v head_dim一样，那么就是

layer * 9 * head_num * head_dim * size_of_dtype

Mamba cache state保存逻辑

核心代码

python\sglang\srt\mem_cache\mamba_radix_cache.py

python\sglang\srt\mem_cache\unified_cache_components\mamba_component.py

对于dense layer的kv cache，是按token存储的，每个token存储一个固定大小的key和value cache，对于MLA只需要存一个单独的hidden的cache。

而对于linear attention的mamba state部分，是在一个stage上循环累加的，理论上一个请求最少只需要存储一个state张量。

linear attention的mamba state不能像dense layer的kv cache那样逐token甚至page存储，因为单个mamba state的容量很大，存储太密集，导致存储空间需求巨大，而存储太稀疏，可能导致命中率降低。

sglang的mamba sate实际存储逻辑：

• prefill部分，每次chunked prefill完成后存储一次。例如输入13600， chunked prefill=6144，，总共进行3次prefill计算，存3个state。
• decode部分，当(输入长度+输出长度)%mamba_track_interval时更新一次mamba state，但是只offload最终的那一个mamba state。也就是decode部分如果输出很短，不会产生新的mamba state，但是通常足够长，会产生一个mamba state存储。

因此mamba state存储数量最大为(in_len + chunked_prefill_size -1)/chunked_prefill_size + 1个状态数量，或者说in_len // chunked_prefill_size + 2。

例如输入13600, 输出512， chunked prefill=6144，总共进行3次prefill计算，存3个state，然后decode完成存一次（前提是decode相比prefill多一个page的token），总共4个mamba state。

Mamba state驱逐逻辑

LRU驱逐

Kimi-k3的kv cache存储大小计算

与传统 Attention 不同，KDA 并不会长期保存每个历史 token 的 Key 和 Value，而是将历史信息不断递推（recurrent）到一个固定大小的状态中。模型会为每个 channel 学习不同的衰减系数，并利用 Delta Correction 控制新信息写入状态的强度，从而在有限状态容量下保留尽可能丰富的历史信息。同时，为了兼顾局部建模能力，KDA 还维护一个固定长度的 Convolution Window，用于保存最近几个 token 的局部信息。

因此，对于每个 KDA 层而言，真正需要持续维护的历史并不是一长串 KV，而是两部分状态：

• Temporal State：递推更新的历史状态；

• Convolution Window：最近几个 token 的局部窗口。

随着新的 token 到来，这两部分状态都会不断原地更新，而不会像传统 KV Cache 一样持续增长。这种设计最大的优势是推理过程中需要访问的缓存大小不再随上下文长度线性增长。即使面对百万级上下文，KDA 层需要维护的递推状态依然保持固定规模，大幅降低了长上下文推理的显存压力。

kimi-k3 linear: dense=3:1， 93层 = 23模块x (3 linear + 1 MLA) + 最后一层 MLA，也就是69 KDA linear + 24 MLA dense层。

Dense层的kv cache大小

kimi-k3 dense层采用MLA，因此kv cache大小与deepseek v3.1一致，每一层576个元素。

BF16 kv cache每个token的存储大小为

kv_lora_rank + qk_rope_head_dim = 512 + 64 = 576

24（layer）*576*2 (bf16) = 27 KB。

FP8直接所有token直接FP8量化，没有像deepseek v3.2那样部分BF16部分FP8，因此kv cache大小直接减半为每个token 13.5 KB。

对于Dense层，在TP并行，非DP attention/DCP的情况下，同一个请求每个GPU的KV cache是一模一样的。

Linear的kv cache大小具体计算逻辑

模型配置为：

• KDA heads：96
• head dimension：128
• short-conv kernel：4
• KDA layers: 69

TP8的场景，每个GPU的head数为96/8 = 12。

KDA 不需要为每个历史 token 保存 K/V。它把历史压缩进一个固定大小的矩阵状态：

S: [local_heads, V, K] = [12, 128, 128]

同时，KDA 输入前有一个 kernel size 为 4 的短卷积，所以还要保存最近 4-1=3 个 Q/K/V 投影输入。

因此，每个状态 slot、每个 KDA 层有两部分：

conv state: [3, head_num/TP * (q_head_dim + k_head_dim + v_head_dim)]
SSM state:  [head_num/TP, k_head_dim, v_head_dim]

shape分配在mamba2_cache_params函数中调用KimiLinearStateShape.create初始化。

因此TP8并行时，每个GPU的kv cache大小：

SSM大小为69(layer) * 96 / 8 * 128 * 128 * 2 (bf16) = 25.9MB。

conv state大小为69(layer) * 3 * 96 / 8 * 3* 128 * 2 (bf16) = 1.82MB

总的kv需要所有GPU加起来，TP 8总和为(25.9+1.82)*8 = 221.8MB.

Speculative decoding

DSpark部分

参考DFlash推测解码和SGLang支持_dflash sglang-CSDN博客

dspark本身的kv cache需要存储在kv cache存储系统中，并且TP并行，每个TP worker需要存储对应的TP切分部分。

replayssm

--enable-linear-replayssm-spec分配推测解码相关的state kv cache，这部分需要占用kv cache的分配计算，但并不需要存储L2 L3。

也就是，开启了--enable-linear-replayssm-spec，mamba-full-memory-ratio需要开的更大，才能获得同样的mamba state kv cache容量，因为有一部分用于replayssm的cache。

KDA 的 recurrent update 可以概括为：

S' = S · Diag(alpha) + d · kᵀ

d = beta · (v - (S · Diag(alpha)) · k)

KDA 的 alpha 是逐 K-channel 的向量，而不是每个 head 一个标量。

普通 speculative target verify 如果一次验证 D 个候选 token，为了最后只提交 accepted prefix，通常需要保存每一步的完整 SSM state：

[num_layers, requests, D, heads, V, K]

Kimi-K3 的 [heads,V,K] 状态很大，对每个 draft token 保存一份成本非常高。

ReplaySSM spec 改成：

1. verify 时仍然算出每一步输出；
2. 不保存每一步完整的 [V,K] state；
3. 只保存生成这个 state 所需的轻量输入记录；
4. acceptance 结束后，仅把被接受的前缀按原 recurrent 顺序重新播放；
5. 将重放结果写回持久化 SSM checkpoint。

KDA 的实现是"每次 commit 都 exact-fold"，不是普通 decode ReplaySSM 所说的"每 L 步才 flush"。相应说明在 kda_replayssm_spec_decode.py (line 11)(/D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:11)。

--linear-replayssm-cache-len 没有显式设置，因此使用默认并发量16。

所有 ring 都按 69 层、81 个 slot、16 个位置分配。

Buffer	Shape	dtype	含义
d	[69,81,12,16,128]	BF16	修正后的 delta/value 向量
k	[69,81,12,16,128]	BF16	归一化/缩放后的 key
g	[69,81,12,16,128]	FP32	KDA 的逐 K-channel log-decay gate
rawv	[69,81,12,16,128]	BF16	exact-fold 使用的原始 value 输入
rawk	[69,81,12,16,128]	BF16	exact-fold 使用的归一化前 key
beta	[69,81,12,16]	FP32	每个 head、每一步的 delta update 系数

分配代码在 memory_pool.py (line 591)(/D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/srt/mem_cache/memory_pool.py:591)。

d / k

spec 模式下它们采用 conv/activation dtype，即 BF16：

69 × 81 × 12 × 16 × 128 × 2
= 274,710,528 bytes
= 0.255844 GiB
→ 各显示 0.256GB

g

KDA 的 gate 是逐 K-channel 的，所以包含最后一个 128 维，而且强制使用 FP32：

69 × 81 × 12 × 16 × 128 × 4
= 549,421,056 bytes
= 0.511688 GiB
→ 显示 0.512GB

这也直接证明日志虽然写着 "GDN"，实际 shape 是 KDA：

• 真正 GDN 的 g 是每 head 一个标量，shape 为 [69,81,12,16]
• KDA 的 g 是 128 维向量，所以恰好大 128 倍

rawv / rawk

两者 shape 和 d/k 一样，都是 BF16：

各 0.255844 GiB
→ 各显示 0.256GB

verify kernel 写入的是：

• 尚未做 delta correction 的 v
• 尚未做 L2 normalization 的 k
• kernel 内实际形成的 FP32 g
• sigmoid(b) 后的 FP32 beta

对应写入逻辑在 fused_sigmoid_gating_recurrent.py (line 210)(/D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py:210)。

beta

69 × 81 × 12 × 16 × 4
= 4,292,352 bytes
= 0.003998 GiB
→ 显示 0.004GB

Ring 总大小

d + k + g + rawv + rawk + beta
= 1.539063 GiB

每个 slot 跨 69 层：

20,401,920 bytes
≈ 19.457 MiB

需要指出：在 KDA spec 路径中，真正用于 commit exact-fold 的主要是：

rawv + rawk + g + beta

d/k 主要是 GDN/普通 ReplaySSM reconstruction 需要的。当前统一内存池仍为 KDA 分配它们，源码注释也明确说它们对 KDA fold 看起来是"dead weight"，但暂时保留以避免改变 decode dispatch 行为。因此这里约有：

d + k = 0.511688 GiB/卡

属于当前实现的额外开销。

--linear-replayssm-cache-len

linear_replayssm_cache_len默认16。

这里的linear_replayssm_cache_len L=16 不是上下文长度，也不是 KV Cache block size，而是 ReplaySSM 为每个请求、每个 KDA/线性注意力层分配的临时环形缓存深度。

当前代码要求：

1. L 必须是 2 的幂：8、16、32、64...

2. L 必须满足：

   L >= 2 × 最大 speculative verify token 数

也就是建议直接取：

L = next_power_of_two(2 × D)

其中 D 是最终的 max_speculative_num_draft_tokens。

DSpark 中：

# speculative_hook.py _handle_dspark

--speculative-dspark-block-size 7
             ↓
gamma = 7
             ↓
target verify window = 1 anchor + 7 drafts
             ↓
speculative_num_draft_tokens = 8
             ↓
ReplaySSM 要求 L >= 2 × 8
             ↓
linear_replayssm_cache_len 最小为 16

第一个 token 来自上一轮的 bonus token，写入位置在

draft_block_ids = torch.full((bs, gamma), ...)
draft_block_ids[:, 0].copy_(draft_input.bonus_tokens.view(-1))

所以 target verify 的形状是：

verify_ids_2d.shape = [batch_size, gamma + 1]

实际分配样例

# no spark spec
kv_cache_dtype = bfloat16
mamba_ssm_dtype='float32'
Mamba Cache is allocated. max_mamba_cache_size: 51, conv_state size: 0.09GB, ssm_state size: 2.63GB
GDN ReplaySSM ring buffers allocated (L=16): d=0.164GB, k=0.164GB, g=0.328GB rawv=0.164GB, rawk=0.164GB, beta=0.003GB
KV Cache is allocated. dtype: torch.bfloat16, #tokens: 974656, KV size: 25.10 GB

kv_cache_dtype='fp8_e4m3'
mamba_ssm_dtype='bfloat16'
Mamba Cache is allocated. max_mamba_cache_size: 80, conv_state size: 0.14GB, ssm_state size: 2.05GB
GDN ReplaySSM ring buffers allocated (L=16): d=0.256GB, k=0.256GB, g=0.512GB rawv=0.256GB, rawk=0.256GB, beta=0.004GB
KV Cache is allocated. dtype: torch.float8_e4m3fn, #tokens: 1947648, KV size: 25.08 GB

实际分配的slot为size + 1

两种Attention KV大小设置

上面介绍了mamba state存储数量每个请求需要大约为prefill // chunked_prefill_size + 2个mamba state状态数量。单个token的dense kv cache和单个mamba state的内存占用大小是确定的。

这导致一个后果：不同的请求长度需要设置不同的mamba-full-memory-ratio，短输入需要设置比较大的值，而长输入需要设置小的值。

设置不合理，会导致推理的并发量受限于主kv或者是Mamba kv，例如下面这个例子主kv还有很大的空闲，但是mamba state kv已经用满了，导致并发被限制：

# in 16k out 3k
Decode batch, #running-req: 16, 
#full token: 206080, full token usage: 0.11, 
mamba num: 64, mamba usage: 0.80

max_running_requests is capped to 16 by the mamba state cache 
(max_mamba_cache_size=80, 5 state slots per request). 
To raise it: increase --mamba-full-memory-ratio or --max-mamba-cache-size, 
or halve the state size with --mamba-ssm-dtype bfloat16.

resolve_max_num_reqs里面使用分配的总的mamba state除以_calculate_mamba_ratio()计算的每个请求预留的mamba state来计算最大的并发数量，例如总共分配80个slot，每个请求预留5个 （5 = 基础安全容量 3 + overlap ping-pong buffer 2），那么只能并发16。这时候要提升并发，需要增大mamba内存分配比例，或者通过SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK=1设置把预留数量降低。

相关参数设置

--max-mamba-cache-size

人工指定整个 Mamba 状态池最多有多少物理槽位。这个参数什么时候有用：因为系统需要每个请求预留4或者5个slot的mamba state，因此最小的slot数量分配就是期望的decode并发数乘以预留slot数量。

--mamba-full-memory-ratio

设置sglang启动时 Mamba 状态与 Full KV 的显存预算比例，通过full和mamba的kv cache比例来自动计算mamba state槽位。

sglang官方的计算器

Kimi-K3 - SGLang Documentation

这个计算公式也有一些缺陷：

• 没有考虑sglang的chunked prefill存储逻辑，因为这个比例计算跟chunked prefill size有关；
• 如果通过--mamba-max-states-per-path设置了每个请求最大的mamba state数量，这个公式也需要修改。
• 当前每个请求预留了4-5个slot（ 设置SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK=1时为4，否则为5**），还需要根据这个值和最大并发量来确定最优ratio比例。**

只开TP不开启DCP时，mamba:dense_kv因为mamba部分按GPU进行head切分，因此大小/TP数量，因此比例更低，而开启DCP时，没有这个冗余，这个mem ratio要乘以TP数。

--mamba-max-states-per-path

每条 Radix 路径保留多少个历史 Mamba 检查点。

整个请求的Dense KV cache在运行未结束前，是不能进行驱逐的，所有token的kv cache必须保持，但是运行中请求的mamba state只需要保留一份，可以进行驱逐。这种情况的mamba kv cache是从根部往尾部驱逐，而不是从尾部往根部驱逐。而请求之间的前缀匹配，进行驱逐的时候，两者都应该从尾部往根部驱逐。

从上面介绍可以看到，请求越长，分配的mamba state越多，但是前面部分的mamba state可能对缓存命中率的贡献并不是很大，但是却占用大量存储。

--mamba-max-states-per-path可以减少长会话不断延伸时积累的历史状态，例如：

--mamba-max-states-per-path 3

这会让历史路径释放更多槽位，给新请求使用，但代价是：

• 请求从较浅前缀分叉时，可能找不到对应的 Mamba 状态；
• 需要从更早的状态重新计算；
• Mamba prefix-cache 命中效果可能下降；
• 若存在 HiCache host backup，GPU 状态被删除后 host 副本仍保留，但命中时需要重新加载。

overlap schedule时至少每个请求需要2个slot。

新请求进入时，如果空闲 slot 不足，SGLang 会自动从 Radix Cache 中 LRU 驱逐 "未锁定、可驱逐"的历史 Mamba checkpoint；活跃请求正在使用或被锁定的状态不会被驱逐。因此，这个--mamba-max-states-per-path通常并不需要设置。

--mamba-track-interval

控制输出部分的mamba state更新逻辑，当前默认256。

例如 prompt 长度为 1000、interval 为 256，decode checkpoint 会在总长度：

1024、1280、1536......

也就是分别 decode 约 24、280、536...... 个已处理 token 后触发，而不是固定在 decode 输出长度 256、512...... 时触发。最终 HiCache 通常只 offload 最新边界对应的一个 Mamba state slot。

值越小：缓存粒度更细、前缀命中后需要重算的 token 更少，但状态保存更频繁，显存和执行开销可能增加。

值越大：保存开销更低，但缓存粒度更粗，前缀复用效果可能降低。

mamba_track_interval 的核心存储链路分为四步：

1. 判断是否到达存储边界

在 schedule_batch.py (line 2999)(/D:/codes/open_engine/sglang/sglang_github/python/sglang/srt/managers/schedule_batch.py:2999)：

mamba_track_interval = get_exec().mamba.mamba_track_interval

self.mamba_track_mask = (
    self.seq_lens_cpu % mamba_track_interval == 0
)

只有序列长度为 interval 整数倍的请求，mamba_track_mask 才为 True。

1. 确定快照目标槽位

在 schedule_batch.py (line 1796)(/D:/codes/open_engine/sglang/sglang_github/python/sglang/srt/managers/schedule_batch.py:1796) 的 set_mamba_track_indices_from_reqs() 中，根据请求的 ping-pong buffer 生成：

batch.mamba_track_indices

它表示当前 Mamba 状态应该写入 Mamba state pool 的哪个槽位。

1. 真正复制 Mamba 状态

主要入口在 hybrid_linear_attn_backend.py (line 711)(/D:/codes/open_engine/sglang/sglang_github/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py:711)：

track_mamba_states_if_needed(
    conv_states,
    ssm_states,
    cache_indices,                  # 当前运行状态
    forward_batch.mamba_track_mask, # 是否到达 interval
    self.forward_metadata.mamba_track_indices,  # 快照目标槽位
    ...
)

真正执行复制的是 Triton kernel：

mamba_state_scatter_triton.py (line 13)(/D:/codes/open_engine/sglang/sglang_github/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py:13)

其核心逻辑等价于：

if mamba_track_mask[i]:
    conv_states[track_slot] = conv_states[active_slot]
    ssm_states[track_slot] = ssm_states[active_slot]

也就是说，存储的是两部分：

• convolution state
• SSM/recurrent state

它们存进 Mamba state pool 的额外 tracking slot，而不是普通 token KV Cache。

1. 更新 checkpoint 元数据并插入 Radix Cache

forward 完成后，batch_result_processor.py (line 1074)(/D:/codes/open_engine/sglang/sglang_github/python/sglang/srt/managers/scheduler_components/batch_result_processor.py:1074) 会记录：

req.mamba_last_track_seqlen = track_seqlen

非 lazy 策略还会切换 ping-pong 槽位：

req.mamba_next_track_idx = other_idx

请求完成或中途缓存时，mamba_radix_cache.py (line 544)(/D:/codes/open_engine/sglang/sglang_github/python/sglang/srt/mem_cache/mamba_radix_cache.py:544) 使用：

cache_len = req.mamba_last_track_seqlen
mamba_value = src_active.clone()
self.insert(...)

把 tracking slot 对应的状态与 token 前缀一起挂到 Radix Tree 节点上。

整体流程是：

seq_len 到达 interval 边界
    → mamba_track_mask=True
    → 把当前 conv/SSM 状态复制到 ping-pong tracking slot
    → 记录 mamba_last_track_seqlen
    → 请求缓存时将该 slot 插入 Radix Cache

--mamba-radix-cache-strategy

可选项：

MAMBA_RADIX_CACHE_STRATEGY_CHOICES = [
    "auto",
    "no_buffer",
    "extra_buffer",
    "extra_buffer_lazy",
]

默认为auto也就是extra_buffer模式。

--mamba-radix-cache-strategy 决定混合 Attention + Mamba/GDN/KDA 模型如何保存线性注意力的循环状态。它主要在三件事之间权衡：

1. 是否启用 overlap scheduler。
2. 能否缓存 Radix Tree 分叉点上的 Mamba 状态。
3. 每个运行中请求要预留多少 Mamba state slot，从而影响最大并发。

四种策略对比

策略	Overlap scheduler	分叉点状态缓存	每请求容量预留	适用场景
auto	自动决定	自动决定	取决于解析结果	通常首选
no_buffer	不支持	未实现	3 slots	显存紧张、兼容性优先、ReplaySSM
extra_buffer	支持	支持	overlap 开启时 5 slots	吞吐优先、稳定生产配置
extra_buffer_lazy	必须开启	支持	4 slots	Mamba state 容量成为瓶颈时

当前sglang给每个请求预留了固定数量的slot，这在mamba ratio比较小，slot数量少的时候会限制并发数量。extra_buffe模式预留5个slot，而extra_buffer_lazy模式只预留4个，而使用SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK=1可以再降低1个。

auto

开启 overlap schedule和page_size>1默认设置为extra_buffer。

no_buffer

不支持overlap schedule，只支持page_size==1。

extra_buffer

普通 extra_buffer 在 overlap 开启时，为每个请求预先分配两个 track slot：

track slot A：CPU/Radix Cache 可以安全读取的旧快照
track slot B：GPU forward 正在写入的新快照

下一轮两者交换，也就是 ping-pong：

第 t 轮：   读取 A，写入 B
第 t+1 轮：读取 B，写入 A

之所以需要两个，是因为 overlap scheduler 允许：

CPU：处理上一轮结果、更新 Radix Cache
GPU：同时执行下一轮 forward

只用一个快照槽时，CPU 读取状态和 GPU 覆盖状态可能发生竞争。

代码直接定义：

self.mamba_ping_pong_track_buffer_size = (
    2 if enable_overlap_schedule else 1
)

见 memory_pool.py (line 1142)(D:/codes/open_engine/sglang/sglang_github/python/sglang/srt/mem_cache/memory_pool.py:1142)。

普通 extra_buffer 会在请求进入时一次性申请全部两个 slot，见 memory_pool.py (line 1372)(D:/codes/open_engine/sglang/sglang_github/python/sglang/srt/mem_cache/memory_pool.py:1372)。

因此最终容量系数是：

base              = 3
overlap ping-pong = 2
---------------------
总计              = 5

对应代码：

MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3
MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2

见 kv_cache_configurator.py (line 108)(D:/codes/open_engine/sglang/sglang_github/python/sglang/srt/mem_cache/kv_cache_configurator.py:108)。

extra_buffer_lazy

extra_buffer_lazy unsupported under PD disaggregation;

SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK=1

frees one more slot per request (experimental, under validation)

Unified Memory Pool

实现路径：python\sglang\srt\mem_cache\unified_memory_pool.py

当前问题：与多项技术不兼容

• --enable-unified-memory is not yet compatible with PD " "disaggregation."
• --enable-unified-memory is not yet compatible with speculative decoding.
• --enable-unified-memory is not yet compatible with hierarchical host-tiered KV cache
• --enable-unified-memory is not yet compatible with decode context parallelism (--dcp-size > 1)

针对设置不合理的--mamba-full-memory-ratio会导致无法自适应不同的业务场景问题，sglang社区提出了主kv和mamba state共用一份kv cache的方案：

--enable-unified-memory

Replace the statically-partitioned hybrid-model pools (full-attn KV + SWA/Mamba state) with one byte buffer split dynamically between sub-pools. Requires the Triton attention / linear-attn / Mamba backends; not yet compatible with PD disaggregation or speculative decoding.

它让 Full KV Cache 和 Mamba Cache 共用同一块显存，两边从相反方向动态增长：

• 长请求多：更多显存用于 Full KV token。
• 短请求多：空闲的 Full KV 显存可动态转成更多 Mamba 槽位。
• mamba_full_memory_ratio 只参与确定启动时总预算，不再固定运行时的两边分界。

开启 enable_unified_memory 后：

低地址                                                   高地址
| Mamba/SWA → → →       动态空闲区       ← ← ← Full KV |
                   边界随运行负载变化

Full KV 从高地址向下增长，Mamba/SWA 从低地址向上增长，直到两边相遇。

底层物理存储

系统只分配一块 GPU 字节缓冲区：

self._raw = torch.empty(total_bytes, dtype=torch.uint8, device=device)

然后在它上面构造不同的 Tensor view：

• MHA：K/V view
• MLA：每层 dense view
• Mamba：conv state 和 temporal state view
• SWA：Full KV 和 SWA KV view

这些 view 指向同一块物理显存，但分配器保证两边实际占用的字节区域不重叠。实现见class UnifiedKVPool。