commit d566b300f25f13b1f7ce3659af0e892c0a408312 Author: Sun Ruoxi Date: Fri Oct 30 17:29:26 2026 +0800 first commit diff --git a/README.md b/README.md new file mode 100644 index 0000000..f93941c --- /dev/null +++ b/README.md @@ -0,0 +1,112 @@ +# DeepSeek-R1 昆仑芯 INT8 转换 + +针对完整 DeepSeek-R1/V3 架构检查点。 +依据已安装的 `vllm==0.25.1`、`vllm-kunlun==0.25.1.dev0` 实现 + +## 输出策略 + +| 权重 | 输出类型 | 原因 | +|---|---|---| +| MLA q_a/q_b/kv_a/o 投影 | 逐输出通道 INT8 | 对接 KunlunScaledMMLinearKernel | +| 58 层 MoE 的 256 个路由专家 gate/up/down | 逐输出通道 INT8 | 对接 KunlunCompressedTensorsW8A8Int8MoEMethod | +| MLA kv_b_proj | 反量化 BF16 | 对接 MLA 的 BF16 W_UK_T/W_UV 拆分路径 | +| 前三层 Dense MLP | 反量化 BF16 | 真实样本逐通道 INT8 误差较大 | +| 共享专家 gate/up/down | 反量化 BF16 | 真实样本逐通道 INT8 误差较大 | +| Embedding、lm_head、norm、router、routing bias | 原样保留 | 保持源类型和值 | +| 第 62 层 MTP，即 model.layers.61 | 不包含 | 本脚本输出普通推理版本，num_nextn_predict_layers=0 | + +输出 `compressed-tensors / int-quantized`，权重逐通道静态对称 INT8，激活逐 token 动态对称 INT8；无需校准集。BF16 例外通过量化配置的 ignore 明确声明，包括 vLLM 融合后的 gate_up_proj 名称。保持源张量名称，由 vLLM 完成 Q/KV A 投影、gate/up 的融合与专家权重切分。 + +```text +W[i,j] = W_fp8[i,j] * weight_scale_inv[i//128,j//128] +s[i,0] = max(abs(W[i,:])) / 127 +W_int8[i,j] = clamp(round(W[i,j]/s[i,0]), -127, 127) +``` + +`weight_scale_inv` 保存的值用于乘法反量化，不能取倒数。 +输出 INT8 scale 为 `weight_scale`，FP32，形状 `[out_features,1]`。 +磁盘存 `absmax/127`，插件加载后自行乘 127，转换时不能提前乘。 +BF16 例外输出 `BF16(FP8 * block_scale)`，删除原 block scale，不产生 INT8 scale。 +全零 INT8 行使用 `q=0,s=1`。 + +脚本按源索引定位跨分片的 scale，逐分片、逐矩阵、分行处理；不加载整个模型，不分配 XPU，不修改原始模型。 +预检逐项验证前 61 层的全部 MLA/Dense/共享及路由专家投影、矩阵形状、FP8 格式和 scale 网格。 + +## 使用 + +脚本已部署到： + +```text +/mnt/data1/quant_tools/deepseek_r1_kunlun_int8/ +``` + +预检不创建输出目录、不转换权重： + +```bash +/opt/vllm_kunlun/bin/python \ + /mnt/data1/quant_tools/deepseek_r1_kunlun_int8/convert_deepseek_fp8_to_int8.py \ + --source /mnt/data1/deepseek-r1 \ + --output /mnt/data1/deepseek-r1-INT8 \ + --dry-run +``` + +开始全量转换： + +```bash +/opt/vllm_kunlun/bin/python \ + /mnt/data1/quant_tools/deepseek_r1_kunlun_int8/convert_deepseek_fp8_to_int8.py \ + --source /mnt/data1/deepseek-r1 \ + --output /mnt/data1/deepseek-r1-INT8 \ + --threads 8 +``` + +中断后在相同命令末尾添加 `--resume`。输出目录必须是独立空目录，或具有匹配的 `conversion_state.json`。默认 `--chunk-rows 1024`，恢复时必须与首次运行一致。每片先写临时文件，检查成功后重命名，再记录完成状态。完成时打印 `Complete` 并生成新配置、索引和 `conversion_report.json`，之后才可部署。 + +恢复时核实源配置和索引、源分片大小与时间戳；核实已完成输出分片大小、张量名称、类型、形状及 safetensors 文件结构。恢复检查不计算完整权重文件校验和。 + +2026-10-29 预检结果：45,032 个 FP8 矩阵，其中 44,788 个转 INT8，244 个转 BF16；160 个输出分片，预计张量数据 **631.84 GiB**。原始目录按真实张量计算为 641.29 GiB（含 MTP）；源索引的 total_size 元数据偏大，脚本不使用该值估算转换空间。 + +## 首次部署参数 + +转换后可运行提供的启动示例： + +```bash +bash /mnt/data1/quant_tools/deepseek_r1_kunlun_int8/serve_int8.sh +``` + +示例使用 TP=8、BF16、eager 模式、4096 上下文、单并发，每卡 KV cache 显式分配 4 GiB。 +`VLLM_KUNLUN_ENABLE_INT8_BMM=0` 与 BF16 kv_b_proj 一致，不能设置为 1。 +首测使用 TP，不沿用原命令中的 `--enable-expert-parallel`；当前 INT8 单体 MoE 路径的 EP 路由和专家重映射未验证。 +这是准备好的初始参数，尚未完成完整模型启动或生成测试。 + +## 已完成的验证与范围 + +```bash +cd /mnt/data1/quant_tools/deepseek_r1_kunlun_int8 +/opt/vllm_kunlun/bin/python test_conversion.py +/opt/vllm_kunlun/bin/python validate_kunlun_contract.py +/opt/vllm_kunlun/bin/python validate_sample_weights.py \ + --source /mnt/data1/deepseek-r1 \ + --report sample_validation_mixed.json +``` + +- CPU 数值测试与合成检查点转换通过：包括跨分片 scale、非整块维度、全零行、BF16 例外、MTP 排除、配置/索引、断点恢复、源文件保护、拒绝覆盖和缺失投影检查。 +- 当前插件的 CPU 配置/参数约定检查通过：MLA 融合投影的 INT8 内核、路由专家 INT8 MoE、scale 乘 127、BF16 ignore 配置和 MLA BF16 权重拆分。测试模拟 TP rank/size，不是完整引擎启动或 XPU 算子执行。 +- 九个真实矩阵抽样通过：覆盖五类 MLA 投影、Dense MLP、共享专家、路由专家 gate/down。 + +初始 absmax INT8 基线中，Dense gate 样本 relative RMSE 为 10.10%，共享专家 gate 为 7.18%。将这两类投影改为 BF16 后，对应样本误差为 0.142% 和 0.152%；最终 INT8 样本误差约为 0.94%–1.99%，满足半个逐行量化步长的舍入误差界。相对误差均以 FP8 乘 block scale 后的 FP32 权重为参考；这些指标不代表任务准确率。 + +服务器上的 `sample_validation.json` 保存初始基线，`sample_validation_mixed.json` 保存最终策略结果。 +这是从已量化 FP8 再转换的基线方案；全量转换、P800 推理及实际任务质量仍需分别验证。 + +## 2026-10-29 MLA 初始化问题及绕过 + +完整部署发现当前 FlashMLAImpl 缺少 vLLM 所要求的 forward_mha/forward_mqa，无法实例化。此前 CPU 权重准备检查未覆盖这个实例化步骤。权重无需重新转换。 + +新增 serve_int8_no_mla.sh，通过 VLLM_MLA_DISABLE=1 选择 DeepseekV2Attention 和昆仑芯普通 Attention，绕过不完整的 FlashMLA 接口；保持 TP=8、不启用 EP、eager、单并发、4096 上下文和每卡 4 GiB KV cache。关闭 MLA 增加 KV cache 开销，可能降低吞吐。该路径尚未经过完整模型启动和输出验证。 + +在容器内运行： + +```bash +bash /mnt/data1/quant_tools/deepseek_r1_kunlun_int8/serve_int8_no_mla.sh +``` diff --git a/convert_deepseek_fp8_to_int8.py b/convert_deepseek_fp8_to_int8.py new file mode 100644 index 0000000..18ec31e --- /dev/null +++ b/convert_deepseek_fp8_to_int8.py @@ -0,0 +1,394 @@ +#!/usr/bin/env python3 +"""DeepSeek R1/V3 block-FP8 -> Kunlun compressed-tensors dynamic W8A8. + +CPU only; no model instantiation, calibration, or accelerator allocation. +Checkpoint weight_scale = absmax / 127 (NOT the Kunlun runtime's absmax). +MLA kv_b_proj, dense MLPs and shared experts are dequantized to BF16; +routed experts and remaining MLA projections use INT8. MTP is omitted. +""" +import argparse +from contextlib import ExitStack +import hashlib +import json +import math +from pathlib import Path +import re +import shutil +import struct +import time + +VERSION = 1 +PROJECTION = re.compile( + r"model\.layers\.\d+\.(?:self_attn\.(?:q_a_proj|q_b_proj|kv_a_proj_with_mqa|kv_b_proj|o_proj)|" + r"mlp\.(?:(?:experts\.\d+|shared_experts)\.)?(?:gate_proj|up_proj|down_proj))\.weight$" +) +DTYPE_BYTES = {"F8_E4M3": 1, "BF16": 2, "F16": 2, "F32": 4, "I8": 1} + + +def read_json(path): + return json.loads(path.read_text()) + + +def write_json(path, value): + tmp = path.with_name(path.name + ".tmp") + tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n") + tmp.replace(path) + + +def read_header(path): + with path.open("rb") as stream: + raw = stream.read(8) + if len(raw) != 8: + raise ValueError(f"Truncated safetensors: {path}") + size = struct.unpack(" 100 * 1024 * 1024: + raise ValueError(f"Invalid safetensors header: {path}") + header = json.loads(stream.read(size)) + header.pop("__metadata__", None) + return header + + +def quantization_config(): + return { + "quant_method": "compressed-tensors", + "format": "int-quantized", + "quantization_status": "compressed", + "config_groups": {"group_0": { + "targets": ["Linear"], + "weights": {"num_bits": 8, "type": "int", "symmetric": True, + "strategy": "channel", "dynamic": False}, + "input_activations": {"num_bits": 8, "type": "int", "symmetric": True, + "strategy": "token", "dynamic": True}, + "output_activations": None, + }}, + "ignore": ["lm_head", "re:.*\\.gate$", "re:.*\\.self_attn\\.kv_b_proj$", + "re:.*\\.mlp\\.(gate_up_proj|gate_proj|up_proj|down_proj)$", + "re:.*\\.mlp\\.shared_experts\\.(gate_up_proj|gate_proj|up_proj|down_proj)$"], + "kv_cache_scheme": None, + } + + +def plan_source(source): + config = read_json(source / "config.json") + quant = config.get("quantization_config", {}) + if (config.get("model_type") != "deepseek_v3" + or quant.get("quant_method") != "fp8" + or quant.get("fmt") not in {"e4m3", "float8_e4m3fn"} + or quant.get("activation_scheme") != "dynamic"): + raise ValueError("Expected DeepSeek R1/V3 dynamic block FP8 E4M3 checkpoint") + if config.get("moe_layer_freq", 1) != 1 or not config.get("q_lora_rank"): + raise ValueError("Only DeepSeek V3/R1 with Q-LoRA and moe_layer_freq=1 is supported") + block = quant.get("weight_block_size") + if block != [128, 128]: + raise ValueError(f"Expected FP8 block size [128, 128], got {block}") + index_path = source / "model.safetensors.index.json" + index = read_json(index_path) + mapping = index["weight_map"] + specs = {} + source_files = sorted(set(mapping.values())) + digest = hashlib.sha256() + for path in (source / "config.json", index_path): + digest.update(path.read_bytes()) + for filename in source_files: + path = (source / filename).resolve() + if not path.is_relative_to(source): + raise ValueError(f"Shard outside source: {filename}") + header = read_header(path) + stat = path.stat() + digest.update(f"{filename}:{stat.st_size}:{stat.st_mtime_ns}".encode()) + for key, spec in header.items(): + if mapping.get(key) == filename: + specs[key] = spec + if set(specs) != set(mapping): + raise ValueError("Source index contains missing tensors") + layers = config["num_hidden_layers"] + mtp_layers = config.get("num_nextn_predict_layers", 0) + active = {} + for key, spec in specs.items(): + match = re.match(r"model\.layers\.(\d+)\.", key) + if match and int(match[1]) >= layers: + if int(match[1]) >= layers + mtp_layers: + raise ValueError(f"Unexpected layer outside decoder/MTP range: {key}") + continue + active[key] = spec + specs = active + expected_shapes = expected_projection_shapes(config) + actual_projections = {key for key in specs if PROJECTION.fullmatch(key)} + if actual_projections != set(expected_shapes): + missing = sorted(set(expected_shapes) - actual_projections)[:5] + extra = sorted(actual_projections - set(expected_shapes))[:5] + raise ValueError(f"Projection coverage mismatch: missing={missing}, extra={extra}") + scales = set() + fp8_count = 0 + output_bytes = 0 + for key, spec in specs.items(): + dtype, shape = spec["dtype"], spec["shape"] + if dtype == "F8_E4M3": + if not PROJECTION.fullmatch(key) or len(shape) != 2: + raise ValueError(f"Unexpected FP8 tensor: {key} {shape}") + if shape != expected_shapes[key]: + raise ValueError(f"Projection shape mismatch: {key}: {shape}, expected {expected_shapes[key]}") + scale_key = key.removesuffix("weight") + "weight_scale_inv" + expected = [math.ceil(shape[0] / block[0]), math.ceil(shape[1] / block[1])] + scale_spec = specs.get(scale_key, {}) + if scale_spec.get("shape") != expected or scale_spec.get("dtype") != "F32": + raise ValueError(f"Missing/invalid block scale: {scale_key}") + scales.add(scale_key) + fp8_count += 1 + output_bytes += (math.prod(shape) * 2 if is_bf16_projection(key) + else math.prod(shape) + shape[0] * 4) + elif key.endswith("weight_scale_inv"): + continue + else: + if PROJECTION.fullmatch(key): + raise ValueError(f"Expected FP8 projection: {key}, got {dtype}") + if dtype not in DTYPE_BYTES: + raise ValueError(f"Unsupported source dtype: {key}: {dtype}") + if len(shape) == 2 and key not in {"model.embed_tokens.weight", "lm_head.weight"} and not re.fullmatch(r"model\.layers\.\d+\.mlp\.gate\.weight", key): + raise ValueError(f"Unexpected unquantized matrix: {key}") + output_bytes += math.prod(shape) * DTYPE_BYTES[dtype] + if scales != {key for key in specs if key.endswith("weight_scale_inv")}: + raise ValueError("Orphan block scales in checkpoint") + expected_count = len(expected_shapes) + if fp8_count != expected_count: + raise ValueError(f"Expected {expected_count} FP8 projections, found {fp8_count}") + groups = {} + for key in sorted(specs): + if key not in scales: + groups.setdefault(mapping[key], []).append(key) + groups = sorted(groups.items()) + config["num_nextn_predict_layers"] = 0 + return config, mapping, specs, groups, digest.hexdigest(), output_bytes, fp8_count + + +def is_bf16_projection(key): + return (key.endswith(".self_attn.kv_b_proj.weight") or bool(re.fullmatch( + r"model\.layers\.\d+\.mlp\.(?:shared_experts\.)?(?:gate_proj|up_proj|down_proj)\.weight", key))) + + +def expected_projection_shapes(config): + hidden = config["hidden_size"] + heads = config["num_attention_heads"] + q_rank, kv_rank = config["q_lora_rank"], config["kv_lora_rank"] + q_head = config["qk_nope_head_dim"] + config["qk_rope_head_dim"] + result = {} + for index in range(config["num_hidden_layers"]): + prefix = f"model.layers.{index}." + for name, shape in { + "q_a_proj": [q_rank, hidden], "q_b_proj": [heads * q_head, q_rank], + "kv_a_proj_with_mqa": [kv_rank + config["qk_rope_head_dim"], hidden], + "kv_b_proj": [heads * (config["qk_nope_head_dim"] + config["v_head_dim"]), kv_rank], + "o_proj": [hidden, heads * config["v_head_dim"]], + }.items(): + result[prefix + "self_attn." + name + ".weight"] = shape + if index < config["first_k_dense_replace"]: + mlps = [("mlp", config["intermediate_size"])] + else: + mlps = [(f"mlp.experts.{e}", config["moe_intermediate_size"]) + for e in range(config["n_routed_experts"])] + if config.get("n_shared_experts"): + mlps.append(("mlp.shared_experts", config["moe_intermediate_size"] * config["n_shared_experts"])) + for mlp, intermediate in mlps: + for name, shape in {"gate_proj": [intermediate, hidden], "up_proj": [intermediate, hidden], + "down_proj": [hidden, intermediate]}.items(): + result[prefix + mlp + "." + name + ".weight"] = shape + return result + + +def dequantize_bf16(weight, block_scale, chunk_rows=1024): + import torch + rows, cols = weight.shape + result = torch.empty((rows, cols), dtype=torch.bfloat16, device="cpu") + columns = torch.arange(cols) // 128 + if not torch.isfinite(block_scale).all() or (block_scale < 0).any(): + raise ValueError("Invalid block scales") + for start in range(0, rows, chunk_rows): + end = min(rows, start + chunk_rows) + row_blocks = torch.arange(start, end) // 128 + value = weight[start:end].float() * block_scale[row_blocks[:, None], columns[None, :]] + value = value.to(torch.bfloat16) + if not torch.isfinite(value).all(): + raise ValueError("Non-finite dequantized BF16 weight") + result[start:end] = value + return result + + +def quantize_weight(weight, block_scale, block=(128, 128), chunk_rows=1024): + """Recover W=FP8*block_scale in FP32, then round(W/(row_absmax/127)).""" + import torch + if weight.ndim != 2 or weight.dtype != torch.float8_e4m3fn: + raise ValueError("Expected 2-D FP8 E4M3 weight") + rows, cols = weight.shape + if tuple(block_scale.shape) != (math.ceil(rows / block[0]), math.ceil(cols / block[1])): + raise ValueError("Block scale shape mismatch") + if chunk_rows < 1 or not torch.isfinite(block_scale).all() or (block_scale < 0).any(): + raise ValueError("Invalid chunk size or block scales") + result = torch.empty((rows, cols), dtype=torch.int8, device="cpu") + scales = torch.empty((rows, 1), dtype=torch.float32, device="cpu") + col_blocks = torch.arange(cols, device="cpu") // block[1] + for start in range(0, rows, chunk_rows): + end = min(rows, start + chunk_rows) + row_blocks = torch.arange(start, end, device="cpu") // block[0] + value = weight[start:end].to(torch.float32) + value.mul_(block_scale[row_blocks[:, None], col_blocks[None, :]]) + if not torch.isfinite(value).all(): + raise ValueError("Non-finite dequantized FP8 weight") + scale = value.abs().amax(dim=1, keepdim=True).div_(127.0) + # A zero row uses scale=1, q=0. Avoid zero scale and underflow. + scale.clamp_(min=torch.finfo(torch.float32).tiny) + scale = torch.where(value.abs().amax(dim=1, keepdim=True) == 0, 1.0, scale) + result[start:end] = value.div_(scale).round_().clamp_(-127, 127).to(torch.int8) + scales[start:end] = scale + return result, scales + + +def output_specs(keys, specs): + expected = {} + for key in keys: + spec = specs[key] + dtype = spec["dtype"] + if dtype == "F8_E4M3": + dtype = "BF16" if is_bf16_projection(key) else "I8" + expected[key] = (dtype, spec["shape"]) + if spec["dtype"] == "F8_E4M3" and not is_bf16_projection(key): + expected[key.removesuffix("weight") + "weight_scale"] = ("F32", [spec["shape"][0], 1]) + return expected + + +def validate_shard(path, expected): + header = read_header(path) + if set(header) != set(expected): + raise ValueError(f"Tensor names mismatch: {path}") + for key, (dtype, shape) in expected.items(): + if header[key]["dtype"] != dtype or header[key]["shape"] != shape: + raise ValueError(f"Tensor dtype/shape mismatch: {path}: {key}") + # safe_open validates offsets and file length without reading tensor data. + from safetensors import safe_open + with safe_open(str(path), framework="pt", device="cpu"): + pass + + +def copy_assets(source, destination): + excluded = {"config.json", "model.safetensors.index.json", "quantize_config.json", + "quantization_config.json", "conversion_state.json", "conversion_report.json"} + for path in source.iterdir(): + if (path.is_file() and path.name not in excluded and not path.name.endswith(".index.json") + and path.suffix in {".json", ".py", ".jinja", ".txt", ".model", ".tiktoken"}): + shutil.copy2(path, destination / path.name) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--threads", type=int, default=8) + parser.add_argument("--chunk-rows", type=int, default=1024) + parser.add_argument("--dry-run", action="store_true", help="Check headers, estimate space; write nothing") + parser.add_argument("--resume", action="store_true", help="Resume matching conversion_state.json") + args = parser.parse_args() + if args.threads < 1 or args.chunk_rows < 1: + parser.error("threads and chunk-rows must be positive") + source, output = args.source.resolve(), args.output.resolve() + if source == output or source in output.parents or output in source.parents: + parser.error("source and output must be separate, non-nested directories") + config, mapping, specs, groups, fingerprint, size, count = plan_source(source) + bf16_count = sum(1 for key, spec in specs.items() + if spec["dtype"] == "F8_E4M3" and is_bf16_projection(key)) + ancestor = output + while not ancestor.exists(): + ancestor = ancestor.parent + free = shutil.disk_usage(ancestor).free + print(json.dumps({"fp8_weights": count, "output_shards": len(groups), + "int8_weights": count - bf16_count, + "bf16_weights": bf16_count, + "bf16_policy": "MLA kv_b_proj, dense MLPs, shared experts", + "mtp": "omitted; main-model inference only", + "estimated_output_GiB": round(size / 2**30, 2), + "free_disk_GiB": round(free / 2**30, 2)}, indent=2), flush=True) + if args.dry_run: + return + import torch + from safetensors import safe_open + from safetensors.torch import save_file + torch.set_num_threads(args.threads) + state_path = output / "conversion_state.json" + identity = {"version": VERSION, "source": str(source), "fingerprint": fingerprint, + "chunk_rows": args.chunk_rows, "quantization_config": quantization_config()} + if args.resume: + state = read_json(state_path) + if state["identity"] != identity: + raise ValueError("Source/settings changed; use a new output directory") + else: + if output.exists() and any(output.iterdir()): + raise ValueError("Output must be empty; use --resume for an interrupted conversion") + state = {"identity": identity, "completed": {}} + completed_bytes = sum((output / name).stat().st_size for name in state["completed"]) + if free < max(0, size - completed_bytes) + 2 * 2**30: + raise ValueError("Insufficient output disk space (including 2 GiB reserve)") + output.mkdir(parents=True, exist_ok=True) + write_json(state_path, state) + output_mapping = {} + total_size = 0 + started = time.monotonic() + for number, (filename, keys) in enumerate(groups, 1): + name = f"model-{number:05d}-of-{len(groups):05d}.safetensors" + path = output / name + expected = output_specs(keys, specs) + if name in state["completed"]: + if path.stat().st_size != state["completed"][name]: + raise ValueError(f"Completed shard changed: {path}") + validate_shard(path, expected) + print(f"[{number}/{len(groups)}] resume: {name}", flush=True) + else: + tensors = {} + with ExitStack() as stack: + readers = {} + + def get(key): + shard = mapping[key] + if shard not in readers: + readers[shard] = stack.enter_context(safe_open( + str(source / shard), framework="pt", device="cpu")) + return readers[shard].get_tensor(key) + + with torch.inference_mode(): + for key in keys: + if specs[key]["dtype"] == "F8_E4M3": + scale_key = key.removesuffix("weight") + "weight_scale_inv" + if is_bf16_projection(key): + tensors[key] = dequantize_bf16(get(key), get(scale_key), chunk_rows=args.chunk_rows) + else: + tensors[key], tensors[key.removesuffix("weight") + "weight_scale"] = ( + quantize_weight(get(key), get(scale_key), chunk_rows=args.chunk_rows)) + else: + tensors[key] = get(key).contiguous() + tmp = path.with_name(path.name + ".tmp") + save_file(tensors, str(tmp), metadata={"format": "pt"}) + validate_shard(tmp, expected) + tmp.replace(path) + del tensors + state["completed"][name] = path.stat().st_size + write_json(state_path, state) + print(f"[{number}/{len(groups)}] wrote {name}; elapsed={time.monotonic()-started:.1f}s", flush=True) + output_mapping.update(dict.fromkeys(expected, name)) + total_size += sum(math.prod(shape) * DTYPE_BYTES[dtype] for dtype, shape in expected.values()) + copy_assets(source, output) + config["quantization_config"] = quantization_config() + write_json(output / "config.json", config) + write_json(output / "model.safetensors.index.json", { + "metadata": {"total_size": total_size}, "weight_map": output_mapping}) + write_json(output / "conversion_report.json", { + "source": str(source), "source_fingerprint": fingerprint, + "fp8_weights_processed": count, "int8_weights": count - bf16_count, + "bf16_weights": bf16_count, "bf16_policy": "MLA kv_b_proj, dense MLPs, shared experts", + "mtp": "omitted", + "output_tensors": len(output_mapping), + "tensor_bytes": total_size, "seconds_this_run": time.monotonic() - started, + "method": "FP8 * block_scale -> FP32 -> per-output-channel INT8 or BF16 (see bf16_policy)", + "calibration": "none; dynamic per-token activation quantization at inference", + }) + print(f"Complete: {output}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/image_build/.dockerignore b/image_build/.dockerignore new file mode 100644 index 0000000..215c725 --- /dev/null +++ b/image_build/.dockerignore @@ -0,0 +1,6 @@ +* +!Dockerfile +!apply_patches.py +!routing_compat.py +!test_route_p800.py +!serve_deepseek.sh diff --git a/image_build/Dockerfile b/image_build/Dockerfile new file mode 100644 index 0000000..88f03ad --- /dev/null +++ b/image_build/Dockerfile @@ -0,0 +1,11 @@ +ARG BASE_IMAGE=xc-llm-kunlun:v0.25.1 +FROM ${BASE_IMAGE} +USER root +COPY apply_patches.py routing_compat.py test_route_p800.py serve_deepseek.sh /opt/kunlun-deepseek-fixes/ +RUN /opt/vllm_kunlun/bin/python /opt/kunlun-deepseek-fixes/apply_patches.py \ + && /opt/vllm_kunlun/bin/python /opt/kunlun-deepseek-fixes/apply_patches.py --check \ + && bash -n /opt/kunlun-deepseek-fixes/serve_deepseek.sh +ENV PATH="/opt/vllm_kunlun/bin:${PATH}" +WORKDIR /workspace +ENTRYPOINT [] +CMD ["/bin/bash"] diff --git a/image_build/README.md b/image_build/README.md new file mode 100644 index 0000000..5f371e4 --- /dev/null +++ b/image_build/README.md @@ -0,0 +1,91 @@ +# 昆仑芯 DeepSeek INT8 修复镜像 + +基础镜像：用户已提交的 `xc-llm-kunlun:v0.25.1`。 +输出镜像：`xc-llm-kunlun:v0.25.1-int8-fixed`。 + +## 修改内容 + +1. `vllm_kunlun/quantization/compressed_tensors/compressed_tensors_moe.py`： + INT8 MoE 从 RoutedExperts 实例读取 scoring_func、num_expert_group、topk_group、e_score_correction_bias 和 routed_scaling_factor，替代旧接口的默认值。避免 DeepSeek sigmoid 分组路由被错误地执行成 softmax；尊重外层 runner 的缩放分工，不硬编码 2.5。 +2. `vllm_kunlun/ops/_custom_ops.py`： + 两个 sigmoid 路由包装函数向 kunlun_ops 传参时，将 `block_static=block_static` 改为 `block_statistic=block_static`。保留上层 Torch 注册接口的 block_static 名称。 + +补丁可重复执行。构建时检查版本、语法、路由参数与底层函数签名，不导入 torch/vllm、不使用 XPU。文件原本已修复时只验证，不重复插入。 + +MLA 接口并未修复，启动脚本使用 `VLLM_MLA_DISABLE=1` 绕过当前 FlashMLAImpl 缺少 forward_mha/forward_mqa 的问题。使用 `VLLM_KUNLUN_ENABLE_INT8_BMM=0`，与已转换权重中的 BF16 kv_b_proj 一致。首测使用 TP=8、不启用 EP、eager、单并发、4096 上下文、每卡 4 GiB KV cache。 + +模型权重不打入镜像。当前转换结果中，路由专家和部分 MLA 投影为 INT8，kv_b_proj、Dense MLP 与共享专家为 BF16；MTP 已排除。 + +## 构建 + +在 172.28.4.201 宿主机执行： + +```bash +cd /mnt/disk1/quant_tools/deepseek_r1_image +docker build --pull=false \ + --build-arg BASE_IMAGE=xc-llm-kunlun:v0.25.1 \ + -t xc-llm-kunlun:v0.25.1-int8-fixed . +``` + +仅针对内置文件验证补丁，不挂卡、不挂模型： + +```bash +docker run --rm xc-llm-kunlun:v0.25.1-int8-fixed \ + /opt/vllm_kunlun/bin/python /opt/kunlun-deepseek-fixes/apply_patches.py --check +``` + +## 启动与验证 + +切换服务前先停止旧服务，释放设备和端口。若旧服务由 srx_test 承载，可停止旧容器后再启动新容器： + +```bash +docker stop srx_test +cd /mnt/disk1/quant_tools/deepseek_r1_image +bash run_container.sh xc-llm-kunlun:v0.25.1-int8-fixed deepseek_int8_fixed +``` + +run_container.sh 挂入八张 XPU、xpuctrl、32 GiB 共享内存，使用 host 网络，沿用实际宿主机映射： + +```text +/mnt/disk0 -> /mnt/data0 +/mnt/disk1 -> /mnt/data1 +``` + +执行真实 P800 路由数值检查（小张量，不加载模型）： + +```bash +docker exec deepseek_int8_fixed \ + /opt/vllm_kunlun/bin/python /opt/kunlun-deepseek-fixes/test_route_p800.py +``` + +启动完整服务（该命令前台运行，测试请求另开终端）： + +```bash +docker exec -it deepseek_int8_fixed \ + bash /opt/kunlun-deepseek-fixes/serve_deepseek.sh +``` + +确认模型文件存在于宿主机 `/mnt/disk1/deepseek-r1-INT8`。镜像内脚本使用容器路径 `/mnt/data1/deepseek-r1-INT8`。 + +服务启动后： + +```bash +curl -sS http://127.0.0.1:8000/v1/models +curl -sS --max-time 300 http://127.0.0.1:8000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"/mnt/data1/deepseek-r1-INT8","messages":[{"role":"user","content":"北京是哪个国家的首都？请简短回答。"}],"temperature":0,"max_tokens":2048}' +``` + +查看 content、reasoning 和 finish_reason。完整输出质量需实际业务问题验证，算子数值检查不能代替该步骤。 + +## 导出与回滚 + +```bash +docker save -o /mnt/disk1/xc-llm-kunlun-v0.25.1-int8-fixed.tar \ + xc-llm-kunlun:v0.25.1-int8-fixed +# 另一台同架构、驱动兼容的 P800 机器： +docker load -i /path/to/xc-llm-kunlun-v0.25.1-int8-fixed.tar +``` + +镜像归档不包含挂载模型；迁移时模型目录需要单独准备。 +回滚时停止新服务/容器，再使用原镜像或启动原容器；旧镜像标签保持不变。 diff --git a/image_build/apply_patches.py b/image_build/apply_patches.py new file mode 100644 index 0000000..30f03e7 --- /dev/null +++ b/image_build/apply_patches.py @@ -0,0 +1,81 @@ +"""Idempotent source patches and CPU checks. No vLLM/torch imports or XPU needed.""" +import argparse +import ast +import copy +import hashlib +import inspect +import json +from importlib import metadata +from pathlib import Path +from types import SimpleNamespace + +from routing_compat import patch_source, regression_tests, INSERT + +parser = argparse.ArgumentParser() +parser.add_argument('--site-packages', type=Path, + default=Path('/opt/vllm_kunlun/lib/python3.10/site-packages')) +parser.add_argument('--check', action='store_true', help='Verify without changing installed files') +args = parser.parse_args() +root = args.site_packages +route = root / 'vllm_kunlun/quantization/compressed_tensors/compressed_tensors_moe.py' +wrapper = root / 'vllm_kunlun/ops/_custom_ops.py' +lower = root / 'kunlun_ops/_xflashinfer_moe.py' +versions = {p: metadata.version(p) for p in ('vllm', 'vllm-kunlun', 'kunlun-ops')} +if not versions['vllm'].startswith('0.25.1') or not versions['vllm-kunlun'].startswith('0.25.1'): + raise RuntimeError(f'Unsupported environment: {versions}') + +before_route = route.read_text() +after_route = patch_source(before_route) +# Reproduce the legacy bug and exercise patched method dispatch. +regression_tests(after_route.replace(INSERT, ''), after_route) +lower_fn = next(n for n in ast.parse(lower.read_text()).body + if isinstance(n, ast.FunctionDef) and n.name == 'moe_sigmoid_group_topk_norm') +signature = inspect.Signature([inspect.Parameter(a.arg, inspect.Parameter.POSITIONAL_OR_KEYWORD) + for a in lower_fn.args.args]) +assert 'block_statistic' in signature.parameters + +before_wrapper = wrapper.read_text() +lines = before_wrapper.splitlines(keepends=True) +names = {'moe_sigmoid_group_topk_norm', 'moe_sigmoid_group_topk_norm_cuda'} +functions = [n for n in ast.parse(before_wrapper).body if isinstance(n, ast.FunctionDef) and n.name in names] +assert len(functions) == 2 +for fn in functions: + for call in ast.walk(fn): + if isinstance(call, ast.Call) and ast.unparse(call.func) == 'kunlun_ops.moe_sigmoid_group_topk_norm': + for kw in call.keywords: + if kw.arg == 'block_static': + idx = kw.lineno - 1 + assert 'block_static=block_static' in lines[idx] + lines[idx] = lines[idx].replace('block_static=block_static', 'block_statistic=block_static') +after_wrapper = ''.join(lines) +for fn in [n for n in ast.parse(after_wrapper).body if isinstance(n, ast.FunctionDef) and n.name in names]: + node = copy.deepcopy(fn) + node.decorator_list = []; node.returns = None + for a in node.args.args: + a.annotation = None + calls = [] + def spy(*a, **kw): + calls.append(signature.bind(*a, **kw).arguments) + scope = {'kunlun_ops': SimpleNamespace(moe_sigmoid_group_topk_norm=spy)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '', 'exec'), scope) + values = {k: object() for k in ('x', 'topk_index', 'norm_score', 'block_static', 'bias')} + scope[fn.name](**values, scale=1.0, n_group=8, topk_group=4) + assert len(calls) == 1 and calls[0]['block_statistic'] is values['block_static'] + assert calls[0]['bias'] is values['bias'] + print('WRAPPER_PASS', fn.name) + +for path, before, after in [(route, before_route, after_route), (wrapper, before_wrapper, after_wrapper)]: + compile(after, str(path), 'exec') + if args.check: + if before != after: + raise RuntimeError(f'Missing patch: {path}') + elif before != after: + backup = path.with_name(path.name + '.bak-image-build') + if backup.exists(): + raise RuntimeError(f'Backup already exists: {backup}') + backup.write_text(before) + temp = path.with_name(path.name + '.image-build.tmp') + temp.write_text(after); temp.chmod(path.stat().st_mode & 0o777); temp.replace(path) + print('SHA256', hashlib.sha256(after.encode()).hexdigest(), path) +print('VERSIONS', json.dumps(versions)) +print('ALL_PATCH_CHECKS_PASS') diff --git a/image_build/routing_compat.py b/image_build/routing_compat.py new file mode 100644 index 0000000..4fc9453 --- /dev/null +++ b/image_build/routing_compat.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""Patch the paused/running container's installed Kunlun INT8 MoE, with backup. + +Run on the Docker host. Tests exercise the actual method through routing dispatch +with tensor allocation/kernel calls mocked; they do not execute XPU kernels. +""" +import ast +import copy +import datetime +import hashlib +import io +import subprocess +import tarfile +from types import SimpleNamespace as NS + +CONTAINER = 'srx_test' +DIRECTORY = '/opt/vllm_kunlun/lib/python3.10/site-packages/vllm_kunlun/quantization/compressed_tensors' +NAME = 'compressed_tensors_moe.py' +CLASS = 'KunlunCompressedTensorsW8A8Int8MoEMethod' +MARKER = '# Kunlun vLLM 0.25 routing compatibility fix' +INSERT = ''' # Kunlun vLLM 0.25 routing compatibility fix + # RoutedExperts.forward_monolithic passes only layer/x/logits/input_ids. + # Read routing from the layer, not legacy method defaults. In particular, + # routed_scaling_factor may be 1 because the runner applies it externally. + scoring_func = layer.scoring_func + num_expert_group = layer.num_expert_group + topk_group = layer.topk_group + e_score_correction_bias = layer.e_score_correction_bias + routed_scaling_factor = layer.routed_scaling_factor + if scoring_func not in ("softmax", "sigmoid"): + raise NotImplementedError(f"Unsupported Kunlun INT8 MoE routing: {scoring_func}") + +''' + + +def method_node(source): + module = ast.parse(source) + cls = next(n for n in module.body if isinstance(n, ast.ClassDef) and n.name == CLASS) + return next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == 'apply_monolithic') + + +def patch_source(source): + if MARKER in source: + return source + node = method_node(source) + first = node.body[0] + if ast.unparse(first) != 'hidden_states = x': + raise RuntimeError('Unexpected installed method; refusing to patch') + lines = source.splitlines(keepends=True) + lines.insert(first.lineno - 1, INSERT) + patched = ''.join(lines) + compile(patched, NAME, 'exec') + return patched + + +class Captured(Exception): + pass + + +def dispatch(source, scoring, groups, top_groups, scale, bias): + node = copy.deepcopy(method_node(source)) + for arg in node.args.posonlyargs + node.args.args + node.args.kwonlyargs: + arg.annotation = None + node.returns = None + calls = [] + + def capture(kind): + def callback(**kwargs): + calls.append((kind, kwargs)) + raise Captured() + return callback + + ops = NS(moe_softmax_topk_norm=capture('softmax'), + moe_sigmoid_group_topk_norm=capture('sigmoid')) + torch = NS(float32='float32', int32='int32', + empty=lambda *a, **k: object(), zeros=lambda *a, **k: object(), ops=NS(_C=ops)) + scope = {'torch': torch} + exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '', 'exec'), scope) + layer = NS(scoring_func=scoring, num_expert_group=groups, topk_group=top_groups, + routed_scaling_factor=scale, e_score_correction_bias=bias, + w13_weight=NS(shape=(256, 512, 128)), w2_weight=NS(shape=(256, 128, 256))) + try: + scope['apply_monolithic'](NS(moe=NS(experts_per_token=8)), layer, + NS(shape=(1, 128), device='cpu'), NS(float=lambda: object()), + input_ids=None) + except Captured: + pass + assert len(calls) == 1 + return calls[0] + + +def regression_tests(original, patched): + bias = object() + if MARKER not in original: + kind, _ = dispatch(original, 'sigmoid', 8, 4, 1.0, bias) + assert kind == 'softmax', 'Expected to reproduce the old routing bug' + print('REPRODUCED: old method used softmax for DeepSeek sigmoid routing') + for scale in (1.0, 2.5): + kind, kwargs = dispatch(patched, 'sigmoid', 8, 4, scale, bias) + assert kind == 'sigmoid' + assert kwargs['n_group'] == 8 and kwargs['topk_group'] == 4 + assert kwargs['bias'] is bias and kwargs['scale'] == scale + assert dispatch(patched, 'softmax', None, None, 1.0, None)[0] == 'softmax' + kind, kwargs = dispatch(patched, 'sigmoid', None, None, 1.0, bias) + assert kind == 'sigmoid' and kwargs['bias'] is bias + print('PASS: sigmoid groups/bias, external/internal scaling, softmax and ungrouped sigmoid') + + +def read_installed(): + data = subprocess.check_output(['docker', 'cp', f'{CONTAINER}:{DIRECTORY}/{NAME}', '-']) + with tarfile.open(fileobj=io.BytesIO(data)) as archive: + entry = next(m for m in archive.getmembers() if m.isfile()) + return entry, archive.extractfile(entry).read() + + +def main(): + entry, original_bytes = read_installed() + original = original_bytes.decode() + patched = patch_source(original) + regression_tests(original, patched) + if patched == original: + print('ALREADY_PATCHED') + return + backup = NAME + '.bak-routing-' + datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f') + output = io.BytesIO() + with tarfile.open(fileobj=output, mode='w') as archive: + for name, content in [(backup, original_bytes), (NAME, patched.encode())]: + info = copy.copy(entry) + info.name, info.size = name, len(content) + archive.addfile(info, io.BytesIO(content)) + subprocess.run(['docker', 'cp', '-', f'{CONTAINER}:{DIRECTORY}'], input=output.getvalue(), check=True) + _, installed = read_installed() + assert installed == patched.encode(), 'Installed content differs from patch' + print('INSTALLED', DIRECTORY + '/' + NAME) + print('BACKUP', DIRECTORY + '/' + backup) + print('SHA256', hashlib.sha256(installed).hexdigest()) + print('Restart the vLLM service to load the fix; existing workers retain old code.') + + +if __name__ == '__main__': + main() diff --git a/image_build/run_container.sh b/image_build/run_container.sh new file mode 100644 index 0000000..787da0d --- /dev/null +++ b/image_build/run_container.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail +image=${1:-vllm-kunlun:deepseek-int8-fixed-20261029} +name=${2:-deepseek_int8_fixed} +devices=(--device=/dev/xpuctrl:/dev/xpuctrl) +for i in {0..7}; do devices+=(--device="/dev/xpu${i}:/dev/xpu${i}"); done +exec docker run -dit --name "$name" --network host --shm-size 32g \ + "${devices[@]}" \ + -v /mnt/disk0:/mnt/data0 \ + -v /mnt/disk1:/mnt/data1 \ + "$image" /bin/bash diff --git a/image_build/serve_deepseek.sh b/image_build/serve_deepseek.sh new file mode 100644 index 0000000..84592b8 --- /dev/null +++ b/image_build/serve_deepseek.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Run inside srx_test, after conversion has completed. +set -euo pipefail +export XPU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +# Avoid the incomplete MLA interface in the installed Kunlun plugin. +export VLLM_MLA_DISABLE=1 +# kv_b_proj remains BF16; keep INT8 BMM disabled. +export VLLM_KUNLUN_ENABLE_INT8_BMM=0 +exec /opt/vllm_kunlun/bin/vllm serve /mnt/data1/deepseek-r1-INT8 \ + --trust-remote-code \ + --dtype bfloat16 \ + --quantization compressed-tensors \ + --tensor-parallel-size 8 \ + --distributed-executor-backend mp \ + --max-model-len 4096 \ + --max-num-seqs 1 \ + --max-num-batched-tokens 4096 \ + --kv-cache-memory-bytes 4294967296 \ + --enforce-eager \ + --no-enable-prefix-caching \ + --no-enable-chunked-prefill \ + --enable-auto-tool-choice \ + --tool-call-parser deepseek_v3 \ + --reasoning-parser deepseek_r1 \ + "$@" diff --git a/image_build/test_route_p800.py b/image_build/test_route_p800.py new file mode 100644 index 0000000..6519795 --- /dev/null +++ b/image_build/test_route_p800.py @@ -0,0 +1,28 @@ +"""Tiny real P800 routing test, run after docker build with an exposed XPU.""" +import torch +from vllm.platforms import current_platform +current_platform.pre_register_and_update() +for count in (2, 128): + for scale in (1.0, 2.5): + torch.manual_seed(17) + logits = torch.randn(count, 256, dtype=torch.float32) + bias = torch.randn(256, dtype=torch.float32) * 0.1 + scores = logits.sigmoid() + adjusted = scores + bias + groups = adjusted.view(count, 8, 32).topk(2, dim=-1).values.sum(-1).topk(4, dim=-1).indices + mask = torch.zeros(count, 8, dtype=torch.bool).scatter_(1, groups, True).repeat_interleave(32, dim=1) + reference_ids = adjusted.masked_fill(~mask, float('-inf')).topk(8, dim=-1).indices + ids = torch.empty(count, 8, dtype=torch.int32, device='cuda') + weights = torch.empty(count, 8, dtype=torch.float32, device='cuda') + stats = torch.zeros(12, 256, dtype=torch.int32, device='cuda') + torch.ops._C.moe_sigmoid_group_topk_norm( + x=logits.cuda(), topk_index=ids, norm_score=weights, block_static=stats, + bias=bias.cuda(), scale=scale, n_group=8, topk_group=4) + torch.cuda.synchronize() + actual_ids = ids.cpu().long() + assert torch.equal(actual_ids.sort(-1).values, reference_ids.sort(-1).values) + expected = scores.gather(1, actual_ids) + expected = expected / expected.sum(-1, keepdim=True) * scale + error = (weights.cpu() - expected).abs().max().item() + assert error < 1e-5, error + print('P800_ROUTE_PASS', count, scale, error, flush=True) diff --git a/routing_fix/patch_moe_routing.py b/routing_fix/patch_moe_routing.py new file mode 100644 index 0000000..4fc9453 --- /dev/null +++ b/routing_fix/patch_moe_routing.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""Patch the paused/running container's installed Kunlun INT8 MoE, with backup. + +Run on the Docker host. Tests exercise the actual method through routing dispatch +with tensor allocation/kernel calls mocked; they do not execute XPU kernels. +""" +import ast +import copy +import datetime +import hashlib +import io +import subprocess +import tarfile +from types import SimpleNamespace as NS + +CONTAINER = 'srx_test' +DIRECTORY = '/opt/vllm_kunlun/lib/python3.10/site-packages/vllm_kunlun/quantization/compressed_tensors' +NAME = 'compressed_tensors_moe.py' +CLASS = 'KunlunCompressedTensorsW8A8Int8MoEMethod' +MARKER = '# Kunlun vLLM 0.25 routing compatibility fix' +INSERT = ''' # Kunlun vLLM 0.25 routing compatibility fix + # RoutedExperts.forward_monolithic passes only layer/x/logits/input_ids. + # Read routing from the layer, not legacy method defaults. In particular, + # routed_scaling_factor may be 1 because the runner applies it externally. + scoring_func = layer.scoring_func + num_expert_group = layer.num_expert_group + topk_group = layer.topk_group + e_score_correction_bias = layer.e_score_correction_bias + routed_scaling_factor = layer.routed_scaling_factor + if scoring_func not in ("softmax", "sigmoid"): + raise NotImplementedError(f"Unsupported Kunlun INT8 MoE routing: {scoring_func}") + +''' + + +def method_node(source): + module = ast.parse(source) + cls = next(n for n in module.body if isinstance(n, ast.ClassDef) and n.name == CLASS) + return next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == 'apply_monolithic') + + +def patch_source(source): + if MARKER in source: + return source + node = method_node(source) + first = node.body[0] + if ast.unparse(first) != 'hidden_states = x': + raise RuntimeError('Unexpected installed method; refusing to patch') + lines = source.splitlines(keepends=True) + lines.insert(first.lineno - 1, INSERT) + patched = ''.join(lines) + compile(patched, NAME, 'exec') + return patched + + +class Captured(Exception): + pass + + +def dispatch(source, scoring, groups, top_groups, scale, bias): + node = copy.deepcopy(method_node(source)) + for arg in node.args.posonlyargs + node.args.args + node.args.kwonlyargs: + arg.annotation = None + node.returns = None + calls = [] + + def capture(kind): + def callback(**kwargs): + calls.append((kind, kwargs)) + raise Captured() + return callback + + ops = NS(moe_softmax_topk_norm=capture('softmax'), + moe_sigmoid_group_topk_norm=capture('sigmoid')) + torch = NS(float32='float32', int32='int32', + empty=lambda *a, **k: object(), zeros=lambda *a, **k: object(), ops=NS(_C=ops)) + scope = {'torch': torch} + exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '', 'exec'), scope) + layer = NS(scoring_func=scoring, num_expert_group=groups, topk_group=top_groups, + routed_scaling_factor=scale, e_score_correction_bias=bias, + w13_weight=NS(shape=(256, 512, 128)), w2_weight=NS(shape=(256, 128, 256))) + try: + scope['apply_monolithic'](NS(moe=NS(experts_per_token=8)), layer, + NS(shape=(1, 128), device='cpu'), NS(float=lambda: object()), + input_ids=None) + except Captured: + pass + assert len(calls) == 1 + return calls[0] + + +def regression_tests(original, patched): + bias = object() + if MARKER not in original: + kind, _ = dispatch(original, 'sigmoid', 8, 4, 1.0, bias) + assert kind == 'softmax', 'Expected to reproduce the old routing bug' + print('REPRODUCED: old method used softmax for DeepSeek sigmoid routing') + for scale in (1.0, 2.5): + kind, kwargs = dispatch(patched, 'sigmoid', 8, 4, scale, bias) + assert kind == 'sigmoid' + assert kwargs['n_group'] == 8 and kwargs['topk_group'] == 4 + assert kwargs['bias'] is bias and kwargs['scale'] == scale + assert dispatch(patched, 'softmax', None, None, 1.0, None)[0] == 'softmax' + kind, kwargs = dispatch(patched, 'sigmoid', None, None, 1.0, bias) + assert kind == 'sigmoid' and kwargs['bias'] is bias + print('PASS: sigmoid groups/bias, external/internal scaling, softmax and ungrouped sigmoid') + + +def read_installed(): + data = subprocess.check_output(['docker', 'cp', f'{CONTAINER}:{DIRECTORY}/{NAME}', '-']) + with tarfile.open(fileobj=io.BytesIO(data)) as archive: + entry = next(m for m in archive.getmembers() if m.isfile()) + return entry, archive.extractfile(entry).read() + + +def main(): + entry, original_bytes = read_installed() + original = original_bytes.decode() + patched = patch_source(original) + regression_tests(original, patched) + if patched == original: + print('ALREADY_PATCHED') + return + backup = NAME + '.bak-routing-' + datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f') + output = io.BytesIO() + with tarfile.open(fileobj=output, mode='w') as archive: + for name, content in [(backup, original_bytes), (NAME, patched.encode())]: + info = copy.copy(entry) + info.name, info.size = name, len(content) + archive.addfile(info, io.BytesIO(content)) + subprocess.run(['docker', 'cp', '-', f'{CONTAINER}:{DIRECTORY}'], input=output.getvalue(), check=True) + _, installed = read_installed() + assert installed == patched.encode(), 'Installed content differs from patch' + print('INSTALLED', DIRECTORY + '/' + NAME) + print('BACKUP', DIRECTORY + '/' + backup) + print('SHA256', hashlib.sha256(installed).hexdigest()) + print('Restart the vLLM service to load the fix; existing workers retain old code.') + + +if __name__ == '__main__': + main() diff --git a/routing_fix/patch_sigmoid_wrapper.py b/routing_fix/patch_sigmoid_wrapper.py new file mode 100644 index 0000000..7745466 --- /dev/null +++ b/routing_fix/patch_sigmoid_wrapper.py @@ -0,0 +1,62 @@ +"""Run inside srx_test: fix lower-level keyword without changing Torch schema.""" +import ast +import copy +import datetime +import inspect +from pathlib import Path +from types import SimpleNamespace + +import kunlun_ops + +path = Path('/opt/vllm_kunlun/lib/python3.10/site-packages/vllm_kunlun/ops/_custom_ops.py') +original = path.read_text() +signature = inspect.signature(kunlun_ops.moe_sigmoid_group_topk_norm) +assert 'block_statistic' in signature.parameters +names = {'moe_sigmoid_group_topk_norm', 'moe_sigmoid_group_topk_norm_cuda'} +tree = ast.parse(original) +functions = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names] +assert len(functions) == 2 +lines = original.splitlines(keepends=True) +changed = 0 +for fn in functions: + for call in ast.walk(fn): + if not isinstance(call, ast.Call) or ast.unparse(call.func) != 'kunlun_ops.moe_sigmoid_group_topk_norm': + continue + for keyword in call.keywords: + if keyword.arg == 'block_static': + index = keyword.lineno - 1 + assert 'block_static=block_static' in lines[index] + lines[index] = lines[index].replace('block_static=block_static', 'block_statistic=block_static') + changed += 1 +assert changed in (0, 2), f'Unexpected matching calls: {changed}' +patched = ''.join(lines) +compile(patched, str(path), 'exec') +# Execute both actual wrapper bodies against the installed lower-level signature. +for fn in [n for n in ast.parse(patched).body if isinstance(n, ast.FunctionDef) and n.name in names]: + node = copy.deepcopy(fn) + node.decorator_list = [] + node.returns = None + for arg in node.args.args: + arg.annotation = None + calls = [] + def spy(*args, **kwargs): + bound = signature.bind(*args, **kwargs) + calls.append(bound.arguments) + scope = {'kunlun_ops': SimpleNamespace(moe_sigmoid_group_topk_norm=spy)} + exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '', 'exec'), scope) + sentinels = {name: object() for name in ('x', 'topk_index', 'norm_score', 'block_static', 'bias')} + scope[fn.name](**sentinels, scale=1.0, n_group=8, topk_group=4) + assert calls[0]['block_statistic'] is sentinels['block_static'] + assert calls[0]['bias'] is sentinels['bias'] + print('WRAPPER_SIGNATURE_PASS', fn.name) +if changed: + backup = path.with_name(path.name + '.bak-sigmoid-keyword-' + datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f')) + backup.write_text(original) + temp = path.with_name(path.name + '.routing-fix.tmp') + temp.write_text(patched) + temp.chmod(path.stat().st_mode & 0o777) + temp.replace(path) + print('BACKUP', backup) + print('PATCHED', path) +else: + print('ALREADY_PATCHED') diff --git a/serve_int8.sh b/serve_int8.sh new file mode 100644 index 0000000..277936a --- /dev/null +++ b/serve_int8.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail +export XPU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +# kv_b_proj is BF16: use the MLA BF16 BMM path. +export VLLM_KUNLUN_ENABLE_INT8_BMM=0 +exec /opt/vllm_kunlun/bin/vllm serve /mnt/data1/deepseek-r1-INT8 \ + --trust-remote-code \ + --dtype bfloat16 \ + --quantization compressed-tensors \ + --tensor-parallel-size 8 \ + --distributed-executor-backend mp \ + --max-model-len 4096 \ + --max-num-seqs 1 \ + --max-num-batched-tokens 4096 \ + --kv-cache-memory-bytes 4294967296 \ + --enforce-eager \ + --no-enable-prefix-caching \ + --no-enable-chunked-prefill \ + --enable-auto-tool-choice \ + --tool-call-parser deepseek_v3 \ + --reasoning-parser deepseek_r1 \ + "$@" diff --git a/serve_int8_no_mla.sh b/serve_int8_no_mla.sh new file mode 100644 index 0000000..7c75f4b --- /dev/null +++ b/serve_int8_no_mla.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail +export XPU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +# Avoid the incomplete MLA interface in the installed Kunlun plugin. +export VLLM_MLA_DISABLE=1 +# kv_b_proj remains BF16; keep INT8 BMM disabled. +export VLLM_KUNLUN_ENABLE_INT8_BMM=0 +exec /opt/vllm_kunlun/bin/vllm serve /mnt/data1/deepseek-r1-INT8 \ + --trust-remote-code \ + --dtype bfloat16 \ + --quantization compressed-tensors \ + --tensor-parallel-size 8 \ + --distributed-executor-backend mp \ + --max-model-len 4096 \ + --max-num-seqs 1 \ + --max-num-batched-tokens 4096 \ + --kv-cache-memory-bytes 4294967296 \ + --enforce-eager \ + --no-enable-prefix-caching \ + --no-enable-chunked-prefill \ + --enable-auto-tool-choice \ + --tool-call-parser deepseek_v3 \ + --reasoning-parser deepseek_r1 \ + "$@" diff --git a/test_conversion.py b/test_conversion.py new file mode 100644 index 0000000..78a0535 --- /dev/null +++ b/test_conversion.py @@ -0,0 +1,123 @@ +"""CPU tests for numerical conversion, MLA, shared experts, MTP and resume.""" +import hashlib +import json +from pathlib import Path +import subprocess +import sys +import tempfile +import unittest + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +import convert_deepseek_fp8_to_int8 as convert + + +def fixture(source): + source.mkdir() + config = { + "model_type": "deepseek_v3", "architectures": ["DeepseekV3ForCausalLM"], + "num_hidden_layers": 2, "num_nextn_predict_layers": 1, "hidden_size": 128, + "q_lora_rank": 128, "kv_lora_rank": 128, "num_attention_heads": 8, + "qk_nope_head_dim": 16, "qk_rope_head_dim": 64, "v_head_dim": 16, + "first_k_dense_replace": 1, "intermediate_size": 256, + "moe_intermediate_size": 128, "n_routed_experts": 2, "n_shared_experts": 1, + "moe_layer_freq": 1, "torch_dtype": "bfloat16", + "quantization_config": {"quant_method": "fp8", "fmt": "e4m3", + "activation_scheme": "dynamic", "weight_block_size": [128, 128]}, + } + (source / "config.json").write_text(json.dumps(config)) + (source / "tokenizer_config.json").write_text('{}') + weights, scales = {}, {} + torch.manual_seed(12) + for name, shape in convert.expected_projection_shapes(config).items(): + weights[name] = torch.randn(*shape).clamp(-5, 5).to(torch.float8_e4m3fn) + scales[name.removesuffix("weight") + "weight_scale_inv"] = ( + torch.rand((shape[0] + 127) // 128, (shape[1] + 127) // 128) * 0.1) + weights["model.embed_tokens.weight"] = torch.ones(8, 128, dtype=torch.bfloat16) + weights["lm_head.weight"] = torch.ones(8, 128, dtype=torch.bfloat16) + weights["model.layers.1.mlp.gate.weight"] = torch.ones(2, 128, dtype=torch.bfloat16) + weights["model.layers.1.mlp.gate.e_score_correction_bias"] = torch.ones(2) + weights["model.norm.weight"] = torch.ones(128, dtype=torch.bfloat16) + # MTP has no destination in the main model and must not leak into output. + weights["model.layers.2.eh_proj.weight"] = torch.ones(128, 256, dtype=torch.bfloat16) + weights["model.layers.2.self_attn.q_a_proj.weight"] = torch.ones(128, 128).to(torch.float8_e4m3fn) + scales["model.layers.2.self_attn.q_a_proj.weight_scale_inv"] = torch.ones(1, 1) + save_file(weights, str(source / "weights.safetensors")) + save_file(scales, str(source / "scales.safetensors")) + mapping = {key: "weights.safetensors" for key in weights} + mapping.update({key: "scales.safetensors" for key in scales}) + (source / "model.safetensors.index.json").write_text(json.dumps({ + "metadata": {"total_size": 999999999999999}, "weight_map": mapping})) + return config, weights, scales + + +class ConversionTests(unittest.TestCase): + def test_numeric(self): + torch.manual_seed(42) + weight = torch.randn(129, 257).to(torch.float8_e4m3fn) + weight[0] = torch.zeros(257).to(torch.float8_e4m3fn) + block_scale = torch.tensor([[0.1, 0.2, 0.3], [0.5, 0.6, 0.7]]) + reference = weight.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:129, :257] + q, scales = convert.quantize_weight(weight, block_scale, chunk_rows=37) + self.assertTrue(torch.all((q.float() * scales - reference).abs() <= scales / 2 + 1e-5)) + self.assertEqual(q[0].count_nonzero().item(), 0) + self.assertEqual(scales[0].item(), 1) + self.assertTrue(torch.equal(convert.dequantize_bf16(weight, block_scale, 37), reference.bfloat16())) + with self.assertRaises(ValueError): + convert.quantize_weight(weight, block_scale * float('nan')) + + def test_checkpoint(self): + with tempfile.TemporaryDirectory() as temp: + source, output = Path(temp) / "source", Path(temp) / "output" + config, weights, scales = fixture(source) + original = {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in source.iterdir()} + command = [sys.executable, str(Path(convert.__file__)), "--source", str(source), + "--output", str(output), "--threads", "1"] + subprocess.run(command + ["--dry-run"], check=True) + self.assertFalse(output.exists()) + subprocess.run(command, check=True) + mapping = json.loads((output / "model.safetensors.index.json").read_text())["weight_map"] + self.assertFalse(any(k.startswith('model.layers.2.') for k in mapping)) + self.assertFalse(any(k.endswith('weight_scale_inv') for k in mapping)) + for key, filename in mapping.items(): + with safe_open(str(output / filename), framework="pt") as reader: + value = reader.get_tensor(key) + if key.endswith("weight_scale"): + self.assertEqual(value.dtype, torch.float32) + self.assertEqual(value.shape[-1], 1) + elif weights[key].dtype != torch.float8_e4m3fn: + self.assertTrue(torch.equal(value, weights[key])) + elif convert.is_bf16_projection(key): + self.assertEqual(value.dtype, torch.bfloat16) + self.assertNotIn(key.removesuffix('weight') + 'weight_scale', mapping) + self.assertTrue(torch.equal(value, convert.dequantize_bf16( + weights[key], scales[key.removesuffix('weight') + 'weight_scale_inv']))) + else: + self.assertEqual(value.dtype, torch.int8) + final_config = json.loads((output / 'config.json').read_text()) + self.assertEqual(final_config['num_nextn_predict_layers'], 0) + self.assertEqual(final_config['quantization_config'], convert.quantization_config()) + for name in ('config.json', 'model.safetensors.index.json', 'conversion_report.json'): + (output / name).unlink() + subprocess.run(command + ['--resume'], check=True) + self.assertTrue((output / 'conversion_report.json').exists()) + self.assertEqual(original, {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in source.iterdir()}) + self.assertNotEqual(subprocess.run(command, capture_output=True).returncode, 0) + + def test_missing_projection_rejected(self): + with tempfile.TemporaryDirectory() as temp: + source = Path(temp) / 'source' + fixture(source) + path = source / 'model.safetensors.index.json' + index = json.loads(path.read_text()) + del index['weight_map']['model.layers.1.mlp.experts.0.up_proj.weight'] + path.write_text(json.dumps(index)) + with self.assertRaisesRegex(ValueError, 'Projection coverage'): + convert.plan_source(source) + + +if __name__ == '__main__': + torch.set_num_threads(1) + unittest.main() diff --git a/validate_kunlun_contract.py b/validate_kunlun_contract.py new file mode 100644 index 0000000..28e76d5 --- /dev/null +++ b/validate_kunlun_contract.py @@ -0,0 +1,82 @@ +"""CPU checks of installed Kunlun INT8 loader and BF16 MLA weight preparation. + +Uses synthetic layers and mocked TP rank/size, not a full engine startup. +""" +import copy +import os +from types import SimpleNamespace +from unittest.mock import patch + +os.environ["VLLM_KUNLUN_ENABLE_INT8_BMM"] = "0" + +import torch +from convert_deepseek_fp8_to_int8 import quantization_config + + +def main(): + from vllm.platforms import current_platform + current_platform.pre_register_and_update() + from vllm_kunlun.quantization.compressed_tensors.compressed_tensors import KunlunCompressedTensorsConfig + from vllm_kunlun.quantization.compressed_tensors.compressed_tensors_moe import KunlunCompressedTensorsMoEMethod + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + cfg = KunlunCompressedTensorsConfig.from_config(copy.deepcopy(quantization_config())) + cfg.packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"], + } + for name, widths in [ + ("model.layers.0.self_attn.fused_qkv_a_proj", [128, 64]), + ("model.layers.0.self_attn.q_b_proj", [128]), + ("model.layers.0.self_attn.o_proj", [128]), + ]: + layer = torch.nn.Linear(128, sum(widths), bias=False) + scheme = cfg.get_scheme(layer, name) + del layer._parameters['weight'] + with patch('vllm.model_executor.parameter.get_tensor_model_parallel_rank', return_value=0), \ + patch('vllm.model_executor.parameter.get_tensor_model_parallel_world_size', return_value=1): + scheme.create_weights(layer, widths, 128, torch.bfloat16, lambda *a, **k: None) + assert type(scheme.kernel).__name__ == 'KunlunScaledMMLinearKernel' + assert layer.weight.dtype == torch.int8 + assert tuple(layer.weight_scale.shape) == (sum(widths), 1) + print('LINEAR_PASS', name, type(scheme.kernel).__name__) + + class RoutedExperts(torch.nn.Module): + pass + + experts = RoutedExperts() + experts.moe_config = SimpleNamespace(is_act_and_mul=True) + method = KunlunCompressedTensorsMoEMethod.get_moe_method(cfg, experts, 'model.layers.3.mlp.experts') + assert type(method).__name__ == 'KunlunCompressedTensorsW8A8Int8MoEMethod' + method.create_weights(experts, 2, 128, 128, torch.bfloat16) + assert tuple(experts.w13_weight_scale.shape) == (2, 256, 1) + with torch.no_grad(): + experts.w13_weight_scale.fill_(0.25) + experts.w2_weight_scale.fill_(0.25) + method.process_weights_after_loading(experts) + assert experts.w13_weight_scale.flatten()[0].item() == 31.75 + print('MOE_SCALE_PASS', type(method).__name__) + for name in ['lm_head', 'model.layers.3.mlp.gate', 'model.layers.0.self_attn.kv_b_proj', + 'model.layers.0.mlp.gate_up_proj', 'model.layers.0.mlp.down_proj', + 'model.layers.3.mlp.shared_experts.gate_up_proj', + 'model.layers.3.mlp.shared_experts.down_proj']: + assert cfg.get_scheme(torch.nn.Linear(128, 128), name) is None + print('IGNORE_PASS: head, router, MLA kv_b_proj, dense MLPs, shared experts') + + from vllm_kunlun.v1.attention.backends.mla.common import MLACommonImpl + projection = torch.nn.Linear(128, 256, bias=False, dtype=torch.bfloat16) + projection.quant_method = UnquantizedLinearMethod() + impl = SimpleNamespace(kv_b_proj=projection, kv_lora_rank=128, + num_heads=2, qk_nope_head_dim=64, v_head_dim=64) + with torch.inference_mode(): + MLACommonImpl.process_weights_after_loading(impl, torch.bfloat16) + assert tuple(impl.W_UK_T.shape) == (2, 64, 128) + assert tuple(impl.W_UV.shape) == (2, 128, 64) + assert impl.W_UK_T.dtype == impl.W_UV.dtype == torch.bfloat16 + expected = projection.weight.T.reshape(128, 2, 128) + assert torch.equal(impl.W_UK_T, expected[:, :, :64].permute(1, 2, 0)) + assert torch.equal(impl.W_UV, expected[:, :, 64:].transpose(0, 1)) + print('MLA_BF16_PREPARATION_PASS') + + +if __name__ == '__main__': + main() diff --git a/validate_sample_weights.py b/validate_sample_weights.py new file mode 100644 index 0000000..f2fb3ed --- /dev/null +++ b/validate_sample_weights.py @@ -0,0 +1,51 @@ +"""Check real DeepSeek layer-0 MLA/Dense and layer-3 shared/routed experts.""" +import argparse +import json +from pathlib import Path + +import torch +from safetensors import safe_open +from convert_deepseek_fp8_to_int8 import quantize_weight, dequantize_bf16, is_bf16_projection + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--source', required=True, type=Path) + parser.add_argument('--report', required=True, type=Path) + args = parser.parse_args() + torch.set_num_threads(8) + mapping = json.loads((args.source / 'model.safetensors.index.json').read_text())['weight_map'] + names = [f'model.layers.0.self_attn.{name}.weight' for name in + ['q_a_proj', 'q_b_proj', 'kv_a_proj_with_mqa', 'kv_b_proj', 'o_proj']] + names += ['model.layers.0.mlp.gate_proj.weight', + 'model.layers.3.mlp.shared_experts.gate_proj.weight', + 'model.layers.3.mlp.experts.0.gate_proj.weight', + 'model.layers.3.mlp.experts.0.down_proj.weight'] + report = [] + for name in names: + scale_name = name.removesuffix('weight') + 'weight_scale_inv' + with safe_open(str(args.source / mapping[name]), framework='pt') as reader: + weight = reader.get_tensor(name) + with safe_open(str(args.source / mapping[scale_name]), framework='pt') as reader: + block_scale = reader.get_tensor(scale_name) + with torch.inference_mode(): + reference = weight.float() * block_scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:weight.shape[0], :weight.shape[1]] + if is_bf16_projection(name): + value = dequantize_bf16(weight, block_scale) + assert torch.equal(value, reference.bfloat16()) + error = value.float() - reference + check = 'exact BF16 rounding' + else: + q, scales = quantize_weight(weight, block_scale) + error = q.float() * scales - reference + assert torch.all(error.abs() <= scales / 2 + scales * 1e-4 + 1e-7), name + check = 'half-scale rounding bound' + item = {'name': name, 'shape': list(weight.shape), 'check': check, + 'relative_rmse': (error.square().mean() / reference.square().mean()).sqrt().item()} + report.append(item) + print(json.dumps(item), flush=True) + args.report.write_text(json.dumps(report, indent=2) + '\n') + + +if __name__ == '__main__': + main()