企业官网建设流程全解析

Ref

Kimi K3 Tech Blog: Open Frontier Intelligence

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

https://pytorch.org/blog/hybrid-models-meet-sglang-more-than-full-attention/

Mamba模型简介

mamba-ish可以理解为：

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

其中-ish是英语后缀，表示“类似……的”。它不是一个严格的论文模型名称，而是 SGLang 内部使用的工程分类。

采用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 保存一份。

核心特征

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

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

缓存随着序列长度增长：

KV cache memory ∝ sequence length

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

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

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

state memory ≈ 固定大小 / request

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

状态通常包含什么

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

conv state + temporal/SSM state

Conv state

保存短卷积所需的最近几个输入：

[channels, conv_kernel - 1]

它类似一个滑动窗口。

Temporal/SSM state

保存长期递归状态。

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

[HV, V, K]

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

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

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

Mamba KV cache保存逻辑

核心代码

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

所以每个 TP rank 上：

local_heads = 96 / TP # head per GPU K = V = 128 # head_dim conv_history = kernel_size - 1 = 3 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, Q_local + K_local + V_local] SSM state: [head_num/TP, q_head_dim, v_head_dim]

其中每个 Q/K/V 的本地宽度为：

12 heads × 128 = 1536

所以：

conv state shape = [3, 1536 × 3] = [3, 4608] SSM state shape = [12, 128, 128]

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

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

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

conv state大小为3*1536*3*69*2 = 1.82MB

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

Speculative decoding

--enable-linear-replayssm-spec分配推测解码相关的state kv 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 个位置分配。

分配代码在 [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 内实际形成的 FP32g
• 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/卡

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

实际分配样例

# 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，从而影响最大并发。

四种策略对比

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。