first commit

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