[Feature] Add MiniCPM-SALA support (#30360)

Co-authored-by: Alex Nails <alex.nails@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/kernels/jit/csrc/minicpm_sala/get_block_table.cuh

@@ -0,0 +1,212 @@
+// MiniCPM-SALA sparse attention: build the per-token sparse block table.
+//
+// Migrated from `3rdparty/sparse_kernel/get_table_kernel.cu`. The original
+// CUDA kernels are kept almost verbatim; only the host-side wrappers are
+// rewritten from the torch::Tensor + pybind interface to the jit_kernel
+// tvm::ffi::TensorView + TensorMatcher/LaunchKernel convention.
+//
+// The Python wrapper compiles and caches one module per sparse layout.
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace sglang::minicpm_sala {
+
+constexpr int kTopkPerBlock = 16;
+
+// topk_idx:        [head_group, token_num, kSparseTopK]  int32
+// block_table:     [batch_size, seqlen_q_max]           int32
+// token_to_bs:     [token_num]                           int32
+// token_pos_in_bs: [token_num]                           int32
+// seqlen_q:        [batch_size]                          int32
+// out_block_table: [token_num, head_group, kSparseTopK * kSparseBlockSize] int32
+
+// 1 thread calc 64 element of out_block_table.
+// This allows topk_idx to be read once and all corresponding
+// out_block_table elements calculated, reducing memory access.
+template <int kSparseTopK, int kHeadGroup, int kSparseBlockSize>
+__global__ void get_block_table_cuda_blockwise(
+    const int* topk_idx,
+    const int* block_table,
+    const int* token_to_bs,
+    const int* token_pos_in_bs,
+    const int* seqlen_q,
+    int* out_block_table,
+    const int seqlen_q_max,
+    const int token_num) {
+  int token_idx = (blockIdx.x * blockDim.x + threadIdx.x) / (kSparseTopK * kHeadGroup);
+  if (token_idx >= token_num) return;
+  int head_group_idx = ((blockIdx.x * blockDim.x + threadIdx.x) / kSparseTopK) % kHeadGroup;
+  int topk_idx_in_head = (blockIdx.x * blockDim.x + threadIdx.x) % kSparseTopK;
+  int bs = token_to_bs[token_idx];
+  int pos_in_bs = token_pos_in_bs[token_idx];
+  int seqlen_q_bs = seqlen_q[bs];
+  int sparse_block_idx =
+      topk_idx[head_group_idx * token_num * kSparseTopK + token_idx * kSparseTopK + topk_idx_in_head];
+
+  auto out_view = reinterpret_cast<int (*)[kHeadGroup][kSparseTopK][kSparseBlockSize]>(out_block_table);
+  for (int i = 0; i < kSparseBlockSize; i++) {
+    int token_idx_in_batch = sparse_block_idx * kSparseBlockSize + i;
+    if (sparse_block_idx >= 0 && token_idx_in_batch < seqlen_q_bs && token_idx_in_batch < pos_in_bs) {
+      out_view[token_idx][head_group_idx][topk_idx_in_head][i] =
+          kHeadGroup * block_table[bs * seqlen_q_max + token_idx_in_batch] + head_group_idx;
+    } else {
+      out_view[token_idx][head_group_idx][topk_idx_in_head][i] = 0;
+    }
+  }
+}
+
+// 1 thread calculates 1 element of out_block_table. A 1024-thread block
+// expands 16 selected blocks in parallel when kSparseBlockSize is 64.
+template <int kSparseTopK, int kHeadGroup, int kSparseBlockSize>
+__global__ void get_block_table_cuda_elementwise(
+    const int* topk_idx,
+    const int* block_table,
+    const int* token_to_bs,
+    const int* token_pos_in_bs,
+    const int* seqlen_q,
+    int* out_block_table,
+    const int seqlen_q_max,
+    const int token_num) {
+  constexpr int kBlockPerTokenHead = kSparseTopK / kTopkPerBlock;
+  // calc 16 topk -> 1024 output
+  __shared__ int topk_idx_share[kTopkPerBlock];
+  const int tidx = threadIdx.x;
+  const int bidx = blockIdx.x;
+
+  if (threadIdx.x < kTopkPerBlock) {
+    topk_idx_share[tidx] = topk_idx[bidx * kTopkPerBlock + tidx];
+  }
+
+  __syncthreads();
+
+  const int head_group_idx = (bidx / kBlockPerTokenHead) / token_num;
+  const int token_idx = (bidx / kBlockPerTokenHead) % token_num;
+  const int topk_idx_in_head = bidx % kBlockPerTokenHead * kTopkPerBlock + tidx / kSparseBlockSize;
+
+  const int sparse_block_idx = topk_idx_share[tidx / kSparseBlockSize];
+
+  const int token_idx_src = sparse_block_idx * kSparseBlockSize + tidx % kSparseBlockSize;
+  const int token_idx_dst = token_idx * kHeadGroup * kSparseTopK * kSparseBlockSize +
+                            head_group_idx * kSparseTopK * kSparseBlockSize + topk_idx_in_head * kSparseBlockSize +
+                            tidx % kSparseBlockSize;
+  if (sparse_block_idx < 0) {
+    out_block_table[token_idx_dst] = 0;
+    return;
+  }
+
+  const int bs = token_to_bs[token_idx];
+  const int pos_in_bs = token_pos_in_bs[token_idx];
+  const int seqlen_q_bs = seqlen_q[bs];
+
+  if (token_idx_src < seqlen_q_bs && token_idx_src < pos_in_bs) {
+    out_block_table[token_idx_dst] = kHeadGroup * block_table[bs * seqlen_q_max + token_idx_src] + head_group_idx;
+  } else {
+    out_block_table[token_idx_dst] = 0;
+  }
+}
+
+// Validate all inputs that are shared across the two kernel variants and
+// bind the symbolic dims (token_num / batch_size / seqlen_q_max). The output
+// tensor is pre-allocated and fully initialized by the selected kernel.
+template <int kSparseTopK, int kHeadGroup, int kSparseBlockSize>
+void verify_inputs(
+    tvm::ffi::TensorView out,
+    tvm::ffi::TensorView topk_idx,
+    tvm::ffi::TensorView block_table,
+    tvm::ffi::TensorView token_to_bs,
+    tvm::ffi::TensorView token_pos_in_bs,
+    tvm::ffi::TensorView seqlen_q,
+    host::SymbolicSize& token_num,
+    host::SymbolicSize& batch_size,
+    host::SymbolicSize& seqlen_q_max,
+    host::SymbolicDevice& device) {
+  using namespace host;
+  constexpr int64_t kOutLastDim = static_cast<int64_t>(kSparseTopK) * kSparseBlockSize;
+
+  // topk_idx: [kHeadGroup, token_num, kSparseTopK]
+  TensorMatcher({static_cast<int64_t>(kHeadGroup), token_num, static_cast<int64_t>(kSparseTopK)})
+      .with_dtype<int32_t>()
+      .with_device<kDLCUDA>(device)
+      .verify(topk_idx);
+  // block_table: [batch_size, seqlen_q_max]
+  TensorMatcher({batch_size, seqlen_q_max})  //
+      .with_dtype<int32_t>()
+      .with_device<kDLCUDA>(device)
+      .verify(block_table);
+  // token_to_bs / token_pos_in_bs: [token_num]
+  TensorMatcher({token_num})  //
+      .with_dtype<int32_t>()
+      .with_device<kDLCUDA>(device)
+      .verify(token_to_bs)
+      .verify(token_pos_in_bs);
+  // seqlen_q: [batch_size]
+  TensorMatcher({batch_size})  //
+      .with_dtype<int32_t>()
+      .with_device<kDLCUDA>(device)
+      .verify(seqlen_q);
+  // out: [token_num, kHeadGroup, kSparseTopK * kSparseBlockSize]
+  TensorMatcher({token_num, static_cast<int64_t>(kHeadGroup), kOutLastDim})
+      .with_dtype<int32_t>()
+      .with_device<kDLCUDA>(device)
+      .verify(out);
+}
+
+template <bool kElementwise, int kSparseTopK, int kHeadGroup, int kSparseBlockSize>
+void get_block_table(
+    tvm::ffi::TensorView out,
+    tvm::ffi::TensorView topk_idx,
+    tvm::ffi::TensorView block_table,
+    tvm::ffi::TensorView token_to_bs,
+    tvm::ffi::TensorView token_pos_in_bs,
+    tvm::ffi::TensorView seqlen_q) {
+  using namespace host;
+  SymbolicSize token_num{"token_num"}, batch_size{"batch_size"}, seqlen_q_max{"seqlen_q_max"};
+  SymbolicDevice device;
+  device.set_options<kDLCUDA>();
+  verify_inputs<kSparseTopK, kHeadGroup, kSparseBlockSize>(
+      out, topk_idx, block_table, token_to_bs, token_pos_in_bs, seqlen_q, token_num, batch_size, seqlen_q_max, device);
+
+  const int n_token = static_cast<int>(token_num.unwrap());
+  const int s_q_max = static_cast<int>(seqlen_q_max.unwrap());
+  const DLDevice dev = device.unwrap();
+
+  constexpr int kThreadsPerBlock = 1024;
+  constexpr int kElementsPerEntry = kElementwise ? kSparseBlockSize : 1;
+  const int64_t total = static_cast<int64_t>(n_token) * kHeadGroup * kSparseTopK * kElementsPerEntry;
+  const int64_t num_blocks = (total + kThreadsPerBlock - 1) / kThreadsPerBlock;
+
+  if constexpr (!kElementwise) {
+    LaunchKernel(num_blocks, kThreadsPerBlock, dev)(
+        get_block_table_cuda_blockwise<kSparseTopK, kHeadGroup, kSparseBlockSize>,
+        static_cast<const int*>(topk_idx.data_ptr()),
+        static_cast<const int*>(block_table.data_ptr()),
+        static_cast<const int*>(token_to_bs.data_ptr()),
+        static_cast<const int*>(token_pos_in_bs.data_ptr()),
+        static_cast<const int*>(seqlen_q.data_ptr()),
+        static_cast<int*>(out.data_ptr()),
+        s_q_max,
+        n_token);
+  } else {
+    LaunchKernel(num_blocks, kThreadsPerBlock, dev)(
+        get_block_table_cuda_elementwise<kSparseTopK, kHeadGroup, kSparseBlockSize>,
+        static_cast<const int*>(topk_idx.data_ptr()),
+        static_cast<const int*>(block_table.data_ptr()),
+        static_cast<const int*>(token_to_bs.data_ptr()),
+        static_cast<const int*>(token_pos_in_bs.data_ptr()),
+        static_cast<const int*>(seqlen_q.data_ptr()),
+        static_cast<int*>(out.data_ptr()),
+        s_q_max,
+        n_token);
+  }
+}
+
+}  // namespace sglang::minicpm_sala

python/sglang/kernels/jit/minicpm_sala/__init__.py

@@ -0,0 +1,3 @@
+from sglang.kernels.jit.minicpm_sala.get_block_table import get_block_table
+
+__all__ = ["get_block_table"]

python/sglang/kernels/jit/minicpm_sala/get_block_table.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_get_block_table_module(
+    topk: int, head_group_num: int, block_size: int
+) -> Module:
+    """Compile and cache the JIT module for a given sparse topk value.
+
+    One module is built per topk value, replacing the original runtime
+    ``VALUE_SPLITS_SWITCH(topk, ...)`` dispatch with a compile-time template
+    argument ``kSparseTopK``.
+    """
+    args = make_cpp_args(topk, head_group_num, block_size)
+    wrappers = [
+        (
+            "get_block_table_blockwise",
+            f"minicpm_sala::get_block_table<false, {args}>",
+        ),
+    ]
+    if block_size == 64 and topk % 16 == 0:
+        wrappers.append(
+            (
+                "get_block_table_elementwise",
+                f"minicpm_sala::get_block_table<true, {args}>",
+            )
+        )
+    return load_jit(
+        f"get_block_table_strategies_topk{topk}_g{head_group_num}_b{block_size}",
+        *args,
+        cuda_files=["minicpm_sala/get_block_table.cuh"],
+        cuda_wrappers=wrappers,
+    )
+
+
+def get_block_table(
+    topk_idx: torch.Tensor,
+    block_table: torch.Tensor,
+    token_to_bs: torch.Tensor,
+    token_pos_in_bs: torch.Tensor,
+    seqlen_q: torch.Tensor,
+    head_group_num: int = 2,
+    block_size: int = 64,
+    *,
+    elementwise: bool,
+) -> torch.Tensor:
+    if topk_idx.dim() != 3:
+        raise RuntimeError(
+            f"topk_idx must be 3D [head_group, token_num, topk], got shape {tuple(topk_idx.shape)}"
+        )
+    token_num = topk_idx.shape[1]
+    topk = topk_idx.shape[2]
+    if topk <= 0 or block_size <= 0:
+        raise RuntimeError(
+            f"topk and block_size must be positive, got {topk=} and {block_size=}"
+        )
+    kernel_name = (
+        "get_block_table_elementwise"
+        if elementwise and block_size == 64 and topk % 16 == 0
+        else "get_block_table_blockwise"
+    )
+
+    out = torch.empty(
+        (token_num, head_group_num, topk * block_size),
+        dtype=torch.int32,
+        device=topk_idx.device,
+    )
+    module = _jit_get_block_table_module(topk, head_group_num, block_size)
+    getattr(module, kernel_name)(
+        out, topk_idx, block_table, token_to_bs, token_pos_in_bs, seqlen_q
+    )
+    return out

python/sglang/srt/arg_groups/overrides.py

@@ -1217,6 +1217,54 @@ def _moss_vl_overrides(server_args: Any, hf_config: Any) -> dict:
     return overrides
 
 
+@_register_for("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM")
+def _minicpm_sala_overrides(server_args: Any, hf_config: Any) -> dict:
+    if server_args.enable_dp_attention:
+        raise ValueError("MiniCPM does not support DP attention")
+    has_sparse_attention = getattr(hf_config, "has_minicpm_sparse_attention", False)
+    has_hybrid_attention = has_sparse_attention or getattr(
+        hf_config, "has_lightning_layers", False
+    )
+    overrides: Dict[str, Any] = {}
+    if has_hybrid_attention:
+        if server_args.enable_hierarchical_cache:
+            raise ValueError("MiniCPM SALA does not support hierarchical cache")
+        overrides["disable_radix_cache"] = True
+    if envs.SGLANG_MINICPM_FORCE_DENSE.get():
+        dense_backends = {
+            "minicpm_flashattn": ("fa4" if is_blackwell_supported() else "fa3"),
+            "minicpm_flashinfer": "flashinfer",
+        }
+        for backend_field in (
+            "attention_backend",
+            "prefill_attention_backend",
+            "decode_attention_backend",
+        ):
+            dense_backend = dense_backends.get(getattr(server_args, backend_field))
+            if dense_backend is not None:
+                overrides[backend_field] = dense_backend
+    elif has_sparse_attention:
+        uses_sparse_backend = server_args.is_attention_backend_not_set() or any(
+            backend in ("minicpm_flashattn", "minicpm_flashinfer")
+            for backend in (
+                server_args.attention_backend,
+                server_args.prefill_attention_backend,
+                server_args.decode_attention_backend,
+            )
+        )
+        if uses_sparse_backend and server_args.disaggregation_mode != "null":
+            raise ValueError(
+                "MiniCPM sparse attention does not support PD disaggregation"
+            )
+        if server_args.is_attention_backend_not_set():
+            overrides["attention_backend"] = (
+                "minicpm_flashinfer"
+                if is_blackwell_supported()
+                else "minicpm_flashattn"
+            )
+    return overrides
+
+
 @_register_for("MiniCPMV4_6ForConditionalGeneration")
 def _minicpm_v4_6_overrides(server_args: Any, hf_config: Any) -> dict:
     if is_sm100_supported() and server_args.attention_backend is None:

python/sglang/srt/configs/__init__.py

@@ -36,6 +36,7 @@ from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig
 from sglang.srt.configs.lfm2_vl import Lfm2VlConfig
 from sglang.srt.configs.locate_anything import LocateAnythingConfig
 from sglang.srt.configs.longcat_flash import LongcatFlashConfig
+from sglang.srt.configs.minicpm import MiniCPMHybridConfig
 from sglang.srt.configs.minicpmv4_6 import MiniCPMV4_6Config, MiniCPMV4_6VisionConfig
 from sglang.srt.configs.minimax_vl import MiniMaxM3VLConfig
 from sglang.srt.configs.muse_glimmer import (
@@ -113,6 +114,7 @@ __all__ = [
     "NemotronH_Nano_Omni_Reasoning_V3_Config",
     "JetNemotronConfig",
     "JetVLMConfig",
+    "MiniCPMHybridConfig",
     "Step3p5Config",
     "MiniMaxM3VLConfig",
     "Step3p7Config",

python/sglang/srt/configs/hybrid_arch.py

@@ -15,6 +15,7 @@ from sglang.srt.configs import (
     Lfm2Config,
     Lfm2MoeConfig,
     Lfm2VlConfig,
+    MiniCPMHybridConfig,
     NemotronH_Nano_VL_V2_Config,
     NemotronHConfig,
     Qwen3_5Config,
@@ -42,6 +43,8 @@ def hybrid_lightning_config(model_config: ModelConfig):
     config = model_config.hf_config
     if isinstance(config, BailingHybridConfig):
         return config
+    if isinstance(config, MiniCPMHybridConfig) and config.has_lightning_layers:
+        return config
     return None
 
 

python/sglang/srt/configs/minicpm.py

@@ -0,0 +1,194 @@
+from transformers import PretrainedConfig
+
+from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
+from sglang.srt.runtime_context import get_parallel
+
+_MIXER_TYPE_ALIASES = {
+    "minicpm4": "minicpm4",
+    "minicpm": "minicpm4",
+    "standard": "minicpm4",
+    "attention": "minicpm4",
+    "attn": "minicpm4",
+    "lightning": "lightning-attn",
+    "lightning_attn": "lightning-attn",
+    "lightning-attn": "lightning-attn",
+}
+
+
+class MiniCPMHybridConfig(PretrainedConfig):
+    """
+    Configuration class for hybrid MiniCPM models.
+
+    This config extends PretrainedConfig to match the pattern used by other
+    hybrid/linear attention models (Falcon H1, Nemotron H, Kimi Linear, etc.)
+    and provides cache parameters for the Simple GLA attention mechanism.
+    """
+
+    model_type = "minicpm_sala"
+
+    def __init__(
+        self,
+        # Base model config fields
+        vocab_size=150528,
+        hidden_size=4096,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        head_dim=None,
+        hidden_act="silu",
+        intermediate_size=14336,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        max_position_embeddings=32768,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        scale_emb=12,
+        scale_depth=1.4,
+        dim_model_base=256,
+        # MiniCPM-specific hybrid config fields
+        mixer_types=None,
+        lightning_nh=None,
+        lightning_nkv=None,
+        lightning_head_dim=None,
+        lightning_scale="1/sqrt(d)",
+        lightning_layerwise_decay=False,
+        lightning_use_rope=True,
+        use_output_gate=False,
+        attention_bias=False,
+        use_output_norm=False,
+        qk_norm=True,
+        attn_use_rope=True,
+        attn_use_output_gate=False,
+        sparse_config=None,
+        **kwargs,
+    ):
+        for unused_field in ("minicpm4", "lightning", "sparse_use_nope"):
+            kwargs.pop(unused_field, None)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = (
+            head_dim if head_dim is not None else hidden_size // num_attention_heads
+        )
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.scale_emb = scale_emb
+        self.scale_depth = scale_depth
+        self.dim_model_base = dim_model_base
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        # Hybrid config fields
+        if not mixer_types:
+            mixer_types = ["minicpm4"]
+        elif len(mixer_types) > num_hidden_layers:
+            raise ValueError(f"Invalid number of mixer types: {len(mixer_types)}")
+        try:
+            mixer_types = [
+                _MIXER_TYPE_ALIASES[mixer_type] for mixer_type in mixer_types
+            ]
+        except KeyError as exc:
+            raise ValueError(f"Unsupported mixer type: {exc.args[0]}") from exc
+        repeats = (num_hidden_layers + len(mixer_types) - 1) // len(mixer_types)
+        self.mixer_types = (mixer_types * repeats)[:num_hidden_layers]
+        self.lightning_nh = (
+            lightning_nh if lightning_nh is not None else num_attention_heads
+        )
+        self.lightning_nkv = (
+            lightning_nkv if lightning_nkv is not None else num_key_value_heads
+        )
+        self.lightning_head_dim = (
+            lightning_head_dim if lightning_head_dim is not None else self.head_dim
+        )
+        if (
+            "lightning-attn" in self.mixer_types
+            and self.lightning_nh != self.lightning_nkv
+        ):
+            raise ValueError(
+                "MiniCPM Lightning attention requires equal query and KV head "
+                "counts because the seg_la backend does not support GQA: "
+                f"lightning_nh={self.lightning_nh}, "
+                f"lightning_nkv={self.lightning_nkv}"
+            )
+        self.lightning_scale = lightning_scale
+        self.lightning_layerwise_decay = lightning_layerwise_decay
+        self.lightning_use_rope = lightning_use_rope
+        self.use_output_gate = use_output_gate
+        self.attention_bias = attention_bias
+        self.use_output_norm = use_output_norm
+        self.qk_norm = qk_norm
+        self.attn_use_rope = attn_use_rope
+        self.attn_use_output_gate = attn_use_output_gate
+        self.sparse_config = sparse_config
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def num_linear_key_value_heads(self) -> int:
+        return self.lightning_nkv
+
+    @property
+    def mamba2_cache_params(self):
+        """Return linear-attention cache parameters for lightning layers."""
+        lightning_layer_ids = self.lightning_layer_ids
+
+        if (
+            not lightning_layer_ids
+            or not self.lightning_nkv
+            or not self.lightning_head_dim
+        ):
+            return None
+
+        shape = Mamba2StateShape.create(
+            tp_world_size=get_parallel().attn_tp_size,
+            intermediate_size=0,
+            n_groups=0,
+            num_heads=self.lightning_nkv,
+            head_dim=self.lightning_head_dim,
+            state_size=self.lightning_head_dim,
+            conv_kernel=1,
+        )
+
+        return Mamba2CacheParams(shape=shape, layers=lightning_layer_ids)
+
+    @property
+    def full_attention_layer_ids(self):
+        return [
+            i
+            for i, mixer_type in enumerate(self.mixer_types)
+            if mixer_type == "minicpm4"
+        ]
+
+    @property
+    def has_minicpm_sparse_attention(self) -> bool:
+        """Check if this config has MiniCPM sparse attention layers."""
+        return self.sparse_config is not None and any(
+            mt == "minicpm4" for mt in self.mixer_types
+        )
+
+    @property
+    def has_lightning_layers(self) -> bool:
+        """Check if this config has lightning attention layers."""
+        return any(mt == "lightning-attn" for mt in self.mixer_types)
+
+    @property
+    def lightning_layer_ids(self) -> list:
+        """Get the indices of layers with lightning attention."""
+        return [i for i, mt in enumerate(self.mixer_types) if mt == "lightning-attn"]

python/sglang/srt/disaggregation/decode.py

@@ -162,6 +162,7 @@ class DecodeReqToTokenPool:
         # here: HybridMambaDecodeReqToTokenPool borrows this __init__ while
         # inheriting ReqToTokenPool.alloc, which bumps it.
         self.req_generation = torch.zeros(self._alloc_size, dtype=torch.int64)
+        self._aux_cache: Any = None
 
     def write(self, indices, values):
         self.req_to_token[indices] = values
@@ -169,6 +170,20 @@ class DecodeReqToTokenPool:
     def available_size(self):
         return len(self.free_slots)
 
+    def reset_aux_cache_allocator(self) -> None:
+        pass
+
+    def schedulable_token_capacity(self, physical_capacity: int) -> int:
+        return physical_capacity
+
+    def alloc_aux_to_lengths(
+        self,
+        *,
+        req_pool_indices_cpu: torch.Tensor,
+        target_seq_lens_cpu: torch.Tensor,
+    ) -> None:
+        pass
+
     def alloc(self, reqs: List[Req]) -> Optional[List[int]]:
         # Indices of reqs that already have a req_pool_idx and will reuse
         # their existing slot (e.g. chunked prefill continuing across chunks).

python/sglang/srt/environ.py

@@ -1074,6 +1074,11 @@ class Envs:
     # ===================================================================
     # Kernel selection and fused backends
     # ===================================================================
+    # MiniCPM sparse attention developer switches
+    SGLANG_MINICPM_FUSE_TOPK = EnvBool(False)
+    SGLANG_MINICPM_DENSE_AS_SPARSE = EnvBool(False)
+    SGLANG_MINICPM_FORCE_DENSE = EnvBool(False)
+
     SGLANG_USE_SGL_FA3_KERNEL = EnvBool(True)
     # Force every sglang.kernels BaseFusedOp onto one backend (a KernelBackend
     # value, e.g. "torch" / "torch_compile" / "triton" / "aot"); unset =

python/sglang/srt/layers/attention/attention_registry.py

@@ -321,6 +321,20 @@ def attn_backend_wrapper_for_draft_decode(runner: "ModelRunner", backend):
     return backend
 
 
+@register_attention_backend("minicpm_flashattn")
+def create_minicpm_flashattn_backend(runner):
+    from sglang.srt.layers.attention.minicpm.backend import MiniCPMSparseBackend
+
+    return MiniCPMSparseBackend(runner, use_flashinfer=False)
+
+
+@register_attention_backend("minicpm_flashinfer")
+def create_minicpm_flashinfer_backend(runner):
+    from sglang.srt.layers.attention.minicpm.backend import MiniCPMSparseBackend
+
+    return MiniCPMSparseBackend(runner, use_flashinfer=True)
+
+
 def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBackend"):
     """
     Wrapper for special models like hybrid GDN, so we don't

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -1170,6 +1170,61 @@ class FlashAttentionBackend(AttentionBackend):
 
         self.forward_metadata = metadata
 
+    def get_paged_mha_kv_cache(
+        self,
+        layer: RadixAttention,
+        *,
+        head_group_num: int = 1,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        key_cache, value_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+        return (
+            key_cache.view(
+                -1,
+                self.page_size,
+                layer.tp_k_head_num // head_group_num,
+                layer.head_dim,
+            ),
+            value_cache.view(
+                -1,
+                self.page_size,
+                layer.tp_v_head_num // head_group_num,
+                layer.v_head_dim,
+            ),
+        )
+
+    def prepare_paged_mha_query(
+        self,
+        q: torch.Tensor,
+        q_rope: Optional[torch.Tensor],
+        k_rope: Optional[torch.Tensor],
+        layer: RadixAttention,
+        *,
+        logical_batch_size: int,
+        kv_head_num: int,
+        is_prefill: bool,
+    ) -> tuple[
+        torch.Tensor,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+    ]:
+        k_descale = v_descale = None
+        if (
+            self.kv_cache_dtype_str != "auto"
+            and layer.head_dim <= 256
+            and not self.kv_cache_is_mxfp8
+            and (not is_prefill or self.fa_impl_ver != 4)
+        ):
+            if layer.k_scale is not None:
+                descale_shape = (logical_batch_size, kv_head_num)
+                k_descale = layer.k_scale.expand(descale_shape)
+                v_descale = layer.v_scale.expand(descale_shape)
+            q = q.to(self.kv_cache_dtype)
+            q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None
+            k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None
+        return q, q_rope, k_rope, k_descale, v_descale
+
     def forward_extend(
         self,
         q: torch.Tensor,
@@ -1282,24 +1337,15 @@ class FlashAttentionBackend(AttentionBackend):
             if is_swa_layer
             else (-1, -1)
         )
-        fa_k_descale, fa_v_descale = None, None
+        q, q_rope, k_rope, fa_k_descale, fa_v_descale = self.prepare_paged_mha_query(
-        # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
+            q,
-        # has corresponding quantization method so that layer.k_scale is not None,
+            q_rope,
-        # 3) layer.head_dim <= 256 since fa3 kernel require fp16 and bf16 data type in this case,
+            k_rope,
-        # 4) fa_impl_ver != 4 since fa4 does not currently support fp8 queries and keys.
+            layer,
-        if (
+            logical_batch_size=forward_batch.batch_size,
-            self.kv_cache_dtype_str != "auto"
+            kv_head_num=layer.tp_k_head_num,
-            and layer.head_dim <= 256
+            is_prefill=True,
-            and self.fa_impl_ver != 4
+        )
-            and not self.kv_cache_is_mxfp8
-        ):
-            if layer.k_scale is not None:
-                descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
-                fa_k_descale = layer.k_scale.expand(descale_shape)
-                fa_v_descale = layer.v_scale.expand(descale_shape)
-            q = q.to(self.kv_cache_dtype)
-            q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None
-            k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None
         # Check if we should use local attention
         use_local_attn = (
             self.has_local_attention
@@ -1385,13 +1431,8 @@ class FlashAttentionBackend(AttentionBackend):
         # Use Flash Attention for prefill
         if not self.use_mla:
             # Do multi-head attention
-            key_cache, value_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+            key_cache, value_cache = self.get_paged_mha_kv_cache(
-
+                layer,
-            key_cache = key_cache.view(
-                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
-            )
-            value_cache = value_cache.view(
-                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
             )
             if layer.is_cross_attention:
                 page_table = metadata.encoder_page_table
@@ -1869,34 +1910,23 @@ class FlashAttentionBackend(AttentionBackend):
             else None
         )
 
-        fa_k_descale, fa_v_descale = None, None
+        q, q_rope, k_rope, fa_k_descale, fa_v_descale = self.prepare_paged_mha_query(
-        # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
+            q,
-        # has corresponding quantization method so that layer.k_scale is not None,
+            q_rope,
-        # 3) layer.head_dim <= 256 since fa3 kernel require fp16 and bf16 data type in this case.
+            k_rope,
-        if (
+            layer,
-            self.kv_cache_dtype_str != "auto"
+            logical_batch_size=forward_batch.batch_size,
-            and layer.head_dim <= 256
+            kv_head_num=layer.tp_k_head_num,
-            and not self.kv_cache_is_mxfp8
+            is_prefill=False,
-        ):
+        )
-            if layer.k_scale is not None:
-                descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
-                fa_k_descale = layer.k_scale.expand(descale_shape)
-                fa_v_descale = layer.v_scale.expand(descale_shape)
-            q = q.to(self.kv_cache_dtype)
-            q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None
-            k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None
         if fa_k_descale is not None:
             kwargs["k_descale"] = fa_k_descale
             kwargs["v_descale"] = fa_v_descale
         if not self.use_mla:
             # Do multi-head attention
 
-            key_cache, value_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+            key_cache, value_cache = self.get_paged_mha_kv_cache(
-            key_cache = key_cache.view(
+                layer,
-                -1, self.page_size, layer.tp_k_head_num, layer.head_dim
-            )
-            value_cache = value_cache.view(
-                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
             )
 
             if layer.is_cross_attention:

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -1061,6 +1061,23 @@ class FlashInferAttnBackend(AttentionBackend):
             for i in range(self.num_wrappers)
         ]
 
+    def get_cuda_graph_decode_wrappers(
+        self,
+        *,
+        bs: int,
+        num_tokens: int,
+    ) -> list:
+        wrappers = self.decode_cuda_graph_metadata.get(bs)
+        if wrappers is None:
+            self._prepare_cuda_graph_metadata(
+                bs,
+                num_tokens,
+                ForwardMode.DECODE,
+                spec_info=None,
+            )
+            wrappers = self.decode_cuda_graph_metadata[bs]
+        return wrappers
+
     def _create_prefill_wrappers(self, bs: int, use_custom_mask: bool = False) -> list:
         # FlashInfer's prefill wrapper decides mask mode based on whether
         # `custom_mask_buf` is initialized (not whether a custom mask is provided).

python/sglang/srt/layers/attention/linear/lightning_backend.py

@@ -74,14 +74,20 @@ class LightningAttentionBackend(MambaAttnBackendBase):
             if hasattr(model_runner.model_config, "block")
             else 256
         )
-        total_num_heads = model_runner.model_config.hf_config.num_attention_heads
+        config = model_runner.model_config.hf_config
-        num_hidden_layers = model_runner.model_config.hf_config.num_hidden_layers
+        total_num_heads = getattr(
+            config, "num_linear_key_value_heads", config.num_attention_heads
+        )
+        layerwise_decay = getattr(config, "lightning_layerwise_decay", True)
+        assert total_num_heads % get_parallel().attn_tp_size == 0
+        num_hidden_layers = config.num_hidden_layers
         self.tp_slope = LightningAttentionBackend._build_slope_tensor(
-            total_num_heads, num_hidden_layers, self.device
+            total_num_heads,
-        )
+            num_hidden_layers,
-        self.linear_backend = getattr(
+            self.device,
-            model_runner.model_config.hf_config, "linear_backend", "seg_la"
+            layerwise_decay=layerwise_decay,
         )
+        self.linear_backend = getattr(config, "linear_backend", "seg_la")
         logger.info(
             f"linear_backend for linear attention in hybrid_linear_backend: {self.linear_backend}"
         )
@@ -130,7 +136,10 @@ class LightningAttentionBackend(MambaAttnBackendBase):
 
     @staticmethod
     def _build_slope_tensor(
-        n_attention_heads: int, num_hidden_layers: int, device="cuda"
+        n_attention_heads: int,
+        num_hidden_layers: int,
+        device="cuda",
+        layerwise_decay: bool = True,
     ):
         def get_slopes(n):
             def get_slopes_power_of_2(n):
@@ -153,7 +162,9 @@ class LightningAttentionBackend(MambaAttnBackendBase):
 
         tp_heads = n_attention_heads // get_parallel().attn_tp_size
         tp_rank = get_parallel().attn_tp_rank
-        if num_hidden_layers <= 1:
+        if not layerwise_decay:
+            slope_rate_list = [slopes] * num_hidden_layers
+        elif num_hidden_layers <= 1:
             slope_rate_list = [slopes * (1 + 1e-5)]
         else:
             slope_rate_list = [
@@ -285,6 +296,7 @@ class LightningAttentionBackend(MambaAttnBackendBase):
             cache_indices=intermediate_state_indices,
             track_lens=track_lens,
             track_state_indices=track_state_indices,
+            softmax_scale=layer.scaling,
             decouple=True,
         )
         return hidden

python/sglang/srt/layers/attention/minicpm/attention_adapter.py

@@ -0,0 +1,288 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.ops.attention.flash_attention import flash_attn_with_kvcache
+from sglang.kernels.ops.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
+from sglang.srt.layers.attention.minicpm.sparse_utils import (
+    MiniCPMSparseMetadata,
+)
+from sglang.srt.utils import is_flashinfer_available
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+class MiniCPMFlashAttentionAdapter:
+    def __init__(self, flash_attn_backend: FlashAttentionBackend):
+        self.flash_attn_backend = flash_attn_backend
+
+    def prepare_forward(
+        self,
+        metadata: MiniCPMSparseMetadata,
+        *,
+        is_prefill: bool,
+        graph: bool,
+    ) -> None:
+        pass
+
+    def init_cuda_graph_state(self, max_num_tokens: int) -> None:
+        pass
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        metadata: MiniCPMSparseMetadata,
+        layer: RadixAttention,
+        *,
+        is_prefill: bool,
+        k_descale: Optional[torch.Tensor] = None,
+        v_descale: Optional[torch.Tensor] = None,
+        sinks: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        kwargs = {}
+        if sinks is not None:
+            kwargs["sinks"] = sinks
+        if k_descale is not None:
+            kwargs["k_descale"] = k_descale
+            kwargs["v_descale"] = v_descale
+
+        return flash_attn_with_kvcache(
+            q=q,
+            k_cache=key_cache,
+            v_cache=value_cache,
+            page_table=metadata.sparse_page_table,
+            cache_seqlens=metadata.sparse_cache_seqlens_int32,
+            cu_seqlens_q=metadata.sparse_cu_seqlens_q,
+            cu_seqlens_k_new=metadata.sparse_cu_seqlens_k,
+            max_seqlen_q=(
+                metadata.sparse_max_seq_len_q
+                if is_prefill
+                else metadata.base.max_seq_len_q
+            ),
+            softmax_scale=layer.scaling,
+            causal=True,
+            window_size=(-1, -1),
+            softcap=layer.logit_cap,
+            num_splits=self.flash_attn_backend.num_splits,
+            ver=self.flash_attn_backend.fa_impl_ver,
+            **kwargs,
+        )
+
+
+class MiniCPMFlashInferAdapter:
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        *,
+        head_group_num: int,
+        heads_per_group: int,
+        head_dim: int,
+        page_size: int,
+        max_kv_tokens_per_row: int,
+    ):
+        if not is_flashinfer_available():
+            raise RuntimeError("minicpm_flashinfer requires the flashinfer package.")
+
+        from sglang.srt.layers.attention.flashinfer_backend import (
+            FlashInferAttnBackend,
+        )
+
+        self.device = model_runner.device
+        self.head_group_num = head_group_num
+        self.num_qo_heads = heads_per_group
+        self.num_kv_heads = 1
+        self.head_dim = head_dim
+        self.page_size = page_size
+        self.max_kv_tokens_per_row = max_kv_tokens_per_row
+        self.q_dtype = model_runner.dtype
+        self.kv_dtype = model_runner.kv_cache_dtype
+
+        max_sparse_bs = model_runner.req_to_token_pool.size * head_group_num
+        self.kv_indptr = torch.zeros(
+            max_sparse_bs + 1,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.kv_indices = torch.zeros(
+            max_sparse_bs * max_kv_tokens_per_row,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.kv_last_page_len = torch.ones(
+            max_sparse_bs,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.rows = torch.arange(
+            max_sparse_bs,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.flashinfer_backend = FlashInferAttnBackend(
+            model_runner,
+            skip_prefill=False,
+            kv_indptr_buf=self.kv_indptr,
+            kv_last_page_len_buf=self.kv_last_page_len,
+        )
+        self.active_wrapper = None
+        self.active_kv_indptr = None
+        self.active_kv_indices = None
+        self.active_rows = None
+        self.prefill_planned = False
+
+    def prepare_forward(
+        self,
+        metadata: MiniCPMSparseMetadata,
+        *,
+        is_prefill: bool,
+        graph: bool,
+    ) -> None:
+        if is_prefill:
+            self.prefill_planned = False
+        else:
+            self._prepare(
+                metadata,
+                is_prefill=False,
+                graph=graph,
+            )
+
+    def init_cuda_graph_state(self, max_num_tokens: int) -> None:
+        self.flashinfer_backend.init_cuda_graph_state(
+            max_num_tokens,
+            max_num_tokens,
+            kv_indices_buf=self.kv_indices,
+        )
+
+    def _prepare(
+        self,
+        metadata: MiniCPMSparseMetadata,
+        *,
+        is_prefill: bool,
+        graph: bool = False,
+    ) -> None:
+        cache_seqlens = metadata.sparse_cache_seqlens_int32
+        sparse_bs = cache_seqlens.numel()
+        if sparse_bs == 0:
+            self.active_wrapper = None
+            return
+
+        if is_prefill:
+            kv_indptr = metadata.sparse_cu_seqlens_k
+            kv_indices = torch.empty(
+                metadata.sparse_page_table.numel(),
+                dtype=torch.int32,
+                device=self.device,
+            )
+            kv_last_page_len = (cache_seqlens > 0).to(torch.int32)
+            rows = torch.arange(
+                sparse_bs,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            wrapper = self.flashinfer_backend.prefill_wrappers_paged[0]
+            wrapper.begin_forward(
+                metadata.sparse_cu_seqlens_q,
+                kv_indptr,
+                kv_indices,
+                kv_last_page_len,
+                self.num_qo_heads,
+                self.num_kv_heads,
+                self.head_dim,
+                self.page_size,
+                causal=True,
+                q_data_type=self.q_dtype,
+                kv_data_type=self.kv_dtype,
+                non_blocking=True,
+            )
+        else:
+            kv_indptr = self.kv_indptr[: sparse_bs + 1]
+            kv_indptr[0] = 0
+            torch.cumsum(cache_seqlens, dim=0, out=kv_indptr[1:])
+            kv_indices = self.kv_indices[: sparse_bs * self.max_kv_tokens_per_row]
+            kv_last_page_len = self.kv_last_page_len[:sparse_bs]
+            kv_last_page_len.copy_((cache_seqlens > 0).to(torch.int32))
+            rows = self.rows[:sparse_bs]
+            if graph:
+                graph_bs = sparse_bs // self.head_group_num
+                wrapper = self.flashinfer_backend.get_cuda_graph_decode_wrappers(
+                    bs=graph_bs,
+                    num_tokens=sparse_bs,
+                )[0]
+            else:
+                wrapper = self.flashinfer_backend.decode_wrappers[0]
+
+            wrapper.begin_forward(
+                kv_indptr,
+                kv_indices,
+                kv_last_page_len,
+                self.num_qo_heads,
+                self.num_kv_heads,
+                self.head_dim,
+                self.page_size,
+                q_data_type=self.q_dtype,
+                kv_data_type=self.kv_dtype,
+                non_blocking=True,
+            )
+
+        self.active_wrapper = wrapper
+        self.active_kv_indptr = kv_indptr
+        self.active_kv_indices = kv_indices
+        self.active_rows = rows
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        metadata: MiniCPMSparseMetadata,
+        layer: RadixAttention,
+        *,
+        is_prefill: bool,
+        k_descale: Optional[torch.Tensor] = None,
+        v_descale: Optional[torch.Tensor] = None,
+        sinks: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if sinks is not None:
+            raise NotImplementedError(
+                "minicpm_flashinfer does not support attention sinks"
+            )
+        if is_prefill and not self.prefill_planned:
+            self._prepare(metadata, is_prefill=True)
+            self.prefill_planned = True
+
+        cache_seqlens = metadata.sparse_cache_seqlens_int32
+        sparse_bs = cache_seqlens.numel()
+        create_flashinfer_kv_indices_triton[(sparse_bs,)](
+            metadata.sparse_page_table,
+            self.active_rows,
+            cache_seqlens,
+            self.active_kv_indptr,
+            None,
+            self.active_kv_indices,
+            metadata.sparse_page_table.stride(0),
+        )
+        kwargs = {
+            "sm_scale": layer.scaling,
+            "logits_soft_cap": layer.logit_cap or None,
+            "k_scale": layer.k_scale_float,
+            "v_scale": layer.v_scale_float,
+        }
+        if is_prefill:
+            return self.active_wrapper.forward(
+                q,
+                (key_cache, value_cache),
+                causal=True,
+                **kwargs,
+            )
+        return self.active_wrapper.forward(
+            q,
+            (key_cache, value_cache),
+            **kwargs,
+        )

python/sglang/srt/layers/attention/minicpm/cache.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+
+
+class MiniCPMCompressedCache:
+    def __init__(
+        self,
+        pool: ReqToTokenPool,
+        allocator: BaseTokenToKVPoolAllocator,
+        *,
+        kernel_size: int,
+        kernel_stride: int,
+        enable_memory_saver: bool,
+    ):
+        self.pool = pool
+        self.allocator = allocator
+        self.kernel_size = kernel_size
+        self.kernel_stride = kernel_stride
+        saver = TorchMemorySaverAdapter.create(enable=enable_memory_saver)
+        with saver.region(GPU_MEMORY_TYPE_KV_CACHE):
+            k1_size = (pool.max_context_len - kernel_size) // kernel_stride + 1
+            k2_size = (pool.max_context_len - kernel_size * 4) // (
+                kernel_stride * 4
+            ) + 1
+            pool.req_to_sparse_k1_token = torch.zeros(
+                (pool._alloc_size, k1_size), dtype=torch.int32, device=pool.device
+            )
+            pool.req_to_sparse_k2_token = torch.zeros(
+                (pool._alloc_size, k2_size), dtype=torch.int32, device=pool.device
+            )
+        self.allocated_lens = [
+            [0] * pool._alloc_size,
+            [0] * pool._alloc_size,
+        ]
+        self.reserved_slots = torch.empty(0, dtype=torch.int64, device=pool.device)
+        self.free_slots = self.reserved_slots
+        self.reset_allocator()
+
+    def reset_allocator(self) -> None:
+        """Reserve K1/K2 capacity after the backing allocator is cleared."""
+        if self.allocator.page_size != 1:
+            raise ValueError("MiniCPM sparse attention requires page_size=1")
+
+        total_slots = self.allocator.available_size()
+        # K1 and K2 consume at most 1/s and 1/(4s) slots per dense token.
+        denominator = 4 * self.kernel_stride + 5
+        reserve_size = (5 * total_slots + denominator - 1) // denominator
+        reserved_slots = self.allocator.alloc(reserve_size)
+        if reserved_slots is None:
+            raise RuntimeError(
+                f"Unable to reserve {reserve_size} MiniCPM compressed-cache slots"
+            )
+
+        self.dense_capacity = total_slots - reserve_size
+        self.reserved_slots = reserved_slots
+        self.free_slots = reserved_slots
+        self.clear()
+
+    def _alloc_reserved(self, size: int) -> torch.Tensor:
+        if size > len(self.free_slots):
+            raise RuntimeError(
+                "MiniCPM compressed cache is out of reserved slots: "
+                f"requested={size}, available={len(self.free_slots)}"
+            )
+        slots = self.free_slots[:size]
+        self.free_slots = self.free_slots[size:]
+        return slots
+
+    def _free_reserved(self, slots: torch.Tensor) -> None:
+        self.free_slots = torch.cat(
+            (self.free_slots, slots.to(self.reserved_slots.dtype))
+        )
+
+    def _sparse_len(self, length: int, scale: int) -> int:
+        kernel_size = self.kernel_size * scale
+        if length < kernel_size:
+            return 0
+        return (length - kernel_size) // (self.kernel_stride * scale) + 1
+
+    def alloc_to_lengths(
+        self,
+        *,
+        req_pool_indices_cpu: torch.Tensor,
+        target_seq_lens_cpu: torch.Tensor,
+    ) -> None:
+        req_indices = req_pool_indices_cpu.tolist()
+        seq_lens = target_seq_lens_cpu.tolist()
+        tables = (
+            self.pool.req_to_sparse_k1_token,
+            self.pool.req_to_sparse_k2_token,
+        )
+        plans = []
+        for level, (table, scale) in enumerate(zip(tables, (1, 4))):
+            targets = {
+                req_idx: self._sparse_len(seq_len, scale)
+                for req_idx, seq_len in zip(req_indices, seq_lens)
+            }
+            rows = [
+                (req_idx, self.allocated_lens[level][req_idx], target)
+                for req_idx, target in targets.items()
+                if target > self.allocated_lens[level][req_idx]
+            ]
+            plans.append((table, rows, sum(end - start for _, start, end in rows)))
+
+        allocated = []
+        try:
+            for _, _, size in plans:
+                allocated.append(self._alloc_reserved(size) if size > 0 else None)
+
+            for (table, rows, _), locs in zip(plans, allocated):
+                if locs is None:
+                    continue
+                offset = 0
+                for req_idx, start, end in rows:
+                    count = end - start
+                    table[req_idx, start:end] = locs[offset : offset + count].to(
+                        torch.int32
+                    )
+                    offset += count
+            for level, (_, rows, _) in enumerate(plans):
+                for req_idx, _, end in rows:
+                    self.allocated_lens[level][req_idx] = end
+        except Exception:
+            for locs in allocated:
+                if locs is not None:
+                    self._free_reserved(locs)
+            raise
+
+    def free(self, req_pool_idx: int) -> None:
+        allocated = []
+        for table, lengths in zip(
+            (
+                self.pool.req_to_sparse_k1_token,
+                self.pool.req_to_sparse_k2_token,
+            ),
+            self.allocated_lens,
+        ):
+            length = lengths[req_pool_idx]
+            if length > 0:
+                allocated.append(table[req_pool_idx, :length].clone())
+                table[req_pool_idx, :length].zero_()
+                lengths[req_pool_idx] = 0
+
+        if allocated:
+            self._free_reserved(torch.cat(allocated))
+
+    def clear(self) -> None:
+        self.pool.req_to_sparse_k1_token.zero_()
+        self.pool.req_to_sparse_k2_token.zero_()
+        for lengths in self.allocated_lens:
+            lengths[:] = [0] * len(lengths)
+        self.free_slots = self.reserved_slots
+
+
+def attach_compressed_cache(
+    pool: ReqToTokenPool,
+    allocator: BaseTokenToKVPoolAllocator,
+    *,
+    kernel_size: int,
+    kernel_stride: int,
+    enable_memory_saver: bool,
+) -> ReqToTokenPool:
+    if isinstance(pool._aux_cache, MiniCPMCompressedCache):
+        return pool
+
+    pool.attach_aux_cache(
+        MiniCPMCompressedCache(
+            pool,
+            allocator,
+            kernel_size=kernel_size,
+            kernel_stride=kernel_stride,
+            enable_memory_saver=enable_memory_saver,
+        )
+    )
+    return pool

python/sglang/srt/layers/attention/minicpm/fuse_kernel.py

@@ -0,0 +1,425 @@
+import math
+from functools import partial
+
+import tilelang
+import tilelang.language as T
+import tilelang.math
+
+_pass_configs = {
+    tilelang.PassConfigKey.TL_DISABLE_THREAD_STORAGE_SYNC: True,
+    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+}
+
+
+@tilelang.jit(pass_configs=_pass_configs)
+def _fused_attn_pooling_online_topk(
+    batch_size: int,
+    groups: int,
+    heads: int,
+    dim: int,
+    topk: int,
+    max_seqlen_q_grid: int,  # Static param for grid (use bucketing)
+    pooled_k_len: int,  # Static param (use bucketing) = ceil(max_seqlen_k / block_size)
+    is_causal: bool,
+    dense_len: int = 0,
+    m_block_dim: int = 16,
+    block_M: int = 16,
+    block_N: int = 64,
+    # infllmv2 pooling parameters:
+    # block_stride = block_size // kernel_stride = 64 // 16 = 4
+    # pad_len = kernel_size // kernel_stride - 1 = 32 // 16 - 1 = 1
+    # num_offs = kernel_size // kernel_stride + block_size // kernel_stride - 1 = 2 + 4 - 1 = 5
+    block_stride: int = 4,  # pool output block stride
+    pad_len: int = 1,  # padding for pool blocks
+    num_offs: int = 5,  # number of k positions each pool block reads
+    kernel_stride: int = 16,
+    block_size: int = 64,  # block size for q/k block computation
+    init_blocks: int = 0,
+    local_blocks: int = 0,
+    num_stages: int = 0,
+    threads: int = 128,
+    dtype_str: str = "bfloat16",
+):
+    """
+    Fused Attention + Max Pooling + Online TopK for prefill and decode.
+
+    Chunk prefill support:
+    - cache_lens: tensor of shape [batch_size], cache length for each batch
+    - When cache_lens[i] = 0, it's standard prefill
+    - When cache_lens[i] > 0, it's chunk prefill (continuing from cached state)
+
+    Pooling logic aligned with infllmv2_cuda_impl:
+    - For each pool block b, it aggregates k scores in range [b * block_stride - pad_len, b * block_stride - pad_len + num_offs)
+    - block_stride = 4, pad_len = 1, num_offs = 5
+    - pool block 0: k in [0-1, 0-1+5) = [-1, 4) -> [0, 4)
+    - pool block 1: k in [3, 8)
+    - pool block 2: k in [7, 12)
+    - etc.
+    """
+    assert topk == tilelang.math.next_power_of_2(topk), "topk must be power of 2"
+
+    scale = (1.0 / dim) ** 0.5 * 1.44269504
+    head_kv = heads // groups
+
+    # Dynamic dimensions - inferred from tensor shapes at runtime
+    UQ = T.dynamic("UQ")
+    UKV = T.dynamic("UKV")
+
+    q_shape = [UQ * groups, head_kv, dim]
+    kv_shape = [UKV, head_kv, dim]
+    topk_indices_shape = [head_kv, UQ, topk]
+    topk_values_shape = [head_kv, UQ, topk]
+
+    dtype = dtype_str
+    accum_dtype = "float"
+
+    N = 2 * topk
+    num_sort_iters = int(round(math.log2(N)))
+    block_P = topk
+
+    @T.macro
+    def bitonic_sort(
+        topk_index_shared: T.Buffer([N], "int32"),
+        topk_value_shared: T.Buffer([N], "float32"),
+    ):
+        T.sync_threads()
+        for i1 in T.serial(num_sort_iters):
+            for i2 in T.serial(i1 + 1):
+                for i in T.Parallel(N):
+                    ascending = (i & (1 << (i1 + 1))) != 0
+                    j = i ^ (1 << (i1 - i2))
+                    if i < j and (
+                        (ascending and topk_value_shared[i] > topk_value_shared[j])
+                        or (
+                            not ascending
+                            and topk_value_shared[i] < topk_value_shared[j]
+                        )
+                    ):
+                        val = topk_value_shared[i]
+                        topk_value_shared[i] = topk_value_shared[j]
+                        topk_value_shared[j] = val
+                        idx = topk_index_shared[i]
+                        topk_index_shared[i] = topk_index_shared[j]
+                        topk_index_shared[j] = idx
+                T.sync_threads()
+
+    @T.prim_func
+    def main(
+        Q_unpad: T.Tensor(q_shape, dtype),
+        K_unpad: T.Tensor(kv_shape, dtype),
+        cu_seqlens_q: T.Tensor([batch_size + 1], "int32"),
+        cu_seqlens_k: T.Tensor([batch_size + 1], "int32"),
+        cache_lens: T.Tensor(
+            [batch_size], "int32"
+        ),  # Per-batch cache length for chunk prefill
+        TopkIndices: T.Tensor(topk_indices_shape, "int32"),
+        TopkValues: T.Tensor(topk_values_shape, "float32"),
+    ):
+        with T.Kernel(max_seqlen_q_grid, head_kv, batch_size, threads=threads) as (
+            bx,
+            by,
+            bz,
+        ):
+            Q_shared = T.alloc_shared([block_M, dim], dtype)
+            K_shared = T.alloc_shared([block_N, dim], dtype)
+            topk_index_shared = T.alloc_shared([N], "int32")
+            topk_value_shared = T.alloc_shared([N], "float32")
+            pool_max_shared = T.alloc_shared([block_P], "float32")
+
+            acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)
+            scores_max = T.alloc_fragment([block_M], accum_dtype)
+            scores_max_prev = T.alloc_fragment([block_M], accum_dtype)
+            scores_scale = T.alloc_fragment([block_M], accum_dtype)
+            scores_sum = T.alloc_fragment([block_M], accum_dtype)
+            logsum = T.alloc_fragment([block_M], accum_dtype)
+            acc_output = T.alloc_fragment([block_N], accum_dtype)
+
+            batch_idx = bz
+            kv_head_idx = by
+            original_q_idx = bx
+
+            q_start_idx = cu_seqlens_q[batch_idx]
+            k_start_idx = cu_seqlens_k[batch_idx]
+            q_end_idx = cu_seqlens_q[batch_idx + 1]
+            k_end_idx = cu_seqlens_k[batch_idx + 1]
+
+            q_current_seqlen = T.alloc_var("int32", init=q_end_idx - q_start_idx)
+            k_current_seqlen = T.alloc_var("int32", init=k_end_idx - k_start_idx)
+
+            # Chunk prefill: cache_len from tensor (0 for standard prefill, >0 for chunk prefill)
+            cache_len = cache_lens[batch_idx]
+            if not is_causal:
+                active = cache_len + 1 >= dense_len
+                q_current_seqlen = T.if_then_else(active, q_current_seqlen, 0)
+                k_current_seqlen = T.if_then_else(active, k_current_seqlen, 0)
+            if is_causal:
+                actual_pooled_k_len = (
+                    k_current_seqlen - 1 + pad_len
+                ) // block_stride + 1
+            else:
+                actual_pooled_k_len = (1 + cache_len + block_size - 1) // block_size
+            effective_pooled_k_len = T.min(actual_pooled_k_len, pooled_k_len)
+
+            T.fill(topk_index_shared, -1)
+            T.fill(topk_value_shared, float("-inf"))
+            T.sync_threads()
+
+            # Use q_end_idx to avoid out-of-bounds access for Q
+            q_copy_end = T.min(
+                q_start_idx * groups + (bx + 1) * block_M, q_end_idx * groups
+            )
+            T.copy(
+                Q_unpad[
+                    q_start_idx * groups + bx * block_M : q_copy_end, kv_head_idx, :
+                ],
+                Q_shared,
+            )
+            for i, d in T.Parallel(block_M, dim):
+                if original_q_idx >= q_current_seqlen:
+                    Q_shared[i, d] = 0
+
+            T.fill(logsum, 0)
+            T.fill(scores_max, -T.infinity(accum_dtype))
+
+            loop_range_k = T.ceildiv(k_current_seqlen, block_N)
+
+            for k in T.Pipelined(loop_range_k, num_stages=num_stages):
+                # Use k_end_idx to avoid out-of-bounds access for last block
+                k_copy_end = T.min(k_start_idx + (k + 1) * block_N, k_end_idx)
+                T.copy(
+                    K_unpad[k_start_idx + k * block_N : k_copy_end, kv_head_idx, :],
+                    K_shared,
+                )
+                for i, d in T.Parallel(block_N, dim):
+                    if k * block_N + i >= k_current_seqlen:
+                        K_shared[i, d] = 0
+
+                for i, j in T.Parallel(block_M, block_N):
+                    k_idx = k * block_N + j
+                    boundary_mask = (original_q_idx >= q_current_seqlen) or (
+                        k_idx >= k_current_seqlen
+                    )
+                    if is_causal:
+                        row_idx = original_q_idx * block_M + i + cache_len * block_M
+                        orig_row_idx = row_idx // m_block_dim
+                        orig_seqlen_q = (
+                            (q_current_seqlen + cache_len) * block_M
+                        ) // m_block_dim
+                        compressed_seqlen_q = (
+                            orig_seqlen_q - kernel_stride + 1
+                        ) // kernel_stride
+                        offset_row_idx = T.max(
+                            0,
+                            (orig_row_idx + 1) // kernel_stride
+                            - 1
+                            + k_current_seqlen
+                            - compressed_seqlen_q,
+                        )
+                        q_compress_clamped = T.min(k_current_seqlen, offset_row_idx)
+                        causal_mask = k_idx > q_compress_clamped
+                        acc_s[i, j] = T.if_then_else(
+                            boundary_mask or causal_mask, -1e9, 0
+                        )
+                    else:
+                        acc_s[i, j] = T.if_then_else(boundary_mask, -1e9, 0)
+
+                T.gemm(
+                    Q_shared,
+                    K_shared,
+                    acc_s,
+                    transpose_B=True,
+                    policy=T.GemmWarpPolicy.FullRow,
+                )
+
+                T.copy(scores_max, scores_max_prev)
+                T.fill(scores_max, -T.infinity(accum_dtype))
+                T.reduce_max(acc_s, scores_max, dim=1, clear=False)
+                for i in T.Parallel(block_M):
+                    scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
+
+                for i in T.Parallel(block_M):
+                    scores_scale[i] = T.exp2(
+                        scores_max_prev[i] * scale - scores_max[i] * scale
+                    )
+
+                for i, j in T.Parallel(block_M, block_N):
+                    acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
+
+                T.reduce_sum(acc_s, scores_sum, dim=1)
+                for i in T.Parallel(block_M):
+                    logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
+
+            loop_range_pool = T.ceildiv(effective_pooled_k_len, block_P)
+
+            for p_block in T.serial(loop_range_pool):
+                T.fill(pool_max_shared, float("-inf"))
+                T.sync_threads()
+
+                for k in T.serial(loop_range_k):
+                    # Use k_end_idx to avoid out-of-bounds access for last block
+                    k_copy_end = T.min(k_start_idx + (k + 1) * block_N, k_end_idx)
+                    T.copy(
+                        K_unpad[k_start_idx + k * block_N : k_copy_end, kv_head_idx, :],
+                        K_shared,
+                    )
+                    for i, d in T.Parallel(block_N, dim):
+                        if k * block_N + i >= k_current_seqlen:
+                            K_shared[i, d] = 0
+
+                    for i, j in T.Parallel(block_M, block_N):
+                        k_idx = k * block_N + j
+                        boundary_mask = (original_q_idx >= q_current_seqlen) or (
+                            k_idx >= k_current_seqlen
+                        )
+                        if is_causal:
+                            row_idx = original_q_idx * block_M + i + cache_len * block_M
+                            orig_row_idx = row_idx // m_block_dim
+                            orig_seqlen_q = (
+                                (q_current_seqlen + cache_len) * block_M
+                            ) // m_block_dim
+                            compressed_seqlen_q = (
+                                orig_seqlen_q - kernel_stride + 1
+                            ) // kernel_stride
+                            offset_row_idx = T.max(
+                                0,
+                                (orig_row_idx + 1) // kernel_stride
+                                - 1
+                                + k_current_seqlen
+                                - compressed_seqlen_q,
+                            )
+                            q_compress_clamped = T.min(k_current_seqlen, offset_row_idx)
+                            causal_mask = k_idx > q_compress_clamped
+                            acc_s[i, j] = T.if_then_else(
+                                boundary_mask or causal_mask, -1e9, 0
+                            )
+                        else:
+                            acc_s[i, j] = T.if_then_else(boundary_mask, -1e9, 0)
+
+                    T.gemm(
+                        Q_shared,
+                        K_shared,
+                        acc_s,
+                        transpose_B=True,
+                        policy=T.GemmWarpPolicy.FullRow,
+                    )
+
+                    # Normalize and handle NaN/Inf (when logsum is 0 or very small)
+                    for i, j in T.Parallel(block_M, block_N):
+                        normalized = (
+                            T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
+                            / logsum[i]
+                        )
+                        # Handle NaN/Inf: if logsum is very small or result is invalid, set to 0
+                        acc_s[i, j] = T.if_then_else(
+                            (logsum[i] > 1e-10)
+                            and (normalized >= 0)
+                            and (normalized <= 1e10),
+                            normalized,
+                            T.Cast(accum_dtype, 0.0),
+                        )
+
+                    T.fill(acc_output, 0)
+                    T.reduce_sum(acc_s, acc_output, dim=0)
+
+                    # infllmv2 block-based pooling:
+                    # For pool block b, it aggregates k scores in range [b * block_stride - pad_len, b * block_stride - pad_len + num_offs)
+                    # For k_idx, it contributes to pool block b if:
+                    #   b * block_stride - pad_len <= k_idx < b * block_stride - pad_len + num_offs
+                    # So:
+                    #   start_b = max(0, ceil((k_idx - num_offs + 1 + pad_len) / block_stride))
+                    #   end_b = floor((k_idx + pad_len) / block_stride)
+                    for j in T.Parallel(block_N):
+                        k_idx = k * block_N + j
+                        if (
+                            original_q_idx < q_current_seqlen
+                            and k_idx < k_current_seqlen
+                        ):
+                            # Calculate which pool blocks this k_idx contributes to
+                            # start_b = ceil((k_idx - num_offs + 1 + pad_len) / block_stride)
+                            #         = ceil((k_idx - 5 + 1 + 1) / 4) = ceil((k_idx - 3) / 4)
+                            start_pool = T.max(
+                                0,
+                                (k_idx - num_offs + 1 + pad_len + block_stride - 1)
+                                // block_stride,
+                            )
+                            end_pool = T.min(
+                                effective_pooled_k_len - 1,
+                                (k_idx + pad_len) // block_stride,
+                            )
+
+                            pool_block_start = p_block * block_P
+                            pool_block_end = T.min(
+                                (p_block + 1) * block_P, effective_pooled_k_len
+                            )
+
+                            for p_off in T.serial(
+                                num_offs
+                            ):  # at most num_offs pool blocks per k
+                                p_idx = start_pool + p_off
+                                if (
+                                    p_idx >= pool_block_start
+                                    and p_idx < pool_block_end
+                                    and p_idx <= end_pool
+                                ):
+                                    local_p_idx = p_idx - pool_block_start
+                                    T.atomic_max(
+                                        pool_max_shared[local_p_idx], acc_output[j]
+                                    )
+                    T.sync_threads()
+
+                for p_off in T.Parallel(block_P):
+                    p_idx = p_block * block_P + p_off
+                    if (
+                        p_idx < effective_pooled_k_len
+                        and original_q_idx < q_current_seqlen
+                    ):
+                        off_bq = (original_q_idx + cache_len) // block_size
+                        off_bk = p_idx
+
+                        # Match Torch implementation exactly:
+                        # if init_blocks > 0 and off_bk < init_blocks:
+                        #     should_mask_inf = True
+                        # elif local_blocks > 0:
+                        #     if (off_bq >= off_bk) and (off_bq <= off_bk + local_blocks):
+                        #         should_mask_inf = True
+                        is_init_masked = (init_blocks > 0) and (off_bk < init_blocks)
+                        is_local_masked = (
+                            (local_blocks > 0)
+                            and (off_bq >= off_bk)
+                            and (off_bq <= off_bk + local_blocks)
+                        )
+                        # Use elif logic: local_blocks check only when not init_masked
+                        is_masked = T.if_then_else(
+                            is_init_masked, 1, T.if_then_else(is_local_masked, 1, 0)
+                        )
+
+                        topk_index_shared[topk + p_off] = p_idx
+                        # Use inf for masked blocks to force selection
+                        # Compare only index sets, not order
+                        topk_value_shared[topk + p_off] = T.if_then_else(
+                            is_masked == 1,
+                            T.Cast("float32", float("inf")),
+                            pool_max_shared[p_off],
+                        )
+                T.sync_threads()
+
+                bitonic_sort(topk_index_shared, topk_value_shared)
+
+            for i in T.Parallel(topk):
+                if original_q_idx < q_current_seqlen:
+                    global_q_idx = q_start_idx + original_q_idx
+                    TopkIndices[kv_head_idx, global_q_idx, i] = topk_index_shared[i]
+                    TopkValues[kv_head_idx, global_q_idx, i] = topk_value_shared[i]
+
+    return main
+
+
+fused_attn_pooling_online_topk_prefill = partial(
+    _fused_attn_pooling_online_topk, is_causal=True
+)
+fused_attn_pooling_online_topk_decode = partial(
+    _fused_attn_pooling_online_topk,
+    max_seqlen_q_grid=1,
+    is_causal=False,
+)

python/sglang/srt/layers/attention/minicpm/sparse_kernels.py

@@ -0,0 +1,198 @@
+import triton
+import triton.language as tl
+
+
+# TODO. Now only page size == 1 is supported. Consider extend to page size > 1
+@triton.jit
+def compress_k_complete_kernel_new(
+    key_cache_ptr,
+    token_table_ptr,
+    cu_new_k_token_nums_ptr,
+    history_compress_k_token_nums_ptr,
+    compressed_k_table_ptr,
+    cu_total_compress_k_token_nums_ptr,
+    full_compressed_k_ptr,
+    batch_size,
+    max_chunks_per_seq,
+    token_table_cols,
+    compressed_k_table_cols,
+    head_num_k: tl.constexpr,
+    head_dim: tl.constexpr,
+    kernel_size: tl.constexpr,
+    kernel_stride: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    max_grid_chunks: tl.constexpr,
+):
+    """
+    Single-kernel implementation that fuses k computation, key compression,
+    key_cache write, and full_compressed_k read for ALL chunks (history + new).
+
+    Grid: (batch_size, min(max_total_chunks, max_grid_chunks), head_num_k)
+    where max_total_chunks = max_chunks_per_seq + max_history_chunks
+    - chunk_in_seq in [0, history_chunks_in_seq): process HISTORY chunks
+    - chunk_in_seq in [history_chunks_in_seq, total_chunks_in_seq): process NEW chunks
+
+    If total_chunks > max_grid_chunks, each thread block loops to handle multiple chunks.
+
+    Each program processes one (batch, chunk_in_seq, head) combination.
+    """
+    batch_idx = tl.program_id(0)
+    grid_chunk_idx = tl.program_id(1)
+    head_idx = tl.program_id(2)
+
+    # Total number of chunks this thread block needs to process
+    chunk_stride = max_grid_chunks
+
+    if batch_idx >= batch_size or head_idx >= head_num_k:
+        return
+
+    # ====================================================================
+    # PHASE 0: Determine chunk type and boundaries
+    # ====================================================================
+
+    history_compress = tl.load(history_compress_k_token_nums_ptr + batch_idx)
+
+    # Compute how many NEW chunks this sequence actually has
+    cu_new_k_start = tl.load(cu_new_k_token_nums_ptr + batch_idx)
+    cu_new_k_end = tl.load(cu_new_k_token_nums_ptr + batch_idx + 1)
+    new_k_count = cu_new_k_end - cu_new_k_start
+    new_chunks_in_seq = tl.where(
+        new_k_count >= kernel_size, (new_k_count - kernel_size) // kernel_stride + 1, 0
+    )
+
+    # Total chunks = history + new
+    history_chunks_in_seq = history_compress
+    total_chunks_in_seq = history_chunks_in_seq + new_chunks_in_seq
+
+    output_start = tl.load(cu_total_compress_k_token_nums_ptr + batch_idx)
+
+    # ====================================================================
+    # LOOP: Handle multiple chunks per thread block if needed
+    # ====================================================================
+
+    # Iterate over all chunks assigned to this thread block
+    chunk_in_seq = grid_chunk_idx
+
+    while chunk_in_seq < total_chunks_in_seq:
+        # Determine if processing history or new chunks
+        is_history_chunk = chunk_in_seq < history_chunks_in_seq
+
+        if is_history_chunk:
+            # ====================================================================
+            # PHASE 1: Process HISTORY chunks
+            # ====================================================================
+
+            # chunk_in_seq in [0, history_compress) -> history chunk index
+            history_chunk_idx = chunk_in_seq
+
+            global_full_idx = output_start + history_chunk_idx
+
+            # Read from compressed_k_table: indices at y = history_chunk_idx
+            full_compressed_idx = tl.load(
+                compressed_k_table_ptr
+                + batch_idx * compressed_k_table_cols
+                + history_chunk_idx
+            ).to(tl.int32)
+
+            head_offset = (
+                full_compressed_idx * head_num_k * head_dim + head_idx * head_dim
+            )
+            x = tl.load(
+                key_cache_ptr + head_offset + tl.arange(0, BLOCK_SIZE),
+                mask=tl.arange(0, BLOCK_SIZE) < head_dim,
+                other=0.0,
+            )
+            out_offset = global_full_idx * head_num_k * head_dim + head_idx * head_dim
+            tl.store(
+                full_compressed_k_ptr + out_offset + tl.arange(0, BLOCK_SIZE),
+                x,
+                mask=tl.arange(0, BLOCK_SIZE) < head_dim,
+            )
+
+        else:
+            # ====================================================================
+            # PHASE 2: Process NEW chunks
+            # ====================================================================
+
+            # chunk_in_seq in [history_compress, total_chunks_in_seq) -> new chunk index
+            new_chunk_idx = chunk_in_seq - history_chunks_in_seq
+
+            # Compute y index in token_table for this new chunk
+            # y = new_chunk_idx * kernel_stride + history_compress * kernel_stride
+            y = (new_chunk_idx + history_compress) * kernel_stride
+
+            # Use nested if instead of continue (Triton doesn't support continue)
+            if y < token_table_cols:
+                # Compute y index in compressed_k_table for new_compressed_k_indices
+                # y = new_chunk_idx + history_compress
+                compressed_table_y = new_chunk_idx + history_compress
+
+                if compressed_table_y < compressed_k_table_cols:
+                    # Read new_compressed_k_indices from compressed_k_table
+                    new_compressed_k_indices = tl.load(
+                        compressed_k_table_ptr
+                        + batch_idx * compressed_k_table_cols
+                        + compressed_table_y
+                    ).to(tl.int32)
+
+                    # ====================================================================
+                    # PHASE 3: Perform mean pooling compression on k
+                    # ====================================================================
+
+                    # Accumulate over all tokens in this chunk
+                    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
+
+                    for token_offset in range(kernel_size):
+                        # Compute k_indices for this token
+                        token_y = (
+                            new_chunk_idx * kernel_stride + token_offset
+                        ) + history_compress * kernel_stride
+
+                        # Read k_indices from token_table
+                        if token_y < token_table_cols:
+                            token_k_indices = tl.load(
+                                token_table_ptr + batch_idx * token_table_cols + token_y
+                            ).to(tl.int32)
+                        else:
+                            token_k_indices = 0
+
+                        # Load k from key_cache: key_cache[token_k_indices, head_idx, :]
+                        key_base_offset = (
+                            token_k_indices * head_num_k * head_dim
+                            + head_idx * head_dim
+                        )
+
+                        # Vectorized load of head_dim values
+                        x = tl.load(
+                            key_cache_ptr + key_base_offset + tl.arange(0, BLOCK_SIZE),
+                            mask=tl.arange(0, BLOCK_SIZE) < head_dim,
+                            other=0.0,
+                        ).to(tl.float32)
+
+                        acc += x
+
+                    # Compute mean over the chunk
+                    acc = acc / kernel_size
+
+                    head_offset = (
+                        new_compressed_k_indices * head_num_k * head_dim
+                        + head_idx * head_dim
+                    )
+                    tl.store(
+                        key_cache_ptr + head_offset + tl.arange(0, BLOCK_SIZE),
+                        acc,
+                        mask=tl.arange(0, BLOCK_SIZE) < head_dim,
+                    )
+
+                    global_full_idx = output_start + history_compress + new_chunk_idx
+                    out_offset = (
+                        global_full_idx * head_num_k * head_dim + head_idx * head_dim
+                    )
+                    tl.store(
+                        full_compressed_k_ptr + out_offset + tl.arange(0, BLOCK_SIZE),
+                        acc,
+                        mask=tl.arange(0, BLOCK_SIZE) < head_dim,
+                    )
+
+        # Move to next chunk for this thread block
+        chunk_in_seq += chunk_stride

python/sglang/srt/layers/attention/minicpm/sparse_utils.py

@@ -0,0 +1,719 @@
+"""Sparse attention utilities for MiniCPM models.
+
+This module provides sparse attention helpers and utilities for MiniCPM models,
+combining both backend-agnostic sparse attention components and kernel utilities.
+"""
+
+from __future__ import annotations
+
+from itertools import accumulate
+from typing import TYPE_CHECKING, Optional
+
+import msgspec
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.layers.attention.flashattention_backend import (
+    FlashAttentionMetadata,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+import triton
+from sgl_kernel import infllmv2_attn_stage1, max_pooling_1d_varlen
+
+from sglang.srt.layers.attention.minicpm.sparse_kernels import (
+    compress_k_complete_kernel_new,
+)
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+
+
+def batched_gather(a, lengths_cpu, select):
+    offsets = [0, *accumulate(map(int, lengths_cpu))]
+    return torch.cat([a[offsets[i] : offsets[i + 1]] for i in select])
+
+
+def compress_k_core_new(
+    full_compressed_k,  # output
+    batch,
+    key_cache,
+    token_table,
+    compressed_k_table,
+    cu_new_k_token_nums,
+    history_compress_k_token_nums,
+    cu_total_compress_k_token_nums,
+    kernel_size,
+    kernel_stride,
+    max_context_length,
+):
+    head_num_k = key_cache.shape[1]
+    head_dim = key_cache.shape[2]
+
+    # ==============================================================================
+    # BUFFER ALLOCATION
+    # ==============================================================================
+
+    # Use provided explicit parameters for buffer allocation
+    # max_chunks_per_seq is already the maximum possible chunks for any sequence
+    # given max_context_length, kernel_size, and kernel_stride
+    max_chunks_per_seq = max(0, (max_context_length - kernel_size) // kernel_stride + 1)
+
+    # ==============================================================================
+    # Launch kernel for ALL chunks (history + new)
+    # ==============================================================================
+    # Grid: (batch, max_chunks_per_seq, head_num_k)
+    # - chunk_in_seq in [0, history_compress): process HISTORY chunks
+    # - chunk_in_seq in [history_compress, total_chunks_in_seq): process NEW chunks
+    #
+    # max_chunks_per_seq is already the maximum possible chunks for any sequence,
+    # so it's sufficient for both history and new chunks.
+    #
+    # All operations are in a single kernel, CUDA graph compatible.
+
+    # Limit grid size to avoid too many thread blocks
+    # If max_chunks_per_seq > max_grid_chunks, kernel will loop to handle remaining chunks
+    MAX_GRID_CHUNKS = 1024  # Adjustable limit for grid dimension
+    max_grid_chunks = min(max_chunks_per_seq, MAX_GRID_CHUNKS)
+
+    BLOCK_SIZE = triton.next_power_of_2(head_dim)
+    # Grid size is now limited, kernel uses loop to handle all chunks
+    grid = (batch, max_grid_chunks, head_num_k)
+
+    compress_k_complete_kernel_new[grid](
+        key_cache,
+        token_table,
+        cu_new_k_token_nums,
+        history_compress_k_token_nums,
+        compressed_k_table,
+        cu_total_compress_k_token_nums,
+        full_compressed_k,
+        batch,
+        max_chunks_per_seq,
+        token_table.shape[1],
+        compressed_k_table.shape[1],
+        head_num_k,
+        head_dim,
+        kernel_size,
+        kernel_stride,
+        BLOCK_SIZE,
+        max_grid_chunks,  # Pass the limit to kernel for loop control
+    )
+
+    return
+
+
+def get_compress_k_v2(
+    layer,
+    forward_batch,
+    metadata: MiniCPMSparseMetadata,
+    full_compressed_k1,
+    full_compressed_k2,
+    max_context_length,
+    k1_kernel_size,
+    k1_kernel_stride,
+    k2_kernel_size,
+    k2_kernel_stride,
+):
+    batch = len(forward_batch.req_pool_indices)
+    key_cache = get_token_to_kv_pool().get_key_buffer(layer.layer_id)
+    key_cache = key_cache.view(-1, layer.tp_k_head_num, layer.head_dim)
+
+    for full_compressed_k, level, kernel_size, kernel_stride in (
+        (
+            full_compressed_k1,
+            metadata.k1,
+            k1_kernel_size,
+            k1_kernel_stride,
+        ),
+        (
+            full_compressed_k2,
+            metadata.k2,
+            k2_kernel_size,
+            k2_kernel_stride,
+        ),
+    ):
+        compress_k_core_new(
+            full_compressed_k,
+            batch,
+            key_cache,
+            metadata.base.page_table,
+            level.table,
+            level.cu_new_token_nums,
+            level.history_compress_token_nums,
+            level.cu_total_compress_token_nums,
+            kernel_size,
+            kernel_stride,
+            max_context_length,
+        )
+
+
+def allocate_and_compress_keys(
+    layer,
+    forward_batch,
+    metadata: MiniCPMSparseMetadata,
+    k1_token_nums: int,
+    k2_token_nums: int,
+    k1_kernel_size: int,
+    k1_kernel_stride: int,
+    k2_kernel_size: int,
+    k2_kernel_stride: int,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device = None,
+    max_context_length: int = 32768,
+):
+    """Allocate compressed key tensors and run compression.
+
+    Args:
+        layer: Model layer with head configuration
+        forward_batch: Forward batch info
+        metadata: MiniCPM sparse metadata
+        k1_token_nums: Number of k1 tokens to allocate
+        k2_token_nums: Number of k2 tokens to allocate
+        k1_kernel_size: K1 compression window
+        k1_kernel_stride: K1 compression stride
+        k2_kernel_size: K2 compression window
+        k2_kernel_stride: K2 compression stride
+        dtype: Tensor data type (default: bfloat16)
+        device: Tensor device (default: layer device)
+        max_context_length: Maximum context length for the model (default: 32768)
+
+    Returns:
+        Tuple of (full_compressed_k1, full_compressed_k2)
+    """
+    if device is None:
+        device = forward_batch.input_ids.device
+
+    full_compressed_k1 = torch.full(
+        (k1_token_nums, layer.tp_k_head_num, layer.head_dim),
+        dtype=dtype,
+        device=device,
+        fill_value=float("-inf"),
+    )
+    full_compressed_k2 = torch.full(
+        (k2_token_nums, layer.tp_k_head_num, layer.head_dim),
+        dtype=dtype,
+        device=device,
+        fill_value=float("-inf"),
+    )
+
+    get_compress_k_v2(
+        layer,
+        forward_batch,
+        metadata,
+        full_compressed_k1,
+        full_compressed_k2,
+        max_context_length=max_context_length,
+        k1_kernel_size=k1_kernel_size,
+        k1_kernel_stride=k1_kernel_stride,
+        k2_kernel_size=k2_kernel_size,
+        k2_kernel_stride=k2_kernel_stride,
+    )
+
+    return full_compressed_k1, full_compressed_k2
+
+
+def compressed_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    k2: torch.Tensor,
+    kernel_stride: int,
+    block_size: int,
+    topk: int,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    cu_seqlens_k2: torch.Tensor,
+    max_seqlen_q: int,
+    max_context_len: int,
+    init_blocks: int = 1,
+    local_blocks: int = 2,
+    cache_lens: Optional[torch.Tensor] = None,
+    cu_seqlens_q_adjusted: Optional[torch.Tensor] = None,
+    max_seqlen_q_adjusted: Optional[int] = None,
+) -> torch.Tensor:
+    """Compressed attention computation for sparse attention.
+
+    Computes attention scores between query and compressed keys (k and k2),
+    then performs max pooling and selects top-k blocks.
+
+    Args:
+        q: Query tensor, shape (total_q_len, num_heads, head_dim)
+        k: Compressed key tensor k1, shape (total_k_len, num_heads, head_dim)
+        k2: Compressed key tensor k2, shape (total_k_len, num_heads, head_dim)
+        kernel_stride: Stride of compression kernel
+        block_size: Size of attention blocks
+        topk: Number of top blocks to select
+        cu_seqlens_q: Cumulative sequence lengths for query, shape (batch_size + 1)
+        cu_seqlens_k: Cumulative sequence lengths for k, shape (batch_size + 1)
+        cu_seqlens_k2: Cumulative sequence lengths for k2, shape (batch_size + 1)
+        max_seqlen_q: Maximum sequence length in query
+        init_blocks: Number of initial blocks to always attend to
+        local_blocks: Number of local blocks to consider
+        cache_lens: Cache lengths for each batch (optional)
+        cu_seqlens_q_adjusted: Adjusted cumulative sequence lengths for query (for stage1 optimization)
+        max_seqlen_q_adjusted: Adjusted maximum sequence length for query (for stage1 optimization)
+
+    Returns:
+        Top-k block indices, shape (num_heads, total_q_len, topk)
+    """
+    with torch.no_grad():
+        batch_size = cu_seqlens_q.shape[0] - 1
+
+        is_prefilling = max_seqlen_q > 1
+
+        if is_prefilling:
+            if cache_lens is None:
+                cache_lens = torch.zeros(batch_size, dtype=torch.int32, device=q.device)
+
+        score = infllmv2_attn_stage1(
+            q.contiguous(),
+            k.contiguous(),
+            k2.contiguous(),
+            cu_seqlens_q=cu_seqlens_q_adjusted,
+            cu_seqlens_k=cu_seqlens_k,
+            cu_seqlens_v=cu_seqlens_k2,
+            max_seqlen_q=max_seqlen_q_adjusted,
+            max_seqlen_k=max_context_len // kernel_stride,
+            causal=is_prefilling,
+        )
+
+        block_score = max_pooling_1d_varlen(
+            score.contiguous(),
+            cu_seqlens_q,
+            cu_seqlens_k,
+            cache_lens,
+            max_seqlen_q,
+            max_context_len,
+            local_blocks=local_blocks,
+            init_blocks=init_blocks,
+            block_size=block_size,
+            stride=kernel_stride,
+        )
+
+        topk_idx = block_score.topk(topk, dim=-1).indices.sort(-1).values
+        topk_idx = topk_idx.to(torch.int32)
+
+    return topk_idx
+
+
+def compressed_attention_tilelang(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    block_size: int,
+    topk: int,
+    kernel_topk: int,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    cache_lens=None,
+    fused_kernel=None,
+    max_cache_len=-1,
+) -> torch.Tensor:
+    """
+    使用 tilelang online topk kernel 计算 compressed attention topk indices
+    """
+    with torch.no_grad():
+        batch_size = cu_seqlens_q.shape[0] - 1
+
+        total_q_len = q.shape[0]
+        num_kv_heads = k.shape[1]
+        head_dim = k.shape[2]
+
+        num_heads = q.shape[1]
+        groups = num_heads // num_kv_heads
+        q_kernel = q.view(total_q_len, num_kv_heads, groups, head_dim)
+        q_kernel = (
+            q_kernel.transpose(1, 2)
+            .reshape(total_q_len * groups, num_kv_heads, head_dim)
+            .contiguous()
+        )
+
+        k_kernel = k.contiguous()
+
+        pooled_k_len = (max_cache_len + block_size - 1) // block_size
+
+        assert fused_kernel is not None, "fused_kernel is not initialized"
+
+        # Compute actual output topk (same as original: min(topk, num_blocks))
+        output_topk = min(topk, pooled_k_len)
+
+        # Allocate output tensors
+        topk_indices = torch.full(
+            (num_kv_heads, total_q_len, kernel_topk),
+            -1,
+            dtype=torch.int32,
+            device=q.device,
+        )
+        topk_values = torch.full(
+            (num_kv_heads, total_q_len, kernel_topk),
+            float("-inf"),
+            dtype=torch.float32,
+            device=q.device,
+        )
+
+        if cache_lens is None:
+            cache_lens_tensor = torch.zeros(
+                batch_size, dtype=torch.int32, device=q.device
+            )
+        else:
+            cache_lens_tensor = cache_lens.to(torch.int32)
+
+        fused_kernel(
+            q_kernel,
+            k_kernel,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            cache_lens_tensor,
+            topk_indices,
+            topk_values,
+        )
+
+        # Note: q_idx masking is handled inside the kernel via causal_mask
+        # which sets scores to -1e9 for K blocks beyond the causal boundary.
+        # These blocks won't be selected in topk due to their low scores.
+
+        # Sort with -1 values at the end (match original behavior)
+        # Replace -1 with large value, sort, then replace back
+        large_val = pooled_k_len + 1000  # Any value larger than max valid index
+        topk_for_sort = topk_indices.clone()
+        topk_for_sort[topk_for_sort == -1] = large_val
+        topk_idx = topk_for_sort.sort(-1).values
+        topk_idx[topk_idx == large_val] = -1
+
+        # Truncate to output_topk (same as original: min(topk, num_blocks))
+        topk_idx = topk_idx[:, :, :output_topk].contiguous()
+
+        return topk_idx
+
+
+class CompressionLevelMetadata(msgspec.Struct):
+    """Metadata for a single compression level (k1 or k2).
+
+    This struct groups all metadata fields for one compression level,
+    reducing duplication and making the code more maintainable.
+    """
+
+    # Cumulative sequence lengths for compressed cache
+    cu_seqlens: Optional[torch.Tensor] = None
+    cu_seqlens_cpu: Optional[list[int]] = None
+
+    # Token mapping table (request pool indices -> compressed cache tokens)
+    table: Optional[torch.Tensor] = None
+
+    # Compressed cache metadata
+    history_compress_token_nums: Optional[torch.Tensor] = None
+    cu_new_token_nums: Optional[torch.Tensor] = None
+    cu_total_compress_token_nums: Optional[torch.Tensor] = None
+
+
+class MiniCPMSparseMetadata(msgspec.Struct):
+    base: FlashAttentionMetadata
+    k1: Optional[CompressionLevelMetadata] = None
+    k2: Optional[CompressionLevelMetadata] = None
+    sparse_bs_list: Optional[list[int]] = None
+    sparse_idx: Optional[list[int]] = None
+    dense_layout: Optional[list[tuple[int, int, int, int]]] = None
+    seqlen_k_sparse_bs_tensor: Optional[torch.Tensor] = None
+    token_to_bs: Optional[torch.Tensor] = None
+    token_pos_in_bs: Optional[torch.Tensor] = None
+    sparse_page_table: Optional[torch.Tensor] = None
+    sparse_cache_seqlens_int32: Optional[torch.Tensor] = None
+    sparse_cu_seqlens_q: Optional[torch.Tensor] = None
+    sparse_cu_seqlens_k: Optional[torch.Tensor] = None
+    sparse_max_seq_len_q: int = 1
+    cache_seqlens_int32_stage1: Optional[torch.Tensor] = None
+    cu_seqlens_q_adjusted: Optional[torch.Tensor] = None
+    max_seqlen_q_adjusted: int = 1
+    topk_cu_seqlens_q: Optional[torch.Tensor] = None
+    topk_cu_seqlens_k: Optional[torch.Tensor] = None
+    topk_max_seqlen_q: int = 1
+    topk_max_seqlen_k: int = 1
+
+
+def _compute_single_compression_metadata(
+    seq_lens_cpu: torch.Tensor,
+    token_nums: torch.Tensor,
+    history_lens: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    req_to_sparse_token: torch.Tensor,
+    kernel_size: int,
+    kernel_stride: int,
+) -> CompressionLevelMetadata:
+    seqlen_cpu = torch.clamp(
+        (seq_lens_cpu - kernel_size) // kernel_stride + 1,
+        min=0,
+    )
+
+    cu_seqlens_cpu = F.pad(
+        torch.cumsum(seqlen_cpu, dim=0, dtype=torch.int32), (1, 0)
+    ).tolist()
+    cu_seqlens = F.pad(
+        torch.cumsum(seqlen_cpu.to(device=token_nums.device), dim=0, dtype=torch.int32),
+        (1, 0),
+    )
+    token_table = req_to_sparse_token[req_pool_indices]
+    history_compress_token_nums = torch.clamp(
+        (history_lens - kernel_size) // kernel_stride + 1,
+        min=0,
+    )
+    new_token_nums = token_nums - history_compress_token_nums * kernel_stride
+    cu_new_token_nums = F.pad(
+        torch.cumsum(new_token_nums, dim=0, dtype=torch.int32), (1, 0)
+    )
+    new_compress_token_nums = torch.clamp(
+        (new_token_nums - kernel_size) // kernel_stride + 1,
+        min=0,
+    )
+    total_compress_token_nums = history_compress_token_nums + new_compress_token_nums
+    cu_total_compress_token_nums = F.pad(
+        torch.cumsum(total_compress_token_nums, dim=0, dtype=torch.int32), (1, 0)
+    )
+
+    return CompressionLevelMetadata(
+        cu_seqlens=cu_seqlens,
+        cu_seqlens_cpu=cu_seqlens_cpu,
+        table=token_table,
+        history_compress_token_nums=history_compress_token_nums,
+        cu_new_token_nums=cu_new_token_nums,
+        cu_total_compress_token_nums=cu_total_compress_token_nums,
+    )
+
+
+def _build_k1_k2_compression_metadata(
+    forward_batch: ForwardBatch,
+    base_metadata: FlashAttentionMetadata,
+    req_to_sparse_k1_token: torch.Tensor,
+    req_to_sparse_k2_token: torch.Tensor,
+    k1_kernel_size: int,
+    k1_kernel_stride: int,
+    k2_kernel_size: int,
+    k2_kernel_stride: int,
+    cu_seqlens_q: torch.Tensor,
+) -> tuple[CompressionLevelMetadata, CompressionLevelMetadata]:
+    bs = forward_batch.batch_size
+    seq_lens_cpu = torch.as_tensor(
+        forward_batch.seq_lens_cpu,
+        dtype=base_metadata.cu_seqlens_q.dtype,
+        device="cpu",
+    )
+    token_nums = (
+        base_metadata.cu_seqlens_k[1 : bs + 1] - base_metadata.cu_seqlens_k[:bs]
+    )
+    input_lens = cu_seqlens_q[1 : bs + 1] - cu_seqlens_q[:bs]
+    history_lens = token_nums - input_lens
+
+    return tuple(
+        _compute_single_compression_metadata(
+            seq_lens_cpu,
+            token_nums,
+            history_lens,
+            forward_batch.req_pool_indices,
+            req_to_sparse_token,
+            kernel_size,
+            kernel_stride,
+        )
+        for req_to_sparse_token, kernel_size, kernel_stride in (
+            (req_to_sparse_k1_token, k1_kernel_size, k1_kernel_stride),
+            (req_to_sparse_k2_token, k2_kernel_size, k2_kernel_stride),
+        )
+    )
+
+
+def _get_sparse_cache_lens(
+    seq_lens: torch.Tensor,
+    sparse_capacity: int,
+    block_size: int,
+) -> torch.Tensor:
+    remainder = seq_lens % block_size
+    sparse_lens = torch.where(
+        remainder == 0,
+        sparse_capacity,
+        sparse_capacity - block_size + remainder,
+    )
+    return torch.where(seq_lens <= sparse_capacity, seq_lens, sparse_lens)
+
+
+def _plan_sparse_prefill(
+    forward_batch: ForwardBatch,
+    metadata: MiniCPMSparseMetadata,
+    head_group_num: int,
+    heads_per_group: int,
+    dense_len: int,
+    sparse_topk: int,
+    block_size: int,
+) -> None:
+    device = metadata.base.cu_seqlens_q.device
+    sparse_capacity = sparse_topk * block_size
+    sparse_bs_list = []
+    sparse_idx = []
+    dense_layout = []
+    row_q_lens = []
+    sparse_cache_seqlens = []
+    token_to_bs = []
+    token_pos_in_bs = []
+    sparse_q_lens = []
+    sparse_k_lens = []
+    dense_q_lens = []
+    max_sparse_cache_len = 0
+    query_group_start = 0
+
+    for batch_idx in range(forward_batch.batch_size):
+        seq_len = int(forward_batch.seq_lens_cpu[batch_idx])
+        query_len = int(forward_batch.extend_seq_lens_cpu[batch_idx])
+        prefix_len = int(forward_batch.extend_prefix_lens_cpu[batch_idx])
+        row_start = len(row_q_lens)
+        if seq_len >= dense_len:
+            sparse_batch_idx = len(sparse_bs_list)
+            sparse_bs_list.append(batch_idx)
+            sparse_q_lens.append(query_len)
+            sparse_k_lens.append(seq_len)
+            sparse_idx.extend(range(row_start, row_start + query_len * head_group_num))
+            row_q_lens.extend([1] * (query_len * head_group_num))
+            token_to_bs.extend([sparse_batch_idx] * query_len)
+            token_pos_in_bs.extend(range(prefix_len + 1, prefix_len + query_len + 1))
+            token_seq_lens = torch.arange(
+                prefix_len + 1,
+                prefix_len + query_len + 1,
+                dtype=torch.int32,
+            )
+            sparse_cache_seqlens.extend(
+                _get_sparse_cache_lens(token_seq_lens, sparse_capacity, block_size)
+                .repeat_interleave(head_group_num)
+                .tolist()
+            )
+            max_sparse_cache_len = max(max_sparse_cache_len, sparse_capacity)
+        else:
+            dense_layout.append((batch_idx, row_start, query_group_start, query_len))
+            dense_q_lens.append(query_len)
+            row_q_lens.extend([query_len] * head_group_num)
+            sparse_cache_seqlens.extend([seq_len] * head_group_num)
+            max_sparse_cache_len = max(max_sparse_cache_len, seq_len)
+        query_group_start += query_len * head_group_num
+
+    metadata.sparse_bs_list = sparse_bs_list
+    metadata.sparse_idx = sparse_idx
+    metadata.dense_layout = dense_layout
+    metadata.token_to_bs = torch.tensor(token_to_bs, dtype=torch.int32, device=device)
+    metadata.token_pos_in_bs = torch.tensor(
+        token_pos_in_bs, dtype=torch.int32, device=device
+    )
+    metadata.seqlen_k_sparse_bs_tensor = torch.tensor(
+        sparse_k_lens, dtype=torch.int32, device=device
+    )
+    metadata.sparse_page_table = torch.zeros(
+        (len(row_q_lens), max_sparse_cache_len),
+        dtype=metadata.base.page_table.dtype,
+        device=metadata.base.page_table.device,
+    )
+    row_q_lens_tensor = torch.tensor(
+        row_q_lens, dtype=metadata.base.cu_seqlens_q.dtype, device=device
+    )
+    metadata.sparse_cu_seqlens_q = F.pad(
+        torch.cumsum(row_q_lens_tensor, dim=0, dtype=torch.int32), (1, 0)
+    )
+    metadata.sparse_max_seq_len_q = max(dense_q_lens, default=1)
+    metadata.sparse_cache_seqlens_int32 = torch.tensor(
+        sparse_cache_seqlens,
+        dtype=torch.int32,
+        device=device,
+    )
+    metadata.sparse_cu_seqlens_k = F.pad(
+        torch.cumsum(metadata.sparse_cache_seqlens_int32, dim=0, dtype=torch.int32),
+        (1, 0),
+    )
+    metadata.cache_seqlens_int32_stage1 = (
+        metadata.base.cache_seqlens_int32[sparse_bs_list] - 1
+    )
+
+    if sparse_bs_list:
+        sparse_q_lens_tensor = torch.tensor(
+            sparse_q_lens, dtype=torch.int32, device=device
+        )
+        metadata.topk_cu_seqlens_q = F.pad(
+            torch.cumsum(sparse_q_lens_tensor, dim=0, dtype=torch.int32), (1, 0)
+        )
+        metadata.topk_cu_seqlens_k = F.pad(
+            torch.cumsum(metadata.seqlen_k_sparse_bs_tensor, dim=0, dtype=torch.int32),
+            (1, 0),
+        )
+        metadata.topk_max_seqlen_q = max(sparse_q_lens)
+        metadata.topk_max_seqlen_k = max(sparse_k_lens)
+        metadata.cu_seqlens_q_adjusted = metadata.topk_cu_seqlens_q * heads_per_group
+        metadata.max_seqlen_q_adjusted = metadata.topk_max_seqlen_q * heads_per_group
+    else:
+        metadata.cu_seqlens_q_adjusted = metadata.base.cu_seqlens_q * heads_per_group
+        metadata.max_seqlen_q_adjusted = metadata.base.max_seq_len_q * heads_per_group
+
+
+def _plan_sparse_decode(
+    forward_batch: ForwardBatch,
+    metadata: MiniCPMSparseMetadata,
+    head_group_num: int,
+    dense_len: int,
+    sparse_topk: int,
+    block_size: int,
+) -> None:
+    base_metadata = metadata.base
+    bs = forward_batch.batch_size
+    cache_seqlens = base_metadata.cache_seqlens_int32
+    page_table = base_metadata.page_table
+    seq_lens_cpu = torch.as_tensor(
+        forward_batch.seq_lens_cpu, dtype=cache_seqlens.dtype, device="cpu"
+    )
+    sparse_capacity = sparse_topk * block_size
+    cache_lens_cpu = torch.where(
+        seq_lens_cpu >= dense_len,
+        _get_sparse_cache_lens(seq_lens_cpu, sparse_capacity, block_size),
+        seq_lens_cpu,
+    )
+    sparse_mask_cpu = seq_lens_cpu >= dense_len
+    sparse_bs_list = sparse_mask_cpu.nonzero().flatten().tolist()
+    dense_bs_list = (~sparse_mask_cpu).nonzero().flatten().tolist()
+    sparse_idx = [
+        row
+        for batch_idx in sparse_bs_list
+        for row in range(batch_idx * head_group_num, (batch_idx + 1) * head_group_num)
+    ]
+    max_sparse_cache_len = max(
+        int(cache_lens_cpu.max()),
+        sparse_capacity if sparse_bs_list else 0,
+    )
+    sparse_cache_seqlens_cpu = cache_lens_cpu.repeat_interleave(head_group_num)
+
+    sparse_cache_seqlens_int32 = sparse_cache_seqlens_cpu.to(
+        device=cache_seqlens.device
+    )
+    sparse_cu_seqlens_k = F.pad(
+        torch.cumsum(sparse_cache_seqlens_int32, dim=0, dtype=torch.int32), (1, 0)
+    )
+    sparse_cu_seqlens_q = torch.arange(
+        0,
+        bs * head_group_num + 1,
+        dtype=torch.int32,
+        device=base_metadata.cu_seqlens_q.device,
+    )
+    token_to_bs = torch.arange(
+        0, len(sparse_bs_list), dtype=torch.int32, device=page_table.device
+    )
+    sparse_page_table = torch.zeros(
+        (head_group_num * bs, max_sparse_cache_len),
+        dtype=page_table.dtype,
+        device=page_table.device,
+    )
+
+    metadata.sparse_cache_seqlens_int32 = sparse_cache_seqlens_int32
+    metadata.sparse_cu_seqlens_k = sparse_cu_seqlens_k
+    metadata.sparse_cu_seqlens_q = sparse_cu_seqlens_q
+    metadata.sparse_page_table = sparse_page_table
+    metadata.sparse_bs_list = sparse_bs_list
+    metadata.sparse_idx = sparse_idx
+    metadata.dense_layout = [
+        (batch_idx, batch_idx * head_group_num, batch_idx * head_group_num, 1)
+        for batch_idx in dense_bs_list
+    ]
+    metadata.token_to_bs = token_to_bs
+    metadata.topk_cu_seqlens_q = torch.arange(
+        0,
+        len(sparse_bs_list) + 1,
+        dtype=torch.int32,
+        device=base_metadata.cu_seqlens_q.device,
+    )

python/sglang/srt/managers/scheduler.py

@@ -4386,6 +4386,7 @@ class Scheduler(
             self.tree_cache.reset()
             self.req_to_token_pool.clear()
             self.token_to_kv_pool_allocator.clear()
+            self.req_to_token_pool.reset_aux_cache_allocator()
             self.grammar_manager.clear()
             self.metrics_reporter.reset_metrics()
 

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -98,7 +98,9 @@ class SchedulerInvariantChecker:
             else:
                 protected = self.tree_cache.protected_size()
             session_held = self.pool_stats_observer.session_held_tokens()
-            total = self.token_to_kv_pool_allocator.size
+            total = self.req_to_token_pool.schedulable_token_capacity(
+                self.token_to_kv_pool_allocator.size
+            )
         else:
             protected = self.tree_cache.protected_size()
             session_held = self.pool_stats_observer.session_held_tokens()
@@ -164,12 +166,20 @@ class SchedulerInvariantChecker:
                 return leak, msg
             free_full_pages = set(free_pages.tolist() + release_pages.tolist())
             cached_full_pages = set(self.tree_cache.all_values_flatten().tolist())
-            expected_full_pages = set(
+            full_page_msg = ""
-                range(1, self.token_to_kv_pool_allocator.size + 1)
+            if (
-            )
+                self.req_to_token_pool.schedulable_token_capacity(
-            leaked_full_pages = (
+                    self.token_to_kv_pool_allocator.size
-                expected_full_pages - free_full_pages - cached_full_pages
+                )
-            )
+                == self.token_to_kv_pool_allocator.size
+            ):
+                expected_full_pages = set(
+                    range(1, self.token_to_kv_pool_allocator.size + 1)
+                )
+                leaked_full_pages = (
+                    expected_full_pages - free_full_pages - cached_full_pages
+                )
+                full_page_msg = f", leaked_full_pages={leaked_full_pages or None}"
             mamba_allocator = self.req_to_token_pool.mamba_allocator
             free_mamba_pages = set(mamba_allocator.free_slots.tolist())
             cached_mamba_pages = set(
@@ -179,10 +189,8 @@ class SchedulerInvariantChecker:
             leaked_mamba_pages = (
                 expected_mamba_pages - free_mamba_pages - cached_mamba_pages
             )
-            msg += (
+            msg += full_page_msg
-                f", leaked_full_pages={leaked_full_pages or None}"
+            msg += f", leaked_mamba_pages={leaked_mamba_pages or None}"
-                f", leaked_mamba_pages={leaked_mamba_pages or None}"
-            )
         return leak, msg
 
     def _check_mamba_pool_with_int8(self, ps: PoolStats, ckpt_pool) -> Tuple[bool, str]:

python/sglang/srt/managers/scheduler_components/pool_stats_observer.py

@@ -261,13 +261,14 @@ class SchedulerPoolStatsObserver:
             if (is_mamba_radix_cache and not has_int8_ckpt)
             else 0
         )
-        full_num_used = self.token_to_kv_pool_allocator.size - (
+        full_capacity = self.req_to_token_pool.schedulable_token_capacity(
-            full_available_size + full_evictable_size
+            self.token_to_kv_pool_allocator.size
         )
+        full_num_used = full_capacity - (full_available_size + full_evictable_size)
         mamba_num_used = self.req_to_token_pool.mamba_pool.size - (
             mamba_available_size + mamba_evictable_size
         )
-        full_token_usage = full_num_used / self.token_to_kv_pool_allocator.size
+        full_token_usage = full_num_used / full_capacity
         mamba_usage = mamba_num_used / self.req_to_token_pool.mamba_pool.size
 
         return PoolStats(

python/sglang/srt/managers/tp_worker.py

@@ -536,7 +536,9 @@ class TpModelWorker(BaseTpWorker):
             - 1,
         )
         return (
-            self.model_runner.max_total_num_tokens,
+            self.model_runner.req_to_token_pool.schedulable_token_capacity(
+                self.model_runner.max_total_num_tokens
+            ),
             get_schedule().max_prefill_tokens,
             self.model_runner.max_running_requests,
             get_schedule().max_queued_requests,

python/sglang/srt/mem_cache/allocation.py

@@ -359,6 +359,14 @@ def alloc_for_extend(
         prefix_tensors,
         batch.req_to_token_pool,
     )
+    try:
+        batch.req_to_token_pool.alloc_aux_to_lengths(
+            req_pool_indices_cpu=req_pool_indices_cpu,
+            target_seq_lens_cpu=batch.seq_lens_cpu,
+        )
+    except Exception:
+        batch.tree_cache.token_to_kv_pool_allocator.free(out_cache_loc)
+        raise
 
     # DSV4-NPU hook: no-op on non-DSV4 paths.
     if _is_npu:
@@ -559,6 +567,15 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
             token_per_req,
         )
 
+    try:
+        batch.req_to_token_pool.alloc_aux_to_lengths(
+            req_pool_indices_cpu=batch.req_pool_indices_cpu,
+            target_seq_lens_cpu=batch.seq_lens_cpu + token_per_req,
+        )
+    except Exception:
+        batch.tree_cache.token_to_kv_pool_allocator.free(out_cache_loc)
+        raise
+
     for req in batch.reqs:
         req.kv.kv_allocated_len += token_per_req
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -281,6 +281,7 @@ class ReqToTokenPool:
             )
         self.free_slots = list(range(1, self._alloc_size))
         self.req_generation = torch.zeros(self._alloc_size, dtype=torch.int64)
+        self._aux_cache: Any = None
 
     def write(self, indices, values):
         self.req_to_token[indices] = values
@@ -324,12 +325,41 @@ class ReqToTokenPool:
 
     def free(self, req: Req):
         assert req.req_pool_idx is not None, "request must have req_pool_idx"
+        if self._aux_cache is not None:
+            self._aux_cache.free(req.req_pool_idx)
         self.free_slots.append(req.req_pool_idx)
         req.req_pool_idx = None
 
     def clear(self):
         self.free_slots = list(range(1, self._alloc_size))
         self.req_generation.zero_()
+        if self._aux_cache is not None:
+            self._aux_cache.clear()
+
+    def attach_aux_cache(self, aux_cache: Any) -> None:
+        assert self._aux_cache is None
+        self._aux_cache = aux_cache
+
+    def reset_aux_cache_allocator(self) -> None:
+        if self._aux_cache is not None:
+            self._aux_cache.reset_allocator()
+
+    def schedulable_token_capacity(self, physical_capacity: int) -> int:
+        if self._aux_cache is None:
+            return physical_capacity
+        return self._aux_cache.dense_capacity
+
+    def alloc_aux_to_lengths(
+        self,
+        *,
+        req_pool_indices_cpu: torch.Tensor,
+        target_seq_lens_cpu: torch.Tensor,
+    ) -> None:
+        if self._aux_cache is not None:
+            self._aux_cache.alloc_to_lengths(
+                req_pool_indices_cpu=req_pool_indices_cpu,
+                target_seq_lens_cpu=target_seq_lens_cpu,
+            )
 
 
 class MambaPool:

python/sglang/srt/model_executor/model_runner.py

@@ -1298,9 +1298,12 @@ class ModelRunner:
     def effective_max_total_num_tokens(self):
         """Return the max token pool size considering hybrid swa settings."""
         if self.is_hybrid_swa:
-            return self.full_max_total_num_tokens or self.swa_max_total_num_tokens
+            capacity = self.full_max_total_num_tokens or self.swa_max_total_num_tokens
         else:
-            return self.max_total_num_tokens
+            capacity = self.max_total_num_tokens
+        if (req_to_token_pool := getattr(self, "req_to_token_pool", None)) is not None:
+            return req_to_token_pool.schedulable_token_capacity(capacity)
+        return capacity
 
     @property
     def max_token_pool_size(self):

python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py

@@ -96,6 +96,7 @@ def compute_post_capture_kv_resize(
     )
     pool.finalize_backing(config)
     model_runner.token_to_kv_pool_allocator.resize(config)
+    model_runner.req_to_token_pool.reset_aux_cache_allocator()
 
     capped_max_running_requests = None
     if model_runner.max_running_requests is not None:

python/sglang/srt/models/minicpm.py

@@ -17,11 +17,14 @@ import math
 from typing import Any, Dict, Iterable, Optional, Tuple
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 
+from sglang.srt.configs.minicpm import MiniCPMHybridConfig
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
+    ColumnParallelLinear,
     MergedColumnParallelLinear,
     QKVParallelLinear,
     RowParallelLinear,
@@ -35,9 +38,12 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
-from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.model_loader.weight_utils import (
+    default_weight_loader,
+    sharded_weight_loader,
+)
 from sglang.srt.runtime_context import get_parallel
-from sglang.srt.utils import add_prefix
+from sglang.srt.utils import add_prefix, set_weight_attrs
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
 
@@ -85,11 +91,15 @@ class MiniCPMAttention(nn.Module):
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
+        head_dim: Optional[int] = None,
         layer_id: int = 0,
         rope_theta: float = 10000,
         rope_scaling: Optional[Dict[str, Any]] = None,
         max_position_embeddings: int = 8192,
         quant_config: Optional[QuantizationConfig] = None,
+        attn_use_rope: bool = True,
+        use_output_gate: bool = False,
+        attention_bias: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -108,37 +118,42 @@ class MiniCPMAttention(nn.Module):
             # the KV heads across multiple tensor parallel GPUs.
             assert tp_size % self.total_num_kv_heads == 0
         self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
-        self.head_dim = hidden_size // self.total_num_heads
+        self.head_dim = (
+            head_dim if head_dim is not None else hidden_size // self.total_num_heads
+        )
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
         self.rope_theta = rope_theta
         self.max_position_embeddings = max_position_embeddings
+        self.attn_use_rope = attn_use_rope
+        self.use_output_gate = use_output_gate
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
             self.head_dim,
             self.total_num_heads,
             self.total_num_kv_heads,
-            bias=False,
+            bias=attention_bias,
             quant_config=quant_config,
             prefix=add_prefix("qkv_proj", prefix),
         )
         self.o_proj = RowParallelLinear(
             self.total_num_heads * self.head_dim,
             hidden_size,
-            bias=False,
+            bias=attention_bias,
             quant_config=quant_config,
             prefix=add_prefix("o_proj", prefix),
         )
 
-        self.rotary_emb = get_rope(
+        if self.attn_use_rope:
-            self.head_dim,
+            self.rotary_emb = get_rope(
-            rotary_dim=self.head_dim,
+                self.head_dim,
-            max_position=max_position_embeddings,
+                rotary_dim=self.head_dim,
-            base=rope_theta,
+                max_position=max_position_embeddings,
-            rope_scaling=rope_scaling,
+                base=rope_theta,
-        )
+                rope_scaling=rope_scaling,
+            )
         self.attn = RadixAttention(
             self.num_heads,
             self.head_dim,
@@ -149,21 +164,194 @@ class MiniCPMAttention(nn.Module):
             prefix=add_prefix("attn", prefix),
         )
 
+        if self.use_output_gate:
+            self.o_gate = ColumnParallelLinear(
+                hidden_size,
+                self.total_num_heads * self.head_dim,
+                bias=attention_bias,
+                quant_config=quant_config,
+                prefix=add_prefix("o_gate", prefix),
+            )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        if self.attn_use_rope:
+            orig_dtype = q.dtype
+            q, k = q.float(), k.float()
+            q, k = self.rotary_emb(positions, q, k)
+            q, k = q.to(orig_dtype), k.to(orig_dtype)
+
+        attn_output = self.attn(q, k, v, forward_batch)
+
+        if self.use_output_gate:
+            o_gate_output, _ = self.o_gate(hidden_states)
+            attn_output = attn_output * F.sigmoid(o_gate_output)
+
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class MiniCPMLightningMixer(nn.Module):
+    """Lightning attention mixer backed by the shared linear-attention backend.
+
+    This is a wrapper that prepares inputs for the backend and handles
+    the QKV projection, normalization, RoPE, and output processing,
+    while delegating the recurrent computation through RadixAttention.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        layer_id: int = 0,
+        rope_theta: float = 10000,
+        rope_scaling: Optional[Dict[str, Any]] = None,
+        max_position_embeddings: int = 8192,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+        use_rope: bool = True,
+        use_output_gate: bool = False,
+        attention_bias: bool = False,
+        rms_norm_eps: float = 1e-6,
+        use_output_norm: bool = False,
+        qk_norm: bool = True,
+        scale: str | float = "1/sqrt(d)",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_parallel().tp_size
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim
+        if scale == "1/sqrt(d)":
+            scaling = self.head_dim ** (-0.5)
+        elif scale == "1/d":
+            scaling = self.head_dim ** (-1.0)
+        elif isinstance(scale, (int, float)):
+            scaling = float(scale)
+        else:
+            raise ValueError(f"Unsupported lightning scale: {scale}")
+        self.use_output_gate = use_output_gate
+        self.attention_bias = attention_bias
+        self.rms_norm_eps = rms_norm_eps
+        self.use_rope = use_rope
+        self.qk_norm = qk_norm
+        self.use_output_norm = use_output_norm
+
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=self.attention_bias,
+            quant_config=quant_config,
+            prefix=add_prefix("qkv_proj", prefix),
+        )
+
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=self.attention_bias,
+            quant_config=quant_config,
+            prefix=add_prefix("o_proj", prefix),
+        )
+
+        if self.use_output_norm:
+            self.o_norm = RMSNorm(self.num_heads * self.head_dim, eps=self.rms_norm_eps)
+            set_weight_attrs(
+                self.o_norm.weight, {"weight_loader": sharded_weight_loader(0)}
+            )
+
+        if self.use_output_gate:
+            self.z_proj = ColumnParallelLinear(
+                self.hidden_size,
+                self.total_num_heads * self.head_dim,
+                bias=self.attention_bias,
+                quant_config=quant_config,
+                prefix=add_prefix("z_proj", prefix),
+            )
+
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, eps=self.rms_norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, eps=self.rms_norm_eps)
+
+        if self.use_rope:
+            self.rotary_emb = get_rope(
+                self.head_dim,
+                rotary_dim=self.head_dim,
+                max_position=max_position_embeddings,
+                base=rope_theta,
+                rope_scaling=rope_scaling,
+            )
+
+        self.attn = RadixAttention(
+            self.num_heads,
+            self.head_dim,
+            scaling,
+            num_kv_heads=self.num_kv_heads,
+            layer_id=layer_id,
+            quant_config=quant_config,
+            prefix=add_prefix("attn", prefix),
+        )
+
     def forward(
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
+        if forward_batch.forward_mode.is_idle():
+            return hidden_states.new_empty(hidden_states.shape[0], self.hidden_size)
+
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        orig_dtype = q.dtype
+
-        q, k = q.float(), k.float()
+        if self.qk_norm:
-        q, k = self.rotary_emb(positions, q, k)
+            q = self.q_norm(q.reshape(-1, self.head_dim))
-        q, k = q.to(orig_dtype), k.to(orig_dtype)
+            k = self.k_norm(k.reshape(-1, self.head_dim))
-        attn_output = self.attn(q, k, v, forward_batch)
+
-        output, _ = self.o_proj(attn_output)
+        if self.use_rope:
-        return output
+            q = q.reshape(-1, self.num_heads * self.head_dim)
+            k = k.reshape(-1, self.num_kv_heads * self.head_dim)
+            orig_dtype = q.dtype
+            q, k = q.float(), k.float()
+            q, k = self.rotary_emb(positions, q, k)
+            q, k = q.to(orig_dtype), k.to(orig_dtype)
+
+        q = q.reshape(-1, self.num_heads, self.head_dim)
+        k = k.reshape(-1, self.num_kv_heads, self.head_dim)
+        v = v.reshape(-1, self.num_kv_heads, self.head_dim)
+
+        o = self.attn(q, k, v, forward_batch)
+
+        if self.use_output_norm:
+            o = self.o_norm(o)
+
+        if self.use_output_gate:
+            z, _ = self.z_proj(hidden_states)
+            o = o * F.sigmoid(z)
+
+        y, _ = self.o_proj(o)
+        return y
 
 
 class MiniCPMDecoderLayer(nn.Module):
@@ -176,20 +364,59 @@ class MiniCPMDecoderLayer(nn.Module):
     ) -> None:
         super().__init__()
         self.config = config
+        self.layer_id = layer_id
         self.hidden_size = config.hidden_size
+        if isinstance(config, MiniCPMHybridConfig):
+            self.mixer_type = config.mixer_types[layer_id]
+            attn_use_rope = config.attn_use_rope
+            attn_use_output_gate = config.attn_use_output_gate
+            attention_bias = config.attention_bias
+        else:
+            self.mixer_type = "minicpm4"
+            attn_use_rope = True
+            attn_use_output_gate = False
+            attention_bias = False
+
         rope_theta, rope_scaling = get_rope_config(config)
         max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
-        self.self_attn = MiniCPMAttention(
+        if self.mixer_type == "minicpm4":
-            hidden_size=self.hidden_size,
+            self.self_attn = MiniCPMAttention(
-            num_heads=config.num_attention_heads,
+                hidden_size=self.hidden_size,
-            num_kv_heads=config.num_key_value_heads,
+                num_heads=config.num_attention_heads,
-            layer_id=layer_id,
+                num_kv_heads=config.num_key_value_heads,
-            rope_theta=rope_theta,
+                head_dim=getattr(config, "head_dim", None),
-            rope_scaling=rope_scaling,
+                layer_id=layer_id,
-            max_position_embeddings=max_position_embeddings,
+                rope_theta=rope_theta,
-            quant_config=quant_config,
+                rope_scaling=rope_scaling,
-            prefix=add_prefix("self_attn", prefix),
+                max_position_embeddings=max_position_embeddings,
-        )
+                quant_config=quant_config,
+                attn_use_rope=attn_use_rope,
+                use_output_gate=attn_use_output_gate,
+                attention_bias=attention_bias,
+                prefix=add_prefix("self_attn", prefix),
+            )
+        elif self.mixer_type == "lightning-attn":
+            self.self_attn = MiniCPMLightningMixer(
+                hidden_size=self.hidden_size,
+                num_heads=config.lightning_nh,
+                num_kv_heads=config.lightning_nkv,
+                head_dim=config.lightning_head_dim,
+                layer_id=layer_id,
+                rope_theta=rope_theta,
+                rope_scaling=rope_scaling,
+                max_position_embeddings=max_position_embeddings,
+                quant_config=quant_config,
+                use_rope=config.lightning_use_rope,
+                use_output_gate=config.use_output_gate,
+                attention_bias=config.attention_bias,
+                rms_norm_eps=config.rms_norm_eps,
+                use_output_norm=config.use_output_norm,
+                qk_norm=config.qk_norm,
+                scale=config.lightning_scale,
+                prefix=add_prefix("self_attn", prefix),
+            )
+        else:
+            raise ValueError(f"Unsupported mixer type: {self.mixer_type}")
         self.mlp = MiniCPMMLP(
             hidden_size=self.hidden_size,
             intermediate_size=config.intermediate_size,
@@ -287,7 +514,7 @@ class MiniCPMModel(nn.Module):
         return hidden_states
 
 
-class MiniCPMForCausalLM(nn.Module):
+class MiniCPMSALAForCausalLM(nn.Module):
     def __init__(
         self,
         config,
@@ -396,4 +623,8 @@ class MiniCPMForCausalLM(nn.Module):
                     weight_loader(param, loaded_weight)
 
 
-EntryClass = MiniCPMForCausalLM
+class MiniCPMForCausalLM(MiniCPMSALAForCausalLM):
+    """Alias for MiniCPM checkpoints whose config uses the HF architecture name."""
+
+
+EntryClass = [MiniCPMSALAForCausalLM, MiniCPMForCausalLM]

python/sglang/srt/server_args.py

@@ -201,6 +201,8 @@ ATTENTION_BACKEND_CHOICES = [
     "trtllm_mha",
     "dual_chunk_flash_attn",
     "hpc_ops",  # HPC-Ops (https://github.com/Tencent/hpc-ops), Hopper (SM90) only, requires --page-size 64
+    "minicpm_flashattn",
+    "minicpm_flashinfer",
     # AMD specific
     "aiter",
     "wave",
@@ -969,7 +971,12 @@ class ServerArgs:
         NS("schedule"),
     ] = False
     disable_radix_cache: A[
-        bool, "Disable RadixAttention for prefix caching.", NS("memory")
+        bool,
+        Arg(
+            help="Disable RadixAttention for prefix caching.",
+            resolvable=True,
+        ),
+        NS("memory"),
     ] = False
     enable_page_major_kv_layout: A[
         bool,

python/sglang/srt/session/streaming_session.py

@@ -465,7 +465,7 @@ class StreamingSession(BasePrefixCache):
                     slot.req_pool_idx, start:end
                 ]
                 self.token_to_kv_pool_allocator.free(kv_indices)
-            self.req_to_token_pool.free_slots.append(slot.req_pool_idx)
+            self.req_to_token_pool.free(slot)
 
         self._free_slot_mamba(slot)
 

python/sglang/srt/utils/hf_transformers/common.py

@@ -49,6 +49,7 @@ from sglang.srt.configs import (
     LagunaConfig,
     LocateAnythingConfig,
     LongcatFlashConfig,
+    MiniCPMHybridConfig,
     MiniCPMV4_6Config,
     MiniCPMV4_6VisionConfig,
     MiniMaxM3VLConfig,
@@ -131,6 +132,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
         KimiK25Config,
         Step3p5Config,
         Step3p7Config,
+        MiniCPMHybridConfig,
         MiniCPMV4_6Config,
         MiniCPMV4_6VisionConfig,
         InklingModelConfig,

test/registered/kernels/benchmark/bench_get_block_table.py

@@ -0,0 +1,49 @@
+import torch
+
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.minicpm_sala import get_block_table
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=20, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+_HEAD_GROUP = 2
+_SPARSE_BLOCK_SIZE = 64
+_TOPK = 96
+
+
+def _make_valid_inputs(token_num: int, topk: int, device: str = "cuda"):
+    """Well-formed inputs shared by both expansion strategies.
+
+    ``seqlen_q_max`` is tied to ``token_num`` so the per-token causal position
+    (``token_pos_in_bs``) never indexes past ``block_table``.
+    """
+    seqlen_q_max = token_num
+    num_blocks = max(1, seqlen_q_max // _SPARSE_BLOCK_SIZE)
+    torch.manual_seed(0)
+    topk_idx = torch.randint(
+        0, num_blocks, (_HEAD_GROUP, token_num, topk), dtype=torch.int32, device=device
+    )
+    block_table = torch.arange(
+        1, seqlen_q_max + 1, dtype=torch.int32, device=device
+    ).reshape(1, seqlen_q_max)
+    token_to_bs = torch.zeros((token_num,), dtype=torch.int32, device=device)
+    token_pos_in_bs = torch.arange(1, token_num + 1, dtype=torch.int32, device=device)
+    seqlen_q = torch.tensor([seqlen_q_max], dtype=torch.int32, device=device)
+    return topk_idx, block_table, token_to_bs, token_pos_in_bs, seqlen_q
+
+
+@marker.parametrize("token_num", [2**n for n in range(9, 15)], [512, 4096])
+@marker.benchmark("provider", ["blockwise", "elementwise"])
+def benchmark(token_num: int, provider: str):
+    inputs = _make_valid_inputs(token_num, _TOPK)
+
+    def fn(*args):
+        return get_block_table(*args, elementwise=provider == "elementwise")
+
+    return marker.do_bench(fn, input_args=inputs)
+
+
+if __name__ == "__main__":
+    benchmark.run()

test/registered/kernels/ops/test_get_block_table.py

@@ -0,0 +1,153 @@
+import pytest
+import torch
+
+from sglang.kernels.jit.minicpm_sala.get_block_table import get_block_table
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+_HEAD_GROUP = 2
+_SPARSE_BLOCK_SIZE = 64
+
+
+def _make_inputs(token_num, seqlen_q_max, topk, batch_size=1, device="cuda"):
+    """Build the same kind of inputs as the original CUDA kernel test."""
+    topk_idx = torch.full(
+        (_HEAD_GROUP, token_num, topk), -1, dtype=torch.int32, device=device
+    )
+    # Plant a few valid blocks at fixed positions, like the original UT.
+    topk_idx[0, 32, 0:2] = torch.tensor([0, 1], dtype=torch.int32, device=device)
+    topk_idx[1, 32, 0:2] = torch.tensor([0, 1], dtype=torch.int32, device=device)
+    topk_idx[1, 64, 0:2] = torch.tensor([0, 1], dtype=torch.int32, device=device)
+    topk_idx[0, 1000, 0:10] = torch.tensor(
+        [0, 1, 5, 11, 14, 16, 17, 25, 26, 27], dtype=torch.int32, device=device
+    )
+
+    block_table = torch.arange(
+        1, seqlen_q_max * batch_size + 1, dtype=torch.int32, device=device
+    ).reshape(batch_size, seqlen_q_max)
+    token_to_bs = torch.zeros((token_num,), dtype=torch.int32, device=device)
+    token_pos_in_bs = torch.arange(1, token_num + 1, dtype=torch.int32, device=device)
+    seqlen_q = torch.tensor([seqlen_q_max], dtype=torch.int32, device=device)
+    return topk_idx, block_table, token_to_bs, token_pos_in_bs, seqlen_q
+
+
+def _make_valid_inputs(
+    token_num,
+    seqlen_q_max,
+    topk,
+    batch_size=1,
+    head_group=_HEAD_GROUP,
+    block_size=_SPARSE_BLOCK_SIZE,
+    device="cuda",
+):
+    """Build inputs with only non-negative block indices."""
+    num_blocks = seqlen_q_max // block_size
+    torch.manual_seed(0)
+    topk_idx = torch.randint(
+        0, num_blocks, (head_group, token_num, topk), dtype=torch.int32, device=device
+    )
+    block_table = torch.arange(
+        1, seqlen_q_max * batch_size + 1, dtype=torch.int32, device=device
+    ).reshape(batch_size, seqlen_q_max)
+    token_to_bs = torch.zeros((token_num,), dtype=torch.int32, device=device)
+    token_pos_in_bs = torch.arange(1, token_num + 1, dtype=torch.int32, device=device)
+    seqlen_q = torch.tensor([seqlen_q_max], dtype=torch.int32, device=device)
+    return topk_idx, block_table, token_to_bs, token_pos_in_bs, seqlen_q
+
+
+def _get_block_table_reference(
+    topk_idx,
+    block_table,
+    token_to_bs,
+    token_pos_in_bs,
+    seqlen_q,
+    block_size=_SPARSE_BLOCK_SIZE,
+):
+    head_group = topk_idx.shape[0]
+    token_num = topk_idx.shape[1]
+    source = topk_idx.permute(1, 0, 2).unsqueeze(-1) * block_size + torch.arange(
+        block_size, device=topk_idx.device
+    )
+    valid = (source >= 0) & (
+        source
+        < torch.minimum(seqlen_q[token_to_bs], token_pos_in_bs).view(token_num, 1, 1, 1)
+    )
+    gathered = torch.gather(
+        block_table[token_to_bs],
+        1,
+        source.reshape(token_num, -1).clamp(0, block_table.shape[1] - 1),
+    ).view_as(source)
+    heads = torch.arange(head_group, device=topk_idx.device).view(1, -1, 1, 1)
+    return torch.where(valid, gathered * head_group + heads, 0).flatten(2)
+
+
+def test_get_block_table_supports_tp_local_head_group():
+    inputs = _make_valid_inputs(64, 64, 96, head_group=1)
+    expected = _get_block_table_reference(*inputs)
+    actual = get_block_table(*inputs, head_group_num=1, elementwise=False)
+    assert torch.equal(expected, actual)
+
+
+def _golden_check_blockwise(out_block_table, block_table, token_num):
+    """The assertions ported verbatim from the original kernel test."""
+    # check token 32
+    assert (out_block_table[32, 0] != 0).sum().item() == 33
+    assert (out_block_table[32, 1] != 0).sum().item() == 33
+    assert torch.equal(out_block_table[32, 0, 0:33], block_table[0][:33] * 2)
+    assert torch.equal(out_block_table[32, 1, 0:33], block_table[0][:33] * 2 + 1)
+
+    # check token 64
+    assert (out_block_table[64, 1] != 0).sum().item() == 65
+    assert torch.equal(out_block_table[64, 1, 0:65], block_table[0][:65] * 2 + 1)
+
+    # check token 1000
+    topk_blocks = [0, 1, 5, 11, 14, 16, 17, 25, 26, 27]
+    tokens = []
+    for b in topk_blocks:
+        tokens.extend(range(b * _SPARSE_BLOCK_SIZE, (b + 1) * _SPARSE_BLOCK_SIZE))
+    tokens = [t for t in tokens if t < token_num and t < 1001]
+    assert (out_block_table[1000, 0] != 0).sum().item() == len(tokens)
+    assert torch.equal(
+        out_block_table[1000, 0, : len(tokens)], block_table[0][tokens] * 2
+    )
+
+
+@pytest.mark.parametrize("topk", [96, 128])
+def test_get_block_table_blockwise_golden(topk):
+    token_num, seqlen_q_max = 8192, 8192
+    inputs = _make_inputs(token_num, seqlen_q_max, topk)
+    out = get_block_table(*inputs, elementwise=False)
+    assert out.shape == (token_num, _HEAD_GROUP, topk * _SPARSE_BLOCK_SIZE)
+    _golden_check_blockwise(out, inputs[1], token_num)
+
+
+@pytest.mark.parametrize("topk", [96, 128])
+def test_get_block_table_strategies_match_reference(topk):
+    """Both expansion strategies match the Torch reference, including -1."""
+    token_num, seqlen_q_max = 2048, 2048
+    inputs = _make_inputs(token_num, seqlen_q_max, topk)
+    expected = _get_block_table_reference(*inputs)
+    assert torch.equal(expected, get_block_table(*inputs, elementwise=False))
+    assert torch.equal(expected, get_block_table(*inputs, elementwise=True))
+
+
+@pytest.mark.parametrize(("topk", "block_size"), [(10, 32), (7, 128)])
+def test_get_block_table_supports_configured_layout(topk, block_size):
+    token_num = seqlen_q_max = 256
+    inputs = _make_valid_inputs(
+        token_num,
+        seqlen_q_max,
+        topk,
+        block_size=block_size,
+    )
+    expected = _get_block_table_reference(*inputs, block_size=block_size)
+    kwargs = {"head_group_num": _HEAD_GROUP, "block_size": block_size}
+    assert torch.equal(expected, get_block_table(*inputs, **kwargs, elementwise=False))
+    assert torch.equal(expected, get_block_table(*inputs, **kwargs, elementwise=True))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/test_minicpm_compress_k.py

@@ -0,0 +1,83 @@
+import pytest
+import torch
+
+from sglang.srt.layers.attention.minicpm.fuse_kernel import (
+    fused_attn_pooling_online_topk_decode,
+)
+from sglang.srt.layers.attention.minicpm.sparse_utils import compress_k_core_new
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def test_compress_k_writes_each_head_once():
+    """Each compressed output must be produced once even with multiple KV heads."""
+    key_cache = torch.arange(
+        9 * 2 * 6,
+        dtype=torch.float32,
+        device="cuda",
+    ).reshape(9, 2, 6)
+    original = key_cache.clone()
+    token_table = torch.tensor([[0, 0, 0, 1, 2, 3]], dtype=torch.int32, device="cuda")
+    compressed_table = torch.tensor([[6, 7, 8]], dtype=torch.int32, device="cuda")
+    full_compressed = torch.empty((3, 2, 6), device="cuda")
+
+    compress_k_core_new(
+        full_compressed,
+        1,
+        key_cache,
+        token_table,
+        compressed_table,
+        torch.tensor([0, 4], dtype=torch.int32, device="cuda"),
+        torch.tensor([1], dtype=torch.int32, device="cuda"),
+        torch.tensor([0, 3], dtype=torch.int32, device="cuda"),
+        2,
+        2,
+        6,
+    )
+
+    expected = torch.stack(
+        (
+            original[6],
+            original[0:2].mean(dim=0),
+            original[2:4].mean(dim=0),
+        )
+    )
+    torch.testing.assert_close(full_compressed, expected)
+    torch.testing.assert_close(key_cache[7:9], expected[1:])
+
+
+def test_fused_decode_topk_skips_dense_rows():
+    kernel = fused_attn_pooling_online_topk_decode(
+        batch_size=2,
+        groups=16,
+        heads=16,
+        dim=128,
+        topk=8,
+        pooled_k_len=8,
+        dense_len=5,
+        dtype_str="bfloat16",
+    )
+    topk_indices = torch.full((1, 2, 8), -1, dtype=torch.int32, device="cuda")
+    topk_values = torch.full(
+        (1, 2, 8), float("-inf"), dtype=torch.float32, device="cuda"
+    )
+
+    kernel(
+        torch.randn(32, 1, 128, dtype=torch.bfloat16, device="cuda"),
+        torch.randn(4, 1, 128, dtype=torch.bfloat16, device="cuda"),
+        torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda"),
+        torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda"),
+        torch.tensor([3, 7], dtype=torch.int32, device="cuda"),
+        topk_indices,
+        topk_values,
+    )
+
+    assert torch.all(topk_indices[:, 0] == -1)
+    assert torch.any(topk_indices[:, 1] >= 0)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/unit/configs/test_minicpm_config.py

@@ -0,0 +1,349 @@
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.configs.hybrid_arch import (
+    hybrid_lightning_config,
+    mambaish_config,
+)
+from sglang.srt.configs.linear_attn_model_registry import (
+    get_linear_attn_config,
+    get_linear_attn_spec_by_arch,
+)
+from sglang.srt.configs.mamba_utils import Mamba2CacheParams
+from sglang.srt.configs.minicpm import MiniCPMHybridConfig
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    MambaAttnBackendBase,
+)
+from sglang.srt.layers.attention.linear.lightning_backend import (
+    LightningAttentionBackend,
+)
+from sglang.srt.models import minicpm as minicpm_module
+from sglang.srt.models.minicpm import (
+    MiniCPMAttention,
+    MiniCPMDecoderLayer,
+    MiniCPMLightningMixer,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def test_minicpm_lightning_config_defaults_are_complete():
+    """A checkpoint missing optional SALA fields must still define every model input."""
+    config = MiniCPMHybridConfig()
+
+    assert config.scale_emb == 12
+    assert config.scale_depth == 1.4
+    assert config.dim_model_base == 256
+    assert config.lightning_use_rope is True
+    assert config.use_output_gate is False
+    assert config.attention_bias is False
+    assert config.use_output_norm is False
+    assert config.qk_norm is True
+
+
+def test_minicpm_empty_mixer_types_default_to_full_attention():
+    config = MiniCPMHybridConfig(num_hidden_layers=3, mixer_types=[])
+
+    assert config.mixer_types == ["minicpm4", "minicpm4", "minicpm4"]
+    assert config.full_attention_layer_ids == [0, 1, 2]
+
+
+def test_minicpm_sparse_config_uses_nested_fields_only():
+    sparse_config = {
+        "block_size": 64,
+        "dense_len": 8192,
+        "init_blocks": 1,
+        "kernel_size": 32,
+        "kernel_stride": 16,
+        "topk": 64,
+        "window_size": 2048,
+    }
+    config = MiniCPMHybridConfig(sparse_config=sparse_config)
+
+    assert config.has_minicpm_sparse_attention
+    assert config.sparse_config == sparse_config
+    assert not hasattr(config, "sparse_dense_len")
+
+
+def test_minicpm_short_mixer_pattern_repeats_to_layer_count():
+    config = MiniCPMHybridConfig(
+        num_hidden_layers=5,
+        mixer_types=["minicpm4", "lightning-attn"],
+        lightning_nkv=32,
+    )
+
+    assert config.mixer_types == [
+        "minicpm4",
+        "lightning-attn",
+        "minicpm4",
+        "lightning-attn",
+        "minicpm4",
+    ]
+    assert config.full_attention_layer_ids == [0, 2, 4]
+    assert config.lightning_layer_ids == [1, 3]
+
+
+def test_minicpm_mixer_aliases_are_canonicalized():
+    config = MiniCPMHybridConfig(
+        num_hidden_layers=4,
+        mixer_types=["attention", "lightning_attn"],
+        lightning_nkv=32,
+    )
+
+    assert config.mixer_types == [
+        "minicpm4",
+        "lightning-attn",
+        "minicpm4",
+        "lightning-attn",
+    ]
+
+
+def test_minicpm_rejects_more_mixer_types_than_layers():
+    with pytest.raises(ValueError, match="Invalid number of mixer types: 3"):
+        MiniCPMHybridConfig(
+            num_hidden_layers=2,
+            mixer_types=["minicpm4", "lightning", "minicpm4"],
+        )
+
+
+def test_minicpm_lightning_dimensions_fall_back_to_base_attention():
+    config = MiniCPMHybridConfig(
+        hidden_size=96,
+        num_attention_heads=6,
+        num_key_value_heads=3,
+        head_dim=None,
+        lightning_nh=None,
+        lightning_nkv=None,
+        lightning_head_dim=None,
+    )
+
+    assert config.head_dim == 16
+    assert config.lightning_nh == 6
+    assert config.lightning_nkv == 3
+    assert config.lightning_head_dim == 16
+
+
+def test_minicpm_rejects_lightning_gqa():
+    with pytest.raises(ValueError, match="seg_la backend does not support GQA"):
+        MiniCPMHybridConfig(
+            num_attention_heads=6,
+            num_key_value_heads=3,
+            mixer_types=["lightning-attn"],
+        )
+
+
+def test_minicpm_lightning_idle_batch_returns_empty_output():
+    """An idle DP rank must return empty output instead of reducing empty tensors."""
+    mixer = MiniCPMLightningMixer.__new__(MiniCPMLightningMixer)
+    torch.nn.Module.__init__(mixer)
+    mixer.hidden_size = 8
+    forward_batch = SimpleNamespace(forward_mode=SimpleNamespace(is_idle=lambda: True))
+
+    output = mixer.forward(
+        positions=torch.empty(0, dtype=torch.int64),
+        hidden_states=torch.empty(0, 4),
+        forward_batch=forward_batch,
+    )
+
+    assert output.shape == (0, 8)
+
+
+def test_minicpm_lightning_attention_bias_applies_to_every_projection():
+    """Enabling attention bias must cover every Lightning projection."""
+    with get_parallel().override(tp_size=1, tp_rank=0):
+        mixer = MiniCPMLightningMixer(
+            hidden_size=8,
+            num_heads=2,
+            num_kv_heads=2,
+            head_dim=4,
+            use_rope=False,
+            use_output_gate=True,
+            attention_bias=True,
+            qk_norm=False,
+        )
+
+    assert mixer.qkv_proj.bias is not None
+    assert mixer.o_proj.bias is not None
+    assert mixer.z_proj.bias is not None
+
+
+def test_minicpm_lightning_rejects_unknown_scale():
+    with (
+        get_parallel().override(tp_size=1, tp_rank=0),
+        pytest.raises(ValueError, match="Unsupported lightning scale"),
+    ):
+        MiniCPMLightningMixer(
+            hidden_size=8,
+            num_heads=2,
+            num_kv_heads=2,
+            head_dim=4,
+            use_rope=False,
+            qk_norm=False,
+            scale="unknown",
+        )
+
+
+def test_minicpm_full_attention_bias_applies_to_every_projection():
+    """Enabling attention bias must cover every full-attention projection."""
+    with get_parallel().override(tp_size=1, tp_rank=0):
+        mixer = MiniCPMAttention(
+            hidden_size=8,
+            num_heads=2,
+            num_kv_heads=2,
+            attn_use_rope=False,
+            use_output_gate=True,
+            attention_bias=True,
+        )
+
+    assert mixer.qkv_proj.bias is not None
+    assert mixer.o_proj.bias is not None
+    assert mixer.o_gate.bias is not None
+
+
+def test_minicpm_full_attention_uses_configured_head_dim(monkeypatch):
+    monkeypatch.setattr(minicpm_module, "SiluAndMul", torch.nn.Identity)
+    config = MiniCPMHybridConfig(
+        hidden_size=16,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=6,
+        intermediate_size=32,
+        attn_use_rope=False,
+    )
+
+    with get_parallel().override(tp_size=1, tp_rank=0):
+        layer = MiniCPMDecoderLayer(config)
+
+    assert layer.self_attn.head_dim == 6
+    assert layer.self_attn.q_size == 12
+    assert layer.self_attn.kv_size == 12
+
+
+def test_minicpm_lightning_reuses_shared_backend_and_cache_shape():
+    config = MiniCPMHybridConfig(
+        num_hidden_layers=2,
+        mixer_types=["lightning", "minicpm4"],
+        lightning_nh=4,
+        lightning_nkv=4,
+        lightning_head_dim=64,
+    )
+
+    model_config = SimpleNamespace(
+        hf_config=config,
+        linear_attn_registry_result=get_linear_attn_config(config),
+    )
+    assert hybrid_lightning_config(model_config) is config
+    assert mambaish_config(model_config) is config
+
+    with get_parallel().override(attn_tp_size=1):
+        cache = config.mamba2_cache_params
+    assert isinstance(cache, Mamba2CacheParams)
+    assert cache.layers == [0]
+    assert cache.shape.conv == [(0, 0)]
+    assert cache.shape.temporal == (4, 64, 64)
+    assert config.num_linear_key_value_heads == 4
+
+    with get_parallel().override(attn_tp_size=1, attn_tp_rank=0):
+        slopes = LightningAttentionBackend._build_slope_tensor(
+            4, 2, device="cpu", layerwise_decay=False
+        )
+    assert len(slopes) == 2
+    assert slopes[0].equal(slopes[1])
+
+
+def test_non_lightning_minicpm_is_not_classified_as_linear_attention():
+    config = MiniCPMHybridConfig(
+        num_hidden_layers=1,
+        mixer_types=["minicpm4"],
+        sparse_config={},
+    )
+    model_config = SimpleNamespace(
+        hf_config=config,
+        linear_attn_registry_result=get_linear_attn_config(config),
+    )
+
+    assert hybrid_lightning_config(model_config) is None
+    assert mambaish_config(model_config) is None
+    for architecture in ("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM"):
+        assert get_linear_attn_spec_by_arch(architecture) is None
+
+
+def test_lightning_backend_reads_structural_linear_config(monkeypatch):
+    def fake_base_init(self, model_runner):
+        self.topk = 1
+
+    monkeypatch.setattr(MambaAttnBackendBase, "__init__", fake_base_init)
+    config = SimpleNamespace(
+        num_attention_heads=8,
+        num_linear_key_value_heads=4,
+        num_hidden_layers=2,
+        lightning_layerwise_decay=False,
+    )
+    model_runner = SimpleNamespace(
+        req_to_token_pool=SimpleNamespace(
+            mamba_pool=SimpleNamespace(
+                mamba_cache=SimpleNamespace(conv=[torch.empty(0)])
+            )
+        ),
+        sliding_window_size=None,
+        model_config=SimpleNamespace(
+            hf_config=config,
+            is_encoder_decoder=False,
+            context_len=128,
+            block=256,
+        ),
+        device="cpu",
+        kv_cache_dtype=torch.float32,
+        kv_cache_dtype_str="float32",
+    )
+
+    with get_parallel().override(attn_tp_size=1, attn_tp_rank=0):
+        backend = LightningAttentionBackend(model_runner)
+
+    assert [slope.shape for slope in backend.tp_slope] == [(4, 1, 1), (4, 1, 1)]
+    assert backend.tp_slope[0].equal(backend.tp_slope[1])
+
+
+def test_lightning_backend_uses_layer_scale(monkeypatch):
+    """Each layer's attention scale must reach the linear-attention computation."""
+    captured = {}
+
+    def fake_seg_la_fwd(**kwargs):
+        captured.update(kwargs)
+        return kwargs["q"]
+
+    monkeypatch.setattr(
+        "sglang.srt.layers.attention.linear.lightning_backend.seg_la_fwd",
+        fake_seg_la_fwd,
+    )
+    backend = LightningAttentionBackend.__new__(LightningAttentionBackend)
+    backend.tp_slope = [torch.ones(1, 1, 1)]
+    layer = SimpleNamespace(layer_id=0, scaling=0.25)
+    metadata = SimpleNamespace(
+        batch_size=1,
+        query_start_loc=torch.tensor([0, 1]),
+        has_initial_states=torch.tensor([False]),
+    )
+    q = torch.ones(1, 1, 1)
+
+    backend._linear_attention_entry(
+        q=q,
+        k=q,
+        v=q,
+        kv_cache=torch.zeros(1, 1, 1, 1),
+        state_indices_tensor=torch.tensor([0]),
+        metadata=metadata,
+        layer=layer,
+    )
+
+    assert captured["softmax_scale"] == 0.25
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/disaggregation/test_decode_req_to_token_pool.py

@@ -0,0 +1,55 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.srt.disaggregation.decode import (
+    DecodeReqToTokenPool,
+    HybridMambaDecodeReqToTokenPool,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+def _init_decode_pool(pool):
+    DecodeReqToTokenPool.__init__(
+        pool,
+        size=1,
+        max_context_len=4,
+        device="cpu",
+        enable_memory_saver=False,
+        pre_alloc_size=1,
+    )
+    return pool
+
+
+def test_decode_pool_reports_physical_capacity():
+    pool = _init_decode_pool(DecodeReqToTokenPool.__new__(DecodeReqToTokenPool))
+
+    assert pool.schedulable_token_capacity(17) == 17
+
+
+def test_decode_pool_supports_noop_aux_cache_contract():
+    pool = _init_decode_pool(DecodeReqToTokenPool.__new__(DecodeReqToTokenPool))
+    req_to_token = pool.req_to_token.clone()
+
+    pool.alloc_aux_to_lengths(
+        req_pool_indices_cpu=torch.tensor([1]),
+        target_seq_lens_cpu=torch.tensor([3]),
+    )
+    pool.reset_aux_cache_allocator()
+
+    assert torch.equal(pool.req_to_token, req_to_token)
+
+
+def test_hybrid_decode_pool_initializes_aux_cache_contract():
+    pool = _init_decode_pool(
+        HybridMambaDecodeReqToTokenPool.__new__(HybridMambaDecodeReqToTokenPool)
+    )
+
+    assert pool.schedulable_token_capacity(17) == 17
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/layers/test_flashattention_paged_mha.py

@@ -0,0 +1,91 @@
+import sys
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock, patch
+
+import torch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+with patch.dict(
+    sys.modules,
+    {
+        module: MagicMock()
+        for module in (
+            "sgl_kernel",
+            "sgl_kernel.quantization",
+            "sgl_kernel.scalar_type",
+        )
+    },
+):
+    from sglang.srt.layers.attention.flashattention_backend import (
+        FlashAttentionBackend,
+    )
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def _backend():
+    backend = FlashAttentionBackend.__new__(FlashAttentionBackend)
+    backend.page_size = 1
+    backend.kv_cache_dtype = torch.float16
+    backend.kv_cache_dtype_str = "float8_e4m3fn"
+    backend.kv_cache_is_mxfp8 = False
+    backend.fa_impl_ver = 3
+    backend.num_splits = 4
+    return backend
+
+
+class TestFlashAttentionPagedMHA(unittest.TestCase):
+    def test_get_paged_mha_kv_cache_supports_head_groups(self):
+        backend = _backend()
+        backend.token_to_kv_pool = SimpleNamespace(
+            get_kv_buffer=Mock(
+                return_value=(
+                    torch.empty(8, 2, 16),
+                    torch.empty(8, 2, 16),
+                )
+            )
+        )
+        layer = SimpleNamespace(
+            layer_id=3,
+            tp_k_head_num=2,
+            tp_v_head_num=2,
+            head_dim=16,
+            v_head_dim=16,
+        )
+
+        key_cache, value_cache = backend.get_paged_mha_kv_cache(
+            layer,
+            head_group_num=2,
+        )
+
+        self.assertEqual(key_cache.shape, (16, 1, 1, 16))
+        self.assertEqual(value_cache.shape, (16, 1, 1, 16))
+
+    def test_prepare_paged_mha_query_reuses_fa_scaling_policy(self):
+        backend = _backend()
+        layer = SimpleNamespace(
+            head_dim=16,
+            k_scale=torch.tensor(2.0),
+            v_scale=torch.tensor(4.0),
+        )
+        q = torch.ones(2, 16, dtype=torch.bfloat16)
+
+        q, _, _, k_descale, v_descale = backend.prepare_paged_mha_query(
+            q,
+            None,
+            None,
+            layer,
+            logical_batch_size=2,
+            kv_head_num=1,
+            is_prefill=True,
+        )
+
+        self.assertEqual(q.dtype, torch.float16)
+        self.assertEqual(k_descale.shape, (2, 1))
+        self.assertEqual(v_descale.shape, (2, 1))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/test_minicpm_attention_adapter.py

@@ -0,0 +1,214 @@
+import sys
+import unittest
+from types import ModuleType, SimpleNamespace
+from unittest.mock import MagicMock, Mock, patch
+
+import torch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+with patch.dict(
+    sys.modules,
+    {
+        module: MagicMock()
+        for module in (
+            "sgl_kernel",
+            "sgl_kernel.quantization",
+            "sgl_kernel.scalar_type",
+        )
+    },
+):
+    from sglang.srt.layers.attention.minicpm import attention_adapter as adapter_module
+    from sglang.srt.layers.attention.minicpm.attention_adapter import (
+        MiniCPMFlashAttentionAdapter,
+        MiniCPMFlashInferAdapter,
+    )
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def _metadata(rows=1):
+    return SimpleNamespace(
+        sparse_page_table=torch.zeros((rows, 4), dtype=torch.int32),
+        sparse_cache_seqlens_int32=torch.full(
+            (rows,),
+            4,
+            dtype=torch.int32,
+        ),
+        sparse_cu_seqlens_q=torch.arange(rows + 1, dtype=torch.int32),
+        sparse_cu_seqlens_k=torch.arange(
+            0,
+            (rows + 1) * 4,
+            4,
+            dtype=torch.int32,
+        ),
+        sparse_max_seq_len_q=1,
+        max_seq_len_q=1,
+    )
+
+
+class TestMiniCPMAttentionAdapter(unittest.TestCase):
+    def test_flashattention_adapter_owns_kernel_arguments(self):
+        expected = torch.ones(1, 1, 1)
+        flash_attn_backend = SimpleNamespace(
+            num_splits=4,
+            fa_impl_ver=3,
+        )
+        adapter = MiniCPMFlashAttentionAdapter(flash_attn_backend)
+        metadata = _metadata()
+        layer = SimpleNamespace(scaling=0.125, logit_cap=0.0)
+        k_descale = torch.tensor([[2.0]])
+        v_descale = torch.tensor([[4.0]])
+
+        with patch.object(
+            adapter_module,
+            "flash_attn_with_kvcache",
+            return_value=expected,
+        ) as kernel:
+            result = adapter.forward(
+                torch.ones(1, 1, 1),
+                torch.ones(4, 1, 1, 1),
+                torch.ones(4, 1, 1, 1),
+                metadata,
+                layer,
+                is_prefill=True,
+                k_descale=k_descale,
+                v_descale=v_descale,
+            )
+
+        self.assertIs(result, expected)
+        kwargs = kernel.call_args.kwargs
+        self.assertIs(kwargs["page_table"], metadata.sparse_page_table)
+        self.assertIs(kwargs["k_descale"], k_descale)
+        self.assertIs(kwargs["v_descale"], v_descale)
+        self.assertEqual(kwargs["num_splits"], 4)
+        self.assertEqual(kwargs["ver"], 3)
+
+    def test_flashinfer_prefill_plans_once_and_executes_each_layer(self):
+        adapter = MiniCPMFlashInferAdapter.__new__(MiniCPMFlashInferAdapter)
+        adapter.prefill_planned = False
+        adapter._prepare = Mock()
+        adapter.active_rows = torch.tensor([0], dtype=torch.int32)
+        adapter.active_kv_indptr = torch.tensor([0, 4], dtype=torch.int32)
+        adapter.active_kv_indices = torch.empty(4, dtype=torch.int32)
+        expected = torch.ones(1, 1, 1)
+        adapter.active_wrapper = SimpleNamespace(forward=Mock(return_value=expected))
+        metadata = _metadata()
+        layer = SimpleNamespace(
+            scaling=0.125,
+            logit_cap=0.0,
+            k_scale_float=1.0,
+            v_scale_float=1.0,
+        )
+
+        with patch.object(
+            adapter_module,
+            "create_flashinfer_kv_indices_triton",
+        ) as index_kernel:
+            first = adapter.forward(
+                torch.ones(1, 1, 1),
+                torch.ones(4, 1, 1, 1),
+                torch.ones(4, 1, 1, 1),
+                metadata,
+                layer,
+                is_prefill=True,
+            )
+            second = adapter.forward(
+                torch.ones(1, 1, 1),
+                torch.ones(4, 1, 1, 1),
+                torch.ones(4, 1, 1, 1),
+                metadata,
+                layer,
+                is_prefill=True,
+            )
+
+        self.assertIs(first, expected)
+        self.assertIs(second, expected)
+        adapter._prepare.assert_called_once_with(metadata, is_prefill=True)
+        self.assertEqual(index_kernel.__getitem__.return_value.call_count, 2)
+        self.assertEqual(adapter.active_wrapper.forward.call_count, 2)
+
+    def test_flashinfer_graph_uses_backend_wrapper_cache(self):
+        adapter = MiniCPMFlashInferAdapter.__new__(MiniCPMFlashInferAdapter)
+        adapter.device = torch.device("cpu")
+        adapter.head_group_num = 2
+        adapter.num_qo_heads = 4
+        adapter.num_kv_heads = 1
+        adapter.head_dim = 16
+        adapter.page_size = 1
+        adapter.max_kv_tokens_per_row = 4
+        adapter.q_dtype = torch.float16
+        adapter.kv_dtype = torch.float16
+        adapter.kv_indptr = torch.zeros(3, dtype=torch.int32)
+        adapter.kv_indices = torch.zeros(8, dtype=torch.int32)
+        adapter.kv_last_page_len = torch.ones(2, dtype=torch.int32)
+        adapter.rows = torch.arange(2, dtype=torch.int32)
+        wrapper = SimpleNamespace(begin_forward=Mock())
+        adapter.flashinfer_backend = SimpleNamespace(
+            get_cuda_graph_decode_wrappers=Mock(return_value=[wrapper]),
+        )
+        metadata = _metadata(rows=2)
+
+        adapter.prepare_forward(
+            metadata,
+            is_prefill=False,
+            graph=True,
+        )
+
+        adapter.flashinfer_backend.get_cuda_graph_decode_wrappers.assert_called_once_with(
+            bs=1,
+            num_tokens=2,
+        )
+        wrapper.begin_forward.assert_called_once()
+        self.assertIs(adapter.active_wrapper, wrapper)
+
+    def test_flashinfer_decode_indices_cover_dense_rows(self):
+        wrapper = SimpleNamespace(begin_forward=Mock())
+        flashinfer_backend = SimpleNamespace(decode_wrappers=[wrapper])
+        flashinfer_backend_module = ModuleType(
+            "sglang.srt.layers.attention.flashinfer_backend"
+        )
+        flashinfer_backend_module.FlashInferAttnBackend = Mock(
+            return_value=flashinfer_backend
+        )
+        model_runner = SimpleNamespace(
+            device=torch.device("cpu"),
+            dtype=torch.float16,
+            kv_cache_dtype=torch.float16,
+            req_to_token_pool=SimpleNamespace(size=1),
+        )
+
+        with (
+            patch.object(adapter_module, "is_flashinfer_available", return_value=True),
+            patch.dict(
+                sys.modules,
+                {
+                    "sglang.srt.layers.attention.flashinfer_backend": (
+                        flashinfer_backend_module
+                    ),
+                },
+            ),
+        ):
+            adapter = MiniCPMFlashInferAdapter(
+                model_runner,
+                head_group_num=2,
+                heads_per_group=16,
+                head_dim=128,
+                page_size=1,
+                max_kv_tokens_per_row=7,
+            )
+
+        metadata = _metadata(rows=2)
+        metadata.sparse_cache_seqlens_int32.fill_(7)
+        adapter.prepare_forward(
+            metadata,
+            is_prefill=False,
+            graph=False,
+        )
+
+        self.assertEqual(adapter.kv_indices.numel(), 14)
+        self.assertEqual(adapter.active_kv_indices.numel(), 14)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/test_minicpm_sparse_cache.py

@@ -0,0 +1,344 @@
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.minicpm.cache import (
+    attach_compressed_cache,
+)
+from sglang.srt.managers.scheduler_components.invariant_checker import (
+    SchedulerInvariantChecker,
+)
+from sglang.srt.managers.scheduler_components.pool_stats_observer import (
+    SchedulerPoolStatsObserver,
+)
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class RecordingAllocator:
+    def __init__(self, capacity: int):
+        self.capacity = capacity
+        self.page_size = 1
+        self.next_slot = 1
+        self.live: set[int] = set()
+
+    @property
+    def size(self):
+        return self.capacity
+
+    def alloc(self, size: int):
+        if size > self.available_size():
+            return None
+        slots = torch.arange(self.next_slot, self.next_slot + size, dtype=torch.int64)
+        self.next_slot += size
+        self.live.update(slots.tolist())
+        return slots
+
+    def free(self, slots: torch.Tensor):
+        self.live.difference_update(slots.tolist())
+
+    def available_size(self):
+        return self.capacity - len(self.live)
+
+    def clear(self):
+        self.next_slot = 1
+        self.live.clear()
+
+
+def make_pool_and_req(capacity: int = 64):
+    allocator = RecordingAllocator(capacity)
+    pool = ReqToTokenPool(
+        size=2,
+        max_context_len=64,
+        device="cpu",
+        enable_memory_saver=False,
+    )
+    attach_compressed_cache(
+        pool,
+        allocator,
+        kernel_size=4,
+        kernel_stride=2,
+        enable_memory_saver=False,
+    )
+    req = SimpleNamespace(
+        req_pool_idx=None,
+        inflight_middle_chunks=0,
+        kv_committed_len=0,
+    )
+    req_pool_idx = pool.alloc([req])[0]
+    return pool, req, req_pool_idx, allocator
+
+
+def alloc_extend(pool, req_pool_idx: int, seq_len: int):
+    pool.alloc_aux_to_lengths(
+        req_pool_indices_cpu=torch.tensor([req_pool_idx], dtype=torch.int64),
+        target_seq_lens_cpu=torch.tensor([seq_len], dtype=torch.int64),
+    )
+
+
+def test_extend_allocates_at_sparse_boundaries():
+    pool, _, req_pool_idx, allocator = make_pool_and_req()
+    cache = pool._aux_cache
+
+    alloc_extend(pool, req_pool_idx, seq_len=3)
+    assert allocator.available_size() == 39
+    assert len(cache.free_slots) == 25
+
+    alloc_extend(pool, req_pool_idx, seq_len=4)
+    assert allocator.available_size() == 39
+    assert len(cache.free_slots) == 24
+
+    alloc_extend(pool, req_pool_idx, seq_len=16)
+    assert allocator.available_size() == 39
+    assert len(cache.free_slots) == 17
+
+
+def test_chunk_reuse_only_allocates_new_sparse_slots():
+    pool, _, req_pool_idx, _ = make_pool_and_req()
+    cache = pool._aux_cache
+
+    alloc_extend(pool, req_pool_idx, seq_len=8)
+    assert len(cache.free_slots) == 22
+
+    alloc_extend(pool, req_pool_idx, seq_len=12)
+    assert len(cache.free_slots) == 20
+
+    alloc_extend(pool, req_pool_idx, seq_len=12)
+    assert len(cache.free_slots) == 20
+
+
+def test_decode_does_not_duplicate_sparse_slots():
+    """Retrying the same decode position must not allocate duplicate cache slots."""
+    pool, _, req_pool_idx, _ = make_pool_and_req()
+    cache = pool._aux_cache
+    alloc_extend(pool, req_pool_idx, seq_len=15)
+
+    pool.alloc_aux_to_lengths(
+        req_pool_indices_cpu=torch.tensor([req_pool_idx], dtype=torch.int64),
+        target_seq_lens_cpu=torch.tensor([16], dtype=torch.int64),
+    )
+    available_after_first_decode = len(cache.free_slots)
+
+    pool.alloc_aux_to_lengths(
+        req_pool_indices_cpu=torch.tensor([req_pool_idx], dtype=torch.int64),
+        target_seq_lens_cpu=torch.tensor([16], dtype=torch.int64),
+    )
+    assert available_after_first_decode == 17
+    assert len(cache.free_slots) == available_after_first_decode
+
+
+def test_reserve_leaves_only_dense_capacity_visible():
+    allocator = RecordingAllocator(capacity=69)
+    pool = ReqToTokenPool(
+        size=2,
+        max_context_len=64,
+        device="cpu",
+        enable_memory_saver=False,
+    )
+    attach_compressed_cache(
+        pool,
+        allocator,
+        kernel_size=32,
+        kernel_stride=16,
+        enable_memory_saver=False,
+    )
+
+    assert allocator.available_size() == 64
+    assert len(pool._aux_cache.reserved_slots) == 5
+    assert pool.schedulable_token_capacity(69) == 64
+
+
+def test_reserved_slots_are_excluded_from_full_pool_invariant():
+    pool, _, _, allocator = make_pool_and_req(capacity=69)
+    checker = SchedulerInvariantChecker(
+        is_hybrid_swa=False,
+        is_hybrid_ssm=True,
+        disaggregation_mode=None,
+        page_size=1,
+        full_tokens_per_layer=None,
+        swa_tokens_per_layer=None,
+        max_total_num_tokens=64,
+        tree_cache=SimpleNamespace(
+            supports_mamba=lambda: False,
+            protected_size=lambda: 0,
+        ),
+        token_to_kv_pool_allocator=allocator,
+        req_to_token_pool=pool,
+        pool_stats_observer=SimpleNamespace(session_held_tokens=lambda: 0),
+        get_last_batch=lambda: None,
+        get_running_batch=lambda: None,
+    )
+
+    leak, message = checker._check_full_pool(
+        SimpleNamespace(
+            full_available_size=allocator.available_size(), full_evictable_size=0
+        )
+    )
+
+    assert not leak, message
+
+
+def test_hybrid_pool_stats_exclude_reserved_slots():
+    pool, _, _, allocator = make_pool_and_req(capacity=69)
+    pool.mamba_allocator = SimpleNamespace(available_size=lambda: 1)
+    pool.mamba_pool = SimpleNamespace(size=1)
+    observer = SchedulerPoolStatsObserver(
+        tree_cache=SimpleNamespace(supports_mamba=lambda: False),
+        token_to_kv_pool_allocator=allocator,
+        req_to_token_pool=pool,
+        session_controller=None,
+        hisparse_coordinator=None,
+        is_hybrid_swa=False,
+        is_hybrid_ssm=True,
+        enable_hisparse=False,
+        full_tokens_per_layer=None,
+        swa_tokens_per_layer=None,
+        max_total_num_tokens=42,
+        get_last_batch=lambda: None,
+        get_running_batch=lambda: None,
+    )
+
+    stats = observer._get_mamba_token_info()
+
+    assert stats.full_num_used == 0
+    assert stats.full_token_usage == 0
+
+
+def test_streaming_session_release_frees_compressed_slots():
+    pool, _, req_pool_idx, allocator = make_pool_and_req()
+    alloc_extend(pool, req_pool_idx, seq_len=16)
+    dense_slots = allocator.alloc(16)
+    pool.req_to_token[req_pool_idx, :16] = dense_slots.to(torch.int32)
+    compressed_cache = pool._aux_cache
+    assert len(compressed_cache.free_slots) < len(compressed_cache.reserved_slots)
+
+    session = StreamingSession(
+        SimpleNamespace(
+            req_to_token_pool=pool,
+            token_to_kv_pool_allocator=allocator,
+            page_size=1,
+        )
+    )
+    session.slots["session-a"] = SessionSlot(
+        req_pool_idx=req_pool_idx,
+        kv=SimpleNamespace(kv_allocated_len=16),
+    )
+
+    session.release_session("session-a")
+
+    assert req_pool_idx in pool.free_slots
+    assert len(compressed_cache.free_slots) == len(compressed_cache.reserved_slots)
+
+
+def test_mamba_leak_diagnostic_does_not_report_reserved_slots():
+    pool, _, _, allocator = make_pool_and_req(capacity=69)
+    allocator.free_pages = torch.arange(6, 70, dtype=torch.int64)
+    allocator.release_pages = torch.empty(0, dtype=torch.int64)
+    pool.mamba_pool = SimpleNamespace(size=1)
+    pool.mamba_allocator = SimpleNamespace(
+        size=1,
+        free_slots=torch.empty(0, dtype=torch.int64),
+    )
+    checker = SchedulerInvariantChecker(
+        is_hybrid_swa=False,
+        is_hybrid_ssm=True,
+        disaggregation_mode=None,
+        page_size=1,
+        full_tokens_per_layer=None,
+        swa_tokens_per_layer=None,
+        max_total_num_tokens=64,
+        tree_cache=SimpleNamespace(
+            mamba_protected_size=lambda: 0,
+            all_values_flatten=lambda: torch.empty(0, dtype=torch.int64),
+            all_mamba_values_flatten=lambda: torch.empty(0, dtype=torch.int64),
+        ),
+        token_to_kv_pool_allocator=allocator,
+        req_to_token_pool=pool,
+        pool_stats_observer=SimpleNamespace(
+            session_held_mamba_slots=lambda: 0,
+        ),
+        get_last_batch=lambda: None,
+        get_running_batch=lambda: None,
+    )
+
+    leak, message = checker._check_mamba_pool(
+        SimpleNamespace(mamba_available_size=0, mamba_evictable_size=0)
+    )
+
+    assert leak
+    assert "leaked_full_pages" not in message
+    assert "leaked_mamba_pages={1}" in message
+
+
+def test_partial_failure_rolls_back_and_free_releases_every_slot():
+    """A failed cache-level allocation must release slots allocated for other levels."""
+    pool, req, req_pool_idx, allocator = make_pool_and_req(capacity=18)
+    cache = pool._aux_cache
+
+    with pytest.raises(RuntimeError, match="out of reserved slots"):
+        alloc_extend(pool, req_pool_idx, seq_len=16)
+
+    assert allocator.available_size() == 11
+    assert len(cache.free_slots) == 7
+
+    allocator.capacity = 20
+    allocator.clear()
+    pool.reset_aux_cache_allocator()
+    alloc_extend(pool, req_pool_idx, seq_len=16)
+    assert allocator.available_size() == 12
+    assert len(cache.free_slots) == 0
+
+    pool.free(req)
+    assert req.req_pool_idx is None
+    assert allocator.available_size() == 12
+    assert len(cache.free_slots) == 8
+
+
+def test_allocator_reset_rebuilds_reserve():
+    pool, _, _, allocator = make_pool_and_req()
+
+    allocator.clear()
+    assert allocator.available_size() == 64
+
+    pool.reset_aux_cache_allocator()
+    assert allocator.available_size() == 39
+    assert len(pool._aux_cache.free_slots) == 25
+
+
+def test_attach_compressed_cache_is_idempotent():
+    allocator = RecordingAllocator(capacity=64)
+    pool = ReqToTokenPool(
+        size=2,
+        max_context_len=64,
+        device="cpu",
+        enable_memory_saver=False,
+    )
+    attach_compressed_cache(
+        pool,
+        allocator,
+        kernel_size=4,
+        kernel_stride=2,
+        enable_memory_saver=False,
+    )
+    cache = pool._aux_cache
+    k1_table = pool.req_to_sparse_k1_token
+
+    attach_compressed_cache(
+        pool,
+        allocator,
+        kernel_size=4,
+        kernel_stride=2,
+        enable_memory_saver=False,
+    )
+
+    assert pool._aux_cache is cache
+    assert pool.req_to_sparse_k1_token is k1_table
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/mem_cache/test_streaming_session_unit.py

@@ -18,6 +18,16 @@ class _FakeAllocator:
         self.freed.append(free_index.clone())
 
 
+class _FakeReqToTokenPool:
+    def __init__(self, req_to_token):
+        self.req_to_token = req_to_token
+        self.free_slots = []
+
+    def free(self, req):
+        self.free_slots.append(req.req_pool_idx)
+        req.req_pool_idx = None
+
+
 class _FakeInnerCache:
     def __init__(self, req_to_token_pool, allocator, page_size, match_results=None):
         self.req_to_token_pool = req_to_token_pool
@@ -85,7 +95,7 @@ def test_preabort_detaches_session_and_preserves_slot():
     """Pre-aborted req (to_finish set before match_prefix) is detached from
     the session: session=None, abort_req() called. Slot stays intact."""
     req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
-    req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
+    req_to_token_pool = _FakeReqToTokenPool(req_to_token)
     allocator = _FakeAllocator()
     inner = _FakeInnerCache(
         req_to_token_pool,
@@ -133,7 +143,7 @@ def test_first_mid_abort_nukes_ephemeral_slot():
     slot is created from req state and nuked via release_session."""
     page_size = 1
     req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
-    req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
+    req_to_token_pool = _FakeReqToTokenPool(req_to_token)
     allocator = _FakeAllocator()
     inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
     tree_cache = StreamingSession(inner)
@@ -159,7 +169,7 @@ def test_nth_mid_abort_nukes_session_slot():
     in req_nodes for next turn's re-prefill."""
     page_size = 1
     req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
-    req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
+    req_to_token_pool = _FakeReqToTokenPool(req_to_token)
     allocator = _FakeAllocator()
     inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
     tree_cache = StreamingSession(inner)
@@ -197,7 +207,7 @@ def test_release_session_threads_mamba_skip_ids():
     from sglang.srt.mem_cache.unified_cache.components import ComponentType
 
     req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
-    req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
+    req_to_token_pool = _FakeReqToTokenPool(req_to_token)
     allocator = _FakeAllocator()
     inner = _FakeInnerCache(req_to_token_pool, allocator, page_size=1)
     tree_cache = StreamingSession(inner)
@@ -235,7 +245,7 @@ def test_trim_overshoot_postcondition():
     """
     page_size = 1
     req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
-    req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
+    req_to_token_pool = _FakeReqToTokenPool(req_to_token)
     allocator = _FakeAllocator()
     tree_cache = StreamingSession(
         _FakeInnerCache(req_to_token_pool, allocator, page_size)

test/registered/unit/model_executor/model_runner_components/test_attention_backend_setup.py

@@ -80,5 +80,42 @@ def test_split_full_attention_applies_model_wrapper_once():
         override.restore()
 
 
+def test_equal_resolved_backends_ignore_stale_global_backend():
+    runner = SimpleNamespace(
+        server_args=SimpleNamespace(
+            attention_backend="global-test",
+            speculative_attention_mode="prefill",
+        ),
+        kv_cache_dtype=None,
+        token_to_kv_pool=object(),
+        req_to_token_pool=object(),
+        init_new_workspace=None,
+    )
+    constructors = {
+        "global-test": lambda _runner: _FakeBackend("global"),
+        "resolved-test": lambda _runner: _FakeBackend("resolved"),
+    }
+    resolved = ResolvedAttentionBackendStr(
+        decode="resolved-test",
+        prefill="resolved-test",
+    )
+
+    with (
+        patch.dict(attention_backend_setup.ATTENTION_BACKENDS, constructors),
+        patch.object(
+            attention_backend_setup,
+            "attn_backend_wrapper",
+            side_effect=lambda _runner, backend: backend,
+        ),
+    ):
+        result = attention_backend_setup._build_resolved_backend(
+            model_runner=runner,
+            resolved=resolved,
+            init_new_workspace=False,
+        )
+
+    assert result.name == "resolved"
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/test_model_overrides.py

@@ -22,6 +22,7 @@ from sglang.srt.arg_groups.overrides import (
     register_model_override,
     validate_declarations,
 )
+from sglang.srt.configs.minicpm import MiniCPMHybridConfig
 from sglang.srt.environ import envs
 from sglang.srt.runtime_context import (
     get_context,
@@ -77,6 +78,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
                     "dcp_comm_backend",
                     "dcp_replicate_q_proj",
                     "disable_overlap_schedule",
+                    "disable_radix_cache",
                     "uses_mamba_radix_cache",
                     "mamba_radix_cache_strategy",
                     "mamba_full_memory_ratio",
@@ -294,6 +296,235 @@ class TestGoldenModelOverrides(_IsolatedPublish):
         "v_head_dim": 16,
     }
 
+    @staticmethod
+    def _minicpm_overrides(
+        architecture,
+        *,
+        sparse_attention=False,
+        lightning_attention=False,
+        attention_backend=None,
+        prefill_attention_backend=None,
+        decode_attention_backend=None,
+        disaggregation_mode="null",
+        enable_dp_attention=False,
+        enable_hierarchical_cache=False,
+    ):
+        args = SimpleNamespace(
+            attention_backend=attention_backend,
+            prefill_attention_backend=prefill_attention_backend,
+            decode_attention_backend=decode_attention_backend,
+            disaggregation_mode=disaggregation_mode,
+            enable_dp_attention=enable_dp_attention,
+            enable_hierarchical_cache=enable_hierarchical_cache,
+        )
+        args.is_attention_backend_not_set = lambda: all(
+            backend is None
+            for backend in (
+                args.attention_backend,
+                args.prefill_attention_backend,
+                args.decode_attention_backend,
+            )
+        )
+        mixer_types = []
+        if sparse_attention:
+            mixer_types.append("minicpm4")
+        if lightning_attention:
+            mixer_types.append("lightning-attn")
+        if not mixer_types:
+            mixer_types.append("minicpm4")
+        declarations = collect_model_override_declarations(
+            architecture,
+            args,
+            hf_config=MiniCPMHybridConfig(
+                num_hidden_layers=len(mixer_types),
+                num_attention_heads=1,
+                num_key_value_heads=1,
+                mixer_types=mixer_types,
+                sparse_config={} if sparse_attention else None,
+            ),
+        )
+        return {
+            field: value
+            for _, declaration in declarations
+            for field, value in declaration.items()
+        }
+
+    def test_minicpm_disables_radix_cache_only_for_hybrid_layers(self):
+        for architecture in ("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM"):
+            with self.subTest(architecture=architecture):
+                self.assertNotIn(
+                    "disable_radix_cache",
+                    self._minicpm_overrides(architecture),
+                )
+                self.assertTrue(
+                    self._minicpm_overrides(architecture, sparse_attention=True)[
+                        "disable_radix_cache"
+                    ]
+                )
+                self.assertTrue(
+                    self._minicpm_overrides(architecture, lightning_attention=True)[
+                        "disable_radix_cache"
+                    ]
+                )
+
+    def test_minicpm_rejects_dp_attention(self):
+        for architecture in ("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM"):
+            with self.subTest(architecture=architecture):
+                with self.assertRaisesRegex(
+                    ValueError,
+                    "MiniCPM does not support DP attention",
+                ):
+                    self._minicpm_overrides(
+                        architecture,
+                        enable_dp_attention=True,
+                    )
+
+    def test_minicpm_rejects_hierarchical_cache_for_hybrid_models(self):
+        for capability in ("sparse_attention", "lightning_attention"):
+            with self.subTest(capability=capability):
+                with self.assertRaisesRegex(
+                    ValueError,
+                    "MiniCPM SALA does not support hierarchical cache",
+                ):
+                    self._minicpm_overrides(
+                        "MiniCPMSALAForCausalLM",
+                        enable_hierarchical_cache=True,
+                        **{capability: True},
+                    )
+
+    def test_sparse_minicpm_defaults_to_sparse_attention_backend(self):
+        with patch.object(
+            overrides_module,
+            "is_blackwell_supported",
+            return_value=False,
+        ):
+            for architecture in ("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM"):
+                with self.subTest(architecture=architecture):
+                    self.assertEqual(
+                        self._minicpm_overrides(
+                            architecture,
+                            sparse_attention=True,
+                        )["attention_backend"],
+                        "minicpm_flashattn",
+                    )
+
+    def test_minicpm_overrides_use_config_capabilities(self):
+        args = SimpleNamespace(
+            attention_backend=None,
+            prefill_attention_backend=None,
+            decode_attention_backend=None,
+            disaggregation_mode="null",
+            enable_dp_attention=False,
+            enable_hierarchical_cache=False,
+            is_attention_backend_not_set=lambda: True,
+        )
+        config = SimpleNamespace(
+            has_minicpm_sparse_attention=True,
+            has_lightning_layers=False,
+        )
+
+        with patch.object(
+            overrides_module, "is_blackwell_supported", return_value=False
+        ):
+            overrides = overrides_module._minicpm_sala_overrides(args, config)
+
+        self.assertTrue(overrides["disable_radix_cache"])
+        self.assertEqual(overrides["attention_backend"], "minicpm_flashattn")
+
+    def test_sparse_minicpm_defaults_to_flashinfer_on_blackwell(self):
+        with patch.object(
+            overrides_module,
+            "is_blackwell_supported",
+            return_value=True,
+        ):
+            self.assertEqual(
+                self._minicpm_overrides(
+                    "MiniCPMSALAForCausalLM",
+                    sparse_attention=True,
+                )["attention_backend"],
+                "minicpm_flashinfer",
+            )
+
+    def test_minicpm_preserves_explicit_attention_backend(self):
+        overrides = self._minicpm_overrides(
+            "MiniCPMSALAForCausalLM",
+            sparse_attention=True,
+            attention_backend="fa3",
+        )
+        self.assertNotIn("attention_backend", overrides)
+
+    def test_sparse_minicpm_rejects_pd_disaggregation(self):
+        for disaggregation_mode in ("prefill", "decode"):
+            with self.subTest(disaggregation_mode=disaggregation_mode):
+                with self.assertRaisesRegex(
+                    ValueError,
+                    "MiniCPM sparse attention does not support PD disaggregation",
+                ):
+                    self._minicpm_overrides(
+                        "MiniCPMSALAForCausalLM",
+                        sparse_attention=True,
+                        disaggregation_mode=disaggregation_mode,
+                    )
+        for backend_field in (
+            "prefill_attention_backend",
+            "decode_attention_backend",
+        ):
+            with self.subTest(backend_field=backend_field):
+                with self.assertRaisesRegex(
+                    ValueError,
+                    "MiniCPM sparse attention does not support PD disaggregation",
+                ):
+                    self._minicpm_overrides(
+                        "MiniCPMSALAForCausalLM",
+                        sparse_attention=True,
+                        disaggregation_mode="decode",
+                        **{backend_field: "minicpm_flashattn"},
+                    )
+
+    def test_minicpm_force_dense_uses_stock_attention_backend(self):
+        with envs.SGLANG_MINICPM_FORCE_DENSE.override(True):
+            self.assertNotIn(
+                "attention_backend",
+                self._minicpm_overrides(
+                    "MiniCPMSALAForCausalLM",
+                    sparse_attention=True,
+                ),
+            )
+            self.assertEqual(
+                self._minicpm_overrides(
+                    "MiniCPMSALAForCausalLM",
+                    sparse_attention=True,
+                    attention_backend="minicpm_flashinfer",
+                )["attention_backend"],
+                "flashinfer",
+            )
+            with patch.object(
+                overrides_module,
+                "is_blackwell_supported",
+                return_value=True,
+            ):
+                self.assertEqual(
+                    self._minicpm_overrides(
+                        "MiniCPMSALAForCausalLM",
+                        sparse_attention=True,
+                        attention_backend="minicpm_flashattn",
+                    )["attention_backend"],
+                    "fa4",
+                )
+            with patch.object(
+                overrides_module,
+                "is_blackwell_supported",
+                return_value=False,
+            ):
+                split_overrides = self._minicpm_overrides(
+                    "MiniCPMSALAForCausalLM",
+                    sparse_attention=True,
+                    prefill_attention_backend="minicpm_flashattn",
+                    decode_attention_backend="minicpm_flashattn",
+                )
+                self.assertEqual(split_overrides["prefill_attention_backend"], "fa3")
+                self.assertEqual(split_overrides["decode_attention_backend"], "fa3")
+
     def _construct(self, arch, model_type, config_extra=None, **server_kwargs):
         from sglang.srt.server_args import ServerArgs