first commit

validate_kunlun_contract.py

@@ -0,0 +1,62 @@
+"""Check installed Kunlun loader's configuration and CPU parameter contracts.
+
+This uses synthetic layers and a mocked TP rank, not a full engine startup.
+"""
+import copy
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from convert_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
+
+    cfg = KunlunCompressedTensorsConfig.from_config(copy.deepcopy(quantization_config()))
+    cfg.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
+    linear = torch.nn.Linear(128, 128, bias=False)
+    scheme = cfg.get_scheme(linear, "model.layers.0.self_attn.qkv_proj")
+    # CPU-only parameter allocation; no distributed processes or accelerator work.
+    del linear._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(linear, [128], 128, torch.bfloat16, lambda *a, **k: None)
+    assert type(scheme.kernel).__name__ == "KunlunScaledMMLinearKernel"
+    assert linear.weight.dtype == torch.int8
+    assert tuple(linear.weight_scale.shape) == (128, 1)
+    print("LINEAR_PASS", type(scheme.kernel).__name__)
+
+    class RoutedExperts(torch.nn.Module):
+        pass
+
+    experts = RoutedExperts()
+    experts.ckpt_gate_proj_name = "w1"
+    experts.ckpt_down_proj_name = "w2"
+    experts.ckpt_up_proj_name = "w3"
+    experts.moe_config = SimpleNamespace(is_act_and_mul=True)
+    method = KunlunCompressedTensorsMoEMethod.get_moe_method(
+        cfg, experts, "model.layers.0.block_sparse_moe.experts")
+    assert type(method).__name__ == "KunlunCompressedTensorsW8A8Int8MoEMethod"
+    method.create_weights(experts, 2, 128, 128, torch.bfloat16)
+    assert tuple(experts.w13_weight.shape) == (2, 256, 128)
+    assert tuple(experts.w13_weight_scale.shape) == (2, 256, 1)
+    assert tuple(experts.w2_weight_scale.shape) == (2, 128, 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_PASS", type(method).__name__)
+    print("SCALE_PASS: checkpoint scale * 127 = runtime absmax")
+    assert cfg.get_scheme(torch.nn.Linear(128, 128), "lm_head") is None
+    assert cfg.get_scheme(torch.nn.Linear(128, 128), "model.layers.0.block_sparse_moe.gate") is None
+    print("IGNORE_PASS: lm_head and router gate")
+
+
+if __name__ == "__main__":
+    main()