0813-add-unit-tests

Author: SangChengC

lightllm/models/qwen2_5_vl/qwen2_5_visual.py

@@ -209,8 +209,6 @@ def __init__(
         self.processor = Qwen2VLImageProcessor(**processor_config_dict)
 
         self._init_datatype()
-        self.load_model(kvargs["weight_dir"])
-        self.cuda()
 
     def _init_datatype(self):
         if isinstance(self.data_type, torch.dtype):

unit_tests/models/qwen2_vl/test_rotary_pos_emb.py

@@ -0,0 +1,55 @@
+import math
+import torch
+import pytest
+
+from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(tensor: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    orig_dtype = tensor.dtype
+    tensor = tensor.float()
+    cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+    sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+    output = (tensor * cos) + (rotate_half(tensor) * sin)
+    output = output.to(orig_dtype)
+    return output
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (16, 1296, 64, 80),
+        (2, 1024, 2, 192),
+        (1, 1024, 1, 256),
+        (2, 1024, 3, 160),
+    ],
+)
+def test_triton_matches_reference(shape):
+    B, L, H, D = shape
+    assert D % 2 == 0
+
+    torch.manual_seed(0)
+
+    freqs = torch.randn(L, D // 2, device="cuda", dtype=torch.bfloat16)
+    cos = freqs.cos()
+    sin = freqs.sin()
+
+    tensor = torch.randn(B, L, H, D, device="cuda", dtype=torch.bfloat16)
+
+    ref = apply_rotary_pos_emb_vision(tensor, cos, sin)
+    out = apply_rotary_pos_emb_triton(tensor, cos, sin)
+
+    assert out.dtype == tensor.dtype, "输出 dtype 应与输入一致"
+    assert out.shape == tensor.shape, "输出形状应与输入一致"
+    assert torch.allclose(out, ref, rtol=1e-2, atol=1e-2), "Triton 与参考实现不一致"
+
+
+if __name__ == "__main__":
+    pytest.main()