minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Kimi] Support kimi-k3 (#32541)

Browse Source 
Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Ziyi Xu <ziyi.xu@radixark.ai>
Co-authored-by: Zijie Xia <37504505+zijiexia@users.noreply.github.com>
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
Co-authored-by: zhangxiaohao <1024393531@qq.com>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: Julien Lin <jullin@nvidia.com>
Co-authored-by: Hao Phan <htphan@nvidia.com>
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
Co-authored-by: RolaoDenthu <xinyisong0111@gmail.com>
Co-authored-by: pigeonsoup <32922982+pigeonsoup@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com>
Co-authored-by: Lee Nau <lee.nau@gmail.com>
Co-authored-by: HMING <126185151+Hearum@users.noreply.github.com>
Co-authored-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
Co-authored-by: Byron Hsu <byronhsu1230@gmail.com>
Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com>
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
Co-authored-by: BBuf <xiaoyu.zhang@radixark.ai>
Co-authored-by: Hanming Lu <hanminglu@meta.com>
Co-authored-by: Xinyi Song <xinyis10@illinois.edu>
This commit is contained in:
Liangsheng Yin
2026-08-04 13:22:49 -07:00
committed by GitHub
co-authored by DarkSharpness Xiaoyu Zhang Mick Yuhao Yang Cheng Wan Ke Bao Baizhou Zhang Chunan Zeng Khoa Pham Ziyi Xu Zijie Xia Yuwei An zhangxiaohao Yangmin Li Julien Lin Hao Phan Thomas Wang RolaoDenthu pigeonsoup HaiShaw Xinyuan Tong Pranjal Shankhdhar Lee Nau HMING elvischenv Byron Hsu Byron Hsu Claude Opus 5 Thomas Wang Xinyi Song Mohammad Miadh Angkad Cheng Wan BBuf Hanming Lu Xinyi Song
parent 0753663b8e
commit abddb1c7e9
139 changed files with 15414 additions and 911 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool.py
+51 -23
View File
@@ -344,7 +344,7 @@ class MambaPool:
# replayssm_rawv: [num_layers, num_slots, HV, L, V] (conv/activation dtype)
# replayssm_rawk: [num_layers, num_slots, H, L, K] (conv/activation dtype)
# replayssm_beta: [num_layers, num_slots, HV, L] (fp32)
# The raw rings + beta exist only under --enable-gdn-replayssm-spec: the
# The raw rings + beta exist only under --enable-linear-replayssm-spec: the
# closed-loop exact fold sequentially replays them through the recurrent
# update at flush -- bit-identical to the recurrent baseline -- instead
# of folding the chunked `d` records open-loop (which accumulates error
@@ -380,7 +380,7 @@ class MambaPool:
@dataclass(frozen=True, kw_only=True)
class SpeculativeState(State):
# None under --enable-gdn-replayssm-spec: the spec ring owns rollback
# None under --enable-linear-replayssm-spec: the spec ring owns rollback
# (verify writes ring records, commit moves cursors), so the per-draft
# full-state snapshots are never produced or consumed.
intermediate_ssm: Optional[torch.Tensor]
@@ -468,7 +468,7 @@ class MambaPool:
enable_linear_replayssm: bool = False,
linear_replayssm_cache_len: int = 16,
envelope_layout: bool = False,
enable_gdn_replayssm_spec: bool = False,
enable_linear_replayssm_spec: bool = False,
):
conv_state_shape = cache_params.shape.conv
temporal_state_shape = cache_params.shape.temporal
@@ -487,12 +487,13 @@ class MambaPool:
self.linear_replayssm_cache_len = linear_replayssm_cache_len
# ReplaySSM: the decode ring (--enable-linear-replayssm) allocates the
# chunked (d, k) records + write_pos; the spec-verify flag
# (--enable-gdn-replayssm-spec) always uses fold-every-commit and
# (--enable-linear-replayssm-spec) always uses fold-every-commit and
# allocates only the raw (v, k, g, beta) window -- no chunked records,
# no cursors. The shared g allocation gates on `_replayssm_on`.
self.enable_gdn_replayssm_spec = enable_gdn_replayssm_spec
self.replayssm_spec_fold = bool(enable_gdn_replayssm_spec)
_replayssm_on = enable_linear_replayssm or enable_gdn_replayssm_spec
# no cursors (KDA additionally keeps d/k, see the allocation below).
# The shared g allocation gates on `_replayssm_on`.
self.enable_linear_replayssm_spec = enable_linear_replayssm_spec
self.replayssm_spec_fold = bool(enable_linear_replayssm_spec)
_replayssm_on = enable_linear_replayssm or enable_linear_replayssm_spec
# for disagg with nvlink
self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
@@ -570,7 +571,7 @@ class MambaPool:
# flag is on; otherwise left as None so the legacy State is
# byte-identical. temporal_state_shape == (HV, V, K). Either the decode
# ring (--enable-linear-replayssm) or the spec-verify ring
# (--enable-gdn-replayssm-spec) shares this allocation.
# (--enable-linear-replayssm-spec) shares this allocation.
replayssm_d = replayssm_k = replayssm_g = None
replayssm_rawv = replayssm_rawk = replayssm_beta = None
if _replayssm_on:
@@ -580,16 +581,23 @@ class MambaPool:
num_slots = size + 1
# Ring dtype. DECODE ring (--enable-linear-replayssm): records
# follow the SSM dtype -- its flush folds `d` directly into the
# state. SPEC-verify ring (--enable-gdn-replayssm-spec): d/k feed
# state. SPEC-verify ring (--enable-linear-replayssm-spec): d/k feed
# ONLY the one-shot output reconstruction (the closed-loop exact
# fold replays the raw rings for state instead), so their
# quantization noise stays below the bf16 output cast; keep them
# in the conv/activation dtype instead of the (fp32-enforced)
# SSM dtype to halve the ring traffic. g stays fp32 everywhere
# (exact-fold input). The two flags are mutually exclusive.
ring_dtype = conv_dtype if enable_gdn_replayssm_spec else ssm_dtype
# Fold-every-commit: one verify window, no chunked (d, k) records.
if self.replayssm_spec_fold:
ring_dtype = conv_dtype if enable_linear_replayssm_spec else ssm_dtype
# Fold-every-commit: one verify window, no chunked (d, k)
# records. KDA is the exception on both counts: its window
# stays L-sized (the fused verify ring-write drops
# absorb-inflated rows past L), and d/k stay allocated --
# forward_decode routes on `replayssm_d is None` (fused vs
# decode-ring), so skipping them would flip KDA decode to the
# fused path, a behavior change needing its own validation
# (memory follow-up).
if self.replayssm_spec_fold and not cache_params.is_kda:
record_len = (
speculative_num_draft_tokens
if speculative_num_draft_tokens is not None
@@ -597,6 +605,7 @@ class MambaPool:
)
else:
record_len = L
if not self.replayssm_spec_fold or cache_params.is_kda:
replayssm_d = torch.zeros(
size=(num_mamba_layers, num_slots, hv, L, v_dim),
dtype=ring_dtype,
@@ -626,7 +635,22 @@ class MambaPool:
# flush replays these through the recurrent update sequentially
# (bit-identical to the recurrent baseline) instead of folding
# the chunked `d` records open-loop.
if enable_gdn_replayssm_spec:
if enable_linear_replayssm_spec:
if cache_params.is_kda:
# Backstop for the KDA ring invariants; this pool is
# sized with the final adaptive-aware draft maximum.
if L & (L - 1) != 0:
raise ValueError(
f"spec-verify ring length must be a power of two, got {L}"
)
if (
speculative_num_draft_tokens is not None
and L < 2 * speculative_num_draft_tokens
):
raise ValueError(
f"spec-verify ring too small: {L} < "
f"2 * {speculative_num_draft_tokens} (early-flush margin)"
)
replayssm_rawv = torch.zeros(
size=(num_mamba_layers, num_slots, hv, record_len, v_dim),
dtype=conv_dtype,
@@ -662,7 +686,11 @@ class MambaPool:
# The recurrent-verify fallback cannot be reached under the flag
# (GDN + linear chain + triton enforced in server_args; the
# backend asserts loudly if it ever is).
if enable_gdn_replayssm_spec:
# ReplaySSM skips this dominant scratch (~9GB @ K3 dspark γ=7): the
# KDA verify kernel takes intermediate_states_buffer=None (skips the
# per-step write, CACHE_INTERMEDIATE_STATES=False) and the commit
# replays the ring into the checkpoint instead. This is the memory win.
if enable_linear_replayssm_spec:
intermediate_ssm_state_cache = None
else:
intermediate_ssm_state_cache = torch.zeros(
@@ -795,7 +823,7 @@ class MambaPool:
f"rawv={get_tensor_size_bytes(replayssm_rawv) / GB:.3f}GB, "
f"rawk={get_tensor_size_bytes(replayssm_rawk) / GB:.3f}GB, "
f"beta={get_tensor_size_bytes(replayssm_beta) / GB:.3f}GB "
if enable_gdn_replayssm_spec
if enable_linear_replayssm_spec
else ""
)
)
@@ -818,12 +846,12 @@ class MambaPool:
# all GDN layers of one verify step; advanced by commit_gdn_replayssm_spec.
self.replayssm_cache_base = (
torch.zeros((size + 1,), dtype=torch.int32, device=device)
if enable_gdn_replayssm_spec and not self.replayssm_spec_fold
if enable_linear_replayssm_spec and not self.replayssm_spec_fold
else None
)
self.replayssm_is_flush = (
torch.zeros((size + 1,), dtype=torch.int8, device=device)
if enable_gdn_replayssm_spec and not self.replayssm_spec_fold
if enable_linear_replayssm_spec and not self.replayssm_spec_fold
else None
)
mem_usage_bytes = self.mamba_cache.mem_usage_bytes()
@@ -1130,7 +1158,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
enable_linear_replayssm: bool = False,
linear_replayssm_cache_len: int = 16,
mamba_envelope_layout: bool = False,
enable_gdn_replayssm_spec: bool = False,
enable_linear_replayssm_spec: bool = False,
):
super().__init__(
size=size,
@@ -1157,7 +1185,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
enable_linear_replayssm=enable_linear_replayssm,
linear_replayssm_cache_len=linear_replayssm_cache_len,
mamba_envelope_layout=mamba_envelope_layout,
enable_gdn_replayssm_spec=enable_gdn_replayssm_spec,
enable_linear_replayssm_spec=enable_linear_replayssm_spec,
)
def _init_mamba_pool(
@@ -1173,7 +1201,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
enable_linear_replayssm: bool = False,
linear_replayssm_cache_len: int = 16,
mamba_envelope_layout: bool = False,
enable_gdn_replayssm_spec: bool = False,
enable_linear_replayssm_spec: bool = False,
):
self.mamba_pool = self.mamba_pool_cls(
size=mamba_size,
@@ -1187,7 +1215,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
enable_linear_replayssm=enable_linear_replayssm,
linear_replayssm_cache_len=linear_replayssm_cache_len,
envelope_layout=mamba_envelope_layout,
enable_gdn_replayssm_spec=enable_gdn_replayssm_spec,
enable_linear_replayssm_spec=enable_linear_replayssm_spec,
)
self.mamba_allocator = MambaSlotAllocator(
size=mamba_size,
@@ -3663,7 +3691,7 @@ class HybridLinearKVPool(KVCache):
def get_kv_buffer_shape(self) -> Tuple[torch.Size, torch.Size]:
# Hybrid layer ids are global model-layer ids, while the backing pool
# is dense over only full-attention layers. Shape discovery does not
# is dense over only full-attention layers. Shape discovery does not
# need a global layer lookup, so delegate it to that backing pool.
return self.full_kv_pool.get_kv_buffer_shape()