[Kimi] Support kimi-k3 (#32541)

Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Ziyi Xu <ziyi.xu@radixark.ai>
Co-authored-by: Zijie Xia <37504505+zijiexia@users.noreply.github.com>
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
Co-authored-by: zhangxiaohao <1024393531@qq.com>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: Julien Lin <jullin@nvidia.com>
Co-authored-by: Hao Phan <htphan@nvidia.com>
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
Co-authored-by: RolaoDenthu <xinyisong0111@gmail.com>
Co-authored-by: pigeonsoup <32922982+pigeonsoup@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com>
Co-authored-by: Lee Nau <lee.nau@gmail.com>
Co-authored-by: HMING <126185151+Hearum@users.noreply.github.com>
Co-authored-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
Co-authored-by: Byron Hsu <byronhsu1230@gmail.com>
Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com>
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
Co-authored-by: BBuf <xiaoyu.zhang@radixark.ai>
Co-authored-by: Hanming Lu <hanminglu@meta.com>
Co-authored-by: Xinyi Song <xinyis10@illinois.edu>

python/pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
   "scipy",
   "sentencepiece",
   "setproctitle",
-  "sgl-deep-gemm==0.1.5",
+  "sgl-deep-gemm==0.1.5.post1",
   "sglang-kernel==0.4.5",
   "smg-grpc-servicer>=0.5.0",
   "soundfile==0.13.1",

python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh

@@ -18,9 +18,11 @@
 #include <tvm/ffi/container/tensor.h>
 
 #include <cstdint>
-#include <cuda_fp8.h>
 #include <limits>
 #include <type_traits>
+#ifndef USE_ROCM
+#include <cuda_fp8.h>
+#endif
 
 namespace sglang {
 

python/sglang/kernels/jit/csrc/moe/align_single_token.cuh

@@ -0,0 +1,107 @@
+// Tiny moe_align_block_size for M == 1 decode: one warp replaces the
+// moe_align_block_size + count_and_sort_expert_tokens kernel pair (~4.4us)
+// with a single ~1.5us launch.
+//
+// For a single token the top-k expert ids are distinct, so the aligned
+// layout is exactly: experts sorted ascending, one block per expert,
+// slot i of block b = flat topk index for b's expert, remaining block
+// slots padded with numel (= topk). num_tokens_post_padded = topk * block.
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/type.cuh>   // For device::cast
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+struct AlignSingleTokenParams {
+  const int32_t* __restrict__ topk_ids;  // [1, topk]
+  int32_t* __restrict__ sorted_ids;      // [topk * block_size]
+  int32_t* __restrict__ expert_ids;      // [topk]
+  int32_t* __restrict__ num_post;        // [1]
+  uint32_t topk;
+  uint32_t block_size;
+};
+
+template <bool kUsePDL>
+__global__ void align_single_token_kernel(const AlignSingleTokenParams __grid_constant__ params) {
+  using namespace device;
+  const uint32_t lane = threadIdx.x;  // one warp
+  const uint32_t topk = params.topk;
+  const uint32_t bs = params.block_size;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  int32_t my_id = (lane < topk) ? params.topk_ids[lane] : INT32_MAX;
+
+  // Rank of my expert id among the topk (ids are distinct for one token;
+  // tie-break on lane keeps this robust anyway).
+  uint32_t rank = 0;
+  for (uint32_t j = 0; j < topk; ++j) {
+    int32_t other = __shfl_sync(0xffffffff, my_id, j);
+    if (other < my_id || (other == my_id && j < lane)) {
+      rank++;
+    }
+  }
+
+  if (lane < topk) {
+    params.expert_ids[rank] = my_id;
+    // block `rank`: first slot is my flat index (token 0, slot `lane`),
+    // rest padded with numel (= topk).
+    params.sorted_ids[rank * bs] = static_cast<int32_t>(lane);
+  }
+  // Fill padding cooperatively: positions not equal to a block start.
+  for (uint32_t p = lane; p < topk * bs; p += 32) {
+    if (p % bs != 0) {
+      params.sorted_ids[p] = static_cast<int32_t>(topk);
+    }
+  }
+  if (lane == 0) {
+    params.num_post[0] = static_cast<int32_t>(topk * bs);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <bool kUsePDL>
+struct AlignSingleTokenKernel {
+  static constexpr auto kernel = align_single_token_kernel<kUsePDL>;
+
+  static void
+  run(const tvm::ffi::TensorView topk_ids,
+      const tvm::ffi::TensorView sorted_ids,
+      const tvm::ffi::TensorView expert_ids,
+      const tvm::ffi::TensorView num_post,
+      int64_t block_size) {
+    using namespace host;
+
+    auto One_ = SymbolicSize{"one"};
+    auto K_ = SymbolicSize{"topk"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({One_, K_}).with_dtype<int32_t>().with_device(device).verify(topk_ids);
+
+    const auto topk = static_cast<uint32_t>(K_.unwrap());
+    RuntimeCheck(One_.unwrap() == 1, "moe_align_single_token requires M == 1");
+    RuntimeCheck(topk <= 32, "moe_align_single_token requires topk <= 32");
+
+    const auto params = AlignSingleTokenParams{
+        .topk_ids = static_cast<const int32_t*>(topk_ids.data_ptr()),
+        .sorted_ids = static_cast<int32_t*>(sorted_ids.data_ptr()),
+        .expert_ids = static_cast<int32_t*>(expert_ids.data_ptr()),
+        .num_post = static_cast<int32_t*>(num_post.data_ptr()),
+        .topk = topk,
+        .block_size = static_cast<uint32_t>(block_size),
+    };
+
+    LaunchKernel(dim3(1), 32, device.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/kernels/ops/moe/moe_align_single_token.py

@@ -0,0 +1,49 @@
+"""CUDA JIT single-warp moe_align_block_size for M == 1 decode batches."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_align_single_token_module() -> Module:
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        "moe_align_single_token",
+        *args,
+        cuda_files=["moe/align_single_token.cuh"],
+        cuda_wrappers=[("run", f"AlignSingleTokenKernel<{args}>::run")],
+        extra_cuda_cflags=["-O3"],
+    )
+
+
+def moe_align_single_token(
+    topk_ids: torch.Tensor, block_size: int
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """moe_align_block_size for a single token (distinct expert ids).
+
+    Returns (sorted_token_ids, expert_ids, num_tokens_post_padded) with the
+    same layout as moe_align_block_size: experts ascending, one block per
+    expert, padding value = numel.
+    """
+    topk = topk_ids.shape[1]
+    device = topk_ids.device
+    sorted_ids = torch.empty((topk * block_size,), dtype=torch.int32, device=device)
+    expert_ids = torch.empty((topk,), dtype=torch.int32, device=device)
+    num_post = torch.empty((1,), dtype=torch.int32, device=device)
+    _jit_align_single_token_module().run(
+        topk_ids, sorted_ids, expert_ids, num_post, block_size
+    )
+    return sorted_ids, expert_ids, num_post

python/sglang/kernels/ops/sampling/renorm_triton.py

@@ -0,0 +1,172 @@
+"""ROCm-compatible top-k / top-p probability renormalization fallbacks."""
+
+from __future__ import annotations
+
+from typing import Union
+
+import torch
+import triton
+import triton.language as tl
+
+_BLOCK_SIZE = 1024
+
+
+@triton.jit
+def _mask_and_partial_sum_kernel(
+    probs_ptr,
+    pivots_ptr,
+    out_ptr,
+    partial_sums_ptr,
+    vocab_size: tl.constexpr,
+    num_chunks: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    chunk = tl.program_id(1)
+    offsets = chunk * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < vocab_size
+    row_offsets = row * vocab_size + offsets
+
+    probs = tl.load(probs_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32)
+    pivot = tl.load(pivots_ptr + row)
+    kept = tl.where(mask & (probs >= pivot), probs, 0.0)
+
+    tl.store(out_ptr + row_offsets, kept, mask=mask)
+    tl.store(partial_sums_ptr + row * num_chunks + chunk, tl.sum(kept, axis=0))
+
+
+@triton.jit
+def _normalize_kernel(
+    out_ptr,
+    row_sums_ptr,
+    numel,
+    vocab_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < numel
+    row = offsets // vocab_size
+    values = tl.load(out_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+    denominator = tl.load(row_sums_ptr + row, mask=mask, other=1.0)
+    tl.store(out_ptr + offsets, values / denominator, mask=mask)
+
+
+def _prepare_probs(probs: torch.Tensor) -> torch.Tensor:
+    if probs.ndim != 2:
+        raise ValueError(f"probs must be 2D, got shape={tuple(probs.shape)}")
+    if not probs.is_cuda:
+        raise ValueError("renorm kernels require a CUDA/HIP tensor")
+    return probs.float().contiguous()
+
+
+def _renorm_from_pivots(probs_fp32: torch.Tensor, pivots: torch.Tensor) -> torch.Tensor:
+    batch_size, vocab_size = probs_fp32.shape
+    num_chunks = triton.cdiv(vocab_size, _BLOCK_SIZE)
+    out = torch.empty_like(probs_fp32)
+    partial_sums = torch.empty(
+        (batch_size, num_chunks), device=probs_fp32.device, dtype=torch.float32
+    )
+    _mask_and_partial_sum_kernel[(batch_size, num_chunks)](
+        probs_fp32,
+        pivots,
+        out,
+        partial_sums,
+        vocab_size=vocab_size,
+        num_chunks=num_chunks,
+        BLOCK_SIZE=_BLOCK_SIZE,
+        num_warps=8,
+    )
+
+    row_sums = partial_sums.sum(dim=1)
+    _normalize_kernel[(triton.cdiv(out.numel(), _BLOCK_SIZE),)](
+        out,
+        row_sums,
+        out.numel(),
+        vocab_size=vocab_size,
+        BLOCK_SIZE=_BLOCK_SIZE,
+        num_warps=8,
+    )
+    return out
+
+
+def top_p_renorm_probs_triton(
+    probs: torch.Tensor, top_p: Union[torch.Tensor, float]
+) -> torch.Tensor:
+    """Apply exact top-p thresholding and renormalize each probability row.
+
+    Sorting and prefix sums use PyTorch's device kernels because a vocabulary-sized
+    in-register Triton sort does not scale to 100K+ vocabularies. Triton performs
+    the bandwidth-heavy masking, partial reduction, and normalization.
+    """
+    probs_fp32 = _prepare_probs(probs)
+    batch_size, vocab_size = probs_fp32.shape
+    if batch_size == 0 or vocab_size == 0:
+        return probs_fp32
+
+    if isinstance(top_p, torch.Tensor):
+        top_ps = top_p.to(device=probs.device, dtype=torch.float32).reshape(-1)
+        if top_ps.numel() == 1:
+            top_ps = top_ps.expand(batch_size)
+        elif top_ps.numel() != batch_size:
+            raise ValueError(
+                f"top_p must be scalar or have one value per row, got "
+                f"{top_ps.numel()} values for {batch_size} rows"
+            )
+    else:
+        if not 0.0 < float(top_p) <= 1.0:
+            raise ValueError("top_p values must be in (0, 1]")
+        top_ps = torch.full(
+            (batch_size,), float(top_p), device=probs.device, dtype=torch.float32
+        )
+
+    # Match FlashInfer's threshold semantics: sort ascending, discard the prefix
+    # whose cumulative mass is below 1 - p, and retain all ties at the pivot.
+    sorted_probs = torch.sort(probs_fp32, dim=-1).values
+    cdf = torch.cumsum(sorted_probs, dim=-1)
+    cutoff = torch.searchsorted(cdf, (1.0 - top_ps).unsqueeze(1), right=False).squeeze(
+        1
+    )
+    cutoff.clamp_(max=vocab_size - 1)
+    pivots = sorted_probs.gather(1, cutoff.unsqueeze(1)).squeeze(1).contiguous()
+
+    return _renorm_from_pivots(probs_fp32, pivots)
+
+
+def top_k_renorm_probs_triton(
+    probs: torch.Tensor, top_k: Union[torch.Tensor, int]
+) -> torch.Tensor:
+    """Apply exact top-k thresholding and renormalize each probability row.
+
+    Sorting uses PyTorch's device kernels because a vocabulary-sized in-register
+    Triton sort does not scale to 100K+ vocabularies. Triton performs the
+    bandwidth-heavy masking, partial reduction, and normalization.
+    """
+    probs_fp32 = _prepare_probs(probs)
+    batch_size, vocab_size = probs_fp32.shape
+    if batch_size == 0 or vocab_size == 0:
+        return probs_fp32
+
+    if isinstance(top_k, torch.Tensor):
+        top_ks = top_k.to(device=probs.device, dtype=torch.int64).reshape(-1)
+        if top_ks.numel() == 1:
+            top_ks = top_ks.expand(batch_size)
+        elif top_ks.numel() != batch_size:
+            raise ValueError(
+                f"top_k must be scalar or have one value per row, got "
+                f"{top_ks.numel()} values for {batch_size} rows"
+            )
+    else:
+        top_ks = torch.full(
+            (batch_size,), int(top_k), device=probs.device, dtype=torch.int64
+        )
+
+    # Match FlashInfer's threshold semantics: sort descending, keep the k highest
+    # probabilities, and retain all ties at the pivot.
+    sorted_probs = torch.sort(probs_fp32, dim=-1, descending=True).values
+    cutoff = (top_ks - 1).clamp_(min=0, max=vocab_size - 1)
+    pivots = sorted_probs.gather(1, cutoff.unsqueeze(1)).squeeze(1).contiguous()
+
+    return _renorm_from_pivots(probs_fp32, pivots)
+
+
+__all__ = ["top_k_renorm_probs_triton", "top_p_renorm_probs_triton"]

python/sglang/srt/arg_groups/kimi_k3_hook.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+def apply_kimi_k3_spec_backend_defaults(server_args: ServerArgs) -> None:
+    """Apply speculative backend defaults for Kimi hybrid models."""
+    from sglang.srt.utils import is_sm100_supported
+
+    if server_args.speculative_algorithm is None:
+        return
+
+    # Use the fused Kimi-K3/DSPARK CuTeDSL kernel for KDA target verification.
+    # Decode is left free (its bf16-ssm SM100+ flashinfer default is fine -- the
+    # target only verifies under spec); the verify backend is pinned directly.
+    if server_args.linear_attn_verify_backend is None:
+        server_args.linear_attn_verify_backend = "nv_cutedsl"
+        logger.info(
+            "Kimi hybrid model with speculative decoding: pinning "
+            "--linear-attn-verify-backend to nv_cutedsl (uses the fused "
+            "Kimi-K3/DSPARK CuTeDSL kernel)."
+        )
+
+    # dspark's draft is dense MQA; trtllm_mha avoids flashinfer's blocking
+    # per-step host plan. DSPARK-only: other spec algos use MLA-family drafts.
+    if (
+        server_args.speculative_algorithm == "DSPARK"
+        and server_args.speculative_draft_attention_backend is None
+        and is_sm100_supported()
+    ):
+        server_args.speculative_draft_attention_backend = "trtllm_mha"
+        logger.info(
+            "Kimi hybrid DSPARK: defaulting "
+            "--speculative-draft-attention-backend to trtllm_mha."
+        )
+
+
+def disable_kimi_k3_symm_mem(server_args: ServerArgs) -> None:
+    """Turn `--enable-symm-mem` back off unless every phase runs eager.
+
+    Symm-mem allocations are per-forward, so an address captured into a graph is
+    neither reserved for its lifetime nor at the same offset on every rank. Under
+    capture that corrupts spec decode: accept collapses to 1.000, or the server
+    silently emits garbage with accept pinned at the ceiling. Prefill counts too --
+    the same allocation sits in any captured RowParallelLinear.
+
+    Gates on the arch itself: this runs from cuda-graph resolution, which is earlier
+    than the model-specific hook block.
+    """
+    from sglang.srt.connector import ConnectorType
+    from sglang.srt.model_executor.cuda_graph_config import Backend
+    from sglang.srt.utils import parse_connector_type
+
+    if not server_args.enable_symm_mem:
+        return
+    if parse_connector_type(server_args.model_path) == ConnectorType.INSTANCE:
+        return
+    if server_args.get_model_config().hf_config.architectures[0] not in (
+        "KimiLinearForCausalLM",
+        "KimiK3ForConditionalGeneration",
+    ):
+        return
+    graph = server_args.cuda_graph_config
+    if (
+        graph.decode.backend == Backend.DISABLED
+        and graph.prefill.backend == Backend.DISABLED
+    ):
+        return
+    server_args.enable_symm_mem = False
+    logger.warning(
+        "Kimi hybrid model: ignoring --enable-symm-mem because CUDA graphs are on. "
+        "The symmetric-memory pool's per-forward allocations are not valid for the "
+        "lifetime of a captured graph, which corrupts speculative decoding and can "
+        "silently produce wrong output. The auto-probed K3 fused all-reduce is faster "
+        "anyway. Disable capture on every phase "
+        "(--cuda-graph-backend-decode=disabled --cuda-graph-backend-prefill=disabled) "
+        "if you genuinely need symmetric memory."
+    )
+
+
+def apply_kimi_k3_linear_attn_defaults(server_args: ServerArgs) -> None:
+    """KDA decode-fallback default for Kimi hybrid models (spec-independent)."""
+    from sglang.srt.utils import is_sm100_supported
+
+    # Preempts the generic SM100+bf16 flashinfer switch (a GDN default): on
+    # KDA shapes the triton packed decode measures ~35% faster than
+    # recurrent_kda across bs 1-256, and ReplaySSM requires triton.
+    if (
+        server_args.linear_attn_decode_backend is None
+        and server_args.mamba_ssm_dtype == "bfloat16"
+        and is_sm100_supported()
+    ):
+        server_args.linear_attn_decode_backend = "triton"
+        logger.info(
+            "Kimi hybrid model with bf16 SSM state: defaulting "
+            "--linear-attn-decode-backend to triton."
+        )

python/sglang/srt/arg_groups/overrides.py

@@ -30,6 +30,7 @@ Two declaration forms, keyed on ``hf_config.architectures[0]``:
 from __future__ import annotations
 
 import dataclasses
+import inspect
 import logging
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
@@ -39,6 +40,7 @@ from sglang.srt.model_executor.cuda_graph_config import Backend
 from sglang.srt.utils.common import (
     cpu_has_amx_support,
     get_device_capability,
+    get_device_name,
     get_device_sm,
     get_nvidia_driver_version,
     get_quantization_config,
@@ -48,6 +50,7 @@ from sglang.srt.utils.common import (
     is_flashinfer_available,
     is_gfx95_supported,
     is_hip,
+    is_mnnvl_fabric_device,
     is_musa,
     is_npu,
     is_sm90_supported,
@@ -315,6 +318,234 @@ def _register_for(*architectures: str):
     return decorator
 
 
+def _dspark_verify_on_decode_backend(
+    backend: Optional[str], q_len: int, kv_cache_dtype: Optional[str]
+) -> bool:
+    """Whether the MLA decode backend can serve a q_len-wide target verify."""
+    if backend == "trtllm_mla":
+        return True
+    if backend == "tokenspeed_mla":
+        return kv_cache_dtype == "fp8_e4m3" and q_len <= 8
+    if backend == "cutedsl_mla":
+        # The cute-dsl kernel rejects q_len >= 5 with no fallback.
+        return q_len <= 4
+    return False
+
+
+_KIMI_K3_DCP_PATCH_URL = (
+    "https://github.com/sgl-project/sglang/blob/"
+    "b701464720ca22aa1851d5dda7144e84a410f2c7/"
+    "docker/kimi_k3/kimi_k3_cu13.Dockerfile#L116-L123"
+)
+
+
+def _require_kimi_k3_cutedsl_dcp_support() -> None:
+    try:
+        from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
+
+        parameters = inspect.signature(trtllm_batch_decode_with_kv_cache_mla).parameters
+    except (ImportError, TypeError, ValueError) as exc:
+        raise RuntimeError(
+            "Kimi-K3 DCP with decode_attention_backend='cutedsl_mla' requires "
+            "a DCP-patched FlashInfer "
+            "trtllm_batch_decode_with_kv_cache_mla exposing enable_dcp in its "
+            f"signature. Apply the patch as shown in {_KIMI_K3_DCP_PATCH_URL}."
+        ) from exc
+
+    if "enable_dcp" not in parameters:
+        raise RuntimeError(
+            "Kimi-K3 DCP with decode_attention_backend='cutedsl_mla' requires "
+            "enable_dcp in the signature of "
+            "flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla. Apply "
+            f"the FlashInfer DCP patch as shown in {_KIMI_K3_DCP_PATCH_URL}."
+        )
+
+
+@_register_for("KimiK3ForConditionalGeneration")
+def _kimi_k3_overrides(server_args: Any, hf_config: Any) -> dict:
+    if server_args.dcp_size > 1:
+        overrides = {}
+        if server_args.enable_symm_mem:
+            logger.warning(
+                "Kimi-K3 DCP disables --enable-symm-mem due to decode CUDA "
+                "graph correctness issues."
+            )
+            overrides["enable_symm_mem"] = False
+
+        if server_args.speculative_algorithm == "DSPARK":
+            from sglang.srt.speculative.ragged_verify import (
+                RaggedVerifyMode,
+                read_ragged_verify_mode,
+            )
+
+            ragged_mode = read_ragged_verify_mode()
+            if ragged_mode is not RaggedVerifyMode.STATIC:
+                raise ValueError(
+                    "Kimi-K3 DCP + DSPARK currently requires "
+                    "SGLANG_RAGGED_VERIFY_MODE=static; compact/cap-accept are "
+                    f"not validated under DCP (got {ragged_mode.value!r})."
+                )
+
+            # DSPARK target-verify + draft-extend must run on the decode
+            # (cutedsl_mla) backend, whose _run_decode_kernel implements the DCP
+            # signature (causal_seqs / cp_world / cp_rank). The default
+            # "prefill" routes verify to trtllm_mla, whose base _run_decode_kernel
+            # lacks that DCP path (TypeError: unexpected kwarg 'causal_seqs').
+            overrides["speculative_attention_mode"] = "decode"
+
+        prefill_backend, decode_backend = attention_backends_of(server_args)
+        if decode_backend == "cutedsl_mla" or decode_backend is None:
+            _require_kimi_k3_cutedsl_dcp_support()
+            logger.info(
+                "Kimi-K3 DCP keeps decode attention backend 'cutedsl_mla' "
+                f"(prefill={prefill_backend!r} -> 'trtllm_mla')."
+            )
+            overrides.update(
+                prefill_attention_backend="trtllm_mla",
+                decode_attention_backend="cutedsl_mla",
+            )
+        elif decode_backend == "tokenspeed_mla":
+            logger.info(
+                "Kimi-K3 DCP overrides attention backends: "
+                f"prefill={prefill_backend!r}, decode={decode_backend!r} -> "
+                "'tokenspeed_mla'."
+            )
+            logger.info(
+                "Kimi-K3 DCP with tokenspeed mla backend overrides KV cache dtype: "
+                f"{server_args.kv_cache_dtype!r} -> 'fp8_e4m3'."
+            )
+            overrides.update(
+                prefill_attention_backend="tokenspeed_mla",
+                decode_attention_backend="tokenspeed_mla",
+                kv_cache_dtype="fp8_e4m3",
+            )
+        else:
+            raise AssertionError(
+                f"Decode attention backend for Kimi-K3 DCP must be 'cutedsl_mla' or 'tokenspeed_mla', got {decode_backend!r}."
+            )
+
+        if server_args.dcp_replicate_q_proj is None:
+            logger.info("Kimi-K3 DCP enables replicated Q projection by default.")
+            overrides["dcp_replicate_q_proj"] = True
+
+        device_name = get_device_name()
+        dcp_comm_backend = "fi_a2a" if is_mnnvl_fabric_device() else "a2a"
+        logger.info(
+            "Kimi-K3 DCP selects communication backend on "
+            f"{device_name!r}: {server_args.dcp_comm_backend!r} -> "
+            f"{dcp_comm_backend!r}."
+        )
+        overrides["dcp_comm_backend"] = dcp_comm_backend
+        return overrides
+
+    if not (is_sm100_supported() and get_device_sm() in (100, 103)):
+        return {}
+    backends_unset = server_args.is_attention_backend_not_set()
+    if server_args.speculative_algorithm != "DSPARK":
+        if not backends_unset:
+            return {}
+        logger.info(
+            "Use trtllm_mla as the default prefill and decode attention "
+            "backend for Kimi-K3 on SM100/SM103."
+        )
+        return {
+            "decode_attention_backend": "trtllm_mla",
+            "prefill_attention_backend": "trtllm_mla",
+        }
+    # DSPARK: verify runs on the decode backend (mode=decode below), so this
+    # picks the verify kernel -- mode=prefill routes it to flashinfer, which is
+    # slow and syncs, while plain decode is cold under dspark.
+    q_len = server_args.speculative_num_draft_tokens or (
+        server_args.speculative_dspark_block_size + 1
+        if server_args.speculative_dspark_block_size is not None
+        # Checkpoint auto-infer happens after overrides; K3 draft uses block 7.
+        else 8
+    )
+    overrides = {}
+    if backends_unset:
+        backend = "trtllm_mla"
+        overrides["decode_attention_backend"] = backend
+        overrides["prefill_attention_backend"] = "trtllm_mla"
+    else:
+        # Explicit backend knobs keep priority, but the mode is a separate knob
+        # that still needs declaring -- else verify stays on the prefill backend,
+        # whose host-side plan (flashinfer by default) forces a per-step D2H.
+        _, backend = attention_backends_of(server_args)
+    if _dspark_verify_on_decode_backend(backend, q_len, server_args.kv_cache_dtype):
+        overrides["speculative_attention_mode"] = "decode"
+        logger.info(
+            "Kimi-K3 DSPARK on SM100/SM103: decode/verify attention backend "
+            f"{backend} (speculative_attention_mode=decode)."
+        )
+    else:
+        logger.warning(
+            f"Kimi-K3 DSPARK: decode attention backend {backend!r} cannot serve "
+            f"target verify at q_len={q_len}, so verify runs on the prefill "
+            "backend (speculative_attention_mode=prefill). A host-plan prefill "
+            "backend costs a per-step seq_lens D2H sync; leave the attention "
+            "backend knobs unset for the sync-free default."
+        )
+    return overrides
+
+
+def _is_mxfp4_pack_quantized(hf_config: Any) -> bool:
+    qc = getattr(
+        getattr(hf_config, "text_config", hf_config), "quantization_config", None
+    )
+    if not isinstance(qc, dict):
+        return False
+    groups = qc.get("config_groups") or {}
+    return any(
+        "mxfp4" in str(g.get("format", ""))
+        for g in groups.values()
+        if isinstance(g, dict)
+    )
+
+
+@_register_for("KimiK3ForConditionalGeneration")
+def _kimi_k3_moe_runner_overrides(server_args: Any, hf_config: Any) -> dict:
+    # MoE runner default, independent of the attention-backend gate above.
+    # trtllm-gen fused MoE (flashinfer_mxfp4) beats marlin on both the decode
+    # (M=bs) and the target-verify (M=bs*(gamma+1)) regimes on SM100/SM103;
+    # it hard-requires the SiTU cubin pool on the box (K3's SiTU activation has
+    # no public cubins). Do not silently trade W4A8 for Marlin W4A16 when the
+    # default cannot start; explicit non-FlashInfer runner choices still win.
+    if server_args.moe_runner_backend not in ("auto", "flashinfer_mxfp4"):
+        return {}
+    if not (is_sm100_supported() and get_device_sm() in (100, 103)):
+        return {}
+    if not _is_mxfp4_pack_quantized(hf_config):
+        return {}
+    from sglang.kernels.ops.moe.trtllm_gen_moe import available as _trtllm_gen_moe_ok
+
+    if not _trtllm_gen_moe_ok():
+        raise RuntimeError(
+            "Kimi-K3 on Blackwell with moe_runner_backend='auto' or "
+            "'flashinfer_mxfp4' requires a valid "
+            "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL. Install it with:\n"
+            "wget https://github.com/sgl-project/whl/releases/download/"
+            "trtllm_gen_moe_cubin_20260617/"
+            "trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip\n"
+            "sudo mkdir -p /opt/trtllm_gen_moe_cubin_pool\n"
+            "sudo unzip -q "
+            "trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip -d "
+            "/opt/trtllm_gen_moe_cubin_pool\n"
+            "export "
+            "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/opt/trtllm_gen_moe_cubin_pool/"
+            "trtllm_gen_moe_cubin_pool_20260617_v0613rc1\n"
+            "To use Marlin "
+            "instead, set --moe-runner-backend marlin explicitly."
+        )
+
+    if server_args.moe_runner_backend == "auto":
+        logger.info(
+            "Kimi-K3 on SM100/SM103: moe_runner_backend=flashinfer_mxfp4 "
+            "(trtllm-gen SiTU cubin pool found)."
+        )
+        return {"moe_runner_backend": "flashinfer_mxfp4"}
+    return {}
+
+
 @_register_for(
     "DeepseekV3ForCausalLM",
     "DeepseekV32ForCausalLM",
@@ -1175,6 +1406,7 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict:
 _MAMBA_RADIX_CACHE_ARCHS = frozenset(
     {
         "KimiLinearForCausalLM",
+        "KimiK3ForConditionalGeneration",
         "BailingMoeV2_5ForCausalLM",
         "Qwen3NextForCausalLM",
         "Qwen3_5MoeForConditionalGeneration",
@@ -1208,6 +1440,10 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset(
         "GraniteMoeHybridForCausalLM",
         "NemotronHForCausalLM",
         "NemotronHPuzzleForCausalLM",
+        # KDA-based: same MambaPool ping-pong machinery as GDN; requires the
+        # KDA backend's track-snapshot writes (decode + extend) so donated
+        # slots hold real states for prefix-cache restores.
+        "KimiK3ForConditionalGeneration",
     }
 )
 

python/sglang/srt/configs/__init__.py

@@ -19,6 +19,7 @@ from sglang.srt.configs.interns2preview import InternS2PreviewConfig
 from sglang.srt.configs.janus_pro import MultiModalityConfig
 from sglang.srt.configs.jet_nemotron import JetNemotronConfig
 from sglang.srt.configs.jet_vlm import JetVLMConfig
+from sglang.srt.configs.kimi_k3 import KimiK3Config
 from sglang.srt.configs.kimi_k25 import KimiK25Config
 from sglang.srt.configs.kimi_linear import KimiLinearConfig
 from sglang.srt.configs.kimi_vl import KimiVLConfig
@@ -71,6 +72,7 @@ __all__ = [
     "Step3VisionEncoderConfig",
     "Olmo3Config",
     "KimiLinearConfig",
+    "KimiK3Config",
     "KimiK25Config",
     "LagunaConfig",
     "Qwen3NextConfig",

python/sglang/srt/configs/hybrid_arch.py

@@ -102,6 +102,9 @@ def kimi_linear_config(model_config: ModelConfig):
     config = model_config.hf_config
     if isinstance(config, KimiLinearConfig):
         return config
+    text_config = getattr(config, "text_config", None)
+    if isinstance(text_config, KimiLinearConfig):
+        return text_config
     return None
 
 

python/sglang/srt/configs/kimi_k3.py

@@ -0,0 +1,124 @@
+from transformers.configuration_utils import PretrainedConfig
+
+from sglang.srt.configs.kimi_linear import KimiLinearConfig
+
+
+class KimiK3VisionConfig(PretrainedConfig):
+    model_type = "kimi_k3_vision"
+
+    def __init__(
+        self,
+        patch_size: int = 14,
+        init_pos_emb_height: int = 64,
+        init_pos_emb_width: int = 64,
+        init_pos_emb_time: int = 4,
+        pos_emb_type: str = "divided_fixed",
+        vt_num_attention_heads: int = 12,
+        vt_num_hidden_layers: int = 27,
+        vt_hidden_size: int = 1024,
+        vt_intermediate_size: int = 4096,
+        merge_kernel_size: tuple[int, int] = (2, 2),
+        video_attn_type: str = "spatial_temporal",
+        merge_type: str = "sd2_tpool",
+        _attn_implementation: str = "flash_attention_2",
+        mm_projector_type: str = "patchmergerv2",
+        mm_hidden_size: int | None = None,
+        projector_hidden_act: str = "gelu",
+        projector_ln_eps: float = 1e-5,
+        qkv_hidden_size: int = 1536,
+        norm_type: str = "rmsnorm",
+        attn_bias: bool = False,
+        patch_embed_proj_bias: bool = False,
+        mlp_type: str = "mlp2",
+        linear_bias: bool = False,
+        activation_func: str = "gelu_pytorch_tanh",
+        pos_emb_interpolation_mode: str = "bilinear",
+        text_hidden_size: int = 2304,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.patch_size = patch_size
+        self.init_pos_emb_height = init_pos_emb_height
+        self.init_pos_emb_width = init_pos_emb_width
+        self.init_pos_emb_time = init_pos_emb_time
+        self.pos_emb_type = pos_emb_type
+        self.vt_num_attention_heads = vt_num_attention_heads
+        self.vt_num_hidden_layers = vt_num_hidden_layers
+        self.vt_hidden_size = vt_hidden_size
+        self.vt_intermediate_size = vt_intermediate_size
+        self.merge_kernel_size = tuple(merge_kernel_size)
+        self.video_attn_type = video_attn_type
+        self.merge_type = merge_type
+        self._attn_implementation = _attn_implementation
+
+        self.mm_projector_type = mm_projector_type
+        self.mm_hidden_size = (
+            mm_hidden_size if mm_hidden_size is not None else vt_hidden_size
+        )
+        self.projector_hidden_act = projector_hidden_act
+        self.projector_ln_eps = projector_ln_eps
+        self.text_hidden_size = text_hidden_size
+
+        self.qkv_hidden_size = qkv_hidden_size
+        self.norm_type = norm_type
+        self.attn_bias = attn_bias
+        self.patch_embed_proj_bias = patch_embed_proj_bias
+        self.mlp_type = mlp_type
+        self.linear_bias = linear_bias
+        self.activation_func = activation_func
+        self.pos_emb_interpolation_mode = pos_emb_interpolation_mode
+
+        # Aliases consumed by the K2.5 vision implementation.
+        self.num_attention_heads = vt_num_attention_heads
+        self.num_hidden_layers = vt_num_hidden_layers
+        self.hidden_size = vt_hidden_size
+        self.intermediate_size = vt_intermediate_size
+
+
+class KimiK3Config(PretrainedConfig):
+    model_type = "kimi_k3"
+
+    def __init__(
+        self,
+        text_config: dict | KimiLinearConfig | None = None,
+        vision_config: dict | KimiK3VisionConfig | None = None,
+        ignore_index: int = -100,
+        media_placeholder_token_id: int = 163605,
+        pad_token_id: int = 0,
+        image_placeholder: str = "<|kimi_image_placeholder|>",
+        **kwargs,
+    ):
+        if text_config is None:
+            self.text_config = KimiLinearConfig()
+        elif isinstance(text_config, dict):
+            self.text_config = KimiLinearConfig(**text_config)
+        else:
+            self.text_config = text_config
+
+        if vision_config is None:
+            self.vision_config = KimiK3VisionConfig()
+        elif isinstance(vision_config, dict):
+            self.vision_config = KimiK3VisionConfig(**vision_config)
+        else:
+            self.vision_config = vision_config
+
+        if self.vision_config.text_hidden_size != self.text_config.hidden_size:
+            self.vision_config.text_hidden_size = self.text_config.hidden_size
+
+        self.ignore_index = ignore_index
+        self.media_placeholder_token_id = media_placeholder_token_id
+        self.image_placeholder = image_placeholder
+
+        if getattr(self.text_config, "quantization_config", None) is not None:
+            self.quantization_config = self.text_config.quantization_config
+
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+    @property
+    def hidden_size(self) -> int:
+        return self.text_config.hidden_size
+
+    @property
+    def vocab_size(self) -> int:
+        return self.text_config.vocab_size

python/sglang/srt/configs/kimi_linear.py

@@ -43,6 +43,7 @@ class KimiLinearConfig(PretrainedConfig):
         use_grouped_topk: bool = True,
         num_expert_group: int = 1,
         topk_group: int = 1,
+        topk_method: str = "noaux_tc",
         q_lora_rank: int | None = None,
         kv_lora_rank: int | None = None,
         qk_nope_head_dim: int | None = None,
@@ -51,6 +52,13 @@ class KimiLinearConfig(PretrainedConfig):
         mla_use_nope: bool | None = False,
         num_nextn_predict_layers: int = 0,
         linear_attn_config: dict | None = None,
+        attn_res_block_size: int | None = None,
+        routed_expert_hidden_size: int | None = None,
+        latent_moe_use_norm: bool = False,
+        activation_situ_beta: float | None = None,
+        activation_situ_linear_beta: float | None = None,
+        mla_use_output_gate: bool = False,
+        max_position_embeddings: int = 4096,
         **kwargs,
     ):
         self.model_type = model_type
@@ -83,6 +91,7 @@ class KimiLinearConfig(PretrainedConfig):
         self.mla_use_nope = mla_use_nope
         # moe config
         self.n_routed_experts = self.num_experts = num_experts
+        self.topk_method = topk_method
         self.num_experts_per_token = num_experts_per_token
         self.moe_renormalize = moe_renormalize
         self.num_shared_experts = num_shared_experts
@@ -97,6 +106,14 @@ class KimiLinearConfig(PretrainedConfig):
         self.topk_group = topk_group
         self.num_nextn_predict_layers = num_nextn_predict_layers
 
+        self.attn_res_block_size = attn_res_block_size
+        self.routed_expert_hidden_size = routed_expert_hidden_size
+        self.latent_moe_use_norm = latent_moe_use_norm
+        self.activation_situ_beta = activation_situ_beta
+        self.activation_situ_linear_beta = activation_situ_linear_beta
+        self.mla_use_output_gate = mla_use_output_gate
+        self.max_position_embeddings = max_position_embeddings
+
         if linear_attn_config is not None:
             assert linear_attn_config["kda_layers"] is not None
             assert linear_attn_config["full_attn_layers"] is not None

python/sglang/srt/configs/mamba_utils.py

@@ -128,18 +128,26 @@ class BaseLinearStateParams(ABC):
         """Per-slot bytes of the ReplaySSM spec-verify fold window (all
         layers). Not part of ``mamba_cache_per_req``, so the memory solver
         must charge it separately. MUST mirror the ``MambaPool`` allocation:
-        raw v/k in the conv dtype + fp32 g and beta. GDN scalar-g layout
-        only."""
-        assert not self.is_kda, "replayssm ring accounting supports GDN only"
+        raw v/k in the conv dtype + fp32 beta, plus the fp32 gate ring
+        (per-head scalar for GDN, per-K vector for KDA). KDA additionally
+        keeps the chunked d/k rings under spec (its forward_decode routes on
+        their presence), also in the conv dtype."""
         hv, v_dim, k_dim = self.shape.temporal
         h_k = self.shape.num_k_heads_per_tp
         conv_b = self.dtype.conv.itemsize
         fp32_b = 4
         per_layer = (
-            hv * record_len * v_dim * conv_b
-            + h_k * record_len * k_dim * conv_b
-            + 2 * hv * record_len * fp32_b
+            hv * record_len * v_dim * conv_b  # rawv
+            + h_k * record_len * k_dim * conv_b  # rawk
+            + hv * record_len * fp32_b  # beta (fp32)
+            # g (fp32): GDN per-head scalar, KDA per-K vector
+            + hv * record_len * (k_dim if self.is_kda else 1) * fp32_b
         )
+        if self.is_kda:
+            per_layer += (
+                hv * record_len * v_dim * conv_b  # d
+                + h_k * record_len * k_dim * conv_b  # k
+            )
         return per_layer * len(self.layers)
 
     @property

python/sglang/srt/configs/model_config.py

@@ -121,6 +121,10 @@ def is_deepseek_dsa(config) -> bool:
     )
 
 
+def is_kimi_k3(config) -> bool:
+    return _hf_arch(config) == "KimiK3ForConditionalGeneration"
+
+
 def is_deepseek_v4(config) -> bool:
     return _hf_arch(config) in (
         "DeepseekV4ForCausalLM",
@@ -484,6 +488,9 @@ class ModelConfig:
         self.is_multimodal_breakable_cuda_graph_supported = enable_multimodal and (
             is_multimodal_breakable_cuda_graph_supported(self.hf_config.architectures)
         )
+        self.is_mla_breakable_cuda_graph_supported = (
+            is_mla_breakable_cuda_graph_supported(self.hf_config.architectures)
+        )
         self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
 
         # Derive context length and model shapes
@@ -912,18 +919,20 @@ class ModelConfig:
             self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
             self.v_head_dim = self.hf_text_config.v_head_dim
             self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
-        elif "KimiLinearForCausalLM" in self.hf_config.architectures:
+        elif (
+            "KimiLinearForCausalLM" in self.hf_config.architectures
+            or "KimiK3ForConditionalGeneration" in self.hf_config.architectures
+        ):
+            tc = self.hf_text_config
             self.head_dim = 72
             self.attention_arch = AttentionArch.MLA
-            self.kv_lora_rank = self.hf_config.kv_lora_rank
-            self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
-            self.v_head_dim = self.hf_config.v_head_dim
-            self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim
+            self.kv_lora_rank = tc.kv_lora_rank
+            self.qk_rope_head_dim = tc.qk_rope_head_dim
+            self.v_head_dim = tc.v_head_dim
+            self.qk_nope_head_dim = tc.qk_nope_head_dim
             self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
-            if self.hf_config.rope_scaling:
-                self.scaling = compute_mla_mscale_scaling(
-                    self.hf_config.rope_scaling, self.scaling
-                )
+            if getattr(tc, "rope_scaling", None):
+                self.scaling = compute_mla_mscale_scaling(tc.rope_scaling, self.scaling)
         elif (
             "BailingMoeV2_5ForCausalLM" in self.hf_config.architectures
             or "BailingMoeForCausalLMNextN" in self.hf_config.architectures
@@ -1841,6 +1850,15 @@ multimodal_breakable_cuda_graph_supported_model_archs = [
     "Qwen3_5MoeForConditionalGeneration",
 ]
 
+# MLA archs validated to run breakable CUDA graph when it is explicitly
+# requested (--cuda-graph-backend-prefill=breakable bypasses the ServerArgs
+# disable rules). Dispatch pins the absorbed MLA path inside capture/replay
+# for these archs, so the prefill runner's MHA-companion prefix restrictions
+# do not apply (see PrefillCudaGraphRunner.mla_pinned_under_bcg).
+mla_breakable_cuda_graph_supported_model_archs = [
+    "KimiK3ForConditionalGeneration",
+]
+
 if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get():
     multimodal_model_archs.append(external_mm_model_arch)
 
@@ -1915,6 +1933,14 @@ def is_multimodal_breakable_cuda_graph_supported(model_architectures: List[str])
     )
 
 
+def is_mla_breakable_cuda_graph_supported(model_architectures: List[str]):
+    """Whether an MLA arch may keep prefill breakable CUDA graph enabled."""
+    return any(
+        arch in mla_breakable_cuda_graph_supported_model_archs
+        for arch in model_architectures
+    )
+
+
 # SequenceClassification models that use CrossEncodingPooler
 _cross_encoding_pooler_archs = [
     "BertForSequenceClassification",
@@ -1936,8 +1962,13 @@ def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float
     """Compute MLA attention scaling factor from rope_scaling with mscale.
 
     Used by DeepSeek, BailingMoe, SarvamMLA and similar MLA models.
-    Warns if 'factor' is missing from rope_scaling (common in v5 configs).
+    Transformers v5 also exposes the default RoPE parameters through
+    ``rope_scaling``. Those parameters do not request any scaling.
     """
+    rope_type = rope_scaling.get("rope_type") or rope_scaling.get("type")
+    if rope_type == "default":
+        return base_scaling
+
     if not rope_scaling.get("apply_yarn_scaling", True) or not rope_scaling.get(
         "apply_scale", True
     ):

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

@@ -206,6 +206,7 @@ class CommonKVManager(BaseKVManager):
         self.request_status: Dict[int, KVPoll] = {}
         self._socket_cache: Dict[str, zmq.Socket] = {}
         self._monitor_cache: Dict[str, zmq.Socket] = {}
+        self._socket_send_locks: Dict[str, threading.Lock] = {}
         self._socket_lock = threading.Lock()
         self.failure_records: Dict[int, str] = {}
         self.failure_lock = threading.Lock()
@@ -776,8 +777,18 @@ class CommonKVManager(BaseKVManager):
             self._monitor_cache[endpoint] = sock.get_monitor_socket(
                 zmq.EVENT_DISCONNECTED
             )
+            self._socket_send_locks.setdefault(endpoint, threading.Lock())
             return sock
 
+    def _send_multipart_locked(
+        self, endpoint: str, parts: List[bytes], is_ipv6: bool = False
+    ):
+        # Cached sockets are shared across sender threads and zmq sockets are
+        # not thread-safe; serialize sends per endpoint.
+        sock = self._connect(endpoint, is_ipv6=is_ipv6)
+        with self._socket_send_locks[endpoint]:
+            sock.send_multipart(parts)
+
     def get_mha_kv_ptrs_with_pp(
         self, src_kv_ptrs: List[int], dst_kv_ptrs: List[int]
     ) -> Tuple[List[int], List[int], List[int], List[int], int]:

python/sglang/srt/disaggregation/decode.py

@@ -211,7 +211,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
         speculative_eagle_topk: Optional[int] = None,
         linear_replayssm_cache_len: int = 16,
         mamba_envelope_layout: bool = False,
-        enable_gdn_replayssm_spec: bool = False,
+        enable_linear_replayssm_spec: bool = False,
     ):
         DecodeReqToTokenPool.__init__(
             self,
@@ -258,7 +258,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
             speculative_eagle_topk=speculative_eagle_topk,
             linear_replayssm_cache_len=linear_replayssm_cache_len,
             mamba_envelope_layout=mamba_envelope_layout,
-            enable_gdn_replayssm_spec=enable_gdn_replayssm_spec,
+            enable_linear_replayssm_spec=enable_linear_replayssm_spec,
         )
 
     def clear(self):

python/sglang/srt/disaggregation/encode_receiver.py

@@ -481,6 +481,7 @@ _GENERAL_VIDEO_META_ATTRS = (
     "video_timestamps",
     "second_per_grid_ts",
 )
+_GENERAL_IMAGE_META_ATTRS = ("original_image_sizes",)
 # MiMo-VL audio-in-video fields; appended only when model_type is MiMo.
 _MIMO_VIDEO_AUDIO_META_ATTRS = (
     "video_audio_feature_lens",
@@ -551,6 +552,8 @@ class MultiModalEmbeddingData(EmbeddingData):
         self.img_grid_thw = [None] * num_parts
         self.video_grid_thw = [None] * num_parts
         self.audio_feature_lens = [None] * num_parts
+        for attr in _GENERAL_IMAGE_META_ATTRS:
+            setattr(self, attr, [None] * num_parts)
         self.modality_list = [
             modality if part_idx == i else None for i in range(num_parts)
         ]
@@ -567,6 +570,18 @@ class MultiModalEmbeddingData(EmbeddingData):
         self._set_part_grid(part_idx, modality, self.get_grid())
         if modality == Modality.VIDEO:
             self._set_video_meta_for_part(part_idx, kwargs)
+        if modality == Modality.IMAGE:
+            self._set_image_meta_for_part(part_idx, kwargs)
+
+    def _set_image_meta_for_part(self, part_idx, source):
+        for attr_name in _GENERAL_IMAGE_META_ATTRS:
+            val = (
+                source.get(attr_name)
+                if isinstance(source, dict)
+                else getattr(source, attr_name, None)
+            )
+            if val is not None:
+                getattr(self, attr_name)[part_idx] = val
 
     def _set_part_grid(self, part_idx, modality, grid):
         """Set the grid for one part according to modality (IMAGE/VIDEO/AUDIO)."""
@@ -601,6 +616,10 @@ class MultiModalEmbeddingData(EmbeddingData):
             val = getattr(embedding_data, attr, None)
             if val is not None:
                 extra[attr] = val
+        for attr in _GENERAL_IMAGE_META_ATTRS:
+            val = getattr(embedding_data, attr, None)
+            if val is not None:
+                extra[attr] = val
         mm_data = cls(
             part_idx=embedding_data.part_idx,
             num_parts=embedding_data.num_parts,
@@ -651,6 +670,10 @@ class MultiModalEmbeddingData(EmbeddingData):
                 kwargs[attr] = torch.cat(valid, dim=0)
             else:
                 kwargs[attr] = list(itertools.chain(*valid))
+        for attr in _GENERAL_IMAGE_META_ATTRS:
+            valid = [value for value in getattr(self, attr) if value is not None]
+            if valid:
+                kwargs[attr] = list(itertools.chain(*valid))
         return kwargs
 
     def add(self, embedding_data: EmbeddingData):
@@ -669,6 +692,8 @@ class MultiModalEmbeddingData(EmbeddingData):
         self._set_part_grid(pid, embedding_data.modality, embedding_data.get_grid())
         if embedding_data.modality == Modality.VIDEO:
             self._set_video_meta_for_part(pid, embedding_data)
+        if embedding_data.modality == Modality.IMAGE:
+            self._set_image_meta_for_part(pid, embedding_data)
 
 
 class WaitingImageRequestStatus(IntEnum):
@@ -678,6 +703,13 @@ class WaitingImageRequestStatus(IntEnum):
     TIMEOUT = -2
 
 
+def _select_mm_processor_prompt(recv_req, mm_processor):
+    """Mirror tokenizer-side prompt selection for scheduler-side EPD rebuilds."""
+    if mm_processor.prefer_tokenized_input and recv_req.input_ids is not None:
+        return list(recv_req.input_ids)
+    return recv_req.input_text or recv_req.input_ids
+
+
 def create_part_req_id(original_req_id: str, part_idx: int) -> str:
     """Create a unique part request ID by appending part index suffix."""
     return f"{original_req_id}_local_part_{part_idx}"
@@ -724,6 +756,8 @@ class WaitingImageRequest:
         model_type,
         host_name,
         receive_count,
+        zmq_context=None,
+        embedding_port=None,
     ):
         self.rid = rid
         self.recv_req = recv_req
@@ -736,9 +770,16 @@ class WaitingImageRequest:
         self.host_name = host_name
         self.receive_count = receive_count
         self.num_items_assigned = recv_req.num_items_assigned
-        self.embedding_port, self.recv_socket = get_zmq_socket_on_host(
-            zmq.Context(), zmq.PULL, host=host_name
-        )
+        self.zmq_context = zmq_context
+        if embedding_port is None:
+            if self.zmq_context is None:
+                raise ValueError("zmq_context is required for a per-request socket")
+            self.embedding_port, self.recv_socket = get_zmq_socket_on_host(
+                self.zmq_context, zmq.PULL, host=host_name
+            )
+        else:
+            self.embedding_port = embedding_port
+            self.recv_socket = None
         logger.info(f"Waiting for input {self.embedding_port = }")
         self.recv_embedding_data = None
         # ok=1 pending=0 fail=-1
@@ -835,71 +876,78 @@ class WaitingImageRequest:
     def _try_recv_mm_data(self):
         if self.status != WaitingImageRequestStatus.PENDING:
             return
-        while self.recv_embedding_data is None or not self.recv_embedding_data.ready:
+        if self.recv_socket is None:
+            return
+        while self.status == WaitingImageRequestStatus.PENDING:
             try:
                 parts = self.recv_socket.recv_multipart(flags=zmq.NOBLOCK, copy=False)
             except zmq.Again:
                 # No data available yet, wait a bit and retry
                 return
-            try:
-                recv_obj: EmbeddingData = safe_pickle_loads(parts[0])
-                if getattr(recv_obj, "error_msg", None) is not None:
-                    logger.warning(
-                        f"Received error signal from encoder for {self.rid}: {recv_obj.error_msg} {recv_obj.error_code = }"
-                    )
-                    self.error_msg = recv_obj.error_msg
-                    self.error_code = recv_obj.error_code
-                    self.status = WaitingImageRequestStatus.FAIL
-                    self.recv_socket.close()
-                    return
+            self.consume_parts(parts)
 
-                # Extract original req_id from part_req_id and drop stale payloads
-                # that may arrive on a reused ZMQ port after a prior request aborted.
-                original_req_id = extract_original_req_id(recv_obj.req_id)
-                if original_req_id != self.recv_req.rid:
-                    logger.warning(
-                        f"Dropping stale embedding data: expected rid={self.recv_req.rid}, "
-                        f"got rid={recv_obj.req_id} (likely from ZMQ port reuse)"
-                    )
-                    continue
-                recv_obj.req_id = original_req_id
+    def consume_parts(self, parts):
+        if self.status != WaitingImageRequestStatus.PENDING:
+            return
 
-                buffer = parts[1].buffer if hasattr(parts[1], "buffer") else parts[1]
-                recv_obj.embedding = (
-                    torch.frombuffer(buffer, dtype=recv_obj.dtype)
-                    .reshape(recv_obj.shape)
-                    .clone()
+        try:
+            recv_obj: EmbeddingData = safe_pickle_loads(parts[0])
+            if getattr(recv_obj, "error_msg", None) is not None:
+                logger.warning(
+                    f"Received error signal from encoder for {self.rid}: {recv_obj.error_msg} {recv_obj.error_code = }"
                 )
-
-                if self.recv_embedding_data is None:
-                    self.recv_embedding_data = (
-                        MultiModalEmbeddingData.from_embedding_data(
-                            recv_obj, model_type=self.model_type
-                        )
-                    )
-                else:
-                    self.recv_embedding_data.add(recv_obj)
-            except Exception as e:
-                # A message the scheduler cannot decode (blocked unpickle,
-                # bad shape/dtype, ...) must fail this request, not crash the
-                # scheduler event loop; FAIL still reaches the TP-wide status
-                # all-reduce in _process_waiting_requests.
-                logger.exception(
-                    "Failed to decode embedding message for rid=%s", self.rid
-                )
-                self.error_msg = f"Failed to decode embedding message: {e}"
+                self.error_msg = recv_obj.error_msg
+                self.error_code = recv_obj.error_code
                 self.status = WaitingImageRequestStatus.FAIL
-                self._cleanup_gpu_buffer()
-                self.recv_socket.close()
+                self.close_recv_socket()
                 return
 
+            # Extract original req_id from part_req_id and drop stale payloads
+            # that may arrive on a reused ZMQ port after a prior request aborted.
+            original_req_id = extract_original_req_id(recv_obj.req_id)
+            if original_req_id != self.recv_req.rid:
+                logger.warning(
+                    f"Dropping stale embedding data: expected rid={self.recv_req.rid}, "
+                    f"got rid={recv_obj.req_id} (likely from ZMQ port reuse)"
+                )
+                return
+            recv_obj.req_id = original_req_id
+
+            buffer = parts[1].buffer if hasattr(parts[1], "buffer") else parts[1]
+            recv_obj.embedding = (
+                torch.frombuffer(buffer, dtype=recv_obj.dtype)
+                .reshape(recv_obj.shape)
+                .clone()
+            )
+
+            if self.recv_embedding_data is None:
+                self.recv_embedding_data = MultiModalEmbeddingData.from_embedding_data(
+                    recv_obj, model_type=self.model_type
+                )
+            else:
+                self.recv_embedding_data.add(recv_obj)
+        except Exception as e:
+            # A message the scheduler cannot decode (blocked unpickle,
+            # bad shape/dtype, ...) must fail this request, not crash the
+            # scheduler event loop; FAIL still reaches the TP-wide status
+            # all-reduce in _process_waiting_requests.
+            logger.exception("Failed to decode embedding message for rid=%s", self.rid)
+            self.error_msg = f"Failed to decode embedding message: {e}"
+            self.status = WaitingImageRequestStatus.FAIL
+            self._cleanup_gpu_buffer()
+            self.close_recv_socket()
+            return
+
+        if not self.recv_embedding_data.ready:
+            return
+
         # Assemble mm_inputs. Wrapped so an assembly failure still reaches the
         # TP-wide status all-reduce in _process_waiting_requests instead of
         # raising past it.
         try:
             recv_embedding = self.recv_embedding_data.get_embedding(is_concat=True)
             mm_inputs = self.mm_processor.get_mm_data(
-                self.recv_req.input_text,
+                _select_mm_processor_prompt(self.recv_req, self.mm_processor),
                 recv_embedding,
                 **self.recv_embedding_data.get_mm_extra_meta(),
             )
@@ -913,7 +961,12 @@ class WaitingImageRequest:
             self.status = WaitingImageRequestStatus.FAIL
             self.error_msg = f"Failed to assemble multimodal inputs: {e}"
             self._cleanup_gpu_buffer()
-        self.recv_socket.close()
+        self.close_recv_socket()
+
+    def close_recv_socket(self):
+        if self.recv_socket is not None:
+            self.recv_socket.close()
+            self.recv_socket = None
 
     def _cleanup_gpu_buffer(self):
         pass
@@ -975,11 +1028,13 @@ class WaitingImageRDMARequest(WaitingImageRequest):
         encoder_urls,
         host_name,
         receive_count,
+        zmq_context,
         embeddings_engine,
         dtype,
         gpu_id=0,
         model_type: Optional[str] = None,
         embedding_pool=None,
+        embedding_port=None,
     ):
         super().__init__(
             rid=rid,
@@ -989,6 +1044,8 @@ class WaitingImageRDMARequest(WaitingImageRequest):
             model_type=model_type,
             host_name=host_name,
             receive_count=receive_count,
+            zmq_context=zmq_context,
+            embedding_port=embedding_port,
         )
         self.embeddings_engine = embeddings_engine
         self.dtype = dtype
@@ -1256,7 +1313,7 @@ class WaitingImageRDMARequest(WaitingImageRequest):
         else:
             recv_embedding = self.recv_embedding_data.get_embedding(is_concat=True)
         mm_inputs = self.mm_processor.get_mm_data(
-            self.recv_req.input_text,
+            _select_mm_processor_prompt(self.recv_req, self.mm_processor),
             recv_embedding,
             **self.recv_embedding_data.get_mm_extra_meta(),
         )
@@ -1511,6 +1568,10 @@ class MMReceiverBase(ABC):
         encode_urls: Optional[List[str]] = None,
     ):
         self.context = zmq.asyncio.Context(20)
+        # Scheduler-side receive is polled synchronously. Keep one regular ZMQ
+        # context alive for the process instead of creating a temporary context
+        # whose destruction also closes its per-request socket.
+        self.scheduler_context = zmq.Context()
         self.encoder_transfer_backend = server_args.encoder_transfer_backend
         # When ``encode_urls`` is shared with an :class:`EncoderBootstrapServer`
         # (tokenizer manager process), it grows / shrinks in place as encoders
@@ -1529,6 +1590,24 @@ class MMReceiverBase(ABC):
         self.nnodes = server_args.nnodes
         self.hostname = get_local_ip_auto()
         self.waiting_list: List[WaitingImageRequest] = []
+        self.waiting_by_rid: Dict[str, WaitingImageRequest] = {}
+        self.scheduler_embedding_port = None
+        self.scheduler_recv_socket = None
+        if (
+            self.encoder_transfer_backend == "zmq_to_scheduler"
+            and scheduler is not None
+        ):
+            (
+                self.scheduler_embedding_port,
+                self.scheduler_recv_socket,
+            ) = get_zmq_socket_on_host(
+                self.scheduler_context, zmq.PULL, host=self.hostname
+            )
+            logger.info(
+                "Scheduler TP rank %s reuses ZMQ embedding port %s",
+                self.tp_rank,
+                self.scheduler_embedding_port,
+            )
         self.scheduler = scheduler
         self.gpu_id = scheduler.ps.gpu_id if scheduler is not None else 0
         self.wait_timeout = envs.SGLANG_ENCODER_RECV_TIMEOUT.get()
@@ -1864,6 +1943,28 @@ class MMReceiverBase(ABC):
         waiting_req.error_code = best_code
 
     # For zmq_to_scheduler
+    def _drain_scheduler_embeddings(self):
+        if self.scheduler_recv_socket is None:
+            return
+
+        while True:
+            try:
+                parts = self.scheduler_recv_socket.recv_multipart(
+                    flags=zmq.NOBLOCK, copy=False
+                )
+            except zmq.Again:
+                return
+
+            recv_obj: EmbeddingData = safe_pickle_loads(parts[0])
+            rid = extract_original_req_id(recv_obj.req_id)
+            waiting_req = self.waiting_by_rid.get(rid)
+            if waiting_req is None:
+                logger.warning(
+                    "Dropping embedding data for inactive request %s", recv_obj.req_id
+                )
+                continue
+            waiting_req.consume_parts(parts)
+
     def _process_waiting_requests(self, recv_reqs, waiting_cls, **extra_kwargs):
         new_recv_reqs = []
         for recv_req in recv_reqs:
@@ -1886,8 +1987,16 @@ class MMReceiverBase(ABC):
                     model_type=self.model_type,
                     host_name=self.hostname,
                     receive_count=self.tp_size,
+                    zmq_context=(
+                        None
+                        if self.scheduler_recv_socket is not None
+                        else self.scheduler_context
+                    ),
+                    embedding_port=self.scheduler_embedding_port,
                     **extra_kwargs,
                 )
+                if self.scheduler_recv_socket is not None:
+                    self.waiting_by_rid[waiting_req.rid] = waiting_req
                 waiting_req.send_encode_request()
                 self.waiting_list.append(waiting_req)
             else:
@@ -1896,14 +2005,16 @@ class MMReceiverBase(ABC):
         if len(self.waiting_list) == 0:
             return new_recv_reqs, []
 
+        self._drain_scheduler_embeddings()
         current_time = time.time()
         local_status = []
         for waiting_req in self.waiting_list:
-            waiting_req._try_recv_mm_data()
+            if self.scheduler_recv_socket is None:
+                waiting_req._try_recv_mm_data()
             if current_time - waiting_req.start_time > self.wait_timeout:
                 waiting_req.status = WaitingImageRequestStatus.TIMEOUT
                 waiting_req._cleanup_gpu_buffer()
-                waiting_req.recv_socket.close()
+                waiting_req.close_recv_socket()
             local_status.append(waiting_req.status)
 
         local_status = torch.tensor(local_status, device="cpu", dtype=torch.int32)
@@ -1945,6 +2056,8 @@ class MMReceiverBase(ABC):
                 )
             else:  # status_value == WaitingImageRequestStatus.PENDING
                 new_waiting.append(waiting_req)
+                continue
+            self.waiting_by_rid.pop(waiting_req.rid, None)
 
         self.waiting_list = new_waiting
         return new_recv_reqs, abort_reqs

python/sglang/srt/disaggregation/encode_server.py

@@ -55,6 +55,9 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
 from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
+from sglang.srt.model_executor.model_runner_components.load_model_utils import (
+    maybe_precompile_model_kernels_after_loading,
+)
 from sglang.srt.model_loader import get_model
 from sglang.srt.multimodal.processors.qwen_vl import preprocess_video
 from sglang.srt.observability.metrics_collector import EncoderMetricsCollector
@@ -114,6 +117,7 @@ mooncake_send_done_count: Dict[str, int] = dict()
 use_image_processor_gpu = envs.SGLANG_ENCODER_IMAGE_PROCESSOR_USE_GPU.get()
 
 ENCODER_MAX_BATCH_SIZE = envs.SGLANG_ENCODER_MAX_BATCH_SIZE.get()
+ENCODER_MAX_BATCH_SIZE_EXPLICIT = envs.SGLANG_ENCODER_MAX_BATCH_SIZE.is_set()
 # Watchdog: max time to wait for a batched /encode result. Bounds HTTP latency
 # if the batch worker stalls (NCCL hang, dead worker proc, etc.).
 ENCODER_REQ_TIMEOUT = envs.SGLANG_ENCODER_REQ_TIMEOUT.get()
@@ -186,7 +190,7 @@ def _convert(data):
 
 
 _mm_grid_attrs = {
-    # Kimi K2.5 HF processor uses grid_thws (see base_processor.ATTR_NAME_TO_MODALITY).
+    # Kimi K2.5/K3 HF processors use grid_thws (see base_processor.ATTR_NAME_TO_MODALITY).
     Modality.IMAGE: ["image_grid_thw", "image_grid_hws", "grid_thws"],
     Modality.VIDEO: ["video_grid_thw"],
     Modality.AUDIO: ["audio_feature_lens_raw"],
@@ -200,14 +204,18 @@ _mm_feature_attrs = {
 
 
 def _get_mm_grid_dim(mm_inputs, modality, model_type: Optional[str] = None):
+    # Kimi K2.5/K3 vision processors only emit `grid_thws`; prefer it over generic keys
+    # so we never pick a mis-typed or stale `image_grid_hws` field from kwargs.
     attrs = _mm_grid_attrs[modality]
     model_type = (model_type or "").lower()
     if modality == Modality.IMAGE:
-        # Kimi K2.5 emits grid_thws, while Kimi-VL emits image_grid_hws.
-        if model_type == "kimi_k25":
+        # Kimi K2.5/K3 emit grid_thws, while Kimi-VL emits image_grid_hws.
+        # Other model types keep the generic attr order above.
+        if model_type in ("kimi_k25", "kimi_k3"):
             attrs = ("grid_thws", "image_grid_thw", "image_grid_hws")
         elif model_type == "kimi_vl":
             attrs = ("image_grid_hws", "image_grid_thw", "grid_thws")
+
     for attr in attrs:
         if attr in mm_inputs and mm_inputs[attr] is not None:
             return _convert(mm_inputs[attr])
@@ -249,10 +257,29 @@ def _normalize_aux_value(val):
 
 def _build_mm_aux_data(mm_inputs, model_type=None):
     # Video aux metadata, scoped to model_type's video-meta attrs.
-    return {
+    aux = {
         attr: _normalize_aux_value(mm_inputs.get(attr))
         for attr in video_meta_attrs_for(model_type)
     }
+    if model_type == "kimi_k3":
+        aux["original_image_sizes"] = _normalize_aux_value(
+            mm_inputs.get("original_image_sizes")
+        )
+    return aux
+
+
+def _get_original_image_size(image):
+    """Return an image's original (width, height) before encoder preprocessing."""
+    if isinstance(image, dict):
+        image = image.get("image")
+    if isinstance(image, torch.Tensor):
+        if image.ndim < 2:
+            raise ValueError(f"Invalid image tensor shape: {tuple(image.shape)}")
+        return [int(image.shape[-1]), int(image.shape[-2])]
+    if hasattr(image, "size"):
+        width, height = image.size
+        return [int(width), int(height)]
+    raise TypeError(f"Cannot determine original image size from {type(image)}")
 
 
 class MMEncoder:
@@ -320,9 +347,12 @@ class MMEncoder:
             load_config=self.load_config,
             device_config=self.device_config,
         )
+        maybe_precompile_model_kernels_after_loading(self.model, self.device)
 
         self.context = zmq.asyncio.Context(2)
         self.sync_context = zmq.Context()  # Reuse sync context for thread pool
+        self.scheduler_send_sockets = {}
+        self.scheduler_send_locks = {}
         self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=10)
         # Dedicated executor for image preprocessing (resize/normalize).
         # Separate from self.executor (ZMQ sends) to avoid contention under high concurrency.
@@ -799,7 +829,7 @@ class MMEncoder:
             return self._get_feat_extract_output_lengths(input_length)
         else:
             if (
-                self.model_type in ["kimi_k25", "kimi_vl"]
+                self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"]
                 and modality == Modality.IMAGE
             ):
                 return self._kimi_tokens_from_patch_grid(grid)
@@ -855,6 +885,9 @@ class MMEncoder:
         """
         if grid_thw is None:
             grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
+        split_kimi_k3_images = (
+            self.model_type == "kimi_k3" and modality == Modality.IMAGE
+        )
 
         # Audio features are per-item (list of mels for mimo_v2, or batched
         # N x n_mels x T_max for qwen2_audio); slice by item index and keep
@@ -874,31 +907,50 @@ class MMEncoder:
                 offsets.append(curr)
             for idx in indices:
                 sub_feature_list.append(mm_feature[offsets[idx] : offsets[idx + 1]])
-            sub_feature = torch.cat(sub_feature_list, dim=0)
+            if not split_kimi_k3_images:
+                sub_feature = torch.cat(sub_feature_list, dim=0)
 
-        mm_item = MultimodalDataItem.from_dict(
-            {
-                "modality": modality,
-                "feature": (
-                    sub_feature
-                    if isinstance(sub_feature, list)
-                    else _convert(sub_feature)
-                ),
-            }
-        )
+        if split_kimi_k3_images:
+            mm_items = [
+                MultimodalDataItem.from_dict(
+                    {
+                        "modality": modality,
+                        "feature": _convert(feature),
+                    }
+                )
+                for feature in sub_feature_list
+            ]
+        else:
+            mm_items = [
+                MultimodalDataItem.from_dict(
+                    {
+                        "modality": modality,
+                        "feature": (
+                            sub_feature
+                            if isinstance(sub_feature, list)
+                            else _convert(sub_feature)
+                        ),
+                    }
+                )
+            ]
 
         for k, v in mm_inputs.items():
             if k in _mm_feature_attrs.get(modality, []):
                 continue
             val = _convert(v)
             if k in _mm_grid_attrs.get(modality, []):
-                mm_item.set(k, val[indices])
+                if split_kimi_k3_images:
+                    for mm_item, idx in zip(mm_items, indices):
+                        mm_item.set(k, val[idx : idx + 1])
+                else:
+                    mm_items[0].set(k, val[indices])
             else:
-                mm_item.set(k, val)
+                for mm_item in mm_items:
+                    mm_item.set(k, val)
 
         forward_start = time.perf_counter()
         with torch.inference_mode():
-            new_embeddings = get_feature_fn([mm_item])
+            new_embeddings = get_feature_fn(mm_items)
             if not keep_on_gpu:
                 new_embeddings = new_embeddings.cpu()
             if new_embeddings.ndim != 2:
@@ -1473,12 +1525,20 @@ class MMEncoder:
     def _grid_count_per_leaf(self, leaves: List, modality: Modality) -> List[int]:
         """Number of grid entries each leaf produces under the model's processor.
 
-        Most processors map 1 leaf → 1 grid. Kimi-VL/K25 image processors expand
+        Most processors map 1 leaf → 1 grid. Kimi-VL/K2.5/K3 image processors expand
         a leaf shaped {"type": "image", "image": [pil1, pil2, ...]} into N grids
         (see _normalize_kimi_encoder_images). Cross-request batching needs these
         counts to keep per-request boundaries aligned with grid_dim.
         """
-        if self.model_type not in ("kimi_k25", "kimi_vl") or modality != Modality.IMAGE:
+        if (
+            self.model_type
+            not in (
+                "kimi_k25",
+                "kimi_k3",
+                "kimi_vl",
+            )
+            or modality != Modality.IMAGE
+        ):
             return [1] * len(leaves)
 
         def count(leaf):
@@ -1527,7 +1587,7 @@ class MMEncoder:
                     normalized.append(img)
             return normalized
 
-        # Kimi-K2.5 vision processor expects media dicts.
+        # Kimi-K2.5/K3 vision processors expect media dicts.
         normalized = []
         for img in images:
             wrapped = wrap_one(img)
@@ -1580,12 +1640,16 @@ class MMEncoder:
         if model_preprocessor:
             return model_preprocessor(images, Modality.IMAGE, self.vision_config)
         image_config = self.vision_config.get("image", {})
-        if self.model_type in ["kimi_k25", "kimi_vl"]:
+        original_image_sizes = [_get_original_image_size(item) for item in images]
+        if self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"]:
             images = self._normalize_kimi_encoder_images(images)
-        return await asyncio.get_running_loop().run_in_executor(
+        processor_input = await asyncio.get_running_loop().run_in_executor(
             self.preproc_executor,
             functools.partial(self.image_processor, images=images, **image_config),
         )
+        if self.model_type == "kimi_k3":
+            processor_input["original_image_sizes"] = original_image_sizes
+        return processor_input
 
     async def _process_video_items(self, mm_items, model_preprocessor):
         if model_preprocessor:
@@ -1893,20 +1957,79 @@ class MMEncoder:
                 serialized_data = pickle.dumps(new_mm_data)
                 buffer = embedding_tensor.__buffer__()
 
-        # Use thread pool executor for parallel ZMQ send operations
+        _zmq_xfer_start = time.perf_counter()
+        if (
+            self.server_args.encoder_transfer_backend == "zmq_to_scheduler"
+            and url is not None
+        ):
+            lock = self.scheduler_send_locks.get(endpoint)
+            if lock is None:
+                lock = asyncio.Lock()
+                self.scheduler_send_locks[endpoint] = lock
+
+            async with lock:
+                sock = self.scheduler_send_sockets.get(endpoint)
+                if sock is None:
+                    sock = self.context.socket(zmq.PUSH)
+                    config_socket(sock, zmq.PUSH)
+                    sock.setsockopt(zmq.IMMEDIATE, 1)
+                    sock.setsockopt(zmq.SNDTIMEO, int(self.send_timeout * 1000))
+                    sock.connect(endpoint)
+                    self.scheduler_send_sockets[endpoint] = sock
+                try:
+                    frames = (
+                        [serialized_data, buffer]
+                        if buffer is not None
+                        else [serialized_data]
+                    )
+                    tracker = await sock.send_multipart(frames, copy=False, track=True)
+                except Exception:
+                    if self.scheduler_send_sockets.get(endpoint) is sock:
+                        self.scheduler_send_sockets.pop(endpoint, None)
+                    sock.close(linger=0)
+                    raise
+
+            # MessageTracker.wait() protects the zero-copy source buffer; it
+            # is not a receiver acknowledgement. Waiting under the per-peer
+            # lock serialized every large embedding on that TCP connection.
+            # Queue sends in order under the lock, then wait for buffer
+            # ownership independently so libzmq can pipeline the connection.
+            try:
+                await asyncio.to_thread(tracker.wait, self.send_timeout)
+            except Exception:
+                if self.scheduler_send_sockets.get(endpoint) is sock:
+                    self.scheduler_send_sockets.pop(endpoint, None)
+                    sock.close(linger=0)
+                raise
+
+            if encoder_metrics_collector is not None:
+                encoder_metrics_collector.observe_transfer(
+                    time.perf_counter() - _zmq_xfer_start,
+                    backend=self.server_args.encoder_transfer_backend,
+                )
+            return
+
+        # Per-request sockets remain for zmq_to_tokenizer and legacy direct
+        # scheduler sends. Scheduler URL sends use persistent sockets above.
         def send_with_socket():
             sock = self.sync_context.socket(zmq.PUSH)
             config_socket(sock, zmq.PUSH)
+            sock.setsockopt(zmq.IMMEDIATE, 1)
+            sock.setsockopt(zmq.SNDTIMEO, int(self.send_timeout * 1000))
             try:
                 sock.connect(endpoint)
                 if buffer is not None:
-                    sock.send_multipart([serialized_data, buffer], copy=False)
+                    tracker = sock.send_multipart(
+                        [serialized_data, buffer], copy=False, track=True
+                    )
                 else:
-                    sock.send_multipart([serialized_data], copy=False)
+                    tracker = sock.send_multipart(
+                        [serialized_data], copy=False, track=True
+                    )
+                tracker.wait(timeout=self.send_timeout)
             finally:
                 sock.close(linger=5000)
 
-        _zmq_xfer_start = time.perf_counter()
         await asyncio.get_event_loop().run_in_executor(self.executor, send_with_socket)
         if (
             encoder_metrics_collector is not None
@@ -2077,7 +2200,7 @@ class MMEncoder:
         try:
             mm_inputs, get_feature_fn = await self._process_mm_items(mm_items, modality)
             grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
-            aux_data = _build_mm_aux_data(mm_inputs)
+            aux_data = _build_mm_aux_data(mm_inputs, self.model_type)
 
             # Setup metadata and event management
             nbytes, total_tokens, embedding_dim, event = (
@@ -2178,7 +2301,7 @@ class MMEncoder:
         """Cross-request encoder fusion (image/audio). No cache path."""
         # items_per_req counts grid entries (post-expansion) so per-request
         # slicing of grid_dim/final_slices stays aligned for processors that
-        # expand one leaf into multiple grids (e.g. Kimi-VL/K25 dict-of-images).
+        # expand one leaf into multiple grids (e.g. Kimi-VL/K2.5/K3 dict-of-images).
         flat_items, items_per_req = [], []
         for req in requests:
             leaves = MMEncoder._flatten_nested_items(req["mm_items"])
@@ -2235,14 +2358,17 @@ class MMEncoder:
             if self.profiler is not None:
                 for _ in requests:
                     self.profiler.step()
-            # No aux_data here: batch_encode only handles IMAGE/AUDIO
-            # (_BATCHABLE_MODALITIES), and _build_mm_aux_data only extracts
-            # video-meta fields — which never appear in image/audio mm_inputs.
+            aux_data = _build_mm_aux_data(mm_inputs, self.model_type)
             results = []
             offset = 0
             for req, n in zip(requests, items_per_req):
                 slices = final_slices[offset : offset + n]
                 emb = slices[0] if n == 1 else torch.cat(slices, dim=0)
+                req_aux_data = {}
+                if aux_data.get("original_image_sizes") is not None:
+                    req_aux_data["original_image_sizes"] = aux_data[
+                        "original_image_sizes"
+                    ][offset : offset + n]
                 if self.rank == 0:
                     self.embedding_to_send[req["req_id"]] = EmbeddingData(
                         req["req_id"],
@@ -2251,6 +2377,7 @@ class MMEncoder:
                         grid_dim[offset : offset + n],
                         modality,
                         emb,
+                        **req_aux_data,
                     )
                 results.append((emb.nbytes, emb.shape[0], emb.shape[1], None, None))
                 offset += n
@@ -2462,6 +2589,20 @@ class PendingRequest:
 # VIDEO excluded: per-video preprocess kwargs (do_sample_frames, video_metadata)
 # vary per request and can't merge into one HF processor call.
 _BATCHABLE_MODALITIES = {Modality.IMAGE, Modality.AUDIO}
+_KIMI_K3_DEFAULT_ENCODER_MAX_BATCH_SIZE = 2
+
+
+def _resolve_encoder_batch_policy(
+    model_type: str,
+    configured_max_batch_size: int,
+    max_batch_size_is_explicit: bool,
+) -> Tuple[int, bool]:
+    """Return effective batch size and same-turn coalescing policy."""
+    max_batch_size = max(1, int(configured_max_batch_size))
+    coalesce_same_turn = model_type == "kimi_k3"
+    if coalesce_same_turn and not max_batch_size_is_explicit:
+        max_batch_size = min(max_batch_size, _KIMI_K3_DEFAULT_ENCODER_MAX_BATCH_SIZE)
+    return max_batch_size, coalesce_same_turn
 
 
 class EncoderScheduler:
@@ -2472,11 +2613,13 @@ class EncoderScheduler:
         encoder: "MMEncoder",
         send_sockets: List[zmq.Socket],
         max_batch_size: int,
+        coalesce_same_turn: bool = False,
         request_timeout: float = ENCODER_REQ_TIMEOUT,
     ):
         self.encoder = encoder
         self.send_sockets = send_sockets
         self.max_batch_size = max(1, int(max_batch_size))
+        self.coalesce_same_turn = bool(coalesce_same_turn)
         self.request_timeout = max(1.0, float(request_timeout))
         self.pending_queue: asyncio.Queue[PendingRequest] = asyncio.Queue()
         self._worker_task: Optional[asyncio.Task] = None
@@ -2485,7 +2628,9 @@ class EncoderScheduler:
         if self._worker_task is None:
             self._worker_task = asyncio.create_task(self._batch_worker())
             logger.info(
-                f"EncoderScheduler started with max_batch_size={self.max_batch_size}"
+                "EncoderScheduler started with "
+                f"max_batch_size={self.max_batch_size}, "
+                f"coalesce_same_turn={self.coalesce_same_turn}"
             )
 
     async def stop(self) -> None:
@@ -2520,6 +2665,17 @@ class EncoderScheduler:
 
     async def _collect_batch(self) -> List[PendingRequest]:
         batch = [await self.pending_queue.get()]
+        first_modality = Modality.from_str(batch[0].request.get("modality", "image"))
+        should_yield = (
+            self.coalesce_same_turn
+            and self.max_batch_size > 1
+            and first_modality in _BATCHABLE_MODALITIES
+        )
+        if should_yield:
+            # Let HTTP handlers that arrived in the same event-loop turn enqueue
+            # before dispatch. Unlike a fixed sleep, this adds no millisecond-scale
+            # tax to an isolated request.
+            await asyncio.sleep(0)
         while len(batch) < self.max_batch_size:
             try:
                 batch.append(self.pending_queue.get_nowait())
@@ -2603,23 +2759,29 @@ class EncoderScheduler:
                 encoder_metrics_collector.observe_queue_wait(
                     max(0.0, start - p.submit_time), modality=modality_str
                 )
-        for sock in self.send_sockets:
-            sock_send(
-                sock,
-                wrap_as_pickle(
-                    {
-                        "type": "batch_encode",
-                        "modality": modality.name,
-                        "requests": requests,
-                        "enter_time": start,
-                    }
-                ),
-            )
-
-        logger.info(f"Dispatching batch of {len(group)} {modality.name} requests")
-
         try:
-            results = await self.encoder.batch_encode(requests, modality)
+            # The scheduler is the sole owner of batched dispatch order. Keep
+            # the collective broadcast and rank-0 execution under the same
+            # lock, while allowing concurrent HTTP handlers to enqueue before
+            # waiting on their individual futures.
+            async with self.encoder.encode_dispatch_lock:
+                for sock in self.send_sockets:
+                    sock_send(
+                        sock,
+                        wrap_as_pickle(
+                            {
+                                "type": "batch_encode",
+                                "modality": modality.name,
+                                "requests": requests,
+                                "enter_time": start,
+                            }
+                        ),
+                    )
+
+                logger.info(
+                    f"Dispatching batch of {len(group)} {modality.name} requests"
+                )
+                results = await self.encoder.batch_encode(requests, modality)
             if len(group) > 1:
                 logger.info(
                     f"Batch of {len(group)} {modality.name} requests completed in "
@@ -3376,8 +3538,16 @@ async def run_dp_worker(
         encoder_metrics_collector = EncoderMetricsCollector(labels)
         enc.dp_rank = dp_rank
 
+    max_batch_size, coalesce_same_turn = _resolve_encoder_batch_policy(
+        enc.model_type,
+        ENCODER_MAX_BATCH_SIZE,
+        ENCODER_MAX_BATCH_SIZE_EXPLICIT,
+    )
     sched = EncoderScheduler(
-        encoder=enc, send_sockets=[], max_batch_size=ENCODER_MAX_BATCH_SIZE
+        encoder=enc,
+        send_sockets=[],
+        max_batch_size=max_batch_size,
+        coalesce_same_turn=coalesce_same_turn,
     )
 
     ctx = zmq.asyncio.Context(2)
@@ -3386,12 +3556,12 @@ async def run_dp_worker(
     send_lock = asyncio.Lock()
     inflight: Set[asyncio.Task] = set()
     # Acquire-before-recv → back-pressure propagates to the dispatcher
-    # PUSH buffer. Must be ≥ ENCODER_MAX_BATCH_SIZE or batching degrades.
+    # PUSH buffer. Must be at least max_batch_size or batching degrades.
     max_inflight = envs.SGLANG_ENCODER_DP_WORKER_MAX_INFLIGHT.get()
-    if max_inflight < ENCODER_MAX_BATCH_SIZE:
+    if max_inflight < max_batch_size:
         logger.warning(
             f"SGLANG_ENCODER_DP_WORKER_MAX_INFLIGHT={max_inflight} is below "
-            f"ENCODER_MAX_BATCH_SIZE={ENCODER_MAX_BATCH_SIZE}; the encoder "
+            f"the effective encoder max_batch_size={max_batch_size}; the encoder "
             f"will never assemble a full batch."
         )
     inflight_sem = asyncio.Semaphore(max_inflight)
@@ -3472,8 +3642,16 @@ async def _lifespan(app: FastAPI):
         yield
         return
     if encoder is not None:
+        max_batch_size, coalesce_same_turn = _resolve_encoder_batch_policy(
+            encoder.model_type,
+            ENCODER_MAX_BATCH_SIZE,
+            ENCODER_MAX_BATCH_SIZE_EXPLICIT,
+        )
         encoder_scheduler = EncoderScheduler(
-            encoder, send_sockets, max_batch_size=ENCODER_MAX_BATCH_SIZE
+            encoder,
+            send_sockets,
+            max_batch_size=max_batch_size,
+            coalesce_same_turn=coalesce_same_turn,
         )
         encoder_scheduler.start()
     try:
@@ -3937,8 +4115,10 @@ async def handle_encode_request(request: dict):
                     },
                 )
         else:
-            # Lock direct dispatch together with rank0 await so its NCCL launch
-            # order matches the ZMQ dispatch order rank>0 sees.
+            # Non-batched requests still own their collective dispatch order
+            # directly; batched requests take this lock in _dispatch_group.
+            # Locking direct dispatch together with the rank0 await keeps its
+            # NCCL launch order matching the ZMQ dispatch order rank>0 sees.
             async with encoder.encode_dispatch_lock:
                 for socket in send_sockets:
                     sock_send(socket, wrap_as_pickle(request))
@@ -4157,12 +4337,6 @@ async def health_generate():
     if encoder is None:
         return Response(status_code=503)
 
-    # Skip the dummy encode when real requests are already in flight — the
-    # ongoing traffic already proves liveness, matching the scheduler's
-    # `is_fully_idle`-based health-check skip pattern.
-    if encoder.embedding_to_send:
-        return Response(status_code=200)
-
     # Pick the first available modality for the dummy encode
     if encoder.image_processor is not None:
         mm_items = [f"data:image/png;base64,{MINIMUM_PNG_PICTURE_BASE64}"]
@@ -4186,21 +4360,26 @@ async def health_generate():
             "part_idx": 0,
         }
 
-        # Broadcast to other TP ranks so distributed ops stay in sync
-        for socket in send_sockets:
-            sock_send(socket, wrap_as_pickle(dummy_request))
+        # A health encode participates in the same TP collectives as a real
+        # request. Serialize its broadcast and rank-0 forward with every other
+        # collective dispatch, then recheck whether traffic made the probe
+        # unnecessary while it waited for the lock.
+        async with encoder.encode_dispatch_lock:
+            if encoder.embedding_to_send:
+                return Response(status_code=200)
+            for socket in send_sockets:
+                sock_send(socket, wrap_as_pickle(dummy_request))
 
-        # Run encode on rank 0 with timeout
-        _, _, _, error_msg, _ = await asyncio.wait_for(
-            encoder.encode(
-                mm_items=mm_items,
-                modality=modality,
-                req_id=req_id,
-                num_parts=1,
-                part_idx=0,
-            ),
-            timeout=HEALTH_CHECK_TIMEOUT,
-        )
+            _, _, _, error_msg, _ = await asyncio.wait_for(
+                encoder.encode(
+                    mm_items=mm_items,
+                    modality=modality,
+                    req_id=req_id,
+                    num_parts=1,
+                    part_idx=0,
+                ),
+                timeout=HEALTH_CHECK_TIMEOUT,
+            )
 
         # Clean up stored embedding
         encoder.embedding_to_send.pop(req_id, None)

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

@@ -404,10 +404,8 @@ class MooncakeKVManager(CommonKVManager):
     def _send_chunk_ready(self, req, chunk_idx, kv_chunk, prefill_unique_rank):
         """Notify decode that a non-last staging chunk RDMA is complete."""
         na = NetworkAddress(req.endpoint, req.dst_port)
-        self._connect(
+        self._send_multipart_locked(
             na.to_tcp(),
-            is_ipv6=na.is_ipv6,
-        ).send_multipart(
             [
                 b"CHUNK_READY",
                 str(req.room).encode("ascii"),
@@ -416,7 +414,8 @@ class MooncakeKVManager(CommonKVManager):
                 str(len(kv_chunk.prefill_kv_indices)).encode("ascii"),
                 req.mooncake_session_id.encode("ascii"),
                 str(prefill_unique_rank).encode("ascii"),
-            ]
+            ],
+            is_ipv6=na.is_ipv6,
         )
 
     def _do_staging_transfer(
@@ -1091,9 +1090,8 @@ class MooncakeKVManager(CommonKVManager):
         data: bytes,
     ):
         na = NetworkAddress(remote, dst_port)
-        socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6)
-
-        socket.send_multipart(
+        self._send_multipart_locked(
+            na.to_tcp(),
             [
                 MooncakeKVManager.AUX_DATA_HEADER,
                 str(room).encode("ascii"),
@@ -1101,7 +1099,8 @@ class MooncakeKVManager(CommonKVManager):
                 str(aux_index).encode("ascii"),
                 struct.pack(">I", len(data)),
                 data,
-            ]
+            ],
+            is_ipv6=na.is_ipv6,
         )
 
     def _handle_aux_data(self, msg: List[bytes]):
@@ -1502,12 +1501,14 @@ class MooncakeKVManager(CommonKVManager):
         self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int
     ):
         na = NetworkAddress(remote, dst_port)
-        self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart(
+        self._send_multipart_locked(
+            na.to_tcp(),
             [
                 str(room).encode("ascii"),
                 str(status).encode("ascii"),
                 str(prefill_rank).encode("ascii"),
-            ]
+            ],
+            is_ipv6=na.is_ipv6,
         )
 
     def transfer_worker(
@@ -1736,12 +1737,34 @@ class MooncakeKVManager(CommonKVManager):
 
                         if kv_chunk.is_last_chunk:
                             if kv_chunk.state_indices and not skip_state:
-                                self.maybe_send_extra(
+                                state_rc = self.maybe_send_extra(
                                     req,
                                     kv_chunk.state_indices,
                                     executor,
                                     target_rank_registration_info,
                                 )
+                                if state_rc != 0:
+                                    with self.session_lock:
+                                        self.session_failures[
+                                            req.mooncake_session_id
+                                        ] += 1
+                                        self.failed_sessions.add(
+                                            req.mooncake_session_id
+                                        )
+                                    self.record_failure(
+                                        kv_chunk.room,
+                                        f"Failed to send state components of {kv_chunk.room} to "
+                                        f"{NetworkAddress(req.endpoint, req.dst_port).to_host_port_str()}",
+                                    )
+                                    self.update_status(kv_chunk.room, KVPoll.Failed)
+                                    self.sync_status_to_decode_endpoint(
+                                        req.endpoint,
+                                        req.dst_port,
+                                        req.room,
+                                        KVPoll.Failed,
+                                        prefill_unique_rank,
+                                    )
+                                    break
 
                             # Only the last chunk we need to send the aux data
                             ret = self.send_aux(
@@ -1856,11 +1879,13 @@ class MooncakeKVManager(CommonKVManager):
                     # Send ACK back to decode endpoint
                     try:
                         na = NetworkAddress(decode_ip, decode_port)
-                        self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart(
+                        self._send_multipart_locked(
+                            na.to_tcp(),
                             [
                                 b"ABORT_ACK",
                                 str(room_to_be_aborted).encode("ascii"),
-                            ]
+                            ],
+                            is_ipv6=na.is_ipv6,
                         )
                         logger.debug(
                             f"Sent ABORT_ACK for room {room_to_be_aborted} to "

python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py

@@ -58,6 +58,10 @@ _SEMAPHORE_BYTES = 128
 _MAX_GRAPH_INPUTS = 131072
 # resolved once at import time; explicit constructor sizes take precedence
 _DEFAULT_MAX_SIZE = envs.SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB.get() * 1024
+# forced per-direction workspace sizes; highest priority, override both the
+# tuned config and explicit constructor sizes (None = not forced)
+_FORCE_PULL_SIZE_KB = envs.SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB.get()
+_FORCE_PUSH_SIZE_KB = envs.SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB.get()
 
 
 class _PullMode(enum.Enum):
@@ -109,6 +113,10 @@ class CustomAllReduceV2:
                               the tuned size and ``max_size``.
         :param max_push_size: explicit per-buffer push workspace size;
                               overrides both the tuned size and ``max_size``.
+
+        ``SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB`` /
+        ``SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB`` take the highest
+        priority and override all of the size parameters above.
         """
         self.disabled = True
         if not can_use_custom_all_reduce_v2(group=group, device=device):
@@ -123,6 +131,25 @@ class CustomAllReduceV2:
             max_pull_size = min(base_config.max_pull_bytes, max_size)
         if max_push_size is None:
             max_push_size = min(base_config.max_push_bytes, max_size)
+        if _FORCE_PULL_SIZE_KB is not None:
+            max_pull_size = int(_FORCE_PULL_SIZE_KB) * 1024
+        if _FORCE_PUSH_SIZE_KB is not None:
+            max_push_size = int(_FORCE_PUSH_SIZE_KB) * 1024
+
+        def force_thresholds(heuristic):
+            # forced sizes bypass the tuned NCCL-crossover heuristics: lift
+            # each direction's ceiling to the forced workspace capacity (the
+            # clip() below only ever lowers thresholds)
+            if _FORCE_PULL_SIZE_KB is not None:
+                heuristic = heuristic._replace(two_shot_pull_threshold=max_pull_size)
+            if _FORCE_PUSH_SIZE_KB is not None:
+                heuristic = heuristic._replace(one_shot_push_threshold=max_push_size)
+            return heuristic
+
+        base_config = base_config._replace(
+            graph=force_thresholds(base_config.graph),
+            eager=force_thresholds(base_config.eager),
+        )
         # a minimal workspace keeps the Communicator valid even when a caller
         # only uses one direction (e.g. push-only fused qk-norm instances)
         self.max_pull_size = _ceil_align(max(max_pull_size, _ALIGN_BYTES), _ALIGN_BYTES)
@@ -138,8 +165,12 @@ class CustomAllReduceV2:
             max_push_bytes=self.max_push_size, max_pull_bytes=self.max_pull_size
         )._replace(num_pull_blocks=num_pull_blocks, num_push_blocks=num_push_blocks)
         self.override_algo: Optional[AllReduceAlgo] = None
-        self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
-
+        # On a multi-node (MNNVL) group the symm-mem workspace plane works
+        # across nodes, but graph zero-copy input registration (cudaIpc /
+        # node-local VMM remap) does not — force eager pull inside graphs.
+        is_multinode = not all(in_the_same_node_as(group, source_rank=0))
+        tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
+        self._is_graph_mode_supported = not tms_cudagraph and not is_multinode
         # device-side pointer table: one row of world_size pointers per
         # graph-captured all-reduce input
         self.graph_params = torch.zeros(
@@ -342,7 +373,7 @@ class CustomAllReduceV2:
             yield
             return
         try:
-            self._graph_mode_allowed = not self.tms_cudagraph
+            self._graph_mode_allowed = self._is_graph_mode_supported
             yield
         finally:
             self._graph_mode_allowed = False
@@ -426,6 +457,22 @@ def can_use_custom_all_reduce_v2(
     group: ProcessGroup,
     device: torch.device,
 ) -> bool:
+    # Multi-node (MNNVL): the node-local NVLink/P2P topology checks below
+    # are meaningless across nodes; the torch symm-mem rendezvous (fabric
+    # handles) is the real capability gate there.
+    if envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get() and not all(
+        in_the_same_node_as(group, source_rank=0)
+    ):
+        world_size = dist.get_world_size(group=group)
+        if world_size in range(2, 9):
+            logger.warning(
+                "CustomAllReduceV2 enabled on a multi-node group "
+                "(world_size=%d); graph zero-copy is disabled.",
+                world_size,
+            )
+            return True
+        return False
+
     supported = list(range(2, 17))
     if dist.get_world_size(group=group) not in supported:
         return False

python/sglang/srt/distributed/parallel_state.py

@@ -695,10 +695,22 @@ class GroupCoordinator:
             self.pymscclpp_comm is not None
             and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
         )
+        # With the MNNVL opt-in, let CustomAllReduceV2 take eligible (small)
+        # inputs ahead of the symm-mem pynccl fast path; otherwise pynccl
+        # would absorb every all-reduce whenever --enable-symm-mem is on and
+        # v2 never runs. Large inputs fail should_custom_ar and still go to
+        # the symm-mem path below.
+        _ca_takes_input = (
+            _CA_V2_MULTINODE
+            and self.ca_comm is not None
+            and not self.ca_comm.disabled
+            and self.ca_comm.should_custom_ar(input_)
+        )
         if (
             self.pynccl_comm is not None
             and self.is_symmetric_memory_enabled()
             and not should_use_pymscclpp_allreduce
+            and not _ca_takes_input
         ):
             self.debug_check_symmetric_mempool(self, {"input": input_}, "all_reduce")
             with self.pynccl_comm.change_state(enable=True):
@@ -1984,6 +1996,9 @@ logger = logging.getLogger(__name__)
 _ENABLE_CUSTOM_ALL_REDUCE = True
 _ENABLE_MSCCLPP_ALL_REDUCE = False
 _ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = False
+# Read once at import: whether CustomAllReduceV2 is opted in on a multi-node
+# (MNNVL) group. Used on the all_reduce hot path (see GroupCoordinator).
+_CA_V2_MULTINODE = envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get()
 
 
 def set_custom_all_reduce(enable: bool):

python/sglang/srt/entrypoints/engine.py

@@ -107,6 +107,7 @@ from sglang.srt.utils import (
     configure_logger,
     get_bool_env_var,
     is_cuda,
+    is_mnnvl_fabric_device,
     kill_process_tree,
     launch_dummy_health_check_server,
     maybe_reindex_device_id,
@@ -1572,6 +1573,12 @@ class Engine(EngineScoreMixin, EngineBase):
 
 def _set_envs_and_config(server_args: ServerArgs):
     # Set global environments
+    # MNNVL fabric (GB200/GB300) multi-node: cross-node NVLink needs NCCL's
+    # cuMem-based buffers and MNNVL transport. Default them on (user-set
+    # values win; the symm-mem override below only fires when unset).
+    if server_args.nnodes > 1 and is_mnnvl_fabric_device():
+        os.environ.setdefault("NCCL_CUMEM_ENABLE", "1")
+        os.environ.setdefault("NCCL_MNNVL_ENABLE", "1")
     if "NCCL_CUMEM_ENABLE" not in os.environ or server_args.enable_symm_mem:
         os.environ["NCCL_CUMEM_ENABLE"] = str(int(server_args.enable_symm_mem))
     if (

python/sglang/srt/entrypoints/openai/serving_chat.py

@@ -194,6 +194,7 @@ class OpenAIServingChat(OpenAIServingBase):
     """Handler for /v1/chat/completions requests"""
 
     _default_sampling_params_logged = False
+    _KIMI_K3_GENERATION_STUB_TOKENS = 3
 
     def __init__(
         self,
@@ -650,6 +651,14 @@ class OpenAIServingChat(OpenAIServingBase):
             content["meta_info"].get("cached_tokens", 0)
         )
 
+    def _reported_prompt_tokens(self, meta_info: Dict[str, Any]) -> int:
+        prompt_tokens = meta_info.get("prompt_tokens", 0)
+        if self.chat_encoding_spec == "kimi_k3":
+            # K3's three-token assistant generation stub is model input, but the
+            # reference API excludes it from billed/reported prompt tokens.
+            prompt_tokens = max(0, prompt_tokens - self._KIMI_K3_GENERATION_STUB_TOKENS)
+        return prompt_tokens
+
     async def _generate_stream_content(
         self,
         content: Dict[str, Any],
@@ -1493,7 +1502,9 @@ class OpenAIServingChat(OpenAIServingBase):
             ):
                 index = content.get("index", 0)
 
-                prompt_tokens[index] = content["meta_info"].get("prompt_tokens", 0)
+                prompt_tokens[index] = self._reported_prompt_tokens(
+                    content["meta_info"]
+                )
                 completion_tokens[index] = content["meta_info"].get(
                     "completion_tokens", 0
                 )
@@ -1720,6 +1731,20 @@ class OpenAIServingChat(OpenAIServingBase):
         created: int,
     ) -> Union[ChatCompletionResponse, ORJSONResponse]:
         """Build chat completion response from generation results"""
+        if self.chat_encoding_spec == "kimi_k3":
+            ret = [
+                {
+                    **item,
+                    "meta_info": {
+                        **item["meta_info"],
+                        "prompt_tokens": self._reported_prompt_tokens(
+                            item["meta_info"]
+                        ),
+                    },
+                }
+                for item in ret
+            ]
+
         choices = []
 
         # Build sglext at response level (from first ret_item, as these are per-request)
@@ -2378,7 +2403,7 @@ class OpenAIServingChat(OpenAIServingBase):
 
             # Add usage stats if continuous_usage_stats is enabled
             if continuous_usage_stats:
-                prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
+                prompt_tokens = self._reported_prompt_tokens(content["meta_info"])
                 completion_tokens = content["meta_info"].get("completion_tokens", 0)
                 reasoning_tokens = content["meta_info"].get("reasoning_tokens", 0)
                 chunk.usage = UsageProcessor.calculate_token_usage(
@@ -2431,7 +2456,7 @@ class OpenAIServingChat(OpenAIServingBase):
 
             # Add usage stats if continuous_usage_stats is enabled
             if continuous_usage_stats:
-                prompt_tokens = content["meta_info"].get("prompt_tokens", 0)
+                prompt_tokens = self._reported_prompt_tokens(content["meta_info"])
                 completion_tokens = content["meta_info"].get("completion_tokens", 0)
                 reasoning_tokens = content["meta_info"].get("reasoning_tokens", 0)
                 chunk.usage = UsageProcessor.calculate_token_usage(

python/sglang/srt/environ.py

@@ -808,6 +808,39 @@ class Envs:
     # Triton two_dot variant, 1.16-1.38x faster across GLM/DS shapes).
     SGLANG_OPT_Q8KV8_QPREP_VARIANT = EnvStr("auto")
 
+    # TRT-LLM-gen fused MoE (SiTU) via sglang JIT: path to an unpacked SiTU
+    # cubin pool (cubins + flat ABI headers + overlay/; distributed as a
+    # single downloadable archive). Needs the public flashinfer package
+    # installed for the unmodified JIT sources. Unset = feature off.
+    SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL = EnvStr(None)
+
+    # MNNVL fused all-reduce (bf16, TP8): zero-copy 1shot multicast-push for
+    # small messages and in-place NVLS 2shot on symmetric-memory tensors for
+    # large ones, with an optional fused residual add. Covers the KDA o_proj
+    # output and the latent|shared MoE reduce; everything else falls back to
+    # the regular all-reduce path. Auto-enabled on SM100/SM103 when
+    # CustomAllReduceV2 with multicast is available; set 0/1 to override in
+    # either direction. See srt/layers/k3_ar_fusion.py.
+    SGLANG_K3_AR_FUSION = EnvBool(False)
+    # K3 SP-MoE fused residual + reduce-scatter and matching all-gather over
+    # CustomAllReduceV2's MNNVL push workspace. Auto-probed for the validated
+    # TP8 GB300 configuration; set 0/1 to override. See
+    # srt/layers/k3_sp_collective.py.
+    SGLANG_K3_SP_COLLECTIVE = EnvBool(False)
+    # Keep K3's post-MoE residual stream token-sharded between consecutive
+    # SP-MoE layers. The next attention-residual aggregation and snapshot
+    # bank write run on the local shard, then only the normalized attention
+    # input is all-gathered. Requires SGLANG_K3_SP_COLLECTIVE.
+    SGLANG_K3_SP_ATTN_RES = EnvBool(False)
+    # Fused o_proj GEMM + all-reduce (bf16, TP 2..8, SM100+): one
+    # kernel computes the TP-local o_proj partial and the cross-rank sum over
+    # a P2P comm region, replacing the GEMM + NCCL AR pair at M <= 512.
+    SGLANG_K3_GEMM_AR = EnvBool(False)
+    # Merge the router gate and routed_expert_down_proj weights so the K3 MoE
+    # front reads hidden_states once, and run the top-k plus the bf16 cast in one
+    # epilogue kernel. See kernels/ops/moe/moe_front.py. Default on.
+    SGLANG_K3_FUSED_FRONT = EnvBool(True)
+
     # sgl-kernel
     SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
 
@@ -848,6 +881,16 @@ class Envs:
     # Default per-direction workspace cap for CustomAllReduceV2; explicit
     # constructor sizes take precedence over this.
     SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB = EnvInt(16 * 1024)
+    SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB = EnvInt(None)
+    SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB = EnvInt(None)
+    # Allow CustomAllReduceV2 on a process group that spans nodes (MNNVL
+    # fabric). Requires torch symmetric memory to rendezvous across nodes
+    # (fabric handles + IMEX). Graph zero-copy input registration is not
+    # supported in this mode and is disabled; all-reduce inside CUDA graphs
+    # falls back to eager pull from the symm workspace. Auto-enabled on
+    # MNNVL-fabric devices (GB200/GB300) when nnodes > 1; set 0/1 to
+    # override in either direction.
+    SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE = EnvBool(False)
     SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096)
     SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048)
     SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096)
@@ -903,6 +946,9 @@ class Envs:
     SGLANG_MM_BUFFER_SIZE_MB = EnvInt(0)
     SGLANG_MM_PRECOMPUTE_HASH = EnvBool(False)
     SGLANG_VIT_ENABLE_CUDA_GRAPH = EnvBool(False)
+    SGLANG_KIMI_K3_VIT_CUDA_GRAPH_CACHE_CAPACITY = EnvInt(2)
+    SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MIN_HITS = EnvInt(2)
+    SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MAX_SEQLEN = EnvInt(6144)
     # Use the fully-vectorized ViT position-embedding interpolation (no per-image
     # Python loop / CPU<->GPU sync). Bit-exact with the legacy implementation;
     # set False to fall back to the per-image loop.
@@ -914,7 +960,9 @@ class Envs:
 
     # VLM Item CUDA IPC Transport
     SGLANG_USE_CUDA_IPC_TRANSPORT = EnvBool(False)
-    SGLANG_USE_IPC_POOL_HANDLE_CACHE = EnvBool(False)
+    # Reuse the mapping for the already-allocated bounded CUDA IPC pool. This
+    # has no effect unless CUDA IPC feature transport is explicitly selected.
+    SGLANG_USE_IPC_POOL_HANDLE_CACHE = EnvBool(True)
     SGLANG_MM_FEATURE_CACHE_MB = EnvInt(1 * 1024)
     SGLANG_MM_ITEM_MEM_POOL_RECYCLE_INTERVAL_SEC = EnvFloat(0.05)
 
@@ -929,7 +977,6 @@ class Envs:
     # mamba pool ratio accordingly. Frees one resident slot per running request,
     # raising max_running_requests. Off = original locking + ratio (escape hatch).
     SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK = EnvBool(False)
-
     # Unified Radix Tree
     SGLANG_ENABLE_UNIFIED_RADIX_TREE = EnvBool(False)
     # Registered TreeCore backend serving the unified radix cache.

python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py

@@ -85,6 +85,7 @@ def fused_sigmoid_gating_delta_rule_update(
         dt_bias=dt_bias,
         softplus_beta=softplus_beta,
         softplus_threshold=softplus_threshold,
+        lower_bound=0.0,
         q=q,
         k=k,
         v=v,
@@ -118,6 +119,7 @@ def fused_sigmoid_gating_delta_rule_update(
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
         IS_VARLEN=cu_seqlens is not None,
         IS_KDA=is_kda,
+        USE_LOWER_BOUND=False,
         DISABLE_STATE_UPDATE=disable_state_update,
         CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None,
         HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None,

python/sglang/srt/layers/activation.py

@@ -181,6 +181,37 @@ class SiluAndMul(MultiPlatformOp):
         return self._musa_swish_glu(x)
 
 
+class SituAndMul(MultiPlatformOp):
+    """SituGLU activation used by Kimi K3.
+
+    Computes beta * tanh(gate / beta) * sigmoid(gate) * up.
+    When linear_beta is set, up is softly clipped:
+        up = linear_beta * tanh(up / linear_beta).
+    """
+
+    def __init__(self, beta: float = 1.0, linear_beta: float | None = None):
+        super().__init__()
+        self.beta = float(beta)
+        self.linear_beta = None if linear_beta is None else float(linear_beta)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        gate = x[..., :d].float()
+        up = x[..., d:].float()
+        gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate)
+        if self.linear_beta is not None:
+            up = self.linear_beta * torch.tanh(up / self.linear_beta)
+        return (gate * up).to(x.dtype)
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        from sglang.kernels.ops.kimi_k3.activation import situ_and_mul
+
+        return situ_and_mul(x, None, self.beta, self.linear_beta)
+
+    def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+
 class GeluAndMul(MultiPlatformOp):
     def __init__(self, approximate="tanh"):
         super().__init__()

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

@@ -70,6 +70,18 @@ def create_flashinfer_backend(runner):
 def create_trtllm_mla_backend(runner):
     if not runner.use_mla_backend:
         raise ValueError("trtllm_mla backend can only be used with MLA models.")
+    if (
+        runner.server_args.dcp_size > 1
+        and runner.server_args.speculative_algorithm is not None
+    ):
+        _, decode_backend = runner.server_args.get_attention_backends()
+        if decode_backend == "trtllm_mla":
+            raise ValueError(
+                "trtllm_mla cannot serve decode context parallelism with speculative "
+                "decoding: it does not forward the cyclic DCP metadata to its decode "
+                "kernel and returns no rank-local LSE for the cross-rank merge. "
+                "Select cutedsl_mla or tokenspeed_mla."
+            )
     from sglang.srt.layers.attention.trtllm_mla_backend import TRTLLMMLABackend
 
     return TRTLLMMLABackend(runner)

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

@@ -178,15 +178,18 @@ class CuteDslMLABackend(TRTLLMMLABackend):
             bs, num_tokens, forward_mode, seq_lens, device
         )
         if get_parallel().dcp_enabled:
-            if forward_mode.is_target_verify():
-                self.forward_decode_metadata.global_seq_lens_k = torch.zeros_like(
-                    self.forward_decode_metadata.seq_lens_k
-                )
-            self.forward_decode_metadata.max_seq_len_k = (
-                self._get_dcp_local_max_seq_len(
-                    self.max_context_len
-                    + (self.num_draft_tokens if forward_mode.is_target_verify() else 0)
+            metadata = self.forward_decode_metadata
+            if metadata.global_seq_lens_k is None:
+                # Plain decode under DCP also keeps the int32 GLOBAL lens in a
+                # capture-stable buffer (super allocates it only for verify):
+                # the DCP kernel consumes both the rank-local and the global
+                # lens every MLA layer, so both are maintained once per step.
+                metadata.global_seq_lens_k = torch.zeros(
+                    (bs,), dtype=torch.int32, device=device
                 )
+            metadata.max_seq_len_k = self._get_dcp_local_max_seq_len(
+                self.max_context_len
+                + (self.num_draft_tokens if forward_mode.is_target_verify() else 0)
             )
 
     def _apply_cuda_graph_metadata(
@@ -224,7 +227,13 @@ class CuteDslMLABackend(TRTLLMMLABackend):
             local_seq_lens = self._get_dcp_local_seq_lens(seq_lens)
         else:
             seq_lens = seq_lens[:bs]
-            local_seq_lens = self._get_dcp_local_seq_lens(seq_lens)
+            # Hoist: refresh the int32 global + rank-local lens once per step
+            # into the capture-stable buffers; forward_decode reads them
+            # instead of recomputing get_dcp_lens + two int32 casts per MLA
+            # layer.
+            metadata.global_seq_lens_k.copy_(seq_lens)
+            metadata.seq_lens_k.copy_(self._get_dcp_local_seq_lens(seq_lens))
+            local_seq_lens = metadata.seq_lens_k
 
         self._fill_dcp_block_kv_indices(
             metadata.block_kv_indices,
@@ -249,6 +258,19 @@ class CuteDslMLABackend(TRTLLMMLABackend):
                 metadata.seq_lens_k = self._get_dcp_local_seq_lens(
                     metadata.global_seq_lens_k
                 )
+            elif (
+                forward_batch.forward_mode.is_decode_or_idle()
+                and self.forward_decode_metadata.seq_lens_k is not None
+            ):
+                # Same hoist as verify: the parent stored the int32 GLOBAL
+                # lens in seq_lens_k; keep it as global_seq_lens_k and derive
+                # the rank-local view once per step (forward_decode consumes
+                # both every MLA layer).
+                metadata = self.forward_decode_metadata
+                metadata.global_seq_lens_k = metadata.seq_lens_k
+                metadata.seq_lens_k = self._get_dcp_local_seq_lens(
+                    metadata.global_seq_lens_k
+                )
             self.forward_decode_metadata.max_seq_len_k = (
                 self._get_dcp_local_max_seq_len(
                     self.forward_decode_metadata.max_seq_len_k
@@ -349,12 +371,30 @@ class CuteDslMLABackend(TRTLLMMLABackend):
         # Query / KV preparation mirrors the base cute-dsl decode (both FP16 and
         # FP8 KV), then swaps to the DCP kernel call + rank-local return.
         merge_query = q_rope is not None
+        query = None
         if self.data_type == torch.float8_e4m3fn:
             assert q_rope is not None and k_rope is not None
             if cos_sin_cache is None:
-                q, k, k_rope = mla_quantize_without_rope_for_fp8(
-                    q, q_rope, k.squeeze(1), k_rope.squeeze(1)
-                )
+                if (
+                    save_kv_cache
+                    and self._fused_set_kv_concat_q_fp8
+                    and not self._unified_mla
+                ):
+                    # Static pool: out_cache_loc is already the physical loc.
+                    # Fused: bf16->fp8 quantize + KV scatter + q concat in one
+                    # launch; None when not covered.
+                    query = self._set_kv_and_concat_q_fp8_fused(
+                        layer=layer,
+                        loc=forward_batch.out_cache_loc,
+                        q=q,
+                        q_rope=q_rope,
+                        k=k,
+                        k_rope=k_rope,
+                    )
+                if query is None:
+                    q, k, k_rope = mla_quantize_without_rope_for_fp8(
+                        q, q_rope, k.squeeze(1), k_rope.squeeze(1)
+                    )
             else:
                 q, k, k_rope = mla_quantize_and_rope_for_fp8(
                     q,
@@ -369,13 +409,15 @@ class CuteDslMLABackend(TRTLLMMLABackend):
                 )
             merge_query = False
 
-        if save_kv_cache:
+        if query is None and save_kv_cache:
             assert k is not None and k_rope is not None
             self.token_to_kv_pool.set_mla_kv_buffer(
                 layer, forward_batch.out_cache_loc, k, k_rope
             )
 
-        if merge_query:
+        if query is not None:
+            pass  # fused fp8 path already built the query and wrote KV
+        elif merge_query:
             q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
             q_rope_reshaped = q_rope.view(
                 -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
@@ -404,8 +446,14 @@ class CuteDslMLABackend(TRTLLMMLABackend):
             self.init_forward_metadata(forward_batch)
             metadata = forward_batch.decode_trtllm_mla_metadata
 
-        global_seq_lens = forward_batch.seq_lens[: forward_batch.batch_size]
-        local_seq_lens = self._get_dcp_local_seq_lens(global_seq_lens)
+        if metadata.seq_lens_k is not None and metadata.global_seq_lens_k is not None:
+            # Hoisted path: int32 rank-local + global lens maintained once per
+            # step by metadata init / graph replay-prep.
+            local_seq_lens = metadata.seq_lens_k[: forward_batch.batch_size]
+            global_seq_lens = metadata.global_seq_lens_k[: forward_batch.batch_size]
+        else:
+            global_seq_lens = forward_batch.seq_lens[: forward_batch.batch_size]
+            local_seq_lens = self._get_dcp_local_seq_lens(global_seq_lens)
         raw_out, lse = self._run_decode_kernel(
             query=query,
             kv_cache=kv_cache,

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

@@ -30,6 +30,9 @@ from sglang.srt.layers.dcp import (
 )
 from sglang.srt.layers.dcp.planner import plan_dcp_decode_metadata
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph 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,
 )
@@ -205,6 +208,10 @@ class FlashInferMhaChunkKVRunner:
 class FlashInferMLAAttnBackend(AttentionBackend):
     """Flashinfer attention kernels."""
 
+    # Verify metadata is ragged-layout aware via generate_attn_arg_prefill;
+    # graphs key their wrappers by token tier (_verify_graph_key).
+    supports_ragged_verify_graph: bool = True
+
     def __init__(
         self,
         model_runner: ModelRunner,
@@ -352,7 +359,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                     kv_len_arr=self.cuda_graph_kv_lens[:bs],
                     backend="auto",
                 )
-                self.prefill_cuda_graph_metadata[bs] = prefill_wrapper
+                self.prefill_cuda_graph_metadata[
+                    self._verify_graph_key(bs, spec_info)
+                ] = prefill_wrapper
                 self.forward_metadata = PrefillMetadata(prefill_wrapper, False)
             else:
                 raise ValueError(f"Invalid mode: {forward_mode=}")
@@ -404,8 +413,11 @@ class FlashInferMLAAttnBackend(AttentionBackend):
             use_ragged = (
                 not get_exec().kernel.flashinfer_mla_disable_ragged
                 and extend_no_prefix
-                # Piecewise cuda graph should use paged prefill to be compatible with prefix cache
+                # Captured prefill (tc_piecewise or breakable) must use paged
+                # prefill: it stays compatible with prefix cache, and the ragged
+                # wrapper rejects the absorbed-MLA head dims (qk=576, vo=512).
                 and not is_in_tc_piecewise_cuda_graph()
+                and not is_in_breakable_cuda_graph()
             )
 
             self.indices_updater_prefill.update(
@@ -452,6 +464,15 @@ class FlashInferMLAAttnBackend(AttentionBackend):
             "kv_indices": self.cuda_graph_kv_indices,
         }
 
+    @staticmethod
+    def _verify_graph_key(bs: int, spec_info: Optional[SpecInput]):
+        """bs for uniform graphs; token tier for ragged (tiers share slot
+        counts but each graph must replay its own recorded plan buffers)."""
+        layout = spec_info.ragged_verify_layout if spec_info is not None else None
+        if layout is None:
+            return bs
+        return ("ragged", layout.graph_num_tokens)
+
     def _apply_cuda_graph_metadata(
         self,
         bs: int,
@@ -496,7 +517,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                 seq_lens[:bs],
                 seq_lens_sum,
                 prefix_lens=None,
-                prefill_wrapper_paged=self.prefill_cuda_graph_metadata[bs],
+                prefill_wrapper_paged=self.prefill_cuda_graph_metadata[
+                    self._verify_graph_key(bs, spec_info)
+                ],
                 use_ragged=False,
                 spec_info=spec_info,
             )

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

@@ -36,14 +36,23 @@ class HybridAttnBackend(AttentionBackend):
         self.spec_attn_is_prefill = (
             model_runner.server_args.speculative_attention_mode == "prefill"
         )
-        # decide_needs_cpu_seq_lens ORs this flag across backends; without the
-        # delegation the base-class default (True) forces a per-step seq_lens
-        # D2H + host sync even when both sub-backends opted out.
-        self.needs_cpu_seq_lens = (
-            prefill_backend.needs_cpu_seq_lens or decode_backend.needs_cpu_seq_lens
+        # Gates the FutureMap's per-step seq_lens D2H (decide_needs_cpu_seq_lens
+        # ORs it across backends). Count only what runs in the spec decode loop:
+        # decode always, prefill only when mode=prefill routes verify to it --
+        # else a cpu-lens prefill backend forces the D2H on steps it never serves.
+        self.needs_cpu_seq_lens = decode_backend.needs_cpu_seq_lens or (
+            self.spec_attn_is_prefill and prefill_backend.needs_cpu_seq_lens
         )
         self.max_context_len = model_runner.model_config.context_len
 
+    @property
+    def supports_ragged_verify_graph(self) -> bool:
+        # Ragged verify is TARGET_VERIFY-only; delegate to its executor.
+        backend = (
+            self.decode_backend if self.spec_attn_is_decode else self.prefill_backend
+        )
+        return backend.supports_ragged_verify_graph
+
     def _select_backend(self, forward_mode: ForwardMode) -> AttentionBackend:
         """
         Select the appropriate attention backend based on the forward mode.
@@ -113,6 +122,19 @@ class HybridAttnBackend(AttentionBackend):
     def get_cuda_graph_seq_len_fill_value(self):
         return self.decode_backend.get_cuda_graph_seq_len_fill_value()
 
+    def init_mha_chunk_metadata(
+        self, forward_batch: ForwardBatch, disable_flashinfer_ragged: bool = False
+    ):
+        # Chunked-prefix / one-shot MHA metadata is a prefill concern. Without
+        # this delegation the MLA MHA path silently skips (re)planning its
+        # ragged wrappers when the full-attn backend is this prefill/decode
+        # split (e.g. --decode-attention-backend trtllm_mla), and any
+        # prefix-cache-hit extend batch then runs against a stale plan:
+        #   ValueError: q.shape[0] (...) does not match qo_indptr[-1] (...)
+        init = getattr(self.prefill_backend, "init_mha_chunk_metadata", None)
+        if init is not None:
+            init(forward_batch, disable_flashinfer_ragged)
+
     @property
     def verify_mask(self) -> Optional[VerifyMask]:
         return self._select_backend(ForwardMode.TARGET_VERIFY).verify_mask

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

@@ -11,6 +11,7 @@ from sglang.kernels.ops.mamba.mamba_state_indices_triton import (
 )
 from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
     scatter_mamba_states_after_mtp_verify,
+    track_mamba_states_all_layers,
     track_mamba_states_if_needed,
 )
 from sglang.srt.configs.hybrid_arch import mamba2_config
@@ -118,7 +119,7 @@ class MambaAttnBackendBase(AttentionBackend):
             # The ring cursor is a per-slot decode counter shared by all GDN layers;
             # manage it once here (snapshot, hand to layers, advance mod L), not per-layer.
             # Gate on the linear_replayssm FLAG, not on cursor-tensor presence: the
-            # spec-verify ring (--enable-gdn-replayssm-spec) shares the write_pos
+            # spec-verify ring (--enable-linear-replayssm-spec) shares the write_pos
             # allocation but owns it exclusively via commit_gdn_replayssm_spec
             # (advance-by-accept-count once per verify step). Advancing it here as
             # well inserts one phantom/stale ring entry per step and cumulatively
@@ -178,13 +179,18 @@ class MambaAttnBackendBase(AttentionBackend):
                 # skip mamba metadata.
                 query_start_loc = None
             elif forward_batch.forward_mode.is_target_verify():
-                query_start_loc = torch.arange(
-                    0,
-                    forward_batch.input_ids.shape[0] + 1,
-                    step=forward_batch.spec_info.draft_token_num,
-                    dtype=torch.int32,
-                    device=forward_batch.input_ids.device,
-                )
+                ragged_layout = forward_batch.spec_info.ragged_verify_layout
+                if ragged_layout is not None:
+                    # Compact ragged verify: variable per-request verify lens.
+                    query_start_loc = ragged_layout.qo_indptr_device
+                else:
+                    query_start_loc = torch.arange(
+                        0,
+                        forward_batch.input_ids.shape[0] + 1,
+                        step=forward_batch.spec_info.draft_token_num,
+                        dtype=torch.int32,
+                        device=forward_batch.input_ids.device,
+                    )
 
                 if self.topk > 1:
                     retrieve_next_token = forward_batch.spec_info.retrieve_next_token
@@ -371,7 +377,7 @@ class MambaAttnBackendBase(AttentionBackend):
         """True iff --enable-linear-replayssm is on for this pool.
 
         Gate on the FLAG, not on ``replayssm_write_pos is not None``: the
-        spec-verify ring (--enable-gdn-replayssm-spec) also allocates the
+        spec-verify ring (--enable-linear-replayssm-spec) also allocates the
         cursor tensor but owns it exclusively via commit_gdn_replayssm_spec.
         The decode-ring metadata machinery gated here (per-bs static cursor
         buffers, the per-replay snapshot + advance-by-one in _replay_metadata,
@@ -485,9 +491,16 @@ class MambaAttnBackendBase(AttentionBackend):
                 self.cached_cuda_graph_decode_query_start_loc[: bs + 1]
             )
         elif forward_mode.is_target_verify():
-            self.query_start_loc_list[bs - 1].copy_(
-                self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
+            ragged_layout = (
+                spec_info.ragged_verify_layout if spec_info is not None else None
             )
+            if ragged_layout is not None:
+                # Ragged capture: qsl from the runner's synthetic layout.
+                self.query_start_loc_list[bs - 1].copy_(ragged_layout.qo_indptr_device)
+            else:
+                self.query_start_loc_list[bs - 1].copy_(
+                    self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
+                )
         else:
             raise ValueError(f"Invalid forward mode: {forward_mode=}")
         mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
@@ -654,7 +667,19 @@ class MambaAttnBackendBase(AttentionBackend):
                     bs - num_padding
                 )
         elif forward_mode.is_target_verify():
-            if num_padding == 0:
+            ragged_layout = (
+                spec_info.ragged_verify_layout if spec_info is not None else None
+            )
+            if ragged_layout is not None:
+                # Mamba kernels index dense [bs, N] scratch, so they need the
+                # capped variant (see padded_to_bucket). Padding rows carry
+                # mamba slot -1 and are skipped.
+                if ragged_layout.bs != bs or ragged_layout.cap is None:
+                    ragged_layout = ragged_layout.padded_to_bucket(
+                        padded_bs=bs, cap=spec_info.draft_token_num
+                    )
+                self.query_start_loc_list[bs - 1].copy_(ragged_layout.qo_indptr_device)
+            elif num_padding == 0:
                 self.query_start_loc_list[bs - 1].copy_(
                     self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
                 )
@@ -705,26 +730,73 @@ class MambaAttnBackendBase(AttentionBackend):
     def get_cpu_graph_seq_len_fill_value(self):
         return 1
 
+    def _track_pools(self):
+        """Full [num_layers, pool_size, ...] conv/ssm pools plus the pool index
+        of the last mamba layer, for the fused all-layers track launch. None if
+        the pool shape is not the expected layout."""
+        cached = getattr(self, "_track_pools_cache", False)
+        if cached is False:
+            pools = None
+            try:
+                mamba_cache = self.req_to_token_pool.mamba_pool.mamba_cache
+                conv_pool = mamba_cache.conv[0]
+                ssm_pool = mamba_cache.temporal
+                last_pool_idx = max(self.req_to_token_pool.mamba_map.values())
+                if (
+                    conv_pool.dim() >= 3
+                    and ssm_pool.dim() >= 3
+                    and conv_pool.shape[0] == ssm_pool.shape[0] == last_pool_idx + 1
+                ):
+                    pools = (conv_pool, ssm_pool, last_pool_idx)
+            except (AttributeError, IndexError):
+                pools = None
+            self._track_pools_cache = pools
+            cached = pools
+        return cached
+
     def _track_mamba_state_decode(
         self,
         forward_batch: ForwardBatch,
         conv_states: torch.Tensor,
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,
+        layer_id: Optional[int] = None,
     ):
         """Copy decode conv/SSM states to track slots for prefix caching. Track
         dests come from the metadata (under cuda-graph: the static buffer), so the
-        InputBuffer registry slot is never mutated."""
-        if forward_batch.mamba_track_mask is not None:
-            track_mamba_states_if_needed(
-                conv_states,
-                ssm_states,
-                cache_indices,
-                forward_batch.mamba_track_mask,
-                self.forward_metadata.mamba_track_indices,
-                forward_batch.batch_size,
-                check_freed_slots=self.enable_unified_memory,
-            )
+        InputBuffer registry slot is never mutated.
+
+        With a known layer_id, the per-layer launches collapse into ONE
+        all-layers launch fired at the last mamba layer: the mask/src/dst
+        indices are shared across layers and every layer's state is final by
+        then, so the result is identical."""
+        if forward_batch.mamba_track_mask is None:
+            return
+        if layer_id is not None:
+            pools = self._track_pools()
+            if pools is not None:
+                conv_pool, ssm_pool, last_pool_idx = pools
+                if self.req_to_token_pool.mamba_map.get(layer_id) != last_pool_idx:
+                    return
+                track_mamba_states_all_layers(
+                    conv_pool,
+                    ssm_pool,
+                    cache_indices,
+                    forward_batch.mamba_track_mask,
+                    self.forward_metadata.mamba_track_indices,
+                    forward_batch.batch_size,
+                    check_freed_slots=self.enable_unified_memory,
+                )
+                return
+        track_mamba_states_if_needed(
+            conv_states,
+            ssm_states,
+            cache_indices,
+            forward_batch.mamba_track_mask,
+            self.forward_metadata.mamba_track_indices,
+            forward_batch.batch_size,
+            check_freed_slots=self.enable_unified_memory,
+        )
 
     def _track_mamba_state_extend(
         self,
@@ -893,8 +965,18 @@ class HybridLinearAttnBackend(AttentionBackend):
 
     @property
     def data_type(self):
+        # KV-cache dtype readers (e.g. the trtllm_mla fused-rope check) reach the
+        # wrapper since split backends are wrapped once (#31439); the full-attn
+        # side owns the KV cache, so its dtype is authoritative.
         return self.full_attn_backend.data_type
 
+    @property
+    def supports_ragged_verify_graph(self) -> bool:
+        return (
+            self.full_attn_backend.supports_ragged_verify_graph
+            and self.linear_attn_backend.supports_ragged_verify_graph
+        )
+
     def _is_full_attn(
         self, layer: Optional[RadixAttention], layer_id: Optional[int] = None
     ) -> bool:
@@ -1121,9 +1203,32 @@ class HybridLinearAttnBackend(AttentionBackend):
             ]
         )
 
-        mamba_caches = (
-            self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers()
-        )
+        req_pool = self.linear_attn_backend.req_to_token_pool
+        mamba_caches = req_pool.get_speculative_mamba2_params_all_layers()
+
+        # ReplaySSM-KDA: the accepted drafts live in the per-slot ring (written
+        # during verify); no intermediate_ssm is allocated. Replay the accepted
+        # prefix into `temporal` instead of scattering an intermediate state.
+        # dspark/dflash call this method directly; the generic spec_utils commit
+        # handles replayssm before reaching here (returns early), so this branch is
+        # only hit by the direct callers. Chain layout only (topk <= 1), so
+        # accept_lens == last_correct_step_indices + 1.
+        mamba_pool = req_pool.mamba_pool
+        if getattr(mamba_pool, "replayssm_is_kda", False):
+            from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import (
+                commit_kda_replayssm_after_verify,
+            )
+
+            commit_kda_replayssm_after_verify(
+                spec_state=mamba_caches,
+                state_batch_indices=state_indices_tensor,
+                accept_lens=last_correct_step_indices + 1,
+                last_correct_step_indices=last_correct_step_indices,
+                mamba_track_indices=mamba_track_indices,
+                mamba_steps_to_track=mamba_steps_to_track,
+                null_block_id=-1,
+            )
+            return
 
         scatter_mamba_states_after_mtp_verify(
             mamba_caches,

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

@@ -12,6 +12,7 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBack
 from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
 from sglang.srt.layers.attention.linear.utils import (
     LinearAttnKernelBackend,
+    build_verify_intermediate_state_indices,
     get_linear_attn_decode_backend,
     get_linear_attn_prefill_backend,
 )
@@ -346,8 +347,13 @@ class GDNAttnBackend(MambaAttnBackendBase):
         decode_backend = get_linear_attn_decode_backend()
         prefill_backend = get_linear_attn_prefill_backend()
         self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend)
-        self.verify_intermediate_state_indices = torch.arange(
-            self.req_to_token_pool.size, dtype=torch.int32, device=model_runner.device
+        # Sized past the pool for attn_tp-padded warmup/MLP-sync batches (see helper).
+        self.verify_intermediate_state_indices = (
+            build_verify_intermediate_state_indices(
+                self.req_to_token_pool.size,
+                model_runner.server_args,
+                model_runner.device,
+            )
         )
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
@@ -427,7 +433,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
                 replayssm_force_flush=replayssm_force_flush,
             )
             self._track_mamba_state_decode(
-                forward_batch, conv_states, ssm_states, cache_indices
+                forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id
             )
             return core_attn_out
 
@@ -456,7 +462,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
         )
 
         self._track_mamba_state_decode(
-            forward_batch, conv_states, ssm_states, cache_indices
+            forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id
         )
 
         return core_attn_out
@@ -633,13 +639,13 @@ class GDNAttnBackend(MambaAttnBackendBase):
                 )
             else:
                 # The recurrent fallback needs the per-draft snapshots, which
-                # the pool gates OFF under --enable-gdn-replayssm-spec (the
+                # the pool gates OFF under --enable-linear-replayssm-spec (the
                 # same flag that makes `use_replayssm_spec` true above), so
                 # this branch is unreachable with a None buffer by
                 # construction -- keep it loud rather than silently frozen.
                 assert intermediate_state_cache is not None, (
                     "recurrent target_verify fallback requires intermediate_ssm, "
-                    "which is not allocated under --enable-gdn-replayssm-spec"
+                    "which is not allocated under --enable-linear-replayssm-spec"
                 )
                 core_attn_out = self.kernel_dispatcher.target_verify(
                     A_log=layer.A_log,

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

@@ -1,7 +1,9 @@
+import importlib.util
 from typing import Optional, Tuple, Union
 
 import torch
 
+from sglang.kernels.ops.attention import kda_fused_decode
 from sglang.kernels.ops.mamba.causal_conv1d_triton import (
     causal_conv1d_fn,
     causal_conv1d_update,
@@ -10,8 +12,10 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBack
 from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
 from sglang.srt.layers.attention.linear.utils import (
     LinearAttnKernelBackend,
+    build_verify_intermediate_state_indices,
     get_linear_attn_decode_backend,
     get_linear_attn_prefill_backend,
+    get_linear_attn_verify_backend,
 )
 from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.utils import is_cpu, is_cuda, is_npu
@@ -39,7 +43,9 @@ class KDAKernelDispatcher:
         self,
         decode_backend: LinearAttnKernelBackend,
         prefill_backend: LinearAttnKernelBackend,
+        verify_backend: LinearAttnKernelBackend,
     ):
+        self.verify_backend = verify_backend
         triton_kernel = TritonKDAKernel()
 
         if decode_backend.is_triton():
@@ -68,15 +74,39 @@ class KDAKernelDispatcher:
                 "KDA supports 'triton', 'cutedsl', or 'flashinfer'."
             )
 
-        # target_verify (MTP / speculative decode) kernel: each decode backend
-        # verifies with its own kernel. FlashInfer decode uses recurrent_kda (SM100,
-        # chain only); Triton -- and CuTe DSL, which has no verify of its own -- use
-        # the Triton fused KDA verify, which handles chain + tree
-        # (retrieve_parent_token) and per-step checkpointing and is the reference the
-        # KDA backend correctness tests assert against.
-        self.verify_kernel = (
-            self.decode_kernel if decode_backend.is_flashinfer() else triton_kernel
-        )
+        # target_verify kernel, selected via --linear-attn-verify-backend (defaults
+        # to follow decode: flashinfer -> recurrent_kda, else triton).
+        #   triton: fused chain + tree (retrieve_parent_token) verify; the reference
+        #     the KDA correctness tests assert against.
+        #   flashinfer: recurrent_kda (SM100, chain only); reuses the decode kernel
+        #     when decode is also flashinfer.
+        #   nv_cutedsl: fused Kimi-K3/DSpARK dense verify.
+        if verify_backend.is_triton():
+            self.verify_kernel = triton_kernel
+        elif verify_backend.is_flashinfer():
+            if decode_backend.is_flashinfer():
+                self.verify_kernel = self.decode_kernel
+            else:
+                if not is_cuda():
+                    raise ValueError("KDA FlashInfer verify backend requires CUDA")
+                from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import (
+                    FlashInferKDAKernel,
+                )
+
+                self.verify_kernel = FlashInferKDAKernel()
+        elif verify_backend.is_nv_cutedsl():
+            self.verify_kernel = triton_kernel
+        elif verify_backend.is_custom():
+            # Future custom KDA verify kernel plugs in here.
+            raise NotImplementedError(
+                "--linear-attn-verify-backend custom: no custom KDA verify kernel "
+                "is registered yet."
+            )
+        else:
+            raise ValueError(
+                f"Unsupported KDA verify backend: {verify_backend}. "
+                "KDA verify supports 'triton', 'nv_cutedsl', or 'flashinfer'."
+            )
 
         if prefill_backend.is_triton():
             self.extend_kernel = triton_kernel
@@ -104,11 +134,40 @@ class KDAKernelDispatcher:
                 rank0_log(
                     "KDA cutedsl prefill needs SM100; falling back to Triton extend."
                 )
+        elif prefill_backend.is_ptx_kda():
+            if not is_cuda():
+                raise ValueError("PTX KDA prefill backend requires CUDA")
+            from sglang.srt.layers.attention.linear.kernels.kda_ptx import PtxKDAKernel
+
+            ptx_kda_kernel = PtxKDAKernel()
+            if ptx_kda_kernel.supports_prefill:
+                self.extend_kernel = ptx_kda_kernel
+            else:
+                self.extend_kernel = triton_kernel
+                rank0_log(
+                    "PTX KDA prefill needs SM103 (GB300); falling back to Triton "
+                    "extend."
+                )
+        elif prefill_backend.is_nvidia_kda():
+            if not is_cuda():
+                raise ValueError("NVIDIA KDA prefill backend requires CUDA")
+            from sglang.srt.layers.attention.linear.kernels.kda_nvidia import (
+                NvidiaKDAKernel,
+            )
+
+            nvidia_kda_kernel = NvidiaKDAKernel()
+            if nvidia_kda_kernel.supports_prefill:
+                self.extend_kernel = nvidia_kda_kernel
+            else:
+                self.extend_kernel = triton_kernel
+                rank0_log(
+                    "NVIDIA KDA prefill needs SM100; falling back to Triton extend."
+                )
         else:
             raise ValueError(
                 f"Unsupported KDA prefill backend: {prefill_backend}. "
-                "KDA supports 'triton', 'flashkda', or 'cutedsl' "
-                "(cutedsl prefill needs SM100)."
+                "KDA supports 'triton', 'flashkda', 'cutedsl', 'nvidia_kda', or "
+                "'ptx_kda' (cutedsl/nvidia_kda prefill need SM100, ptx_kda SM103)."
             )
 
         self.supports_packed_decode = getattr(
@@ -201,6 +260,7 @@ class KDAKernelDispatcher:
         intermediate_state_indices: torch.Tensor,
         cache_steps: int,
         retrieve_parent_token: torch.Tensor,
+        lower_bound: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         """MTP / speculative-decode verify, routed to ``self.verify_kernel``
@@ -221,6 +281,9 @@ class KDAKernelDispatcher:
             intermediate_state_indices=intermediate_state_indices,
             cache_steps=cache_steps,
             retrieve_parent_token=retrieve_parent_token,
+            lower_bound=lower_bound,
+            # Forward extras (e.g. the fused ring-write cache_ring/replayssm_*).
+            **kwargs,
         )
 
     def extend(
@@ -249,19 +312,48 @@ class KDAKernelDispatcher:
         )
 
 
+def ragged_verify_dense_scatter_indices(
+    *,
+    query_start_loc: torch.Tensor,
+    seq_len: int,
+    draft_token_num: int,
+) -> torch.Tensor:
+    """Dense [bs, draft_token_num] slot index per packed ragged-verify token.
+
+    Rows never exceed draft_token_num under either layout variant (cap for
+    graph replay, planner construction for eager -- see
+    RaggedVerifyLayout.padded_to_bucket), so in-row offsets stay in-row;
+    tokens past the layout's coverage collapse into one ghost row at index
+    bs * draft_token_num.
+    """
+    batch_size = query_start_loc.shape[0] - 1
+    token_pos = torch.arange(seq_len, device=query_start_loc.device, dtype=torch.int32)
+    token_slots = torch.searchsorted(query_start_loc[1:], token_pos, right=True)
+    return (
+        token_slots * draft_token_num
+        + (token_pos - query_start_loc[token_slots]).to(torch.int64)
+    ).clamp_(max=batch_size * draft_token_num)
+
+
 class KDAAttnBackend(MambaAttnBackendBase):
     """Attention backend for KDA (Kimi Delta Attention) linear attention."""
 
-    # Same GPU-only contract as GDNAttnBackend / Mamba2AttnBackend: KDA metadata
-    # never reads the spec-v2 seq_lens_cpu mirror (replay padding comes from
-    # forward_batch.num_padding, and the replayssm track-flush mask paths are
-    # gated `not is_kda`), so don't force FutureMap's blocking per-step
-    # seq_lens D2H (~0.5 ms/step host stall in bs=1 MTP decode).
+    # The verify kernel is varlen and the conv path scatters ragged tokens
+    # 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
+    # (graph replay already passes seq_lens_cpu=None), extend reads
+    # 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):
         super().__init__(model_runner)
-        # mamba_cache.conv is [..., kernel-1, dim] while conv_states_shape expects the window length (kernel-1) at shape[-1], hence the transpose.
+        # Needed by the extra_buffer track path: _init_track_conv_indices reads
+        # conv_states_shape[-1] as the conv window length (kernel_size - 1).
+        # The KDA pool stores conv states as [kernel-1, dim] — transposed vs
+        # Mamba2/GDN's [dim, kernel-1] — so expose the transposed shape here.
         self.conv_states_shape = (
             model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0]
             .transpose(-1, -2)
@@ -269,22 +361,35 @@ class KDAAttnBackend(MambaAttnBackendBase):
         )
         decode_backend = get_linear_attn_decode_backend()
         prefill_backend = get_linear_attn_prefill_backend()
-        # KDA FlashInfer speculative decode (target_verify) is linear-chain only --
-        # recurrent_kda has no tree-ancestor traversal. Reject EAGLE tree verify
-        # (topk > 1) early at setup instead of deep in the per-step verify call.
-        # (The kernel keeps a per-call retrieve_parent_token guard as a backstop; it
-        # also covers ngram tree, which this topk field does not.)
+        verify_backend = get_linear_attn_verify_backend()
+        # KDA FlashInfer target_verify (recurrent_kda) is chain-only (no tree-ancestor
+        # traversal). Reject EAGLE tree verify (topk > 1) early at setup, keyed on the
+        # verify backend (not decode). The kernel keeps a per-call
+        # retrieve_parent_token backstop that also covers ngram tree.
         speculative_topk = model_runner.server_args.speculative_eagle_topk or 1
-        if decode_backend.is_flashinfer() and speculative_topk > 1:
+        if verify_backend.is_flashinfer() and speculative_topk > 1:
             raise ValueError(
                 "KDA FlashInfer speculative decoding only supports topk=1 "
                 "(EAGLE tree verify / retrieve_parent_token is unsupported)."
             )
-        self.kernel_dispatcher = KDAKernelDispatcher(decode_backend, prefill_backend)
+        self.kernel_dispatcher = KDAKernelDispatcher(
+            decode_backend, prefill_backend, verify_backend
+        )
+        # One-shot; emitted at the first fused-decode interception below.
+        self._fused_override_notice = (
+            "K3 fused KDA decode engaged: --linear-attn-decode-backend "
+            f"{decode_backend} only picks the fallback kernel for shapes "
+            "the fused kernel does not cover."
+        )
         # Per-request row index into the speculative `intermediate_ssm` scratch,
-        # used by the MTP / target_verify path (mirrors GDNAttnBackend).
-        self.verify_intermediate_state_indices = torch.arange(
-            self.req_to_token_pool.size, dtype=torch.int32, device=model_runner.device
+        # used by the MTP / target_verify path (mirrors GDNAttnBackend). Sized
+        # past the pool for attn_tp-padded warmup/MLP-sync batches (see helper).
+        self.verify_intermediate_state_indices = (
+            build_verify_intermediate_state_indices(
+                self.req_to_token_pool.size,
+                model_runner.server_args,
+                model_runner.device,
+            )
         )
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
@@ -326,6 +431,77 @@ class KDAAttnBackend(MambaAttnBackendBase):
         replayssm_k = layer_cache.replayssm_k
         replayssm_g = layer_cache.replayssm_g
 
+        # Fully fused decode step: conv1d update + delta-rule recurrence +
+        # gated RMSNorm in one kernel. Engages only when the model handed off
+        # the output-norm gate for this forward (attempt-and-verify stash,
+        # see kimi_k3.py) and the shapes are covered; the model applies the
+        # norm itself whenever the stash is left unconsumed.
+        if replayssm_d is None:
+            fused_static = getattr(layer, "_k3_fused_decode_args", None)
+            onorm_gate = getattr(layer, "_k3_onorm_gate", None)
+            if (
+                fused_static is not None
+                and onorm_gate is not None
+                and mixed_qkv.shape[0] == cache_indices.shape[0]
+                and b.ndim == 3
+                and kda_fused_decode.covered(
+                    mixed_qkv,
+                    a,
+                    b[0],
+                    conv_states,
+                    ssm_states,
+                    cache_indices,
+                    onorm_gate,
+                )
+            ):
+                if self._fused_override_notice is not None:
+                    rank0_log(self._fused_override_notice)
+                    self._fused_override_notice = None
+                w_q_t, w_k_t, w_v_t, conv_bias, a_log, onorm_w, onorm_eps = fused_static
+                core_attn_out = kda_fused_decode.kda_fused_decode(
+                    mixed_qkv,
+                    a,
+                    b[0],
+                    conv_states,
+                    w_q_t,
+                    w_k_t,
+                    w_v_t,
+                    conv_bias,
+                    a_log,
+                    layer.dt_bias,
+                    onorm_gate,
+                    onorm_w,
+                    ssm_states,
+                    cache_indices,
+                    scale=layer.head_k_dim**-0.5,
+                    onorm_eps=onorm_eps,
+                    lower_bound=layer.lower_bound,
+                )
+                layer._k3_onorm_consumed = True
+                self._track_mamba_state_decode(
+                    forward_batch,
+                    conv_states,
+                    ssm_states,
+                    cache_indices,
+                    layer.layer_id,
+                )
+                return core_attn_out
+            elif fused_static is not None and onorm_gate is not None:
+                # One-shot diagnostics: the model offered the handoff but the
+                # runtime shapes were rejected (decode stays on the unfused
+                # chain, which is correct but slower).
+                if not getattr(KDAAttnBackend, "_fused_reject_logged", False):
+                    KDAAttnBackend._fused_reject_logged = True
+                    rank0_log(
+                        "KDA fused decode rejected by covered(): "
+                        f"mixed_qkv {tuple(mixed_qkv.shape)}/{mixed_qkv.dtype} "
+                        f"stride {mixed_qkv.stride()}, "
+                        f"conv_states {tuple(conv_states.shape)} "
+                        f"stride {conv_states.stride()}, "
+                        f"ssm_states {tuple(ssm_states.shape)}/{ssm_states.dtype}, "
+                        f"b {tuple(b.shape)}, indices {cache_indices.dtype}"
+                    )
+
         qkv = causal_conv1d_update(
             mixed_qkv,
             conv_states.transpose(-1, -2),
@@ -353,6 +529,7 @@ class KDAAttnBackend(MambaAttnBackendBase):
                 cache_indices=cache_indices,
                 num_v_heads=layer.num_v_heads,
                 head_v_dim=layer.head_v_dim,
+                lower_bound=layer.lower_bound,
                 replayssm_d=replayssm_d,
                 replayssm_k=replayssm_k,
                 replayssm_g=replayssm_g,
@@ -360,7 +537,7 @@ class KDAAttnBackend(MambaAttnBackendBase):
                 replayssm_force_flush=replayssm_force_flush,
             )
             self._track_mamba_state_decode(
-                forward_batch, conv_states, ssm_states, cache_indices
+                forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id
             )
             return core_attn_out
 
@@ -380,10 +557,11 @@ class KDAAttnBackend(MambaAttnBackendBase):
             ssm_states=ssm_states,
             cache_indices=cache_indices,
             query_start_loc=query_start_loc,
+            lower_bound=layer.lower_bound,
         )
 
         self._track_mamba_state_decode(
-            forward_batch, conv_states, ssm_states, cache_indices
+            forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id
         )
 
         return core_attn_out
@@ -418,6 +596,11 @@ class KDAAttnBackend(MambaAttnBackendBase):
         has_initial_state = forward_batch.extend_prefix_lens > 0
 
         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
+            # source). The KDA pool stores conv states as [kernel-1, dim], so
+            # rows of the raw [tokens, dim] pre-conv input index in directly
+            # (GDN needs a transpose here; KDA does not).
             mamba_cache_params.conv[0][
                 self.forward_metadata.conv_states_mask_indices
             ] = mixed_qkv[self.forward_metadata.track_conv_indices]
@@ -483,14 +666,23 @@ class KDAAttnBackend(MambaAttnBackendBase):
             query_start_loc=query_start_loc,
             A_log=layer.A_log,
             dt_bias=layer.dt_bias,
-            lower_bound=getattr(layer, "lower_bound", None),
+            lower_bound=layer.lower_bound,
             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.
             is_spec_decode=forward_batch.forward_mode.is_draft_extend_v2(),
             return_intermediate_states=track_ssm,
+            # Which global chunk rows of h the track snapshot will read; lets
+            # kernels that cannot materialize per-chunk states (NVIDIA KDA) take the
+            # fast path when the snapshot only needs the final state.
+            track_ssm_h_src=(
+                self.forward_metadata.track_ssm_h_src if track_ssm else None
+            ),
         )
         if track_ssm:
+            # Snapshot the SSM state at the last track-aligned chunk boundary
+            # from the kernel's per-chunk states (h) / final states into the
+            # ping-pong track slots (see _init_track_ssm_indices).
             core_attn_out, h = core_attn_out
             self._track_mamba_state_extend(
                 forward_batch, h, ssm_states, self.forward_metadata
@@ -525,7 +717,12 @@ class KDAAttnBackend(MambaAttnBackendBase):
         conv_states = mamba_cache_params.conv[0]
         ssm_states = mamba_cache_params.temporal
         intermediate_state_cache = getattr(mamba_cache_params, "intermediate_ssm", None)
-        if intermediate_state_cache is None:
+        # ReplaySSM: intermediate_ssm is intentionally None (the ring + commit-time
+        # fold replace the per-step snapshots). Pass None to the verify kernel so it
+        # skips the write (CACHE_INTERMEDIATE_STATES=False); the output is unaffected.
+        replayssm_rawv = getattr(mamba_cache_params, "replayssm_rawv", None)
+        replayssm_on = replayssm_rawv is not None
+        if intermediate_state_cache is None and not replayssm_on:
             raise RuntimeError(
                 "KDA target_verify requires a speculative mamba cache "
                 "(MambaPool.SpeculativeState); none found."
@@ -534,13 +731,66 @@ class KDAAttnBackend(MambaAttnBackendBase):
         intermediate_state_indices = self.verify_intermediate_state_indices
 
         draft_token_num = forward_batch.spec_info.draft_token_num
-        batch_size = seq_len // draft_token_num
+        ragged_layout = forward_batch.spec_info.ragged_verify_layout
+        if self._can_run_dspark_cutedsl_mtp(
+            layer=layer,
+            mixed_qkv=mixed_qkv,
+            a=a,
+            b=b,
+            draft_token_num=draft_token_num,
+            ragged_layout=ragged_layout,
+            conv_states=conv_states,
+            ssm_states=ssm_states,
+            intermediate_state_cache=intermediate_state_cache,
+            intermediate_conv_window_cache=intermediate_conv_window_cache,
+            retrieve_parent_token=retrieve_parent_token,
+            replayssm_rawv=replayssm_rawv,
+        ):
+            return self._run_dspark_cutedsl_mtp(
+                layer=layer,
+                mixed_qkv=mixed_qkv,
+                a=a,
+                b=b,
+                conv_states=conv_states,
+                ssm_states=ssm_states,
+                intermediate_state_cache=intermediate_state_cache,
+                intermediate_conv_window_cache=intermediate_conv_window_cache,
+                intermediate_state_indices=intermediate_state_indices,
+                cache_indices=cache_indices,
+                query_start_loc=query_start_loc,
+                replayssm_rawv=replayssm_rawv,
+                replayssm_rawk=mamba_cache_params.replayssm_rawk,
+                replayssm_g=mamba_cache_params.replayssm_g,
+                replayssm_beta=mamba_cache_params.replayssm_beta,
+            )
+        if ragged_layout is None:
+            batch_size = seq_len // draft_token_num
+            dense_token_indices = None
+            mixed_qkv_dense = mixed_qkv.view(batch_size, draft_token_num, -1)
+        else:
+            # Conv update and its per-step scratch want the dense
+            # [bs, draft_token_num] layout: scatter ragged tokens to their
+            # step slots, gather back after (pad steps are never committed).
+            # Uncovered tier-leftover tokens land in the ghost row (see
+            # ragged_verify_dense_scatter_indices); their values are pad
+            # garbage, discarded downstream (ghost collisions are
+            # value-irrelevant).
+            batch_size = query_start_loc.shape[0] - 1
+            num_dense_tokens = batch_size * draft_token_num
+            dense_token_indices = ragged_verify_dense_scatter_indices(
+                query_start_loc=query_start_loc,
+                seq_len=seq_len,
+                draft_token_num=draft_token_num,
+            )
+            dense = mixed_qkv.new_zeros(num_dense_tokens + 1, mixed_qkv.shape[-1])
+            dense.index_copy_(0, dense_token_indices, mixed_qkv)
+            mixed_qkv_dense = dense[:num_dense_tokens].view(
+                batch_size, draft_token_num, -1
+            )
 
         # causal_conv1d_update expects [.., dim, width]. KDA keeps dense conv-window
         # scratch because the deduplicated overlapping layout cannot be transposed.
-        mixed_qkv_reshaped = mixed_qkv.view(batch_size, draft_token_num, -1).transpose(
-            1, 2
-        )
+        mixed_qkv_reshaped = mixed_qkv_dense.transpose(1, 2)
         mixed_qkv_processed = causal_conv1d_update(
             mixed_qkv_reshaped,
             conv_states.transpose(-1, -2),
@@ -554,14 +804,41 @@ class KDAAttnBackend(MambaAttnBackendBase):
             retrieve_next_sibling=retrieve_next_sibling,
             retrieve_parent_token=retrieve_parent_token,
         )
-        mixed_qkv = mixed_qkv_processed.transpose(1, 2).reshape(seq_len, -1)
+        mixed_qkv_flat = mixed_qkv_processed.transpose(1, 2).reshape(
+            batch_size * draft_token_num, -1
+        )
+        if dense_token_indices is None:
+            mixed_qkv = mixed_qkv_flat
+        else:
+            # Ghost row (zeros) so uncovered tail tokens gather finite values.
+            padded_flat = mixed_qkv_flat.new_zeros(
+                batch_size * draft_token_num + 1, mixed_qkv_flat.shape[-1]
+            )
+            padded_flat[: batch_size * draft_token_num] = mixed_qkv_flat
+            mixed_qkv = padded_flat[dense_token_indices]
 
         q, k, v = mixed_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)
         v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0)
 
-        return self.kernel_dispatcher.target_verify(
+        # ReplaySSM: the ring-write is fused into the triton verify kernel
+        # (CACHE_RING). Ragged layouts work natively -- step_idx is the
+        # within-row step under varlen, so row i writes
+        # ring[slot][0..verify_lens[i]) and commit folds at most commit_lens
+        # of them (absorb overflow is bounded in-kernel). ring_kwargs stays
+        # empty for non-triton verify kernels, which never see replayssm.
+        ring_kwargs = {}
+        if replayssm_rawv is not None:
+            ring_kwargs = dict(
+                cache_ring=True,
+                replayssm_rawv=replayssm_rawv,
+                replayssm_rawk=mamba_cache_params.replayssm_rawk,
+                replayssm_g=mamba_cache_params.replayssm_g,
+                replayssm_beta=mamba_cache_params.replayssm_beta,
+            )
+
+        core_attn_out = self.kernel_dispatcher.target_verify(
             A_log=layer.A_log,
             dt_bias=layer.dt_bias,
             q=q,
@@ -576,4 +853,193 @@ class KDAAttnBackend(MambaAttnBackendBase):
             intermediate_state_indices=intermediate_state_indices,
             cache_steps=draft_token_num,
             retrieve_parent_token=retrieve_parent_token,
+            lower_bound=layer.lower_bound,
+            **ring_kwargs,
         )
+        if dense_token_indices is not None:
+            # Kernel output is empty-allocated and the capped qsl skips the
+            # uncovered tail rows; zero them so discarded pad hidden states
+            # stay finite. Uncovered == clamped-to-ghost.
+            covered = dense_token_indices < (batch_size * draft_token_num)
+            core_attn_out = torch.where(covered.view(1, -1, 1, 1), core_attn_out, 0.0)
+        return core_attn_out
+
+    def _can_run_dspark_cutedsl_mtp(
+        self,
+        *,
+        layer: RadixLinearAttention,
+        mixed_qkv: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        draft_token_num: int,
+        ragged_layout,
+        conv_states: torch.Tensor,
+        ssm_states: torch.Tensor,
+        intermediate_state_cache: torch.Tensor,
+        intermediate_conv_window_cache: torch.Tensor,
+        retrieve_parent_token: Optional[torch.Tensor],
+        replayssm_rawv: Optional[torch.Tensor],
+    ) -> bool:
+        """Return whether the fixed Kimi-K3/DSpARK CuTe contract is satisfied."""
+        if not self.kernel_dispatcher.verify_backend.is_nv_cutedsl() or not is_cuda():
+            return False
+        if importlib.util.find_spec("cutlass") is None:
+            return False
+        if torch.cuda.get_device_capability()[0] != 10:
+            return False
+        if ragged_layout is not None or retrieve_parent_token is not None:
+            return False
+        # draft_token_num = 1 bonus + dspark block size; the CuTe kernel is
+        # specialized per block size and capped at 8 by shared-memory growth.
+        if not 2 <= draft_token_num <= 8:
+            return False
+        if layer.bias is not None or layer.lower_bound is None:
+            return False
+        if (
+            layer.head_q_dim != 128
+            or layer.head_k_dim != 128
+            or layer.head_v_dim != 128
+        ):
+            return False
+        if (
+            layer.num_q_heads != layer.num_k_heads
+            or layer.num_k_heads != layer.num_v_heads
+        ):
+            return False
+        if (
+            mixed_qkv.dtype != torch.bfloat16
+            or a.dtype != torch.bfloat16
+            or b.dtype != torch.bfloat16
+        ):
+            return False
+        if layer.conv_weights is None or tuple(layer.conv_weights.shape) != (
+            layer.q_dim + layer.k_dim + layer.v_dim,
+            4,
+        ):
+            return False
+        if layer.conv_weights.dtype != torch.float32:
+            return False
+        if (
+            ssm_states.dtype != torch.float32
+            or tuple(ssm_states.shape[-3:])
+            != (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim)
+            or tuple(ssm_states.stride()[-3:])
+            != (
+                layer.head_v_dim * layer.head_k_dim,
+                layer.head_k_dim,
+                1,
+            )
+            or ssm_states.stride(0) % 4 != 0
+            or ssm_states.storage_offset() % 4 != 0
+        ):
+            return False
+        if conv_states.shape[-2] != 3 or intermediate_conv_window_cache.shape[-2] != 3:
+            return False
+        if intermediate_state_cache is None:
+            # ReplaySSM: the ring replaces the per-step snapshots. The kernel
+            # wrapper validates ring layout/dtypes and raises loudly (there is
+            # no recurrent fallback without intermediate_ssm).
+            return replayssm_rawv is not None
+        if intermediate_state_cache.shape[1] < draft_token_num:
+            return False
+        if tuple(intermediate_state_cache.shape[-3:]) != (
+            layer.num_v_heads,
+            layer.head_v_dim,
+            layer.head_k_dim,
+        ):
+            return False
+        return True
+
+    @staticmethod
+    def _run_dspark_cutedsl_mtp(
+        *,
+        layer: RadixLinearAttention,
+        mixed_qkv: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        conv_states: torch.Tensor,
+        ssm_states: torch.Tensor,
+        intermediate_state_cache: torch.Tensor,
+        intermediate_conv_window_cache: torch.Tensor,
+        intermediate_state_indices: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        replayssm_rawv: Optional[torch.Tensor] = None,
+        replayssm_rawk: Optional[torch.Tensor] = None,
+        replayssm_g: Optional[torch.Tensor] = None,
+        replayssm_beta: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        from sglang.kernels.ops.kimi_k3.kda_decode_mtp import (
+            fused_kda_decode_mtp_dspark,
+        )
+
+        seq_len = mixed_qkv.shape[0]
+        h = layer.num_v_heads
+        x_q, x_k, x_v = mixed_qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)
+        x_q = x_q.reshape(1, seq_len, h, layer.head_q_dim)
+        x_k = x_k.reshape(1, seq_len, h, layer.head_k_dim)
+        x_v = x_v.reshape(1, seq_len, h, layer.head_v_dim)
+        w_q, w_k, w_v = layer.conv_weights.split(
+            [layer.q_dim, layer.k_dim, layer.v_dim], dim=0
+        )
+        cs_q, cs_k, cs_v = conv_states.split(
+            [layer.q_dim, layer.k_dim, layer.v_dim], dim=-1
+        )
+        cs_q = cs_q.transpose(-1, -2)
+        cs_k = cs_k.transpose(-1, -2)
+        cs_v = cs_v.transpose(-1, -2)
+        intermediate_conv_window_cache = intermediate_conv_window_cache.transpose(
+            -1, -2
+        )
+        ic_q, ic_k, ic_v = intermediate_conv_window_cache.split(
+            [layer.q_dim, layer.k_dim, layer.v_dim], dim=-2
+        )
+        # The kernel owns all 128 output channels and can fold gated RMSNorm
+        # into the recurrence.
+        onorm_gate = getattr(layer, "_k3_onorm_gate", None)
+        fused_static = getattr(layer, "_k3_fused_decode_args", None)
+        apply_onorm = onorm_gate is not None and fused_static is not None
+        if apply_onorm:
+            onorm_weight = fused_static[5]
+            onorm_eps = fused_static[6]
+            onorm_gate = onorm_gate.reshape(1, seq_len, h, layer.head_v_dim)
+        else:
+            onorm_weight = None
+            onorm_eps = None
+            onorm_gate = None
+
+        out = fused_kda_decode_mtp_dspark(
+            x_q=x_q,
+            x_k=x_k,
+            x_v=x_v,
+            w_q=w_q,
+            w_k=w_k,
+            w_v=w_v,
+            cs_q=cs_q,
+            cs_k=cs_k,
+            cs_v=cs_v,
+            g=a,
+            beta=b,
+            A_log=layer.A_log.reshape(-1),
+            dt_bias=layer.dt_bias.reshape(-1),
+            recurrent_state=ssm_states,
+            intermediate_ssm=intermediate_state_cache,
+            intermediate_state_indices=intermediate_state_indices,
+            intermediate_conv_q=ic_q,
+            intermediate_conv_k=ic_k,
+            intermediate_conv_v=ic_v,
+            ssm_state_indices=cache_indices.to(torch.int32),
+            cu_seqlens=query_start_loc.to(torch.int32),
+            lower_bound=float(layer.lower_bound),
+            scale=layer.head_q_dim**-0.5,
+            replayssm_rawv=replayssm_rawv,
+            replayssm_rawk=replayssm_rawk,
+            replayssm_g=replayssm_g,
+            replayssm_beta=replayssm_beta,
+            onorm_gate=onorm_gate,
+            onorm_weight=onorm_weight,
+            onorm_eps=onorm_eps,
+        )
+        if apply_onorm:
+            layer._k3_onorm_consumed = True
+        return out

python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py

@@ -75,6 +75,11 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
         query_start_loc: torch.Tensor,
         **kwargs,
     ) -> torch.Tensor:
+        if kwargs.get("lower_bound") is not None:
+            raise NotImplementedError(
+                "KDA safe gate (lower_bound) is not implemented in the CuTe DSL "
+                "decode kernel; use --linear-attn-decode-backend triton."
+            )
         return cutedsl_fused_sigmoid_gating_kda_update(
             A_log=A_log,
             dt_bias=dt_bias,
@@ -108,6 +113,10 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
         **kwargs,
     ) -> torch.Tensor:
         if kwargs.get("return_intermediate_states"):
+            # The mamba radix extra_buffer track path needs per-chunk states
+            # (h), which chunk_kda_cutedsl does not expose. Refuse instead of
+            # silently skipping the snapshot (that corrupts prefix-cache
+            # restores). Use the Triton prefill backend with extra_buffer.
             raise NotImplementedError(
                 "CuteDSLKDAKernel.extend cannot return intermediate chunk "
                 "states required by mamba_radix_cache_strategy=extra_buffer; "

python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py

@@ -78,8 +78,51 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
         # Cache the constant per-(row-map, batch, T) verify scatter indices
         # (ssm_state_indices), which never change across verify calls.
         self._verify_idx_cache: dict = {}
+        # State pools whose stride layout has been validated against the
+        # recurrent_kda contract (per-layer views are pool-stable, so id() is
+        # a stable key — same lifetime argument as _gate_cache).
+        self._state_contract_ok: set = set()
         logger.info("Using FlashInfer KDA kernel")
 
+    def _check_state_stride_contract(self, ssm_states: torch.Tensor) -> None:
+        """One-time (per pool view) check that ``ssm_states`` matches the
+        layout ``recurrent_kda`` was compiled for.
+
+        The kernel's state argument is a CuTe fake tensor of shape
+        ``[N, HV, V, K]`` with stride ``(sym_int64(divisibility=16), V*K, K, 1)``
+        and ``assumed_align=32`` (flashinfer ``kda_kernels/recurrent_kda.py``):
+        the slot stride is free — which is what lets the envelope-strided pools
+        (unified memory / page-major layout, slot stride = per-slot envelope
+        pitch) be passed in and updated IN PLACE on the cu_seqlens path — but
+        the inner strides are compiled-in constants and the divisibility /
+        alignment are hard assumptions. A pool violating them would mis-address
+        state in-kernel without any error; fail loudly here instead.
+        """
+        key = id(ssm_states)
+        if key in self._state_contract_ok:
+            return
+        if ssm_states.dim() != 4:
+            raise ValueError(
+                f"recurrent_kda needs a [N, HV, V, K] state pool; got "
+                f"shape {tuple(ssm_states.shape)}"
+            )
+        _, hv, v, k = ssm_states.shape
+        if ssm_states.stride()[1:] != (v * k, k, 1):
+            raise ValueError(
+                "recurrent_kda state inner strides must be compact "
+                f"(V*K, K, 1)=({v * k}, {k}, 1); got {ssm_states.stride()[1:]} "
+                "(only the slot stride may be non-compact)"
+            )
+        base_bytes = ssm_states.storage_offset() * ssm_states.element_size()
+        if ssm_states.stride(0) % 16 != 0 or base_bytes % 32 != 0:
+            raise ValueError(
+                "recurrent_kda state pool breaks the compiled stride contract: "
+                f"slot stride {ssm_states.stride(0)} elements must be a multiple "
+                f"of 16 and the base byte offset {base_bytes} a multiple of 32 "
+                "(sym_int64(divisibility=16) / assumed_align=32)"
+            )
+        self._state_contract_ok.add(key)
+
     # ---- gate / beta normalization (shared by decode + verify) ----
 
     def _prep_gate_params(self, A_log: torch.Tensor, dt_bias: torch.Tensor):
@@ -118,6 +161,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,
         query_start_loc: torch.Tensor,
+        lower_bound: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         batch_size = cache_indices.shape[0]
@@ -126,6 +170,11 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
         num_v_heads = v.shape[2]
         head_v_dim = v.shape[3]
 
+        # The committed pool goes into the kernel as-is (in-place update); under
+        # unified memory / page-major it is an envelope-strided view, which the
+        # cu_seqlens path supports — verify the compiled contract once per pool.
+        self._check_state_stride_contract(ssm_states)
+
         # Pack each request as a length-1 sequence ([1, B, ...] + cu_seqlens) so
         # recurrent_kda indexes the committed pool IN-KERNEL via ssm_state_indices.
         # The plain [B, 1, ...] path (no cu_seqlens) instead python-gathers
@@ -144,9 +193,8 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
 
         A_log_fi, dt_bias_fi = self._prep_gate_params(A_log, dt_bias)
 
-        # Softplus gate (lower_bound=None) to match the Triton KDA decode path;
-        # in-place state update into the committed pool (no rollback for decode).
-        # query_start_loc is the decode cu_seqlens (one token per request).
+        # Gate contract matches the Triton decode path (safe gate when
+        # lower_bound set); in-place state update, no rollback for decode.
         output_fi, _ = self._recurrent_kda(
             q=query_fi,
             k=key_fi,
@@ -160,7 +208,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
             output_final_state=False,
             use_qk_l2norm_in_kernel=True,
             use_gate_in_kernel=True,
-            lower_bound=None,
+            lower_bound=lower_bound,
             cu_seqlens=query_start_loc.to(torch.int32),
             ssm_state_indices=cache_indices.to(torch.int32),
         )
@@ -186,6 +234,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
         intermediate_state_indices: torch.Tensor,
         cache_steps: int,
         retrieve_parent_token: torch.Tensor,
+        lower_bound: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         if retrieve_parent_token is not None:
@@ -272,7 +321,7 @@ class FlashInferKDAKernel(LinearAttnKernelBase):
             output_final_state=False,
             use_qk_l2norm_in_kernel=True,
             use_gate_in_kernel=True,
-            lower_bound=None,
+            lower_bound=lower_bound,
             cu_seqlens=query_start_loc.to(torch.int32),
             ssm_state_indices=ssm_state_indices,
             num_spec_tokens=num_spec_tokens,

python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py

@@ -117,6 +117,10 @@ class FlashKDAKernel(LinearAttnKernelBase):
         return_intermediate_states: bool = False,
         **kwargs,
     ) -> torch.Tensor:
+        # The fused kernel cannot expose per-chunk states (h), which the mamba
+        # radix extra_buffer track path needs; route tracked batches through
+        # the Triton chunk_kda fallback instead of silently skipping the
+        # snapshot (that would corrupt prefix-cache restores).
         if return_intermediate_states or self._should_fall_back(
             lower_bound, is_spec_decode, query_start_loc, extend_seq_lens_cpu
         ):

python/sglang/srt/layers/attention/linear/kernels/kda_nvidia.py

@@ -0,0 +1,495 @@
+# SPDX-License-Identifier: Apache-2.0
+"""NVIDIA split KDA prefill backend (K1-K4 CuTe/Triton/cuTile pipeline).
+
+Wraps the vendored ``kda_nvidia_prefill`` package through its FLA-compatible
+``chunk_kda_fwd`` interface. SGLang owns the boundary conversion from its
+V-major ``[B,H,V,K]`` cache to the vendor's K-major ``[B,H,K,V]`` state and
+back; the vendored split kernels keep their original internal layouts.
+
+The serving wrapper repacks ordinary packed prefill batches into bounded
+equal-length groups (B <= 8) and pads every sequence to one of
+2k/4k/8k/16k. This makes short multi-sequence prefill use the same NVIDIA KDA
+pipeline while bounding compile variants and transient workspaces. Everything
+that is not an ordinary state-committing prefill stays on its dedicated path:
+
+- single-sequence prefill, including chunks of one long request, where Triton
+  retains the neutral BS=1 performance;
+- track batches needing an *interior* chunk snapshot (mamba extra_buffer
+  with a non-boundary track point): the pipeline never materializes
+  per-chunk states. Boundary-aligned track batches reuse the final state;
+- spec-decode extends, which must stay rollback-able.
+
+Inputs are staged into per-bucket persistent buffers (2k/4k/8k/16k):
+stable pointers keep the vendored pipeline's pointer-keyed wrapper and
+fast-path caches hot, and bound both the workspace set and the compile
+variants. Pad rows are state-neutral (k/v/beta zero => no rank-1 update;
+g sentinel -1000 => decay exactly 1).
+
+Any unexpected kernel failure falls back to Triton for that batch with a
+loud log (attempt-and-verify, same stance as the fused decode kernel).
+"""
+
+import logging
+from typing import Optional
+
+import torch
+
+from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
+from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
+    LinearAttnKernelBase,
+)
+
+logger = logging.getLogger(__name__)
+
+_BUCKETS = (2048, 4096, 8192, 16384)
+_MAX_NVIDIA_KDA_BATCH = 8
+
+
+def _to_nvidia_kda_state_layout(
+    state: torch.Tensor, *, head_k_dim: int, head_v_dim: int
+) -> torch.Tensor:
+    """Materialize SGLang [B,H,V,K] state as vendor [B,H,K,V]."""
+    expected = (head_v_dim, head_k_dim)
+    if state.ndim != 4 or tuple(state.shape[-2:]) != expected:
+        raise ValueError(
+            "SGLang KDA state must be [B,H,V,K] with "
+            f"(V,K)={expected}, got {tuple(state.shape)}"
+        )
+    return state.transpose(-1, -2).float().contiguous()
+
+
+def _from_nvidia_kda_state_layout(
+    state: torch.Tensor,
+    *,
+    head_k_dim: int,
+    head_v_dim: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Materialize vendor [B,H,K,V] state as SGLang [B,H,V,K]."""
+    expected = (head_k_dim, head_v_dim)
+    if state.ndim != 4 or tuple(state.shape[-2:]) != expected:
+        raise ValueError(
+            "NVIDIA KDA state must be [B,H,K,V] with "
+            f"(K,V)={expected}, got {tuple(state.shape)}"
+        )
+    return state.transpose(-1, -2).to(dtype=dtype).contiguous()
+
+
+class NvidiaKDAKernel(LinearAttnKernelBase):
+    def __init__(self):
+        # This kernel uses tcgen05 + TMEM, which are available on datacenter
+        # Blackwell (SM100/SM103, reported as capability major 10), but not on
+        # Blackwell desktop SM120 even though its capability number is larger.
+        self.supports_prefill = torch.cuda.is_available() and (
+            torch.cuda.get_device_capability()[0] == 10
+        )
+        self._fwd = None
+        self._l2norm = None
+        self._triton = TritonKDAKernel()
+        # Stable detached fp32 views of the (frozen) gate params: nn.Parameters
+        # trip cute's from_dlpack (grad-tracking) and per-call views defeat
+        # the pointer-keyed wrapper caches. Keyed per source tensor -- this
+        # kernel instance is shared by every KDA layer (one dispatcher per
+        # attention backend) and each layer has its own A_log/dt_bias.
+        self._param_flat = {}
+        # One-shot engagement log per (batch size, bucket) (observability).
+        self._engaged_logged = set()
+        self._unsupported_logged = set()
+        # (batch size, bucket, shape, device) -> staging dict.
+        self._staging = {}
+
+    def _ensure_loaded(self):
+        if self._fwd is None:
+            from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
+            from sglang.kernels.ops.attention.linear.kda_nvidia_prefill import (
+                chunk_kda_fwd,
+            )
+
+            self._fwd = chunk_kda_fwd
+            self._l2norm = l2norm_fwd
+            logger.info("Using NVIDIA chunked KDA prefill (Blackwell)")
+
+    def decode(self, *args, **kwargs):
+        raise NotImplementedError("NvidiaKDAKernel is prefill-only")
+
+    def target_verify(self, *args, **kwargs):
+        raise NotImplementedError("NvidiaKDAKernel does not support target_verify")
+
+    def _triton_extend(
+        self,
+        q,
+        k,
+        v,
+        g,
+        beta,
+        ssm_states,
+        cache_indices,
+        query_start_loc,
+        A_log,
+        dt_bias,
+        lower_bound,
+        return_intermediate_states,
+        kwargs,
+    ):
+        return self._triton.extend(
+            q,
+            k,
+            v,
+            g,
+            beta,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            lower_bound=lower_bound,
+            return_intermediate_states=return_intermediate_states,
+            **kwargs,
+        )
+
+    def _flat_param(self, t: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        if t is None:
+            return None
+        key = (t.data_ptr(), t.dtype, tuple(t.shape))
+        v = self._param_flat.get(key)
+        if v is None:
+            v = t.detach().reshape(-1).float().contiguous()
+            self._param_flat[key] = v
+        return v
+
+    def _get_staging(self, batch_size, bucket, num_heads, head_k_dim, head_v_dim, dev):
+        key = (
+            batch_size,
+            bucket,
+            num_heads,
+            head_k_dim,
+            head_v_dim,
+            dev.index,
+        )
+        st = self._staging.get(key)
+        if st is None:
+            st = {
+                "q": torch.zeros(
+                    batch_size,
+                    bucket,
+                    num_heads,
+                    head_k_dim,
+                    dtype=torch.bfloat16,
+                    device=dev,
+                ),
+                "k": torch.zeros(
+                    batch_size,
+                    bucket,
+                    num_heads,
+                    head_k_dim,
+                    dtype=torch.bfloat16,
+                    device=dev,
+                ),
+                "v": torch.zeros(
+                    batch_size,
+                    bucket,
+                    num_heads,
+                    head_v_dim,
+                    dtype=torch.bfloat16,
+                    device=dev,
+                ),
+                "g": torch.full(
+                    (batch_size, bucket, num_heads, head_k_dim),
+                    -1000.0,
+                    dtype=torch.bfloat16,
+                    device=dev,
+                ),
+                "beta": torch.zeros(
+                    batch_size,
+                    bucket,
+                    num_heads,
+                    dtype=torch.bfloat16,
+                    device=dev,
+                ),
+                "s0": torch.zeros(
+                    batch_size,
+                    num_heads,
+                    head_k_dim,
+                    head_v_dim,
+                    dtype=torch.float32,
+                    device=dev,
+                ),
+            }
+            self._staging[key] = st
+        return st
+
+    def extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        *,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        A_log: Optional[torch.Tensor] = None,
+        dt_bias: Optional[torch.Tensor] = None,
+        lower_bound: Optional[float] = None,
+        return_intermediate_states: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        num_tokens = q.shape[1]
+        seq_lens_cpu = kwargs.get("extend_seq_lens_cpu")
+        track_h_src = kwargs.get("track_ssm_h_src")
+        supported_shape = (
+            q.ndim == 4
+            and k.ndim == 4
+            and v.ndim == 4
+            and q.shape[-1] == 128
+            and k.shape[-1] == 128
+            and v.shape[-1] == 128
+            and q.dtype == torch.bfloat16
+            and k.dtype == torch.bfloat16
+            and v.dtype == torch.bfloat16
+            and g.dtype == torch.bfloat16
+            and beta.dtype == torch.float32
+            and A_log is not None
+        )
+        needs_interior_snapshot = (
+            return_intermediate_states
+            and track_h_src is not None
+            and track_h_src.numel() > 0
+        )
+        eligible = (
+            not needs_interior_snapshot
+            and not kwargs.get("is_spec_decode")
+            and seq_lens_cpu is not None
+            and len(seq_lens_cpu) > 1
+            and supported_shape
+        )
+        if not eligible:
+            if (
+                seq_lens_cpu is not None
+                and len(seq_lens_cpu) > 1
+                and not kwargs.get("is_spec_decode")
+                and not needs_interior_snapshot
+                and not supported_shape
+            ):
+                unsupported_key = (
+                    q.dtype,
+                    k.dtype,
+                    v.dtype,
+                    g.dtype,
+                    beta.dtype,
+                    q.shape[-1],
+                    k.shape[-1],
+                    v.shape[-1],
+                    A_log is not None,
+                )
+                if unsupported_key not in self._unsupported_logged:
+                    self._unsupported_logged.add(unsupported_key)
+                    logger.warning(
+                        "NVIDIA KDA prefill only supports BF16 q/k/v/g, FP32 beta, "
+                        "K=V=128 with A_log; "
+                        "got q/k/v/g/beta=%s/%s/%s/%s/%s, Kq/Kk/V=%d/%d/%d, "
+                        "A_log=%s. Falling back to Triton.",
+                        q.dtype,
+                        k.dtype,
+                        v.dtype,
+                        g.dtype,
+                        beta.dtype,
+                        q.shape[-1],
+                        k.shape[-1],
+                        v.shape[-1],
+                        A_log is not None,
+                    )
+            return self._triton_extend(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                ssm_states,
+                cache_indices,
+                query_start_loc,
+                A_log,
+                dt_bias,
+                lower_bound,
+                return_intermediate_states,
+                kwargs,
+            )
+
+        self._ensure_loaded()
+
+        seq_lens = [int(length) for length in seq_lens_cpu]
+        if (
+            any(length <= 0 or length > _BUCKETS[-1] for length in seq_lens)
+            or sum(seq_lens) != num_tokens
+            or len(seq_lens) != cache_indices.numel()
+        ):
+            logger.warning(
+                "NVIDIA KDA cannot repack prefill shape T=%d seq_lens=%s; falling "
+                "back to Triton",
+                num_tokens,
+                seq_lens,
+            )
+            return self._triton_extend(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                ssm_states,
+                cache_indices,
+                query_start_loc,
+                A_log,
+                dt_bias,
+                lower_bound,
+                return_intermediate_states,
+                kwargs,
+            )
+
+        num_heads = q.shape[2]
+        head_k_dim = q.shape[-1]
+        head_v_dim = v.shape[-1]
+
+        alog_flat = self._flat_param(A_log)
+        dtb_flat = self._flat_param(dt_bias)
+
+        all_slot_indices = torch.where(
+            cache_indices >= 0, cache_indices, ssm_states.shape[0] - 1
+        ).to(torch.int64)
+        state_backup = ssm_states.index_select(0, all_slot_indices).clone()
+        packed_output = torch.empty_like(v)
+
+        try:
+            seq_start = 0
+            token_start = 0
+            while seq_start < len(seq_lens):
+                group_lens = seq_lens[seq_start : seq_start + _MAX_NVIDIA_KDA_BATCH]
+                group_size = len(group_lens)
+                group_tokens = sum(group_lens)
+                bucket = next(b for b in _BUCKETS if b >= max(group_lens))
+                engage_key = (group_size, bucket)
+                if engage_key not in self._engaged_logged:
+                    self._engaged_logged.add(engage_key)
+                    logger.info(
+                        "NVIDIA KDA prefill engaged: sequences=%d max_T=%d -> B=%d "
+                        "bucket=%d",
+                        len(seq_lens),
+                        max(group_lens),
+                        group_size,
+                        bucket,
+                    )
+
+                st = self._get_staging(
+                    group_size,
+                    bucket,
+                    num_heads,
+                    head_k_dim,
+                    head_v_dim,
+                    q.device,
+                )
+                st["q"].zero_()
+                st["k"].zero_()
+                st["v"].zero_()
+                st["g"].fill_(-1000.0)
+                st["beta"].zero_()
+
+                row_token_start = token_start
+                for row, length in enumerate(group_lens):
+                    row_token_end = row_token_start + length
+                    st["q"][row, :length].copy_(
+                        self._l2norm(q[0, row_token_start:row_token_end].contiguous())
+                    )
+                    st["k"][row, :length].copy_(
+                        self._l2norm(k[0, row_token_start:row_token_end].contiguous())
+                    )
+                    st["v"][row, :length].copy_(v[0, row_token_start:row_token_end])
+                    g_in = g[0, row_token_start:row_token_end]
+                    if g_in.dim() == 2:
+                        g_in = g_in.view(length, num_heads, head_k_dim)
+                    st["g"][row, :length].copy_(g_in)
+                    st["beta"][row, :length].copy_(
+                        beta[0, row_token_start:row_token_end]
+                    )
+                    row_token_start = row_token_end
+
+                group_slots = all_slot_indices[seq_start : seq_start + group_size]
+                st["s0"].copy_(
+                    _to_nvidia_kda_state_layout(
+                        ssm_states.index_select(0, group_slots),
+                        head_k_dim=head_k_dim,
+                        head_v_dim=head_v_dim,
+                    )
+                )
+                res = self._fwd(
+                    st["q"],
+                    st["k"],
+                    st["v"],
+                    st["g"],
+                    st["beta"],
+                    scale=head_k_dim**-0.5,
+                    initial_state=st["s0"],
+                    output_final_state=True,
+                    cu_seqlens=None,
+                    safe_gate=lower_bound is not None,
+                    lower_bound=lower_bound,
+                    use_gate_in_kernel=True,
+                    A_log=alog_flat,
+                    dt_bias=dtb_flat,
+                )
+                group_output, final_state = res[0], res[1]
+
+                row_token_start = token_start
+                for row, length in enumerate(group_lens):
+                    row_token_end = row_token_start + length
+                    packed_output[0, row_token_start:row_token_end].copy_(
+                        group_output[row, :length]
+                    )
+                    row_token_start = row_token_end
+
+                ssm_states.index_copy_(
+                    0,
+                    group_slots,
+                    _from_nvidia_kda_state_layout(
+                        final_state,
+                        head_k_dim=head_k_dim,
+                        head_v_dim=head_v_dim,
+                        dtype=ssm_states.dtype,
+                    ),
+                )
+                seq_start += group_size
+                token_start += group_tokens
+        except Exception:
+            ssm_states.index_copy_(0, all_slot_indices, state_backup)
+            logger.warning(
+                "NVIDIA KDA prefill failed (T=%d sequences=%d); restored states "
+                "and fell back to Triton for this batch",
+                num_tokens,
+                len(seq_lens),
+                exc_info=True,
+            )
+            return self._triton_extend(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                ssm_states,
+                cache_indices,
+                query_start_loc,
+                A_log,
+                dt_bias,
+                lower_bound,
+                return_intermediate_states,
+                kwargs,
+            )
+
+        if return_intermediate_states:
+            # Boundary-aligned track: the snapshot comes from the final state
+            # already scattered into the pool (track_ssm_final_src); h is
+            # never indexed (track_ssm_h_src is empty), a zero-row stand-in
+            # keeps the (out, h) contract.
+            h_empty = q.new_empty(
+                (1, 0) + tuple(ssm_states.shape[1:]), dtype=torch.float32
+            )
+            return packed_output, h_empty
+        return packed_output

python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py

@@ -0,0 +1,349 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PTX/tcgen05 KDA chunked-prefill backend (``--linear-attn-prefill-backend ptx_kda``).
+
+Wraps the vendored hand-CUDA ``kda_ptx_prefill`` kernel (GB300 / sm_103a) through
+its FLA-compatible ``chunk_kda_fwd`` interface. The kernel selects a fused
+long-sequence route or a two-launch high-head/many-sequence route at the KDA
+serving shape: K = V = 128, chunk 64.
+
+Scope: ordinary extend batches satisfying the kernel's fixed tensor contract.
+Correctness-sensitive cases stay on Triton:
+
+- track batches receive dense intermediate SSM states directly from the kernel
+  when the cache checkpoint stride is also 64 tokens. Other interior snapshots
+  stay on Triton; boundary-only tracking can still use the final state;
+- spec-decode extends, which must stay rollback-able.
+
+Single-sequence token counts that are not a multiple of the kernel's 64-token
+chunk are padded up to a bucket (1k/2k/4k/8k/16k/32k) in a persistent staging
+buffer, which bounds the resident workspace set. Pad rows are state-neutral:
+k/v/beta zero => no rank-1 update; raw gate -1000 => transformed decay of
+exactly 1. Multi-sequence batches go through the kernel's own varlen grid
+(real cu_seqlens, no padding), so their shapes are whatever the scheduler
+produces and each distinct shape can retain another workspace.
+
+The kernel caches one workspace per distinct shape for the process lifetime;
+callers should account for that resident-memory behavior when sending highly
+variable packed batches. Any unexpected kernel failure restores the state and
+falls back to Triton for that batch with a loud log.
+"""
+
+import logging
+from typing import Optional
+
+import torch
+
+from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
+from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
+    LinearAttnKernelBase,
+)
+
+logger = logging.getLogger(__name__)
+
+_CHUNK = 64
+_PAD_BUCKETS = (1024, 2048, 4096, 8192, 16384, 32768)
+# Raw-gate sentinel for pad rows: softplus(-1000) ~ 0 and
+# sigmoid(-1000) ~ 0, so either gate transform yields glog 0 => decay 1.
+_PAD_GATE = -1000.0
+
+
+class PtxKDAKernel(LinearAttnKernelBase):
+    def __init__(self):
+        # tcgen05 + TMEM with sm_103a-only encodings: GB300 (SM103) only.
+        self.supports_prefill = torch.cuda.is_available() and (
+            torch.cuda.get_device_capability() == (10, 3)
+        )
+        self._fwd = None
+        self._triton = TritonKDAKernel()
+        # Stable detached fp32 views of the (frozen) gate params, keyed per
+        # source tensor: this kernel instance is shared by every KDA layer.
+        self._param_flat = {}
+        self._unsupported_logged = False
+        # (bucket, H, K, V, device) -> staging dict for ragged token counts.
+        self._staging = {}
+
+    def _ensure_loaded(self):
+        if self._fwd is None:
+            from sglang.kernels.ops.attention.linear.kda_ptx_prefill import (
+                chunk_kda_fwd,
+                load_ext,
+            )
+
+            logger.info("Building the PTX KDA prefill extension (first use, ~1-2 min)")
+            load_ext()
+            self._fwd = chunk_kda_fwd
+            logger.info("Using PTX KDA chunked prefill (GB300 / sm_103a)")
+
+    def decode(self, *args, **kwargs):
+        raise NotImplementedError("PtxKDAKernel is prefill-only")
+
+    def target_verify(self, *args, **kwargs):
+        raise NotImplementedError("PtxKDAKernel does not support target_verify")
+
+    def _flat_param(self, t: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        if t is None:
+            return None
+        key = (t.data_ptr(), t.dtype, tuple(t.shape))
+        flat = self._param_flat.get(key)
+        if flat is None:
+            flat = t.detach().reshape(-1).float().contiguous()
+            self._param_flat[key] = flat
+        return flat
+
+    def _get_staging(self, bucket, num_heads, head_k_dim, head_v_dim, dev):
+        key = (bucket, num_heads, head_k_dim, head_v_dim, dev.index)
+        st = self._staging.get(key)
+        if st is None:
+            st = {
+                "q": torch.zeros(
+                    1, bucket, num_heads, head_k_dim, dtype=torch.bfloat16, device=dev
+                ),
+                "k": torch.zeros(
+                    1, bucket, num_heads, head_k_dim, dtype=torch.bfloat16, device=dev
+                ),
+                "v": torch.zeros(
+                    1, bucket, num_heads, head_v_dim, dtype=torch.bfloat16, device=dev
+                ),
+                "g": torch.full(
+                    (1, bucket, num_heads, head_k_dim),
+                    _PAD_GATE,
+                    dtype=torch.bfloat16,
+                    device=dev,
+                ),
+                "beta": torch.zeros(
+                    1, bucket, num_heads, dtype=torch.bfloat16, device=dev
+                ),
+            }
+            self._staging[key] = st
+        return st
+
+    def _triton_extend(
+        self,
+        q,
+        k,
+        v,
+        g,
+        beta,
+        ssm_states,
+        cache_indices,
+        query_start_loc,
+        A_log,
+        dt_bias,
+        lower_bound,
+        return_intermediate_states,
+        kwargs,
+    ):
+        return self._triton.extend(
+            q,
+            k,
+            v,
+            g,
+            beta,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            lower_bound=lower_bound,
+            return_intermediate_states=return_intermediate_states,
+            **kwargs,
+        )
+
+    def extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        *,
+        ssm_states: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        A_log: Optional[torch.Tensor] = None,
+        dt_bias: Optional[torch.Tensor] = None,
+        lower_bound: Optional[float] = None,
+        return_intermediate_states: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        num_tokens = q.shape[1]
+        seq_lens_cpu = kwargs.get("extend_seq_lens_cpu")
+        track_h_src = kwargs.get("track_ssm_h_src")
+        supported_shape = (
+            q.ndim == 4
+            and k.ndim == 4
+            and v.ndim == 4
+            and q.shape[-1] == 128
+            and k.shape[-1] == 128
+            and v.shape[-1] == 128
+            and q.dtype == torch.bfloat16
+            and k.dtype == torch.bfloat16
+            and v.dtype == torch.bfloat16
+            and g.dtype == torch.bfloat16
+            and A_log is not None
+        )
+        needs_interior_snapshot = (
+            return_intermediate_states
+            and track_h_src is not None
+            and track_h_src.numel() > 0
+        )
+        intermediate_stride_supported = True
+        if needs_interior_snapshot:
+            # The CUDA kernel hardcodes BT=64 and h[c] is the state before
+            # global chunk c. SGLang indexes h using mamba_cache_chunk_size,
+            # which may be larger (for example 256 for Nemotron-H). Returning
+            # the raw 64-token rows in that case would silently select the
+            # wrong boundary and compute wrong offsets for later sequences.
+            from sglang.srt.runtime_context import get_server_args
+
+            intermediate_stride_supported = (
+                get_server_args().mamba_cache_chunk_size == _CHUNK
+            )
+        seq_lens = (
+            [int(length) for length in seq_lens_cpu] if seq_lens_cpu is not None else []
+        )
+        shape_known = (
+            len(seq_lens) >= 1
+            and len(seq_lens) == cache_indices.numel()
+            and sum(seq_lens) == num_tokens
+            and min(seq_lens, default=0) > 0
+        )
+        # Only the single-sequence path pads (staging buffers keep the resident
+        # workspace set bounded); varlen batches go in as-is. A ragged single
+        # sequence beyond the largest bucket uses the kernel's ragged grid.
+        padded_tokens = num_tokens
+        if len(seq_lens) == 1 and num_tokens % _CHUNK != 0:
+            padded_tokens = next(
+                (bucket for bucket in _PAD_BUCKETS if bucket >= num_tokens),
+                num_tokens,
+            )
+        eligible = (
+            not kwargs.get("is_spec_decode")
+            and intermediate_stride_supported
+            and shape_known
+            and supported_shape
+        )
+        if not eligible:
+            if shape_known and not supported_shape and not self._unsupported_logged:
+                self._unsupported_logged = True
+                logger.warning(
+                    "PTX KDA prefill only supports BF16 q/k/v/g with K=V=128 and "
+                    "A_log; got q/k/v/g=%s/%s/%s/%s, Kq/Kk/V=%d/%d/%d, A_log=%s. "
+                    "Falling back to Triton.",
+                    q.dtype,
+                    k.dtype,
+                    v.dtype,
+                    g.dtype,
+                    q.shape[-1],
+                    k.shape[-1],
+                    v.shape[-1],
+                    A_log is not None,
+                )
+            return self._triton_extend(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                ssm_states,
+                cache_indices,
+                query_start_loc,
+                A_log,
+                dt_bias,
+                lower_bound,
+                return_intermediate_states,
+                kwargs,
+            )
+
+        self._ensure_loaded()
+
+        num_heads = q.shape[2]
+        head_k_dim = q.shape[-1]
+        head_v_dim = v.shape[-1]
+
+        slot = torch.where(
+            cache_indices >= 0, cache_indices, ssm_states.shape[0] - 1
+        ).to(torch.int64)
+        state_backup = ssm_states.index_select(0, slot).clone()
+
+        try:
+            qn = q
+            kn = k
+            g4 = g.reshape(1, num_tokens, num_heads, head_k_dim)
+            b3 = beta.reshape(1, num_tokens, num_heads)
+            if padded_tokens != num_tokens:
+                st = self._get_staging(
+                    padded_tokens, num_heads, head_k_dim, head_v_dim, q.device
+                )
+                for name, src in (("q", qn), ("k", kn), ("v", v), ("g", g4)):
+                    st[name][0, :num_tokens].copy_(src[0])
+                st["beta"][0, :num_tokens].copy_(b3[0])
+                st["q"][0, num_tokens:].zero_()
+                st["k"][0, num_tokens:].zero_()
+                st["v"][0, num_tokens:].zero_()
+                st["g"][0, num_tokens:].fill_(_PAD_GATE)
+                st["beta"][0, num_tokens:].zero_()
+                qn, kn, v_in, g4, b3 = st["q"], st["k"], st["v"], st["g"], st["beta"]
+            else:
+                v_in = v
+
+            # Multi-sequence packed batch, or a single sequence beyond the
+            # padding buckets -> the kernel's varlen grid. Host
+            # cu_seqlens avoids the D2H sync its piece table would otherwise do.
+            cu_kwargs = {}
+            if len(seq_lens) > 1 or padded_tokens % _CHUNK != 0:
+                cu_cpu = torch.zeros(len(seq_lens) + 1, dtype=torch.int32)
+                cu_cpu[1:] = torch.tensor(seq_lens, dtype=torch.int32).cumsum(0)
+                cu_kwargs = dict(cu_seqlens=query_start_loc, cu_seqlens_cpu=cu_cpu)
+
+            result = self._fwd(
+                qn,
+                kn,
+                v_in,
+                g4,
+                b3,
+                scale=head_k_dim**-0.5,
+                initial_state=ssm_states.index_select(0, slot),
+                output_final_state=True,
+                **cu_kwargs,
+                # SGLang keeps the recurrent state V-major ([slots, H, V, K]);
+                # the kernel is K-major, so let the wrapper transpose both ways.
+                state_v_first=True,
+                safe_gate=lower_bound is not None,
+                lower_bound=lower_bound,
+                use_gate_in_kernel=True,
+                A_log=self._flat_param(A_log),
+                dt_bias=self._flat_param(dt_bias),
+                return_intermediate_states=return_intermediate_states,
+                use_qk_l2norm_in_kernel=True,
+            )
+            out, final_state, h = result[0], result[1], result[10]
+            ssm_states.index_copy_(0, slot, final_state.to(ssm_states.dtype))
+        except Exception:
+            ssm_states.index_copy_(0, slot, state_backup)
+            logger.warning(
+                "PTX KDA prefill failed (T=%d); restored state and fell back to "
+                "Triton for this batch",
+                num_tokens,
+                exc_info=True,
+            )
+            return self._triton_extend(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                ssm_states,
+                cache_indices,
+                query_start_loc,
+                A_log,
+                dt_bias,
+                lower_bound,
+                return_intermediate_states,
+                kwargs,
+            )
+
+        packed_output = out[:, :num_tokens].contiguous()
+        if return_intermediate_states:
+            return packed_output, h
+        return packed_output

python/sglang/srt/layers/attention/linear/kernels/kda_triton.py

@@ -38,6 +38,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
         cache_indices: torch.Tensor,
         num_v_heads: int,
         head_v_dim: int,
+        lower_bound: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         """Packed decode fast path: feed the conv-1d output ``mixed_qkv``
@@ -67,6 +68,11 @@ class TritonKDAKernel(LinearAttnKernelBase):
             and replayssm_g is not None
             and replayssm_write_pos is not None
         ):
+            if lower_bound is not None:
+                raise NotImplementedError(
+                    "KDA safe gate (lower_bound) is not implemented in the "
+                    "ReplaySSM decode kernel; disable --enable-linear-replayssm."
+                )
             K = ssm_states.shape[-1]  # ssm_states: [num_slots, HV, V, K]
             fused_recurrent_linear_replayssm_decode(
                 mixed_qkv=mixed_qkv,
@@ -105,6 +111,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
             out=out,
             ssm_state_indices=cache_indices,
             use_qk_l2norm_in_kernel=True,
+            lower_bound=lower_bound,
         )
         # [B, 1, HV, V] -> [1, B, HV, V] view to match existing decode layout.
         return out.transpose(0, 1)
@@ -122,6 +129,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,
         query_start_loc: torch.Tensor,
+        lower_bound: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         return fused_sigmoid_gating_delta_rule_update(
@@ -139,6 +147,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
             softplus_beta=1.0,
             softplus_threshold=20.0,
             is_kda=True,
+            lower_bound=lower_bound,
         )
 
     def target_verify(
@@ -158,6 +167,13 @@ class TritonKDAKernel(LinearAttnKernelBase):
         intermediate_state_indices: torch.Tensor,
         cache_steps: int,
         retrieve_parent_token: torch.Tensor,
+        lower_bound: Optional[float] = None,
+        # fused ReplaySSM ring-write (dense verify only; off elsewhere).
+        cache_ring: bool = False,
+        replayssm_rawv: Optional[torch.Tensor] = None,
+        replayssm_rawk: Optional[torch.Tensor] = None,
+        replayssm_g: Optional[torch.Tensor] = None,
+        replayssm_beta: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
         # KDA MTP / speculative-decode verify via the fused KDA kernel (IS_KDA=True),
@@ -186,6 +202,12 @@ class TritonKDAKernel(LinearAttnKernelBase):
             intermediate_state_indices=intermediate_state_indices,
             cache_steps=cache_steps,
             retrieve_parent_token=retrieve_parent_token,
+            lower_bound=lower_bound,
+            cache_ring=cache_ring,
+            replayssm_rawv=replayssm_rawv,
+            replayssm_rawk=replayssm_rawk,
+            replayssm_g=replayssm_g,
+            replayssm_beta=replayssm_beta,
         )
 
     def extend(

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

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 from enum import Enum
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Dict, Optional
 
 from sglang.srt.utils.common import rank0_log
 
@@ -15,8 +15,11 @@ logger = logging.getLogger(__name__)
 class LinearAttnKernelBackend(Enum):
     TRITON = "triton"
     CUTEDSL = "cutedsl"
+    NV_CUTEDSL = "nv_cutedsl"
     FLASHINFER = "flashinfer"
     FLASHKDA = "flashkda"
+    NVIDIA_KDA = "nvidia_kda"
+    PTX_KDA = "ptx_kda"
     CUSTOM = "custom"
 
     @classmethod
@@ -29,51 +32,109 @@ class LinearAttnKernelBackend(Enum):
     def is_cutedsl(self):
         return self == LinearAttnKernelBackend.CUTEDSL
 
+    def is_nv_cutedsl(self):
+        return self == LinearAttnKernelBackend.NV_CUTEDSL
+
     def is_flashinfer(self):
         return self == LinearAttnKernelBackend.FLASHINFER
 
     def is_flashkda(self):
         return self == LinearAttnKernelBackend.FLASHKDA
 
+    def is_nvidia_kda(self):
+        return self == LinearAttnKernelBackend.NVIDIA_KDA
+
+    def is_ptx_kda(self):
+        return self == LinearAttnKernelBackend.PTX_KDA
+
     def is_custom(self):
         return self == LinearAttnKernelBackend.CUSTOM
 
 
-LINEAR_ATTN_DECODE_BACKEND: Optional[LinearAttnKernelBackend] = None
-LINEAR_ATTN_PREFILL_BACKEND: Optional[LinearAttnKernelBackend] = None
+_BACKENDS: Dict[str, Optional[LinearAttnKernelBackend]] = {
+    "decode": None,
+    "prefill": None,
+    "verify": None,
+}
 
 
 def initialize_linear_attn_config(
     server_args: ServerArgs, prefill_default: Optional[str] = None
 ):
-    global LINEAR_ATTN_DECODE_BACKEND
-    global LINEAR_ATTN_PREFILL_BACKEND
-
     base = server_args.linear_attn_backend
     decode = server_args.linear_attn_decode_backend or base
     prefill = server_args.linear_attn_prefill_backend or prefill_default or base
 
-    LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend(decode)
-    LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend(prefill)
+    _BACKENDS["decode"] = LinearAttnKernelBackend(decode)
+    _BACKENDS["prefill"] = LinearAttnKernelBackend(prefill)
 
-    rank0_log(f"Linear attention kernel backend: decode={decode}, prefill={prefill}")
+    # Verify backend. Unset -> follow decode (flashinfer -> its recurrent kernel,
+    # else triton), preserving historical behavior.
+    verify = server_args.linear_attn_verify_backend
+    if verify is None:
+        verify = decode if _BACKENDS["decode"].is_flashinfer() else "triton"
+    _BACKENDS["verify"] = LinearAttnKernelBackend(verify)
+
+    rank0_log(
+        f"Linear attention kernel backend: decode={decode}, prefill={prefill}, "
+        f"verify={verify}"
+    )
+
+
+def _get_backend(phase: str) -> LinearAttnKernelBackend:
+    backend = _BACKENDS[phase]
+    if backend is None:
+        logger.warning(
+            "linear-attn %s backend is not initialized, using triton backend", phase
+        )
+        backend = _BACKENDS[phase] = LinearAttnKernelBackend.TRITON
+    return backend
 
 
 def get_linear_attn_decode_backend() -> LinearAttnKernelBackend:
-    global LINEAR_ATTN_DECODE_BACKEND
-    if LINEAR_ATTN_DECODE_BACKEND is None:
-        logger.warning(
-            "LINEAR_ATTN_DECODE_BACKEND is not initialized, using triton backend"
-        )
-        LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend.TRITON
-    return LINEAR_ATTN_DECODE_BACKEND
+    return _get_backend("decode")
 
 
 def get_linear_attn_prefill_backend() -> LinearAttnKernelBackend:
-    global LINEAR_ATTN_PREFILL_BACKEND
-    if LINEAR_ATTN_PREFILL_BACKEND is None:
-        logger.warning(
-            "LINEAR_ATTN_PREFILL_BACKEND is not initialized, using triton backend"
+    return _get_backend("prefill")
+
+
+def get_linear_attn_verify_backend() -> LinearAttnKernelBackend:
+    return _get_backend("verify")
+
+
+def build_verify_intermediate_state_indices(
+    pool_size: int, server_args: ServerArgs, device
+):
+    """Per-request row index into the speculative intermediate scratch
+    (`intermediate_ssm` / `intermediate_conv_window`) for the MTP /
+    target_verify path: request slot i owns scratch row i.
+
+    The scratch is allocated with one extra padding row (the `+1` in
+    MambaPool.SpeculativeState, index `pool_size`). Warmup and MLP-sync
+    batches can be padded past the pool capacity — under DP attention
+    `get_eager_max_batch_size` ceil-aligns the eager warmup bs to attn_tp —
+    and the verify kernels index this table positionally up to that padded
+    bs. Size the table to the padded maximum and clamp every out-of-pool row
+    onto the padding row: pad rows race onto one discard row, which is
+    value-irrelevant (same convention as the ragged-verify ghost row).
+    """
+    import torch
+
+    from sglang.srt.utils.common import get_eager_max_batch_size
+
+    padded_bs = max(get_eager_max_batch_size(server_args, pool_size), pool_size)
+    indices = torch.arange(pool_size, dtype=torch.int32, device=device)
+    if padded_bs > pool_size:
+        indices = torch.cat(
+            [
+                indices,
+                torch.full(
+                    (padded_bs - pool_size,),
+                    pool_size,
+                    dtype=torch.int32,
+                    device=device,
+                ),
+            ]
         )
-        LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend.TRITON
-    return LINEAR_ATTN_PREFILL_BACKEND
+    return indices

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

@@ -55,6 +55,7 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
     TRTLLMMLAMultiStepDraftBackend,
 )
 from sglang.srt.layers.dcp.layout import get_dcp_lens
+from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import is_flashinfer_available, is_tokenspeed_mla_available
 
@@ -75,7 +76,8 @@ logger = logging.getLogger(__name__)
 
 # Workspace upper bound for tokenspeed_mla_decode:
 #   num_sms * num_heads * max_q_len * (kv_lora_rank + 1) * sizeof(float32)
-# MAX_Q_LEN=8 covers EAGLE3 num_draft_tokens=4 plus headroom.
+# MAX_Q_LEN=8 covers EAGLE3 num_draft_tokens=4 plus headroom. Larger
+# speculative widths grow the bound at backend initialization.
 _TOKENSPEED_MAX_Q_LEN = 8
 
 
@@ -401,7 +403,10 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
     ):
         if not get_parallel().dcp_enabled:
             return super()._apply_cuda_graph_metadata(
-                bs, req_pool_indices, seq_lens, forward_mode
+                bs,
+                req_pool_indices,
+                seq_lens,
+                forward_mode,
             )
 
         metadata = self.decode_cuda_graph_metadata[bs]
@@ -511,6 +516,25 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
         llama_4_scaling: Optional[torch.Tensor] = None,
     ):
         parallel = get_parallel()
+        # FlashInfer autotunes MoE kernels with a synthetic full-model decode
+        # and discards the attention/logits result.  On multi-node GB300, the
+        # synthetic full-head DCP metadata can make both the TokenSpeed and
+        # TRTLLM decode kernels surface cudaErrorNvlinkUncorrectable.  Skip
+        # attention only inside that explicitly scoped dummy pass.  Real
+        # requests and CUDA graph capture continue through TokenSpeed below.
+        if parallel.dcp_enabled and get_in_autotune_dummy_run():
+            output = torch.zeros(
+                (q.shape[0], layer.tp_q_head_num * layer.v_head_dim),
+                dtype=self.q_data_type,
+                device=q.device,
+            )
+            lse = torch.zeros(
+                (q.shape[0], layer.tp_q_head_num),
+                dtype=torch.float32,
+                device=q.device,
+            )
+            return output, lse
+
         if not parallel.dcp_enabled:
             return super().forward_decode(
                 q,
@@ -615,7 +639,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
         # (prepare_prefill_qkv / pack_prefix_chunk_kv); no quantize here.
         # Hybrid MLA models resolve the model-side hook through the outer
         # HybridLinearAttnBackend, so their fallback MHA path can pass V as a
-        # last-dimension slice of kv_b_proj (stride(-2) > size(-1)). The
+        # last-dimension slice of kv_b_proj (stride(-2) > size(-1)).  The
         # TokenSpeed prefill kernel requires dense Q/K/V layouts even though
         # the public wrapper accepts arbitrary torch tensors.
         q = q.contiguous()

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

@@ -465,11 +465,19 @@ class TritonAttnBackend(AttentionBackend):
         spec_info,
     ):
         """Fill all cuda-graph buffers for target_verify mode."""
+        # Prefer the spec_info's per-request query length (DSpark draft propose
+        # uses gamma < verify window); fall back to the configured verify window.
+        num_draft_tokens = self.num_draft_tokens
+        if (
+            spec_info is not None
+            and getattr(spec_info, "draft_token_num", None) is not None
+        ):
+            num_draft_tokens = int(spec_info.draft_token_num)
         qo_indptr = self.qo_indptr[: bs + 1]
         qo_indptr[: bs + 1] = torch.arange(
             0,
-            (1 + bs) * self.num_draft_tokens,
-            step=self.num_draft_tokens,
+            (1 + bs) * num_draft_tokens,
+            step=num_draft_tokens,
             dtype=torch.int32,
             device=self.device,
         )
@@ -506,7 +514,7 @@ class TritonAttnBackend(AttentionBackend):
             custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
         else:
             custom_mask = None
-        seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
+        seq_mask_len = num_draft_tokens * (seq_lens + num_draft_tokens)
         mask_indptr = self.mask_indptr[: bs + 1]
         mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
         return (
@@ -799,10 +807,18 @@ class TritonAttnBackend(AttentionBackend):
             max_extend_len = None
         elif forward_batch.forward_mode.is_target_verify():
             bs = len(forward_batch.req_pool_indices)
+            # self.num_draft_tokens is the verify window (gamma + 1), while
+            # DSpark draft propose runs a gamma-token TARGET_VERIFY forward.
+            num_draft_tokens = self.num_draft_tokens
+            if (
+                spec_info is not None
+                and getattr(spec_info, "draft_token_num", None) is not None
+            ):
+                num_draft_tokens = int(spec_info.draft_token_num)
             qo_indptr = torch.arange(
                 0,
-                (1 + bs) * self.num_draft_tokens,
-                step=self.num_draft_tokens,
+                (1 + bs) * num_draft_tokens,
+                step=num_draft_tokens,
                 dtype=torch.int32,
                 device=self.device,
             )
@@ -839,13 +855,13 @@ class TritonAttnBackend(AttentionBackend):
                 )
 
             custom_mask = spec_info.custom_mask
-            seq_mask_len = self.num_draft_tokens * (
-                forward_batch.seq_lens + self.num_draft_tokens
+            seq_mask_len = num_draft_tokens * (
+                forward_batch.seq_lens + num_draft_tokens
             )
             mask_indptr = self.mask_indptr
             mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0)
             mask_indptr = mask_indptr[: bs + 1]
-            max_extend_len = self.num_draft_tokens
+            max_extend_len = num_draft_tokens
             num_kv_splits = None
             attn_logits = None
             attn_lse = None
@@ -1093,10 +1109,16 @@ class TritonAttnBackend(AttentionBackend):
                 and getattr(spec_info, "custom_mask", None) is not None
                 else None
             )
+            max_extend_len = self.num_draft_tokens
+            if (
+                spec_info is not None
+                and getattr(spec_info, "draft_token_num", None) is not None
+            ):
+                max_extend_len = int(spec_info.draft_token_num)
             return ForwardMetadata(
                 attn_logits=None,
                 attn_lse=None,
-                max_extend_len=self.num_draft_tokens,
+                max_extend_len=max_extend_len,
                 num_kv_splits=None,
                 kv_indptr=self.kv_indptr[: bs + 1],
                 kv_indices=self.cuda_graph_kv_indices,

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

@@ -19,6 +19,20 @@ from sglang.kernels.ops.attention.pad import (
 from sglang.kernels.ops.attention.pad import (
     unpad_draft_extend_output as unpad_draft_extend_output_triton,
 )
+from sglang.kernels.ops.attention.set_mla_kv_concat_q import (
+    can_use_set_mla_kv_concat_q,
+    can_use_set_mla_kv_concat_q_fp8,
+)
+from sglang.kernels.ops.attention.set_mla_kv_concat_q import (
+    covered as set_mla_kv_concat_q_covered,
+)
+from sglang.kernels.ops.attention.set_mla_kv_concat_q import (
+    covered_fp8 as set_mla_kv_concat_q_fp8_covered,
+)
+from sglang.kernels.ops.attention.set_mla_kv_concat_q import (
+    set_mla_kv_concat_q,
+    set_mla_kv_concat_q_fp8,
+)
 from sglang.kernels.ops.attention.utils import (
     concat_mla_absorb_q_general,
     mla_quantize_and_rope_for_fp8,
@@ -146,6 +160,13 @@ class TRTLLMMLAPrefillMetadata:
     fallback_to_flashinfer_impl: bool = False
 
 
+from sglang.kernels.jit.utils import is_arch_support_pdl
+
+# Arm PDL on the trtllm-gen decode launch so its prolog overlaps the tail of
+# the query-prep kernels (which already trigger their PDL secondary).
+_ENABLE_PDL = is_arch_support_pdl()
+
+
 @dataclass
 class TRTLLMMLADecodeMetadata:
     """Metadata for TRTLLM MLA decode operations."""
@@ -167,6 +188,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
     # read seq_lens_cpu / seq_lens_sum; opt out of the D2H sync.
     needs_cpu_seq_lens: bool = False
 
+    # Ragged verify: the packed query is front-aligned into the dense
+    # [bs, draft_token_num] layout in forward_extend; metadata stays uniform.
+    supports_ragged_verify_graph: bool = True
+
     def __init__(
         self,
         model_runner: ModelRunner,
@@ -266,6 +291,27 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         # write translates through the pool's _full_translate hook instead).
         self._decode_dense_loc: Optional[torch.Tensor] = None
         self.cuda_graph_out_cache_loc_dense: Optional[torch.Tensor] = None
+        # Fused KV-scatter + q-concat on the decode dense-loc path (one launch
+        # instead of set_mla_kv_buffer + concat_mla_absorb_q). Disabled under
+        # async asserts: the fused path writes the pool directly and would
+        # skip the pool's OOB probe.
+        self._fused_set_kv_concat_q = (
+            self.data_type == torch.bfloat16
+            and not envs.SGLANG_ENABLE_ASYNC_ASSERT.get()
+            and can_use_set_mla_kv_concat_q(
+                self.kv_lora_rank * 2, self.qk_rope_head_dim * 2
+            )
+        )
+        # fp8 sibling: quantize + KV scatter + q concat in one launch
+        # (replaces mla_quantize_without_rope_for_fp8's concat + three aten
+        # casts plus the KV-row write on the fp8 decode path).
+        self._fused_set_kv_concat_q_fp8 = (
+            self.data_type == torch.float8_e4m3fn
+            and not envs.SGLANG_ENABLE_ASYNC_ASSERT.get()
+            and self.kv_lora_rank == 512
+            and self.qk_rope_head_dim == 64
+            and can_use_set_mla_kv_concat_q_fp8()
+        )
 
     def _calc_padded_blocks(self, max_seq_len: int) -> int:
         """
@@ -415,6 +461,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         device: torch.device,
     ):
         """Allocate persistent metadata buffers for CUDA graph capture."""
+        if forward_mode.is_target_verify() and bs in self.decode_cuda_graph_metadata:
+            # Token tiers at the same slot count must share one per-bs buffer
+            # set (each graph bakes in the tensors it captured).
+            self.forward_decode_metadata = self.decode_cuda_graph_metadata[bs]
+            return
         metadata = TRTLLMMLADecodeMetadata()
 
         if forward_mode.is_target_verify():
@@ -438,6 +489,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             )
             metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device)
 
+        if metadata.seq_lens_k is None:
+            # Plain decode: static int32 seq_lens buffer, refreshed by the
+            # capture+replay body below (same pattern as target-verify).
+            metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device)
+
         # Capture with full width so future longer sequences are safe during replay.
         max_blocks_per_seq = self._calc_padded_blocks(self.max_context_len)
         block_kv_indices = self.decode_cuda_graph_kv_indices[:bs, :max_blocks_per_seq]
@@ -476,6 +532,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             metadata.sum_seq_lens_q = num_tokens_per_req * bs
             seq_lens = seq_lens[:bs]
             metadata.seq_lens_k.copy_(seq_lens)
+        elif metadata.seq_lens_k is not None:
+            # Plain decode: int64 -> int32 downcast copy into the static
+            # buffer (once per step, replacing the per-layer conversion).
+            metadata.seq_lens_k.copy_(seq_lens[:bs])
 
         # Update block indices for new sequences.
         create_flashmla_kv_indices_triton[
@@ -643,6 +703,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 self.forward_decode_metadata.global_seq_lens_k = (
                     self.forward_decode_metadata.seq_lens_k
                 )
+            elif forward_batch.forward_mode.is_decode_or_idle():
+                # One int32 conversion per step; forward_decode reads it back
+                # so the per-layer .to(int32) in _run_decode_kernel stays a
+                # no-op (24 elementwise copies/step otherwise).
+                self.forward_decode_metadata.seq_lens_k = seq_lens.to(torch.int32)
             elif forward_batch.forward_mode.is_draft_extend_v2():
                 sum_seq_lens_q = sum(forward_batch.extend_seq_lens_cpu)
                 max_seq_len_q = max(forward_batch.extend_seq_lens_cpu)
@@ -699,14 +764,32 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         cu_seqlens_q: torch.Tensor,
         seq_lens_q: torch.Tensor,
         sum_seq_lens_q: int,
+        zero_uncovered: bool = False,
     ) -> torch.Tensor:
-        """Unpad draft extended output using Triton kernel."""
+        """Unpad draft extended output using Triton kernel.
+
+        zero_uncovered: ragged verify's clamped rows leave output positions
+        unwritten; zero the destination so those discarded rows stay finite
+        (draft_extend writes every position and does not need this).
+        """
+        output_buffer = self.unpad_output_buffer
+        if zero_uncovered:
+            if output_buffer is not None:
+                output_buffer[:sum_seq_lens_q].zero_()
+            else:
+                # No persistent buffer without cuda graph state; the triton
+                # wrapper's dynamic fallback is torch.empty.
+                output_buffer = torch.zeros(
+                    (sum_seq_lens_q, raw_out.shape[2], raw_out.shape[3]),
+                    dtype=raw_out.dtype,
+                    device=raw_out.device,
+                )
         return unpad_draft_extend_output_triton(
             raw_out,
             cu_seqlens_q,
             seq_lens_q,
             sum_seq_lens_q,
-            self.unpad_output_buffer,
+            output_buffer,
         )
 
     def _compute_decode_bmm1_scale(self, layer: RadixAttention) -> float:
@@ -738,8 +821,25 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         seq_lens: torch.Tensor,
         max_seq_len: int,
         layer: RadixAttention,
+        *,
+        causal_seqs: Optional[torch.Tensor] = None,
+        cp_world: int = 1,
+        cp_rank: int = 0,
+        return_lse: bool = False,
     ) -> torch.Tensor:
-        """Hook for subclasses to swap the decode/spec-verify kernel."""
+        """Hook for subclasses to swap the decode/spec-verify kernel.
+
+        The DCP arguments belong to the hook contract because forward_extend
+        passes them on the DCP target-verify path. This implementation does not
+        forward them to the kernel and returns no LSE, so only the DCP-capable
+        subclasses serve them."""
+        if cp_world > 1 or return_lse:
+            raise NotImplementedError(
+                "trtllm_mla does not forward the cyclic DCP metadata to its "
+                "decode kernel and returns no rank-local LSE for the cross-rank "
+                "merge; select cutedsl_mla or tokenspeed_mla for a DCP "
+                "target-verify run"
+            )
 
         # Scale computation for TRTLLM MLA kernel BMM1 operation:
         # The final BMM1 scale is computed as: q_scale * k_scale * softmax_scale
@@ -761,6 +861,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         return flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
             query=query,
             kv_cache=kv_cache,
+            enable_pdl=_ENABLE_PDL,
             workspace_buffer=self.workspace_buffer,
             qk_nope_head_dim=self.qk_nope_head_dim,
             kv_lora_rank=self.kv_lora_rank,
@@ -818,6 +919,101 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(),
         )
 
+    def _set_kv_and_concat_q_fused(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        k: torch.Tensor,
+        k_rope: torch.Tensor,
+        q: torch.Tensor,
+        q_rope: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        """Decode: scatter the KV row at ``loc`` (already physical — the
+        dense-loc buffer on the unified pool, or out_cache_loc on the static
+        pool where ``_full_translate`` is identity) and build the
+        [q_nope | q_rope] fmha query in one kernel launch (saves one launch
+        per MLA layer and keeps the PDL chain intact).
+
+        Returns the concatenated query, or None when the fused kernel does
+        not cover the inputs (caller falls back to the two-kernel path).
+        """
+        k_nope_2d = k.view(k.shape[0], -1)
+        k_rope_2d = k_rope.view(k_rope.shape[0], -1)
+        q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+        q_rope_3d = q_rope.view(
+            -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+        )
+        # Same raw per-layer buffer the decode kernel reads below (bf16-only
+        # gate means store_dtype == dtype, so no view); get_key_buffer applies
+        # the hybrid pool's full-attention layer-id mapping.
+        kv_raw = self.token_to_kv_pool.get_key_buffer(layer.layer_id)
+        kv_2d = kv_raw.view(kv_raw.shape[0], -1) if kv_raw.dim() != 2 else kv_raw
+        if not set_mla_kv_concat_q_covered(
+            kv_buffer=kv_2d,
+            loc=loc,
+            k_nope=k_nope_2d,
+            k_rope=k_rope_2d,
+            q_nope=q_nope,
+            q_rope=q_rope_3d,
+        ):
+            return None
+        return set_mla_kv_concat_q(
+            kv_buffer=kv_2d,
+            loc=loc,
+            cache_k_nope=k_nope_2d,
+            cache_k_rope=k_rope_2d,
+            q_nope=q_nope,
+            q_rope=q_rope_3d,
+        )
+
+    def _set_kv_and_concat_q_fp8_fused(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        q: torch.Tensor,
+        q_rope: torch.Tensor,
+        k: torch.Tensor,
+        k_rope: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        """fp8-KV decode: quantize + scatter the KV row at ``loc`` (already
+        physical) and build the fp8 [q_nope | q_rope] query in one launch.
+
+        Returns the fp8 query, or None when the fused kernel does not cover
+        the inputs (caller falls back to the aten quantize chain).
+        """
+        k_nope_2d = k.view(k.shape[0], -1)
+        k_rope_2d = k_rope.view(k_rope.shape[0], -1)
+        q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+        q_rope_3d = q_rope.view(
+            -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+        )
+        # fp8 view of the pool's uint8 store; same buffer the decode kernel
+        # reads below.
+        kv_raw = self.token_to_kv_pool.get_key_buffer(layer.layer_id)
+        kv_2d = kv_raw.view(kv_raw.shape[0], -1) if kv_raw.dim() != 2 else kv_raw
+        if not set_mla_kv_concat_q_fp8_covered(
+            kv_buffer=kv_2d,
+            loc=loc,
+            k_nope=k_nope_2d,
+            k_rope=k_rope_2d,
+            q_nope=q_nope,
+            q_rope=q_rope_3d,
+        ):
+            return None
+        parallel = get_parallel()
+        return set_mla_kv_concat_q_fp8(
+            kv_buffer=kv_2d,
+            loc=loc,
+            cache_k_nope=k_nope_2d,
+            cache_k_rope=k_rope_2d,
+            q_nope=q_nope,
+            q_rope=q_rope_3d,
+            # DCP cyclic KV sharding: virtual loc -> owner mask + loc//world
+            # (identity when attn_dcp_size == 1).
+            dcp_world_size=parallel.attn_dcp_size,
+            dcp_rank=parallel.attn_dcp_rank,
+        )
+
     def forward_decode(
         self,
         q: torch.Tensor,  # q_nope
@@ -834,12 +1030,33 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
     ) -> torch.Tensor:
         """Run forward for decode using TRTLLM MLA kernel."""
         merge_query = q_rope is not None
+        fused_fp8_query = None
         if self.data_type == torch.float8_e4m3fn:
             assert q_rope is not None and k_rope is not None
             if cos_sin_cache is None:
-                q, k, k_rope = mla_quantize_without_rope_for_fp8(
-                    q, q_rope, k.squeeze(1), k_rope.squeeze(1)
-                )
+                if save_kv_cache and self._fused_set_kv_concat_q_fp8:
+                    loc = (
+                        self._decode_dense_loc
+                        if self._decode_dense_loc is not None
+                        else (
+                            None if self._unified_mla else forward_batch.out_cache_loc
+                        )
+                    )
+                    if loc is not None:
+                        # Fused: bf16->fp8 quantize + KV scatter + q concat
+                        # in one launch; None when not covered.
+                        fused_fp8_query = self._set_kv_and_concat_q_fp8_fused(
+                            layer=layer,
+                            loc=loc,
+                            q=q,
+                            q_rope=q_rope,
+                            k=k,
+                            k_rope=k_rope,
+                        )
+                if fused_fp8_query is None:
+                    q, k, k_rope = mla_quantize_without_rope_for_fp8(
+                        q, q_rope, k.squeeze(1), k_rope.squeeze(1)
+                    )
             else:
                 q, k, k_rope = mla_quantize_and_rope_for_fp8(
                     q,
@@ -854,34 +1071,65 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 )
             merge_query = False
 
-        # Save KV cache if requested
-        if save_kv_cache:
+        # Save KV cache if requested (the fused fp8 path already wrote it)
+        query = fused_fp8_query
+        if query is None and save_kv_cache:
             assert (
                 k is not None and k_rope is not None
             ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
             if self._decode_dense_loc is not None:
                 # cuda-graph path: dense write loc precomputed out-of-graph, so
                 # the in-graph write captures no translate allocation.
-                self.token_to_kv_pool.set_mla_kv_buffer(
-                    layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True
-                )
+                if merge_query and self._fused_set_kv_concat_q:
+                    # Fused: KV scatter + [q_nope | q_rope] concat in one
+                    # launch; None when the inputs are not covered.
+                    query = self._set_kv_and_concat_q_fused(
+                        layer=layer,
+                        loc=self._decode_dense_loc,
+                        k=k,
+                        k_rope=k_rope,
+                        q=q,
+                        q_rope=q_rope,
+                    )
+                if query is None:
+                    self.token_to_kv_pool.set_mla_kv_buffer(
+                        layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True
+                    )
             else:
                 # eager (or static pool): the pool's _full_translate handles it.
-                self.token_to_kv_pool.set_mla_kv_buffer(
-                    layer, forward_batch.out_cache_loc, k, k_rope
-                )
+                if (
+                    merge_query
+                    and self._fused_set_kv_concat_q
+                    and not self._unified_mla
+                ):
+                    # Static pool: _full_translate is identity, so
+                    # out_cache_loc is already the physical write loc.
+                    query = self._set_kv_and_concat_q_fused(
+                        layer=layer,
+                        loc=forward_batch.out_cache_loc,
+                        k=k,
+                        k_rope=k_rope,
+                        q=q,
+                        q_rope=q_rope,
+                    )
+                if query is None:
+                    self.token_to_kv_pool.set_mla_kv_buffer(
+                        layer, forward_batch.out_cache_loc, k, k_rope
+                    )
 
-        # Prepare query tensor inline
-        if merge_query:
-            # For FP16 path, we merge the query and rope parts into a single tensor
-            q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
-            q_rope_reshaped = q_rope.view(
-                -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
-            )
-            query = concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
-        else:
-            # For FP8 path, we already have the query and rope parts merged because of the quantize_and_rope_for_fp8 function
-            query = q.view(-1, layer.tp_q_head_num, layer.head_dim)
+        # Prepare query tensor inline (already built when the fused save-KV
+        # path ran)
+        if query is None:
+            if merge_query:
+                # For FP16 path, we merge the query and rope parts into a single tensor
+                q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+                q_rope_reshaped = q_rope.view(
+                    -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                )
+                query = concat_mla_absorb_q_general(q_nope, q_rope_reshaped)
+            else:
+                # For FP8 path, we already have the query and rope parts merged because of the quantize_and_rope_for_fp8 function
+                query = q.view(-1, layer.tp_q_head_num, layer.head_dim)
 
         # Apply llama 4 scaling if provided
         if llama_4_scaling is not None:
@@ -915,7 +1163,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             query=query,
             kv_cache=kv_cache,
             block_tables=metadata.block_kv_indices,
-            seq_lens=forward_batch.seq_lens,
+            seq_lens=(
+                metadata.seq_lens_k
+                if metadata.seq_lens_k is not None
+                else forward_batch.seq_lens
+            ),
             max_seq_len=metadata.max_seq_len_k,
             layer=layer,
         )
@@ -1038,9 +1290,38 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 max_seq_len = metadata.max_seq_len_k + (
                     0 if dcp_enabled else draft_token_num
                 )
-                # For target_verify, all sequences have the same number of draft tokens
-                q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
-                needs_unpad = False
+                ragged_layout = forward_batch.spec_info.ragged_verify_layout
+                if ragged_layout is None or dcp_enabled:
+                    q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
+                    needs_unpad = False
+                else:
+                    if ragged_layout.bs != bs:
+                        # Capped variant: the dense [bs, draft_token_num] q
+                        # buffer below cannot take a row the full-coverage pad
+                        # may inflate past the verify window, and it keeps
+                        # qo_indptr consistent with the clamped lens (same
+                        # contract as the KDA dense path).
+                        ragged_layout = ragged_layout.padded_to_bucket(
+                            padded_bs=bs, cap=draft_token_num
+                        )
+                    total_tokens = q.shape[0]
+                    seq_lens_q = torch.clamp(
+                        ragged_layout.verify_lens, max=draft_token_num
+                    )
+                    cu_seqlens_q = ragged_layout.qo_indptr_device
+                    padded_q = torch.zeros(
+                        (bs, draft_token_num, layer.tp_q_head_num, layer.head_dim),
+                        dtype=q.dtype,
+                        device=q.device,
+                    )
+                    q = self.pad_draft_extend_query(
+                        q, padded_q, seq_lens_q, cu_seqlens_q
+                    )
+                    needs_unpad = True
+                    unpad_zero_uncovered = True
+                    unpad_seq_lens_q = seq_lens_q
+                    unpad_cu_seqlens_q = cu_seqlens_q
+                    unpad_sum_seq_lens_q = total_tokens
             else:
                 # draft_extend: handle varying num_correct_drafts_per_req. If total_tokens % bs == 0,
                 # we can directly reshape q; otherwise, pad to max_seq_len_q.
@@ -1084,6 +1365,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                         q, padded_q, actual_seq_lens_q, actual_cu_seqlens_q
                     )
                     needs_unpad = True
+                    unpad_zero_uncovered = False
                     unpad_seq_lens_q = actual_seq_lens_q
                     unpad_cu_seqlens_q = actual_cu_seqlens_q
                     unpad_sum_seq_lens_q = total_tokens
@@ -1145,6 +1427,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                     unpad_cu_seqlens_q,
                     unpad_seq_lens_q,
                     unpad_sum_seq_lens_q,
+                    zero_uncovered=unpad_zero_uncovered,
                 )
                 output = output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
             else:

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

@@ -169,6 +169,43 @@ def prepare_vision_attention_metadata(
     )
 
 
+def prepare_flashinfer_cudnn_vision_attention_metadata(
+    cu_seqlens: torch.Tensor,
+    device: torch.device,
+    *,
+    elem_per_token: int,
+) -> VisionAttentionMetadata:
+    cu_seqlens = cu_seqlens.to(device=device, dtype=torch.int32, non_blocking=True)
+    seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+    batch_size = int(seq_lens.numel())
+    padded_batch_size = next(
+        (size for size in BATCH_BUCKETS if size >= batch_size),
+        math.ceil(batch_size / BATCH_BUCKETS[0]) * BATCH_BUCKETS[0],
+    )
+    if padded_batch_size != batch_size:
+        pad_size = padded_batch_size - batch_size
+        padded_indptrs = torch.cat([cu_seqlens, cu_seqlens[-1].expand(pad_size)])
+        padded_seq_lens = torch.cat([seq_lens, seq_lens.new_zeros(pad_size)])
+    else:
+        padded_indptrs = cu_seqlens
+        padded_seq_lens = seq_lens
+
+    elem_indptrs = padded_indptrs * elem_per_token
+    real_max_seqlen = int(seq_lens.max().item())
+    max_seqlen = next(
+        (size for size in FLASHINFER_MAX_SEQLEN_BUCKETS if size >= real_max_seqlen),
+        math.ceil(real_max_seqlen / FLASHINFER_MAX_SEQLEN_BUCKETS[-1])
+        * FLASHINFER_MAX_SEQLEN_BUCKETS[-1],
+    )
+    return prepare_vision_attention_metadata(
+        cu_seqlens,
+        device=device,
+        packed_indptrs=torch.cat([elem_indptrs] * 3),
+        sequence_lengths=padded_seq_lens.view(-1, 1, 1, 1),
+        flashinfer_max_seqlen=max_seqlen,
+    )
+
+
 # TODO: requires real seqlens from images
 @functools.lru_cache(maxsize=128)
 def _get_cu_seqlens_for_shape(batch_size: int, seqlen: int, device) -> torch.Tensor:

python/sglang/srt/layers/attn_residual.py

@@ -0,0 +1,505 @@
+# SPDX-License-Identifier: Apache-2.0
+# Kimi-K3 Attention Residual: snapshot bank + aggregation.
+#
+# The public API is the AttnResidual class (constructed once per forward pass).
+# It owns the frozen snapshot bank [T, NB, H] and the valid-row counter, and
+# dispatches each aggregation point (score rows → softmax → weighted sum →
+# RMSNorm) by hardware capability:
+#   fast  — warp-specialized TMA kernel: cp.async.bulk producer +
+#           online-softmax consumers over a double-buffered chunk ring, out
+#           norm fused, per-nvb tuned launch config, one persistent CTA per
+#           SM. Taken on SM100+ with H=7168.
+#   fused — Triton 2-kernel pipeline with full H-parallelism; the fallback
+#           everywhere the fast kernel does not apply.
+# aggregate_stream_torch is the eager reference (tests and the
+# H % _BLOCK_H != 0 shape fallback of aggregate_stream).
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.linear import ReplicatedLinear
+
+_BLOCK_H: int = 1024  # H = 7168 = 7 x 1024
+_MAX_ROWS: int = 16  # next_pow2(8 + 1), K3 has <= 8 snapshots
+
+_FAST_SUPPORTED = None
+
+
+def _use_fast(hidden_size: int) -> bool:
+    """The TMA kernel needs SM100+ (tcgen05, cp.async.bulk) and its H=7168
+    template instantiation; everything else takes the triton pipeline."""
+    global _FAST_SUPPORTED
+    if _FAST_SUPPORTED is None:
+        major, _ = torch.cuda.get_device_capability()
+        _FAST_SUPPORTED = major >= 10
+    return _FAST_SUPPORTED and hidden_size == 7168
+
+
+def get_cw(
+    proj: ReplicatedLinear,
+    norm: RMSNorm,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Cached product norm_weight ⊙ proj_weight (both [H]) in `dtype`.
+
+    Cached per dtype: the fast kernel consumes bf16 while the triton path
+    consumes fp32, and a shared slot would hand one path the other's dtype."""
+    cache = getattr(proj, "_attn_res_cw_cache", None)
+    if cache is None:
+        cache = {}
+        proj._attn_res_cw_cache = cache
+    cw = cache.get(dtype)
+    if cw is None:
+        cw = (norm.weight.float() * proj.weight.squeeze().float()).contiguous()
+        cw = cache[dtype] = cw.to(dtype)
+    return cw
+
+
+def _aggregate_fast(
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    nvb: int,
+    score_proj: ReplicatedLinear,
+    score_norm: RMSNorm,
+    out_norm: RMSNorm,
+    write_bank_row: bool = False,
+) -> torch.Tensor:
+    """Warp-specialized TMA kernel: online softmax over row chunks with the
+    output RMSNorm fused, one persistent CTA per SM, per-nvb tuned launch
+    config (GB300 benchmark winner across nvb). With write_bank_row the kernel
+    also snapshots the prefix row into bank[:, nvb, :] (bit-exact, zero extra
+    reads — the row streams through the score pass anyway)."""
+    from sglang.kernels.ops.kimi_k3.attn_res import attn_res_fused_tma
+
+    # The kernel applies one eps to both the score norm and the output norm.
+    assert score_norm.variance_epsilon == out_norm.variance_epsilon
+
+    cw = get_cw(score_proj, score_norm, dtype=torch.bfloat16)
+    out = torch.empty_like(prefix_sum)
+    attn_res_fused_tma(
+        prefix_sum,
+        bank,
+        cw,
+        out_norm.weight,
+        out,
+        nvb,
+        score_norm.variance_epsilon,
+        write_prefix=write_bank_row,
+    )
+    return out
+
+
+@triton.jit
+def _score_kernel(
+    prefix_ptr,  # [T, H]
+    bank_ptr,  # [T, NB_total, H]
+    cw_ptr,  # [H] fp32
+    scores_ptr,  # [T, MAX_ROWS] fp32
+    NVB,
+    eps,
+    stride_pm,
+    stride_bm,
+    stride_bb,
+    stride_sm,
+    H: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    """One CTA per (token, row): scan H, output one scalar score."""
+    pid_t = tl.program_id(0)
+    j = tl.program_id(1)
+    if j > NVB:
+        return
+    sumsq = 0.0
+    dotv = 0.0
+    for h0 in tl.static_range(0, H, BLOCK_H):
+        offs_h = h0 + tl.arange(0, BLOCK_H)
+        if j < NVB:
+            v = tl.load(bank_ptr + pid_t * stride_bm + j * stride_bb + offs_h).to(
+                tl.float32
+            )
+        else:
+            v = tl.load(prefix_ptr + pid_t * stride_pm + offs_h).to(tl.float32)
+        cw = tl.load(cw_ptr + offs_h)
+        sumsq += tl.sum(v * v)
+        dotv += tl.sum(v * cw)
+    rrms = 1.0 / tl.sqrt(sumsq / H + eps)
+    tl.store(scores_ptr + pid_t * stride_sm + j, dotv * rrms)
+
+
+@triton.jit
+def _combine_kernel(
+    prefix_ptr,
+    bank_ptr,
+    scores_ptr,  # [T, MAX_ROWS] fp32
+    out_ptr,  # [T, H]
+    NVB,
+    stride_pm,
+    stride_bm,
+    stride_bb,
+    stride_sm,
+    stride_om,
+    BLOCK_H: tl.constexpr,
+    MAX_ROWS: tl.constexpr,
+):
+    """One CTA per (token, H-chunk): softmax(scores) → weighted sum → write chunk.
+
+    Softmax is redundantly computed by each H-chunk CTA (≤16 elements, trivial).
+    This gives full H-parallelism: 7 CTAs for H=7168/1024.
+    """
+    pid_t = tl.program_id(0)
+    pid_h = tl.program_id(1)
+    h0 = pid_h * BLOCK_H
+
+    # Softmax (redundant per chunk, 16 fp32 ops)
+    offs_b = tl.arange(0, MAX_ROWS)
+    mask_b = offs_b <= NVB
+    raw = tl.load(
+        scores_ptr + pid_t * stride_sm + offs_b, mask=mask_b, other=float("-inf")
+    )
+    m = tl.max(raw, axis=0)
+    e = tl.where(mask_b, tl.exp(raw - m), 0.0)
+    p = e / tl.sum(e, axis=0)
+
+    # Weighted sum for this H chunk
+    offs_h = h0 + tl.arange(0, BLOCK_H)
+    acc = tl.zeros([BLOCK_H], tl.float32)
+    for j in range(0, NVB + 1):
+        if j < NVB:
+            v = tl.load(bank_ptr + pid_t * stride_bm + j * stride_bb + offs_h).to(
+                tl.float32
+            )
+        else:
+            v = tl.load(prefix_ptr + pid_t * stride_pm + offs_h).to(tl.float32)
+        p_j = tl.sum(tl.where(offs_b == j, p, 0.0), axis=0)
+        acc += p_j * v
+    tl.store(
+        out_ptr + pid_t * stride_om + offs_h,
+        acc.to(out_ptr.dtype.element_ty),
+    )
+
+
+def _mix_fused(
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    nvb: int,
+    score_proj: ReplicatedLinear,
+    score_norm: RMSNorm,
+) -> torch.Tensor:
+    """Triton score + combine pair: returns the pre-norm mixture."""
+    T, H = prefix_sum.shape
+    cw = get_cw(score_proj, score_norm)
+    n_h_blocks = H // _BLOCK_H
+
+    # Step 1: score each row (2D grid, full row-parallelism)
+    scores = torch.empty((T, _MAX_ROWS), dtype=torch.float32, device=prefix_sum.device)
+    _score_kernel[(T, nvb + 1)](
+        prefix_sum,
+        bank,
+        cw,
+        scores,
+        nvb,
+        score_norm.variance_epsilon,
+        prefix_sum.stride(0),
+        bank.stride(0),
+        bank.stride(1),
+        scores.stride(0),
+        H=H,
+        BLOCK_H=_BLOCK_H,
+        num_warps=8,
+    )
+
+    # Step 2: softmax + weighted sum (2D grid, full H-parallelism)
+    out = torch.empty_like(prefix_sum)
+    _combine_kernel[(T, n_h_blocks)](
+        prefix_sum,
+        bank,
+        scores,
+        out,
+        nvb,
+        prefix_sum.stride(0),
+        bank.stride(0),
+        bank.stride(1),
+        scores.stride(0),
+        out.stride(0),
+        BLOCK_H=_BLOCK_H,
+        MAX_ROWS=_MAX_ROWS,
+        num_warps=4,
+    )
+    return out
+
+
+def _aggregate_fused(
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    nvb: int,
+    score_proj: ReplicatedLinear,
+    score_norm: RMSNorm,
+    out_norm: RMSNorm,
+) -> torch.Tensor:
+    # Step 3: standard RMSNorm (sglang's optimized kernel)
+    return out_norm(_mix_fused(prefix_sum, bank, nvb, score_proj, score_norm))
+
+
+def aggregate_stream_torch(
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    nvb: int,
+    score_proj: ReplicatedLinear,
+    score_norm: RMSNorm,
+) -> torch.Tensor:
+    """Eager reference for aggregate_stream (materializes [T, R, H])."""
+    if nvb == 0:
+        return prefix_sum
+    T, H = prefix_sum.shape
+    # rows = [bank[0..nvb-1], prefix_sum]  shape [T, nvb+1, H]
+    rows = torch.cat([bank[:, :nvb, :], prefix_sum.unsqueeze(1)], dim=1)
+    R = nvb + 1
+    normed = score_norm(rows.reshape(T * R, H))
+    scores = score_proj(normed)[0].reshape(T, R)
+    probs = torch.softmax(scores.float(), dim=-1)
+    mixed = (probs.unsqueeze(-1) * rows.float()).sum(dim=1)
+    return mixed.to(prefix_sum.dtype)
+
+
+def aggregate_stream(
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    nvb: int,
+    score_proj: ReplicatedLinear,
+    score_norm: RMSNorm,
+) -> torch.Tensor:
+    """Pre-norm aggregated stream value (softmax mixture, no output norm):
+    the K3 analogue of the residual stream, for dspark aux capture -- the
+    raw wire only carries the current block's running prefix."""
+    if nvb == 0:
+        return prefix_sum
+    if prefix_sum.shape[1] % _BLOCK_H != 0:
+        return aggregate_stream_torch(prefix_sum, bank, nvb, score_proj, score_norm)
+    return _mix_fused(prefix_sum, bank, nvb, score_proj, score_norm)
+
+
+def _aggregate_fused_add(
+    prefix_a: torch.Tensor,
+    prefix_b: torch.Tensor,
+    bank: torch.Tensor,
+    nvb: int,
+    score_proj: ReplicatedLinear,
+    score_norm: RMSNorm,
+    out_norm: RMSNorm,
+    write_bank_row: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Aggregation point with a pending upstream residual add: materialize
+    prefix = prefix_a + prefix_b, then aggregate. Returns (normed, prefix).
+    write_bank_row rides _aggregate (fast path only)."""
+    prefix = prefix_a + prefix_b
+    return (
+        _aggregate(
+            prefix,
+            bank,
+            nvb,
+            score_proj,
+            score_norm,
+            out_norm,
+            write_bank_row=write_bank_row,
+        ),
+        prefix,
+    )
+
+
+def _aggregate(
+    prefix_sum: torch.Tensor,
+    bank: torch.Tensor,
+    nvb: int,
+    score_proj: ReplicatedLinear,
+    score_norm: RMSNorm,
+    out_norm: RMSNorm,
+    write_bank_row: bool = False,
+) -> torch.Tensor:
+    """Single aggregation point: score → softmax → mix → norm.
+
+    Caller handles nvb == 0 (layer 0 attn side: just out_norm(prefix_sum)).
+    write_bank_row is fast-path only (in-kernel snapshot of the prefix row
+    into bank[:, nvb, :]); the triton path keeps the standalone .write() copy —
+    the caller (AttnResidual.forward) owns that fallback.
+    """
+    if _use_fast(prefix_sum.shape[1]):
+        return _aggregate_fast(
+            prefix_sum,
+            bank,
+            nvb,
+            score_proj,
+            score_norm,
+            out_norm,
+            write_bank_row=write_bank_row,
+        )
+    assert not write_bank_row, "fused bank write is fast-path only"
+    return _aggregate_fused(prefix_sum, bank, nvb, score_proj, score_norm, out_norm)
+
+
+class AttnResidual:
+    """Snapshot bank + aggregation of one K3 attention-residual stream,
+    backed by the capability-dispatched kernels above.
+
+    One instance lives for one model forward pass.
+    """
+
+    def __init__(
+        self,
+        hidden_states: torch.Tensor,
+        block_num: int,
+        block_residual: Optional[torch.Tensor] = None,
+    ) -> None:
+        num_tokens, hidden_size = hidden_states.shape
+        # Frozen snapshot rows [T, NB, H]; raw tensor for PP transfer and the
+        # legacy kernel path.
+        self.block_residual = hidden_states.new_empty(
+            (num_tokens, block_num, hidden_size)
+        )
+        self.num_valid_blocks = 0
+        if block_residual is not None:  # inherited from the previous PP rank
+            self.num_valid_blocks = block_residual.size(1)
+            self.block_residual[:, : self.num_valid_blocks, :].copy_(block_residual)
+
+    def write(self, prefix_sum: torch.Tensor, rows: Optional[slice] = None) -> None:
+        """Snapshot the pre-attention prefix into the next bank row.
+
+        Under SP attention-residual carry each rank owns a disjoint token
+        slice, so only that slice is written and subsequently read locally.
+        """
+        bank = self.block_residual if rows is None else self.block_residual[rows]
+        bank[:, self.num_valid_blocks, :].copy_(prefix_sum)
+        self.num_valid_blocks += 1
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        prefix_sum: Optional[torch.Tensor],
+        score_proj: ReplicatedLinear,
+        score_norm: RMSNorm,
+        out_norm: RMSNorm,
+        rows: Optional[slice] = None,
+        write: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Aggregate; with write=True also snapshot the aggregated prefix
+        (the second return value) into the next bank row — fused into the
+        fast kernel (the row streams through its score pass anyway), a
+        standalone .write() copy on every other path."""
+        nvb = self.num_valid_blocks
+        # Layer 0 attention side: nothing banked yet
+        if nvb == 0:
+            assert prefix_sum is None
+            if write:
+                self.write(hidden_states, rows)
+            return out_norm(hidden_states), hidden_states
+
+        # SP-MoE: the caller holds only its token shard; align the banked
+        # residual rows to it (dim-0 slice of a contiguous buffer stays
+        # contiguous for the jit kernels).
+        block_residual = (
+            self.block_residual if rows is None else self.block_residual[rows]
+        )
+
+        fused_write = write and _use_fast(hidden_states.shape[1])
+        if prefix_sum is None:
+            # hidden_states already is the whole head (PP entry or a
+            # block-boundary restart).
+            normed = _aggregate(
+                hidden_states,
+                block_residual,
+                nvb,
+                score_proj,
+                score_norm,
+                out_norm,
+                write_bank_row=fused_write,
+            )
+            prefix = hidden_states
+        else:
+            # Pending add: materialize the prefix, then aggregate.
+            normed, prefix = _aggregate_fused_add(
+                prefix_sum,
+                hidden_states,
+                block_residual,
+                nvb,
+                score_proj,
+                score_norm,
+                out_norm,
+                write_bank_row=fused_write,
+            )
+        if fused_write:
+            self.num_valid_blocks += 1  # row nvb written in-kernel
+        elif write:
+            self.write(prefix, rows)
+        return normed, prefix
+
+    def forward_sp_all_gather(
+        self,
+        hidden_states: torch.Tensor,
+        prefix_sum: Optional[torch.Tensor],
+        score_proj: ReplicatedLinear,
+        score_norm: RMSNorm,
+        out_norm: RMSNorm,
+        rows: slice,
+        write: bool = False,
+    ) -> Optional[tuple[torch.Tensor, torch.Tensor]]:
+        """Fuse a local aggregation point and the following row all-gather."""
+        nvb = self.num_valid_blocks
+        if nvb == 0:
+            return None
+        if prefix_sum is not None and prefix_sum.shape != hidden_states.shape:
+            prefix_sum = prefix_sum[rows]
+        prefix = hidden_states if prefix_sum is None else prefix_sum.add(hidden_states)
+        bank = self.block_residual[rows]
+        cw = get_cw(score_proj, score_norm, dtype=torch.bfloat16)
+        assert score_norm.variance_epsilon == out_norm.variance_epsilon
+        from sglang.srt.layers import k3_sp_collective
+
+        normed = k3_sp_collective.attn_res_all_gather(
+            prefix,
+            bank,
+            cw,
+            out_norm.weight,
+            nvb,
+            score_norm.variance_epsilon,
+            write_prefix=write,
+        )
+        if normed is None:
+            return None
+        if write:
+            self.num_valid_blocks += 1
+        return normed, prefix
+
+    def forward_sp_reduce_scatter(
+        self,
+        hidden_states: torch.Tensor,
+        prefix_sum: Optional[torch.Tensor],
+        score_proj: ReplicatedLinear,
+        score_norm: RMSNorm,
+        out_norm: RMSNorm,
+        rows: slice,
+    ) -> Optional[tuple[torch.Tensor, torch.Tensor]]:
+        """Fuse o_proj RS, the pending local prefix add, and aggregation."""
+        nvb = self.num_valid_blocks
+        if nvb == 0:
+            return None
+        local_tokens = rows.stop - rows.start
+        residual = prefix_sum
+        if residual is not None and residual.shape[0] != local_tokens:
+            residual = residual[rows]
+        bank = self.block_residual[rows]
+        cw = get_cw(score_proj, score_norm, dtype=torch.bfloat16)
+        assert score_norm.variance_epsilon == out_norm.variance_epsilon
+        from sglang.srt.layers import k3_sp_collective
+
+        return k3_sp_collective.reduce_scatter_attn_res(
+            hidden_states,
+            residual,
+            bank,
+            cw,
+            out_norm.weight,
+            nvb,
+            score_norm.variance_epsilon,
+        )

python/sglang/srt/layers/k3_ar_fusion.py

@@ -0,0 +1,386 @@
+"""K3 MNNVL fused all-reduce dispatch.
+
+Auto-enabled on SM100/SM103 when CustomAllReduceV2 with multicast is
+available; ``SGLANG_K3_AR_FUSION`` overrides in either direction (0 = off,
+1 = attempt anywhere and warn when unavailable).
+
+Glue between the model and the ``kernels.ops.kimi_k3.all_reduce`` kernels. Small
+messages take the 1shot multicast-push through the group's CustomAllReduceV2
+workspace; large ones take the in-place NVLS 2shot, which needs its input to be a
+:func:`symm_buffer` slice and otherwise falls back to the regular all-reduce.
+
+Call-site contract when the fusion is active: ``enabled()`` was
+checked once (which initializes the state), ``x`` is bf16 and contiguous,
+and the residual is identical on every rank (a fully reduced tensor such
+as the attn-res prefix sum) or ``None``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, List, NamedTuple, Optional
+
+import torch
+
+import sglang.srt.runtime_context as ctx
+from sglang.kernels.jit.utils import cache_once
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+        CustomAllReduceV2,
+    )
+
+logger = logging.getLogger(__name__)
+
+# push wins below ~0.5 MB on B200x8/GB300 (bs<=32 @ H=7168); NVLS 2shot wins
+# above (measured: push 8.5us vs pull 11.0us at 448KB, 15.5 vs 11.0 at 896KB)
+_PUSH_MAX_BYTES = 512 * 1024
+
+# Latent width of the K3 latent|shared MoE buffer the fused-norm AR expects
+# ([N, NORM_DIM] latent then [N, 2*NORM_DIM] shared). MUST match kNormDim in
+# csrc/kimi_k3/comm/ar_fusion.cuh — the mod hardcodes this row width.
+NORM_DIM = 3584
+
+# :func:`symm_buffer` names: o_proj's [rows, hidden_size] output, and the MoE
+# [latent | shared] pair flattened to rows x (moe_hidden_size + hidden_size).
+ATTN_O_PROJ = "attn_o_proj"
+MOE_LATENT_SHARED = "moe_latent_shared"
+
+
+class _State(NamedTuple):
+    comm: CustomAllReduceV2
+    world_size: int
+    group_name: str
+
+
+@cache_once
+def _get_state() -> Optional[_State]:
+    explicit = envs.SGLANG_K3_AR_FUSION.is_set()
+    if explicit:
+        if not envs.SGLANG_K3_AR_FUSION.get():
+            return None
+    else:
+        # Auto-probe: only SM100/SM103, and not under --enable-symm-mem or EP
+        # a2a (their allocator contexts misroute the buffers the pull needs).
+        from sglang.srt.utils.common import get_device_sm
+
+        if get_device_sm() not in (100, 103):
+            return None
+        server_args = ctx.get_server_args()
+        if server_args.enable_symm_mem or server_args.moe_a2a_backend != "none":
+            logger.info(
+                "K3 all-reduce fusion auto-probe: skipping "
+                "(enable_symm_mem=%s, moe_a2a_backend=%s; under symm-mem the "
+                "allocator contexts conflict, and under EP a2a the model's "
+                "symm-pool allocation contract does not hold on every AR "
+                "call-site. Set SGLANG_K3_AR_FUSION=1 to force.)",
+                server_args.enable_symm_mem,
+                server_args.moe_a2a_backend,
+            )
+            return None
+    from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+        CustomAllReduceV2,
+    )
+    from sglang.srt.distributed.parallel_state import get_tp_group
+    from sglang.srt.runtime_context import get_parallel
+
+    if get_parallel().tp_size <= 1:
+        return None
+    group = get_tp_group()
+    comm = group.ca_comm
+    if (
+        not isinstance(comm, CustomAllReduceV2)
+        or comm.disabled
+        or comm.mc_base_ptr == 0
+    ):
+        if explicit:
+            logger.warning(
+                "SGLANG_K3_AR_FUSION requested but CustomAllReduceV2 with multicast "
+                "is unavailable; falling back to the regular all-reduce path."
+            )
+        else:
+            logger.info(
+                "K3 all-reduce fusion auto-probe: CustomAllReduceV2 with "
+                "multicast is unavailable; using the regular all-reduce path."
+            )
+        return None
+    from sglang.kernels.ops.kimi_k3 import all_reduce as mod
+
+    mod.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
+    logger.info("K3 all-reduce fusion enabled (world_size=%d)", comm.world_size)
+    return _State(comm, comm.world_size, group.cpu_group.group_name)
+
+
+def enabled() -> bool:
+    return _get_state() is not None
+
+
+class _Buffer(NamedTuple):
+    region: torch.Tensor  # flat [max_rows * width]
+    base: int  # local VA of region[0]
+    mc_base: int  # multicast VA of region[0]
+    max_rows: int
+    width: int
+
+
+# Reverse pointer -> multicast lookup for get_mc_ptr().
+_BUFS: List[_Buffer] = []
+
+
+@cache_once
+def _max_buffer_rows() -> int:
+    """Rows to reserve per buffer: the largest batch the server args allow."""
+    server_args = ctx.get_server_args()
+    chunked = server_args.chunked_prefill_size
+    if chunked is not None and chunked > 0:
+        return int(chunked)
+    return int(server_args.max_prefill_tokens or 0)
+
+
+def _create_buffer(
+    name: str, width: int, dtype: torch.dtype, group_name: str
+) -> _Buffer:
+    """One symmetric buffer, rendezvoused once. ``empty_strided_p2p`` skips the
+    caching allocator, so this is right for a persistent buffer only."""
+    from torch._C._distributed_c10d import _SymmetricMemory
+
+    max_rows = _max_buffer_rows()
+    # outside inference mode on purpose: created inside it the buffer would be an
+    # inference tensor, and the in-place write from a no_grad path (CUDA graph
+    # capture) is then rejected
+    with torch.inference_mode(False):
+        region = _SymmetricMemory.empty_strided_p2p(
+            (max_rows * width,),
+            [1],
+            dtype,
+            torch.device("cuda", torch.cuda.current_device()),
+            group_name,
+        ).view(max_rows, width)
+    handle = _SymmetricMemory.rendezvous(region)
+    assert handle is not None and handle.multicast_ptr != 0
+    buf = _Buffer(
+        region=region,
+        base=int(region.data_ptr()),
+        mc_base=int(handle.multicast_ptr),
+        max_rows=max_rows,
+        width=width,
+    )
+    _BUFS.append(buf)
+    logger.info(
+        "K3 symm buffer %r: %d rows x %d, %.1f MB (rendezvoused once on first "
+        "use; no per-forward symmetric allocation)",
+        name,
+        max_rows,
+        width,
+        region.numel() * region.element_size() / (1024 * 1024),
+    )
+    return buf
+
+
+def symm_buffer(
+    name: str,
+    rows: int,
+    width: int,
+    dtype: torch.dtype,
+    group_name: Optional[str] = None,
+):
+    """``[rows, width]`` view of a named persistent symmetric buffer.
+
+    The pull reduces in place on the input's multicast alias, so every rank must
+    resolve the same ``(buffer, offset)``; a per-forward symm-pool allocation does
+    not, because torch's caching allocator breaks ties on the absolute address and
+    ``rendezvous`` validates sizes only.
+
+    ``group_name`` defaults to the TP group, the one the fused all-reduce reduces
+    over; a caller reducing over another group (SP-MoE's attention TP group) passes
+    its own. Each name belongs to one group, so the name alone identifies it.
+    """
+    if group_name is None:
+        from sglang.srt.distributed.parallel_state import get_tp_group
+
+        group_name = get_tp_group().cpu_group.group_name
+    buf: _Buffer = ctx.get_buffer(
+        f"k3_symm:{name}", lambda: _create_buffer(name, width, dtype, group_name)
+    )
+    assert 0 < rows <= buf.max_rows and width == buf.width and dtype == buf.region.dtype
+    return buf.region[:rows]
+
+
+def find_mc_ptr(x: torch.Tensor) -> Optional[int]:
+    """Multicast VA of ``x`` if it lies inside a named buffer, else None."""
+    ptr = int(x.data_ptr())
+    end = ptr + x.numel() * x.element_size()
+    for buf in _BUFS:
+        if buf.base <= ptr and end <= buf.base + buf.region.nbytes:
+            return buf.mc_base + (ptr - buf.base)
+    return None
+
+
+def get_mc_ptr(x: torch.Tensor) -> int:
+    """Multicast VA of ``x``, which must be a :func:`symm_buffer` slice."""
+    mc = find_mc_ptr(x)
+    assert mc is not None, "K3 fused pull input is not a symm_buffer slice"
+    return mc
+
+
+def all_reduce(
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """In-place ``x = allreduce(x) [+ residual]``; returns ``x``."""
+    from sglang.kernels.ops.kimi_k3 import all_reduce as mod
+
+    state = _get_state()
+    assert state is not None
+    if x.shape[0] == 0:
+        return x if residual is None else x + residual
+    nbytes = x.numel() * 2
+    if nbytes <= min(_PUSH_MAX_BYTES, state.comm.max_push_size):
+        return mod.all_reduce_push_res(
+            state.world_size, x, residual, ws_mc_base=state.comm.mc_base_ptr
+        )
+    return mod.all_reduce_pull_res(
+        state.world_size, x, residual, input_mc_ptr=get_mc_ptr(x)
+    )
+
+
+def all_reduce_low_sm(
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    *,
+    num_blocks: Optional[int] = None,
+    unroll: Optional[int] = None,
+) -> torch.Tensor:
+    """In-place ``x = allreduce(x) [+ residual]``, pinned to the low-SM NVLS pull.
+
+    For side-stream use: push would fan out over many blocks and steal SMs from
+    the GEMMs it overlaps. ``num_blocks`` / ``unroll`` default to the tuned tables.
+    """
+    from sglang.kernels.ops.kimi_k3 import all_reduce as mod
+
+    state = _get_state()
+    assert state is not None
+    if x.shape[0] == 0:
+        return x if residual is None else x + residual
+    return mod.all_reduce_pull_res(
+        state.world_size,
+        x,
+        residual,
+        input_mc_ptr=get_mc_ptr(x),
+        num_blocks=num_blocks,
+        unroll=unroll,
+    )
+
+
+def all_reduce_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    num_tokens: int,
+) -> torch.Tensor:
+    """In-place ``x = allreduce(x)`` with a fused RMSNorm over the first
+    ``num_tokens`` rows of the 2D ``[rows, NORM_DIM]`` input."""
+    from sglang.kernels.ops.kimi_k3 import all_reduce as mod
+
+    state = _get_state()
+    assert state is not None
+    if x.shape[0] == 0:
+        return x
+    nbytes = x.numel() * 2
+    if nbytes <= min(_PUSH_MAX_BYTES, state.comm.max_push_size):
+        return mod.all_reduce_push_norm(
+            state.world_size,
+            x,
+            weight,
+            eps,
+            num_norm_rows=num_tokens,
+            ws_mc_base=state.comm.mc_base_ptr,
+        )
+    return mod.all_reduce_pull_norm(
+        state.world_size,
+        x,
+        weight,
+        eps,
+        num_norm_rows=num_tokens,
+        input_mc_ptr=get_mc_ptr(x),
+    )
+
+
+def finalize_push_fits(num_tokens: int) -> bool:
+    """Whether a [num_tokens, NORM_DIM] latent fits the push window; the
+    finalize-fused AR is push-only."""
+    state = _get_state()
+    assert state is not None
+    nbytes = num_tokens * NORM_DIM * 2
+    return nbytes <= min(_PUSH_MAX_BYTES, state.comm.max_push_size)
+
+
+def gemm_ag_up_fits(num_tokens: int) -> bool:
+    """Whether :func:`gemm_ag_up_proj` covers this decode batch: TP8, within the
+    kernel's win range, and the staging slice fits a push slot."""
+    from sglang.kernels.ops.kimi_k3 import gemm_ag as mod
+
+    state = _get_state()
+    assert state is not None
+    comm = state.comm
+    return (
+        0 < num_tokens <= mod.MAX_TOKENS
+        and state.world_size == 8
+        and num_tokens * (2 * NORM_DIM // 8) * 2 <= comm.max_push_size
+        # The GEMV producer reads a single phase counter, so its grid is not
+        # bound to the counter array; the spin consumer still needs one block
+        # per counter slot plus the cleanup block.
+        and comm.config.num_push_blocks >= 2
+    )
+
+
+def gemm_ag_up_proj(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    b: torch.Tensor,
+    c: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """``up_proj(x) + b (+ c)`` via column-parallel GEMV + multicast all-gather +
+    fused add3. ``weight`` is the FULL replicated up_proj (the kernel slices this
+    rank's rows). Caller checked :func:`gemm_ag_up_fits`."""
+    from sglang.kernels.ops.kimi_k3 import gemm_ag as mod
+
+    state = _get_state()
+    assert state is not None
+    return mod.gemm_ag_up_proj(
+        state.world_size,
+        x,
+        weight,
+        b,
+        c,
+        torch.empty_like(b),
+        ws_mc_base=state.comm.mc_base_ptr,
+    )
+
+
+def finalize_all_reduce_push_norm(
+    out: torch.Tensor,
+    gemm2_out: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    expert_weights: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Deferred MoE finalize fused into the 1shot push AR + RMSNorm; ``out`` is
+    output-only. Caller checked :func:`finalize_push_fits`."""
+    from sglang.kernels.ops.kimi_k3 import all_reduce as mod
+
+    state = _get_state()
+    assert state is not None
+    return mod.finalize_all_reduce_push_norm(
+        state.world_size,
+        out,
+        gemm2_out,
+        expanded_idx_to_permuted_idx,
+        expert_weights,
+        weight,
+        eps,
+        ws_mc_base=state.comm.mc_base_ptr,
+    )

python/sglang/srt/layers/k3_gemm_ar.py

@@ -0,0 +1,95 @@
+"""K3 fused o_proj GEMM+AR dispatch (``SGLANG_K3_GEMM_AR``).
+
+Glue between the model and ``kernels.ops.kimi_k3.gemm_ar``: lazily allocates
+the P2P comm region on the TP group and swaps the o_proj RowParallelLinear
+forward for the single fused GEMM+all-reduce kernel at decode shapes
+(falling through to the regular GEMM + AR path whenever the input doesn't
+fit — prefill, capture, non-2D input). Eager-mode only; see
+kernels/ops/kimi_k3/GEMM_AR_README.md.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.linear import RowParallelLinear
+
+logger = logging.getLogger(__name__)
+
+_INITIALIZED = False
+_ENABLED = False
+
+
+def _init() -> bool:
+    global _INITIALIZED, _ENABLED
+    if _INITIALIZED:
+        return _ENABLED
+    _INITIALIZED = True
+    if not envs.SGLANG_K3_GEMM_AR.get():
+        return False
+    from sglang.srt.runtime_context import get_parallel
+
+    world_size = get_parallel().tp_size
+    if not (2 <= world_size <= 8):
+        logger.warning("SGLANG_K3_GEMM_AR requires 2 <= TP <= 8; disabled.")
+        return False
+    if torch.cuda.get_device_capability() < (10, 0):
+        logger.warning("SGLANG_K3_GEMM_AR requires SM100+; disabled.")
+        return False
+    _ENABLED = True
+    logger.info("K3 fused o_proj GEMM+AR enabled (world_size=%d)", world_size)
+    return True
+
+
+def maybe_wrap_o_proj(o_proj: RowParallelLinear) -> None:
+    """SGLANG_K3_GEMM_AR: route decode-shaped o_proj calls through the fused
+    GEMM+all-reduce kernel; everything else falls through to the original
+    forward. The comm region + JIT compile happen HERE (model build, weight
+    shape known, well before CUDA-graph capture) — capture must only see the
+    ready-to-launch path."""
+    if not _init():
+        return
+    from sglang.kernels.ops.kimi_k3 import gemm_ar as mod
+    from sglang.srt.distributed.parallel_state import get_tp_group
+    from sglang.srt.runtime_context import get_parallel
+
+    parallel = get_parallel()
+    world_size = parallel.tp_size
+    if not o_proj.reduce_results or o_proj.tp_size != world_size:
+        return
+    weight = o_proj.weight
+    if not (
+        isinstance(weight, torch.Tensor)
+        and weight.shape[0] == mod.N
+        and weight.shape[1] % 128 == 0
+    ):
+        return
+
+    mod.init(
+        world_size=world_size,
+        rank=parallel.tp_rank,
+        group=get_tp_group().cpu_group,
+        k=weight.shape[1],
+    )
+    # per-K compile + base-address stash, pre-capture
+    mod._module_with_bases(weight.shape[1], world_size)
+
+    inner = o_proj.forward
+
+    def _gemm_ar_forward(x, *args, **kwargs):
+        weight = o_proj.weight
+        if (
+            isinstance(weight, torch.Tensor)
+            and weight.dtype == torch.bfloat16
+            and mod.fits(x)
+        ):
+            return mod.o_proj_gemm_ar(x, weight), None
+        return inner(x, *args, **kwargs)
+
+    o_proj.forward = _gemm_ar_forward

python/sglang/srt/layers/k3_sp_collective.py

@@ -0,0 +1,470 @@
+"""K3 SP-MoE MNNVL reduce-scatter/all-gather dispatch.
+
+The reduce-scatter folds the pending attention residual into its reduction
+epilogue.  The matching all-gather reassembles the token shards after MoE.
+Both reuse CustomAllReduceV2's push workspace and fall back as a pair outside
+the checked-in GB300 tuning envelope.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers import k3_ar_fusion
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+        CustomAllReduceV2,
+    )
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
+
+logger = logging.getLogger(__name__)
+
+_HIDDEN_SIZE = 7168
+_SUPPORTED_WORLD_SIZES = (4, 8)
+
+# Named persistent symmetric buffers, one per NVLS-aliased tensor. Every rank
+# must resolve the same (buffer, offset) for these, which a per-forward
+# allocation does not give -- see k3_ar_fusion.symm_buffer.
+_O_PROJ = "sp_o_proj"  # TP-partial o_proj output, read by the pull RS
+_ALL_GATHER = "sp_all_gather"  # full-batch output of the direct AG
+_ATTN_RES_AG = "sp_attn_res_ag"  # ... and of its attention-residual fusion
+
+
+class _State:
+    def __init__(self, group: GroupCoordinator, comm: CustomAllReduceV2):
+        self.group = group
+        self.comm = comm
+
+
+_STATE: Optional[_State] = None
+_INITIALIZED = False
+_O_PROJ_RESULT_BUFFERS: dict[int, torch.Tensor] = {}
+_O_PROJ_OUTPUT_ROWS: ContextVar[Optional[int]] = ContextVar(
+    "k3_sp_o_proj_output_rows", default=None
+)
+
+
+def _init_state() -> Optional[_State]:
+    global _STATE, _INITIALIZED
+    if _INITIALIZED:
+        return _STATE
+    _INITIALIZED = True
+
+    explicit = envs.SGLANG_K3_SP_COLLECTIVE.is_set()
+    if explicit and not envs.SGLANG_K3_SP_COLLECTIVE.get():
+        return None
+
+    from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+        CustomAllReduceV2,
+    )
+    from sglang.srt.runtime_context import get_parallel, get_server_args
+    from sglang.srt.utils.common import get_device_sm
+
+    server_args = get_server_args()
+    a2a = server_args.moe_a2a_backend
+    group = get_parallel().attn_tp_group
+    comm = group.ca_comm
+    if (
+        get_device_sm() != 103
+        or group.world_size not in _SUPPORTED_WORLD_SIZES
+        or a2a not in ("megamoe", "deepep")
+        or not isinstance(comm, CustomAllReduceV2)
+        or comm.disabled
+        or comm.mc_base_ptr == 0
+    ):
+        message = (
+            "K3 SP collective requires SM103, TP4/TP8, MegaMoE/DeepEP, and "
+            "CustomAllReduceV2 with multicast; using NCCL."
+        )
+        (logger.warning if explicit else logger.info)(message)
+        return None
+
+    from sglang.kernels.ops.kimi_k3 import attn_res, sp_collective
+
+    # Refuse to enable without a checked-in table for this exact device.
+    if (
+        sp_collective.get_dispatch(
+            "reduce_scatter",
+            group.world_size,
+            _HIDDEN_SIZE,
+            1,
+            torch.device("cuda"),
+        )
+        is None
+    ):
+        (logger.warning if explicit else logger.info)(
+            "K3 SP collective has no tuning table for %s; using NCCL.",
+            torch.cuda.get_device_name(),
+        )
+        return None
+
+    sp_collective.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
+    attn_res.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
+    _STATE = _State(group, comm)
+    logger.info(
+        "K3 SP collective enabled (TP%d, fused RS residual + AG)",
+        group.world_size,
+    )
+    if envs.SGLANG_K3_SP_ATTN_RES.get():
+        logger.info(
+            "K3 SP attention-residual carry enabled "
+            "(local agg/bank write + normalized AG)"
+        )
+    return _STATE
+
+
+def enabled() -> bool:
+    return _init_state() is not None
+
+
+def _symm_buffer(
+    state: _State, name: str, rows: int, width: int, dtype: torch.dtype
+) -> torch.Tensor:
+    """Named persistent symmetric buffer over the group this module reduces on."""
+    return k3_ar_fusion.symm_buffer(
+        name, rows, width, dtype, group_name=state.group.cpu_group.group_name
+    )
+
+
+def requires_symmetric_rs(num_tokens: int, device: torch.device) -> bool:
+    """Whether standalone or fused RS reads o_proj through its NVLS alias."""
+    state = _init_state()
+    if state is None:
+        return False
+    from sglang.kernels.ops.kimi_k3 import sp_collective
+
+    dispatch = sp_collective.get_dispatch(
+        "reduce_scatter",
+        state.group.world_size,
+        _HIDDEN_SIZE,
+        num_tokens,
+        device,
+    )
+    if dispatch is not None and dispatch.strategy == "pull":
+        return True
+    if not envs.SGLANG_K3_SP_ATTN_RES.get():
+        return False
+    fusion = sp_collective.get_fusion_dispatch(
+        "reduce_scatter_attn_res",
+        state.group.world_size,
+        _HIDDEN_SIZE,
+        num_tokens,
+        device,
+    )
+    return fusion is not None and fusion.strategy == "fused_pull"
+
+
+def get_o_proj_output_buffer(
+    num_tokens: int, dtype: torch.dtype, hidden_size: int = _HIDDEN_SIZE
+) -> torch.Tensor:
+    """Persistent symmetric storage for a TP-partial o_proj output."""
+    state = _init_state()
+    assert state is not None, "K3 SP collective is not initialized"
+    return _symm_buffer(state, _O_PROJ, num_tokens, hidden_size, dtype)
+
+
+@contextmanager
+def o_proj_output_rows(num_rows: int):
+    """Tell o_proj to target a padded full-batch symmetric output."""
+    token = _O_PROJ_OUTPUT_ROWS.set(num_rows)
+    try:
+        yield
+    finally:
+        _O_PROJ_OUTPUT_ROWS.reset(token)
+
+
+def get_o_proj_output_rows(default: int) -> int:
+    num_rows = _O_PROJ_OUTPUT_ROWS.get()
+    return default if num_rows is None else num_rows
+
+
+def register_o_proj_output(result: torch.Tensor, output: torch.Tensor) -> None:
+    """Associate a model-visible o_proj result/view with its symmetric backing."""
+    if (
+        result.ndim != output.ndim
+        or result.shape[1:] != output.shape[1:]
+        or result.shape[0] > output.shape[0]
+        or result.dtype != output.dtype
+    ):
+        raise RuntimeError(
+            "K3 o_proj result does not match its persistent symmetric output"
+        )
+    _O_PROJ_RESULT_BUFFERS[result.data_ptr()] = output
+
+
+def finish_padded_o_proj_output(
+    result: torch.Tensor, num_padded: int
+) -> Optional[torch.Tensor]:
+    """Zero the padding tail and return the full persistent symmetric buffer."""
+    output = _O_PROJ_RESULT_BUFFERS.get(result.data_ptr())
+    if output is None or output.shape[0] != num_padded:
+        return None
+    output[result.shape[0] :].zero_()
+    _O_PROJ_RESULT_BUFFERS[output.data_ptr()] = output
+    return output
+
+
+def _resolve_symmetric_o_proj_input(tensor: torch.Tensor) -> tuple[torch.Tensor, int]:
+    """Recover the persistent o_proj buffer if a model wrapper rebound its view.
+
+    The second element is 0 when neither is symmetric, and the caller falls back.
+    """
+    input_mc_ptr = k3_ar_fusion.find_mc_ptr(tensor)
+    if input_mc_ptr is not None:
+        return tensor, input_mc_ptr
+    output = _O_PROJ_RESULT_BUFFERS.get(tensor.data_ptr())
+    if output is None:
+        return tensor, 0
+    return output, k3_ar_fusion.find_mc_ptr(output) or 0
+
+
+def _eligible(
+    state: _State,
+    tensor: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+) -> bool:
+    if (
+        tensor.dtype != torch.bfloat16
+        or not tensor.is_contiguous()
+        or tensor.ndim != 2
+        or tensor.shape[1] != _HIDDEN_SIZE
+        or tensor.shape[0] <= 0
+        or tensor.shape[0] % state.group.world_size != 0
+    ):
+        return False
+    if residual is not None:
+        local_shape = (tensor.shape[0] // state.group.world_size, tensor.shape[1])
+        if (
+            residual.shape not in (tensor.shape, local_shape)
+            or residual.dtype != tensor.dtype
+            or not residual.is_contiguous()
+        ):
+            return False
+    local_bytes = tensor.numel() * tensor.element_size() // state.group.world_size
+    return local_bytes <= state.comm.max_push_size
+
+
+def reduce_scatter_res(
+    tensor: torch.Tensor, residual: Optional[torch.Tensor]
+) -> Optional[torch.Tensor]:
+    """Return a fused local shard, or None when the NCCL fallback should run."""
+    state = _init_state()
+    if state is None or not _eligible(state, tensor, residual):
+        return None
+    from sglang.kernels.ops.kimi_k3 import sp_collective
+
+    dispatch = sp_collective.get_dispatch(
+        "reduce_scatter",
+        state.group.world_size,
+        tensor.shape[1],
+        tensor.shape[0],
+        tensor.device,
+    )
+    if dispatch is None:
+        return None
+    output = torch.empty(
+        (tensor.shape[0] // state.group.world_size, tensor.shape[1]),
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    if dispatch.strategy == "push":
+        return sp_collective.reduce_scatter_res(
+            state.group.world_size,
+            tensor,
+            output,
+            residual,
+            tuning=dispatch.tuning,
+        )
+    if dispatch.strategy == "pull":
+        tensor, input_mc_ptr = _resolve_symmetric_o_proj_input(tensor)
+        if input_mc_ptr == 0:
+            logger.warning("K3 pull RS input is not symmetric; using NCCL.")
+            return None
+        return sp_collective.reduce_scatter_pull(
+            state.group.world_size,
+            tensor,
+            output,
+            residual,
+            input_mc_ptr=input_mc_ptr,
+            tuning=dispatch.tuning,
+        )
+    raise AssertionError(f"unknown K3 reduce-scatter strategy: {dispatch.strategy}")
+
+
+def reduce_scatter_attn_res(
+    tensor: torch.Tensor,
+    residual: Optional[torch.Tensor],
+    bank: torch.Tensor,
+    cw: torch.Tensor,
+    ow: torch.Tensor,
+    nvb: int,
+    eps: float,
+) -> Optional[tuple[torch.Tensor, torch.Tensor]]:
+    """Fused NVLS pull RS + local residual + attention-residual aggregation."""
+    state = _init_state()
+    if (
+        state is None
+        or not envs.SGLANG_K3_SP_ATTN_RES.get()
+        or not _eligible(state, tensor, residual)
+    ):
+        return None
+    local_tokens = tensor.shape[0] // state.group.world_size
+    if (
+        bank.shape[0] != local_tokens
+        or bank.ndim != 3
+        or bank.shape[2] != _HIDDEN_SIZE
+        or not bank.is_contiguous()
+        or cw.shape != (_HIDDEN_SIZE,)
+        or ow.shape != (_HIDDEN_SIZE,)
+    ):
+        return None
+    from sglang.kernels.ops.kimi_k3 import attn_res, sp_collective
+
+    dispatch = sp_collective.get_fusion_dispatch(
+        "reduce_scatter_attn_res",
+        state.group.world_size,
+        tensor.shape[1],
+        tensor.shape[0],
+        tensor.device,
+    )
+    if dispatch is None or dispatch.strategy != "fused_pull":
+        return None
+    tensor, input_mc_ptr = _resolve_symmetric_o_proj_input(tensor)
+    if input_mc_ptr == 0:
+        logger.warning("K3 fused pull RS input is not symmetric; using separate path.")
+        return None
+    normed = torch.empty(
+        (local_tokens, tensor.shape[1]), dtype=tensor.dtype, device=tensor.device
+    )
+    prefix = torch.empty_like(normed)
+    attn_res.attn_res_fused_pull_rs(
+        state.group.world_size,
+        tensor,
+        residual,
+        bank,
+        cw,
+        ow,
+        normed,
+        prefix,
+        nvb,
+        eps,
+        input_mc_ptr=input_mc_ptr,
+        max_blocks=dispatch.max_blocks,
+    )
+    return normed, prefix
+
+
+def all_gather(tensor: torch.Tensor) -> Optional[torch.Tensor]:
+    """Return the reassembled full batch, or None for the NCCL fallback."""
+    state = _init_state()
+    if state is None:
+        return None
+    global_tokens = tensor.shape[0] * state.group.world_size
+    if (
+        tensor.dtype != torch.bfloat16
+        or not tensor.is_contiguous()
+        or tensor.ndim != 2
+        or tensor.shape[1] != _HIDDEN_SIZE
+        or tensor.shape[0] <= 0
+        or tensor.numel() * tensor.element_size() > state.comm.max_push_size
+    ):
+        return None
+    from sglang.kernels.ops.kimi_k3 import sp_collective
+
+    dispatch = sp_collective.get_dispatch(
+        "all_gather",
+        state.group.world_size,
+        tensor.shape[1],
+        global_tokens,
+        tensor.device,
+    )
+    if dispatch is None:
+        return None
+    output_shape = (global_tokens, tensor.shape[1])
+    if dispatch.strategy == "push":
+        output = torch.empty(output_shape, dtype=tensor.dtype, device=tensor.device)
+        return sp_collective.all_gather(
+            state.group.world_size,
+            tensor,
+            output,
+            ws_mc_base=state.comm.mc_base_ptr,
+            tuning=dispatch.tuning,
+        )
+    if dispatch.strategy == "direct":
+        output = _symm_buffer(
+            state, _ALL_GATHER, global_tokens, tensor.shape[1], tensor.dtype
+        )
+        return sp_collective.all_gather_direct(
+            state.group.world_size,
+            tensor,
+            output,
+            output_mc_ptr=k3_ar_fusion.get_mc_ptr(output),
+            tuning=dispatch.tuning,
+        )
+    raise AssertionError(f"unknown K3 all-gather strategy: {dispatch.strategy}")
+
+
+def attn_res_all_gather(
+    prefix: torch.Tensor,
+    bank: torch.Tensor,
+    cw: torch.Tensor,
+    ow: torch.Tensor,
+    nvb: int,
+    eps: float,
+    *,
+    write_prefix: bool = False,
+) -> Optional[torch.Tensor]:
+    """Fused local attention-residual aggregation + direct multicast AG."""
+    state = _init_state()
+    if (
+        state is None
+        or not envs.SGLANG_K3_SP_ATTN_RES.get()
+        or prefix.dtype != torch.bfloat16
+        or not prefix.is_contiguous()
+        or prefix.ndim != 2
+        or prefix.shape[1] != _HIDDEN_SIZE
+        or prefix.shape[0] <= 0
+        or bank.ndim != 3
+        or bank.shape[0] != prefix.shape[0]
+        or bank.shape[2] != _HIDDEN_SIZE
+        or not bank.is_contiguous()
+        or cw.shape != (_HIDDEN_SIZE,)
+        or ow.shape != (_HIDDEN_SIZE,)
+    ):
+        return None
+    from sglang.kernels.ops.kimi_k3 import attn_res, sp_collective
+
+    global_tokens = prefix.shape[0] * state.group.world_size
+    dispatch = sp_collective.get_fusion_dispatch(
+        "attn_res_all_gather",
+        state.group.world_size,
+        prefix.shape[1],
+        global_tokens,
+        prefix.device,
+    )
+    if dispatch is None or dispatch.strategy != "fused_direct":
+        return None
+    output = _symm_buffer(
+        state, _ATTN_RES_AG, global_tokens, prefix.shape[1], prefix.dtype
+    )
+    attn_res.attn_res_fused_direct_ag(
+        state.group.world_size,
+        prefix,
+        bank,
+        cw,
+        ow,
+        output,
+        nvb,
+        eps,
+        output_mc_ptr=k3_ar_fusion.get_mc_ptr(output),
+        max_blocks=dispatch.max_blocks,
+        write_prefix=write_prefix,
+    )
+    return output

python/sglang/srt/layers/linear.py

@@ -1560,7 +1560,13 @@ class RowParallelLinear(LinearBase):
                 # Fallback for parameters that don't accept additional args
                 param.load_row_parallel_weight(loaded_weight)
 
-    def forward(self, input_, skip_all_reduce=False, forward_batch=None):
+    def forward(
+        self,
+        input_,
+        skip_all_reduce=False,
+        forward_batch=None,
+        output_tensor=None,
+    ):
         if self.input_is_parallel:
             input_parallel = input_
         else:
@@ -1581,7 +1587,20 @@ class RowParallelLinear(LinearBase):
                 get_tp_group(), disabled=not is_allocation_symmetric()
             )
         with symm_ctx:
-            output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
+            if output_tensor is None:
+                output_parallel = self.quant_method.apply(
+                    self, input_parallel, bias=bias_
+                )
+            else:
+                apply_into = getattr(self.quant_method, "apply_into", None)
+                if apply_into is None:
+                    raise RuntimeError(
+                        f"{type(self.quant_method).__name__} cannot write into "
+                        "caller-owned linear output"
+                    )
+                output_parallel = apply_into(
+                    self, input_parallel, output_tensor, bias=bias_
+                )
 
         # skip_all_reduce: explicit call-site override. Also honor
         # ForwardFlags (fuse_mlp_allreduce / mlp_reduce_scatter) published by

python/sglang/srt/layers/moe/ep_moe/layer.py

@@ -125,6 +125,15 @@ class DeepEPMoE(FusedMoE):
             in ("modelopt_fp4", "modelopt_mixed", "nvfp4_online")
         ):
             self.deprecate_flag = True
+        elif (
+            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+            and get_moe_runner_backend().is_deep_gemm()
+            and quant_config is not None
+            and quant_config.get_name() == "mxfp4"
+        ):
+            # MXFP4 experts (e.g. Kimi K3) on the DeepGEMM fp8_fp4 W4A8 path:
+            # route through the modern FusedMoE runner (Mxfp4MoEMethod.apply).
+            self.deprecate_flag = True
         elif (
             quant_config is None
             and self.w13_weight.dtype == torch.bfloat16
@@ -322,7 +331,7 @@ class DeepEPMoE(FusedMoE):
         self,
         dispatch_output: DeepEPNormalDispatchOutput,
     ):
-        assert self.moe_runner_config.activation == "silu"
+        assert self.moe_runner_config.activation in ("silu", "situ")
         assert isinstance(self.quant_method, W4AFp8MoEMethod)
         return self.quant_method.apply_deepep_normal(
             layer=self,
@@ -333,7 +342,7 @@ class DeepEPMoE(FusedMoE):
         self,
         dispatch_output: DeepEPLLDispatchOutput,
     ):
-        assert self.moe_runner_config.activation == "silu"
+        assert self.moe_runner_config.activation in ("silu", "situ")
         assert isinstance(self.quant_method, W4AFp8MoEMethod)
         return self.quant_method.apply_deepep_ll(
             layer=self,

python/sglang/srt/layers/moe/fused_moe_native.py

@@ -6,7 +6,7 @@ It is based on https://github.com/pytorch-labs/gpt-fast/blob/32971d3129541c5bfb4
 import torch
 from torch.nn import functional as F
 
-from sglang.srt.layers.activation import GeluAndMul, SiluAndMul
+from sglang.srt.layers.activation import GeluAndMul, SiluAndMul, SituAndMul
 from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
 from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
     swiglu_gpt_oss_sigmoid_alpha,
@@ -40,6 +40,14 @@ def fused_moe_forward_native(
         x1 = F.silu(x1)
     elif moe_runner_config.activation == "gelu":
         x1 = F.gelu(x1)
+    elif moe_runner_config.activation == "situ":
+        beta = (
+            moe_runner_config.gemm1_alpha
+            if moe_runner_config.gemm1_alpha is not None
+            else 4.0
+        )
+        x1 = beta * torch.tanh(x1.float() / beta) * torch.sigmoid(x1.float())
+        x1 = x1.to(x.dtype)
     else:
         raise ValueError(f"Unsupported activation: {moe_runner_config.activation=}")
     x3 = torch.einsum("ti, taoi -> tao", x, w3_weights)
@@ -77,6 +85,14 @@ def moe_forward_native(
         act = SiluAndMul()
     elif moe_runner_config.activation == "gelu":
         act = GeluAndMul()
+    elif moe_runner_config.activation == "situ":
+        situ_beta = (
+            moe_runner_config.gemm1_alpha
+            if moe_runner_config.gemm1_alpha is not None
+            else 4.0
+        )
+        situ_linear_beta = moe_runner_config.gemm1_clamp_limit
+        act = SituAndMul(beta=situ_beta, linear_beta=situ_linear_beta)
     else:
         raise ValueError(f"Unsupported activation: {moe_runner_config.activation=}")
 

python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py

@@ -2,7 +2,10 @@ from typing import Optional
 
 import torch
 import torch.nn.functional as F
+import triton
+import triton.language as tl
 
+from sglang.srt.layers import zero_copy_context
 from sglang.srt.utils import is_cuda
 from sglang.srt.utils.custom_op import register_custom_op
 
@@ -15,6 +18,69 @@ if _is_cuda:
     from sglang.kernels.ops.moe.moe_wna16_marlin import moe_wna16_marlin_gemm
 
 
+@triton.jit
+def _tl_tanh(x):
+    return 2.0 * tl.sigmoid(2.0 * x) - 1.0
+
+
+@triton.jit
+def _situ_and_mul_kernel(
+    x_ptr,  # [M, 2N] gate;up halves (non-interleaved)
+    out_ptr,  # [M, N]
+    N,
+    situ_beta,
+    linear_beta,
+    stride_xm,
+    stride_om,
+    BLOCK_N: tl.constexpr,
+    HAS_LINEAR_BETA: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = offs < N
+    base = x_ptr + pid_m * stride_xm
+    gate = tl.load(base + offs, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(base + N + offs, mask=mask, other=0.0).to(tl.float32)
+    gate = situ_beta * _tl_tanh(gate / situ_beta) * tl.sigmoid(gate)
+    if HAS_LINEAR_BETA:
+        up = linear_beta * _tl_tanh(up / linear_beta)
+    out = gate * up
+    tl.store(
+        out_ptr + pid_m * stride_om + offs,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def situ_and_mul(
+    output: torch.Tensor,
+    x: torch.Tensor,
+    situ_beta: float,
+    linear_beta: Optional[float],
+) -> None:
+    """SiTU gated activation (Kimi K3), fused into one elementwise kernel:
+    out = situ_beta*tanh(gate/situ_beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta)
+    where x = [gate; up] halves along the last dim.
+    """
+    M, N2 = x.shape
+    N = N2 // 2
+    assert output.shape == (M, N)
+    BLOCK_N = 1024
+    grid = (M, triton.cdiv(N, BLOCK_N))
+    _situ_and_mul_kernel[grid](
+        x,
+        output,
+        N,
+        float(situ_beta),
+        float(linear_beta) if linear_beta is not None else 0.0,
+        x.stride(0),
+        output.stride(0),
+        BLOCK_N=BLOCK_N,
+        HAS_LINEAR_BETA=linear_beta is not None,
+    )
+
+
 def get_scalar_type(
     num_bits: int,
     has_zp: bool,
@@ -177,9 +243,25 @@ def fused_marlin_moe(
 
     if global_num_experts == -1:
         global_num_experts = E
-    sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
-        topk_ids, block_size_m, global_num_experts
-    )
+    if (
+        M == 1
+        and topk <= 32
+        and expert_map is None
+        # The JIT kernel is int32-only; torch-native topk emits int64 -- let
+        # that (test-only) shape take the generic path instead of casting.
+        and topk_ids.dtype == torch.int32
+    ):
+        # Single-token decode: top-k ids are distinct, so alignment is a
+        # single-warp sort instead of the align + count_and_sort kernel pair.
+        from sglang.kernels.ops.moe.moe_align_single_token import moe_align_single_token
+
+        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_single_token(
+            topk_ids, block_size_m
+        )
+    else:
+        sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+            topk_ids, block_size_m, global_num_experts
+        )
 
     if workspace is None:
         max_workspace_size = (max(2 * N, K) // 64) * (
@@ -266,6 +348,13 @@ def fused_marlin_moe(
         )
     elif activation == "silu" and is_gated:
         silu_and_mul(intermediate_cache1.view(-1, gemm1_n), intermediate_cache2)
+    elif activation == "situ" and is_gated:
+        situ_and_mul(
+            intermediate_cache2,
+            intermediate_cache1.view(-1, gemm1_n),
+            situ_beta=gemm1_alpha if gemm1_alpha is not None else 4.0,
+            linear_beta=clamp_limit,
+        )
     elif activation == "silu" and not is_gated:
         intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N))
     elif activation == "relu2" and not is_gated:
@@ -305,10 +394,28 @@ def fused_marlin_moe(
         is_zp_float=False,
     ).view(-1, topk, K)
 
-    output = hidden_states if inplace else torch.empty_like(hidden_states)
+    output = zero_copy_context.get_moe_output(hidden_states)
+    if output is None:
+        output = hidden_states if inplace else torch.empty_like(hidden_states)
 
     if is_mxfp4_marlin:
-        return torch.sum(intermediate_cache3, dim=1, out=output)
+        # Top-k weights (incl. routed scaling) are already applied above via
+        # mul_topk_weights, so this is a plain sum over the topk dim. The JIT
+        # vectorized pass (~1.5us at decode shapes) beats sgl_kernel's
+        # moe_sum_reduce_kernel_general (~5.7us) and the generic at::native
+        # reduce_kernel torch.sum dispatches to (~6.7us).
+        if (
+            intermediate_cache3.dtype == torch.bfloat16
+            and intermediate_cache3.is_contiguous()
+            and output.is_contiguous()
+            and intermediate_cache3.shape[-1] % 8 == 0
+        ):
+            from sglang.kernels.ops.moe.moe_topk_sum import moe_topk_sum
+
+            moe_topk_sum(intermediate_cache3, output)
+        else:
+            moe_sum_reduce(intermediate_cache3, output, 1.0)
+        return output
     else:
         if routed_scaling_factor is None:
             routed_scaling_factor = 1.0

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -658,13 +658,32 @@ class FusedMoE(torch.nn.Module):
                 if not is_bias and self.use_triton_kernels:
                     # do not transpose for bias
                     loaded_weight = loaded_weight.transpose(-2, -1)
+                # When the buffer is padded (e.g., MXFP4 SM100 rounds
+                # intermediate up to 256), shard_size from the buffer may
+                # exceed the checkpoint's per-TP slice.  Derive the actual
+                # shard size from the loaded weight so we index correctly.
+                loaded_shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size
                 loaded_weight = loaded_weight.narrow(
-                    shard_dim, shard_size * tp_rank, shard_size
+                    shard_dim, loaded_shard_size * tp_rank, loaded_shard_size
                 )
 
             expert_data = expert_data.narrow(shard_dim, start, shard_size)
+
         loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight)
-        expert_data.copy_(loaded_weight)
+        # loaded_weight may be smaller than expert_data along shard_dim when
+        # the buffer is padded.  Copy into the leading slice and leave the
+        # trailing padding as zeros.  Rank-mismatched tensors (bias / scalar
+        # scales) don't carry shard_dim; they keep the plain broadcast copy.
+        if (
+            loaded_weight.dim() == expert_data.dim()
+            and shard_dim < expert_data.dim()
+            and loaded_weight.shape[shard_dim] < expert_data.shape[shard_dim]
+        ):
+            expert_data.narrow(shard_dim, 0, loaded_weight.shape[shard_dim]).copy_(
+                loaded_weight
+            )
+        else:
+            expert_data.copy_(loaded_weight)
 
     def _load_w2(
         self,
@@ -729,13 +748,28 @@ class FusedMoE(torch.nn.Module):
             if not is_bias and not self.use_presharded_weights:
                 if self.use_triton_kernels:
                     loaded_weight = loaded_weight.transpose(-2, -1)
+                # Derive shard size from the loaded weight so padded buffers
+                # do not cause out-of-bounds indexing into the checkpoint.
+                loaded_shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size
                 loaded_weight = loaded_weight.narrow(
-                    shard_dim, shard_size * tp_rank, shard_size
+                    shard_dim, loaded_shard_size * tp_rank, loaded_shard_size
                 )
 
         # w2, down_proj: Load into only logical weight of w2.
         loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight)
-        expert_data.copy_(loaded_weight)
+        # loaded_weight may be smaller than expert_data along shard_dim when
+        # the buffer is padded.  Copy into the leading slice only.  See the
+        # rank-mismatch note in _load_w13.
+        if (
+            loaded_weight.dim() == expert_data.dim()
+            and shard_dim < expert_data.dim()
+            and loaded_weight.shape[shard_dim] < expert_data.shape[shard_dim]
+        ):
+            expert_data.narrow(shard_dim, 0, loaded_weight.shape[shard_dim]).copy_(
+                loaded_weight
+            )
+        else:
+            expert_data.copy_(loaded_weight)
 
     def _maybe_load_fp8_shared_expert_as_fp4(
         self,
@@ -873,7 +907,7 @@ class FusedMoE(torch.nn.Module):
         # if expert_id is None, then
         # all the experts are loaded at the same time
         if (
-            not expert_id
+            expert_id is None
             and self.quant_config is not None
             and self.quant_config.get_name() == "mxfp4"
             and self.quant_config.is_static_cfg()

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

@@ -91,7 +91,11 @@ class AiterRunnerOutput(RunnerOutput):
         return MoeRunnerBackend.AITER
 
 
-_AITER_ACTIVATIONS = {"silu": "Silu", "swiglu": "Swiglu"}
+_AITER_ACTIVATIONS = {
+    "silu": "Silu",
+    "swiglu": "Swiglu",
+    "situ": "Situv2",
+}
 
 
 def _aiter_activation(activation: str):
@@ -162,7 +166,15 @@ class AiterRunnerCore(MoeRunnerCore):
             extra["num_local_tokens"] = runner_input.num_local_tokens
         if runner_input.output_dtype is not None:
             extra["dtype"] = runner_input.output_dtype
-        if quant_info.swiglu_limit > 0:
+        if self.config.activation == "situ":
+            from aiter.ops.flydsl.moe_common import GateMode
+
+            extra["gate_mode"] = GateMode.SEPARATED.value
+            if self.config.gemm1_alpha is not None:
+                extra["beta"] = float(self.config.gemm1_alpha)
+            if self.config.gemm1_clamp_limit is not None:
+                extra["linear_beta"] = float(self.config.gemm1_clamp_limit)
+        elif quant_info.swiglu_limit > 0:
             # GateMode is only needed for the gpt-oss MXFP4 swiglu_limit path.
             # Import lazily so models that don't use it (e.g. DeepSeek-V3 fp8,
             # swiglu_limit==0) still run on aiter builds where this module

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

@@ -6,6 +6,8 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple
 
 import einops
 import torch
+import triton
+import triton.language as tl
 
 from sglang.kernels.ops.attention.dsv4 import silu_and_mul_masked_post_quant
 from sglang.kernels.ops.quantization import per_token_group_quant
@@ -165,7 +167,9 @@ class DeepGemmMoeQuantInfo(MoeQuantInfo):
 class DeepGemmRunnerCore(MoeRunnerCore):
     def __init__(self, config: MoeRunnerConfig):
         super().__init__(config)
-        assert self.config.activation == "silu"
+        # SiTU (Kimi K3) is applied outside the GEMMs in python, so it only
+        # needs the masked-gemm activation site to branch (see _run_masked_gemm).
+        assert self.config.activation in ("silu", "situ")
         assert self.config.is_gated
         self.swiglu_limit = self.config.swiglu_limit
         self.use_swizzle = False
@@ -256,7 +260,61 @@ class DeepGemmRunnerCore(MoeRunnerCore):
         dispose_tensor(hidden_states)
         dispose_tensor(hidden_states_scale)
 
-        if envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get():
+        if self.config.activation == "situ":
+            situ_beta = self.config.gemm1_alpha
+            situ_linear_beta = self.config.gemm1_clamp_limit
+            assert situ_beta is not None and situ_linear_beta is not None
+            if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+                # Fused SiTU + per-group fp8 quant over the compacted rows,
+                # then the proven round-up e8m0 cast (mn-major packed layout).
+                rows = gateup_output.shape[0]
+                half_n = N // 2
+                kg = half_n // scale_block_size
+                down_input_fp8 = torch.empty(
+                    (rows, half_n),
+                    device=gateup_output.device,
+                    dtype=torch.float8_e4m3fn,
+                )
+                s = torch.empty(
+                    (rows, kg), device=gateup_output.device, dtype=torch.float32
+                )
+                _situ_mul_quant_contig_kernel[(rows,)](
+                    gateup_output,
+                    down_input_fp8,
+                    s,
+                    half_n,
+                    kg,
+                    situ_beta,
+                    situ_linear_beta,
+                    GROUP=scale_block_size,
+                    KG_POW2=triton.next_power_of_2(kg),
+                    num_warps=8,
+                )
+                del gateup_output
+                down_input_scale = _cast_to_e8m0_with_rounding_up(
+                    s.unsqueeze(0)
+                ).squeeze(0)
+            else:
+                from sglang.kernels.ops.quantization.fp8_kernel import (
+                    sglang_per_token_group_quant_fp8,
+                )
+
+                gate = gateup_output[:, : N // 2].float()
+                up = gateup_output[:, N // 2 :].float()
+                gate = situ_beta * torch.tanh(gate / situ_beta) * torch.sigmoid(gate)
+                up = situ_linear_beta * torch.tanh(up / situ_linear_beta)
+                down_input = (gate * up).to(torch.bfloat16)
+                del gateup_output
+
+                down_input_fp8, down_input_scale = sglang_per_token_group_quant_fp8(
+                    down_input,
+                    scale_block_size,
+                    column_major_scales=False,
+                    scale_tma_aligned=False,
+                    scale_ue8m0=False,
+                )
+                del down_input
+        elif envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get():
             swiglu_limit_arg: Optional[float] = self.swiglu_limit
 
             down_input_fp8 = torch.empty(
@@ -519,28 +577,38 @@ class DeepGemmRunnerCore(MoeRunnerCore):
                 )
 
         # Act.
-        topk_ids_rs = running_state.get("topk_ids")
-        num_real_tokens = (
-            topk_ids_rs.shape[0]
-            if (
-                use_mxfp8
-                and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
-                and topk_ids_rs is not None
-                and "src2dst" in running_state
+        if self.config.activation == "situ":
+            down_input, down_input_scale = _varlen_deep_gemm_situ_mul_quant(
+                gateup_output,
+                masked_m,
+                group_size=128,
+                topk=self.config.top_k,
+                beta=self.config.gemm1_alpha,
+                linear_beta=self.config.gemm1_clamp_limit,
+            )
+        else:
+            topk_ids_rs = running_state.get("topk_ids")
+            num_real_tokens = (
+                topk_ids_rs.shape[0]
+                if (
+                    use_mxfp8
+                    and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
+                    and topk_ids_rs is not None
+                    and "src2dst" in running_state
+                )
+                else None
+            )
+            down_input, down_input_scale = _varlen_deep_gemm_silu_mul_quant(
+                gateup_output,
+                masked_m,
+                group_size=scale_block_size,
+                topk=self.config.top_k,
+                swiglu_limit=swiglu_limit_arg,
+                swizzle=self.use_swizzle,
+                gemm1_alpha=self.config.gemm1_alpha,
+                gemm1_clamp_limit=self.config.gemm1_clamp_limit,
+                num_real_tokens=num_real_tokens,
             )
-            else None
-        )
-        down_input, down_input_scale = _varlen_deep_gemm_silu_mul_quant(
-            gateup_output,
-            masked_m,
-            group_size=scale_block_size,
-            topk=self.config.top_k,
-            swiglu_limit=swiglu_limit_arg,
-            swizzle=self.use_swizzle,
-            gemm1_alpha=self.config.gemm1_alpha,
-            gemm1_clamp_limit=self.config.gemm1_clamp_limit,
-            num_real_tokens=num_real_tokens,
-        )
         del gateup_output
 
         # Down activation is quantised locally at scale_block_size (never DeepEP-LL),
@@ -885,10 +953,11 @@ def post_permute_deep_gemm_to_standard(
     hidden_states_shape = running_state["hidden_states_shape"]
     hidden_states_dtype = running_state["hidden_states_dtype"]
     hidden_states_device = running_state["hidden_states_device"]
-    src2dst = running_state["src2dst"]
     topk_ids = running_state["topk_ids"]
     topk_weights = running_state["topk_weights"]
 
+    src2dst = running_state["src2dst"]
+
     with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
         output = torch.empty(
             hidden_states_shape, dtype=hidden_states_dtype, device=hidden_states_device
@@ -975,7 +1044,7 @@ def pre_permute_deepep_normal_to_deep_gemm(
         topk_weights,
         num_recv_tokens_per_expert,
     ) = dispatch_output
-    assert runner_config.activation == "silu"
+    assert runner_config.activation in ("silu", "situ")
 
     all_tokens = sum(num_recv_tokens_per_expert)
     running_state["all_tokens"] = all_tokens
@@ -1090,6 +1159,53 @@ def post_permute_deep_gemm_to_deepep_normal(
     )
 
 
+def _varlen_deep_gemm_situ_mul_quant(
+    gateup_output: torch.Tensor,
+    masked_m: torch.Tensor,
+    group_size: int,
+    topk: int,
+    beta: float,
+    linear_beta: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Fused SiTU activation + per-group fp8 quant via CUDA JIT kernel."""
+    from sglang.kernels.ops.kimi_k3 import situ_and_mul_masked_post_quant
+
+    E, N, D_2 = gateup_output.shape
+    D = D_2 // 2
+    G = D // group_size
+    packed_ue8m0 = deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
+
+    down_input = torch.empty(
+        (E, N, D), device=gateup_output.device, dtype=torch.float8_e4m3fn
+    )
+    if packed_ue8m0:
+        down_input_scale = torch.empty(
+            (E, G // 4, N), device=gateup_output.device, dtype=torch.int32
+        )
+    else:
+        down_input_scale = torch.empty(
+            (E, N, G), device=gateup_output.device, dtype=torch.float32
+        )
+
+    situ_and_mul_masked_post_quant(
+        gateup_output,
+        down_input,
+        down_input_scale,
+        group_size,
+        masked_m,
+        beta=beta,
+        linear_beta=linear_beta,
+        scale_ue8m0=packed_ue8m0,
+        topk=topk,
+        transposed=packed_ue8m0,
+    )
+
+    if packed_ue8m0:
+        down_input_scale = down_input_scale.transpose(-1, -2)
+
+    return down_input, down_input_scale
+
+
 def _varlen_deep_gemm_silu_mul_quant(
     gateup_output: torch.Tensor,
     masked_m: Optional[torch.Tensor],
@@ -1202,6 +1318,37 @@ def _varlen_deep_gemm_silu_mul_quant(
     )
 
 
+@triton.jit
+def _situ_mul_quant_contig_kernel(
+    g_ptr,  # [rows, 2N] bf16, non-interleaved [gate; up] halves
+    q_ptr,  # [rows, N] fp8 out
+    s_ptr,  # [rows, KG] fp32 scales out
+    N,
+    KG,
+    situ_beta,
+    situ_linear_beta,
+    GROUP: tl.constexpr,
+    KG_POW2: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    rows2d = tl.arange(0, KG_POW2)[:, None]
+    cols = tl.arange(0, GROUP)[None, :]
+    offs = rows2d * GROUP + cols
+    mask = rows2d < KG
+    gate = tl.load(g_ptr + row * 2 * N + offs, mask=mask, other=0.0).to(tl.float32)
+    up = tl.load(g_ptr + row * 2 * N + N + offs, mask=mask, other=0.0).to(tl.float32)
+    # tanh(x) == 2*sigmoid(2x) - 1 (avoids a libdevice dependency)
+    gate_t = 2.0 * tl.sigmoid(2.0 * gate / situ_beta) - 1.0
+    gate = situ_beta * gate_t * tl.sigmoid(gate)
+    up_t = 2.0 * tl.sigmoid(2.0 * up / situ_linear_beta) - 1.0
+    y = gate * situ_linear_beta * up_t
+    amax = tl.clamp(tl.max(tl.abs(y), axis=1), min=1e-10, max=float("inf"))
+    q = (y * (448.0 / amax)[:, None]).to(tl.float8e4nv)
+    tl.store(q_ptr + row * N + offs, q, mask=mask)
+    srow = tl.arange(0, KG_POW2)
+    tl.store(s_ptr + row * KG + srow, amax / 448.0, mask=srow < KG)
+
+
 def _apply_swiglu_limit(
     gateup_output: torch.Tensor, swiglu_limit: float
 ) -> torch.Tensor:

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

@@ -140,7 +140,10 @@ def fused_experts_none_to_marlin(
         )
 
     if runner_config.is_gated:
-        assert runner_config.activation == "silu", "Only gated SiLU is supported."
+        assert runner_config.activation in {
+            "silu",
+            "situ",
+        }, f"Only gated SiLU/SiTU is supported, got {runner_config.activation}."
     elif runner_config.activation not in {"silu", "relu2"}:
         raise ValueError(
             f"Unsupported Marlin MoE activation: {runner_config.activation}"

python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py

@@ -685,6 +685,17 @@ def _fused_moe_kernel_sequence(
                 x = intermediate_cache1.view(-1, N)
                 d = x.shape[-1] // 2
                 intermediate_cache2.copy_(F.silu(x[..., :d]) * x[..., d:])
+    elif activation == "situ" and is_gated:
+        d = N // 2
+        x = intermediate_cache1.view(-1, N)
+        gate = x[..., :d].float()
+        up = x[..., d:].float()
+        situ_beta = gemm1_alpha if gemm1_alpha is not None else 4.0
+        gate = situ_beta * torch.tanh(gate / situ_beta) * torch.sigmoid(gate)
+        situ_linear_beta = gemm1_limit
+        if situ_linear_beta is not None:
+            up = situ_linear_beta * torch.tanh(up / situ_linear_beta)
+        intermediate_cache2.copy_((gate * up).to(intermediate_cache1.dtype))
     elif activation == "gelu" and is_gated:
         assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu"
         assert gemm1_limit is None, "gemm1_limit is not supported for gelu"

python/sglang/srt/layers/moe/route_quant_handoff.py

@@ -0,0 +1,130 @@
+"""Attempt-and-verify handoff for the fused K3 MoE-front prep launch.
+
+At decode batch sizes the chain between the K3 fused-front GEMM and the
+trtllm-gen SiTU MoE op is three tiny back-to-back kernels on the critical
+path — route_radix (top-16 of 896), the triton ``(id << 16) | bf16(weight)``
+pack, and per_token_group_quant (mxfp8, ``[T, 3584]``) — about 7.5 us busy
+plus two extra launches per MoE layer, each leaving the SMs near idle. The
+fused kernel (kernels/ops/moe/moe_route_quant_fused.py) runs all three in one
+launch: routing CTAs and one quant CTA per token, concurrently.
+
+The inputs live in different modules (router logits reach the router through
+TopK, activations through the MoE runner), so the fusion is wired as a
+consume-once stash instead of new signatures:
+
+    KimiK3MoE._forward_routed*      stage(x) before self.topk, clear() after
+                                    the experts call
+    biased_grouped_topk_gpu         try_route_quant_fused() replaces the
+                                    moe_fused_gate call on a hit
+    Mxfp4MoEMethod.apply (situ)     take(x) skips the quant and the pack
+
+Every step is fallback-safe: if the routing dispatch never consumes the staged
+request (different model, uncovered shape, triton fallback) or the runner's
+``take`` misses (activations re-viewed or copied), the unfused chain runs as
+before. The stage/clear bracket in the model layer guarantees a published
+entry can never leak into another layer whose allocator reused the same
+activation address.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import msgspec
+import torch
+
+
+class _Handoff(msgspec.Struct):
+    # staged by the model layer: the activation rows the runner will quantize
+    request_x: Optional[torch.Tensor] = None
+    # published by the routing dispatch, keyed by the staged activations
+    produced_x: Optional[torch.Tensor] = None
+    packed: Optional[torch.Tensor] = None
+    x_q: Optional[torch.Tensor] = None
+    x_s: Optional[torch.Tensor] = None
+
+
+_handoff = _Handoff()
+
+
+def stage(x: torch.Tensor) -> None:
+    """Publish the routed activations for the upcoming topk call. Caller pairs
+    this with clear() after the experts call (try/finally)."""
+    _handoff.request_x = x
+    _handoff.produced_x = None
+
+
+def clear() -> None:
+    _handoff.request_x = None
+    _handoff.produced_x = None
+    _handoff.packed = None
+    _handoff.x_q = None
+    _handoff.x_s = None
+
+
+def try_route_quant_fused(
+    gating_output: torch.Tensor,
+    correction_bias: torch.Tensor,
+    topk: int,
+    *,
+    num_fused_shared_experts: int,
+    renormalize: bool,
+    routed_scaling_factor: Optional[float],
+    apply_routed_scaling_factor_on_output: bool,
+) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+    """Fused replacement for the ungrouped-sigmoid moe_fused_gate call when a
+    staged request covers it. Returns (weights, ids) on a hit, None otherwise
+    (caller falls through to the unfused router)."""
+    x = _handoff.request_x
+    if x is None or num_fused_shared_experts != 0:
+        return None
+
+    from sglang.kernels.ops.moe import moe_route_quant_fused
+
+    if (
+        not moe_route_quant_fused.covered(gating_output, correction_bias, topk, x)
+        or not moe_route_quant_fused.available()
+    ):
+        return None
+
+    weights, ids, packed, x_q, x_s = moe_route_quant_fused.route_quant_fused(
+        gating_output,
+        correction_bias,
+        x,
+        topk,
+        renormalize=renormalize,
+        routed_scaling_factor=(
+            routed_scaling_factor if routed_scaling_factor is not None else 1.0
+        ),
+        apply_scale=apply_routed_scaling_factor_on_output,
+    )
+    _handoff.request_x = None
+    _handoff.produced_x = x
+    _handoff.packed = packed
+    _handoff.x_q = x_q
+    _handoff.x_s = x_s
+    return weights, ids
+
+
+def take(
+    x: torch.Tensor,
+) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    """Consume the published (packed_topk, x_q, x_s int32) for these exact
+    activation rows, or None. Storage identity is verified so a re-viewed or
+    copied tensor simply misses."""
+    produced = _handoff.produced_x
+    if produced is None:
+        return None
+    if (
+        produced.data_ptr() != x.data_ptr()
+        or produced.shape != x.shape
+        or produced.dtype != x.dtype
+        or produced.stride() != x.stride()
+    ):
+        return None
+    out = (_handoff.packed, _handoff.x_q, _handoff.x_s)
+    _handoff.produced_x = None
+    _handoff.packed = None
+    _handoff.x_q = None
+    _handoff.x_s = None
+    return out

python/sglang/srt/layers/moe/topk.py

@@ -1574,7 +1574,22 @@ def biased_grouped_topk_gpu(
         return topk_weights, topk_ids
     else:
         num_experts = gating_output.shape[1]
-        if _is_cuda and num_experts == 384 and num_expert_group == 1:
+        # The JIT triton router (single fused kernel: scoring + bias + topk +
+        # renorm) handles the ungrouped case with arbitrary num_experts/topk.
+        # Original user: Kimi K2 (384 experts). Also dispatch shapes the other
+        # fused kernels cannot cover, e.g. Kimi K3 (896 experts, top-16):
+        # fused_topk_deepseek needs pow2 experts + topk<=8, jit_grouped_topk
+        # needs experts<=512 + topk<=8.
+        _jit_gate_ok = (
+            _is_cuda
+            and num_expert_group == 1
+            and (topk_group is None or topk_group == 1)
+            and (
+                num_experts == 384
+                or (num_experts <= 1024 and (num_experts > 512 or topk > 8))
+            )
+        )
+        if _jit_gate_ok:
             # ===== TO BE REFACTORED ====
             _use_jit_bf16_gate = False
             if _SGLANG_EXPERIMENTAL_LORA_OPTI:
@@ -1601,9 +1616,36 @@ def biased_grouped_topk_gpu(
             # ===== END TO BE REFACTORED ====
             from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate as jit_gate
 
+            # Pass BF16 logits through untouched for the automatic radix-select
+            # fast path. BF16 -> FP32 is exact, and unsupported shapes fall back
+            # inside moe_fused_gate.
+            _gating = (
+                gating_output
+                if gating_output.dtype == torch.bfloat16
+                else gating_output.to(dtype=torch.float32)
+            )
+            # K3 staged fusion: when the model layer staged the routed
+            # activations (route_quant_handoff), the radix route, the trtllm id
+            # pack and the mxfp8 quant run as one launch. Bit-identical
+            # (weights, ids); a miss falls through to the unfused router.
+            from sglang.srt.layers.moe import route_quant_handoff
+
+            fused = route_quant_handoff.try_route_quant_fused(
+                _gating,
+                correction_bias.to(dtype=torch.float32),
+                topk,
+                num_fused_shared_experts=num_fused_shared_experts,
+                renormalize=renormalize,
+                routed_scaling_factor=routed_scaling_factor,
+                apply_routed_scaling_factor_on_output=bool(
+                    apply_routed_scaling_factor_on_output
+                ),
+            )
+            if fused is not None:
+                return fused
             return jit_gate(
-                gating_output.to(dtype=torch.float32),
-                correction_bias,
+                _gating,
+                correction_bias.to(dtype=torch.float32),
                 topk=topk,
                 scoring_func="sigmoid",
                 num_fused_shared_experts=num_fused_shared_experts,
@@ -2255,6 +2297,49 @@ def select_experts(
     return StandardTopKOutput(topk_weights, topk_ids, router_logits)
 
 
+def precomputed_topk_postprocess_is_noop(
+    topk_config: TopKConfig,
+    num_token_non_padded: Optional[torch.Tensor] = None,
+    expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
+) -> bool:
+    """Whether :func:`build_precomputed_topk_output` can stand in for
+    :func:`select_experts`' post-processing.
+
+    A router that produces (weights, ids) itself -- e.g. K3's fused gate+top-k
+    kernel -- skips select_experts entirely, so it must not skip the work
+    select_experts does *after* the top-k: the EPLB logical->physical remap, the
+    padded-region mask, and the shared-expert append.  This returns True only
+    when all of those reduce to no-ops, leaving just the capture hook and the
+    distribution recorder (which the builder below still runs).
+    """
+    return (
+        _is_cuda
+        and topk_config.num_fused_shared_experts == 0
+        and num_token_non_padded is None
+        and expert_location_dispatch_info is None
+        and not envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get()
+        and not envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get()
+    )
+
+
+def build_precomputed_topk_output(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk_config: TopKConfig,
+    layer_id: int,
+) -> StandardTopKOutput:
+    """Wrap a router's own (weights, ids) as a STANDARD top-k output, running the
+    capture hook and the expert-distribution recorder that select_experts would.
+
+    Only valid when :func:`precomputed_topk_postprocess_is_noop` holds.
+    """
+    capture_routed_experts_if_allowed(topk_config, layer_id, topk_ids)
+    get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
+    # router_logits is only read by the BYPASSED formats and by the
+    # shared-expert append (excluded above); STANDARD consumers take ids/weights.
+    return StandardTopKOutput(topk_weights, topk_ids, None)
+
+
 # NOTE: the AOT sgl_kernel::moe_fused_gate and sgl_kernel::kimi_k2_moe_fused_gate
 # ops (and their torch.compile fake impls) were retired here — both CUDA gate
 # paths now route through the unified Triton router (kernels/ops/moe/moe_fused_gate.py),

python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py

@@ -178,6 +178,17 @@ class CompressedTensorsConfig(QuantizationConfig):
         from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
 
         if isinstance(layer, FusedMoE):
+            # Detect MXFP4 before the scheme-based path: MXFP4 uses a
+            # dedicated FusedMoEMethodBase (Mxfp4MoEMethod) that already
+            # handles all MoE backends, bypassing the scheme abstraction.
+            if self._is_mxfp4_moe(layer_name=prefix):
+                from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
+
+                logger.info_once(
+                    "Using Mxfp4MoEMethod for MXFP4 compressed-tensors MoE"
+                )
+                return Mxfp4MoEMethod(prefix=prefix)
+
             layer.scheme = self.get_moe_scheme(layer=layer, layer_name=prefix)
             if layer.scheme is None:  # ignored layer
                 use_triton_kernels = get_moe_runner_backend().is_triton_kernels()
@@ -559,6 +570,28 @@ class CompressedTensorsConfig(QuantizationConfig):
 
         return is_mxint4 and input_quant_none and is_symmetric and is_static
 
+    def _is_mxfp4_moe(self, layer_name: str) -> bool:
+        """Detect MXFP4-quantized MoE from global format or target scheme."""
+        if "mxfp4" in (self.quant_format or ""):
+            return True
+        self._add_fused_moe_to_target_scheme_map()
+        for key in ["FusedMoE", "Linear"]:
+            scheme = self.target_scheme_map.get(key)
+            if scheme is None:
+                continue
+            wq = scheme.get("weights")
+            if wq is None:
+                continue
+            if (
+                wq.num_bits == 4
+                and wq.type == QuantizationType.FLOAT
+                and wq.strategy == QuantizationStrategy.GROUP.value
+                and wq.group_size == 32
+                and wq.symmetric
+            ):
+                return True
+        return False
+
     def _is_dynamic_token_w4(
         self, weight_quant: BaseModel, input_quant: BaseModel
     ) -> bool:

python/sglang/srt/layers/quantization/marlin_utils.py

@@ -250,7 +250,7 @@ def check_moe_marlin_supports_layer(
     # apply_router_weight_on_input is not supported for moe marlin
     supports_router_weight = not layer.moe_runner_config.apply_router_weight_on_input
     if layer.moe_runner_config.is_gated:
-        supports_activation = layer.moe_runner_config.activation == "silu"
+        supports_activation = layer.moe_runner_config.activation in {"silu", "situ"}
     else:
         supports_activation = layer.moe_runner_config.activation in {
             "silu",

python/sglang/srt/layers/quantization/mxfp4.py

@@ -33,13 +33,13 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
 )
 from sglang.srt.environ import envs
+from sglang.srt.layers import zero_copy_context
 from sglang.srt.layers.amx_utils import (
     CPUQuantMethod,
     _amx_process_weight_after_loading,
 )
 from sglang.srt.layers.dp_attention import is_allocation_symmetric
 from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
-from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
 from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend, get_moe_runner_backend
 from sglang.srt.layers.quantization.base_config import (
@@ -145,6 +145,7 @@ if TYPE_CHECKING:
 _is_cpu = is_cpu()
 _is_hip = is_hip()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+_aiter_k3_opt = _use_aiter and get_bool_env_var("SGLANG_AITER_K3_OPT")
 _is_shuffle_moe_mxfp4 = is_gfx95_supported()
 _is_cpu_amx_available = cpu_has_amx_support()
 
@@ -332,6 +333,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         self.with_bias = False
         self.use_flashinfer = get_moe_runner_backend().is_flashinfer_mxfp4()
         self.use_marlin = get_moe_runner_backend().is_marlin()
+        # True W4A8: DeepGEMM fp8_fp4 grouped GEMM (SM100 MXF8F6F4 UMMA).
+        # Weights stay MXFP4 (e2m1 + ue8m0 g32, zero requantization);
+        # activations are quantized to fp8 per-token-group-128.
+        self.use_deep_gemm = get_moe_runner_backend().is_deep_gemm()
         self.flashinfer_mxfp4_moe_precision = (
             get_exec().moe.flashinfer_mxfp4_moe_precision
         )
@@ -392,12 +397,21 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 intermediate_size_per_partition_after_pad
                 - layer.intermediate_size_per_partition
             )
+        elif self.use_deep_gemm:
+            # DeepGEMM fp8_fp4 grouped GEMM consumes the checkpoint layout
+            # directly (packed e2m1 K-major + ue8m0 g32 scales); no padding.
+            pass
         elif is_sm100_supported():
             if self.use_flashinfer:
+                # FlashInfer trtllm-gen FP4 kernel actual alignment:
+                # intermediate: scale shuffle needs M%128==0 → intermediate%64==0
+                # hidden: finalize kernel needs K%32==0, block quant needs K%32==0
+                # Using 128 alignment (not 256) since 256 was overly conservative
+                # and causes 60GB/GPU waste for models like K3 (384 intermediate).
                 intermediate_size_per_partition_after_pad = round_up(
-                    intermediate_size_per_partition, 256
+                    intermediate_size_per_partition, 128
                 )
-                hidden_size = round_up(hidden_size, 256)
+                hidden_size = round_up(hidden_size, 128)
             else:
                 intermediate_size_per_partition_after_pad = round_up(
                     intermediate_size_per_partition, triton_kernels_padding_alignment
@@ -422,9 +436,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             # naive-copy fast path is correct.
             intermediate_size_per_partition_after_pad = intermediate_size_per_partition
         elif _use_aiter:
-
+            # Expert intermediate is padded to 128 or 256. On K3 (TP8) the 256
+            # default costs +33% mem (3072/8=384 -> 512). K3 FlyDSL MoE works with
+            # 128 padding, so relaxing to 128 saves that memory and lets the TP8
+            # decode CUDA graph range grow from 20 to 120.
+            _inter_align = 128 if _aiter_k3_opt else 256
             intermediate_size_per_partition_after_pad = round_up(
-                intermediate_size_per_partition, 256
+                intermediate_size_per_partition, _inter_align
             )
 
             hidden_size = round_up(hidden_size, 256)
@@ -465,17 +483,20 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         )
         layer.register_parameter("w13_weight_scale", w13_weight_scale)
         set_weight_attrs(w13_weight_scale, extra_weight_attrs)
+        w13_weight_scale.quant_method = "group"
 
-        w13_weight_bias = torch.nn.Parameter(
-            torch.zeros(
-                layer.num_local_experts,
-                2 * intermediate_size_per_partition_after_pad,
-                dtype=torch.bfloat16,
-            ),
-            requires_grad=False,
-        )
-        layer.register_parameter("w13_weight_bias", w13_weight_bias)
-        set_weight_attrs(w13_weight_bias, extra_weight_attrs)
+        create_bias = with_bias or not _is_hip
+        if create_bias:
+            w13_weight_bias = torch.nn.Parameter(
+                torch.zeros(
+                    layer.num_local_experts,
+                    2 * intermediate_size_per_partition_after_pad,
+                    dtype=torch.bfloat16,
+                ),
+                requires_grad=False,
+            )
+            layer.register_parameter("w13_weight_bias", w13_weight_bias)
+            set_weight_attrs(w13_weight_bias, extra_weight_attrs)
 
         # down_proj (row parallel)
         w2_weight = torch.nn.Parameter(
@@ -501,13 +522,15 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         )
         layer.register_parameter("w2_weight_scale", w2_weight_scale)
         set_weight_attrs(w2_weight_scale, extra_weight_attrs)
+        w2_weight_scale.quant_method = "group"
 
-        w2_weight_bias = torch.nn.Parameter(
-            torch.zeros(layer.num_local_experts, hidden_size, dtype=torch.bfloat16),
-            requires_grad=False,
-        )
-        layer.register_parameter("w2_weight_bias", w2_weight_bias)
-        set_weight_attrs(w2_weight_bias, extra_weight_attrs)
+        if create_bias:
+            w2_weight_bias = torch.nn.Parameter(
+                torch.zeros(layer.num_local_experts, hidden_size, dtype=torch.bfloat16),
+                requires_grad=False,
+            )
+            layer.register_parameter("w2_weight_bias", w2_weight_bias)
+            set_weight_attrs(w2_weight_bias, extra_weight_attrs)
 
     def process_weights_after_loading(self, layer):
         if self.use_marlin:
@@ -519,19 +542,73 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 prepare_moe_mxfp4_layer_for_marlin,
             )
 
-            if not is_sm90_supported() and not is_sm120_supported():
-                raise RuntimeError("MXFP4 Marlin requires SM90 or SM120.")
+            if (
+                not is_sm90_supported()
+                and not is_sm100_supported()
+                and not is_sm120_supported()
+            ):
+                raise RuntimeError("MXFP4 Marlin requires SM90+.")
             if not check_moe_marlin_supports_layer(layer, 32, allow_tile_padding=True):
                 raise RuntimeError(
                     "Current MXFP4 MoE layer is not supported by Marlin."
                 )
 
-            if self.moe_runner_config.gemm1_alpha is not None:
+            if self.moe_runner_config.gemm1_alpha is not None and getattr(
+                self.moe_runner_config, "gate_up_interleaved", True
+            ):
                 deinterleave_moe_mxfp4_w13_for_marlin(layer)
             prepare_moe_mxfp4_layer_for_marlin(layer)
             layer._mxfp4_backend = "marlin"
             return
 
+        if self.use_deep_gemm:
+            from deep_gemm import transform_sf_into_required_layout
+
+            # Packed fp4 (e2m1 x2 per byte) weights: DeepGEMM expects int8.
+            layer.w13_weight.data = layer.w13_weight.data.view(torch.int8)
+            layer.w2_weight.data = layer.w2_weight.data.view(torch.int8)
+            # Checkpoint scales are uint8 e8m0 (biased exponents). DeepGEMM
+            # SM100 needs them in packed-UE8M0 TMA-aligned MN-major layout.
+            # Round-trip through fp32 is exact (values are powers of two).
+            for scale_name, weight in (
+                ("w13_weight_scale", layer.w13_weight),
+                ("w2_weight_scale", layer.w2_weight),
+            ):
+                scale = getattr(layer, scale_name)
+                num_experts, n, _ = scale.data.shape
+                k = weight.shape[2] * 2
+                scale_f32 = scale.data.view(torch.float8_e8m0fnu).to(torch.float32)
+                scale.data = transform_sf_into_required_layout(
+                    scale_f32,
+                    mn=n,
+                    k=k,
+                    recipe=(1, 32),
+                    num_groups=num_experts,
+                    disable_ue8m0_cast=False,
+                )
+            if get_moe_a2a_backend().is_megamoe():
+                # MegaMoE consumes the same transformed sf, plus its own
+                # interleaved/UTCCP weight layout. K3 routes EVERY batch
+                # through mega (the megamoe backend has no a2a fallback), so
+                # the contig-layout originals are dead weight — repoint the
+                # params at the mega tensors to reclaim ~1.9GB/layer/rank
+                # (keeping both layouts OOMs: 92 layers double the experts).
+                from deep_gemm import transform_weights_for_mega_moe
+
+                l1_pair, l2_pair = transform_weights_for_mega_moe(
+                    (layer.w13_weight.data, layer.w13_weight_scale.data),
+                    (layer.w2_weight.data, layer.w2_weight_scale.data),
+                )
+                layer.mega_l1_weights = l1_pair
+                layer.mega_l2_weights = l2_pair
+                layer.w13_weight.data = l1_pair[0]
+                layer.w13_weight_scale.data = l1_pair[1]
+                layer.w2_weight.data = l2_pair[0]
+                layer.w2_weight_scale.data = l2_pair[1]
+                layer._mega_moe_weights_built = True
+            layer._mxfp4_backend = "deep_gemm"
+            return
+
         if self._fi_kernel == "cutlass_sm90":
             self._process_weights_for_sm90_cutlass(layer)
             return
@@ -539,38 +616,42 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             self._process_weights_for_sm120_cutlass(layer)
             return
         if self.use_flashinfer:
-            # TODO: these values are hardcoded for now, we need to get them from the model
+            # Per-expert buffers are local (create_weights uses num_local_experts);
+            # the global self.num_experts here breaks EP>1. Mirrors the SM90 path.
+            E = layer.num_local_experts
+            _alpha = getattr(layer.moe_runner_config, "gemm1_alpha", None) or 1.702
+            _limit = getattr(layer.moe_runner_config, "gemm1_clamp_limit", None) or 7.0
             layer.gemm1_alpha = Parameter(
-                torch.tensor([1.702] * self.num_experts, dtype=torch.float32).cuda(),
+                torch.tensor([_alpha] * E, dtype=torch.float32).cuda(),
                 requires_grad=False,
             )
             layer.gemm1_beta = Parameter(
-                torch.tensor([1.0] * self.num_experts, dtype=torch.float32).cuda(),
+                torch.tensor([1.0] * E, dtype=torch.float32).cuda(),
                 requires_grad=False,
             )
             layer.gemm1_clamp_limit = Parameter(
-                torch.tensor([7.0] * self.num_experts, dtype=torch.float32).cuda(),
+                torch.tensor([_limit] * E, dtype=torch.float32).cuda(),
                 requires_grad=False,
             )
             sf_block_size = 32  # mxfp4 block size
 
             assert (
                 layer.w13_weight.dim() == 3
-                and layer.w13_weight.shape[0] == self.num_experts
+                and layer.w13_weight.shape[0] == E
                 and layer.w13_weight.shape[1]
                 == self.intermediate_size_per_partition * 2
                 and layer.w13_weight.shape[2] == self.hidden_size // 2
             )
             assert (
                 layer.w13_weight_scale.dim() == 3
-                and layer.w13_weight_scale.shape[0] == self.num_experts
+                and layer.w13_weight_scale.shape[0] == E
                 and layer.w13_weight_scale.shape[1]
                 == self.intermediate_size_per_partition * 2
                 and layer.w13_weight_scale.shape[2] == self.hidden_size // sf_block_size
             )
             assert (
                 layer.w2_weight.dim() == 3
-                and layer.w2_weight.shape[0] == self.num_experts
+                and layer.w2_weight.shape[0] == E
                 and layer.w2_weight.shape[1] == self.hidden_size
                 and layer.w2_weight.shape[2]
                 == self.intermediate_size_per_partition // 2
@@ -583,13 +664,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             )
             assert (
                 layer.w13_weight_bias.dim() == 2
-                and layer.w13_weight_bias.shape[0] == self.num_experts
+                and layer.w13_weight_bias.shape[0] == E
                 and layer.w13_weight_bias.shape[1]
                 == self.intermediate_size_per_partition * 2
             )
             assert (
                 layer.w2_weight_bias.dim() == 2
-                and layer.w2_weight_bias.shape[0] == self.num_experts
+                and layer.w2_weight_bias.shape[0] == E
                 and layer.w2_weight_bias.shape[1] == self.hidden_size
             )
 
@@ -618,9 +699,26 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 new_shape = list(shape)
                 return x.reshape(*new_shape)
 
-            w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2)
-            w13_weight = swap_every_two_rows(w13_weight, -2)
-            w13_bias = swap_every_two_rows(w13_bias, -1)
+            if getattr(layer.moe_runner_config, "gate_up_interleaved", True):
+                w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2)
+                w13_weight = swap_every_two_rows(w13_weight, -2)
+                w13_bias = swap_every_two_rows(w13_bias, -1)
+            else:
+                # Non-interleaved layout (e.g. K3 Latent MoE): first half = gate
+                # (w1), second half = up (w3). The trtllm-gen fused gated-act
+                # epilogue wants rows interleaved as (up_i, gate_i) PAIRS -
+                # the layout get_reorder_rows_for_gated_act_gemm_row_indices
+                # produces from [linear; gate] halves, and what the
+                # interleaved branch above arrives at via swap_every_two_rows.
+                # A plain halves swap keeps rows blocked and pairs up_i with
+                # up_{i+1}, which scrambles the gated activation.
+                half = w13_weight.shape[-2] // 2
+                pair_idx = torch.empty(2 * half, dtype=torch.long)
+                pair_idx[0::2] = torch.arange(half) + half  # up (w3)
+                pair_idx[1::2] = torch.arange(half)  # gate (w1)
+                w13_weight = w13_weight[..., pair_idx, :].contiguous()
+                w13_weight_scale = w13_weight_scale[..., pair_idx, :].contiguous()
+                w13_bias = w13_bias[..., pair_idx].contiguous()
 
             # Shuffle weights and scaling factors for transposed mma output
             gemm1_weights_mxfp4_shuffled = []
@@ -658,7 +756,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 epilogue_tile_m,
             )
 
-            for i in range(self.num_experts):
+            for i in range(E):
                 gemm1_weights_mxfp4_shuffled.append(
                     w13_weight[i]
                     .view(torch.uint8)[w13_weight_permute_indices]
@@ -699,7 +797,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             w13_weight_scale = (
                 torch.stack(gemm1_scales_mxfp4_shuffled)
                 .reshape(
-                    self.num_experts,
+                    E,
                     2 * self.intermediate_size_per_partition,
                     self.hidden_size // sf_block_size,
                 )
@@ -710,7 +808,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             w2_weight_scale = (
                 torch.stack(gemm2_scales_mxfp4_shuffled)
                 .reshape(
-                    self.num_experts,
+                    E,
                     self.hidden_size,
                     self.intermediate_size_per_partition // sf_block_size,
                 )
@@ -722,44 +820,57 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             layer.w2_weight = Parameter(w2_weight, requires_grad=False)
             layer.w2_weight_scale = Parameter(w2_weight_scale, requires_grad=False)
             layer.w13_weight_bias = Parameter(
-                torch.stack(gemm1_bias_shuffled).reshape(self.num_experts, -1),
+                torch.stack(gemm1_bias_shuffled).reshape(E, -1),
                 requires_grad=False,
             )
             layer.w2_weight_bias = Parameter(
-                torch.stack(gemm2_bias_shuffled).reshape(self.num_experts, -1),
+                torch.stack(gemm2_bias_shuffled).reshape(E, -1),
                 requires_grad=False,
             )
             return
         if _use_aiter:
-            if layer.w13_weight_bias is not None:
+            if getattr(layer, "w13_weight_bias", None) is not None:
                 layer.w13_weight_bias.data = layer.w13_weight_bias.data.to(
                     torch.float32
                 )
-            if layer.w2_weight_bias is not None:
+            if getattr(layer, "w2_weight_bias", None) is not None:
                 layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32)
 
             e, n, k = layer.w13_weight.shape
-            layer.w13_weight.view(torch.uint8).copy_(
-                layer.w13_weight.data.view(torch.uint8)
-                .view(e, n // 2, 2, k)
-                .permute(0, 2, 1, 3)
-                .contiguous()
-                .view(e, n, k)
-            )
-            layer.w13_weight_scale.data = (
-                layer.w13_weight_scale.data.view(e, n // 2, 2, -1)
-                .permute(0, 2, 1, 3)
-                .contiguous()
-                .view(e, n, -1)
-            )
-            layer.w13_weight_bias.data = (
-                layer.w13_weight_bias.data.view(-1, n // 2, 2)
-                .permute(0, 2, 1)
-                .contiguous()
-                .view(-1, n)
+            gate_up_interleaved = getattr(
+                layer.moe_runner_config, "gate_up_interleaved", True
             )
 
-            if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get():
+            if gate_up_interleaved:
+                layer.w13_weight.view(torch.uint8).copy_(
+                    layer.w13_weight.data.view(torch.uint8)
+                    .view(e, n // 2, 2, k)
+                    .permute(0, 2, 1, 3)
+                    .contiguous()
+                    .view(e, n, k)
+                )
+                layer.w13_weight_scale.data = (
+                    layer.w13_weight_scale.data.view(e, n // 2, 2, -1)
+                    .permute(0, 2, 1, 3)
+                    .contiguous()
+                    .view(e, n, -1)
+                )
+                if getattr(layer, "w13_weight_bias", None) is not None:
+                    layer.w13_weight_bias.data = (
+                        layer.w13_weight_bias.data.view(-1, n // 2, 2)
+                        .permute(0, 2, 1)
+                        .contiguous()
+                        .view(-1, n)
+                    )
+
+            k3_situ_a8w4 = (
+                os.environ.get("AITER_SITUV2_A8W4", "0") == "1"
+                and getattr(layer.moe_runner_config, "activation", None) == "situ"
+            )
+            use_aiter_gu_interleave = k3_situ_a8w4 or (
+                envs.SGLANG_USE_AITER_MOE_GU_ITLV.get() and gate_up_interleaved
+            )
+            if use_aiter_gu_interleave:
                 layer.w13_weight.data = shuffle_weight_a16w4(layer.w13_weight, 16, True)
                 shuffled_w13_scale = shuffle_scale_a16w4(
                     layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
@@ -925,11 +1036,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         # half along its row dim (N) from N_un to N_pad with zeros, and along
         # its last dim (K) from K_un (or K_un / sf_block_size) to K_pad.
 
+        _interleaved = getattr(layer.moe_runner_config, "gate_up_interleaved", True)
+
         def _stack_up_gate_w13(unpadded_w13, last_pad, last_un):
             # unpadded_w13: [E, 2*N_un, last_un]
             # Returns: [E, 2*N_pad, last_pad] in [up_padded; gate_padded] order.
-            gate_rows = unpadded_w13[:, 0::2, :]  # [E, N_un, last_un]
-            up_rows = unpadded_w13[:, 1::2, :]  # [E, N_un, last_un]
+            if _interleaved:
+                gate_rows = unpadded_w13[:, 0::2, :]  # [E, N_un, last_un]
+                up_rows = unpadded_w13[:, 1::2, :]  # [E, N_un, last_un]
+            else:
+                # Non-interleaved: first half=gate, second half=up
+                gate_rows = unpadded_w13[:, :N_un, :]
+                up_rows = unpadded_w13[:, N_un:, :]
             out = torch.zeros(
                 E, 2 * N_pad, last_pad, dtype=unpadded_w13.dtype, device=device
             )
@@ -948,8 +1066,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             K_un // sf_block_size,
         )
         # Bias: same de-interleave on dim=-1.
-        w13_bias_gate = layer.w13_weight_bias.data[:, 0::2]  # [E, N_un]
-        w13_bias_up = layer.w13_weight_bias.data[:, 1::2]  # [E, N_un]
+        if _interleaved:
+            w13_bias_gate = layer.w13_weight_bias.data[:, 0::2]  # [E, N_un]
+            w13_bias_up = layer.w13_weight_bias.data[:, 1::2]  # [E, N_un]
+        else:
+            w13_bias_gate = layer.w13_weight_bias.data[:, :N_un]
+            w13_bias_up = layer.w13_weight_bias.data[:, N_un:]
         w13_bias_padded = torch.zeros(E, 2 * N_pad, dtype=bias_dtype, device=device)
         w13_bias_padded[:, :N_un] = w13_bias_up
         w13_bias_padded[:, N_pad : N_pad + N_un] = w13_bias_gate
@@ -971,9 +1093,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         w2_bias_padded = torch.zeros(E, K_pad, dtype=bias_dtype, device=device)
         w2_bias_padded[:, :K_un] = layer.w2_weight_bias.data
 
-        # ---- Per-expert SwiGLU scalars (GPT-OSS defaults) ------------------
+        # ---- Per-expert SwiGLU scalars (read from runner config, fallback to GPT-OSS defaults)
+        _sm90_alpha = getattr(layer.moe_runner_config, "gemm1_alpha", None) or 1.702
+        _sm90_limit = getattr(layer.moe_runner_config, "gemm1_clamp_limit", None) or 7.0
         layer.swiglu_alpha = Parameter(
-            torch.full((E,), 1.702, dtype=torch.float32, device=device),
+            torch.full((E,), _sm90_alpha, dtype=torch.float32, device=device),
             requires_grad=False,
         )
         layer.swiglu_beta = Parameter(
@@ -981,7 +1105,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             requires_grad=False,
         )
         layer.swiglu_limit = Parameter(
-            torch.full((E,), 7.0, dtype=torch.float32, device=device),
+            torch.full((E,), _sm90_limit, dtype=torch.float32, device=device),
             requires_grad=False,
         )
 
@@ -1121,14 +1245,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 moe_runner_backend = MoeRunnerBackend.TRITON
 
         if moe_runner_backend.is_aiter():
-            # MXFP4 hard-codes Swiglu in the AITER kernel path.
-            self.runner = MoeRunner(
-                moe_runner_backend, replace(moe_runner_config, activation="swiglu")
-            )
+            # MXFP4 hard-codes Swiglu in the AITER kernel path, so the
+            # checkpoint's "silu" has to be translated. K3's SiTU is the one
+            # activation the kernel selects on its own -- leave it alone.
+            aiter_config = moe_runner_config
+            if aiter_config.activation != "situ":
+                aiter_config = replace(aiter_config, activation="swiglu")
+            self.runner = MoeRunner(moe_runner_backend, aiter_config)
         elif (
             moe_runner_backend.is_triton_kernels()
             or moe_runner_backend.is_triton()
             or moe_runner_backend.is_marlin()
+            or moe_runner_backend.is_deep_gemm()
         ):
             self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
         elif moe_runner_backend.is_flashinfer_mxfp4() and self._fi_kernel in (
@@ -1172,6 +1300,28 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         )
         return self.runner.run(dispatch_output, quant_info)
 
+    def _apply_marlin(self, layer, dispatch_output):
+        """MXFP4 x BF16 MoE via the Marlin runner. The quant_info (incl. the
+        dispatcher's EP mapping) is shared with
+        :class:`~sglang.srt.layers.quantization.mxfp4_marlin_moe.Mxfp4MarlinMoEMethod`;
+        only the input padding is local (``hidden_size`` is pre-rounded to the
+        Marlin-padded width at create_weights)."""
+        from sglang.srt.layers.quantization.mxfp4_marlin_moe import (
+            build_marlin_moe_quant_info,
+        )
+
+        x = dispatch_output.hidden_states
+        if x.shape[-1] == self.hidden_size:
+            x_padded = x
+        else:
+            x_padded = torch.nn.functional.pad(
+                x, (0, self.hidden_pad), mode="constant", value=0.0
+            )
+        quant_info = build_marlin_moe_quant_info(layer)
+        return self.runner.run(
+            dispatch_output._replace(hidden_states=x_padded), quant_info
+        )
+
     def _apply_sm120_cutlass(self, layer, dispatch_output):
         """SM120 GPT-OSS MXFP8 x MXFP4 MoE via FlashInfer CUTLASS."""
         from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
@@ -1206,6 +1356,24 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
         from sglang.srt.layers.moe.topk import TopKOutputChecker
 
+        if self.use_deep_gemm:
+            # Handles standard AND deepep_ll/deepep_normal dispatch formats via
+            # the runner's registered pre/post-permute functions. Must run
+            # before the `.topk_output` unpack below: deepep dispatch outputs
+            # carry topk_ids/topk_weights directly and have no `.topk_output`.
+            from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo
+
+            quant_info = DeepGemmMoeQuantInfo(
+                w13_weight=layer.w13_weight,
+                w2_weight=layer.w2_weight,
+                use_fp8=True,
+                w13_scale=layer.w13_weight_scale,
+                w2_scale=layer.w2_weight_scale,
+                block_shape=[128, 128],
+                is_fp4_experts=True,
+            )
+            return self.runner.run(dispatch_output, quant_info)
+
         x = dispatch_output.hidden_states
         topk_output = dispatch_output.topk_output
         if _is_cpu:
@@ -1248,27 +1416,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
 
         if self.use_marlin:
             assert TopKOutputChecker.format_is_standard(topk_output)
-            if x.shape[-1] == self.hidden_size:
-                x_padded = x
-            else:
-                x_padded = torch.nn.functional.pad(
-                    x, (0, self.hidden_pad), mode="constant", value=0.0
-                )
-            quant_info = MarlinMoeQuantInfo(
-                w13_qweight=layer.w13_weight,
-                w2_qweight=layer.w2_weight,
-                w13_scales=layer.w13_weight_scale,
-                w2_scales=layer.w2_weight_scale,
-                w13_g_idx_sort_indices=None,
-                w2_g_idx_sort_indices=None,
-                weight_bits=4,
-                is_k_full=True,
-                w13_bias=getattr(layer, "w13_weight_bias", None),
-                w2_bias=getattr(layer, "w2_weight_bias", None),
-            )
-            return self.runner.run(
-                dispatch_output._replace(hidden_states=x_padded), quant_info
-            )
+            return self._apply_marlin(layer, dispatch_output)
 
         if self._fi_kernel == "cutlass_sm90":
             return self._apply_sm90_cutlass(layer, dispatch_output)
@@ -1279,6 +1427,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations,
             # which can theoretically improve performance
             origin_hidden_states_dim = x.shape[-1]
+            # Filled by the staged K3 route+pack+quant fusion below; the pack
+            # site further down falls back to PackTopkIds when it is None.
+            prepared_packed_topk = None
             if self.flashinfer_mxfp4_moe_precision == "bf16":
                 assert x.dtype == torch.bfloat16
                 x_quant = x
@@ -1293,31 +1444,200 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                         value=0.0,
                     )
             elif self.flashinfer_mxfp4_moe_precision == "default":
-                from sglang.srt.layers.quantization.fp8_utils import (
-                    flashinfer_mxfp8_quantize,
-                )
+                if x.shape[-1] == self.hidden_size:
+                    if x.dim() > 2:
+                        x = x.view(-1, x.shape[-1])
+                    # K3 staged fusion (route_quant_handoff): the routing
+                    # dispatch already quantized these rows and packed the
+                    # topk ids in the fused route launch — consume both and
+                    # skip the two standalone kernels. Identity-verified;
+                    # a miss runs the unfused chain below.
+                    from sglang.srt.layers.moe import route_quant_handoff
 
-                x_quant, x_scale = flashinfer_mxfp8_quantize(
-                    x, False, alignment=self.hidden_size
-                )
-                x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1)
+                    prepared = route_quant_handoff.take(x)
+                    if prepared is not None:
+                        prepared_packed_topk, x_quant, x_scale = prepared
+                        x_scale = x_scale.view(torch.float8_e4m3fn)
+                    else:
+                        from sglang.kernels.ops.quantization.per_token_group_quant import (
+                            per_token_group_quant,
+                        )
+
+                        x_quant, x_scale = per_token_group_quant(
+                            x, group_size=32, scale_ue8m0=True
+                        )
+                        x_scale = x_scale.view(torch.float8_e4m3fn)
+                else:
+                    from sglang.srt.layers.quantization.fp8_utils import (
+                        flashinfer_mxfp8_quantize,
+                    )
+
+                    x_quant, x_scale = flashinfer_mxfp8_quantize(
+                        x, False, alignment=self.hidden_size
+                    )
+                    x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
+                        *x.shape[:-1], -1
+                    )
             else:
                 raise NotImplementedError()
 
             assert x_quant.shape[-1] == self.hidden_size
-            assert TopKOutputChecker.format_is_bypassed(topk_output)
+            is_standard = TopKOutputChecker.format_is_standard(topk_output)
+            # The situ path accepts precomputed (standard) routing; the
+            # public path below is logits-only.
+            assert is_standard or TopKOutputChecker.format_is_bypassed(
+                topk_output
+            ), f"unsupported topk format: {topk_output.format}"
+            if is_standard:
+                assert (
+                    self.moe_runner_config.activation == "situ"
+                ), "standard topk output only wired for the situ path"
+                top_k = topk_output.topk_ids.shape[1]
+                router_logits = None
+            else:
+                top_k = topk_output.topk_config.top_k
+                router_logits = topk_output.router_logits
 
-            top_k = topk_output.topk_config.top_k
-            router_logits = topk_output.router_logits
+            num_tokens = x_quant.shape[0]
+            hidden_size = origin_hidden_states_dim
+            # The K3 fused-front path publishes its [latent | shared] buffer
+            # slice as the output destination (zero_copy_context); writing
+            # the finalize output there directly skips this allocation and
+            # the copy_ back in _forward_fused. The slice lives in the same
+            # symmetric buffer the caller all-reduces.
+            symm_output = zero_copy_context.get_moe_output_spec(
+                torch.Size((num_tokens, hidden_size)),
+                torch.bfloat16,
+                x_quant.device,
+            )
+            if symm_output is None:
+                with use_symmetric_memory(
+                    get_tp_group(), disabled=not is_allocation_symmetric()
+                ):
+                    symm_output = torch.empty(
+                        num_tokens,
+                        hidden_size,
+                        dtype=torch.bfloat16,
+                        device=x_quant.device,
+                    )
 
-            with use_symmetric_memory(
-                get_tp_group(), disabled=not is_allocation_symmetric()
-            ):
-                num_tokens = x_quant.shape[0]
-                hidden_size = origin_hidden_states_dim
-                symm_output = torch.empty(
-                    num_tokens, hidden_size, dtype=torch.bfloat16, device=x_quant.device
+            if self.moe_runner_config.activation == "situ":
+                # SiTU is only in the private trtllm-gen cubin pool (the
+                # public artifact bakes swiglu into the fused-act cubins and
+                # silently computes the wrong activation). Routing must also
+                # be noaux_tc (sigmoid + correction bias, DeepSeekV3 method),
+                # not the renormalize-softmax default below.
+                from sglang.kernels.ops.moe import trtllm_gen_moe as situ_moe
+
+                if not situ_moe.available():
+                    raise RuntimeError(
+                        "activation='situ' with the flashinfer_mxfp4 runner "
+                        "needs the SiTU cubin pool: set "
+                        "SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL (see "
+                        "sglang/kernels/ops/moe/trtllm_gen_moe.py)."
+                    )
+                # EP is cubin-internal: each rank computes its local expert slice
+                # [offset, +num_local) and the caller all-reduces. ep=1 -> TP path.
+                local_expert_offset = layer.moe_ep_rank * layer.num_local_experts
+                if TopKOutputChecker.format_is_standard(topk_output):
+                    # Precomputed routing (radix router upstream): skip the
+                    # in-op routing kernels entirely. At small T the in-op
+                    # single-CTA routing costs ~22 us/layer vs ~6 us for the
+                    # external radix router.
+                    if prepared_packed_topk is not None:
+                        packed_topk = prepared_packed_topk
+                    else:
+                        from sglang.kernels.ops.moe.pack_topk_ids import PackTopkIds
+
+                        packed_topk = PackTopkIds.execute(
+                            topk_output.topk_ids, topk_output.topk_weights
+                        )
+                    # Deferred finalize (K3 forward_deferred_finalize): return
+                    # the finalize inputs instead of the finalized output.
+                    from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+                        _deferred_finalize_enabled,
+                    )
+
+                    defer_finalize = _deferred_finalize_enabled.get()
+                    result = situ_moe.trtllm_fp4_block_scale_routed_moe(
+                        packed_topk_ids=packed_topk,
+                        hidden_states=x_quant,
+                        hidden_states_scale=x_scale,
+                        gemm1_weights=layer.w13_weight,
+                        gemm1_weights_scale=layer.w13_weight_scale,
+                        gemm1_alpha=layer.gemm1_alpha,
+                        # SiTuGlu: gatedActBeta is the linear-half tanh
+                        # clip; K3 stores it in gemm1_clamp_limit.
+                        gemm1_beta=layer.gemm1_clamp_limit,
+                        gemm2_weights=layer.w2_weight,
+                        gemm2_weights_scale=layer.w2_weight_scale,
+                        output1_scale_scalar=None,
+                        output1_scale_gate_scalar=None,
+                        output2_scale_scalar=None,
+                        num_experts=layer.num_experts,
+                        top_k=packed_topk.shape[1],
+                        intermediate_size=self.intermediate_size_per_partition,
+                        activation_type=situ_moe.ACTIVATION_SITU,
+                        local_expert_offset=local_expert_offset,
+                        local_num_experts=layer.num_local_experts,
+                        output=symm_output,
+                        do_finalize=not defer_finalize,
+                    )
+                    if defer_finalize:
+                        from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
+                            FlashInferTrtllmDeferredFinalizeOutput,
+                        )
+
+                        gemm2_out, topk_weights, expanded_idx = result
+                        result = FlashInferTrtllmDeferredFinalizeOutput(
+                            gemm2_out=gemm2_out,
+                            expert_weights=topk_weights,
+                            expanded_idx_to_permuted_idx=expanded_idx,
+                            top_k=packed_topk.shape[1],
+                        )
+                    return StandardCombineInput(hidden_states=result)
+
+                # Bypassed topk: route from logits inside the op.
+                correction_bias = topk_output.topk_config.correction_bias
+                bias_bf16 = getattr(layer, "_situ_routing_bias_bf16", None)
+                if bias_bf16 is None and correction_bias is not None:
+                    bias_bf16 = correction_bias.to(torch.bfloat16)
+                    layer._situ_routing_bias_bf16 = bias_bf16
+                situ_moe.trtllm_fp4_block_scale_moe(
+                    # router_logits is a row-strided slice of the K3 fused
+                    # front GEMM output; the FFI reads it as dense.
+                    routing_logits=router_logits.to(torch.bfloat16).contiguous(),
+                    routing_bias=bias_bf16,
+                    hidden_states=x_quant,
+                    hidden_states_scale=x_scale,
+                    gemm1_weights=layer.w13_weight,
+                    gemm1_weights_scale=layer.w13_weight_scale,
+                    gemm1_alpha=layer.gemm1_alpha,
+                    # SiTuGlu: gatedActBeta is the linear-half tanh clip;
+                    # K3 stores it in gemm1_clamp_limit (situ_linear_beta).
+                    gemm1_beta=layer.gemm1_clamp_limit,
+                    gemm2_weights=layer.w2_weight,
+                    gemm2_weights_scale=layer.w2_weight_scale,
+                    output1_scale_scalar=None,
+                    output1_scale_gate_scalar=None,
+                    output2_scale_scalar=None,
+                    num_experts=layer.num_experts,
+                    top_k=top_k,
+                    n_group=topk_output.topk_config.num_expert_group,
+                    topk_group=topk_output.topk_config.topk_group,
+                    intermediate_size=self.intermediate_size_per_partition,
+                    routed_scaling_factor=(
+                        topk_output.topk_config.routed_scaling_factor or 1.0
+                    ),
+                    routing_method_type=situ_moe.ROUTING_DEEPSEEK_V3,
+                    activation_type=situ_moe.ACTIVATION_SITU,
+                    norm_topk_prob=topk_output.topk_config.renormalize,
+                    local_expert_offset=local_expert_offset,
+                    local_num_experts=layer.num_local_experts,
+                    output=symm_output,
                 )
+                return StandardCombineInput(hidden_states=symm_output)
+
             trtllm_gen_output = trtllm_fp4_block_scale_moe(
                 router_logits.to(torch.bfloat16),
                 None,  # routing_bias
@@ -1385,22 +1705,23 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 quant_type=AiterQuantType.PER_1X32,
                 w13_scale=layer.w13_weight_scale,
                 w2_scale=layer.w2_weight_scale,
-                b13=layer.w13_weight_bias,
-                b2=layer.w2_weight_bias,
+                b13=layer.w13_weight_bias if self.with_bias else None,
+                b2=layer.w2_weight_bias if self.with_bias else None,
                 expert_mask=layer.dispatcher.expert_mask_gpu,
                 doweight_stage1=self.moe_runner_config.apply_router_weight_on_input,
                 hidden_pad=self.hidden_pad,
                 intermediate_pad=self.intermediate_pad,
-                # Triggers aiter's INTERLEAVE gate_mode dispatch (required for our
-                # preshuffled gate/up-interleaved weight layout) and applies the
-                # model's swiglu clamp. Models populate the same scalar under
-                # different MoeRunnerConfig fields: gpt-oss uses `gemm1_clamp_limit`
-                # (renamed in `models/gpt_oss.py` from `config.swiglu_limit`); DSv4
-                # / FP8 uses `swiglu_limit` directly. Accept either.
+                # Applies swiglu clamp for GPT-OSS-style activations. K3 SiTU
+                # uses gemm1_clamp_limit as linear_beta, which is forwarded by
+                # the AITER runner and must not be treated as swiglu_limit.
                 swiglu_limit=(
-                    self.moe_runner_config.gemm1_clamp_limit
-                    or self.moe_runner_config.swiglu_limit
-                    or 0.0
+                    0.0
+                    if self.moe_runner_config.activation == "situ"
+                    else (
+                        self.moe_runner_config.gemm1_clamp_limit
+                        or self.moe_runner_config.swiglu_limit
+                        or 0.0
+                    )
                 ),
             )
             return self.runner.run(

python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py

@@ -17,6 +17,31 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def build_marlin_moe_quant_info(layer: Module) -> MarlinMoeQuantInfo:
+    """Build the Marlin quant_info for an MXFP4 MoE layer.
+
+    Single source for the runner inputs shared by the marlin path of
+    ``Mxfp4MoEMethod.apply`` and :class:`Mxfp4MarlinMoEMethod`, including
+    the dispatcher's EP mapping (global -> local expert ids) when EP is on.
+    """
+    expert_map = getattr(layer.dispatcher, "local_expert_mapping", None)
+    global_num_experts = layer.dispatcher.num_experts if expert_map is not None else -1
+    return MarlinMoeQuantInfo(
+        w13_qweight=layer.w13_weight,
+        w2_qweight=layer.w2_weight,
+        w13_scales=layer.w13_weight_scale,
+        w2_scales=layer.w2_weight_scale,
+        w13_g_idx_sort_indices=None,
+        w2_g_idx_sort_indices=None,
+        weight_bits=4,
+        is_k_full=True,
+        w13_bias=getattr(layer, "w13_weight_bias", None),
+        w2_bias=getattr(layer, "w2_weight_bias", None),
+        expert_map=expert_map,
+        global_num_experts=global_num_experts,
+    )
+
+
 class Mxfp4MarlinMoEMethod:
     """MXFP4 (E8M0 scales) MoE quantization method using the Marlin backend."""
 
@@ -162,18 +187,7 @@ class Mxfp4MarlinMoEMethod:
                 value=0.0,
             )
 
-        quant_info = MarlinMoeQuantInfo(
-            w13_qweight=layer.w13_weight,
-            w2_qweight=layer.w2_weight,
-            w13_scales=layer.w13_weight_scale,
-            w2_scales=layer.w2_weight_scale,
-            w13_g_idx_sort_indices=None,
-            w2_g_idx_sort_indices=None,
-            weight_bits=4,
-            is_k_full=True,
-            w13_bias=getattr(layer, "w13_weight_bias", None),
-            w2_bias=getattr(layer, "w2_weight_bias", None),
-        )
+        quant_info = build_marlin_moe_quant_info(layer)
         runner_output = self.runner.run(
             dispatch_output._replace(hidden_states=hidden_states_padded),
             quant_info=quant_info,

python/sglang/srt/layers/quantization/unquant.py

@@ -254,6 +254,48 @@ class UnquantizedLinearMethod(LinearMethodBase):
 
         return F.linear(x, layer.weight, bias)
 
+    def apply_into(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        output: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Run an inference-only BF16 linear into caller-owned storage."""
+        if (
+            get_bf16_gemm_backend().is_cutedsl()
+            and x.is_cuda
+            and x.ndim == 2
+            and x.dtype == torch.bfloat16
+            and layer.weight.dtype == torch.bfloat16
+            and output.dtype == torch.bfloat16
+            and output.is_contiguous()
+            and output.shape == (x.shape[0], layer.weight.shape[0])
+            and (bias is None or bias.dtype == torch.bfloat16)
+            and not layer.weight.requires_grad
+            and (bias is None or not bias.requires_grad)
+            and _use_cutedsl_bf16_gemm(
+                x.shape[0], layer.weight.shape[0], layer.weight.shape[1]
+            )
+        ):
+            from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import (
+                cutedsl_bf16_gemm_out,
+            )
+
+            return cutedsl_bf16_gemm_out(x, layer.weight, output, bias)
+
+        if x.ndim != 2:
+            raise ValueError("caller-owned linear output currently requires a 2D input")
+        if output.shape != (x.shape[0], layer.weight.shape[0]):
+            raise ValueError(
+                f"linear output has shape {output.shape}, expected "
+                f"{(x.shape[0], layer.weight.shape[0])}"
+            )
+        torch.mm(x, layer.weight.t(), out=output)
+        if bias is not None:
+            output.add_(bias)
+        return output
+
 
 class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
     """MoE method without quantization."""

python/sglang/srt/layers/radix_linear_attention.py

@@ -74,6 +74,7 @@ class RadixLinearAttention(nn.Module):
 
         self.A_log = A_log
         self.dt_bias = dt_bias
+        self.lower_bound = None
 
     def forward(
         self,

python/sglang/srt/layers/zero_copy_context.py

@@ -0,0 +1,56 @@
+import contextlib
+from typing import Iterator, Optional
+
+import msgspec
+import torch
+
+
+class ZeroCopyContext(msgspec.Struct, frozen=True):
+    moe_output: Optional[torch.Tensor] = None
+
+
+ctx = ZeroCopyContext()
+
+
+@contextlib.contextmanager
+def set_moe_output(out: torch.Tensor) -> Iterator[None]:
+    """Publish `out` as the MoE runner's output destination for the block."""
+    old_output = ctx.moe_output
+    msgspec.structs.force_setattr(ctx, "moe_output", out)
+    try:
+        yield
+    finally:
+        msgspec.structs.force_setattr(ctx, "moe_output", old_output)
+
+
+def get_moe_output(ref: torch.Tensor) -> Optional[torch.Tensor]:
+    """The published destination iff it can stand in for empty_like(ref)."""
+    return get_moe_output_spec(ref.shape, ref.dtype, ref.device)
+
+
+def get_moe_output_spec(
+    shape: torch.Size, dtype: torch.dtype, device: torch.device
+) -> Optional[torch.Tensor]:
+    """Spec form of get_moe_output for callers that would otherwise have to
+    materialize a reference tensor just for the match (e.g. the trtllm-gen
+    runner, whose activation input is fp4-packed and shaped differently
+    from its output)."""
+    out = ctx.moe_output
+    if (
+        out is not None
+        and out.shape == shape
+        and out.dtype == dtype
+        and out.device == device
+        and out.is_contiguous()
+    ):
+        return out
+    return None
+
+
+__all__ = [
+    "ZeroCopyContext",
+    "ctx",
+    "set_moe_output",
+    "get_moe_output",
+    "get_moe_output_spec",
+]

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

@@ -5,7 +5,7 @@ import logging
 import os
 import threading
 import time
-from queue import Queue
+from queue import Empty, Queue
 from typing import TYPE_CHECKING, Any, Callable, List, Optional
 
 import torch
@@ -675,14 +675,62 @@ class HybridCacheController(BaseHiCacheController):
         operation.pool_transfers_done = True
 
     def _page_backup(self, operation):
-        # Backup extra pools
-        if operation.pool_transfers:
+        # MLA KV is replicated across TP ranks and should still be written only
+        # by TP0. On follower ranks, only the rank-sharded Mamba/KDA pool is
+        # owned by the rank and must be written here. Do not replicate other
+        # sidecar pools (for example SWA or indexer state) accidentally.
+        backup_transfers = operation.pool_transfers
+        if self.backup_skip:
+            backup_transfers = [
+                transfer
+                for transfer in operation.pool_transfers or []
+                if transfer.name == PoolName.MAMBA
+            ]
+
+        if backup_transfers:
             self._resolve_sidecar_derived_pool_transfers(operation)
-            results = self.storage_backend.batch_set_v2(operation.pool_transfers)
+            results = self.storage_backend.batch_set_v2(backup_transfers)
             operation.pool_storage_result.update_extra_pool_hit_pages(results)
 
-        # Backup kv pools
-        super()._page_backup(operation)
+        if not self.backup_skip:
+            super()._page_backup(operation)
+        else:
+            sidecar_ok = bool(backup_transfers)
+            if sidecar_ok:
+                for transfer in backup_transfers:
+                    result = results.get(transfer.name)
+                    if result is None:
+                        result = results.get(transfer.name.value)
+                    expected = len(transfer.keys or [])
+                    if expected == 0 and transfer.host_indices is not None:
+                        expected = int(transfer.host_indices.numel())
+                    if (
+                        not isinstance(result, (list, tuple))
+                        or len(result) != expected
+                        or not all(bool(ok) for ok in result)
+                    ):
+                        sidecar_ok = False
+                        break
+            operation.completed_tokens = (
+                len(operation.hash_value) * self.page_size if sidecar_ok else 0
+            )
+
+    def backup_thread_func(self):
+        """Back up rank-sharded sidecars on every TP rank.
+
+        The base implementation skips the entire operation on non-zero MLA TP
+        ranks. That optimization is valid for replicated MLA KV, but not for
+        hybrid rank-sharded pools such as Kimi-K3 Mamba state.
+        """
+        while not self.storage_stop_event.is_set():
+            try:
+                operation = self.backup_queue.get(block=True, timeout=1)
+                if operation is None:
+                    continue
+                self._page_backup(operation)
+                self.ack_backup_queue.put(operation)
+            except Empty:
+                continue
 
     def _resolve_sidecar_derived_pool_transfers(self, operation):
         for transfer in operation.pool_transfers:

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

@@ -27,7 +27,7 @@ from sglang.srt.mem_cache.pool_host.mha import (
     get_mha_host_pool_cls,
 )
 from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
-from sglang.srt.mem_cache.unified_cache.components import ComponentType
+from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
 from sglang.srt.runtime_context import get_parallel
 
 if TYPE_CHECKING:

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -8,7 +8,11 @@ from typing import TYPE_CHECKING, Any, Optional
 import msgspec
 import torch
 
-from sglang.srt.configs.hybrid_arch import hybrid_gdn_config, mambaish_config
+from sglang.srt.configs.hybrid_arch import (
+    hybrid_gdn_config,
+    kimi_linear_config,
+    mambaish_config,
+)
 from sglang.srt.configs.model_config import (
     ModelConfig,
     get_dsa_index_head_dim,
@@ -702,9 +706,16 @@ class KVCacheConfigurator:
             start_layer=self.layer_info.start_layer,
             linear_replayssm_cache_len=get_exec().mamba.linear_replayssm_cache_len,
             mamba_envelope_layout=get_memory().enable_page_major_kv_layout,
-            enable_gdn_replayssm_spec=(
-                get_exec().mamba.enable_gdn_replayssm_spec
-                and self.hybrid_gdn_config is not None
+            # ReplaySSM spec-verify is for linear-attn models (GDN fold or KDA
+            # fold); activate the pool machinery only for those, so any other
+            # mamba-ish model (Mamba2/Nemotron, lightning, ...) run with the
+            # flag set stays byte-identical to flag-off.
+            enable_linear_replayssm_spec=(
+                get_exec().mamba.enable_linear_replayssm_spec
+                and (
+                    self.hybrid_gdn_config is not None
+                    or kimi_linear_config(self.model_config) is not None
+                )
             ),
         )
         return req_to_token_pool
@@ -738,6 +749,18 @@ class KVCacheConfigurator:
         max_num_reqs: int,
         extra_max_context_len: int,
     ) -> ReqToTokenPool:
+        # DSPARK/DFLASH commit routes through the backend fold (KDA-only); a
+        # non-KDA model there would scatter a None intermediate_ssm and crash.
+        _algo = (self.server_args.speculative_algorithm or "").upper()
+        if (
+            get_exec().mamba.enable_linear_replayssm_spec
+            and _algo in ("DSPARK", "DFLASH")
+            and kimi_linear_config(self.model_config) is None
+        ):
+            raise ValueError(
+                "--enable-linear-replayssm-spec with DSPARK/DFLASH requires a KDA "
+                "(kimi_linear) model; got a non-KDA model."
+            )
         req_to_token_pool = HybridReqToTokenPool(
             size=max_num_reqs,
             mamba_size=get_schedule().max_mamba_cache_size,
@@ -762,9 +785,16 @@ class KVCacheConfigurator:
             enable_linear_replayssm=get_exec().mamba.enable_linear_replayssm,
             linear_replayssm_cache_len=get_exec().mamba.linear_replayssm_cache_len,
             mamba_envelope_layout=get_memory().enable_page_major_kv_layout,
-            enable_gdn_replayssm_spec=(
-                get_exec().mamba.enable_gdn_replayssm_spec
-                and self.hybrid_gdn_config is not None
+            # ReplaySSM spec-verify is for linear-attn models (GDN fold or KDA
+            # fold); activate the pool machinery only for those, so any other
+            # mamba-ish model (Mamba2/Nemotron, lightning, ...) run with the
+            # flag set stays byte-identical to flag-off.
+            enable_linear_replayssm_spec=(
+                get_exec().mamba.enable_linear_replayssm_spec
+                and (
+                    self.hybrid_gdn_config is not None
+                    or kimi_linear_config(self.model_config) is not None
+                )
             ),
         )
         return req_to_token_pool
@@ -1688,11 +1718,23 @@ class KVCacheConfigurator:
             requested_per_worker = None
             max_num_reqs = min(estimated, token_capacity // 2)
 
+        capped_by_mamba = False
         if self.mambaish_config is not None:
             ratio = self._calculate_mamba_ratio()
-            max_num_reqs = min(
-                max_num_reqs, get_schedule().max_mamba_cache_size // ratio
-            )
+            mamba_cap = get_schedule().max_mamba_cache_size // ratio
+            if mamba_cap < max_num_reqs:
+                capped_by_mamba = True
+                logger.warning(
+                    "max_running_requests is capped to %d by the mamba state "
+                    "cache (max_mamba_cache_size=%d, %d state slots per "
+                    "request). To raise it: increase --mamba-full-memory-ratio "
+                    "or --max-mamba-cache-size, or halve the state size with "
+                    "--mamba-ssm-dtype bfloat16.",
+                    mamba_cap,
+                    get_schedule().max_mamba_cache_size,
+                    ratio,
+                )
+            max_num_reqs = min(max_num_reqs, mamba_cap)
 
             if max_num_reqs <= 0:
                 raise RuntimeError(
@@ -1703,7 +1745,11 @@ class KVCacheConfigurator:
                     f"(2) increase --mem-fraction-static, or "
                     f"(3) use GPUs with more memory."
                 )
-        if requested_per_worker is not None and max_num_reqs < requested_per_worker:
+        if (
+            requested_per_worker is not None
+            and max_num_reqs < requested_per_worker
+            and not capped_by_mamba
+        ):
             logger.warning(
                 "max_running_requests was reduced from the requested %d to %d "
                 "(per dp worker) due to the available KV cache capacity.",
@@ -1760,18 +1806,25 @@ class KVCacheConfigurator:
         assert config is not None
 
         has_spec_dec = not self.spec_algorithm.is_none()
+        # ReplaySSM drops the per-step intermediate_ssm scratch, so its mamba budget
+        # no longer reserves the (1 + D/ratio) intermediate factor -- the whole
+        # budget goes to persistent slots (K sized like non-spec), which is how the
+        # freed ~9GB turns into higher max_running.
         # The ring is allocated per slot but is not part of mamba_cache_per_req;
         # the solve must charge it too or num_slots is over-provisioned.
-        replayssm_active = (
-            get_exec().mamba.enable_gdn_replayssm_spec
-            and self.hybrid_gdn_config is not None
+        replayssm_active = get_exec().mamba.enable_linear_replayssm_spec and (
+            self.hybrid_gdn_config is not None
+            or kimi_linear_config(self.model_config) is not None
         )
         if replayssm_active:
-            record_len = (
-                server_args.max_speculative_num_draft_tokens
-                if server_args.max_speculative_num_draft_tokens is not None
-                else get_exec().mamba.linear_replayssm_cache_len
-            )
+            # GDN sizes the fold window to the draft maximum; the KDA ring
+            # stays --linear-replayssm-cache-len long (mirrors MambaPool).
+            if kimi_linear_config(self.model_config) is not None:
+                record_len = get_exec().mamba.linear_replayssm_cache_len
+            elif server_args.max_speculative_num_draft_tokens is not None:
+                record_len = server_args.max_speculative_num_draft_tokens
+            else:
+                record_len = get_exec().mamba.linear_replayssm_cache_len
             replayssm_ring_per_req = (
                 config.mamba2_cache_params.replayssm_ring_bytes_per_req(
                     record_len=record_len
@@ -1790,7 +1843,9 @@ class KVCacheConfigurator:
                 max_mamba_cache_size=get_schedule().max_mamba_cache_size
                 // self.ps.attn_dp_size,
             )
-            # Reserve intermediate memory based on capped max_num_reqs (+1 padding slot)
+            # Reserve intermediate memory based on capped max_num_reqs (+1: the
+            # pool's padding slot, see memory_pool.py). Skipped under replayssm
+            # (no intermediate_ssm allocated).
             if has_spec_dec and not replayssm_active:
                 ratio = self._calculate_mamba_ratio()
                 capped_reqs = min(
@@ -1813,7 +1868,8 @@ class KVCacheConfigurator:
                 max_mamba_cache_size=get_schedule().max_running_requests
                 // self.ps.attn_dp_size,
             )
-            # Reserve intermediate memory based on capped max_num_reqs (+1 padding slot)
+            # Reserve intermediate memory based on capped max_num_reqs (+1: the
+            # pool's padding slot). Skipped under replayssm.
             if has_spec_dec and not replayssm_active:
                 intermediate_size = (
                     config.mamba2_cache_params.mamba_cache_per_req
@@ -1880,7 +1936,9 @@ class KVCacheConfigurator:
                 f"(4) use GPUs with more memory."
             )
 
-        # +1: the pool's padding slot
+        # +1: the pool's padding slot is allocated alongside the request slots.
+        # ReplaySSM ring rides on every slot too (replayssm_ring_per_req is 0 when
+        # the ring is not allocated).
         mamba_state_memory = (
             (get_schedule().max_mamba_cache_size + 1)
             * (config.mamba2_cache_params.mamba_cache_per_req + replayssm_ring_per_req)

python/sglang/srt/mem_cache/memory_pool.py

@@ -344,7 +344,7 @@ class MambaPool:
         #   replayssm_rawv: [num_layers, num_slots, HV, L, V]  (conv/activation dtype)
         #   replayssm_rawk: [num_layers, num_slots, H,  L, K]  (conv/activation dtype)
         #   replayssm_beta: [num_layers, num_slots, HV, L]     (fp32)
-        # The raw rings + beta exist only under --enable-gdn-replayssm-spec: the
+        # The raw rings + beta exist only under --enable-linear-replayssm-spec: the
         # closed-loop exact fold sequentially replays them through the recurrent
         # update at flush -- bit-identical to the recurrent baseline -- instead
         # of folding the chunked `d` records open-loop (which accumulates error
@@ -380,7 +380,7 @@ class MambaPool:
 
     @dataclass(frozen=True, kw_only=True)
     class SpeculativeState(State):
-        # None under --enable-gdn-replayssm-spec: the spec ring owns rollback
+        # None under --enable-linear-replayssm-spec: the spec ring owns rollback
         # (verify writes ring records, commit moves cursors), so the per-draft
         # full-state snapshots are never produced or consumed.
         intermediate_ssm: Optional[torch.Tensor]
@@ -468,7 +468,7 @@ class MambaPool:
         enable_linear_replayssm: bool = False,
         linear_replayssm_cache_len: int = 16,
         envelope_layout: bool = False,
-        enable_gdn_replayssm_spec: bool = False,
+        enable_linear_replayssm_spec: bool = False,
     ):
         conv_state_shape = cache_params.shape.conv
         temporal_state_shape = cache_params.shape.temporal
@@ -487,12 +487,13 @@ class MambaPool:
         self.linear_replayssm_cache_len = linear_replayssm_cache_len
         # ReplaySSM: the decode ring (--enable-linear-replayssm) allocates the
         # chunked (d, k) records + write_pos; the spec-verify flag
-        # (--enable-gdn-replayssm-spec) always uses fold-every-commit and
+        # (--enable-linear-replayssm-spec) always uses fold-every-commit and
         # allocates only the raw (v, k, g, beta) window -- no chunked records,
-        # no cursors. The shared g allocation gates on `_replayssm_on`.
-        self.enable_gdn_replayssm_spec = enable_gdn_replayssm_spec
-        self.replayssm_spec_fold = bool(enable_gdn_replayssm_spec)
-        _replayssm_on = enable_linear_replayssm or enable_gdn_replayssm_spec
+        # no cursors (KDA additionally keeps d/k, see the allocation below).
+        # The shared g allocation gates on `_replayssm_on`.
+        self.enable_linear_replayssm_spec = enable_linear_replayssm_spec
+        self.replayssm_spec_fold = bool(enable_linear_replayssm_spec)
+        _replayssm_on = enable_linear_replayssm or enable_linear_replayssm_spec
 
         # for disagg with nvlink
         self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
@@ -570,7 +571,7 @@ class MambaPool:
             # flag is on; otherwise left as None so the legacy State is
             # byte-identical. temporal_state_shape == (HV, V, K). Either the decode
             # ring (--enable-linear-replayssm) or the spec-verify ring
-            # (--enable-gdn-replayssm-spec) shares this allocation.
+            # (--enable-linear-replayssm-spec) shares this allocation.
             replayssm_d = replayssm_k = replayssm_g = None
             replayssm_rawv = replayssm_rawk = replayssm_beta = None
             if _replayssm_on:
@@ -580,16 +581,23 @@ class MambaPool:
                 num_slots = size + 1
                 # Ring dtype. DECODE ring (--enable-linear-replayssm): records
                 # follow the SSM dtype -- its flush folds `d` directly into the
-                # state. SPEC-verify ring (--enable-gdn-replayssm-spec): d/k feed
+                # state. SPEC-verify ring (--enable-linear-replayssm-spec): d/k feed
                 # ONLY the one-shot output reconstruction (the closed-loop exact
                 # fold replays the raw rings for state instead), so their
                 # quantization noise stays below the bf16 output cast; keep them
                 # in the conv/activation dtype instead of the (fp32-enforced)
                 # SSM dtype to halve the ring traffic. g stays fp32 everywhere
                 # (exact-fold input). The two flags are mutually exclusive.
-                ring_dtype = conv_dtype if enable_gdn_replayssm_spec else ssm_dtype
-                # Fold-every-commit: one verify window, no chunked (d, k) records.
-                if self.replayssm_spec_fold:
+                ring_dtype = conv_dtype if enable_linear_replayssm_spec else ssm_dtype
+                # Fold-every-commit: one verify window, no chunked (d, k)
+                # records. KDA is the exception on both counts: its window
+                # stays L-sized (the fused verify ring-write drops
+                # absorb-inflated rows past L), and d/k stay allocated --
+                # forward_decode routes on `replayssm_d is None` (fused vs
+                # decode-ring), so skipping them would flip KDA decode to the
+                # fused path, a behavior change needing its own validation
+                # (memory follow-up).
+                if self.replayssm_spec_fold and not cache_params.is_kda:
                     record_len = (
                         speculative_num_draft_tokens
                         if speculative_num_draft_tokens is not None
@@ -597,6 +605,7 @@ class MambaPool:
                     )
                 else:
                     record_len = L
+                if not self.replayssm_spec_fold or cache_params.is_kda:
                     replayssm_d = torch.zeros(
                         size=(num_mamba_layers, num_slots, hv, L, v_dim),
                         dtype=ring_dtype,
@@ -626,7 +635,22 @@ class MambaPool:
                 # flush replays these through the recurrent update sequentially
                 # (bit-identical to the recurrent baseline) instead of folding
                 # the chunked `d` records open-loop.
-                if enable_gdn_replayssm_spec:
+                if enable_linear_replayssm_spec:
+                    if cache_params.is_kda:
+                        # Backstop for the KDA ring invariants; this pool is
+                        # sized with the final adaptive-aware draft maximum.
+                        if L & (L - 1) != 0:
+                            raise ValueError(
+                                f"spec-verify ring length must be a power of two, got {L}"
+                            )
+                        if (
+                            speculative_num_draft_tokens is not None
+                            and L < 2 * speculative_num_draft_tokens
+                        ):
+                            raise ValueError(
+                                f"spec-verify ring too small: {L} < "
+                                f"2 * {speculative_num_draft_tokens} (early-flush margin)"
+                            )
                     replayssm_rawv = torch.zeros(
                         size=(num_mamba_layers, num_slots, hv, record_len, v_dim),
                         dtype=conv_dtype,
@@ -662,7 +686,11 @@ class MambaPool:
                 # The recurrent-verify fallback cannot be reached under the flag
                 # (GDN + linear chain + triton enforced in server_args; the
                 # backend asserts loudly if it ever is).
-                if enable_gdn_replayssm_spec:
+                # ReplaySSM skips this dominant scratch (~9GB @ K3 dspark γ=7): the
+                # KDA verify kernel takes intermediate_states_buffer=None (skips the
+                # per-step write, CACHE_INTERMEDIATE_STATES=False) and the commit
+                # replays the ring into the checkpoint instead. This is the memory win.
+                if enable_linear_replayssm_spec:
                     intermediate_ssm_state_cache = None
                 else:
                     intermediate_ssm_state_cache = torch.zeros(
@@ -795,7 +823,7 @@ class MambaPool:
                         f"rawv={get_tensor_size_bytes(replayssm_rawv) / GB:.3f}GB, "
                         f"rawk={get_tensor_size_bytes(replayssm_rawk) / GB:.3f}GB, "
                         f"beta={get_tensor_size_bytes(replayssm_beta) / GB:.3f}GB "
-                        if enable_gdn_replayssm_spec
+                        if enable_linear_replayssm_spec
                         else ""
                     )
                 )
@@ -818,12 +846,12 @@ class MambaPool:
             # all GDN layers of one verify step; advanced by commit_gdn_replayssm_spec.
             self.replayssm_cache_base = (
                 torch.zeros((size + 1,), dtype=torch.int32, device=device)
-                if enable_gdn_replayssm_spec and not self.replayssm_spec_fold
+                if enable_linear_replayssm_spec and not self.replayssm_spec_fold
                 else None
             )
             self.replayssm_is_flush = (
                 torch.zeros((size + 1,), dtype=torch.int8, device=device)
-                if enable_gdn_replayssm_spec and not self.replayssm_spec_fold
+                if enable_linear_replayssm_spec and not self.replayssm_spec_fold
                 else None
             )
             mem_usage_bytes = self.mamba_cache.mem_usage_bytes()
@@ -1130,7 +1158,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
         enable_linear_replayssm: bool = False,
         linear_replayssm_cache_len: int = 16,
         mamba_envelope_layout: bool = False,
-        enable_gdn_replayssm_spec: bool = False,
+        enable_linear_replayssm_spec: bool = False,
     ):
         super().__init__(
             size=size,
@@ -1157,7 +1185,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
             enable_linear_replayssm=enable_linear_replayssm,
             linear_replayssm_cache_len=linear_replayssm_cache_len,
             mamba_envelope_layout=mamba_envelope_layout,
-            enable_gdn_replayssm_spec=enable_gdn_replayssm_spec,
+            enable_linear_replayssm_spec=enable_linear_replayssm_spec,
         )
 
     def _init_mamba_pool(
@@ -1173,7 +1201,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
         enable_linear_replayssm: bool = False,
         linear_replayssm_cache_len: int = 16,
         mamba_envelope_layout: bool = False,
-        enable_gdn_replayssm_spec: bool = False,
+        enable_linear_replayssm_spec: bool = False,
     ):
         self.mamba_pool = self.mamba_pool_cls(
             size=mamba_size,
@@ -1187,7 +1215,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
             enable_linear_replayssm=enable_linear_replayssm,
             linear_replayssm_cache_len=linear_replayssm_cache_len,
             envelope_layout=mamba_envelope_layout,
-            enable_gdn_replayssm_spec=enable_gdn_replayssm_spec,
+            enable_linear_replayssm_spec=enable_linear_replayssm_spec,
         )
         self.mamba_allocator = MambaSlotAllocator(
             size=mamba_size,
@@ -3663,7 +3691,7 @@ class HybridLinearKVPool(KVCache):
 
     def get_kv_buffer_shape(self) -> Tuple[torch.Size, torch.Size]:
         # Hybrid layer ids are global model-layer ids, while the backing pool
-        # is dense over only full-attention layers. Shape discovery does not
+        # is dense over only full-attention layers.  Shape discovery does not
         # need a global layer lookup, so delegate it to that backing pool.
         return self.full_kv_pool.get_kv_buffer_shape()
 

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -710,13 +710,14 @@ class UnifiedMambaPool(MambaPool):
         self.enable_custom_mem_pool = False
         self.custom_mem_pool = None
         self.num_mamba_layers = spec.layer_num
-        # GDN/KDA ReplaySSM unsupported; replicate parent's disabled-state attrs so
-        # paths guarded by `replayssm_write_pos is not None` don't AttributeError.
+        # GDN/KDA ReplaySSM / spec unsupported; replicate parent's disabled-state
+        # attrs so unconditional reads (e.g. `replayssm_cache_base is not None` in
+        # the req-slot alloc path) and `... is not None` guards don't AttributeError.
         self.enable_linear_replayssm = False
         self.linear_replayssm_cache_len = 16
         self.replayssm_write_pos = None
         self.replayssm_is_kda = False
-        self.enable_gdn_replayssm_spec = False
+        self.enable_linear_replayssm_spec = False
         self.replayssm_spec_fold = False
         self.replayssm_cache_base = None
         self.replayssm_is_flush = None
@@ -941,10 +942,10 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
         mamba_envelope_layout: bool = False,
         enable_linear_replayssm: bool = False,
         linear_replayssm_cache_len: int = 16,
-        enable_gdn_replayssm_spec: bool = False,
+        enable_linear_replayssm_spec: bool = False,
     ):
         # mamba_envelope_layout / speculative_eagle_topk / enable_linear_replayssm /
-        # linear_replayssm_cache_len / enable_gdn_replayssm_spec: accepted to match
+        # linear_replayssm_cache_len / enable_linear_replayssm_spec: accepted to match
         # the parent signature but NOT forwarded — the shared pool's conv/temporal
         # state are fixed-shape views (replayssm/spec are gated off under unified).
         assert mamba_size == self._shared_mamba_size, (

python/sglang/srt/model_executor/input_buffers.py

@@ -66,12 +66,35 @@ def share_input_buffers_in(obj) -> None:
             setattr(obj, name, share_input_buffer(name, buffer))
 
 
+# Values that index the rope table, the KV pool, req_to_token, or the mamba
+# state pool, so stale content is unsafe to execute.
+_INDEX_SEMANTIC_BUFFERS = frozenset(
+    {
+        "positions",
+        "mrope_positions",
+        "out_cache_loc",
+        "req_pool_indices",
+        "mamba_track_indices",
+        "mamba_track_mask",
+    }
+)
+
+
 @dataclass
 class ForwardInputBuffers:
 
     def _share_one_buffer(self, name: str, new_buffer: torch.Tensor) -> torch.Tensor:
         return share_input_buffer(name, new_buffer)
 
+    def reset_index_buffers(self) -> None:
+        """Zero the index-semantic buffers this set declares."""
+        for f in fields(self):
+            if f.name not in _INDEX_SEMANTIC_BUFFERS:
+                continue
+            buffer = getattr(self, f.name)
+            if buffer is not None:
+                buffer.zero_()
+
     def share_buffers(self):
         # disable share input buffer on npu due to accuracy issue
         if is_npu():

python/sglang/srt/model_executor/model_runner.py

@@ -134,6 +134,7 @@ from sglang.srt.model_executor.model_runner_components.load_model_utils import (
     load_model_with_memory_saver,
     maybe_downgrade_dtype_for_legacy_gpu,
     maybe_enable_ipc_weight_cache,
+    maybe_precompile_model_kernels_after_loading,
     maybe_register_debug_tensor_dump_hook,
     maybe_trigger_remote_instance_nccl_send_group,
     report_online_quantization,
@@ -736,6 +737,14 @@ class ModelRunner:
             start_layer=self.layer_info.start_layer,
         )
 
+    def get_pp_proxy_residual_num_blocks(self) -> Optional[int]:
+        return misc_utils.resolve_pp_proxy_residual_num_blocks(
+            model_config=self.model_config,
+            pp_size=self.ps.pp_size,
+            pp_rank=self.ps.pp_rank,
+            start_layer=self.layer_info.start_layer,
+        )
+
     def decode_num_tokens_per_req(
         self, *, num_draft_tokens: Optional[int] = None
     ) -> int:
@@ -1060,6 +1069,8 @@ class ModelRunner:
         if not self.is_draft_worker:
             get_offloader().post_init()
 
+        self.maybe_precompile_model_kernels_after_loading()
+
         # Register model for layerwise NVTX profiling if enabled
         if get_exec().comm.enable_layerwise_nvtx_marker:
             pyt_hooks = PytHooks()
@@ -1124,6 +1135,9 @@ class ModelRunner:
             is_ep_joiner=self.server_args.is_ep_joiner,
         )
 
+    def maybe_precompile_model_kernels_after_loading(self) -> None:
+        maybe_precompile_model_kernels_after_loading(self.model, self.device)
+
     def maybe_init_dwdp(self):
         if self.is_draft_worker:
             return
@@ -1190,7 +1204,9 @@ class ModelRunner:
                 model_dtype=getattr(self, "dtype", torch.bfloat16),
                 is_draft_worker=getattr(self, "is_draft_worker", False),
                 is_dflash=(
-                    spec_algorithm.is_dflash() if spec_algorithm is not None else False
+                    spec_algorithm.is_dflash_family()
+                    if spec_algorithm is not None
+                    else False
                 ),
                 speculative_draft_attention_backend=getattr(
                     self.server_args, "speculative_draft_attention_backend", None

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

@@ -24,6 +24,7 @@ from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
     RemoteInstanceWeightLoaderBackend,
     trigger_init_weights_send_group_for_remote_instance_request,
 )
+from sglang.srt.platforms import current_platform
 from sglang.srt.utils.common import is_npu
 from sglang.srt.utils.network import NetworkAddress
 
@@ -40,6 +41,20 @@ _is_npu = is_npu()
 UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480  # leave more time for post data processing
 
 
+def maybe_precompile_model_kernels_after_loading(model, device: str) -> None:
+    precompile = getattr(model, "precompile_kernels_after_loading", None)
+    if precompile is None:
+        return
+
+    if device == "cuda":
+        current_platform.synchronize()
+        current_platform.empty_cache()
+    precompile()
+    if device == "cuda":
+        current_platform.synchronize()
+        current_platform.empty_cache()
+
+
 class LoadedModel(msgspec.Struct, frozen=True, kw_only=True):
     loader: Any
     model: Any

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

@@ -3,7 +3,11 @@ from __future__ import annotations
 import logging
 from typing import TYPE_CHECKING, Any, Optional
 
-from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa
+from sglang.srt.configs.model_config import (
+    dsa_layer_skips_topk,
+    is_deepseek_dsa,
+    is_kimi_k3,
+)
 from sglang.srt.runtime_context import get_context, get_exec, get_schedule
 from sglang.srt.server_args import CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS
 
@@ -68,3 +72,16 @@ def resolve_pp_proxy_topk_size(
     ):
         return None
     return getattr(hf_config, "index_topk", None)
+
+
+def resolve_pp_proxy_residual_num_blocks(
+    *, model_config: ModelConfig, pp_size: int, pp_rank: int, start_layer: int
+) -> Optional[int]:
+    """Return the inherited Kimi K3 attention-residual bank width."""
+    if pp_size <= 1 or pp_rank == 0 or not is_kimi_k3(model_config.hf_config):
+        return None
+
+    block_size = getattr(model_config.hf_text_config, "attn_res_block_size", None)
+    if block_size is None:
+        return None
+    return (start_layer + block_size - 1) // block_size

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

@@ -81,6 +81,7 @@ def _allocate_decode_buffers(
     ne_token_table: Optional[torch.Tensor] = None,
     hc_hidden_size: Optional[int] = None,
     pp_proxy_topk_size: Optional[int] = None,
+    pp_proxy_residual_num_blocks: Optional[int] = None,
 ) -> SimpleNamespace:
     """Allocate the FB-shared decode buffers."""
     with torch.device(device):
@@ -115,9 +116,14 @@ def _allocate_decode_buffers(
                 "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
             }
             if not is_mhc:
-                pp_proxy_tensors["residual"] = torch.zeros(
-                    (max_bs, hidden_size), dtype=dtype
+                # Only Kimi K3 supplies num_blocks: its PP bank is token-major
+                # [T, blocks, H]. Other models keep the legacy [max_bs, H].
+                residual_shape = (
+                    (max_num_token, pp_proxy_residual_num_blocks, hidden_size)
+                    if pp_proxy_residual_num_blocks is not None
+                    else (max_bs, hidden_size)
                 )
+                pp_proxy_tensors["residual"] = torch.zeros(residual_shape, dtype=dtype)
             if pp_proxy_topk_size is not None:
                 pp_proxy_tensors["topk_indices"] = torch.zeros(
                     (max_num_token, pp_proxy_topk_size), dtype=torch.int32
@@ -338,6 +344,7 @@ class BaseRunner(ABC):
             ),
             hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None),
             pp_proxy_topk_size=mr.get_pp_proxy_topk_size(),
+            pp_proxy_residual_num_blocks=mr.get_pp_proxy_residual_num_blocks(),
         )
 
     def _dummy_run(

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

@@ -297,6 +297,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             else None
         )
         self._ragged_graph_size = 0
+        # Per-tier capture layouts; their verify_lens / qo_indptr tensors are
+        # baked into the captured graphs and refreshed in place each replay.
+        self._captured_ragged_layouts: dict[int, object] = {}
         if self.ragged_verify_mode and (
             self.enable_two_batch_overlap
             or model_runner.server_args.enable_lora
@@ -377,6 +380,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 self.model_runner.model_config, "hc_hidden_size", None
             ),
             pp_proxy_topk_size=self.model_runner.get_pp_proxy_topk_size(),
+            pp_proxy_residual_num_blocks=(
+                self.model_runner.get_pp_proxy_residual_num_blocks()
+            ),
         )
         self.buffers.share_buffers()
         # FB-shared slot registry adopting DecodeInputBuffers storage (same
@@ -495,11 +501,27 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             num_slots=self._ragged_capture_slots(num_tokens),
             num_draft_tokens=self.captured_req_width,
         )
-        return RaggedVerifyLayout.from_verify_lens(
+        layout = RaggedVerifyLayout.from_verify_lens(
             verify_lens_cpu=verify_lens_cpu,
             device=self.device,
             grid=self.capture_num_tokens,
         )
+        self._captured_ragged_layouts[num_tokens] = layout
+        return layout
+
+    def _stage_ragged_verify_layout(self, ragged_layout, graph_size_key: int) -> None:
+        # Without this refresh every replay reuses the capture-time synthetic
+        # verify_lens / qo_indptr and mis-slices the packed q rows.
+        cap_layout = self._captured_ragged_layouts.get(graph_size_key)
+        if cap_layout is None:
+            return
+        live = ragged_layout
+        if live.bs != cap_layout.bs or live.cap is None:
+            live = live.padded_to_bucket(
+                padded_bs=cap_layout.bs, cap=self.captured_req_width
+            )
+        cap_layout.verify_lens.copy_(live.verify_lens)
+        cap_layout.qo_indptr_device.copy_(live.qo_indptr_device)
 
     def can_run_graph(self, forward_batch: ForwardBatch):
         # Disable for token embedding overrides (dynamic per-request)
@@ -823,6 +845,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         # captured graph needs before capturing.
         self.buffers.seq_lens.fill_(self.seq_len_fill_value)
         self.buffers.seq_lens_cpu.fill_(self.seq_len_fill_value)
+        # Capture runs real forwards, so a mid-serving recapture would index --
+        # and write KV -- through the previous batch's live values. Replay is
+        # already covered by the registry's padding policy.
+        self.buffers.reset_index_buffers()
 
         # Trigger CUDA graph capture for specific shapes.
         # Capture the large shapes first so that the smaller shapes
@@ -1036,6 +1062,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     f"stale ragged raw_num_token {self.raw_num_token} != "
                     f"{ragged_layout.graph_num_tokens}"
                 )
+                self._stage_ragged_verify_layout(ragged_layout, graph_size_key)
             self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
             self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
             if (
@@ -1072,6 +1099,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 f"{raw_bs}; the planner must reject this batch before replay"
             )
             padded_num_tokens = graph_size_key
+            self._stage_ragged_verify_layout(ragged_layout, graph_size_key)
         else:
             raw_num_token = raw_bs * self.captured_req_width
             if self.require_mlp_tp_gather:

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

@@ -337,6 +337,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         self.has_mha_companion_layers = any(
             layer is not None for layer in self.mha_companion_layers
         )
+        # Archs on the MLA-BCG allowlist pin the absorbed MLA path inside
+        # capture/replay (attention_backend_handler), so the MHA companion is
+        # never captured and the MHA-prefix restrictions below don't apply.
+        self.mla_pinned_under_bcg = (
+            self.model_runner.model_config.is_mla_breakable_cuda_graph_supported
+        )
         self.moe_layers = self.model_runner.moe_layers
         self.moe_fusions = self.model_runner.moe_fusions
         self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None)
@@ -1052,11 +1058,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             return False
         # A prefix forces the MHA companion path, whose captured state is
         # frozen prefix-free; DSA models are exempt (capture/replay force
-        # the sparse path, which takes any prefix via device metadata).
+        # the sparse path, which takes any prefix via device metadata), as
+        # are archs on the MLA-BCG allowlist (they pin the absorbed MLA path
+        # inside capture/replay, so the MHA companion is never captured).
         if (
             self.prefill_backend_name == Backend.BREAKABLE
             and self.has_mha_companion_layers
             and not self.dsa_sparse_prefill_forced
+            and not self.mla_pinned_under_bcg
             and prefix_lens is not None
             and any(prefix_lens)
         ):
@@ -1536,6 +1545,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         if (
             isinstance(self.backend, BreakableCudaGraphBackend)
             and self.has_mha_companion_layers
+            and not self.mla_pinned_under_bcg
         ):
             self._restore_mha_capture_state(static_forward_batch)
 

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -105,6 +105,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
         ne_token_table: Optional[torch.Tensor] = None,
         hc_hidden_size: Optional[int] = None,
         pp_proxy_topk_size: Optional[int] = None,
+        pp_proxy_residual_num_blocks: Optional[int] = None,
     ) -> DecodeInputBuffers:
         with torch.device(device):
             input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
@@ -135,8 +136,15 @@ class DecodeInputBuffers(ForwardInputBuffers):
                     "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
                 }
                 if not is_mhc:
+                    # Only Kimi K3 supplies num_blocks: its PP bank is token-major
+                    # [T, blocks, H]. Other models keep the legacy [max_bs, H].
+                    residual_shape = (
+                        (max_num_token, pp_proxy_residual_num_blocks, hidden_size)
+                        if pp_proxy_residual_num_blocks is not None
+                        else (max_bs, hidden_size)
+                    )
                     pp_proxy_tensors["residual"] = torch.zeros(
-                        (max_bs, hidden_size), dtype=dtype
+                        residual_shape, dtype=dtype
                     )
                 if pp_proxy_topk_size is not None:
                     pp_proxy_tensors["topk_indices"] = torch.zeros(

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

@@ -80,7 +80,10 @@ def _support_mha_one_shot(attn, forward_batch, backend_name):
 
 
 def _handle_attention_backend(attn, forward_batch, backend_name):
-    if is_in_tc_piecewise_cuda_graph():
+    # Captured prefill (tc_piecewise or breakable) must keep a single attention
+    # path: pin the absorbed MLA method — MHA one-shot/chunked shapes vary with
+    # kv-len and cannot be captured.
+    if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph():
         return AttnForwardMethod.MLA
 
     # MLA prefill CP forces absorbed MLA regardless of prefix length: the
@@ -145,7 +148,7 @@ def handle_attention_fa4(attn, forward_batch):
 
 
 def handle_attention_trtllm_mla(attn, forward_batch):
-    if is_in_tc_piecewise_cuda_graph():
+    if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph():
         return AttnForwardMethod.MLA
 
     sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
@@ -190,7 +193,7 @@ def handle_attention_dsa(attn, forward_batch):
 
 
 def handle_attention_triton(attn, forward_batch):
-    if is_in_tc_piecewise_cuda_graph():
+    if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph():
         return AttnForwardMethod.MLA
 
     # when deterministic inference is enabled, use MLA

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

@@ -27,6 +27,7 @@ from sglang.srt.layers.dcp import (
     cp_lse_ag_out_rs_mla,
     dcp_a2a_lse_reduce,
 )
+from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
 from sglang.srt.layers.quantization.fp8_utils import (
     materialize_bpreshuffle_fp8_scale_tuple,
 )
@@ -68,7 +69,7 @@ from sglang.srt.models.deepseek_common.utils import (
     _use_aiter_bpreshuffle_gfx95,
     _use_aiter_gfx95,
 )
-from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args, get_spec
+from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args
 from sglang.srt.state_capturer.indexer_topk import (
     maybe_capture_indexer_topk,
 )
@@ -92,23 +93,12 @@ class MlaBmmFusionPlan:
     attn_output_buf: torch.Tensor
 
 
-def _is_dcp_mla_decode_phase(forward_batch: ForwardBatch) -> bool:
-    if not get_parallel().dcp_enabled:
-        return False
-    if forward_batch.forward_mode.is_decode():
-        return True
-    if not forward_batch.forward_mode.is_target_verify() or not _is_cuda:
-        return False
-
-    server_args = get_server_args()
-    decode_backend = (
-        server_args.decode_attention_backend or server_args.attention_backend
-    )
-    return (
-        get_spec().speculative_algorithm == "DSPARK"
-        and get_spec().speculative_attention_mode == "decode"
-        and decode_backend in ("tokenspeed_mla", "cutedsl_mla")
-    )
+def _select_local_dcp_heads_for_autotune(
+    attn_output: torch.Tensor, num_local_heads: int
+) -> torch.Tensor:
+    """Select this rank's head shard without communicating dummy outputs."""
+    rank = get_parallel().attn_dcp_rank
+    return attn_output.narrow(1, rank * num_local_heads, num_local_heads)
 
 
 def _is_mla_dcp_lse_base_on_e(attention_backend: Optional[str]) -> bool:
@@ -366,7 +356,11 @@ class DeepseekMLAForwardMixin:
         # weights and skip the per-layer Q all-gather (bf16 decode absorb only).
         q_replicate_active = (
             get_parallel().dcp_replicate_q_proj
-            and _is_dcp_mla_decode_phase(forward_batch)
+            and get_parallel().dcp_enabled
+            and (
+                forward_batch.forward_mode.is_decode()
+                or forward_batch.forward_mode.is_target_verify()
+            )
             and not self.use_deep_gemm_bmm
             and self.w_kc_qrep is not None
             and self.q_b_proj_qrep_weight is not None
@@ -802,7 +796,10 @@ class DeepseekMLAForwardMixin:
 
         # all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch.
         if get_parallel().dcp_enabled:
-            if _is_dcp_mla_decode_phase(forward_batch):
+            if (
+                forward_batch.forward_mode.is_decode()
+                or forward_batch.forward_mode.is_target_verify()
+            ):
                 if not q_replicate_active:
                     q_nope_out, q_pe = all_gather_q_for_mla_decode(
                         q_nope_out=q_nope_out,
@@ -955,7 +952,10 @@ class DeepseekMLAForwardMixin:
                         topk_indices=topk_indices,
                     )
                     attn_output = fusion_plan.attn_output_buf
-                elif _is_dcp_mla_decode_phase(forward_batch):
+                elif (
+                    forward_batch.forward_mode.is_decode()
+                    or forward_batch.forward_mode.is_target_verify()
+                ) and get_parallel().dcp_enabled:
                     # set return_lse=True to correct attn_output
                     attn_output, lse = self.attn_mqa_for_dcp_decode(
                         q_nope_out,
@@ -1029,31 +1029,44 @@ class DeepseekMLAForwardMixin:
             )
 
         # correct attn_output with respect to lse from other ranks
-        if _is_dcp_mla_decode_phase(forward_batch):
+        if (
+            forward_batch.forward_mode.is_decode()
+            or forward_batch.forward_mode.is_target_verify()
+        ) and get_parallel().dcp_enabled:
             attn_output = attn_output.view(
                 -1,
                 self.num_local_heads * get_parallel().attn_dcp_size,
                 self.kv_lora_rank,
             )
-            dcp_comm_backend = get_parallel().dcp_comm_backend
-            is_lse_base_on_e = _is_mla_dcp_lse_base_on_e(self.current_attention_backend)
-            if dcp_comm_backend in ("a2a", "fi_a2a"):
-                # A2A exchange of head partials + LSE, then local Triton combine.
-                attn_output = dcp_a2a_lse_reduce(
-                    attn_output.contiguous(),
-                    lse.contiguous(),
-                    get_parallel().dcp_group,
-                    is_lse_base_on_e=is_lse_base_on_e,
-                    comm_backend=dcp_comm_backend,
+            if get_in_autotune_dummy_run():
+                # The synthetic FlashInfer MoE autotune pass discards model
+                # outputs. Avoid an unnecessary cross-node MNNVL exchange of
+                # zero attention partials.
+                attn_output = _select_local_dcp_heads_for_autotune(
+                    attn_output, self.num_local_heads
                 )
             else:
-                attn_output = cp_lse_ag_out_rs_mla(
-                    attn_output,
-                    lse,
-                    get_parallel().dcp_group,
-                    is_lse_base_on_e=is_lse_base_on_e,
+                dcp_comm_backend = get_parallel().dcp_comm_backend
+                is_lse_base_on_e = _is_mla_dcp_lse_base_on_e(
+                    self.current_attention_backend
                 )
-                attn_output = attn_output.transpose(0, 1)
+                if dcp_comm_backend in ("a2a", "fi_a2a"):
+                    # A2A exchange of head partials + LSE, then local Triton combine.
+                    attn_output = dcp_a2a_lse_reduce(
+                        attn_output.contiguous(),
+                        lse.contiguous(),
+                        get_parallel().dcp_group,
+                        is_lse_base_on_e=is_lse_base_on_e,
+                        comm_backend=dcp_comm_backend,
+                    )
+                else:
+                    attn_output = cp_lse_ag_out_rs_mla(
+                        attn_output,
+                        lse,
+                        get_parallel().dcp_group,
+                        is_lse_base_on_e=is_lse_base_on_e,
+                    )
+                    attn_output = attn_output.transpose(0, 1)
         attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
 
         _kvb_v = None

python/sglang/srt/models/dspark.py

@@ -596,11 +596,17 @@ class DSparkDraftMixin:
 
 class DSparkDraftModel(DSparkDraftMixin, DFlashDraftModel):
 
-    pass
+    def prune_to_ctx_kv_injection(self) -> None:
+        self.markov_head = None
+        self.confidence_head = None
+        for layer in self.layers:
+            layer.mlp = None
+            layer.self_attn.o_proj = None
+        torch.cuda.empty_cache()
 
 
 class Qwen3DSparkModel(DSparkDraftModel):
     pass
 
 
-EntryClass = [Qwen3DSparkModel]
+EntryClass = [Qwen3DSparkModel, DSparkDraftModel]

python/sglang/srt/models/inkling_common/dense_mlp.py

@@ -217,6 +217,7 @@ class InklingBatchDenseMLP(nn.Module, FusedMoELoadingMixin):
             self.quant_method,
             self.moe_runner_config,
             self.moe_tp_rank,
+            self.moe_tp_size,
         )
         self.quant_method.create_weights(
             layer=self,

python/sglang/srt/models/inkling_common/util.py

@@ -88,6 +88,7 @@ class FusedMoELoadingMixin(abc.ABC):
         quant_method: UnquantizedFusedMoEMethod,
         moe_runner_config: MoeRunnerConfig,
         moe_tp_rank: int,
+        moe_tp_size: int,
     ) -> None:
         super().__init__()
         helper = FusedMoE.__new__(FusedMoE)
@@ -97,6 +98,7 @@ class FusedMoELoadingMixin(abc.ABC):
         helper.moe_runner_config = moe_runner_config
         helper.use_triton_kernels = False
         helper.moe_tp_rank = moe_tp_rank
+        helper.moe_tp_size = moe_tp_size
         helper.use_presharded_weights = False
         helper.use_flashinfer_trtllm_moe = False
         # Keep this parameterless loading helper out of the module tree so

python/sglang/srt/models/kimi_k3_vl.py

@@ -0,0 +1,937 @@
+"""Kimi K3 vision tower (MoonViT3d) and projector.
+
+Faithful port of the checkpoint reference implementation
+(modeling_kimi_k3.py). Dedicated to K3 — do not share with Kimi K2.5:
+K3 uses qkv_hidden_size != hidden_size (head_dim = qkv_hidden_size //
+num_heads), RMSNorm encoder norms, bias-free linears, and the
+PatchMergerMLPV2 projector (no pre-norm, post RMSNorm), all of which
+differ from the K2.5 vision code.
+
+Weight names match the checkpoint exactly (wqkv/wo, mlp.fc0/fc1,
+mm_projector.proj.0/proj.2, mm_projector.post_norm), so loading needs no
+renames.
+"""
+
+from dataclasses import dataclass, replace
+from typing import List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from sglang.kernels.ops.attention.vision_rope import (
+    apply_fused_qk_complex_rope,
+    precompile_fused_qk_complex_rope,
+)
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.vision import (
+    FLASHINFER_WORKSPACE_SIZE_BYTES,
+    QKV_BACKEND_IMPL,
+    VisionAttentionMetadata,
+    prepare_flashinfer_cudnn_vision_attention_metadata,
+    prepare_vision_attention_metadata,
+)
+from sglang.srt.models.kimi_vl_moonvit import tpool_patch_merger
+from sglang.srt.multimodal.mm_utils import concat_or_single
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.utils import get_bool_env_var, is_hip, print_info_once
+
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+
+if _use_aiter:
+    from aiter.ops.triton.conv.conv2d import conv2d as aiter_conv2d
+
+_SM103_TRITON_MAX_SEQLEN = 1536
+_SM103_FA4_MIN_ATTENTION_WORK = 3_000_000
+
+GridTHW = Tuple[int, int, int]
+SegmentBounds = Tuple[Tuple[int, int], ...]
+
+
+def _resolve_grid_thw_list(
+    grid_thws: torch.Tensor, grid_thw_list: Optional[Sequence[Sequence[int]]] = None
+) -> Tuple[GridTHW, ...]:
+    values = grid_thws.tolist() if grid_thw_list is None else grid_thw_list
+    return tuple((int(t), int(h), int(w)) for t, h, w in values)
+
+
+def _get_mm_attention_backend() -> str:
+    try:
+        server_args = get_server_args()
+    except ValueError:
+        return "auto"
+    return server_args.mm_attention_backend or "auto"
+
+
+def _is_fa4_available() -> bool:
+    try:
+        from sglang.kernels.ops.attention.flash_attention_v4 import (
+            is_flash_attention_v4_available,
+        )
+    except ImportError:
+        return False
+    return is_flash_attention_v4_available()
+
+
+def _resolve_mm_attention_backend(
+    configured_backend: str,
+    *,
+    max_seqlen: int,
+    total_tokens: int,
+    device: torch.device,
+    fa4_available: Optional[bool] = None,
+) -> str:
+    if configured_backend != "auto":
+        return configured_backend
+    if device.type != "cuda":
+        return "sdpa"
+    if torch.cuda.get_device_capability(device) != (10, 3):
+        return "sdpa"
+
+    use_fa4 = (
+        max_seqlen > _SM103_TRITON_MAX_SEQLEN
+        or max_seqlen * total_tokens >= _SM103_FA4_MIN_ATTENTION_WORK
+    )
+    if use_fa4:
+        if fa4_available is None:
+            fa4_available = _is_fa4_available()
+        return "fa4" if fa4_available else "sdpa"
+    return "triton_attn"
+
+
+def apply_rope(
+    xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    freqs_cis = freqs_cis.unsqueeze(-2)
+    xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().view(*xk.shape[:-1], -1, 2))
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+def _can_use_fused_rope(hidden_states: torch.Tensor, freqs_cis: torch.Tensor) -> bool:
+    return _can_use_fused_rope_for_shape(
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+        freqs_cis=freqs_cis,
+    )
+
+
+def _can_use_fused_rope_for_shape(
+    *,
+    dtype: torch.dtype,
+    device: torch.device,
+    freqs_cis: torch.Tensor,
+) -> bool:
+    if not (
+        device.type == "cuda"
+        and freqs_cis.is_cuda
+        and device == freqs_cis.device
+        and dtype in (torch.bfloat16, torch.float16)
+        and freqs_cis.dtype == torch.complex64
+    ):
+        return False
+    major, _ = torch.cuda.get_device_capability(device)
+    return major >= 9
+
+
+def sdpa_varlen_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    *,
+    segment_bounds: Optional[SegmentBounds] = None,
+) -> torch.Tensor:
+    is_single_segment = (
+        len(segment_bounds) == 1
+        if segment_bounds is not None
+        else cu_seqlens.numel() == 2
+    )
+    if is_single_segment:
+        out = F.scaled_dot_product_attention(
+            q.transpose(0, 1).unsqueeze(0),
+            k.transpose(0, 1).unsqueeze(0),
+            v.transpose(0, 1).unsqueeze(0),
+        )
+        return out.squeeze(0).transpose(0, 1).flatten(start_dim=-2)
+    outputs = []
+    if segment_bounds is None:
+        bounds = cu_seqlens.tolist()
+        segment_bounds = tuple(zip(bounds[:-1], bounds[1:]))
+    for start, end in segment_bounds:
+        seg_q = q[start:end].transpose(0, 1).unsqueeze(0)
+        seg_k = k[start:end].transpose(0, 1).unsqueeze(0)
+        seg_v = v[start:end].transpose(0, 1).unsqueeze(0)
+        out = F.scaled_dot_product_attention(seg_q, seg_k, seg_v)
+        outputs.append(out.squeeze(0).transpose(0, 1))
+    return torch.cat(outputs).flatten(start_dim=-2)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    return np.concatenate([emb_sin, emb_cos], axis=1)
+
+
+def get_1d_sincos_pos_embed(embed_dim: int, t_size: int) -> np.ndarray:
+    grid_t = np.arange(t_size, dtype=np.float32)
+    return get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t)
+
+
+def interpolate_pos_emb(
+    weight: torch.Tensor, interpolation_mode: str, shape: Tuple[int, int]
+) -> torch.Tensor:
+    return (
+        F.interpolate(
+            weight.permute((2, 0, 1)).contiguous().unsqueeze(0),
+            size=shape,
+            mode=interpolation_mode,
+        )
+        .squeeze(0)
+        .permute((1, 2, 0))
+        .flatten(end_dim=1)
+    )
+
+
+class Learnable2DInterpPosEmbDividedFixed(nn.Module):
+    def __init__(
+        self,
+        height: int,
+        width: int,
+        num_frames: int,
+        dim: int,
+        interpolation_mode: str = "bicubic",
+    ) -> None:
+        super().__init__()
+        self.height = height
+        self.width = width
+        self.num_frames = num_frames
+        self.dim = dim
+        self.interpolation_mode = interpolation_mode
+        self.weight = nn.Parameter(torch.empty(height, width, dim))
+        self.register_buffer(
+            "time_weight",
+            torch.from_numpy(get_1d_sincos_pos_embed(dim, num_frames))
+            .float()
+            .unsqueeze(1),
+            persistent=False,
+        )
+
+    def position_embeddings(
+        self,
+        grid_thws: torch.Tensor,
+        *,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+    ) -> torch.Tensor:
+        pos_embs = []
+        for t, h, w in _resolve_grid_thw_list(grid_thws, grid_thw_list):
+            assert t <= self.num_frames, f"t:{t} > num_frames:{self.num_frames}"
+            if (h, w) == self.weight.shape[:-1]:
+                pos_emb_2d = self.weight.flatten(end_dim=1)
+            else:
+                pos_emb_2d = interpolate_pos_emb(
+                    self.weight, self.interpolation_mode, (h, w)
+                )
+
+            if t == 1:
+                pos_emb_3d = pos_emb_2d
+            else:
+                pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[
+                    0:t
+                ].to(pos_emb_2d.dtype)
+
+            pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1]))
+
+        return torch.cat(pos_embs)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        grid_thws: torch.Tensor,
+        *,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+        position_embeddings: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if position_embeddings is None:
+            position_embeddings = self.position_embeddings(
+                grid_thws, grid_thw_list=grid_thw_list
+            )
+        return x + position_embeddings
+
+
+class MoonVision3dPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        out_dim: int,
+        in_dim: int = 3,
+        patch_size: Union[int, Tuple[int, int]] = (14, 14),
+        pos_emb_height: int = 14,
+        pos_emb_width: int = 14,
+        pos_emb_time: int = 4,
+        pos_emb_type: str = "divided_fixed",
+        pos_emb_interpolation_mode: str = "bicubic",
+        patch_embed_proj_bias: bool = True,
+    ):
+        super().__init__()
+        if isinstance(patch_size, int):
+            patch_size = (patch_size, patch_size)
+        self.patch_size = patch_size
+
+        self.proj = nn.Conv2d(
+            in_dim,
+            out_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=patch_embed_proj_bias,
+        )
+
+        if pos_emb_type != "divided_fixed":
+            raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}")
+        self.pos_emb = Learnable2DInterpPosEmbDividedFixed(
+            height=pos_emb_height,
+            width=pos_emb_width,
+            num_frames=pos_emb_time,
+            dim=out_dim,
+            interpolation_mode=pos_emb_interpolation_mode,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        grid_thws: torch.Tensor,
+        *,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+        position_embeddings: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # MIOpen can overflow grid_size for some patch shapes. Prefer AITER's
+        # Triton convolution on AMD, with an equivalent linear fallback.
+        if _use_aiter:
+            x = aiter_conv2d(
+                x,
+                self.proj.weight,
+                self.proj.bias,
+                stride=self.patch_size,
+                padding=(0, 0),
+                dilation=(1, 1),
+            ).view(x.size(0), -1)
+        elif _is_hip:
+            x = F.linear(x.flatten(1), self.proj.weight.flatten(1), self.proj.bias)
+        else:
+            x = self.proj(x).view(x.size(0), -1)
+        return self.pos_emb(
+            x,
+            grid_thws,
+            grid_thw_list=grid_thw_list,
+            position_embeddings=position_embeddings,
+        )
+
+
+class Rope2DPosEmbRepeated(nn.Module):
+    def __init__(
+        self, dim: int, max_height: int, max_width: int, theta_base: float = 10000
+    ):
+        super().__init__()
+        assert dim % 4 == 0, "dim must be divisible by 4"
+        self.dim = dim
+        self.max_height = max_height
+        self.max_width = max_width
+        self.theta_base = theta_base
+
+    def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor:
+        N = self.max_height * self.max_width
+        flat_pos = torch.arange(0, N).float().to(device)
+        x_pos = flat_pos % self.max_width
+        y_pos = flat_pos // self.max_width
+        dim_range = torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device)
+        freqs = 1.0 / (self.theta_base ** (dim_range / self.dim))
+        x_freqs = torch.outer(x_pos, freqs).float()
+        y_freqs = torch.outer(y_pos, freqs).float()
+        x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
+        y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
+        freqs_cis = torch.cat(
+            [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1
+        )
+        return freqs_cis.reshape(self.max_height, self.max_width, -1)
+
+    def get_freqs_cis(
+        self,
+        grid_thws: torch.Tensor,
+        device: torch.device,
+        *,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+    ) -> torch.Tensor:
+        if not hasattr(self, "freqs_cis"):
+            self.register_buffer(
+                "freqs_cis", self._precompute_freqs_cis(device), persistent=False
+            )
+
+        shapes = _resolve_grid_thw_list(grid_thws, grid_thw_list)
+        assert all(
+            1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes
+        ), (shapes, self.max_height, self.max_width)
+        return torch.cat(
+            [
+                self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1)
+                for t, h, w in shapes
+            ],
+            dim=0,
+        )
+
+
+class MLP2(nn.Module):
+    def __init__(self, dims: List[int], activation, bias: bool = True):
+        super().__init__()
+        assert len(dims) == 3
+        self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
+        self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
+        self.activation = activation
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc1(self.activation(self.fc0(x)))
+
+
+def _make_norm(norm_type: str, dim: int) -> nn.Module:
+    if norm_type == "layernorm":
+        return nn.LayerNorm(dim)
+    if norm_type == "rmsnorm":
+        return nn.RMSNorm(dim)
+    raise NotImplementedError(f"Not support norm_type: {norm_type}")
+
+
+class MoonViTEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        qkv_hidden_size: Optional[int] = None,
+        norm_type: str = "layernorm",
+        *,
+        activation=F.gelu,
+        attn_bias: bool = False,
+        linear_bias: bool = True,
+        attention_backend: str = "sdpa",
+        attention_workspace: Optional[torch.Tensor] = None,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.qkv_hidden_size = (
+            hidden_dim if qkv_hidden_size is None else qkv_hidden_size
+        )
+        self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads
+
+        self.norm0 = _make_norm(norm_type, hidden_dim)
+        self.norm1 = _make_norm(norm_type, hidden_dim)
+        self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation, bias=linear_bias)
+        self.wqkv = nn.Linear(hidden_dim, self.qkv_hidden_size * 3, bias=attn_bias)
+        self.wo = nn.Linear(self.qkv_hidden_size, hidden_dim, bias=attn_bias)
+        self.attention_backend = attention_backend
+        if attention_backend == "auto":
+            if not torch.cuda.is_available():
+                implementation_backends = ()
+            elif _is_hip:
+                implementation_backends = ("triton_attn",)
+            else:
+                implementation_backends = ("triton_attn", "fa4")
+        elif attention_backend == "sdpa":
+            implementation_backends = ()
+        else:
+            implementation_backends = (attention_backend,)
+        self.attention_backend_impls = nn.ModuleDict(
+            {
+                backend: QKV_BACKEND_IMPL[backend](
+                    use_data_parallel=True,
+                    workspace_buffer=attention_workspace,
+                )
+                for backend in implementation_backends
+            }
+        )
+
+    def _attention(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        segment_bounds: SegmentBounds,
+        rope_freqs_cis: torch.Tensor,
+        forward_metadata: VisionAttentionMetadata,
+        use_fused_rope: bool,
+        selected_attention_backend: str,
+    ) -> torch.Tensor:
+        xqkv = self.wqkv(x)
+        qkv_shape = xqkv.size()[:-1] + (
+            3,
+            self.num_heads,
+            self.hidden_size_per_attention_head,
+        )
+        xqkv = xqkv.view(*qkv_shape)
+        xq, xk, xv = torch.unbind(xqkv, dim=-3)
+
+        if use_fused_rope:
+            xq, xk = apply_fused_qk_complex_rope(xq, xk, rope_freqs_cis)
+        else:
+            xq, xk = apply_rope(xq, xk, rope_freqs_cis)
+
+        if selected_attention_backend == "sdpa":
+            attn_out = sdpa_varlen_attention(
+                xq,
+                xk,
+                xv,
+                cu_seqlens,
+                segment_bounds=segment_bounds,
+            )
+        else:
+            if selected_attention_backend == "flashinfer_cudnn":
+                xv = xv.contiguous()
+            attn_out = self.attention_backend_impls[selected_attention_backend](
+                xq,
+                xk,
+                xv,
+                cu_seqlens=cu_seqlens,
+                bsz=1,
+                seq_len=xq.shape[0],
+                forward_metadata=forward_metadata,
+            ).flatten(start_dim=-2)
+        return self.wo(attn_out)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        segment_bounds: SegmentBounds,
+        rope_freqs_cis: torch.Tensor,
+        forward_metadata: VisionAttentionMetadata,
+        use_fused_rope: bool,
+        selected_attention_backend: str,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.norm0(hidden_states)
+        hidden_states = self._attention(
+            hidden_states,
+            cu_seqlens,
+            segment_bounds,
+            rope_freqs_cis,
+            forward_metadata,
+            use_fused_rope,
+            selected_attention_backend,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
+
+
+@dataclass(frozen=True)
+class KimiK3VisionForwardMetadata:
+    grid_thw_list: Tuple[GridTHW, ...]
+    segment_bounds: SegmentBounds
+    rope_freqs_cis: torch.Tensor
+    attention: VisionAttentionMetadata
+    use_fused_rope: bool
+    selected_attention_backend: str
+    position_embeddings: Optional[torch.Tensor] = None
+
+
+class MoonViT3dEncoder(nn.Module):
+    def __init__(self, hidden_dim: int, num_layers: int, block_cfg: dict) -> None:
+        super().__init__()
+        qkv_hidden_size = block_cfg.get("qkv_hidden_size") or block_cfg["hidden_dim"]
+        attention_backend = _get_mm_attention_backend()
+        if attention_backend != "auto" and attention_backend not in QKV_BACKEND_IMPL:
+            raise ValueError(
+                f"Unsupported Kimi-K3 vision attention backend: {attention_backend}"
+            )
+        attention_workspace = None
+        if attention_backend == "flashinfer_cudnn" and torch.cuda.is_available():
+            attention_workspace = torch.empty(
+                FLASHINFER_WORKSPACE_SIZE_BYTES,
+                dtype=torch.uint8,
+                device=torch.device("cuda", torch.cuda.current_device()),
+            )
+        self.attention_backend = attention_backend
+        if attention_backend == "auto":
+            print_info_once(
+                "Kimi-K3 vision attention uses shape-aware auto selection on "
+                "B300/GB300 (Triton for small workloads, FA4 otherwise)."
+            )
+        self.attention_width = qkv_hidden_size
+        self.rope_2d = Rope2DPosEmbRepeated(
+            qkv_hidden_size // block_cfg["num_heads"], 512, 512
+        )
+        self.blocks = nn.ModuleList(
+            [
+                MoonViTEncoderLayer(
+                    **block_cfg,
+                    attention_backend=attention_backend,
+                    attention_workspace=attention_workspace,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_layernorm = _make_norm(
+            block_cfg.get("norm_type", "layernorm"), hidden_dim
+        )
+
+    def precompile_fused_rope(self, dtype: torch.dtype, device: torch.device) -> bool:
+        if not self.blocks:
+            return False
+        return precompile_fused_qk_complex_rope(
+            num_heads=self.blocks[0].num_heads,
+            head_dim=self.blocks[0].hidden_size_per_attention_head,
+            dtype=dtype,
+            device=device,
+        )
+
+    def precompile_attention_backend(
+        self, dtype: torch.dtype, device: torch.device
+    ) -> bool:
+        if (
+            self.attention_backend != "auto"
+            or not self.blocks
+            or device.type != "cuda"
+            or torch.cuda.get_device_capability(device) != (10, 3)
+            or not _is_fa4_available()
+        ):
+            return False
+
+        block = self.blocks[0]
+        if "fa4" not in block.attention_backend_impls:
+            return False
+
+        num_tokens = 256
+        packed_qkv = torch.zeros(
+            (
+                num_tokens,
+                3,
+                block.num_heads,
+                block.hidden_size_per_attention_head,
+            ),
+            dtype=dtype,
+            device=device,
+        )
+        q = packed_qkv[:, 0].contiguous()
+        k = packed_qkv[:, 1].contiguous()
+        v = packed_qkv[:, 2]
+        cu_seqlens = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
+        metadata = prepare_vision_attention_metadata(cu_seqlens, device=device)
+        with torch.inference_mode():
+            block.attention_backend_impls["fa4"](
+                q,
+                k,
+                v,
+                cu_seqlens=cu_seqlens,
+                bsz=1,
+                seq_len=num_tokens,
+                forward_metadata=metadata,
+            )
+        torch.cuda.synchronize(device)
+        return True
+
+    def prepare_forward_metadata(
+        self,
+        *,
+        grid_thws: torch.Tensor,
+        total_tokens: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+    ) -> KimiK3VisionForwardMetadata:
+        shapes = _resolve_grid_thw_list(grid_thws, grid_thw_list)
+        rope_freqs_cis = self.rope_2d.get_freqs_cis(
+            grid_thws=grid_thws,
+            device=device,
+            grid_thw_list=shapes,
+        )
+        cumulative_lengths = [0]
+        for t, h, w in shapes:
+            cumulative_lengths.append(cumulative_lengths[-1] + t * h * w)
+        cu_seqlens = torch.tensor(cumulative_lengths, dtype=torch.int32, device=device)
+
+        if self.attention_backend == "flashinfer_cudnn":
+            attention = prepare_flashinfer_cudnn_vision_attention_metadata(
+                cu_seqlens,
+                device=device,
+                elem_per_token=self.attention_width,
+            )
+        else:
+            attention = prepare_vision_attention_metadata(cu_seqlens, device=device)
+
+        selected_attention_backend = _resolve_mm_attention_backend(
+            self.attention_backend,
+            max_seqlen=attention.max_seqlen,
+            total_tokens=total_tokens,
+            device=device,
+        )
+        return KimiK3VisionForwardMetadata(
+            grid_thw_list=shapes,
+            segment_bounds=tuple(zip(cumulative_lengths[:-1], cumulative_lengths[1:])),
+            rope_freqs_cis=rope_freqs_cis,
+            attention=attention,
+            use_fused_rope=_can_use_fused_rope_for_shape(
+                dtype=dtype,
+                device=device,
+                freqs_cis=rope_freqs_cis,
+            ),
+            selected_attention_backend=selected_attention_backend,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        grid_thws: torch.Tensor,
+        *,
+        forward_metadata: Optional[KimiK3VisionForwardMetadata] = None,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+    ) -> torch.Tensor:
+        if forward_metadata is None:
+            forward_metadata = self.prepare_forward_metadata(
+                grid_thws=grid_thws,
+                total_tokens=hidden_states.shape[0],
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+                grid_thw_list=grid_thw_list,
+            )
+            forward_metadata = replace(
+                forward_metadata,
+                use_fused_rope=_can_use_fused_rope(
+                    hidden_states, forward_metadata.rope_freqs_cis
+                ),
+            )
+        attention = forward_metadata.attention
+        cu_seqlens = attention.cu_seqlens
+        rope_freqs_cis = forward_metadata.rope_freqs_cis
+
+        for block in self.blocks:
+            hidden_states = block(
+                hidden_states,
+                cu_seqlens,
+                forward_metadata.segment_bounds,
+                rope_freqs_cis,
+                attention,
+                forward_metadata.use_fused_rope,
+                forward_metadata.selected_attention_backend,
+            )
+
+        return self.final_layernorm(hidden_states)
+
+
+class KimiK3VisionTower(nn.Module):
+    def __init__(self, vision_config, **kwargs):
+        super().__init__()
+        config = vision_config
+        self.config = config
+        self.merge_kernel_size = tuple(config.merge_kernel_size)
+        self.patch_size = config.patch_size
+        self.merge_type = config.merge_type
+        if self.merge_type != "sd2_tpool":
+            raise NotImplementedError(f"Not support merge_type: {self.merge_type}")
+
+        hidden_size = getattr(config, "vt_hidden_size", None) or config.hidden_size
+        num_heads = (
+            getattr(config, "vt_num_attention_heads", None)
+            or config.num_attention_heads
+        )
+        num_layers = (
+            getattr(config, "vt_num_hidden_layers", None) or config.num_hidden_layers
+        )
+        intermediate_size = (
+            getattr(config, "vt_intermediate_size", None) or config.intermediate_size
+        )
+
+        self.patch_embed = MoonVision3dPatchEmbed(
+            out_dim=hidden_size,
+            patch_size=config.patch_size,
+            pos_emb_height=config.init_pos_emb_height,
+            pos_emb_width=config.init_pos_emb_width,
+            pos_emb_time=config.init_pos_emb_time,
+            pos_emb_type=config.pos_emb_type,
+            pos_emb_interpolation_mode=config.pos_emb_interpolation_mode,
+            patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True),
+        )
+
+        activation_func = getattr(config, "activation_func", "gelu_pytorch_tanh")
+        if activation_func == "gelu_pytorch_tanh":
+            activation = lambda x: F.gelu(x, approximate="tanh")
+        elif activation_func == "gelu":
+            activation = F.gelu
+        else:
+            raise NotImplementedError(f"Not support activation_func: {activation_func}")
+
+        self.encoder = MoonViT3dEncoder(
+            hidden_dim=hidden_size,
+            num_layers=num_layers,
+            block_cfg={
+                "num_heads": num_heads,
+                "hidden_dim": hidden_size,
+                "qkv_hidden_size": getattr(config, "qkv_hidden_size", None),
+                "mlp_dim": intermediate_size,
+                "norm_type": getattr(config, "norm_type", "layernorm"),
+                "activation": activation,
+                "attn_bias": getattr(config, "attn_bias", True),
+                "linear_bias": getattr(config, "linear_bias", True),
+            },
+        )
+        self.cuda_graph_runner = None
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.patch_embed.proj.weight.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.patch_embed.proj.weight.device
+
+    def precompile_fused_rope(self) -> bool:
+        return self.encoder.precompile_fused_rope(self.dtype, self.device)
+
+    def precompile_attention_backend(self) -> bool:
+        return self.encoder.precompile_attention_backend(self.dtype, self.device)
+
+    def prepare_forward_metadata(
+        self,
+        grid_thws: torch.Tensor,
+        *,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+        total_tokens: int,
+        dtype: Optional[torch.dtype] = None,
+    ) -> KimiK3VisionForwardMetadata:
+        metadata = self.encoder.prepare_forward_metadata(
+            grid_thws=grid_thws,
+            total_tokens=total_tokens,
+            dtype=self.dtype if dtype is None else dtype,
+            device=self.device,
+            grid_thw_list=grid_thw_list,
+        )
+        return replace(
+            metadata,
+            position_embeddings=self.patch_embed.pos_emb.position_embeddings(
+                grid_thws, grid_thw_list=metadata.grid_thw_list
+            ),
+        )
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thws: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+        *,
+        grid_hw: Optional[torch.Tensor] = None,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+        forward_metadata: Optional[KimiK3VisionForwardMetadata] = None,
+    ) -> List[torch.Tensor]:
+        # run_dp_sharded_mrope_vision_model calls rope_2d towers with
+        # grid_hw=/max_seqlen= keywords (#30878); K3 grids are (t, h, w) and the
+        # encoder derives its own varlen metadata, so max_seqlen is unused.
+        if grid_thws is None:
+            grid_thws = grid_hw
+        assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, grid_thws.shape
+        if (
+            forward_metadata is None
+            and pixel_values.is_cuda
+            and envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get()
+        ):
+            if grid_thw_list is None:
+                grid_thw_list = _resolve_grid_thw_list(grid_thws)
+            else:
+                grid_thw_list = _resolve_grid_thw_list(grid_thws, grid_thw_list)
+            if self.cuda_graph_runner is None:
+                from sglang.srt.multimodal.kimi_k3_vit_cuda_graph_runner import (
+                    KimiK3ViTCudaGraphRunner,
+                )
+
+                self.cuda_graph_runner = KimiK3ViTCudaGraphRunner(
+                    self,
+                    capacity=(envs.SGLANG_KIMI_K3_VIT_CUDA_GRAPH_CACHE_CAPACITY.get()),
+                    min_hits=envs.SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MIN_HITS.get(),
+                    max_seqlen=(envs.SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MAX_SEQLEN.get()),
+                )
+            return self.cuda_graph_runner.run(pixel_values, grid_thws, grid_thw_list)
+        return self._forward_eager(
+            pixel_values,
+            grid_thws,
+            grid_thw_list=grid_thw_list,
+            forward_metadata=forward_metadata,
+        )
+
+    def _forward_eager(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thws: torch.Tensor,
+        *,
+        grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
+        forward_metadata: Optional[KimiK3VisionForwardMetadata] = None,
+    ) -> List[torch.Tensor]:
+        if forward_metadata is not None:
+            grid_thw_list = forward_metadata.grid_thw_list
+        hidden_states = self.patch_embed(
+            pixel_values,
+            grid_thws,
+            grid_thw_list=grid_thw_list,
+            position_embeddings=(
+                None
+                if forward_metadata is None
+                else forward_metadata.position_embeddings
+            ),
+        )
+        hidden_states = self.encoder(
+            hidden_states,
+            grid_thws,
+            forward_metadata=forward_metadata,
+            grid_thw_list=grid_thw_list,
+        )
+        return tpool_patch_merger(
+            hidden_states,
+            grid_thws,
+            merge_kernel_size=self.merge_kernel_size,
+            grid_thw_list=grid_thw_list,
+        )
+
+
+class KimiK3MultiModalProjector(nn.Module):
+    """PatchMergerMLPV2: bias-free two-layer MLP over merged patches with a
+    post RMSNorm; K3 has no pre-norm, unlike the K2.5 projector."""
+
+    def __init__(self, vision_config):
+        super().__init__()
+        config = vision_config
+        mm_hidden_size = (
+            getattr(config, "mm_hidden_size", None) or config.vt_hidden_size
+        )
+        merge_h, merge_w = config.merge_kernel_size
+        self.hidden_size = mm_hidden_size * merge_h * merge_w
+        text_hidden_size = getattr(config, "text_hidden_size", None) or getattr(
+            config, "hidden_size"
+        )
+        eps = config.projector_ln_eps
+
+        self.proj = nn.Sequential(
+            nn.Linear(self.hidden_size, self.hidden_size, bias=False),
+            nn.GELU(),
+            nn.Linear(self.hidden_size, text_hidden_size, bias=False),
+        )
+        self.post_norm = nn.RMSNorm(text_hidden_size, eps=eps)
+
+    def forward(
+        self, image_features: Union[torch.Tensor, List[torch.Tensor]]
+    ) -> torch.Tensor:
+        if isinstance(image_features, (list, tuple)):
+            x = concat_or_single(
+                [item.reshape(item.shape[0], -1) for item in image_features]
+            )
+        else:
+            x = image_features.reshape(image_features.shape[0], -1)
+        return self.post_norm(self.proj(x))

python/sglang/srt/models/kimi_linear.py

@@ -113,6 +113,9 @@ class KimiMoE(nn.Module):
             layer_id=self.layer_idx,
             quant_config=quant_config,
             routed_scaling_factor=self.routed_scaling_factor,
+            activation=config.hidden_act,
+            gemm1_alpha=config.activation_situ_beta,
+            gemm1_clamp_limit=config.activation_situ_linear_beta,
             prefix=add_prefix("experts", prefix),
         )
 
@@ -405,7 +408,10 @@ class KimiDeltaAttention(nn.Module):
             forget_gate = forget_gate.unflatten(
                 -1, (-1, self.head_dim)
             )  # [T, H*K] -> [T, H, K]
-            beta = beta.float().sigmoid()
+            if not forward_batch.forward_mode.is_target_verify():
+                # Only chunk_kda (extend) wants pre-activated beta; the verify
+                # kernel sigmoids it in-kernel like decode.
+                beta = beta.float().sigmoid()
             forget_gate = forget_gate.unsqueeze(0)
         beta = beta.unsqueeze(0)
 
@@ -866,6 +872,12 @@ class KimiLinearForCausalLM(nn.Module):
                     weight_loader(param, loaded_weight, **kwargs)
             loaded_params.add(name)
 
+        self.post_load_weights()
+
+    def post_load_weights(self):
+        # Derive the absorbed MLA weights. Also called by the dummy loader
+        # (`_post_load_weights`), which never runs `load_weights` — keeping the
+        # derivation only there left `w_kc` None under --load-format dummy.
         for layer_id in self.config.full_attention_layer_ids:
             if not self.model.start_layer <= layer_id < self.model.end_layer:
                 continue

python/sglang/srt/models/kimi_vl_moonvit.py

@@ -570,13 +570,21 @@ def tpool_patch_merger(
     x: torch.Tensor,
     grid_thws: torch.Tensor,
     merge_kernel_size: tuple[int, int] = (2, 2),
+    *,
+    grid_thw_list: Optional[Sequence[Sequence[int]]] = None,
 ) -> List[torch.Tensor]:
-    """Group spatial patches and average only across real video frames."""
+    """Group spatial patches and average only across real video frames.
+
+    ``grid_thw_list`` lets a graph-aware tower pass the host grid it already
+    has instead of paying a device sync for ``grid_thws.tolist()``.
+    """
 
     d_model = x.size(-1)
     outputs = []
     pre_sum = 0
-    for t, h, w in grid_thws.tolist():
+    shapes = grid_thws.tolist() if grid_thw_list is None else grid_thw_list
+    for t, h, w in shapes:
+        t, h, w = int(t), int(h), int(w)
         seq = x[pre_sum : pre_sum + t * h * w]
         kernel_height, kernel_width = merge_kernel_size
         new_height, new_width = h // kernel_height, w // kernel_width

python/sglang/srt/multimodal/kimi_k3_vit_cuda_graph_runner.py

@@ -0,0 +1,211 @@
+"""Bounded full-tower CUDA graph runner for Kimi K3 vision."""
+
+from __future__ import annotations
+
+import logging
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Hashable, List, Tuple
+
+import torch
+
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.models.kimi_k3_vl import (
+        GridTHW,
+        KimiK3VisionForwardMetadata,
+        KimiK3VisionTower,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _CapturedVisionGraph:
+    graph: torch.cuda.CUDAGraph
+    input_buffer: torch.Tensor
+    outputs: Tuple[torch.Tensor, ...]
+    metadata: KimiK3VisionForwardMetadata
+
+
+class KimiK3ViTCudaGraphRunner:
+    def __init__(
+        self,
+        tower: KimiK3VisionTower,
+        *,
+        capacity: int,
+        min_hits: int,
+        max_seqlen: int | None = None,
+    ) -> None:
+        if capacity <= 0:
+            raise ValueError(f"capacity must be positive, got {capacity}")
+        if min_hits <= 0:
+            raise ValueError(f"min_hits must be positive, got {min_hits}")
+        if max_seqlen is not None and max_seqlen <= 0:
+            raise ValueError(f"max_seqlen must be positive, got {max_seqlen}")
+        self.tower = tower
+        self.capacity = capacity
+        self.min_hits = min_hits
+        self.max_seqlen = max_seqlen
+        self.graphs: dict[Hashable, _CapturedVisionGraph] = {}
+        self.seen: OrderedDict[Hashable, int] = OrderedDict()
+        self.failed_keys: set[Hashable] = set()
+        self._graph_memory_pool: Any = None
+        self._capacity_logged = False
+        self._max_seqlen_logged = False
+        logger.info(
+            "Kimi-K3 ViT CUDA graph: capacity=%d min_hits=%d max_seqlen=%s",
+            capacity,
+            min_hits,
+            max_seqlen,
+        )
+
+    @staticmethod
+    def graph_key(grid_thw_list: Tuple[GridTHW, ...]) -> Hashable:
+        return grid_thw_list
+
+    def _record_hit(self, key: Hashable) -> int:
+        count = self.seen.pop(key, 0) + 1
+        self.seen[key] = count
+        seen_capacity = max(self.capacity * 8, self.capacity)
+        if len(self.seen) > seen_capacity:
+            self.seen.popitem(last=False)
+        return count
+
+    def _run_eager(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thws: torch.Tensor,
+        grid_thw_list: Tuple[GridTHW, ...],
+    ) -> tuple[List[torch.Tensor], KimiK3VisionForwardMetadata]:
+        metadata = self.tower.prepare_forward_metadata(
+            grid_thws,
+            grid_thw_list=grid_thw_list,
+            total_tokens=pixel_values.shape[0],
+            dtype=pixel_values.dtype,
+        )
+        outputs = self.tower._forward_eager(
+            pixel_values,
+            grid_thws,
+            grid_thw_list=grid_thw_list,
+            forward_metadata=metadata,
+        )
+        return outputs, metadata
+
+    def _capture(
+        self,
+        key: Hashable,
+        pixel_values: torch.Tensor,
+        grid_thws: torch.Tensor,
+        grid_thw_list: Tuple[GridTHW, ...],
+        metadata: KimiK3VisionForwardMetadata,
+    ) -> _CapturedVisionGraph:
+        allocated_before = torch.cuda.memory_allocated(pixel_values.device)
+        reserved_before = torch.cuda.memory_reserved(pixel_values.device)
+        input_buffer = torch.empty_like(pixel_values)
+        input_buffer.copy_(pixel_values)
+        graph = torch.cuda.CUDAGraph()
+        if self._graph_memory_pool is None:
+            self._graph_memory_pool = get_or_create_global_graph_memory_pool(torch.cuda)
+        with torch.cuda.graph(graph, pool=self._graph_memory_pool):
+            outputs = self.tower._forward_eager(
+                input_buffer,
+                grid_thws,
+                grid_thw_list=grid_thw_list,
+                forward_metadata=metadata,
+            )
+        torch.cuda.synchronize(pixel_values.device)
+        entry = _CapturedVisionGraph(
+            graph=graph,
+            input_buffer=input_buffer,
+            outputs=tuple(outputs),
+            metadata=metadata,
+        )
+        position_embeddings = getattr(metadata, "position_embeddings", None)
+        position_cache_bytes = (
+            position_embeddings.numel() * position_embeddings.element_size()
+            if position_embeddings is not None
+            else 0
+        )
+        logger.info(
+            "Captured Kimi-K3 ViT CUDA graph: key=%s cache=%d/%d "
+            "allocated_delta_mib=%.1f reserved_delta_mib=%.1f "
+            "position_cache_mib=%.1f allocated_total_mib=%.1f "
+            "reserved_total_mib=%.1f",
+            key,
+            len(self.graphs) + 1,
+            self.capacity,
+            (torch.cuda.memory_allocated(pixel_values.device) - allocated_before)
+            / 2**20,
+            (torch.cuda.memory_reserved(pixel_values.device) - reserved_before) / 2**20,
+            position_cache_bytes / 2**20,
+            torch.cuda.memory_allocated(pixel_values.device) / 2**20,
+            torch.cuda.memory_reserved(pixel_values.device) / 2**20,
+        )
+        return entry
+
+    def run(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thws: torch.Tensor,
+        grid_thw_list: Tuple[GridTHW, ...],
+    ) -> List[torch.Tensor]:
+        key = self.graph_key(grid_thw_list)
+        entry = self.graphs.get(key)
+        if entry is not None:
+            entry.input_buffer.copy_(pixel_values)
+            entry.graph.replay()
+            return list(entry.outputs)
+
+        max_seqlen = max(t * h * w for t, h, w in grid_thw_list)
+        if self.max_seqlen is not None and max_seqlen > self.max_seqlen:
+            if not self._max_seqlen_logged:
+                logger.info(
+                    "Kimi-K3 ViT CUDA graph uses eager fallback above "
+                    "max_seqlen=%d to avoid low-value captures and HBM use",
+                    self.max_seqlen,
+                )
+                self._max_seqlen_logged = True
+            outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list)
+            return outputs
+
+        if key in self.failed_keys or self._record_hit(key) < self.min_hits:
+            outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list)
+            return outputs
+
+        if len(self.graphs) >= self.capacity:
+            if not self._capacity_logged:
+                logger.warning(
+                    "Kimi-K3 ViT CUDA graph cache reached capacity=%d; "
+                    "new shapes use eager fallback to bound HBM",
+                    self.capacity,
+                )
+                self._capacity_logged = True
+            outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list)
+            return outputs
+
+        metadata = self.tower.prepare_forward_metadata(
+            grid_thws,
+            grid_thw_list=grid_thw_list,
+            total_tokens=pixel_values.shape[0],
+            dtype=pixel_values.dtype,
+        )
+        try:
+            entry = self._capture(key, pixel_values, grid_thws, grid_thw_list, metadata)
+        except Exception:
+            self.failed_keys.add(key)
+            logger.exception(
+                "Kimi-K3 ViT CUDA graph capture failed for key=%s; "
+                "using eager fallback",
+                key,
+            )
+            outputs, _ = self._run_eager(pixel_values, grid_thws, grid_thw_list)
+            return outputs
+
+        self.graphs[key] = entry
+        entry.input_buffer.copy_(pixel_values)
+        entry.graph.replay()
+        return list(entry.outputs)

python/sglang/srt/multimodal/mm_utils.py

@@ -557,9 +557,11 @@ def run_dp_sharded_mrope_vision_model(
     grid_thw_list: list,
     *,
     rope_type: Literal["rope_3d", "rope_2d", "rope_2d_packed"],
+    pool_temporal_dimension: bool = False,
     load_local_pixel_values: Optional[Callable[[list[int]], torch.Tensor]] = None,
     pixel_values_device: Optional[torch.device] = None,
     pixel_values_dtype: Optional[torch.dtype] = None,
+    pass_grid_thw_list: bool = False,
 ):
     """Run a vision model with data parallelism (DP) sharding.
     The function will shard the input image tensor on the
@@ -576,6 +578,11 @@ def run_dp_sharded_mrope_vision_model(
                    "rope_2d" for packed 2D rope outputs (e.g., Kimi-VL)
                    "rope_2d_packed" for packed 2D rope outputs that accept
                    ``grid_thws`` positionally (e.g., Kimi-K2.5/K2.7)
+        pool_temporal_dimension: Whether the vision model pools away the temporal
+                   grid dimension. Its output length is then h * w divided by
+                   the spatial merge area instead of t * h * w divided by it.
+        pass_grid_thw_list: Forward the existing host grid list to the vision
+            model so graph-aware towers do not materialize it from a CUDA tensor.
     Returns:
         torch.Tensor: Output image embeddings
 
@@ -614,11 +621,13 @@ def run_dp_sharded_mrope_vision_model(
             device=pixel_values.device if rope_type == "rope_2d" else None,
         )
         if rope_type == "rope_2d":
-            image_embeds = vision_model(
-                pixel_values,
-                grid_hw=grid_thw,
-                max_seqlen=max(math.prod(grid) for grid in grid_thw_list),
-            )
+            kwargs = {
+                "grid_hw": grid_thw,
+                "max_seqlen": max(math.prod(grid) for grid in grid_thw_list),
+            }
+            if pass_grid_thw_list:
+                kwargs["grid_thw_list"] = grid_thw_list
+            image_embeds = vision_model(pixel_values, **kwargs)
             # MoonViT returns one tensor per image. The multi-GPU path below
             # already concatenates these tensors before returning, so keep the
             # TP=1 DP-encoder path on the same projector-facing contract.
@@ -687,7 +696,9 @@ def run_dp_sharded_mrope_vision_model(
         )
 
     output_tokens_per_image = [
-        math.prod(grid) // embed_dim_reduction_factor for grid in grid_thw_list
+        math.prod(grid[1:] if pool_temporal_dimension else grid)
+        // embed_dim_reduction_factor
+        for grid in grid_thw_list
     ]
     grouped_output_lengths = []
     assignment_offset = 0
@@ -717,11 +728,13 @@ def run_dp_sharded_mrope_vision_model(
                 device=(pixel_values_local.device if rope_type == "rope_2d" else None),
             )
             if rope_type == "rope_2d":
-                image_embeds_local = vision_model(
-                    pixel_values_local,
-                    grid_hw=local_grid_thw,
-                    max_seqlen=max(math.prod(grid) for grid in local_grid_thw_list),
-                )
+                kwargs = {
+                    "grid_hw": local_grid_thw,
+                    "max_seqlen": max(math.prod(grid) for grid in local_grid_thw_list),
+                }
+                if pass_grid_thw_list:
+                    kwargs["grid_thw_list"] = local_grid_thw_list
+                image_embeds_local = vision_model(pixel_values_local, **kwargs)
             else:
                 image_embeds_local = vision_model(pixel_values_local, local_grid_thw)
             if isinstance(image_embeds_local, list):

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

@@ -1,7 +1,7 @@
 import math
 import re
 from collections import defaultdict
-from typing import Dict, List, Union
+from typing import Callable, Dict, List, Optional, Union
 
 import numpy as np
 import torch
@@ -194,15 +194,24 @@ def _process_single_image(
     image_scale: torch.Tensor,
     image_bias: torch.Tensor,
     patch_size: int,
+    to_chw: Callable[[Union[torch.Tensor, Image.Image]], torch.Tensor] = _to_cuda_chw,
+    post_resize: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
 ) -> torch.Tensor:
-    """Process a single image on GPU: resize -> pad -> normalize -> patchify."""
-    image = _to_cuda_chw(image)
+    """Process a single image on GPU: resize -> pad -> normalize -> patchify.
+
+    ``to_chw`` converts the input to a CUDA CHW tensor (a model may keep an
+    alpha channel here); ``post_resize`` runs on the resized ``(B, C, H, W)``
+    batch before patchify (K3 composites transparent backgrounds there).
+    """
+    image = to_chw(image)
 
     new_h, new_w = config["new_height"], config["new_width"]
     padded_h = new_h + config["pad_height"]
     padded_w = new_w + config["pad_width"]
 
     x = _resize_bicubic_if_needed(image.unsqueeze(0), new_h, new_w)
+    if post_resize is not None:
+        x = post_resize(x)
 
     return normalize_and_patchify(
         x, image_scale, image_bias, patch_size, padded_h, padded_w
@@ -250,6 +259,8 @@ def _gpu_preprocess_images(
     image_scale: torch.Tensor,
     image_bias: torch.Tensor,
     patch_size: int,
+    to_chw: Callable[[Union[torch.Tensor, Image.Image]], torch.Tensor] = _to_cuda_chw,
+    post_resize: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """GPU preprocessing pipeline for a batch of images.
 
@@ -278,21 +289,30 @@ def _gpu_preprocess_images(
         if len(group) == 1:
             idx, image, config = group[0]
             patches = _process_single_image(
-                image, config, image_scale, image_bias, patch_size
+                image,
+                config,
+                image_scale,
+                image_bias,
+                patch_size,
+                to_chw=to_chw,
+                post_resize=post_resize,
             )
             all_patches[idx] = patches
             all_grids[idx] = _grid_thw_from_resize_config(config, patch_size)
         else:
-            indexed_images = [(idx, _to_cuda_chw(image)) for idx, image, _ in group]
+            indexed_images = [(idx, to_chw(image)) for idx, image, _ in group]
 
             # One NaViT target group can include several original resolutions.
             # Batch only source-compatible images, which removes redundant
             # bicubic launches for common multi-image requests without padding
             # random-size inputs to a larger source resolution.
-            batch = torch.cat(
-                _resize_images_by_source_shape(indexed_images, target_h, target_w),
-                dim=0,
-            )
+            resized = _resize_images_by_source_shape(indexed_images, target_h, target_w)
+            if post_resize is not None:
+                # Before the concat: a hook may change the channel count (K3
+                # composites RGBA onto a background and returns RGB), and mixed
+                # 3/4-channel sources cannot be concatenated first.
+                resized = [post_resize(part) for part in resized]
+            batch = torch.cat(resized, dim=0)
 
             T = 1
             gh, gw = padded_h // patch_size, padded_w // patch_size

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

@@ -0,0 +1,438 @@
+"""Kimi K3 multimodal processor.
+
+GPU image preprocessing dedicated to K3: unlike the K2.5 wrapper it keeps
+the alpha channel through the bicubic resize and then composites RGBA
+images onto the checkpoint-configured background
+(``transparent_bg_config`` with ``transparent_bg_fill_stage ==
+"after_resize"`` in preprocessor_config.json), instead of dropping alpha
+at load time.
+"""
+
+import re
+from typing import Dict, List, Union
+
+import numpy as np
+import torch
+from PIL import Image
+
+from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput
+from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
+from sglang.srt.multimodal.processors.base_processor import (
+    BaseMultimodalProcessor as SGLangBaseProcessor,
+)
+from sglang.srt.multimodal.processors.base_processor import (
+    MultimodalSpecialTokens,
+)
+from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin
+from sglang.srt.multimodal.processors.kimi_k25 import (
+    KimiGPUProcessorWrapper,
+    _get_image_dimensions,
+    _gpu_preprocess_images,
+    navit_resize_config,
+)
+from sglang.srt.utils.cuda_ipc_transport_utils import (
+    DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
+)
+
+
+def _encode_k3_special_tokens(tokenizer, text: str) -> list[int]:
+    """Encode K3 control tokens without allowing them to be BPE-split."""
+    try:
+        return list(tokenizer.encode(text, allowed_special="all"))
+    except TypeError:
+        # Keep the helper usable with lightweight tokenizer stubs in CPU tests.
+        return list(tokenizer.encode(text))
+
+
+def _expand_k3_image_prompt_token_ids(
+    input_ids: Union[List[int], torch.Tensor],
+    image_token_id: int,
+    image_token_counts: List[int],
+    image_sizes: List[tuple[int, int]],
+    tokenizer,
+) -> torch.Tensor:
+    """Expand K3 image placeholders into the checkpoint's media contract.
+
+    K3 requires each image feature span to be enclosed by its original uploaded
+    dimensions.  The chat template deliberately emits one ``media_pad`` per
+    image; after decode, insert the surrounding control tokens and expand that
+    one placeholder to the NaViT feature count.
+    """
+    if len(image_token_counts) != len(image_sizes):
+        raise ValueError("Expected one original size for each K3 image.")
+
+    if isinstance(input_ids, torch.Tensor):
+        input_ids = input_ids.detach().flatten().cpu().numpy()
+    input_ids = np.asarray(input_ids, dtype=np.int64)
+
+    placeholder_count = np.count_nonzero(input_ids == image_token_id)
+    if placeholder_count != len(image_token_counts):
+        raise ValueError(
+            f"Expected {len(image_token_counts)} image placeholder token(s), "
+            f"found {placeholder_count}."
+        )
+
+    output = []
+    image_index = 0
+    for token_id in input_ids:
+        if token_id != image_token_id:
+            output.append(int(token_id))
+            continue
+
+        width, height = image_sizes[image_index]
+        output.extend(
+            _encode_k3_special_tokens(
+                tokenizer,
+                f"<|media_begin|>image {width}x{height}<|media_content|>",
+            )
+        )
+        output.extend([image_token_id] * image_token_counts[image_index])
+        output.extend(_encode_k3_special_tokens(tokenizer, "<|media_end|>"))
+        image_index += 1
+
+    return torch.tensor(output, dtype=torch.long).unsqueeze(0)
+
+
+def _expand_k3_image_prompt_text(
+    input_text: str,
+    image_token: str,
+    image_token_counts: List[int],
+    image_sizes: List[tuple[int, int]],
+) -> str:
+    """Render the K3 media framing for the CPU HF-processor fallback."""
+    parts = input_text.split(image_token)
+    if len(parts) - 1 != len(image_token_counts):
+        raise ValueError(
+            f"Expected {len(image_token_counts)} image placeholder(s), "
+            f"found {len(parts) - 1}."
+        )
+
+    output = [parts[0]]
+    for image_token_count, (width, height), suffix in zip(
+        image_token_counts, image_sizes, parts[1:]
+    ):
+        output.extend(
+            (
+                f"<|media_begin|>image {width}x{height}<|media_content|>",
+                image_token * image_token_count,
+                "<|media_end|>",
+                suffix,
+            )
+        )
+    return "".join(output)
+
+
+def _k3_to_cuda_chw(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor:
+    if isinstance(image, Image.Image):
+        # The checkpoint's fill_transparent_bg_with() returns RGB-mode images
+        # untouched before it ever inspects the alpha bands, so an RGB image
+        # carrying a stray "transparency" info key must NOT be promoted to
+        # RGBA here.
+        has_alpha = image.mode != "RGB" and (
+            "A" in image.getbands() or "transparency" in image.info
+        )
+        arr = np.asarray(image.convert("RGBA" if has_alpha else "RGB"))
+        return torch.from_numpy(arr).permute(2, 0, 1).cuda()
+
+    image = image.cuda()
+    if image.dim() == 2:
+        image = image.unsqueeze(0)
+    if image.shape[0] == 1:
+        image = image.repeat(3, 1, 1)
+    return image
+
+
+def _chessboard_background(
+    height: int, width: int, cfg: dict, device: torch.device
+) -> torch.Tensor:
+    square = cfg.get("chessboard_square_size", 16)
+    white = float(cfg.get("chessboard_white_value", 255))
+    gray = float(cfg.get("chessboard_gray_value", 200))
+    on_top_left = cfg.get("chessboard_square_on_top_left", True)
+
+    ys = torch.arange(height, device=device) // square
+    xs = torch.arange(width, device=device) // square
+    parity = (ys.unsqueeze(1) + xs.unsqueeze(0)) % 2
+    gray_parity = 1 if on_top_left else 0
+    bg = torch.where(parity == gray_parity, gray, white)
+    return bg.unsqueeze(0).expand(3, height, width)
+
+
+def _fill_transparent_bg(x: torch.Tensor, bg_cfg: Union[dict, None]) -> torch.Tensor:
+    """Composite a resized (1, 4, H, W) float image in [0, 255] onto the
+    configured background; 3-channel input passes through."""
+    if x.shape[1] == 3:
+        return x
+    rgb = x[:, :3]
+    if bg_cfg is None:
+        return rgb
+
+    _, _, height, width = x.shape
+    pattern = bg_cfg.get("pattern", "black")
+    if pattern == "chessboard":
+        bg = _chessboard_background(height, width, bg_cfg, x.device)
+    elif pattern == "white":
+        bg = torch.full((3, height, width), 255.0, device=x.device)
+    elif pattern == "black":
+        bg = torch.zeros(3, height, width, device=x.device)
+    elif pattern == "gray":
+        bg = torch.full((3, height, width), 128.0, device=x.device)
+    else:
+        raise ValueError(f"Invalid background pattern: {pattern}")
+
+    alpha = (x[:, 3:4] / 255.0).clamp(0.0, 1.0)
+    # The checkpoint processor casts the composited float result back with
+    # numpy's astype(np.uint8), which truncates; floor matches that exactly
+    # (a composite of [0, 255] inputs is always non-negative).
+    return (alpha * rgb + (1.0 - alpha) * bg).clamp(0.0, 255.0).floor_()
+
+
+class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper):
+    def __init__(self, *args, transparent_bg_config=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._transparent_bg_config = transparent_bg_config
+
+    def _prepare_input_ids(
+        self, input_text, resize_configs, original_input_ids, image_sizes
+    ):
+        image_token_counts = [config["num_tokens"] for config in resize_configs]
+        if original_input_ids is None:
+            original_input_ids = _encode_k3_special_tokens(
+                self._hf_processor.tokenizer, input_text
+            )
+        return _expand_k3_image_prompt_token_ids(
+            original_input_ids,
+            self._image_token_id,
+            image_token_counts,
+            image_sizes,
+            self._hf_processor.tokenizer,
+        )
+
+    def __call__(self, text=None, images=None, **kwargs):
+        images = images or kwargs.pop("images", None)
+        original_input_ids = kwargs.pop("sglang_original_input_ids", None)
+        if images and torch.cuda.is_available():
+            return self._gpu_call(text, images, original_input_ids)
+        return self._cpu_call(text, images, original_input_ids, **kwargs)
+
+    def _gpu_call(self, text, images, original_input_ids=None):
+        input_text = text[0] if isinstance(text, list) else text
+
+        resize_configs = []
+        image_sizes = []
+        for image in images:
+            w, h = _get_image_dimensions(image)
+            image_sizes.append((w, h))
+            resize_configs.append(
+                navit_resize_config(
+                    w,
+                    h,
+                    self._patch_size,
+                    self._merge_kernel_size,
+                    self._in_patch_limit,
+                    self._patch_limit_on_one_side,
+                    self._fixed_output_tokens,
+                )
+            )
+
+        input_ids = self._prepare_input_ids(
+            input_text, resize_configs, original_input_ids, image_sizes
+        )
+
+        image_scale, image_bias = self._get_gpu_norm_tensors()
+        # Shared source-compatible batched pipeline (same as K2.5): RGBA
+        # inputs land in their own source-shape groups, and the
+        # transparent-background compositing runs on each resized batch
+        # before patchify -- identical order to the previous per-image path.
+        pixel_values, grid_thws = _gpu_preprocess_images(
+            images,
+            resize_configs,
+            image_scale,
+            image_bias,
+            self._patch_size,
+            to_chw=_k3_to_cuda_chw,
+            post_resize=lambda x: _fill_transparent_bg(x, self._transparent_bg_config),
+        )
+
+        return {
+            "input_ids": input_ids,
+            "pixel_values": pixel_values,
+            "image_grid_thw": grid_thws,
+        }
+
+    def _cpu_call(self, text, images, original_input_ids=None, **kwargs):
+        """HF fallback with the same K3 media framing as the GPU path."""
+        input_text = text[0] if isinstance(text, list) else text
+        if not images:
+            return self._hf_processor(text=[input_text], **kwargs)
+
+        image_sizes = [_get_image_dimensions(image) for image in images]
+        image_token_counts = [
+            self._hf_processor.media_processor.media_tokens_calculator(
+                {"type": "image", "image": image}
+            )
+            for image in images
+        ]
+        expanded_text = _expand_k3_image_prompt_text(
+            input_text,
+            self._image_token,
+            image_token_counts,
+            image_sizes,
+        )
+        kwargs["medias"] = [{"type": "image", "image": image} for image in images]
+        out = self._hf_processor(text=[expanded_text], **kwargs)
+        out["input_ids"] = self._prepare_input_ids(
+            input_text,
+            [{"num_tokens": count} for count in image_token_counts],
+            original_input_ids,
+            image_sizes,
+        )
+        grid_thws = out.pop("grid_thws", None)
+        if grid_thws is not None:
+            out["image_grid_thw"] = grid_thws
+        return out
+
+
+class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
+    models = [KimiK3ForConditionalGeneration]
+    gpu_image_decode = True
+    prefer_tokenized_input = True
+    precompute_hash_before_cpu_transfer = True
+    auto_mm_processor_worker_num = 2
+    auto_mm_io_worker_num = 16
+    supports_mm_processor_concurrency = True
+    preserve_processor_input_ids = True
+
+    def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
+        mm_tokens = MultimodalSpecialTokens(
+            image_token="<|media_pad|>",
+            image_token_id=hf_config.media_placeholder_token_id,
+            image_token_regex=re.compile(r"(?:<\|media_pad\|>)+"),
+        ).build(_processor)
+
+        media_proc_cfg = _processor.media_processor.media_proc_cfg
+
+        processor = KimiK3GPUProcessorWrapper(
+            _processor,
+            image_token=mm_tokens.image_token,
+            image_token_id=mm_tokens.image_token_id,
+            patch_size=media_proc_cfg["patch_size"],
+            merge_kernel_size=media_proc_cfg["merge_kernel_size"],
+            in_patch_limit=media_proc_cfg["in_patch_limit"],
+            patch_limit_on_one_side=media_proc_cfg["patch_limit_on_one_side"],
+            fixed_output_tokens=media_proc_cfg.get("fixed_output_tokens"),
+            image_mean=media_proc_cfg["image_mean"],
+            image_std=media_proc_cfg["image_std"],
+            transparent_bg_config=media_proc_cfg.get("transparent_bg_config"),
+        )
+        super().__init__(hf_config, server_args, processor, *args, **kwargs)
+        self.mm_tokens = mm_tokens
+
+    async def process_mm_data_async(
+        self,
+        image_data: List[Union[str, bytes, Dict]],
+        input_text,
+        request_obj,
+        *args,
+        **kwargs,
+    ):
+        if getattr(request_obj, "video_data", None) or kwargs.get("audio_data"):
+            raise ValueError("Kimi-K3 supports image input only")
+
+        expected_image_count = len(image_data or [])
+        placeholder_count = self.count_image_placeholders(
+            input_text, self.mm_tokens.image_token_id
+        )
+        if placeholder_count is not None:
+            if placeholder_count != expected_image_count:
+                raise ValueError(
+                    "Kimi image placeholders must map one-to-one to image data: "
+                    f"expected {expected_image_count}, found {placeholder_count} token(s)"
+                )
+            # Keep structural media tokens distinct from user text that happens to
+            # spell ``<|media_pad|>``. Decoding the whole prompt and matching the
+            # resulting string would lose that distinction and could bind an image
+            # to user-provided text instead of the renderer-inserted token.
+            base_output = await self.fast_load_mm_data(
+                prompt=input_text,
+                image_data=image_data,
+                multimodal_tokens=self.mm_tokens,
+                discard_alpha_channel=False,
+                # Unlike load_mm_data, fast_load_mm_data does not derive
+                # input_ids from the prompt. Without this the wrapper falls
+                # back to re-encoding the decoded string, which is the loss of
+                # the structural/user distinction described above.
+                input_ids=input_text,
+            )
+        else:
+            base_output = await self.load_mm_data(
+                prompt=input_text,
+                image_data=image_data,
+                multimodal_tokens=self.mm_tokens,
+                discard_alpha_channel=False,
+            )
+
+        if len(base_output.images) != expected_image_count:
+            raise ValueError(
+                "Kimi image placeholders must map one-to-one to image data: "
+                f"expected {expected_image_count}, loaded {len(base_output.images)}"
+            )
+
+        mm_items, input_ids, _ = await self.process_and_combine_mm_data_async(
+            base_output,
+            self.mm_tokens,
+            sglang_original_input_ids=base_output.input_ids,
+        )
+
+        # K3's tower is unconditionally image-wise data-parallel (each image
+        # is consumed by exactly one TP rank), so keep IPC proxies lazy until
+        # that assignment is known: one tokenizer/scheduler crossing per
+        # image instead of one per rank. K2.5 gates this on
+        # --mm-enable-dp-encoder; K3 needs no flag.
+        if getattr(self, "use_cuda_ipc", False):
+            for item in mm_items:
+                item.model_specific_data[DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY] = (
+                    True
+                )
+
+        return MultimodalProcessorOutput(
+            input_ids=input_ids.tolist(),
+            mm_items=mm_items,
+            im_token_id=self.mm_tokens.image_token_id,
+        )
+
+    def get_mm_data(self, prompt, embeddings, **kwargs):
+        img_grid_thw = kwargs.get("img_grid_thw", None)
+        output = self._build_kimi_mm_data_from_grids(
+            prompt=prompt,
+            embeddings=embeddings,
+            image_token_id=self.mm_tokens.image_token_id,
+            img_grid_thw=img_grid_thw,
+        )
+        image_sizes = kwargs.get("original_image_sizes")
+        if image_sizes is None:
+            return output
+
+        counts = [self._num_image_tokens_from_grid(grid) for grid in img_grid_thw]
+        if len(image_sizes) != len(counts):
+            raise ValueError(
+                "Expected one original image size for each K3 encoder grid."
+            )
+        output.input_ids = (
+            _expand_k3_image_prompt_token_ids(
+                prompt,
+                self.mm_tokens.image_token_id,
+                counts,
+                [tuple(size) for size in image_sizes],
+                self._tokenizer,
+            )
+            .flatten()
+            .tolist()
+        )
+
+        search_start = 0
+        for item, count in zip(output.mm_items, counts):
+            start = output.input_ids.index(self.mm_tokens.image_token_id, search_start)
+            item.offsets = [(start, start + count - 1)]
+            search_start = start + count
+        return output

python/sglang/srt/server_args.py

@@ -78,6 +78,7 @@ from sglang.srt.utils.common import (
     is_hip,
     is_hopper_with_cuda_12_3,
     is_host_cpu_arm64,
+    is_mnnvl_fabric_device,
     is_mps,
     is_musa,
     is_no_spec_infer_or_topk_one,
@@ -313,7 +314,23 @@ RL_ON_POLICY_TARGET_CHOICES = ["fsdp"]
 
 LORA_BACKEND_CHOICES = ["triton", "csgmv", "ascend", "torch_native"]
 
-ENCODER_TRANSFER_BACKEND_CHOICES = ["zmq_to_scheduler", "zmq_to_tokenizer", "mooncake"]
+ENCODER_TRANSFER_BACKEND_CHOICES = [
+    "auto",
+    "zmq_to_scheduler",
+    "zmq_to_tokenizer",
+    "mooncake",
+]
+
+
+def resolve_encoder_transfer_backend(
+    backend: str, model_arch: str, tp_size: int
+) -> str:
+    if backend != "auto":
+        return backend
+    if model_arch == "KimiK3ForConditionalGeneration" and tp_size > 1:
+        return "zmq_to_tokenizer"
+    return "zmq_to_scheduler"
+
 
 DSA_PREFILL_CP_SPLIT_CHOICES = ["in-seq-split", "round-robin-split"]
 NSA_PREFILL_CP_SPLIT_CHOICES = DSA_PREFILL_CP_SPLIT_CHOICES  # deprecated alias
@@ -348,7 +365,14 @@ MAMBA_RADIX_CACHE_STRATEGY_CHOICES = [
 
 MAMBA_BACKEND_CHOICES = ["triton", "flashinfer"]
 
-LINEAR_ATTN_KERNEL_BACKEND_CHOICES = ["triton", "cutedsl", "flashinfer", "flashkda"]
+LINEAR_ATTN_KERNEL_BACKEND_CHOICES = [
+    "triton",
+    "cutedsl",
+    "flashinfer",
+    "flashkda",
+    "nvidia_kda",
+    "ptx_kda",
+]
 
 
 # Allow external code to add more choices
@@ -1866,7 +1890,12 @@ class ServerArgs:
         NS("exec.comm"),
     ] = False
     enable_symm_mem: A[
-        bool, "Enable NCCL symmetric memory for fast collectives.", NS("exec.comm")
+        bool,
+        Arg(
+            help="Enable NCCL symmetric memory for fast collectives.",
+            resolvable=True,
+        ),
+        NS("exec.comm"),
     ] = False
     triton_attention_reduce_in_fp32: A[
         bool,
@@ -2097,6 +2126,7 @@ class ServerArgs:
         Arg(
             help="Attention backend for speculative decoding operations (both target verify and draft extend). Can be one of 'prefill' (default) or 'decode'.",
             choices=["prefill", "decode"],
+            resolvable=True,
         ),
         NS("spec"),
     ] = "prefill"
@@ -2523,6 +2553,14 @@ class ServerArgs:
         ),
         NS("exec.mamba"),
     ] = None
+    linear_attn_verify_backend: A[
+        Optional[str],
+        Arg(
+            help="Override the kernel backend for linear attention speculative target-verify. If not set, follows the decode backend (flashinfer decode -> flashinfer verify, otherwise triton). KDA supports triton, nv_cutedsl, and flashinfer verify backends.",
+            choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES + ["nv_cutedsl"],
+        ),
+        NS("exec.mamba"),
+    ] = None
     # ReplaySSM buffered output-only linear-attn decode (GDN + KDA): per-slot
     # ring + periodic flush to cut per-step HBM state traffic.
     enable_linear_replayssm: A[
@@ -2542,14 +2580,15 @@ class ServerArgs:
         "Ring-buffer length L for ReplaySSM linear-attn decode. The full recurrent state is flushed to HBM every L decode steps.",
         NS("exec.mamba"),
     ] = 16
-    # ReplaySSM spec-verify (Part B of RFC #28511): GDN linear-chain target-verify
-    # via fold-every-commit instead of per-draft full-state snapshots -- the
-    # verify stores each draft step's raw inputs into a per-slot window and the
-    # commit replays the accepted prefix into the fp32 checkpoint. GDN only;
-    # linear-chain (topk <= 1) only.
-    enable_gdn_replayssm_spec: A[
+    # ReplaySSM spec-verify (Part B of RFC #28511): linear-attn target-verify via
+    # fold-every-commit instead of per-draft full-state snapshots -- the verify
+    # stores each draft step's raw inputs into a per-slot window and the commit
+    # replays the accepted prefix into the fp32 checkpoint. GDN sizes the window
+    # to the draft maximum; KDA folds a (raw v, pre-norm k, gate, beta) ring of
+    # length --linear-replayssm-cache-len. Linear-chain (topk <= 1) only.
+    enable_linear_replayssm_spec: A[
         bool,
-        "Enable the ReplaySSM GDN spec-verify (Part B of RFC #28511): fold-every-commit -- a per-slot raw-input window sized to the draft maximum replaces the recurrent verify's per-draft full-state snapshots. GDN only, linear-chain (--speculative-eagle-topk in {None, 1}) only.",
+        "Enable the ReplaySSM spec-verify: fold-every-commit -- a per-slot raw-input window replaces the recurrent verify's per-draft full-state snapshots. GDN or KDA hybrid linear-attn models, linear-chain (--speculative-eagle-topk in {None, 1}) only.",
         NS("exec.mamba"),
     ] = False
 
@@ -2704,7 +2743,10 @@ class ServerArgs:
     mm_feature_transport: A[
         Optional[Literal["cpu", "cuda_ipc"]],
         "Transport multimodal features through CPU memory or a bounded CUDA IPC pool. "
-        "The default is CPU transport; CUDA IPC reserves GPU memory on the base GPU.",
+        "Unset resolves automatically: single-node CUDA deployments (without "
+        "disaggregation) use cuda_ipc, everything else uses cpu. CUDA IPC reserves "
+        "SGLANG_MM_FEATURE_CACHE_MB (default 1024 MiB) on the base GPU and falls "
+        "back to CPU transport per tensor when the pool is full.",
         NS("mm"),
     ] = None
     keep_mm_feature_on_device: A[
@@ -3005,7 +3047,7 @@ class ServerArgs:
     encoder_transfer_backend: A[
         str,
         Arg(
-            help="The backend for encoder disaggregation transfer. Default is zmq_to_scheduler.",
+            help="The backend for encoder disaggregation transfer. Auto selects a model- and TP-aware backend.",
             choices=ENCODER_TRANSFER_BACKEND_CHOICES,
         ),
         NS("disagg"),
@@ -3550,8 +3592,6 @@ class ServerArgs:
 
         handle_speculative_decoding(self)
 
-        # Needs the draft-token count derived just above.
-
         # Validate the CuteDSL A2A token budget now that num_tokens_per_req is final.
         self._validate_cutedsl_a2a_token_budget()
 
@@ -3753,16 +3793,20 @@ class ServerArgs:
             )
 
     def _handle_model_source_paths(self):
-        """Resolve model/tokenizer paths backed by remote object stores."""
-        if is_runai_obj_uri(self.model_path):
-            ObjectStorageModel.download_and_get_path(self.model_path)
-
-        if (
-            self.tokenizer_path is not None
-            and is_runai_obj_uri(self.tokenizer_path)
-            and self.tokenizer_path != self.model_path
+        """Prepare metadata for model paths backed by remote object stores."""
+        seen_paths = set()
+        for model_path in (
+            self.model_path,
+            self.tokenizer_path,
+            self.speculative_draft_model_path,
         ):
-            ObjectStorageModel.download_and_get_path(self.tokenizer_path)
+            if (
+                model_path is not None
+                and model_path not in seen_paths
+                and is_runai_obj_uri(model_path)
+            ):
+                ObjectStorageModel.download_and_get_path(model_path)
+                seen_paths.add(model_path)
 
     def _handle_pd_disaggregation(self):
         from sglang.srt.arg_groups.pd_disaggregation_hook import (
@@ -3772,9 +3816,6 @@ class ServerArgs:
         handle_pd_disaggregation(self)
 
     def _handle_dcp_validation(self):
-        # Decode context parallel (DCP) is currently implemented and validated
-        # only on AMD HIP/ROCm. Reject invalid or unverified configurations
-        # early instead of letting them fail deeper in model initialization.
         if self.dcp_size < 1:
             raise ValueError(
                 "Decode context parallel size (--dcp-size / "
@@ -3805,48 +3846,6 @@ class ServerArgs:
                     "communication backend (it removes the head-dim Q all-gather); "
                     f"got --dcp-comm-backend={self.dcp_comm_backend}."
                 )
-        if not self.dcp_size > 1:
-            return
-        if is_hip():
-            return
-        elif is_cuda():
-            if self.speculative_algorithm is not None:
-                model_arches = self.get_model_config().hf_config.architectures
-                decode_backend = self.decode_attention_backend or self.attention_backend
-                kimi_linear_dspark = (
-                    self.speculative_algorithm == "DSPARK"
-                    and "KimiLinearForCausalLM" in model_arches
-                    and self.speculative_attention_mode == "decode"
-                    and decode_backend in ("tokenspeed_mla", "cutedsl_mla")
-                )
-                if kimi_linear_dspark:
-                    ragged_verify_mode = envs.SGLANG_RAGGED_VERIFY_MODE.get()
-                    if ragged_verify_mode != "static":
-                        raise ValueError(
-                            "Kimi Linear DCP + DSPARK currently requires "
-                            "SGLANG_RAGGED_VERIFY_MODE=static, but got "
-                            f"{ragged_verify_mode!r}."
-                        )
-                else:
-                    raise ValueError(
-                        "Decode context parallel (--dcp-size / "
-                        "--decode-context-parallel-size > 1) with speculative "
-                        "decoding on CUDA is supported only for Kimi Linear + "
-                        "DSPARK + --speculative-attention-mode decode + "
-                        "tokenspeed_mla, or experimental cutedsl_mla, but got "
-                        f"architectures={model_arches}, "
-                        f"speculative_algorithm={self.speculative_algorithm!r}, "
-                        "speculative_attention_mode="
-                        f"{self.speculative_attention_mode!r}, "
-                        f"decode_attention_backend={decode_backend!r}."
-                    )
-        else:
-            raise ValueError(
-                "Decode context parallel (--dcp-size / "
-                "--decode-context-parallel-size > 1) is currently only "
-                f"supported on the AMD HIP platform, but got dcp_size="
-                f"{self.dcp_size} on a non-HIP platform."
-            )
 
     def _handle_load_balance_method(self):
         if self.disaggregation_mode not in ("null", "prefill", "decode"):
@@ -4231,7 +4230,12 @@ class ServerArgs:
             self.cuda_graph_backend_prefill = Backend.FULL
 
     def _handle_cuda_graph_config(self):
+        from sglang.srt.arg_groups.kimi_k3_hook import disable_kimi_k3_symm_mem
+
         self._parse_cuda_graph_config()
+        # Reads the resolved per-phase backends; must precede the compat rules
+        # below and _handle_gpu_memory_settings, which key off enable_symm_mem.
+        disable_kimi_k3_symm_mem(self)
         self._apply_cuda_graph_compatibility()
         self._apply_deepep_adjustments()
         self._apply_cuda_graph_disaggregation_roles()
@@ -5088,6 +5092,18 @@ class ServerArgs:
         )
         validate_declarations(self, self._resolved_overrides)
 
+        if model_arch in (
+            "KimiLinearForCausalLM",
+            "KimiK3ForConditionalGeneration",
+        ):
+            from sglang.srt.arg_groups.kimi_k3_hook import (
+                apply_kimi_k3_linear_attn_defaults,
+                apply_kimi_k3_spec_backend_defaults,
+            )
+
+            apply_kimi_k3_linear_attn_defaults(self)
+            apply_kimi_k3_spec_backend_defaults(self)
+
         if model_arch in [
             "DeepseekV4ForCausalLM",
         ]:
@@ -5582,7 +5598,9 @@ class ServerArgs:
                 "--mamba-radix-cache-strategy extra_buffer."
             )
             algo = (view.speculative_algorithm or "").upper()
-            assert algo not in ("DFLASH", "DSPARK"), (
+            # dspark verifies through prepare_mamba_track_for_verify (lazy plan
+            # wired); dflash bypasses that hook, so it stays unsupported.
+            assert algo != "DFLASH", (
                 f"extra_buffer_lazy unsupported with {view.speculative_algorithm}; "
                 "use --mamba-radix-cache-strategy extra_buffer."
             )
@@ -6018,6 +6036,10 @@ class ServerArgs:
             self.linear_attn_decode_backend is None
             and is_sm100_supported()
             and self.mamba_ssm_dtype == "bfloat16"
+            # Stage 4: flashinfer's recurrent_kda compiles the state slot stride
+            # as a free int64, so it reads the page-major/unified envelope-strided
+            # state correctly — the unified-memory skip is no longer needed (the
+            # page-major gate now allows flashinfer for linear-attn decode).
         ):
             self.linear_attn_decode_backend = "flashinfer"
             logger.info(
@@ -6059,6 +6081,21 @@ class ServerArgs:
                 f"got {self.mamba_ssm_dtype!r}"
             )
 
+        verify = self.linear_attn_verify_backend
+        if verify is None and decode == "flashinfer":
+            verify = "flashinfer"
+        if (
+            verify == "flashinfer"
+            and self.mamba_ssm_dtype != "bfloat16"
+            and is_cuda()
+            and torch.cuda.get_device_capability()[0] >= 10
+        ):
+            raise ValueError(
+                "--linear-attn-verify-backend flashinfer on SM100+ requires "
+                "--mamba-ssm-dtype bfloat16, "
+                f"got {self.mamba_ssm_dtype!r}"
+            )
+
         # SM100+ FlashInfer GDN prefill requires CUDA 13+ (CuTe DSL kernel)
         # for correctness and best performance.
         prefill = self.linear_attn_prefill_backend or self.linear_attn_backend
@@ -6123,20 +6160,20 @@ class ServerArgs:
                     f"{self.linear_replayssm_cache_len}."
                 )
 
-        # ReplaySSM spec-verify (Part B of #28511): GDN-only, linear-chain target
-        # verify. Reuses the `linear_replayssm` ring (replayssm_d/k/g + write_pos)
-        # plus two extra per-slot cursors (cache_base, is_flush) and the chunked
-        # (I+A)^-1 reconstruction verify kernel. The intra-window interaction uses a
-        # strictly-lower causal mask, so it is valid ONLY for a linear draft chain
-        # (speculative_eagle_topk in {None, 1}, i.e. NEXTN / MTP); EAGLE tree verify
-        # (topk > 1) must fall back to the recurrent verify. GDN-only is enforced at
-        # runtime (KDA routes through kda_backend, which never enters this path; the
-        # pool gate also checks `not cache_params.is_kda`). The ring length reuses
-        # --linear-replayssm-cache-len (no separate flag).
-        if self.enable_gdn_replayssm_spec:
+        # ReplaySSM spec-verify (Part B of #28511): linear-chain target verify via
+        # fold-every-commit -- the verify stores each draft step's raw inputs into
+        # the per-slot (rawv, rawk, g, beta) window and the commit replays the
+        # accepted prefix into the fp32 checkpoint. The intra-window interaction
+        # uses a strictly-lower causal mask, so it is valid ONLY for a linear
+        # draft chain (speculative_eagle_topk in {None, 1}, i.e. NEXTN / MTP);
+        # EAGLE tree verify (topk > 1) must fall back to the recurrent verify.
+        # GDN sizes the window to the draft maximum; KDA (kda_backend) keeps a
+        # --linear-replayssm-cache-len window and folds via its own fused
+        # verify ring-write + commit_kda_replayssm_after_verify.
+        if self.enable_linear_replayssm_spec:
             if self.speculative_eagle_topk not in (None, 1):
                 raise ValueError(
-                    "--enable-gdn-replayssm-spec requires a linear draft chain "
+                    "--enable-linear-replayssm-spec requires a linear draft chain "
                     "(--speculative-eagle-topk in {None, 1}); the chunked verify "
                     "kernel uses a strictly-lower causal mask and is invalid for "
                     "EAGLE tree verify. Got "
@@ -6144,45 +6181,61 @@ class ServerArgs:
                 )
             if decode not in ("triton", "flashinfer"):
                 raise ValueError(
-                    "--enable-gdn-replayssm-spec requires the triton or "
+                    "--enable-linear-replayssm-spec requires the triton or "
                     "flashinfer linear-attn decode backend, got "
                     f"--linear-attn-decode-backend={decode!r}."
                 )
-            # The auto->extra_buffer strategy resolution is still a declaration
-            # here, so read it through the resolved view.
-            view = self._resolved()
-            if (
-                view.disable_radix_cache is False
-                and view.mamba_radix_cache_strategy == "extra_buffer_lazy"
-            ):
-                raise ValueError(
-                    "--enable-gdn-replayssm-spec is not validated with "
-                    "--mamba-radix-cache-strategy extra_buffer_lazy yet; "
-                    "use extra_buffer."
-                )
+            from sglang.srt.speculative.ragged_verify import (
+                RaggedVerifyMode,
+                read_ragged_verify_mode,
+            )
+
+            ragged_mode = read_ragged_verify_mode()
+            if ragged_mode is not RaggedVerifyMode.STATIC:
+                # Ragged ring-writes need the KDA fold-every-commit family
+                # (DSPARK/DFLASH) + the triton verify kernel (nv_cutedsl falls
+                # back to it for ragged layouts). The GDN ring-write kernels do
+                # not take the ragged layout and the flashinfer verify kernel
+                # never writes the ring -> a stale ring would be folded; keep
+                # refusing those combinations.
+                _algo = (self.speculative_algorithm or "").upper()
+                verify = self.linear_attn_verify_backend
+                if _algo not in ("DSPARK", "DFLASH") or verify not in (
+                    "triton",
+                    "nv_cutedsl",
+                ):
+                    raise ValueError(
+                        "--enable-linear-replayssm-spec with "
+                        f"SGLANG_RAGGED_VERIFY_MODE={ragged_mode.value} requires the "
+                        "KDA fold-every-commit family (DSPARK/DFLASH) and a "
+                        "ring-writing verify kernel (--linear-attn-verify-backend "
+                        "triton or nv_cutedsl); got "
+                        f"algorithm={self.speculative_algorithm!r}, "
+                        f"verify={verify!r}. Use SGLANG_RAGGED_VERIFY_MODE=static."
+                    )
             if self.disaggregation_mode == "prefill":
                 raise ValueError(
-                    "--enable-gdn-replayssm-spec is not supported on a PD "
+                    "--enable-linear-replayssm-spec is not supported on a PD "
                     "prefill server: the ring is spec-verify-only scratch and "
                     "the prefill server never runs spec verify."
                 )
             if self.enable_linear_replayssm:
                 raise ValueError(
-                    "--enable-gdn-replayssm-spec and --enable-linear-replayssm are "
+                    "--enable-linear-replayssm-spec and --enable-linear-replayssm are "
                     "mutually exclusive: they share the ring storage but drive it "
                     "with incompatible cursor protocols (per-decode-forward vs "
                     "per-verify-commit advance)."
                 )
             if self.mamba_ssm_dtype is None:
                 logger.info(
-                    "--enable-gdn-replayssm-spec: setting --mamba-ssm-dtype "
+                    "--enable-linear-replayssm-spec: setting --mamba-ssm-dtype "
                     "float32 (the closed-loop exact fold keeps the SSM checkpoint "
                     "bit-identical to the recurrent baseline)."
                 )
                 self.mamba_ssm_dtype = "float32"
             elif self.mamba_ssm_dtype != "float32":
                 logger.warning(
-                    "--enable-gdn-replayssm-spec with --mamba-ssm-dtype=%s: the "
+                    "--enable-linear-replayssm-spec with --mamba-ssm-dtype=%s: the "
                     "closed-loop fold re-quantizes the committed state each "
                     "commit/flush (fp32 keeps it bit-exact to the fp32 recurrent "
                     "baseline), so it may drift over long sequences. Validate "
@@ -7212,6 +7265,13 @@ class ServerArgs:
         elif is_remote_url(self.model_path):
             self.load_format = "remote"
 
+        if (
+            self.speculative_draft_model_path is not None
+            and is_runai_obj_uri(self.speculative_draft_model_path)
+            and self.speculative_draft_load_format is None
+        ):
+            self.speculative_draft_load_format = "runai_streamer"
+
         if self.custom_weight_loader is None:
             self.custom_weight_loader = []
 
@@ -7361,6 +7421,17 @@ class ServerArgs:
         # Validate model type for encoder disaggregation
         hf_config = self.get_model_config().hf_config
         model_arch = hf_config.architectures[0]
+        if self.encoder_transfer_backend == "auto":
+            self.encoder_transfer_backend = resolve_encoder_transfer_backend(
+                self.encoder_transfer_backend, model_arch, self.tp_size
+            )
+            if self.encoder_only or self.language_only:
+                logger.info(
+                    "Encoder transfer backend auto-resolved to %s for %s at TP%d.",
+                    self.encoder_transfer_backend,
+                    model_arch,
+                    self.tp_size,
+                )
         if (self.encoder_only or self.language_only) and model_arch not in [
             "Qwen2VLForConditionalGeneration",
             "Qwen3VLForConditionalGeneration",
@@ -7374,6 +7445,7 @@ class ServerArgs:
             "Qwen2_5OmniForConditionalGeneration",
             "KimiVLForConditionalGeneration",
             "KimiK25ForConditionalGeneration",
+            "KimiK3ForConditionalGeneration",
             "MiMoV2ForCausalLM",
         ]:
             raise ValueError(
@@ -7523,6 +7595,26 @@ class ServerArgs:
                     "--mm-feature-transport=%s instead.",
                     requested_transport,
                 )
+            elif self.encoder_only:
+                requested_transport = "cpu"
+                logger.info(
+                    "Multimodal feature transport auto-resolved to cpu for "
+                    "encoder-only serving; encoder outputs use "
+                    "--encoder-transfer-backend instead."
+                )
+            elif is_cuda() and self.nnodes == 1 and self.disaggregation_mode == "null":
+                # Auto policy: single-node CUDA serving defaults to the bounded
+                # CUDA-IPC pool. The pool is only allocated when a multimodal
+                # processor exists, so text-only deployments are unaffected; a
+                # full pool degrades to CPU transport per tensor. Multi-node
+                # (IPC handles are intra-node) and PD-disaggregated deployments
+                # keep CPU transport.
+                requested_transport = "cuda_ipc"
+                logger.info(
+                    "Multimodal feature transport auto-resolved to cuda_ipc "
+                    "(single-node CUDA). Pass --mm-feature-transport=cpu to "
+                    "opt out."
+                )
             else:
                 requested_transport = "cpu"
         elif legacy_ipc_is_set and legacy_ipc_enabled != (
@@ -7535,6 +7627,14 @@ class ServerArgs:
                 int(legacy_ipc_enabled),
             )
 
+        if self.encoder_only and requested_transport == "cuda_ipc":
+            logger.warning(
+                "--mm-feature-transport=cuda_ipc does not control encoder-only "
+                "output transfer; using cpu for this inactive transport. Select "
+                "--encoder-transfer-backend for encoder outputs."
+            )
+            requested_transport = "cpu"
+
         if requested_transport == "cuda_ipc":
             if not is_cuda():
                 raise ValueError(
@@ -7554,6 +7654,16 @@ class ServerArgs:
                 self.base_gpu_id,
                 self.tokenizer_worker_num,
             )
+            logger.info(
+                "CUDA IPC pool-handle caching is %s. It reuses mappings to the "
+                "existing bounded pool without reserving another pool; set "
+                "SGLANG_USE_IPC_POOL_HANDLE_CACHE=0 to disable it.",
+                (
+                    "enabled"
+                    if envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.get()
+                    else "disabled"
+                ),
+            )
 
         self.mm_feature_transport = requested_transport
         # The bounded IPC pool owns device residency. Do not retain unpooled
@@ -7563,6 +7673,41 @@ class ServerArgs:
             "1" if requested_transport == "cuda_ipc" else "0"
         )
 
+    def _handle_custom_all_reduce_v2_multinode(self):
+        # Custom all-reduce v2's graph zero-copy path uses IPC handles and is
+        # intra-node only. On MNNVL-fabric devices (GB200/GB300) the eager pull
+        # path works across nodes via the symm-mem workspace, so opt into the
+        # multinode mode automatically (a failed fabric rendezvous falls back
+        # to the legacy path at init). Elsewhere force-disable v2 on
+        # multi-node so the dispatch falls back to the legacy CustomAllreduce
+        # path, unless the MNNVL opt-in is set explicitly.
+        if self.nnodes <= 1 or not envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
+            return
+        if (
+            not envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.is_set()
+            and is_mnnvl_fabric_device()
+            # CustomAllReduceV2 supports world sizes 2..8 only
+            # (can_use_custom_all_reduce_v2 rejects larger groups); don't
+            # auto-opt-in a TP16+ launch just to fall back downstream.
+            and self.tp_size <= 8
+        ):
+            logger.info(
+                "MNNVL fabric device detected with nnodes=%d: enabling "
+                "custom all-reduce v2 multinode mode "
+                "(SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE=1; set it "
+                "to 0 to opt out).",
+                self.nnodes,
+            )
+            envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.set("1")
+        if not envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get():
+            if envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.is_set():
+                logger.warning(
+                    "Disabling SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 because nnodes=%d "
+                    "(custom all-reduce v2 is intra-node only).",
+                    self.nnodes,
+                )
+            envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.set("0")
+
     def _handle_environment_variables(self):
         self._handle_multimodal_feature_transport()
         envs.SGLANG_ENABLE_TORCH_COMPILE.set("1" if self.enable_torch_compile else "0")
@@ -7574,6 +7719,7 @@ class ServerArgs:
         envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set(
             "1" if self.enable_deterministic_inference else "0"
         )
+        self._handle_custom_all_reduce_v2_multinode()
         if self.debug_cuda_graph:
             if not (is_cuda() or is_hip()):
                 logger.warning(
@@ -7815,19 +7961,43 @@ class ServerArgs:
             f"paged MLA backends); got {sorted(backends)}, allowed "
             f"{sorted(allowed_full)}. Pass a compatible --attention-backend."
         )
-        # The Mamba state is stored in envelope-strided views; only the
-        # stride-aware Triton causal-conv / SSM kernels read them correctly.
-        linear_backends = {
-            self.linear_attn_backend,
-            self.linear_attn_decode_backend,
-            self.linear_attn_prefill_backend,
-            self.mamba_backend,
-        }
-        linear_backends.discard(None)
-        assert linear_backends <= {"triton"}, (
-            "--enable-page-major-kv-layout requires the Triton linear-attention / "
-            f"Mamba kernels for the strided conv/SSM state; got "
-            f"{sorted(linear_backends)}. Pass --linear-attn-backend triton and "
+        # The Mamba/KDA state is stored in envelope-strided views; only
+        # stride-audited kernels may read it (Stage 4 audit, per slot):
+        # - decode: triton; flashinfer (recurrent_kda compiles the state slot
+        #   stride as a free int64 — natively strided); cutedsl (KDA fused
+        #   sigmoid-gating update made stride-safe) on KDA-hybrid models only —
+        #   cutedsl_gdn still compiles h0 against a contiguous dummy.
+        # - prefill: triton; flashkda (wrapper gathers/scatters a contiguous
+        #   per-slot copy, external kernel never sees the pool); cutedsl
+        #   (kernel_h compiles h0/ht with dynamic int64 strides), same
+        #   KDA-only caveat.
+        # - mamba (mamba2/short-conv state): triton only.
+        # use_mla_backend() distinguishes the KDA-hybrid family (K3/KimiLinear
+        # are MLA-hybrid) from GDN models (GQA-hybrid) for the cutedsl caveat.
+        decode_allowed = {"triton", "flashinfer"}
+        prefill_allowed = {"triton", "flashkda"}
+        if self.use_mla_backend():
+            decode_allowed.add("cutedsl")
+            prefill_allowed.add("cutedsl")
+        resolved_linear_decode = (
+            self.linear_attn_decode_backend or self.linear_attn_backend
+        )
+        resolved_linear_prefill = (
+            self.linear_attn_prefill_backend or self.linear_attn_backend
+        )
+        assert resolved_linear_decode in decode_allowed | {None}, (
+            "--enable-page-major-kv-layout: linear-attention DECODE backend must "
+            f"be one of {sorted(decode_allowed)} for the strided conv/SSM state; "
+            f"got {resolved_linear_decode!r}."
+        )
+        assert resolved_linear_prefill in prefill_allowed | {None}, (
+            "--enable-page-major-kv-layout: linear-attention PREFILL backend must "
+            f"be one of {sorted(prefill_allowed)} for the strided conv/SSM state; "
+            f"got {resolved_linear_prefill!r}."
+        )
+        assert self.mamba_backend in (None, "triton"), (
+            "--enable-page-major-kv-layout requires the Triton Mamba kernels for "
+            f"the strided conv/SSM state; got {self.mamba_backend!r}. Pass "
             "--mamba-backend triton."
         )
 
@@ -8226,6 +8396,13 @@ class ServerArgs:
             new_flag="--enable-prefill-cp",
             help="[Deprecated] Use --enable-prefill-cp instead.",
         )
+        parser.add_argument(
+            "--enable-gdn-replayssm-spec",
+            dest="enable_linear_replayssm_spec",
+            action=DeprecatedStoreTrueAction,
+            new_flag="--enable-linear-replayssm-spec",
+            help="[Deprecated] Use --enable-linear-replayssm-spec instead.",
+        )
         parser.add_argument(
             "--enable-prefill-context-parallel",
             dest="enable_prefill_context_parallel",

python/sglang/srt/speculative/dflash_info.py

@@ -105,6 +105,9 @@ class DFlashVerifyInput(SpecInput):
         bs = len(req_pool_indices)
 
         layout = self.ragged_verify_layout
+        if layout is not None and layout.bs != bs:
+            # Graph replay pads the batch to the captured slots; match it.
+            layout = layout.padded_to_bucket(padded_bs=bs)
 
         if layout is None:
             qo_indptr = torch.arange(

python/sglang/srt/speculative/dflash_utils.py

@@ -13,7 +13,7 @@ from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 from sglang.srt.layers.sampler import apply_custom_logit_processor
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.speculative.spec_utils import _sample_simulated_acc_len
-from sglang.srt.utils import is_cuda, is_musa
+from sglang.srt.utils import is_cuda, is_hip, is_musa
 
 DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>"
 
@@ -46,6 +46,15 @@ if is_cuda() or is_musa():
         top_k_renorm_prob = None
         top_p_renorm_prob = None
         tree_speculative_sampling_target_only = None
+elif is_hip():
+    from sglang.kernels.ops.sampling.renorm_triton import (
+        top_k_renorm_probs_triton as top_k_renorm_prob,
+    )
+    from sglang.kernels.ops.sampling.renorm_triton import (
+        top_p_renorm_probs_triton as top_p_renorm_prob,
+    )
+
+    _DFLASH_SAMPLING_VERIFY_AVAILABLE = True
 else:
     top_k_renorm_prob = None
     top_p_renorm_prob = None

python/sglang/srt/speculative/draft_worker_common.py

@@ -22,13 +22,16 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+# trtllm_mha: decode-only dense-MQA drafts (dspark). DFLASH excludes it
+# earlier, at arg resolution (speculative_hook.py) -- its draft path needs
+# per-layer DFlash attention -- so it never reaches this gate with it.
 _SUPPORTED_DRAFT_BACKENDS = (
     "flashinfer",
     "fa3",
     "fa4",
     "triton",
-    "trtllm_mha",
     "ascend",
+    "trtllm_mha",
 )
 
 

python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py

@@ -92,8 +92,15 @@ class DSparkWorkerV2(BaseSpecWorker):
         self._draft_dp_context_enabled = (
             server_args.enable_dp_attention and not self._draft_is_moe
         )
-        attn_tp_size = server_args.tp_size // max(server_args.dp_size, 1)
-        if server_args.enable_dp_attention and self._draft_is_moe and attn_tp_size > 1:
+        self._is_pd_prefill = server_args.disaggregation_mode == "prefill"
+        self._decode_graph_allowed = (
+            not server_args.disable_cuda_graph and not self._is_pd_prefill
+        )
+        if (
+            server_args.enable_dp_attention
+            and self._draft_is_moe
+            and ps.attn_tp_size > 1
+        ):
             raise ValueError(
                 "DSpark + dp attention with a DeepSeek-V4 (MoE) draft requires "
                 "attn_tp == 1 (set --dp-size == --tp). attn_tp > 1 corrupts the "
@@ -173,7 +180,7 @@ class DSparkWorkerV2(BaseSpecWorker):
             server_args.enable_dp_attention
             and not self._draft_is_moe
             and self._verify_planner.is_compact_mode
-            and not server_args.disable_cuda_graph
+            and self._decode_graph_allowed
         ):
             raise ValueError(
                 "DSpark dense-draft compact verify under --enable-dp-attention does not "
@@ -201,7 +208,7 @@ class DSparkWorkerV2(BaseSpecWorker):
         self._verify_epilogue = None
         if (
             self._verify_planner.is_compact_mode
-            and not server_args.disable_cuda_graph
+            and self._decode_graph_allowed
             and is_cuda()
         ):
             self._verify_epilogue = DsparkVerifyEpilogue(
@@ -262,6 +269,9 @@ class DSparkWorkerV2(BaseSpecWorker):
             simulate_acc_len=self._simulate_acc_len,
         )
 
+        if self._is_pd_prefill and not self._draft_is_moe:
+            self.draft_model.prune_to_ctx_kv_injection()
+
     def _resolve_target_embed_tokens(self, target_model):
         if hasattr(target_model, "get_input_embeddings"):
             return target_model.get_input_embeddings()
@@ -311,7 +321,7 @@ class DSparkWorkerV2(BaseSpecWorker):
         )
 
     def init_cuda_graphs(self):
-        capture_decode_cuda_graph = not get_exec().graph.disable_cuda_graph
+        capture_decode_cuda_graph = self._decode_graph_allowed
         if is_cuda() and capture_decode_cuda_graph:
             available_mem = get_available_gpu_memory(self.device, self.gpu_id)
             if available_mem < 1.0:
@@ -420,6 +430,8 @@ class DSparkWorkerV2(BaseSpecWorker):
         if batch.out_cache_loc is None:
             raise RuntimeError("DSpark prefill expected out_cache_loc, but got None.")
 
+        # Must inject before prefill returns: the scheduler may update radix
+        # afterward, invalidating out_cache_loc.
         device = next_token_ids.device
         ctx_lens = torch.tensor(batch.extend_lens, dtype=torch.int32, device=device)
         draft_seq_lens = torch.tensor(
@@ -436,6 +448,7 @@ class DSparkWorkerV2(BaseSpecWorker):
             cache_loc=batch.out_cache_loc,
             positions=positions,
         )
+        # Avoid copying large hidden-state buffers to CPU in overlap scheduling.
         logits_output.hidden_states = None
 
         batch_output.next_draft_input = make_next_draft_input(

python/sglang/srt/speculative/dspark_disaggregation.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.managers.overlap_utils import RelayPayload
+from sglang.srt.speculative.dspark_components.dspark_draft import make_next_draft_input
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.overlap_utils import FutureMap
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.server_args import ServerArgs
+    from sglang.srt.speculative.spec_info import SpecInput
+
+
+def build_dspark_disagg_draft_input(
+    batch: ScheduleBatch,
+    server_args: ServerArgs,
+    last_tokens_tensor: torch.Tensor,
+    future_map: FutureMap,
+) -> SpecInput:
+    spec_info = make_next_draft_input(
+        bonus_tokens=last_tokens_tensor,
+        new_seq_lens=batch.seq_lens,
+    )
+    if batch.enable_overlap:
+        spec_info.future_dsa_topk_indices_available = False
+        spec_info.future_indices = batch.req_pool_indices
+        future_map.publish(spec_info.future_indices, batch.seq_lens)
+        future_map.stash(
+            spec_info.future_indices, RelayPayload.from_draft_input(spec_info)
+        )
+    return spec_info

python/sglang/srt/speculative/ragged_verify.py

@@ -55,6 +55,9 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True):
     kv_lens_host: Optional[torch.Tensor] = None
     max_q_len: Optional[int] = None
     max_kv_len: Optional[int] = None
+    # Per-row upper bound (capped padded variant); rows never exceed it, so
+    # dense [bs, cap] consumers stay in bounds. None = full-coverage variant.
+    cap: Optional[int] = None
 
     def __post_init__(self) -> None:
         if self.verify_lens_cpu is None:
@@ -66,6 +69,11 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True):
                 f"every request must verify the anchor (verify_len >= 1), got "
                 f"{self.verify_lens_cpu}"
             )
+        if self.cap is not None and max(self.verify_lens_cpu) > self.cap:
+            raise ValueError(
+                f"capped layout has a row exceeding cap={self.cap}: "
+                f"{self.verify_lens_cpu}"
+            )
         if self.total_verify_tokens != sum(self.verify_lens_cpu):
             raise ValueError(
                 f"total_verify_tokens {self.total_verify_tokens} != "
@@ -89,6 +97,7 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True):
         graph_num_tokens: int,
         verify_lens_cpu: Optional[list[int]] = None,
         total_verify_tokens: Optional[int] = None,
+        cap: Optional[int] = None,
     ) -> RaggedVerifyLayout:
         from sglang.kernels.ops.speculative.ragged_verify_kernels import (
             BuildQoIndptr,
@@ -103,6 +112,7 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True):
             qo_indptr_device=indptr.qo_indptr,
             verify_lens_cpu=verify_lens_cpu,
             total_verify_tokens=total_verify_tokens,
+            cap=cap,
         )
 
     @classmethod
@@ -154,7 +164,16 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True):
             device=device,
         )
 
-    def padded_to_bucket(self, *, padded_bs: int) -> RaggedVerifyLayout:
+    def padded_to_bucket(
+        self, *, padded_bs: int, cap: Optional[int] = None
+    ) -> RaggedVerifyLayout:
+        """Pad to the captured slot count. The full-coverage variant (cap=None)
+        keeps sum(verify_lens) == graph_num_tokens, so padding rows (or, with
+        no padding rows, the last real row) can exceed the per-request verify
+        window. Dense-layout consumers (mamba/KDA) must pass cap=N: rows are
+        clamped to N and the layout no longer covers the tier's leftover
+        tokens -- those collapse into the consumer's ghost slot. Real tokens
+        map identically under both variants."""
         from sglang.kernels.ops.speculative.ragged_verify_kernels import (
             PaddedToBucket,
         )
@@ -165,11 +184,14 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True):
             bs=self.bs,
             padded_bs=padded_bs,
         )
+        if cap is not None:
+            padded = torch.clamp(padded, max=cap)
 
         return RaggedVerifyLayout._assemble_device(
             verify_lens=padded,
             graph_num_tokens=self.graph_num_tokens,
-            total_verify_tokens=self.graph_num_tokens,
+            total_verify_tokens=None if cap is not None else self.graph_num_tokens,
+            cap=cap,
         )
 
 

python/sglang/srt/speculative/spec_info.py

@@ -184,6 +184,14 @@ class SpeculativeAlgorithm(Enum):
             return build_eagle_disagg_draft_input(
                 batch, server_args, last_tokens_tensor, future_map
             )
+        if self.is_dspark():
+            from sglang.srt.speculative.dspark_disaggregation import (
+                build_dspark_disagg_draft_input,
+            )
+
+            return build_dspark_disagg_draft_input(
+                batch, server_args, last_tokens_tensor, future_map
+            )
         return None
 
     def need_topk(self) -> bool:

python/sglang/srt/speculative/spec_utils.py

@@ -935,10 +935,72 @@ def commit_mamba_states_after_verify(
         # NOTE: radix mamba prefix-caching (mamba_track / extra_buffer) would need
         # a device-side force-flush so `temporal` reflects the ring before a
         # snapshot; not wired for Part B (server_args forbids extra_buffer with
-        # --enable-gdn-replayssm-spec), so the per-track scatters are intentionally
+        # --enable-linear-replayssm-spec), so the per-track scatters are intentionally
         # skipped here.
         return
 
+    # KDA ReplaySSM (fold-every-commit): KDA keeps its own recurrent verify kernel
+    # for the OUTPUT, so we replay the accepted window into the fp32 checkpoint
+    # (`temporal`) here on commit -- `temporal` is always the current committed
+    # state. The draft window's raw inputs were written to the ring during verify
+    # by the KDA backend. Gate on the fold flag + is_kda (the cursor tensors are
+    # never allocated under fold, so they cannot serve as the signal).
+    if (
+        mamba_pool is not None
+        and getattr(mamba_pool, "replayssm_spec_fold", False)
+        and getattr(mamba_pool, "replayssm_is_kda", False)
+    ):
+        if batch.forward_mode.is_idle() or accept_index.numel() == 0:
+            return
+        from sglang.kernels.ops.attention.fla.kda_replayssm_spec_decode import (
+            commit_kda_replayssm_after_verify,
+        )
+
+        spec_state = req_pool.get_speculative_mamba2_params_all_layers()
+        bs = accept_lens.shape[0]
+        state_batch_indices = req_pool.get_mamba_indices(batch.req_pool_indices)
+        accept_indices_offset = torch.arange(
+            0,
+            bs * draft_token_num,
+            step=draft_token_num,
+            dtype=accept_lens.dtype,
+            device=accept_lens.device,
+        )
+        req_idx = torch.arange(bs, dtype=torch.int64, device=accept_lens.device)
+        last_correct_step_indices = (
+            accept_index[req_idx, (accept_lens - 1).to(torch.int64)]
+            - accept_indices_offset
+        )
+        # extra_buffer: the interval-crossing step whose state must snapshot into
+        # the track ping-pong slot (mirrors the regular commit's
+        # mamba_steps_to_track); commit_kda_replayssm_spec folds it in one pass, so
+        # `temporal` stays current and no device-side force-flush is needed.
+        mamba_track_indices = batch.mamba_track_indices
+        mamba_steps_to_track = None
+        if mamba_track_indices is not None:
+            ti = get_exec().mamba.mamba_track_interval
+            seq_pre = batch.seq_lens
+            seq_post = batch.seq_lens + accept_lens
+            to_track_mask = seq_pre // ti != seq_post // ti
+            tracking_point = seq_post // ti * ti
+            to_track_ith = torch.clamp(tracking_point - seq_pre - 1, min=0).to(
+                torch.int64
+            )
+            candidate = accept_index[req_idx, to_track_ith] - accept_indices_offset
+            mamba_steps_to_track = torch.where(
+                to_track_mask, candidate, torch.full_like(candidate, -1)
+            )
+        commit_kda_replayssm_after_verify(
+            spec_state=spec_state,
+            state_batch_indices=state_batch_indices,
+            accept_lens=accept_lens,  # incl. bonus token
+            last_correct_step_indices=last_correct_step_indices,
+            mamba_track_indices=mamba_track_indices,
+            mamba_steps_to_track=mamba_steps_to_track,
+            null_block_id=-1,  # SGLang: valid slots >= 0, padding == -1
+        )
+        return
+
     attn_backend = model_runner.attn_backend
 
     bs = accept_lens.shape[0]