[NPU] Support Kimi-K3 DCP with shared compact communication

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

@@ -12,11 +12,13 @@
     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
@@ -32,8 +34,10 @@ def _jit_mla_output_gate_module() -> Module:
 
 
 def covered(x: torch.Tensor, gate: torch.Tensor) -> bool:
+    # transfer_to_npu aliases Tensor.is_cuda to Tensor.is_npu; this JIT is CUDA-only.
     return (
-        x.is_cuda
+        not _is_npu
+        and x.is_cuda
         and gate.is_cuda
         and x.dtype == torch.bfloat16
         and gate.dtype == torch.bfloat16

python/sglang/srt/arg_groups/model_hook.py

@@ -146,6 +146,7 @@ def handle_model_specific_adjustments(server_args: Any):
         NONCAUSAL_FULL_ATTENTION,
         get_mimo_v2_fused_qkv_expected_tp_size,
         is_deepseek_dsa,
+        is_kimi_k3,
     )
 
     if cfg.enable_deterministic_inference:
@@ -169,11 +170,15 @@ def handle_model_specific_adjustments(server_args: Any):
     hf_config = model_config.hf_config
     model_arch = hf_config.architectures[0]
 
-    if get_platform().is_npu and cfg.dcp_size > 1 and not is_deepseek_dsa(hf_config):
+    if (
+        get_platform().is_npu
+        and cfg.dcp_size > 1
+        and not (is_deepseek_dsa(hf_config) or is_kimi_k3(hf_config))
+    ):
         raise ValueError(
             "NPU decode context parallelism is currently implemented only for "
-            "DeepSeek DSA models; got "
-            f"{model_arch}. Set --decode-context-parallel-size=1 or use a DSA model."
+            "DeepSeek DSA and Kimi-K3 models; got "
+            f"{model_arch}. Set --decode-context-parallel-size=1 or use a supported model."
         )
 
     if model_arch == "InternS2MobiusForConditionalGeneration":

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

@@ -48,11 +48,52 @@ def _require_kimi_k3_aiter_gluon_dcp_support() -> None:
         )
 
 
-@_register_for("KimiK3ForConditionalGeneration")
+@_register_for("KimiK3ForConditionalGeneration", "KimiK3LinearForCausalLM")
 def _kimi_k3_overrides(server_args: Any, hf_config: Any) -> dict:
     cfg = resolving_view(server_args)
     if cfg.dcp_size > 1:
         overrides = {}
+        if get_platform().is_npu:
+            if cfg.speculative_algorithm not in (None, "DSPARK"):
+                raise ValueError(
+                    "Kimi-K3 NPU DCP currently supports speculative decoding "
+                    "only with DSPARK."
+                )
+            if cfg.speculative_algorithm == "DSPARK":
+                from sglang.srt.speculative.ragged_verify import (
+                    RaggedVerifyMode,
+                    read_ragged_verify_mode,
+                )
+
+                ragged_mode = read_ragged_verify_mode()
+                if ragged_mode is not RaggedVerifyMode.STATIC:
+                    raise ValueError(
+                        "Kimi-K3 NPU DCP + DSPARK currently requires "
+                        "SGLANG_RAGGED_VERIFY_MODE=static; compact/cap-accept "
+                        f"are not validated under DCP (got {ragged_mode.value!r})."
+                    )
+                overrides["speculative_attention_mode"] = "decode"
+            if cfg.dcp_replicate_q_proj:
+                raise ValueError(
+                    "Kimi-K3 NPU DCP does not support --dcp-replicate-q-proj "
+                    "yet; the Ascend path uses an HCCL Q AllGather."
+                )
+            if cfg.enable_symm_mem:
+                logger.warning(
+                    "Kimi-K3 NPU DCP disables --enable-symm-mem; the initial "
+                    "Ascend path uses ordinary HCCL collectives."
+                )
+                overrides["enable_symm_mem"] = False
+            logger.info(
+                "Kimi-K3 NPU DCP selects the HCCL A2A communication path and "
+                "disables replicated Q projection."
+            )
+            overrides.update(
+                dcp_comm_backend="a2a",
+                dcp_replicate_q_proj=False,
+            )
+            return overrides
+
         if cfg.enable_symm_mem:
             logger.warning(
                 "Kimi-K3 DCP disables --enable-symm-mem due to decode CUDA "