From 2d8484740d5e248dca6b66c471206de1d81af597 Mon Sep 17 00:00:00 2001 From: hujianmin <76199145+beyondHJM@users.noreply.github.com> Date: Wed, 26 Aug 2026 10:05:12 +0800 Subject: [PATCH] Support deepseek v4 and kimi k3 on ssd (#35314) Co-authored-by: 1BIN4 <1741738350@qq.com> Co-authored-by: L-Ark Co-authored-by: Chikati Co-authored-by: mengzili --- .../deepseek_v4/benchmark_deepseek_5090.py | 492 +++++++++ .../runtime/kimi_k3/benchmark_kimi_k3_5090.py | 573 +++++++++++ .../kernels/jit/csrc/moe/expert_pack_mxfp4.cu | 565 +++++++++++ .../kernels/jit/csrc/ngram_corpus/result.h | 1 + .../kernels/ops/attention/dsv4/compress.py | 7 + python/sglang/kernels/ops/kimi_k3/attn_res.py | 7 +- .../kernels/ops/moe/expert_pack_mxfp4.py | 113 +++ .../sglang/srt/arg_groups/expert_pack_hook.py | 198 ++++ python/sglang/srt/configs/load_config.py | 1 + python/sglang/srt/configs/model_config.py | 6 +- .../srt/layers/attention/dsv4/compressor.py | 21 +- .../srt/layers/attention/dsv4/indexer.py | 8 +- python/sglang/srt/layers/linear.py | 8 + python/sglang/srt/layers/moe/expert_pack.py | 931 ++++++++++++++++++ .../srt/layers/quantization/expert_pack.py | 322 ++++++ python/sglang/srt/layers/quantization/gguf.py | 21 +- .../mxfp4_flashinfer_trtllm_moe.py | 2 + .../srt/lora/deepseek_mla_correction.py | 3 +- .../deepseek_mla_correction.py | 3 +- .../sglang/srt/model_loader/deepseek4_gguf.py | 181 ++++ .../srt/model_loader/expert_pack_config.py | 41 + .../srt/model_loader/expert_pack_loader.py | 320 ++++++ .../srt/model_loader/expert_pack_runtime.py | 575 +++++++++++ .../sglang/srt/model_loader/kimi_k3_gguf.py | 203 ++++ python/sglang/srt/model_loader/loader.py | 5 + .../attention_forward_methods/forward_mla.py | 28 +- .../forward_mla_fused_rope_cpu.py | 11 +- python/sglang/srt/models/deepseek_v2.py | 42 +- python/sglang/srt/models/deepseek_v4.py | 187 ++-- python/sglang/srt/models/kimi_k3.py | 99 +- python/sglang/srt/server_args.py | 19 + .../expert_pack/test_expert_pack_mxfp4.py | 360 +++++++ .../expert_pack/test_expert_pack_runtime.py | 444 +++++++++ .../expert_pack/test_kimi_k3_gguf.py | 115 +++ .../layers/test_minicpm_sparse_metadata.py | 51 +- .../unit/lora/test_deepseek_mla_correction.py | 24 + test/registered/unit/test_model_overrides.py | 5 +- ...test_supplied_instance_exposure_ratchet.py | 6 + tools/expert_pack/__init__.py | 1 + tools/expert_pack/build.py | 635 ++++++++++++ tools/expert_pack/format.py | 313 ++++++ tools/expert_pack/kimi_ggml.py | 738 ++++++++++++++ tools/expert_pack/prepare_deepseek_pack.py | 264 +++++ tools/expert_pack/prepare_kimi_manifest.py | 71 ++ tools/expert_pack/prepare_kimi_pack.py | 125 +++ tools/expert_pack/validate.py | 335 +++++++ 46 files changed, 8347 insertions(+), 133 deletions(-) create mode 100755 examples/runtime/deepseek_v4/benchmark_deepseek_5090.py create mode 100755 examples/runtime/kimi_k3/benchmark_kimi_k3_5090.py create mode 100644 python/sglang/kernels/jit/csrc/moe/expert_pack_mxfp4.cu create mode 100644 python/sglang/kernels/ops/moe/expert_pack_mxfp4.py create mode 100644 python/sglang/srt/arg_groups/expert_pack_hook.py create mode 100644 python/sglang/srt/layers/moe/expert_pack.py create mode 100644 python/sglang/srt/layers/quantization/expert_pack.py create mode 100644 python/sglang/srt/model_loader/deepseek4_gguf.py create mode 100644 python/sglang/srt/model_loader/expert_pack_config.py create mode 100644 python/sglang/srt/model_loader/expert_pack_loader.py create mode 100644 python/sglang/srt/model_loader/expert_pack_runtime.py create mode 100644 python/sglang/srt/model_loader/kimi_k3_gguf.py create mode 100644 test/registered/expert_pack/test_expert_pack_mxfp4.py create mode 100644 test/registered/expert_pack/test_expert_pack_runtime.py create mode 100644 test/registered/expert_pack/test_kimi_k3_gguf.py create mode 100644 test/registered/unit/lora/test_deepseek_mla_correction.py create mode 100644 tools/expert_pack/__init__.py create mode 100755 tools/expert_pack/build.py create mode 100755 tools/expert_pack/format.py create mode 100755 tools/expert_pack/kimi_ggml.py create mode 100755 tools/expert_pack/prepare_deepseek_pack.py create mode 100755 tools/expert_pack/prepare_kimi_manifest.py create mode 100755 tools/expert_pack/prepare_kimi_pack.py create mode 100755 tools/expert_pack/validate.py diff --git a/examples/runtime/deepseek_v4/benchmark_deepseek_5090.py b/examples/runtime/deepseek_v4/benchmark_deepseek_5090.py new file mode 100755 index 000000000..5171f6cc5 --- /dev/null +++ b/examples/runtime/deepseek_v4/benchmark_deepseek_5090.py @@ -0,0 +1,492 @@ +#!/usr/bin/env python3 +"""Run a one-shot DeepSeek V4 Flash expert-pack benchmark on one RTX 5090.""" + +from __future__ import annotations + +import argparse +import fcntl +import hashlib +import json +import os +import signal +import socket +import subprocess +import sys +import time +import urllib.request +from pathlib import Path +from typing import Any + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_PROMPT = "Please introduce Shenzhen" +CHAT_PREFIX = "<\uff5cbegin\u2581of\u2581sentence\uff5c>You are a helpful assistant.<\uff5cUser\uff5c>" +CHAT_SUFFIX = "<\uff5cAssistant\uff5c>" +DEFAULT_LOCK = Path("/tmp/sglang-deepseek-v4-5090-benchmark.lock") +METADATA_FORMAT_VERSION = 4 +ACTIVE_MOE_LAYERS = tuple(range(43)) + + +def cache_root() -> Path: + return Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() + + +def artifact_dir_for_source(path: Path) -> Path: + stat = path.stat() + fingerprint = hashlib.sha256( + f"{path.resolve()}:{stat.st_size}:{stat.st_mtime_ns}:{METADATA_FORMAT_VERSION}".encode() + ).hexdigest()[:20] + return cache_root() / "sglang-expert-pack" / "deepseek-v4-flash" / fingerprint + + +def find_sglang_repo() -> Path: + configured = os.environ.get("SGLANG_REPO") + if configured: + return Path(configured).expanduser().resolve() + for candidate in (SCRIPT_DIR, *SCRIPT_DIR.parents): + if (candidate / "python" / "sglang").is_dir() and ( + candidate / "tools" / "expert_pack" + ).is_dir(): + return candidate + raise RuntimeError("could not locate the SGLang repository") + + +def format_prompt(prompt: str) -> str: + return f"{CHAT_PREFIX}{prompt}{CHAT_SUFFIX}" + + +def server_url(args: argparse.Namespace) -> str: + return f"http://{args.host}:{args.port}" + + +def port_in_use(host: str, port: int, timeout: float = 0.5) -> bool: + try: + with socket.create_connection((host, port), timeout=timeout): + return True + except OSError: + return False + + +def detect_rtx_5090() -> str: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=index,name", "--format=csv,noheader"], + check=True, + capture_output=True, + text=True, + ) + rows = [line.strip() for line in result.stdout.splitlines() if line.strip()] + gpu_zero = next( + (line.split(",", 1)[1].strip() for line in rows if line.startswith("0,")), + None, + ) + if gpu_zero is None or "5090" not in gpu_zero: + raise RuntimeError( + f"CUDA device 0 must be an RTX 5090; detected: {', '.join(rows) or 'none'}" + ) + return gpu_zero + + +def build_server_command(args: argparse.Namespace) -> list[str]: + extra_config = { + "cache_vram_mib": args.expert_cache_mib, + "cache_vram_reserve_mib": args.expert_cache_reserve_mib, + "stage_slots": args.stage_slots, + "read_splits": args.read_splits, + "direct_io": args.direct_io, + "stats_flush_interval": len(ACTIVE_MOE_LAYERS), + "stats_path": str(args.stats_path), + } + return [ + sys.executable, + "-m", + "sglang.launch_server", + "--model-path", + str(args.gguf), + "--trust-remote-code", + "--load-format", + "expert_pack", + "--model-loader-extra-config", + json.dumps(extra_config, separators=(",", ":")), + "--attention-backend", + "dsv4", + "--tp-size", + "1", + "--ep-size", + "1", + "--disable-flashinfer-autotune", + "--skip-server-warmup", + "--context-length", + str(args.context_length), + "--max-total-tokens", + str(args.max_total_tokens), + "--max-running-requests", + "1", + "--mem-fraction-static", + str(args.mem_fraction_static), + "--watchdog-timeout", + str(args.watchdog_timeout), + "--host", + args.host, + "--port", + str(args.port), + ] + + +def start_server(args: argparse.Namespace) -> subprocess.Popen: + if port_in_use(args.host, args.port): + raise RuntimeError(f"server address is already in use: {server_url(args)}") + args.server_log.parent.mkdir(parents=True, exist_ok=True) + log = args.server_log.open("wb", buffering=0) + env = os.environ.copy() + python_path = [str(args.sglang_repo), str(args.sglang_repo / "python")] + if env.get("PYTHONPATH"): + python_path.append(env["PYTHONPATH"]) + env["CUDA_VISIBLE_DEVICES"] = "0" + env["PYTHONPATH"] = os.pathsep.join(python_path) + env.setdefault("SGLANG_OPT_USE_TILELANG_INDEXER", "1") + conda_lib = str(Path(sys.prefix) / "lib") + cuda_root = Path("/usr/local/cuda") + if (cuda_root / "bin" / "nvcc").is_file(): + env["CUDA_HOME"] = str(cuda_root) + env["CUDA_PATH"] = str(cuda_root) + env["PATH"] = os.pathsep.join((str(cuda_root / "bin"), env.get("PATH", ""))) + env["LD_LIBRARY_PATH"] = os.pathsep.join( + value + for value in ( + conda_lib, + str(cuda_root / "lib64") if (cuda_root / "lib64").is_dir() else None, + env.get("LD_LIBRARY_PATH"), + ) + if value + ) + command = build_server_command(args) + print( + f"SERVICE_STARTING url={server_url(args)} timeout={args.startup_timeout:.0f}s " + f"log={args.server_log}", + flush=True, + ) + process = subprocess.Popen( + command, + cwd=args.sglang_repo, + stdin=subprocess.DEVNULL, + stdout=log, + stderr=subprocess.STDOUT, + env=env, + start_new_session=True, + ) + process._benchmark_log = log # type: ignore[attr-defined] + try: + deadline = time.monotonic() + args.startup_timeout + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"SGLang exited during startup with code {process.returncode}; " + f"see {args.server_log}" + ) + if port_in_use(args.host, args.port): + print( + f"SERVICE_READY pid={process.pid} url={server_url(args)}", + flush=True, + ) + return process + time.sleep(2) + raise TimeoutError( + f"SGLang did not become ready within {args.startup_timeout:.0f}s; " + f"see {args.server_log}" + ) + except BaseException: + stop_server(process, args) + raise + + +def stop_server(process: subprocess.Popen | None, args: argparse.Namespace) -> None: + if process is None: + return + log = getattr(process, "_benchmark_log", None) + try: + if process.poll() is None: + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + pass + try: + process.wait(timeout=45) + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=15) + deadline = time.monotonic() + 10 + while ( + port_in_use(args.host, args.port, timeout=0.2) + and time.monotonic() < deadline + ): + time.sleep(0.2) + print(f"SERVICE_STOPPED pid={process.pid} url={server_url(args)}", flush=True) + finally: + if log is not None: + log.close() + + +def generate( + args: argparse.Namespace, + prompt: str, + max_new_tokens: int, + *, + stream_output: bool, +) -> dict[str, Any]: + payload = { + "text": format_prompt(prompt), + "sampling_params": { + "temperature": args.temperature, + "top_p": args.top_p, + "sampling_seed": args.seed, + "max_new_tokens": max_new_tokens, + "ignore_eos": True, + }, + "stream": True, + } + request = urllib.request.Request( + server_url(args) + "/generate", + data=json.dumps(payload, ensure_ascii=False).encode("utf-8"), + headers={"Content-Type": "application/json"}, + method="POST", + ) + started = time.perf_counter_ns() + first_token = last_token = None + completion_tokens = 0 + prompt_tokens = None + output = "" + finish_reason = None + if stream_output: + print(f"prompt: {prompt}", flush=True) + print("output: ", end="", flush=True) + + with urllib.request.urlopen(request, timeout=args.request_timeout) as response: + for raw_line in response: + now = time.perf_counter_ns() + line = raw_line.decode("utf-8").strip() + if not line: + continue + if line.startswith("data: "): + line = line[6:] + if line == "[DONE]": + continue + event = json.loads(line) + meta = event.get("meta_info") or {} + current_tokens = int(meta.get("completion_tokens", 0)) + if current_tokens > completion_tokens: + first_token = first_token or now + last_token = now + completion_tokens = current_tokens + if meta.get("prompt_tokens") is not None: + prompt_tokens = int(meta["prompt_tokens"]) + event_output = event.get("text") + if event_output is not None: + if stream_output and event_output != output: + if event_output.startswith(output): + print(event_output[len(output) :], end="", flush=True) + else: + print(f"\n[output revised]\n{event_output}", end="", flush=True) + output = event_output + finish_reason = meta.get("finish_reason", finish_reason) + if stream_output: + print(flush=True) + if first_token is None or last_token is None or prompt_tokens is None: + raise RuntimeError( + "SGLang response did not contain complete token timing metadata" + ) + ttft_s = (first_token - started) / 1e9 + decode_span_s = (last_token - first_token) / 1e9 + total_s = (time.perf_counter_ns() - started) / 1e9 + decode_intervals = max(0, completion_tokens - 1) + return { + "prompt": prompt, + "output": output, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "finish_reason": finish_reason, + "ttft_ms": ttft_s * 1000, + "prefill_token_rate": prompt_tokens / ttft_s if ttft_s > 0 else None, + "decode_token_rate": ( + decode_intervals / decode_span_s if decode_span_s > 0 else None + ), + "tpot_ms": ( + decode_span_s * 1000 / decode_intervals if decode_intervals else None + ), + "total_elapsed_s": total_s, + "end_to_end_token_rate": completion_tokens / total_s if total_s > 0 else None, + } + + +def run_benchmark(args: argparse.Namespace) -> dict[str, Any]: + process = None + try: + if args.stats_path.exists(): + args.stats_path.unlink() + process = start_server(args) + return generate(args, args.prompt, args.max_new_tokens, stream_output=True) + finally: + stop_server(process, args) + + +def read_stats(path: Path) -> dict[str, Any]: + if not path.is_file(): + raise FileNotFoundError(f"expert-pack stats were not written: {path}") + return json.loads(path.read_text(encoding="utf-8")) + + +def audit_routes(stats: dict[str, Any], expected_tokens: int) -> None: + token_counts = stats.get("route_tokens_by_layer") or [] + call_counts = stats.get("route_calls_by_layer") or [] + if len(token_counts) != len(ACTIVE_MOE_LAYERS) or len(call_counts) != len( + ACTIVE_MOE_LAYERS + ): + raise RuntimeError("DeepSeek Expert Pack stats have an unexpected layer count") + for layer in ACTIVE_MOE_LAYERS: + if call_counts[layer] <= 0 or token_counts[layer] != expected_tokens: + raise RuntimeError( + f"layer {layer} routed {token_counts[layer]} tokens in " + f"{call_counts[layer]} calls; expected {expected_tokens} tokens" + ) + if int(stats.get("fallback_count", 0)) != 0: + raise RuntimeError("the request used an expert fallback") + if int(stats.get("io_errors", 0)) != 0: + raise RuntimeError("the request encountered Expert Pack I/O errors") + + +def _git_sha(repo: Path) -> str | None: + result = subprocess.run( + ["git", "-C", str(repo), "rev-parse", "HEAD"], + capture_output=True, + text=True, + ) + return result.stdout.strip() if result.returncode == 0 else None + + +def write_report(args: argparse.Namespace, gpu: str, result: dict[str, Any]) -> None: + report = { + "format": "SGLANG-DEEPSEEK-V4-FLASH-EXPERT-PACK-BENCHMARK-v1", + "git_sha": _git_sha(args.sglang_repo), + "gpu": gpu, + "source_path": str(args.gguf), + "result": result, + "expert_pack_stats": read_stats(args.stats_path), + "server_log": str(args.server_log), + } + temporary = args.report_path.with_suffix(args.report_path.suffix + ".tmp") + temporary.write_text( + json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(args.report_path) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--gguf", + type=Path, + required=True, + help="DeepSeek-V4-Flash source GGUF; server startup derives all artifacts", + ) + parser.add_argument("--max-new-tokens", type=int, default=200) + parser.set_defaults( + prompt=DEFAULT_PROMPT, + temperature=0.0, + top_p=0.95, + seed=20260810, + host="127.0.0.1", + port=30000, + startup_timeout=1200, + request_timeout=3600, + watchdog_timeout=1800, + context_length=32768, + max_total_tokens=32768, + mem_fraction_static=0.96, + expert_cache_mib=21 * 1024, + expert_cache_reserve_mib=2 * 1024, + stage_slots=12, + read_splits=4, + direct_io=True, + ) + args = parser.parse_args(argv) + if args.max_new_tokens < 1: + parser.error("--max-new-tokens must be positive") + if not 1 <= args.port <= 65535: + parser.error("--port must be between 1 and 65535") + for name in ( + "expert_cache_mib", + "expert_cache_reserve_mib", + "stage_slots", + "read_splits", + ): + if getattr(args, name) < 1: + parser.error(f"--{name.replace('_', '-')} must be positive") + args.gguf = args.gguf.expanduser().resolve(strict=True) + args.sglang_repo = find_sglang_repo() + args.artifact_dir = artifact_dir_for_source(args.gguf).resolve() + args.server_log = args.artifact_dir / "deepseek-v4-5090-server.log" + args.stats_path = args.artifact_dir / "deepseek-v4-expert-pack.stats.json" + args.report_path = args.artifact_dir / "deepseek-v4-5090-benchmark.json" + return args + + +def handle_termination(signum: int, _frame: object) -> None: + raise KeyboardInterrupt(f"received signal {signum}") + + +def print_result(result: dict[str, Any], gpu: str) -> None: + print(f"gpu: {gpu}") + print(f"prompt_tokens: {result['prompt_tokens']}") + print(f"completion_tokens: {result['completion_tokens']}") + for name, suffix in ( + ("ttft_ms", ""), + ("prefill_token_rate", " tok/s"), + ("decode_token_rate", " tok/s"), + ("tpot_ms", " ms/token"), + ("end_to_end_token_rate", " tok/s"), + ): + value = result[name] + print(f"{name}: {'n/a' if value is None else f'{value:.3f}'}{suffix}") + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + signal.signal(signal.SIGTERM, handle_termination) + signal.signal(signal.SIGHUP, handle_termination) + args.artifact_dir.mkdir(parents=True, exist_ok=True) + lock = DEFAULT_LOCK.open("w") + try: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + print( + f"error: another benchmark is running (lock: {DEFAULT_LOCK})", + file=sys.stderr, + ) + return 2 + lock.write(f"{os.getpid()}\n") + lock.flush() + try: + gpu = detect_rtx_5090() + print(f"MODEL_INPUT_READY gpu={gpu} gguf={args.gguf}", flush=True) + result = run_benchmark(args) + stats = read_stats(args.stats_path) + audit_routes(stats, result["prompt_tokens"] + result["completion_tokens"]) + write_report(args, gpu, result) + print_result(result, gpu) + print(f"report: {args.report_path}") + print(f"server_log: {args.server_log}") + return 0 + except KeyboardInterrupt: + print("error: benchmark interrupted", file=sys.stderr) + return 130 + except Exception as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + finally: + lock.close() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/runtime/kimi_k3/benchmark_kimi_k3_5090.py b/examples/runtime/kimi_k3/benchmark_kimi_k3_5090.py new file mode 100755 index 000000000..d22375334 --- /dev/null +++ b/examples/runtime/kimi_k3/benchmark_kimi_k3_5090.py @@ -0,0 +1,573 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Run a one-shot Kimi K3 expert-pack benchmark on one RTX 5090.""" + +from __future__ import annotations + +import argparse +import fcntl +import hashlib +import json +import os +import signal +import socket +import subprocess +import sys +import time +import urllib.request +from pathlib import Path + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_PROMPT = "请介绍深圳" +DEFAULT_SERVER_LOG = SCRIPT_DIR / "logs/kimi-k3-5090-benchmark-server.log" +DEFAULT_LOCK = "/tmp/sglang-kimi-k3-5090-benchmark.lock" +METADATA_FORMAT_VERSION = 3 +ACTIVE_MOE_LAYERS = tuple(range(1, 93)) +IMMUTABLE_TOP_K = 16 + + +def cache_root() -> Path: + return Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() + + +def artifact_dir_for_source(gguf: Path) -> Path: + stat = gguf.stat() + fingerprint = hashlib.sha256( + f"{gguf.parent.resolve()}:{stat.st_size}:{stat.st_mtime_ns}:" + f"{METADATA_FORMAT_VERSION}".encode() + ).hexdigest()[:20] + return cache_root() / "sglang-expert-pack" / "kimi-k3" / fingerprint + + +def find_sglang_repo() -> Path: + configured = os.environ.get("SGLANG_REPO") + if configured: + return Path(configured).expanduser().resolve() + for candidate in (SCRIPT_DIR, *SCRIPT_DIR.parents): + if (candidate / "python" / "sglang").is_dir() and ( + candidate / "tools" / "expert_pack" + ).is_dir(): + return candidate + raise RuntimeError("could not locate the SGLang repository") + + +def write_json_atomic(path: Path, value: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + f".{os.getpid()}.tmp") + temporary.write_text( + json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def detect_rtx_5090() -> str: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=index,name", "--format=csv,noheader"], + check=True, + capture_output=True, + text=True, + ) + rows = [line.strip() for line in result.stdout.splitlines() if line.strip()] + gpu_zero = next( + (line.split(",", 1)[1].strip() for line in rows if line.startswith("0,")), + None, + ) + if gpu_zero is None or "5090" not in gpu_zero: + raise RuntimeError( + f"CUDA device 0 must be an RTX 5090; detected: {', '.join(rows)}" + ) + return gpu_zero + + +def server_url(args: argparse.Namespace) -> str: + return f"http://{args.host}:{args.port}" + + +def port_in_use(host: str, port: int, timeout: float = 0.5) -> bool: + try: + with socket.create_connection((host, port), timeout=timeout): + return True + except OSError: + return False + + +def build_server_command(args: argparse.Namespace) -> list[str]: + extra_config = { + "cache_vram_mib": args.expert_cache_mib, + "cache_vram_reserve_mib": args.expert_cache_reserve_mib, + "stage_slots": args.stage_slots, + "read_splits": args.read_splits, + "direct_io": args.direct_io, + "stats_flush_interval": len(ACTIVE_MOE_LAYERS), + "stats_path": str(args.stats_path), + } + return [ + sys.executable, + "-m", + "sglang.launch_server", + "--model-path", + str(args.gguf), + "--trust-remote-code", + "--load-format", + "expert_pack", + "--model-loader-extra-config", + json.dumps(extra_config, separators=(",", ":")), + "--tp-size", + "1", + "--ep-size", + "1", + "--disable-radix-cache", + "--mamba-radix-cache-strategy", + "no_buffer", + "--disable-overlap-schedule", + "--skip-server-warmup", + "--context-length", + str(args.context_length), + "--max-total-tokens", + str(args.max_total_tokens), + "--chunked-prefill-size", + str(args.chunked_prefill_size), + "--watchdog-timeout", + str(args.watchdog_timeout), + "--max-running-requests", + "1", + "--mem-fraction-static", + str(args.mem_fraction_static), + "--host", + args.host, + "--port", + str(args.port), + ] + + +def start_server(args: argparse.Namespace) -> subprocess.Popen: + if port_in_use(args.host, args.port): + raise RuntimeError(f"server address is already in use: {server_url(args)}") + args.server_log.parent.mkdir(parents=True, exist_ok=True) + log = args.server_log.open("wb", buffering=0) + env = os.environ.copy() + python_path = [str(args.sglang_repo), str(args.sglang_repo / "python")] + if env.get("PYTHONPATH"): + python_path.append(env["PYTHONPATH"]) + env["CUDA_VISIBLE_DEVICES"] = "0" + env["PYTHONPATH"] = os.pathsep.join(python_path) + conda_lib = str(Path(sys.prefix) / "lib") + cuda_root = Path("/usr/local/cuda") + if (cuda_root / "bin" / "nvcc").is_file(): + env["CUDA_HOME"] = str(cuda_root) + env["CUDA_PATH"] = str(cuda_root) + env["PATH"] = os.pathsep.join((str(cuda_root / "bin"), env.get("PATH", ""))) + env["LD_LIBRARY_PATH"] = os.pathsep.join( + value + for value in ( + conda_lib, + str(cuda_root / "lib64") if (cuda_root / "lib64").is_dir() else None, + env.get("LD_LIBRARY_PATH"), + ) + if value + ) + command = build_server_command(args) + print( + f"SERVICE_STARTING url={server_url(args)} timeout={args.startup_timeout:.0f}s " + f"log={args.server_log}", + flush=True, + ) + process = subprocess.Popen( + command, + stdin=subprocess.DEVNULL, + stdout=log, + stderr=subprocess.STDOUT, + env=env, + start_new_session=True, + ) + process._benchmark_log = log # type: ignore[attr-defined] + try: + deadline = time.monotonic() + args.startup_timeout + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"SGLang exited during startup with code {process.returncode}; " + f"see {args.server_log}" + ) + if port_in_use(args.host, args.port): + print( + f"SERVICE_READY pid={process.pid} url={server_url(args)}", + flush=True, + ) + return process + time.sleep(2) + raise TimeoutError( + f"SGLang did not become ready within {args.startup_timeout:.0f}s; " + f"see {args.server_log}" + ) + except BaseException: + stop_server(process, server_url(args)) + raise + + +def stop_server(process: subprocess.Popen | None, url: str) -> None: + if process is None: + return + log = getattr(process, "_benchmark_log", None) + try: + if process.poll() is None: + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + pass + try: + process.wait(timeout=45) + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=15) + deadline = time.monotonic() + 10 + while ( + port_in_use(*server_address(url), timeout=0.2) + and time.monotonic() < deadline + ): + time.sleep(0.2) + print(f"SERVICE_STOPPED pid={process.pid} url={url}", flush=True) + finally: + if log is not None: + log.close() + + +def server_address(url: str) -> tuple[str, int]: + without_scheme = url.removeprefix("http://") + host, port = without_scheme.rsplit(":", 1) + return host, int(port) + + +def make_prompt(model_dir: Path, prompt: str) -> tuple[list[int], str]: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + model_dir, trust_remote_code=True, local_files_only=True + ) + input_ids = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=True, + add_generation_prompt=True, + ) + prompt_text = tokenizer.decode(input_ids, skip_special_tokens=False) + if not input_ids: + raise ValueError("Kimi chat template produced an empty prompt") + return [int(value) for value in input_ids], prompt_text + + +def generate( + url: str, + input_ids: list[int], + prompt: str, + max_new_tokens: int, + temperature: float, + top_p: float, + seed: int, + *, + stream_output: bool, +) -> dict: + payload = { + "input_ids": input_ids, + "sampling_params": { + "temperature": temperature, + "top_p": top_p, + "sampling_seed": seed, + "max_new_tokens": max_new_tokens, + "ignore_eos": True, + }, + "return_logprob": True, + "stream": True, + } + request = urllib.request.Request( + url.rstrip("/") + "/generate", + data=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json"}, + method="POST", + ) + started = time.perf_counter_ns() + first_token = None + last_token = None + completion_tokens = 0 + prompt_tokens = None + output = "" + finish_reason = None + output_token_ids: dict[int, int] = {} + + if stream_output: + print(f"prompt: {prompt}", flush=True) + print("output: ", end="", flush=True) + with urllib.request.urlopen(request, timeout=3600) as response: + for raw_line in response: + now = time.perf_counter_ns() + line = raw_line.decode("utf-8").strip() + if not line: + continue + if line.startswith("data: "): + line = line[6:] + if line == "[DONE]": + continue + event = json.loads(line) + meta = event.get("meta_info") or {} + current_tokens = int(meta.get("completion_tokens", 0)) + if current_tokens > completion_tokens: + first_token = first_token or now + last_token = now + completion_tokens = current_tokens + if meta.get("prompt_tokens") is not None: + prompt_tokens = int(meta["prompt_tokens"]) + logprobs = meta.get("output_token_logprobs") or [] + logprob_length = int( + meta.get("output_token_logprobs_length", current_tokens) + ) + offset = logprob_length - len(logprobs) + for index, item in enumerate(logprobs): + output_token_ids[offset + index] = int(item[1]) + event_output = event.get("text") + if event_output is not None: + if stream_output and event_output != output: + if event_output.startswith(output): + print(event_output[len(output) :], end="", flush=True) + else: + print(f"\n[output revised]\n{event_output}", end="", flush=True) + output = event_output + finish_reason = meta.get("finish_reason", finish_reason) + if stream_output: + print(flush=True) + + if first_token is None or last_token is None or prompt_tokens is None: + raise RuntimeError("SGLang response omitted token timing metadata") + if prompt_tokens != len(input_ids): + raise RuntimeError( + f"server prompt token count {prompt_tokens} != tokenizer count {len(input_ids)}" + ) + ordered_token_ids = [output_token_ids[index] for index in sorted(output_token_ids)] + if len(ordered_token_ids) != completion_tokens: + raise RuntimeError( + "SGLang response omitted output token IDs: " + f"{len(ordered_token_ids)} != {completion_tokens}" + ) + ttft_s = (first_token - started) / 1e9 + decode_span_s = (last_token - first_token) / 1e9 + total_s = (time.perf_counter_ns() - started) / 1e9 + return { + "output": output, + "output_token_ids": ordered_token_ids, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "finish_reason": finish_reason, + "ttft_ms": ttft_s * 1000, + "prefill_token_rate": prompt_tokens / ttft_s if ttft_s > 0 else None, + "decode_token_rate": ( + (completion_tokens - 1) / decode_span_s + if completion_tokens > 1 and decode_span_s > 0 + else None + ), + "tpot_ms": ( + decode_span_s * 1000 / (completion_tokens - 1) + if completion_tokens > 1 + else None + ), + "total_elapsed_s": total_s, + "end_to_end_token_rate": completion_tokens / total_s if total_s > 0 else None, + } + + +def read_stats(path: Path) -> dict: + if not path.is_file(): + raise FileNotFoundError(f"expert-pack stats were not written: {path}") + return json.loads(path.read_text(encoding="utf-8")) + + +def audit_routes(stats: dict, expected_tokens: int) -> None: + token_counts = stats["route_tokens_by_layer"] + call_counts = stats["route_calls_by_layer"] + for layer in ACTIVE_MOE_LAYERS: + if ( + token_counts[layer] != expected_tokens + or call_counts[layer] != expected_tokens + ): + raise RuntimeError( + f"layer {layer} routed {token_counts[layer]} tokens in " + f"{call_counts[layer]} calls; expected {expected_tokens} exact Top-16 calls" + ) + if any( + token_counts[layer] + for layer in set(range(len(token_counts))) - set(ACTIVE_MOE_LAYERS) + ): + raise RuntimeError("routed experts outside model layers 1..92") + if int(stats.get("fallback_count", 0)) != 0: + raise RuntimeError("the request used an expert fallback") + + +def git_commit(repo: Path) -> str: + return subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--gguf", + type=Path, + required=True, + help="one Kimi-K3 GGUF shard; sibling shards and serving artifacts are derived", + ) + parser.add_argument("--prompt", default=DEFAULT_PROMPT) + parser.add_argument("--max-new-tokens", type=int, default=200) + parser.set_defaults( + host="127.0.0.1", + port=30001, + temperature=0.0, + top_p=0.95, + seed=20260813, + startup_timeout=1200, + context_length=384, + max_total_tokens=512, + chunked_prefill_size=64, + watchdog_timeout=1800, + mem_fraction_static=0.98, + expert_cache_mib=5120, + expert_cache_reserve_mib=1536, + stage_slots=16, + read_splits=1, + direct_io=True, + ) + args = parser.parse_args(argv) + if args.max_new_tokens < 1: + parser.error("--max-new-tokens must be positive") + if not 1 <= args.port <= 65535: + parser.error("--port must be between 1 and 65535") + for name in ( + "expert_cache_mib", + "expert_cache_reserve_mib", + "stage_slots", + "read_splits", + ): + if getattr(args, name) < 1: + parser.error(f"--{name.replace('_', '-')} must be positive") + args.gguf = args.gguf.expanduser().resolve(strict=True) + args.sglang_repo = find_sglang_repo() + args.artifact_dir = artifact_dir_for_source(args.gguf).resolve() + args.server_log = args.artifact_dir / DEFAULT_SERVER_LOG.name + args.stats_path = args.artifact_dir / "kimi-k3-expert-pack.stats.json" + args.report_path = args.artifact_dir / "kimi-k3-5090-benchmark.json" + return args + + +def print_result(result: dict, stats: dict) -> None: + def value(name: str) -> str: + item = result[name] + return "n/a" if item is None else f"{item:.3f}" + + print(f"prompt_tokens: {result['prompt_tokens']}") + print(f"completion_tokens: {result['completion_tokens']}") + print(f"ttft_ms: {value('ttft_ms')}") + print(f"prefill_token_rate: {value('prefill_token_rate')} tok/s") + print(f"decode_token_rate: {value('decode_token_rate')} tok/s") + print(f"tpot_ms: {value('tpot_ms')} ms/token") + print(f"end_to_end_token_rate: {value('end_to_end_token_rate')} tok/s") + print( + f"expert_cache: hits={stats['cache_hits']} misses={stats['cache_misses']} " + f"evictions={stats['cache_evictions']} reads={stats['pack_reads']} " + f"read_bytes={stats['pack_read_bytes']} h2d_bytes={stats['h2d_bytes']}" + ) + + +def main() -> int: + args = parse_args() + signal.signal(signal.SIGTERM, lambda *_: (_ for _ in ()).throw(KeyboardInterrupt())) + signal.signal(signal.SIGHUP, lambda *_: (_ for _ in ()).throw(KeyboardInterrupt())) + args.artifact_dir.mkdir(parents=True, exist_ok=True) + lock = Path(DEFAULT_LOCK).open("w") + try: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + print( + f"error: another benchmark is running (lock: {DEFAULT_LOCK})", + file=sys.stderr, + ) + return 2 + lock.write(f"{os.getpid()}\n") + lock.flush() + + process = None + try: + gpu = detect_rtx_5090() + prepared = { + "gpu": gpu, + "source_gguf": str(args.gguf), + "top_k": IMMUTABLE_TOP_K, + } + print( + f"MODEL_INPUT_READY gpu={gpu} top_k={IMMUTABLE_TOP_K} gguf={args.gguf}", + flush=True, + ) + + if args.stats_path.exists(): + args.stats_path.unlink() + process = start_server(args) + prompt_ids, prompt_text = make_prompt( + args.artifact_dir / "model-meta", args.prompt + ) + result = generate( + server_url(args), + prompt_ids, + args.prompt, + args.max_new_tokens, + args.temperature, + args.top_p, + args.seed, + stream_output=True, + ) + model_tokens = result["prompt_tokens"] + result["completion_tokens"] - 1 + stop_server(process, server_url(args)) + process = None + stats = read_stats(args.stats_path) + audit_routes(stats, model_tokens) + + report = { + **prepared, + "status": "passed", + "sglang_commit": git_commit(args.sglang_repo), + "python": sys.version, + "command": sys.argv, + "server_url": server_url(args), + "server_log": str(args.server_log), + "stats_path": str(args.stats_path), + "prompt": args.prompt, + "formatted_prompt": prompt_text, + "result": result, + "expert_pack_stats": stats, + "route_audit": { + "active_moe_layers": list(ACTIVE_MOE_LAYERS), + "immutable_top_k": IMMUTABLE_TOP_K, + "model_tokens_per_layer": model_tokens, + "fallback_count": int(stats.get("fallback_count", 0)), + }, + } + write_json_atomic(args.report_path, report) + print_result(result, stats) + print(f"report: {args.report_path}") + print(f"server_log: {args.server_log}") + return 0 + except KeyboardInterrupt: + print("error: benchmark interrupted", file=sys.stderr) + return 130 + except Exception as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + finally: + stop_server(process, server_url(args)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/sglang/kernels/jit/csrc/moe/expert_pack_mxfp4.cu b/python/sglang/kernels/jit/csrc/moe/expert_pack_mxfp4.cu new file mode 100644 index 000000000..5fece0b19 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/moe/expert_pack_mxfp4.cu @@ -0,0 +1,565 @@ +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +#include +#include +#include +#include + +namespace { + +constexpr int kQuantBlock = 32; +constexpr int kBlockBytes = 17; +constexpr int kWarpsPerBlock = 4; +constexpr int kRowsPerWarp = 4; +constexpr int kMarlinTileK = 16; +constexpr int kMarlinTileN = 64; +constexpr int kMarlinTileWords = 128; + +__device__ __forceinline__ float fp4_value(uint8_t value) { + constexpr float table[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, 0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; + return table[value & 0x0f]; +} + +template +__device__ __forceinline__ float load_scalar(const scalar_t* input, int index); + +template <> +__device__ __forceinline__ float load_scalar<__nv_bfloat16>(const __nv_bfloat16* input, int index) { + return __bfloat162float(input[index]); +} + +template <> +__device__ __forceinline__ float load_scalar(const half* input, int index) { + return __half2float(input[index]); +} + +template +__device__ __forceinline__ scalar_t store_scalar(float value); + +template <> +__device__ __forceinline__ __nv_bfloat16 store_scalar<__nv_bfloat16>(float value) { + return __float2bfloat16_rn(value); +} + +template <> +__device__ __forceinline__ half store_scalar(float value) { + return __float2half_rn(value); +} + +template +__global__ void mxfp4_matvec_kernel( + const scalar_t* __restrict__ input, + const uint8_t* __restrict__ cache, + int64_t cache_stride, + const int32_t* __restrict__ slot_ids, + int64_t role_offset, + int input_size, + int output_size, + int records, + int records_per_input, + scalar_t* __restrict__ output) { + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int output_row_base = (blockIdx.x * kWarpsPerBlock + warp) * kRowsPerWarp; + const int record = blockIdx.y; + if (record >= records || output_row_base >= output_size) { + return; + } + + const int input_row = record / records_per_input; + const scalar_t* input_ptr = input + static_cast(input_row) * input_size; + const int blocks_per_row = input_size / kQuantBlock; + const int64_t row_bytes = static_cast(blocks_per_row) * kBlockBytes; + const int32_t slot = slot_ids[record]; + const uint8_t* weight_base = cache + static_cast(slot) * cache_stride + role_offset; + + float sums[kRowsPerWarp] = {}; + for (int block = lane; block < blocks_per_row; block += 32) { + const int input_base = block * kQuantBlock; + float block_sums[kRowsPerWarp] = {}; +#pragma unroll + for (int index = 0; index < 16; ++index) { + const float input_low = load_scalar(input_ptr, input_base + index); + const float input_high = load_scalar(input_ptr, input_base + index + 16); +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const int output_row = output_row_base + row; + if (output_row < output_size) { + const uint8_t* quant = + weight_base + static_cast(output_row) * row_bytes + static_cast(block) * kBlockBytes; + const uint8_t packed = quant[index + 1]; + block_sums[row] = fmaf(input_low, fp4_value(packed), block_sums[row]); + block_sums[row] = fmaf(input_high, fp4_value(packed >> 4), block_sums[row]); + } + } + } +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const int output_row = output_row_base + row; + if (output_row < output_size) { + const uint8_t* quant = + weight_base + static_cast(output_row) * row_bytes + static_cast(block) * kBlockBytes; + const int exponent = static_cast(quant[0]) - 127; + sums[row] = fmaf(block_sums[row], ldexpf(1.0f, exponent), sums[row]); + } + } + } + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + sums[row] += __shfl_down_sync(0xffffffffu, sums[row], offset); + } + } + if (lane == 0) { +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const int output_row = output_row_base + row; + if (output_row < output_size) { + output[static_cast(record) * output_size + output_row] = store_scalar(sums[row]); + } + } + } +} + +// Compute gate and up together so the hidden-state vector is loaded once. +template +__global__ void mxfp4_matvec_dual_kernel( + const scalar_t* __restrict__ input, + const uint8_t* __restrict__ cache, + int64_t cache_stride, + const int32_t* __restrict__ slot_ids, + int64_t role_offset_a, + int64_t role_offset_b, + int input_size, + int output_size, + int records, + int records_per_input, + scalar_t* __restrict__ output_a, + scalar_t* __restrict__ output_b) { + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int output_row_base = (blockIdx.x * kWarpsPerBlock + warp) * kRowsPerWarp; + const int record = blockIdx.y; + if (record >= records || output_row_base >= output_size) { + return; + } + + const int input_row = record / records_per_input; + const scalar_t* input_ptr = input + static_cast(input_row) * input_size; + const int blocks_per_row = input_size / kQuantBlock; + const int64_t row_bytes = static_cast(blocks_per_row) * kBlockBytes; + const int32_t slot = slot_ids[record]; + const uint8_t* weight_base_a = cache + static_cast(slot) * cache_stride + role_offset_a; + const uint8_t* weight_base_b = cache + static_cast(slot) * cache_stride + role_offset_b; + + float sums_a[kRowsPerWarp] = {}; + float sums_b[kRowsPerWarp] = {}; + for (int block = lane; block < blocks_per_row; block += 32) { + const int input_base = block * kQuantBlock; + float block_sums_a[kRowsPerWarp] = {}; + float block_sums_b[kRowsPerWarp] = {}; +#pragma unroll + for (int index = 0; index < 16; ++index) { + const float input_low = load_scalar(input_ptr, input_base + index); + const float input_high = load_scalar(input_ptr, input_base + index + 16); +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const int output_row = output_row_base + row; + if (output_row < output_size) { + const uint8_t* quant_a = + weight_base_a + static_cast(output_row) * row_bytes + static_cast(block) * kBlockBytes; + const uint8_t* quant_b = + weight_base_b + static_cast(output_row) * row_bytes + static_cast(block) * kBlockBytes; + const uint8_t packed_a = quant_a[index + 1]; + const uint8_t packed_b = quant_b[index + 1]; + block_sums_a[row] = fmaf(input_low, fp4_value(packed_a), block_sums_a[row]); + block_sums_a[row] = fmaf(input_high, fp4_value(packed_a >> 4), block_sums_a[row]); + block_sums_b[row] = fmaf(input_low, fp4_value(packed_b), block_sums_b[row]); + block_sums_b[row] = fmaf(input_high, fp4_value(packed_b >> 4), block_sums_b[row]); + } + } + } +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const int output_row = output_row_base + row; + if (output_row < output_size) { + const uint8_t* quant_a = + weight_base_a + static_cast(output_row) * row_bytes + static_cast(block) * kBlockBytes; + const uint8_t* quant_b = + weight_base_b + static_cast(output_row) * row_bytes + static_cast(block) * kBlockBytes; + const int exponent_a = static_cast(quant_a[0]) - 127; + const int exponent_b = static_cast(quant_b[0]) - 127; + sums_a[row] = fmaf(block_sums_a[row], ldexpf(1.0f, exponent_a), sums_a[row]); + sums_b[row] = fmaf(block_sums_b[row], ldexpf(1.0f, exponent_b), sums_b[row]); + } + } + } + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + sums_a[row] += __shfl_down_sync(0xffffffffu, sums_a[row], offset); + sums_b[row] += __shfl_down_sync(0xffffffffu, sums_b[row], offset); + } + } + if (lane == 0) { +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const int output_row = output_row_base + row; + if (output_row < output_size) { + output_a[static_cast(record) * output_size + output_row] = store_scalar(sums_a[row]); + output_b[static_cast(record) * output_size + output_row] = store_scalar(sums_b[row]); + } + } + } +} + +__device__ __forceinline__ uint32_t load_raw_word( + const uint8_t* raw, int64_t cache_stride, int slot, int role_offset, int row, int blocks_per_row, int packed_word) { + const int block = packed_word / 4; + const int word_in_block = packed_word & 3; + const int64_t row_bytes = static_cast(blocks_per_row) * kBlockBytes; + const uint8_t* ptr = raw + static_cast(slot) * cache_stride + role_offset + + static_cast(row) * row_bytes + static_cast(block) * kBlockBytes + 1 + + word_in_block * 4; + return static_cast(ptr[0]) | (static_cast(ptr[1]) << 8) | (static_cast(ptr[2]) << 16) | + (static_cast(ptr[3]) << 24); +} + +__device__ __forceinline__ uint8_t load_raw_scale( + const uint8_t* raw, int64_t cache_stride, int slot, int role_offset, int row, int blocks_per_row, int block) { + const int64_t row_bytes = static_cast(blocks_per_row) * kBlockBytes; + const uint8_t* ptr = raw + static_cast(slot) * cache_stride + role_offset + + static_cast(row) * row_bytes + static_cast(block) * kBlockBytes; + return *ptr; +} + +__device__ __forceinline__ uint8_t marlin_scale_perm(int index) { + constexpr int local_perm[4] = {0, 2, 1, 3}; + const int interleaved = (index / 4) * 4 + local_perm[index & 3]; + return static_cast(((interleaved & 7) * 8) + (interleaved >> 3)); +} + +__device__ __forceinline__ uint8_t marlin_nibble(uint32_t word, int value_index) { + return static_cast((word >> ((value_index & 7) * 4)) & 0x0f); +} + +__global__ void mxfp4_marlin_repack_weight_kernel( + const uint8_t* __restrict__ raw, + int64_t raw_stride, + const int32_t* __restrict__ source_slots, + const int32_t* __restrict__ target_slots, + int64_t role_bytes, + int input_size, + int output_size, + bool gate_up, + int32_t* __restrict__ output, + int64_t output_stride) { + const int batch = blockIdx.y; + const int64_t total_words = static_cast(input_size / kMarlinTileK) * (output_size * 2); + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (batch >= gridDim.y || index >= total_words) return; + + const int64_t tile_span = static_cast(output_size / kMarlinTileN) * kMarlinTileWords; + const int tile_k = static_cast(index / tile_span); + const int64_t tile_rem = index % tile_span; + const int tile_n = static_cast(tile_rem / kMarlinTileWords); + const int local = static_cast(tile_rem % kMarlinTileWords); + const int warp = local & 3; + const int thread = local >> 2; + const int cur_n = warp * 16 + thread / 4; + const int tc_row = (thread & 3) * 2; + constexpr int offsets[4] = {0, 1, 8, 9}; + constexpr int pack_index[8] = {0, 2, 4, 6, 1, 3, 5, 7}; + + const int source_slot = source_slots[batch]; + const int target_slot = target_slots[batch]; + const int rows_per_role = gate_up ? output_size / 2 : output_size; + const int blocks_per_row = input_size / kQuantBlock; + const int role0_offset = 0; + const int role1_offset = static_cast(role_bytes); + const int role2_offset = static_cast(2 * role_bytes); + uint8_t values[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int value_index = tc_row + offsets[i]; + const int source_row = tile_n * kMarlinTileN + cur_n; + const int role = gate_up && source_row >= rows_per_role ? 1 : (gate_up ? 0 : 2); + const int row = gate_up ? source_row % rows_per_role : source_row; + const int role_offset = role == 0 ? role0_offset : (role == 1 ? role1_offset : role2_offset); + const uint32_t word = + load_raw_word(raw, raw_stride, source_slot, role_offset, row, blocks_per_row, tile_k * 2 + value_index / 8); + values[i] = marlin_nibble(word, value_index); + const int high_source_row = tile_n * kMarlinTileN + cur_n + 8; + const int high_role = gate_up && high_source_row >= rows_per_role ? 1 : (gate_up ? 0 : 2); + const int high_role_offset = high_role == 0 ? role0_offset : (high_role == 1 ? role1_offset : role2_offset); + const int high_row = gate_up ? high_source_row % rows_per_role : high_source_row; + const uint32_t high_word = load_raw_word( + raw, raw_stride, source_slot, high_role_offset, high_row, blocks_per_row, tile_k * 2 + value_index / 8); + values[4 + i] = marlin_nibble(high_word, value_index); + } + + uint32_t packed = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) { + packed |= static_cast(values[pack_index[i]]) << (i * 4); + } + output[static_cast(target_slot) * output_stride + index] = static_cast(packed); +} + +__global__ void mxfp4_marlin_repack_scale_kernel( + const uint8_t* __restrict__ raw, + int64_t raw_stride, + const int32_t* __restrict__ source_slots, + const int32_t* __restrict__ target_slots, + int64_t role_bytes, + int input_size, + int output_size, + bool gate_up, + uint8_t* __restrict__ output, + int64_t output_stride) { + const int batch = blockIdx.y; + const int groups = input_size / kQuantBlock; + const int64_t total = static_cast(groups) * output_size; + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (batch >= gridDim.y || index >= total) return; + + const int group = static_cast(index / output_size); + const int column = static_cast(index % output_size); + const int source_column = (column / 64) * 64 + marlin_scale_perm(column & 63); + const int rows_per_role = gate_up ? output_size / 2 : output_size; + const int role = gate_up && source_column >= rows_per_role ? 1 : (gate_up ? 0 : 2); + const int row = gate_up ? source_column % rows_per_role : source_column; + const int role_offset = role == 0 ? 0 : (role == 1 ? static_cast(role_bytes) : static_cast(2 * role_bytes)); + const uint8_t value = load_raw_scale(raw, raw_stride, source_slots[batch], role_offset, row, groups, group); + output[static_cast(target_slots[batch]) * output_stride + index] = value; +} + +void mxfp4_marlin_repack( + torch::Tensor raw, + torch::Tensor source_slots, + torch::Tensor target_slots, + int64_t role_bytes, + int64_t hidden_size, + int64_t intermediate_size, + torch::Tensor w13, + torch::Tensor w2, + torch::Tensor w13_scale, + torch::Tensor w2_scale) { + TORCH_CHECK(raw.is_cuda() && source_slots.is_cuda() && target_slots.is_cuda(), "repack inputs must be CUDA tensors"); + TORCH_CHECK(raw.scalar_type() == at::kByte && raw.dim() == 2, "raw cache must be a uint8 matrix"); + TORCH_CHECK( + source_slots.scalar_type() == at::kInt && target_slots.scalar_type() == at::kInt, "slot ids must be int32"); + TORCH_CHECK(source_slots.numel() == target_slots.numel(), "slot id size mismatch"); + TORCH_CHECK(w13.scalar_type() == at::kInt && w2.scalar_type() == at::kInt, "Marlin weights must be int32"); + TORCH_CHECK( + w13_scale.scalar_type() == at::kByte && w2_scale.scalar_type() == at::kByte, + "Marlin scales must be uint8 storage"); + TORCH_CHECK(hidden_size % 32 == 0 && intermediate_size % 32 == 0, "MXFP4 dimensions must be divisible by 32"); + const int batch = static_cast(source_slots.numel()); + if (batch == 0) return; + const int threads = 256; + const auto stream = at::cuda::getCurrentCUDAStream(); + const int w13_n = static_cast(2 * intermediate_size); + const int w2_n = static_cast(hidden_size); + const int w13_k = static_cast(hidden_size); + const int w2_k = static_cast(intermediate_size); + const int64_t w13_words = static_cast(w13_k / kMarlinTileK) * w13_n * 2; + const int64_t w2_words = static_cast(w2_k / kMarlinTileK) * w2_n * 2; + const int64_t w13_scales = static_cast(w13_k / kQuantBlock) * w13_n; + const int64_t w2_scales = static_cast(w2_k / kQuantBlock) * w2_n; + mxfp4_marlin_repack_weight_kernel<<>>( + raw.data_ptr(), + raw.stride(0), + source_slots.data_ptr(), + target_slots.data_ptr(), + role_bytes, + w13_k, + w13_n, + true, + w13.data_ptr(), + w13.stride(0)); + mxfp4_marlin_repack_weight_kernel<<>>( + raw.data_ptr(), + raw.stride(0), + source_slots.data_ptr(), + target_slots.data_ptr(), + role_bytes, + w2_k, + w2_n, + false, + w2.data_ptr(), + w2.stride(0)); + mxfp4_marlin_repack_scale_kernel<<>>( + raw.data_ptr(), + raw.stride(0), + source_slots.data_ptr(), + target_slots.data_ptr(), + role_bytes, + w13_k, + w13_n, + true, + w13_scale.data_ptr(), + w13_scale.stride(0)); + mxfp4_marlin_repack_scale_kernel<<>>( + raw.data_ptr(), + raw.stride(0), + source_slots.data_ptr(), + target_slots.data_ptr(), + role_bytes, + w2_k, + w2_n, + false, + w2_scale.data_ptr(), + w2_scale.stride(0)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +torch::Tensor mxfp4_matvec( + torch::Tensor input, + torch::Tensor cache, + torch::Tensor slot_ids, + int64_t role_offset, + int64_t role_bytes, + int64_t input_size, + int64_t output_size, + int64_t records_per_input) { + TORCH_CHECK( + input.is_cuda() && cache.is_cuda() && slot_ids.is_cuda(), "input, cache, and slot_ids must be CUDA tensors"); + TORCH_CHECK( + input.is_contiguous() && cache.is_contiguous() && slot_ids.is_contiguous(), + "input, cache, and slot_ids must be contiguous"); + TORCH_CHECK(input.scalar_type() == at::kBFloat16 || input.scalar_type() == at::kHalf, "input must be BF16 or FP16"); + TORCH_CHECK(cache.scalar_type() == at::kByte && cache.dim() == 2, "cache must be a two-dimensional uint8 tensor"); + TORCH_CHECK( + slot_ids.scalar_type() == at::kInt && slot_ids.dim() == 1, "slot_ids must be a one-dimensional int32 tensor"); + TORCH_CHECK(input.dim() == 2 && input.size(1) == input_size, "input shape does not match input_size"); + TORCH_CHECK(input_size > 0 && input_size % kQuantBlock == 0, "input_size must be divisible by 32"); + TORCH_CHECK(records_per_input > 0, "records_per_input must be positive"); + TORCH_CHECK( + slot_ids.numel() == input.size(0) * records_per_input, + "slot count does not match input rows and records_per_input"); + const int64_t expected_role_bytes = output_size * (input_size / kQuantBlock) * kBlockBytes; + TORCH_CHECK(role_bytes == expected_role_bytes, "role byte count does not match matrix dimensions"); + TORCH_CHECK(role_offset >= 0 && role_offset + role_bytes <= cache.size(1), "role range is outside each cache slot"); + + const auto records = slot_ids.numel(); + auto output = torch::empty({records, output_size}, input.options()); + const dim3 block(kWarpsPerBlock * 32); + const dim3 grid((output_size + kWarpsPerBlock * kRowsPerWarp - 1) / (kWarpsPerBlock * kRowsPerWarp), records); + const auto stream = at::cuda::getCurrentCUDAStream(); + if (input.scalar_type() == at::kBFloat16) { + mxfp4_matvec_kernel<<>>( + reinterpret_cast(input.data_ptr()), + cache.data_ptr(), + cache.stride(0), + slot_ids.data_ptr(), + role_offset, + input_size, + output_size, + records, + records_per_input, + reinterpret_cast<__nv_bfloat16*>(output.data_ptr())); + } else { + mxfp4_matvec_kernel<<>>( + reinterpret_cast(input.data_ptr()), + cache.data_ptr(), + cache.stride(0), + slot_ids.data_ptr(), + role_offset, + input_size, + output_size, + records, + records_per_input, + reinterpret_cast(output.data_ptr())); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +std::tuple mxfp4_matvec_dual( + torch::Tensor input, + torch::Tensor cache, + torch::Tensor slot_ids, + int64_t role_offset_a, + int64_t role_offset_b, + int64_t role_bytes, + int64_t input_size, + int64_t output_size, + int64_t records_per_input) { + TORCH_CHECK( + input.is_cuda() && cache.is_cuda() && slot_ids.is_cuda(), "input, cache, and slot_ids must be CUDA tensors"); + TORCH_CHECK( + input.is_contiguous() && cache.is_contiguous() && slot_ids.is_contiguous(), + "input, cache, and slot_ids must be contiguous"); + TORCH_CHECK(input.scalar_type() == at::kBFloat16 || input.scalar_type() == at::kHalf, "input must be BF16 or FP16"); + TORCH_CHECK(cache.scalar_type() == at::kByte && cache.dim() == 2, "cache must be a two-dimensional uint8 tensor"); + TORCH_CHECK( + slot_ids.scalar_type() == at::kInt && slot_ids.dim() == 1, "slot_ids must be a one-dimensional int32 tensor"); + TORCH_CHECK(input.dim() == 2 && input.size(1) == input_size, "input shape does not match input_size"); + TORCH_CHECK(input_size > 0 && input_size % kQuantBlock == 0, "input_size must be divisible by 32"); + TORCH_CHECK(records_per_input > 0, "records_per_input must be positive"); + TORCH_CHECK( + slot_ids.numel() == input.size(0) * records_per_input, + "slot count does not match input rows and records_per_input"); + const int64_t expected_role_bytes = output_size * (input_size / kQuantBlock) * kBlockBytes; + TORCH_CHECK(role_bytes == expected_role_bytes, "role byte count does not match matrix dimensions"); + TORCH_CHECK( + role_offset_a >= 0 && role_offset_a + role_bytes <= cache.size(1), "gate role range is outside each cache slot"); + TORCH_CHECK( + role_offset_b >= 0 && role_offset_b + role_bytes <= cache.size(1), "up role range is outside each cache slot"); + + const auto records = slot_ids.numel(); + auto output_a = torch::empty({records, output_size}, input.options()); + auto output_b = torch::empty({records, output_size}, input.options()); + const dim3 block(kWarpsPerBlock * 32); + const dim3 grid((output_size + kWarpsPerBlock * kRowsPerWarp - 1) / (kWarpsPerBlock * kRowsPerWarp), records); + const auto stream = at::cuda::getCurrentCUDAStream(); + if (input.scalar_type() == at::kBFloat16) { + mxfp4_matvec_dual_kernel<<>>( + reinterpret_cast(input.data_ptr()), + cache.data_ptr(), + cache.stride(0), + slot_ids.data_ptr(), + role_offset_a, + role_offset_b, + input_size, + output_size, + records, + records_per_input, + reinterpret_cast<__nv_bfloat16*>(output_a.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(output_b.data_ptr())); + } else { + mxfp4_matvec_dual_kernel<<>>( + reinterpret_cast(input.data_ptr()), + cache.data_ptr(), + cache.stride(0), + slot_ids.data_ptr(), + role_offset_a, + role_offset_b, + input_size, + output_size, + records, + records_per_input, + reinterpret_cast(output_a.data_ptr()), + reinterpret_cast(output_b.data_ptr())); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return std::make_tuple(output_a, output_b); +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("mxfp4_matvec", &mxfp4_matvec, "GGUF MXFP4 matrix-vector multiply"); + module.def("mxfp4_matvec_dual", &mxfp4_matvec_dual, "GGUF MXFP4 gate/up matrix-vector multiply"); + module.def("mxfp4_marlin_repack", &mxfp4_marlin_repack, "Repack raw GGUF MXFP4 objects to Marlin layout"); +} diff --git a/python/sglang/kernels/jit/csrc/ngram_corpus/result.h b/python/sglang/kernels/jit/csrc/ngram_corpus/result.h index 8a480c8df..a04d70cc5 100644 --- a/python/sglang/kernels/jit/csrc/ngram_corpus/result.h +++ b/python/sglang/kernels/jit/csrc/ngram_corpus/result.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include diff --git a/python/sglang/kernels/ops/attention/dsv4/compress.py b/python/sglang/kernels/ops/attention/dsv4/compress.py index bb3540076..99b9597c9 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress.py @@ -409,6 +409,11 @@ def compress_forward( else: fn = module.decode if plan.is_decode else module.prefill + # C4/C128 kernels use the same InputFloat type for APE and kv_score_input. + # Keep the model parameter in FP32 but convert it to the kernel input dtype + # at the fused-kernel boundary. + if ape.dtype != kv_score_input.dtype: + ape = ape.to(dtype=kv_score_input.dtype) fn(kv_score_buffer, kv_score_input, out, ape, *plan[1:3]) return out @@ -448,6 +453,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 != kv.dtype: + norm_weight = norm_weight.to(dtype=kv.dtype) fn( kv, plan[1], diff --git a/python/sglang/kernels/ops/kimi_k3/attn_res.py b/python/sglang/kernels/ops/kimi_k3/attn_res.py index 76a83bac5..33b936975 100644 --- a/python/sglang/kernels/ops/kimi_k3/attn_res.py +++ b/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, diff --git a/python/sglang/kernels/ops/moe/expert_pack_mxfp4.py b/python/sglang/kernels/ops/moe/expert_pack_mxfp4.py new file mode 100644 index 000000000..dda8c3454 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/arg_groups/expert_pack_hook.py b/python/sglang/srt/arg_groups/expert_pack_hook.py new file mode 100644 index 000000000..b7ee90484 --- /dev/null +++ b/python/sglang/srt/arg_groups/expert_pack_hook.py @@ -0,0 +1,198 @@ +# 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.arg_groups.overrides import declare_resolution +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 + ) + declare_resolution( + server_args, + "handle_expert_pack", + 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}") + + declare_resolution( + server_args, + "handle_expert_pack", + disable_cuda_graph=True, + 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." + ) diff --git a/python/sglang/srt/configs/load_config.py b/python/sglang/srt/configs/load_config.py index 9b63ab964..5a0e06d02 100644 --- a/python/sglang/srt/configs/load_config.py +++ b/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" diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index cbd054056..32806be68 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -132,7 +132,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_dspark_draft(config) -> bool: @@ -990,6 +993,7 @@ class ModelConfig: 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 diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 054b781e0..31fec7b84 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -40,6 +40,7 @@ _is_npu = is_npu() if TYPE_CHECKING: from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.deepseek_v4_backend import DeepseekV4AttnBackend + from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.rotary_embedding import RotaryEmbedding from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -342,6 +343,7 @@ class Compressor(BaseFusedOp): head_dim: int, rotate: bool = False, prefix: str = "", + quant_config: Optional[QuantizationConfig] = None, rotary_emb: Optional[RotaryEmbedding] = None, ) -> None: super().__init__() @@ -365,7 +367,7 @@ class Compressor(BaseFusedOp): self.dim, 2 * coff * self.head_dim, bias=False, - quant_config=None, + quant_config=quant_config, prefix=add_prefix("wkv_gate", prefix), params_dtype=wkv_gate_dtype, ) @@ -425,7 +427,7 @@ class Compressor(BaseFusedOp): comm_stream = getattr(forward_batch, "_cp_prefetch_comm_stream", None) if comm_stream is None or not dsa_use_prefill_cp(forward_batch): return - kv_score = linear_bf16_fp32(x, self.wkv_gate.weight) + kv_score = self._compute_wkv_gate(x) # Keyed by forward_batch: each TBO ubatch carries its own, so the two # ubatches cannot collect each other's gather. pending = forward_batch.__dict__.setdefault("_cp_pending_gathers", {}) @@ -440,7 +442,7 @@ class Compressor(BaseFusedOp): if handle is not None: return cp_all_gather_rerange_finish(handle) - kv_score = linear_bf16_fp32(x, self.wkv_gate.weight) + kv_score = self._compute_wkv_gate(x) # CUDA path: delegate to backend if dsa_use_prefill_cp(forward_batch): @@ -451,6 +453,19 @@ class Compressor(BaseFusedOp): ) return kv_score + def _compute_wkv_gate(self, x: torch.Tensor) -> torch.Tensor: + weight = getattr(self.wkv_gate, "weight", None) + if weight is not None: + return linear_bf16_fp32(x, weight) + + from sglang.srt.layers.quantization.gguf import fused_mul_mat_gguf + + return fused_mul_mat_gguf( + x, + self.wkv_gate.qweight, + self.wkv_gate.qweight_type.weight_type, + ) + def forward_native( self, x: torch.Tensor, diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 983171efb..774b2729b 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -901,11 +901,16 @@ class C4Indexer(nn.Module): params_dtype=torch.bfloat16, prefix=add_prefix("wq_b", prefix), ) + expert_pack_quant_config = ( + quant_config + if quant_config is not None and quant_config.get_name() == "expert_pack" + else None + ) self.weights_proj = ReplicatedLinear( self.dim, self.n_heads, bias=False, - quant_config=None, + quant_config=expert_pack_quant_config, params_dtype=torch.bfloat16, prefix=add_prefix("weights_proj", prefix), ) @@ -918,6 +923,7 @@ class C4Indexer(nn.Module): head_dim=self.head_dim, rotate=True, prefix=add_prefix("compressor", prefix), + quant_config=expert_pack_quant_config, rotary_emb=rotary_emb, ) self.rotary_emb = rotary_emb diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 6cf18264b..c91458423 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -267,6 +267,14 @@ class ReplicatedLinear(LinearBase): if len(loaded_weight.shape) == 0: loaded_weight = loaded_weight.reshape(1) + is_gguf_weight = getattr(param, "is_gguf_weight", False) + is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False) + if is_gguf_weight_type: + param.weight_type = loaded_weight.item() + + if is_gguf_weight and isinstance(param, UninitializedParameter): + param.materialize(tuple(loaded_weight.shape), dtype=loaded_weight.dtype) + # The per-tensor quant-scale must be 1 dimension if _is_npu: if param.size() != loaded_weight.size() and param.size(0) == 1: diff --git a/python/sglang/srt/layers/moe/expert_pack.py b/python/sglang/srt/layers/moe/expert_pack.py new file mode 100644 index 000000000..bdbc353a2 --- /dev/null +++ b/python/sglang/srt/layers/moe/expert_pack.py @@ -0,0 +1,931 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Runtime reader and VRAM cache for SGLANG-EXPERTPACK-v1.""" + +from __future__ import annotations + +import atexit +import concurrent.futures +import hashlib +import json +import logging +import os +import re +import struct +import threading +import time +from collections import OrderedDict +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + +MAGIC = b"SGLANG-EXPERTPACK-v1\0\0\0\0" +ROLE_NAMES = ("gate", "up", "down") +HEADER_STRUCT = struct.Struct("<24sIIIIQQQIIII32s32s32s") +ENTRY_STRUCT = struct.Struct("\d+)\.ffn_(?Pup|gate|down)_exps\.weight$" +) + + +def _fixed_string(value: bytes) -> str: + return value.split(b"\0", 1)[0].decode("utf-8") + + +def _sha256_file(path: Path, chunk_bytes: int = 16 * 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb", buffering=0) as stream: + while chunk := stream.read(chunk_bytes): + digest.update(chunk) + return digest.hexdigest() + + +@dataclass(frozen=True) +class ExpertPackHeader: + flags: int + index_count: int + data_start: int + alignment: int + num_layers: int + num_experts: int + top_k: int + role_count: int + model_identity_sha256: str + source_blob_sha256: str + config_sha256: str + + @classmethod + def read(cls, stream) -> ExpertPackHeader: + raw = stream.read(HEADER_STRUCT.size) + if len(raw) != HEADER_STRUCT.size: + raise ValueError("expert-pack header is truncated") + values = HEADER_STRUCT.unpack(raw) + if values[0] != MAGIC or values[1] != 1: + raise ValueError("expert-pack magic or version does not match") + if values[2] != HEADER_STRUCT.size or values[3] != ENTRY_STRUCT.size: + raise ValueError("expert-pack struct sizes do not match") + header = cls( + flags=values[4], + index_count=values[5], + data_start=values[6], + alignment=values[7], + num_layers=values[8], + num_experts=values[9], + top_k=values[10], + role_count=values[11], + model_identity_sha256=values[12].hex(), + source_blob_sha256=values[13].hex(), + config_sha256=values[14].hex(), + ) + expected = header.num_layers * header.num_experts * len(ROLE_NAMES) + if header.index_count != expected or header.role_count != len(ROLE_NAMES): + raise ValueError("expert-pack header coverage is inconsistent") + if header.flags & REQUIRED_FLAGS != REQUIRED_FLAGS: + raise ValueError("expert-pack is not identity triplet layout") + if header.alignment <= 0 or header.alignment & (header.alignment - 1): + raise ValueError("expert-pack alignment is invalid") + minimum = HEADER_STRUCT.size + header.index_count * ENTRY_STRUCT.size + if header.data_start < minimum or header.data_start % header.alignment: + raise ValueError("expert-pack data offset is invalid") + return header + + +@dataclass(frozen=True) +class ExpertPackEntry: + layer: int + expert: int + role_id: int + dtype_id: int + dtype: str + tensor_name: str + source_slice_offset: int + source_slice_nbytes: int + pack_offset: int + pack_nbytes: int + checksum: str + shape: tuple[int, ...] + quant_scheme: str + transform_id: str + block_size: int + generation: int + + @classmethod + def read(cls, stream) -> ExpertPackEntry: + raw = stream.read(ENTRY_STRUCT.size) + if len(raw) != ENTRY_STRUCT.size: + raise ValueError("expert-pack index is truncated") + values = ENTRY_STRUCT.unpack(raw) + role_id, rank = values[2], values[3] + if role_id >= len(ROLE_NAMES) or not 1 <= rank <= 4: + raise ValueError("expert-pack index role or rank is invalid") + return cls( + layer=values[0], + expert=values[1], + role_id=role_id, + dtype_id=values[4], + dtype=_fixed_string(values[5]), + tensor_name=_fixed_string(values[6]), + source_slice_offset=values[9], + source_slice_nbytes=values[10], + pack_offset=values[11], + pack_nbytes=values[12], + checksum=values[15].hex(), + shape=tuple(values[16 : 16 + rank]), + quant_scheme=_fixed_string(values[20]), + transform_id=_fixed_string(values[21]), + block_size=values[22], + generation=values[23], + ) + + +@dataclass +class _CacheSlot: + key: tuple[int, int] | None = None + generation: int = 0 + frequency: int = 0 + last_use: torch.cuda.Event | None = None + ready: torch.cuda.Event | None = None + + +def _initialize_runtime_state( + store, + *, + cache_vram_mib: int, + cache_vram_reserve_mib: int, + stage_slots: int, + read_splits: int, + direct_io: bool, + stats_flush_interval: int, + stats_path: str | os.PathLike[str] | None, +) -> None: + store.cache_vram_mib = int(cache_vram_mib) + store.cache_vram_reserve_mib = int(cache_vram_reserve_mib) + store.kernel_backend = "custom" + store.stage_slot_count = int(stage_slots) + store.read_splits = int(read_splits) + store.direct_io = bool(direct_io) + store.stats_flush_interval = int(stats_flush_interval) + if ( + store.cache_vram_mib <= 0 + or store.cache_vram_reserve_mib <= 0 + or store.stage_slot_count <= 0 + or store.read_splits <= 0 + ): + raise ValueError( + "expert cache and staging budgets, and read splits, must be positive" + ) + if store.stats_flush_interval < 0: + raise ValueError("expert-pack stats flush interval cannot be negative") + if store.direct_io and not hasattr(os, "O_DIRECT"): + raise ValueError("expert-pack direct I/O is unavailable on this platform") + open_flags = os.O_RDONLY | (os.O_DIRECT if store.direct_io else 0) + store._fd = os.open(store.path, open_flags) + store._lock = threading.RLock() + store._cache = None + store._cache_slots = [] + store._key_to_slot = {} + store._key_frequency = {} + store._lru = OrderedDict() + store._staging = [] + store._stage_events = [] + store._stage_cursor = 0 + store._transfer_stream = None + store._read_executor = None + store._active_keys = set() + store._route_calls_by_layer = [0] * store.header.num_layers + store._route_tokens_by_layer = [0] * store.header.num_layers + store.stats_path = Path(stats_path).resolve() if stats_path else None + store._last_stats_flush_calls = 0 + store.stats = { + "pack_path": str(store.path), + "pack_entries": len(store.entries), + "pack_reads": 0, + "pack_read_bytes": 0, + "pack_read_ns": 0, + "read_splits": store.read_splits, + "direct_io": store.direct_io, + "cache_hits": 0, + "cache_misses": 0, + "cache_evictions": 0, + "cache_policy": "reuse-lfu-lru-v2", + "kernel_backend": "custom", + "resident_experts": 0, + "resident_bytes": 0, + "h2d_bytes": 0, + "cache_vram_reserve_mib": store.cache_vram_reserve_mib, + "fallback_count": 0, + "io_errors": 0, + } + atexit.register(store.close) + + +class ExpertPackStore: + """Validated pack index plus generation-aware GPU and host caches.""" + + def __init__( + self, + pack_path: str | os.PathLike[str], + *, + manifest_path: str | os.PathLike[str] | None = None, + expected_layers: int, + expected_experts: int, + expected_top_k: int, + expected_source_sha256: str | None = None, + expected_model_identity_sha256: str | None = None, + expected_config_sha256: str | None = None, + cache_vram_mib: int = 20 * 1024, + cache_vram_reserve_mib: int = 3 * 1024, + stage_slots: int = 8, + read_splits: int = READ_SPLITS, + direct_io: bool = False, + stats_flush_interval: int = 0, + verify_pack_sha256: bool = False, + stats_path: str | os.PathLike[str] | None = None, + ) -> None: + self.path = Path(pack_path).resolve() + self.manifest_path = Path( + manifest_path or str(self.path) + ".manifest.json" + ).resolve() + if not self.path.is_file() or not self.manifest_path.is_file(): + raise FileNotFoundError( + f"expert-pack or manifest is missing: {self.path}, {self.manifest_path}" + ) + self.manifest = json.loads(self.manifest_path.read_text(encoding="utf-8")) + if not self.manifest.get("complete"): + raise ValueError("expert-pack manifest is not complete") + + with self.path.open("rb", buffering=0) as stream: + self.header = ExpertPackHeader.read(stream) + entries = [ + ExpertPackEntry.read(stream) for _ in range(self.header.index_count) + ] + + expected_dimensions = (expected_layers, expected_experts, expected_top_k) + actual_dimensions = ( + self.header.num_layers, + self.header.num_experts, + self.header.top_k, + ) + if actual_dimensions != expected_dimensions: + raise ValueError( + f"expert-pack dimensions {actual_dimensions} != {expected_dimensions}" + ) + expected_digests = { + "source_blob_sha256": expected_source_sha256, + "model_identity_sha256": expected_model_identity_sha256, + "config_sha256": expected_config_sha256, + } + for field, expected in expected_digests.items(): + if expected and getattr(self.header, field) != expected: + raise ValueError( + f"expert-pack {field} does not match configured digest" + ) + + self.entries: dict[tuple[int, int, int], ExpertPackEntry] = {} + role_bytes: int | None = None + object_generations: dict[tuple[int, int], int] = {} + for entry in entries: + key = (entry.layer, entry.expert, entry.role_id) + if key in self.entries: + raise ValueError(f"duplicate expert-pack entry {key}") + if entry.dtype != "MXFP4" or entry.quant_scheme != "MXFP4": + raise ValueError(f"unsupported expert dtype for {key}: {entry.dtype}") + if entry.transform_id != "identity-v1" or entry.block_size != 32: + raise ValueError(f"unsupported expert transform for {key}") + if entry.pack_nbytes != entry.source_slice_nbytes: + raise ValueError(f"non-identity expert payload size for {key}") + if role_bytes is None: + role_bytes = entry.pack_nbytes + elif role_bytes != entry.pack_nbytes: + raise ValueError("expert-pack roles are not fixed size") + object_key = key[:2] + generation = object_generations.setdefault(object_key, entry.generation) + if generation != entry.generation: + raise ValueError(f"mixed generation in expert object {object_key}") + self.entries[key] = entry + + assert role_bytes is not None + self.role_bytes = role_bytes + self.object_payload_bytes = role_bytes * len(ROLE_NAMES) + self.object_stride = int(self.manifest["object_stride"]) + if self.object_stride < self.object_payload_bytes: + raise ValueError("expert-pack object stride is smaller than its payload") + expected_size = self.header.data_start + ( + expected_layers * expected_experts * self.object_stride + ) + if self.path.stat().st_size != expected_size: + raise ValueError("expert-pack file size does not match its index") + self.object_offsets: dict[tuple[int, int], int] = {} + self.active_moe_layer_ids = frozenset(range(expected_layers)) + for layer in range(expected_layers): + for expert in range(expected_experts): + object_offset = ( + self.header.data_start + + (layer * expected_experts + expert) * self.object_stride + ) + self.object_offsets[(layer, expert)] = object_offset + for role_id in range(len(ROLE_NAMES)): + entry = self.entries[(layer, expert, role_id)] + if entry.pack_offset != object_offset + role_id * role_bytes: + raise ValueError( + f"expert-pack object layout mismatch at {(layer, expert, role_id)}" + ) + + manifest_pack_sha = self.manifest.get("pack_sha256") + if verify_pack_sha256: + actual_pack_sha = _sha256_file(self.path) + if actual_pack_sha != manifest_pack_sha: + raise ValueError("expert-pack SHA-256 does not match its manifest") + self.pack_sha256 = str(manifest_pack_sha) + self.role_offsets = { + role: role_id * self.role_bytes for role_id, role in enumerate(ROLE_NAMES) + } + self.role_nbytes = {role: self.role_bytes for role in ROLE_NAMES} + _initialize_runtime_state( + self, + cache_vram_mib=cache_vram_mib, + cache_vram_reserve_mib=cache_vram_reserve_mib, + stage_slots=stage_slots, + read_splits=read_splits, + direct_io=direct_io, + stats_flush_interval=stats_flush_interval, + stats_path=stats_path, + ) + + def initialize_device_cache(self, device: torch.device | str) -> None: + if self._cache is not None: + return + device = torch.device(device) + if device.type != "cuda": + raise ValueError("expert-pack runtime currently requires CUDA") + requested = self.cache_vram_mib * 1024 * 1024 + free_bytes, _ = torch.cuda.mem_get_info(device) + reserve = self.cache_vram_reserve_mib * 1024 * 1024 + budget = min(requested, max(0, free_bytes - reserve)) + slot_count = max(0, budget // self.object_payload_bytes) + if slot_count < self.header.top_k: + raise MemoryError( + "insufficient free VRAM for one top-k expert working set: " + f"free={free_bytes}, object={self.object_payload_bytes}" + ) + cache_bytes = slot_count * self.object_payload_bytes + self._cache = torch.empty( + (slot_count, self.object_payload_bytes), + dtype=torch.uint8, + device=device, + ) + self._cache_slots = [_CacheSlot() for _ in range(slot_count)] + self._staging = [ + torch.empty(self.object_payload_bytes, dtype=torch.uint8, pin_memory=True) + for _ in range(self.stage_slot_count) + ] + if self.direct_io: + alignment = 4096 + ranges = self._object_read_ranges() + if any(offset % alignment for offset in self.object_offsets.values()): + raise ValueError("expert-pack direct I/O requires aligned objects") + if any(start % alignment or length % alignment for start, length in ranges): + raise ValueError("expert-pack direct I/O requires aligned read ranges") + if any(staging.data_ptr() % alignment for staging in self._staging): + raise ValueError("expert-pack direct I/O requires aligned staging") + self._stage_events = [None] * self.stage_slot_count + self._transfer_stream = torch.cuda.Stream(device=device) + self._read_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=self.stage_slot_count * self.read_splits, + thread_name_prefix="expert-pack-read", + ) + self.stats["cache_capacity_experts"] = slot_count + self.stats["cache_capacity_bytes"] = cache_bytes + staged_bytes = self.stage_slot_count * self.object_payload_bytes + self.stats["staged_bytes"] = staged_bytes + logger.info( + "Expert pack ready: entries=%d resident_experts=0 dense_bytes=external " + "staged_bytes=%d cache_capacity_experts=%d cache_capacity_bytes=%d", + len(self.entries), + staged_bytes, + slot_count, + cache_bytes, + ) + + def _read_object( + self, layer: int, expert: int, staging: torch.Tensor + ) -> tuple[int, int]: + return self._read_object_range( + layer, expert, staging, start=0, length=self.object_payload_bytes + ) + + def _read_object_range( + self, + layer: int, + expert: int, + staging: torch.Tensor, + *, + start: int, + length: int, + ) -> tuple[int, int]: + try: + object_offset = self.object_offsets[(layer, expert)] + except KeyError as exc: + raise ValueError(f"no expert-pack object for {(layer, expert)}") from exc + offset = object_offset + start + view = memoryview(staging.numpy()).cast("B")[start : start + length] + started = time.perf_counter_ns() + read_bytes = os.preadv(self._fd, [view], offset) + elapsed = time.perf_counter_ns() - started + if read_bytes != length: + raise OSError( + f"short expert-pack read for {(layer, expert, start, length)}: " + f"{read_bytes} != {length}" + ) + return read_bytes, elapsed + + def _object_read_ranges(self) -> list[tuple[int, int]]: + split_count = self.read_splits + alignment = 4096 if self.object_payload_bytes >= split_count * 4096 else 1 + boundaries = [ + self.object_payload_bytes * part // split_count // alignment * alignment + for part in range(split_count) + ] + [self.object_payload_bytes] + return [ + (boundaries[part], boundaries[part + 1] - boundaries[part]) + for part in range(split_count) + ] + + def _victim_slot( + self, + protected: set[tuple[int, int]], + *, + preserve_oldest: bool = False, + ) -> int: + for index, slot in enumerate(self._cache_slots): + if slot.key is None: + return index + keys = reversed(self._lru) if preserve_oldest else iter(self._lru) + victim_index = None + victim_frequency = None + for key in keys: + if key in protected: + continue + slot_index = self._key_to_slot[key] + frequency = self._cache_slots[slot_index].frequency + if victim_frequency is None or frequency < victim_frequency: + victim_index = slot_index + victim_frequency = frequency + if victim_index is not None: + return victim_index + raise RuntimeError("expert cache cannot evict the active top-k working set") + + def _install_staging( + self, + staging: torch.Tensor, + slot_index: int, + stream: torch.cuda.Stream, + ) -> torch.cuda.Event: + """Publish one host object into the custom GPU cache.""" + with torch.cuda.stream(stream): + assert self._cache is not None + self._cache[slot_index].copy_(staging, non_blocking=True) + ready = torch.cuda.Event() + ready.record(stream) + return ready + + def _record_read(self, read_bytes: int, elapsed: int) -> None: + self.stats["pack_reads"] = int(self.stats["pack_reads"]) + 1 + self.stats["pack_read_bytes"] = int(self.stats["pack_read_bytes"]) + read_bytes + self.stats["pack_read_ns"] = int(self.stats["pack_read_ns"]) + elapsed + + def acquire( + self, layer: int, topk_ids: torch.Tensor, *, is_prefill: bool | None = None + ) -> tuple[torch.Tensor, list[int]]: + if ( + self._cache is None + or self._transfer_stream is None + or self._read_executor is None + ): + raise RuntimeError("expert device cache is not initialized") + if layer not in self.active_moe_layer_ids: + raise ValueError(f"layer {layer} is not an active routed MoE layer") + if topk_ids.ndim != 2 or topk_ids.shape[-1] != self.header.top_k: + raise ValueError( + f"runtime top-k must be exactly {self.header.top_k}; " + f"received shape {tuple(topk_ids.shape)}" + ) + route_ids = [int(value) for value in topk_ids.detach().cpu().reshape(-1)] + if is_prefill is None: + is_prefill = topk_ids.shape[0] > 1 + if any(expert < 0 or expert >= self.header.num_experts for expert in route_ids): + raise ValueError("route contains an out-of-range expert id") + requested = {(layer, expert) for expert in route_ids} + events: list[torch.cuda.Event] = [] + + with self._lock: + for key in requested: + self._key_frequency[key] = min(self._key_frequency.get(key, 0) + 1, 255) + self._active_keys.update(requested) + self._route_calls_by_layer[layer] += 1 + self._route_tokens_by_layer[layer] += int(topk_ids.shape[0]) + pending: list[tuple[tuple[int, int], int]] = [] + requested_keys = sorted( + dict.fromkeys((layer, expert) for expert in route_ids) + ) + for key in requested_keys: + slot_index = self._key_to_slot.get(key) + generation = self.entries[(key[0], key[1], 0)].generation + if ( + slot_index is not None + and self._cache_slots[slot_index].generation == generation + ): + self.stats["cache_hits"] = int(self.stats["cache_hits"]) + 1 + self._lru.move_to_end(key) + slot = self._cache_slots[slot_index] + slot.frequency = self._key_frequency[key] + if slot.ready is not None: + events.append(slot.ready) + continue + + self.stats["cache_misses"] = int(self.stats["cache_misses"]) + 1 + if slot_index is None: + slot_index = self._victim_slot( + requested, preserve_oldest=is_prefill + ) + slot = self._cache_slots[slot_index] + if slot.key is not None: + self.stats["cache_evictions"] = ( + int(self.stats["cache_evictions"]) + 1 + ) + self._key_to_slot.pop(slot.key, None) + self._lru.pop(slot.key, None) + slot.key = key + slot.generation = generation + slot.frequency = self._key_frequency[key] + self._key_to_slot[key] = slot_index + self._lru[key] = None + pending.append((key, slot_index)) + + for batch_start in range(0, len(pending), len(self._staging)): + batch = pending[batch_start : batch_start + len(self._staging)] + jobs = [] + for key, slot_index in batch: + stage_index = self._stage_cursor + self._stage_cursor = (self._stage_cursor + 1) % len(self._staging) + stage_event = self._stage_events[stage_index] + if stage_event is not None: + stage_event.synchronize() + staging = self._staging[stage_index] + futures = tuple( + self._read_executor.submit( + self._read_object_range, + key[0], + key[1], + staging, + start=start, + length=length, + ) + for start, length in self._object_read_ranges() + ) + jobs.append((futures, key, stage_index, staging, slot_index)) + + for futures, key, stage_index, staging, slot_index in jobs: + if futures: + try: + results = [future.result() for future in futures] + except OSError: + self.stats["io_errors"] = int(self.stats["io_errors"]) + 1 + raise + read_bytes = sum(result[0] for result in results) + elapsed = max(result[1] for result in results) + self._record_read(read_bytes, elapsed) + slot = self._cache_slots[slot_index] + if slot.ready is not None: + self._transfer_stream.wait_event(slot.ready) + if slot.last_use is not None: + self._transfer_stream.wait_event(slot.last_use) + ready = self._install_staging( + staging, + slot_index, + self._transfer_stream, + ) + self._stage_events[stage_index] = ready + events.append(ready) + slot.last_use = None + slot.ready = ready + self.stats["h2d_bytes"] = int(self.stats["h2d_bytes"]) + ( + self.object_payload_bytes + ) + + cache_device = self._cache.device + current_stream = torch.cuda.current_stream(cache_device) + for event in events: + current_stream.wait_event(event) + slots = [self._key_to_slot[(layer, expert)] for expert in route_ids] + self.stats["resident_experts"] = len(self._key_to_slot) + self.stats["resident_bytes"] = ( + len(self._key_to_slot) * self.object_payload_bytes + ) + return ( + torch.tensor(slots, dtype=torch.int32, device=cache_device), + slots, + ) + + def mark_used(self, slot_indices: list[int]) -> None: + event = torch.cuda.Event() + event.record(torch.cuda.current_stream()) + with self._lock: + for slot_index in set(slot_indices): + slot = self._cache_slots[slot_index] + slot.last_use = event + slot.ready = None + if slot.key is not None: + self._active_keys.discard(slot.key) + route_calls = sum(self._route_calls_by_layer) + if ( + self.stats_path is not None + and self.stats_flush_interval + and route_calls - self._last_stats_flush_calls + >= self.stats_flush_interval + ): + self._write_stats() + self._last_stats_flush_calls = route_calls + + @property + def device_cache(self) -> torch.Tensor: + if self._cache is None: + raise RuntimeError("raw expert device cache is not initialized") + return self._cache + + def snapshot(self) -> dict[str, Any]: + value = dict(self.stats) + reads = int(value["pack_reads"]) + value["mean_read_ms"] = ( + int(value["pack_read_ns"]) / reads / 1e6 if reads else 0.0 + ) + value["route_calls_by_layer"] = list(self._route_calls_by_layer) + value["route_tokens_by_layer"] = list(self._route_tokens_by_layer) + return value + + def _write_stats(self) -> None: + assert self.stats_path is not None + self.stats_path.parent.mkdir(parents=True, exist_ok=True) + temporary = self.stats_path.with_name( + self.stats_path.name + f".{os.getpid()}.tmp" + ) + temporary.write_text( + json.dumps(self.snapshot(), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(self.stats_path) + + def close(self) -> None: + if self._read_executor is not None: + self._read_executor.shutdown(wait=True) + self._read_executor = None + if self.stats_path is not None: + self._write_stats() + if getattr(self, "_fd", -1) >= 0: + os.close(self._fd) + self._fd = -1 + + +class KimiGGMLExpertPackStore(ExpertPackStore): + """Runtime cache for the audited, zero-copy Kimi GGMLMOEPACKv1 layout.""" + + def __init__( + self, + pack_path: str | os.PathLike[str], + *, + manifest_path: str | os.PathLike[str], + expected_layers: int, + expected_experts: int, + expected_top_k: int, + cache_vram_mib: int = 18 * 1024, + cache_vram_reserve_mib: int = 3 * 1024, + stage_slots: int = 16, + read_splits: int = READ_SPLITS, + direct_io: bool = False, + stats_flush_interval: int = 0, + verify_pack_sha256: bool = False, + stats_path: str | os.PathLike[str] | None = None, + ) -> None: + self.path = Path(pack_path).resolve() + self.manifest_path = Path(manifest_path).resolve() + if not self.path.is_file() or not self.manifest_path.is_file(): + raise FileNotFoundError( + f"Kimi expert-pack or manifest is missing: " + f"{self.path}, {self.manifest_path}" + ) + self.manifest = json.loads(self.manifest_path.read_text(encoding="utf-8")) + if ( + not self.manifest.get("complete") + or self.manifest.get("format") != KIMI_FORMAT + ): + raise ValueError("Kimi expert-pack manifest is incomplete or unsupported") + + constraints = self.manifest.get("hard_constraints", {}) + if constraints != { + "all_selected_experts_must_execute": True, + "expert_pruning_allowed": False, + "requantization_allowed": False, + "top_k": 16, + "top_k_is_immutable": True, + }: + raise ValueError( + "Kimi manifest hard constraints do not match runtime policy" + ) + model = self.manifest["model"] + dimensions = ( + int(model["num_hidden_layers"]), + int(model["num_experts"]), + int(model["num_experts_per_token"]), + ) + expected_dimensions = (expected_layers, expected_experts, expected_top_k) + if dimensions != expected_dimensions or expected_top_k != 16: + raise ValueError( + f"Kimi expert-pack dimensions {dimensions} != {expected_dimensions}; " + "Top-K is immutable at 16" + ) + active_layers = tuple(int(value) for value in model["active_moe_layer_ids"]) + if active_layers != tuple(range(1, 93)): + raise ValueError("Kimi active routed MoE layers must be exactly 1..92") + + pack_manifest = self.manifest["expert_pack"] + if Path(pack_manifest["path"]).resolve() != self.path: + raise ValueError("Kimi manifest expert-pack path does not match pack_path") + if int(pack_manifest["size"]) != self.path.stat().st_size: + raise ValueError("Kimi expert-pack size does not match its manifest") + if pack_manifest.get("physical_role_order") != list(KIMI_PHYSICAL_ROLES): + raise ValueError("Kimi expert-pack physical role order is unsupported") + roles = pack_manifest["roles"] + expected_roles = { + "up": ("Q2_K", 10), + "gate": ("Q2_K", 10), + "down": ("Q3_K", 11), + } + for role, (dtype, dtype_id) in expected_roles.items(): + if ( + roles[role]["dtype"] != dtype + or int(roles[role]["dtype_id"]) != dtype_id + ): + raise ValueError(f"Kimi expert-pack {role} quant type is unsupported") + + expected_entry_count = len(active_layers) * expected_experts * len(ROLE_NAMES) + index_digest = hashlib.sha256() + self.entries: dict[tuple[int, int, int], ExpertPackEntry] = {} + self.object_offsets: dict[tuple[int, int], int] = {} + object_payload_bytes = int(pack_manifest["object_bytes"]) + previous_end = int(pack_manifest["data_start"]) + role_offsets: dict[str, int] = {} + role_nbytes = { + role: int(roles[role]["expert_bytes"]) for role in KIMI_PHYSICAL_ROLES + } + running_role_offset = 0 + for role in KIMI_PHYSICAL_ROLES: + role_offsets[role] = running_role_offset + running_role_offset += role_nbytes[role] + if running_role_offset != object_payload_bytes: + raise ValueError("Kimi expert-pack role sizes do not match object bytes") + + with self.path.open("rb", buffering=0) as stream: + raw_header = stream.read(GGML_PACK_HEADER.size) + if len(raw_header) != GGML_PACK_HEADER.size: + raise ValueError("Kimi expert-pack header is truncated") + index_digest.update(raw_header) + magic, version, header_size, index_count, data_start = ( + GGML_PACK_HEADER.unpack(raw_header) + ) + if ( + magic != GGML_PACK_MAGIC + or version != 1 + or header_size != GGML_PACK_HEADER.size + or index_count != expected_entry_count + or data_start != int(pack_manifest["data_start"]) + ): + raise ValueError("Kimi expert-pack header does not match its manifest") + + for index in range(index_count): + raw_entry = stream.read(GGML_PACK_ENTRY.size) + if len(raw_entry) != GGML_PACK_ENTRY.size: + raise ValueError("Kimi expert-pack index is truncated") + index_digest.update(raw_entry) + name_raw, expert, reserved, offset, nbytes = GGML_PACK_ENTRY.unpack( + raw_entry + ) + object_index, physical_role_id = divmod(index, len(KIMI_PHYSICAL_ROLES)) + layer_index, expected_expert = divmod(object_index, expected_experts) + layer = active_layers[layer_index] + physical_role = KIMI_PHYSICAL_ROLES[physical_role_id] + name = _fixed_string(name_raw) + match = KIMI_EXPERT_RE.fullmatch(name) + if ( + match is None + or int(match.group("layer")) != layer + or match.group("role") != physical_role + or expert != expected_expert + or reserved != 0 + ): + raise ValueError( + f"Kimi expert-pack identity mismatch at index {index}" + ) + expected_nbytes = role_nbytes[physical_role] + if ( + nbytes != expected_nbytes + or offset % int(pack_manifest["alignment"]) + or offset < previous_end + or offset + nbytes > self.path.stat().st_size + ): + raise ValueError( + f"Kimi expert-pack range mismatch at index {index}" + ) + previous_end = offset + nbytes + object_key = (layer, expert) + if physical_role_id == 0: + self.object_offsets[object_key] = offset + expected_offset = ( + self.object_offsets[object_key] + role_offsets[physical_role] + ) + if offset != expected_offset: + raise ValueError( + f"Kimi expert object is not contiguous at index {index}" + ) + generation_bytes = hashlib.sha256( + f"{pack_manifest['index_sha256']}:{layer}:{expert}".encode("ascii") + ).digest() + generation = int.from_bytes(generation_bytes[:8], "little") or 1 + logical_role_id = ROLE_NAMES.index(physical_role) + logical_shape = tuple( + int(value) for value in roles[physical_role]["logical_shape"] + ) + self.entries[(layer, expert, logical_role_id)] = ExpertPackEntry( + layer=layer, + expert=expert, + role_id=logical_role_id, + dtype_id=int(roles[physical_role]["dtype_id"]), + dtype=str(roles[physical_role]["dtype"]), + tensor_name=name, + source_slice_offset=0, + source_slice_nbytes=nbytes, + pack_offset=offset, + pack_nbytes=nbytes, + checksum="", + shape=logical_shape, + quant_scheme=str(roles[physical_role]["dtype"]), + transform_id="identity-v1", + block_size=256, + generation=generation, + ) + + if previous_end != self.path.stat().st_size: + raise ValueError("Kimi expert-pack file has trailing or missing bytes") + actual_index_sha256 = index_digest.hexdigest() + if actual_index_sha256 != pack_manifest["index_sha256"]: + raise ValueError("Kimi expert-pack index SHA-256 does not match manifest") + if verify_pack_sha256: + expected_sha256 = pack_manifest.get("sha256") + if not expected_sha256: + raise ValueError( + "full pack verification requested, but manifest has no full SHA-256" + ) + if _sha256_file(self.path) != expected_sha256: + raise ValueError("Kimi expert-pack SHA-256 does not match manifest") + + self.header = ExpertPackHeader( + flags=REQUIRED_FLAGS, + index_count=expected_entry_count, + data_start=int(pack_manifest["data_start"]), + alignment=int(pack_manifest["alignment"]), + num_layers=expected_layers, + num_experts=expected_experts, + top_k=expected_top_k, + role_count=len(ROLE_NAMES), + model_identity_sha256="0" * 64, + source_blob_sha256=str(self.manifest["source"]["inventory_sha256"]), + config_sha256=str(model["config_sha256"]), + ) + self.active_moe_layer_ids = frozenset(active_layers) + self.role_offsets = role_offsets + self.role_nbytes = role_nbytes + self.role_bytes = role_nbytes["gate"] + self.object_payload_bytes = object_payload_bytes + self.object_stride = object_payload_bytes + self.pack_sha256 = str(pack_manifest.get("sha256") or actual_index_sha256) + _initialize_runtime_state( + self, + cache_vram_mib=cache_vram_mib, + cache_vram_reserve_mib=cache_vram_reserve_mib, + stage_slots=stage_slots, + read_splits=read_splits, + direct_io=direct_io, + stats_flush_interval=stats_flush_interval, + stats_path=stats_path, + ) + self.stats["pack_format"] = KIMI_FORMAT + self.stats["active_moe_layers"] = len(active_layers) diff --git a/python/sglang/srt/layers/quantization/expert_pack.py b/python/sglang/srt/layers/quantization/expert_pack.py new file mode 100644 index 000000000..2bc8a4408 --- /dev/null +++ b/python/sglang/srt/layers/quantization/expert_pack.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GGUF dense weights plus streamed GGUF-MXFP4 routed experts.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from sgl_kernel.quantization import ggml_moe_a8_vec + +from sglang.kernels.ops.moe.expert_pack_mxfp4 import ( + mxfp4_matvec, + mxfp4_matvec_dual, +) +from sglang.srt.layers.linear import LinearBase +from sglang.srt.layers.moe.expert_pack import ( + ExpertPackStore, + KimiGGMLExpertPackStore, +) +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.gguf import ( + GGUFConfig, + GGUFEmbeddingMethod, + GGUFLinearMethod, +) + + +def _clamped_swiglu( + gate: torch.Tensor, up: torch.Tensor, limit: float | None +) -> torch.Tensor: + if limit is not None: + if gate.is_cuda: + from sglang.kernels.ops.attention.dsv4 import silu_and_mul_clamp + + gate_up = torch.cat((gate, up), dim=-1) + output = torch.empty_like(gate) + silu_and_mul_clamp(gate_up, output, float(limit)) + return output + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + return F.silu(gate) * up + + +class ExpertPackConfig(GGUFConfig): + """Use regular GGUF methods except for routed FusedMoE layers.""" + + is_fp4_experts = True + supports_kimi_k3_quantized_latent_projections = True + + def __init__(self, store: ExpertPackStore) -> None: + super().__init__() + self.store = store + + def get_name(self) -> str: + return "expert_pack" + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding + + if isinstance(layer, FusedMoE): + return ExpertPackMoEMethod(self.store, prefix) + if isinstance(layer, LinearBase): + return GGUFLinearMethod(self) + if isinstance(layer, VocabParallelEmbedding): + return GGUFEmbeddingMethod(self) + return None + + +class ExpertPackMoEMethod(FusedMoEMethodBase): + def __init__(self, store: ExpertPackStore, prefix: str) -> None: + self.store = store + self.prefix = prefix + self.layer_id: int | None = None + self.hidden_size: int | None = None + self.intermediate_size: int | None = None + self.activation = "silu" + self.swiglu_limit: float | None = None + self.situ_beta: float | None = None + self.situ_linear_beta: float | None = None + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del extra_weight_attrs + if layer.num_fused_shared_experts: + raise ValueError( + "expert-pack requires --disable-shared-experts-fusion so the " + "shared expert remains on the dense GGUF path" + ) + if layer.moe_ep_size != 1 or layer.moe_tp_size != 1: + raise ValueError("expert-pack v1 supports only single-GPU TP=EP=1") + if num_experts != self.store.header.num_experts: + raise ValueError("FusedMoE expert count does not match expert-pack") + if params_dtype not in (torch.bfloat16, torch.float16): + raise ValueError("expert-pack kernel requires BF16 or FP16 activations") + gate_shape = self.store.entries[(layer.layer_id, 0, 0)].shape + down_shape = self.store.entries[(layer.layer_id, 0, 2)].shape + if gate_shape != (hidden_size, intermediate_size_per_partition): + raise ValueError( + f"expert-pack gate shape {gate_shape} does not match " + f"{(hidden_size, intermediate_size_per_partition)}" + ) + if down_shape != (intermediate_size_per_partition, hidden_size): + raise ValueError( + f"expert-pack down shape {down_shape} does not match " + f"{(intermediate_size_per_partition, hidden_size)}" + ) + self.layer_id = layer.layer_id + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size_per_partition + + # An empty, non-persistent marker makes the absence of eager expert + # parameters visible in module/state audits without reserving VRAM. + layer.register_buffer( + "expert_pack_marker", + torch.empty(0, dtype=torch.uint8), + persistent=False, + ) + + def create_moe_runner(self, layer, moe_runner_config) -> None: + del layer + if isinstance(self.store, KimiGGMLExpertPackStore): + if moe_runner_config.activation != "situ": + raise ValueError("Kimi-K3 expert-pack requires SiTU experts") + if ( + float(moe_runner_config.gemm1_alpha or 0.0), + float(moe_runner_config.gemm1_clamp_limit or 0.0), + ) != (4.0, 25.0): + raise ValueError("Kimi-K3 SiTU constants must be exactly 4.0 and 25.0") + self.situ_beta = 4.0 + self.situ_linear_beta = 25.0 + elif moe_runner_config.activation != "silu": + raise ValueError("DeepSeek expert-pack requires SiLU experts") + self.activation = moe_runner_config.activation + self.swiglu_limit = moe_runner_config.swiglu_limit + + @staticmethod + def _kimi_vec( + inputs: torch.Tensor, + weights: torch.Tensor, + expert_ids: torch.Tensor, + *, + top_k: int, + weight_type: int, + output_size: int, + ) -> torch.Tensor: + return ggml_moe_a8_vec( + inputs, + weights, + expert_ids, + top_k, + weight_type, + output_size, + inputs.shape[0], + ) + + def _apply_kimi(self, hidden_states, topk_ids, topk_weights, slots): + if self.intermediate_size is None or self.hidden_size is None: + raise RuntimeError("Kimi-K3 expert-pack dimensions are unavailable") + if self.situ_beta != 4.0 or self.situ_linear_beta != 25.0: + raise RuntimeError("Kimi-K3 SiTU constants were not initialized") + + cache = self.store.device_cache + top_k = topk_ids.shape[-1] + role_types = {"gate": 10, "up": 10, "down": 11} + row_bytes = { + "gate": self.store.role_nbytes["gate"] // self.intermediate_size, + "up": self.store.role_nbytes["up"] // self.intermediate_size, + "down": self.store.role_nbytes["down"] // self.hidden_size, + } + + gate_start = self.store.role_offsets["gate"] + up_start = self.store.role_offsets["up"] + down_start = self.store.role_offsets["down"] + if hidden_states.shape[0] != 1: + raise ValueError("Kimi expert-pack compact kernel expects one token") + slot_indices = slots.long() + compact_ids = torch.arange( + top_k, dtype=torch.int32, device=hidden_states.device + ).view(1, top_k) + gate_weights = torch.index_select( + cache[:, gate_start : gate_start + self.store.role_nbytes["gate"]], + 0, + slot_indices, + ).view(top_k, self.intermediate_size, row_bytes["gate"]) + up_weights = torch.index_select( + cache[:, up_start : up_start + self.store.role_nbytes["up"]], + 0, + slot_indices, + ).view(top_k, self.intermediate_size, row_bytes["up"]) + gate = self._kimi_vec( + hidden_states, + gate_weights, + compact_ids, + top_k=top_k, + weight_type=role_types["gate"], + output_size=self.intermediate_size, + ) + up = self._kimi_vec( + hidden_states, + up_weights, + compact_ids, + top_k=top_k, + weight_type=role_types["up"], + output_size=self.intermediate_size, + ) + gate_fp32 = gate.float() + gate = self.situ_beta * torch.tanh(gate_fp32 / self.situ_beta) + gate = gate * torch.sigmoid(gate_fp32) + up = self.situ_linear_beta * torch.tanh(up.float() / self.situ_linear_beta) + activated = (gate * up).to(hidden_states.dtype) + down_weights = torch.index_select( + cache[:, down_start : down_start + self.store.role_nbytes["down"]], + 0, + slot_indices, + ).view(top_k, self.hidden_size, row_bytes["down"]) + down = self._kimi_vec( + activated, + down_weights, + compact_ids.reshape(-1, 1), + top_k=1, + weight_type=role_types["down"], + output_size=self.hidden_size, + ) + down = down.view(hidden_states.shape[0], top_k, self.hidden_size) + output = torch.zeros( + (hidden_states.shape[0], self.hidden_size), + dtype=torch.float32, + device=hidden_states.device, + ) + for route_index in range(top_k): + output.add_( + down[:, route_index].float() + * topk_weights[:, route_index].float().unsqueeze(-1) + ) + return output.to(hidden_states.dtype) + + def apply(self, layer, dispatch_output): + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + if self.layer_id is None or self.hidden_size is None: + raise RuntimeError("expert-pack MoE method was not initialized") + hidden_states = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + if topk_ids.shape[-1] != self.store.header.top_k: + raise ValueError( + f"runtime top-k {topk_ids.shape[-1]} does not match expert-pack " + f"top-k {self.store.header.top_k}" + ) + if hidden_states.shape[0] == 0: + return StandardCombineInput(hidden_states=torch.empty_like(hidden_states)) + + if isinstance(self.store, KimiGGMLExpertPackStore): + token_outputs = [] + for token_index in range(topk_ids.shape[0]): + slots, host_slots = self.store.acquire( + self.layer_id, + topk_ids[token_index : token_index + 1], + is_prefill=topk_ids.shape[0] > 1, + ) + try: + token_outputs.append( + self._apply_kimi( + hidden_states[token_index : token_index + 1], + topk_ids[token_index : token_index + 1], + topk_weights[token_index : token_index + 1], + slots, + ) + ) + finally: + self.store.mark_used(host_slots) + return StandardCombineInput(hidden_states=torch.cat(token_outputs, dim=0)) + slots, host_slots = self.store.acquire( + self.layer_id, + topk_ids, + is_prefill=topk_ids.shape[0] > 1, + ) + cache = self.store.device_cache + records_per_input = topk_ids.shape[-1] + gate, up = mxfp4_matvec_dual( + hidden_states, + cache, + slots, + gate_role_offset=0, + up_role_offset=self.store.role_bytes, + role_bytes=self.store.role_bytes, + input_size=self.hidden_size, + output_size=self.intermediate_size, + records_per_input=records_per_input, + ) + intermediate = _clamped_swiglu(gate, up, self.swiglu_limit) + down = mxfp4_matvec( + intermediate, + cache, + slots, + role_offset=2 * self.store.role_bytes, + role_bytes=self.store.role_bytes, + input_size=self.intermediate_size, + output_size=self.hidden_size, + records_per_input=1, + ) + output = ( + down.view(hidden_states.shape[0], records_per_input, self.hidden_size) + * topk_weights.unsqueeze(-1).to(down.dtype) + ).sum(dim=1) + self.store.mark_used(host_slots) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py index 98784000d..cc4ed6e1e 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -71,6 +71,17 @@ else: logger = logging.getLogger(__name__) +def _ordered_gguf_shard_ids(shard_ids: list) -> list: + """Return checkpoint shards in the fused layer's logical output order.""" + if len(shard_ids) == 3 and set(shard_ids) == {"q", "k", "v"}: + return ["q", "k", "v"] + if all(isinstance(shard_id, int) for shard_id in shard_ids) and set( + shard_ids + ) == set(range(len(shard_ids))): + return sorted(shard_ids) + return list(shard_ids) + + class GGUFConfig(QuantizationConfig): """Config class for GGUF.""" @@ -424,16 +435,20 @@ class GGUFLinearMethod(LinearMethodBase): ) # (dim0_start, dim0_end, dim1_size) shard_offset_map = dict[str, tuple[int, int, int]]() - for idx in shard_id: + ordered_shard_ids = _ordered_gguf_shard_ids(shard_id) + cursor = 0 + for idx in ordered_shard_ids: id_in_container = shard_id_map[idx] - start = sum(x.size(0) for x in data_container[:id_in_container]) + start = cursor end = start + data_container[id_in_container].size(0) size = data_container[id_in_container].size(1) padded_data[start:end, :size] = data_container[id_in_container] shard_offset_map[idx] = (start, end, size) + cursor = end qweight.data_container.clear() padded_param = Parameter(padded_data, requires_grad=False) set_weight_attrs(padded_param, vars(qweight)) + padded_param.shard_id = ordered_shard_ids set_weight_attrs(padded_param, {"shard_offset_map": shard_offset_map}) layer.register_parameter("qweight", padded_param) @@ -447,7 +462,7 @@ class GGUFLinearMethod(LinearMethodBase): if shard_id: # dequantize shard weights respectively - shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id + shard_id = _ordered_gguf_shard_ids(shard_id) qweight = layer.qweight result = [] for idx in shard_id: diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index f4ec27acc..9b7c45c31 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -382,6 +382,7 @@ def maybe_fuse_routed_scale_and_shared_add( # alpha=scale)`. With no shared output, the missing scale is applied # in-place. Otherwise `routed` is already scale-final and we just add # `shared` (or pass through if there is none). + from sglang.srt.layers.quantization.expert_pack import ExpertPackMoEMethod from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import ( Mxfp4FlashinferCutlassMoEMethod, ) @@ -395,6 +396,7 @@ def maybe_fuse_routed_scale_and_shared_add( Mxfp4FlashinferTrtllmMoEMethod, Mxfp4FlashinferCutlassMoEMethod, Mxfp4MarlinMoEMethod, + ExpertPackMoEMethod, ), ) if fused: diff --git a/python/sglang/srt/lora/deepseek_mla_correction.py b/python/sglang/srt/lora/deepseek_mla_correction.py index f8250ad02..ddc1a378a 100644 --- a/python/sglang/srt/lora/deepseek_mla_correction.py +++ b/python/sglang/srt/lora/deepseek_mla_correction.py @@ -35,7 +35,8 @@ def is_kv_b_lora_active(attn_module: DeepseekV2AttentionMLA) -> bool: """Cheap precondition check used at call sites in the attention forward to skip the entire LoRA-correction path when no ``kv_b_proj`` adapter is wrapped on this module (the common case).""" - return getattr(attn_module.kv_b_proj, "set_lora", False) + kv_b_proj = getattr(attn_module, "kv_b_proj", None) + return getattr(kv_b_proj, "set_lora", False) def _get_state( diff --git a/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py b/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py index 7ae19430d..dc4d5a46c 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py @@ -54,7 +54,8 @@ def is_kv_b_lora_active(attn_module: DeepseekV2AttentionMLA) -> bool: """Cheap precondition check used at call sites in the attention forward to skip the entire LoRA-correction path when no ``kv_b_proj`` adapter is wrapped on this module (the common case).""" - return getattr(attn_module.kv_b_proj, "set_lora", False) + kv_b_proj = getattr(attn_module, "kv_b_proj", None) + return getattr(kv_b_proj, "set_lora", False) def _get_state( diff --git a/python/sglang/srt/model_loader/deepseek4_gguf.py b/python/sglang/srt/model_loader/deepseek4_gguf.py new file mode 100644 index 000000000..ea23c76d2 --- /dev/null +++ b/python/sglang/srt/model_loader/deepseek4_gguf.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact GGUF tensor-name mapping for DeepSeek-V4 checkpoints. + +DeepSeek-V4 GGUF files use the ``deepseek4`` architecture label, while the +current gguf Python package only exposes the closely related DeepSeek-V2 name +map. The shared entries are sufficient for most tensors; V4-only attention +compressor and mHC tensors are handled explicitly below. +""" + +from __future__ import annotations + +import re +from collections import defaultdict +from typing import Any, Iterable + +_ROUTED_EXPERT_RE = re.compile( + r"^blk\.(?P\d+)\.ffn_(?Pgate|up|down)_exps\.weight$" +) +_TENSOR_SUFFIXES = frozenset(("weight", "bias", "scale", "tid2eid")) + + +def routed_expert_tensor(name: str) -> tuple[int, str] | None: + """Return ``(layer, role)`` for an aggregated routed-expert tensor.""" + + match = _ROUTED_EXPERT_RE.fullmatch(name) + if match is None: + return None + return int(match.group("layer")), match.group("role") + + +def _split_suffix(name: str) -> tuple[str, str]: + base, separator, suffix = name.rpartition(".") + if separator and suffix in _TENSOR_SUFFIXES: + return base, suffix + return name, "" + + +def _v4_checkpoint_name(name: str) -> str | None: + base, suffix = _split_suffix(name) + suffix_part = f".{suffix}" if suffix else "" + + top_level = { + "token_embd": "embed", + "output": "head", + "output_norm": "norm", + } + if base in top_level: + return top_level[base] + suffix_part + + match = re.fullmatch(r"output_hc_(base|fn|scale)", base) + if match: + # These are direct nn.Parameters. llama.cpp adds the .weight alias + # when reading converted four-expert files, but SGLang does not. + return f"hc_head_{match.group(1)}" + + match = re.fullmatch(r"blk\.(\d+)\.(.+)", base) + if match: + layer, tensor = match.groups() + + direct_parameter = { + "attn_sinks": "attn.attn_sink", + "ffn_gate_tid2eid": "ffn.gate.tid2eid", + "hc_attn_base": "hc_attn_base", + "hc_attn_fn": "hc_attn_fn", + "hc_attn_scale": "hc_attn_scale", + "hc_ffn_base": "hc_ffn_base", + "hc_ffn_fn": "hc_ffn_fn", + "hc_ffn_scale": "hc_ffn_scale", + } + if tensor in direct_parameter: + return f"layers.{layer}.{direct_parameter[tensor]}" + + linear_or_norm = { + "attn_kv": "attn.wkv", + "attn_kv_a_norm": "attn.kv_norm", + "attn_norm": "attn_norm", + "attn_output_a": "attn.wo_a", + "attn_output_b": "attn.wo_b", + "attn_q_a": "attn.wq_a", + "attn_q_a_norm": "attn.q_norm", + "attn_q_b": "attn.wq_b", + "ffn_down_exps": "ffn.experts.w2", + "ffn_down_shexp": "ffn.shared_experts.w2", + "ffn_gate_exps": "ffn.experts.w1", + "ffn_gate_inp": "ffn.gate", + "ffn_gate_shexp": "ffn.shared_experts.w1", + "ffn_norm": "ffn_norm", + "ffn_up_exps": "ffn.experts.w3", + "ffn_up_shexp": "ffn.shared_experts.w3", + "indexer.attn_q_b": "attn.indexer.wq_b", + "indexer.proj": "attn.indexer.weights_proj", + } + if tensor in linear_or_norm: + return f"layers.{layer}.{linear_or_norm[tensor]}{suffix_part}" + + match = re.fullmatch(r"(attn|indexer)_compressor_(ape|gate|kv|norm)", tensor) + if match: + owner, component = match.groups() + owner_part = "attn" if owner == "attn" else "attn.indexer" + if component == "ape": + # Compressor.ape is a direct nn.Parameter. + return f"layers.{layer}.{owner_part}.compressor.ape" + component = {"gate": "wgate", "kv": "wkv"}.get(component, component) + return f"layers.{layer}.{owner_part}.compressor.{component}{suffix_part}" + + if tensor == "exp_probs_b": + return f"layers.{layer}.ffn.gate{suffix_part}" + + return None + + +def _candidate_score(alias: str) -> tuple[int, int, str]: + # DeepSeek's native checkpoint aliases use layers.N.attn/ffn. Selecting + # them keeps the downstream DeepSeek-V4 remapper authoritative. + if alias.startswith("layers.") and (".attn." in alias or ".ffn." in alias): + priority = 0 + elif alias.startswith("layers."): + priority = 1 + elif alias.startswith("model.layers."): + priority = 2 + else: + priority = 3 + return priority, len(alias), alias + + +def build_deepseek4_checkpoint_name_map( + gguf_module: Any, + tensor_names: Iterable[str], + num_layers: int, +) -> dict[str, str]: + """Map every source GGUF tensor to a DeepSeek checkpoint tensor name. + + The function fails closed if a source tensor has no deterministic mapping + or if two source tensors would load the same checkpoint tensor. + """ + + try: + arch = gguf_module.MODEL_ARCH.DEEPSEEK2 + except AttributeError as exc: + raise RuntimeError( + "gguf package does not provide the DeepSeek name map" + ) from exc + + name_map = gguf_module.get_tensor_name_map(arch, num_layers) + aliases_by_gguf_base: dict[str, list[str]] = defaultdict(list) + for alias, mapping in name_map.mapping.items(): + aliases_by_gguf_base[mapping[1]].append(alias) + + result: dict[str, str] = {} + reverse: dict[str, str] = {} + missing: list[str] = [] + for tensor_name in tensor_names: + checkpoint_name = _v4_checkpoint_name(tensor_name) + if checkpoint_name is None: + base, suffix = _split_suffix(tensor_name) + candidates = aliases_by_gguf_base.get(base, ()) + if candidates: + alias = min(candidates, key=_candidate_score) + checkpoint_name = alias + if suffix and not alias.endswith(f".{suffix}"): + checkpoint_name += f".{suffix}" + + if checkpoint_name is None: + missing.append(tensor_name) + continue + if checkpoint_name in reverse: + other = reverse[checkpoint_name] + raise RuntimeError( + "DeepSeek-V4 GGUF mapping collision: " + f"{other!r} and {tensor_name!r} -> {checkpoint_name!r}" + ) + result[tensor_name] = checkpoint_name + reverse[checkpoint_name] = tensor_name + + if missing: + preview = ", ".join(repr(name) for name in missing[:8]) + raise RuntimeError( + f"No DeepSeek-V4 checkpoint mapping for {len(missing)} GGUF tensors: " + f"{preview}" + ) + return result diff --git a/python/sglang/srt/model_loader/expert_pack_config.py b/python/sglang/srt/model_loader/expert_pack_config.py new file mode 100644 index 000000000..60bc36ebf --- /dev/null +++ b/python/sglang/srt/model_loader/expert_pack_config.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Lightweight model constraints shared by expert-pack startup and loading.""" + +from __future__ import annotations + +from typing import Any + +DEEPSEEK_V4_MODEL_TYPE = "deepseek_v4" +KIMI_K3_MODEL_TYPE = "kimi_linear" + +KIMI_K3_REQUIRED_CONFIG = { + "num_hidden_layers": 93, + "num_experts": 896, + "num_experts_per_token": 16, + "first_k_dense_replace": 1, + "routed_expert_hidden_size": 3584, + "moe_intermediate_size": 3072, + "num_shared_experts": 2, + "hidden_act": "situ", + "activation_situ_beta": 4.0, + "activation_situ_linear_beta": 25.0, +} + + +def validate_expert_pack_model_config(hf_config: Any) -> tuple[str | None, list[str]]: + """Return the supported model kind and every violated hard constraint.""" + model_type = getattr(hf_config, "model_type", None) + if model_type == DEEPSEEK_V4_MODEL_TYPE: + return DEEPSEEK_V4_MODEL_TYPE, [] + if model_type != KIMI_K3_MODEL_TYPE: + return None, [ + "model_type must be 'deepseek_v4' or the text-only Kimi-K3 " + f"'kimi_linear' config, got {model_type!r}" + ] + + errors = [] + for field, expected in KIMI_K3_REQUIRED_CONFIG.items(): + actual = getattr(hf_config, field, None) + if actual != expected: + errors.append(f"Kimi-K3 {field} must be {expected!r}, got {actual!r}") + return KIMI_K3_MODEL_TYPE, errors diff --git a/python/sglang/srt/model_loader/expert_pack_loader.py b/python/sglang/srt/model_loader/expert_pack_loader.py new file mode 100644 index 000000000..02f69d95d --- /dev/null +++ b/python/sglang/srt/model_loader/expert_pack_loader.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SSD expert-pack loader for deepseek-v4-flash and text-only kimi-k3. + +Only these two language-model paths are currently supported. The multimodal +kimi-k3 model is outside the scope of this loader. +""" + +from __future__ import annotations + +import logging +import os +import time +from pathlib import Path +from typing import Generator, Tuple + +import numpy as np +import torch +from torch import nn + +from sglang.kernels.ops.moe.expert_pack_mxfp4 import prewarm_mxfp4_extension +from sglang.srt.layers.moe.expert_pack import ( + ExpertPackStore, + KimiGGMLExpertPackStore, +) +from sglang.srt.layers.quantization.expert_pack import ( + ExpertPackConfig, + _clamped_swiglu, +) +from sglang.srt.model_loader.deepseek4_gguf import ( + build_deepseek4_checkpoint_name_map, + routed_expert_tensor, +) +from sglang.srt.model_loader.expert_pack_config import ( + KIMI_K3_MODEL_TYPE, + validate_expert_pack_model_config, +) +from sglang.srt.model_loader.kimi_k3_gguf import kimi_k3_nonexpert_weights_iterator +from sglang.srt.model_loader.loader import ( + BaseModelLoader, + _initialize_model, + device_loading_context, +) +from sglang.srt.model_loader.utils import set_default_torch_dtype +from sglang.srt.runtime_context import get_exec, get_parallel + +logger = logging.getLogger(__name__) + + +def _bf16_tensor(data: np.ndarray) -> torch.Tensor: + raw = np.asarray(data) + if raw.dtype != np.uint8 or raw.shape[-1] % 2: + raise ValueError("GGUF BF16 payload does not have a byte-pair layout") + values = raw.view(np.uint16).reshape(*raw.shape[:-1], raw.shape[-1] // 2) + return torch.from_numpy(values.copy()).view(torch.bfloat16) + + +def _compressor_component(source_name: str) -> str | None: + if "_compressor_kv.weight" in source_name: + return "kv" + if "_compressor_gate.weight" in source_name: + return "gate" + return None + + +def _fused_compressor_name(checkpoint_name: str) -> str: + result = checkpoint_name.replace(".wkv.weight", ".wkv_gate.weight") + result = result.replace(".wgate.weight", ".wkv_gate.weight") + if result == checkpoint_name: + raise ValueError(f"invalid compressor checkpoint name: {checkpoint_name}") + return result + + +def deepseek4_nonexpert_weights_iterator( + source_path: str | os.PathLike[str], + num_layers: int, +) -> Generator[Tuple[str, torch.Tensor], None, None]: + """Yield exact non-routed tensors without materializing routed experts.""" + + import gguf + + reader = gguf.GGUFReader(str(source_path), mode="r") + names = [tensor.name for tensor in reader.tensors] + mapping = build_deepseek4_checkpoint_name_map(gguf, names, num_layers) + tensors = {tensor.name: tensor for tensor in reader.tensors} + + # GGUF quant methods must know the type before the raw qweight arrives. + for tensor in reader.tensors: + if routed_expert_tensor(tensor.name) is not None: + continue + weight_type = tensor.tensor_type + if weight_type.name == "Q8_0": + component = _compressor_component(tensor.name) + if component == "gate": + continue + checkpoint_name = ( + _fused_compressor_name(mapping[tensor.name]) + if component == "kv" + else mapping[tensor.name] + ) + if not checkpoint_name.endswith(".weight"): + raise ValueError( + f"quantized tensor maps to a non-weight parameter: {tensor.name}" + ) + yield ( + checkpoint_name.removesuffix("weight") + "qweight_type", + torch.tensor(int(weight_type), dtype=torch.uint8), + ) + + for tensor in reader.tensors: + if routed_expert_tensor(tensor.name) is not None: + continue + checkpoint_name = mapping[tensor.name] + weight_type = tensor.tensor_type + if weight_type.name == "Q8_0": + component = _compressor_component(tensor.name) + if component == "gate": + continue + if component == "kv": + gate_name = tensor.name.replace("_compressor_kv", "_compressor_gate") + gate = tensors.get(gate_name) + if gate is None or gate.tensor_type != weight_type: + raise ValueError( + f"missing matching compressor gate tensor: {gate_name}" + ) + checkpoint_name = _fused_compressor_name(checkpoint_name) + raw_weight = torch.cat( + (torch.tensor(tensor.data), torch.tensor(gate.data)), dim=0 + ) + else: + raw_weight = torch.tensor(tensor.data) + yield checkpoint_name.removesuffix("weight") + "qweight", raw_weight + elif weight_type.name == "BF16": + yield checkpoint_name, _bf16_tensor(tensor.data) + elif weight_type.name in ("F32", "I32"): + yield checkpoint_name, torch.tensor(tensor.data) + else: + raise ValueError( + f"unsupported non-routed GGUF type {weight_type.name} for {tensor.name}" + ) + + +class ExpertPackModelLoader(BaseModelLoader): + def __init__(self, load_config) -> None: + super().__init__(load_config) + config = dict(load_config.model_loader_extra_config or {}) + pack_path = config.get("pack_path") or os.getenv("SGLANG_EXPERT_PACK_PATH") + if not pack_path: + raise ValueError( + "expert_pack load format requires pack_path or SGLANG_EXPERT_PACK_PATH" + ) + self.config = config + self.pack_path = Path(pack_path).resolve() + self.manifest_path = ( + Path(config["manifest_path"]).resolve() + if config.get("manifest_path") + else None + ) + self.source_path = ( + Path(config["source_path"]).resolve() if config.get("source_path") else None + ) + + def download_model(self, model_config) -> None: + if not Path(model_config.model_path).is_dir(): + raise ValueError( + "expert_pack model_path must be the verified tokenizer/config directory" + ) + + def load_model(self, *, model_config, device_config) -> nn.Module: + hf_config = model_config.hf_config + model_kind, model_errors = validate_expert_pack_model_config(hf_config) + if model_errors: + details = "\n".join(f"- {error}" for error in model_errors) + raise ValueError(f"Invalid expert_pack model configuration:\n{details}") + is_kimi = model_kind == KIMI_K3_MODEL_TYPE + if is_kimi: + if self.manifest_path is None or not self.manifest_path.is_file(): + raise FileNotFoundError("Kimi-K3 expert_pack requires manifest_path") + + parallel = get_parallel() + exec_config = get_exec() + if ( + parallel.tp_size != 1 + or parallel.moe_dp_size != 1 + or parallel.moe_ep_size != 1 + or not exec_config.graph.disable_cuda_graph + or not exec_config.moe.disable_shared_experts_fusion + ): + raise RuntimeError( + "expert_pack ServerArgs invariants were not applied before model load" + ) + + if is_kimi: + stats_path = self.config.get("stats_path") + store = KimiGGMLExpertPackStore( + self.pack_path, + manifest_path=self.manifest_path, + expected_layers=93, + expected_experts=896, + expected_top_k=16, + cache_vram_mib=int(self.config.get("cache_vram_mib", 4 * 1024)), + cache_vram_reserve_mib=int( + self.config.get("cache_vram_reserve_mib", 2 * 1024) + ), + stage_slots=int(self.config.get("stage_slots", 16)), + read_splits=int(self.config.get("read_splits", 1)), + direct_io=bool(self.config.get("direct_io", True)), + stats_flush_interval=int(self.config.get("stats_flush_interval", 0)), + stats_path=stats_path, + ) + weights = kimi_k3_nonexpert_weights_iterator(self.manifest_path) + else: + stats_path = self.config.get("stats_path") or os.getenv( + "SGLANG_EXPERT_PACK_STATS_PATH" + ) + required = ( + "source_path", + "source_sha256", + "model_identity_sha256", + "config_sha256", + ) + missing = [name for name in required if not self.config.get(name)] + if missing: + raise ValueError( + "expert_pack loader config is missing: " + + ", ".join(sorted(missing)) + ) + if self.source_path is None or not self.source_path.is_file(): + raise FileNotFoundError("DeepSeek source GGUF is missing") + store = ExpertPackStore( + self.pack_path, + manifest_path=self.manifest_path, + expected_layers=int(hf_config.num_hidden_layers), + expected_experts=int(hf_config.n_routed_experts), + expected_top_k=int(hf_config.num_experts_per_tok), + expected_source_sha256=self.config["source_sha256"], + expected_model_identity_sha256=self.config["model_identity_sha256"], + expected_config_sha256=self.config["config_sha256"], + cache_vram_mib=int( + self.config.get( + "cache_vram_mib", + os.getenv("SGLANG_EXPERT_CACHE_VRAM_MIB", 20 * 1024), + ) + ), + cache_vram_reserve_mib=int( + self.config.get("cache_vram_reserve_mib", 3 * 1024) + ), + stage_slots=int( + self.config.get( + "stage_slots", os.getenv("SGLANG_EXPERT_STAGE_SLOTS", 8) + ) + ), + read_splits=int(self.config.get("read_splits", 1)), + direct_io=bool(self.config.get("direct_io", True)), + stats_flush_interval=int(self.config.get("stats_flush_interval", 0)), + stats_path=stats_path, + ) + weights = deepseek4_nonexpert_weights_iterator( + self.source_path, int(hf_config.num_hidden_layers) + ) + quant_config = ExpertPackConfig(store) + target_device = torch.device(device_config.device) + with set_default_torch_dtype(model_config.dtype): + with target_device: + model = _initialize_model(model_config, self.load_config, quant_config) + loaded_params = model.load_weights(weights) + if is_kimi: + if loaded_params is None: + raise RuntimeError( + "Kimi-K3 load_weights did not return its parameter coverage" + ) + expected_params = {name for name, _ in model.named_parameters()} + missing_params = sorted(expected_params - set(loaded_params)) + if missing_params: + preview = ", ".join(missing_params[:16]) + raise RuntimeError( + "Kimi-K3 GGUF did not initialize all model parameters: " + f"missing={len(missing_params)} [{preview}]" + ) + logger.info( + "Kimi-K3 parameter coverage complete: loaded=%d expected=%d", + len(set(loaded_params) & expected_params), + len(expected_params), + ) + for _, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None: + with device_loading_context(module, target_device): + quant_method.process_weights_after_loading(module) + + store.initialize_device_cache(target_device) + if not is_kimi: + prewarm_started = time.monotonic() + prewarm_mxfp4_extension() + activation_input = torch.zeros( + (1, int(hf_config.moe_intermediate_size)), + dtype=model_config.dtype, + device=target_device, + ) + _clamped_swiglu(activation_input, activation_input, hf_config.swiglu_limit) + torch.cuda.synchronize(target_device) + del activation_input + torch.cuda.empty_cache() + logger.info( + "Expert-pack CUDA extension and clamped SwiGLU prewarmed in %.3fs", + time.monotonic() - prewarm_started, + ) + dense_bytes = sum( + value.numel() * value.element_size() + for value in list(model.parameters()) + list(model.buffers()) + ) + store.stats["dense_bytes"] = dense_bytes + model.expert_pack_store = store + logger.info( + "Loaded verified DeepSeek expert-pack model: source_sha256=%s " + "pack_sha256=%s dense_bytes=%d resident_experts=0", + store.header.source_blob_sha256, + store.pack_sha256, + dense_bytes, + ) + return model.eval() diff --git a/python/sglang/srt/model_loader/expert_pack_runtime.py b/python/sglang/srt/model_loader/expert_pack_runtime.py new file mode 100644 index 000000000..173d890b0 --- /dev/null +++ b/python/sglang/srt/model_loader/expert_pack_runtime.py @@ -0,0 +1,575 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Internal preparation of source assets for the expert-pack load format.""" + +from __future__ import annotations + +import fcntl +import hashlib +import importlib.metadata +import json +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +METADATA_FORMAT_VERSION = 3 +GGUF_SHARD_SUFFIX_RE = re.compile(r"-\d{5}-of-\d{5}\.gguf$") +DEEPSEEK_METADATA_FORMAT_VERSION = 4 + + +def cache_root() -> Path: + return Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() + + +def artifact_dir_for_source(gguf: Path) -> Path: + stat = gguf.stat() + fingerprint = hashlib.sha256( + f"{gguf.parent.resolve()}:{stat.st_size}:{stat.st_mtime_ns}:" + f"{METADATA_FORMAT_VERSION}".encode() + ).hexdigest()[:20] + return cache_root() / "sglang-expert-pack" / "kimi-k3" / fingerprint + + +def _tokenizer_candidate(path: Path) -> bool: + return ( + path.is_dir() + and (path / "config.json").is_file() + and any( + (path / name).is_file() + for name in ("tokenizer.json", "tiktoken.model", "tokenizer_config.json") + ) + ) + + +def resolve_kimi_tokenizer(gguf: Path, explicit: str | None = None) -> Path: + if explicit: + candidate = Path(explicit).expanduser().resolve() + if not _tokenizer_candidate(candidate): + raise ValueError(f"Kimi tokenizer directory is invalid: {candidate}") + return candidate + + candidates = [gguf.parent / "tokenizer", gguf.parent.parent / "kimi-k3-tokenizer"] + candidates.extend( + sorted(path for path in gguf.parent.parent.glob("*tokenizer*") if path.is_dir()) + ) + tokenizers = [] + for path in candidates: + path = path.resolve() + if path not in tokenizers and _tokenizer_candidate(path): + tokenizers.append(path) + if len(tokenizers) != 1: + names = ", ".join(str(path) for path in tokenizers) or "none" + raise RuntimeError( + f"could not uniquely derive Kimi tokenizer beside {gguf.parent}; " + f"candidates: {names}" + ) + return tokenizers[0] + + +def _write_json_atomic(path: Path, value: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + f".{os.getpid()}.tmp") + temporary.write_text( + json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def prepare_kimi_model_metadata(tokenizer_dir: Path, artifact_dir: Path) -> Path: + tokenizer_dir = tokenizer_dir.resolve(strict=True) + source_config = json.loads( + (tokenizer_dir / "config.json").read_text(encoding="utf-8") + ) + if "text_config" not in source_config: + raise ValueError("Kimi tokenizer config does not contain text_config") + config = dict(source_config["text_config"]) + config["architectures"] = ["KimiK3LinearForCausalLM"] + config["model_type"] = "kimi_linear" + config.pop("auto_map", None) + config.pop("quantization_config", None) + + output_dir = artifact_dir / "model-meta" + output_dir.mkdir(parents=True, exist_ok=True) + for source in tokenizer_dir.iterdir(): + if source.is_file() and source.name != "config.json": + destination = output_dir / source.name + if ( + not destination.is_file() + or destination.stat().st_size != source.stat().st_size + or destination.stat().st_mtime_ns != source.stat().st_mtime_ns + ): + shutil.copy2(source, destination) + _write_json_atomic(output_dir / "config.json", config) + return output_dir + + +def _expert_pack_path(gguf: Path) -> Path: + match = GGUF_SHARD_SUFFIX_RE.search(gguf.name) + if match is None: + raise ValueError(f"Kimi GGUF is not a numbered shard: {gguf}") + return gguf.parent / f"{gguf.name[: match.start()]}.expert-major.pack" + + +def _repo_root() -> Path: + for candidate in Path(__file__).resolve().parents: + if (candidate / "tools" / "expert_pack" / "prepare_kimi_pack.py").is_file(): + return candidate + raise RuntimeError( + "expert_pack cannot auto-build Kimi artifacts from an installed package; " + "run from an SGLang source checkout" + ) + + +def ensure_kimi_assets( + gguf: Path, + *, + tokenizer_dir: str | None = None, +) -> dict[str, Path]: + """Build or reuse Kimi artifacts and return the internal serving paths.""" + gguf = gguf.expanduser().resolve(strict=True) + if not gguf.is_file() or gguf.suffix != ".gguf": + raise ValueError(f"expert_pack expects a local Kimi GGUF shard, got {gguf}") + if "KIMI-K3" not in gguf.name.upper(): + raise ValueError( + "raw GGUF auto-preparation currently supports only Kimi-K3; " + "provide the model metadata and loader artifacts for other models" + ) + gguf_dir = gguf.parent + artifact_dir = artifact_dir_for_source(gguf).resolve() + pack = _expert_pack_path(gguf).resolve() + manifest = artifact_dir / "kimi-k3-expert-pack.manifest.json" + tokenizer = resolve_kimi_tokenizer(gguf, tokenizer_dir) + lock_path = pack.with_name(pack.name + ".startup.lock") + repo = _repo_root() + with lock_path.open("w") as lock: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX) + model_dir = prepare_kimi_model_metadata(tokenizer, artifact_dir) + subprocess.run( + [ + sys.executable, + str(repo / "tools" / "expert_pack" / "prepare_kimi_pack.py"), + "--gguf", + str(gguf), + "--model-config", + str(model_dir / "config.json"), + ], + cwd=repo, + check=True, + ) + subprocess.run( + [ + sys.executable, + str(repo / "tools" / "expert_pack" / "prepare_kimi_manifest.py"), + "--gguf-dir", + str(gguf_dir), + "--expert-pack", + str(pack), + "--model-config", + str(model_dir / "config.json"), + "--tokenizer-dir", + str(tokenizer), + "--output", + str(manifest), + "--payload-samples", + "6", + ], + cwd=repo, + check=True, + ) + return { + "gguf": gguf, + "gguf_dir": gguf_dir, + "tokenizer_dir": tokenizer, + "model_dir": model_dir, + "pack_path": pack, + "manifest_path": manifest, + "stats_path": artifact_dir / "kimi-k3-expert-pack.stats.json", + "artifact_dir": artifact_dir, + } + + +def prepare_raw_kimi_server_args( + server_args: Any, loader_config: dict[str, Any] +) -> None: + """Resolve a raw GGUF model path into the normal loader inputs.""" + model_path = Path(server_args.model_path).expanduser() + if not model_path.is_file() or model_path.suffix.lower() != ".gguf": + return + tokenizer_path = server_args.tokenizer_path + if tokenizer_path and Path(tokenizer_path).expanduser() == model_path: + tokenizer_path = None + assets = ensure_kimi_assets( + model_path, + tokenizer_dir=tokenizer_path, + ) + server_args._declare( + "prepare_raw_kimi_server_args", + model_path=str(assets["model_dir"]), + tokenizer_path=str(assets["model_dir"]), + ) + for key in ("pack_path", "manifest_path", "stats_path"): + loader_config.setdefault(key, str(assets[key])) + loader_config.setdefault("source_path", str(assets["gguf"])) + + +def _deepseek_cache_root() -> Path: + return cache_root() / "sglang-expert-pack" / "deepseek-v4-flash" + + +def _deepseek_artifact_dir_for_source(source: Path) -> Path: + stat = source.stat() + fingerprint = hashlib.sha256( + f"{source.resolve()}:{stat.st_size}:{stat.st_mtime_ns}:" + f"{DEEPSEEK_METADATA_FORMAT_VERSION}".encode() + ).hexdigest()[:20] + return _deepseek_cache_root() / fingerprint + + +def _deepseek_gguf_value(reader: object, name: str) -> object: + fields = getattr(reader, "fields") + if name not in fields: + raise ValueError(f"GGUF metadata is missing required field: {name}") + return fields[name].contents() + + +def _deepseek_model_config_from_gguf(reader: object) -> dict[str, Any]: + def value(name: str) -> object: + return _deepseek_gguf_value(reader, f"deepseek4.{name}") + + architecture = _deepseek_gguf_value(reader, "general.architecture") + if architecture != "deepseek4": + raise ValueError(f"expected GGUF architecture deepseek4, got {architecture!r}") + if int(value("expert_gating_func")) != 4: + raise ValueError("unsupported deepseek4.expert_gating_func; expected 4") + swiglu_limits = [float(item) for item in value("swiglu_clamp_exp")] + if not swiglu_limits or any(item != swiglu_limits[0] for item in swiglu_limits): + raise ValueError("deepseek4.swiglu_clamp_exp must be constant") + tokens = _deepseek_gguf_value(reader, "tokenizer.ggml.tokens") + return { + "architectures": ["DeepseekV4ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": int( + _deepseek_gguf_value(reader, "tokenizer.ggml.bos_token_id") + ), + "eos_token_id": int( + _deepseek_gguf_value(reader, "tokenizer.ggml.eos_token_id") + ), + "expert_dtype": "fp4", + "hc_eps": float(value("hyper_connection.epsilon")), + "hc_mult": int(value("hyper_connection.count")), + "hc_sinkhorn_iters": int(value("hyper_connection.sinkhorn_iterations")), + "head_dim": int(value("attention.key_length")), + "hidden_act": "silu", + "hidden_size": int(value("embedding_length")), + "index_head_dim": int(value("attention.indexer.key_length")), + "index_n_heads": int(value("attention.indexer.head_count")), + "index_topk": int(value("attention.indexer.top_k")), + "initializer_range": 0.02, + "max_position_embeddings": int(value("context_length")), + "model_type": "deepseek_v4", + "moe_intermediate_size": int(value("expert_feed_forward_length")), + "n_routed_experts": int(value("expert_count")), + "n_shared_experts": int(value("expert_shared_count")), + "norm_topk_prob": bool(value("expert_weights_norm")), + "num_attention_heads": int(value("attention.head_count")), + "num_experts_per_tok": int(value("expert_used_count")), + "num_hidden_layers": int(value("block_count")), + "num_hash_layers": int(value("hash_layer_count")), + "num_key_value_heads": int(value("attention.head_count_kv")), + "num_nextn_predict_layers": 1, + "o_groups": int(value("attention.output_group_count")), + "o_lora_rank": int(value("attention.output_lora_rank")), + "q_lora_rank": int(value("attention.q_lora_rank")), + "qk_rope_head_dim": int(value("rope.dimension_count")), + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + }, + "rms_norm_eps": float(value("attention.layer_norm_rms_epsilon")), + "rope_scaling": { + "beta_fast": float(value("rope.scaling.yarn_beta_fast")), + "beta_slow": float(value("rope.scaling.yarn_beta_slow")), + "factor": float(value("rope.scaling.factor")), + "original_max_position_embeddings": int( + value("rope.scaling.original_context_length") + ), + "type": str(value("rope.scaling.type")), + }, + "rope_theta": float(value("rope.freq_base")), + "routed_scaling_factor": float(value("expert_weights_scale")), + "scoring_func": "sqrtsoftplus", + "sliding_window": int(value("attention.sliding_window")), + "swiglu_limit": swiglu_limits[0], + "tie_word_embeddings": False, + "topk_method": "noaux_tc", + "torch_dtype": "bfloat16", + "transformers_version": importlib.metadata.version("transformers"), + "use_cache": True, + "vocab_size": len(tokens), + "compress_rope_theta": float(value("attention.compress_rope_freq_base")), + "compress_ratios": [int(item) for item in value("attention.compress_ratios")], + } + + +def _write_deepseek_tokenizer_from_gguf( + reader: object, output_dir: Path, config: dict[str, Any] +) -> None: + from tokenizers import AddedToken, Regex, normalizers, pre_tokenizers + from transformers.integrations.ggml import convert_gguf_tokenizer + + tokenizer_type = _deepseek_gguf_value(reader, "tokenizer.ggml.model") + pre_tokenizer_type = _deepseek_gguf_value(reader, "tokenizer.ggml.pre") + if tokenizer_type != "gpt2" or pre_tokenizer_type != "joyai-llm": + raise ValueError( + f"unsupported GGUF tokenizer: model={tokenizer_type!r} " + f"pre={pre_tokenizer_type!r}" + ) + tokens = list(_deepseek_gguf_value(reader, "tokenizer.ggml.tokens")) + token_types = list(_deepseek_gguf_value(reader, "tokenizer.ggml.token_type")) + tokenizer_data = { + "tokenizer_type": tokenizer_type, + "tokens": tokens, + "token_type": token_types, + "merges": list(_deepseek_gguf_value(reader, "tokenizer.ggml.merges")), + "bos_token_id": config["bos_token_id"], + "eos_token_id": config["eos_token_id"], + "pad_token_id": int( + _deepseek_gguf_value(reader, "tokenizer.ggml.padding_token_id") + ), + } + tokenizer, _ = convert_gguf_tokenizer("gpt2", tokenizer_data) + tokenizer.add_special_tokens( + [ + AddedToken(token, normalized=False, special=True) + for token, token_type in zip(tokens, token_types) + if token_type in (3, 4) + ] + ) + tokenizer.normalizer = normalizers.Sequence([]) + tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split(Regex(r"\p{N}{1,3}"), behavior="isolated"), + pre_tokenizers.Split( + Regex(r"[\u4e00-\u9fa5\u3040-\u309f\u30a0-\u30ff]+"), + behavior="isolated", + ), + pre_tokenizers.Split( + Regex( + r"[!\"#$%&'()*+,\-./:;<=>?@\[\\\]^_`{|}~][A-Za-z]+|" + r"[^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+|" + r" ?[\p{P}\p{S}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" + ), + behavior="isolated", + ), + pre_tokenizers.ByteLevel( + add_prefix_space=False, trim_offsets=True, use_regex=False + ), + ] + ) + tokenizer.save(str(output_dir / "tokenizer.json")) + + def token(token_id: int) -> dict[str, object]: + return { + "__type": "AddedToken", + "content": tokens[token_id], + "lstrip": False, + "normalized": False, + "rstrip": False, + "single_word": False, + } + + tokenizer_config = { + "add_bos_token": bool( + _deepseek_gguf_value(reader, "tokenizer.ggml.add_bos_token") + ), + "add_eos_token": bool( + _deepseek_gguf_value(reader, "tokenizer.ggml.add_eos_token") + ), + "bos_token": token(config["bos_token_id"]), + "chat_template": _deepseek_gguf_value(reader, "tokenizer.chat_template"), + "clean_up_tokenization_spaces": False, + "eos_token": token(config["eos_token_id"]), + "model_max_length": config["max_position_embeddings"], + "pad_token": token( + int(_deepseek_gguf_value(reader, "tokenizer.ggml.padding_token_id")) + ), + "tokenizer_class": "PreTrainedTokenizerFast", + "unk_token": None, + } + (output_dir / "tokenizer_config.json").write_text( + json.dumps(tokenizer_config, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +def _prepare_deepseek_model_metadata(source: Path, artifact_dir: Path) -> Path: + output_dir = artifact_dir / "model-meta" + marker = output_dir / "metadata.json" + stat = source.stat() + if ( + marker.is_file() + and (output_dir / "config.json").is_file() + and (output_dir / "tokenizer.json").is_file() + ): + try: + metadata = json.loads(marker.read_text(encoding="utf-8")) + if ( + metadata.get("size") == stat.st_size + and metadata.get("mtime_ns") == stat.st_mtime_ns + ): + return output_dir / "config.json" + except (OSError, ValueError): + pass + try: + import gguf + except ImportError as exc: + raise RuntimeError( + "the gguf Python package is required to prepare DeepSeek metadata" + ) from exc + output_dir.mkdir(parents=True, exist_ok=True) + reader = gguf.GGUFReader(str(source), "r") + config = _deepseek_model_config_from_gguf(reader) + (output_dir / "config.json").write_text( + json.dumps(config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + ) + generation_config = { + "_from_model_config": True, + "bos_token_id": config["bos_token_id"], + "eos_token_id": config["eos_token_id"], + "do_sample": True, + "temperature": float(_deepseek_gguf_value(reader, "general.sampling.temp")), + "top_p": float(_deepseek_gguf_value(reader, "general.sampling.top_p")), + } + (output_dir / "generation_config.json").write_text( + json.dumps(generation_config, indent=2) + "\n", encoding="utf-8" + ) + _write_deepseek_tokenizer_from_gguf(reader, output_dir, config) + marker.write_text( + json.dumps( + { + "format_version": DEEPSEEK_METADATA_FORMAT_VERSION, + "gguf": str(source), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + return output_dir / "config.json" + + +def _deepseek_digest(value: object, field: str) -> str: + digest = str(value or "").lower() + if len(digest) != 64 or any(char not in "0123456789abcdef" for char in digest): + raise ValueError(f"manifest {field} is not a SHA-256 digest") + return digest + + +def _prepare_deepseek_pack( + source: Path, model_config: Path, repo: Path +) -> tuple[Path, Path]: + tool = repo / "tools" / "expert_pack" / "prepare_deepseek_pack.py" + if not tool.is_file(): + raise FileNotFoundError(f"missing DeepSeek Expert Pack preparer: {tool}") + subprocess.run( + [ + sys.executable, + str(tool), + "--gguf", + str(source), + "--model-config", + str(model_config), + ], + cwd=repo, + check=True, + ) + return ( + source.parent / "DeepSeek-V4-Flash.expert-pack", + source.parent / "DeepSeek-V4-Flash.expert-pack.manifest.json", + ) + + +def prepare_raw_deepseek_server_args( + server_args: Any, loader_config: dict[str, Any] +) -> None: + """Resolve a raw DeepSeek V4 GGUF into metadata and Expert Pack inputs.""" + source = Path(server_args.model_path).expanduser().resolve(strict=True) + if not source.is_file(): + return + repo = _repo_root() + artifact_dir = _deepseek_artifact_dir_for_source(source).resolve() + lock_path = artifact_dir / "deepseek-v4-startup.lock" + artifact_dir.mkdir(parents=True, exist_ok=True) + with lock_path.open("w") as lock: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX) + model_config = _prepare_deepseek_model_metadata(source, artifact_dir) + pack, manifest = _prepare_deepseek_pack(source, model_config, repo) + manifest_value = json.loads(manifest.read_text(encoding="utf-8")) + source_value = manifest_value.get("source") or {} + model_value = manifest_value.get("model") or {} + source_sha256 = _deepseek_digest(source_value.get("sha256"), "source.sha256") + model_identity_sha256 = _deepseek_digest( + model_value.get("model_identity_sha256"), + "model.model_identity_sha256", + ) + config_sha256 = _deepseek_digest( + model_value.get("config_sha256"), "model.config_sha256" + ) + server_args._declare( + "prepare_raw_deepseek_server_args", + model_path=str(model_config.parent), + tokenizer_path=str(model_config.parent), + ) + for key, value in { + "pack_path": pack, + "manifest_path": manifest, + "source_path": source, + "source_sha256": source_sha256, + "model_identity_sha256": model_identity_sha256, + "config_sha256": config_sha256, + "stats_path": artifact_dir / "deepseek-v4-expert-pack.stats.json", + }.items(): + loader_config.setdefault(key, str(value) if isinstance(value, Path) else value) + + +def prepare_raw_expert_pack_server_args( + server_args: Any, loader_config: dict[str, Any] +) -> None: + """Dispatch a raw GGUF to the model-specific expert-pack preparation path.""" + source = Path(server_args.model_path).expanduser() + if not source.is_file(): + return + name = source.name.upper() + if "KIMI" in name: + prepare_raw_kimi_server_args(server_args, loader_config) + return + if "DEEPSEEK" in name: + prepare_raw_deepseek_server_args(server_args, loader_config) + return + try: + import gguf + + reader = gguf.GGUFReader(str(source), "r") + architecture = _deepseek_gguf_value(reader, "general.architecture") + except Exception as exc: + raise ValueError( + "raw GGUF auto-preparation currently supports DeepSeek-V4 and Kimi-K3; " + f"could not identify {source}: {exc}" + ) from exc + if architecture == "deepseek4": + prepare_raw_deepseek_server_args(server_args, loader_config) + return + raise ValueError( + "raw GGUF auto-preparation currently supports DeepSeek-V4 and Kimi-K3; " + f"detected architecture {architecture!r}" + ) diff --git a/python/sglang/srt/model_loader/kimi_k3_gguf.py b/python/sglang/srt/model_loader/kimi_k3_gguf.py new file mode 100644 index 000000000..876c73c78 --- /dev/null +++ b/python/sglang/srt/model_loader/kimi_k3_gguf.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact non-routed GGUF loader for the Kimi-K3 expert-pack runtime.""" + +from __future__ import annotations + +import json +import os +import re +from collections import defaultdict +from collections.abc import Generator +from pathlib import Path + +import torch + +_LAYER_RE = re.compile(r"^blk\.(?P\d+)\.(?P.+)$") +_ROUTED_EXPERT_RE = re.compile(r"^blk\.\d+\.ffn_(?:gate|up|down)_exps\.weight$") + +_TOP_LEVEL_NAMES = { + "token_embd.weight": "model.embed_tokens.weight", + "output.weight": "lm_head.weight", + "output_norm.weight": "model.norm.weight", +} + +_COMMON_LAYER_NAMES = { + "attn_norm.weight": "input_layernorm.weight", + "ffn_norm.weight": "post_attention_layernorm.weight", + "attn_output.weight": "self_attn.o_proj.weight", + "ffn_gate.weight": "mlp.gate_proj.weight", + "ffn_up.weight": "mlp.up_proj.weight", + "ffn_down.weight": "mlp.down_proj.weight", + "exp_probs_b.bias": "mlp.gate.e_score_correction_bias", + "ffn_gate_inp.weight": "mlp.gate.weight", + "ffn_routed_down.weight": "mlp.routed_expert_down_proj.weight", + "ffn_routed_norm.weight": "mlp.routed_expert_norm.weight", + "ffn_routed_up.weight": "mlp.routed_expert_up_proj.weight", + "ffn_gate_shexp.weight": "mlp.shared_experts.gate_proj.weight", + "ffn_up_shexp.weight": "mlp.shared_experts.up_proj.weight", + "ffn_down_shexp.weight": "mlp.shared_experts.down_proj.weight", +} + +_KDA_NAMES = { + "attn_q.weight": "self_attn.q_proj.weight", + "attn_k.weight": "self_attn.k_proj.weight", + "attn_v.weight": "self_attn.v_proj.weight", + "ssm_g.weight": "self_attn.g_proj.weight", + "ssm_beta.weight": "self_attn.b_proj.weight", + "ssm_f_a.weight": "self_attn.f_a_proj.weight", + "ssm_f_b.weight": "self_attn.f_b_proj.weight", + "ssm_conv1d_q.weight": "self_attn.q_conv1d.weight", + "ssm_conv1d_k.weight": "self_attn.k_conv1d.weight", + "ssm_conv1d_v.weight": "self_attn.v_conv1d.weight", + "ssm_a": "self_attn.A_log", + "ssm_dt.bias": "self_attn.dt_bias", + "ssm_norm.weight": "self_attn.o_norm.weight", +} + +_MLA_NAMES = { + "attn_q_a.weight": "self_attn.q_a_proj.weight", + "attn_q_a_norm.weight": "self_attn.q_a_layernorm.weight", + "attn_q_b.weight": "self_attn.q_b_proj.weight", + "attn_kv_a_mqa.weight": "self_attn.kv_a_proj_with_mqa.weight", + "attn_kv_a_norm.weight": "self_attn.kv_a_layernorm.weight", + "attn_gate.weight": "self_attn.g_proj.weight", + # K and V use different GGUF types and must remain separate. + "attn_k_b.weight": "self_attn.k_b_qweight", + "attn_v_b.weight": "self_attn.v_b_qweight", +} + + +def routed_expert_tensor(name: str) -> bool: + return _ROUTED_EXPERT_RE.fullmatch(name) is not None + + +def kimi_k3_checkpoint_targets(source_name: str) -> tuple[str, ...]: + """Map one llama.cpp Kimi-K3 tensor to exact SGLang parameter names.""" + + if source_name == "output_res_score.weight": + return ( + "model.output_attn_res_proj.weight", + "model.output_attn_res_norm.weight", + ) + if source_name in _TOP_LEVEL_NAMES: + return (_TOP_LEVEL_NAMES[source_name],) + + match = _LAYER_RE.fullmatch(source_name) + if match is None: + raise KeyError(f"unsupported Kimi-K3 GGUF tensor name: {source_name}") + layer = int(match.group("layer")) + suffix = match.group("suffix") + prefix = f"model.layers.{layer}." + if suffix == "attn_res_score.weight": + return ( + prefix + "self_attention_res_proj.weight", + prefix + "self_attention_res_norm.weight", + ) + if suffix == "ffn_res_score.weight": + return ( + prefix + "mlp_res_proj.weight", + prefix + "mlp_res_norm.weight", + ) + target = _COMMON_LAYER_NAMES.get(suffix) + if target is None: + target = _KDA_NAMES.get(suffix) + if target is None: + target = _MLA_NAMES.get(suffix) + if target is None: + raise KeyError(f"unsupported Kimi-K3 GGUF tensor name: {source_name}") + return (prefix + target,) + + +def _runtime_name(checkpoint_name: str, quantized: bool) -> str: + if not quantized or not checkpoint_name.endswith(".weight"): + return checkpoint_name + return checkpoint_name.removesuffix("weight") + "qweight" + + +def _residual_target_value(raw: torch.Tensor, target_index: int) -> torch.Tensor: + if raw.ndim != 1: + raise ValueError( + f"Kimi-K3 attention-residual score must be a vector, got {tuple(raw.shape)}" + ) + if target_index == 0: + return raw.unsqueeze(0) + if target_index == 1: + return torch.ones_like(raw) + raise ValueError(f"invalid Kimi-K3 attention-residual target {target_index}") + + +def _kda_a_log_target_value(raw: torch.Tensor) -> torch.Tensor: + """Undo llama.cpp's GGUF-time ``A_log -> -exp(A_log)`` transform.""" + if not raw.is_floating_point() or not torch.isfinite(raw).all(): + raise ValueError("Kimi-K3 GGUF ssm_a must contain finite floating values") + if not torch.all(raw < 0): + raise ValueError("Kimi-K3 GGUF ssm_a must contain only -exp(A_log) values") + return torch.log(-raw) + + +def kimi_k3_nonexpert_weights_iterator( + manifest_path: str | os.PathLike[str], +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Stream non-routed tensors shard by shard without reading routed payloads.""" + + import gguf + + manifest_file = Path(manifest_path).resolve() + manifest = json.loads(manifest_file.read_text(encoding="utf-8")) + if manifest.get("format") != "SGLANG-KIMI-GGMLMOEPACK-ADAPTER-v1": + raise ValueError("Kimi-K3 manifest format is unsupported") + if not manifest.get("complete"): + raise ValueError("Kimi-K3 manifest is incomplete") + + records_by_shard: dict[int, list[dict]] = defaultdict(list) + for record in manifest["source"]["tensors"]: + records_by_shard[int(record["shard_index"])].append(record) + + emitted: set[str] = set() + for shard in manifest["source"]["shards"]: + shard_index = int(shard["index"]) + shard_path = Path(shard["path"]).resolve() + if not shard_path.is_file() or shard_path.stat().st_size != int(shard["size"]): + raise FileNotFoundError( + f"Kimi-K3 GGUF shard is missing or changed: {shard_path}" + ) + reader = gguf.GGUFReader(str(shard_path), mode="r") + tensors = {tensor.name: tensor for tensor in reader.tensors} + expected = {record["name"]: record for record in records_by_shard[shard_index]} + if set(tensors) != set(expected): + raise ValueError(f"Kimi-K3 GGUF shard inventory changed: {shard_path}") + + for source_name, tensor in tensors.items(): + record = expected[source_name] + if tensor.tensor_type.name != record["dtype"]: + raise ValueError(f"Kimi-K3 GGUF tensor type changed: {source_name}") + if routed_expert_tensor(source_name): + continue + + targets = kimi_k3_checkpoint_targets(source_name) + quantized = tensor.tensor_type.name not in ("F32", "F16", "BF16") + raw = torch.tensor(tensor.data) + for target_index, checkpoint_name in enumerate(targets): + if source_name.endswith(".ssm_a"): + value = _kda_a_log_target_value(raw) + elif len(targets) == 2: + value = _residual_target_value(raw, target_index) + else: + value = raw + runtime_name = _runtime_name(checkpoint_name, quantized) + if runtime_name in emitted: + raise ValueError( + f"duplicate Kimi-K3 target parameter: {runtime_name}" + ) + if quantized: + type_name = runtime_name.removesuffix("qweight") + "qweight_type" + if type_name in emitted: + raise ValueError( + f"duplicate Kimi-K3 target parameter: {type_name}" + ) + emitted.add(type_name) + yield type_name, torch.tensor( + int(tensor.tensor_type), dtype=torch.uint8 + ) + emitted.add(runtime_name) + yield runtime_name, value diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 23efee357..85be7f2b5 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -4302,6 +4302,11 @@ def get_model_loader( if load_config.load_format == LoadFormat.GGUF: return GGUFModelLoader(load_config) + if load_config.load_format == LoadFormat.EXPERT_PACK: + from sglang.srt.model_loader.expert_pack_loader import ExpertPackModelLoader + + return ExpertPackModelLoader(load_config) + if load_config.load_format == LoadFormat.LAYERED: return LayeredModelLoader(load_config) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index e140a117d..7dad2fb6b 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -148,6 +148,8 @@ class DeepseekMLAForwardMixin: def _can_fuse_bmm_into_attention( self: DeepseekV2AttentionMLA, forward_batch: ForwardBatch ) -> bool: + if getattr(self, "_kimi_split_gguf_kv_b", False): + return False # Shared activation surface with the DSA indexer graph dispatch # (in piecewise/breakable graph + non-speculative extend). Like the indexer # dispatch, this fusion is on by default on that surface. @@ -477,6 +479,17 @@ class DeepseekMLAForwardMixin: .transpose(0, 1) .contiguous() ) + elif getattr(self, "_kimi_split_gguf_kv_b", False): + from sglang.srt.layers.quantization.gguf import fused_mul_mat_gguf + + k_type = int(self.k_b_qweight_type.weight_type) + q_nope_out = torch.stack( + [ + fused_mul_mat_gguf(q_nope[:, head], self.k_b_qweight[head], k_type) + for head in range(self.num_local_heads) + ], + dim=1, + ) elif fusion_plan is not None: # The composite split op fills q_nope_out_buf and attention reads # this transposed alias directly. @@ -797,7 +810,20 @@ class DeepseekMLAForwardMixin: _kvb_v = kv_b_lora_v_prepare(self, attn_output) - if self.use_deep_gemm_bmm: + if getattr(self, "_kimi_split_gguf_kv_b", False): + from sglang.srt.layers.quantization.gguf import fused_mul_mat_gguf + + v_type = int(self.v_b_qweight_type.weight_type) + attn_bmm_output = torch.stack( + [ + fused_mul_mat_gguf( + attn_output[:, head], self.v_b_qweight[head], v_type + ) + for head in range(self.num_local_heads) + ], + dim=1, + ).flatten(1, 2) + elif self.use_deep_gemm_bmm: ( attn_output_val, attn_output_scale, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py index fa781b39c..40420c5fe 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py @@ -29,15 +29,22 @@ class DeepseekMLACpuForwardMixin: weight_names=["w_kc", "w_vc"], transpose_dims=[[1, 2], [1, 2]] ) + fused_qkv_weight = ( + getattr(self.fused_qkv_a_proj_with_mqa, "weight", None) + if self.has_fused_proj + else None + ) self.qkv_proj_with_rope_is_int8 = ( self.has_fused_proj and not self.is_packed_weight - and self.fused_qkv_a_proj_with_mqa.weight.dtype == torch.int8 + and fused_qkv_weight is not None + and fused_qkv_weight.dtype == torch.int8 ) self.qkv_proj_with_rope_is_fp8 = ( self.has_fused_proj and not self.is_packed_weight - and self.fused_qkv_a_proj_with_mqa.weight.dtype == torch.float8_e4m3fn + and fused_qkv_weight is not None + and fused_qkv_weight.dtype == torch.float8_e4m3fn ) self.weight_block_size = None diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 570fa2bba..1259dcd2f 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -357,6 +357,7 @@ class DeepseekV2MLP(nn.Module): if ( gemm_output_zero_allocator is not None and x.shape[0] <= 256 + and getattr(self.gate_up_proj, "weight", None) is not None and self.gate_up_proj.weight.dtype == torch.uint8 ): y = gemm_output_zero_allocator.allocate( @@ -371,6 +372,7 @@ class DeepseekV2MLP(nn.Module): if ( self.swiglu_limit is not None and not self.down_proj.reduce_results + and getattr(self.down_proj, "weight", None) is not None and self.down_proj.weight.dtype == torch.uint8 and hasattr(self.down_proj, "weight_scale_inv") ): @@ -490,9 +492,12 @@ class MoEGate(nn.Module): "quark", ): correction_bias_dtype = torch.bfloat16 - self.e_score_correction_bias = nn.Parameter( - torch.empty((config.n_routed_experts), dtype=correction_bias_dtype) + correction_bias = torch.empty( + (config.n_routed_experts), dtype=correction_bias_dtype ) + if quant_config is not None and quant_config.get_name() == "expert_pack": + correction_bias.zero_() + self.e_score_correction_bias = nn.Parameter(correction_bias) else: self.e_score_correction_bias = None if _is_cpu and _is_cpu_amx_available: @@ -785,13 +790,23 @@ class DeepseekV2MoE(nn.Module): "awq_marlin", "moe_wna16", } + shared_gate_up_weight = getattr( + self.shared_experts.gate_up_proj, "weight", None + ) + if shared_gate_up_weight is None: + shared_gate_up_weight = getattr( + self.shared_experts.gate_up_proj, "qweight", None + ) + if shared_gate_up_weight is None: + raise ValueError( + "shared expert gate/up projection has no weight storage" + ) self.shared_experts_is_int8 = ( - not is_packed_weight - and self.shared_experts.gate_up_proj.weight.dtype == torch.int8 + not is_packed_weight and shared_gate_up_weight.dtype == torch.int8 ) self.shared_experts_is_fp8 = ( not is_packed_weight - and self.shared_experts.gate_up_proj.weight.dtype == torch.float8_e4m3fn + and shared_gate_up_weight.dtype == torch.float8_e4m3fn ) if self.shared_experts_is_fp8: if ( @@ -1959,8 +1974,11 @@ class DeepseekV2AttentionMLA( self.has_q_b_proj = hasattr(self, "q_b_proj") q_b_proj_verified_shapes = {(2048, 2048), (4096, 2048)} - self._q_b_proj_verified_shape = self.has_q_b_proj and ( - tuple(self.q_b_proj.weight.shape) in q_b_proj_verified_shapes + q_b_weight = ( + getattr(self.q_b_proj, "weight", None) if self.has_q_b_proj else None + ) + self._q_b_proj_verified_shape = q_b_weight is not None and ( + tuple(q_b_weight.shape) in q_b_proj_verified_shapes ) self._use_min_latency_q_b_gemm: bool | None = None @@ -2075,7 +2093,7 @@ class DeepseekV2AttentionMLA( llama_4_scaling: Optional[torch.Tensor] = None, prev_topk_indices: Optional[torch.Tensor] = None, ): - if self.attn_mha.kv_b_proj is None: + if self.attn_mha.kv_b_proj is None and hasattr(self, "kv_b_proj"): self.attn_mha.kv_b_proj = self.kv_b_proj # when hidden_states is a tuple of tensors, the tuple will include quantized weight and scale tensor @@ -2215,6 +2233,14 @@ class DeepseekV2AttentionMLA( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch ): assert self.q_lora_rank is not None + if hasattr(self, "q_a_proj"): + return torch.cat( + ( + self.q_a_proj(hidden_states)[0], + self.kv_a_proj_with_mqa(hidden_states)[0], + ), + dim=-1, + ) if self._use_min_latency_fused_a_gemm is None: self._use_min_latency_fused_a_gemm = ( self.has_fused_proj diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 74eba43ef..247c05908 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -185,38 +185,22 @@ class MhcOps(NamedTuple): hc_split_sinkhorn: Callable[..., Any] mhc_fused_post_pre: Optional[Callable[..., Any]] npu_hc_pre: Optional[Callable[..., Any]] - mhc_pre: Optional[Callable[..., Any]] - mhc_post: Optional[Callable[..., Any]] - fused_hc_head: Optional[Callable[..., Any]] @functools.cache def _get_mhc_ops() -> MhcOps: """Load MHC kernels only when a DeepSeek-V4 layer needs them. - Model modules are imported eagerly by the registry. Importing + Model modules are imported eagerly by the registry. Importing ``sglang.kernels.ops.layernorm.mhc`` owns TileLang-backed MHC kernels. Import it only when a DeepSeek-V4 layer executes so registry discovery cannot initialize an optional CUDA runtime before unrelated models set up - their communication workspaces. DeepSeek-V4 is the sole consumer here. + their communication workspaces. DeepSeek-V4 is the sole consumer here. """ if _is_xpu: - from sgl_kernel import ( - fused_hc_head, - hc_post, - hc_split_sinkhorn, - mhc_fused_post_pre, - mhc_pre, - ) + from sgl_kernel import hc_split_sinkhorn - return MhcOps( - hc_split_sinkhorn=hc_split_sinkhorn, - mhc_fused_post_pre=mhc_fused_post_pre, - npu_hc_pre=None, - mhc_pre=mhc_pre, - mhc_post=hc_post, - fused_hc_head=fused_hc_head, - ) + return MhcOps(hc_split_sinkhorn, None, None) from sglang.kernels.ops.layernorm.mhc import ( hc_split_sinkhorn, @@ -224,14 +208,7 @@ def _get_mhc_ops() -> MhcOps: npu_hc_pre, ) - return MhcOps( - hc_split_sinkhorn=hc_split_sinkhorn, - mhc_fused_post_pre=mhc_fused_post_pre, - npu_hc_pre=npu_hc_pre, - mhc_pre=None, - mhc_post=None, - fused_hc_head=None, - ) + return MhcOps(hc_split_sinkhorn, mhc_fused_post_pre, npu_hc_pre) logger = logging.getLogger(__name__) @@ -245,13 +222,6 @@ DEEPSEEK_V4_STACKED_PARAMS_MAPPING: List[Tuple[str, str, int]] = [ ] -def _is_fused_mhc_post_pre_enabled_xpu() -> bool: - if _is_xpu: - return envs.SGLANG_OPT_FUSE_MHC_POST_PRE.get() - - return False - - # FlashInfer's mhc_pre_big_fuse only accepts these split-K counts. _FLASHINFER_MHC_PRE_SPLITS = (1, 2, 4, 8, 16) @@ -506,6 +476,31 @@ def _freqs_cis_to_cos_sin( return cos, sin +def _apply_gguf_grouped_wo_a( + o: torch.Tensor, + qweight: torch.Tensor, + qweight_type: int, + o_lora_rank: int, + matmul_fn: Optional[Callable] = None, +) -> torch.Tensor: + if matmul_fn is None: + from sglang.srt.layers.quantization.gguf import fused_mul_mat_gguf + + matmul_fn = fused_mul_mat_gguf + + group_outputs = [] + for group_id in range(o.shape[1]): + start = group_id * o_lora_rank + group_outputs.append( + matmul_fn( + o[:, group_id, :].contiguous(), + qweight[start : start + o_lora_rank], + qweight_type, + ) + ) + return torch.stack(group_outputs, dim=1) + + if TYPE_CHECKING: from sglang.srt.layers.attention.deepseek_v4_backend import ( DeepseekV4AttnBackend, @@ -638,8 +633,11 @@ class MqaAttentionBase(nn.Module): else wo_b_reduce_results ) if wo_a_keeps_quant_config is None: + keep_source_quant = ( + quant_config is not None and quant_config.get_name() == "expert_pack" + ) wo_a_quant_config: Optional[QuantizationConfig] = ( - quant_config if fp8 else None + quant_config if fp8 or keep_source_quant else None ) elif wo_a_keeps_quant_config: wo_a_quant_config = quant_config @@ -857,6 +855,11 @@ class MQALayer(MqaAttentionBase): self.compressor = None self.indexer = None if self.compress_ratio in (4, 128): + expert_pack_quant_config = ( + quant_config + if quant_config is not None and quant_config.get_name() == "expert_pack" + else None + ) self.compressor = Compressor( config, layer_id=self.layer_id, @@ -866,6 +869,7 @@ class MQALayer(MqaAttentionBase): head_dim=self.head_dim, rotate=False, prefix=add_prefix("compressor", prefix), + quant_config=expert_pack_quant_config, rotary_emb=self.rotary_emb, ) if self.compress_ratio == 4: @@ -1723,10 +1727,19 @@ class MQALayer(MqaAttentionBase): ) o = output else: - wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) - o = _apply_wo_a_bf16_matmul( - o, wo_a, is_decode=forward_batch.forward_mode.is_decode() - ) + wo_a_weight = getattr(self.wo_a, "weight", None) + if wo_a_weight is not None: + wo_a = wo_a_weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = _apply_wo_a_bf16_matmul( + o, wo_a, is_decode=forward_batch.forward_mode.is_decode() + ) + else: + o = _apply_gguf_grouped_wo_a( + o, + self.wo_a.qweight, + self.wo_a.qweight_type.weight_type, + self.o_lora_rank, + ) o, _ = self.wo_b(o.flatten(1)) if self.attn_tp_size > 1 and self.attn_tp_size < get_parallel().tp_size: @@ -1814,9 +1827,7 @@ class DeepseekV4DecoderLayer(nn.Module): ) = make_hc_mixing_params(hc_mult, config.hidden_size) self.rms_norm_eps = config.rms_norm_eps self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() - self.use_fused_mhc_post_pre = ( - is_cross_layer_mhc_fusion_enabled() or _is_fused_mhc_post_pre_enabled_xpu() - ) + self.use_fused_mhc_post_pre = is_cross_layer_mhc_fusion_enabled() self._input_layernorm_weight_bf16 = None self._post_attention_layernorm_weight_bf16 = None @@ -1892,26 +1903,6 @@ class DeepseekV4DecoderLayer(nn.Module): ) return y, post, comb, False - if _is_xpu: - norm_kwargs = {} - if norm is not None: - norm_kwargs["norm_weight"] = norm.weight.data - norm_kwargs["norm_eps"] = norm.variance_epsilon - - post, comb, y = _get_mhc_ops().mhc_pre( - residual=x, - fn=hc_fn, - hc_scale=hc_scale, - hc_base=hc_base, - rms_eps=self.rms_norm_eps, - hc_pre_eps=self.hc_eps, - hc_sinkhorn_eps=self.hc_eps, - hc_post_mult_value=_MHC_POST_MULT_VALUE, - sinkhorn_repeat=self.hc_sinkhorn_iters, - **norm_kwargs, - ) - return y, post, comb, norm is not None - if envs.SGLANG_OPT_USE_FLASHINFER_MHC.get(): y, post, comb = _flashinfer_hc_pre( x, @@ -2017,9 +2008,6 @@ class DeepseekV4DecoderLayer(nn.Module): if _is_npu: return torch.ops.custom.npu_hc_post(x, residual, post, comb) - if _is_xpu: - return _get_mhc_ops().mhc_post(x, residual, post, comb) - if envs.SGLANG_OPT_USE_FLASHINFER_MHC.get(): from flashinfer.mhc import mhc_post @@ -2726,10 +2714,17 @@ class DeepseekV4Model(nn.Module): self.pp_group = get_pp_group() self.hidden_size = config.hidden_size if self.pp_group.is_first_rank: + embedding_quant_config = ( + quant_config + if quant_config is not None and quant_config.get_name() == "expert_pack" + else None + ) self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, enable_tp=not is_dp_attention_enabled(), + quant_config=embedding_quant_config, + prefix=add_prefix("embed_tokens", prefix), ) else: self.embed_tokens = PPMissingLayer() @@ -2798,15 +2793,6 @@ class DeepseekV4Model(nn.Module): hc_base: torch.Tensor, ): if x.numel() > 0: - if _is_xpu: - return _get_mhc_ops().fused_hc_head( - x.contiguous(), - hc_fn, - hc_scale, - hc_base, - norm_eps=self.norm_eps, - hc_eps=self.hc_eps, - ) from sglang.kernels.ops.layernorm.mhc_head import fused_hc_head return fused_hc_head( @@ -3441,10 +3427,10 @@ class DeepseekV4ForCausalLM(nn.Module): is_nextn: bool = False, num_hidden_layers: Optional[int] = None, ) -> str: - if name == "embed.weight": - return "model.embed_tokens.weight" - if name == "head.weight": - return "lm_head.weight" + if name.startswith("embed."): + return "model.embed_tokens." + name.removeprefix("embed.") + if name.startswith("head."): + return "lm_head." + name.removeprefix("head.") if name == "norm.weight": return "model.norm.weight" if name.startswith("hc_head_"): @@ -3508,7 +3494,7 @@ class DeepseekV4ForCausalLM(nn.Module): if self._mhc_prewarmed_at_load: return self._mhc_prewarmed_at_load = True - if _is_npu or _is_xpu or not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get(): + if _is_npu or not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get(): return layer = next( (m for m in self.model.layers if isinstance(m, DeepseekV4DecoderLayer)), @@ -3583,7 +3569,7 @@ class DeepseekV4ForCausalLM(nn.Module): raise ValueError("num_nextn_predict_layers is not in the config") if not _FP8_WO_A_GEMM: - weights = _dequant_fp8_wo_a_streaming(weights) + weights = _prepare_deepseek_v4_weights(weights, self.quant_config) stacked_params_mapping = DEEPSEEK_V4_STACKED_PARAMS_MAPPING @@ -3780,7 +3766,7 @@ class DeepseekV4ForCausalLM(nn.Module): or name == "lm_head.weight" ) and not self.pp_group.is_last_rank: continue - elif COMPRESSOR_PART in name: + elif COMPRESSOR_PART in name and ".wkv_gate." not in name: is_kv = name.endswith(".wkv.weight") is_wgate = name.endswith(".wgate.weight") assert is_kv != is_wgate @@ -3817,6 +3803,10 @@ class DeepseekV4ForCausalLM(nn.Module): or name.endswith(".wq_a.weight_scale_inv") or name.endswith(".wkv.weight") or name.endswith(".wkv.weight_scale_inv") + or name.endswith(".wq_a.qweight") + or name.endswith(".wkv.qweight") + or name.endswith(".wq_a.qweight_type") + or name.endswith(".wkv.qweight_type") ): is_q = ".wq_a." in name param_name = name.replace( @@ -3831,8 +3821,8 @@ class DeepseekV4ForCausalLM(nn.Module): loaded_weight ) if len(bucket) == 2: - fused_weight = torch.cat( - [bucket["q"], bucket["kv"]], dim=0 + fused_weight = _fuse_deepseek_v4_wqkv_a_pair( + param_name, bucket ) param = params_dict[param_name] weight_loader = auto_weight_loader(param) @@ -4027,3 +4017,32 @@ def _dequant_fp8_wo_a( yield name, _dequant_fp8(weight, scale) yield from weights_dict.items() + + +def _prepare_deepseek_v4_weights( + weights: Iterable[Tuple[str, torch.Tensor]], + quant_config: Optional[QuantizationConfig], +) -> Iterable[Tuple[str, torch.Tensor]]: + """Keep Expert Pack GGUF weights on the streaming load path.""" + + if quant_config is not None and quant_config.get_name() == "expert_pack": + logger.info("Keep Expert Pack GGUF weights on the streaming load path") + return weights + return _dequant_fp8_wo_a_streaming(weights) + + +def _fuse_deepseek_v4_wqkv_a_pair( + param_name: str, bucket: dict[str, torch.Tensor] +) -> torch.Tensor: + """Fuse Q/KV rows while preserving their common GGUF type scalar.""" + + q = bucket["q"] + kv = bucket["kv"] + if param_name.endswith(".qweight_type"): + if q.numel() != 1 or kv.numel() != 1 or q.item() != kv.item(): + raise ValueError( + f"cannot fuse different GGUF qweight types for {param_name}: " + f"q={q.tolist()} kv={kv.tolist()}" + ) + return q + return torch.cat([q, kv], dim=0) diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 45e83672d..3e0d30aba 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -100,6 +100,9 @@ from sglang.srt.model_loader.weight_utils import ( maybe_remap_kv_scale_name, sharded_weight_loader, ) +from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods import ( + AttnForwardMethod, +) from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, MoEGate from sglang.srt.models.kimi_k3_vl import ( KimiK3MultiModalProjector, @@ -664,7 +667,9 @@ class KimiK3MoE(nn.Module): self.fuse_ar_norm and self.tp_size == 8 and self.routed_expert_up_proj is not None - and isinstance(self.routed_expert_up_proj.weight, torch.Tensor) + and isinstance( + getattr(self.routed_expert_up_proj, "weight", None), torch.Tensor + ) and self.routed_expert_up_proj.weight.dtype == torch.bfloat16 and self.routed_expert_up_proj.weight.is_contiguous() ) @@ -705,6 +710,8 @@ class KimiK3MoE(nn.Module): mods = [self.gate, self.routed_expert_down_proj] else: return + if any(getattr(module, "weight", None) is None for module in mods): + return dtypes = {m.weight.dtype for m in mods} if len(dtypes) != 1 or dtypes.pop() not in (torch.bfloat16, torch.float16): return @@ -1712,6 +1719,8 @@ class KimiK3DeltaAttention(nn.Module): self._bfa_w = torch.cat(weights, dim=0).contiguous() self._bfa_f_b_w = _get_k3_dense_weight(self.f_b_proj).contiguous() else: + if any(getattr(mod, "weight", None) is None for mod in mods): + return self._bfa_w, sizes = _merge_weights_as_views(mods, pad_rows_to=8) self._bfa_f_b_w = self.f_b_proj.weight self._bfa_fa_size, self._bfa_b_size = sizes @@ -1891,6 +1900,9 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): alt_stream: Optional[torch.cuda.Stream] = None, gate_alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: + split_gguf_kv_b = getattr( + quant_config, "supports_kimi_k3_quantized_latent_projections", False + ) 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 @@ -1912,6 +1924,51 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): reduce_results=not self.all_reduce_fusion, alt_stream=alt_stream, ) + if split_gguf_kv_b: + del self.fused_qkv_a_proj_with_mqa + del self.kv_b_proj + self.q_a_proj = ReplicatedLinear( + config.hidden_size, + config.q_lora_rank, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_a_proj", + ) + self.kv_a_proj_with_mqa = ReplicatedLinear( + config.hidden_size, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + self.has_fused_proj = False + + from sglang.srt.layers.quantization.gguf import GGUFUninitializedParameter + + for role in ("k", "v"): + qweight = GGUFUninitializedParameter(requires_grad=False) + set_weight_attrs( + qweight, + { + "is_gguf_weight": True, + "weight_loader": self._split_kv_b_weight_loader, + }, + ) + self.register_parameter(f"{role}_b_qweight", qweight) + qweight_type = nn.Parameter( + torch.empty(1, dtype=torch.uint8), requires_grad=False + ) + set_weight_attrs( + qweight_type, + { + "is_gguf_weight_type": True, + "weight_type": 0, + "ignore_warning": True, + "weight_loader": self._split_kv_b_weight_loader, + }, + ) + self.register_parameter(f"{role}_b_qweight_type", qweight_type) + self._kimi_split_gguf_kv_b = True # Installed before the output-gate wrap below so the gate multiply is # applied to x before the fused GEMM+AR sees it. if self.all_reduce_fusion and not _o_proj_takes_output(self.o_proj): @@ -2016,6 +2073,24 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): self.o_proj.forward = _gated_o_proj_forward + @staticmethod + def _split_kv_b_weight_loader(param, loaded_weight) -> None: + from torch.nn.parameter import UninitializedParameter + + if getattr(param, "is_gguf_weight_type", False): + param.weight_type = int(loaded_weight.item()) + param.data.copy_(loaded_weight.reshape_as(param)) + return + if isinstance(param, UninitializedParameter): + param.materialize(tuple(loaded_weight.shape), dtype=loaded_weight.dtype) + param.data.copy_(loaded_weight) + + def dispatch_attn_forward_method(self, forward_batch) -> AttnForwardMethod: + method = super().dispatch_attn_forward_method(forward_batch) + if getattr(self, "_kimi_split_gguf_kv_b", False): + return AttnForwardMethod.MLA + return method + def _precompute_output_gate(self, hidden_states: torch.Tensor) -> None: """Issue the output-gate GEMM on the alt stream so it overlaps the attention core; the lazy path in the o_proj wrap otherwise computes @@ -2543,9 +2618,15 @@ class KimiK3LinearModel(nn.Module): self._trim_padded_attn = require_mlp_sync(get_server_args()) if self.pp_group.is_first_rank: + embedding_quant_config = ( + quant_config + if quant_config is not None and quant_config.get_name() == "expert_pack" + else None + ) self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, + quant_config=embedding_quant_config, prefix=f"{prefix}.embed_tokens", # Under DP attention each rank embeds only its local tokens: # reduce within the attention-TP group, not the full TP group. @@ -3053,6 +3134,7 @@ class KimiK3LinearForCausalLM(nn.Module): loaded_params.add(name) self.post_load_weights() + return loaded_params def post_load_weights(self): # Also invoked by loader post-load hooks (DummyModelLoader, @@ -3067,6 +3149,13 @@ class KimiK3LinearForCausalLM(nn.Module): if isinstance(layer, PPMissingLayer): continue self_attn = layer.self_attn + if getattr(self_attn, "_kimi_split_gguf_kv_b", False): + if int(self_attn.k_b_qweight_type.weight_type) != 2: + raise ValueError("Kimi-K3 MLA K projection must remain GGUF Q4_0") + if int(self_attn.v_b_qweight_type.weight_type) != 10: + raise ValueError("Kimi-K3 MLA V projection must remain GGUF Q2_K") + self_attn.use_deep_gemm_bmm = False + continue kv_b_weight = _get_k3_dense_weight(self_attn.kv_b_proj) w_kc, w_vc = kv_b_weight.unflatten( 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) @@ -3124,9 +3213,13 @@ class KimiK3LinearForCausalLM(nn.Module): precompile_k3_recompute_w_u_kernel, ) + o_proj_weight = getattr(layer.self_attn.o_proj, "weight", None) + if o_proj_weight is None: + o_proj_weight = layer.self_attn.o_proj.qweight if precompile_k3_recompute_w_u_kernel( num_heads=layer.self_attn.local_num_heads, - dtype=layer.self_attn.o_proj.params_dtype, + dtype=getattr(layer.self_attn.o_proj, "params_dtype", None) + or o_proj_weight.dtype, device=layer.self_attn.dt_bias.device, ): rank0_log("Precompiled the Kimi-K3 KDA prefill kernel.") @@ -3538,4 +3631,4 @@ class KimiK3ForConditionalGeneration(nn.Module): pass -EntryClass = [KimiK3ForConditionalGeneration] +EntryClass = [KimiK3ForConditionalGeneration, KimiK3LinearForCausalLM] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6cc87be0c..1ee1ca36d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -125,6 +125,12 @@ LOAD_FORMAT_CHOICES = [ "sharded_state", "presharded", "gguf", + # Experimental and intentionally narrow: expert_pack is validated only for + # DeepSeek-V4-Flash-0731 MXFP4 GGUF (MXFP4 experts, FP8 dense weights) + # and KIMI-K3-MXP4-DERISKED-Q2_K-*.gguf (Q2_K gate/up, Q3_K down weights): + # https://huggingface.co/unsloth/DeepSeek-V4-Flash-0731-GGUF + # https://huggingface.co/Blackfrost-AI/KIMI-K3-Q2_K-GGUF-ABLITERATED + "expert_pack", "bitsandbytes", "mistral", "layered", @@ -573,6 +579,9 @@ class ServerArgs: '"dummy" will initialize the weights with random values, ' "which is mainly for profiling." '"gguf" will load the weights in the gguf format. ' + '"expert_pack" is experimental and loads only the validated ' + "DeepSeek-V4-Flash-0731 MXFP4 or text-only Kimi-K3 Q2_K GGUF " + "model with routed experts stored in an SSD expert pack. " '"bitsandbytes" will load the weights using bitsandbytes ' "quantization." '"layered" loads weights layer by layer so that one can quantize a ' @@ -3839,6 +3848,11 @@ class ServerArgs: # Set missing default values. self._handle_missing_default_values() + # expert_pack may replace a raw GGUF input with its generated local + # model metadata before any model-specific handler calls get_model_config. + # It also establishes eager-only invariants before CUDA graph parsing. + self._handle_expert_pack() + # Validate PD disaggregation flags before CUDA graph config. self._handle_pd_disaggregation() @@ -8079,6 +8093,11 @@ class ServerArgs: speculative_draft_model_path=resolved_draft, ) + def _handle_expert_pack(self): + from sglang.srt.arg_groups.expert_pack_hook import handle_expert_pack + + handle_expert_pack(self) + def _handle_load_format(self): # The quantization side of the gguf coupling moved to the pipeline # (arg_groups/overrides.py: _gguf_quantization); load_format itself is diff --git a/test/registered/expert_pack/test_expert_pack_mxfp4.py b/test/registered/expert_pack/test_expert_pack_mxfp4.py new file mode 100644 index 000000000..8018b32b3 --- /dev/null +++ b/test/registered/expert_pack/test_expert_pack_mxfp4.py @@ -0,0 +1,360 @@ +"""CUDA unit tests for the MXFP4 expert-pack kernels.""" + +import unittest + +import torch + +from sglang.kernels.ops.moe.expert_pack_mxfp4 import ( + mxfp4_marlin_repack, + mxfp4_matvec, + mxfp4_matvec_dual, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=90, stage="base-b", runner_config="1-gpu-small") + + +_FP4_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + 0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) +_PACK_INDEX = (0, 2, 4, 6, 1, 3, 5, 7) + + +def _matvec_reference( + input_cpu, + cache_cpu, + slots_cpu, + role_offset, + input_size, + output_size, + records_per_input, +): + blocks = input_size // 32 + row_bytes = blocks * 17 + result = torch.zeros((slots_cpu.numel(), output_size), dtype=torch.float32) + for record, slot in enumerate(slots_cpu.tolist()): + input_row = record // records_per_input + for output_row in range(output_size): + total = 0.0 + row_offset = ( + int(slot) * cache_cpu.stride(0) + role_offset + output_row * row_bytes + ) + for block in range(blocks): + quant_offset = row_offset + block * 17 + scale = 2.0 ** (int(cache_cpu.view(-1)[quant_offset]) - 127) + block_sum = 0.0 + for index in range(16): + packed = int(cache_cpu.view(-1)[quant_offset + index + 1]) + block_sum += ( + float(input_cpu[input_row, block * 32 + index]) + * _FP4_VALUES[packed & 0xF] + ) + block_sum += ( + float(input_cpu[input_row, block * 32 + index + 16]) + * _FP4_VALUES[packed >> 4] + ) + total += block_sum * scale + result[record, output_row] = total + return result + + +def _fill_mxfp4_cache(slots, role_bytes, input_size, device): + blocks = input_size // 32 + cache_cpu = torch.zeros((slots, 3 * role_bytes), dtype=torch.uint8) + for slot in range(slots): + for role in range(3): + role_rows = role_bytes // (blocks * 17) + for row in range(role_rows): + for block in range(blocks): + offset = ( + slot * cache_cpu.stride(0) + + role * role_bytes + + row * blocks * 17 + + block * 17 + ) + cache_cpu.view(-1)[offset] = 126 + ((slot + role + row + block) % 3) + for byte in range(16): + low = (slot + role * 3 + row + block + byte) % 16 + high = (15 + slot + role + row + block - byte) % 16 + cache_cpu.view(-1)[offset + byte + 1] = low | (high << 4) + return cache_cpu.to(device=device) + + +def _load_raw_word(raw_cpu, slot, role_offset, row, blocks_per_row, packed_word): + block = packed_word // 4 + word_in_block = packed_word & 3 + offset = ( + slot * raw_cpu.stride(0) + + role_offset + + row * blocks_per_row * 17 + + block * 17 + + 1 + + word_in_block * 4 + ) + word = 0 + for byte in range(4): + word |= int(raw_cpu.view(-1)[offset + byte]) << (8 * byte) + return word + + +def _marlin_nibble(word, value_index): + return (word >> ((value_index & 7) * 4)) & 0xF + + +def _marlin_scale_perm(index): + local_perm = (0, 2, 1, 3) + interleaved = (index // 4) * 4 + local_perm[index & 3] + return ((interleaved & 7) * 8) + (interleaved >> 3) + + +def _repack_reference( + raw_cpu, source_slot, role_bytes, input_size, output_size, gate_up +): + blocks = input_size // 32 + total_words = (input_size // 16) * output_size * 2 + output = torch.empty(total_words, dtype=torch.int32) + tile_span = (output_size // 64) * 128 + rows_per_role = output_size // 2 if gate_up else output_size + for index in range(total_words): + tile_k, tile_rem = divmod(index, tile_span) + tile_n, local = divmod(tile_rem, 128) + warp, thread = local & 3, local >> 2 + cur_n = warp * 16 + thread // 4 + tc_row = (thread & 3) * 2 + values = [] + for high in (False, True): + source_row = tile_n * 64 + cur_n + (8 if high else 0) + role = ( + 1 if gate_up and source_row >= rows_per_role else (0 if gate_up else 2) + ) + row = source_row % rows_per_role if gate_up else source_row + for offset in (0, 1, 8, 9): + value_index = tc_row + offset + word = _load_raw_word( + raw_cpu, + source_slot, + role * role_bytes, + row, + blocks, + tile_k * 2 + value_index // 8, + ) + values.append(_marlin_nibble(word, value_index)) + packed = 0 + for output_index, value_index in enumerate(_PACK_INDEX): + packed |= values[value_index] << (output_index * 4) + # The CUDA kernel stores the packed uint32 bit pattern in int32. + output[index] = packed if packed < (1 << 31) else packed - (1 << 32) + return output + + +def _scale_reference( + raw_cpu, source_slot, role_bytes, input_size, output_size, gate_up +): + blocks = input_size // 32 + output = torch.empty(blocks * output_size, dtype=torch.uint8) + rows_per_role = output_size // 2 if gate_up else output_size + for index in range(output.numel()): + group, column = divmod(index, output_size) + source_column = (column // 64) * 64 + _marlin_scale_perm(column & 63) + role = ( + 1 if gate_up and source_column >= rows_per_role else (0 if gate_up else 2) + ) + row = source_column % rows_per_role if gate_up else source_column + offset = ( + source_slot * raw_cpu.stride(0) + + role * role_bytes + + row * blocks * 17 + + group * 17 + ) + output[index] = raw_cpu.view(-1)[offset] + return output + + +@unittest.skipUnless(torch.cuda.is_available(), "MXFP4 kernel tests require CUDA") +class TestExpertPackMxfp4(unittest.TestCase): + def test_matvec_fp16_and_bf16(self): + input_size, output_size, records_per_input = 64, 5, 2 + blocks = input_size // 32 + role_bytes = output_size * blocks * 17 + cache = _fill_mxfp4_cache(2, role_bytes, input_size, "cuda") + slots = torch.tensor([1, 0, 1, 0], dtype=torch.int32, device="cuda") + for dtype in (torch.float16, torch.bfloat16): + if dtype is torch.bfloat16 and not torch.cuda.is_bf16_supported(): + continue + input_tensor = torch.arange( + 2 * input_size, dtype=torch.float32, device="cuda" + ).reshape(2, input_size) + input_tensor = ((input_tensor % 19) - 9).to(dtype) + output = mxfp4_matvec( + input_tensor, + cache, + slots, + role_offset=role_bytes, + role_bytes=role_bytes, + input_size=input_size, + output_size=output_size, + records_per_input=records_per_input, + ) + reference = _matvec_reference( + input_tensor.cpu(), + cache.cpu(), + slots.cpu(), + role_bytes, + input_size, + output_size, + records_per_input, + ).to(dtype) + torch.testing.assert_close(output.cpu(), reference, rtol=0.03, atol=0.25) + + def test_matvec_dual_matches_two_roles(self): + input_size, output_size, records_per_input = 64, 17, 2 + blocks = input_size // 32 + role_bytes = output_size * blocks * 17 + cache = _fill_mxfp4_cache(2, role_bytes, input_size, "cuda") + slots = torch.tensor([0, 1, 1, 0], dtype=torch.int32, device="cuda") + input_tensor = (torch.randn(2, input_size, device="cuda") * 0.5).to( + torch.float16 + ) + output_a, output_b = mxfp4_matvec_dual( + input_tensor, + cache, + slots, + gate_role_offset=0, + up_role_offset=role_bytes, + role_bytes=role_bytes, + input_size=input_size, + output_size=output_size, + records_per_input=records_per_input, + ) + expected_a = mxfp4_matvec( + input_tensor, + cache, + slots, + role_offset=0, + role_bytes=role_bytes, + input_size=input_size, + output_size=output_size, + records_per_input=records_per_input, + ) + expected_b = mxfp4_matvec( + input_tensor, + cache, + slots, + role_offset=role_bytes, + role_bytes=role_bytes, + input_size=input_size, + output_size=output_size, + records_per_input=records_per_input, + ) + torch.testing.assert_close(output_a, expected_a) + torch.testing.assert_close(output_b, expected_b) + + def test_marlin_repack_weights_and_scales(self): + hidden_size, intermediate_size = 64, 32 + w13_n, w2_n = 2 * intermediate_size, hidden_size + role_bytes = intermediate_size * (hidden_size // 32) * 17 + raw_cpu = torch.zeros((4, 3 * role_bytes), dtype=torch.uint8) + for slot in range(4): + for role, rows, size in ( + (0, intermediate_size, hidden_size), + (1, intermediate_size, hidden_size), + (2, hidden_size, intermediate_size), + ): + groups = size // 32 + for row in range(rows): + for group in range(groups): + offset = ( + slot * raw_cpu.stride(0) + + role * role_bytes + + row * groups * 17 + + group * 17 + ) + raw_cpu.view(-1)[offset] = ( + 100 + slot * 7 + role * 11 + row + group + ) % 256 + for byte in range(16): + raw_cpu.view(-1)[offset + byte + 1] = ( + (byte + row + role) % 16 + ) | (((15 - byte + slot + group) % 16) << 4) + raw = raw_cpu.cuda() + source_slots = torch.tensor([1, 0], dtype=torch.int32, device="cuda") + target_slots = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + w13_words = (hidden_size // 16) * w13_n * 2 + w2_words = (intermediate_size // 16) * w2_n * 2 + w13_scales = (hidden_size // 32) * w13_n + w2_scales = (intermediate_size // 32) * w2_n + w13 = torch.full((4, w13_words), -1, dtype=torch.int32, device="cuda") + w2 = torch.full((4, w2_words), -1, dtype=torch.int32, device="cuda") + w13_scale = torch.full((4, w13_scales), 255, dtype=torch.uint8, device="cuda") + w2_scale = torch.full((4, w2_scales), 255, dtype=torch.uint8, device="cuda") + mxfp4_marlin_repack( + raw, + source_slots, + target_slots, + role_bytes=role_bytes, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + ) + torch.cuda.synchronize() + for source_slot, target_slot in zip( + source_slots.cpu().tolist(), target_slots.cpu().tolist() + ): + expected_w13 = _repack_reference( + raw_cpu, source_slot, role_bytes, hidden_size, w13_n, True + ) + expected_w2 = _repack_reference( + raw_cpu, source_slot, role_bytes, intermediate_size, w2_n, False + ) + expected_w13_scale = _scale_reference( + raw_cpu, source_slot, role_bytes, hidden_size, w13_n, True + ) + expected_w2_scale = _scale_reference( + raw_cpu, source_slot, role_bytes, intermediate_size, w2_n, False + ) + torch.testing.assert_close(w13[target_slot].cpu(), expected_w13) + torch.testing.assert_close(w2[target_slot].cpu(), expected_w2) + torch.testing.assert_close(w13_scale[target_slot].cpu(), expected_w13_scale) + torch.testing.assert_close(w2_scale[target_slot].cpu(), expected_w2_scale) + self.assertTrue(torch.all(w13[0] == -1)) + self.assertTrue(torch.all(w2[0] == -1)) + self.assertTrue(torch.all(w13_scale[0] == 255)) + self.assertTrue(torch.all(w2_scale[0] == 255)) + + def test_invalid_dimensions_are_rejected(self): + input_tensor = torch.zeros((1, 31), dtype=torch.float16, device="cuda") + cache = torch.zeros((1, 17), dtype=torch.uint8, device="cuda") + slots = torch.zeros((1,), dtype=torch.int32, device="cuda") + with self.assertRaisesRegex(RuntimeError, "divisible by 32"): + mxfp4_matvec( + input_tensor, + cache, + slots, + role_offset=0, + role_bytes=17, + input_size=31, + output_size=1, + records_per_input=1, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/expert_pack/test_expert_pack_runtime.py b/test/registered/expert_pack/test_expert_pack_runtime.py new file mode 100644 index 000000000..7b3ae20a8 --- /dev/null +++ b/test/registered/expert_pack/test_expert_pack_runtime.py @@ -0,0 +1,444 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import sys +import tempfile +import unittest +from dataclasses import replace +from pathlib import Path + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + +ROOT = Path(__file__).resolve().parents[3] +TOOLS = ROOT / "tools" / "expert_pack" +sys.path.insert(0, str(TOOLS)) + +from format import ( # noqa: E402 + ENTRY_STRUCT, + FLAG_IDENTITY_PAYLOAD, + FLAG_TRIPLET_OBJECTS, + HEADER_STRUCT, + IndexEntry, + PackHeader, + align_up, + read_header, + read_index, +) + +from sglang.srt.layers.moe.expert_pack import ( # noqa: E402 + ExpertPackStore, + _CacheSlot, +) + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _make_pack(directory: Path) -> tuple[Path, Path, dict[str, str]]: + layers, experts, top_k = 1, 2, 1 + role_bytes = 17 + object_stride = 4096 + index_count = layers * experts * 3 + data_start = align_up(HEADER_STRUCT.size + index_count * ENTRY_STRUCT.size, 4096) + digests = { + "model_identity": hashlib.sha256(b"model-identity").hexdigest(), + "source": hashlib.sha256(b"source").hexdigest(), + "config": hashlib.sha256(b"config").hexdigest(), + } + header = PackHeader( + flags=FLAG_IDENTITY_PAYLOAD | FLAG_TRIPLET_OBJECTS, + index_count=index_count, + data_start=data_start, + alignment=4096, + num_layers=layers, + num_experts=experts, + top_k=top_k, + role_count=3, + model_identity_sha256=digests["model_identity"], + source_blob_sha256=digests["source"], + config_sha256=digests["config"], + ) + entries = [] + payloads = [] + for expert in range(experts): + object_offset = data_start + expert * object_stride + generation = expert + 100 + for role_id, role in enumerate(("gate", "up", "down")): + payload = bytes([expert * 3 + role_id]) * role_bytes + payload_hash = hashlib.sha256(payload).hexdigest() + entries.append( + IndexEntry( + layer=0, + expert=expert, + role=role, + dtype_id=39, + dtype="MXFP4", + tensor_name=f"blk.0.ffn_{role}_exps.weight", + source_tensor_offset=0, + source_tensor_nbytes=role_bytes * experts, + source_slice_offset=expert * role_bytes, + source_slice_nbytes=role_bytes, + pack_offset=object_offset + role_id * role_bytes, + pack_nbytes=role_bytes, + source_tensor_sha256=payload_hash, + source_slice_sha256=payload_hash, + checksum=payload_hash, + shape=(32, 1), + quant_scheme="MXFP4", + transform_id="identity-v1", + block_size=32, + generation=generation, + ) + ) + payloads.append((object_offset + role_id * role_bytes, payload)) + + pack = directory / "runtime.expert-pack" + with pack.open("w+b") as stream: + stream.write(header.pack()) + for entry in entries: + stream.write(entry.pack()) + stream.truncate(data_start + experts * object_stride) + for offset, payload in payloads: + stream.seek(offset) + stream.write(payload) + manifest = directory / "runtime.expert-pack.manifest.json" + manifest.write_text( + json.dumps( + { + "complete": True, + "object_stride": object_stride, + "pack_sha256": _sha256(pack), + } + ), + encoding="utf-8", + ) + return pack, manifest, digests + + +def _refresh_manifest_hash(pack: Path, manifest: Path) -> None: + value = json.loads(manifest.read_text(encoding="utf-8")) + value["pack_sha256"] = _sha256(pack) + manifest.write_text(json.dumps(value), encoding="utf-8") + + +class TestExpertPackRuntime(unittest.TestCase): + def test_victim_prefers_low_frequency_then_lru(self): + store = object.__new__(ExpertPackStore) + keys = [(0, 0), (0, 1), (0, 2)] + store._cache_slots = [ + _CacheSlot(key=keys[0], frequency=5), + _CacheSlot(key=keys[1], frequency=1), + _CacheSlot(key=keys[2], frequency=1), + ] + store._key_to_slot = {key: index for index, key in enumerate(keys)} + store._lru = dict.fromkeys(keys) + + self.assertEqual(store._victim_slot(set()), 1) + self.assertEqual(store._victim_slot({keys[1]}), 2) + self.assertEqual(store._victim_slot(set(), preserve_oldest=True), 2) + with self.assertRaisesRegex(RuntimeError, "active top-k"): + store._victim_slot(set(keys)) + + def test_zero_staging_slots_are_rejected(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + with self.assertRaisesRegex(ValueError, "staging budgets"): + ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=0, + ) + + def test_full_pack_verification_is_opt_in(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + store = ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + ) + self.assertEqual(len(store.entries), 6) + self.assertEqual(store.object_payload_bytes, 51) + staging = torch.empty(51, dtype=torch.uint8) + read_bytes, elapsed_ns = store._read_object(0, 1, staging) + self.assertEqual(read_bytes, 51) + self.assertGreaterEqual(elapsed_ns, 0) + self.assertEqual(staging.tolist(), [3] * 17 + [4] * 17 + [5] * 17) + store.close() + + with pack.open("r+b") as stream: + stream.seek(-1, 2) + stream.write(b"\x01") + + store = ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + ) + store.close() + + with self.assertRaisesRegex(ValueError, "SHA-256"): + ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + verify_pack_sha256=True, + ) + + def test_split_read_ranges_reconstruct_exact_object(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + store = ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + ) + ranges = store._object_read_ranges() + self.assertEqual(len(ranges), 4) + self.assertEqual(ranges[0][0], 0) + self.assertEqual(sum(length for _, length in ranges), 51) + self.assertTrue(all(length > 0 for _, length in ranges)) + self.assertTrue( + all( + ranges[index][0] + ranges[index][1] == ranges[index + 1][0] + for index in range(len(ranges) - 1) + ) + ) + + staging = torch.empty(51, dtype=torch.uint8) + for start, length in ranges: + read_bytes, elapsed_ns = store._read_object_range( + 0, 1, staging, start=start, length=length + ) + self.assertEqual(read_bytes, length) + self.assertGreaterEqual(elapsed_ns, 0) + self.assertEqual(staging.tolist(), [3] * 17 + [4] * 17 + [5] * 17) + store.close() + + def test_read_splits_follow_runtime_configuration(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + store = ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + read_splits=2, + stats_flush_interval=7, + ) + ranges = store._object_read_ranges() + self.assertEqual(store.stats["read_splits"], 2) + self.assertEqual(store.stats_flush_interval, 7) + self.assertEqual(len(ranges), 2) + self.assertEqual(ranges[0][0], 0) + self.assertEqual(sum(length for _, length in ranges), 51) + store.close() + + def test_short_object_read_raises(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + store = ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + ) + second_object = store.object_offsets[(0, 1)] + with pack.open("r+b", buffering=0) as stream: + stream.truncate(second_object + 10) + staging = torch.empty(store.object_payload_bytes, dtype=torch.uint8) + with self.assertRaisesRegex(OSError, "short expert-pack read"): + store._read_object(0, 1, staging) + store.close() + + def test_duplicate_entry_is_rejected_even_with_valid_pack_hash(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + with pack.open("r+b", buffering=0) as stream: + header = read_header(stream) + entries = read_index(stream, header) + duplicate = replace( + entries[1], + layer=entries[0].layer, + expert=entries[0].expert, + role=entries[0].role, + ) + stream.seek(HEADER_STRUCT.size + ENTRY_STRUCT.size) + stream.write(duplicate.pack()) + _refresh_manifest_hash(pack, manifest) + with self.assertRaisesRegex(ValueError, "duplicate expert-pack entry"): + ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + ) + + def test_out_of_range_entry_is_rejected_even_with_valid_pack_hash(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + with pack.open("r+b", buffering=0) as stream: + header = read_header(stream) + entries = read_index(stream, header) + invalid = replace( + entries[0], pack_offset=pack.stat().st_size + header.alignment + ) + stream.seek(HEADER_STRUCT.size) + stream.write(invalid.pack()) + _refresh_manifest_hash(pack, manifest) + with self.assertRaisesRegex(ValueError, "object layout mismatch"): + ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + ) + + def test_source_identity_mismatch_is_rejected(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + with self.assertRaisesRegex(ValueError, "source_blob_sha256"): + ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256="0" * 64, + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + ) + + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA pinned memory") + def test_direct_io_rejects_unaligned_object_ranges(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + store = ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + cache_vram_reserve_mib=1, + stage_slots=1, + direct_io=True, + ) + with self.assertRaisesRegex(ValueError, "aligned read ranges"): + store.initialize_device_cache("cuda") + store.close() + + def test_close_flushes_stats_atomically(self): + with tempfile.TemporaryDirectory() as value: + root = Path(value) + pack, manifest, digests = _make_pack(root) + stats_path = root / "stats.json" + store = ExpertPackStore( + pack, + manifest_path=manifest, + expected_layers=1, + expected_experts=2, + expected_top_k=1, + expected_source_sha256=digests["source"], + expected_model_identity_sha256=digests["model_identity"], + expected_config_sha256=digests["config"], + cache_vram_mib=1, + stage_slots=1, + stats_path=stats_path, + ) + store._route_calls_by_layer[0] = 2 + store._route_tokens_by_layer[0] = 3 + store.stats["pack_reads"] = 4 + store.stats["pack_read_bytes"] = 204 + store.close() + stats = json.loads(stats_path.read_text(encoding="utf-8")) + self.assertEqual(stats["pack_reads"], 4) + self.assertEqual(stats["pack_read_bytes"], 204) + self.assertEqual(stats["route_calls_by_layer"], [2]) + self.assertEqual(stats["route_tokens_by_layer"], [3]) + self.assertFalse(list(root.glob("stats.json.*.tmp"))) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/expert_pack/test_kimi_k3_gguf.py b/test/registered/expert_pack/test_kimi_k3_gguf.py new file mode 100644 index 000000000..ce530b4c9 --- /dev/null +++ b/test/registered/expert_pack/test_kimi_k3_gguf.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.layers.quantization.gguf import ( + GGUFLinearMethod, + _ordered_gguf_shard_ids, +) +from sglang.srt.model_loader.kimi_k3_gguf import ( + _kda_a_log_target_value, + _residual_target_value, + kimi_k3_checkpoint_targets, + routed_expert_tensor, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class TestKimiK3GGUFMapping(unittest.TestCase): + def test_maps_dense_kda_mla_moe_and_residual_tensors(self) -> None: + cases = { + "token_embd.weight": ("model.embed_tokens.weight",), + "blk.0.ffn_gate.weight": ("model.layers.0.mlp.gate_proj.weight",), + "blk.1.ssm_g.weight": ("model.layers.1.self_attn.g_proj.weight",), + "blk.3.attn_q_b.weight": ("model.layers.3.self_attn.q_b_proj.weight",), + "blk.3.attn_k_b.weight": ("model.layers.3.self_attn.k_b_qweight",), + "blk.3.attn_v_b.weight": ("model.layers.3.self_attn.v_b_qweight",), + "blk.1.ffn_routed_down.weight": ( + "model.layers.1.mlp.routed_expert_down_proj.weight", + ), + "blk.1.ffn_gate_shexp.weight": ( + "model.layers.1.mlp.shared_experts.gate_proj.weight", + ), + "blk.2.attn_res_score.weight": ( + "model.layers.2.self_attention_res_proj.weight", + "model.layers.2.self_attention_res_norm.weight", + ), + "output_res_score.weight": ( + "model.output_attn_res_proj.weight", + "model.output_attn_res_norm.weight", + ), + } + for source, expected in cases.items(): + with self.subTest(source=source): + self.assertEqual(kimi_k3_checkpoint_targets(source), expected) + + def test_only_routed_aggregate_tensors_are_skipped(self) -> None: + self.assertTrue(routed_expert_tensor("blk.92.ffn_up_exps.weight")) + self.assertTrue(routed_expert_tensor("blk.1.ffn_down_exps.weight")) + self.assertFalse(routed_expert_tensor("blk.1.ffn_up_shexp.weight")) + self.assertFalse(routed_expert_tensor("blk.1.ffn_routed_up.weight")) + + def test_unknown_tensor_fails_closed(self) -> None: + with self.assertRaisesRegex(KeyError, "unsupported Kimi-K3"): + kimi_k3_checkpoint_targets("blk.7.unexpected.weight") + + def test_residual_score_preserves_exact_combined_weight(self) -> None: + source = torch.tensor([0.5, -1.25, 3.0], dtype=torch.float32) + projection = _residual_target_value(source, 0) + norm = _residual_target_value(source, 1) + self.assertEqual(tuple(projection.shape), (1, 3)) + self.assertEqual(tuple(norm.shape), (3,)) + torch.testing.assert_close(norm * projection.squeeze(0), source) + + def test_residual_score_rejects_non_vector_source(self) -> None: + with self.assertRaisesRegex(ValueError, "must be a vector"): + _residual_target_value(torch.ones(1, 3), 0) + + def test_restores_kda_a_log_from_gguf_transform(self) -> None: + original = torch.tensor([-0.75, 0.0, 1.5], dtype=torch.float32) + stored = -torch.exp(original) + torch.testing.assert_close(_kda_a_log_target_value(stored), original) + with self.assertRaisesRegex(ValueError, "only -exp"): + _kda_a_log_target_value(torch.tensor([-1.0, 0.0])) + with self.assertRaisesRegex(ValueError, "finite"): + _kda_a_log_target_value(torch.tensor([-1.0, float("nan")])) + + def test_merged_gguf_output_uses_logical_shard_order(self) -> None: + qweight = torch.tensor([[30], [0], [20], [10]], dtype=torch.uint8) + qweight.shard_id = [3, 0, 2, 1] + qweight.shard_offset_map = { + 3: (0, 1, 1), + 0: (1, 2, 1), + 2: (2, 3, 1), + 1: (3, 4, 1), + } + qweight.gguf_prefix = "" + layer = SimpleNamespace( + qweight=qweight, + qweight_type=SimpleNamespace(shard_weight_type={0: 0, 1: 0, 2: 0, 3: 0}), + ) + method = object.__new__(GGUFLinearMethod) + + def fake_matmul(_x, weight, _weight_type): + return weight[:, 0].float().unsqueeze(0) + + with patch( + "sglang.srt.layers.quantization.gguf.fused_mul_mat_gguf", + side_effect=fake_matmul, + ): + output = method.apply(layer, torch.zeros(1, 1)) + torch.testing.assert_close(output, torch.tensor([[0.0, 10.0, 20.0, 30.0]])) + + def test_unknown_gguf_shard_layouts_preserve_checkpoint_order(self) -> None: + self.assertEqual(_ordered_gguf_shard_ids(["q", "k"]), ["q", "k"]) + self.assertEqual(_ordered_gguf_shard_ids([4, 2]), [4, 2]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/test_minicpm_sparse_metadata.py b/test/registered/unit/layers/test_minicpm_sparse_metadata.py index 6e501b1c9..80146bec5 100644 --- a/test/registered/unit/layers/test_minicpm_sparse_metadata.py +++ b/test/registered/unit/layers/test_minicpm_sparse_metadata.py @@ -1,7 +1,7 @@ import subprocess import sys import unittest -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import MagicMock, Mock, patch import torch @@ -279,7 +279,18 @@ class TestMiniCPMSparseMetadata(CustomTestCase): model_runner.server_args.attention_backend = "minicpm_flashinfer" flashinfer_adapter = object() + fake_fuse_kernel = ModuleType("sglang.srt.layers.attention.minicpm.fuse_kernel") + fake_fuse_kernel.fused_attn_pooling_online_topk_prefill = Mock( + return_value="prefill" + ) + fake_fuse_kernel.fused_attn_pooling_online_topk_decode = Mock( + return_value="decode" + ) with ( + patch.dict( + sys.modules, + {"sglang.srt.layers.attention.minicpm.fuse_kernel": fake_fuse_kernel}, + ), patch.object(backend_module, "MiniCPMHybridConfig", SimpleNamespace), patch.object(backend_module, "is_blackwell_supported", return_value=True), patch.object( @@ -297,11 +308,6 @@ class TestMiniCPMSparseMetadata(CustomTestCase): "get_parallel", return_value=SimpleNamespace(attn_tp_size=1), ), - patch( - "sglang.srt.layers.attention.minicpm.fuse_kernel." - "fused_attn_pooling_online_topk_prefill", - return_value="prefill", - ), patch.object(backend_module, "attach_compressed_cache"), ): backend = MiniCPMSparseBackend(model_runner, use_flashinfer=True) @@ -1079,10 +1085,16 @@ class TestMiniCPMSparseMetadata(CustomTestCase): backend._get_fused_topk_kernel.assert_called_once_with(1, is_prefill=False) def test_fused_topk_prefill_kernels_compile_for_all_batches_at_startup(self): - with patch( - "sglang.srt.layers.attention.minicpm.fuse_kernel." - "fused_attn_pooling_online_topk_prefill", - side_effect=lambda **kwargs: f"prefill-{kwargs['batch_size']}", + fake_fuse_kernel = ModuleType("sglang.srt.layers.attention.minicpm.fuse_kernel") + fake_fuse_kernel.fused_attn_pooling_online_topk_prefill = Mock( + side_effect=lambda **kwargs: f"prefill-{kwargs['batch_size']}" + ) + fake_fuse_kernel.fused_attn_pooling_online_topk_decode = Mock( + side_effect=lambda **kwargs: f"decode-{kwargs['batch_size']}" + ) + with patch.dict( + sys.modules, + {"sglang.srt.layers.attention.minicpm.fuse_kernel": fake_fuse_kernel}, ): backend, *_ = _construct_sparse_backend( max_running_requests=3, @@ -1104,17 +1116,14 @@ class TestMiniCPMSparseMetadata(CustomTestCase): backend.fused_kernel_kwargs = {"topk": 8} backend.prefill_kernel_max_seqlen_q_grid = 64 - with ( - patch( - "sglang.srt.layers.attention.minicpm.fuse_kernel." - "fused_attn_pooling_online_topk_prefill", - return_value="prefill", - ) as prefill, - patch( - "sglang.srt.layers.attention.minicpm.fuse_kernel." - "fused_attn_pooling_online_topk_decode", - return_value="decode", - ) as decode, + fake_fuse_kernel = ModuleType("sglang.srt.layers.attention.minicpm.fuse_kernel") + prefill = Mock(return_value="prefill") + decode = Mock(return_value="decode") + fake_fuse_kernel.fused_attn_pooling_online_topk_prefill = prefill + fake_fuse_kernel.fused_attn_pooling_online_topk_decode = decode + with patch.dict( + sys.modules, + {"sglang.srt.layers.attention.minicpm.fuse_kernel": fake_fuse_kernel}, ): self.assertEqual( backend._get_fused_topk_kernel(3, is_prefill=True), "prefill" diff --git a/test/registered/unit/lora/test_deepseek_mla_correction.py b/test/registered/unit/lora/test_deepseek_mla_correction.py new file mode 100644 index 000000000..7c20fe2c2 --- /dev/null +++ b/test/registered/unit/lora/test_deepseek_mla_correction.py @@ -0,0 +1,24 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.lora.deepseek_mla_correction import is_kv_b_lora_active +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestDeepseekMLACorrection(unittest.TestCase): + def test_kv_b_lora_probe(self): + self.assertFalse(is_kv_b_lora_active(SimpleNamespace())) + self.assertFalse( + is_kv_b_lora_active(SimpleNamespace(kv_b_proj=SimpleNamespace())) + ) + self.assertTrue( + is_kv_b_lora_active( + SimpleNamespace(kv_b_proj=SimpleNamespace(set_lora=True)) + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 5c3f38bcc..c7d748f6a 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -636,7 +636,10 @@ class TestGoldenModelOverrides(_IsolatedPublish): ) # Thor (SM110) and other architectures keep the existing auto behavior. - with patch.object(overrides_module, "is_sm100_supported", return_value=False): + with ( + patch.object(overrides_module, "is_sm100_supported", return_value=False), + patch.object(overrides_module, "is_sm120_supported", return_value=False), + ): non_sm10x = self._construct( "MiniMaxM2ForCausalLM", "llama", quantization="modelopt_fp4" ) diff --git a/test/registered/unit/test_supplied_instance_exposure_ratchet.py b/test/registered/unit/test_supplied_instance_exposure_ratchet.py index 333855ed4..a52354d79 100644 --- a/test/registered/unit/test_supplied_instance_exposure_ratchet.py +++ b/test/registered/unit/test_supplied_instance_exposure_ratchet.py @@ -139,7 +139,9 @@ _EXPOSED = { ("dllm/config.py", "model_path"), ("multimodal/processors/base_processor.py", "image_processor_backend"), ("speculative/spec_registry.py", "disable_overlap_schedule"), + ("disaggregation/encoder/server.py", "model_loader_extra_config"), ("layers/moe/utils.py", "deepep_mode"), + ("layers/moe/utils.py", "disable_shared_experts_fusion"), ("layers/moe/utils.py", "moe_a2a_backend"), ("layers/moe/utils.py", "moe_runner_backend"), ("layers/moe/utils.py", "quantization"), @@ -171,6 +173,8 @@ _EXPOSED = { ("layers/flashinfer_comm_fusion.py", "flashinfer_allreduce_fusion_backend"), ("lora/lora_manager.py", "enable_lora_overlap_loading"), ("lora/marlin_lora_temp/policy.py", "lora_paths"), + ("model_loader/expert_pack_runtime.py", "model_path"), + ("model_loader/expert_pack_runtime.py", "tokenizer_path"), ("parser/template_detection.py", "model_path"), ("speculative/adaptive_spec_params.py", "speculative_algorithm"), ("speculative/adaptive_spec_params.py", "speculative_eagle_topk"), @@ -189,6 +193,7 @@ _EXPOSED = { ("weight_cache/daemon.py", "enable_dp_lm_head"), ("weight_cache/daemon.py", "ep_size"), ("weight_cache/daemon.py", "load_format"), + ("weight_cache/daemon.py", "model_loader_extra_config"), ("weight_cache/daemon.py", "model_path"), ("weight_cache/daemon.py", "moe_a2a_backend"), ("weight_cache/daemon.py", "moe_dense_tp_size"), @@ -212,6 +217,7 @@ _OVERRIDDEN_AND_READ = { ("dllm/config.py", "model_path"), ("entrypoints/engine.py", "reasoning_parser"), ("entrypoints/engine.py", "tool_call_parser"), + ("model_loader/expert_pack_runtime.py", "model_path"), ("weight_cache/daemon.py", "dp_size"), ("weight_cache/daemon.py", "dtype"), ("weight_cache/daemon.py", "ep_size"), diff --git a/tools/expert_pack/__init__.py b/tools/expert_pack/__init__.py new file mode 100644 index 000000000..58eb7c0b4 --- /dev/null +++ b/tools/expert_pack/__init__.py @@ -0,0 +1 @@ +"""DeepSeek expert-pack build, inspection, and validation tools.""" diff --git a/tools/expert_pack/build.py b/tools/expert_pack/build.py new file mode 100755 index 000000000..a61cd1e66 --- /dev/null +++ b/tools/expert_pack/build.py @@ -0,0 +1,635 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shutil +import subprocess +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + +import gguf + +try: + from .format import ( + ENTRY_STRUCT, + FLAG_IDENTITY_PAYLOAD, + FLAG_TRIPLET_OBJECTS, + HEADER_STRUCT, + ROLE_NAMES, + IndexEntry, + PackHeader, + align_up, + inspect_pack, + sha256_file, + write_index, + ) +except ImportError: + from format import ( # type: ignore[no-redef] + ENTRY_STRUCT, + FLAG_IDENTITY_PAYLOAD, + FLAG_TRIPLET_OBJECTS, + HEADER_STRUCT, + ROLE_NAMES, + IndexEntry, + PackHeader, + align_up, + inspect_pack, + sha256_file, + write_index, + ) + + +FORMAT = "SGLANG-EXPERTPACK-v1" +EXPERT_RE = re.compile( + r"^blk\.(?P\d+)\.ffn_(?Pgate|up|down)_exps\.weight$" +) +COPY_CHUNK_BYTES = 16 * 1024 * 1024 + + +def now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def write_json_atomic(path: Path, value: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + ".tmp") + with temporary.open("w", encoding="utf-8", newline="\n") as stream: + json.dump(value, stream, ensure_ascii=False, indent=2, sort_keys=True) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + directory_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + + +def hash_range(stream, offset: int, nbytes: int) -> str: + digest = hashlib.sha256() + stream.seek(offset) + remaining = nbytes + while remaining: + chunk = stream.read(min(remaining, COPY_CHUNK_BYTES)) + if not chunk: + raise EOFError( + f"short read at source offset {offset}, {remaining} bytes remain" + ) + digest.update(chunk) + remaining -= len(chunk) + return digest.hexdigest() + + +def copy_range(source_fd: int, output, offset: int, nbytes: int) -> str: + digest = hashlib.sha256() + copied = 0 + while copied < nbytes: + chunk = os.pread( + source_fd, min(COPY_CHUNK_BYTES, nbytes - copied), offset + copied + ) + if not chunk: + raise EOFError(f"short read at source offset {offset + copied}") + output.write(chunk) + digest.update(chunk) + copied += len(chunk) + return digest.hexdigest() + + +def load_inventory( + path: Path, source: Path, expected_sha256: str +) -> tuple[dict, list[dict], dict]: + rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] + headers = [row for row in rows if row.get("kind") == "source"] + summaries = [row for row in rows if row.get("kind") == "summary"] + tensors = [row for row in rows if row.get("kind") == "tensor"] + if len(headers) != 1 or len(summaries) != 1: + raise ValueError( + "inventory must contain exactly one source and one summary record" + ) + header = headers[0] + if Path(header["path"]).resolve() != source.resolve(): + raise ValueError("inventory source path does not match --source") + if header.get("source_sha256") != expected_sha256 or not header.get( + "payload_hashes" + ): + raise ValueError( + "inventory is not a full-hash inventory for the requested source" + ) + if int(header["size"]) != source.stat().st_size: + raise ValueError("inventory source size does not match the source file") + if len(tensors) != int(header["tensor_count"]): + raise ValueError("inventory tensor count does not match its header") + if any(not row.get("sha256") for row in tensors): + raise ValueError("inventory contains a tensor without a payload hash") + return header, tensors, summaries[0] + + +def create_inventory(source: Path, source_sha256: str) -> tuple[dict, list[dict], dict]: + reader = gguf.GGUFReader(source, "r") + tensors = [] + inventory_digest = hashlib.sha256() + with source.open("rb", buffering=0) as stream: + for tensor in sorted(reader.tensors, key=lambda item: item.name): + record = { + "kind": "tensor", + "name": tensor.name, + "shape": [int(value) for value in tensor.shape.tolist()], + "type": tensor.tensor_type.name, + "type_id": int(tensor.tensor_type), + "offset": int(tensor.data_offset), + "nbytes": int(tensor.n_bytes), + "sha256": hash_range( + stream, int(tensor.data_offset), int(tensor.n_bytes) + ), + } + encoded = json.dumps(record, sort_keys=True).encode("utf-8") + b"\n" + inventory_digest.update(encoded) + tensors.append(record) + header = { + "kind": "source", + "path": str(source.resolve()), + "size": source.stat().st_size, + "source_sha256": source_sha256, + "gguf_data_offset": int(reader.data_offset), + "gguf_alignment": int(reader.alignment), + "tensor_count": len(tensors), + "metadata_count": len(reader.fields), + "payload_hashes": True, + } + summary = { + "kind": "summary", + "tensor_count": len(tensors), + "inventory_sha256": inventory_digest.hexdigest(), + } + return header, tensors, summary + + +def validate_inventory_against_reader(source: Path, tensors: list[dict]) -> None: + reader = gguf.GGUFReader(source, "r") + actual = { + tensor.name: { + "shape": [int(value) for value in tensor.shape.tolist()], + "type": tensor.tensor_type.name, + "type_id": int(tensor.tensor_type), + "offset": int(tensor.data_offset), + "nbytes": int(tensor.n_bytes), + } + for tensor in reader.tensors + } + recorded = {row["name"]: row for row in tensors} + if set(actual) != set(recorded): + raise ValueError("inventory tensor names do not match the GGUF reader") + for name, value in actual.items(): + if any(value[field] != recorded[name].get(field) for field in value): + raise ValueError(f"inventory metadata mismatch for tensor {name}") + ordered = sorted(tensors, key=lambda row: int(row["offset"])) + previous_end = 0 + for row in ordered: + offset = int(row["offset"]) + end = offset + int(row["nbytes"]) + if offset < previous_end or end > source.stat().st_size: + raise ValueError(f"invalid or overlapping source range for {row['name']}") + previous_end = end + + +def generation(model_digest: str, source_digest: str, layer: int, expert: int) -> int: + value = hashlib.sha256( + f"{model_digest}:{source_digest}:{layer}:{expert}".encode("ascii") + ).digest() + return int.from_bytes(value[:8], "little") or 1 + + +def tool_sha256() -> str: + digest = hashlib.sha256() + for path in (Path(__file__), Path(__file__).with_name("format.py")): + digest.update(path.name.encode("ascii") + b"\0") + digest.update(path.read_bytes()) + return digest.hexdigest() + + +def git_sha() -> str: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).resolve().parents[2], + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + +def build(args: argparse.Namespace) -> dict[str, object]: + started_at = now() + started_monotonic = time.monotonic() + source = args.source.resolve(strict=True) + output = args.output.resolve() + manifest_path = args.manifest.resolve() + checkpoint_path = args.checkpoint.resolve() + partial_path = output.with_name(output.name + ".partial") + output.parent.mkdir(parents=True, exist_ok=True) + manifest_path.parent.mkdir(parents=True, exist_ok=True) + + if output.exists(): + raise ValueError(f"completed output already exists: {output}") + if args.config_blob is not None: + if sha256_file(args.config_blob.resolve(strict=True)) != args.config_sha256: + raise ValueError("DeepSeek config hash does not match its digest") + + actual_source_sha256 = sha256_file(source) + if actual_source_sha256 != args.source_sha256: + raise ValueError( + f"source SHA-256 mismatch: expected {args.source_sha256}, got {actual_source_sha256}" + ) + + if args.inventory is None: + inventory_header, tensors, inventory_summary = create_inventory( + source, actual_source_sha256 + ) + else: + inventory_header, tensors, inventory_summary = load_inventory( + args.inventory.resolve(strict=True), source, actual_source_sha256 + ) + validate_inventory_against_reader(source, tensors) + + expert_tensors: dict[tuple[int, str], dict] = {} + for row in tensors: + match = EXPERT_RE.fullmatch(row["name"]) + if match is not None: + key = int(match.group("layer")), match.group("role") + if key in expert_tensors: + raise ValueError(f"duplicate routed-expert tensor {key}") + expert_tensors[key] = row + expected_tensor_keys = { + (layer, role) for layer in range(args.num_layers) for role in ROLE_NAMES + } + if set(expert_tensors) != expected_tensor_keys: + missing = sorted(expected_tensor_keys - set(expert_tensors)) + extra = sorted(set(expert_tensors) - expected_tensor_keys) + raise ValueError( + f"routed-expert tensor coverage mismatch: missing={missing[:8]} extra={extra[:8]}" + ) + + role_bytes = set() + for (layer, role), row in expert_tensors.items(): + shape = [int(value) for value in row["shape"]] + if len(shape) != 3 or shape[-1] != args.num_experts: + raise ValueError( + f"unexpected expert shape for layer={layer} role={role}: {shape}" + ) + if int(row["nbytes"]) % args.num_experts: + raise ValueError(f"expert tensor is not evenly sliceable: {row['name']}") + slice_bytes = int(row["nbytes"]) // args.num_experts + block_size, type_size = gguf.GGML_QUANT_SIZES[ + gguf.GGMLQuantizationType(int(row["type_id"])) + ] + logical_elements = 1 + for dimension in shape[:-1]: + logical_elements *= dimension + expected_bytes = logical_elements // block_size * type_size + if logical_elements % block_size or slice_bytes != expected_bytes: + raise ValueError( + f"quantized slice size mismatch for {row['name']}: {slice_bytes} != {expected_bytes}" + ) + role_bytes.add(slice_bytes) + if len(role_bytes) != 1: + raise ValueError( + f"triplet v1 requires uniform role sizes, got {sorted(role_bytes)}" + ) + role_nbytes = role_bytes.pop() + + object_count = args.num_layers * args.num_experts + object_payload_bytes = role_nbytes * len(ROLE_NAMES) + object_stride = align_up(object_payload_bytes, args.alignment) + index_count = object_count * len(ROLE_NAMES) + data_start = align_up( + HEADER_STRUCT.size + index_count * ENTRY_STRUCT.size, args.alignment + ) + expected_pack_bytes = data_start + object_count * object_stride + header = PackHeader( + flags=FLAG_IDENTITY_PAYLOAD | FLAG_TRIPLET_OBJECTS, + index_count=index_count, + data_start=data_start, + alignment=args.alignment, + num_layers=args.num_layers, + num_experts=args.num_experts, + top_k=args.top_k, + role_count=len(ROLE_NAMES), + model_identity_sha256=args.model_identity_sha256, + source_blob_sha256=actual_source_sha256, + config_sha256=args.config_sha256, + ) + header_raw = header.pack() + + existing_bytes = ( + partial_path.stat().st_size if args.resume and partial_path.exists() else 0 + ) + remaining_bytes = max(expected_pack_bytes - existing_bytes, 0) + free_bytes = shutil.disk_usage(output.parent).free + safety_bytes = int(args.safety_margin_gib * 1024**3) + if free_bytes < remaining_bytes + safety_bytes: + raise OSError( + f"insufficient free space: free={free_bytes}, remaining_pack={remaining_bytes}, " + f"safety={safety_bytes}" + ) + + if args.resume: + checkpoint = json.loads(checkpoint_path.read_text(encoding="utf-8")) + if checkpoint.get("status") != "in_progress": + raise ValueError("resume checkpoint is not in progress") + for field, expected in ( + ("source_sha256", actual_source_sha256), + ("model_identity_sha256", args.model_identity_sha256), + ("config_sha256", args.config_sha256), + ("tool_sha256", tool_sha256()), + ("expected_pack_bytes", expected_pack_bytes), + ): + if checkpoint.get(field) != expected: + raise ValueError(f"resume checkpoint {field} mismatch") + completed_layers = [int(value) for value in checkpoint["completed_layers"]] + if completed_layers != list(range(len(completed_layers))): + raise ValueError("completed layers in checkpoint are not a prefix") + entries = [IndexEntry.from_dict(value) for value in checkpoint["entries"]] + stream = partial_path.open("r+b", buffering=0) + if stream.read(len(header_raw)) != header_raw: + raise ValueError("partial pack header does not match the requested build") + pack_end = int(checkpoint["pack_end"]) + if partial_path.stat().st_size < pack_end: + raise ValueError("partial pack is shorter than its checkpoint") + stream.truncate(pack_end) + stream.seek(pack_end) + else: + if partial_path.exists() or checkpoint_path.exists() or manifest_path.exists(): + raise ValueError( + "build outputs already exist; use --resume for an in-progress build" + ) + completed_layers = [] + entries: list[IndexEntry] = [] + stream = partial_path.open("x+b", buffering=0) + stream.write(header_raw) + stream.truncate(data_start) + stream.seek(data_start) + checkpoint = { + "format": FORMAT + "-checkpoint", + "version": 1, + "status": "in_progress", + "started_at": started_at, + "source_sha256": actual_source_sha256, + "model_identity_sha256": args.model_identity_sha256, + "config_sha256": args.config_sha256, + "tool_sha256": tool_sha256(), + "expected_pack_bytes": expected_pack_bytes, + "completed_layers": [], + "pack_end": data_start, + "entries": [], + } + write_json_atomic(checkpoint_path, checkpoint) + + source_fd = os.open(source, os.O_RDONLY) + try: + for layer in range(len(completed_layers), args.num_layers): + layer_entries = [] + for expert in range(args.num_experts): + object_ordinal = layer * args.num_experts + expert + object_offset = data_start + object_ordinal * object_stride + if stream.tell() != object_offset: + raise ValueError( + f"pack cursor mismatch: {stream.tell()} != {object_offset}" + ) + object_generation = generation( + args.model_identity_sha256, actual_source_sha256, layer, expert + ) + for role in ROLE_NAMES: + tensor = expert_tensors[(layer, role)] + source_slice_offset = int(tensor["offset"]) + expert * role_nbytes + pack_offset = stream.tell() + slice_sha256 = copy_range( + source_fd, stream, source_slice_offset, role_nbytes + ) + dtype_id = int(tensor["type_id"]) + block_size = int( + gguf.GGML_QUANT_SIZES[gguf.GGMLQuantizationType(dtype_id)][0] + ) + entry = IndexEntry( + layer=layer, + expert=expert, + role=role, + dtype_id=dtype_id, + dtype=str(tensor["type"]), + tensor_name=str(tensor["name"]), + source_tensor_offset=int(tensor["offset"]), + source_tensor_nbytes=int(tensor["nbytes"]), + source_slice_offset=source_slice_offset, + source_slice_nbytes=role_nbytes, + pack_offset=pack_offset, + pack_nbytes=role_nbytes, + source_tensor_sha256=str(tensor["sha256"]), + source_slice_sha256=slice_sha256, + checksum=slice_sha256, + shape=tuple(int(value) for value in tensor["shape"][:-1]), + quant_scheme=str(tensor["type"]), + transform_id="identity-v1", + block_size=block_size, + generation=object_generation, + ) + entry.pack() + layer_entries.append(entry) + padding = object_stride - object_payload_bytes + if padding: + stream.write(bytes(padding)) + stream.flush() + os.fsync(stream.fileno()) + entries.extend(layer_entries) + checkpoint["completed_layers"].append(layer) + checkpoint["pack_end"] = stream.tell() + checkpoint["entries"] = [entry.to_dict() for entry in entries] + write_json_atomic(checkpoint_path, checkpoint) + print( + f"completed layer {layer}/{args.num_layers - 1}: pack_end={stream.tell()}", + file=sys.stderr, + flush=True, + ) + + if stream.tell() != expected_pack_bytes: + raise ValueError( + f"final pack size mismatch: {stream.tell()} != {expected_pack_bytes}" + ) + index_sha256 = write_index(stream, header, entries) + stream.flush() + os.fsync(stream.fileno()) + finally: + os.close(source_fd) + stream.close() + + pack_sha256 = sha256_file(partial_path) + reader = gguf.GGUFReader(source, "r") + with source.open("rb", buffering=0) as source_stream: + source_metadata_sha256 = hash_range(source_stream, 0, int(reader.data_offset)) + + routed_names = {row["name"] for row in expert_tensors.values()} + tensor_manifest = [] + for tensor in sorted(tensors, key=lambda row: row["name"]): + routed = tensor["name"] in routed_names + tensor_manifest.append( + { + "name": tensor["name"], + "shape": tensor["shape"], + "type": tensor["type"], + "type_id": tensor["type_id"], + "source_offset": tensor["offset"], + "source_nbytes": tensor["nbytes"], + "source_payload_sha256": tensor["sha256"], + "category": "routed_expert" if routed else "non_routed", + "mapping": "expert_pack_identity" if routed else "gguf_direct_identity", + "scale_storage": ( + "inline_quant_block" + if tensor["type"] == "MXFP4" + else "tensor_native" + ), + } + ) + + manifest = { + "format": FORMAT, + "version": 1, + "complete": True, + "created_at": now(), + "layout": "triplet_identity", + "role_order": list(ROLE_NAMES), + "alignment": args.alignment, + "pack_path": str(output), + "pack_size": expected_pack_bytes, + "pack_sha256": pack_sha256, + "header_bytes": header.header_bytes, + "index_count": index_count, + "index_entry_bytes": header.entry_bytes, + "index_sha256": index_sha256, + "data_start": data_start, + "object_count": object_count, + "object_payload_bytes": object_payload_bytes, + "object_stride": object_stride, + "role_bytes": role_nbytes, + "payload_bytes": index_count * role_nbytes, + "padding_bytes": expected_pack_bytes - data_start - index_count * role_nbytes, + "model": { + "ref": args.model_ref, + "model_identity_sha256": args.model_identity_sha256, + "config_sha256": args.config_sha256, + "num_layers": args.num_layers, + "num_routed_experts": args.num_experts, + "top_k": args.top_k, + "single_gpu": True, + }, + "source": { + "path": str(source), + "size": source.stat().st_size, + "sha256": actual_source_sha256, + "gguf_data_offset": int(reader.data_offset), + "gguf_metadata_sha256": source_metadata_sha256, + "inventory_path": str(args.inventory.resolve()) if args.inventory else None, + "inventory_sha256": inventory_summary.get("inventory_sha256"), + "tensor_count": len(tensors), + }, + "coverage": { + "layers": list(range(args.num_layers)), + "experts_per_layer": args.num_experts, + "roles": list(ROLE_NAMES), + "routed_tensor_count": len(routed_names), + "non_routed_tensor_count": len(tensors) - len(routed_names), + }, + "transform": { + "id": "identity-v1", + "description": "Contiguous source bytes; no dequantization, requantization, or value transform", + "reversible": True, + "tool_sha256": tool_sha256(), + }, + "builder": { + "git_sha": git_sha(), + "python": sys.version, + "command": " ".join(sys.argv), + "started_at": started_at, + "elapsed_s": time.monotonic() - started_monotonic, + }, + "tensors": tensor_manifest, + } + + os.replace(partial_path, output) + write_json_atomic(manifest_path, manifest) + checkpoint["status"] = "complete" + checkpoint["completed_at"] = now() + checkpoint["manifest_sha256"] = sha256_file(manifest_path) + checkpoint["pack_sha256"] = pack_sha256 + write_json_atomic(checkpoint_path, checkpoint) + return manifest + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Build an auditable SGLang expert pack from GGUF" + ) + parser.add_argument("--source", type=Path) + parser.add_argument("--source-sha256") + parser.add_argument("--inventory", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--manifest", type=Path) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--model-ref", default="deepseek-v4-flash") + parser.add_argument("--model-identity-sha256") + parser.add_argument("--config-blob", type=Path) + parser.add_argument("--config-sha256") + parser.add_argument("--num-layers", type=int, default=43) + parser.add_argument("--num-experts", type=int, default=256) + parser.add_argument("--top-k", type=int, default=6) + parser.add_argument("--alignment", type=int, default=4096) + parser.add_argument("--safety-margin-gib", type=float, default=16.0) + parser.add_argument("--resume", action="store_true") + parser.add_argument("--inspect", action="store_true") + parser.add_argument("--limit", type=int, default=12) + args = parser.parse_args() + if args.inspect: + return args + for name in ("source", "source_sha256", "model_identity_sha256", "config_sha256"): + if getattr(args, name) is None: + parser.error(f"--{name.replace('_', '-')} is required when building") + args.manifest = args.manifest or args.output.with_name( + args.output.name + ".manifest.json" + ) + args.checkpoint = args.checkpoint or args.output.with_name( + args.output.name + ".checkpoint.json" + ) + if args.num_layers <= 0 or args.num_experts <= 0 or args.top_k <= 0: + parser.error("model dimensions and top-k must be positive") + align_up(0, args.alignment) + return args + + +def main() -> int: + args = parse_args() + if args.inspect: + inspect_pack(args.output, args.limit) + return 0 + manifest = build(args) + print( + json.dumps( + { + "pack": manifest["pack_path"], + "pack_sha256": manifest["pack_sha256"], + "manifest": str(args.manifest.resolve()), + "objects": manifest["object_count"], + "entries": manifest["index_count"], + "pack_size": manifest["pack_size"], + }, + sort_keys=True, + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/expert_pack/format.py b/tools/expert_pack/format.py new file mode 100755 index 000000000..401b05ff8 --- /dev/null +++ b/tools/expert_pack/format.py @@ -0,0 +1,313 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import hashlib +import json +import struct +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO, Iterable + +MAGIC = b"SGLANG-EXPERTPACK-v1\0\0\0\0" +VERSION = 1 +ROLE_NAMES = ("gate", "up", "down") +ROLE_IDS = {name: index for index, name in enumerate(ROLE_NAMES)} +FLAG_IDENTITY_PAYLOAD = 1 << 0 +FLAG_TRIPLET_OBJECTS = 1 << 1 + +# magic, version, header bytes, entry bytes, flags, index count, data start, +# alignment, layer count, expert count, top-k, role count, and three digests. +HEADER_STRUCT = struct.Struct("<24sIIIIQQQIIII32s32s32s") + +# layer, expert, role, rank, GGML dtype id, dtype, tensor name, six ranges, +# source tensor/slice and pack hashes, logical role shape, quant/transform, +# quant block size, and generation. +ENTRY_STRUCT = struct.Struct(" int: + if alignment <= 0 or alignment & (alignment - 1): + raise ValueError("alignment must be a positive power of two") + return (value + alignment - 1) // alignment * alignment + + +def sha256_file(path: Path, chunk_bytes: int = 16 * 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb", buffering=0) as stream: + while chunk := stream.read(chunk_bytes): + digest.update(chunk) + return digest.hexdigest() + + +def parse_sha256(value: str, field: str) -> bytes: + if len(value) != 64: + raise ValueError(f"{field} must be a 64-character SHA-256 digest") + try: + result = bytes.fromhex(value) + except ValueError as exc: + raise ValueError(f"{field} must be a hexadecimal SHA-256 digest") from exc + if len(result) != 32: + raise ValueError(f"{field} must decode to 32 bytes") + return result + + +def encode_fixed(value: str, size: int, field: str) -> bytes: + encoded = value.encode("utf-8") + if len(encoded) >= size: + raise ValueError(f"{field} is too long for its {size}-byte field: {value!r}") + return encoded + bytes(size - len(encoded)) + + +def decode_fixed(value: bytes) -> str: + return value.split(b"\0", 1)[0].decode("utf-8") + + +@dataclass(frozen=True) +class PackHeader: + flags: int + index_count: int + data_start: int + alignment: int + num_layers: int + num_experts: int + top_k: int + role_count: int + model_identity_sha256: str + source_blob_sha256: str + config_sha256: str + + @property + def header_bytes(self) -> int: + return HEADER_STRUCT.size + + @property + def entry_bytes(self) -> int: + return ENTRY_STRUCT.size + + def pack(self) -> bytes: + if self.role_count != len(ROLE_NAMES): + raise ValueError(f"role_count must be {len(ROLE_NAMES)}") + expected_entries = self.num_layers * self.num_experts * self.role_count + if self.index_count != expected_entries: + raise ValueError( + f"index_count {self.index_count} does not match {expected_entries}" + ) + minimum_data_start = self.header_bytes + self.index_count * self.entry_bytes + if self.data_start < minimum_data_start or self.data_start % self.alignment: + raise ValueError("data_start is too small or is not aligned") + return HEADER_STRUCT.pack( + MAGIC, + VERSION, + self.header_bytes, + self.entry_bytes, + self.flags, + self.index_count, + self.data_start, + self.alignment, + self.num_layers, + self.num_experts, + self.top_k, + self.role_count, + parse_sha256(self.model_identity_sha256, "model_identity_sha256"), + parse_sha256(self.source_blob_sha256, "source_blob_sha256"), + parse_sha256(self.config_sha256, "config_sha256"), + ) + + @classmethod + def unpack(cls, raw: bytes) -> PackHeader: + if len(raw) != HEADER_STRUCT.size: + raise ValueError("expert-pack header is truncated") + ( + magic, + version, + header_bytes, + entry_bytes, + flags, + index_count, + data_start, + alignment, + num_layers, + num_experts, + top_k, + role_count, + model_identity_digest, + source_digest, + config_digest, + ) = HEADER_STRUCT.unpack(raw) + if magic != MAGIC: + raise ValueError("expert-pack magic does not match") + if version != VERSION: + raise ValueError(f"unsupported expert-pack version {version}") + if header_bytes != HEADER_STRUCT.size or entry_bytes != ENTRY_STRUCT.size: + raise ValueError( + "expert-pack struct sizes do not match this implementation" + ) + result = cls( + flags=flags, + index_count=index_count, + data_start=data_start, + alignment=alignment, + num_layers=num_layers, + num_experts=num_experts, + top_k=top_k, + role_count=role_count, + model_identity_sha256=model_identity_digest.hex(), + source_blob_sha256=source_digest.hex(), + config_sha256=config_digest.hex(), + ) + result.pack() + return result + + +@dataclass(frozen=True) +class IndexEntry: + layer: int + expert: int + role: str + dtype_id: int + dtype: str + tensor_name: str + source_tensor_offset: int + source_tensor_nbytes: int + source_slice_offset: int + source_slice_nbytes: int + pack_offset: int + pack_nbytes: int + source_tensor_sha256: str + source_slice_sha256: str + checksum: str + shape: tuple[int, ...] + quant_scheme: str + transform_id: str + block_size: int + generation: int + + @property + def key(self) -> tuple[int, int, int]: + return self.layer, self.expert, ROLE_IDS[self.role] + + def pack(self) -> bytes: + if self.role not in ROLE_IDS: + raise ValueError(f"unknown expert role {self.role!r}") + if not 0 <= self.layer <= 0xFFFF or not 0 <= self.expert <= 0xFFFF: + raise ValueError("layer/expert does not fit the pack index") + if not 0 <= self.dtype_id <= 0xFFFF: + raise ValueError("dtype_id does not fit the pack index") + if not 1 <= len(self.shape) <= 4 or any(value <= 0 for value in self.shape): + raise ValueError(f"invalid role shape {self.shape}") + dims = self.shape + (0,) * (4 - len(self.shape)) + return ENTRY_STRUCT.pack( + self.layer, + self.expert, + ROLE_IDS[self.role], + len(self.shape), + self.dtype_id, + encode_fixed(self.dtype, 16, "dtype"), + encode_fixed(self.tensor_name, 80, "tensor_name"), + self.source_tensor_offset, + self.source_tensor_nbytes, + self.source_slice_offset, + self.source_slice_nbytes, + self.pack_offset, + self.pack_nbytes, + parse_sha256(self.source_tensor_sha256, "source_tensor_sha256"), + parse_sha256(self.source_slice_sha256, "source_slice_sha256"), + parse_sha256(self.checksum, "checksum"), + *dims, + encode_fixed(self.quant_scheme, 16, "quant_scheme"), + encode_fixed(self.transform_id, 16, "transform_id"), + self.block_size, + self.generation, + ) + + @classmethod + def unpack(cls, raw: bytes) -> IndexEntry: + if len(raw) != ENTRY_STRUCT.size: + raise ValueError("expert-pack index entry is truncated") + values = ENTRY_STRUCT.unpack(raw) + role_id = values[2] + rank = values[3] + if role_id >= len(ROLE_NAMES) or not 1 <= rank <= 4: + raise ValueError("expert-pack index contains an invalid role or rank") + shape = tuple(values[16 : 16 + rank]) + return cls( + layer=values[0], + expert=values[1], + role=ROLE_NAMES[role_id], + dtype_id=values[4], + dtype=decode_fixed(values[5]), + tensor_name=decode_fixed(values[6]), + source_tensor_offset=values[7], + source_tensor_nbytes=values[8], + source_slice_offset=values[9], + source_slice_nbytes=values[10], + pack_offset=values[11], + pack_nbytes=values[12], + source_tensor_sha256=values[13].hex(), + source_slice_sha256=values[14].hex(), + checksum=values[15].hex(), + shape=shape, + quant_scheme=decode_fixed(values[20]), + transform_id=decode_fixed(values[21]), + block_size=values[22], + generation=values[23], + ) + + def to_dict(self) -> dict[str, object]: + value = dict(self.__dict__) + value["shape"] = list(self.shape) + return value + + @classmethod + def from_dict(cls, value: dict[str, object]) -> IndexEntry: + fields = dict(value) + fields["shape"] = tuple(int(item) for item in fields["shape"]) + return cls(**fields) + + +def read_header(stream: BinaryIO) -> PackHeader: + stream.seek(0) + return PackHeader.unpack(stream.read(HEADER_STRUCT.size)) + + +def read_index(stream: BinaryIO, header: PackHeader) -> list[IndexEntry]: + stream.seek(header.header_bytes) + entries = [] + for _ in range(header.index_count): + entries.append(IndexEntry.unpack(stream.read(header.entry_bytes))) + return entries + + +def write_index( + stream: BinaryIO, header: PackHeader, entries: Iterable[IndexEntry] +) -> str: + ordered = sorted(entries, key=lambda entry: entry.key) + if len(ordered) != header.index_count: + raise ValueError( + f"expected {header.index_count} index entries, got {len(ordered)}" + ) + digest = hashlib.sha256() + stream.seek(header.header_bytes) + for entry in ordered: + raw = entry.pack() + stream.write(raw) + digest.update(raw) + return digest.hexdigest() + + +def inspect_pack(path: Path, limit: int = 12) -> dict[str, object]: + with path.open("rb", buffering=0) as stream: + header = read_header(stream) + entries = read_index(stream, header) + summary = { + "path": str(path.resolve()), + "size": path.stat().st_size, + "header": header.__dict__, + "role_counts": { + role: sum(entry.role == role for entry in entries) for role in ROLE_NAMES + }, + "payload_bytes": sum(entry.pack_nbytes for entry in entries), + "entries": [entry.to_dict() for entry in entries[:limit]], + } + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary diff --git a/tools/expert_pack/kimi_ggml.py b/tools/expert_pack/kimi_ggml.py new file mode 100755 index 000000000..d87252077 --- /dev/null +++ b/tools/expert_pack/kimi_ggml.py @@ -0,0 +1,738 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Structural inventory and adapter manifest for Kimi K3 GGUF assets.""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import struct +import subprocess +from collections.abc import Callable, Iterable +from contextlib import ExitStack +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, BinaryIO + +FORMAT = "SGLANG-KIMI-GGMLMOEPACK-ADAPTER-v1" +PACK_MAGIC = b"GGMLMOEPACKv1\0\0\0" +PACK_VERSION = 1 +PACK_HEADER = struct.Struct("<16sIIQQ") +PACK_ENTRY = struct.Struct("<128siIQQ") +PACK_ALIGNMENT = 4096 +ROLE_ORDER = ("up", "gate", "down") +EXPERT_RE = re.compile( + r"^blk\.(?P\d+)\.ffn_(?Pup|gate|down)_exps\.weight$" +) +SHARD_RE = re.compile(r"-(?P\d{5})-of-(?P\d{5})\.gguf$") +COPY_CHUNK_BYTES = 16 * 1024 * 1024 + + +@dataclass(frozen=True) +class KimiK3Spec: + num_hidden_layers: int + first_k_dense_replace: int + num_experts: int + top_k: int + num_shared_experts: int + hidden_size: int + routed_expert_hidden_size: int + moe_intermediate_size: int + hidden_act: str + active_moe_layer_ids: tuple[int, ...] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> KimiK3Spec: + text_config = config.get("text_config", config) + num_hidden_layers = int(text_config["num_hidden_layers"]) + first_dense = int(text_config["first_k_dense_replace"]) + active_layers = tuple(range(first_dense, num_hidden_layers)) + result = cls( + num_hidden_layers=num_hidden_layers, + first_k_dense_replace=first_dense, + num_experts=int(text_config["num_experts"]), + top_k=int(text_config["num_experts_per_token"]), + num_shared_experts=int(text_config["num_shared_experts"]), + hidden_size=int(text_config["hidden_size"]), + routed_expert_hidden_size=int(text_config["routed_expert_hidden_size"]), + moe_intermediate_size=int(text_config["moe_intermediate_size"]), + hidden_act=str(text_config["hidden_act"]), + active_moe_layer_ids=active_layers, + ) + result.validate_kimi_k3() + return result + + def validate_kimi_k3(self) -> None: + expected = { + "num_hidden_layers": 93, + "first_k_dense_replace": 1, + "num_experts": 896, + "top_k": 16, + "num_shared_experts": 2, + "hidden_size": 7168, + "routed_expert_hidden_size": 3584, + "moe_intermediate_size": 3072, + "hidden_act": "situ", + } + actual = { + "num_hidden_layers": self.num_hidden_layers, + "first_k_dense_replace": self.first_k_dense_replace, + "num_experts": self.num_experts, + "top_k": self.top_k, + "num_shared_experts": self.num_shared_experts, + "hidden_size": self.hidden_size, + "routed_expert_hidden_size": self.routed_expert_hidden_size, + "moe_intermediate_size": self.moe_intermediate_size, + "hidden_act": self.hidden_act, + } + if actual != expected: + raise ValueError( + "Kimi K3 model invariants do not match the audited model: " + f"expected={expected}, actual={actual}" + ) + expected_layers = tuple(range(1, 93)) + if self.active_moe_layer_ids != expected_layers: + raise ValueError("Kimi K3 active MoE layers must be exactly 1..92") + + +@dataclass(frozen=True) +class TensorRecord: + name: str + shape: tuple[int, ...] + dtype: str + dtype_id: int + shard_index: int + shard_path: str + data_offset: int + nbytes: int + + @property + def expert_key(self) -> tuple[int, str] | None: + match = EXPERT_RE.fullmatch(self.name) + if match is None: + return None + return int(match.group("layer")), match.group("role") + + +@dataclass(frozen=True) +class PackEntryRecord: + tensor_name: str + expert: int + offset: int + nbytes: int + + +def _field_value(reader: Any, name: str) -> Any: + field = reader.fields.get(name) + if field is None: + raise ValueError(f"GGUF metadata is missing required field {name!r}") + return field.contents() + + +def _sha256_range(stream: BinaryIO, offset: int, nbytes: int) -> str: + digest = hashlib.sha256() + stream.seek(offset) + remaining = nbytes + while remaining: + chunk = stream.read(min(COPY_CHUNK_BYTES, remaining)) + if not chunk: + raise EOFError(f"short read at offset {offset}; {remaining} bytes remain") + digest.update(chunk) + remaining -= len(chunk) + return digest.hexdigest() + + +def sha256_file(path: Path) -> str: + with path.open("rb", buffering=0) as stream: + return _sha256_range(stream, 0, path.stat().st_size) + + +def canonical_sha256(value: Any) -> str: + encoded = json.dumps( + value, ensure_ascii=False, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def write_json_atomic(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(path.name + ".tmp") + with temporary.open("w", encoding="utf-8", newline="\n") as stream: + json.dump(value, stream, ensure_ascii=False, indent=2, sort_keys=True) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + + +def discover_gguf_shards(directory: Path) -> list[Path]: + directory = directory.resolve(strict=True) + candidates = sorted(directory.glob("*.gguf")) + if not candidates: + raise ValueError(f"no GGUF files found in {directory}") + numbered: list[tuple[int, int, Path]] = [] + for path in candidates: + match = SHARD_RE.search(path.name) + if match is None: + raise ValueError(f"GGUF shard name does not contain NNNNN-of-NNNNN: {path}") + numbered.append((int(match.group("number")), int(match.group("count")), path)) + counts = {count for _, count, _ in numbered} + if len(counts) != 1: + raise ValueError(f"GGUF shard filenames disagree on split count: {counts}") + count = counts.pop() + numbers = [number for number, _, _ in numbered] + if count != len(numbered) or sorted(numbers) != list(range(1, count + 1)): + raise ValueError( + f"GGUF shard set is incomplete: count={count}, numbers={sorted(numbers)}" + ) + return [path for _, _, path in sorted(numbered)] + + +def _git_sha(repo: Path | None) -> str: + if repo is None: + return "unknown" + try: + return subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + except (OSError, subprocess.CalledProcessError): + return "unknown" + + +def scan_gguf_shards( + paths: list[Path], *, full_source_hashes: bool = False +) -> tuple[list[dict[str, Any]], list[TensorRecord], dict[str, Any]]: + import gguf + + shard_records: list[dict[str, Any]] = [] + tensors: list[TensorRecord] = [] + seen_names: set[str] = set() + split_count: int | None = None + split_tensor_count: int | None = None + architecture: str | None = None + for shard_index, path in enumerate(paths): + reader = gguf.GGUFReader(str(path), mode="r") + shard_split_count = int(_field_value(reader, "split.count")) + shard_split_no = int(_field_value(reader, "split.no")) + shard_tensor_count = int(_field_value(reader, "split.tensors.count")) + shard_architecture = str(_field_value(reader, "general.architecture")) + if shard_split_no != shard_index: + raise ValueError( + f"GGUF split.no mismatch for {path}: {shard_split_no} != {shard_index}" + ) + if split_count is None: + split_count = shard_split_count + split_tensor_count = shard_tensor_count + architecture = shard_architecture + elif ( + split_count != shard_split_count + or split_tensor_count != shard_tensor_count + or architecture != shard_architecture + ): + raise ValueError(f"GGUF split metadata mismatch at {path}") + + metadata_nbytes = int(reader.data_offset) + with path.open("rb", buffering=0) as stream: + metadata_sha256 = _sha256_range(stream, 0, metadata_nbytes) + shard_record: dict[str, Any] = { + "index": shard_index, + "path": str(path.resolve()), + "size": path.stat().st_size, + "metadata_nbytes": metadata_nbytes, + "metadata_sha256": metadata_sha256, + "tensor_count": len(reader.tensors), + } + if full_source_hashes: + shard_record["sha256"] = sha256_file(path) + shard_records.append(shard_record) + + for tensor in reader.tensors: + if tensor.name in seen_names: + raise ValueError(f"duplicate GGUF tensor across shards: {tensor.name}") + seen_names.add(tensor.name) + tensors.append( + TensorRecord( + name=tensor.name, + shape=tuple(int(value) for value in tensor.shape.tolist()), + dtype=tensor.tensor_type.name, + dtype_id=int(tensor.tensor_type), + shard_index=shard_index, + shard_path=str(path.resolve()), + data_offset=int(tensor.data_offset), + nbytes=int(tensor.n_bytes), + ) + ) + del reader + + if split_count != len(paths): + raise ValueError(f"GGUF split.count={split_count}, found {len(paths)} files") + if split_tensor_count != len(tensors): + raise ValueError( + f"GGUF split.tensors.count={split_tensor_count}, found {len(tensors)}" + ) + if architecture != "kimi-k3": + raise ValueError(f"expected GGUF architecture 'kimi-k3', got {architecture!r}") + summary = { + "architecture": architecture, + "shard_count": len(paths), + "tensor_count": len(tensors), + "total_bytes": sum(item["size"] for item in shard_records), + "full_source_hashes": full_source_hashes, + } + return shard_records, tensors, summary + + +def validate_expert_tensors( + tensors: Iterable[TensorRecord], spec: KimiK3Spec +) -> dict[tuple[int, str], TensorRecord]: + expert_tensors: dict[tuple[int, str], TensorRecord] = {} + for tensor in tensors: + key = tensor.expert_key + if key is None: + continue + if key in expert_tensors: + raise ValueError(f"duplicate routed expert tensor {key}") + expert_tensors[key] = tensor + expected_keys = { + (layer, role) for layer in spec.active_moe_layer_ids for role in ROLE_ORDER + } + if set(expert_tensors) != expected_keys: + missing = sorted(expected_keys - set(expert_tensors)) + extra = sorted(set(expert_tensors) - expected_keys) + raise ValueError( + f"routed expert tensor coverage mismatch: missing={missing[:8]}, " + f"extra={extra[:8]}" + ) + expected_layout = { + "up": ((spec.routed_expert_hidden_size, spec.moe_intermediate_size), "Q2_K"), + "gate": ( + (spec.routed_expert_hidden_size, spec.moe_intermediate_size), + "Q2_K", + ), + "down": ( + (spec.moe_intermediate_size, spec.routed_expert_hidden_size), + "Q3_K", + ), + } + for (layer, role), tensor in expert_tensors.items(): + expected_shape, expected_dtype = expected_layout[role] + if tensor.shape != (*expected_shape, spec.num_experts): + raise ValueError( + f"unexpected expert shape for {(layer, role)}: {tensor.shape}" + ) + if tensor.dtype != expected_dtype: + raise ValueError( + f"unexpected expert dtype for {(layer, role)}: {tensor.dtype}" + ) + if tensor.nbytes % spec.num_experts: + raise ValueError(f"expert tensor is not evenly sliceable: {tensor.name}") + return expert_tensors + + +def _decode_name(raw: bytes) -> str: + try: + return raw.split(b"\0", 1)[0].decode("utf-8") + except UnicodeDecodeError as exc: + raise ValueError("expert-pack tensor name is not valid UTF-8") from exc + + +def _read_pack_entry(stream: BinaryIO, digest: hashlib._Hash) -> PackEntryRecord: + raw = stream.read(PACK_ENTRY.size) + if len(raw) != PACK_ENTRY.size: + raise ValueError("GGML expert-pack index is truncated") + digest.update(raw) + name, expert, reserved, offset, nbytes = PACK_ENTRY.unpack(raw) + if reserved != 0: + raise ValueError("GGML expert-pack reserved entry field must be zero") + return PackEntryRecord(_decode_name(name), expert, offset, nbytes) + + +def _compare_ranges( + pack_stream: BinaryIO, + pack_entry: PackEntryRecord, + tensor: TensorRecord, + expert_bytes: int, +) -> None: + source_offset = tensor.data_offset + pack_entry.expert * expert_bytes + with Path(tensor.shard_path).open("rb", buffering=0) as source_stream: + source_stream.seek(source_offset) + pack_stream.seek(pack_entry.offset) + remaining = expert_bytes + while remaining: + size = min(COPY_CHUNK_BYTES, remaining) + source_chunk = source_stream.read(size) + pack_chunk = pack_stream.read(size) + if source_chunk != pack_chunk or len(source_chunk) != size: + raise ValueError( + "expert-pack payload does not match GGUF source for " + f"{pack_entry.tensor_name} expert {pack_entry.expert}" + ) + remaining -= size + + +def validate_ggml_moe_pack( + path: Path, + expert_tensors: dict[tuple[int, str], TensorRecord], + spec: KimiK3Spec, + *, + payload_samples: int = 6, + full_pack_hash: bool = False, +) -> dict[str, Any]: + path = path.resolve(strict=True) + expected_count = len(spec.active_moe_layer_ids) * spec.num_experts * len(ROLE_ORDER) + file_size = path.stat().st_size + index_digest = hashlib.sha256() + role_summary: dict[str, dict[str, Any]] = {} + sample_indices = set() + if payload_samples > 0: + if payload_samples == 1: + sample_indices.add(0) + else: + sample_indices.update( + round(index * (expected_count - 1) / (payload_samples - 1)) + for index in range(payload_samples) + ) + sampled_entries: list[tuple[PackEntryRecord, TensorRecord, int]] = [] + object_bytes: int | None = None + object_start: int | None = None + previous_end = 0 + + with path.open("rb", buffering=0) as stream: + raw_header = stream.read(PACK_HEADER.size) + if len(raw_header) != PACK_HEADER.size: + raise ValueError("GGML expert-pack header is truncated") + index_digest.update(raw_header) + magic, version, header_size, index_count, data_start = PACK_HEADER.unpack( + raw_header + ) + if magic != PACK_MAGIC or version != PACK_VERSION: + raise ValueError("GGML expert-pack magic or version does not match") + if header_size != PACK_HEADER.size: + raise ValueError("GGML expert-pack header size does not match") + if index_count != expected_count: + raise ValueError( + f"GGML expert-pack has {index_count} entries; expected {expected_count}" + ) + minimum_data_start = PACK_HEADER.size + index_count * PACK_ENTRY.size + if data_start < minimum_data_start or data_start % PACK_ALIGNMENT: + raise ValueError("GGML expert-pack data_start is invalid or unaligned") + previous_end = data_start + + for index in range(index_count): + entry = _read_pack_entry(stream, index_digest) + object_index, physical_role_id = divmod(index, len(ROLE_ORDER)) + active_layer_index, expected_expert = divmod(object_index, spec.num_experts) + expected_layer = spec.active_moe_layer_ids[active_layer_index] + expected_role = ROLE_ORDER[physical_role_id] + match = EXPERT_RE.fullmatch(entry.tensor_name) + if match is None: + raise ValueError(f"invalid expert tensor name: {entry.tensor_name!r}") + actual_key = (int(match.group("layer")), match.group("role")) + expected_key = (expected_layer, expected_role) + if actual_key != expected_key or entry.expert != expected_expert: + raise ValueError( + "GGML expert-pack is not complete expert-major up/gate/down " + f"layout at index {index}: expected={(*expected_key, expected_expert)}, " + f"actual={(*actual_key, entry.expert)}" + ) + tensor = expert_tensors[expected_key] + expert_bytes = tensor.nbytes // spec.num_experts + if entry.nbytes != expert_bytes: + raise ValueError( + f"expert-pack byte size mismatch for index {index}: " + f"{entry.nbytes} != {expert_bytes}" + ) + if entry.offset % PACK_ALIGNMENT: + raise ValueError(f"expert-pack entry {index} is not 4 KiB aligned") + if entry.offset < previous_end or entry.offset + entry.nbytes > file_size: + raise ValueError( + f"expert-pack entry {index} overlaps or is out of range" + ) + previous_end = entry.offset + entry.nbytes + summary = role_summary.setdefault( + expected_role, + { + "dtype": tensor.dtype, + "dtype_id": tensor.dtype_id, + "logical_shape": list(tensor.shape[:2]), + "expert_bytes": expert_bytes, + "entry_count": 0, + "payload_bytes": 0, + }, + ) + if summary["expert_bytes"] != expert_bytes: + raise ValueError(f"variable expert bytes for role {expected_role}") + summary["entry_count"] += 1 + summary["payload_bytes"] += entry.nbytes + + if physical_role_id == 0: + object_start = entry.offset + elif physical_role_id == len(ROLE_ORDER) - 1: + assert object_start is not None + span = entry.offset + entry.nbytes - object_start + if object_bytes is None: + object_bytes = span + elif object_bytes != span: + raise ValueError("expert-pack object spans are not fixed size") + if index in sample_indices: + sampled_entries.append((entry, tensor, expert_bytes)) + + if previous_end != file_size: + raise ValueError( + f"expert-pack has unexplained trailing bytes: {file_size - previous_end}" + ) + for entry, tensor, expert_bytes in sampled_entries: + _compare_ranges(stream, entry, tensor, expert_bytes) + + assert object_bytes is not None + result: dict[str, Any] = { + "path": str(path), + "size": file_size, + "magic": PACK_MAGIC.rstrip(b"\0").decode("ascii"), + "version": PACK_VERSION, + "header_bytes": PACK_HEADER.size, + "entry_bytes": PACK_ENTRY.size, + "index_count": expected_count, + "data_start": data_start, + "alignment": PACK_ALIGNMENT, + "index_sha256": index_digest.hexdigest(), + "physical_role_order": list(ROLE_ORDER), + "active_moe_layer_ids": list(spec.active_moe_layer_ids), + "num_experts": spec.num_experts, + "top_k": spec.top_k, + "object_bytes": object_bytes, + "roles": role_summary, + "payload_samples_verified": len(sampled_entries), + "full_pack_hash": full_pack_hash, + } + if full_pack_hash: + result["sha256"] = sha256_file(path) + return result + + +def _align_up(value: int, alignment: int = PACK_ALIGNMENT) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _pack_layout( + expert_tensors: dict[tuple[int, str], TensorRecord], spec: KimiK3Spec +) -> tuple[list[tuple[PackEntryRecord, TensorRecord]], int, int]: + index_count = len(spec.active_moe_layer_ids) * spec.num_experts * len(ROLE_ORDER) + data_start = _align_up(PACK_HEADER.size + index_count * PACK_ENTRY.size) + offset = data_start + entries: list[tuple[PackEntryRecord, TensorRecord]] = [] + for layer in spec.active_moe_layer_ids: + for expert in range(spec.num_experts): + for role in ROLE_ORDER: + tensor = expert_tensors[(layer, role)] + expert_bytes = tensor.nbytes // spec.num_experts + offset = _align_up(offset) + entries.append( + ( + PackEntryRecord(tensor.name, expert, offset, expert_bytes), + tensor, + ) + ) + offset += expert_bytes + return entries, data_start, offset + + +def estimate_ggml_moe_pack_size( + expert_tensors: dict[tuple[int, str], TensorRecord], spec: KimiK3Spec +) -> int: + return _pack_layout(expert_tensors, spec)[2] + + +def _copy_tensor_slice( + source: BinaryIO, output: BinaryIO, offset: int, nbytes: int +) -> None: + source.seek(offset) + remaining = nbytes + while remaining: + chunk = source.read(min(COPY_CHUNK_BYTES, remaining)) + if not chunk: + raise EOFError( + f"short GGUF read at offset {offset}; {remaining} bytes remain" + ) + output.write(chunk) + remaining -= len(chunk) + + +def write_ggml_moe_pack( + path: Path, + expert_tensors: dict[tuple[int, str], TensorRecord], + spec: KimiK3Spec, + *, + progress: Callable[[int, int], None] | None = None, +) -> int: + entries, data_start, final_size = _pack_layout(expert_tensors, spec) + path = path.resolve() + partial = path.with_name(path.name + ".partial") + if path.exists(): + raise FileExistsError(f"refusing to overwrite existing Expert Pack: {path}") + if partial.exists(): + raise FileExistsError(f"partial Expert Pack already exists: {partial}") + path.parent.mkdir(parents=True, exist_ok=True) + + with ExitStack() as stack: + sources = { + source_path: stack.enter_context(Path(source_path).open("rb", buffering=0)) + for source_path in {tensor.shard_path for _, tensor in entries} + } + output = stack.enter_context(partial.open("xb", buffering=0)) + output.write( + PACK_HEADER.pack( + PACK_MAGIC, + PACK_VERSION, + PACK_HEADER.size, + len(entries), + data_start, + ) + ) + for entry, _ in entries: + encoded_name = entry.tensor_name.encode("utf-8") + if len(encoded_name) >= 128: + raise ValueError( + f"expert tensor name is too long for the Pack index: {entry.tensor_name}" + ) + output.write( + PACK_ENTRY.pack( + encoded_name.ljust(128, b"\0"), + entry.expert, + 0, + entry.offset, + entry.nbytes, + ) + ) + output.write(bytes(data_start - output.tell())) + + total = len(entries) + for index, (entry, tensor) in enumerate(entries, start=1): + padding = entry.offset - output.tell() + if padding < 0: + raise RuntimeError("Expert Pack layout moved backwards") + if padding: + output.write(bytes(padding)) + source_offset = tensor.data_offset + entry.expert * ( + tensor.nbytes // spec.num_experts + ) + _copy_tensor_slice( + sources[tensor.shard_path], output, source_offset, entry.nbytes + ) + if progress is not None and (index % 1024 == 0 or index == total): + progress(index, total) + output.flush() + os.fsync(output.fileno()) + + if partial.stat().st_size != final_size: + raise RuntimeError( + f"generated Expert Pack size {partial.stat().st_size} != {final_size}" + ) + os.replace(partial, path) + return final_size + + +def create_manifest( + *, + gguf_dir: Path, + expert_pack: Path, + model_config: Path, + tokenizer_dir: Path, + payload_samples: int = 6, + full_source_hashes: bool = False, + full_pack_hash: bool = False, + repo: Path | None = None, +) -> dict[str, Any]: + config = json.loads(model_config.read_text(encoding="utf-8")) + spec = KimiK3Spec.from_config(config) + shard_paths = discover_gguf_shards(gguf_dir) + shard_records, tensors, source_summary = scan_gguf_shards( + shard_paths, full_source_hashes=full_source_hashes + ) + expert_tensors = validate_expert_tensors(tensors, spec) + pack = validate_ggml_moe_pack( + expert_pack, + expert_tensors, + spec, + payload_samples=payload_samples, + full_pack_hash=full_pack_hash, + ) + + tokenizer_files = [] + for path in sorted(tokenizer_dir.resolve(strict=True).iterdir()): + if path.is_file(): + tokenizer_files.append( + { + "name": path.name, + "size": path.stat().st_size, + "sha256": sha256_file(path), + } + ) + tensor_records = [ + { + "name": tensor.name, + "shape": list(tensor.shape), + "dtype": tensor.dtype, + "dtype_id": tensor.dtype_id, + "shard_index": tensor.shard_index, + "data_offset": tensor.data_offset, + "nbytes": tensor.nbytes, + } + for tensor in sorted(tensors, key=lambda item: item.name) + ] + source_inventory = { + "summary": source_summary, + "shards": shard_records, + "tensors": tensor_records, + } + model = { + "config_path": str(model_config.resolve()), + "config_sha256": sha256_file(model_config), + "architecture": "KimiLinearForCausalLM", + "num_hidden_layers": spec.num_hidden_layers, + "active_moe_layer_ids": list(spec.active_moe_layer_ids), + "num_experts": spec.num_experts, + "num_experts_per_token": spec.top_k, + "num_shared_experts": spec.num_shared_experts, + "hidden_size": spec.hidden_size, + "routed_expert_hidden_size": spec.routed_expert_hidden_size, + "moe_intermediate_size": spec.moe_intermediate_size, + "hidden_act": spec.hidden_act, + } + return { + "complete": True, + "format": FORMAT, + "created_at": datetime.now(timezone.utc).isoformat(), + "sglang_git_sha": _git_sha(repo), + "hard_constraints": { + "top_k": 16, + "top_k_is_immutable": True, + "all_selected_experts_must_execute": True, + "expert_pruning_allowed": False, + "requantization_allowed": False, + }, + "model": model, + "source": { + **source_inventory, + "inventory_sha256": canonical_sha256(source_inventory), + }, + "expert_pack": pack, + "tokenizer": { + "path": str(tokenizer_dir.resolve()), + "files": tokenizer_files, + "inventory_sha256": canonical_sha256(tokenizer_files), + }, + "verification": { + "structure": "complete", + "payload_samples": payload_samples, + "full_source_hashes": full_source_hashes, + "full_pack_hash": full_pack_hash, + }, + } diff --git a/tools/expert_pack/prepare_deepseek_pack.py b/tools/expert_pack/prepare_deepseek_pack.py new file mode 100755 index 000000000..f8c2c5075 --- /dev/null +++ b/tools/expert_pack/prepare_deepseek_pack.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +"""Validate or build the DeepSeek expert-pack used by the RTX 5090 benchmark.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import sys +from pathlib import Path + +try: + from .format import read_header +except ImportError: + from format import read_header # type: ignore[no-redef] + + +FORMAT = "SGLANG-EXPERTPACK-v1" +EXPERT_PACK_FILENAME = "DeepSeek-V4-Flash.expert-pack" + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def write_json_atomic(path: Path, value: object) -> None: + temporary = path.with_name(path.name + ".tmp") + temporary.write_text( + json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def normalize_manifest_identity(manifest_path: Path, manifest: dict) -> None: + model = manifest["model"] + normalized_model = { + "ref": model["ref"], + "model_identity_sha256": model["model_identity_sha256"], + "config_sha256": model["config_sha256"], + "num_layers": model["num_layers"], + "num_routed_experts": model["num_routed_experts"], + "top_k": model["top_k"], + "single_gpu": model["single_gpu"], + } + if model != normalized_model: + manifest["model"] = normalized_model + write_json_atomic(manifest_path, manifest) + + +def artifact_paths(source: Path) -> tuple[Path, Path, Path]: + pack = source.parent / EXPERT_PACK_FILENAME + manifest = source.parent / f"{EXPERT_PACK_FILENAME}.manifest.json" + checkpoint = source.parent / f"{EXPERT_PACK_FILENAME}.checkpoint.json" + return pack, manifest, checkpoint + + +def validate_pack( + pack: Path, manifest_path: Path, expected_source: Path | None = None +) -> tuple[bool, str, dict | None]: + try: + pack = pack.resolve(strict=True) + manifest_path = manifest_path.resolve(strict=True) + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + if manifest.get("format") != FORMAT or manifest.get("complete") is not True: + raise ValueError( + "manifest is not a completed SGLANG-EXPERTPACK-v1 manifest" + ) + if Path(manifest["pack_path"]).resolve() != pack: + raise ValueError( + "manifest pack_path does not match the fixed expert-pack path" + ) + if pack.stat().st_size != int(manifest["pack_size"]): + raise ValueError("pack size does not match manifest") + + source = Path(manifest["source"]["path"]).resolve(strict=True) + if expected_source is not None and source != expected_source.resolve( + strict=True + ): + raise ValueError("manifest source path does not match --gguf") + if source.stat().st_size != int(manifest["source"]["size"]): + raise ValueError("source GGUF size does not match manifest") + + with pack.open("rb", buffering=0) as stream: + header = read_header(stream) + stream.seek(header.header_bytes) + raw_index = stream.read(header.index_count * header.entry_bytes) + if len(raw_index) != header.index_count * header.entry_bytes: + raise ValueError("pack index is truncated") + if hashlib.sha256(raw_index).hexdigest() != manifest["index_sha256"]: + raise ValueError("pack index SHA-256 does not match manifest") + + model = manifest["model"] + model_identity_sha256 = model.get( + "model_identity_sha256", header.model_identity_sha256 + ) + model["model_identity_sha256"] = model_identity_sha256 + expected = ( + (header.index_count, manifest["index_count"], "index count"), + (header.data_start, manifest["data_start"], "data start"), + (header.num_layers, model["num_layers"], "layer count"), + (header.num_experts, model["num_routed_experts"], "expert count"), + (header.top_k, model["top_k"], "top-k"), + ( + header.source_blob_sha256, + manifest["source"]["sha256"], + "source digest", + ), + ( + header.model_identity_sha256, + model_identity_sha256, + "model identity digest", + ), + (header.config_sha256, model["config_sha256"], "config digest"), + ) + for actual, wanted, label in expected: + if actual != wanted: + raise ValueError(f"pack header {label} does not match manifest") + return True, "header, index, source path and sizes are valid", manifest + except Exception as exc: + return False, str(exc), None + + +def load_build_inputs(args: argparse.Namespace) -> dict[str, object]: + source = args.gguf.resolve(strict=True) + expert_pack, expert_pack_manifest, checkpoint = artifact_paths(source) + model_config_path = args.model_config.resolve(strict=True) + model_config = json.loads(model_config_path.read_text(encoding="utf-8")) + source_sha256 = sha256_file(source) + config_sha256 = sha256_file(model_config_path) + model_identity = hashlib.sha256( + f"sglang-deepseek-expert-pack-v1:{source_sha256}:{config_sha256}".encode( + "ascii" + ) + ).hexdigest() + return { + "source": source, + "source_sha256": source_sha256, + "expert_pack": expert_pack, + "expert_pack_manifest": expert_pack_manifest, + "checkpoint": checkpoint, + "config_blob": model_config_path, + "config_sha256": config_sha256, + "model_identity_sha256": model_identity, + "num_layers": int(model_config["num_hidden_layers"]), + "num_experts": int(model_config["n_routed_experts"]), + "top_k": int(model_config["num_experts_per_tok"]), + } + + +def remove_invalid_outputs(pack: Path, manifest: Path, checkpoint: Path) -> None: + for path in (pack, manifest, pack.with_name(pack.name + ".partial"), checkpoint): + if path.exists(): + print(f"EXPERT_PACK_REMOVE_INVALID path={path}", flush=True) + path.unlink() + + +def build_pack(args: argparse.Namespace, inputs: dict[str, object]) -> None: + build_script = Path(__file__).with_name("build.py") + expert_pack = Path(inputs["expert_pack"]) + manifest = Path(inputs["expert_pack_manifest"]) + checkpoint = Path(inputs["checkpoint"]) + partial = expert_pack.with_name(expert_pack.name + ".partial") + resume = partial.is_file() and checkpoint.is_file() and not expert_pack.exists() + if not resume: + remove_invalid_outputs(expert_pack, manifest, checkpoint) + + command = [ + sys.executable, + str(build_script), + "--source", + str(inputs["source"]), + "--source-sha256", + str(inputs["source_sha256"]), + "--output", + str(expert_pack), + "--manifest", + str(manifest), + "--checkpoint", + str(checkpoint), + "--model-ref", + args.model_ref, + "--model-identity-sha256", + str(inputs["model_identity_sha256"]), + "--config-blob", + str(inputs["config_blob"]), + "--config-sha256", + str(inputs["config_sha256"]), + "--num-layers", + str(inputs["num_layers"]), + "--num-experts", + str(inputs["num_experts"]), + "--top-k", + str(inputs["top_k"]), + "--alignment", + str(args.alignment), + "--safety-margin-gib", + str(args.safety_margin_gib), + ] + if args.inventory and args.inventory.is_file(): + command.extend(("--inventory", str(args.inventory.resolve()))) + if resume: + command.append("--resume") + print( + f"EXPERT_PACK_BUILD_START output={expert_pack} resume={str(resume).lower()}", + flush=True, + ) + subprocess.run(command, check=True) + print(f"EXPERT_PACK_BUILD_COMPLETE output={expert_pack}", flush=True) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--gguf", type=Path, required=True) + parser.add_argument("--model-config", type=Path, required=True) + parser.add_argument("--inventory", type=Path) + parser.add_argument("--model-ref", default="deepseek-v4-flash") + parser.add_argument("--alignment", type=int, default=4096) + parser.add_argument("--safety-margin-gib", type=float, default=16.0) + parser.add_argument("--check-only", action="store_true") + args = parser.parse_args() + return args + + +def main() -> int: + args = parse_args() + source = args.gguf.resolve(strict=True) + expert_pack, expert_pack_manifest, _ = artifact_paths(source) + valid, reason, manifest = validate_pack(expert_pack, expert_pack_manifest, source) + if valid and manifest is not None: + normalize_manifest_identity(expert_pack_manifest, manifest) + print(f"EXPERT_PACK_VALID path={expert_pack} detail={reason}", flush=True) + return 0 + print(f"EXPERT_PACK_INVALID path={expert_pack} detail={reason}", flush=True) + if args.check_only: + return 1 + + existing_valid, existing_reason, existing_manifest = validate_pack( + expert_pack, expert_pack_manifest + ) + if existing_valid and existing_manifest is not None: + existing_source = Path(existing_manifest["source"]["path"]) + raise RuntimeError( + f"the fixed expert-pack already belongs to a different GGUF: {existing_source}; " + f"move the requested GGUF to its own directory instead of overwriting {expert_pack}" + ) + + inputs = load_build_inputs(args) + build_pack(args, inputs) + valid, reason, manifest = validate_pack(expert_pack, expert_pack_manifest, source) + if not valid or manifest is None: + raise RuntimeError(f"generated expert-pack failed validation: {reason}") + normalize_manifest_identity(expert_pack_manifest, manifest) + print(f"EXPERT_PACK_READY path={expert_pack} detail={reason}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/expert_pack/prepare_kimi_manifest.py b/tools/expert_pack/prepare_kimi_manifest.py new file mode 100755 index 000000000..f614a3a4b --- /dev/null +++ b/tools/expert_pack/prepare_kimi_manifest.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Prepare a small, zero-copy Kimi K3 GGUF/expert-pack adapter manifest.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +try: + from .kimi_ggml import create_manifest, write_json_atomic +except ImportError: + from kimi_ggml import create_manifest, write_json_atomic # type: ignore[no-redef] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--gguf-dir", type=Path, required=True) + parser.add_argument("--expert-pack", type=Path, required=True) + parser.add_argument("--model-config", type=Path, required=True) + parser.add_argument("--tokenizer-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--payload-samples", + type=int, + default=6, + help="Evenly spaced pack entries compared byte-for-byte with GGUF (default: 6).", + ) + parser.add_argument("--full-source-hashes", action="store_true") + parser.add_argument("--full-pack-hash", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.payload_samples < 0: + raise ValueError("--payload-samples must be non-negative") + repo = Path(__file__).resolve().parents[2] + manifest = create_manifest( + gguf_dir=args.gguf_dir, + expert_pack=args.expert_pack, + model_config=args.model_config.resolve(strict=True), + tokenizer_dir=args.tokenizer_dir, + payload_samples=args.payload_samples, + full_source_hashes=args.full_source_hashes, + full_pack_hash=args.full_pack_hash, + repo=repo, + ) + write_json_atomic(args.output.resolve(), manifest) + summary = { + "manifest": str(args.output.resolve()), + "format": manifest["format"], + "source_shards": manifest["source"]["summary"]["shard_count"], + "source_tensors": manifest["source"]["summary"]["tensor_count"], + "pack_entries": manifest["expert_pack"]["index_count"], + "pack_index_sha256": manifest["expert_pack"]["index_sha256"], + "top_k": manifest["hard_constraints"]["top_k"], + "payload_samples_verified": manifest["expert_pack"]["payload_samples_verified"], + } + print(json.dumps(summary, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as exc: + print(f"error: {exc}", file=sys.stderr) + raise SystemExit(1) diff --git a/tools/expert_pack/prepare_kimi_pack.py b/tools/expert_pack/prepare_kimi_pack.py new file mode 100755 index 000000000..d925f728c --- /dev/null +++ b/tools/expert_pack/prepare_kimi_pack.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Validate or build the Kimi K3 Expert Pack derived from GGUF shards.""" + +from __future__ import annotations + +import argparse +import fcntl +import json +import shutil +from pathlib import Path + +try: + from .kimi_ggml import ( + SHARD_RE, + KimiK3Spec, + discover_gguf_shards, + estimate_ggml_moe_pack_size, + scan_gguf_shards, + validate_expert_tensors, + validate_ggml_moe_pack, + write_ggml_moe_pack, + ) +except ImportError: + from kimi_ggml import ( # type: ignore[no-redef] + SHARD_RE, + KimiK3Spec, + discover_gguf_shards, + estimate_ggml_moe_pack_size, + scan_gguf_shards, + validate_expert_tensors, + validate_ggml_moe_pack, + write_ggml_moe_pack, + ) + + +def expert_pack_path(gguf: Path) -> Path: + match = SHARD_RE.search(gguf.name) + if match is None: + raise ValueError(f"Kimi GGUF name is not a numbered shard: {gguf}") + return gguf.parent / f"{gguf.name[: match.start()]}.expert-major.pack" + + +def validate_pack(pack: Path, expert_tensors: dict, spec: KimiK3Spec) -> str: + result = validate_ggml_moe_pack( + pack, expert_tensors, spec, payload_samples=6, full_pack_hash=False + ) + return ( + f"entries={result['index_count']} size={result['size']} " + f"samples={result['payload_samples_verified']}" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--gguf", type=Path, required=True) + parser.add_argument("--model-config", type=Path, required=True) + parser.add_argument("--safety-margin-gib", type=float, default=2.0) + parser.add_argument("--check-only", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + gguf = args.gguf.expanduser().resolve(strict=True) + model_config = args.model_config.expanduser().resolve(strict=True) + if args.safety_margin_gib < 0: + raise ValueError("--safety-margin-gib must be non-negative") + + shards = discover_gguf_shards(gguf.parent) + if gguf not in shards: + raise ValueError(f"--gguf is not part of the discovered shard set: {gguf}") + config = json.loads(model_config.read_text(encoding="utf-8")) + spec = KimiK3Spec.from_config(config) + _, tensors, _ = scan_gguf_shards(shards) + expert_tensors = validate_expert_tensors(tensors, spec) + pack = expert_pack_path(gguf) + partial = pack.with_name(pack.name + ".partial") + lock_path = pack.with_name(pack.name + ".lock") + + with lock_path.open("w") as lock: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX) + try: + detail = validate_pack(pack, expert_tensors, spec) + except (OSError, ValueError) as exc: + print(f"EXPERT_PACK_INVALID path={pack} detail={exc}", flush=True) + if args.check_only: + return 1 + else: + print(f"EXPERT_PACK_VALID path={pack} {detail}", flush=True) + return 0 + + for path in (pack, partial): + if path.exists(): + print(f"EXPERT_PACK_REMOVE_INVALID path={path}", flush=True) + path.unlink() + + estimated_size = estimate_ggml_moe_pack_size(expert_tensors, spec) + safety_margin = int(args.safety_margin_gib * 1024**3) + available = shutil.disk_usage(pack.parent).free + if available < estimated_size + safety_margin: + raise OSError( + f"insufficient space for Kimi Expert Pack: available={available}, " + f"required={estimated_size + safety_margin}" + ) + + print( + f"EXPERT_PACK_BUILD_START output={pack} size={estimated_size}", + flush=True, + ) + + def report_progress(completed: int, total: int) -> None: + print( + f"EXPERT_PACK_BUILD_PROGRESS completed={completed} total={total}", + flush=True, + ) + + write_ggml_moe_pack(pack, expert_tensors, spec, progress=report_progress) + detail = validate_pack(pack, expert_tensors, spec) + print(f"EXPERT_PACK_READY path={pack} {detail}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/expert_pack/validate.py b/tools/expert_pack/validate.py new file mode 100755 index 000000000..e6e713786 --- /dev/null +++ b/tools/expert_pack/validate.py @@ -0,0 +1,335 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import hashlib +import json +import random +import time +from pathlib import Path + +try: + from .format import ( + ROLE_NAMES, + IndexEntry, + read_header, + read_index, + sha256_file, + ) +except ImportError: + from format import ( # type: ignore[no-redef] + ROLE_NAMES, + IndexEntry, + read_header, + read_index, + sha256_file, + ) + + +CHUNK_BYTES = 16 * 1024 * 1024 + + +def hash_range(stream, offset: int, nbytes: int) -> str: + digest = hashlib.sha256() + stream.seek(offset) + remaining = nbytes + while remaining: + chunk = stream.read(min(remaining, CHUNK_BYTES)) + if not chunk: + raise EOFError(f"short read at offset {offset}; {remaining} bytes remain") + digest.update(chunk) + remaining -= len(chunk) + return digest.hexdigest() + + +def compare_ranges(source, pack, entry: IndexEntry) -> None: + remaining = entry.pack_nbytes + source_offset = entry.source_slice_offset + pack_offset = entry.pack_offset + while remaining: + length = min(remaining, CHUNK_BYTES) + source.seek(source_offset) + pack.seek(pack_offset) + source_data = source.read(length) + pack_data = pack.read(length) + if len(source_data) != length or len(pack_data) != length: + raise EOFError(f"short source/pack read for entry {entry.key}") + if source_data != pack_data: + raise ValueError(f"source/pack bytes differ for entry {entry.key}") + source_offset += length + pack_offset += length + remaining -= length + + +def validate(args: argparse.Namespace) -> dict[str, object]: + started = time.monotonic() + pack_path = args.pack.resolve(strict=True) + manifest_path = args.manifest.resolve(strict=True) + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + if manifest.get("format") != "SGLANG-EXPERTPACK-v1" or not manifest.get("complete"): + raise ValueError("manifest is not a complete SGLANG-EXPERTPACK-v1 manifest") + if Path(manifest["pack_path"]).resolve() != pack_path: + raise ValueError("manifest pack path does not match --pack") + if pack_path.stat().st_size != int(manifest["pack_size"]): + raise ValueError("pack size does not match manifest") + + source_path = ( + args.source.resolve(strict=True) + if args.source is not None + else Path(manifest["source"]["path"]).resolve(strict=True) + ) + if source_path.stat().st_size != int(manifest["source"]["size"]): + raise ValueError("source size does not match manifest") + + with pack_path.open("rb", buffering=0) as pack: + header = read_header(pack) + index_start = header.header_bytes + pack.seek(index_start) + raw_index = pack.read(header.index_count * header.entry_bytes) + if len(raw_index) != header.index_count * header.entry_bytes: + raise ValueError("pack index is truncated") + if hashlib.sha256(raw_index).hexdigest() != manifest["index_sha256"]: + raise ValueError("pack index SHA-256 does not match manifest") + entries = read_index(pack, header) + + model = manifest["model"] + source = manifest["source"] + for actual, expected, field in ( + ( + header.model_identity_sha256, + model["model_identity_sha256"], + "model identity digest", + ), + (header.source_blob_sha256, source["sha256"], "source digest"), + (header.config_sha256, model["config_sha256"], "config digest"), + (header.num_layers, model["num_layers"], "layer count"), + (header.num_experts, model["num_routed_experts"], "expert count"), + (header.top_k, model["top_k"], "top-k"), + (header.index_count, manifest["index_count"], "index count"), + (header.data_start, manifest["data_start"], "data start"), + (header.alignment, manifest["alignment"], "alignment"), + ): + if actual != expected: + raise ValueError(f"pack header {field} does not match manifest") + + expected_keys = { + (layer, expert, role) + for layer in range(header.num_layers) + for expert in range(header.num_experts) + for role in ROLE_NAMES + } + by_key = {(entry.layer, entry.expert, entry.role): entry for entry in entries} + if len(by_key) != len(entries) or set(by_key) != expected_keys: + raise ValueError("pack index does not have exact layer/expert/role coverage") + + tensor_map = {tensor["name"]: tensor for tensor in manifest["tensors"]} + non_routed = [ + tensor for tensor in manifest["tensors"] if tensor["category"] == "non_routed" + ] + routed = [ + tensor + for tensor in manifest["tensors"] + if tensor["category"] == "routed_expert" + ] + if len(non_routed) != manifest["coverage"]["non_routed_tensor_count"]: + raise ValueError("non-routed tensor coverage does not match manifest summary") + if len(routed) != manifest["coverage"]["routed_tensor_count"]: + raise ValueError("routed tensor coverage does not match manifest summary") + + ranges = [] + object_stride = int(manifest["object_stride"]) + for layer in range(header.num_layers): + for expert in range(header.num_experts): + object_entries = [by_key[(layer, expert, role)] for role in ROLE_NAMES] + expected_object_start = ( + header.data_start + + (layer * header.num_experts + expert) * object_stride + ) + if object_entries[0].pack_offset != expected_object_start: + raise ValueError( + f"object {(layer, expert)} is not at its expected aligned offset" + ) + if expected_object_start % header.alignment: + raise ValueError(f"object {(layer, expert)} is not aligned") + cursor = expected_object_start + generations = set() + for entry in object_entries: + entry.pack() + tensor = tensor_map.get(entry.tensor_name) + if tensor is None or tensor["category"] != "routed_expert": + raise ValueError( + f"entry {entry.key} does not map to a routed tensor" + ) + if ( + entry.pack_offset != cursor + or entry.pack_nbytes != entry.source_slice_nbytes + ): + raise ValueError( + f"entry {entry.key} breaks identity triplet layout" + ) + if ( + entry.transform_id != "identity-v1" + or entry.checksum != entry.source_slice_sha256 + ): + raise ValueError( + f"entry {entry.key} is not an auditable identity transform" + ) + if entry.source_tensor_offset != tensor["source_offset"]: + raise ValueError(f"entry {entry.key} source tensor offset mismatch") + if entry.source_tensor_nbytes != tensor["source_nbytes"]: + raise ValueError(f"entry {entry.key} source tensor size mismatch") + if entry.source_tensor_sha256 != tensor["source_payload_sha256"]: + raise ValueError(f"entry {entry.key} source tensor hash mismatch") + expected_slice_offset = ( + entry.source_tensor_offset + expert * entry.source_slice_nbytes + ) + if entry.source_slice_offset != expected_slice_offset: + raise ValueError(f"entry {entry.key} source slice offset mismatch") + if entry.source_slice_offset + entry.source_slice_nbytes > ( + entry.source_tensor_offset + entry.source_tensor_nbytes + ): + raise ValueError(f"entry {entry.key} source slice is out of bounds") + ranges.append( + ( + entry.pack_offset, + entry.pack_offset + entry.pack_nbytes, + entry.key, + ) + ) + generations.add(entry.generation) + cursor += entry.pack_nbytes + if len(generations) != 1: + raise ValueError( + f"object {(layer, expert)} has inconsistent generations" + ) + if cursor > expected_object_start + object_stride: + raise ValueError(f"object {(layer, expert)} exceeds its stride") + + ranges.sort() + previous_end = header.data_start + for start, end, key in ranges: + if start < previous_end or end > pack_path.stat().st_size: + raise ValueError(f"overlapping or out-of-range pack entry {key}") + previous_end = end + + bytes_hashed = 0 + pack_hash_ok = None + if args.full_pack_hash: + pack_hash_ok = sha256_file(pack_path) == manifest["pack_sha256"] + bytes_hashed += pack_path.stat().st_size + if not pack_hash_ok: + raise ValueError("full pack SHA-256 does not match manifest") + + entry_hash_count = 0 + if args.full_pack_entry_hashes: + with pack_path.open("rb", buffering=0) as pack: + for entry in sorted(entries, key=lambda value: value.pack_offset): + if ( + hash_range(pack, entry.pack_offset, entry.pack_nbytes) + != entry.checksum + ): + raise ValueError( + f"pack payload checksum mismatch for entry {entry.key}" + ) + bytes_hashed += entry.pack_nbytes + entry_hash_count += 1 + + source_tensor_hash_count = 0 + if args.full_source_tensor_hashes: + with source_path.open("rb", buffering=0) as source_stream: + for tensor in sorted( + manifest["tensors"], key=lambda value: value["source_offset"] + ): + digest = hash_range( + source_stream, + int(tensor["source_offset"]), + int(tensor["source_nbytes"]), + ) + if digest != tensor["source_payload_sha256"]: + raise ValueError( + f"source tensor hash mismatch for {tensor['name']}" + ) + bytes_hashed += int(tensor["source_nbytes"]) + source_tensor_hash_count += 1 + + source_file_hash_ok = None + if args.full_source_file_hash: + source_file_hash_ok = sha256_file(source_path) == source["sha256"] + bytes_hashed += source_path.stat().st_size + if not source_file_hash_ok: + raise ValueError("full source file SHA-256 does not match manifest") + + sample_count = min(args.source_range_samples, len(entries)) + sampled_entries = [] + if sample_count: + seed = int(source["sha256"][:16], 16) + sampled_entries = random.Random(seed).sample(entries, sample_count) + with ( + source_path.open("rb", buffering=0) as source_stream, + pack_path.open("rb", buffering=0) as pack_stream, + ): + for entry in sampled_entries: + compare_ranges(source_stream, pack_stream, entry) + bytes_hashed += entry.pack_nbytes * 2 + + elapsed_s = time.monotonic() - started + result = { + "status": "PASS", + "pack": str(pack_path), + "manifest": str(manifest_path), + "source": str(source_path), + "layers": header.num_layers, + "experts_per_layer": header.num_experts, + "top_k": header.top_k, + "index_count": len(entries), + "object_count": header.num_layers * header.num_experts, + "non_routed_tensor_count": len(non_routed), + "routed_tensor_count": len(routed), + "full_pack_hash": pack_hash_ok, + "full_pack_entry_hash_count": entry_hash_count, + "full_source_tensor_hash_count": source_tensor_hash_count, + "full_source_file_hash": source_file_hash_ok, + "source_range_compare_count": len(sampled_entries), + "bytes_verified": bytes_hashed, + "elapsed_s": elapsed_s, + "verified_mib_s": bytes_hashed / 1024**2 / elapsed_s if bytes_hashed else None, + } + if args.report is not None: + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return result + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Validate SGLANG-EXPERTPACK-v1") + parser.add_argument("--pack", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--source", type=Path) + parser.add_argument("--source-range-samples", type=int, default=96) + parser.add_argument("--full-pack-hash", action="store_true") + parser.add_argument("--full-pack-entry-hashes", action="store_true") + parser.add_argument("--full-source-tensor-hashes", action="store_true") + parser.add_argument("--full-source-file-hash", action="store_true") + parser.add_argument("--full", action="store_true") + parser.add_argument("--report", type=Path) + args = parser.parse_args() + if args.source_range_samples < 0: + parser.error("--source-range-samples must be non-negative") + if args.full: + args.full_pack_hash = True + args.full_pack_entry_hashes = True + args.full_source_tensor_hashes = True + args.full_source_file_hash = True + return args + + +def main() -> int: + args = parse_args() + print(json.dumps(validate(args), indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())