minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support deepseek v4 and kimi k3 on ssd (#35314)

Browse Source 
Co-authored-by: 1BIN4 <1741738350@qq.com>
Co-authored-by: L-Ark <fliangae@connect.ust.hk>
Co-authored-by: Chikati <jxudn@connect.ust.hk>
Co-authored-by: mengzili <zilim@ust.hk>
This commit is contained in:
hujianmin
2026-08-26 10:05:12 +08:00
committed by GitHub
co-authored by 1BIN4 L-Ark Chikati mengzili
parent bec6248272
commit 2d8484740d
46 changed files with 8347 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/runtime/deepseek_v4/benchmark_deepseek_5090.py
Executable
+492
View File
@@ -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><think>"
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())
examples/runtime/kimi_k3/benchmark_kimi_k3_5090.py
Executable
+573
View File
@@ -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())
python/sglang/kernels/jit/csrc/moe/expert_pack_mxfp4.cu
+565
View File
@@ -0,0 +1,565 @@
// SPDX-License-Identifier: Apache-2.0
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <cstdint>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <tuple>
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 <typename scalar_t>
__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<half>(const half* input, int index) {
return __half2float(input[index]);
}
template <typename scalar_t>
__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<half>(float value) {
return __float2half_rn(value);
}
template <typename scalar_t>
__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<int64_t>(input_row) * input_size;
const int blocks_per_row = input_size / kQuantBlock;
const int64_t row_bytes = static_cast<int64_t>(blocks_per_row) * kBlockBytes;
const int32_t slot = slot_ids[record];
const uint8_t* weight_base = cache + static_cast<int64_t>(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<int64_t>(output_row) * row_bytes + static_cast<int64_t>(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<int64_t>(output_row) * row_bytes + static_cast<int64_t>(block) * kBlockBytes;
const int exponent = static_cast<int>(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<int64_t>(record) * output_size + output_row] = store_scalar<scalar_t>(sums[row]);
}
}
}
}
// Compute gate and up together so the hidden-state vector is loaded once.
template <typename scalar_t>
__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<int64_t>(input_row) * input_size;
const int blocks_per_row = input_size / kQuantBlock;
const int64_t row_bytes = static_cast<int64_t>(blocks_per_row) * kBlockBytes;
const int32_t slot = slot_ids[record];
const uint8_t* weight_base_a = cache + static_cast<int64_t>(slot) * cache_stride + role_offset_a;
const uint8_t* weight_base_b = cache + static_cast<int64_t>(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<int64_t>(output_row) * row_bytes + static_cast<int64_t>(block) * kBlockBytes;
const uint8_t* quant_b =
weight_base_b + static_cast<int64_t>(output_row) * row_bytes + static_cast<int64_t>(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<int64_t>(output_row) * row_bytes + static_cast<int64_t>(block) * kBlockBytes;
const uint8_t* quant_b =
weight_base_b + static_cast<int64_t>(output_row) * row_bytes + static_cast<int64_t>(block) * kBlockBytes;
const int exponent_a = static_cast<int>(quant_a[0]) - 127;
const int exponent_b = static_cast<int>(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<int64_t>(record) * output_size + output_row] = store_scalar<scalar_t>(sums_a[row]);
output_b[static_cast<int64_t>(record) * output_size + output_row] = store_scalar<scalar_t>(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<int64_t>(blocks_per_row) * kBlockBytes;
const uint8_t* ptr = raw + static_cast<int64_t>(slot) * cache_stride + role_offset +
static_cast<int64_t>(row) * row_bytes + static_cast<int64_t>(block) * kBlockBytes + 1 +
word_in_block * 4;
return static_cast<uint32_t>(ptr[0]) | (static_cast<uint32_t>(ptr[1]) << 8) | (static_cast<uint32_t>(ptr[2]) << 16) |
(static_cast<uint32_t>(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<int64_t>(blocks_per_row) * kBlockBytes;
const uint8_t* ptr = raw + static_cast<int64_t>(slot) * cache_stride + role_offset +
static_cast<int64_t>(row) * row_bytes + static_cast<int64_t>(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<uint8_t>(((interleaved & 7) * 8) + (interleaved >> 3));
}
__device__ __forceinline__ uint8_t marlin_nibble(uint32_t word, int value_index) {
return static_cast<uint8_t>((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<int64_t>(input_size / kMarlinTileK) * (output_size * 2);
const int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (batch >= gridDim.y || index >= total_words) return;
const int64_t tile_span = static_cast<int64_t>(output_size / kMarlinTileN) * kMarlinTileWords;
const int tile_k = static_cast<int>(index / tile_span);
const int64_t tile_rem = index % tile_span;
const int tile_n = static_cast<int>(tile_rem / kMarlinTileWords);
const int local = static_cast<int>(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<int>(role_bytes);
const int role2_offset = static_cast<int>(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<uint32_t>(values[pack_index[i]]) << (i * 4);
}
output[static_cast<int64_t>(target_slot) * output_stride + index] = static_cast<int32_t>(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<int64_t>(groups) * output_size;
const int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (batch >= gridDim.y || index >= total) return;
const int group = static_cast<int>(index / output_size);
const int column = static_cast<int>(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<int>(role_bytes) : static_cast<int>(2 * role_bytes));
const uint8_t value = load_raw_scale(raw, raw_stride, source_slots[batch], role_offset, row, groups, group);
output[static_cast<int64_t>(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<int>(source_slots.numel());
if (batch == 0) return;
const int threads = 256;
const auto stream = at::cuda::getCurrentCUDAStream();
const int w13_n = static_cast<int>(2 * intermediate_size);
const int w2_n = static_cast<int>(hidden_size);
const int w13_k = static_cast<int>(hidden_size);
const int w2_k = static_cast<int>(intermediate_size);
const int64_t w13_words = static_cast<int64_t>(w13_k / kMarlinTileK) * w13_n * 2;
const int64_t w2_words = static_cast<int64_t>(w2_k / kMarlinTileK) * w2_n * 2;
const int64_t w13_scales = static_cast<int64_t>(w13_k / kQuantBlock) * w13_n;
const int64_t w2_scales = static_cast<int64_t>(w2_k / kQuantBlock) * w2_n;
mxfp4_marlin_repack_weight_kernel<<<dim3((w13_words + threads - 1) / threads, batch), threads, 0, stream>>>(
raw.data_ptr<uint8_t>(),
raw.stride(0),
source_slots.data_ptr<int32_t>(),
target_slots.data_ptr<int32_t>(),
role_bytes,
w13_k,
w13_n,
true,
w13.data_ptr<int32_t>(),
w13.stride(0));
mxfp4_marlin_repack_weight_kernel<<<dim3((w2_words + threads - 1) / threads, batch), threads, 0, stream>>>(
raw.data_ptr<uint8_t>(),
raw.stride(0),
source_slots.data_ptr<int32_t>(),
target_slots.data_ptr<int32_t>(),
role_bytes,
w2_k,
w2_n,
false,
w2.data_ptr<int32_t>(),
w2.stride(0));
mxfp4_marlin_repack_scale_kernel<<<dim3((w13_scales + threads - 1) / threads, batch), threads, 0, stream>>>(
raw.data_ptr<uint8_t>(),
raw.stride(0),
source_slots.data_ptr<int32_t>(),
target_slots.data_ptr<int32_t>(),
role_bytes,
w13_k,
w13_n,
true,
w13_scale.data_ptr<uint8_t>(),
w13_scale.stride(0));
mxfp4_marlin_repack_scale_kernel<<<dim3((w2_scales + threads - 1) / threads, batch), threads, 0, stream>>>(
raw.data_ptr<uint8_t>(),
raw.stride(0),
source_slots.data_ptr<int32_t>(),
target_slots.data_ptr<int32_t>(),
role_bytes,
w2_k,
w2_n,
false,
w2_scale.data_ptr<uint8_t>(),
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<<<grid, block, 0, stream>>>(
reinterpret_cast<const __nv_bfloat16*>(input.data_ptr()),
cache.data_ptr<uint8_t>(),
cache.stride(0),
slot_ids.data_ptr<int32_t>(),
role_offset,
input_size,
output_size,
records,
records_per_input,
reinterpret_cast<__nv_bfloat16*>(output.data_ptr()));
} else {
mxfp4_matvec_kernel<<<grid, block, 0, stream>>>(
reinterpret_cast<const half*>(input.data_ptr()),
cache.data_ptr<uint8_t>(),
cache.stride(0),
slot_ids.data_ptr<int32_t>(),
role_offset,
input_size,
output_size,
records,
records_per_input,
reinterpret_cast<half*>(output.data_ptr()));
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
return output;
}
std::tuple<torch::Tensor, torch::Tensor> 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<<<grid, block, 0, stream>>>(
reinterpret_cast<const __nv_bfloat16*>(input.data_ptr()),
cache.data_ptr<uint8_t>(),
cache.stride(0),
slot_ids.data_ptr<int32_t>(),
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<<<grid, block, 0, stream>>>(
reinterpret_cast<const half*>(input.data_ptr()),
cache.data_ptr<uint8_t>(),
cache.stride(0),
slot_ids.data_ptr<int32_t>(),
role_offset_a,
role_offset_b,
input_size,
output_size,
records,
records_per_input,
reinterpret_cast<half*>(output_a.data_ptr()),
reinterpret_cast<half*>(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");
}
python/sglang/kernels/jit/csrc/ngram_corpus/result.h
+1
View File
@@ -1,5 +1,6 @@
#pragma once
#include <cstddef>
#include <cstdint>
#include <unordered_map>
#include <vector>
python/sglang/kernels/ops/attention/dsv4/compress.py
+7
View File
@@ -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],
python/sglang/kernels/ops/kimi_k3/attn_res.py
+4 -3
View File
@@ -34,11 +34,12 @@ def _jit_fused_tma_module(
"""Compile and cache the warp-specialized TMA aggregation kernel (per-row
bulk copies into chunk slots; chunk_rows / occupancy / consumer_regs are
tuning knobs). The smem ring is frozen at 2 chunk slots and PDL is always
on: the kernel targets SM100+, where both are unconditional wins."""
on: the kernel targets SM100+ except SM12x."""
major, minor = torch.cuda.get_device_capability()
if major < 10:
if major < 10 or major == 12:
raise RuntimeError(
"attn_res_fused_tma requires SM100+ (tcgen05, cp.async.bulk)"
"attn_res_fused_tma requires SM100+ excluding SM12x; "
f"SM{major}{minor} is unsupported"
)
args = make_cpp_args(
_DIM,
python/sglang/kernels/ops/moe/expert_pack_mxfp4.py
+113
View File
@@ -0,0 +1,113 @@
# SPDX-License-Identifier: Apache-2.0
"""Lazy-built CUDA kernels used by the expert-pack MoE runtime."""
from __future__ import annotations
import os
from functools import lru_cache
import torch
from torch.utils.cpp_extension import load
from sglang.kernels.jit.utils import KERNEL_PATH
_EXTENSION_NAME = "sglang_expert_pack_mxfp4"
@lru_cache(maxsize=1)
def _extension():
source = KERNEL_PATH / "csrc" / "moe" / "expert_pack_mxfp4.cu"
return load(
name=_EXTENSION_NAME,
sources=[str(source)],
extra_cflags=["-O3"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=os.getenv("SGLANG_EXPERT_PACK_BUILD_VERBOSE", "0") == "1",
)
def mxfp4_matvec(
x: torch.Tensor,
cache: torch.Tensor,
slot_ids: torch.Tensor,
*,
role_offset: int,
role_bytes: int,
input_size: int,
output_size: int,
records_per_input: int,
) -> torch.Tensor:
"""Multiply selected raw GGUF MXFP4 matrices by BF16/FP16 rows."""
return _extension().mxfp4_matvec(
x,
cache,
slot_ids,
role_offset,
role_bytes,
input_size,
output_size,
records_per_input,
)
def mxfp4_matvec_dual(
x: torch.Tensor,
cache: torch.Tensor,
slot_ids: torch.Tensor,
*,
gate_role_offset: int,
up_role_offset: int,
role_bytes: int,
input_size: int,
output_size: int,
records_per_input: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Compute gate and up projections while loading each input row once."""
return _extension().mxfp4_matvec_dual(
x,
cache,
slot_ids,
gate_role_offset,
up_role_offset,
role_bytes,
input_size,
output_size,
records_per_input,
)
def prewarm_mxfp4_extension() -> None:
"""Build and load the extension before the server accepts requests."""
_extension()
def mxfp4_marlin_repack(
raw: torch.Tensor,
source_slots: torch.Tensor,
target_slots: torch.Tensor,
*,
role_bytes: int,
hidden_size: int,
intermediate_size: int,
w13: torch.Tensor,
w2: torch.Tensor,
w13_scale: torch.Tensor,
w2_scale: torch.Tensor,
) -> None:
"""Repack raw GGUF objects into contiguous Marlin SoA cache tensors."""
_extension().mxfp4_marlin_repack(
raw,
source_slots,
target_slots,
role_bytes,
hidden_size,
intermediate_size,
w13,
w2,
w13_scale,
w2_scale,
)
python/sglang/srt/arg_groups/expert_pack_hook.py
+198
View File
@@ -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."
)
python/sglang/srt/configs/load_config.py
+1
View File
@@ -23,6 +23,7 @@ class LoadFormat(str, enum.Enum):
SHARDED_STATE = "sharded_state"
PRESHARDED = "presharded"
GGUF = "gguf"
EXPERT_PACK = "expert_pack"
BITSANDBYTES = "bitsandbytes"
MISTRAL = "mistral"
LAYERED = "layered"
python/sglang/srt/configs/model_config.py
+5 -1
View File
@@ -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
python/sglang/srt/layers/attention/dsv4/compressor.py
+18 -3
View File
@@ -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,
python/sglang/srt/layers/attention/dsv4/indexer.py
+7 -1
View File
@@ -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
python/sglang/srt/layers/linear.py
+8
View File
@@ -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:
python/sglang/srt/layers/moe/expert_pack.py
+931
View File
@@ -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("<HHBBH16s80sQQQQQQ32s32s32s4Q16s16sQQ")
REQUIRED_FLAGS = (1 << 0) | (1 << 1)
READ_SPLITS = 4
KIMI_FORMAT = "SGLANG-KIMI-GGMLMOEPACK-ADAPTER-v1"
GGML_PACK_MAGIC = b"GGMLMOEPACKv1\0\0\0"
GGML_PACK_HEADER = struct.Struct("<16sIIQQ")
GGML_PACK_ENTRY = struct.Struct("<128siIQQ")
KIMI_PHYSICAL_ROLES = ("up", "gate", "down")
KIMI_EXPERT_RE = re.compile(
r"^blk\.(?P<layer>\d+)\.ffn_(?P<role>up|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)
python/sglang/srt/layers/quantization/expert_pack.py
+322
View File
@@ -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)
python/sglang/srt/layers/quantization/gguf.py
+18 -3
View File
@@ -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:
python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py
+2
View File
@@ -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:
python/sglang/srt/lora/deepseek_mla_correction.py
+2 -1
View File
@@ -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(
python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py
+2 -1
View File
@@ -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(
python/sglang/srt/model_loader/deepseek4_gguf.py
+181
View File
@@ -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<layer>\d+)\.ffn_(?P<role>gate|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
python/sglang/srt/model_loader/expert_pack_config.py
+41
View File
@@ -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
python/sglang/srt/model_loader/expert_pack_loader.py
+320
View File
@@ -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()
python/sglang/srt/model_loader/expert_pack_runtime.py
+575
View File
@@ -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}"
)
python/sglang/srt/model_loader/kimi_k3_gguf.py
+203
View File
@@ -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<layer>\d+)\.(?P<suffix>.+)$")
_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
python/sglang/srt/model_loader/loader.py
+5
View File
@@ -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)
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+27 -1
View File
@@ -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,
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py
+9 -2
View File
@@ -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
python/sglang/srt/models/deepseek_v2.py
+34 -8
View File
@@ -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
python/sglang/srt/models/deepseek_v4.py
+103 -84
View File
@@ -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)
python/sglang/srt/models/kimi_k3.py
+96 -3
View File
@@ -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]
python/sglang/srt/server_args.py
+19
View File
@@ -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
test/registered/expert_pack/test_expert_pack_mxfp4.py
+360
View File
@@ -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()
test/registered/expert_pack/test_expert_pack_runtime.py
+444
View File
@@ -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()
test/registered/expert_pack/test_kimi_k3_gguf.py
+115
View File
@@ -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()
test/registered/unit/layers/test_minicpm_sparse_metadata.py
+30 -21
View File
@@ -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"
test/registered/unit/lora/test_deepseek_mla_correction.py
+24
View File
@@ -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()
test/registered/unit/test_model_overrides.py
+4 -1
View File
@@ -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"
)
test/registered/unit/test_supplied_instance_exposure_ratchet.py
+6
View File
@@ -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"),
tools/expert_pack/__init__.py
+1
View File
@@ -0,0 +1 @@
"""DeepSeek expert-pack build, inspection, and validation tools."""
tools/expert_pack/build.py
Executable
+635
View File
@@ -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<layer>\d+)\.ffn_(?P<role>gate|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())
tools/expert_pack/format.py
Executable
+313
View File
@@ -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("<HHBBH16s80sQQQQQQ32s32s32s4Q16s16sQQ")
def align_up(value: int, alignment: int) -> 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
tools/expert_pack/kimi_ggml.py
Executable
+738
View File
@@ -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<layer>\d+)\.ffn_(?P<role>up|gate|down)_exps\.weight$"
)
SHARD_RE = re.compile(r"-(?P<number>\d{5})-of-(?P<count>\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,
},
}
tools/expert_pack/prepare_deepseek_pack.py
Executable
+264
View File
@@ -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())
tools/expert_pack/prepare_kimi_manifest.py
Executable
+71
View File
@@ -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)
tools/expert_pack/prepare_kimi_pack.py
Executable
+125
View File
@@ -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())
tools/expert_pack/validate.py
Executable
+335
View File
@@ -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())