[Kimi-K3][NPU] Support Kimi-K3 on NPU (#33465)

python/sglang/kernels/ops/elementwise/add3.py

@@ -12,6 +12,7 @@ from sglang.kernels.jit.utils import (
     load_jit,
     make_cpp_args,
 )
+from sglang.srt.utils import is_npu
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
@@ -38,7 +39,8 @@ def covered(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> bool:
     """Same-shape contiguous CUDA bf16 tensors, numel a multiple of the
     widest vector (16 elements)."""
     return (
-        a.dtype == b.dtype == c.dtype == torch.bfloat16
+        not is_npu()
+        and a.dtype == b.dtype == c.dtype == torch.bfloat16
         and a.shape == b.shape == c.shape
         and a.is_contiguous()
         and b.is_contiguous()

python/sglang/kernels/ops/kimi_k3/__init__.py

@@ -2,9 +2,13 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Optional
 
+from sglang.srt.utils import is_npu
+
 if TYPE_CHECKING:
     import torch
 
+_is_npu = is_npu()
+
 _K3_N_GEMM_DISPATCH_MAP = {
     (144, 7168): 16,
     (896, 7168): 8,
@@ -65,12 +69,13 @@ def kimi_k3_tiny_gemm(
 
     m, k = x.shape
     n, _ = w.shape
-    if max_num_tokens := _K3_N_GEMM_DISPATCH_MAP.get((n, k)):
-        if 0 < m <= max_num_tokens:
-            return tiny_n_gemm_bf16(x, w)
-    if max_num_tokens := _K3_K_GEMM_DISPATCH_MAP.get((n, k)):
-        if 0 < m <= max_num_tokens:
-            return tiny_k_gemm_bf16(x, w)
+    if not _is_npu:
+        if max_num_tokens := _K3_N_GEMM_DISPATCH_MAP.get((n, k)):
+            if 0 < m <= max_num_tokens:
+                return tiny_n_gemm_bf16(x, w)
+        if max_num_tokens := _K3_K_GEMM_DISPATCH_MAP.get((n, k)):
+            if 0 < m <= max_num_tokens:
+                return tiny_k_gemm_bf16(x, w)
     return torch.nn.functional.linear(x, w)
 
 

python/sglang/kernels/ops/kimi_k3/mla_output_gate.py

@@ -12,11 +12,13 @@ from sglang.kernels.jit.utils import (
     load_jit,
     make_cpp_args,
 )
+from sglang.srt.utils import is_npu
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
 
 _THREADS: int = 256
+_is_npu = is_npu()
 
 
 @cache_once
@@ -33,7 +35,8 @@ def _jit_mla_output_gate_module() -> Module:
 
 def covered(x: torch.Tensor, gate: torch.Tensor) -> bool:
     return (
-        x.dtype == torch.bfloat16
+        not _is_npu
+        and x.dtype == torch.bfloat16
         and gate.dtype == torch.bfloat16
         and x.shape == gate.shape
         and x.is_contiguous()

python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py

@@ -9,6 +9,7 @@ from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_f
 from sglang.kernels.ops.speculative.dspark.dispatch import inputs_on_cuda
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
+from sglang.srt.utils import is_npu
 
 
 class RaggedVerifyWindow(msgspec.Struct, frozen=True):
@@ -797,7 +798,7 @@ def build_commit_inject_layout_triton(
 class BuildOutTokens:
     @classmethod
     def execute(cls, *args, **kwargs) -> torch.Tensor:
-        if inputs_on_cuda(*args, **kwargs):
+        if not is_npu() and inputs_on_cuda(*args, **kwargs):
             return cls.triton(*args, **kwargs)
         return cls.torch(*args, **kwargs)
 

python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py

@@ -52,6 +52,12 @@ class AscendFABackend(AttentionBackend):
     def get_builder_cls() -> type["AttentionMetadataBuilder"]:
         return AscendFAMetadataBuilder
 
+    @classmethod
+    def supports_ring_rotation(cls) -> bool:
+        """Whether this backend can serve as the ring-attention kernel; the
+        per-hop online-softmax merge needs the kernel's softmax LSE."""
+        return True
+
 
 class AscendFAImpl(AttentionImpl):
 

python/sglang/srt/arg_groups/speculative_hook.py

@@ -276,14 +276,17 @@ def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool:
 
 
 def _handle_dspark(server_args: ServerArgs) -> None:
-    if not server_args.device.startswith("cuda"):
-        raise ValueError("DSpark speculative decoding only supports CUDA device.")
+    _is_npu = server_args.device.startswith("npu")
+    if not server_args.device.startswith("cuda") and not _is_npu:
+        raise ValueError(
+            "DSpark speculative decoding only supports CUDA and NPU devices."
+        )
 
     # dp_size==1 with dp_attention is a degenerate flag under DSV4 CP; skip DP-only checks.
     if server_args.enable_dp_attention and server_args.dp_size > 1:
         if not server_args.enable_dp_lm_head:
             raise ValueError("DSpark with dp attention requires --enable-dp-lm-head.")
-        if server_args.moe_a2a_backend != "none":
+        if not _is_npu and server_args.moe_a2a_backend != "none":
             raise ValueError(
                 "DSpark with dp attention only supports the built-in TP MoE "
                 f"(moe_a2a_backend='none'), got {server_args.moe_a2a_backend!r}."
@@ -294,7 +297,8 @@ def _handle_dspark(server_args: ServerArgs) -> None:
                 f"(attn_cp_size={server_args.attn_cp_size})."
             )
         if (
-            server_args.speculative_moe_a2a_backend is not None
+            not _is_npu
+            and server_args.speculative_moe_a2a_backend is not None
             and server_args.speculative_moe_a2a_backend != server_args.moe_a2a_backend
         ):
             raise ValueError(

python/sglang/srt/environ.py

@@ -348,6 +348,9 @@ class Envs:
     SGLANG_DSPARK_FP32_LM_HEAD = EnvBool(False)
     SGLANG_DSPARK_FAST_SAMPLING = EnvBool(True)
     SGLANG_DSPARK_FOLDED_SAMPLING = EnvInt(DsparkFoldedSampling.AUTO)
+    SGLANG_DSPARK_FOLDED_PROPOSAL = EnvBool(True)
+    SGLANG_DSPARK_STACKED_CTX_KV = EnvBool(True)
+    SGLANG_DSPARK_EMBED_IN_GRAPH = EnvBool(True)
     SGLANG_DSPARK_OPT_MARKOV_W2_BF16 = EnvBool(True)
     SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD = EnvBool(True)
     SGLANG_DSPARK_ENABLE_MULTI_STREAM = EnvBool(True)
@@ -670,6 +673,13 @@ class Envs:
     SGLANG_EXPERIMENTAL_LORA_OPTI = EnvBool(False)
     # Enable int4x2 weights loading
     SGLANG_NPU_W4A4_NEW_PACKING = EnvBool(False)
+    # Keep K3 shared experts and dense MLPs sharded over attention TP.
+    SGLANG_K3_SHARED_EXPERTS_ATTN_TP = EnvBool(False)
+    SGLANG_K3_DENSE_MLP_ATTN_TP = EnvBool(False)
+    # Use the graph-safe Triton-Ascend kernel for masked speculative KV commits.
+    SGLANG_NPU_USE_TRITON_PREFIX_KV_CACHE_STORE = EnvBoolWithAlias(
+        False, deprecated_name="SGLANG_NPU_USE_TRITON_KV_CACHE_STORE"
+    )
     # Quantize x to int8 in the dispatch operator
     DEEP_NORMAL_MODE_USE_INT8_QUANT = EnvBool(False)  # This argument is deprecated
     SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1)

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -28,7 +28,11 @@ from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_flags, get_spec
 from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
-from sglang.srt.utils import get_bool_env_var, get_current_device_stream_fast
+from sglang.srt.utils import (
+    get_bool_env_var,
+    get_current_device_stream_fast,
+    next_power_of_2,
+)
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
@@ -337,6 +341,15 @@ class AscendAttnBackend(AttentionBackend):
         self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
         self.enable_torch_compile = get_flags().capture.enable_torch_compile
         self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens
+        if (
+            self.speculative_num_draft_tokens is not None
+            and model_runner.is_draft_worker
+        ):
+            self.speculative_num_draft_tokens = (
+                model_runner.spec_algorithm.get_num_tokens_per_req_for_target_verify(
+                    int(self.speculative_num_draft_tokens), is_draft_worker=True
+                )
+            )
         self.ascend_attn_mask_builder = AscendAttnMaskBuilder(
             model_runner, self.device, self.use_fia, self.use_mla
         )
@@ -430,14 +443,23 @@ class AscendAttnBackend(AttentionBackend):
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Init the metadata for a forward pass."""
         self.forward_metadata = ForwardMetadata()
-        seq_lens_max = forward_batch.seq_lens.max()
         if forward_batch.forward_mode.is_target_verify():
-            seq_lens_max += self.speculative_num_draft_tokens
+            # Overlap scheduling can publish the CPU sequence length one step
+            # ahead of the device tensor. FIA consumes seq_lens_cpu below, so
+            # derive the block-table width from the same source. Otherwise a
+            # page-aligned request can expose KV_S=N while asking FIA for N+1.
+            seq_lens_max = (
+                forward_batch.seq_lens_cpu.max().item()
+                + self.speculative_num_draft_tokens
+            )
         elif (
             forward_batch.forward_mode.is_decode_or_idle()
             and forward_batch.spec_info is not None
         ):
+            seq_lens_max = forward_batch.seq_lens.max()
             seq_lens_max += self.speculative_step_id + 1
+        else:
+            seq_lens_max = forward_batch.seq_lens.max()
         self.forward_metadata.block_tables = (
             self.req_to_token_pool.req_to_token[
                 forward_batch.req_pool_indices, :seq_lens_max
@@ -2781,13 +2803,32 @@ class AscendAttnBackend(AttentionBackend):
                     layer.tp_q_head_num,
                     self.qk_rope_head_dim,
                 )
+                if (layer.tp_q_head_num & (layer.tp_q_head_num - 1)) != 0:
+                    power_of_2_head = next_power_of_2(layer.tp_q_head_num)
+                    padding_head = power_of_2_head - layer.tp_q_head_num
+                    q_padding_tensor = torch.zeros(
+                        [num_tokens, q.shape[1], padding_head, q.shape[-1]],
+                        dtype=q.dtype,
+                        device=q.device,
+                    )
+                    q = torch.cat((q, q_padding_tensor), dim=-2)
+                    q_rope_padding_tensor = torch.zeros(
+                        [num_tokens, q_rope.shape[1], padding_head, q_rope.shape[-1]],
+                        dtype=q_rope.dtype,
+                        device=q_rope.device,
+                    )
+                    q_rope = torch.cat((q_rope, q_rope_padding_tensor), dim=-2)
+                    tp_q_head_num = power_of_2_head
+                else:
+                    tp_q_head_num = layer.tp_q_head_num
+
                 attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
                     q,
                     kv_c,
                     kv_c,
                     query_rope=q_rope,
                     key_rope=k_pe,
-                    num_heads=layer.tp_q_head_num,
+                    num_heads=tp_q_head_num,
                     num_key_value_heads=layer.tp_k_head_num,
                     input_layout="BSND",
                     atten_mask=None,
@@ -2799,6 +2840,7 @@ class AscendAttnBackend(AttentionBackend):
                     block_size=self.page_size,
                     actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_int,
                 )
+                attn_output = attn_output[:, :, : layer.tp_q_head_num, :]
             else:
                 assert (
                     self.graph_mode == False

python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py

@@ -0,0 +1,634 @@
+import math
+from typing import Optional
+
+import torch
+from sgl_kernel_npu.fla.kda_chunk_delta_h import (
+    chunk_gated_delta_rule_fwd_h_npu,
+)
+from sgl_kernel_npu.fla.kda_gate import fused_kda_gate_npu
+from sgl_kernel_npu.fla.kda_prefill import (
+    chunk_gla_fwd_o_gk_npu,
+    recompute_w_u_fwd_npu,
+)
+from sgl_kernel_npu.fla.kda_target_verify import kda_target_verify_npu
+from sgl_kernel_npu.fla.solve_tril import solve_tril_npu
+from sgl_kernel_npu.fla.utils import prepare_chunk_indices
+from sgl_kernel_npu.mamba.causal_conv1d import (
+    causal_conv1d_fn_npu,
+    causal_conv1d_update_npu,
+)
+from sgl_kernel_npu.mamba.causal_conv1d_verify import (
+    causal_conv1d_linear_verify_npu,
+)
+
+from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum
+from sglang.kernels.ops.attention.fla.kda import chunk_kda_scaled_dot_kkt_fwd
+from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
+from sglang.srt.layers.attention.linear.kda_backend import (
+    KDAAttnBackend,
+    ragged_verify_dense_scatter_indices,
+)
+from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+_LOG2_E = math.log2(math.e)
+
+
+class _AscendKDAExtendKernel:
+    """Ascend-only KDA prefill decomposition backed by sgl-kernel-npu."""
+
+    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,
+        return_intermediate_states: bool = False,
+        **kwargs,
+    ):
+        chunk_size = 64
+        q = l2norm_fwd(q.contiguous())
+        k = l2norm_fwd(k.contiguous())
+        v = v.contiguous()
+        beta = beta.contiguous()
+        chunk_indices = prepare_chunk_indices(query_start_loc, chunk_size)
+        g = chunk_local_cumsum(
+            g.contiguous(),
+            chunk_size=chunk_size,
+            scale=_LOG2_E,
+            cu_seqlens=query_start_loc,
+            chunk_indices=chunk_indices,
+        )
+
+        triangular, query_key = chunk_kda_scaled_dot_kkt_fwd(
+            q=q,
+            k=k,
+            gk=g,
+            beta=beta,
+            scale=k.shape[-1] ** -0.5,
+            cu_seqlens=query_start_loc,
+            output_dtype=torch.float32,
+        )
+        triangular = solve_tril_npu(
+            A=triangular,
+            cu_seqlens=query_start_loc,
+            output_dtype=k.dtype,
+        )
+        w, u, gated_k = recompute_w_u_fwd_npu(
+            k=k,
+            v=v,
+            beta=beta,
+            A=triangular,
+            gk=g,
+            cu_seqlens=query_start_loc,
+            chunk_indices=chunk_indices,
+        )
+        del triangular
+        chunk_states, new_values = chunk_gated_delta_rule_fwd_h_npu(
+            k=gated_k,
+            w=w,
+            u=u,
+            gk=g,
+            initial_state=ssm_states,
+            initial_state_indices=cache_indices,
+            cu_seqlens=query_start_loc,
+            chunk_indices=chunk_indices,
+            use_exp2=True,
+        )
+        del w, u, gated_k
+        out = chunk_gla_fwd_o_gk_npu(
+            q=q,
+            v=new_values,
+            g=g,
+            A=query_key,
+            h=chunk_states,
+            out=v,
+            scale=k.shape[-1] ** -0.5,
+            cu_seqlens=query_start_loc,
+            chunk_size=chunk_size,
+            chunk_indices=chunk_indices,
+        )
+        del query_key, new_values
+        if return_intermediate_states:
+            return out, chunk_states.transpose(-1, -2).contiguous()
+        return out
+
+
+class AscendKDAAttnBackend(KDAAttnBackend):
+    """Ascend implementation of Kimi Delta Attention.
+
+    The model, scheduler, metadata, and non-operator control flow stay in the
+    shared KDA backend. This class contains only the layout and operator
+    differences required by Ascend.
+    """
+
+    supports_speculative_conv_state_snapshots: bool = True
+
+    def __init__(self, model_runner):
+        super().__init__(model_runner)
+        # The NPU pool is allocated directly as [pool, channels, window].
+        self.conv_states_shape = (
+            model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
+        )
+        self.kernel_dispatcher.extend_kernel = _AscendKDAExtendKernel()
+
+    def forward_decode(
+        self,
+        layer: RadixLinearAttention,
+        forward_batch: ForwardBatch,
+        mixed_qkv: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,
+    ):
+        """Run KDA decode against the native channel-first Ascend cache."""
+        assert isinstance(mixed_qkv, torch.Tensor)
+        layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
+        conv_states = layer_cache.conv[0]
+        ssm_states = layer_cache.temporal
+        query_start_loc = self.forward_metadata.query_start_loc
+        cache_indices = self.forward_metadata.mamba_cache_indices
+
+        qkv = causal_conv1d_update_npu(
+            mixed_qkv,
+            conv_states,
+            layer.conv_weights,
+            layer.bias,
+            activation="silu",
+            conv_state_indices=cache_indices,
+        )
+
+        if self.kernel_dispatcher.supports_packed_decode:
+            assert qkv.shape[0] == cache_indices.shape[0], (
+                "KDA packed decode requires one token per sequence (T=1): "
+                f"got {qkv.shape[0]} tokens for {cache_indices.shape[0]} requests."
+            )
+            core_attn_out = self.kernel_dispatcher.packed_decode(
+                mixed_qkv=qkv,
+                a=a,
+                b=b,
+                A_log=layer.A_log,
+                dt_bias=layer.dt_bias,
+                scale=layer.head_k_dim**-0.5,
+                ssm_states=ssm_states,
+                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=layer_cache.replayssm_d,
+                replayssm_k=layer_cache.replayssm_k,
+                replayssm_g=layer_cache.replayssm_g,
+                replayssm_write_pos=getattr(
+                    self.forward_metadata, "replayssm_write_pos", None
+                ),
+                replayssm_force_flush=getattr(
+                    self.forward_metadata, "replayssm_force_flush", None
+                ),
+            )
+        else:
+            q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)
+            q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0)
+            k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)
+            v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0)
+            core_attn_out = self.kernel_dispatcher.decode(
+                q=q,
+                k=k,
+                v=v,
+                a=a,
+                b=b,
+                A_log=layer.A_log,
+                dt_bias=layer.dt_bias,
+                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,
+            layer.layer_id,
+        )
+        return core_attn_out
+
+    def forward_extend(
+        self,
+        layer: RadixLinearAttention,
+        forward_batch: ForwardBatch,
+        mixed_qkv: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,
+    ):
+        """Run Ascend prefill without changing the shared KDA backend."""
+        assert isinstance(mixed_qkv, torch.Tensor)
+        if forward_batch.forward_mode.is_target_verify():
+            return self._forward_target_verify(layer, forward_batch, mixed_qkv, a, b)
+
+        query_start_loc = self.forward_metadata.query_start_loc
+        cache_indices = self.forward_metadata.mamba_cache_indices
+        cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
+        conv_states = cache.conv[0]
+        ssm_states = cache.temporal
+
+        if forward_batch.extend_prefix_lens is None:
+            raise RuntimeError(
+                "extend_prefix_lens cannot be None in non-TARGET_VERIFY mode."
+            )
+        has_initial_state = forward_batch.extend_prefix_lens > 0
+
+        if self.forward_metadata.has_mamba_track_mask:
+            conv_states[self.forward_metadata.conv_states_mask_indices] = mixed_qkv[
+                self.forward_metadata.track_conv_indices
+            ].transpose(-1, -2)
+
+        splits = [layer.q_dim, layer.k_dim, layer.v_dim]
+        q, k, v = mixed_qkv.transpose(0, 1).split(splits, dim=0)
+        q_conv_weight, k_conv_weight, v_conv_weight = layer.conv_weights.split(
+            splits, dim=0
+        )
+        q_conv_state, k_conv_state, v_conv_state = conv_states.split(splits, dim=-2)
+        if layer.bias is not None:
+            q_bias, k_bias, v_bias = layer.bias.split(splits, dim=0)
+        else:
+            q_bias, k_bias, v_bias = None, None, None
+
+        conv_kwargs = dict(
+            has_initial_state=has_initial_state,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+        )
+        q = self._causal_conv1d_extend(
+            q, q_conv_weight, q_bias, q_conv_state, **conv_kwargs
+        )
+        k = self._causal_conv1d_extend(
+            k, k_conv_weight, k_bias, k_conv_state, **conv_kwargs
+        )
+        v = self._causal_conv1d_extend(
+            v, v_conv_weight, v_bias, v_conv_state, **conv_kwargs
+        )
+
+        q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0)
+        k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)
+        v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0)
+        g, beta, extend_A_log, extend_dt_bias = self._prepare_extend_gate_inputs(
+            layer, a, b
+        )
+        track_ssm = self.forward_metadata.has_mamba_track_mask
+        core_attn_out = self.kernel_dispatcher.extend(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+            A_log=extend_A_log,
+            dt_bias=extend_dt_bias,
+            lower_bound=layer.lower_bound,
+            extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+            is_spec_decode=forward_batch.forward_mode.is_draft_extend_v2(),
+            return_intermediate_states=track_ssm,
+            track_ssm_h_src=(
+                self.forward_metadata.track_ssm_h_src if track_ssm else None
+            ),
+        )
+        if track_ssm:
+            core_attn_out, h = core_attn_out
+            self._track_mamba_state_extend(
+                forward_batch, h, ssm_states, self.forward_metadata
+            )
+        return core_attn_out
+
+    def _causal_conv1d_extend(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        state: torch.Tensor,
+        *,
+        has_initial_state: torch.Tensor,
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+    ) -> torch.Tensor:
+        # The Ascend varlen kernel pads in the weight dtype. K3 keeps its
+        # weights in FP32 and its persistent convolution cache in BF16, so use
+        # a compact FP32 working set for the active rows and cast it back.
+        local_indices = torch.arange(
+            cache_indices.shape[0],
+            device=cache_indices.device,
+            dtype=cache_indices.dtype,
+        )
+        state_work = state.index_select(0, cache_indices.to(torch.int64))
+        state_work = state_work.to(weight.dtype).contiguous()
+        out = causal_conv1d_fn_npu(
+            x.to(weight.dtype),
+            weight,
+            bias,
+            activation="silu",
+            conv_states=state_work,
+            has_initial_state=has_initial_state,
+            cache_indices=local_indices,
+            query_start_loc=query_start_loc,
+            seq_lens_cpu=seq_lens_cpu,
+        )
+        state.index_copy_(0, cache_indices.to(torch.int64), state_work.to(state.dtype))
+        return out.to(x.dtype).transpose(0, 1)
+
+    def _prepare_extend_gate_inputs(
+        self,
+        layer: RadixLinearAttention,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+    ]:
+        """Apply the Ascend prefill gate contract.
+
+        The checkpoint was validated with FP32 gate activation before
+        ``chunk_kda``. Keeping this platform override here leaves the shared
+        GPU model/backend paths unchanged.
+        """
+        preactivated_g = fused_kda_gate_npu(
+            g.flatten(-2),
+            layer.A_log,
+            layer.head_k_dim,
+            gate_bias=layer.dt_bias,
+            lower_bound=layer.lower_bound,
+        )
+        return preactivated_g, beta, None, None
+
+    def _forward_target_verify(
+        self,
+        layer: RadixLinearAttention,
+        forward_batch: ForwardBatch,
+        mixed_qkv: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run fixed-width DSpark verify with Ascend-native state snapshots."""
+        metadata = self.forward_metadata
+        seq_len = mixed_qkv.shape[0]
+        query_start_loc = metadata.query_start_loc
+        cache_indices = metadata.mamba_cache_indices
+
+        cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
+        intermediate_state = cache.intermediate_ssm
+        if intermediate_state is None:
+            raise RuntimeError(
+                "Ascend KDA target verify requires speculative Mamba scratch."
+            )
+
+        draft_token_num = forward_batch.spec_info.draft_token_num
+        batch_size = query_start_loc.shape[0] - 1
+        num_dense_tokens = batch_size * draft_token_num
+        ragged_layout = forward_batch.spec_info.ragged_verify_layout
+        if ragged_layout is None and seq_len == num_dense_tokens:
+            dense_token_indices = None
+            dense_qkv = mixed_qkv.view(batch_size, draft_token_num, -1)
+            dense_a = a
+            dense_b = b
+        else:
+            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_qkv = self._scatter_tokens_to_dense(
+                mixed_qkv, dense_token_indices, num_dense_tokens
+            ).view(batch_size, draft_token_num, -1)
+            dense_a = self._scatter_gate_to_dense(
+                a, dense_token_indices, num_dense_tokens
+            )
+            dense_b = self._scatter_gate_to_dense(
+                b, dense_token_indices, num_dense_tokens
+            )
+
+        intermediate_indices = self.verify_intermediate_state_indices[:batch_size]
+        processed_qkv = causal_conv1d_linear_verify_npu(
+            dense_qkv.transpose(1, 2).contiguous(),
+            cache.conv[0],
+            layer.conv_weights,
+            layer.bias,
+            cache_indices[:batch_size],
+            cache.intermediate_conv_window[0],
+            intermediate_indices,
+            activation="silu",
+            update_persistent_state=False,
+        )
+        processed_qkv = processed_qkv.transpose(1, 2).reshape(num_dense_tokens, -1)
+        q, k, v = processed_qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)
+        q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0)
+        k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)
+        v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0)
+
+        # Activate the forget gate and beta in FP32 before entering the
+        # recurrent kernel to match the checkpoint's verify contract.
+        # This stays in the Ascend backend so shared/GPU model code is unchanged.
+        preactivated_a = fused_kda_gate_npu(
+            dense_a.flatten(-2),
+            layer.A_log,
+            layer.head_k_dim,
+            gate_bias=layer.dt_bias,
+            lower_bound=layer.lower_bound,
+        )
+        preactivated_b = dense_b.float().sigmoid()
+        out = kda_target_verify_npu(
+            A_log=layer.A_log,
+            dt_bias=layer.dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            a=preactivated_a,
+            b=preactivated_b,
+            initial_state_source=cache.temporal,
+            initial_state_indices=cache_indices[:batch_size],
+            intermediate_states_buffer=intermediate_state,
+            intermediate_state_indices=intermediate_indices,
+            cache_steps=draft_token_num,
+            lower_bound=None,
+            gates_are_preactivated=True,
+        )
+        if dense_token_indices is None:
+            return out
+        padded_out = out.new_zeros(1, num_dense_tokens + 1, *out.shape[2:])
+        padded_out[:, :num_dense_tokens] = out
+        return padded_out[:, dense_token_indices]
+
+    @staticmethod
+    def _scatter_tokens_to_dense(
+        value: torch.Tensor,
+        dense_token_indices: torch.Tensor,
+        num_dense_tokens: int,
+    ) -> torch.Tensor:
+        dense = value.new_zeros((num_dense_tokens + 1, *value.shape[1:]))
+        dense.index_copy_(0, dense_token_indices, value)
+        return dense[:num_dense_tokens]
+
+    @classmethod
+    def _scatter_gate_to_dense(
+        cls,
+        value: torch.Tensor,
+        dense_token_indices: torch.Tensor,
+        num_dense_tokens: int,
+    ) -> torch.Tensor:
+        has_leading_singleton = value.ndim >= 2 and value.shape[0] == 1
+        token_value = value.squeeze(0) if has_leading_singleton else value
+        dense = cls._scatter_tokens_to_dense(
+            token_value, dense_token_indices, num_dense_tokens
+        )
+        return dense.unsqueeze(0) if has_leading_singleton else dense
+
+
+class AscendKDAHybridLinearAttnBackend:
+    """KDA-specific hybrid backend with strided destination state mover.
+
+    ``AscendHybridLinearAttnBackend`` uses ``move_intermediate_cache`` which
+    assumes a contiguous destination layout. KDA's temporal SSM state on NPU
+    is transposed (-1, -2) and requires the strided variant
+    ``move_intermediate_cache_kda`` to preserve correct (V, K) indexing.
+
+    This class overrides only ``update_mamba_state_after_mtp_verify`` to
+    substitute the KDA-aware mover; the rest of the hybrid behaviour is
+    inherited unchanged.
+    """
+
+    def __new__(cls, *args, **kwargs):
+        # Delay importing AscendHybridLinearAttnBackend to avoid circular deps.
+        from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import (
+            AscendHybridLinearAttnBackend as _Base,
+        )
+
+        # Dynamically create a subclass of _Base with our override.
+        class _AscendKDAHybrid(_Base):
+            def update_mamba_state_after_mtp_verify(
+                self,
+                last_correct_step_indices,
+                mamba_track_indices,
+                mamba_steps_to_track,
+                model,
+                req_pool_indices=None,
+            ):
+                from sgl_kernel_npu.mamba.mamba_state_update_triton import (
+                    conv_state_rollback,
+                    move_intermediate_cache_kda,
+                )
+                from sgl_kernel_npu.mamba.speculative_state_scatter import (
+                    speculative_state_scatter_npu,
+                )
+
+                del req_pool_indices
+                request_number = last_correct_step_indices.shape[0]
+
+                state_indices_tensor = (
+                    self.linear_attn_backend.forward_metadata.mamba_cache_indices[
+                        :request_number
+                    ]
+                )
+
+                mamba_caches = (
+                    self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers()
+                )
+
+                conv_states = mamba_caches.conv[0]
+                ssm_states = mamba_caches.temporal
+                intermediate_state_cache = mamba_caches.intermediate_ssm
+                dst_indices_tensor = state_indices_tensor.to(torch.int32)
+                src_indices_tensor = torch.arange(
+                    dst_indices_tensor.shape[0],
+                    device=dst_indices_tensor.device,
+                    dtype=torch.int32,
+                )
+                last_steps = last_correct_step_indices.to(torch.int32)
+
+                move_intermediate_cache_kda(
+                    ssm_states,
+                    intermediate_state_cache,
+                    dst_indices_tensor,
+                    src_indices_tensor,
+                    last_steps,
+                    h_block_size=1,
+                )
+                draft_token_num = intermediate_state_cache.shape[2]
+                has_conv_snapshots = getattr(
+                    self.linear_attn_backend,
+                    "supports_speculative_conv_state_snapshots",
+                    False,
+                )
+                if has_conv_snapshots:
+                    intermediate_conv_window_cache = (
+                        mamba_caches.intermediate_conv_window[0]
+                    )
+                    speculative_state_scatter_npu(
+                        conv_states,
+                        intermediate_conv_window_cache,
+                        dst_indices_tensor,
+                        src_indices_tensor,
+                        last_steps,
+                    )
+                if mamba_track_indices is not None:
+                    assert mamba_steps_to_track is not None
+                    mamba_track_indices = mamba_track_indices.to(torch.int32)
+                    mamba_steps_to_track = mamba_steps_to_track.to(torch.int32)
+
+                    move_intermediate_cache_kda(
+                        ssm_states,
+                        intermediate_state_cache,
+                        mamba_track_indices,
+                        src_indices_tensor,
+                        mamba_steps_to_track,
+                        h_block_size=1,
+                    )
+
+                    if has_conv_snapshots:
+                        speculative_state_scatter_npu(
+                            conv_states,
+                            intermediate_conv_window_cache,
+                            mamba_track_indices,
+                            src_indices_tensor,
+                            mamba_steps_to_track,
+                        )
+                    else:
+                        track_mask = mamba_steps_to_track >= 0
+                        track_indices = mamba_track_indices[track_mask]
+                        if track_indices.numel() > 0:
+                            conv_states[:, track_indices] = conv_states[
+                                :, dst_indices_tensor[track_mask]
+                            ]
+
+                if not has_conv_snapshots:
+                    if dst_indices_tensor.numel() > 0:
+                        conv_state_rollback(
+                            conv_states,
+                            dst_indices_tensor,
+                            last_steps,
+                            draft_token_num,
+                        )
+
+                    if (
+                        mamba_track_indices is not None
+                        and mamba_track_indices.numel() > 0
+                    ):
+                        conv_state_rollback(
+                            conv_states,
+                            mamba_track_indices,
+                            mamba_steps_to_track,
+                            draft_token_num,
+                        )
+
+                return
+
+        return _AscendKDAHybrid(*args, **kwargs)

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

@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Optional
 import torch
 
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.environ import envs
 from sglang.srt.mem_cache.memory_pool import (
     MHATokenToKVPool,
     MLATokenToKVPool,
@@ -20,19 +21,26 @@ if is_npu():
 
 
 def _init_npu_conv_state(
-    conv_state_in, conv_state_shape, speculative_num_draft_tokens: Optional[int] = None
+    conv_state_in,
+    conv_state_shape,
+    speculative_num_draft_tokens: Optional[int] = None,
+    is_kda: bool = False,
 ):
     extra_conv_len = 0
     if speculative_num_draft_tokens is not None:
         extra_conv_len = speculative_num_draft_tokens - 1
 
-    # conv_state shape (layers, pool_size, conv_wind + draft_step, dim) for conv1d ascendc ops require dim as last dim
+    # Mamba shapes are (channels, window), while KDA shapes are
+    # (window, channels). NPU kernels consume KDA state as
+    # [layers, pool, channels, window] and other Mamba state as
+    # [layers, pool, window, channels]. KDA keeps the base window fixed;
+    # speculative per-step windows live in the intermediate cache.
     conv_state = [
         torch.zeros(
             size=(
                 conv_state_in.shape[0],
                 conv_state_in.shape[1],
-                conv_shape[1] + extra_conv_len,
+                conv_shape[1] if is_kda else conv_shape[1] + extra_conv_len,
                 conv_shape[0],
             ),
             dtype=conv_state_in.dtype,
@@ -66,6 +74,9 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
         **kwargs,
     ):
         self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
+        self.use_triton_prefix_kv_cache_store = (
+            envs.SGLANG_NPU_USE_TRITON_PREFIX_KV_CACHE_STORE.get()
+        )
         super().__init__(
             size=size,
             page_size=page_size,
@@ -201,16 +212,36 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
         if self.use_fia:
             k_buffer_layer = self.k_buffer[layer_id - self.start_layer]
             v_buffer_layer = self.v_buffer[layer_id - self.start_layer]
+            num_rows = loc.numel()
+            expected_k_numel = num_rows * self.head_num * self.head_dim
+            expected_v_numel = num_rows * self.head_num * self.v_head_dim
+            if (
+                cache_k.numel() != expected_k_numel
+                or cache_v.numel() != expected_v_numel
+            ):
+                raise ValueError(
+                    "NPU FIA KV scatter row mismatch: "
+                    f"loc_rows={num_rows}, cache_k_shape={tuple(cache_k.shape)}, "
+                    f"cache_v_shape={tuple(cache_v.shape)}, "
+                    f"head_num={self.head_num}, head_dim={self.head_dim}, "
+                    f"v_head_dim={self.v_head_dim}."
+                )
 
+            # aclnnScatterNdUpdate on the deployed CANN rejects the otherwise
+            # valid 4-D [slot, 1, head, dim] update during tiling. Flatten only
+            # the singleton FIA layout axis and scatter through an equivalent
+            # 3-D view; the underlying KV storage and attention layout stay
+            # unchanged.
+            loc_indices = loc.contiguous().view(-1, 1)
             torch_npu.npu_scatter_nd_update_(
-                k_buffer_layer,
-                loc.view(-1, 1),
-                cache_k.view(-1, 1, self.head_num, self.head_dim),
+                k_buffer_layer.view(-1, self.head_num, self.head_dim),
+                loc_indices,
+                cache_k.contiguous().view(num_rows, self.head_num, self.head_dim),
             )
             torch_npu.npu_scatter_nd_update_(
-                v_buffer_layer,
-                loc.view(-1, 1),
-                cache_v.view(-1, 1, self.head_num, self.v_head_dim),
+                v_buffer_layer.view(-1, self.head_num, self.v_head_dim),
+                loc_indices,
+                cache_v.contiguous().view(num_rows, self.head_num, self.v_head_dim),
             )
         else:
             loc = loc.to(torch.int32)
@@ -226,6 +257,80 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
                 slot_indices=loc,
             )
 
+    def set_kv_buffer_prefix_valid(
+        self,
+        layer: "RadixAttention",
+        loc_2d: torch.Tensor,
+        commit_lens: torch.Tensor,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        layer_id_override: Optional[int] = None,
+    ):
+        if not self.use_triton_prefix_kv_cache_store:
+            return super().set_kv_buffer_prefix_valid(
+                layer,
+                loc_2d,
+                commit_lens,
+                cache_k,
+                cache_v,
+                k_scale,
+                v_scale,
+                layer_id_override,
+            )
+
+        if layer_id_override is not None:
+            layer_id = layer_id_override
+        else:
+            layer_id = layer.layer_id
+        if loc_2d.ndim != 2:
+            raise ValueError(f"loc_2d must be rank-2, got {tuple(loc_2d.shape)}")
+
+        num_rows = loc_2d.numel()
+        if (
+            cache_k.numel() != num_rows * self.head_num * self.head_dim
+            or cache_v.numel() != num_rows * self.head_num * self.v_head_dim
+        ):
+            raise ValueError(
+                "dense NPU KV rows must match loc_2d size: "
+                f"cache_k={tuple(cache_k.shape)}, cache_v={tuple(cache_v.shape)}, "
+                f"loc_2d={tuple(loc_2d.shape)}"
+            )
+
+        if cache_k.dtype != self.dtype:
+            if k_scale is not None:
+                cache_k.div_(k_scale)
+            if v_scale is not None:
+                cache_v.div_(v_scale)
+            cache_k = cache_k.to(self.dtype)
+            cache_v = cache_v.to(self.dtype)
+        if self.store_dtype != self.dtype:
+            cache_k = cache_k.contiguous().view(self.store_dtype)
+            cache_v = cache_v.contiguous().view(self.store_dtype)
+
+        k_buffer_layer = self.k_buffer[layer_id - self.start_layer]
+        v_buffer_layer = self.v_buffer[layer_id - self.start_layer]
+        if loc_2d.device != k_buffer_layer.device:
+            loc_2d = loc_2d.to(device=k_buffer_layer.device, non_blocking=True)
+        if commit_lens.device != k_buffer_layer.device:
+            commit_lens = commit_lens.to(
+                device=k_buffer_layer.device, non_blocking=True
+            )
+        self._debug_prefix_valid_backend = "npu_triton"
+        from sgl_kernel_npu.mem_cache.kv_cache_store import (
+            store_kv_cache_prefix_valid_npu_triton,
+        )
+
+        store_kv_cache_prefix_valid_npu_triton(
+            k_buffer_layer.view(-1, self.head_num, self.head_dim),
+            v_buffer_layer.view(-1, self.head_num, self.v_head_dim),
+            cache_k.reshape(num_rows, self.head_num, self.head_dim),
+            cache_v.reshape(num_rows, self.head_num, self.v_head_dim),
+            loc_2d,
+            commit_lens,
+        )
+
     def _chunk_copy_npu_to_cpu(self, buf_of_layers, indices):
         chunk_size = self.cpu_offloading_chunk_size
         out = []
@@ -294,10 +399,10 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
         dtype: torch.dtype,
         kv_lora_rank: int,
         qk_rope_head_dim: int,
-        index_head_dim: Optional[int],
         layer_num: int,
         device: str,
         enable_memory_saver: bool,
+        index_head_dim: Optional[int] = None,
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
     ):

python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py

@@ -203,8 +203,10 @@ def forward_mla_prepare_npu(
                 k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
                 k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
             else:
-                if qkv_latent.shape[0] < 65536 and not dsa_use_prefill_cp(
-                    forward_batch
+                if (
+                    qkv_latent.shape[0] < 65536
+                    and not dsa_use_prefill_cp(forward_batch)
+                    and not getattr(m, "_disable_npu_fused_split_qk_norm", False)
                 ):
                     q, k_nope, k_pe = fused_split_qk_norm(
                         qkv_latent,
@@ -216,6 +218,8 @@ def forward_mla_prepare_npu(
                         eps=m.q_a_layernorm.variance_epsilon,
                     )
                 else:
+                    # The fused split+RMSNorm kernel is not numerically equivalent
+                    # on Ascend. Keep the unfused path for models that opt out.
                     q, latent_cache = qkv_latent.split(
                         [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
                         dim=-1,
@@ -245,7 +249,8 @@ def forward_mla_prepare_npu(
 
         q_nope_out = q_nope_out.transpose(0, 1)
 
-        q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
+        if m.rotary_emb is not None:
+            q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
 
         if dsa_use_prefill_cp(forward_batch):
             # support allgather+rerrange
@@ -297,14 +302,16 @@ def forward_mla_core_npu(
 
     attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
 
-    attn_bmm_output = torch.empty(
-        (attn_output.shape[0], m.num_local_heads, m.v_head_dim),
-        dtype=attn_output.dtype,
-        device=attn_output.device,
-    )
-
     attn_output = attn_output.contiguous()
-    torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output)
+    # torch.ops.npu.batch_matmul_transpose is not numerically equivalent for
+    # Kimi-K3, so use the numerically validated torch_npu implementation.
+    attn_bmm_output = torch_npu.npu_transpose_batchmatmul(
+        attn_output,
+        m.w_vc,
+        perm_x1=(1, 0, 2),
+        perm_x2=(0, 1, 2),
+        perm_y=(1, 0, 2),
+    )
 
     attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
     output, _ = m.o_proj(attn_bmm_output)
@@ -380,8 +387,10 @@ def forward_dsa_prepare_npu(
             if q_event is not None:
                 torch.npu.current_stream().wait_event(q_event)
         else:
-            if fused_qkv_a_proj_out.shape[0] < 65535 and not dsa_use_prefill_cp(
-                forward_batch
+            if (
+                fused_qkv_a_proj_out.shape[0] < 65535
+                and not dsa_use_prefill_cp(forward_batch)
+                and not getattr(m, "_disable_npu_fused_split_qk_norm", False)
             ):
                 q_lora, k_nope, k_pe = fused_split_qk_norm(
                     fused_qkv_a_proj_out,
@@ -393,6 +402,8 @@ def forward_dsa_prepare_npu(
                     eps=m.q_a_layernorm.variance_epsilon,
                 )
             else:
+                # Keep the numerically validated unfused path for models that
+                # explicitly opt out of the fused split and RMSNorm kernel.
                 q, latent_cache = fused_qkv_a_proj_out.split(
                     [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
                 )

python/sglang/srt/hardware_backend/npu/moe/activation.py

@@ -85,6 +85,39 @@ class NPUSwigluDeepEPKernel(BaseActivation):
         return hidden_states, None
 
 
+class NPUSitu(BaseActivation):
+    """SiTU activation and optional INT8 requantization for grouped rows."""
+
+    def __init__(
+        self,
+        *,
+        need_quant: bool,
+        beta: float = 4.0,
+        linear_beta: Optional[float] = 25.0,
+    ):
+        from sgl_kernel_npu.activation.situ import situ
+
+        self.situ = situ
+        self.need_quant = need_quant
+        self.beta = float(beta)
+        self.linear_beta = None if linear_beta is None else float(linear_beta)
+
+    def _apply_activation(
+        self,
+        hidden_states: torch.Tensor,
+        group_list: torch.Tensor,
+        group_list_type: int,
+    ):
+        return self.situ(
+            hidden_states,
+            group_list,
+            group_list_type,
+            need_quant=self.need_quant,
+            beta=self.beta,
+            linear_beta=self.linear_beta,
+        )
+
+
 class NPUGeluAndMul(BaseActivation):
     def __init__(self):
         self._gelu = GeluAndMul()

python/sglang/srt/layers/activation.py

@@ -208,6 +208,11 @@ class SituAndMul(BaseFusedOp):
 
         return situ_and_mul(x, None, self.beta, self.linear_beta)
 
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        from sgl_kernel_npu.activation.situ import situ_and_mul
+
+        return situ_and_mul(x)
+
     def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
         return self.forward_native(x)
 

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

@@ -453,7 +453,16 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
             else:
                 linear_attn_backend = Mamba2AttnBackend(runner)
         elif kimi_linear_config(runner.model_config) is not None:
-            linear_attn_backend = KDAAttnBackend(runner)
+            if _is_npu:
+                from sglang.srt.hardware_backend.npu.attention.ascend_kda_backend import (
+                    AscendKDAAttnBackend,
+                    AscendKDAHybridLinearAttnBackend,
+                )
+
+                linear_attn_backend = AscendKDAAttnBackend(runner)
+                hybrid_backend_cls = AscendKDAHybridLinearAttnBackend
+            else:
+                linear_attn_backend = KDAAttnBackend(runner)
         elif hybrid_lightning_config(runner.model_config) is not None:
             linear_attn_backend = LightningAttentionBackend(runner)
         else:

python/sglang/srt/layers/attn_residual.py

@@ -24,7 +24,7 @@ import triton.language as tl
 
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
-from sglang.srt.utils import is_hip
+from sglang.srt.utils import is_hip, is_npu
 
 _BLOCK_H: int = 1024  # H = 7168 = 7 x 1024
 _MAX_ROWS: int = 16  # next_pow2(8 + 1), K3 has <= 8 snapshots
@@ -37,6 +37,8 @@ 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 is_npu():
+        return False
     if _FAST_SUPPORTED is None:
         major, _ = torch.cuda.get_device_capability()
         _FAST_SUPPORTED = major >= 10
@@ -208,7 +210,19 @@ def _mix_fused(
 ) -> torch.Tensor:
     """Triton score + combine pair: returns the pre-norm mixture."""
     T, H = prefix_sum.shape
+    if T == 0:
+        return prefix_sum
     cw = get_cw(score_proj, score_norm)
+    if is_npu():
+        from sgl_kernel_npu.kimi_k3.attn_residual import mix_fused
+
+        return mix_fused(
+            prefix_sum,
+            bank,
+            nvb,
+            cw,
+            score_norm.variance_epsilon,
+        )
     n_h_blocks = H // _BLOCK_H
 
     # Step 1: score each row (2D grid, full row-parallelism)
@@ -394,6 +408,8 @@ def _aggregate(
     into bank[:, nvb, :]); the triton path keeps the standalone .write() copy —
     the caller (AttnResidual.forward) owns that fallback.
     """
+    if prefix_sum.shape[0] == 0:
+        return prefix_sum
     if _use_fast(prefix_sum.shape[1]):
         return _aggregate_fast(
             prefix_sum,

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

@@ -10,6 +10,7 @@ import torch
 from sglang.srt.hardware_backend.npu.moe.activation import (
     AllGatherActivationWrapper,
     NPUGeluAndMul,
+    NPUSitu,
     NPUSwiglu,
     NPUSwigluDeepEPKernel,
     NPUSwigluOAI,
@@ -101,7 +102,16 @@ class AscendRunnerCore(MoeRunnerCore):
             is_quant_kernel = isinstance(
                 kernel, (NPUW4A8Int8MoEMethod, NPUW8A8Int8MoEMethod)
             )
-            self.activation = NPUSwigluDeepEPKernel(need_quant=is_quant_kernel)
+            if config.activation == "situ":
+                self.activation = NPUSitu(
+                    need_quant=is_quant_kernel,
+                    beta=(
+                        config.gemm1_alpha if config.gemm1_alpha is not None else 4.0
+                    ),
+                    linear_beta=config.gemm1_clamp_limit,
+                )
+            else:
+                self.activation = NPUSwigluDeepEPKernel(need_quant=is_quant_kernel)
         else:
             # Non‑DeepEP (ascend_tp) path
             # 1. Choose the base activation according to the quant method
@@ -169,8 +179,11 @@ class AscendRunnerCore(MoeRunnerCore):
             )
 
             # --- Activation ---
-            # The DeepEP kernel expects extra dispatch metadata
-            if isinstance(self.activation, NPUSwigluDeepEPKernel):
+            # Grouped-row activations require dispatch metadata.
+            if isinstance(
+                self.activation,
+                (NPUSwigluDeepEPKernel, NPUSitu),
+            ):
                 hidden_states, pertoken_scale = self.activation._apply_activation(
                     hidden_states,
                     group_list=expert_tokens,

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

@@ -90,6 +90,8 @@ class ModelSlimConfig(QuantizationConfig):
     Config class for ModelSlim Quantization, a NPU-specific quantization type.
     """
 
+    supports_kimi_k3_quantized_latent_projections = True
+
     def __init__(self, quant_config: Dict[str, Any] = {}):
         super().__init__()
         keys = [k for k in quant_config if isinstance(k, str)]
@@ -127,6 +129,33 @@ class ModelSlimConfig(QuantizationConfig):
     def update_packed_modules_mapping(self, mapping: Dict[str, List[str]]) -> None:
         self.packed_modules_mapping.update(mapping)
 
+    def _quant_prefix_candidates(self, prefix: str) -> List[str]:
+        """Return checkpoint-name variants without copying the large config.
+
+        Kimi-K3's upstream model uses ``mlp`` internally while its ModelSlim
+        checkpoint retains the Hugging Face ``block_sparse_moe`` hierarchy.
+        Some multimodal checkpoints also keep the outer ``language_model``
+        prefix.  Resolve those layout-only differences at the quantization
+        boundary.
+        """
+        candidates = [prefix]
+        if ".mlp." in prefix:
+            candidates.append(prefix.replace(".mlp.", ".block_sparse_moe."))
+
+        for candidate in list(candidates):
+            if candidate.startswith("language_model."):
+                candidates.append(candidate.removeprefix("language_model."))
+            else:
+                candidates.append(f"language_model.{candidate}")
+
+        return list(dict.fromkeys(candidates))
+
+    def _resolve_quant_prefix(self, prefix: str) -> str:
+        for candidate in self._quant_prefix_candidates(prefix):
+            if candidate + ".weight" in self.quant_description:
+                return candidate
+        return prefix
+
     def get_linear_method(self) -> ModelSlimLinearMethod:
         return ModelSlimLinearMethod(self)
 
@@ -176,10 +205,7 @@ class ModelSlimConfig(QuantizationConfig):
                 prefix_in_quant_config = prefix.replace(
                     proj_name, packed_modules_mapping_subset[proj_name][0]
                 )
-                # Verify the remapped prefix exists in quant_description.
-                # If not (e.g. json uses fused name as-is), fall back to original.
-                if prefix_in_quant_config + ".weight" not in self.quant_description:
-                    prefix_in_quant_config = prefix
+            prefix_in_quant_config = self._resolve_quant_prefix(prefix_in_quant_config)
             if self.is_layer_skipped(
                 prefix, packed_modules_mapping_subset
             ) or self.is_layer_skipped(prefix, self.packed_modules_mapping):
@@ -217,6 +243,7 @@ class ModelSlimConfig(QuantizationConfig):
             ("W4A4_MXFP4", ModelSlimMXFP4Scheme),
         ]
 
+        prefix = self._resolve_quant_prefix(prefix)
         quant_schemes = [self.quant_description.get(prefix + ".weight", "")]
 
         for scheme_name, scheme_class in linear_quant_schemes:
@@ -252,46 +279,54 @@ class ModelSlimConfig(QuantizationConfig):
 
         w13_scheme_name = None
         w2_scheme_name = None
-        for gate_name, up_name, down_name in naming_conventions:
-            w13_keys = [
-                f"{prefix}.0.{gate_name}.weight",
-                f"{prefix}.0.{up_name}.weight",
-            ]
-            w2_key = f"{prefix}.0.{down_name}.weight"
-            w13_entries = {
-                key: self.quant_description[key]
-                for key in w13_keys
-                if key in self.quant_description
-            }
-            if w13_entries and w2_key in self.quant_description:
-                w13_names = list(w13_entries.values())
-                # For w13, both projections must agree on the scheme
-                unique_w13 = set(w13_names)
-                if len(unique_w13) > 1:
-                    raise ValueError(
-                        f"Mismatched ModelSlim quantization for W13 in layer {prefix}: "
-                        f"{w13_entries}"
-                    )
-                w13_scheme_name = w13_names[0]
-                w2_scheme_name = self.quant_description[w2_key]
+        resolved_prefix = prefix
+        for candidate in self._quant_prefix_candidates(prefix):
+            for gate_name, up_name, down_name in naming_conventions:
+                w13_keys = [
+                    f"{candidate}.0.{gate_name}.weight",
+                    f"{candidate}.0.{up_name}.weight",
+                ]
+                w2_key = f"{candidate}.0.{down_name}.weight"
+                w13_entries = {
+                    key: self.quant_description[key]
+                    for key in w13_keys
+                    if key in self.quant_description
+                }
+                if w13_entries and w2_key in self.quant_description:
+                    w13_names = list(w13_entries.values())
+                    # For w13, both projections must agree on the scheme
+                    unique_w13 = set(w13_names)
+                    if len(unique_w13) > 1:
+                        raise ValueError(
+                            "Mismatched ModelSlim quantization for W13 in layer "
+                            f"{prefix}: {w13_entries}"
+                        )
+                    w13_scheme_name = w13_names[0]
+                    w2_scheme_name = self.quant_description[w2_key]
+                    resolved_prefix = candidate
+                    break
+            if w13_scheme_name is not None:
                 break
 
         if w13_scheme_name is None:
             # Build a helpful error message listing all attempted key patterns
             all_attempted = []
-            for gate_name, up_name, down_name in naming_conventions:
-                w13_keys = [
-                    f"{prefix}.0.{gate_name}.weight",
-                    f"{prefix}.0.{up_name}.weight",
-                ]
-                w2_key = f"{prefix}.0.{down_name}.weight"
-                w13_found = any(k in self.quant_description for k in w13_keys)
-                w2_found = w2_key in self.quant_description
-                status = (
-                    f"({gate_name}/{up_name}={'found' if w13_found else 'missing'}, "
-                    f"{down_name}={'found' if w2_found else 'missing'})"
-                )
-                all_attempted.append(status)
+            for candidate in self._quant_prefix_candidates(prefix):
+                for gate_name, up_name, down_name in naming_conventions:
+                    w13_keys = [
+                        f"{candidate}.0.{gate_name}.weight",
+                        f"{candidate}.0.{up_name}.weight",
+                    ]
+                    w2_key = f"{candidate}.0.{down_name}.weight"
+                    w13_found = any(k in self.quant_description for k in w13_keys)
+                    w2_found = w2_key in self.quant_description
+                    status = (
+                        f"{candidate} "
+                        f"({gate_name}/{up_name}="
+                        f"{'found' if w13_found else 'missing'}, "
+                        f"{down_name}={'found' if w2_found else 'missing'})"
+                    )
+                    all_attempted.append(status)
             raise ValueError(
                 f"Missing ModelSlim MoE quantization description for layer {prefix}: "
                 + "; ".join(all_attempted)
@@ -306,7 +341,9 @@ class ModelSlimConfig(QuantizationConfig):
         def instantiate(name, weight_group):
             cls = scheme_map.get(name)
             if cls is None:
-                logger.warning(f"Unsupported scheme '{name}' for layer {prefix}")
+                logger.warning(
+                    f"Unsupported scheme '{name}' for layer {resolved_prefix}"
+                )
                 return None
             return cls(self, weight_group)
 
@@ -335,6 +372,7 @@ class ModelSlimConfig(QuantizationConfig):
 
             is_skipped = None
             for shard_prefix in shard_prefixes:
+                shard_prefix = self._resolve_quant_prefix(shard_prefix)
                 is_shard_skipped = (
                     self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT"
                 )
@@ -348,6 +386,7 @@ class ModelSlimConfig(QuantizationConfig):
                         "to have the same precision."
                     )
         else:
+            prefix = self._resolve_quant_prefix(prefix)
             is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT"
 
         assert is_skipped is not None

python/sglang/srt/mem_cache/memory_pool.py

@@ -558,7 +558,10 @@ class MambaPool:
                     )
 
                     conv_state = _init_npu_conv_state(
-                        conv_state[0], conv_state_shape, speculative_num_draft_tokens
+                        conv_state[0],
+                        conv_state_shape,
+                        speculative_num_draft_tokens,
+                        is_kda=cache_params.is_kda,
                     )
 
                 if _is_cpu and _cpu_has_amx_support:
@@ -759,6 +762,14 @@ class MambaPool:
                     # Original dense layout (NPU/CPU, or EAGLE tree verify): one
                     # [dim, K-1] window per draft token.
                     # Shape: [num_layers, size+1, draft_tokens, dim, K-1]
+                    dense_conv_shapes = [
+                        (
+                            (conv_shape[1], conv_shape[0])
+                            if _is_npu and cache_params.is_kda
+                            else conv_shape
+                        )
+                        for conv_shape in conv_state_shape
+                    ]
                     intermediate_conv_window_cache = [
                         torch.zeros(
                             size=(
@@ -771,7 +782,7 @@ class MambaPool:
                             dtype=conv_dtype,
                             device="cuda",
                         )
-                        for conv_shape in conv_state_shape
+                        for conv_shape in dense_conv_shapes
                     ]
                     self._intermediate_conv_window_phys = intermediate_conv_window_cache
                 self.mamba_cache = self.SpeculativeState(

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

@@ -263,7 +263,17 @@ class EagerRunner(BaseRunner):
         if cp_v2_active:
             prepare_cp_forward(forward_batch)
 
-        if forward_batch.needs_forward_metadata_init() or cp_v2_active:
+        # Target verify can arrive with ``forward_metadata_ready`` set by an
+        # upstream/speculative planning step.  That mark does not initialize
+        # the final target hybrid backend, and unlike a graph replay eager has
+        # no static metadata load to fill the gap.  Re-plan target verify from
+        # the final batch every time; eager metadata is intentionally derived
+        # directly from the live ``spec_info`` tensors.
+        if (
+            forward_batch.needs_forward_metadata_init()
+            or cp_v2_active
+            or forward_batch.forward_mode.is_target_verify()
+        ):
             if model_runner.ps.attn_dcp_size > 1 and hasattr(
                 model_runner.model, "prepare_context_parallel_metadata_for_dcp"
             ):

python/sglang/srt/models/dspark.py

@@ -8,6 +8,7 @@ import torch.nn.functional as F
 from torch import nn
 
 from sglang.srt.distributed.communication_op import tensor_model_parallel_all_gather
+from sglang.srt.environ import envs
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.dflash import DFlashDraftModel
 from sglang.srt.speculative.dflash_utils import can_dflash_slice_qkv_weight
@@ -554,6 +555,8 @@ class DSparkDraftMixin:
         Cached; None (per-layer fallback) when a QKV weight cannot be sliced
         (quantized) or layers disagree on norm epsilon / bias presence.
         """
+        if not envs.SGLANG_DSPARK_STACKED_CTX_KV.get():
+            return None
         cached = getattr(self, "_stacked_ctx_kv_cache", False)
         if cached is not False:
             return cached

python/sglang/srt/models/kimi_k3.py

@@ -38,6 +38,8 @@ from sglang.srt.layers.activation import SiluAndMul, SituAndMul
 from sglang.srt.layers.attn_residual import AttnResidual, aggregate_stream, get_cw
 from sglang.srt.layers.dcp.planner import prepare_decode_context_parallel_metadata
 from sglang.srt.layers.dp_attention import (
+    attn_tp_all_gather_into_tensor,
+    attn_tp_reduce_scatter_tensor,
     dp_gather_replicate,
     dp_scatter,
     get_global_dp_buffer,
@@ -117,7 +119,7 @@ from sglang.srt.runtime_context import (
     get_parallel,
     get_server_args,
 )
-from sglang.srt.utils import is_blackwell_supported, is_hip, make_layers
+from sglang.srt.utils import is_blackwell_supported, is_hip, is_npu, make_layers
 from sglang.srt.utils.common import (
     BumpAllocator,
     add_prefix,
@@ -129,7 +131,10 @@ from sglang.srt.utils.common import (
 
 logger = logging.getLogger(__name__)
 _is_hip = is_hip()
+_is_npu = is_npu()
 _aiter_k3_opt = get_bool_env_var("SGLANG_AITER_K3_OPT")
+_k3_shared_experts_attn_tp = envs.SGLANG_K3_SHARED_EXPERTS_ATTN_TP.get()
+_k3_dense_mlp_attn_tp = envs.SGLANG_K3_DENSE_MLP_ATTN_TP.get()
 
 
 def _cdiv(a: int, b: int) -> int:
@@ -207,8 +212,6 @@ def _merge_weights_as_views(
     return merged, sizes
 
 
-# DP attention helpers.
-#
 # K3 cannot use LayerCommunicator: the attn-res aggregation kernels replace
 # input_layernorm / post_attention_layernorm, which the communicator expects
 # to own. Instead the MLP/MoE modules gather/scatter around their own body:
@@ -217,8 +220,6 @@ def _merge_weights_as_views(
 # semantics (its internal all-reduces are unchanged and required — the latent
 # reduce must happen in latent space before the norm), and the delayed
 # prefix_sum add stays local, applied after the scatter back.
-
-
 def _dp_local_buffer_group():
     """Symmetric-memory group for the local DP buffer (mirrors
     CommunicateSummableTensorPairFn._scatter_hidden_states)."""
@@ -267,6 +268,20 @@ class KimiK3MLP(nn.Module):
         tp_size: Optional[int] = None,
     ) -> None:
         super().__init__()
+        # The Ascend path shards the dense MLP inside each attention-TP
+        # replica.  The GPU K3 refactor instead gathers all DP rows and shards
+        # this one dense layer over the full TP group.  Keep the GPU default,
+        # but allow the NPU launcher to retain the proven attention-TP layout
+        # without a device-type branch in shared model code.
+        self._dense_attn_tp = (
+            _k3_dense_mlp_attn_tp
+            and is_dp_attention_enabled()
+            and tp_rank is None
+            and tp_size is None
+        )
+        if self._dense_attn_tp:
+            tp_rank = get_parallel().attn_tp_rank
+            tp_size = get_parallel().attn_tp_size
         _tp_kwargs = (
             dict(tp_rank=tp_rank, tp_size=tp_size) if tp_size is not None else {}
         )
@@ -284,6 +299,7 @@ class KimiK3MLP(nn.Module):
             bias=False,
             quant_config=quant_config,
             reduce_results=reduce_results,
+            use_dp_attention_reduce=self._dense_attn_tp,
             prefix=f"{prefix}.down_proj",
             **_tp_kwargs,
         )
@@ -308,7 +324,9 @@ class KimiK3MLP(nn.Module):
         # DP attention only when driven from the decoder layer (forward_batch
         # given); the shared-experts instance inside KimiK3MoE passes None and
         # runs on the already-gathered buffer.
-        use_dp = self._dp_attention and forward_batch is not None
+        use_dp = (
+            self._dp_attention and forward_batch is not None and not self._dense_attn_tp
+        )
         if use_dp:
             local_hidden_states = hidden_states
             hidden_states = get_global_dp_buffer(get_tp_group())
@@ -342,6 +360,8 @@ def _add3(
         return a + b
     from sglang.kernels.ops.elementwise import add3
 
+    if not add3.covered(a, b, c):
+        return (a + b) + c
     return add3.add3(a, b, c, prefetch_bc=prefetch_bc)
 
 
@@ -439,6 +459,7 @@ class KimiK3MoE(nn.Module):
             use_grouped_topk=True,
             num_expert_group=config.num_expert_group,
             topk_group=config.topk_group,
+            scoring_func=config.moe_router_activation_func,
             correction_bias=self.gate.e_score_correction_bias,
             quant_config=quant_config,
             routed_scaling_factor=self.routed_scaling_factor,
@@ -488,7 +509,11 @@ class KimiK3MoE(nn.Module):
         # (DP attention) — with every global token dispatched exactly once.
         # No DP gather and no TP reduce is needed anywhere in the region.
         _a2a_backend = get_moe_a2a_backend()
-        self._ep_a2a = _a2a_backend.is_megamoe() or _a2a_backend.is_deepep()
+        self._ep_a2a = (
+            _a2a_backend.is_megamoe()
+            or _a2a_backend.is_deepep()
+            or _a2a_backend.is_ascend_fuseep()
+        )
 
         # Defer the trtllm-gen finalize (top-k weighted unpermute) out of the
         # MoE op and fuse it into the push all-reduce's staging pass
@@ -506,7 +531,24 @@ class KimiK3MoE(nn.Module):
         # a2a: the block runs on partial batches (shard / DP-local rows), and
         # a TP-sharded partial sum could never be reduced across ranks that
         # hold different tokens.
-        self._shared_experts_tp1 = self._ep_a2a
+        self._shared_experts_tp1 = self._ep_a2a and not _k3_shared_experts_attn_tp
+        # NPU compatibility mode keeps DeepEP's DP-local token dispatch but
+        # uses the original TP-sharded shared MLP. Gather only that branch's
+        # inputs, then reduce-scatter its output back to the DP-local rows.
+        self._shared_experts_attn_tp_comm = (
+            _k3_shared_experts_attn_tp
+            and self._ep_a2a
+            and self._dp_attention
+            and get_parallel().attn_tp_size > 1
+        )
+        shared_experts_tp_kwargs = {}
+        if self._shared_experts_tp1:
+            shared_experts_tp_kwargs = dict(tp_rank=0, tp_size=1)
+        elif self._shared_experts_attn_tp_comm:
+            shared_experts_tp_kwargs = dict(
+                tp_rank=get_parallel().attn_tp_rank,
+                tp_size=get_parallel().attn_tp_size,
+            )
         if self.num_shared_experts is not None and self.num_shared_experts > 0:
             shared_intermediate_size = moe_intermediate_size * self.num_shared_experts
             self.shared_experts = KimiK3MLP(
@@ -518,7 +560,7 @@ class KimiK3MoE(nn.Module):
                 prefix=f"{prefix}.shared_experts",
                 activation_situ_beta=config.activation_situ_beta,
                 activation_situ_linear_beta=config.activation_situ_linear_beta,
-                **(dict(tp_rank=0, tp_size=1) if self._shared_experts_tp1 else {}),
+                **shared_experts_tp_kwargs,
             )
         else:
             self.shared_experts = None
@@ -538,16 +580,26 @@ class KimiK3MoE(nn.Module):
         # overlap than two streams.
         self._sbo_shared_overlap = (
             self._ep_a2a
+            and not self._shared_experts_attn_tp_comm
             and self.shared_experts is not None
             and self.alt_stream is not None
         )
 
         if self.use_latent_moe:
+            latent_quant_config = (
+                quant_config
+                if getattr(
+                    quant_config,
+                    "supports_kimi_k3_quantized_latent_projections",
+                    False,
+                )
+                else None
+            )
             self.routed_expert_down_proj = ReplicatedLinear(
                 hidden_size,
                 self.moe_hidden_size,
                 bias=False,
-                quant_config=None,
+                quant_config=latent_quant_config,
                 prefix=f"{prefix}.routed_expert_down_proj",
             )
             self.routed_expert_norm = (
@@ -559,7 +611,7 @@ class KimiK3MoE(nn.Module):
                 self.moe_hidden_size,
                 hidden_size,
                 bias=False,
-                quant_config=None,
+                quant_config=latent_quant_config,
                 prefix=f"{prefix}.routed_expert_up_proj",
             )
         else:
@@ -608,6 +660,12 @@ class KimiK3MoE(nn.Module):
         """
         if not self.use_latent_moe:
             return
+        # These merged layouts feed CUDA-only fused front kernels. Keeping the
+        # regular parameters on other devices avoids a large transient copy
+        # during post-load processing and leaves their native kernels in
+        # control of weight layout.
+        if _is_npu:
+            return
         if self.shared_experts is not None and get_moe_a2a_backend().is_none():
             mods = [
                 self.shared_experts.gate_up_proj,
@@ -774,10 +832,11 @@ class KimiK3MoE(nn.Module):
         if cfg.output_format is not TopKOutputFormat.STANDARD:
             return False
         # The kernel implements sigmoid scoring with bias-ranked ungrouped top-k.
-        # Do NOT test cfg.scoring_func: it defaults to "softmax" and TopK
-        # documents it as unused. What actually selects sigmoid is the
-        # grouped-topk-with-correction-bias route (DSv3 noaux_tc), which calls
-        # biased_grouped_topk and hardwires scoring_func="sigmoid".
+        # K3 passes moe_router_activation_func explicitly to TopK. The legacy
+        # GPU biased_grouped_topk path also hardwires sigmoid, but other platform
+        # implementations consume cfg.scoring_func directly.
+        if cfg.scoring_func != "sigmoid":
+            return False
         if not (cfg.use_grouped_topk and cfg.correction_bias is not None):
             return False
         if (cfg.num_expert_group or 1) > 1 or (cfg.topk_group or 1) > 1:
@@ -900,8 +959,26 @@ class KimiK3MoE(nn.Module):
             return self._latent_norm(latent)
         return self._latent_norm(tensor_model_parallel_all_reduce(latent))
 
+    def _forward_shared_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Run TP-sharded shared experts while DeepEP tokens stay scattered."""
+        if not self._shared_experts_attn_tp_comm:
+            return self.shared_experts(hidden_states)
+
+        group = get_parallel().attn_tp_group
+        # SP-MoE presents one contiguous token shard per attention-TP rank;
+        # the DP local buffer is the full reassembled per-replica batch.
+        gathered_hidden_states = get_local_dp_buffer(group)
+        attn_tp_all_gather_into_tensor(gathered_hidden_states, hidden_states)
+        gathered_shared_output = self.shared_experts(gathered_hidden_states)
+        shared_output = torch.empty_like(hidden_states)
+        attn_tp_reduce_scatter_tensor(shared_output, gathered_shared_output)
+        return shared_output
+
     def _forward_unfused(
-        self, hidden_states: torch.Tensor, *, prefix_sum: Optional[torch.Tensor]
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        prefix_sum: Optional[torch.Tensor],
     ) -> torch.Tensor:
         """Front section with three separate GEMMs, each reading
         hidden_states: shared-expert MLP, router gate, latent down-proj."""
@@ -924,10 +1001,10 @@ class KimiK3MoE(nn.Module):
             if self._sbo_shared_overlap:
                 self.alt_stream.wait_stream(torch.cuda.current_stream())
                 with torch.cuda.stream(self.alt_stream):
-                    shared_output = self.shared_experts(hidden_states)
+                    shared_output = self._forward_shared_experts(hidden_states)
                     shared_event = self.alt_stream.record_event()
             else:
-                shared_output = self.shared_experts(hidden_states)
+                shared_output = self._forward_shared_experts(hidden_states)
 
         # Front: gate + TopK (+ latent down-proj when the merged front covers it).
         # The gate and the latent down-proj read the same hidden_states, so the
@@ -945,7 +1022,6 @@ class KimiK3MoE(nn.Module):
             # fp32 logits reach the radix router from moe_fused_gate.
             router_logits = self.gate(hidden_states)
             topk_output = self.topk(hidden_states, router_logits)
-
         issue_shared()
 
         if not self.use_latent_moe:
@@ -968,15 +1044,28 @@ class KimiK3MoE(nn.Module):
             )
 
         if routed_input is None:
-            routed_input, _ = self.routed_expert_down_proj(hidden_states)
+            if hidden_states.shape[0] == 0:
+                # Idle DP ranks must still enter the EP dispatch below so the
+                # active replicas can exchange routed tokens.  Ascend's
+                # quantized matmul does not accept an empty activation, so
+                # materialize its shape-only result without launching GEMM.
+                routed_input = hidden_states.new_empty((0, self.moe_hidden_size))
+            else:
+                routed_input, _ = self.routed_expert_down_proj(hidden_states)
         expert_output = (
             self._forward_mega_experts(routed_input, topk_output)
             if self._use_mega_moe
             else self.experts(routed_input, topk_output)
         )
-        latent = self._reduce_latent(expert_output)
-        # up_proj is replicated, so the routed output is now fully reduced.
-        out, _ = self.routed_expert_up_proj(latent)
+        if expert_output.shape[0] == 0:
+            # The EP combine returns one row per source token.  Keep the
+            # source-side empty result while avoiding empty RMSNorm/up-proj
+            # launches; the collective itself has already completed above.
+            out = hidden_states.new_empty((0, hidden_states.shape[1]))
+        else:
+            latent = self._reduce_latent(expert_output)
+            # up_proj is replicated, so the routed output is now fully reduced.
+            out, _ = self.routed_expert_up_proj(latent)
         if shared_event is not None:
             # SBO join: as late as possible, so the side-stream shared experts
             # get the whole routed a2a + latent tail to hide under.
@@ -984,10 +1073,16 @@ class KimiK3MoE(nn.Module):
         if shared_output is not None:
             # tp1 shared experts (SP-MoE) are complete per-rank; TP-sharded
             # ones need the partial-sum reduction.
-            if self.tp_size > 1 and not self._shared_experts_tp1:
+            if (
+                self.tp_size > 1
+                and not self._shared_experts_tp1
+                and not self._shared_experts_attn_tp_comm
+            ):
                 shared_output = tensor_model_parallel_all_reduce(shared_output)
-            return _add3(out, shared_output, prefix_sum)
-        return out if prefix_sum is None else out + prefix_sum
+            out = _add3(out, shared_output, prefix_sum)
+            return out
+        out = out if prefix_sum is None else out + prefix_sum
+        return out
 
     @cached_property
     def _moe_front_needs_dense_bf16(self) -> bool:
@@ -1581,6 +1676,8 @@ class KimiK3DeltaAttention(nn.Module):
         cuda graph capture)."""
         if not self.use_full_rank_gate:
             return
+        if _is_npu:
+            return
         self._bfa_w, sizes = _merge_weights_as_views(
             [self.f_a_proj, self.b_proj], pad_rows_to=8
         )
@@ -1599,6 +1696,8 @@ class KimiK3DeltaAttention(nn.Module):
             return
         layer = self.attn
         w = layer.conv_weights
+        if _is_npu:
+            return
         seg = 12 * 128  # compiled for H = HV = 12 heads of 128 (TP8)
         if (
             w is None
@@ -1761,6 +1860,9 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
     ) -> None:
         self.all_reduce_fusion = all_reduce_fusion
         self.use_output_gate = getattr(config, "mla_use_output_gate", False)
+        # The fused Ascend split+RMSNorm path is not numerically equivalent for
+        # Kimi-K3. Other MLA models retain the existing fused fast path.
+        self._disable_npu_fused_split_qk_norm = True
         super().__init__(
             layer_id=layer_idx,
             hidden_size=config.hidden_size,
@@ -1966,7 +2068,11 @@ class KimiK3DecoderLayer(nn.Module):
         # token shard.
         _a2a_backend = get_moe_a2a_backend()
         self._sp_moe = (
-            (_a2a_backend.is_megamoe() or _a2a_backend.is_deepep())
+            (
+                _a2a_backend.is_megamoe()
+                or _a2a_backend.is_deepep()
+                or _a2a_backend.is_ascend_fuseep()
+            )
             and self._is_moe_layer
             and get_parallel().attn_tp_group.world_size > 1
         )
@@ -2243,7 +2349,6 @@ class KimiK3DecoderLayer(nn.Module):
         hidden_states, _, _ = self._finish_attn_reduce(
             hidden_states, allow_scatter=False
         )
-
         hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
         hidden_states = self.mlp(hidden_states, forward_batch=forward_batch)
         return hidden_states, residual, False
@@ -2963,6 +3068,8 @@ class KimiK3LinearForCausalLM(nn.Module):
                 layer.self_attn, KimiK3DeltaAttention
             ):
                 continue
+            if _is_npu:
+                continue
             from sglang.kernels.ops.attention.fla.kda import (
                 precompile_k3_recompute_w_u_kernel,
             )
@@ -3016,10 +3123,20 @@ class KimiK3ForConditionalGeneration(nn.Module):
 
         self.language_model = None
         if not config.encoder_only:
+            quant_description = getattr(quant_config, "quant_description", {})
+            uses_wrapper_quant_prefix = any(
+                isinstance(name, str) and name.startswith("language_model.")
+                for name in quant_description
+            )
+            language_prefix = (
+                maybe_prefix(prefix, "language_model")
+                if uses_wrapper_quant_prefix
+                else prefix
+            )
             self.language_model = KimiK3LinearForCausalLM(
                 config.text_config,
                 quant_config,
-                prefix="",
+                prefix=language_prefix,
             )
 
     @property

python/sglang/srt/speculative/dflash_utils.py

@@ -12,10 +12,13 @@ import triton
 import triton.language as tl
 
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
-from sglang.srt.layers.sampler import apply_custom_logit_processor
+from sglang.srt.layers.sampler import (
+    apply_custom_logit_processor,
+    top_p_normalize_probs_torch,
+)
 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_hip, is_musa
+from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu
 
 DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>"
 
@@ -67,6 +70,72 @@ def is_dflash_sampling_verify_available() -> bool:
     return _DFLASH_SAMPLING_VERIFY_AVAILABLE
 
 
+def _dflash_npu_top_k_top_p_renorm_prob(
+    probs: torch.Tensor,
+    *,
+    top_ks: Optional[torch.Tensor] = None,
+    top_ps: Optional[torch.Tensor] = None,
+) -> Optional[torch.Tensor]:
+    if not is_npu() or probs.device.type != "npu":
+        return None
+    try:
+        import torch_npu
+    except ImportError:
+        return None
+    if not hasattr(torch_npu, "npu_top_k_top_p"):
+        return None
+
+    logits = probs.log()
+    npu_top_ps = (
+        top_ps.reshape(-1).to(device=probs.device, dtype=probs.dtype)
+        if top_ps is not None
+        else None
+    )
+    npu_top_ks = (
+        top_ks.reshape(-1).to(device=probs.device, dtype=torch.int32)
+        if top_ks is not None
+        else None
+    )
+    if npu_top_ks is not None and not bool(
+        torch.all((npu_top_ks >= 1) & (npu_top_ks <= 1024)).item()
+    ):
+        return None
+    filtered_logits = torch_npu.npu_top_k_top_p(logits, npu_top_ps, npu_top_ks)
+    return filtered_logits.softmax(dim=-1)
+
+
+def _dflash_top_k_renorm_prob(
+    probs: torch.Tensor, top_ks: torch.Tensor
+) -> torch.Tensor:
+    if top_k_renorm_prob is not None:
+        return top_k_renorm_prob(probs, top_ks)
+
+    npu_probs = _dflash_npu_top_k_top_p_renorm_prob(probs, top_ks=top_ks)
+    if npu_probs is not None:
+        return npu_probs
+
+    vocab_size = probs.shape[-1]
+    top_ks = top_ks.reshape(-1).to(device=probs.device, dtype=torch.int64)
+    top_ks = top_ks.clamp(min=1, max=vocab_size)
+    max_top_k = int(top_ks.max().item())
+    topk_probs, topk_indices = torch.topk(probs, k=max_top_k, dim=-1)
+    ranks = torch.arange(max_top_k, device=probs.device)[None, :]
+    topk_probs.masked_fill_(ranks >= top_ks[:, None], 0.0)
+    topk_probs.div_(topk_probs.sum(dim=-1, keepdim=True))
+    return torch.zeros_like(probs).scatter_(1, topk_indices, topk_probs)
+
+
+def _dflash_top_p_renorm_prob(
+    probs: torch.Tensor, top_ps: torch.Tensor
+) -> torch.Tensor:
+    if top_p_renorm_prob is not None:
+        return top_p_renorm_prob(probs, top_ps)
+    npu_probs = _dflash_npu_top_k_top_p_renorm_prob(probs, top_ps=top_ps)
+    if npu_probs is not None:
+        return npu_probs
+    return top_p_normalize_probs_torch(probs, top_ps)
+
+
 def dflash_draft_cell_size_per_token(
     *,
     draft_model_config: Any,
@@ -879,7 +948,7 @@ def build_dflash_verify_target_probs(
                 repeated_top_ps = torch.repeat_interleave(
                     sampling_info.top_ps, draft_token_num, dim=0
                 )
-                topk_probs = top_p_renorm_prob(topk_probs, repeated_top_ps)
+                topk_probs = _dflash_top_p_renorm_prob(topk_probs, repeated_top_ps)
 
             target_probs = torch.zeros_like(scaled_logits, dtype=topk_probs.dtype)
             target_probs.scatter_(1, topk_indices, topk_probs)
@@ -888,12 +957,12 @@ def build_dflash_verify_target_probs(
     if not sparse_topk_applied:
         target_probs = F.softmax(scaled_logits, dim=-1)
         if need_top_k:
-            target_probs = top_k_renorm_prob(
+            target_probs = _dflash_top_k_renorm_prob(
                 target_probs,
                 torch.repeat_interleave(sampling_info.top_ks, draft_token_num, dim=0),
             )
         if need_top_p:
-            target_probs = top_p_renorm_prob(
+            target_probs = _dflash_top_p_renorm_prob(
                 target_probs,
                 torch.repeat_interleave(sampling_info.top_ps, draft_token_num, dim=0),
             )

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

@@ -212,7 +212,8 @@ class DraftBlockProposer:
         confidence_tap = None
         folded = False
         if (
-            draft_sampler is not None
+            envs.SGLANG_DSPARK_FOLDED_PROPOSAL.get()
+            and draft_sampler is not None
             and fwd.can_run_graph
             and (all_greedy or draft_sampler.folded_sampling)
         ):
@@ -322,7 +323,9 @@ class DraftBlockProposer:
         draft_positions = positions_2d[:, :gamma].reshape(-1)
         draft_cache_loc = verify_cache_loc_2d[:, :gamma].reshape(-1)
 
-        draft_owns_embed = hasattr(self.draft_model, "forward_embed")
+        draft_owns_embed = envs.SGLANG_DSPARK_EMBED_IN_GRAPH.get() and hasattr(
+            self.draft_model, "forward_embed"
+        )
         draft_input_embeds: Optional[torch.Tensor] = None
         if not draft_owns_embed:
             noise_embedding = embed_module(draft_block_ids)
@@ -337,6 +340,7 @@ class DraftBlockProposer:
         else:
             raise RuntimeError("DSpark decode expected batch.seq_lens_cpu, got None")
 
+        draft_num_tokens = bs * gamma
         draft_forward_batch = ForwardBatch(
             forward_mode=ForwardMode.TARGET_VERIFY,
             batch_size=bs,
@@ -351,6 +355,10 @@ class DraftBlockProposer:
             spec_algorithm=SpeculativeAlgorithm.DSPARK,
             spec_info=self._draft_block_spec_info,
             capture_hidden_mode=CaptureHiddenMode.NULL,
+            num_token_non_padded=torch.tensor(
+                draft_num_tokens, dtype=torch.int32, device=device
+            ),
+            num_token_non_padded_cpu=draft_num_tokens,
         )
         self._fill_dp_moe_sync_metadata(draft_forward_batch, batch)
         graph_runner = self.draft_model_runner.decode_cuda_graph_runner
@@ -377,6 +385,9 @@ class DraftBlockProposer:
     def _fill_dp_moe_sync_metadata(
         self, forward_batch: ForwardBatch, batch: ScheduleBatch
     ) -> None:
+        # The dense DSpark draft still reuses the target batch's graph tier.
+        # Set graph eligibility before the DP-MoE-only metadata early return.
+        forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph
         if not self._dp_moe_sync or batch.global_num_tokens is None:
             return
         gnt, gnt_logprob = spec_scale_global_num_tokens(
@@ -400,4 +411,3 @@ class DraftBlockProposer:
         forward_batch.global_num_tokens_for_logprob_gpu = torch.tensor(
             gnt_logprob, dtype=torch.int64
         ).to(device, non_blocking=True)
-        forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph

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

@@ -350,11 +350,18 @@ class DSparkWorkerV2(BaseSpecWorker):
                 )
         with self._draft_context():
             if capture_decode_cuda_graph:
-                self._draft_sampler = self._maybe_build_draft_sampler()
-                if self._draft_sampler is not None:
-                    self.draft_model_runner.capture_tail_hooks.append(
-                        make_draft_sampler_capture_hook(self._draft_sampler)
-                    )
+                # Keep the draft model graph enabled when folded proposal is
+                # disabled, but do not capture the proposal head as a tail
+                # hook. The proposer will compute base logits and the Markov
+                # block eagerly from the graph's hidden states instead. Apart
+                # from being the intended precision fallback, skipping the
+                # unused hook avoids paying for two proposal computations.
+                if envs.SGLANG_DSPARK_FOLDED_PROPOSAL.get():
+                    self._draft_sampler = self._maybe_build_draft_sampler()
+                    if self._draft_sampler is not None:
+                        self.draft_model_runner.capture_tail_hooks.append(
+                            make_draft_sampler_capture_hook(self._draft_sampler)
+                        )
                 self._proposer.attach_draft_sampler(self._draft_sampler)
             self._draft_worker.init_cuda_graphs(
                 capture_decode_cuda_graph=capture_decode_cuda_graph