GLM-5.3-Flash support (#36507)

Co-authored-by: zRzRzRzRzRzRzR <Yuxuan.Zhang2@liverpool.ac.uk>
Co-authored-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
Co-authored-by: zanes-ops <zanes@nvidia.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Jian Chen <jianchen0311@gmail.com>
Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com>
Co-authored-by: andyluo7 <43718156+andyluo7@users.noreply.github.com>
Co-authored-by: Ehsan Akhgari <ehsan.akhgari@gmail.com>
Co-authored-by: kpham-sgl <khoa.pham@radixark.ai>
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Raiden Makoto <81530826+Raiden-Makoto@users.noreply.github.com>

python/sglang/kernels/ops/attention/dsa/quant_k_cache.py

@@ -3,6 +3,57 @@ import triton
 import triton.language as tl
 
 
+def gather_dsa_kv_scales(
+    scale_src,
+    scale_dst,
+    kv_indices,
+    kv_indptr,
+    kv_indptr_idx,
+):
+    _gather_dsa_kv_scales[(32,)](
+        scale_src,
+        scale_dst,
+        kv_indices,
+        kv_indptr,
+        scale_src.stride(0),
+        KV_INDPTR_IDX=kv_indptr_idx,
+        NUM_TILES=scale_src.shape[-1],
+        BLOCK=256,
+    )
+
+
+@triton.jit
+def _gather_dsa_kv_scales(
+    scale_src,
+    scale_dst,
+    kv_indices,
+    kv_indptr,
+    scale_src_stride,
+    KV_INDPTR_IDX: tl.constexpr,
+    NUM_TILES: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    num_programs = tl.num_programs(0)
+    active = tl.load(kv_indptr + KV_INDPTR_IDX)
+    block_start = pid * BLOCK
+    tiles = tl.arange(0, NUM_TILES)
+    while block_start < active:
+        offsets = block_start + tl.arange(0, BLOCK)
+        mask = offsets < active
+        rows = tl.load(kv_indices + offsets, mask=mask, other=0)
+        values = tl.load(
+            scale_src + rows[:, None] * scale_src_stride + tiles[None, :],
+            mask=mask[:, None],
+        )
+        tl.store(
+            scale_dst + rows[:, None] * NUM_TILES + tiles[None, :],
+            values,
+            mask=mask[:, None],
+        )
+        block_start += num_programs * BLOCK
+
+
 def quantize_k_cache(cache_k):
     return _quantize_k_cache_fast_wrapped(cache_k)
 
@@ -22,18 +73,26 @@ def quantize_k_cache_separate(
         k_nope: (num_tokens, dim_nope) or (num_tokens, 1, dim_nope)
                 Must have dim_nope=512 for FP8 MLA quantization
         k_rope: (num_tokens, dim_rope) or (num_tokens, 1, dim_rope)
-                Must have dim_rope=64 for FP8 MLA quantization
+                Must have dim_rope=64 for FP8 MLA quantization, or dim_rope=0
+                for no-PE MLA (empty rope); None is treated
+                the same as an empty rope.
         tile_size: quantization tile size (default 128)
 
     Returns:
         Tuple of (nope_part, rope_part) where:
         - nope_part: (num_tokens, 1, 528) as uint8 view, contains [nope_fp8(512) | scales(16)]
         - rope_part: (num_tokens, 1, 128) as uint8 view, contains [rope_bf16_bytes(128)]
+                     (empty, (num_tokens, 1, 0), when dim_rope=0)
 
         These two tensors can be directly passed to set_mla_kv_buffer_triton(kv_buffer, loc, nope_part, rope_part)
     """
     # Squeeze middle dimension if present
     k_nope_2d = k_nope.squeeze(1) if k_nope.ndim == 3 else k_nope
+    if k_rope is None or k_rope.numel() == 0:
+        k_rope_2d = torch.empty(
+            (k_nope_2d.shape[0], 0), dtype=k_nope_2d.dtype, device=k_nope_2d.device
+        )
+    else:
         k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope
 
     num_tokens = k_nope_2d.shape[0]
@@ -43,8 +102,8 @@ def quantize_k_cache_separate(
     # Validate dimensions for FP8 MLA
     if dim_nope != 512:
         raise ValueError(f"Expected dim_nope=512 for FP8 MLA, got {dim_nope}")
-    if dim_rope != 64:
+    if dim_rope not in (0, 64):
-        raise ValueError(f"Expected dim_rope=64 for FP8 MLA, got {dim_rope}")
+        raise ValueError(f"Expected dim_rope=64 (or 0 for no-PE MLA), got {dim_rope}")
     if k_rope_2d.shape[0] != num_tokens:
         raise ValueError(
             f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}"
@@ -234,7 +293,12 @@ def _quantize_k_cache_fast_separate(k_nope, k_rope, group_size: int = 128):
     # Fixed byte layout for rope_part: [rope_bf16 (dim_rope*2 bytes)]
     nope_q_view = nope_part_u8[:, :dim_nope].view(torch.float8_e4m3fn)
     nope_s_view = nope_part_u8[:, dim_nope:].view(torch.float32)
+    if dim_rope > 0:
         rope_view = rope_part_u8.view(torch.bfloat16)
+    else:
+        rope_view = torch.empty(
+            (num_tokens, 0), dtype=torch.bfloat16, device=k_rope.device
+        )
 
     # Kernel launch parameters
     num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)

python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py

@@ -272,12 +272,13 @@ def sparse_attention_fwd_kernel_v1(
     num_stages=2,
     threads=256,
 ):
-    assert dim == tilelang.math.next_power_of_2(dim), (
+    assert dim == tilelang.math.next_power_of_2(dim) or dim % 64 == 0, (
-        f"haven't check padding correctness yet, dim={dim}"
+        f"dim={dim} must be a power of 2 or a multiple of 64"
     )
-    assert tail_dim == tilelang.math.next_power_of_2(tail_dim), (
+    assert tail_dim == 0 or tail_dim == tilelang.math.next_power_of_2(tail_dim), (
-        f"haven't check padding correctness yet, dim={tail_dim}"
+        f"tail_dim={tail_dim} must be 0 or a power of 2"
     )
+    has_tail = tail_dim > 0
     assert is_causal == True, "non-casual is not supported"
     assert topk % block_I == 0, (
         "otherwise will load some index=0 thus causing wrong kv to be loaded"
@@ -330,8 +331,10 @@ def sparse_attention_fwd_kernel_v1(
             bz,
         ):
             Q_shared = T.alloc_shared([H_per_block, D], dtype)
+            if has_tail:
                 Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype)
             KV_shared = T.alloc_shared([BI, D], dtype)
+            if has_tail:
                 K_tail_shared = T.alloc_shared([BI, D_tail], dtype)
             O_shared = T.alloc_shared([H_per_block, D], dtype)
             mask = T.alloc_fragment([BI], "bool")
@@ -358,6 +361,7 @@ def sparse_attention_fwd_kernel_v1(
             H1 = H0 + H_per_block
 
             T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared)
+            if has_tail:
                 T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared)
 
             for i_i in T.Pipelined(NI, num_stages=num_stages):
@@ -368,9 +372,13 @@ def sparse_attention_fwd_kernel_v1(
                     KV_shared[bi_i, d_i] = KV[
                         b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i
                     ]
+                if has_tail:
                     for bi_i, d_i in T.Parallel(BI, D_tail):
                         K_tail_shared[bi_i, d_i] = KV[
-                        b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i
+                            b_i,
+                            Indices[b_i, s_i, g_i, i_i * BI + bi_i],
+                            g_i,
+                            D + d_i,
                         ]
 
                 for h_i, bi_i in T.Parallel(H_per_block, BI):
@@ -384,6 +392,7 @@ def sparse_attention_fwd_kernel_v1(
                     transpose_B=True,
                     policy=T.GemmWarpPolicy.FullCol,
                 )
+                if has_tail:
                     T.gemm(
                         Q_tail_shared,
                         K_tail_shared,
@@ -1325,7 +1334,7 @@ def tilelang_sparse_fwd(
     dim = q.shape[2]
     tail_dim = dim - d_v
     topk = indices.shape[-1]
-    assert topk == 2048
+    assert topk % 64 == 0, "topk must be padded to a multiple of 64"
 
     if _is_hip:
         is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
@@ -1379,9 +1388,12 @@ def tilelang_sparse_fwd(
         )
         out = kernel_combine(partial_o_batched, partial_lse_batched)
     else:
-        kernel = sparse_attention_fwd_kernel_v2(
+        kernel_factory = (
-            num_heads, d_v, tail_dim, topk, sm_scale=sm_scale
+            sparse_attention_fwd_kernel_v1
+            if tail_dim == 0
+            else sparse_attention_fwd_kernel_v2
         )
+        kernel = kernel_factory(num_heads, d_v, tail_dim, topk, sm_scale=sm_scale)
         out = kernel(q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0))  # type: ignore
     return out
 

python/sglang/kernels/ops/attention/dsa_metadata.py

@@ -90,7 +90,9 @@ def _fused_dsa_decode_metadata_kernel(
     # fused decode CUDA graph drops it and consumes real_page_table alone.
     if HAS_PAGE_TABLE_1:
         tl.store(
-            page_table_1 + row * page_table_stride_0 + offs_n * page_table_stride_1,
+            page_table_1
+            + row.to(tl.int64) * page_table_stride_0
+            + offs_n * page_table_stride_1,
             vals,
             mask=mask,
         )
@@ -100,7 +102,7 @@ def _fused_dsa_decode_metadata_kernel(
         real_cols = offs_n // real_page_size
         tl.store(
             real_page_table
-            + row * real_page_table_stride_0
+            + row.to(tl.int64) * real_page_table_stride_0
             + real_cols * real_page_table_stride_1,
             vals // real_page_size,
             mask=real_mask,
@@ -320,7 +322,9 @@ def _fused_dsa_target_verify_metadata_kernel(
     # fused_dsa_decode_metadata for the optional-page_table_1 contract).
     if HAS_PAGE_TABLE_1:
         tl.store(
-            page_table_1 + out_row * page_table_stride_0 + offs_n * page_table_stride_1,
+            page_table_1
+            + out_row.to(tl.int64) * page_table_stride_0
+            + offs_n * page_table_stride_1,
             vals,
             mask=mask,
         )
@@ -330,7 +334,7 @@ def _fused_dsa_target_verify_metadata_kernel(
         real_cols = offs_n // real_page_size
         tl.store(
             real_page_table
-            + out_row * real_page_table_stride_0
+            + out_row.to(tl.int64) * real_page_table_stride_0
             + real_cols * real_page_table_stride_1,
             vals // real_page_size,
             mask=real_mask,
@@ -592,7 +596,7 @@ def _fused_dsa_draft_extend_metadata_kernel(
     if HAS_PAGE_TABLE_1:
         tl.store(
             page_table_1
-            + out_rows[:, None] * page_table_stride_0
+            + out_rows.to(tl.int64)[:, None] * page_table_stride_0
             + offs_n[None, :] * page_table_stride_1,
             vals[None, :],
             mask=mask,
@@ -603,7 +607,7 @@ def _fused_dsa_draft_extend_metadata_kernel(
         real_cols = offs_n // real_page_size
         tl.store(
             real_page_table
-            + out_rows[:, None] * real_page_table_stride_0
+            + out_rows.to(tl.int64)[:, None] * real_page_table_stride_0
             + real_cols[None, :] * real_page_table_stride_1,
             (vals // real_page_size)[None, :],
             mask=real_mask,

python/sglang/kernels/ops/attention/fla/kda.py

@@ -29,7 +29,6 @@ from sglang.kernels.ops.attention.fla.utils import (
     check_shared_mem,
     is_intel,
     is_nvidia,
-    is_tf32_supported,
 )
 
 if is_intel:
@@ -742,7 +741,7 @@ def recompute_w_u_fwd(
         BT=BT,
         STORE_KG=kg is not None,
         IS_VARLEN=cu_seqlens is not None,
-        DOT_PRECISION="tf32" if is_tf32_supported else "ieee",
+        DOT_PRECISION="ieee",
         **(static_config or {}),
     )
     return w, u, kg
@@ -751,8 +750,8 @@ def recompute_w_u_fwd(
 @triton.autotune(
     configs=[
         triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
-        for BK in [64]
+        for BK in [32, 64]
-        for BV in [64]
+        for BV in [64, 128]
         for num_warps in [2, 4, 8]
         for num_stages in [2, 3, 4]
     ],
@@ -863,7 +862,7 @@ def chunk_gla_fwd_kernel_o(
     # [BT, BT]
     b_A = tl.load(p_A, boundary_check=(0, 1))
     b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
-    b_o += tl.dot(b_A, b_v)
+    b_o += tl.dot(b_A, b_v, allow_tf32=False)
     tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
 
 

python/sglang/srt/arg_groups/attention_hook.py

@@ -564,6 +564,7 @@ def handle_deterministic_inference(server_args: Any):
                     "PixtralForConditionalGeneration",
                     "GlmMoeDsaForCausalLM",
                     "Glm4MoeLiteForCausalLM",
+                    "Glm5NextForConditionalGeneration",
                 ]
             except Exception:
                 pass

python/sglang/srt/arg_groups/cuda_graph_hook.py

@@ -275,10 +275,17 @@ def disable_breakable_cudagraph_if_incompatible(server_args: Any):
     """
 
     cfg = resolving_view(server_args)
-    from sglang.srt.configs.model_config import is_deepseek_v4
+    from sglang.srt.configs.model_config import (
+        is_deepseek_v4,
+        uses_kda_attention,
+    )
     from sglang.srt.layers.cp.bcg import supports_prefill_cp_bcg
 
     rules = [
+        (
+            "KDA hybrid linear attention",
+            lambda: uses_kda_attention(model_config_of(server_args).hf_config),
+        ),
         # DSV4 is BCG-compatible but introduces heavy memory pressure: the
         # c4 indexer scratch is pinned in the capture pool and OOMs. Disable.
         (

python/sglang/srt/arg_groups/model_hook.py

@@ -183,6 +183,7 @@ def handle_model_specific_adjustments(server_args: Any):
         "MistralLarge3ForCausalLM",
         "PixtralForConditionalGeneration",
         "GlmMoeDsaForCausalLM",
+        "Glm5NextForConditionalGeneration",
         "HYV4ForCausalLM",
         "HYV4ForCausalLMNextN",
         "LongcatFlashForCausalLM",

python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py

@@ -1,6 +1,6 @@
 """Config-time override declarations for deepseek_v2.
 
-Architectures: DeepseekV32ForCausalLM, DeepseekV3ForCausalLM, Dots3NoteForCausalLM, GlmMoeDsaForCausalLM, HYV4ForCausalLM, HYV4ForCausalLMNextN, KimiK25ForConditionalGeneration, LongcatFlashForCausalLM, LongcatFlashForCausalLMNextN, MistralLarge3ForCausalLM, PixtralForConditionalGeneration.
+Architectures: DeepseekV32ForCausalLM, DeepseekV3ForCausalLM, Dots3NoteForCausalLM, Glm5NextForConditionalGeneration, GlmMoeDsaForCausalLM, HYV4ForCausalLM, HYV4ForCausalLMNextN, KimiK25ForConditionalGeneration, LongcatFlashForCausalLM, LongcatFlashForCausalLMNextN, MistralLarge3ForCausalLM, PixtralForConditionalGeneration.
 """
 
 import logging
@@ -24,6 +24,7 @@ logger = logging.getLogger(__name__)
     "MistralLarge3ForCausalLM",
     "PixtralForConditionalGeneration",
     "GlmMoeDsaForCausalLM",
+    "Glm5NextForConditionalGeneration",
     "HYV4ForCausalLM",
     "HYV4ForCausalLMNextN",
     "LongcatFlashForCausalLM",

python/sglang/srt/arg_groups/overrides.py

@@ -512,6 +512,7 @@ _MAMBA_RADIX_CACHE_ARCHS = frozenset(
         "Lfm2ForCausalLM",
         "Lfm2MoeForCausalLM",
         "ZayaForCausalLM",
+        "Glm5NextForConditionalGeneration",
     }
 )
 
@@ -533,6 +534,7 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset(
         "BailingMoeV3ForCausalLM",
         "FalconH1ForCausalLM",
         "GraniteMoeHybridForCausalLM",
+        "Glm5NextForConditionalGeneration",
         "NemotronHForCausalLM",
         "NemotronHPuzzleForCausalLM",
         # KDA-based: same MambaPool ping-pong machinery as GDN; requires the
@@ -792,6 +794,7 @@ _DEEPSEEK_FAMILY_ARCHS = frozenset(
         "MistralLarge3ForCausalLM",
         "PixtralForConditionalGeneration",
         "GlmMoeDsaForCausalLM",
+        "Glm5NextForConditionalGeneration",
         "HYV4ForCausalLM",
         "HYV4ForCausalLMNextN",
         "LongcatFlashForCausalLM",

python/sglang/srt/arg_groups/pd_disaggregation_hook.py

@@ -256,9 +256,10 @@ def handle_encoder_disaggregation(server_args: Any):
         "KimiK25ForConditionalGeneration",
         "KimiK3ForConditionalGeneration",
         "MiMoV2ForCausalLM",
+        "Glm5NextForConditionalGeneration",
     ]:
         raise ValueError(
             f"Model type {model_arch} is not supported for encoder disaggregation. "
             f"Supported architectures: Qwen2VL, Qwen3VL, Qwen3.5, InternS2, "
-            f"Qwen2Audio, Qwen2.5Omni, Dots3-Note, Kimi, MiMoV2."
+            f"Qwen2Audio, Qwen2.5Omni, Dots3-Note, Kimi, MiMoV2, GLM5Next."
         )

python/sglang/srt/configs/__init__.py

@@ -16,6 +16,7 @@ from sglang.srt.configs.dots_ocr import DotsOCRConfig
 from sglang.srt.configs.dots_vlm import DotsVLMConfig
 from sglang.srt.configs.exaone import ExaoneConfig
 from sglang.srt.configs.falcon_h1 import FalconH1Config
+from sglang.srt.configs.glm5_next import Glm5NextConfig, Glm5NextTextConfig
 from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
 from sglang.srt.configs.hy_v4 import HYV4Config
 from sglang.srt.configs.inkling import (
@@ -102,6 +103,8 @@ __all__ = [
     "Olmo3Config",
     "MuseGlimmerConfig",
     "MuseGlimmerAssistantConfig",
+    "Glm5NextConfig",
+    "Glm5NextTextConfig",
     "KimiLinearConfig",
     "KimiK3Config",
     "KimiK25Config",

python/sglang/srt/configs/glm5_next.py

@@ -0,0 +1,341 @@
+from typing import List, Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.glm_ocr.configuration_glm_ocr import GlmOcrVisionConfig
+
+from sglang.srt.configs.mamba_utils import KimiLinearCacheParams, KimiLinearStateShape
+from sglang.srt.runtime_context import get_parallel
+
+_GLM5_NEXT_TOP_LEVEL_CONFIG_KEYS = (
+    "architectures",
+    "vocab_size",
+    "hidden_size",
+    "head_dim",
+    "intermediate_size",
+    "moe_intermediate_size",
+    "num_hidden_layers",
+    "num_attention_heads",
+    "num_key_value_heads",
+    "hidden_act",
+    "max_position_embeddings",
+    "rms_norm_eps",
+    "use_cache",
+    "pad_token_id",
+    "bos_token_id",
+    "eos_token_id",
+    "rope_theta",
+    "rope_scaling",
+    "rope_parameters",
+    "partial_rotary_factor",
+    "tie_word_embeddings",
+    "attention_bias",
+    "attention_dropout",
+    "n_routed_experts",
+    "num_experts_per_tok",
+    "n_shared_experts",
+    "n_group",
+    "topk_group",
+    "norm_topk_prob",
+    "routed_scaling_factor",
+    "scoring_func",
+    "topk_method",
+    "first_k_dense_replace",
+    "moe_layer_freq",
+    "q_lora_rank",
+    "kv_lora_rank",
+    "qk_nope_head_dim",
+    "qk_rope_head_dim",
+    "v_head_dim",
+    "swiglu_limit",
+    "mhc",
+    "hc_mult",
+    "hc_sinkhorn_iters",
+    "hc_eps",
+    "num_nextn_predict_layers",
+    "linear_attn_config",
+    "linear_head_dim",
+    "linear_num_heads",
+    "linear_conv_kernel_dim",
+    "linear_lower_bound",
+    "gate_lower_bound",
+    "index_head_dim",
+    "index_topk",
+    "index_kpool",
+    "index_kpool_always_select_tail",
+    "index_kpool_compress",
+    "index_n_heads",
+    "index_topk_freq",
+    "index_topk_pattern",
+    "index_skip_topk_offset",
+    "index_share_for_mtp_iteration",
+    "indexer_rope_interleave",
+    "layer_types",
+    "mlp_layer_types",
+    "quantization_config",
+)
+
+
+class Glm5NextTextConfig(PretrainedConfig):
+    model_type = "glm5_next_text"
+    base_config_key = "text_config"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size: int = 154880,
+        hidden_size: int = 4096,
+        head_dim: Optional[int] = None,
+        intermediate_size: int = 12288,
+        moe_intermediate_size: int = 2048,
+        num_hidden_layers: int = 45,
+        num_attention_heads: int = 64,
+        num_key_value_heads: Optional[int] = None,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 1013760,
+        rms_norm_eps: float = 1e-5,
+        use_cache: bool = True,
+        pad_token_id: Optional[int] = None,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        rope_theta: float = 800000.0,
+        rope_scaling: Optional[dict] = None,
+        rope_parameters: Optional[dict] = None,
+        partial_rotary_factor: float = 1.0,
+        tie_word_embeddings: bool = False,
+        attention_bias: bool = False,
+        attention_dropout: float = 0.0,
+        n_routed_experts: Optional[int] = 288,
+        num_experts_per_tok: int = 7,
+        n_shared_experts: int = 1,
+        n_group: int = 1,
+        topk_group: int = 1,
+        norm_topk_prob: bool = True,
+        routed_scaling_factor: float = 2.5,
+        scoring_func: str = "sigmoid",
+        topk_method: str = "noaux_tc",
+        first_k_dense_replace: int = 3,
+        moe_layer_freq: int = 1,
+        q_lora_rank: Optional[int] = 1536,
+        kv_lora_rank: int = 512,
+        qk_nope_head_dim: int = 256,
+        qk_rope_head_dim: int = 0,
+        v_head_dim: int = 256,
+        swiglu_limit: Optional[float] = None,
+        mhc: bool = False,
+        hc_mult: int = 4,
+        hc_sinkhorn_iters: int = 20,
+        hc_eps: float = 1e-6,
+        num_nextn_predict_layers: int = 1,
+        linear_attn_config: Optional[dict] = None,
+        linear_head_dim: int = 128,
+        linear_num_heads: int = 64,
+        linear_conv_kernel_dim: int = 4,
+        linear_lower_bound: Optional[float] = None,
+        gate_lower_bound: Optional[float] = None,
+        index_head_dim: int | None = None,
+        index_topk: int | None = None,
+        index_n_heads: int | None = None,
+        index_topk_freq: int = 1,
+        index_topk_pattern: Optional[str] = None,
+        index_skip_topk_offset: Optional[int] = None,
+        **kwargs,
+    ):
+        if rope_scaling is None and rope_parameters is not None:
+            rope_scaling = rope_parameters
+        if rope_parameters is not None:
+            rope_theta = rope_parameters.get("rope_theta", rope_theta)
+            partial_rotary_factor = rope_parameters.get(
+                "partial_rotary_factor", partial_rotary_factor
+            )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.head_dim = head_dim
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_intermediate_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.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.partial_rotary_factor = partial_rotary_factor
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.n_routed_experts = n_routed_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_shared_experts = n_shared_experts
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.norm_topk_prob = norm_topk_prob
+        self.routed_scaling_factor = routed_scaling_factor
+        self.scoring_func = scoring_func
+        self.topk_method = topk_method
+        self.first_k_dense_replace = first_k_dense_replace
+        self.moe_layer_freq = moe_layer_freq
+
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.swiglu_limit = swiglu_limit
+        self.mhc = mhc
+        self.hc_mult = hc_mult
+        self.hc_sinkhorn_iters = hc_sinkhorn_iters
+        self.hc_eps = hc_eps
+
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+        self.linear_head_dim = linear_head_dim
+        self.linear_num_heads = linear_num_heads
+        self.linear_conv_kernel_dim = linear_conv_kernel_dim
+        self.linear_lower_bound = linear_lower_bound
+        self.gate_lower_bound = (
+            gate_lower_bound if gate_lower_bound is not None else linear_lower_bound
+        )
+        if linear_attn_config is None:
+            layer_types = kwargs.get("layer_types")
+            if layer_types is None:
+                kda_layers = [
+                    layer_idx
+                    for layer_idx in range(num_hidden_layers)
+                    if layer_idx % 4 != 3
+                ]
+            else:
+                kda_layers = [
+                    layer_idx
+                    for layer_idx, layer_type in enumerate(layer_types)
+                    if layer_type == "linear_attention"
+                ]
+            kda_layer_set = set(kda_layers)
+            linear_attn_config = {
+                "full_attn_layers": [
+                    layer_idx
+                    for layer_idx in range(num_hidden_layers)
+                    if layer_idx not in kda_layer_set
+                ],
+                "head_dim": linear_head_dim,
+                "kda_layers": kda_layers,
+                "num_heads": linear_num_heads,
+                "short_conv_kernel_size": linear_conv_kernel_dim,
+                "gate_lower_bound": self.gate_lower_bound,
+            }
+        self.linear_attn_config = linear_attn_config
+        self.index_head_dim = index_head_dim
+        self.index_topk = index_topk
+        self.index_n_heads = index_n_heads
+        self.index_topk_freq = index_topk_freq
+        self.index_topk_pattern = index_topk_pattern
+        self.index_skip_topk_offset = index_skip_topk_offset
+
+        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,
+        )
+        if rope_parameters is not None or rope_scaling is not None:
+            self.rope_parameters = rope_parameters or rope_scaling
+
+    def is_kda_layer(self, layer_idx: int):
+        return (
+            self.linear_attn_config is not None
+            and layer_idx in self.linear_attn_config["kda_layers"]
+        )
+
+    @property
+    def linear_layer_ids(self):
+        return [i for i in range(self.num_hidden_layers) if self.is_kda_layer(i)]
+
+    @property
+    def nextn_layer_ids(self):
+        num_nextn_layers = self.num_nextn_predict_layers or 0
+        return [self.num_hidden_layers + i for i in range(num_nextn_layers)]
+
+    @property
+    def full_attention_layer_ids(self):
+        return [i for i in range(self.num_hidden_layers) if not self.is_kda_layer(i)]
+
+    @property
+    def mamba2_cache_params(self) -> KimiLinearCacheParams:
+        shape = KimiLinearStateShape.create(
+            tp_world_size=get_parallel().attn_tp_size,
+            num_heads=self.linear_attn_config["num_heads"],
+            head_dim=self.linear_attn_config["head_dim"],
+            conv_kernel_size=self.linear_attn_config["short_conv_kernel_size"],
+        )
+
+        return KimiLinearCacheParams(shape=shape, layers=self.linear_layer_ids)
+
+
+class Glm5NextVisionConfig(GlmOcrVisionConfig):
+    def __init__(
+        self,
+        swiglu_limit: float,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.swiglu_limit = swiglu_limit
+
+
+class Glm5NextConfig(PretrainedConfig):
+    model_type = "glm5_next"
+    sub_configs = {
+        "vision_config": Glm5NextVisionConfig,
+        "text_config": Glm5NextTextConfig,
+    }
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        image_token_id: int = 59280,
+        video_token_id: int = 59281,
+        image_start_token_id: int = 59256,
+        image_end_token_id: int = 59257,
+        video_start_token_id: int = 59258,
+        video_end_token_id: int = 59259,
+        **kwargs,
+    ):
+        top_level_text_config = {
+            key: kwargs[key]
+            for key in _GLM5_NEXT_TOP_LEVEL_CONFIG_KEYS
+            if key in kwargs
+        }
+
+        if isinstance(text_config, dict):
+            text_config = {**top_level_text_config, **text_config}
+            self.text_config = self.sub_configs["text_config"](**text_config)
+        elif text_config is None:
+            self.text_config = self.sub_configs["text_config"](**top_level_text_config)
+        else:
+            self.text_config = text_config
+
+        if vision_config is None:
+            self.vision_config = None
+        else:
+            if isinstance(vision_config, dict):
+                vision_config = dict(vision_config)
+            else:
+                vision_config = vision_config.to_dict()
+            self.vision_config = self.sub_configs["vision_config"](**vision_config)
+
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+        self.image_start_token_id = image_start_token_id
+        self.image_end_token_id = image_end_token_id
+        self.video_start_token_id = video_start_token_id
+        self.video_end_token_id = video_end_token_id
+
+        if getattr(self.text_config, "quantization_config", None) is not None:
+            self.quantization_config = self.text_config.quantization_config
+
+        super().__init__(**kwargs)
+        for key in _GLM5_NEXT_TOP_LEVEL_CONFIG_KEYS:
+            if hasattr(self.text_config, key):
+                setattr(self, key, getattr(self.text_config, key))

python/sglang/srt/configs/hybrid_arch.py

@@ -113,6 +113,16 @@ def kimi_linear_config(model_config: ModelConfig):
     return None
 
 
+def glm5_next_config(model_config: ModelConfig):
+    hf_config = model_config.hf_config
+    if (
+        getattr(hf_config, "model_type", None) == "glm5_next"
+        and not model_config.is_draft_model
+    ):
+        return hf_config.get_text_config()
+    return None
+
+
 def linear_attn_model_spec(model_config: ModelConfig):
     result = _get_linear_attn_registry_result(model_config)
     return result[0] if result else None
@@ -123,6 +133,7 @@ def mambaish_config(model_config: ModelConfig):
         mamba2_config(model_config)
         or hybrid_gdn_config(model_config)
         or kimi_linear_config(model_config)
+        or glm5_next_config(model_config)
         or hybrid_lightning_config(model_config)
     )
     if existing:

python/sglang/srt/configs/model_config.py

@@ -138,6 +138,8 @@ def is_deepseek_dsa(config) -> bool:
             "PixtralForConditionalGeneration",
             "GlmMoeDsaForCausalLM",
             "GlmMoeDsaForCausalLMNextN",
+            "Glm5NextForConditionalGenerationNextN",
+            "Glm5NextForConditionalGeneration",
             "LongcatFlashForCausalLM",
             "LongcatFlashForCausalLMNextN",
             "Dots3NoteForCausalLM",
@@ -156,6 +158,31 @@ def is_kimi_k3(config) -> bool:
     )
 
 
+def uses_kda_attention(config) -> bool:
+    configs = [config]
+    get_text_config = getattr(config, "get_text_config", None)
+    if callable(get_text_config):
+        configs.append(get_text_config())
+    else:
+        text_config = _hf_attr(config, "text_config")
+        if text_config is not None:
+            configs.append(text_config)
+    for config in configs:
+        linear_attn_config = _hf_attr(config, "linear_attn_config")
+        if isinstance(linear_attn_config, dict) and linear_attn_config.get(
+            "kda_layers"
+        ):
+            return True
+        layer_types = _hf_attr(config, "layer_types") or []
+        if (
+            "linear_attention" in layer_types
+            and _hf_attr(config, "linear_num_heads") is not None
+            and _hf_attr(config, "linear_head_dim") is not None
+        ):
+            return True
+    return False
+
+
 def is_dspark_draft(config) -> bool:
     return _hf_arch(config) == "DSparkDraftModel"
 
@@ -285,6 +312,21 @@ def get_dsa_index_n_heads(config: PretrainedConfig) -> int:
     return config.index_n_heads
 
 
+def get_dsa_index_kpool(config: PretrainedConfig) -> int:
+    return getattr(config, "index_kpool", 1)
+
+
+def get_dsa_mtp_topk_width(config: PretrainedConfig) -> int:
+    """MTP seeds include index_topk pooled tokens plus up to index_kpool - 1 tail tokens."""
+    index_kpool = get_dsa_index_kpool(config)
+    assert index_kpool >= 1, f"index_kpool must be positive, got {index_kpool}"
+    return config.index_topk + index_kpool - 1
+
+
+def get_dsa_index_kpool_compress(config: PretrainedConfig) -> bool:
+    return getattr(config, "index_kpool_compress", False)
+
+
 REQUANTIZATION_METHODS = ["quark_mxfp4"]
 
 
@@ -731,6 +773,15 @@ class ModelConfig:
         ):
             self.hf_config.architectures[0] = "Glm4MoeLiteForCausalLMNextN"
 
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "Glm5NextForConditionalGeneration"
+        ):
+            self.hf_config.architectures[0] = "Glm5NextForConditionalGenerationNextN"
+            self.hf_text_config.architectures = list(self.hf_config.architectures)
+            self.hf_text_config.num_nextn_predict_layers = 1
+            self.hf_text_config.linear_attn_config = None
+
         if is_draft_model and self.hf_config.architectures[0] in [
             "GlmOcrForConditionalGeneration",
         ]:
@@ -968,6 +1019,8 @@ class ModelConfig:
             or "Glm4MoeLiteForCausalLMNextN" in self.hf_config.architectures
             or "GlmMoeDsaForCausalLM" in self.hf_config.architectures
             or "GlmMoeDsaForCausalLMNextN" in self.hf_config.architectures
+            or "Glm5NextForConditionalGeneration" in self.hf_config.architectures
+            or "Glm5NextForConditionalGenerationNextN" in self.hf_config.architectures
             or "LongcatFlashForCausalLM" in self.hf_config.architectures
             or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures
             or "HYV4ForCausalLM" in self.hf_config.architectures
@@ -1085,7 +1138,10 @@ class ModelConfig:
             self.v_head_dim = self.hf_config.v_head_dim
             self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim
             self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
-        elif "SarvamMLAForCausalLM" in self.hf_config.architectures:
+        elif (
+            "SarvamMLAForCausalLM" in self.hf_config.architectures
+            or "Glm5NextForConditionalGeneration" in self.hf_config.architectures
+        ):
             self.head_dim = (
                 self.hf_config.qk_nope_head_dim + self.hf_config.qk_rope_head_dim
             )
@@ -1138,6 +1194,16 @@ class ModelConfig:
             self.num_key_value_heads = self.num_attention_heads
         self.hidden_size = self.hf_text_config.hidden_size
         hc_mult = getattr(self.hf_text_config, "hc_mult", 1)
+        is_glm5_next = getattr(self.hf_config, "model_type", None) == "glm5_next" or (
+            getattr(self.hf_text_config, "model_type", None) == "glm5_next_text"
+        )
+        if is_glm5_next and not getattr(self.hf_text_config, "mhc", False):
+            hc_mult = 1
+        if is_glm5_next:
+            # mHC-flattened hidden size; None when not running an mHC model.
+            self.hc_hidden_size = self.hidden_size * hc_mult if hc_mult > 1 else None
+            self.spec_hidden_size = self.hidden_size
+        else:
             self.spec_hidden_size, self.hc_hidden_size = resolve_spec_hidden_size(
                 self.hf_config, self.hidden_size, hc_mult
             )
@@ -1950,6 +2016,7 @@ multimodal_model_archs = [
     "Gemma4UnifiedForConditionalGeneration",
     "Glm4vForConditionalGeneration",
     "Glm4vMoeForConditionalGeneration",
+    "Glm5NextForConditionalGeneration",
     "GlmOcrForConditionalGeneration",
     "GlmAsrForConditionalGeneration",
     "GlmImageForConditionalGeneration",
@@ -2017,6 +2084,7 @@ piecewise_cuda_graph_disabled_model_archs = [
     "DeepseekV4ForCausalLMNextN",
     "DeepseekV4ForCausalLMDSpark",
     "Qwen3NextForCausalLM",
+    "Glm5NextForConditionalGeneration",
     "BailingMoeV2_5ForCausalLM",
     "LLaDAModelLM",
 ]

python/sglang/srt/configs/qwen3_asr.py

@@ -164,5 +164,5 @@ class Qwen3ASRConfig(PretrainedConfig):
         return self.thinker_config.text_config
 
 
-AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
+AutoConfig.register("qwen3_asr", Qwen3ASRConfig, exist_ok=True)
-AutoConfig.register("qwen3_asr_thinker", Qwen3ASRThinkerConfig)
+AutoConfig.register("qwen3_asr_thinker", Qwen3ASRThinkerConfig, exist_ok=True)

python/sglang/srt/disaggregation/base/conn.py

@@ -18,6 +18,9 @@ class StateType(str, enum.Enum):
     MAMBA = "mamba"
     SWA = "swa"
     DSA = "dsa"
+    # DSA kpool-compress tail: one per-request ring row. The indices encode
+    # only the live subrange of that row for the current open pool.
+    DSA_TAIL = "dsa_tail"
     MINIMAX_INDEX_K = "minimax_index_k"
     # DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot
     # (req_pool_idx * ring_stride + pos % ring_stride), needs its own component.

python/sglang/srt/disaggregation/common/conn.py

@@ -1009,7 +1009,7 @@ class CommonKVManager(BaseKVManager):
         returned unchanged.
         """
         start_layer = self.kv_args.prefill_start_layer
-        end_layer = getattr(self.kv_args, "prefill_end_layer", None)
+        end_layer = self.kv_args.prefill_end_layer
         assert end_layer is not None, (
             "KVArgs.prefill_end_layer must be set when using compressed-MLA PD with PP"
         )

python/sglang/srt/disaggregation/decode.py

@@ -54,6 +54,7 @@ from sglang.srt.disaggregation.utils import (
     _is_fake_transfer,
     build_kv_layer_ids,
     build_staging_slot_metadata,
+    get_dsa_tail_state_indices,
     get_dsv4_c128_state_indices,
     get_kv_class,
     is_dsv4_c128_online_enabled,
@@ -1411,6 +1412,13 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 device_page_size = self.token_to_kv_pool.page_size
                 return kv_to_page_indices(kv_indices_full, device_page_size)
 
+            def _dsa_tail_payload():
+                return get_dsa_tail_state_indices(
+                    self.token_to_kv_pool,
+                    decode_req.req.kv.req_pool_idx,
+                    seq_len,
+                )
+
             def _swa_ring_payload():
                 # Mirror of prefill _swa_ring_payload using this side's req_pool_idx.
                 # Same window positions and order -> positional match with prefill.
@@ -1443,6 +1451,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 StateType.MAMBA: _mamba_payload,
                 StateType.SWA: _swa_payload,
                 StateType.DSA: _full_kv_pages_payload,
+                StateType.DSA_TAIL: _dsa_tail_payload,
                 StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
                 StateType.SWA_RING: _swa_ring_payload,
                 StateType.C128_STATE: _c128_state_payload,

python/sglang/srt/disaggregation/encoder/preprocessor.py

@@ -11,6 +11,7 @@ import concurrent.futures
 import functools
 import logging
 from dataclasses import dataclass
+from http import HTTPStatus
 from typing import Callable, List, Optional, Tuple, Union
 
 import numpy as np
@@ -26,6 +27,17 @@ from sglang.srt.multimodal.encoder_preprocessing import (
     EncoderPreprocessOutput,
     invoke_encoder_preprocessor,
 )
+from sglang.srt.multimodal.processors.glm4v import (
+    _glm_effective_presize_budget,
+    glm_budget_kwargs,
+    glm_decode_frames_at,
+    glm_max_image_tokens_from_configs,
+    glm_processor_video_config,
+    glm_sample_and_decode_sync,
+    glm_sample_frame_indices,
+    preprocess_video_frames_sync,
+    split_glm_video_items,
+)
 from sglang.srt.multimodal.processors.qwen_vl import preprocess_video
 from sglang.srt.runtime_context import (
     get_device,
@@ -220,6 +232,9 @@ class EncoderPreprocessor:
                 self.vision_config[modality_str]["device"] = self.device
 
             if modality_str == "video":
+                # GLM reads its own defaults from the HF video processor
+                # (max_frames=2048); applying the Qwen values here would clobber them.
+                if "glm" not in self.model_type:
                     video_defaults = {"fps": 2.0, "max_frames": 768, "min_frames": 4}
                     for k, v in video_defaults.items():
                         self.vision_config["video"].setdefault(k, v)
@@ -334,7 +349,17 @@ class EncoderPreprocessor:
                     }
                 return img
             elif modality == Modality.VIDEO:
-                return load_video(data, frame_count_limit)
+                vid = load_video(data, frame_count_limit)
+                if (
+                    media_metadata
+                    and self.encoder_media_processor_config.preserve_media_metadata
+                ):
+                    return {
+                        "type": "video",
+                        "video": vid,
+                        **media_metadata,
+                    }
+                return vid
             elif modality == Modality.AUDIO:
                 return load_audio(data, self.model_audio_sr)
 
@@ -399,10 +424,109 @@ class EncoderPreprocessor:
     async def _flatten_and_load_images(self, mm_items):
         return await self._flatten_and_load_data_by_modality(mm_items, Modality.IMAGE)
 
+    @staticmethod
+    def _close_video_decoders(video_items) -> None:
+        for video in video_items or []:
+            close = getattr(video, "close", None)
+            if callable(close):
+                close()
+
+    async def _dp_sharded_decode_single_video(
+        self,
+        vr,
+        video_config,
+        *,
+        tp_rank: int,
+        tp_size: int,
+        video_processor_kwargs: dict,
+        precomputed_indices: Optional[List[int]] = None,
+    ):
+        video_config = video_config or {}
+        video_fps = vr.avg_fps
+        duration = len(vr) / video_fps if video_fps else 0
+        global_indices = precomputed_indices or glm_sample_frame_indices(
+            len(vr),
+            video_fps,
+            duration,
+            target_fps=video_config.get("fps"),
+            max_frame_count=video_config.get("max_frames"),
+        )
+        n_units = len(global_indices) // 2
+        base, remainder = divmod(n_units, tp_size)
+        gpu_sample_counts = [
+            base + (1 if rank < remainder else 0) for rank in range(tp_size)
+        ]
+        start = sum(gpu_sample_counts[:tp_rank])
+        count = gpu_sample_counts[tp_rank]
+        local_indices = global_indices[2 * start : 2 * (start + count)]
+
+        local_error = None
+        frames = None
+        try:
+            frames = await asyncio.get_running_loop().run_in_executor(
+                self.io_executor,
+                glm_decode_frames_at,
+                vr,
+                local_indices,
+                video_config,
+            )
+        except Exception as exc:
+            local_error = exc
+
+        # All ranks must either enter the later ViT all-gather or fail before
+        # it.  A rank-local decoder error must therefore be agreed globally.
+        ok = torch.tensor([0 if local_error else 1], dtype=torch.int32)
+        if tp_size > 1:
+            torch.distributed.all_reduce(
+                ok,
+                op=torch.distributed.ReduceOp.MIN,
+                group=get_parallel().attn_tp_group.cpu_group,
+            )
+        if not int(ok.item()):
+            if local_error is not None:
+                raise local_error
+            from sglang.srt.disaggregation.encoder.server import MMError
+
+            raise MMError(
+                "peer encoder rank failed during sharded video decode",
+                code=HTTPStatus.SERVICE_UNAVAILABLE,
+            )
+        if frames is None:
+            height, width = vr.frame_shape
+            frames = np.zeros((0, height, width, 3), dtype=np.uint8)
+
+        video_processor_kwargs["do_sample_frames"] = False
+        video_processor_kwargs["return_metadata"] = True
+        # Preserve the same per-frame spatial budget as the unsharded request.
+        if global_indices and local_indices:
+            budget = video_config.get("max_image_tokens")
+            if budget is None:
+                budget = getattr(self.video_processor, "max_image_tokens", None)
+            if budget is not None:
+                video_processor_kwargs["max_image_tokens"] = max(
+                    1, int(int(budget) * len(local_indices) / len(global_indices))
+                )
+        video_processor_kwargs["_dp_meta"] = {
+            "global_indices": list(global_indices),
+            "fps": video_fps,
+            "n_units": n_units,
+            "gpu_sample_counts": gpu_sample_counts,
+        }
+        return [frames], video_processor_kwargs
+
     async def _flatten_and_load_videos(self, mm_items):
         if not isinstance(mm_items, (list, tuple)):
             mm_items = [mm_items]
 
+        video_configs = [{} for _ in mm_items]
+        if "glm" in self.model_type:
+            mm_items, video_configs = split_glm_video_items(mm_items)
+            defaults = glm_processor_video_config(self.video_processor)
+            defaults.update(self.vision_config.get("video", {}))
+            video_configs = [
+                {**defaults, **dict(config or {})} for config in video_configs
+            ]
+
         futures, _ = self._submit_data_loading_tasks(
             mm_items, [Modality.VIDEO] * len(mm_items)
         )
@@ -422,7 +546,84 @@ class EncoderPreprocessor:
             if video_metadata:
                 video_processor_kwargs["video_metadata"] = video_metadata
             return videos, video_processor_kwargs
+
+        if "glm" in self.model_type:
+            budget_kwargs = glm_budget_kwargs(
+                self.video_processor,
+                user_max_image_tokens=glm_max_image_tokens_from_configs(video_configs),
+                count=len(video_items),
+                split=True,
+            )
+            if budget_kwargs is not None:
+                video_processor_kwargs.update(budget_kwargs)
+                video_configs = [
+                    _glm_effective_presize_budget(
+                        config, budget_kwargs.get("max_image_tokens")
+                    )
+                    for config in video_configs
+                ]
+
+            framed = any(isinstance(video, list) for video in video_items)
+            if framed:
+                processed = await asyncio.gather(
+                    *[
+                        asyncio.get_running_loop().run_in_executor(
+                            self.io_executor, preprocess_video_frames_sync, video
+                        )
+                        for video in video_items
+                    ]
+                )
             else:
+                parallel = get_parallel()
+                tp_size = parallel.attn_tp_size
+                sampled = None
+                if len(video_items) == 1:
+                    vr = video_items[0]
+                    config = video_configs[0]
+                    sampled = glm_sample_frame_indices(
+                        len(vr),
+                        vr.avg_fps,
+                        len(vr) / vr.avg_fps if vr.avg_fps else 0,
+                        target_fps=config.get("fps"),
+                        max_frame_count=config.get("max_frames"),
+                    )
+                if (
+                    self.server_args.mm_enable_dp_encoder
+                    and tp_size > 1
+                    and sampled is not None
+                    and len(sampled) >= max(32, tp_size * 2)
+                ):
+                    result = await self._dp_sharded_decode_single_video(
+                        video_items[0],
+                        video_configs[0],
+                        tp_rank=parallel.attn_tp_rank,
+                        tp_size=tp_size,
+                        video_processor_kwargs=video_processor_kwargs,
+                        precomputed_indices=sampled,
+                    )
+                    self._close_video_decoders(video_items)
+                    return result
+
+                processed = await asyncio.gather(
+                    *[
+                        asyncio.get_running_loop().run_in_executor(
+                            self.io_executor,
+                            glm_sample_and_decode_sync,
+                            video,
+                            video_configs[index],
+                        )
+                        for index, video in enumerate(video_items)
+                    ]
+                )
+            videos, video_metadata = map(list, zip(*processed))
+            video_processor_kwargs["do_sample_frames"] = False
+            video_processor_kwargs["return_metadata"] = True
+            if video_metadata:
+                video_processor_kwargs["video_metadata"] = video_metadata
+            self._close_video_decoders(video_items)
+            return videos, video_processor_kwargs
+
+        self._close_video_decoders(video_items)
         raise NotImplementedError(
             f"Video processing is not supported for {self.model_type} model."
         )

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -47,9 +47,11 @@ from sglang.srt.disaggregation.mooncake.utils import (
 )
 from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
+    build_dsa_tail_transfer_blocks,
     build_transfer_entry_pairs,
     compute_mamba_state_slice_byte_blocks,
     resolve_dcp_dst_entry_indices,
+    slice_dsa_tail_dst_ptrs_for_pp,
 )
 from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
 from sglang.srt.environ import envs
@@ -1406,10 +1408,25 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                         )
                         or rc
                     )
+            elif st == StateType.DSA_TAIL:
+                rc = (
+                    self._send_slot_state(
+                        req,
+                        src_data_ptrs,
+                        src_item_lens,
+                        dst_data_ptrs,
+                        dst_item_lens,
+                        list(indices),
+                        list(dst_indices),
+                        st.value,
+                    )
+                    or rc
+                )
             elif self._is_generic_kvcache_state_type(st):
                 if (
                     target_rank_registration_info is not None
                     and not self.is_mla_backend
+                    and not self.is_hybrid_mla_backend
                     and self.attn_tp_size
                     != target_rank_registration_info.dst_attn_tp_size
                 ):
@@ -1491,6 +1508,43 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                 )
         return rc
 
+    def _send_slot_state(
+        self,
+        req: TransferInfo,
+        src_ptrs: list[int],
+        src_item_lens: list[int],
+        dst_ptrs: list[int],
+        dst_item_lens: list[int],
+        src_indices: list[int],
+        dst_indices: list[int],
+        label: str,
+    ) -> int:
+        try:
+            dst_ptrs = slice_dsa_tail_dst_ptrs_for_pp(
+                src_ptrs,
+                dst_ptrs,
+                self.kv_args.prefill_start_layer,
+                self.kv_args.prefill_end_layer,
+            )
+            dst_item_lens = slice_dsa_tail_dst_ptrs_for_pp(
+                src_ptrs,
+                dst_item_lens,
+                self.kv_args.prefill_start_layer,
+                self.kv_args.prefill_end_layer,
+            )
+            transfer_blocks = build_dsa_tail_transfer_blocks(
+                src_ptrs,
+                src_item_lens,
+                dst_ptrs,
+                src_indices,
+                dst_indices,
+                dst_item_lens,
+            )
+        except ValueError as exc:
+            logger.error("%s: %s", label, exc)
+            return -1
+        return self._transfer_data(req.mooncake_session_id, transfer_blocks)
+
     def _send_mamba_state(
         self,
         req: TransferInfo,
@@ -1549,7 +1603,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
         attn_tp_size, we slice the state accordingly. GDN conv_state is the
         concatenation [query | key | value] with each sub-block head-sharded
         independently, so on the scatter path it is sliced per sub-block via
-        ``src_state_conv_shard_groups`` (see compute_mamba_state_slice_blocks).
+        ``src_state_conv_shard_groups`` (see
+        compute_mamba_state_slice_byte_blocks).
         """
         logger.warning_once(
             "Using Mamba state slice transfer for different TP sizes between prefill and decode. "

python/sglang/srt/disaggregation/nixl/conn.py

@@ -41,9 +41,11 @@ from sglang.srt.disaggregation.common.utils import (
 )
 from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
+    build_dsa_tail_transfer_blocks,
     build_transfer_entry_pairs,
     compute_mamba_state_slice_byte_blocks,
     resolve_dcp_dst_entry_indices,
+    slice_dsa_tail_dst_ptrs_for_pp,
 )
 from sglang.srt.environ import envs
 from sglang.srt.runtime_context import get_parallel, get_schedule
@@ -2086,6 +2088,60 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
             raise Exception("KVSender failed to post transfer")
         return xfer_handle
 
+    def _send_slot_state(
+        self,
+        peer_name: str,
+        src_data_ptrs: list[int],
+        src_item_lens: list[int],
+        dst_data_ptrs: list[int],
+        dst_item_lens: list[int],
+        src_indices: list[int],
+        dst_indices: list[int],
+        dst_gpu_id: int,
+        notif: str,
+    ):
+        dst_data_ptrs = slice_dsa_tail_dst_ptrs_for_pp(
+            src_data_ptrs,
+            dst_data_ptrs,
+            self.kv_args.prefill_start_layer,
+            self.kv_args.prefill_end_layer,
+        )
+        dst_item_lens = slice_dsa_tail_dst_ptrs_for_pp(
+            src_data_ptrs,
+            dst_item_lens,
+            self.kv_args.prefill_start_layer,
+            self.kv_args.prefill_end_layer,
+        )
+        transfer_blocks = build_dsa_tail_transfer_blocks(
+            src_data_ptrs,
+            src_item_lens,
+            dst_data_ptrs,
+            src_indices,
+            dst_indices,
+            dst_item_lens,
+        )
+        if not transfer_blocks:
+            return None
+
+        src_addrs = [
+            (src_addr, length, self.kv_args.gpu_id)
+            for src_addr, _, length in transfer_blocks
+        ]
+        dst_addrs = [
+            (dst_addr, length, dst_gpu_id) for _, dst_addr, length in transfer_blocks
+        ]
+        src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM")
+        dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM")
+        xfer_handle = self.agent.initialize_xfer(
+            "WRITE", src_descs, dst_descs, peer_name, notif.encode("ascii")
+        )
+        if not xfer_handle:
+            raise Exception("KVSender failed to create dsa_tail transfer")
+        state = self.agent.transfer(xfer_handle)
+        if state == "ERR":
+            raise Exception("KVSender failed to post dsa_tail transfer")
+        return xfer_handle
+
     def _send_mamba_state(
         self,
         peer_name: str,
@@ -2169,7 +2225,8 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
         accordingly, mirroring Mooncake's _send_mamba_state_slice. GDN
         conv_state is [query | key | value] with each sub-block head-sharded
         independently, so on the scatter path it is sliced per sub-block via
-        ``src_state_conv_shard_groups`` (see compute_mamba_state_slice_blocks).
+        ``src_state_conv_shard_groups`` (see
+        compute_mamba_state_slice_byte_blocks).
         """
         logger.warning_once(
             "Using Mamba state slice transfer for different TP sizes. "
@@ -2308,7 +2365,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
             src_indices = (
                 prefill_state_indices[i] if i < len(prefill_state_indices) else None
             )
-            if src_indices is None or len(src_indices) == 0:
+            if src_indices is None or (
+                len(src_indices) == 0 and st != StateType.DSA_TAIL
+            ):
                 continue
             src_ptrs = src_state_data_ptrs[i] if i < len(src_state_data_ptrs) else []
             src_lens = src_state_item_lens[i] if i < len(src_state_item_lens) else []
@@ -2369,12 +2428,37 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
                         src_layer_ids=src_lids,
                         dst_layer_ids=dst_lids,
                     )
-            elif st in (
+            elif st == StateType.DSA_TAIL:
-                StateType.SWA,
+                h = self._send_slot_state(
-                StateType.DSA,
+                    peer_name,
-                StateType.SWA_RING,
+                    src_ptrs,
-                StateType.C128_STATE,
+                    src_lens,
-            ):
+                    dst_ptrs,
+                    dst_lens,
+                    list(src_indices),
+                    list(dst_indices),
+                    dst_gpu_id,
+                    comp_notif,
+                )
+            elif st == StateType.DSA:
+                if len(src_indices) != len(dst_indices):
+                    raise RuntimeError(
+                        f"State index length mismatch at component {i}: "
+                        f"prefill={len(src_indices)}, dst={len(dst_indices)}"
+                    )
+                h = self._send_kvcache_generic(
+                    peer_name=peer_name,
+                    src_data_ptrs=src_ptrs,
+                    dst_data_ptrs=dst_ptrs,
+                    item_lens=src_lens,
+                    prefill_data_indices=np.array(src_indices, dtype=np.int32),
+                    dst_data_indices=np.array(dst_indices, dtype=np.int32),
+                    dst_gpu_id=dst_gpu_id,
+                    notif=comp_notif,
+                    state_type=st,
+                    force_flat=True,
+                )
+            elif st in (StateType.SWA, StateType.SWA_RING, StateType.C128_STATE):
                 if not self.is_mla_backend and self.attn_tp_size != decode_tp_size:
                     raise RuntimeError(
                         f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet."

python/sglang/srt/disaggregation/prefill.py

@@ -45,6 +45,7 @@ from sglang.srt.disaggregation.utils import (
     TransferBackend,
     build_kv_layer_ids,
     build_staging_slot_metadata,
+    get_dsa_tail_state_indices,
     get_dsv4_c128_state_indices,
     get_kv_class,
     is_aborted,
@@ -906,6 +907,7 @@ class SchedulerDisaggregationPrefillMixin:
             can_run_cuda_graph=can_run_cuda_graph,
             dp_cooperation_info=batch.dp_cooperation_info,
         )
+        self.maybe_send_health_check_signal()
 
     @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE)
     def process_disagg_prefill_inflight_queue(
@@ -1306,6 +1308,13 @@ class SchedulerDisaggregationPrefillMixin:
                 ]
                 return kv_to_page_indices(kv_indices_full, page_size)
 
+            def _dsa_tail_payload():
+                return get_dsa_tail_state_indices(
+                    self.token_to_kv_pool_allocator.get_kvcache(),
+                    req.kv.req_pool_idx,
+                    seq_len,
+                )
+
             def _swa_ring_payload():
                 # Unified_kv SWA ring rows (req_pool_idx*ring_stride + pos%ring_stride)
                 # for the last `window` positions, in ascending position order so
@@ -1342,6 +1351,7 @@ class SchedulerDisaggregationPrefillMixin:
                 StateType.MAMBA: _mamba_payload,
                 StateType.SWA: _swa_payload,
                 StateType.DSA: _full_kv_pages_payload,
+                StateType.DSA_TAIL: _dsa_tail_payload,
                 StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
                 StateType.SWA_RING: _swa_ring_payload,
                 StateType.C128_STATE: _c128_state_payload,

python/sglang/srt/disaggregation/utils.py

@@ -19,7 +19,7 @@ import numpy as np
 import torch
 import torch.distributed as dist
 
-from sglang.srt.configs.model_config import get_dsa_index_topk
+from sglang.srt.configs.model_config import get_dsa_mtp_topk_width
 from sglang.srt.disaggregation.base import KVPoll
 from sglang.srt.environ import envs
 from sglang.srt.runtime_context import (
@@ -71,7 +71,7 @@ def get_dsa_seed_metadata_dim(hf_config) -> int:
     """Return the model-defined PD seed width, independent of local spec mode."""
     if not getattr(hf_config, "index_share_for_mtp_iteration", False):
         return 0
-    return get_dsa_index_topk(hf_config)
+    return get_dsa_mtp_topk_width(hf_config)
 
 
 def is_dsv4_c128_online_enabled() -> bool:
@@ -1067,6 +1067,191 @@ def append_state_component(
     kv_args.state_layer_ids.append(layer_ids or [])
 
 
+def get_dsa_tail_state_indices(pool, req_pool_idx: int, seq_len: int) -> List[int]:
+    if getattr(pool, "use_dsa", False):
+        pool = pool.full_kv_pool
+    if not pool.kpool_use_compress:
+        return []
+
+    pool_size = int(pool.index_kpool)
+    tail_size = pool_size + int(getattr(pool, "tail_extra_slots", 0))
+    if pool_size <= 1 or tail_size < pool_size:
+        raise ValueError(
+            "DSA kpool-compress requires pool_size > 1 and "
+            f"tail_size >= pool_size, got pool_size={pool_size}, "
+            f"tail_size={tail_size}"
+        )
+
+    n_valid = int(seq_len) % pool_size
+    if n_valid == 0:
+        return []
+    start_phys = (int(seq_len) - n_valid) % tail_size
+    first_n = min(n_valid, tail_size - start_phys)
+    second_n = n_valid - first_n
+    return [
+        int(req_pool_idx),
+        start_phys,
+        first_n,
+        0,
+        second_n,
+        tail_size,
+    ]
+
+
+def slice_dsa_tail_dst_ptrs_for_pp(
+    src_ptrs: List[int],
+    dst_ptrs: List[int],
+    start_layer: int,
+    end_layer: Optional[int],
+) -> List[int]:
+    if len(src_ptrs) == len(dst_ptrs):
+        return list(dst_ptrs)
+    if len(src_ptrs) % 2 != 0 or len(dst_ptrs) % 2 != 0:
+        raise ValueError(
+            "DSA tail pointer lists must contain equal key/score halves, got "
+            f"src={len(src_ptrs)}, dst={len(dst_ptrs)}"
+        )
+
+    src_layers = len(src_ptrs) // 2
+    dst_layers = len(dst_ptrs) // 2
+    expected_end = start_layer + src_layers
+    if end_layer is not None and end_layer - start_layer == src_layers:
+        expected_end = end_layer
+    if start_layer < 0 or expected_end > dst_layers:
+        raise ValueError(
+            "DSA tail pointer count mismatch: "
+            f"src={len(src_ptrs)}, dst={len(dst_ptrs)}, "
+            f"prefill_layers=[{start_layer}, {expected_end})"
+        )
+
+    return list(dst_ptrs[start_layer:expected_end]) + list(
+        dst_ptrs[dst_layers + start_layer : dst_layers + expected_end]
+    )
+
+
+def build_dsa_tail_transfer_blocks(
+    src_ptrs: List[int],
+    src_item_lens: List[int],
+    dst_ptrs: List[int],
+    src_indices: List[int],
+    dst_indices: List[int],
+    dst_item_lens: Optional[List[int]] = None,
+) -> List[Tuple[int, int, int]]:
+    """Remap live DSA tail tokens between rings with different speculative-slot counts."""
+    if not src_indices and not dst_indices:
+        return []
+    if not src_indices or not dst_indices:
+        raise ValueError(
+            f"DSA tail slot index missing: src={src_indices}, dst={dst_indices}"
+        )
+    if len(src_indices) != 6 or len(dst_indices) != 6:
+        raise ValueError(
+            "DSA tail slot indices must be 6-tuples, "
+            f"got src={src_indices}, dst={dst_indices}"
+        )
+    if dst_item_lens is None:
+        dst_item_lens = src_item_lens
+    if not (len(src_ptrs) == len(dst_ptrs) == len(src_item_lens) == len(dst_item_lens)):
+        raise ValueError(
+            "DSA tail pointer metadata mismatch: "
+            f"src_ptrs={len(src_ptrs)}, dst_ptrs={len(dst_ptrs)}, "
+            f"src_item_lens={len(src_item_lens)}, "
+            f"dst_item_lens={len(dst_item_lens)}"
+        )
+
+    src_tail_size = int(src_indices[5])
+    dst_tail_size = int(dst_indices[5])
+    if src_tail_size <= 0 or dst_tail_size <= 0:
+        raise ValueError(
+            "DSA tail ring sizes must be positive: "
+            f"src={src_tail_size}, dst={dst_tail_size}"
+        )
+
+    def parse_segments(indices: List[int], tail_size: int, side: str):
+        segments = []
+        for seg in (1, 2):
+            off = int(indices[seg * 2 - 1])
+            n = int(indices[seg * 2])
+            if min(off, n) < 0:
+                raise ValueError(
+                    f"DSA tail {side} offsets and lengths must be non-negative"
+                )
+            if off + n > tail_size:
+                raise ValueError(
+                    f"DSA tail {side} segment {seg} exceeds ring size "
+                    f"{tail_size}: ({off}, {n})"
+                )
+            if n:
+                segments.append((off, n))
+        return segments
+
+    src_segments = parse_segments(src_indices, src_tail_size, "source")
+    dst_segments = parse_segments(dst_indices, dst_tail_size, "destination")
+    src_count = sum(n for _, n in src_segments)
+    dst_count = sum(n for _, n in dst_segments)
+    if src_count != dst_count:
+        raise ValueError(
+            f"DSA tail live-token count mismatch: src={src_count}, dst={dst_count}"
+        )
+
+    src_idx = int(src_indices[0])
+    dst_idx = int(dst_indices[0])
+    if src_idx < 0 or dst_idx < 0:
+        raise ValueError("DSA tail request row indices must be non-negative")
+
+    transfer_blocks = []
+    for src_ptr, src_row_bytes, dst_ptr, dst_row_bytes in zip(
+        src_ptrs, src_item_lens, dst_ptrs, dst_item_lens
+    ):
+        src_row_bytes = int(src_row_bytes)
+        dst_row_bytes = int(dst_row_bytes)
+        if src_row_bytes == 0 and dst_row_bytes == 0:
+            continue
+        if src_row_bytes <= 0 or src_row_bytes % src_tail_size != 0:
+            raise ValueError(
+                f"DSA source tail row size {src_row_bytes} is not divisible by "
+                f"{src_tail_size}"
+            )
+        if dst_row_bytes <= 0 or dst_row_bytes % dst_tail_size != 0:
+            raise ValueError(
+                f"DSA destination tail row size {dst_row_bytes} is not "
+                f"divisible by {dst_tail_size}"
+            )
+        src_slot_bytes = src_row_bytes // src_tail_size
+        dst_slot_bytes = dst_row_bytes // dst_tail_size
+        if src_slot_bytes != dst_slot_bytes:
+            raise ValueError(
+                "DSA tail slot-size mismatch: "
+                f"src={src_slot_bytes}, dst={dst_slot_bytes}"
+            )
+
+        slot_bytes = src_slot_bytes
+        src_row_base = int(src_ptr) + src_row_bytes * src_idx
+        dst_row_base = int(dst_ptr) + dst_row_bytes * dst_idx
+        src_seg_idx = dst_seg_idx = 0
+        src_consumed = dst_consumed = 0
+        while src_seg_idx < len(src_segments):
+            src_off, src_n = src_segments[src_seg_idx]
+            dst_off, dst_n = dst_segments[dst_seg_idx]
+            n = min(src_n - src_consumed, dst_n - dst_consumed)
+            transfer_blocks.append(
+                (
+                    src_row_base + (src_off + src_consumed) * slot_bytes,
+                    dst_row_base + (dst_off + dst_consumed) * slot_bytes,
+                    n * slot_bytes,
+                )
+            )
+            src_consumed += n
+            dst_consumed += n
+            if src_consumed == src_n:
+                src_seg_idx += 1
+                src_consumed = 0
+            if dst_consumed == dst_n:
+                dst_seg_idx += 1
+                dst_consumed = 0
+    return transfer_blocks
+
+
 def setup_state_kv_args(
     kv_args: KVArgs,
     token_to_kv_pool,
@@ -1100,6 +1285,19 @@ def setup_state_kv_args(
     kv_args.is_hybrid_mla_backend = False
     kv_args.state_conv_shard_groups = []
 
+    def append_dsa_tail(pool) -> None:
+        if not pool.kpool_use_compress:
+            return
+        tail_ptrs, tail_lens, tail_item_lens = pool.get_compress_tail_buf_infos()
+        if tail_ptrs:
+            append_state_component(
+                kv_args,
+                StateType.DSA_TAIL,
+                tail_ptrs,
+                tail_lens,
+                tail_item_lens,
+            )
+
     if isinstance(token_to_kv_pool, MHATokenToKVPoolMXFP8):
         append_state_component(
             kv_args,
@@ -1203,7 +1401,25 @@ def setup_state_kv_args(
                 slice_outer_counts,
                 layer_ids,
             )
+            # Hybrid DSA pools keep their index cache and kpool tail in the
+            # full-attention sub-pool rather than in the Mamba state above.
+            if getattr(token_to_kv_pool, "use_dsa", False):
+                dsa_pool = token_to_kv_pool.full_kv_pool
+                dsa_ptrs, dsa_lens, dsa_item_lens = dsa_pool.get_state_buf_infos()
+                append_state_component(
+                    kv_args,
+                    StateType.DSA,
+                    dsa_ptrs,
+                    dsa_lens,
+                    dsa_item_lens,
+                )
+                append_dsa_tail(dsa_pool)
         elif isinstance(token_to_kv_pool, (DSATokenToKVPool, NPUMLATokenToKVPool)):
+            tail_ptrs, tail_lens, tail_item_lens = [], [], []
+            if isinstance(token_to_kv_pool, DSATokenToKVPool):
+                tail_ptrs, tail_lens, tail_item_lens = (
+                    token_to_kv_pool.get_compress_tail_buf_infos()
+                )
             if draft_token_to_kv_pool is not None and isinstance(
                 draft_token_to_kv_pool, DSATokenToKVPool
             ):
@@ -1215,6 +1431,12 @@ def setup_state_kv_args(
                 data_ptrs = data_ptrs + draft_data_ptrs
                 data_lens = data_lens + draft_data_lens
                 item_lens = item_lens + draft_item_lens
+                draft_tail_ptrs, draft_tail_lens, draft_tail_item_lens = (
+                    draft_token_to_kv_pool.get_compress_tail_buf_infos()
+                )
+                tail_ptrs = tail_ptrs + draft_tail_ptrs
+                tail_lens = tail_lens + draft_tail_lens
+                tail_item_lens = tail_item_lens + draft_tail_item_lens
             if isinstance(token_to_kv_pool, NPUMLATokenToKVPool):
                 kv_args.kv_buf_groups = (
                     len(kv_args.kv_data_ptrs) // token_to_kv_pool.layer_num
@@ -1224,6 +1446,14 @@ def setup_state_kv_args(
                 append_state_component(
                     kv_args, StateType.DSA, data_ptrs, data_lens, item_lens
                 )
+                if tail_ptrs:
+                    append_state_component(
+                        kv_args,
+                        StateType.DSA_TAIL,
+                        tail_ptrs,
+                        tail_lens,
+                        tail_item_lens,
+                    )
 
     if is_npu() and isinstance(token_to_kv_pool, DSV4NPUTokenToKVPool):
         from sglang.srt.disaggregation.ascend.conn import AscendStateType

python/sglang/srt/hardware_backend/npu/memory_pool_npu.py

@@ -355,7 +355,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
     # NPUMHATokenToKVPool stores buffers as
     #   (num_pages, page_size, head_num, head_dim)            # use_fia=False
     #   (num_pages*page_size, 1, head_num, head_dim)          # use_fia=True
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         torch.npu.synchronize()
         buf_of_layers = []
         for local_layer_id in range(self.layer_num):
@@ -370,7 +370,9 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
         torch.npu.synchronize()
         return kv_cache_cpu
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         torch.npu.synchronize()
         chunk_size = self.cpu_offloading_chunk_size
         for local_layer_id in range(self.layer_num):
@@ -743,7 +745,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
             out.append(layer_chunks)
         return out
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         torch.npu.synchronize()
         buf_of_layers = []
         has_ik = self.index_head_dim is not None
@@ -761,7 +763,9 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
         torch.npu.synchronize()
         return kv_cache_cpu
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         torch.npu.synchronize()
         chunk_size = self.cpu_offloading_chunk_size
         has_ik = self.index_head_dim is not None

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

@@ -7,6 +7,7 @@ from sglang.srt.arg_groups.overrides import (
     resolved_view,
 )
 from sglang.srt.configs.hybrid_arch import (
+    glm5_next_config,
     hybrid_gdn_config,
     hybrid_lightning_config,
     kimi_linear_config,
@@ -505,6 +506,8 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
                 hybrid_backend_cls = AscendKDAHybridLinearAttnBackend
             else:
                 linear_attn_backend = KDAAttnBackend(runner)
+        elif glm5_next_config(runner.model_config) is not None:
+            linear_attn_backend = KDAAttnBackend(runner)
         elif hybrid_lightning_config(runner.model_config) is not None:
             linear_attn_backend = LightningAttentionBackend(runner)
         else:

python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.attention.dsa.kpool_plan import (
+    init_kpool_extend_metadata,
+    init_kpool_write_plan,
+    init_kpool_write_plan_capture,
+    init_pooled_paged_mqa_metadata,
+    update_kpool_write_plan,
+    update_pooled_paged_mqa_metadata,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import (
+        PrecomputedMetadata,
+    )
+    from sglang.srt.layers.attention.dsa.dsa_topk_backend import TopkTransformMethod
+    from sglang.srt.layers.attention.dsa_backend import _DSA_IMPL_T, DSAMetadata
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+
+
+@dataclass
+class _KPoolForwardInputs:
+    full_real_page_table: Optional[torch.Tensor] = None
+    full_seqlens_expanded: Optional[torch.Tensor] = None
+
+
+class DeepseekSparseAttnBackendKPoolMixin:
+    """KPool-specific metadata and tail handling for the DSA backend."""
+
+    def _check_kpool_tail_backend(
+        self,
+        topk_indices: Optional[torch.Tensor],
+        dsa_impl: _DSA_IMPL_T,
+        phase: str,
+    ) -> None:
+        if (
+            topk_indices is None
+            or self.dsa_index_kpool <= 1
+            or dsa_impl in ("fa3", "tilelang", "trtllm")
+        ):
+            return
+        raise NotImplementedError(
+            "index_kpool > 1 appends tail tokens to topk_indices and is "
+            f"currently only supported by the FA3/TileLang/TRTLLM DSA {phase} "
+            "backend."
+        )
+
+    def _resolve_kpool_tail_backend(
+        self,
+        topk_indices: Optional[torch.Tensor],
+        dsa_impl: _DSA_IMPL_T,
+    ) -> _DSA_IMPL_T:
+        if (
+            topk_indices is None
+            or self.dsa_index_kpool <= 1
+            or dsa_impl != "flashmla_sparse"
+        ):
+            return dsa_impl
+        if self.device_sm_major >= 10:
+            return "trtllm"
+        if self.device_sm_major == 9:
+            return "fa3"
+        return dsa_impl
+
+    def _kpool_slots_per_page(self) -> int:
+        return getattr(self.token_to_kv_pool, "slots_per_page", self.real_page_size)
+
+    def _build_kpool_paged_mqa_schedule_metadata(self) -> bool:
+        if self.device_sm_major == 9:
+            return self.num_q_heads in (32, 64)
+        return True
+
+    def _init_kpool_metadata(
+        self,
+        metadata: DSAMetadata,
+        forward_batch: ForwardBatch,
+        topk_transform_method: Optional[TopkTransformMethod] = None,
+        kpool_inputs: Optional[_KPoolForwardInputs] = None,
+    ) -> DSAMetadata:
+        if self.dsa_index_kpool <= 1:
+            return metadata
+
+        forward_mode = forward_batch.forward_mode
+        slots_per_page = self._kpool_slots_per_page()
+        build_schedule_metadata = self._build_kpool_paged_mqa_schedule_metadata()
+        if forward_mode.is_extend_without_speculative():
+            assert topk_transform_method is not None
+            assert kpool_inputs is not None
+            return init_kpool_extend_metadata(
+                metadata,
+                forward_batch,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                slots_per_page=slots_per_page,
+                topk_transform_method=topk_transform_method,
+                full_real_page_table=kpool_inputs.full_real_page_table,
+                full_seqlens_expanded=kpool_inputs.full_seqlens_expanded,
+            )
+
+        if forward_mode.is_decode_or_idle():
+            metadata = init_pooled_paged_mqa_metadata(
+                metadata,
+                metadata.cache_seqlens_int32,
+                forward_mode,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                slots_per_page=slots_per_page,
+                build_schedule_metadata=build_schedule_metadata,
+            )
+            return init_kpool_write_plan(
+                metadata,
+                forward_batch,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                real_page_table=metadata.real_page_table,
+                num_draft_tokens=1,
+                write_start=(forward_batch.seq_lens - 1).to(torch.int32),
+                slots_per_page=slots_per_page,
+                build_schedule_metadata=build_schedule_metadata,
+            )
+
+        if forward_mode.is_target_verify():
+            return init_kpool_write_plan(
+                metadata,
+                forward_batch,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                real_page_table=metadata.real_page_table,
+                num_draft_tokens=self.speculative_num_draft_tokens,
+                write_start=forward_batch.seq_lens.to(torch.int32),
+                slots_per_page=slots_per_page,
+                build_schedule_metadata=build_schedule_metadata,
+            )
+
+        if forward_mode.is_draft_extend_v2():
+            spec_info = forward_batch.spec_info
+            effective_n_per_batch = (
+                spec_info.num_accept_tokens
+                if spec_info is not None
+                and getattr(spec_info, "num_accept_tokens", None) is not None
+                else None
+            )
+            return init_kpool_write_plan(
+                metadata,
+                forward_batch,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                real_page_table=metadata.real_page_table,
+                num_draft_tokens=self.speculative_num_draft_tokens,
+                write_start=(
+                    forward_batch.seq_lens - self.speculative_num_draft_tokens
+                ).to(torch.int32),
+                slots_per_page=slots_per_page,
+                effective_n_per_batch=effective_n_per_batch,
+                build_schedule_metadata=build_schedule_metadata,
+            )
+
+        return metadata
+
+    def _init_kpool_metadata_capture(
+        self, metadata: DSAMetadata, bs: int, forward_mode: ForwardMode
+    ) -> DSAMetadata:
+        if self.dsa_index_kpool <= 1:
+            return metadata
+
+        slots_per_page = self._kpool_slots_per_page()
+        build_schedule_metadata = self._build_kpool_paged_mqa_schedule_metadata()
+        if forward_mode.is_decode_or_idle():
+            metadata = init_pooled_paged_mqa_metadata(
+                metadata,
+                metadata.cache_seqlens_int32,
+                forward_mode,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                slots_per_page=slots_per_page,
+                build_schedule_metadata=build_schedule_metadata,
+            )
+
+        if (
+            forward_mode.is_decode_or_idle()
+            or forward_mode.is_target_verify()
+            or forward_mode.is_draft_extend_v2()
+        ):
+            is_decode = forward_mode.is_decode_or_idle()
+            is_v2 = forward_mode.is_draft_extend_v2()
+            metadata = init_kpool_write_plan_capture(
+                metadata,
+                max_bs=bs,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                num_draft_tokens=(
+                    1 if is_decode else self.speculative_num_draft_tokens
+                ),
+                device=self.device,
+                is_verify=not is_decode,
+                slots_per_page=slots_per_page,
+                is_v2=is_v2,
+                build_schedule_metadata=build_schedule_metadata,
+            )
+
+        return metadata
+
+    def _update_kpool_metadata_replay(
+        self,
+        metadata: DSAMetadata,
+        seq_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        forward_mode: ForwardMode,
+        effective_n_per_batch: Optional[torch.Tensor] = None,
+    ) -> None:
+        if self.dsa_index_kpool <= 1:
+            return
+
+        slots_per_page = self._kpool_slots_per_page()
+        build_schedule_metadata = self._build_kpool_paged_mqa_schedule_metadata()
+        if forward_mode.is_decode_or_idle():
+            update_pooled_paged_mqa_metadata(
+                metadata,
+                metadata.cache_seqlens_int32,
+                forward_mode,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                slots_per_page=slots_per_page,
+                build_schedule_metadata=build_schedule_metadata,
+            )
+
+        if not (
+            forward_mode.is_decode_or_idle()
+            or forward_mode.is_target_verify()
+            or forward_mode.is_draft_extend_v2()
+        ):
+            return
+
+        is_decode = forward_mode.is_decode_or_idle()
+        is_v2 = forward_mode.is_draft_extend_v2()
+        if is_decode:
+            write_start = seq_lens.to(torch.int32) - 1
+        elif is_v2:
+            write_start = seq_lens.to(torch.int32) - self.speculative_num_draft_tokens
+        else:
+            # Target verify: write_start == seq_lens exactly; the plan kernel
+            # casts on load, so skip the per-replay int32 alloc + conversion.
+            write_start = seq_lens
+        update_kpool_write_plan(
+            metadata,
+            write_start=write_start,
+            req_pool_indices=req_pool_indices,
+            real_page_table=metadata.real_page_table,
+            pool_size=self.dsa_index_kpool,
+            real_page_size=self.real_page_size,
+            num_draft_tokens=(1 if is_decode else self.speculative_num_draft_tokens),
+            forward_mode=forward_mode,
+            slots_per_page=slots_per_page,
+            effective_n_per_batch=effective_n_per_batch,
+        )
+
+    def _update_kpool_metadata_from_precomputed(
+        self,
+        metadata: DSAMetadata,
+        precomputed: PrecomputedMetadata,
+        forward_mode: ForwardMode,
+    ) -> None:
+        if self.dsa_index_kpool <= 1:
+            return
+
+        slots_per_page = self._kpool_slots_per_page()
+        build_schedule_metadata = self._build_kpool_paged_mqa_schedule_metadata()
+        if forward_mode.is_decode_or_idle():
+            update_pooled_paged_mqa_metadata(
+                metadata,
+                precomputed.cache_seqlens,
+                forward_mode,
+                pool_size=self.dsa_index_kpool,
+                real_page_size=self.real_page_size,
+                slots_per_page=slots_per_page,
+                build_schedule_metadata=build_schedule_metadata,
+            )
+
+        if not (forward_mode.is_decode_or_idle() or forward_mode.is_target_verify()):
+            return
+
+        is_verify = forward_mode.is_target_verify()
+        write_start = precomputed.cache_seqlens.to(torch.int32)
+        write_start = (
+            write_start - self.speculative_num_draft_tokens
+            if is_verify
+            else write_start - 1
+        )
+        update_kpool_write_plan(
+            metadata,
+            write_start=write_start,
+            req_pool_indices=precomputed.req_pool_indices,
+            real_page_table=metadata.real_page_table,
+            pool_size=self.dsa_index_kpool,
+            real_page_size=self.real_page_size,
+            num_draft_tokens=self.speculative_num_draft_tokens if is_verify else 1,
+            forward_mode=forward_mode,
+            slots_per_page=slots_per_page,
+        )

python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py

@@ -35,6 +35,7 @@ class PrecomputedMetadata:
     # Basic seqlens
     cache_seqlens: torch.Tensor  # int32, [bs]
     cu_seqlens_k: torch.Tensor  # int32, [bs+1]
+    req_pool_indices: torch.Tensor  # int64, [bs]
 
     # Page table
     page_indices: torch.Tensor  # int32, [bs, max_len] or [expanded_bs, max_len]
@@ -121,7 +122,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
         """Precompute metadata for normal decode mode."""
         max_len = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
 
-        if _is_cuda and not _is_hip:
+        if _is_cuda and not _is_hip and self.dsa_index_kpool <= 1:
             from sglang.kernels.ops.attention.dsa_metadata import (
                 fused_dsa_decode_metadata,
             )
@@ -173,6 +174,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
             return PrecomputedMetadata(
                 cache_seqlens=cache_seqlens,
                 cu_seqlens_k=cu_seqlens_k,
+                req_pool_indices=req_pool_indices,
                 page_indices=page_indices,
                 real_page_table=real_page_table,
                 seqlens_expanded=seqlens_expanded,
@@ -193,7 +195,9 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
 
         # Compute DSA seqlens
         dsa_cache_seqlens = compute_dsa_seqlens(
-            cache_seqlens, dsa_index_topk=self.dsa_index_topk
+            cache_seqlens,
+            dsa_index_topk=self.dsa_index_topk,
+            index_kpool=self.dsa_index_kpool,
         )
         seqlens_expanded = cache_seqlens
         seqlens_expanded_size = seqlens_expanded.shape[0]
@@ -218,6 +222,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
         return PrecomputedMetadata(
             cache_seqlens=cache_seqlens,
             cu_seqlens_k=cu_seqlens_k,
+            req_pool_indices=req_pool_indices,
             page_indices=page_indices,
             real_page_table=real_page_table,
             seqlens_expanded=seqlens_expanded,
@@ -240,7 +245,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
         max_seqlen_k = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
         seqlens_expanded_size = bs * self.speculative_num_draft_tokens
 
-        if _is_cuda and not _is_hip:
+        if _is_cuda and not _is_hip and self.dsa_index_kpool <= 1:
             from sglang.kernels.ops.attention.dsa_metadata import (
                 fused_dsa_target_verify_metadata,
             )
@@ -305,6 +310,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
             return PrecomputedMetadata(
                 cache_seqlens=cache_seqlens,
                 cu_seqlens_k=cu_seqlens_k,
+                req_pool_indices=req_pool_indices,
                 page_indices=page_indices,
                 real_page_table=real_page_table,
                 seqlens_expanded=seqlens_expanded,
@@ -342,7 +348,11 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
         )
 
         # Compute DSA seqlens
-        dsa_cache_seqlens = compute_dsa_seqlens(seqlens_expanded, self.dsa_index_topk)
+        dsa_cache_seqlens = compute_dsa_seqlens(
+            seqlens_expanded,
+            self.dsa_index_topk,
+            index_kpool=self.dsa_index_kpool,
+        )
         seqlens_expanded_size = seqlens_expanded.shape[0]
 
         # DSA cumsum
@@ -365,6 +375,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
         return PrecomputedMetadata(
             cache_seqlens=cache_seqlens,
             cu_seqlens_k=cu_seqlens_k,
+            req_pool_indices=req_pool_indices,
             page_indices=page_indices,
             real_page_table=real_page_table,
             seqlens_expanded=seqlens_expanded,

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -486,7 +486,6 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
                     dim=-1,
                 )
             with torch.cuda.stream(self.alt_stream):
-                # TODO we should also put DeepGEMM half SM here?
                 if self.use_dsa_indexer_fusion:
                     key, weights_raw = self._fused_k_weights(x)
                 else:
@@ -876,10 +875,18 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
 
         if use_dg_native:
             seqlens_32_2d = ctx_2d
-        elif seqlens_32.dim() == 2:
+        elif ctx_2d is not None:
-            seqlens_32_2d = seqlens_32
+            if ctx_2d.size(1) == 1:
+                seqlens_32_2d = ctx_2d
             else:
-            seqlens_32_2d = seqlens_32.unsqueeze(-1)
+                seqlens_32_2d = ctx_2d.reshape(-1).contiguous().view(-1, 1)
+        elif seqlens_32.dim() == 2:
+            if seqlens_32.size(1) == 1:
+                seqlens_32_2d = seqlens_32.contiguous()
+            else:
+                seqlens_32_2d = seqlens_32.reshape(-1).contiguous().view(-1, 1)
+        else:
+            seqlens_32_2d = seqlens_32.contiguous().view(-1, 1)
         if _is_cuda:
             if schedule_metadata is None:
                 schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
@@ -896,6 +903,52 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
         assert len(weights.shape) == 3
         weights = weights.squeeze(2)
 
+        # SM100 DeepGEMM paged MQA requires batch_size <= num_sms; chunk larger batches.
+        def _chunked_fp8_paged_mqa_logits(
+            q: torch.Tensor,
+            kv_cache: torch.Tensor,
+            w: torch.Tensor,
+            context_lens: torch.Tensor,
+            block_table: torch.Tensor,
+            mqa_schedule_metadata: torch.Tensor,
+            max_len: int,
+            clean_logits: bool = False,
+        ) -> torch.Tensor:
+            batch_size, chunk_next_n = q.shape[:2]
+            if batch_size == 0:
+                return torch.empty((0, max_len), dtype=torch.float32, device=q.device)
+            if batch_size <= self.sm_count:
+                return deep_gemm.fp8_paged_mqa_logits(
+                    q,
+                    kv_cache,
+                    w,
+                    context_lens,
+                    block_table,
+                    mqa_schedule_metadata,
+                    max_len,
+                    clean_logits=clean_logits,
+                )
+            logits_chunks = []
+            for start in range(0, batch_size, self.sm_count):
+                end = min(start + self.sm_count, batch_size)
+                chunk_context_lens = context_lens[start:end]
+                chunk_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
+                    chunk_context_lens, blocksize, self.sm_count
+                )
+                logits_chunks.append(
+                    deep_gemm.fp8_paged_mqa_logits(
+                        q[start:end],
+                        kv_cache,
+                        w[start * chunk_next_n : end * chunk_next_n],
+                        chunk_context_lens,
+                        block_table[start:end],
+                        chunk_schedule_metadata,
+                        max_len,
+                        clean_logits=clean_logits,
+                    )
+                )
+            return torch.cat(logits_chunks, dim=0)
+
         if self.paged_mqa_logits_backend.is_aiter():
             logits = aiter_paged_mqa_logits(
                 q_fp8,
@@ -928,7 +981,7 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
             )
         elif use_dg_native:
             logits = deepgemm_paged_mqa_logits_native(
-                deep_gemm.fp8_paged_mqa_logits,
+                _chunked_fp8_paged_mqa_logits,
                 q_fp8,
                 kv_cache_fp8,
                 weights,
@@ -942,7 +995,7 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
             )
         else:
             logits = deepgemm_paged_mqa_logits_split(
-                deep_gemm.fp8_paged_mqa_logits,
+                _chunked_fp8_paged_mqa_logits,
                 q_fp8,
                 kv_cache_fp8,
                 weights,

python/sglang/srt/layers/attention/dsa/kpool_plan.py

@@ -731,7 +731,6 @@ def update_kpool_write_plan(
     forward_mode: ForwardMode,
     slots_per_page: int,
     effective_n_per_batch: Optional[torch.Tensor] = None,
-    include_deep_gemm_schedule: bool = True,
 ) -> None:
     if not _is_kpool_layout_enabled(pool_size, real_page_size) or not is_cuda():
         return
@@ -767,9 +766,7 @@ def update_kpool_write_plan(
             effective_n_per_batch.to(torch.int32)
         )
 
-    # In-graph replay updates plan lengths too late for host schedule construction;
+    if plan.pool_schedule_metadata is not None:
-    # the caller rebuilds the schedule from raw seq_lens out of graph.
-    if include_deep_gemm_schedule and plan.pool_schedule_metadata is not None:
         new_schedule = _compute_pool_schedule_metadata(
             plan.pool_seqlens_per_q,
             slots_per_page=slots_per_page,
@@ -778,25 +775,6 @@ def update_kpool_write_plan(
             plan.pool_schedule_metadata.copy_(new_schedule)
 
 
-def refresh_kpool_pool_schedule_from(
-    metadata: DSAMetadata,
-    pool_seqlens_per_q: torch.Tensor,
-    *,
-    slots_per_page: int,
-) -> None:
-    """Use an explicit source because the captured plan buffer remains stale
-    until replay."""
-    plan = metadata.kpool_write_plan
-    if plan is None or plan.pool_schedule_metadata is None:
-        return
-    new_schedule = _compute_pool_schedule_metadata(
-        pool_seqlens_per_q,
-        slots_per_page=slots_per_page,
-    )
-    if new_schedule is not None:
-        plan.pool_schedule_metadata.copy_(new_schedule)
-
-
 def init_kpool_write_plan(
     metadata: DSAMetadata,
     forward_batch: ForwardBatch,

python/sglang/srt/layers/attention/dsa/utils.py

@@ -71,9 +71,19 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 
-def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int):
+def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int, index_kpool: int = 1):
+    if index_kpool <= 1:
         return original_seq_lens.clamp(max=dsa_index_topk)
 
+    # Clamp only complete pools; the unfinished tail must remain selectable
+    # outside the pooled top-k budget.
+    full_pool_tokens = (
+        torch.div(original_seq_lens, index_kpool, rounding_mode="floor") * index_kpool
+    )
+    selected_history_tokens = full_pool_tokens.clamp(max=dsa_index_topk)
+    tail_tokens = original_seq_lens - full_pool_tokens
+    return selected_history_tokens + tail_tokens
+
 
 def should_remap_pd_dsa_seed_to_local_slots() -> bool:
     """Whether a PD seed should enter the allocator-local fused TopK domain."""

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

@@ -31,18 +31,34 @@ from sglang.kernels.ops.attention.dsa.dequant_k_cache import (
 )
 from sglang.kernels.ops.attention.dsa.quant_k_cache import quantize_k_cache
 from sglang.kernels.ops.attention.dsa.transform_index import (
+    prepare_trtllm_nope_sparse_metadata,
     transform_index_page_table_decode,
     transform_index_page_table_prefill,
 )
+from sglang.kernels.ops.attention.dsa_metadata import (
+    fused_dsa_decode_metadata,
+    fused_dsa_draft_extend_metadata,
+    fused_dsa_target_verify_metadata,
+)
 from sglang.kernels.ops.attention.utils import (
     concat_mla_absorb_q_general,
     mla_quantize_and_rope_for_fp8,
+    mla_quantize_for_fp8_no_rope,
     q8kv8_topk_length_from_indices,
     seqlens_expand_triton,
 )
 from sglang.kernels.ops.kvcache.cache_ops import concat_and_cast_q_fp8_pad
+from sglang.srt.configs.model_config import (
+    get_dsa_index_kpool,
+    get_dsa_index_topk,
+    is_deepseek_dsa,
+)
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.dsa.dsa_backend_kpool import (
+    DeepseekSparseAttnBackendKPoolMixin,
+    _KPoolForwardInputs,
+)
 from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import (
     DeepseekSparseAttnBackendMTPPrecomputeMixin,
     PrecomputedMetadata,
@@ -53,6 +69,10 @@ from sglang.srt.layers.attention.dsa.dsa_topk_backend import (
     DSATopKBackend,
     TopkTransformMethod,
 )
+from sglang.srt.layers.attention.dsa.kpool_plan import (
+    KPoolExtendPlan,
+    KPoolWritePlan,
+)
 from sglang.srt.layers.attention.dsa.utils import (
     can_dsa_prefill_cp_round_robin_split,
     compute_dsa_seqlens,
@@ -75,6 +95,7 @@ from sglang.srt.layers.utils.cp_utils import (
     cp_split_and_rebuild_position,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.runtime_context import get_buffer, get_exec, get_parallel, get_spec
 from sglang.srt.utils import (
     get_bool_env_var,
     is_cuda,
@@ -83,6 +104,8 @@ from sglang.srt.utils import (
     print_warning_once,
 )
 
+logger = logging.getLogger(__name__)
+
 # Opt-in (default off): route the fp8 sparse-MLA prefill path through the Triton
 # per-query flash kernel instead of TileLang. Validated on gfx950 (GLM-5.1 @
 # TP4: 16 heads, d_v=512, tail=64). Reads q_nope/q_rope directly (skips the
@@ -253,6 +276,13 @@ class DSAMetadata:
     # batch index for each token.
     token_to_batch_idx: Optional[torch.Tensor] = None
 
+    pooled_index_kpool: int = 1
+    pooled_cache_seqlens_int32: Optional[torch.Tensor] = None
+    pooled_real_page_table: Optional[torch.Tensor] = None
+    pooled_paged_mqa_schedule_metadata: Optional[torch.Tensor] = None
+    kpool_extend_plan: Optional[KPoolExtendPlan] = None
+    kpool_write_plan: Optional[KPoolWritePlan] = None
+
 
 @torch.compile
 def _compiled_cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor:
@@ -287,7 +317,9 @@ _DSA_IMPL_T: TypeAlias = Literal[
 
 
 class DeepseekSparseAttnBackend(
-    DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend
+    DeepseekSparseAttnBackendKPoolMixin,
+    DeepseekSparseAttnBackendMTPPrecomputeMixin,
+    AttentionBackend,
 ):
     # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch
     # can never carry more seqs than the pool.
@@ -314,12 +346,15 @@ class DeepseekSparseAttnBackend(
         self.num_splits = (
             1 if get_exec().deterministic.enable_deterministic_inference else 0
         )
-        self.use_dsa = is_deepseek_dsa(model_runner.model_config.hf_config)
+        hf_config = model_runner.model_config.hf_config
+        self.use_dsa = is_deepseek_dsa(hf_config)
         assert self.use_dsa, "DSA backend only supports DeepSeek DSA"
         self.dsa_kv_cache_store_fp8 = (
             model_runner.token_to_kv_pool.dsa_kv_cache_store_fp8
         )
-        self.dsa_index_topk = get_dsa_index_topk(model_runner.model_config.hf_config)
+        self.dsa_index_topk = get_dsa_index_topk(hf_config)
+        self.dsa_index_kpool = get_dsa_index_kpool(hf_config)
+        self.needs_cpu_seq_lens = self.dsa_index_kpool > 1
         self.max_context_len = model_runner.model_config.context_len
         self.num_q_heads = (
             model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
@@ -350,6 +385,23 @@ class DeepseekSparseAttnBackend(
         self.enable_auto_select_prefill_impl = self.dsa_prefill_impl == "flashmla_auto"
         self._sink_pad_cache: dict[tuple[int, int], torch.Tensor] = {}
 
+        # Hoisted per-call imports of set_dsa_prefill_impl. Module-scope
+        # imports would cycle through model_executor (which imports the
+        # attention backends); backend init runs after those modules are
+        # fully imported, so binding the function refs here is cycle-safe.
+        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+            is_in_breakable_cuda_graph,
+        )
+        from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+            is_in_tc_piecewise_cuda_graph,
+        )
+        from sglang.srt.utils import get_device_sm, is_blackwell
+
+        self._is_in_breakable_cuda_graph = is_in_breakable_cuda_graph
+        self._is_in_tc_piecewise_cuda_graph = is_in_tc_piecewise_cuda_graph
+        self._get_device_sm = get_device_sm
+        self._is_blackwell = is_blackwell
+
         self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32)
 
         if _is_hip:
@@ -654,6 +706,21 @@ class DeepseekSparseAttnBackend(
             "num_kv_splits": self.aiter_dsa_max_split_per_batch,
         }
 
+    def _pad_trtllm_sparse_page_table(
+        self, page_table_1: torch.Tensor
+    ) -> Tuple[torch.Tensor, int]:
+        sparse_mla_top_k = page_table_1.shape[1]
+        padded_top_k = ((sparse_mla_top_k + 3) // 4) * 4
+        if padded_top_k == sparse_mla_top_k:
+            return page_table_1, sparse_mla_top_k
+        padding = torch.full(
+            (page_table_1.shape[0], padded_top_k - sparse_mla_top_k),
+            -1,
+            dtype=page_table_1.dtype,
+            device=page_table_1.device,
+        )
+        return torch.cat([page_table_1, padding], dim=1), padded_top_k
+
     def _build_paged_mqa_schedule_2d_ctx_lens(
         self,
         forward_mode: ForwardMode,
@@ -737,6 +804,9 @@ class DeepseekSparseAttnBackend(
             )
         return self._arange_buf[:length]
 
+    def update_verify_buffers_to_fill_after_draft(self, *args, **kwargs):
+        return None
+
     def _graph_page_table_width(self, metadata: DSAMetadata) -> int:
         """Column count to scan req_to_token during graph replay. Reads the wide
         page_table_1 width when present, else req_to_token's width (the wide table
@@ -826,6 +896,16 @@ class DeepseekSparseAttnBackend(
         # seq_len_cpu of selected sequences
         indexer_seq_lens_cpu = forward_batch.seq_lens_cpu
         indexer_seq_lens = forward_batch.seq_lens
+        use_kpool = self.dsa_index_kpool > 1
+        if use_kpool:
+            assert (
+                self.real_page_size == 64
+                and self.real_page_size % self.dsa_index_kpool == 0
+            ), (
+                f"kpool path requires page_size == 64 and page_size % pool_size == 0; "
+                f"got page_size={self.real_page_size}, pool_size={self.dsa_index_kpool}."
+            )
+        kpool_inputs = _KPoolForwardInputs()
 
         if forward_batch.forward_mode.is_decode_or_idle():
             extend_seq_lens_cpu = [1] * batch_size
@@ -925,6 +1005,11 @@ class DeepseekSparseAttnBackend(
                     )
                 ]
             )
+            if use_kpool:
+                kpool_inputs.full_real_page_table = self._transform_table_1_to_real(
+                    page_table
+                )
+                kpool_inputs.full_seqlens_expanded = seqlens_expanded
 
             if can_dsa_prefill_cp_round_robin_split(forward_batch):
                 if is_cp_v2_active(forward_batch):
@@ -1011,6 +1096,7 @@ class DeepseekSparseAttnBackend(
         dsa_cache_seqlens_int32 = compute_dsa_seqlens(
             original_seq_lens=seqlens_expanded,
             dsa_index_topk=self.dsa_index_topk,
+            index_kpool=self.dsa_index_kpool,
         )
         dsa_cache_seqlens_int32 = pad_dsa_cache_seqlens(
             forward_batch, dsa_cache_seqlens_int32
@@ -1072,6 +1158,12 @@ class DeepseekSparseAttnBackend(
             token_to_batch_idx=token_to_batch_idx,
             topk_v2_plan=self._build_topk_v2_plan(seqlens_expanded),
         )
+        metadata = self._init_kpool_metadata(
+            metadata,
+            forward_batch,
+            topk_transform_method=topk_transform_method,
+            kpool_inputs=kpool_inputs,
+        )
         self.forward_metadata = metadata
 
     def _cal_indexer_k_start_end(
@@ -1183,6 +1275,8 @@ class DeepseekSparseAttnBackend(
             and self.real_page_size > 1
             and self.hisparse_coordinator is None
             and not self.speculative_num_draft_tokens
+            # kpool's PAGED fused-topk mapping still reads page_table_1.
+            and self.dsa_index_kpool <= 1
             and self.use_fused_topk
             and self.dsa_topk_backend.should_use_topk_v2()
             and self.dsa_index_topk is not None
@@ -1275,7 +1369,9 @@ class DeepseekSparseAttnBackend(
             # NOTE(dark): this is always arange, since we are decoding
             cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][: bs + 1]
             dsa_cache_seqlens_int32 = compute_dsa_seqlens(
-                cache_seqlens_int32, dsa_index_topk=self.dsa_index_topk
+                cache_seqlens_int32,
+                dsa_index_topk=self.dsa_index_topk,
+                index_kpool=self.dsa_index_kpool,
             )
 
             seqlens_expanded = cache_seqlens_int32
@@ -1338,7 +1434,9 @@ class DeepseekSparseAttnBackend(
                 ]
             )
             dsa_cache_seqlens_int32 = compute_dsa_seqlens(
-                seqlens_expanded, dsa_index_topk=self.dsa_index_topk
+                seqlens_expanded,
+                dsa_index_topk=self.dsa_index_topk,
+                index_kpool=self.dsa_index_kpool,
             )
             dsa_extend_seq_lens_list = [1] * bs * self.speculative_num_draft_tokens
 
@@ -1400,6 +1498,7 @@ class DeepseekSparseAttnBackend(
             dsa_extend_seq_lens_list=dsa_extend_seq_lens_list,
             topk_v2_plan=self._build_topk_v2_plan(seqlens_expanded),
         )
+        metadata = self._init_kpool_metadata_capture(metadata, bs, forward_mode)
         self.decode_cuda_graph_metadata[bs] = metadata
         self.forward_metadata = metadata
 
@@ -1434,24 +1533,22 @@ class DeepseekSparseAttnBackend(
             )
             return
 
+        metadata: DSAMetadata = self.decode_cuda_graph_metadata[bs]
+
         self.set_dsa_prefill_impl(forward_batch=None)
 
         seq_lens = seq_lens[:bs]
         req_pool_indices = req_pool_indices[:bs]
 
         # Normal Decode
-        metadata: DSAMetadata = self.decode_cuda_graph_metadata[bs]
         used_fused_metadata_generation = False
         target_verify_ctx_lens_written = False
+
         if forward_mode.is_decode_or_idle():
             # Normal Decode
             max_len = self._graph_page_table_width(metadata)
 
-            if is_cuda() and not _is_hip:
+            if is_cuda() and not _is_hip and self.dsa_index_kpool <= 1:
-                from sglang.kernels.ops.attention.dsa_metadata import (
-                    fused_dsa_decode_metadata,
-                )
-
                 fused_dsa_decode_metadata(
                     seq_lens=seq_lens,
                     req_pool_indices=req_pool_indices,
@@ -1482,18 +1579,16 @@ class DeepseekSparseAttnBackend(
                 page_indices = self.req_to_token[req_pool_indices, :max_len]
                 metadata.page_table_1[:, :max_len].copy_(page_indices)
                 dsa_cache_seqlens = compute_dsa_seqlens(
-                    cache_seqlens, dsa_index_topk=self.dsa_index_topk
+                    cache_seqlens,
+                    dsa_index_topk=self.dsa_index_topk,
+                    index_kpool=self.dsa_index_kpool,
                 )
                 metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens)
                 seqlens_expanded = cache_seqlens
         elif forward_mode.is_target_verify():
             max_seqlen_k = self._graph_page_table_width(metadata)
 
-            if is_cuda() and not _is_hip:
+            if is_cuda() and not _is_hip and self.dsa_index_kpool <= 1:
-                from sglang.kernels.ops.attention.dsa_metadata import (
-                    fused_dsa_target_verify_metadata,
-                )
-
                 paged_mqa_ctx_lens_2d = None
                 if (
                     self.speculative_num_draft_tokens >= 2
@@ -1565,7 +1660,9 @@ class DeepseekSparseAttnBackend(
                 )
                 metadata.dsa_seqlens_expanded.copy_(seqlens_expanded)
                 dsa_cache_seqlens = compute_dsa_seqlens(
-                    seqlens_expanded, self.dsa_index_topk
+                    seqlens_expanded,
+                    self.dsa_index_topk,
+                    index_kpool=self.dsa_index_kpool,
                 )
                 metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens)
         elif forward_mode.is_draft_extend_v2():
@@ -1587,11 +1684,7 @@ class DeepseekSparseAttnBackend(
                 device=self.device,
             )
 
-            if is_cuda() and not _is_hip:
+            if is_cuda() and not _is_hip and self.dsa_index_kpool <= 1:
-                from sglang.kernels.ops.attention.dsa_metadata import (
-                    fused_dsa_draft_extend_metadata,
-                )
-
                 fused_dsa_draft_extend_metadata(
                     seq_lens=seq_lens,
                     extend_seq_lens=extend_seq_lens,
@@ -1642,7 +1735,9 @@ class DeepseekSparseAttnBackend(
                     seqlens_expanded
                 )
                 dsa_cache_seqlens = compute_dsa_seqlens(
-                    seqlens_expanded, self.dsa_index_topk
+                    seqlens_expanded,
+                    self.dsa_index_topk,
+                    index_kpool=self.dsa_index_kpool,
                 )
                 metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens)
 
@@ -1696,6 +1791,19 @@ class DeepseekSparseAttnBackend(
         else:
             assert metadata.real_page_table is metadata.page_table_1
 
+        effective_n_per_batch = None
+        if forward_mode.is_draft_extend_v2() and spec_info is not None:
+            effective_n_per_batch = getattr(spec_info, "num_accept_tokens", None)
+            if effective_n_per_batch is not None:
+                effective_n_per_batch = effective_n_per_batch[:bs]
+        self._update_kpool_metadata_replay(
+            metadata,
+            seq_lens,
+            req_pool_indices,
+            forward_mode,
+            effective_n_per_batch=effective_n_per_batch,
+        )
+
         if self.dsa_decode_impl == "flashmla_kv":
             flashmla_metadata = metadata.flashmla_metadata.slice(
                 slice(0, seqlens_expanded_size + 1)
@@ -1848,10 +1956,8 @@ class DeepseekSparseAttnBackend(
                 flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1))
                 flashmla_metadata.copy_(precomputed.flashmla_metadata)
 
-        # Refresh DeepGEMM paged MQA schedule metadata for the actual seqlens of
+        # Refresh the schedule because stale shape decomposition can deadlock
-        # this replay (the captured graph holds stale data otherwise, which can
+        # DeepGEMM paged MQA.
-        # deadlock the kernel when the runtime work decomposition diverges from
-        # the captured one).
         if is_cuda():
             if forward_mode.is_decode_or_idle():
                 seqlens_32_2d = _to_2d_context_lens(metadata.cache_seqlens_int32, bs)
@@ -1869,6 +1975,10 @@ class DeepseekSparseAttnBackend(
             else:
                 metadata.paged_mqa_ctx_lens_2d.copy_(seqlens_32_2d)
 
+        self._update_kpool_metadata_from_precomputed(
+            metadata, precomputed, forward_mode
+        )
+
         self.forward_metadata = metadata
 
     def forward_extend(
@@ -1906,6 +2016,17 @@ class DeepseekSparseAttnBackend(
                 f"Learnable attention sinks require flashmla_sparse, got {dsa_impl}"
             )
 
+        phase = (
+            "decode"
+            if (
+                forward_batch.forward_mode.is_target_verify()
+                or forward_batch.forward_mode.is_draft_extend_v2()
+            )
+            else "prefill"
+        )
+        dsa_impl = self._resolve_kpool_tail_backend(topk_indices, dsa_impl)
+        self._check_kpool_tail_backend(topk_indices, dsa_impl, phase)
+
         if dsa_impl == "trtllm" and not self.use_mha:
             return self._forward_trtllm(
                 q,
@@ -1962,7 +2083,9 @@ class DeepseekSparseAttnBackend(
         if q_rope is not None:
             q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
             q_rope = q_rope.view(
-                -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                q_nope.shape[0],
+                layer.tp_q_head_num,
+                layer.head_dim - layer.v_head_dim,
             )
         else:
             q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
@@ -2209,13 +2332,15 @@ class DeepseekSparseAttnBackend(
         metadata = self.forward_metadata
         assert causal, "DSA is causal only"
 
-        if attn_sink is not None and self.dsa_decode_impl != "flashmla_sparse":
+        dsa_impl = self._resolve_kpool_tail_backend(topk_indices, self.dsa_decode_impl)
+        self._check_kpool_tail_backend(topk_indices, dsa_impl, "decode")
+
+        if attn_sink is not None and dsa_impl != "flashmla_sparse":
             raise RuntimeError(
-                "Learnable attention sinks require flashmla_sparse, got "
+                f"Learnable attention sinks require flashmla_sparse, got {dsa_impl}"
-                f"{self.dsa_decode_impl}"
             )
 
-        if self.dsa_decode_impl == "trtllm":
+        if dsa_impl == "trtllm":
             return self._forward_trtllm(
                 q,
                 k,
@@ -2252,7 +2377,9 @@ class DeepseekSparseAttnBackend(
         if q_rope is not None:
             q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
             q_rope = q_rope.view(
-                -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                q_nope.shape[0],
+                layer.tp_q_head_num,
+                layer.head_dim - layer.v_head_dim,
             )
             # Caller passed split q_nope / q_rope; we'll need to concat below if
             # the chosen impl wants q_all.
@@ -2285,7 +2412,7 @@ class DeepseekSparseAttnBackend(
                 page_size=1,
             )
 
-        if self.dsa_decode_impl == "flashmla_sparse":
+        if dsa_impl == "flashmla_sparse":
             if q_rope is not None:
                 q_all = concat_mla_absorb_q_general(q_nope, q_rope)
             return self._forward_flashmla_sparse(
@@ -2297,7 +2424,7 @@ class DeepseekSparseAttnBackend(
                 topk_length=metadata.dsa_cache_seqlens_int32,
                 attn_sink=attn_sink,
             )
-        elif self.dsa_decode_impl == "flashinfer_sparse_mla":
+        elif dsa_impl == "flashinfer_sparse_mla":
             if q_all is None:
                 q_all = concat_mla_absorb_q_general(q_nope, q_rope)
             return self._forward_flashinfer_sparse_mla(
@@ -2308,7 +2435,7 @@ class DeepseekSparseAttnBackend(
                 sm_scale=layer.scaling,
                 skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
             )
-        elif self.dsa_decode_impl == "flashmla_kv":
+        elif dsa_impl == "flashmla_kv":
             if q_rope is not None:
                 q_all = concat_mla_absorb_q_general(q_nope, q_rope)
             return self._forward_flashmla_kv(
@@ -2321,7 +2448,7 @@ class DeepseekSparseAttnBackend(
                 metadata=metadata,
                 page_table_1=page_table_1,
             )
-        elif self.dsa_decode_impl == "tilelang":
+        elif dsa_impl == "tilelang":
             # Cat-skip (HIP-only): when caller passes q_rope=None on HIP, q_all
             # has already been set to a zero-copy view of q in the else branch
             # above and we can reuse it directly. The `not _is_hip` clause keeps
@@ -2335,7 +2462,7 @@ class DeepseekSparseAttnBackend(
                 sm_scale=layer.scaling,
                 v_head_dim=layer.v_head_dim,
             )
-        elif self.dsa_decode_impl == "fa3":
+        elif dsa_impl == "fa3":
             return self._forward_fa3(
                 q_rope=q_rope,
                 kv_cache=kv_cache,
@@ -2350,7 +2477,7 @@ class DeepseekSparseAttnBackend(
                 logit_cap=layer.logit_cap,
                 page_size=1,
             )
-        elif self.dsa_decode_impl == "aiter":
+        elif dsa_impl == "aiter":
             if q_all is None or not _is_hip:
                 q_all = torch.cat([q_nope, q_rope], dim=-1)
             return self._forward_aiter(
@@ -2363,7 +2490,7 @@ class DeepseekSparseAttnBackend(
             )
 
         else:
-            assert False, f"Unsupported {self.dsa_decode_impl = }"
+            assert False, f"Unsupported {dsa_impl = }"
 
     def _forward_fa3(
         self,
@@ -2383,13 +2510,21 @@ class DeepseekSparseAttnBackend(
         k_rope_cache = kv_cache[:, :, v_head_dim:]
         c_kv_cache = kv_cache[:, :, :v_head_dim]
         qk_rope_dim = k_rope_cache.shape[-1]
-        k_rope_cache = k_rope_cache.view(-1, page_size, 1, qk_rope_dim)
+        num_blocks = kv_cache.shape[0] // page_size
-        c_kv_cache = c_kv_cache.view(-1, page_size, 1, v_head_dim)
+        only_qv = qk_rope_dim == 0
+        if only_qv:
+            k_rope_cache = None
+        else:
+            k_rope_cache = k_rope_cache.view(num_blocks, page_size, 1, qk_rope_dim)
+        c_kv_cache = c_kv_cache.view(num_blocks, page_size, 1, v_head_dim)
+        if self.dsa_index_kpool > 1:
+            page_table = page_table.clamp(min=0)
         o = flash_attn_with_kvcache(
-            q=q_rope,
+            q=None if only_qv else q_rope,
             k_cache=k_rope_cache,
             v_cache=c_kv_cache,
             qv=q_nope,
+            only_qv=only_qv,
             page_table=page_table,
             cache_seqlens=cache_seqlens,
             cu_seqlens_q=cu_seqlens_q,
@@ -2955,6 +3090,19 @@ class DeepseekSparseAttnBackend(
     ) -> torch.Tensor:
         from sglang.kernels.ops.attention.dsa.tilelang_kernel import tilelang_sparse_fwd
 
+        # KPool appends up to index_kpool - 1 live tail tokens to the fixed
+        # index_topk columns. TileLang processes indices in 64-column blocks,
+        # so mask-pad the tail-extended table to the next complete block.
+        padding = (-page_table_1.shape[-1]) % 64
+        if padding:
+            page_table_1 = torch.cat(
+                (
+                    page_table_1,
+                    page_table_1.new_full((*page_table_1.shape[:-1], padding), -1),
+                ),
+                dim=-1,
+            )
+
         return tilelang_sparse_fwd(
             q=q_all,
             kv=kv_cache,
@@ -3154,16 +3302,23 @@ class DeepseekSparseAttnBackend(
 
         metadata = self.forward_metadata
 
-        merge_query = q_rope is not None
+        # The BF16 no-RoPE path passes a zero-width q_rope tensor.
+        merge_query = q_rope is not None and self.qk_rope_head_dim > 0
         if self.kv_cache_dtype == torch.float8_e4m3fn:
             # For FP8 path, we quantize the query and rope parts and merge them into a single tensor
             # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend
             assert q_rope is not None, "For FP8 path q_rope should not be None."
             assert k_rope is not None, "For FP8 path k_rope should not be None."
-            assert cos_sin_cache is not None, (
+            if cos_sin_cache is None:
-                "For FP8 path cos_sin_cache should not be None."
+                q, k, k_rope = mla_quantize_for_fp8_no_rope(
+                    q,
+                    q_rope,
+                    k.squeeze(1),
+                    k_rope.squeeze(1),
+                    self.kv_lora_rank,
+                    self.qk_rope_head_dim,
                 )
-
+            else:
                 rope_positions = forward_batch.positions
                 if dsa_use_prefill_cp(forward_batch):
                     if is_cp_v2_active(forward_batch):
@@ -3192,7 +3347,9 @@ class DeepseekSparseAttnBackend(
                             forward_batch, k, k_rope
                         )
                     else:
-                    k, k_rope = _all_gather_dsa_trtllm_fp8_kv(forward_batch, k, k_rope)
+                        k, k_rope = _all_gather_dsa_trtllm_fp8_kv(
+                            forward_batch, k, k_rope
+                        )
             merge_query = False
 
             # Save KV cache if requested
@@ -3244,6 +3401,12 @@ class DeepseekSparseAttnBackend(
                 topk_indices=topk_indices,
                 page_size=1,
             )
+        page_table_1, sparse_mla_top_k = self._pad_trtllm_sparse_page_table(
+            page_table_1
+        )
+        sparse_mla_top_k_lens = None
+        if self.qk_rope_head_dim == 0:
+            sparse_mla_top_k_lens = prepare_trtllm_nope_sparse_metadata(page_table_1)
 
         q_scale = 1.0
         k_scale = (
@@ -3289,10 +3452,11 @@ class DeepseekSparseAttnBackend(
             block_tables=block_tables,
             seq_lens=seq_lens,
             max_seq_len=metadata.max_seq_len_k,
-            sparse_mla_top_k=self.dsa_index_topk,
+            sparse_mla_top_k=sparse_mla_top_k,
             bmm1_scale=bmm1_scale,
             backend="trtllm-gen",
             skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
+            sparse_mla_top_k_lens=sparse_mla_top_k_lens,
             multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer,
         )
 
@@ -3327,13 +3491,11 @@ class DeepseekSparseAttnBackend(
         """
         Decide all attention prefill dispatch strategies for this batch.
         """
-        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+        # Hoisted in __init__ (import cost is per-call otherwise).
-            is_in_breakable_cuda_graph,
+        is_in_breakable_cuda_graph = self._is_in_breakable_cuda_graph
-        )
+        is_in_tc_piecewise_cuda_graph = self._is_in_tc_piecewise_cuda_graph
-        from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+        get_device_sm = self._get_device_sm
-            is_in_tc_piecewise_cuda_graph,
+        is_blackwell = self._is_blackwell
-        )
-        from sglang.srt.utils import get_device_sm, is_blackwell
 
         # Decide MHA vs MLA
         if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph():
@@ -3613,6 +3775,16 @@ class DeepseekSparseAttnMultiStepBackend:
                     precomputed.seqlens_expanded_size,
                 )
 
+                for backend, metadata in zip(
+                    self.attn_backends[:3],
+                    (metadata0, metadata1, metadata2),
+                    strict=True,
+                ):
+                    backend._update_kpool_metadata_from_precomputed(
+                        metadata, precomputed, ForwardMode.DECODE
+                    )
+                    backend.forward_metadata = metadata
+
                 # Copy remaining backends one by one (if > 3 backends)
                 for i in range(3, self.speculative_num_steps - 1):
                     self.attn_backends[
@@ -3641,7 +3813,7 @@ class DeepseekSparseAttnMultiStepBackend:
                         forward_mode=ForwardMode.DECODE,
                     )
         else:
-            # Less than 3 backends: copy to each backend individually
+            # Copy to each backend and refresh its derived metadata independently.
             for i in range(self.speculative_num_steps - 1):
                 self.attn_backends[
                     i

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

@@ -96,6 +96,9 @@ class MambaAttnBackendBase(AttentionBackend):
         self.cached_cuda_graph_decode_query_start_loc: torch.Tensor = None
         self.cached_cuda_graph_verify_query_start_loc: torch.Tensor = None
         self.conv_states_shape: tuple[int, int] = None
+        # Constant (== 1) for mamba-like backends; hoisted so the replay path
+        # skips the per-cycle method dispatch.
+        self._graph_seq_len_fill_value = self.get_cuda_graph_seq_len_fill_value()
 
     @property
     def mamba_chunk_size(self) -> int:
@@ -210,9 +213,13 @@ class MambaAttnBackendBase(AttentionBackend):
                     new_vals[inv] = next_for_valid.to(write_pos_buf.dtype)
                     write_pos_buf[uniq_slots] = new_vals
         elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
-            if forward_batch.forward_mode.is_draft_extend_v2():
+            has_extend_meta = (
-                # DRAFT_EXTEND_V2 runs only full-attn layers in the draft model;
+                forward_batch.extend_start_loc is not None
-                # skip mamba metadata.
+                and forward_batch.extend_seq_lens is not None
+            )
+            if forward_batch.forward_mode.is_draft_extend_v2() and not has_extend_meta:
+                # Draft-extend-v2 may omit linear metadata when the draft runs only
+                # full-attention layers.
                 query_start_loc = None
             elif forward_batch.forward_mode.is_target_verify():
                 ragged_layout = forward_batch.spec_info.ragged_verify_layout
@@ -623,8 +630,9 @@ class MambaAttnBackendBase(AttentionBackend):
                 num_padding = 0
             else:
                 num_padding = torch.count_nonzero(
-                    seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value()
+                    seq_lens_cpu == self._graph_seq_len_fill_value
                 )
+        num_padding = int(num_padding)
         if self._fused_state_indices_ok and self.replayssm_write_pos_list is None:
             # Single-launch fast path: mapping gather + padding sentinel + store
             # into the static buffer, plus zeroing padded req_pool_indices rows —
@@ -731,6 +739,7 @@ class MambaAttnBackendBase(AttentionBackend):
                         )
                         new_vals[inv] = next_for_valid.to(write_pos_buf.dtype)
                         write_pos_buf[uniq_slots] = new_vals
+        is_target_verify = forward_mode.is_target_verify()
         if forward_mode.is_decode_or_idle():
             if num_padding == 0:
                 self.query_start_loc_list[bs - 1].copy_(
@@ -769,8 +778,9 @@ class MambaAttnBackendBase(AttentionBackend):
                 )
         else:
             raise ValueError(f"Invalid forward mode: {forward_mode=}")
+        qsl_buf = self.query_start_loc_list[bs - 1]
 
-        if forward_mode.is_target_verify() and self.topk > 1:
+        if is_target_verify and self.topk > 1:
             if (
                 spec_info is not None
                 and getattr(spec_info, "retrieve_next_token", None) is not None
@@ -783,7 +793,7 @@ class MambaAttnBackendBase(AttentionBackend):
                     spec_info.retrieve_next_sibling
                 )
             return ForwardMetadata(
-                query_start_loc=self.query_start_loc_list[bs - 1],
+                query_start_loc=qsl_buf,
                 mamba_cache_indices=self.state_indices_list[bs - 1],
                 mamba_track_indices=track_buf,
                 retrieve_next_token=self.retrieve_next_token_list[bs - 1],
@@ -794,7 +804,7 @@ class MambaAttnBackendBase(AttentionBackend):
             )
         else:
             return ForwardMetadata(
-                query_start_loc=self.query_start_loc_list[bs - 1],
+                query_start_loc=qsl_buf,
                 mamba_cache_indices=self.state_indices_list[bs - 1],
                 mamba_track_indices=track_buf,
                 replayssm_write_pos=replayssm_write_pos,
@@ -1067,9 +1077,15 @@ class HybridLinearAttnBackend(AttentionBackend):
             and self.linear_attn_backend.supports_ragged_verify_graph
         )
 
+    @property
+    def use_mha(self) -> bool:
+        return getattr(self.full_attn_backend, "use_mha", False)
+
     @property
     def kv_cache_dtype(self):
-        return self.full_attn_backend.kv_cache_dtype
+        # Expose the full-attention backend's cache dtype because fused DSA/NSA RoPE
+        # reads it from this wrapper.
+        return getattr(self.full_attn_backend, "kv_cache_dtype", None)
 
     def _is_full_attn(
         self, layer: Optional[RadixAttention], layer_id: Optional[int] = None
@@ -1136,6 +1152,9 @@ class HybridLinearAttnBackend(AttentionBackend):
         if init is not None:
             init(forward_batch, disable_flashinfer_ragged)
 
+    def get_indexer_metadata(self, layer_id, forward_batch):
+        return self.full_attn_backend.get_indexer_metadata(layer_id, forward_batch)
+
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         for attn_backend in self.attn_backend_list:
             attn_backend.init_cuda_graph_state(max_bs, max_num_tokens)

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

@@ -151,7 +151,7 @@ class KDAKernelDispatcher:
             )
 
             cutedsl_kernel = CuteDSLKDAKernel()
-            if getattr(cutedsl_kernel, "supports_prefill", False):
+            if cutedsl_kernel.supports_prefill:
                 # SM100 chunk prefill pipeline.
                 self.extend_kernel = cutedsl_kernel
             else:
@@ -255,14 +255,17 @@ class KDAKernelDispatcher:
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,
         query_start_loc: torch.Tensor,
+        lower_bound: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
-        kernel = self.decode_kernel
+        if lower_bound is not None and not isinstance(
-        if kwargs.get("lower_bound") is not None and not getattr(
+            self.decode_kernel, TritonKDAKernel
-            kernel, "supports_safe_gate", True
         ):
-            kernel = self.triton_kernel
+            raise NotImplementedError(
-        return kernel.decode(
+                f"lower_bound (safe gate) is only supported by TritonKDAKernel; "
+                f"got {self.decode_kernel.__class__.__name__}."
+            )
+        return self.decode_kernel.decode(
             q,
             k,
             v,
@@ -273,6 +276,7 @@ class KDAKernelDispatcher:
             ssm_states=ssm_states,
             cache_indices=cache_indices,
             query_start_loc=query_start_loc,
+            lower_bound=lower_bound,
             **kwargs,
         )
 
@@ -292,13 +296,20 @@ class KDAKernelDispatcher:
         intermediate_states_buffer: torch.Tensor,
         intermediate_state_indices: torch.Tensor,
         cache_steps: int,
-        retrieve_parent_token: torch.Tensor,
+        retrieve_parent_token: Optional[torch.Tensor],
         lower_bound: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         """MTP / speculative-decode verify, routed to ``self.verify_kernel``
         (FlashInfer decode -> recurrent_kda; Triton / CuTe DSL decode -> the Triton
         fused KDA verify)."""
+        if lower_bound is not None and not isinstance(
+            self.verify_kernel, TritonKDAKernel
+        ):
+            raise NotImplementedError(
+                "lower_bound (safe gate) target verify is only supported by "
+                f"TritonKDAKernel; got {self.verify_kernel.__class__.__name__}."
+            )
         return self.verify_kernel.target_verify(
             A_log=A_log,
             dt_bias=dt_bias,
@@ -380,10 +391,8 @@ class KDAAttnBackend(MambaAttnBackendBase):
     # to its dense layout, so ragged verify graphs are supported.
     supports_ragged_verify_graph: bool = True
 
-    # Read by decide_needs_cpu_seq_lens. Decode/verify metadata is GPU-only
+    # KDA gets graph padding explicitly and never uses ReplaySSM's host-seqlen
-    # (graph replay already passes seq_lens_cpu=None), extend reads
+    # force-flush path.
-    # extend_seq_lens_cpu from schedule, mamba track indices rebuild from req
-    # objects, and the replayssm seq_lens_cpu force-flush is GDN-only.
     needs_cpu_seq_lens: bool = False
 
     def __init__(self, model_runner: ModelRunner):
@@ -721,9 +730,11 @@ class KDAAttnBackend(MambaAttnBackendBase):
             conv_state_indices=cache_indices,
         )
 
-        # The packed kernel assumes one token per request. Assert the dispatch
+        # The packed kernel assumes one token per request.
-        # invariant before taking the fused path.
+        if (
-        if self.kernel_dispatcher.supports_packed_decode:
+            self.kernel_dispatcher.supports_packed_decode
+            and getattr(layer, "lower_bound", None) is None
+        ):
             assert qkv.shape[0] == cache_indices.shape[0], (
                 "KDA packed decode requires one token per sequence (T=1): "
                 f"got {qkv.shape[0]} tokens for {cache_indices.shape[0]} requests."
@@ -805,6 +816,13 @@ class KDAAttnBackend(MambaAttnBackendBase):
             )
         has_initial_state = forward_batch.extend_prefix_lens > 0
 
+        physical_num_tokens = mixed_qkv.shape[0]
+        logical_num_tokens = int(query_start_loc[-1])
+        if logical_num_tokens < physical_num_tokens:
+            mixed_qkv = mixed_qkv[:logical_num_tokens]
+            a = a[:, :logical_num_tokens]
+            b = b[:, :logical_num_tokens]
+
         if self.forward_metadata.has_mamba_track_mask:
             # Snapshot the conv sliding window at the last track-aligned chunk
             # boundary into the ping-pong track slots (the prefix-cache restore
@@ -815,55 +833,29 @@ class KDAAttnBackend(MambaAttnBackendBase):
                 self.forward_metadata.conv_states_mask_indices
             ] = mixed_qkv[self.forward_metadata.track_conv_indices]
 
-        splits = [layer.q_dim, layer.k_dim, layer.v_dim]
+        # Depthwise conv is channel-independent, so one packed call over the
-        q, k, v = mixed_qkv.transpose(0, 1).split(splits, dim=0)
+        # full qkv width matches the decode path and saves two kernel launches.
-        q_conv_weight, k_conv_weight, v_conv_weight = layer.conv_weights.split(
+        qkv = causal_conv1d_fn(
-            splits, dim=0
+            mixed_qkv.transpose(0, 1),
-        )
+            layer.conv_weights,
-        q_conv_state, k_conv_state, v_conv_state = conv_states.split(splits, dim=-2)
+            layer.bias,
-        if layer.bias is not None:
-            q_bias, k_bias, v_bias = layer.bias.split(splits, dim=0)
-        else:
-            q_bias, k_bias, v_bias = None, None, None
-
-        q = causal_conv1d_fn(
-            q,
-            q_conv_weight,
-            q_bias,
             activation="silu",
-            conv_states=q_conv_state,
+            conv_states=conv_states,
-            has_initial_state=has_initial_state,
-            cache_indices=cache_indices,
-            query_start_loc=query_start_loc,
-            seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
-        ).transpose(0, 1)
-        k = causal_conv1d_fn(
-            k,
-            k_conv_weight,
-            k_bias,
-            activation="silu",
-            conv_states=k_conv_state,
-            has_initial_state=has_initial_state,
-            cache_indices=cache_indices,
-            query_start_loc=query_start_loc,
-            seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
-        ).transpose(0, 1)
-        v = causal_conv1d_fn(
-            v,
-            v_conv_weight,
-            v_bias,
-            activation="silu",
-            conv_states=v_conv_state,
             has_initial_state=has_initial_state,
             cache_indices=cache_indices,
             query_start_loc=query_start_loc,
             seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
         ).transpose(0, 1)
+        q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)
 
         q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
         k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
         v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
 
+        gate_was_flat = a.ndim == 3
+        if gate_was_flat:
+            a = a.unflatten(-1, (-1, layer.head_k_dim))
+
         track_ssm = self.forward_metadata.has_mamba_track_mask
         core_attn_out = self.kernel_dispatcher.extend(
             q=q,
@@ -877,6 +869,7 @@ class KDAAttnBackend(MambaAttnBackendBase):
             A_log=layer.A_log,
             dt_bias=layer.dt_bias,
             lower_bound=layer.lower_bound,
+            beta_is_raw=gate_was_flat,
             extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
             # draft_extend_v2 must stay rollback-able, so kernels that commit state
             # in place (e.g. FlashKDA) must not run for it.
@@ -898,6 +891,13 @@ class KDAAttnBackend(MambaAttnBackendBase):
                 forward_batch, h, ssm_states, self.forward_metadata
             )
 
+        if logical_num_tokens < physical_num_tokens:
+            pad = core_attn_out.new_zeros(
+                (1, physical_num_tokens - logical_num_tokens)
+                + tuple(core_attn_out.shape[2:])
+            )
+            core_attn_out = torch.cat((core_attn_out, pad), dim=1)
+
         if (
             self.accept_lens_pool is not None
             and not forward_batch.forward_mode.is_draft_extend_v2()

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

@@ -683,6 +683,10 @@ class TritonAttnBackend(AttentionBackend):
                     window_num_kv_splits=None,
                     window_kv_offsets=None,
                     swa_attn_logits=self.cuda_graph_swa_attn_logits,
+                    lean_Mp=self.cuda_graph_lean_Mp,
+                    lean_Lp=self.cuda_graph_lean_Lp,
+                    lean_Op=self.cuda_graph_lean_Op,
+                    lean_locks=self.cuda_graph_lean_locks,
                 )
                 return
 

python/sglang/srt/layers/communicator.py

@@ -250,12 +250,17 @@ class AttentionInputs:
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
         qkv_latent_func: Callable,
+        *,
+        is_pre_gathered: bool = False,
     ):
         self.hidden_states_local = hidden_states
         self.forward_batch = forward_batch
         self.qkv_latent_func = qkv_latent_func
         self.hidden_states_ = None
         self.qkv_latent_ = None
+        # When True, hidden_states_local is already attn_tp-gathered upstream
+        # (e.g. by MHC's prepare_attn for DSA). fetch_* must NOT gather again.
+        self.is_pre_gathered = is_pre_gathered
 
     def tp_all_gather_hidden_states(self, hidden_states, forward_batch):
         total_tokens = forward_batch.input_ids.shape[0]
@@ -270,7 +275,7 @@ class AttentionInputs:
         self.qkv_latent_ = self.qkv_latent_func(
             self.hidden_states_local, self.forward_batch
         )
-        if get_attn_tp_context().input_scattered:
+        if get_attn_tp_context().input_scattered and not self.is_pre_gathered:
             self.qkv_latent_ = self.tp_all_gather_hidden_states(
                 self.qkv_latent_, self.forward_batch
             )
@@ -280,7 +285,7 @@ class AttentionInputs:
         if self.hidden_states_ is not None:
             return self.hidden_states_
         self.hidden_states_ = self.hidden_states_local
-        if get_attn_tp_context().input_scattered:
+        if get_attn_tp_context().input_scattered and not self.is_pre_gathered:
             self.hidden_states_ = self.tp_all_gather_hidden_states(
                 self.hidden_states_, self.forward_batch
             )
@@ -292,13 +297,15 @@ class AttnTpContext:
         self.allow_input_scattered = False
         self.is_dsa = False
 
-    def init_context(self, q_lora_rank, is_dsa):
+    def init_context(self, q_lora_rank, is_dsa, is_mhc=False):
+        # Only MHC pre-gathers hidden states before DSA attention, so non-MHC DSA
+        # cannot use scattered inputs.
         self.is_dsa = is_dsa
         self.allow_input_scattered = (
             get_parallel().enable_attn_tp_input_scattered
             and (_is_cuda or _is_npu)
             and q_lora_rank is not None
-            and not is_dsa
+            and (is_mhc or not is_dsa)
             and get_parallel().tp_size > 1
             and not is_dp_attention_enabled()
             and get_moe_a2a_backend().is_none()
@@ -330,6 +337,11 @@ class AttnTpContext:
     def set_attn_inputs(self, attn_inputs: AttentionInputs):
         get_forward().set("attn_inputs", attn_inputs)
 
+    def set_hidden_states_local(self, hidden_states: torch.Tensor) -> None:
+        attn_inputs = get_forward().attn_inputs
+        if attn_inputs is not None:
+            attn_inputs.hidden_states_local = hidden_states
+
     def fetch_qkv_latent(self):
         attn_inputs = get_forward().attn_inputs
         assert attn_inputs is not None
@@ -486,6 +498,26 @@ def enable_dwdp():
     return get_parallel().dwdp_size > 1
 
 
+def tp_reduce_scatter(
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor],
+    context: "CommunicateContext",
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Module-level so MHC communicators can reuse it without a
+    ``LayerCommunicator`` instance."""
+    if hidden_states.shape[0] == 0:
+        return hidden_states, hidden_states
+    assert hidden_states.shape[0] % context.tp_size == 0, (
+        f"Expected total tokens {hidden_states.shape[0]} % tp_size {context.tp_size} to be 0"
+    )
+    local_tokens = hidden_states.shape[0] // context.tp_size
+    output = hidden_states.new_empty(local_tokens, *hidden_states.shape[1:])
+    get_tp_group().reduce_scatter_tensor(output, hidden_states)
+    if residual is not None:
+        residual = residual.tensor_split(context.tp_size)[context.tp_rank]
+    return output, residual
+
+
 class LayerCommunicator:
     def __init__(
         self,
@@ -826,19 +858,7 @@ class LayerCommunicator:
         hidden_states: torch.Tensor,
         residual: torch.Tensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        if hidden_states.shape[0] == 0:
+        return tp_reduce_scatter(hidden_states, residual, self._context)
-            return hidden_states, hidden_states
-        assert hidden_states.shape[0] % self._context.tp_size == 0, (
-            f"Expected total tokens {hidden_states.shape[0]} % tp_size {self._context.tp_size} to be 0"
-        )
-        local_tokens = hidden_states.shape[0] // self._context.tp_size
-        output = hidden_states.new_empty(local_tokens, *hidden_states.shape[1:])
-        get_tp_group().reduce_scatter_tensor(output, hidden_states)
-        if residual is not None:
-            residual = residual.tensor_split(self._context.tp_size)[
-                self._context.tp_rank
-            ]
-        return output, residual
 
     def prepare_mlp(
         self,
@@ -862,6 +882,13 @@ class LayerCommunicator:
             context=self._context,
         )
 
+    def maybe_prefetch_next_full_attention_kv(
+        self,
+        forward_batch: ForwardBatch,
+        next_full_attention_layer_id: Optional[int],
+    ) -> None:
+        return
+
     def postprocess_layer(
         self,
         hidden_states: torch.Tensor,

python/sglang/srt/layers/communicator_mhc.py

@@ -0,0 +1,557 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Optional
+
+import torch
+
+from sglang.kernels.ops.layernorm.mhc import hc_contract, hc_expand
+from sglang.srt.distributed import get_tp_group
+from sglang.srt.distributed.communication_op import (
+    attention_tensor_model_parallel_all_reduce,
+)
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
+from sglang.srt.layers.communicator import (
+    AttentionInputs,
+    CommunicateContext,
+    CommunicateSimpleFn,
+    CommunicateSummableTensorPairFn,
+    CommunicateWithAllReduceAndLayerNormFn,
+    LayerCommunicator,
+    LayerScatterModes,
+    ScatterMode,
+    get_attn_tp_context,
+    tp_reduce_scatter,
+)
+from sglang.srt.layers.dp_attention import (
+    attn_tp_all_gather_into_tensor,
+    attn_tp_reduce_scatter_tensor,
+    dp_gather_replicate,
+    dp_reduce_scatter_tensor,
+    dp_scatter,
+    get_dp_global_num_tokens,
+    get_global_dp_buffer,
+    get_local_dp_buffer_mhc,
+    is_allocation_symmetric,
+)
+from sglang.srt.layers.moe import should_use_dp_reduce_scatterv
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+def tp_all_gather_hidden_states(hidden_states, forward_batch):
+    assert get_attn_tp_context().input_scattered, (
+        "Input scattered guarantees same num tokens in TP group."
+    )
+    total_tokens = forward_batch.input_ids.shape[0]
+    output = hidden_states.new_empty((total_tokens, hidden_states.shape[-1]))
+    get_tp_group().all_gather_into_tensor(output, hidden_states)
+
+    return output
+
+
+@dataclass
+class MHCState:
+    """Parameters belong to the owning layer; this state only holds scratch
+    shared across communication stages."""
+
+    hc_mult: int
+    hc_attn_pre: Callable
+    hc_ffn_pre: Callable
+    hc_post: Callable
+    h_res: Optional[torch.Tensor] = None
+    h_post: Optional[torch.Tensor] = None
+
+    @staticmethod
+    def _resolve_out_norm(out_norm):
+        if out_norm is None:
+            return None, None
+        return out_norm.weight.data, out_norm.variance_epsilon
+
+    def attn_split(self, hidden_states, out_norm: Optional[torch.nn.Module] = None):
+        residual = hidden_states
+        out_norm_weight, out_norm_eps = self._resolve_out_norm(out_norm)
+        hidden_states, self.h_res, self.h_post, norm_fused = self.hc_attn_pre(
+            hidden_states, out_norm_weight, out_norm_eps
+        )
+        if out_norm is not None and not norm_fused and hidden_states.shape[0] != 0:
+            hidden_states = out_norm(hidden_states)
+        return hidden_states, residual
+
+    def attn_to_mlp(
+        self, hidden_states, residual, out_norm: Optional[torch.nn.Module] = None
+    ):
+        hidden_states = self.hc_post(hidden_states, residual, self.h_res, self.h_post)
+        residual = hidden_states
+        out_norm_weight, out_norm_eps = self._resolve_out_norm(out_norm)
+        hidden_states, self.h_res, self.h_post, norm_fused = self.hc_ffn_pre(
+            hidden_states, out_norm_weight, out_norm_eps
+        )
+        if out_norm is not None and not norm_fused and hidden_states.shape[0] != 0:
+            hidden_states = out_norm(hidden_states)
+        return hidden_states, residual
+
+    def mlp_combine(self, hidden_states, residual):
+        return self.hc_post(hidden_states, residual, self.h_res, self.h_post)
+
+    def reset_aux(self):
+        self.h_res = None
+        self.h_post = None
+
+
+class MHCCommunicateWithAllReduceAndLayerNormFn(CommunicateWithAllReduceAndLayerNormFn):
+    @staticmethod
+    def get_fn(
+        hidden_states_input_mode: ScatterMode,
+        residual_input_mode: ScatterMode,
+        hidden_states_output_mode: ScatterMode,
+        residual_output_mode: ScatterMode,
+        context: CommunicateContext,
+    ):
+        fn = CommunicateWithAllReduceAndLayerNormFn.get_fn(
+            hidden_states_input_mode,
+            residual_input_mode,
+            hidden_states_output_mode,
+            residual_output_mode,
+            context,
+        )
+        replacements = {
+            CommunicateWithAllReduceAndLayerNormFn._simple: MHCCommunicateWithAllReduceAndLayerNormFn._simple,
+            CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual: MHCCommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual,
+            CommunicateWithAllReduceAndLayerNormFn._scatter_hidden_states_and_residual: MHCCommunicateWithAllReduceAndLayerNormFn._scatter_hidden_states_and_residual,
+        }
+        if isinstance(fn, partial):
+            return partial(
+                replacements.get(fn.func, fn.func),
+                *fn.args,
+                **(fn.keywords or {}),
+            )
+        return replacements.get(fn, fn)
+
+    @staticmethod
+    def _scatter_hidden_states_and_residual(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        layernorm: torch.nn.Module,
+        context: CommunicateContext,
+        *,
+        residual_input_mode,
+        mhc: MHCState,
+    ):
+        input_hidden_states = hidden_states
+        hidden_states = hidden_states.tensor_split(context.attn_tp_size)[
+            context.attn_tp_rank
+        ]
+        attn_tp_reduce_scatter_tensor(hidden_states, input_hidden_states)
+        if residual_input_mode == ScatterMode.TP_ATTN_FULL:
+            residual = residual.tensor_split(context.attn_tp_size)[context.attn_tp_rank]
+            mhc.h_res = mhc.h_res.tensor_split(context.attn_tp_size)[
+                context.attn_tp_rank
+            ]
+            mhc.h_post = mhc.h_post.tensor_split(context.attn_tp_size)[
+                context.attn_tp_rank
+            ]
+
+        hidden_states, residual = mhc.attn_to_mlp(
+            hidden_states, residual, out_norm=layernorm
+        )
+        return hidden_states, residual
+
+    @staticmethod
+    def _simple(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        layernorm: torch.nn.Module,
+        context: CommunicateContext,
+        *,
+        mhc: MHCState,
+    ):
+        hidden_states, residual = mhc.attn_to_mlp(
+            hidden_states, residual, out_norm=layernorm
+        )
+        return hidden_states, residual
+
+    @staticmethod
+    def _tp_all_reduce_with_scattered_residual(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        layernorm: torch.nn.Module,
+        context: CommunicateContext,
+        *,
+        mhc: MHCState,
+    ):
+        if hidden_states.shape[0] == 0:
+            return hidden_states, hidden_states
+
+        scatter_states = hidden_states.tensor_split(context.tp_size)[context.tp_rank]
+        get_tp_group().reduce_scatter_tensor(scatter_states, hidden_states)
+
+        scatter_states, residual = mhc.attn_to_mlp(
+            scatter_states, residual, out_norm=layernorm
+        )
+
+        attn_tp_all_gather_into_tensor(hidden_states, scatter_states)
+
+        return hidden_states, residual
+
+    @staticmethod
+    def _gather_hidden_states_and_residual(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        layernorm: torch.nn.Module,
+        context: CommunicateContext,
+        *,
+        residual_input_mode,
+        mhc: MHCState,
+    ):
+        if get_attn_tp_context().input_scattered:
+            return MHCCommunicateWithAllReduceAndLayerNormFn._tp_all_reduce_with_scattered_residual(
+                hidden_states,
+                residual,
+                layernorm,
+                context,
+                mhc=mhc,
+            )
+
+        if residual_input_mode == ScatterMode.SCATTERED and context.attn_tp_size > 1:
+            raise NotImplementedError(
+                "Unsupported: h_res/h_post allgather not implemented."
+            )
+
+        hidden_states = attention_tensor_model_parallel_all_reduce(hidden_states)
+        if context.attn_dp_size != 1:
+            if hidden_states.shape[0] != 0:
+                with use_symmetric_memory(
+                    get_tp_group(),
+                    disabled=not is_allocation_symmetric(),
+                ):
+                    hidden_states, residual = mhc.attn_to_mlp(
+                        hidden_states, residual, out_norm=layernorm
+                    )
+            else:
+                hidden_states, residual = mhc.attn_to_mlp(hidden_states, residual)
+
+            hidden_states, local_hidden_states = (
+                get_global_dp_buffer(get_tp_group()),
+                hidden_states,
+            )
+            dp_gather_replicate(hidden_states, local_hidden_states, forward_batch)
+        else:
+            hidden_states, residual = mhc.attn_to_mlp(
+                hidden_states, residual, out_norm=layernorm
+            )
+        return hidden_states, residual
+
+
+class MHCCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
+    @staticmethod
+    def get_fn(
+        hidden_states_input_mode: ScatterMode,
+        residual_input_mode: ScatterMode,
+        output_mode: ScatterMode,
+        context: CommunicateContext,
+    ):
+        fn = CommunicateSummableTensorPairFn.get_fn(
+            hidden_states_input_mode,
+            residual_input_mode,
+            output_mode,
+            context,
+        )
+        replacements = {
+            CommunicateSummableTensorPairFn._trivial: MHCCommunicateSummableTensorPairFn._trivial,
+            CommunicateSummableTensorPairFn._scatter_hidden_states: MHCCommunicateSummableTensorPairFn._scatter_hidden_states,
+            CommunicateSummableTensorPairFn._gather: MHCCommunicateSummableTensorPairFn._gather,
+            CommunicateSummableTensorPairFn._scatter: MHCCommunicateSummableTensorPairFn._scatter,
+        }
+        return replacements.get(fn, fn)
+
+    @staticmethod
+    def _trivial(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        context: CommunicateContext,
+        *,
+        mhc: MHCState,
+        is_last_layer: bool,
+        **kwargs,
+    ):
+        if get_attn_tp_context().input_scattered:
+            hidden_states, _ = tp_reduce_scatter(hidden_states, None, context)
+
+        hidden_states = mhc.mlp_combine(hidden_states, residual)
+        if not is_last_layer:
+            return hidden_states, None
+
+        hidden_states = hc_contract(hidden_states, mhc.hc_mult)
+        if get_attn_tp_context().input_scattered:
+            local_states = hidden_states
+            hidden_states = local_states.new_empty(
+                local_states.shape[0] * context.tp_size, *local_states.shape[1:]
+            )
+            get_tp_group().all_gather_into_tensor(hidden_states, local_states)
+
+        return hidden_states, None
+
+    @staticmethod
+    def _scatter_hidden_states(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        context: CommunicateContext,
+        allow_reduce_scatter: bool = False,
+        *,
+        mhc: MHCState,
+        is_last_layer: bool,
+        **kwargs,
+    ):
+        hidden_states, global_hidden_states = (
+            get_local_dp_buffer_mhc(get_tp_group(), 1),
+            hidden_states,
+        )
+        # MoE skips its post-expert all-reduce with reduce_scatterv, so this
+        # scatter must reduce while combining local-expert partial sums.
+        if should_use_dp_reduce_scatterv():
+            get_tp_group().reduce_scatterv(
+                global_hidden_states,
+                output=hidden_states,
+                sizes=get_dp_global_num_tokens(),
+            )
+        elif allow_reduce_scatter and forward_batch.dp_padding_mode.is_max_len():
+            dp_reduce_scatter_tensor(hidden_states, global_hidden_states)
+        else:
+            dp_scatter(hidden_states, global_hidden_states, forward_batch)
+
+        hidden_states = mhc.mlp_combine(hidden_states, residual)
+        if not is_last_layer:
+            return hidden_states, None
+
+        hidden_states = hc_contract(hidden_states, mhc.hc_mult)
+        return hidden_states, None
+
+    @staticmethod
+    def _gather(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        context: CommunicateContext,
+        *,
+        mhc: MHCState,
+        is_last_layer: bool,
+        **kwargs,
+    ):
+        hidden_states = mhc.mlp_combine(hidden_states, residual)
+        if is_last_layer:
+            hidden_states = hc_contract(hidden_states, mhc.hc_mult)
+
+        hidden_states, local_hidden_states = (
+            get_local_dp_buffer_mhc(
+                get_tp_group(), 1 if is_last_layer else mhc.hc_mult
+            ),
+            hidden_states,
+        )
+
+        attn_tp_all_gather_into_tensor(hidden_states, local_hidden_states)
+        return hidden_states, None
+
+    @staticmethod
+    def _scatter(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        context: CommunicateContext,
+        *,
+        mhc: MHCState,
+        is_last_layer: bool,
+        **kwargs,
+    ):
+        hidden_states = hidden_states.tensor_split(context.attn_tp_size)[
+            context.attn_tp_rank
+        ]
+        residual = residual.tensor_split(context.attn_tp_size)[context.attn_tp_rank]
+
+        hidden_states = mhc.mlp_combine(hidden_states, residual)
+
+        return hidden_states, None
+
+
+class MHCLayerCommunicator(LayerCommunicator):
+    def __init__(
+        self,
+        layer_scatter_modes: LayerScatterModes,
+        input_layernorm: torch.nn.Module,
+        post_attention_layernorm: torch.nn.Module,
+        allow_reduce_scatter: bool = False,
+        is_last_layer: bool = False,
+        qkv_latent_func: Optional[Callable] = None,
+        *,
+        is_first_layer: bool,
+        hc_mult: int,
+        hc_attn_pre: Callable,
+        hc_ffn_pre: Callable,
+        hc_post: Callable,
+    ):
+        self.is_first_layer = is_first_layer
+        self.mhc = MHCState(
+            hc_mult=hc_mult,
+            hc_attn_pre=hc_attn_pre,
+            hc_ffn_pre=hc_ffn_pre,
+            hc_post=hc_post,
+        )
+
+        super().__init__(
+            layer_scatter_modes,
+            input_layernorm,
+            post_attention_layernorm,
+            allow_reduce_scatter,
+            is_last_layer,
+            qkv_latent_func,
+        )
+
+    def _post_init_communicate(self):
+        # Base MOE_FULL callables do not accept ``mhc``, so reject this
+        # combination at construction.
+        if self.layer_scatter_modes.mlp_mode == ScatterMode.MOE_FULL:
+            raise NotImplementedError(
+                "MHCLayerCommunicator does not support MOE_FULL "
+                "(moe_dp_size < attention_context_parallel_size). Increase "
+                "moe_dp_size to match attention_context_parallel_size."
+            )
+        self._communicate_simple_fn = CommunicateSimpleFn.get_fn(
+            input_mode=self.layer_scatter_modes.layer_input_mode,
+            output_mode=self.layer_scatter_modes.attn_mode,
+            context=self._context,
+        )
+        self._communicate_with_all_reduce_and_layer_norm_fn = (
+            MHCCommunicateWithAllReduceAndLayerNormFn.get_fn(
+                hidden_states_input_mode=self.layer_scatter_modes.attn_mode,
+                residual_input_mode=self.layer_scatter_modes.layer_input_mode,
+                hidden_states_output_mode=self.layer_scatter_modes.mlp_mode,
+                residual_output_mode=self.layer_scatter_modes.middle_residual_mode,
+                context=self._context,
+            )
+        )
+        self._communicate_summable_tensor_pair_fn = (
+            MHCCommunicateSummableTensorPairFn.get_fn(
+                hidden_states_input_mode=self.layer_scatter_modes.mlp_mode,
+                residual_input_mode=self.layer_scatter_modes.middle_residual_mode,
+                output_mode=self.layer_scatter_modes.layer_output_mode,
+                context=self._context,
+            )
+        )
+
+    def prepare_attn(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ):
+        if self.is_first_layer:
+            if get_attn_tp_context().input_scattered:
+                hidden_states, _ = tp_reduce_scatter(
+                    hidden_states,
+                    None,
+                    self._context,
+                )
+            hidden_states = hc_expand(hidden_states, self.mhc.hc_mult)
+
+        hidden_states, residual = self.mhc.attn_split(
+            hidden_states, out_norm=self.input_layernorm
+        )
+
+        hidden_states = self._communicate_simple_fn(
+            hidden_states=hidden_states,
+            forward_batch=forward_batch,
+            context=self._context,
+        )
+
+        # DSA and attention without a QKV hook consume full hidden states, so
+        # gather them before attention.
+        ctx = get_attn_tp_context()
+        dsa_pre_gather = ctx.input_scattered and ctx.is_dsa
+        no_qkv_latent_pre_gather = ctx.input_scattered and self.qkv_latent_func is None
+        if dsa_pre_gather or no_qkv_latent_pre_gather:
+            hidden_states = tp_all_gather_hidden_states(hidden_states, forward_batch)
+
+        if self.qkv_latent_func is not None:
+            attn_inputs = AttentionInputs(
+                hidden_states,
+                forward_batch,
+                self.qkv_latent_func,
+                is_pre_gathered=dsa_pre_gather,
+            )
+            ctx.set_attn_inputs(attn_inputs)
+
+        return hidden_states, residual
+
+    def prepare_mlp(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        cache=None,
+    ):
+        if cache is not None:
+            self._context.cache = cache
+
+        hidden_states, residual = self._communicate_with_all_reduce_and_layer_norm_fn(
+            hidden_states=hidden_states,
+            residual=residual,
+            forward_batch=forward_batch,
+            layernorm=self.post_attention_layernorm,
+            context=self._context,
+            mhc=self.mhc,
+        )
+
+        return hidden_states, residual
+
+    def postprocess_layer(self, hidden_states, residual, forward_batch):
+        hidden_states, residual = self._communicate_summable_tensor_pair_fn(
+            hidden_states=hidden_states,
+            residual=residual,
+            forward_batch=forward_batch,
+            context=self._context,
+            allow_reduce_scatter=self.allow_reduce_scatter,
+            mhc=self.mhc,
+            is_last_layer=self.is_last_layer,
+        )
+        self.mhc.reset_aux()
+
+        return hidden_states, residual
+
+    def should_fuse_mlp_allreduce_with_next_layer(self, forward_batch):
+        return False
+
+    def should_use_reduce_scatter(self, forward_batch: ForwardBatch):
+        if not self.allow_reduce_scatter:
+            return False
+        if (
+            self._communicate_summable_tensor_pair_fn
+            is MHCCommunicateSummableTensorPairFn._scatter_hidden_states
+        ):
+            # reduce_scatterv already combines expert outputs; returning False
+            # would make RowParallelLinear perform an extra all-reduce.
+            if should_use_dp_reduce_scatterv():
+                return True
+            if forward_batch.dp_padding_mode.is_max_len():
+                return True
+
+        if get_attn_tp_context().input_scattered:
+            return True
+        return False

python/sglang/srt/layers/dp_attention.py

@@ -213,6 +213,21 @@ class _DpGatheredBufferWrapper:
             )
         return buffer
 
+    @classmethod
+    def get_local_dp_buffer_mhc(
+        cls, group: GroupCoordinator, n: int = 1
+    ) -> torch.Tensor:
+        from sglang.srt.runtime_context import get_flags
+
+        dp = get_flags().dp
+        with use_symmetric_memory(group, disabled=not cls._dp_max_padding):
+            buffer = torch.empty(
+                (cls._local_dp_buffer_len, dp.buffer_hidden_size * n),
+                dtype=dp.buffer_dtype,
+                device=dp.buffer_device,
+            )
+        return buffer
+
     @classmethod
     def get_global_dp_buffer_len(cls) -> int:
         return cls._global_dp_buffer_len
@@ -277,6 +292,10 @@ def get_local_dp_buffer(group: GroupCoordinator) -> torch.Tensor:
     return _DpGatheredBufferWrapper.get_local_dp_buffer(group=group)
 
 
+def get_local_dp_buffer_mhc(group: GroupCoordinator, n: int = 1) -> torch.Tensor:
+    return _DpGatheredBufferWrapper.get_local_dp_buffer_mhc(group=group, n=n)
+
+
 def get_global_dp_buffer_len() -> int:
     return _DpGatheredBufferWrapper.get_global_dp_buffer_len()
 

python/sglang/srt/layers/moe/moe_runner/deep_gemm.py

@@ -1650,6 +1650,7 @@ def _situ_mul_quant_contig_kernel(
 def _apply_swiglu_limit(
     gateup_output: torch.Tensor, swiglu_limit: float
 ) -> torch.Tensor:
+    """Clamp the contiguous runner's owned GEMM workspace in place."""
     assert swiglu_limit == 10
 
     num_tokens, hidden_size_x2 = gateup_output.shape
@@ -1659,12 +1660,12 @@ def _apply_swiglu_limit(
     assert gate.shape == (num_tokens, hidden_size_x2 // 2)
     assert up.shape == (num_tokens, hidden_size_x2 // 2)
 
-    up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit)
+    # Both halves are views of a fresh GEMM output. Avoid separate clamped
-    gate = torch.clamp(gate, max=swiglu_limit)
+    # copies and their concatenation: large compact prefills need that
-
+    # headroom for the activation and down-projection workspaces.
-    out = torch.cat([gate, up], dim=-1)
+    up.clamp_(min=-swiglu_limit, max=swiglu_limit)
-    assert out.shape == (num_tokens, hidden_size_x2)
+    gate.clamp_(max=swiglu_limit)
-    return out
+    return gateup_output
 
 
 @register_pre_permute("deepep_v2", "deep_gemm")

python/sglang/srt/managers/schedule_batch.py

@@ -1877,7 +1877,9 @@ class Req(ReqDllmMixin):
         )
         self.kv.retraction_backup = RetractionBackup(
             cpu_tensors=token_to_kv_pool_allocator.get_cpu_copy(
-                token_indices, mamba_indices=self.kv.mamba_pool_idx
+                token_indices,
+                mamba_indices=self.kv.mamba_pool_idx,
+                req_pool_index=self.kv.req_pool_idx,
             ),
             mamba_cpu=(
                 mamba_pool.get_cpu_copy(self.kv.mamba_pool_idx.unsqueeze(0))
@@ -1901,6 +1903,7 @@ class Req(ReqDllmMixin):
             self.kv.retraction_backup.cpu_tensors,
             token_indices,
             mamba_indices=self.kv.mamba_pool_idx,
+            req_pool_index=self.kv.req_pool_idx,
         )
         self.kv.retraction_backup = None
 

python/sglang/srt/managers/scheduler.py

@@ -16,6 +16,7 @@
 import dataclasses
 import faulthandler
 import logging
+import math
 import os
 import signal
 import sys
@@ -681,6 +682,7 @@ class Scheduler(
 
         # Init prefill kv split size when deterministic inference is enabled with various attention backends
         self.init_deterministic_inference_config()
+        self.init_dsa_kpool_truncation_align()
 
         self.init_weight_updater()
 
@@ -1118,6 +1120,16 @@ class Scheduler(
         if self.server_args.is_startup_weight_load_overlap:
             self.tp_worker.finalize_startup_weight_load()
 
+        # Adaptive/speculative graphs and post-capture KV sizing can consume
+        # the headroom seen by the initial DeepGEMM layout budget. Refresh it
+        # after these allocations, before elastic EP rejoins healthy ranks
+        # that do not participate in this startup collective.
+        from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import (
+            refresh_deep_gemm_layout_memory_budget,
+        )
+
+        refresh_deep_gemm_layout_memory_budget(model_runner, only_if_initialized=True)
+
         if (
             get_exec().moe.elastic_ep_backend is not None
             and get_exec().moe.ep_join_mode == "recover"
@@ -1672,6 +1684,28 @@ class Scheduler(
             get_int_env_var(env_var, default_size) if env_var else None
         )
 
+    def init_dsa_kpool_truncation_align(self):
+        """Kpool compress-write asserts chunked extends start on pool boundaries.
+        Use the LCM to preserve any existing deterministic-inference alignment."""
+        from sglang.srt.configs.model_config import (
+            get_dsa_index_kpool,
+            is_deepseek_dsa,
+        )
+
+        if not is_deepseek_dsa(self.model_config.hf_config):
+            return
+
+        dsa_index_kpool = get_dsa_index_kpool(self.model_config.hf_config)
+        if dsa_index_kpool <= 1:
+            return
+
+        if self.truncation_align_size is None:
+            self.truncation_align_size = dsa_index_kpool
+        else:
+            self.truncation_align_size = math.lcm(
+                self.truncation_align_size, dsa_index_kpool
+            )
+
     def init_request_dispatcher(self):
         self._request_dispatcher = TypeBasedDispatcher(
             [

python/sglang/srt/mem_cache/allocator/base.py

@@ -121,10 +121,12 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
         virtual-id pools must override."""
         return kv_indices
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         raise NotImplementedError()
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         raise NotImplementedError()
 
     def alloc_extend(self, *args, **kwargs):

python/sglang/srt/mem_cache/allocator/paged.py

@@ -340,10 +340,19 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.staged_pages: list[torch.Tensor] = []
         self.num_staged_pages = 0
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
-        return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices)
+        return self._kvcache.get_cpu_copy(
+            indices,
+            mamba_indices=mamba_indices,
+            req_pool_index=req_pool_index,
+        )
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         return self._kvcache.load_cpu_copy(
-            kv_cache_cpu, indices, mamba_indices=mamba_indices
+            kv_cache_cpu,
+            indices,
+            mamba_indices=mamba_indices,
+            req_pool_index=req_pool_index,
         )

python/sglang/srt/mem_cache/allocator/swa.py

@@ -471,12 +471,21 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.free_group = None
         self.swa_free_group = []
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
-        return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices)
+        return self._kvcache.get_cpu_copy(
+            indices,
+            mamba_indices=mamba_indices,
+            req_pool_index=req_pool_index,
+        )
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         return self._kvcache.load_cpu_copy(
-            kv_cache_cpu, indices, mamba_indices=mamba_indices
+            kv_cache_cpu,
+            indices,
+            mamba_indices=mamba_indices,
+            req_pool_index=req_pool_index,
         )
 
 

python/sglang/srt/mem_cache/allocator/token.py

@@ -74,10 +74,19 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         else:
             self.free_group.append(self._copy_for_free_group(free_index))
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
-        return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices)
+        return self._kvcache.get_cpu_copy(
+            indices,
+            mamba_indices=mamba_indices,
+            req_pool_index=req_pool_index,
+        )
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         return self._kvcache.load_cpu_copy(
-            kv_cache_cpu, indices, mamba_indices=mamba_indices
+            kv_cache_cpu,
+            indices,
+            mamba_indices=mamba_indices,
+            req_pool_index=req_pool_index,
         )

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -256,10 +256,12 @@ class HiSparseC4DevicePool(DeepSeekV4SingleKVPool):
         loc = self.translate_loc_to_hisparse_device(loc)
         return super().set_key_buffer_fused(layer_id, loc, cache_k)
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         raise NotImplementedError("HiSparseC4DevicePool does not support get_cpu_copy")
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         raise NotImplementedError("HiSparseC4DevicePool does not support load_cpu_copy")
 
 

python/sglang/srt/mem_cache/dsa_cache_layer_split.py

@@ -551,7 +551,7 @@ class LayerSplitDSATokenToKVPool(DSATokenToKVPool):
 
     # ---- HiCache CPU offload: skip empty (non-owned) layers ---------------
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         from sglang.srt.utils import current_platform
 
         current_platform.synchronize()
@@ -569,9 +569,18 @@ class LayerSplitDSATokenToKVPool(DSATokenToKVPool):
                 kv_cache_cpu[-1].append(kv_cpu)
         current_platform.synchronize()
 
-        return {"kv": kv_cache_cpu, "index_k": self.index_key_cache.cpu_copy(indices)}
+        return {
+            "kv": kv_cache_cpu,
+            "index_k": self.index_key_cache.cpu_copy(indices),
+        }
 
-    def load_cpu_copy(self, kv_cache_cpu_dict, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self,
+        kv_cache_cpu_dict,
+        indices,
+        mamba_indices=None,
+        req_pool_index=None,
+    ):
         from sglang.srt.utils import current_platform
 
         kv_cache_cpu = kv_cache_cpu_dict["kv"]

python/sglang/srt/mem_cache/hisparse_memory_pool.py

@@ -40,6 +40,11 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool):
         kv_cache_dim: int,
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
+        index_kpool: int = 1,
+        index_kpool_compress: bool = False,
+        tail_extra_slots: int = 0,
+        max_running_requests: Optional[int] = None,
+        skip_topk_layers: Optional[list[bool]] = None,
         host_to_device_ratio: int = 2,
     ):
         super().__init__(
@@ -56,6 +61,11 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool):
             start_layer=start_layer,
             end_layer=end_layer,
             index_buf_size=size * host_to_device_ratio,
+            index_kpool=index_kpool,
+            index_kpool_compress=index_kpool_compress,
+            tail_extra_slots=tail_extra_slots,
+            max_running_requests=max_running_requests,
+            skip_topk_layers=skip_topk_layers,
         )
         self.bytes_per_token = self.kv_cache_dim * self.dtype.itemsize
 
@@ -115,8 +125,10 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool):
             num_layers=self.layer_num,
         )
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         raise NotImplementedError("HiSparseDevicePool does not support get_cpu_copy")
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         raise NotImplementedError("HiSparseDevicePool does not support load_cpu_copy")

python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py

@@ -788,8 +788,11 @@ def build_hybrid_mamba_stack(
 ) -> tuple[HostPoolGroup, HybridCacheController]:
     transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping)
     mamba_allocator = params.req_to_token_pool.mamba_allocator
+    from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
+
     mtp_draft_device_pools = tuple(
-        pool.full_kv_pool for pool in params.mtp_draft_device_pools
+        pool.full_kv_pool if isinstance(pool, HybridLinearKVPool) else pool
+        for pool in params.mtp_draft_device_pools
     )
     kv_host_size, mamba_host_size = None, 0
     if get_memory().hicache_size > 0:

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -25,6 +25,7 @@ from typing import TYPE_CHECKING
 
 from sglang.srt.arg_groups.overrides import resolving_view
 from sglang.srt.configs.hybrid_arch import (
+    glm5_next_config,
     hybrid_gdn_config,
     hybrid_lightning_config,
     kimi_linear_config,
@@ -126,6 +127,7 @@ def uses_ssm_state(model_config) -> bool:
         or mamba2_config(model_config) is not None
         or (spec.uses_mamba_radix_cache if spec is not None else False)
         or kimi_linear_config(model_config) is not None
+        or glm5_next_config(model_config) is not None
         or hybrid_lightning_config(model_config) is not None
     )
 

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -19,6 +19,8 @@ from sglang.srt.configs.model_config import (
     ModelConfig,
     dsa_layer_skips_topk,
     get_dsa_index_head_dim,
+    get_dsa_index_kpool,
+    get_dsa_index_kpool_compress,
     get_minimax_sparse_attention_config,
     get_minimax_sparse_disable_value_layer_ids,
     get_minimax_sparse_layer_ids,
@@ -999,6 +1001,18 @@ class KVCacheConfigurator:
             )
         return req_to_token_pool
 
+    def _get_mamba_layer_ids_for_req_pool(self) -> list:
+        mamba_layer_ids = [
+            i
+            for i in self.mambaish_config.mamba2_cache_params.layers
+            if self.layer_info.start_layer <= i < self.layer_info.end_layer
+        ]
+        if max_speculative_num_draft_tokens():
+            for layer_id in getattr(self.mambaish_config, "nextn_layer_ids", []):
+                if layer_id not in mamba_layer_ids:
+                    mamba_layer_ids.append(layer_id)
+        return mamba_layer_ids
+
     def _build_hybrid_mamba_decode_req_pool(
         self,
         *,
@@ -1016,13 +1030,7 @@ class KVCacheConfigurator:
             device=self.device,
             enable_memory_saver=get_exec().features.enable_memory_saver,
             cache_params=self.mambaish_config.mamba2_cache_params,
-            mamba_layer_ids=(
+            mamba_layer_ids=self._get_mamba_layer_ids_for_req_pool(),
-                [
-                    i
-                    for i in self.mambaish_config.mamba2_cache_params.layers
-                    if self.layer_info.start_layer <= i < self.layer_info.end_layer
-                ]
-            ),
             speculative_num_draft_tokens=max_speculative_num_draft_tokens(),
             speculative_eagle_topk=get_spec().speculative_eagle_topk,
             enable_mamba_extra_buffer=mamba_extra_buffer_enabled(),
@@ -1095,13 +1103,7 @@ class KVCacheConfigurator:
             device=self.device,
             enable_memory_saver=get_exec().features.enable_memory_saver,
             cache_params=self.mambaish_config.mamba2_cache_params,
-            mamba_layer_ids=(
+            mamba_layer_ids=self._get_mamba_layer_ids_for_req_pool(),
-                [
-                    i
-                    for i in self.mambaish_config.mamba2_cache_params.layers
-                    if self.layer_info.start_layer <= i < self.layer_info.end_layer
-                ]
-            ),
             enable_mamba_extra_buffer=mamba_extra_buffer_enabled(),
             enable_mamba_extra_buffer_lazy=mamba_extra_buffer_lazy_enabled(),
             # A PD prefill server never runs TARGET_VERIFY, so skip the
@@ -1230,9 +1232,10 @@ class KVCacheConfigurator:
                 swa_max_total_num_tokens=sizes.swa_max_total_num_tokens,
                 is_dsa_model=is_dsa_model,
             )
-        elif self.use_mla_backend and is_dsa_model:
+        elif self.use_mla_backend and is_dsa_model and not self.mambaish_config:
             token_to_kv_pool = self._build_dsa_kv_pool(
                 max_total_num_tokens=sizes.max_total_num_tokens,
+                max_running_requests=sizes.max_running_requests,
             )
         elif self.use_mla_backend and not self.mambaish_config:
             assert not is_dsa_model
@@ -1521,7 +1524,9 @@ class KVCacheConfigurator:
         )
         return token_to_kv_pool
 
-    def _build_dsa_kv_pool(self, *, max_total_num_tokens: int) -> KVCache:
+    def _build_dsa_kv_pool(
+        self, *, max_total_num_tokens: int, max_running_requests: int
+    ) -> KVCache:
         from sglang.srt.layers.cp.utils import get_glm_dsa_cp_layer_shard_info
 
         (
@@ -1570,6 +1575,12 @@ class KVCacheConfigurator:
             start_layer=self.layer_info.start_layer,
             end_layer=self.layer_info.end_layer,
             index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config),
+            index_kpool=get_dsa_index_kpool(self.model_config.hf_config),
+            index_kpool_compress=get_dsa_index_kpool_compress(
+                self.model_config.hf_config
+            ),
+            tail_extra_slots=(max_speculative_num_draft_tokens() or 0),
+            max_running_requests=max_running_requests,
             **pool_kwargs,
         )
         return token_to_kv_pool
@@ -1760,12 +1771,6 @@ class KVCacheConfigurator:
         req_to_token_pool: ReqToTokenPool,
         mha_pool_class: type,
     ) -> KVCache:
-        extra_args = {}
-        if self.use_mla_backend:
-            extra_args = {
-                "kv_lora_rank": self.model_config.kv_lora_rank,
-                "qk_rope_head_dim": self.model_config.qk_rope_head_dim,
-            }
         full_attention_layer_ids = (
             [0]
             if self.is_draft_worker
@@ -1775,6 +1780,39 @@ class KVCacheConfigurator:
                 if self.layer_info.start_layer <= i < self.layer_info.end_layer
             ]
         )
+        extra_args = {}
+        if self.use_mla_backend:
+            extra_args = {
+                "kv_lora_rank": self.model_config.kv_lora_rank,
+                "qk_rope_head_dim": self.model_config.qk_rope_head_dim,
+            }
+            if is_deepseek_dsa(self.model_config.hf_config):
+                dsa_index_kpool = get_dsa_index_kpool(self.model_config.hf_config)
+                extra_args.update(
+                    use_dsa=True,
+                    index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config),
+                    kv_cache_dim=calculate_mla_kv_cache_dim(
+                        model_config=self.model_config,
+                        kv_cache_dtype=self.kv_cache_dtype,
+                    ),
+                    index_kpool=dsa_index_kpool,
+                    index_kpool_compress=get_dsa_index_kpool_compress(
+                        self.model_config.hf_config
+                    ),
+                    skip_topk_layers=(
+                        None
+                        if self.is_draft_worker
+                        else [
+                            dsa_layer_skips_topk(self.model_config.hf_config, layer_id)
+                            for layer_id in full_attention_layer_ids
+                        ]
+                    ),
+                )
+                if dsa_index_kpool > 1:
+                    extra_args.update(
+                        tail_extra_slots=(max_speculative_num_draft_tokens() or 0),
+                        max_running_requests=(req_to_token_pool.req_to_token.shape[0]),
+                    )
         quant_method = self._build_mha_quant_method(
             num_layers=len(full_attention_layer_ids)
         )
@@ -2430,14 +2468,21 @@ def calculate_mla_kv_cache_dim(
     if not is_dsa_model:
         return kv_cache_dim
 
-    # TRTLLM backend does not override kv_cache_dim for MLA kv cache
+    # TRTLLM uses the raw MLA KV layout. In disaggregated serving only the
-    # Assuming dsa prefill and decode backends are the same when using trtllm MLA backend,
+    # backend for the local role determines the local pool layout; the
-    # since it is not compatible for trtllm and other mla attn backend due to the different
+    # inactive role may legitimately have a different default backend.
-    # kv cache layout.
+    disaggregation_mode = get_disagg().disaggregation_mode
-    if (
+    if disaggregation_mode == "decode":
+        uses_trtllm_kv_layout = get_exec().kernel.dsa_decode_backend == "trtllm"
+    elif disaggregation_mode == "prefill":
+        uses_trtllm_kv_layout = get_exec().kernel.dsa_prefill_backend == "trtllm"
+    else:
+        uses_trtllm_kv_layout = (
             get_exec().kernel.dsa_prefill_backend == "trtllm"
             or get_exec().kernel.dsa_decode_backend == "trtllm"
-    ):
+        )
+
+    if uses_trtllm_kv_layout:
         return kv_cache_dim
 
     # On HIP, TileLang and AITER DSA kernels consume the raw MLA KV layout:

python/sglang/srt/mem_cache/memory_pool.py

@@ -1406,6 +1406,19 @@ class HybridReqToTokenPool(ReqToTokenPool):
     def mamba2_layer_cache(self, layer_id: int):
         return self.mamba_pool.mamba2_layer_cache(self.mamba2_layer_index(layer_id))
 
+    def copy_mamba_state(
+        self, src_index: torch.Tensor, dst_index: torch.Tensor
+    ) -> None:
+        if src_index.numel() == 0:
+            return
+        if (
+            self.layer_transfer_counter is not None
+            and self.layer_transfer_counter.consumer_index >= 0
+        ):
+            last_mamba_layer = max(self.mamba_map)
+            self.layer_transfer_counter.wait_until(last_mamba_layer - self.start_layer)
+        self.mamba_pool.copy_from(src_index, dst_index)
+
     def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState:
         return self.mamba_pool.get_speculative_mamba2_params_all_layers()
 
@@ -1554,6 +1567,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
             req.kv.mamba_ping_pong_track_buffer = None
             req.kv.mamba_next_track_idx = None
             req.kv.mamba_last_track_idx = None
+            req.kv.mamba_last_track_seqlen = None
+            req.kv.mamba_cow_src_index = None
+            req.kv.mamba_needs_clear = False
 
     def clear(self):
         logger.info("Reset HybridReqToTokenPool")
@@ -1761,10 +1777,12 @@ class KVCache(abc.ABC):
     def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter):
         self.layer_transfer_counter = layer_transfer_counter
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         raise NotImplementedError()
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         raise NotImplementedError()
 
     def get_kv_cache_quant_method(self) -> Any:
@@ -2265,7 +2283,7 @@ class MHATokenToKVPool(KVCache):
         item_lens = [d.item_len_bytes(self.page_size) for d in self._kv_buffer_descs]
         return ptrs, lens, item_lens
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         assert not self.use_hnd, (
             "CPU KV offload indexes by slot (NHD); HND KV cache "
             "(SGLANG_USE_HND_KVCACHE) is not supported with CPU offload yet."
@@ -2287,7 +2305,9 @@ class MHATokenToKVPool(KVCache):
         current_platform.synchronize()
         return kv_cache_cpu
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         assert not self.use_hnd, (
             "CPU KV offload indexes by slot (NHD); HND KV cache "
             "(SGLANG_USE_HND_KVCACHE) is not supported with CPU offload yet."
@@ -3224,13 +3244,15 @@ class PageMajorMHATokenToKVPool(MHATokenToKVPool):
             "with a page-aware transfer scheme)."
         )
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         raise NotImplementedError(
             "CPU offloading is unsupported under the page-major layout "
             "(TODO: split token ids into page/slot for the 4-D index)."
         )
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         raise NotImplementedError(
             "CPU offloading is unsupported under the page-major layout "
             "(TODO: split token ids into page/slot for the 4-D index)."
@@ -3509,7 +3531,7 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool):
             )
         return self.k_scale_buffer[idx][loc], self.v_scale_buffer[idx][loc]
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         # The scales travel with their fp8 payload; a restored slot dequantizes
         # against mismatched exponents without them.
         assert not self.use_hnd, (
@@ -3539,7 +3561,9 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool):
         current_platform.synchronize()
         return kv_cache_cpu
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         assert not self.use_hnd, (
             "CPU KV offload indexes by slot (NHD); HND KV cache "
             "(SGLANG_USE_HND_KVCACHE) is not supported with CPU offload yet."
@@ -3625,6 +3649,14 @@ class HybridLinearKVPool(KVCache):
         use_mla: bool = False,
         kv_lora_rank: int = None,
         qk_rope_head_dim: int = None,
+        use_dsa: bool = False,
+        index_head_dim: Optional[int] = None,
+        kv_cache_dim: Optional[int] = None,
+        index_kpool: int = 1,
+        index_kpool_compress: bool = False,
+        tail_extra_slots: int = 0,
+        max_running_requests: Optional[int] = None,
+        skip_topk_layers: Optional[List[bool]] = None,
         start_layer: Optional[int] = None,
         full_kv_pool_class: Optional[type] = None,
         quant_method=None,
@@ -3648,6 +3680,7 @@ class HybridLinearKVPool(KVCache):
         # `load_cpu_copy`, the only readers, so those ids never arrive here.
         self._mamba_translate = lambda ids: ids
         self.use_mla = use_mla
+        self.use_dsa = use_dsa
         if full_kv_pool is not None:
             # Shared-KV-pool path: the caller built a UnifiedMHATokenToKVPool
             # aliasing the shared byte buffer.
@@ -3693,6 +3726,30 @@ class HybridLinearKVPool(KVCache):
                 **quant_method_kwarg,
                 **post_capture_kwargs,
             )
+        elif use_dsa:
+            # DSA sparse full-attention layers share the MLA latent layout and
+            # additionally keep a paged index_k cache. Only full-attn layer count
+            # is allocated here; the wrapper translates global layer_id to dense.
+            assert index_head_dim is not None and kv_cache_dim is not None, (
+                "HybridLinearKVPool with use_dsa requires index_head_dim and kv_cache_dim"
+            )
+            self.full_kv_pool = DSATokenToKVPool(
+                size=size,
+                page_size=self.page_size,
+                kv_lora_rank=kv_lora_rank,
+                dtype=dtype,
+                qk_rope_head_dim=qk_rope_head_dim,
+                layer_num=self.full_layer_nums,
+                device=device,
+                index_head_dim=index_head_dim,
+                enable_memory_saver=enable_memory_saver,
+                kv_cache_dim=kv_cache_dim,
+                index_kpool=index_kpool,
+                index_kpool_compress=index_kpool_compress,
+                tail_extra_slots=tail_extra_slots,
+                max_running_requests=max_running_requests,
+                skip_topk_layers=skip_topk_layers,
+            )
         else:
             TokenToKVPoolClass = MLATokenToKVPool
 
@@ -3737,6 +3794,38 @@ class HybridLinearKVPool(KVCache):
         self.full_kv_pool._finalize_backing_tokens(config.max_total_num_tokens)
         self.size = int(config.max_total_num_tokens)
 
+    @property
+    def dsa_kv_cache_store_fp8(self) -> bool:
+        return getattr(self.full_kv_pool, "dsa_kv_cache_store_fp8", False)
+
+    @property
+    def kv_cache_dim(self):
+        return getattr(self.full_kv_pool, "kv_cache_dim", None)
+
+    @property
+    def index_head_dim(self) -> Optional[int]:
+        return getattr(self.full_kv_pool, "index_head_dim", None)
+
+    @property
+    def quant_block_size(self) -> Optional[int]:
+        return getattr(self.full_kv_pool, "quant_block_size", None)
+
+    @property
+    def index_kpool(self) -> int:
+        return getattr(self.full_kv_pool, "index_kpool", 1)
+
+    @property
+    def index_kpool_compress(self) -> bool:
+        return bool(getattr(self.full_kv_pool, "index_kpool_compress", False))
+
+    @property
+    def tail_extra_slots(self) -> int:
+        return getattr(self.full_kv_pool, "tail_extra_slots", 0)
+
+    @property
+    def slots_per_page(self) -> int:
+        return getattr(self.full_kv_pool, "slots_per_page", self.page_size)
+
     def get_kv_size_bytes(self):
         return self.full_kv_pool.get_kv_size_bytes()
 
@@ -3901,8 +3990,8 @@ class HybridLinearKVPool(KVCache):
     def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
         self.full_kv_pool.move_kv_cache(tgt_loc, src_loc)
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
-        kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
+        kv_cpu = self.full_kv_pool.get_cpu_copy(indices, req_pool_index=req_pool_index)
         # mamba_pool stores PHYSICAL ids; translate the (unified-pool virtual) ids first.
         mamba_cpu = (
             self.mamba_pool.get_cpu_copy(self._mamba_translate(mamba_indices))
@@ -3911,9 +4000,11 @@ class HybridLinearKVPool(KVCache):
         )
         return kv_cpu, mamba_cpu
 
-    def load_cpu_copy(self, cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         kv_cpu, mamba_cpu = cache_cpu
-        self.full_kv_pool.load_cpu_copy(kv_cpu, indices)
+        self.full_kv_pool.load_cpu_copy(kv_cpu, indices, req_pool_index=req_pool_index)
         if mamba_cpu is not None and mamba_indices is not None:
             self.mamba_pool.load_cpu_copy(
                 mamba_cpu, self._mamba_translate(mamba_indices)
@@ -3951,6 +4042,145 @@ class HybridLinearKVPool(KVCache):
         with self._transfer_id_context(layer):
             return self.full_kv_pool.get_mla_kv_buffer(layer, loc, dst_dtype)
 
+    def set_index_k_scale_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        index_k: torch.Tensor,
+        index_k_scale: torch.Tensor,
+    ) -> None:
+        assert self.use_dsa, "set_index_k_scale_buffer called when use_dsa is False"
+        layer_id = self._transfer_full_attention_id(layer_id)
+        self.full_kv_pool.set_index_k_scale_buffer(
+            layer_id, loc, index_k, index_k_scale
+        )
+
+    def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor:
+        assert self.use_dsa, (
+            "get_index_k_with_scale_buffer called when use_dsa is False"
+        )
+        self._wait_for_layer(layer_id)
+        layer_id = self._transfer_full_attention_id(layer_id)
+        return self.full_kv_pool.get_index_k_with_scale_buffer(layer_id)
+
+    def get_broadcastable_index_k_with_scale_buffer(
+        self, layer_id: int
+    ) -> torch.Tensor:
+        assert self.use_dsa, (
+            "get_broadcastable_index_k_with_scale_buffer called when use_dsa is False"
+        )
+        self._wait_for_layer(layer_id)
+        layer_id = self._transfer_full_attention_id(layer_id)
+        if hasattr(self.full_kv_pool, "_get_broadcastable_index_buffer"):
+            return self.full_kv_pool._get_broadcastable_index_buffer(layer_id)
+        return self.full_kv_pool.get_index_k_with_scale_buffer(layer_id)
+
+    def invalidate_index_buffer_for_layer(self, layer_id: int) -> None:
+        if not self.use_dsa or not hasattr(
+            self.full_kv_pool, "invalidate_index_buffer_for_layer"
+        ):
+            return
+        layer_id = self._transfer_full_attention_id(layer_id)
+        self.full_kv_pool.invalidate_index_buffer_for_layer(layer_id)
+
+    def get_index_k_continuous(
+        self,
+        layer_id: int,
+        seq_len: int,
+        page_indices: torch.Tensor,
+    ):
+        assert self.use_dsa, "get_index_k_continuous called when use_dsa is False"
+        self._wait_for_layer(layer_id)
+        layer_id = self._transfer_full_attention_id(layer_id)
+        return self.full_kv_pool.get_index_k_continuous(layer_id, seq_len, page_indices)
+
+    def get_index_k_scale_continuous(
+        self,
+        layer_id: int,
+        seq_len: int,
+        page_indices: torch.Tensor,
+    ):
+        assert self.use_dsa, "get_index_k_scale_continuous called when use_dsa is False"
+        self._wait_for_layer(layer_id)
+        layer_id = self._transfer_full_attention_id(layer_id)
+        return self.full_kv_pool.get_index_k_scale_continuous(
+            layer_id, seq_len, page_indices
+        )
+
+    def get_index_k_scale_buffer(
+        self,
+        layer_id: int,
+        seq_len_tensor: torch.Tensor,
+        page_indices: torch.Tensor,
+        seq_len_sum: int,
+        max_seq_len: int,
+    ):
+        assert self.use_dsa, "get_index_k_scale_buffer called when use_dsa is False"
+        self._wait_for_layer(layer_id)
+        layer_id = self._transfer_full_attention_id(layer_id)
+        return self.full_kv_pool.get_index_k_scale_buffer(
+            layer_id, seq_len_tensor, page_indices, seq_len_sum, max_seq_len
+        )
+
+    def get_compress_tail_buffers(
+        self, layer_id: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert self.use_dsa, "get_compress_tail_buffers called when use_dsa is False"
+        layer_id = self._transfer_full_attention_id(layer_id)
+        return self.full_kv_pool.get_compress_tail_buffers(layer_id)
+
+    def kpool_decode_update_index_cache(
+        self,
+        layer_id: int,
+        key: torch.Tensor,
+        slot_score: torch.Tensor,
+        ape: torch.Tensor,
+        block_tables: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        positions: torch.Tensor,
+        seq_lens: torch.Tensor,
+        out_cache_loc: torch.Tensor,
+        round_scale: bool = False,
+    ) -> None:
+        assert self.use_dsa, (
+            "kpool_decode_update_index_cache called when use_dsa is False"
+        )
+        layer_id = self._transfer_full_attention_id(layer_id)
+        self.full_kv_pool.kpool_decode_update_index_cache(
+            layer_id=layer_id,
+            key=key,
+            slot_score=slot_score,
+            ape=ape,
+            block_tables=block_tables,
+            req_pool_indices=req_pool_indices,
+            positions=positions,
+            seq_lens=seq_lens,
+            out_cache_loc=out_cache_loc,
+            round_scale=round_scale,
+        )
+
+    def set_compress_tail_for_request(
+        self,
+        layer_id: int,
+        req_pool_idx: torch.Tensor,
+        key_tail: torch.Tensor,
+        score_tail: torch.Tensor,
+        n_remain: int,
+        dst_logical_start: int,
+    ) -> None:
+        assert self.use_dsa, (
+            "set_compress_tail_for_request called when use_dsa is False"
+        )
+        layer_id = self._transfer_full_attention_id(layer_id)
+        self.full_kv_pool.set_compress_tail_for_request(
+            layer_id=layer_id,
+            req_pool_idx=req_pool_idx,
+            key_tail=key_tail,
+            score_tail=score_tail,
+            n_remain=n_remain,
+            dst_logical_start=dst_logical_start,
+        )
+
 
 class MLATokenToKVPool(KVCache):
     def __init__(
@@ -4163,9 +4393,11 @@ class MLATokenToKVPool(KVCache):
         else:
             if cache_k_nope.dtype != self.dtype:
                 cache_k_nope = cache_k_nope.to(self.dtype)
+                if cache_k_rope is not None and cache_k_rope.numel() > 0:
                     cache_k_rope = cache_k_rope.to(self.dtype)
             if self.store_dtype != self.dtype:
                 cache_k_nope = cache_k_nope.view(self.store_dtype)
+                if cache_k_rope is not None and cache_k_rope.numel() > 0:
                     cache_k_rope = cache_k_rope.view(self.store_dtype)
 
             self._scatter_mla_rows(dst_buffer, loc, cache_k_nope, cache_k_rope)
@@ -4213,6 +4445,9 @@ class MLATokenToKVPool(KVCache):
             dtype=dst_dtype,
             device=kv_buffer.device,
         )
+        if self.qk_rope_head_dim == 0:
+            cache_k_rope = None
+        else:
             cache_k_rope = torch.empty(
                 (loc.shape[0], 1, self.qk_rope_head_dim),
                 dtype=dst_dtype,
@@ -4235,12 +4470,14 @@ class MLATokenToKVPool(KVCache):
         for kv_cache in self.kv_buffer:
             kv_cache[tgt_loc_flat] = kv_cache[src_loc_flat]
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         current_platform.synchronize()
         kv_cache_cpu = []
         chunk_size = self.cpu_offloading_chunk_size
         for layer_id in range(self.layer_num):
             kv_cache_cpu.append([])
+            if self.kv_buffer[layer_id].shape[0] == 0:
+                continue
             for i in range(0, len(indices), chunk_size):
                 chunk_indices = indices[i : i + chunk_size]
                 kv_cpu = self.kv_buffer[layer_id][chunk_indices].to(
@@ -4250,10 +4487,14 @@ class MLATokenToKVPool(KVCache):
         current_platform.synchronize()
         return kv_cache_cpu
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         current_platform.synchronize()
         chunk_size = self.cpu_offloading_chunk_size
         for layer_id in range(self.layer_num):
+            if self.kv_buffer[layer_id].shape[0] == 0:
+                continue
             for i in range(0, len(indices), chunk_size):
                 chunk_indices = indices[i : i + chunk_size]
                 kv_cpu = kv_cache_cpu[layer_id][i // chunk_size]
@@ -4381,12 +4622,17 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
                 cache_k_nope_fp4, cache_k_nope_fp4_sf = (
                     FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_nope)
                 )
+                if cache_k_rope is not None and cache_k_rope.numel() > 0:
                     cache_k_rope_fp4, cache_k_rope_fp4_sf = (
                         FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_rope)
                     )
+                else:
+                    cache_k_rope_fp4 = None
+                    cache_k_rope_fp4_sf = None
 
             if self.store_dtype != self.dtype:
                 cache_k_nope = cache_k_nope.view(self.store_dtype)
+                if cache_k_rope is not None and cache_k_rope.numel() > 0:
                     cache_k_rope = cache_k_rope.view(self.store_dtype)
 
             self._scatter_mla_rows(
@@ -4423,6 +4669,10 @@ class DSATokenToKVPool(MLATokenToKVPool):
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
         index_buf_size: Optional[int] = None,
+        index_kpool: int = 1,
+        index_kpool_compress: bool = False,
+        tail_extra_slots: int = 0,
+        max_running_requests: Optional[int] = None,
         skip_topk_layers: Optional[List[bool]] = None,
     ):
         override_dim = (
@@ -4446,6 +4696,10 @@ class DSATokenToKVPool(MLATokenToKVPool):
         # self.index_k_dtype = torch.float8_e4m3fn
         # self.index_k_scale_dtype = torch.float32
         self.index_head_dim = index_head_dim
+        self.index_kpool = index_kpool
+        self.index_kpool_compress = index_kpool_compress
+        self.tail_extra_slots = tail_extra_slots
+        self.slots_per_page = self.page_size
         if index_buf_size is None:
             index_buf_size = size
         self.index_buf_size = index_buf_size
@@ -4471,19 +4725,172 @@ class DSATokenToKVPool(MLATokenToKVPool):
         else:
             assert self.page_size == 64
         self.index_key_cache = self._create_index_key_cache()
+        self._init_kpool_compress_tail_buffers(
+            index_kpool=index_kpool,
+            index_kpool_compress=index_kpool_compress,
+            tail_extra_slots=tail_extra_slots,
+            index_head_dim=index_head_dim,
+            layer_num=layer_num,
+            device=device,
+            max_running_requests=max_running_requests,
+        )
         self._finalize_allocation_log(size)
 
     def _create_index_key_cache(self) -> IndexKeyCache:
         return IndexKeyCache(self, self.index_buf_size)
 
+    def _should_allocate_index_layer(self, local_layer_idx: int) -> bool:
+        return not self.skip_topk_layers[local_layer_idx]
+
     @property
     def index_k_with_scale_buffer(self):
         # Preserve direct HiCache access while storage lives behind the facade.
         return self.index_key_cache.buffer
 
+    def _init_kpool_compress_tail_buffers(
+        self,
+        index_kpool: int,
+        index_kpool_compress: bool,
+        tail_extra_slots: int,
+        index_head_dim: int,
+        layer_num: int,
+        device: str,
+        max_running_requests: Optional[int],
+    ) -> None:
+        """Keep request tails on the pool so they follow the index-cache lifecycle."""
+        self.kpool_use_compress = index_kpool > 1 and index_kpool_compress
+
+        if not self.kpool_use_compress:
+            self._compress_tail_k = None
+            self._compress_tail_score = None
+            return
+
+        assert max_running_requests is not None, (
+            "DSATokenToKVPool with kpool compress requires max_running_requests"
+        )
+        # +1 mirrors req_to_token_pool.size + 1 used by the indexer to
+        # provide an extra slot for invalid / sentinel req indices.
+        req_pool_size = max_running_requests + 1
+        tail_dtype = torch.bfloat16
+        tail_width = index_kpool + tail_extra_slots
+        with (
+            torch.cuda.use_mem_pool(self.custom_mem_pool)
+            if self.custom_mem_pool
+            else nullcontext()
+        ):
+            self._compress_tail_k: Optional[List[torch.Tensor]] = [
+                torch.zeros(
+                    req_pool_size if self._should_allocate_index_layer(i) else 0,
+                    tail_width,
+                    index_head_dim,
+                    dtype=tail_dtype,
+                    device=device,
+                )
+                for i in range(layer_num)
+            ]
+            self._compress_tail_score: Optional[List[torch.Tensor]] = [
+                torch.zeros(
+                    req_pool_size if self._should_allocate_index_layer(i) else 0,
+                    tail_width,
+                    index_head_dim,
+                    dtype=tail_dtype,
+                    device=device,
+                )
+                for i in range(layer_num)
+            ]
+
+    def get_compress_tail_buffers(
+        self, layer_id: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert self.kpool_use_compress, (
+            "get_compress_tail_buffers called when kpool compress is disabled"
+        )
+        idx = layer_id - self.start_layer
+        return (
+            self._compress_tail_k[idx],
+            self._compress_tail_score[idx],
+        )
+
+    def get_compress_tail_buf_infos(self):
+        if not self.kpool_use_compress:
+            return [], [], []
+        transfer_layer_ids = list(range(self.layer_num))
+        # Keep zero-row indexShare entries in the pointer list so layer offsets
+        # stay aligned across PD peers; item_len=0 makes transfer backends skip them.
+        tail_buffers = [self._compress_tail_k[i] for i in transfer_layer_ids] + [
+            self._compress_tail_score[i] for i in transfer_layer_ids
+        ]
+        data_ptrs = [buf.data_ptr() for buf in tail_buffers]
+        data_lens = [buf.nbytes for buf in tail_buffers]
+        item_lens = [buf[0].nbytes if buf.shape[0] > 0 else 0 for buf in tail_buffers]
+        return data_ptrs, data_lens, item_lens
+
+    def kpool_decode_update_index_cache(
+        self,
+        layer_id: int,
+        key: torch.Tensor,
+        slot_score: torch.Tensor,
+        ape: torch.Tensor,
+        block_tables: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        positions: torch.Tensor,
+        seq_lens: torch.Tensor,
+        out_cache_loc: torch.Tensor,
+        round_scale: bool = False,
+    ) -> None:
+        from sglang.srt.layers.attention.dsa.kpool_fp8_index import (
+            kpool_decode_update_and_maybe_write_cache,
+        )
+
+        assert self.kpool_use_compress, (
+            "kpool_decode_update_index_cache called when kpool compress is disabled"
+        )
+        idx = layer_id - self.start_layer
+        buf = self.get_index_k_with_scale_buffer(layer_id)
+        kpool_decode_update_and_maybe_write_cache(
+            pool=self,
+            buf=buf,
+            tail_k=self._compress_tail_k[idx],
+            tail_score=self._compress_tail_score[idx],
+            key=key,
+            slot_score=slot_score,
+            ape=ape,
+            block_tables=block_tables,
+            req_pool_indices=req_pool_indices,
+            positions=positions,
+            seq_lens=seq_lens,
+            out_cache_loc=out_cache_loc,
+            round_scale=round_scale,
+        )
+
+    def set_compress_tail_for_request(
+        self,
+        layer_id: int,
+        req_pool_idx: torch.Tensor,
+        key_tail: torch.Tensor,
+        score_tail: torch.Tensor,
+        n_remain: int,
+        dst_logical_start: int,
+    ) -> None:
+        """Leave the ring untouched at a pool boundary; no tail carries over."""
+        assert self.kpool_use_compress, (
+            "set_compress_tail_for_request called when kpool compress is disabled"
+        )
+        idx = layer_id - self.start_layer
+        if n_remain > 0:
+            slots = (
+                torch.arange(n_remain, device=key_tail.device, dtype=torch.long)
+                + int(dst_logical_start)
+            ) % self._compress_tail_k[idx].shape[1]
+            self._compress_tail_k[idx][req_pool_idx, slots] = key_tail
+            self._compress_tail_score[idx][req_pool_idx, slots] = score_tail
+
     def _clear_buffers(self):
         super()._clear_buffers()
         self.index_key_cache.clear()
+        if hasattr(self, "_compress_tail_k") and self._compress_tail_k is not None:
+            del self._compress_tail_k
+            del self._compress_tail_score
 
     def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
         """Move latent KV and the DSA indexer cache (key + scale) in lockstep."""
@@ -4532,15 +4939,78 @@ class DSATokenToKVPool(MLATokenToKVPool):
     ) -> None:
         self.index_key_cache.store_quantized(layer_id, loc, index_k, index_k_scale)
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def _get_compress_tail_cpu_copy(self, req_pool_index):
-        kv_cache_cpu = super().get_cpu_copy(indices, mamba_indices=mamba_indices)
+        if not self.kpool_use_compress or req_pool_index is None:
-        return {"kv": kv_cache_cpu, "index_k": self.index_key_cache.cpu_copy(indices)}
+            return None
 
-    def load_cpu_copy(self, kv_cache_cpu_dict, indices, mamba_indices=None):
+        tail_k_cpu = []
+        tail_score_cpu = []
+        for tail_k, tail_score in zip(self._compress_tail_k, self._compress_tail_score):
+            if tail_k.shape[0] == 0:
+                tail_k_cpu.append(None)
+                tail_score_cpu.append(None)
+                continue
+            tail_k_cpu.append(tail_k[req_pool_index].to("cpu", non_blocking=True))
+            tail_score_cpu.append(
+                tail_score[req_pool_index].to("cpu", non_blocking=True)
+            )
+        return tail_k_cpu, tail_score_cpu
+
+    def _load_compress_tail_cpu_copy(self, tail_k_cpu, tail_score_cpu, req_pool_index):
+        if (
+            not self.kpool_use_compress
+            or req_pool_index is None
+            or tail_k_cpu is None
+            or tail_score_cpu is None
+        ):
+            return
+
+        for tail_k, tail_score, saved_k, saved_score in zip(
+            self._compress_tail_k,
+            self._compress_tail_score,
+            tail_k_cpu,
+            tail_score_cpu,
+        ):
+            if tail_k.shape[0] == 0 or saved_k is None or saved_score is None:
+                continue
+            tail_k[req_pool_index] = saved_k.to(tail_k.device, non_blocking=True)
+            tail_score[req_pool_index] = saved_score.to(
+                tail_score.device, non_blocking=True
+            )
+
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
+        # Retraction reuses index-cache pages; offload index/scale with KV so resume cannot read another request's entries.
+        kv_cache_cpu = super().get_cpu_copy(indices, mamba_indices=mamba_indices)
+        cpu_copy = {
+            "kv": kv_cache_cpu,
+            "index_k": self.index_key_cache.cpu_copy(indices),
+        }
+        compress_tail = self._get_compress_tail_cpu_copy(req_pool_index)
+        if compress_tail is not None:
+            cpu_copy["tail_k"], cpu_copy["tail_score"] = compress_tail
+        torch.cuda.synchronize()
+        return cpu_copy
+
+    def load_cpu_copy(
+        self,
+        kv_cache_cpu_dict,
+        indices,
+        mamba_indices=None,
+        req_pool_index=None,
+    ):
         super().load_cpu_copy(
-            kv_cache_cpu_dict["kv"], indices, mamba_indices=mamba_indices
+            kv_cache_cpu_dict["kv"],
+            indices,
+            mamba_indices=mamba_indices,
+            req_pool_index=req_pool_index,
         )
         self.index_key_cache.load_cpu_copy(kv_cache_cpu_dict["index_k"], indices)
+        self._load_compress_tail_cpu_copy(
+            kv_cache_cpu_dict.get("tail_k"),
+            kv_cache_cpu_dict.get("tail_score"),
+            req_pool_index,
+        )
+        torch.cuda.synchronize()
 
     def get_state_buf_infos(self):
         return self.index_key_cache.state_buf_infos()

python/sglang/srt/mem_cache/swa_memory_pool.py

@@ -366,10 +366,12 @@ class SWAKVPool(BaseSWAKVPool):
             filtered.append(filtered_layer)
         return filtered
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         # For SWA, we need to copy KV cache from both full and SWA pools
         # The indices are for the full pool, and we use mapping to get SWA indices
-        full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices)
+        full_kv_cpu = self.full_kv_pool.get_cpu_copy(
+            indices, req_pool_index=req_pool_index
+        )
 
         swa_mask = None
         if self.full_to_swa_index_mapping is not None:
@@ -388,14 +390,18 @@ class SWAKVPool(BaseSWAKVPool):
 
         return {"full": full_kv_cpu, "swa": swa_kv_cpu, "swa_mask": swa_mask}
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         # Load KV cache back from CPU to both full and SWA pools
         # Note: indices here are NEW indices (newly allocated), different from get_cpu_copy indices
         full_kv_cpu = kv_cache_cpu["full"]
         swa_kv_cpu = kv_cache_cpu["swa"]
 
         # Load full KV cache to the new indices
-        self.full_kv_pool.load_cpu_copy(full_kv_cpu, indices)
+        self.full_kv_pool.load_cpu_copy(
+            full_kv_cpu, indices, req_pool_index=req_pool_index
+        )
 
         # Load SWA KV cache if it exists
         if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None:

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -1621,23 +1621,29 @@ class UnifiedSWAKVPool(SWAKVPool):
         phys_pages = allocator.virtual_to_physical[virt_pages]
         return phys_pages * ps + offsets
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         assert self._full_allocator is not None
         assert self._swa_allocator is not None
         # `indices` are virtual TOKEN ids; translate per sub-pool.
         full_phys = self._virt_tokens_to_phys_tokens(indices, self._full_allocator)
         swa_phys = self._virt_tokens_to_phys_tokens(indices, self._swa_allocator)
-        full_cpu = self.full_kv_pool.get_cpu_copy(full_phys)
+        full_cpu = self.full_kv_pool.get_cpu_copy(
+            full_phys, req_pool_index=req_pool_index
+        )
         valid = swa_phys >= 0
         swa_cpu = None
         if bool(valid.any().item()):
             swa_cpu = self.swa_kv_pool.get_cpu_copy(swa_phys[valid])
         return {"full": full_cpu, "swa": swa_cpu}
 
-    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
+    ):
         assert self._full_allocator is not None
         full_phys = self._virt_tokens_to_phys_tokens(indices, self._full_allocator)
-        self.full_kv_pool.load_cpu_copy(kv_cache_cpu["full"], full_phys)
+        self.full_kv_pool.load_cpu_copy(
+            kv_cache_cpu["full"], full_phys, req_pool_index=req_pool_index
+        )
         if kv_cache_cpu.get("swa") is not None:
             assert self._swa_allocator is not None
             swa_phys = self._virt_tokens_to_phys_tokens(indices, self._swa_allocator)

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -656,6 +656,10 @@ def build_decode_registry(
             # init_new -- they leave the GLOBAL None and set the replicated LOCAL
             # count directly, so carry that through.
             if fb.global_num_token_non_padded is None:
+                # DFLASH's dense draft can omit both optional counts, even
+                # when EP on the target enables this slot. Preserve the
+                # registry's skip-missing-field behavior for that path.
+                if fb.num_token_non_padded is not None:
                     buf.copy_(fb.num_token_non_padded)
                 return
             sharded = not enable_prefill_cp and attn_tp_sharded_fn(

python/sglang/srt/model_executor/forward_batch_info.py

@@ -56,6 +56,7 @@ from sglang.srt.runtime_context import (
     get_lora,
     get_parallel,
 )
+from sglang.srt.speculative.spec_info import SpecInputType
 from sglang.srt.utils import (
     is_cpu,
     is_cuda,
@@ -1739,6 +1740,26 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
                     logits_output.hidden_states = logits_output.hidden_states[
                         :num_tokens
                     ]
+            elif (
+                self.spec_info.spec_input_type == SpecInputType.EAGLE_DRAFT_EXTEND
+                and not self.forward_mode.is_draft_extend_v2()
+            ):
+                if self.spec_info.num_correct_drafts is not None:
+                    self.spec_info.num_correct_drafts = (
+                        self.spec_info.num_correct_drafts[:bs]
+                    )
+                if self.spec_info.num_accept_tokens is not None:
+                    self.spec_info.num_accept_tokens = self.spec_info.num_accept_tokens[
+                        :bs
+                    ]
+                if self.extend_seq_lens is not None:
+                    self.extend_seq_lens = self.extend_seq_lens[:bs]
+                if logits_output.next_token_logits is not None:
+                    logits_output.next_token_logits = logits_output.next_token_logits[
+                        :bs
+                    ]
+                if logits_output.hidden_states is not None:
+                    logits_output.hidden_states = logits_output.hidden_states[:bs]
             elif self.forward_mode.is_draft_extend_v2():  # draft extend_v2
                 bs = bs * self.spec_info.num_tokens_per_req
                 if logits_output.next_token_logits is not None:

python/sglang/srt/model_executor/model_runner.py

@@ -1741,7 +1741,7 @@ class ModelRunner:
                 )
             else:
                 # mamba_pool is a pure PHYSICAL store; translate both COW slot ids.
-                pool.mamba_pool.copy_from(
+                pool.copy_mamba_state(
                     pool.translate_mamba_indices(forward_batch.mamba_cow_src_indices),
                     pool.translate_mamba_indices(forward_batch.mamba_cow_dst_indices),
                 )

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

@@ -53,6 +53,7 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.runner.base_runner import BaseRunner
 
 logger = logging.getLogger(__name__)
+_deep_gemm_layout_memory_budget_initialized = False
 
 
 def _align_pipeline_layers(layers: list, layer_model) -> list:
@@ -154,6 +155,75 @@ class CudaGraphsCapture(msgspec.Struct, frozen=True, kw_only=True):
         )
 
 
+def refresh_deep_gemm_layout_memory_budget(
+    model_runner: ModelRunner, *, only_if_initialized: bool = False
+) -> None:
+    """Set the all-rank budget before capture, then refresh after startup."""
+    global _deep_gemm_layout_memory_budget_initialized
+    if (
+        model_runner.device != "cuda"
+        or envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower() != "auto"
+    ):
+        return
+
+    if only_if_initialized:
+        # Target and draft share the budget. Its pre-capture initialization
+        # already used a world-wide collective, so this guard is rank-uniform
+        # and also covers a draft-only DeepGEMM backend outside draft context.
+        if not _deep_gemm_layout_memory_budget_initialized:
+            return
+    else:
+        if model_runner.is_draft_worker:
+            moe_runner_backend = (
+                get_spec().speculative_moe_runner_backend
+                or get_exec().moe.moe_runner_backend
+            )
+            moe_a2a_backend = (
+                get_spec().speculative_moe_a2a_backend or get_exec().moe.moe_a2a_backend
+            )
+        else:
+            moe_runner_backend = get_exec().moe.moe_runner_backend
+            moe_a2a_backend = get_exec().moe.moe_a2a_backend
+
+        uses_deep_gemm_moe_runner = moe_runner_backend == "deep_gemm"
+        if moe_runner_backend == "auto" and model_runner.model_config.quantization in (
+            "fp8",
+            "mxfp8",
+        ):
+            from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend
+            from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod
+
+            uses_deep_gemm_moe_runner = (
+                Fp8MoEMethod.is_deepgemm_moe_runner_backend_enabled(
+                    MoeRunnerBackend(moe_runner_backend),
+                    MoeA2ABackend(moe_a2a_backend),
+                )
+            )
+        if not uses_deep_gemm_moe_runner:
+            return
+
+    from sglang.srt.layers.moe.moe_runner.deep_gemm import (
+        set_masked_standard_layout_memory_budget,
+    )
+
+    world_group = get_world_group()
+    available_memory_gb = get_available_gpu_memory(
+        model_runner.device,
+        model_runner.gpu_id,
+        distributed=world_group.world_size > 1,
+        cpu_group=world_group.cpu_group,
+    )
+    budget_bytes = set_masked_standard_layout_memory_budget(
+        int(available_memory_gb * (1 << 30))
+    )
+    _deep_gemm_layout_memory_budget_initialized = True
+    logger.info(
+        "DeepGEMM masked layout budget: %.2f GiB from %.2f GiB free.",
+        budget_bytes / (1 << 30),
+        available_memory_gb,
+    )
+
+
 def capture_cuda_graphs(
     *, model_runner: ModelRunner, capture_decode_cuda_graph: bool = True
 ) -> CudaGraphsCapture:
@@ -176,56 +246,7 @@ def capture_cuda_graphs(
     # runners point at it) and the eager fallback when a cg runner can't run a
     # batch.
     eager_runner = EagerRunner(model_runner)
-
+    refresh_deep_gemm_layout_memory_budget(model_runner)
-    if model_runner.is_draft_worker:
-        moe_runner_backend = (
-            get_spec().speculative_moe_runner_backend
-            or get_exec().moe.moe_runner_backend
-        )
-        moe_a2a_backend = (
-            get_spec().speculative_moe_a2a_backend or get_exec().moe.moe_a2a_backend
-        )
-    else:
-        moe_runner_backend = get_exec().moe.moe_runner_backend
-        moe_a2a_backend = get_exec().moe.moe_a2a_backend
-
-    uses_deep_gemm_moe_runner = moe_runner_backend == "deep_gemm"
-    if moe_runner_backend == "auto" and model_runner.model_config.quantization in (
-        "fp8",
-        "mxfp8",
-    ):
-        from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend
-        from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod
-
-        uses_deep_gemm_moe_runner = Fp8MoEMethod.is_deepgemm_moe_runner_backend_enabled(
-            MoeRunnerBackend(moe_runner_backend),
-            MoeA2ABackend(moe_a2a_backend),
-        )
-
-    if (
-        model_runner.device == "cuda"
-        and envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower() == "auto"
-        and uses_deep_gemm_moe_runner
-    ):
-        from sglang.srt.layers.moe.moe_runner.deep_gemm import (
-            set_masked_standard_layout_memory_budget,
-        )
-
-        world_group = get_world_group()
-        available_memory_gb = get_available_gpu_memory(
-            model_runner.device,
-            model_runner.gpu_id,
-            distributed=world_group.world_size > 1,
-            cpu_group=world_group.cpu_group,
-        )
-        budget_bytes = set_masked_standard_layout_memory_budget(
-            int(available_memory_gb * (1 << 30))
-        )
-        logger.info(
-            "DeepGEMM masked layout budget: %.2f GiB from %.2f GiB free.",
-            budget_bytes / (1 << 30),
-            available_memory_gb,
-        )
 
     # cuda-graph capture: prefill before decode, so both coalesce onto the
     # eager buffer allocated above. (capture_prefill_graph routes prefill

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

@@ -57,6 +57,7 @@ def resolve_spec_aux_hidden_state_config(
     _resolve_eagle_aux_hidden_state(
         config=config,
         server_args=server_args,
+        model_config=model_config,
         spec_algorithm=spec_algorithm,
         is_draft_worker=is_draft_worker,
     )
@@ -74,15 +75,18 @@ def _resolve_eagle_aux_hidden_state(
     *,
     config: SpecAuxHiddenStateConfig,
     server_args: ServerArgs,
+    model_config: ModelConfig,
     spec_algorithm: SpeculativeAlgorithm,
     is_draft_worker: bool,
 ) -> None:
-    if (
+    if not (
         (spec_algorithm.is_eagle() or spec_algorithm.is_standalone())
         and not is_draft_worker
-        and get_spec().speculative_draft_model_path
     ):
-        # Load draft config to get layer count for KV cache sizing
+        return
+
+    draft_model_config = model_config
+    if get_spec().speculative_draft_model_path:
         draft_model_config = ModelConfig.from_server_args(
             server_args,
             model_path=get_spec().speculative_draft_model_path,
@@ -92,18 +96,17 @@ def _resolve_eagle_aux_hidden_state(
     num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers
     if num_nextn_predict_layers is not None:
         config.eagle_draft_num_layers = int(num_nextn_predict_layers)
-        else:
+    elif get_spec().speculative_draft_model_path:
         config.eagle_draft_num_layers = int(
             max(
                 draft_model_config.num_hidden_layers,
                 draft_model_config.num_attention_layers,
             )
         )
+    else:
+        return
 
-        if (
+    if draft_model_config.is_hybrid_swa and not draft_model_config.is_deepseek_v4_arch:
-            draft_model_config.is_hybrid_swa
-            and not draft_model_config.is_deepseek_v4_arch
-        ):
         config.eagle_draft_swa_num_layers = len(
             draft_model_config.swa_attention_layer_ids
         )
@@ -111,17 +114,14 @@ def _resolve_eagle_aux_hidden_state(
     if spec_algorithm.is_eagle3():
         config.eagle_use_aux_hidden_state = True
         try:
-                eagle_config = getattr(
+            eagle_config = getattr(draft_model_config.hf_config, "eagle_config", None)
-                    draft_model_config.hf_config, "eagle_config", None
-                )
             config.eagle_use_aux_hidden_state = eagle_config.get(
                 "use_aux_hidden_state", True
             )
             config.eagle_aux_hidden_state_layer_ids = eagle_config[
                 "eagle_aux_hidden_state_layer_ids"
             ]
-            except:
+        except Exception:
-                # if there is no aux layer, set to None
             config.eagle_aux_hidden_state_layer_ids = None
 
 

python/sglang/srt/model_executor/pool_configurator.py

@@ -114,6 +114,23 @@ def _dflash_draft_cell_size(kvc: KVCacheConfigurator) -> int:
     return int(cell_size) * get_parallel().attn_dcp_size
 
 
+def _get_dsa_cache_layer_ids(kvc: KVCacheConfigurator, num_layers: int) -> list[int]:
+    """Global layer ids represented by the local DSA pool's dense layer slots."""
+    if kvc.mambaish_config and not kvc.is_draft_worker:
+        layer_ids = [
+            layer_id
+            for layer_id in kvc.mambaish_config.full_attention_layer_ids
+            if kvc.layer_info.start_layer <= layer_id < kvc.layer_info.end_layer
+        ]
+    else:
+        layer_ids = list(range(kvc.layer_info.start_layer, kvc.layer_info.end_layer))
+    # Draft pools and a few platform-specific pools may expose a synthetic layer
+    # count. They do not use indexShare, so only the length matters for sizing.
+    if len(layer_ids) != num_layers:
+        return list(range(num_layers))
+    return layer_ids
+
+
 def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]:
     dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower()
     if dtype_name in ("float32", "fp32"):
@@ -280,8 +297,10 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
             get_glm_dsa_layer_split_effective_num_layers,
         )
 
-        effective_num_layers = get_glm_dsa_layer_split_effective_num_layers(
+        effective_num_layers = (
-            kvc, num_layers
+            num_layers
+            if kvc.server_args.enable_hisparse
+            else get_glm_dsa_layer_split_effective_num_layers(kvc, num_layers)
         )
 
         kv_size = torch._utils._element_size(kv_cache_dtype)
@@ -412,18 +431,13 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
             _should_elide_dsa_index_k,
         )
 
-        if allocate_all_layers or not _should_elide_dsa_index_k(
+        if (
-            is_draft_worker=kvc.is_draft_worker
+            allocate_all_layers
+            or kvc.server_args.enable_hisparse
+            or not _should_elide_dsa_index_k(is_draft_worker=kvc.is_draft_worker)
         ):
             num_indexer_layers = num_layers
         else:
-            active_indexer_layers = [
-                layer_id
-                for layer_id in range(
-                    kvc.layer_info.start_layer, kvc.layer_info.end_layer
-                )
-                if not dsa_layer_skips_topk(kvc.model_config.hf_config, layer_id)
-            ]
             from sglang.srt.layers.cp.utils import (
                 get_glm_dsa_cp_layer_shard_info,
                 get_layer_shard_range,
@@ -431,6 +445,16 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
 
             _, shard_size = get_glm_dsa_cp_layer_shard_info(kvc)
             if shard_size > 1:
+                # Preserve the existing LayerSplit sizing semantics. GLM-5.3
+                # hybrid-layer support is intentionally limited to the normal
+                # (non-LayerSplit) pool below.
+                active_indexer_layers = [
+                    layer_id
+                    for layer_id in range(
+                        kvc.layer_info.start_layer, kvc.layer_info.end_layer
+                    )
+                    if not dsa_layer_skips_topk(kvc.model_config.hf_config, layer_id)
+                ]
                 active_set = set(active_indexer_layers)
                 max_owned = 0
                 for rank in range(shard_size):
@@ -444,7 +468,10 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
                     )
                 num_indexer_layers = max_owned + 1
             else:
-                num_indexer_layers = len(active_indexer_layers)
+                num_indexer_layers = sum(
+                    not dsa_layer_skips_topk(kvc.model_config.hf_config, layer_id)
+                    for layer_id in _get_dsa_cache_layer_ids(kvc, num_layers)
+                )
 
         return int(
             indexer_size_per_token * num_indexer_layers * element_size * indexer_ratio

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -1068,12 +1068,15 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         and stash the returned per-bucket metadata object; otherwise fall
         back to the generic eager init that BCG/TC_PIECEWISE use today."""
         attn_backend = self.model_runner.attn_backend
+        with forward_context(ForwardContext(attn_backend=attn_backend)):
             if not self.use_captured_attn_metadata:
                 attn_backend.init_forward_metadata(forward_batch)
                 return
-        metadata = attn_backend.init_forward_metadata_for_breakable_cuda_graph_capture(
+            metadata = (
+                attn_backend.init_forward_metadata_for_breakable_cuda_graph_capture(
                     forward_batch
                 )
+            )
             assert self.attn_metadata_buffers is not None
             self.attn_metadata_buffers[num_tokens] = metadata
 

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -154,16 +154,18 @@ def _uninstall_wait_stream_hook():
 
 
 def _weak_ref_if_tensor(x):
-    """Return a weak-ref tensor view (shared storage, no refcount) for tensors;
+    """Return a weak-ref view for nonempty accelerator tensors; recurse into
-    recurse into tuples/lists; pass-through for non-tensors. Weak-ref'ing
+    tuples/lists and keep CPU, empty, and non-tensor values unchanged.
-    captured args lets the shared mempool reclaim per-layer intermediates
+    Weak-ref'ing captured args lets the shared mempool reclaim per-layer
-    between segments — storage stays alive for each segment CUDAGraph's
+    intermediates between segments — storage stays alive for each segment
-    lifetime via its pool use_count.
+    CUDAGraph's lifetime via its pool use_count.
 
     weak_ref_tensors is imported lazily because it hard-raises on
     platforms without a CUDA/HIP/NPU backend; we only reach this code during
     an active Breakable capture, which runs only on those backends."""
     if torch.is_tensor(x):
+        if x.numel() == 0 or x.device.type == "cpu":
+            return x
         from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
 
         return weak_ref_tensors(x)

python/sglang/srt/model_loader/loader.py

@@ -259,6 +259,10 @@ def _get_quantization_config(
                 f"method {model_config.quantization}. Supported dtypes: "
                 f"{supported_dtypes}"
             )
+        get_hf_to_sglang_mapper = getattr(model_class, "get_hf_to_sglang_mapper", None)
+        if get_hf_to_sglang_mapper is not None:
+            hf_to_sglang_mapper = get_hf_to_sglang_mapper(model_config.hf_config)
+        else:
             hf_to_sglang_mapper = getattr(model_class, "hf_to_sglang_mapper", None)
         # pass mappings by reference to quant_config
         if hf_to_sglang_mapper is not None and quant_config is not None:

python/sglang/srt/model_loader/weight_utils.py

@@ -1571,11 +1571,18 @@ def row_parallel_weight_loader(
 LoaderFunction = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
 
 
-def sharded_weight_loader(shard_axis: int) -> LoaderFunction:
+def sharded_weight_loader(
+    shard_axis: int,
+    tp_rank_getter=None,
+) -> LoaderFunction:
     """Create a weight loader that shards the weights along the given axis"""
 
     def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
-        tp_rank = get_parallel().attn_tp_rank
+        tp_rank = (
+            tp_rank_getter()
+            if tp_rank_getter is not None
+            else get_parallel().attn_tp_rank
+        )
 
         shard_size = param.data.shape[shard_axis]
         start_idx = tp_rank * shard_size

python/sglang/srt/models/deepseek_common/attention_backend_handler.py

@@ -118,9 +118,9 @@ def _handle_attention_backend(attn, forward_batch, backend_name):
         return _dispatch_mla_subtype(attn, forward_batch)
 
     sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
-    disable_ragged = (
+    disable_ragged = (backend_name in ["flashinfer", "flashmla"]) and (
-        backend_name in ["flashinfer", "flashmla"]
+        attn.flashinfer_mla_disable_ragged or attn.qk_rope_head_dim == 0
-    ) and attn.flashinfer_mla_disable_ragged
+    )
 
     if (
         not disable_ragged

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -677,9 +677,22 @@ class DeepseekMHAForwardMixin:
     def _concat_and_cast_mha_k(
         self: DeepseekV2AttentionMLA,
         k_nope: torch.Tensor,
-        k_pe: torch.Tensor,
+        k_pe: torch.Tensor | None,
         forward_batch: ForwardBatch,
     ):
+        if self.qk_rope_head_dim == 0:
+            assert k_pe is None or k_pe.shape[-1] == 0
+            k = k_nope.contiguous()
+            if (
+                _is_cuda
+                and self.current_attention_backend == "fa3"
+                and self.kv_cache_dtype != "auto"
+            ):
+                # fa3 requires k in the pool dtype when KV cache is fp8; the
+                # concat branch below does the same cast for roped models.
+                k = k.to(get_token_to_kv_pool().dtype)
+            return k
+
         # Temporary for DeepSeek V3/R1 only, but can generalize if needed
         k_shape = (k_nope.shape[0], self.num_local_heads, self.qk_head_dim)
         if (

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -960,6 +960,8 @@ class DeepseekMLAForwardMixin:
         """
         Check if we should skip rope and do fused rope+quantize for TRTLLM MLA decode in fp8_e4m3 path.
         """
+        if self.rotary_emb is None:
+            return False
         if self.current_attention_backend in ("dsa", "nsa"):
             return (
                 get_exec().kernel.dsa_decode_backend == "trtllm"

python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py

@@ -570,7 +570,7 @@ class DeepseekV2WeightLoaderMixin:
                 for name in weight_names:
                     if "kv_b_proj" in name:
                         layer_id = int(name.split(".")[2])
-                        if layer_id < self.config.num_hidden_layers:
+                        if self.model.start_layer <= layer_id < self.model.end_layer:
                             layer_ids.add(layer_id)
 
         for layer_id in layer_ids:
@@ -580,6 +580,9 @@ class DeepseekV2WeightLoaderMixin:
                 else self.model.decoder.self_attn
             )
 
+            if not hasattr(self_attn, "kv_b_proj"):
+                continue
+
             if hasattr(self_attn.kv_b_proj, "qweight"):
                 # awq compatible, dequantize the weight if supported
                 awq_dequantize_f = awq_dequantize_func()

python/sglang/srt/models/deepseek_nextn.py

@@ -182,6 +182,7 @@ class DeepseekModelNextN(nn.Module):
             is_nextn=True,
             prefix=add_prefix(layer_name, prefix),
             alt_stream=self.alt_stream,
+            skip_rope=config.qk_rope_head_dim == 0,
             dsa_enable_prefill_cp=self.dsa_enable_prefill_cp,
             mla_enable_prefill_cp=self.mla_enable_prefill_cp,
         )
@@ -220,8 +221,27 @@ class DeepseekModelNextN(nn.Module):
             )
 
             if input_embeds is None:
-                hidden_states = self.embed_tokens(input_ids)
+                # MM positions in input_ids hold MM_PAD_SHIFT_VALUE+hash sentinels
-            else:
+                # (far above vocab_size). Use target-produced mm_input_embeds for
+                # these positions and only call embed_tokens on the appended
+                # next-token to avoid embed OOB.
+                input_embeds = forward_batch.mm_input_embeds
+                if (
+                    forward_batch.forward_mode.is_extend()
+                    and forward_batch.contains_mm_inputs()
+                    and not forward_batch.forward_mode.is_draft_extend_v2()
+                ):
+                    assert input_embeds is not None
+                    last_indices = (
+                        forward_batch.extend_start_loc
+                        + forward_batch.extend_seq_lens
+                        - 1
+                    ).long()
+                    input_embeds[last_indices] = self.embed_tokens(
+                        input_ids[last_indices]
+                    )
+                if input_embeds is None:
+                    input_embeds = self.embed_tokens(input_ids)
             hidden_states = input_embeds
 
             if hidden_states.shape[0] > 0:
@@ -320,6 +340,10 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
         },
     )
 
+    @classmethod
+    def get_hf_to_sglang_mapper(cls, config) -> WeightsMapper:
+        return cls.hf_to_sglang_mapper
+
     def _resolve_nextn_quant_config(self, config, quant_config):
         if quant_config is None or quant_config.get_name() != "quark":
             return quant_config
@@ -327,7 +351,7 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
         from sglang.srt.layers.quantization.quark.utils import should_ignore_layer
 
         ckpt_prefix = f"model.layers.{config.num_hidden_layers}"
-        mapped_prefix = self.hf_to_sglang_mapper._map_name(ckpt_prefix)
+        mapped_prefix = self.get_hf_to_sglang_mapper(config)._map_name(ckpt_prefix)
         if should_ignore_layer(mapped_prefix, quant_config.exclude_layers):
             return None
         return quant_config

python/sglang/srt/models/deepseek_v2.py

@@ -45,6 +45,7 @@ from sglang.srt.configs.model_config import (
     compute_mla_mscale_scaling,
     dsa_layer_skips_topk,
     get_dsa_index_head_dim,
+    get_dsa_index_kpool,
     get_dsa_index_n_heads,
     get_dsa_index_topk,
     is_deepseek_dsa,
@@ -62,6 +63,7 @@ from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.amx_utils import PackWeightMethod
 from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer
+from sglang.srt.layers.attention.dsa.dsa_indexer_kpool import IndexerKPool
 from sglang.srt.layers.attention.dsa.utils import (
     can_dsa_cp_split,
     dsa_use_prefill_cp,
@@ -1860,7 +1862,10 @@ class DeepseekV2AttentionMLA(
 
             if not self.skip_topk or is_nextn:
                 is_neox_style = not getattr(config, "indexer_rope_interleave", False)
-                self.indexer = Indexer(
+                indexer_cls = (
+                    IndexerKPool if get_dsa_index_kpool(config) > 1 else Indexer
+                )
+                indexer_kwargs = dict(
                     hidden_size=hidden_size,
                     index_n_heads=get_dsa_index_n_heads(config),
                     index_head_dim=get_dsa_index_head_dim(config),
@@ -1879,6 +1884,9 @@ class DeepseekV2AttentionMLA(
                     alt_stream=alt_stream,
                     config=config,
                 )
+                if indexer_cls is IndexerKPool:
+                    indexer_kwargs["skip_rope"] = skip_rope
+                self.indexer = indexer_cls(**indexer_kwargs)
 
         self.kv_b_proj = ColumnParallelLinear(
             self.kv_lora_rank,
@@ -1902,7 +1910,7 @@ class DeepseekV2AttentionMLA(
         )
         self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
 
-        if not skip_rope:
+        if not skip_rope and qk_rope_head_dim > 0:
             is_neox_style = not getattr(config, "rope_interleave", True)
             self.rotary_emb = get_rope_wrapper(
                 qk_rope_head_dim,
@@ -2318,6 +2326,7 @@ class DeepseekV2DecoderLayer(nn.Module):
         is_nextn: bool = False,
         prefix: str = "",
         alt_stream: Optional[torch.cuda.Stream] = None,
+        skip_rope: bool = False,
         dsa_enable_prefill_cp: bool = False,
         mla_enable_prefill_cp: bool = False,
     ) -> None:
@@ -2340,6 +2349,10 @@ class DeepseekV2DecoderLayer(nn.Module):
         self.mla_enable_prefill_cp = mla_enable_prefill_cp
         self.layer_id = layer_id
         self.is_nextn = is_nextn
+        if is_nextn and getattr(config, "mla_nope", False):
+            # The NextN draft must match the NoPE target layers, or its Q/K
+            # and the KV it verifies against live in different spaces.
+            skip_rope = True
         self.self_attn = DeepseekV2AttentionMLA(
             config=config,
             hidden_size=self.hidden_size,
@@ -2359,6 +2372,7 @@ class DeepseekV2DecoderLayer(nn.Module):
             reduce_results=False,
             prefix=add_prefix("self_attn", prefix),
             alt_stream=alt_stream,
+            skip_rope=skip_rope,
             is_nextn=is_nextn,
             dsa_enable_prefill_cp=dsa_enable_prefill_cp,
             mla_enable_prefill_cp=mla_enable_prefill_cp,
@@ -2689,6 +2703,7 @@ class DeepseekV2Model(nn.Module):
                 quant_config=quant_config,
                 prefix=prefix,
                 alt_stream=self.alt_stream,
+                skip_rope=config.qk_rope_head_dim == 0,
                 dsa_enable_prefill_cp=self.dsa_enable_prefill_cp,
                 mla_enable_prefill_cp=self.mla_enable_prefill_cp,
             ),

python/sglang/srt/models/glm5_next_nextn.py

@@ -0,0 +1,80 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import logging
+
+from sglang.srt.models.deepseek_nextn import DeepseekV3ForCausalLMNextN
+from sglang.srt.models.glm5_next import Glm5NextForConditionalGeneration
+from sglang.srt.models.utils import WeightsMapper
+
+logger = logging.getLogger(__name__)
+
+
+class Glm5NextForConditionalGenerationNextN(DeepseekV3ForCausalLMNextN):
+    @classmethod
+    def get_hf_to_sglang_mapper(cls, config) -> WeightsMapper:
+        text_config = getattr(config, "text_config", config)
+        return WeightsMapper(
+            orig_to_new_substr={
+                f"model.layers.{text_config.num_hidden_layers}": "model.decoder",
+            },
+        )
+
+    def _resolve_nextn_quant_config(self, config, quant_config):
+        """Mixed checkpoints list the BF16 NextN block in ``quantization_config.ignore``;
+        inheriting global FP8 quantization would corrupt its QKV weights."""
+        raw_quant_config = getattr(config, "quantization_config", None) or {}
+        if hasattr(raw_quant_config, "to_dict"):
+            raw_quant_config = raw_quant_config.to_dict()
+        ignored = (
+            raw_quant_config.get("ignore", [])
+            if isinstance(raw_quant_config, dict)
+            else []
+        )
+        nextn_layer_pattern = f"model.layers.{config.num_hidden_layers}.*"
+        if nextn_layer_pattern in ignored:
+            logger.warning(
+                "GLM5 NextN layer %s is checkpoint-declared unquantized; "
+                "using BF16 draft modules",
+                nextn_layer_pattern,
+            )
+            return None
+        return super()._resolve_nextn_quant_config(config, quant_config)
+
+    def __init__(self, config, quant_config=None, prefix: str = "") -> None:
+        super().__init__(
+            getattr(config, "text_config", config),
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+    def load_weights(self, weights):
+        if not hasattr(self, "fuse_qkv_a_proj"):
+            self.fuse_qkv_a_proj = getattr(self.config, "q_lora_rank", None) is not None
+        layer_id = self.config.num_hidden_layers
+        layer_prefixes = (
+            f"model.layers.{layer_id}.",
+            f"model.language_model.layers.{layer_id}.",
+        )
+        nextn_weights = (
+            (name, weight)
+            for name, weight in weights
+            if name.startswith(layer_prefixes)
+        )
+        return Glm5NextForConditionalGeneration.load_weights(
+            self, nextn_weights, is_nextn=True
+        )
+
+
+EntryClass = [Glm5NextForConditionalGenerationNextN]

python/sglang/srt/models/glm_ocr.py

@@ -26,7 +26,6 @@ import torch.nn as nn
 from einops import rearrange
 from transformers.models.glm_ocr.configuration_glm_ocr import (
     GlmOcrConfig,
-    GlmOcrTextConfig,
     GlmOcrVisionConfig,
 )
 
@@ -158,7 +157,6 @@ class GlmOcrVisionModel(Glm4vVisionModel):
     def __init__(
         self,
         vision_config: GlmOcrVisionConfig,
-        text_config: GlmOcrTextConfig,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
         use_data_parallel: bool = False,
@@ -209,9 +207,16 @@ class GlmOcrVisionModel(Glm4vVisionModel):
                 for layer_idx in range(depth)
             ]
         )
+        projection_intermediate_size = getattr(
+            vision_config, "projection_intermediate_size", None
+        )
         self.merger = GlmOcrVisionPatchMerger(
             d_model=vision_config.out_hidden_size,
-            context_dim=text_config.intermediate_size,
+            context_dim=(
+                projection_intermediate_size
+                if projection_intermediate_size is not None
+                else vision_config.intermediate_size
+            ),
             quant_config=quant_config,
             bias=False,
             prefix=add_prefix("merger", prefix),
@@ -285,7 +290,6 @@ class GlmOcrForConditionalGeneration(Glm4vForConditionalGeneration):
         self.use_data_parallel = get_mm().mm_enable_dp_encoder
         self.visual = GlmOcrVisionModel(
             vision_config=config.vision_config,
-            text_config=config.text_config,
             quant_config=quant_config,
             prefix=add_prefix("visual", prefix),
             use_data_parallel=self.use_data_parallel,

python/sglang/srt/multimodal/mm_utils.py

@@ -823,3 +823,67 @@ def run_dp_sharded_mrope_vision_model(
             current_idx += count
     out_embeddings = torch.cat(original_order_embeddings, dim=0)
     return out_embeddings
+
+
+def run_dp_presharded_mrope_vision_model(
+    vision_model: torch.nn.Module,
+    pixel_values_local: torch.Tensor,
+    local_grid_thw_list: list,
+    global_grid_thw_list: list,
+    gpu_sample_counts: list,
+) -> torch.Tensor:
+    """Rank-local shards are contiguous, so rank-order concatenation restores global video order."""
+    parallel = get_parallel()
+    tp_size = parallel.attn_tp_size
+    patches_per_unit = [math.prod(grid) for grid in global_grid_thw_list]
+    grouped_patch_counts = []
+    offset = 0
+    for rank in range(tp_size):
+        count = gpu_sample_counts[rank]
+        grouped_patch_counts.append(sum(patches_per_unit[offset : offset + count]))
+        offset += count
+
+    merge_factor = vision_model.spatial_merge_size**2
+    grouped_output_lengths = [
+        patch_count // merge_factor for patch_count in grouped_patch_counts
+    ]
+    max_output_length = max(grouped_output_lengths)
+    try:
+        model_device = vision_model.device
+        model_dtype = vision_model.dtype
+    except AttributeError:
+        parameter = next(vision_model.parameters())
+        model_device, model_dtype = parameter.device, parameter.dtype
+
+    if pixel_values_local.shape[0] > 0:
+        pixel_values_local = pixel_values_local.to(
+            device=model_device, dtype=model_dtype
+        )
+        local_embeddings = vision_model(
+            pixel_values_local,
+            grid_thw=torch.tensor(local_grid_thw_list),
+        )
+    else:
+        local_embeddings = torch.empty(
+            (0, vision_model.out_hidden_size),
+            device=model_device,
+            dtype=model_dtype,
+        )
+
+    if local_embeddings.shape[0] < max_output_length:
+        padding = torch.empty(
+            (
+                max_output_length - local_embeddings.shape[0],
+                local_embeddings.shape[1],
+            ),
+            device=local_embeddings.device,
+            dtype=local_embeddings.dtype,
+        )
+        local_embeddings = torch.cat([local_embeddings, padding], dim=0)
+
+    gathered = parallel.attn_tp_group.all_gather(local_embeddings, dim=0)
+    pieces = []
+    for rank, output_length in enumerate(grouped_output_lengths):
+        start = rank * max_output_length
+        pieces.append(gathered[start : start + output_length])
+    return torch.cat(pieces, dim=0)

python/sglang/srt/multimodal/processors/base_processor.py

@@ -965,11 +965,14 @@ class BaseMultimodalProcessor(ABC):
                 img, _ = load_image(data, cls.gpu_image_decode)
                 if isinstance(img, torch.Tensor):
                     return img  # JPEG already decoded on GPU by nvJPEG
+                # PIL decodes lazily; do it here in the io worker so the decode
+                # doesn't run later on the event-loop thread.
                 if discard_alpha_channel:
                     if cls.smart_rgb_conversion:
                         return smart_to_rgb(img)
                     if img.mode != "RGB":
                         return img.convert("RGB")
+                img.load()
                 return img
             elif modality == Modality.VIDEO:
                 return load_video(data, frame_count_limit)

python/sglang/srt/multimodal/processors/glm4v.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import math
 from typing import List, Tuple, Union
 
@@ -77,14 +78,81 @@ def split_glm_video_items(mm_data):
     return urls, configs
 
 
+def glm_budget_kwargs(processor, user_max_image_tokens=None, count=1, split=False):
+    if processor is None:
+        return None
+    default_max = getattr(processor, "max_image_tokens", None)
+    if not default_max:
+        return None
+
+    if user_max_image_tokens is not None:
+        budget = int(user_max_image_tokens)
+    elif split:
+        budget = int(default_max)
+    else:
+        return None
+
+    count = max(int(count or 1), 1)
+    effective = max(1, budget // count if split and count > 1 else budget)
+    if effective == default_max and user_max_image_tokens is None:
+        return None
+    return {"max_image_tokens": effective}
+
+
+def glm_max_image_tokens_from_configs(configs):
+    values = [
+        int(config["max_image_tokens"])
+        for config in configs or []
+        if isinstance(config, dict) and config.get("max_image_tokens") is not None
+    ]
+    return min(values) if values else None
+
+
 def glm_processor_video_config(processor):
     if processor is None:
         return {}
-    return {
+    config = {
         key: value
         for key in GLM_MEDIA_CONFIG_KEYS
         if (value := getattr(processor, key, None)) is not None
     }
+    budget = _glm_processor_resize_budget(processor)
+    if budget is not None:
+        config["_presize_budget"] = budget
+    return config
+
+
+def _glm_processor_resize_budget(processor):
+    """Use token limits because Glm5NextVideoProcessor.size.longest_edge is only a sentinel."""
+    max_image_tokens = getattr(processor, "max_image_tokens", None)
+    if not max_image_tokens:
+        return None
+    patch_size = getattr(processor, "patch_size", None) or GLM_VIDEO_PATCH_SIZE
+    merge_size = getattr(processor, "merge_size", None) or GLM_VIDEO_MERGE_SIZE
+    expand_factor = getattr(processor, "patch_expand_factor", None) or 1
+    temporal_factor = getattr(processor, "temporal_patch_size", None) or 2
+    pixels_per_token = int(temporal_factor * (patch_size * merge_size) ** 2)
+    return {
+        "factor": int(patch_size * merge_size * expand_factor),
+        "temporal_factor": int(temporal_factor),
+        "pixels_per_token": pixels_per_token,
+        "min_pixels": int(getattr(processor, "min_image_tokens", None) or 0)
+        * pixels_per_token,
+        "max_pixels": int(max_image_tokens) * pixels_per_token,
+        "resize_mode": getattr(processor, "resize_mode", None) or "resize",
+    }
+
+
+def _glm_effective_presize_budget(video_config, effective_max_image_tokens):
+    budget = video_config.get("_presize_budget") if video_config else None
+    if not budget or effective_max_image_tokens is None:
+        return video_config
+    config = dict(video_config)
+    config["_presize_budget"] = {
+        **budget,
+        "max_pixels": int(effective_max_image_tokens) * budget["pixels_per_token"],
+    }
+    return config
 
 
 def _merge_glm_video_configs(default_config, item_configs):
@@ -204,6 +272,122 @@ def _resize_frames_to_max_tokens(frames, max_tokens_per_frame):
     return nchw.permute(0, 2, 3, 1).contiguous()
 
 
+def preprocess_video_frames_sync(frame_list: List[dict]):
+    total_num_frames = len(frame_list)
+    if total_num_frames == 0:
+        raise ValueError("GLM video frame list must not be empty")
+    duration = 0.0
+    if frame_list[0].get("detail") is not None:
+        details = json.loads(frame_list[0]["detail"])
+        duration = float(details.get("video_duration", 0))
+    if duration == 0:
+        base_ts = float(frame_list[0].get("timestamp", 0) or 0)
+        duration = float(frame_list[-1].get("timestamp", base_ts) or base_ts) - base_ts
+    images = [frame["frame_image"] for frame in frame_list]
+    if isinstance(images[0], torch.Tensor):
+        images = torch.stack(images).permute(0, 2, 3, 1).contiguous()
+    else:
+        images = [np.asarray(image) for image in images]
+    fps = total_num_frames / duration if duration else 0
+    return images, _glm_video_metadata(
+        total_num_frames, fps, duration, range(total_num_frames)
+    )
+
+
+GLM_VIDEO_PRE_RESIZE_CHUNK = 64
+
+
+def _vendor_smart_resize_canvas(
+    num_frames, height, width, *, temporal_factor, factor, min_pixels, max_pixels
+):
+    """Replica of the vendor Glm5Next smart_resize (align-ceil + budget search)."""
+
+    def align(value):
+        return math.ceil(value / factor) * factor
+
+    def fit_within_budget(aligned_frames):
+        low, high = 1, height
+        best_height, best_width = factor, factor
+        while low <= high:
+            content_height = (low + high) // 2
+            content_width = max(1, math.floor(width * content_height / height))
+            candidate_height = align(content_height)
+            candidate_width = align(content_width)
+            if aligned_frames * candidate_height * candidate_width <= max_pixels:
+                best_height, best_width = candidate_height, candidate_width
+                low = content_height + 1
+            else:
+                high = content_height - 1
+        return best_height, best_width
+
+    aligned_frames = max(
+        temporal_factor, round(num_frames / temporal_factor) * temporal_factor
+    )
+    canvas_height, canvas_width = align(height), align(width)
+    if aligned_frames * canvas_height * canvas_width > max_pixels:
+        canvas_height, canvas_width = fit_within_budget(aligned_frames)
+    elif aligned_frames * canvas_height * canvas_width < min_pixels:
+        scale = math.sqrt(min_pixels / (num_frames * height * width))
+        canvas_height = align(max(1, math.ceil(height * scale)))
+        canvas_width = align(max(1, math.ceil(width * scale)))
+        if aligned_frames * canvas_height * canvas_width > max_pixels:
+            canvas_height, canvas_width = fit_within_budget(aligned_frames)
+    return canvas_height, canvas_width
+
+
+def _pre_resize_frames_for_processor(
+    frames,
+    *,
+    factor,
+    temporal_factor,
+    pixels_per_token,
+    min_pixels,
+    max_pixels,
+    resize_mode,
+):
+    """Pre-resize in chunks to avoid HF's native-resolution float32 intermediate while preserving its output grid."""
+    import torchvision.transforms.functional as TF
+
+    if not isinstance(frames, torch.Tensor):
+        frames = torch.from_numpy(np.asarray(frames))
+    nchw = frames.permute(0, 3, 1, 2)
+    num_frames, _, height, width = nchw.shape
+    canvas_height, canvas_width = _vendor_smart_resize_canvas(
+        num_frames,
+        height,
+        width,
+        temporal_factor=temporal_factor,
+        factor=factor,
+        min_pixels=min_pixels,
+        max_pixels=max_pixels,
+    )
+    if resize_mode == "resize":
+        content_height, content_width = canvas_height, canvas_width
+    else:
+        scale = min(canvas_height / height, canvas_width / width)
+        if num_frames * height * width >= min_pixels:
+            scale = min(1.0, scale)
+        content_height = max(1, min(canvas_height, math.floor(height * scale)))
+        content_width = max(1, min(canvas_width, math.floor(width * scale)))
+    if (content_height, content_width) != (height, width):
+        nchw = torch.cat(
+            [
+                TF.resize(
+                    chunk,
+                    [content_height, content_width],
+                    interpolation=TF.InterpolationMode.BICUBIC,
+                    antialias=True,
+                )
+                for chunk in nchw.split(GLM_VIDEO_PRE_RESIZE_CHUNK)
+            ]
+        )
+    if (content_height, content_width) != (canvas_height, canvas_width):
+        nchw = torch.nn.functional.pad(
+            nchw, (0, canvas_width - content_width, 0, canvas_height - content_height)
+        )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
 def glm_decode_frames_at(vr, indices, video_config=None):
     indices = list(indices)
     if not indices:
@@ -216,6 +400,8 @@ def glm_decode_frames_at(vr, indices, video_config=None):
     max_tokens_per_frame = video_config.get("max_tokens_per_frame")
     if max_tokens_per_frame is not None:
         frames = _resize_frames_to_max_tokens(frames, max_tokens_per_frame)
+    elif budget := video_config.get("_presize_budget"):
+        frames = _pre_resize_frames_for_processor(frames, **budget)
     return frames
 
 
@@ -372,12 +558,21 @@ class Glm4vImageProcessor(SGLangBaseProcessor):
         )
 
         video_metadata = None
+        videos_kwargs = None
         if base_output.videos and not isinstance(base_output.videos[0], dict):
+            videos_kwargs = glm_budget_kwargs(
+                video_processor,
+                user_max_image_tokens=glm_max_image_tokens_from_configs(video_configs),
+                count=len(base_output.videos),
+                split=True,
+            )
+            effective_max_image_tokens = (videos_kwargs or {}).get("max_image_tokens")
             loop = asyncio.get_running_loop()
             decode_tasks = []
             for index, video in enumerate(base_output.videos):
-                video_config = (
+                video_config = _glm_effective_presize_budget(
-                    video_configs[index] if index < len(video_configs) else {}
+                    video_configs[index] if index < len(video_configs) else {},
+                    effective_max_image_tokens,
                 )
                 if isinstance(video, VideoDecoderWrapper):
                     decode_tasks.append(
@@ -389,6 +584,14 @@ class Glm4vImageProcessor(SGLangBaseProcessor):
                             video_processor,
                         )
                     )
+                elif isinstance(video, list) and (
+                    not video or isinstance(video[0], dict)
+                ):
+                    decode_tasks.append(
+                        loop.run_in_executor(
+                            self.io_executor, preprocess_video_frames_sync, video
+                        )
+                    )
                 else:
                     decode_tasks.append(
                         asyncio.sleep(
@@ -407,18 +610,23 @@ class Glm4vImageProcessor(SGLangBaseProcessor):
                     close = getattr(video, "close", None)
                     if callable(close):
                         close()
-            base_output.videos, video_metadata = map(list, zip(*videos_processed))
+            base_output.videos, metadata = map(list, zip(*videos_processed))
+            if metadata and all(item is not None for item in metadata):
+                video_metadata = metadata
 
         combine_kwargs = {}
         if video_metadata is not None:
             # Skip HF resampling because these frames already carry their original indices.
             combine_kwargs["video_metadata"] = video_metadata
             combine_kwargs["do_sample_frames"] = False
-            combine_kwargs["processor_video_config"] = {
+            processor_video_config = {
                 key: value
                 for key, value in self.video_config.items()
                 if key not in {"fps", "max_frames", "max_tokens_per_frame"}
             }
+            if videos_kwargs is not None:
+                processor_video_config.update(videos_kwargs)
+            combine_kwargs["processor_video_config"] = processor_video_config
 
         mm_items, input_ids, ret = await self.process_and_combine_mm_data_async(
             base_output, self.mm_tokens, **combine_kwargs

python/sglang/srt/speculative/adaptive_spec_params.py

@@ -25,7 +25,7 @@ logger = logging.getLogger(__name__)
 
 DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = {
     "1": {
-        "candidate_steps": [1, 3, 7],
+        "candidate_steps": [1, 3, 5, 7],
         "up_hysteresis": 0.0,
         "down_hysteresis": -0.25,
         "ceiling_coeff": 0,

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -237,7 +237,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
 
         dsa_seed_topk = (
             torch.zeros(
-                (self.max_bs, self.eagle_worker.dsa_index_topk),
+                (self.max_bs, self.eagle_worker.dsa_seed_topk_width),
                 dtype=torch.int32,
                 device=model_runner.device,
             )

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -55,6 +55,18 @@ if TYPE_CHECKING:
     from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
 
 
+def resolve_draft_extend_seq_len_fill_value(
+    attn_backend, captured_req_width: int
+) -> int:
+    """Pad synthetic history past the fixed draft-width subtraction and KPool offset."""
+    fill_value = attn_backend.get_cuda_graph_seq_len_fill_value()
+    full_attn_backend = getattr(attn_backend, "full_attn_backend", attn_backend)
+    dsa_index_kpool = getattr(full_attn_backend, "dsa_index_kpool", 1)
+    if dsa_index_kpool > 1:
+        fill_value = max(fill_value, captured_req_width + dsa_index_kpool)
+    return fill_value
+
+
 @dataclass
 class EagleDraftExtendInputBuffers(ForwardInputBuffers):
     input_ids: torch.Tensor
@@ -142,8 +154,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         self.draft_extend_attn_backend.init_cuda_graph_state(
             self.max_bs, self.max_num_token
         )
-        self.seq_len_fill_value = (
+        self.seq_len_fill_value = resolve_draft_extend_seq_len_fill_value(
-            self.draft_extend_attn_backend.get_cuda_graph_seq_len_fill_value()
+            self.draft_extend_attn_backend, self.captured_req_width
         )
         self.extend_seq_lens_cpu = [self.captured_req_width] * self.max_bs
 
@@ -176,8 +188,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
                 if _hidden_size is not None
                 else None
             )
-            self.seq_len_fill_value = (
+            self.seq_len_fill_value = resolve_draft_extend_seq_len_fill_value(
-                self.draft_extend_attn_backend.get_cuda_graph_seq_len_fill_value()
+                self.draft_extend_attn_backend, self.captured_req_width
             )
             seq_lens = torch.full(
                 (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64
@@ -243,7 +255,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
 
         dsa_seed_topk_capture = (
             torch.full(
-                (self.max_num_token, self.eagle_worker.dsa_index_topk),
+                (self.max_num_token, self.eagle_worker.dsa_seed_topk_width),
                 -1,
                 dtype=torch.int32,
                 device=model_runner.device,

python/sglang/srt/speculative/eagle_utils.py

@@ -507,12 +507,19 @@ def get_draft_recurrent_hidden_state_spec(
     )
 
 
+_PREPARE_FOR_VERIFY_DEPS = None
+
+
 def eagle_prepare_for_verify(
     verify_input: EagleVerifyInput,
     req_to_token_pool: ReqToTokenPool,
     batch: ScheduleBatch,
     target_worker: TpModelWorker,
 ):
+    # Imports must stay lazy (import-cycle safety) but only need to resolve
+    # once, not on every decode cycle of this hot path.
+    global _PREPARE_FOR_VERIFY_DEPS
+    if _PREPARE_FOR_VERIFY_DEPS is None:
         from sglang.kernels.ops.speculative.cache_locs import (
             assign_extend_cache_locs_uniform_func,
         )
@@ -523,6 +530,21 @@ def eagle_prepare_for_verify(
         )
         from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify
 
+        _PREPARE_FOR_VERIFY_DEPS = (
+            assign_extend_cache_locs_uniform_func,
+            CaptureHiddenMode,
+            ForwardBatch,
+            ForwardMode,
+            prepare_mamba_track_for_verify,
+        )
+    (
+        assign_extend_cache_locs_uniform_func,
+        CaptureHiddenMode,
+        ForwardBatch,
+        ForwardMode,
+        prepare_mamba_track_for_verify,
+    ) = _PREPARE_FOR_VERIFY_DEPS
+
     if not batch.forward_mode.is_idle():
         # Assign cache locations
         bs = len(batch.req_pool_indices)

python/sglang/srt/speculative/eagle_worker_common.py

@@ -495,6 +495,10 @@ def run_eagle_verify(
     # Batch 1: Target verify
     # Prepare for target verify in a separate stream
     with plan_stream_ctx:
+        if plan_stream is not None:
+            # Verify prep copies draft-produced tree metadata on the plan stream,
+            # so it must not start before the draft frontier.
+            plan_stream.wait_stream(fwd_stream)
         verify_forward_batch, can_run_cuda_graph = eagle_prepare_for_verify(
             verify_input,
             req_to_token_pool,

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -7,6 +7,7 @@ from typing import List, Optional
 import torch
 
 from sglang.kernels.ops.speculative.topk1 import draft_topk1_postprocess
+from sglang.srt.configs.model_config import get_dsa_mtp_topk_width
 from sglang.srt.distributed import get_pp_group
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
@@ -269,8 +270,13 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         # GLM-5.2 MTP IndexShare: seed reused indexer top-k from draft-extend
         # (last verified token), not draft-decode step 0.
         self.dsa_index_topk = getattr(hf_config, "index_topk", None)
+        self.dsa_seed_topk_width = (
+            get_dsa_mtp_topk_width(hf_config)
+            if self.index_share_for_mtp_iteration and self.dsa_index_topk is not None
+            else None
+        )
         self.seed_dsa_topk_from_draft_extend = (
-            self.index_share_for_mtp_iteration and self.dsa_index_topk is not None
+            self.index_share_for_mtp_iteration and self.dsa_seed_topk_width is not None
         )
 
     def init_token_map(self):
@@ -801,16 +807,29 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         if not batch.forward_mode.is_idle():
             # Chunked-prefill-aware tail tokens (see PR #26329).
             tail_tokens = _eagle_prefill_tail_tokens(batch, next_token_ids)
+
             new_input_ids = torch.empty_like(batch.input_ids)
+            if mm_input_embeds is not None:
+                # Rotate mm embeddings the same way as input_ids: shift left by
+                # one per request so they stay aligned with the rotated ids. The
+                # last position per request is filled by the draft model's own
+                # embed_tokens lookup on next_token_ids (see DeepseekModelNextN).
+                rotated_mm = torch.empty_like(mm_input_embeds)
             pt = 0
             for i, extend_len in enumerate(batch.extend_lens):
                 input_ids = batch.input_ids[pt : pt + extend_len]
                 new_input_ids[pt : pt + extend_len].copy_(
                     torch.cat((input_ids[1:], tail_tokens[i].reshape(1)))
                 )
+                if mm_input_embeds is not None:
+                    rotated_mm[pt : pt + extend_len - 1].copy_(
+                        mm_input_embeds[pt + 1 : pt + extend_len]
+                    )
                 pt += extend_len
             assert pt == batch.input_ids.numel()
             batch.input_ids = new_input_ids
+            if mm_input_embeds is not None:
+                mm_input_embeds = rotated_mm
 
         # Draft-extend spec_info for the extend forward; carries only
         # hidden_states + shape info.
@@ -897,11 +916,10 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         )
 
     def _get_dsa_extend_topk_buf(self, num_tokens: int) -> torch.Tensor:
-        """Lazily-grown int32 [num_tokens, index_topk] eager draft-extend seed buffer."""
         buf = self.dsa_extend_topk_buf
         if buf is None or buf.shape[0] < num_tokens:
             buf = torch.full(
-                (num_tokens, self.dsa_index_topk),
+                (num_tokens, self.dsa_seed_topk_width),
                 -1,
                 dtype=torch.int32,
                 device=self.device,

python/sglang/srt/speculative/spec_info.py

@@ -10,6 +10,7 @@ import torch
 from sglang.srt.arg_groups.overrides import resolving_view
 from sglang.srt.runtime_context import get_spec as get_spec_config
 from sglang.srt.speculative.spec_registry import (
+    _RESERVED_ALIASES,
     CustomSpecAlgo,
     ServerArgsValidator,
     WorkerFactory,
@@ -57,6 +58,8 @@ class SpeculativeAlgorithm(Enum):
             return cls[upper]
         except KeyError:
             pass
+        if upper in _RESERVED_ALIASES:
+            return cls.EAGLE
         spec = _get_registered_spec(upper)
         if spec is not None:
             return spec

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -108,6 +108,7 @@ class StandaloneDraftWorker(EagleDraftWorker):
             and self.topk == 1
         )
         self.dsa_index_topk = None
+        self.dsa_seed_topk_width = None
         self.seed_dsa_topk_from_draft_extend = False
         self.dsa_extend_topk_buf = None
 

python/sglang/srt/utils/common.py

@@ -90,7 +90,7 @@ import torch
 import torch.distributed as dist
 import triton
 from packaging import version as pkg_version
-from PIL import Image, UnidentifiedImageError
+from PIL import Image, ImageOps, UnidentifiedImageError
 from starlette.routing import Mount
 from torch import nn
 from torch.library import Library
@@ -1813,6 +1813,7 @@ def smart_to_rgb(
     if not isinstance(image, Image.Image):
         return image
 
+    image = ImageOps.exif_transpose(image)
     if image.mode in ("RGBA", "LA") or "transparency" in image.info:
         image = image.convert("RGBA")
         width, height = image.size
@@ -1954,6 +1955,8 @@ def load_image(
         image = _load_image(image_file=image_file, gpu_image_decode=gpu_image_decode)
     else:
         raise ValueError(f"Invalid image: {image_file}")
+    if image_size is not None and isinstance(image, Image.Image):
+        image_size = (image.width, image.height)
     return image, image_size
 
 

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

@@ -37,6 +37,8 @@ from sglang.srt.configs import (
     DotsVLMConfig,
     ExaoneConfig,
     FalconH1Config,
+    Glm5NextConfig,
+    Glm5NextTextConfig,
     GraniteMoeHybridConfig,
     HYV4Config,
     InklingAudioConfig,
@@ -120,6 +122,8 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
         MuseGlimmerConfig,
         MuseGlimmerAssistantConfig,
         KimiK3Config,
+        Glm5NextConfig,
+        Glm5NextTextConfig,
         KimiLinearConfig,
         Qwen3NextConfig,
         FalconH1Config,

python/sglang/srt/utils/video_decoder.py

@@ -87,9 +87,11 @@ class VideoDecoderWrapper:
         return len(self._decoder)
 
     def __getitem__(self, idx):
-        """Return single frame as numpy NHWC uint8."""
+        """Return one NHWC uint8 frame (numpy on CPU, tensor on CUDA)."""
         if _BACKEND == "torchcodec":
-            return self._decoder[idx].numpy()
+            frame = self._decoder[idx]
+            data = frame.data if hasattr(frame, "data") else frame
+            return data if data.is_cuda else data.numpy()
         else:
             frame = self._decoder[idx]
             return frame.asnumpy() if hasattr(frame, "asnumpy") else np.array(frame)
@@ -101,11 +103,22 @@ class VideoDecoderWrapper:
         else:
             return self._decoder.get_avg_fps()
 
-    def get_frames_at(self, indices: list) -> np.ndarray:
+    @property
-        """Return frames at given indices as numpy array with shape (N, H, W, C)."""
+    def frame_shape(self) -> tuple[int, int]:
+        if _BACKEND == "torchcodec":
+            metadata = self._decoder.metadata
+            height = getattr(metadata, "height", None)
+            width = getattr(metadata, "width", None)
+            if height and width:
+                return int(height), int(width)
+        shape = self[0].shape
+        return int(shape[-3]), int(shape[-2])
+
+    def get_frames_at(self, indices: list):
+        """Return NHWC uint8 frames (numpy on CPU, tensor on CUDA)."""
         if _BACKEND == "torchcodec":
             batch = self._decoder.get_frames_at(indices)
-            return batch.data.numpy()
+            return batch.data if batch.data.is_cuda else batch.data.numpy()
         else:
             return self._decoder.get_batch(indices).asnumpy()
 
@@ -127,7 +140,7 @@ class VideoDecoderWrapper:
 
         if _BACKEND == "torchcodec":
             batch = self._decoder.get_frames_at(indices)
-            return batch.data.pin_memory()
+            return batch.data if batch.data.is_cuda else batch.data.pin_memory()
         else:
             arr = self._decoder.get_batch(indices).asnumpy()
             return torch.from_numpy(arr).pin_memory()
@@ -141,8 +154,15 @@ class VideoDecoderWrapper:
         chunks = [list(c) for c in np.array_split(indices, num_threads) if len(c) > 0]
         source = self._source
         kwargs = self._tc_kwargs
+        cuda_device = None
+        if kwargs.get("device") == "cuda":
+            cuda_device = torch.cuda.current_device()
 
         def _decode_chunk(chunk):
+            # CUDA's current device is thread-local.  Without this, decoder
+            # workers created by TP rank > 0 silently default to GPU 0.
+            if cuda_device is not None:
+                torch.cuda.set_device(cuda_device)
             d = VideoDecoder(source, **kwargs)
             return d.get_frames_at(chunk).data
 
@@ -156,7 +176,8 @@ class VideoDecoderWrapper:
                 idx = future_to_idx[future]
                 results[idx] = future.result()
 
-        return torch.cat(results, dim=0).pin_memory()
+        output = torch.cat(results, dim=0)
+        return output if output.is_cuda else output.pin_memory()
 
     @property
     def source_bytes(self) -> bytes | None:
@@ -171,8 +192,11 @@ class VideoDecoderWrapper:
         return None
 
     def close(self):
-        """Explicitly clean up temporary files."""
+        self._decoder = None
-        if self._tmp_path is not None:
+        self._source = None
+        self._source_bytes = None
+        self._source_path = None
+        if getattr(self, "_tmp_path", None) is not None:
             if os.path.exists(self._tmp_path):
                 os.unlink(self._tmp_path)
             self._tmp_path = None

test/registered/kernels/ops/attention/test_dsa_indexer.py

@@ -219,6 +219,7 @@ class MockModelRunner:
 
         self.sliding_window_size = None
         self.page_size = self.config["page_size"]
+        self.max_running_requests = max_batch_size
 
         # Create req_to_token_pool
         self.req_to_token_pool = type(
@@ -1240,6 +1241,7 @@ class TestDSAIndexer(CustomTestCase):
                 backend.use_fused_topk = True
                 backend.dsa_topk_backend = topk_backend
                 backend.dsa_index_topk = 2048
+                backend.dsa_index_kpool = 1
                 backend.dsa_decode_impl = "fa3"
                 backend.req_to_token = torch.empty(
                     2, 4096, dtype=torch.int32, device=self.device

test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py

@@ -26,7 +26,7 @@ try:
 except ImportError:
     KERNELS_AVAILABLE = False
 
-register_cuda_ci(est_time=6, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True)
 
 
@@ -234,5 +234,40 @@ def test_mtp_single_step_decode(N: int):
     assert state_fail_rate < 0.01, f"State mismatch: fail_rate={state_fail_rate:.2f}%"
 
 
+@pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernels not available")
+def test_verify_scratch_pitch_uses_allocated_steps():
+    # Gear below the allocated step dim must not spill into the neighbor block.
+    N, T, ALLOCATED = 2, 4, 8
+    H, HV, K, V = 16, 32, 128, 128
+
+    A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
+        N, T, H, HV, K, V
+    )
+    buffer = torch.full(
+        (N + 1, ALLOCATED, HV, V, K), float("nan"), dtype=torch.float32, device="cuda"
+    )
+
+    run_fused_mtp(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state,
+        indices,
+        cu_seqlens,
+        disable_state_update=True,
+        intermediate_states_buffer=buffer,
+        intermediate_state_indices=indices,
+        cache_steps=T,
+    )
+
+    assert not torch.isnan(buffer[:N, :T]).any()
+    assert torch.isnan(buffer[N:]).all()
+    assert torch.isnan(buffer[:N, T:]).all()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/test_lean_attention.py

@@ -197,12 +197,12 @@ def _run_pair_fp8(H_Q, H_KV, D, B, S, fp8_dtype, dev="cuda", seed=0):
 def _run_pair_paged(
     H_Q, H_KV, D, B, S, page_size, dev="cuda", dt=torch.float16, seed=0
 ):
-    """Standard vs Lean on a **paged** 4-D KV buffer ``[num_pages, page_size, head, dim]``.
+    """Standard vs Lean with page-aware addressing over a dense 3-D KV buffer.
 
-    The KV cache is stored in pages and addressed through scattered slot ids in ``kv_indices``
+    The dense ``[max_slots, head, dim]`` cache is addressed through scattered slot ids in
-    (a permutation), so the kernel's page-aware address math (``kv_loc // page_size`` /
+    ``kv_indices`` (a permutation). With ``page_size > 1``, the kernel still exercises its
-    ``kv_loc % page_size``) is genuinely exercised — not the contiguous fast path. Both arms read
+    page-aware address math (``kv_loc // page_size`` / ``kv_loc % page_size``). Both arms read the
-    the identical buffer + indices, so their outputs must agree. Returns (o_std, o_lean).
+    identical buffer + indices, so their outputs must agree. Returns (o_std, o_lean).
     """
     torch.manual_seed(seed)
     D_V = D
@@ -212,11 +212,9 @@ def _run_pair_paged(
     assert tot % page_size == 0, (
         "test setup: total tokens must be a multiple of page_size"
     )
-    num_pages = tot // page_size
+    # Unified memory exposes dense 3-D KV views even when the allocator uses pages.
-
+    k = torch.randn(tot, H_KV, D, dtype=dt, device=dev)
-    # 4-D paged KV buffers [num_pages, page_size, head, dim] (the shared-pool layout).
+    v = torch.randn(tot, H_KV, D_V, dtype=dt, device=dev)
-    k = torch.randn(num_pages, page_size, H_KV, D, dtype=dt, device=dev)
-    v = torch.randn(num_pages, page_size, H_KV, D_V, dtype=dt, device=dev)
 
     kv_indptr = torch.arange(0, (B + 1) * S, step=S, device=dev, dtype=torch.int32)
     # Scatter slots across pages so page_id/tok_in_p vary within every BLOCK_N tile.
@@ -312,9 +310,9 @@ class TestLeanAttentionParity(CustomTestCase):
                         )
 
     def test_paged_kv_parity(self):
-        # Lean must read a paged 4-D KV buffer the same way the standard kernel does. Guards
+        # Lean must apply page-aware address math to dense KV views the same way the standard
-        # the page-aware address math (kv_loc // page_size, kv_loc % page_size); a regression
+        # kernel does. A regression in kv_loc // page_size or kv_loc % page_size would scramble
-        # to the contiguous-only form would scramble reads and drop cos well below 1.
+        # the scattered reads and drop cos well below 1.
         for name, H_Q, H_KV, D in GQA_SHAPES:
             for page_size in (16, 64):
                 with self.subTest(model=name, page_size=page_size):

test/registered/models_e2e/test_glm53_flash_b200.py

@@ -0,0 +1,138 @@
+"""B200 per-commit coverage for the GLM-5.3-Flash serving recipes.
+
+Runs the Low Latency, DFlash2, and High Throughput TP4/EP4 recipes on four
+B200 GPUs. All recipes must retain GSM8K accuracy; the Low Latency recipe also
+checks EAGLE speculative acceptance and single-request decode performance.
+"""
+
+import unittest
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    _wait_for_gpu_idle_in_ci,
+    popen_launch_server,
+    try_cached_model,
+)
+
+register_cuda_ci(est_time=2400, stage="base-c", runner_config="4-gpu-b200")
+
+MODEL_PATH = "zai-org/GLM-5.3-Flash"
+DFLASH2_DRAFT_MODEL_PATH = "incoai/GLM-5.3-Flash-DFlash2"
+SERVER_LAUNCH_TIMEOUT = 3600
+GPU_IDLE_TIMEOUT = 120
+
+COMMON_SERVER_ARGS = [
+    "--tp-size",
+    "4",
+    "--ep-size",
+    "4",
+    "--dsa-prefill-backend",
+    "trtllm",
+    "--dsa-decode-backend",
+    "trtllm",
+    "--kv-cache-dtype",
+    "fp8_e4m3",
+    "--moe-runner-backend",
+    "deep_gemm",
+    "--reasoning-parser",
+    "glm45",
+    "--tool-call-parser",
+    "glm47",
+]
+
+
+def _stop_server(process):
+    if process:
+        kill_process_tree(process.pid)
+        _wait_for_gpu_idle_in_ci(timeout=GPU_IDLE_TIMEOUT)
+
+
+class _GLM53FlashB200Base(CustomTestCase):
+    server_args: list[str]
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = try_cached_model(MODEL_PATH)
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = None
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=cls.server_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        _stop_server(getattr(cls, "process", None))
+
+
+class TestGLM53FlashB200LowLatency(
+    SpecDecodingMixin,
+    GSM8KMixin,
+    _GLM53FlashB200Base,
+):
+    gsm8k_score_threshold = 0.93
+    # Match the established DSA+MTP accuracy workload. The generic 200-question,
+    # 5-shot defaults leave a single question worth 0.5 percentage points and
+    # make this tight quality floor unnecessarily sensitive to kernel numerics.
+    gsm8k_num_examples = 500
+    gsm8k_num_shots = 20
+    accept_length_thres = 4.0
+    bs_1_speed_thres = 250
+    server_args = [
+        *COMMON_SERVER_ARGS,
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "5",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "6",
+        "--speculative-adaptive",
+    ]
+
+
+class TestGLM53FlashB200HighThroughput(
+    GSM8KMixin,
+    _GLM53FlashB200Base,
+):
+    gsm8k_score_threshold = 0.93
+    gsm8k_num_examples = 500
+    gsm8k_num_shots = 20
+    server_args = [
+        *COMMON_SERVER_ARGS,
+        "--enable-dp-attention",
+        "--dp-size",
+        "4",
+        "--moe-a2a-backend",
+        "deepep",
+    ]
+
+
+class TestGLM53FlashB200DFlash2(
+    GSM8KMixin,
+    _GLM53FlashB200Base,
+):
+    gsm8k_score_threshold = 0.93
+    gsm8k_num_examples = 500
+    gsm8k_num_shots = 20
+    server_args = [
+        *COMMON_SERVER_ARGS,
+        "--speculative-algorithm",
+        "DFLASH",
+        "--speculative-draft-model-path",
+        DFLASH2_DRAFT_MODEL_PATH,
+        "--speculative-draft-attention-backend",
+        "fa4",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/models_e2e/test_glm53_flash_h200.py

@@ -0,0 +1,119 @@
+"""H200 per-commit coverage for the GLM-5.3-Flash serving recipes.
+
+Runs the Low Latency and High Throughput TP8/EP8 recipes on eight H200 GPUs.
+Both recipes must retain GSM8K accuracy; the Low Latency recipe also checks
+EAGLE speculative acceptance and single-request decode performance.
+"""
+
+import unittest
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    _wait_for_gpu_idle_in_ci,
+    popen_launch_server,
+    try_cached_model,
+)
+
+register_cuda_ci(est_time=2400, stage="extra-b", runner_config="8-gpu-h200")
+
+MODEL_PATH = "zai-org/GLM-5.3-Flash"
+SERVER_LAUNCH_TIMEOUT = 3600
+GPU_IDLE_TIMEOUT = 120
+
+COMMON_SERVER_ARGS = [
+    "--tp-size",
+    "8",
+    "--ep-size",
+    "8",
+    "--dsa-prefill-backend",
+    "tilelang",
+    "--dsa-decode-backend",
+    "tilelang",
+    "--kv-cache-dtype",
+    "bf16",
+    "--moe-runner-backend",
+    "deep_gemm",
+    "--reasoning-parser",
+    "glm45",
+    "--tool-call-parser",
+    "glm47",
+]
+
+
+def _stop_server(process):
+    if process:
+        kill_process_tree(process.pid)
+        _wait_for_gpu_idle_in_ci(timeout=GPU_IDLE_TIMEOUT)
+
+
+class _GLM53FlashH200Base(CustomTestCase):
+    server_args: list[str]
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = try_cached_model(MODEL_PATH)
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = None
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=cls.server_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        _stop_server(getattr(cls, "process", None))
+
+
+class TestGLM53FlashH200LowLatency(
+    SpecDecodingMixin,
+    GSM8KMixin,
+    _GLM53FlashH200Base,
+):
+    gsm8k_score_threshold = 0.93
+    # Match the established DSA+MTP accuracy workload. The generic 200-question,
+    # 5-shot defaults leave a single question worth 0.5 percentage points and
+    # make this tight quality floor unnecessarily sensitive to kernel numerics.
+    gsm8k_num_examples = 500
+    gsm8k_num_shots = 20
+    accept_length_thres = 4.0
+    bs_1_speed_thres = 200
+    server_args = [
+        *COMMON_SERVER_ARGS,
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "5",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "6",
+        "--speculative-adaptive",
+    ]
+
+
+class TestGLM53FlashH200HighThroughput(
+    GSM8KMixin,
+    _GLM53FlashH200Base,
+):
+    gsm8k_score_threshold = 0.93
+    gsm8k_num_examples = 500
+    gsm8k_num_shots = 20
+    server_args = [
+        *COMMON_SERVER_ARGS,
+        "--enable-dp-attention",
+        "--dp-size",
+        "8",
+        "--moe-a2a-backend",
+        "deepep",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py

@@ -60,6 +60,7 @@ class _Scheduler(SchedulerDisaggregationPrefillMixin):
         self.send_kv_chunk = Mock()
         self.output_streamer = Mock()
         self.metrics_reporter = SimpleNamespace(report_prefill_stats=Mock())
+        self.maybe_send_health_check_signal = Mock()
         self.req_to_metadata_buffer_idx_allocator = Mock()
         self.enable_hicache_storage = True
         self.chunked_req = None

test/registered/unit/hardware_backend/mlx/test_attention_patching.py

@@ -1520,6 +1520,7 @@ if _HAS_MLX:
             self.kv = ReqKvInfo()
             self.mamba_branching_seqlen = None
             self.inflight_middle_chunks = 0
+            self.mamba_branching_seqlen = None
 
     class FakeTpWorker:
         def __init__(self, next_token_ids):

test/registered/unit/layers/attention/test_mla_decode_forced_splits.py

@@ -74,16 +74,10 @@ def _inputs(seq_lens, head_num, page_size, max_kv_splits, seed):
     total = sum(seq_lens)
     n_slots = total + 64
 
-    if page_size == 1:
+    if page_size > 1:
-        pool = torch.randn(
+        n_slots = ((n_slots + page_size - 1) // page_size) * page_size
-            n_slots, 1, LK, dtype=torch.bfloat16, device=dev, generator=gen
+    # Unified memory exposes dense 3-D KV views even when the allocator uses pages.
-        )
+    pool = torch.randn(n_slots, 1, LK, dtype=torch.bfloat16, device=dev, generator=gen)
-    else:
-        n_pages = (n_slots + page_size - 1) // page_size
-        pool = torch.randn(
-            n_pages, page_size, 1, LK, dtype=torch.bfloat16, device=dev, generator=gen
-        )
-        n_slots = n_pages * page_size
 
     kv_indptr = torch.zeros(batch + 1, dtype=torch.int32, device=dev)
     kv_indptr[1:] = torch.cumsum(

test/registered/unit/managers/test_generation_auxiliary_output.py

@@ -476,6 +476,7 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit():
         disagg_prefill_inflight_queue=[],
         send_kv_chunk=Mock(),
         metrics_reporter=SimpleNamespace(report_prefill_stats=Mock()),
+        maybe_send_health_check_signal=Mock(),
     )
 
     with patch("sglang.srt.disaggregation.prefill.maybe_cache_unfinished_req"):
@@ -492,6 +493,7 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit():
         host_output,
         [0],
     )
+    scheduler.maybe_send_health_check_signal.assert_called_once_with()
 
 
 def test_logprob_only_reuses_preprocessing_without_observer_lifecycle():

test/registered/unit/mem_cache/test_retraction_mamba_backup.py

@@ -30,10 +30,12 @@ class _Allocator:
     def get_kvcache(self):
         return self._kv
 
-    def get_cpu_copy(self, indices, mamba_indices=None):
+    def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         return "kv"
 
-    def load_cpu_copy(self, cpu_tensors, indices, mamba_indices=None):
+    def load_cpu_copy(
+        self, cpu_tensors, indices, mamba_indices=None, req_pool_index=None
+    ):
         self.loaded_kv = cpu_tensors
 
 

test/registered/unit/model_executor/test_pool_configurator.py

@@ -124,9 +124,11 @@ def _make_model_runner(
     mc.get_num_kv_heads = lambda tp_size, dcp_size=1: num_kv_heads
     mc.get_swa_num_kv_heads = lambda tp_size: swa_num_kv_heads or num_kv_heads
     mc.hf_config = SimpleNamespace(architectures=["LlamaForCausalLM"])
+    mc.hf_config.model_type = "llama"
     mc.hf_config.get_text_config = lambda: mc.hf_config
     mc.linear_attn_registry_result = None
     mc.context_len = 8192
+    mc.is_draft_model = False
     mr.model_config = mc
     mr.kv_cache_dtype = "fake_bf16"
 

test/registered/unit/models/test_deepseek_nextn_mm_embed.py

@@ -0,0 +1,210 @@
+"""Regression test for sgl-project/sglang#37548.
+
+DeepseekModelNextN.forward must use forward_batch.mm_input_embeds for multimodal
+positions (where input_ids hold MM_PAD_SHIFT_VALUE+hash sentinels far above
+vocab_size) instead of calling embed_tokens on those sentinel values, which
+causes a CUDA index-out-of-bounds gather.
+"""
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from sglang.srt.managers.schedule_batch import MM_PAD_SHIFT_VALUE
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=30, suite="base-a-test-cpu")
+
+VOCAB_SIZE = 154880
+HIDDEN_SIZE = 64  # tiny for CPU test
+
+
+def _make_forward_batch(
+    input_ids: torch.Tensor,
+    mm_input_embeds: torch.Tensor = None,
+    extend_seq_lens: torch.Tensor = None,
+    extend_start_loc: torch.Tensor = None,
+    has_mm: bool = True,
+):
+    """Build a minimal mock ForwardBatch for DeepseekModelNextN.forward."""
+    fb = MagicMock()
+    fb.mm_input_embeds = mm_input_embeds
+    fb.contains_mm_inputs.return_value = has_mm
+    fb.forward_mode.is_extend.return_value = True
+    fb.forward_mode.is_draft_extend_v2.return_value = False
+    fb.forward_mode.is_idle.return_value = False
+    fb.extend_seq_lens = extend_seq_lens
+    fb.extend_start_loc = extend_start_loc
+    fb.spec_info.hidden_states = torch.randn(input_ids.shape[0], HIDDEN_SIZE)
+    return fb
+
+
+def _make_model_nextn(vocab_size: int, hidden_size: int):
+    """Build a mock DeepseekModelNextN with a real embed_tokens layer."""
+    from sglang.srt.models.deepseek_nextn import DeepseekModelNextN
+
+    model = DeepseekModelNextN.__new__(DeepseekModelNextN)
+    torch.nn.Module.__init__(model)
+    # Minimal attributes needed by forward
+    model.vocab_size = vocab_size
+    model.embed_tokens = torch.nn.Embedding(vocab_size, hidden_size)
+    model.enorm = torch.nn.RMSNorm(hidden_size)
+    model.hnorm = torch.nn.RMSNorm(hidden_size)
+    model.eh_proj = torch.nn.Linear(2 * hidden_size, hidden_size, bias=False)
+    model.rot_weight = None
+    model.alt_stream = None
+    model.quant_config = None
+    model.cp_rank = None
+    model.cp_size = None
+    model.dsa_enable_prefill_cp = False
+    model.mla_enable_prefill_cp = False
+    model.mtp_block = MagicMock(side_effect=lambda **kw: (kw["hidden_states"], None))
+    return model
+
+
+class TestDeepseekNextNMmEmbed(CustomTestCase):
+    """DeepseekModelNextN must not call embed_tokens on MM sentinel token ids."""
+
+    def test_mm_sentinel_ids_do_not_cause_oob(self):
+        """input_ids containing MM_PAD_SHIFT_VALUE+hash must not reach embed_tokens."""
+        num_tokens = 10
+        mm_start, mm_end = 3, 7  # MM sentinel positions
+
+        input_ids = torch.arange(num_tokens, dtype=torch.long)
+        # Insert MM sentinel values
+        for i in range(mm_start, mm_end):
+            input_ids[i] = MM_PAD_SHIFT_VALUE + i
+
+        # Build mm_input_embeds matching the target-produced embeddings
+        mm_embeds = torch.randn(num_tokens, HIDDEN_SIZE)
+        extend_seq_lens = torch.tensor([num_tokens])
+        extend_start_loc = torch.tensor([0])
+
+        fb = _make_forward_batch(
+            input_ids,
+            mm_input_embeds=mm_embeds.clone(),
+            extend_seq_lens=extend_seq_lens,
+            extend_start_loc=extend_start_loc,
+        )
+
+        model = _make_model_nextn(VOCAB_SIZE, HIDDEN_SIZE)
+
+        # Use MagicMock to track embed_tokens calls
+        mock_embed = MagicMock(side_effect=model.embed_tokens)
+        object.__setattr__(model, "embed_tokens", mock_embed)
+
+        with (
+            patch(
+                "sglang.srt.models.deepseek_nextn.is_cp_v2_active", return_value=False
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.dsa_use_prefill_cp",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.mla_use_prefill_cp",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.fused_eh_norm",
+                side_effect=lambda h, p, ew, hw, eps: torch.cat(
+                    [model.enorm(h), model.hnorm(p)], dim=-1
+                ),
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.get_global_expert_distribution_recorder"
+            ),
+            patch("sglang.srt.models.deepseek_nextn.is_cuda", False),
+            patch("sglang.srt.models.deepseek_nextn.is_npu", False),
+            patch("sglang.srt.models.deepseek_nextn.envs") as mock_envs,
+            patch("sglang.srt.models.deepseek_nextn.get_model") as mock_get_model,
+            patch("sglang.srt.models.deepseek_nextn.get_parallel") as mock_get_parallel,
+            patch("sglang.srt.models.deepseek_nextn.get_spec") as mock_get_spec,
+        ):
+            mock_envs.SGLANG_NPU_USE_MULTI_STREAM.get.return_value = False
+            mock_get_model.return_value.quantization = None
+
+            positions = torch.arange(num_tokens, dtype=torch.long)
+            try:
+                model.forward(input_ids, positions, fb)
+            except Exception:
+                pass  # We only care about embed_tokens call args
+
+        # embed_tokens should only be called for last_indices (the appended
+        # next-token), not with the full input_ids containing MM sentinels.
+        for call in mock_embed.call_args_list:
+            call_ids = call[0][0]
+            max_id = call_ids.max().item()
+            self.assertLess(
+                max_id,
+                VOCAB_SIZE,
+                f"embed_tokens was called with id {max_id} >= vocab_size "
+                f"{VOCAB_SIZE}. MM sentinel values (MM_PAD_SHIFT_VALUE+hash) "
+                f"must not reach embed_tokens.",
+            )
+
+    def test_no_mm_falls_back_to_embed_tokens(self):
+        """Without mm_input_embeds, embed_tokens is called normally."""
+        num_tokens = 5
+        input_ids = torch.arange(num_tokens, dtype=torch.long)
+
+        fb = _make_forward_batch(
+            input_ids,
+            mm_input_embeds=None,
+            extend_seq_lens=torch.tensor([num_tokens]),
+            extend_start_loc=torch.tensor([0]),
+            has_mm=False,
+        )
+
+        model = _make_model_nextn(VOCAB_SIZE, HIDDEN_SIZE)
+        mock_embed = MagicMock(side_effect=model.embed_tokens)
+        object.__setattr__(model, "embed_tokens", mock_embed)
+        embed_calls = mock_embed.call_args_list
+
+        with (
+            patch(
+                "sglang.srt.models.deepseek_nextn.is_cp_v2_active", return_value=False
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.dsa_use_prefill_cp",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.mla_use_prefill_cp",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.fused_eh_norm",
+                side_effect=lambda h, p, ew, hw, eps: torch.cat(
+                    [model.enorm(h), model.hnorm(p)], dim=-1
+                ),
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.get_global_expert_distribution_recorder"
+            ),
+            patch("sglang.srt.models.deepseek_nextn.is_cuda", False),
+            patch("sglang.srt.models.deepseek_nextn.is_npu", False),
+            patch("sglang.srt.models.deepseek_nextn.envs") as mock_envs,
+            patch("sglang.srt.models.deepseek_nextn.get_model") as mock_get_model,
+            patch("sglang.srt.models.deepseek_nextn.get_parallel") as mock_get_parallel,
+            patch("sglang.srt.models.deepseek_nextn.get_spec") as mock_get_spec,
+        ):
+            mock_envs.SGLANG_NPU_USE_MULTI_STREAM.get.return_value = False
+            mock_get_model.return_value.quantization = None
+
+            positions = torch.arange(num_tokens, dtype=torch.long)
+            try:
+                model.forward(input_ids, positions, fb)
+            except Exception:
+                pass
+
+        # embed_tokens should be called with the full input_ids
+        self.assertTrue(mock_embed.call_count > 0, "embed_tokens should be called")
+        full_ids_call = mock_embed.call_args_list[0][0][0]
+        self.assertEqual(full_ids_call.numel(), num_tokens)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/models/test_glm5_next_dflash_capture.py

@@ -0,0 +1,40 @@
+"""Regression for DFLASH aux-hidden capture on mHC models.
+
+GLM-5.3-Flash runs with mhc=True. MHCLayerCommunicator folds the residual
+into the widened hidden state and returns residual=None, so CUDA-graph
+capture used to crash on `hidden_states + residual`. DFLASH also has to
+contract that widened state back to the draft hidden size; skipping the
+contract is a silent shape/quality bug the crash-guard alone would miss.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from sglang.srt.models.glm5_next import Glm5NextModel
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestGlm5NextDflashCapture(CustomTestCase):
+    def test_dflash_contracts_mhc_hidden_state_without_residual(self):
+        model = Glm5NextModel.__new__(Glm5NextModel)
+        nn.Module.__init__(model)
+        model.config = SimpleNamespace(mhc=True, hc_mult=4)
+        model.dflash_capture = True
+
+        hidden_states = torch.arange(24, dtype=torch.float32).reshape(2, 12)
+
+        actual = model._prepare_aux_hidden_state(hidden_states, None)
+        expected = hidden_states.unflatten(-1, (4, -1)).mean(dim=-2)
+
+        torch.testing.assert_close(actual, expected)
+        self.assertEqual(tuple(actual.shape), (2, 3))
+
+
+if __name__ == "__main__":
+    unittest.main()