Support deepseek v4 and kimi k3 on ssd

python/sglang/kernels/ops/attention/dsv4/compress.py

@@ -409,6 +409,10 @@ def compress_forward(
     else:
         fn = module.decode if plan.is_decode else module.prefill
 
+    # The CUDA C4/C128 kernels use BF16 APE values; keep the model parameter
+    # in FP32 but convert at the fused-kernel boundary.
+    if ape.dtype != torch.bfloat16:
+        ape = ape.to(dtype=torch.bfloat16)
     fn(kv_score_buffer, kv_score_input, out, ape, *plan[1:3])
     return out
 
@@ -448,6 +452,8 @@ def compress_norm_rope_store(
             kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size, bf16_store
         )
         fn = module.forward_fp4 if use_fp4 else module.forward
+        if norm_weight.dtype != torch.bfloat16:
+            norm_weight = norm_weight.to(dtype=torch.bfloat16)
         fn(
             kv,
             plan[1],

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

@@ -34,11 +34,12 @@ def _jit_fused_tma_module(
     """Compile and cache the warp-specialized TMA aggregation kernel (per-row
     bulk copies into chunk slots; chunk_rows / occupancy / consumer_regs are
     tuning knobs). The smem ring is frozen at 2 chunk slots and PDL is always
-    on: the kernel targets SM100+, where both are unconditional wins."""
+    on: the kernel targets SM100+ except SM12x."""
     major, minor = torch.cuda.get_device_capability()
-    if major < 10:
+    if major < 10 or major == 12:
         raise RuntimeError(
-            "attn_res_fused_tma requires SM100+ (tcgen05, cp.async.bulk)"
+            "attn_res_fused_tma requires SM100+ excluding SM12x; "
+            f"SM{major}{minor} is unsupported"
         )
     args = make_cpp_args(
         _DIM,

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

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Lazy-built CUDA kernels used by the expert-pack MoE runtime."""
+
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+
+import torch
+from torch.utils.cpp_extension import load
+
+from sglang.kernels.jit.utils import KERNEL_PATH
+
+_EXTENSION_NAME = "sglang_expert_pack_mxfp4"
+
+
+@lru_cache(maxsize=1)
+def _extension():
+    source = KERNEL_PATH / "csrc" / "moe" / "expert_pack_mxfp4.cu"
+    return load(
+        name=_EXTENSION_NAME,
+        sources=[str(source)],
+        extra_cflags=["-O3"],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.getenv("SGLANG_EXPERT_PACK_BUILD_VERBOSE", "0") == "1",
+    )
+
+
+def mxfp4_matvec(
+    x: torch.Tensor,
+    cache: torch.Tensor,
+    slot_ids: torch.Tensor,
+    *,
+    role_offset: int,
+    role_bytes: int,
+    input_size: int,
+    output_size: int,
+    records_per_input: int,
+) -> torch.Tensor:
+    """Multiply selected raw GGUF MXFP4 matrices by BF16/FP16 rows."""
+
+    return _extension().mxfp4_matvec(
+        x,
+        cache,
+        slot_ids,
+        role_offset,
+        role_bytes,
+        input_size,
+        output_size,
+        records_per_input,
+    )
+
+
+def mxfp4_matvec_dual(
+    x: torch.Tensor,
+    cache: torch.Tensor,
+    slot_ids: torch.Tensor,
+    *,
+    gate_role_offset: int,
+    up_role_offset: int,
+    role_bytes: int,
+    input_size: int,
+    output_size: int,
+    records_per_input: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute gate and up projections while loading each input row once."""
+
+    return _extension().mxfp4_matvec_dual(
+        x,
+        cache,
+        slot_ids,
+        gate_role_offset,
+        up_role_offset,
+        role_bytes,
+        input_size,
+        output_size,
+        records_per_input,
+    )
+
+
+def prewarm_mxfp4_extension() -> None:
+    """Build and load the extension before the server accepts requests."""
+
+    _extension()
+
+
+def mxfp4_marlin_repack(
+    raw: torch.Tensor,
+    source_slots: torch.Tensor,
+    target_slots: torch.Tensor,
+    *,
+    role_bytes: int,
+    hidden_size: int,
+    intermediate_size: int,
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+) -> None:
+    """Repack raw GGUF objects into contiguous Marlin SoA cache tensors."""
+
+    _extension().mxfp4_marlin_repack(
+        raw,
+        source_slots,
+        target_slots,
+        role_bytes,
+        hidden_size,
+        intermediate_size,
+        w13,
+        w2,
+        w13_scale,
+        w2_scale,
+    )

python/sglang/srt/arg_groups/expert_pack_hook.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Server-argument resolution for the public expert-pack load format."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any
+
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    CudaGraphConfig,
+    Phase,
+)
+from sglang.srt.model_loader.expert_pack_config import (
+    DEEPSEEK_V4_MODEL_TYPE,
+    KIMI_K3_MODEL_TYPE,
+    validate_expert_pack_model_config,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def handle_expert_pack(server_args: Any) -> None:
+    """Normalize expert-pack settings and report all startup errors together."""
+    if server_args.load_format != "expert_pack":
+        return
+
+    errors = []
+    parallelism = (
+        ("tensor", "--tp-size", server_args.tp_size),
+        ("data", "--dp-size", server_args.dp_size),
+        ("expert", "--ep-size", server_args.ep_size),
+    )
+    for label, option, size in parallelism:
+        if size != 1:
+            errors.append(f"{label} parallelism ({option}) must be 1, got {size}")
+
+    if server_args.enforce_shared_experts_fusion:
+        errors.append(
+            "--enforce-shared-experts-fusion is incompatible with expert_pack"
+        )
+    if server_args.enable_waterfill:
+        errors.append("--enable-waterfill is incompatible with expert_pack")
+
+    explicit_cuda_graph_backends = {
+        Phase.DECODE: server_args.cuda_graph_backend_decode,
+        Phase.PREFILL: server_args.cuda_graph_backend_prefill,
+    }
+    raw_cuda_graph_config = server_args.cuda_graph_config
+    if isinstance(raw_cuda_graph_config, CudaGraphConfig):
+        raw_cuda_graph_config = raw_cuda_graph_config.to_dict()
+    for phase in Phase.ALL:
+        phase_config = (
+            raw_cuda_graph_config.get(phase, {})
+            if isinstance(raw_cuda_graph_config, dict)
+            else {}
+        )
+        explicit_backend = phase_config.get(
+            "backend", explicit_cuda_graph_backends[phase]
+        )
+        if explicit_backend not in (None, Backend.DISABLED):
+            errors.append(
+                f"expert_pack requires the {phase} CUDA graph backend to be "
+                f"disabled, got {explicit_backend!r}"
+            )
+
+    loader_config = server_args.model_loader_extra_config or {}
+    if isinstance(loader_config, str):
+        try:
+            loader_config = json.loads(loader_config)
+        except (TypeError, json.JSONDecodeError) as exc:
+            errors.append(f"--model-loader-extra-config must be valid JSON: {exc}")
+            loader_config = {}
+    if not isinstance(loader_config, dict):
+        errors.append("--model-loader-extra-config must be a JSON object")
+        loader_config = {}
+
+    # A raw GGUF path is the public input form.  Preparation is performed once
+    # here, before model-config parsing and before the loader is constructed.
+    raw_model_path = Path(server_args.model_path).expanduser()
+    raw_preparation_failed = False
+    if not errors and raw_model_path.is_file():
+        try:
+            from sglang.srt.model_loader import expert_pack_runtime
+
+            model_name = raw_model_path.name.upper()
+            if "KIMI" in model_name:
+                expert_pack_runtime.prepare_raw_kimi_server_args(
+                    server_args, loader_config
+                )
+            elif "DEEPSEEK" in model_name:
+                expert_pack_runtime.prepare_raw_deepseek_server_args(
+                    server_args, loader_config
+                )
+            else:
+                expert_pack_runtime.prepare_raw_expert_pack_server_args(
+                    server_args, loader_config
+                )
+            server_args.model_loader_extra_config = loader_config
+        except Exception as exc:
+            errors.append(f"failed to prepare raw expert_pack GGUF input: {exc}")
+            raw_preparation_failed = True
+
+    def parse_path(label: str, value: Any) -> Path | None:
+        if not value:
+            return None
+        try:
+            return Path(value).expanduser()
+        except TypeError:
+            errors.append(
+                f"{label} must be a filesystem path, got {type(value).__name__}"
+            )
+            return None
+
+    pack_path = None
+    if not raw_preparation_failed:
+        pack_path_value = loader_config.get("pack_path") or os.getenv(
+            "SGLANG_EXPERT_PACK_PATH"
+        )
+        pack_path = parse_path("pack_path", pack_path_value)
+        if pack_path is None:
+            errors.append(
+                "pack_path is required in --model-loader-extra-config or "
+                "SGLANG_EXPERT_PACK_PATH"
+            )
+        elif not pack_path.is_file():
+            errors.append(f"expert-pack file does not exist: {pack_path}")
+
+    model_kind = None
+    model_path = parse_path("--model-path", server_args.model_path)
+    if not raw_preparation_failed:
+        if model_path is None or not model_path.is_dir():
+            errors.append(
+                "--model-path must be a local GGUF shard or tokenizer/config "
+                f"directory for expert_pack, got {server_args.model_path!r}"
+            )
+        else:
+            try:
+                hf_config = server_args.get_model_config().hf_config
+            except Exception as exc:
+                errors.append(f"failed to load expert_pack model config: {exc}")
+            else:
+                model_kind, model_errors = validate_expert_pack_model_config(hf_config)
+                errors.extend(model_errors)
+
+    manifest_path_value = loader_config.get("manifest_path")
+    if model_kind == KIMI_K3_MODEL_TYPE and not manifest_path_value:
+        errors.append("Kimi-K3 requires manifest_path in loader config")
+    if manifest_path_value:
+        manifest_path = parse_path("manifest_path", manifest_path_value)
+    elif model_kind == DEEPSEEK_V4_MODEL_TYPE and pack_path is not None:
+        manifest_path = Path(str(pack_path) + ".manifest.json")
+    else:
+        manifest_path = None
+    if manifest_path is not None and not manifest_path.is_file():
+        errors.append(f"expert-pack manifest does not exist: {manifest_path}")
+
+    if model_kind == DEEPSEEK_V4_MODEL_TYPE:
+        required = (
+            "source_path",
+            "source_sha256",
+            "model_identity_sha256",
+            "config_sha256",
+        )
+        for name in required:
+            if not loader_config.get(name):
+                errors.append(f"deepseek-v4-flash loader config requires {name}")
+        source_path = parse_path("source_path", loader_config.get("source_path"))
+        if source_path is not None and not source_path.is_file():
+            errors.append(
+                f"deepseek-v4-flash source GGUF does not exist: {source_path}"
+            )
+
+    if errors:
+        details = "\n".join(f"- {error}" for error in errors)
+        raise ValueError(f"Invalid expert_pack configuration:\n{details}")
+
+    server_args.disable_cuda_graph = True
+    server_args.disable_shared_experts_fusion = True
+    if model_kind == DEEPSEEK_V4_MODEL_TYPE:
+        envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
+    logger.info(
+        "expert_pack selected: CUDA graph and shared-experts fusion are "
+        "disabled for correctness."
+    )

python/sglang/srt/configs/load_config.py

@@ -23,6 +23,7 @@ class LoadFormat(str, enum.Enum):
     SHARDED_STATE = "sharded_state"
     PRESHARDED = "presharded"
     GGUF = "gguf"
+    EXPERT_PACK = "expert_pack"
     BITSANDBYTES = "bitsandbytes"
     MISTRAL = "mistral"
     LAYERED = "layered"

python/sglang/srt/configs/model_config.py

@@ -122,7 +122,10 @@ def is_deepseek_dsa(config) -> bool:
 
 
 def is_kimi_k3(config) -> bool:
-    return _hf_arch(config) == "KimiK3ForConditionalGeneration"
+    return _hf_arch(config) in (
+        "KimiK3ForConditionalGeneration",
+        "KimiK3LinearForCausalLM",
+    )
 
 
 def is_qwen3_5(config) -> bool:
@@ -956,6 +959,7 @@ def _derive_model_shapes(self):
             self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
         elif (
             "KimiLinearForCausalLM" in self.hf_config.architectures
+            or "KimiK3LinearForCausalLM" in self.hf_config.architectures
             or "KimiK3ForConditionalGeneration" in self.hf_config.architectures
         ):
             tc = self.hf_text_config