RefKimi K3 Tech Blog: Open Frontier Intelligence当 Prefix Cache 遇见 KDAMooncake 如何 Day-0 支持 Kimi K3Hybrid Models Meet SGLang: More than Full Attention – PyTorchQwen3.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等。之所以被称为mamba应该是mamba是比较早适配的linear attention模型架构类似llama一样。Qwen 3.6是GQA Linear而kimi-k3是DeepSeek MLA Linearsglang内部采用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 cacheprefix cache、复制、回滚和 speculative commit 都需要特殊处理。因此它们也被纳入mamba-ish。核心特征普通 Transformer attention 保存每个历史 token 的 KVtoken 1 → K₁, V₁ token 2 → K₂, V₂ ... token T → Kₜ, Vₜ缓存随着序列长度增长KV cache memory ∝ sequence lengthMamba、SSM (State Space Model)和部分线性注意力模型则维护递归状态state_t update(state_{t-1}, input_t)处理完一个 token 后历史信息被压缩进固定形状的状态中不需要为每个历史 token 保存一份 KVstate memory ≈ 固定大小 / requestSGLang 把需要这种“每请求一个递归状态 slot”管理方式的模型统称为mamba-ish。状态通常包含什么SGLang 中一个 Mamba-style cache通常包含两部分conv state temporal/SSM stateTemporal/SSM state保存长期递归状态。Mamba2 中可能是标准 SSM stateGDN/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需要保存的状态它的大小从原来的kt vt的head_num*head_dim变成head_num*k_head_dim*v_head_dim。Conv stateShort Conv对 QKV 投影后的向量做 depthwise causal convolutionkernel_size4这在 DeltaNet、GDN 等线性注意力模型中几乎是标配。其作用可以从 TTTTest-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计算。对于某一个 channelc如果后续代码只保留前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_dtypeKimi-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: dense3: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 57624layer*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 heads96head dimension128short-conv kernel4KDA layers: 69TP8的场景每个GPU的head数为96/8 12。KDA 不需要为每个历史 token 保存 K/V。它把历史压缩进一个固定大小的矩阵状态S: [local_heads, V, K] [12, 128, 128]同时KDA 输入前有一个 kernel size 为 4 的短卷积所以还要保存最近4-13个 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.82MBTP 8所有GPU总和为(25.91.82)*8 221.8MB.Speculative decodingDSpark部分参考DFlash推测解码和SGLang支持_dflash sglang-CSDN博客dspark本身的kv cache需要存储在kv cache存储系统中并且TP并行每个TP worker需要存储对应的TP切分部分。因为当前dspark为MHA的attention因此按照MHA的计算逻辑计算这部分的kv cache即可。这个PR显著优化了draft kv cache存储的一系列问题[HiCache] Support packed and sidecar draft caches for MTP/EAGLE/DSpark by 1e4ves · Pull Request #30393 · sgl-project/sglang · GitHubreplayssm部分--enable-linear-replayssm-spec分配推测解码相关的state cache这部分需要占用kv cache的分配计算但并不需要存储L2 L3。也就是开启了--enable-linear-replayssm-specmamba-full-memory-ratio需要开的更大才能获得同样的mamba state kv cache容量因为有一部分用于replayssm的cache。KDA 的 recurrent update 可以概括为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 改成verify 时仍然算出每一步输出不保存每一步完整的[V,K]state只保存生成这个 state 所需的轻量输入记录acceptance 结束后仅把被接受的前缀按原 recurrent 顺序重新播放将重放结果写回持久化 SSM checkpoint。Kimi K3/Qwen 3.6这种hybrid linear attention的模型开启speculative decoding时需要加上--enable-linear-replayssm-spec参数这会导致几个额外state cache的分配例如# --kv-cache-dtype fp8_e4m3 --mamba-ssm-dtype bfloat16 Mamba Cache is allocated. max_mamba_cache_size: 119, conv_state size: 0.21GB, ssm_state size: 3.03GB intermediate_ssm_state_cache size: 0.00GB intermediate_conv_window_cache size: 0.34GB GDN ReplaySSM ring buffers allocated (L16): d0.379GB, k0.379GB, g0.758GB rawv0.379GB, rawk0.379GB, beta0.006GB这里面提到L16以及6个buffer: d, k, g, rawv, rawk, beta--linear-replayssm-cache-len没有显式设置因此使用默认L16。所有 ring 都按 69 层、head_num 96/TP8 12, L16 ,head_dim128分配假如120 个 slot、张量的shape如下BufferShapedtype含义d[69,120,12,16,128]BF16修正后的 delta/value 向量k[69,120,12,16,128]BF16归一化/缩放后的 keyg[69,120,12,16,128]FP32KDA 的逐 K-channel log-decay gaterawv[69,120,12,16,128]BF16exact-fold 使用的原始 value 输入rawk[69,120,12,16,128]BF16exact-fold 使用的归一化前 keybeta[69,120,12,16]FP32每个 head、每一步的 delta update 系数这 6 个 buffer 可以分成两组d / k / g用来从 checkpoint \(S_0\) 快速重建输出。rawv / rawk / g / beta在 speculative verify 接受 token 后精确重放递推并提交新的 \(S\)。1.d / k / g快速重建状态投影ReplaySSM 不会每步都把完整的 \(S_t\in\mathbb R^{V\times K}\) 写回显存而是保留一个 checkpoint \(S_0\)并把最近 \(L16\) 步的\[ (d_t,k_t,g_t) \]追加进 ring。经过 \(m\) 个 token 后状态可表示为\[ S_m S_0\operatorname{Diag} \left( \prod_{i1}^{m}\alpha_i \right) \sum_{j1}^{m} d_j \left( k_j\odot \prod_{ij1}^{m}\alpha_i \right)^\top. \]因此d保存每一步的 value-side rank-1 向量k保存 key-side rank-1 向量g用于计算每条历史 \(d_jk_j^\top\) 后续还要经历多少 decay。代码中直接将其写成\[ d_{\text{cur}} \beta\left(v-(S\operatorname{Diag}(\alpha))k\right), \qquad S_{\text{new}} S\operatorname{Diag}(\alpha)d_{\text{cur}}k^\top. \]对应代码说明在 [fused_recurrent_linear_replayssm.py (line 27)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py:27)实际d_cur计算在 [fused_recurrent_linear_replayssm.py (line 323)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py:323)。2.rawv / rawk / g / beta精确提交 accepted tokens只保存d存在一个问题d是基于当时重建出来的 state 计算的而且通常以 BF16 保存。反复用这些d做 open-loop fold误差会不断进入后续 state。所以--enable-linear-replayssm-spec额外保存生成 \(d_t\) 所需的原始输入\[ (\mathrm{rawv}_t,\mathrm{rawk}_t,g_t,\beta_t). \]accepted token 提交时从 FP32 checkpoint 开始顺序重放\[ k_t \frac{\mathrm{rawk}_t} {\sqrt{\|\mathrm{rawk}_t\|^2\epsilon}} \]\[ \alpha_t\exp(g_t) \]\[ S\leftarrow S\operatorname{Diag}(\alpha_t) \]\[ d_t\leftarrow \beta_t\left(\mathrm{rawv}_t-Sk_t\right) \]\[ S\leftarrow Sd_tk_t^\top. \]代码几乎逐行对应decay[kda_replayssm_spec_decode.py (line 158)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:158)residual[kda_replayssm_spec_decode.py (line 159)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:159)乘 \(\beta\)[kda_replayssm_spec_decode.py (line 160)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:160)rank-1 update[kda_replayssm_spec_decode.py (line 161)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:161)其中rawv的 “raw” 不是模型 projection 之前的 value而是相对于d/residual_v而言的它是进入 delta-rule 前、尚未执行 \(v-Sk\) 的 value。这个部分的kv cache不需要进行L2 L3 offload当前实现将以下字段列入_NON_TRANSFER_STATE_FIELDSintermediate_ssm intermediate_conv_window replayssm_d replayssm_k replayssm_g replayssm_rawv replayssm_rawk replayssm_beta也就是说Mamba/HiCache transfer 只遍历可持久化状态不传 ReplaySSM ring。Speculative decoding的Conv State部分intermediate_conv_window为spec存储的conv的cache.投机解码一次前向处理多个 draft tokenconv1d kernel 会逐 token 推进并 就地滚动 conv_state 。但 draft token 可能被拒绝不能直接把 conv_state 推到最后一个 draft token——否则被拒部分的状态会污染。所以流程是1. verify 前向 kernel 每处理一个 draft token把当时的窗口 [dim, K-1] 存进 intermediate_conv_window[step] 注释 Save the window state after consuming this token 。2. commit 根据接受的 token 数算出最后一个正确 draft 步 last_correct_step_indices 再用 fused_conv_window_scatter_with_mask 把对应步的快照 回写 进持久 conv_state fused_conv_window_scatter_with_mask( spec_state.conv[0], # dst: 持久 conv_state spec_state.intermediate_conv_window[0], # src: 每步快照 state_batch_indices, last_correct_step_indices, # 取第几个 step 的快照 )见 spec_utils.py:858-863 与 mamba_state_scatter_triton.py:476-484 。换句话说 intermediate_conv_window 是 conv_state 在 spec verify 过程中的 逐步 checkpoint commit 时用它在最后一个被接受 token 处把 conv_state 回滚/前推到正确状态丢弃被拒 token 的影响。 intermediate_ssm 对 SSM 的 temporal state 起完全相同的作用。具体的shape大小前面提到linear Attention的conv state的kv cache shape为[layer, slot_num, window_size-1, qkv_head_dim]而spec的conv state为[layer, spec_state_size 1, speculative_num_draft_tokens, window_size-1, qkv_head_dim]spec_state_size max_running_requests。max_running_requests取决于2方面一个是--cuda-graph-max-bs-decode参数一个是slot除以每个请求的slot大小这两个取最小值。speculative_num_draft_tokens speculative-dspark-block-size 1。通常为8.这两者的比例为(max_running_requests 1) * speculative_num_draft_tokens / slot_num。同样这部分不需要存储L3 cache。形状是verify_ids_2d.shape [batch_size, gamma 1]实际分配样例# no spark spec kv_cache_dtype bfloat16 mamba_ssm_dtypefloat32 Mamba Cache is allocated. max_mamba_cache_size: 51, conv_state size: 0.09GB, ssm_state size: 2.63GB GDN ReplaySSM ring buffers allocated (L16): d0.164GB, k0.164GB, g0.328GB rawv0.164GB, rawk0.164GB, beta0.003GB KV Cache is allocated. dtype: torch.bfloat16, #tokens: 974656, KV size: 25.10 GB kv_cache_dtypefp8_e4m3 mamba_ssm_dtypebfloat16 Mamba Cache is allocated. max_mamba_cache_size: 80, conv_state size: 0.14GB, ssm_state size: 2.05GB GDN ReplaySSM ring buffers allocated (L16): d0.256GB, k0.256GB, g0.512GB rawv0.256GB, rawk0.256GB, beta0.004GB KV Cache is allocated. dtype: torch.float8_e4m3fn, #tokens: 1947648, KV size: 25.08 GB实际分配的slot为size 1SGLang Mamba cache state保存逻辑核心代码python\sglang\srt\mem_cache\mamba_radix_cache.pypython\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 prefill6144总共进行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 prefill6144总共进行3次prefill计算存3个state然后decode完成存一次前提是decode相比prefill多一个page的token总共4个mamba state。除了上面的默认策略还有一些其他的优化策略[Hicache][1/2]Support Mamba branching in Unified Radix Cache with HiCache by Chen-0210 · Pull Request #31181 · sgl-project/sglang · GitHub这个PR在full attention kv的分叉点创建一个mamba 状态保存点而不是请求末尾。[Hicache] Support Mamba branching setting by env by llc-kc · Pull Request #37198 · sgl-project/sglang · GitHub这个PR使得在第一次请求时把mamba状态存的更短一些更利于后续的请求命中。Mamba state驱逐逻辑LRU驱逐两种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.80max_running_requests is capped to 16 by the mamba state cache (max_mamba_cache_size80, 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_LOCK1设置把预留数量降低。--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_LOCK1时为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 backupGPU 状态被删除后 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 为256decode checkpoint 会在总长度1024、1280、1536……也就是分别 decode 约24、280、536……个已处理 token 后触发而不是固定在 decode 输出长度256、512……时触发。最终 HiCache 通常只 offload 最新边界对应的一个 Mamba state slot。值越小缓存粒度更细、前缀命中后需要重算的 token 更少但状态保存更频繁显存和执行开销可能增加。值越大保存开销更低但缓存粒度更粗前缀复用效果可能降低。mamba_track_interval的核心存储链路分为四步判断是否到达存储边界在 [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。确定快照目标槽位在 [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 的哪个槽位。真正复制 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 stateSSM/recurrent state它们存进 Mamba state pool 的额外 tracking slot而不是普通 token KV Cache。更新 checkpoint 元数据并插入 Radix Cacheforward 完成后[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_maskTrue → 把当前 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 模型如何保存线性注意力的循环状态。它主要在三件事之间权衡是否启用 overlap scheduler。能否缓存 Radix Tree 分叉点上的 Mamba 状态。每个运行中请求要预留多少 Mamba state slot从而影响最大并发。四种策略对比策略Overlap scheduler分叉点状态缓存每请求容量预留适用场景auto自动决定自动决定取决于解析结果通常首选no_buffer不支持未实现3 slots显存紧张、兼容性优先、ReplaySSMextra_buffer支持支持overlap 开启时 5 slots吞吐优先、稳定生产配置extra_buffer_lazy必须开启支持4 slotsMamba state 容量成为瓶颈时当前sglang给每个请求预留了固定数量的slot这在mamba ratio比较小slot数量少的时候会限制并发数量。extra_buffe模式预留5个slot而extra_buffer_lazy模式只预留4个而使用SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK1可以再降低1个。auto开启 overlap schedule和page_size1默认设置为extra_buffer。no_buffer不支持overlap schedule只支持page_size1。extra_buffer普通extra_buffer在 overlap 开启时为每个请求预先分配两个 track slottrack slot ACPU/Radix Cache 可以安全读取的旧快照 track slot BGPU forward 正在写入的新快照下一轮两者交换也就是 ping-pong第 t 轮 读取 A写入 B 第 t1 轮读取 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)。优点支持 overlap scheduler通常 decode 吞吐更高。支持page_size1。支持更好的 Radix 分支点 checkpoint。分配和运行路径稳定没有 lazy 边界临时申请失败的问题。PD disaggregation 场景应优先选择它而不是 lazy。缺点overlap 开启时每个请求直接持有1 个 live 2 个 track。Mamba state 很大的模型上会明显降低最大并发请求数。只在特定模型架构、设备和 Triton linear-attention backend 上受支持。extra_buffer_lazy要求 overlap scheduler。不支持 PD disaggregation。不支持 DFLASH speculative algorithm。其他已经接好 lazy planning 的 speculative 路径可以使用。同样要求模型和 Triton linear-attention backend 支持 extra-buffer。因此它适合单体服务、Mamba state 显存成为并发瓶颈、并且愿意采用较新调度路径的场景。SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK1frees one more slot per request (experimental, under validation)--enable-linear-replayssm-spec对于kimi k3/ qwen 3.6这种hybrid linear attention模型在开启speculative decoding时开启节约spec需要的state cache需求。--linear-replayssm-cache-lenlinear_replayssm_cache_len默认16一般不需要额外设置。这里的linear_replayssm_cache_lenL16不是上下文长度也不是 KV Cache block size而是 ReplaySSM 为每个请求、每个 KDA/线性注意力层分配的临时环形缓存深度。当前代码要求L 必须是 2 的幂8、16、32、64...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 的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-memoryReplace 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, dtypetorch.uint8, devicedevice)然后在它上面构造不同的 Tensor viewMHAK/V viewMLA每层 dense viewMambaconv state 和 temporal state viewSWAFull KV 和 SWA KV view这些 view 指向同一块物理显存但分配器保证两边实际占用的字节区域不重叠。实现见class UnifiedKVPool。
SGLang Qwen3/kimi-k3 Hybrid Linear Mamba KV cache管理
RefKimi K3 Tech Blog: Open Frontier Intelligence当 Prefix Cache 遇见 KDAMooncake 如何 Day-0 支持 Kimi K3Hybrid Models Meet SGLang: More than Full Attention – PyTorchQwen3.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等。之所以被称为mamba应该是mamba是比较早适配的linear attention模型架构类似llama一样。Qwen 3.6是GQA Linear而kimi-k3是DeepSeek MLA Linearsglang内部采用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 cacheprefix cache、复制、回滚和 speculative commit 都需要特殊处理。因此它们也被纳入mamba-ish。核心特征普通 Transformer attention 保存每个历史 token 的 KVtoken 1 → K₁, V₁ token 2 → K₂, V₂ ... token T → Kₜ, Vₜ缓存随着序列长度增长KV cache memory ∝ sequence lengthMamba、SSM (State Space Model)和部分线性注意力模型则维护递归状态state_t update(state_{t-1}, input_t)处理完一个 token 后历史信息被压缩进固定形状的状态中不需要为每个历史 token 保存一份 KVstate memory ≈ 固定大小 / requestSGLang 把需要这种“每请求一个递归状态 slot”管理方式的模型统称为mamba-ish。状态通常包含什么SGLang 中一个 Mamba-style cache通常包含两部分conv state temporal/SSM stateTemporal/SSM state保存长期递归状态。Mamba2 中可能是标准 SSM stateGDN/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需要保存的状态它的大小从原来的kt vt的head_num*head_dim变成head_num*k_head_dim*v_head_dim。Conv stateShort Conv对 QKV 投影后的向量做 depthwise causal convolutionkernel_size4这在 DeltaNet、GDN 等线性注意力模型中几乎是标配。其作用可以从 TTTTest-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计算。对于某一个 channelc如果后续代码只保留前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_dtypeKimi-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: dense3: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 57624layer*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 heads96head dimension128short-conv kernel4KDA layers: 69TP8的场景每个GPU的head数为96/8 12。KDA 不需要为每个历史 token 保存 K/V。它把历史压缩进一个固定大小的矩阵状态S: [local_heads, V, K] [12, 128, 128]同时KDA 输入前有一个 kernel size 为 4 的短卷积所以还要保存最近4-13个 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.82MBTP 8所有GPU总和为(25.91.82)*8 221.8MB.Speculative decodingDSpark部分参考DFlash推测解码和SGLang支持_dflash sglang-CSDN博客dspark本身的kv cache需要存储在kv cache存储系统中并且TP并行每个TP worker需要存储对应的TP切分部分。因为当前dspark为MHA的attention因此按照MHA的计算逻辑计算这部分的kv cache即可。这个PR显著优化了draft kv cache存储的一系列问题[HiCache] Support packed and sidecar draft caches for MTP/EAGLE/DSpark by 1e4ves · Pull Request #30393 · sgl-project/sglang · GitHubreplayssm部分--enable-linear-replayssm-spec分配推测解码相关的state cache这部分需要占用kv cache的分配计算但并不需要存储L2 L3。也就是开启了--enable-linear-replayssm-specmamba-full-memory-ratio需要开的更大才能获得同样的mamba state kv cache容量因为有一部分用于replayssm的cache。KDA 的 recurrent update 可以概括为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 改成verify 时仍然算出每一步输出不保存每一步完整的[V,K]state只保存生成这个 state 所需的轻量输入记录acceptance 结束后仅把被接受的前缀按原 recurrent 顺序重新播放将重放结果写回持久化 SSM checkpoint。Kimi K3/Qwen 3.6这种hybrid linear attention的模型开启speculative decoding时需要加上--enable-linear-replayssm-spec参数这会导致几个额外state cache的分配例如# --kv-cache-dtype fp8_e4m3 --mamba-ssm-dtype bfloat16 Mamba Cache is allocated. max_mamba_cache_size: 119, conv_state size: 0.21GB, ssm_state size: 3.03GB intermediate_ssm_state_cache size: 0.00GB intermediate_conv_window_cache size: 0.34GB GDN ReplaySSM ring buffers allocated (L16): d0.379GB, k0.379GB, g0.758GB rawv0.379GB, rawk0.379GB, beta0.006GB这里面提到L16以及6个buffer: d, k, g, rawv, rawk, beta--linear-replayssm-cache-len没有显式设置因此使用默认L16。所有 ring 都按 69 层、head_num 96/TP8 12, L16 ,head_dim128分配假如120 个 slot、张量的shape如下BufferShapedtype含义d[69,120,12,16,128]BF16修正后的 delta/value 向量k[69,120,12,16,128]BF16归一化/缩放后的 keyg[69,120,12,16,128]FP32KDA 的逐 K-channel log-decay gaterawv[69,120,12,16,128]BF16exact-fold 使用的原始 value 输入rawk[69,120,12,16,128]BF16exact-fold 使用的归一化前 keybeta[69,120,12,16]FP32每个 head、每一步的 delta update 系数这 6 个 buffer 可以分成两组d / k / g用来从 checkpoint \(S_0\) 快速重建输出。rawv / rawk / g / beta在 speculative verify 接受 token 后精确重放递推并提交新的 \(S\)。1.d / k / g快速重建状态投影ReplaySSM 不会每步都把完整的 \(S_t\in\mathbb R^{V\times K}\) 写回显存而是保留一个 checkpoint \(S_0\)并把最近 \(L16\) 步的\[ (d_t,k_t,g_t) \]追加进 ring。经过 \(m\) 个 token 后状态可表示为\[ S_m S_0\operatorname{Diag} \left( \prod_{i1}^{m}\alpha_i \right) \sum_{j1}^{m} d_j \left( k_j\odot \prod_{ij1}^{m}\alpha_i \right)^\top. \]因此d保存每一步的 value-side rank-1 向量k保存 key-side rank-1 向量g用于计算每条历史 \(d_jk_j^\top\) 后续还要经历多少 decay。代码中直接将其写成\[ d_{\text{cur}} \beta\left(v-(S\operatorname{Diag}(\alpha))k\right), \qquad S_{\text{new}} S\operatorname{Diag}(\alpha)d_{\text{cur}}k^\top. \]对应代码说明在 [fused_recurrent_linear_replayssm.py (line 27)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py:27)实际d_cur计算在 [fused_recurrent_linear_replayssm.py (line 323)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py:323)。2.rawv / rawk / g / beta精确提交 accepted tokens只保存d存在一个问题d是基于当时重建出来的 state 计算的而且通常以 BF16 保存。反复用这些d做 open-loop fold误差会不断进入后续 state。所以--enable-linear-replayssm-spec额外保存生成 \(d_t\) 所需的原始输入\[ (\mathrm{rawv}_t,\mathrm{rawk}_t,g_t,\beta_t). \]accepted token 提交时从 FP32 checkpoint 开始顺序重放\[ k_t \frac{\mathrm{rawk}_t} {\sqrt{\|\mathrm{rawk}_t\|^2\epsilon}} \]\[ \alpha_t\exp(g_t) \]\[ S\leftarrow S\operatorname{Diag}(\alpha_t) \]\[ d_t\leftarrow \beta_t\left(\mathrm{rawv}_t-Sk_t\right) \]\[ S\leftarrow Sd_tk_t^\top. \]代码几乎逐行对应decay[kda_replayssm_spec_decode.py (line 158)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:158)residual[kda_replayssm_spec_decode.py (line 159)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:159)乘 \(\beta\)[kda_replayssm_spec_decode.py (line 160)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:160)rank-1 update[kda_replayssm_spec_decode.py (line 161)](D:/codes/open_engine/sglang/sglang_kimi_k3/python/sglang/kernels/ops/attention/fla/kda_replayssm_spec_decode.py:161)其中rawv的 “raw” 不是模型 projection 之前的 value而是相对于d/residual_v而言的它是进入 delta-rule 前、尚未执行 \(v-Sk\) 的 value。这个部分的kv cache不需要进行L2 L3 offload当前实现将以下字段列入_NON_TRANSFER_STATE_FIELDSintermediate_ssm intermediate_conv_window replayssm_d replayssm_k replayssm_g replayssm_rawv replayssm_rawk replayssm_beta也就是说Mamba/HiCache transfer 只遍历可持久化状态不传 ReplaySSM ring。Speculative decoding的Conv State部分intermediate_conv_window为spec存储的conv的cache.投机解码一次前向处理多个 draft tokenconv1d kernel 会逐 token 推进并 就地滚动 conv_state 。但 draft token 可能被拒绝不能直接把 conv_state 推到最后一个 draft token——否则被拒部分的状态会污染。所以流程是1. verify 前向 kernel 每处理一个 draft token把当时的窗口 [dim, K-1] 存进 intermediate_conv_window[step] 注释 Save the window state after consuming this token 。2. commit 根据接受的 token 数算出最后一个正确 draft 步 last_correct_step_indices 再用 fused_conv_window_scatter_with_mask 把对应步的快照 回写 进持久 conv_state fused_conv_window_scatter_with_mask( spec_state.conv[0], # dst: 持久 conv_state spec_state.intermediate_conv_window[0], # src: 每步快照 state_batch_indices, last_correct_step_indices, # 取第几个 step 的快照 )见 spec_utils.py:858-863 与 mamba_state_scatter_triton.py:476-484 。换句话说 intermediate_conv_window 是 conv_state 在 spec verify 过程中的 逐步 checkpoint commit 时用它在最后一个被接受 token 处把 conv_state 回滚/前推到正确状态丢弃被拒 token 的影响。 intermediate_ssm 对 SSM 的 temporal state 起完全相同的作用。具体的shape大小前面提到linear Attention的conv state的kv cache shape为[layer, slot_num, window_size-1, qkv_head_dim]而spec的conv state为[layer, spec_state_size 1, speculative_num_draft_tokens, window_size-1, qkv_head_dim]spec_state_size max_running_requests。max_running_requests取决于2方面一个是--cuda-graph-max-bs-decode参数一个是slot除以每个请求的slot大小这两个取最小值。speculative_num_draft_tokens speculative-dspark-block-size 1。通常为8.这两者的比例为(max_running_requests 1) * speculative_num_draft_tokens / slot_num。同样这部分不需要存储L3 cache。形状是verify_ids_2d.shape [batch_size, gamma 1]实际分配样例# no spark spec kv_cache_dtype bfloat16 mamba_ssm_dtypefloat32 Mamba Cache is allocated. max_mamba_cache_size: 51, conv_state size: 0.09GB, ssm_state size: 2.63GB GDN ReplaySSM ring buffers allocated (L16): d0.164GB, k0.164GB, g0.328GB rawv0.164GB, rawk0.164GB, beta0.003GB KV Cache is allocated. dtype: torch.bfloat16, #tokens: 974656, KV size: 25.10 GB kv_cache_dtypefp8_e4m3 mamba_ssm_dtypebfloat16 Mamba Cache is allocated. max_mamba_cache_size: 80, conv_state size: 0.14GB, ssm_state size: 2.05GB GDN ReplaySSM ring buffers allocated (L16): d0.256GB, k0.256GB, g0.512GB rawv0.256GB, rawk0.256GB, beta0.004GB KV Cache is allocated. dtype: torch.float8_e4m3fn, #tokens: 1947648, KV size: 25.08 GB实际分配的slot为size 1SGLang Mamba cache state保存逻辑核心代码python\sglang\srt\mem_cache\mamba_radix_cache.pypython\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 prefill6144总共进行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 prefill6144总共进行3次prefill计算存3个state然后decode完成存一次前提是decode相比prefill多一个page的token总共4个mamba state。除了上面的默认策略还有一些其他的优化策略[Hicache][1/2]Support Mamba branching in Unified Radix Cache with HiCache by Chen-0210 · Pull Request #31181 · sgl-project/sglang · GitHub这个PR在full attention kv的分叉点创建一个mamba 状态保存点而不是请求末尾。[Hicache] Support Mamba branching setting by env by llc-kc · Pull Request #37198 · sgl-project/sglang · GitHub这个PR使得在第一次请求时把mamba状态存的更短一些更利于后续的请求命中。Mamba state驱逐逻辑LRU驱逐两种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.80max_running_requests is capped to 16 by the mamba state cache (max_mamba_cache_size80, 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_LOCK1设置把预留数量降低。--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_LOCK1时为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 backupGPU 状态被删除后 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 为256decode checkpoint 会在总长度1024、1280、1536……也就是分别 decode 约24、280、536……个已处理 token 后触发而不是固定在 decode 输出长度256、512……时触发。最终 HiCache 通常只 offload 最新边界对应的一个 Mamba state slot。值越小缓存粒度更细、前缀命中后需要重算的 token 更少但状态保存更频繁显存和执行开销可能增加。值越大保存开销更低但缓存粒度更粗前缀复用效果可能降低。mamba_track_interval的核心存储链路分为四步判断是否到达存储边界在 [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。确定快照目标槽位在 [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 的哪个槽位。真正复制 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 stateSSM/recurrent state它们存进 Mamba state pool 的额外 tracking slot而不是普通 token KV Cache。更新 checkpoint 元数据并插入 Radix Cacheforward 完成后[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_maskTrue → 把当前 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 模型如何保存线性注意力的循环状态。它主要在三件事之间权衡是否启用 overlap scheduler。能否缓存 Radix Tree 分叉点上的 Mamba 状态。每个运行中请求要预留多少 Mamba state slot从而影响最大并发。四种策略对比策略Overlap scheduler分叉点状态缓存每请求容量预留适用场景auto自动决定自动决定取决于解析结果通常首选no_buffer不支持未实现3 slots显存紧张、兼容性优先、ReplaySSMextra_buffer支持支持overlap 开启时 5 slots吞吐优先、稳定生产配置extra_buffer_lazy必须开启支持4 slotsMamba state 容量成为瓶颈时当前sglang给每个请求预留了固定数量的slot这在mamba ratio比较小slot数量少的时候会限制并发数量。extra_buffe模式预留5个slot而extra_buffer_lazy模式只预留4个而使用SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK1可以再降低1个。auto开启 overlap schedule和page_size1默认设置为extra_buffer。no_buffer不支持overlap schedule只支持page_size1。extra_buffer普通extra_buffer在 overlap 开启时为每个请求预先分配两个 track slottrack slot ACPU/Radix Cache 可以安全读取的旧快照 track slot BGPU forward 正在写入的新快照下一轮两者交换也就是 ping-pong第 t 轮 读取 A写入 B 第 t1 轮读取 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)。优点支持 overlap scheduler通常 decode 吞吐更高。支持page_size1。支持更好的 Radix 分支点 checkpoint。分配和运行路径稳定没有 lazy 边界临时申请失败的问题。PD disaggregation 场景应优先选择它而不是 lazy。缺点overlap 开启时每个请求直接持有1 个 live 2 个 track。Mamba state 很大的模型上会明显降低最大并发请求数。只在特定模型架构、设备和 Triton linear-attention backend 上受支持。extra_buffer_lazy要求 overlap scheduler。不支持 PD disaggregation。不支持 DFLASH speculative algorithm。其他已经接好 lazy planning 的 speculative 路径可以使用。同样要求模型和 Triton linear-attention backend 支持 extra-buffer。因此它适合单体服务、Mamba state 显存成为并发瓶颈、并且愿意采用较新调度路径的场景。SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK1frees one more slot per request (experimental, under validation)--enable-linear-replayssm-spec对于kimi k3/ qwen 3.6这种hybrid linear attention模型在开启speculative decoding时开启节约spec需要的state cache需求。--linear-replayssm-cache-lenlinear_replayssm_cache_len默认16一般不需要额外设置。这里的linear_replayssm_cache_lenL16不是上下文长度也不是 KV Cache block size而是 ReplaySSM 为每个请求、每个 KDA/线性注意力层分配的临时环形缓存深度。当前代码要求L 必须是 2 的幂8、16、32、64...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 的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-memoryReplace 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, dtypetorch.uint8, devicedevice)然后在它上面构造不同的 Tensor viewMHAK/V viewMLA每层 dense viewMambaconv state 和 temporal state viewSWAFull KV 和 SWA KV view这些 view 指向同一块物理显存但分配器保证两边实际占用的字节区域不重叠。实现见class UnifiedKVPool。