first commit

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
+```

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("<Q", raw)[0]
+        if size > 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()

image_build/.dockerignore

@@ -0,0 +1,6 @@
+*
+!Dockerfile
+!apply_patches.py
+!routing_compat.py
+!test_route_p800.py
+!serve_deepseek.sh

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"]

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
+```
+
+镜像归档不包含挂载模型；迁移时模型目录需要单独准备。
+回滚时停止新服务/容器，再使用原镜像或启动原容器；旧镜像标签保持不变。

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=[])), '<wrapper-check>', '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')

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=[])), '<actual-method>', '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()

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

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 \
+  "$@"

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)

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=[])), '<actual-method>', '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()

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=[])), '<wrapper-test>', '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')

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 \
+  "$@"

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 \
+  "$@"

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()

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()

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()