add-unittest-mrope

Author: sangchengmeng

unit_tests/models/qwen2_vl/test_mrope.py

@@ -43,16 +43,23 @@ def test_mrope_triton_correctness(B, H_q, H_k, L, D, mrope_section):
 
     torch.manual_seed(0)
     device = "cuda"
+    HALF = D // 2
 
     q = torch.rand((B, H_q, L, D), dtype=torch.float32, device=device)
     k = torch.rand((B, H_k, L, D), dtype=torch.float32, device=device)
-    cos = torch.rand((3, 1, L, D), dtype=torch.float32, device=device)
-    sin = torch.rand((3, 1, L, D), dtype=torch.float32, device=device)
 
-    ref_q, ref_k = apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1)
+    cos_half = torch.rand((3, L, HALF), dtype=torch.float32, device=device)
+    sin_half = torch.rand((3, L, HALF), dtype=torch.float32, device=device)
 
-    out_q, out_k = mrope_triton(q, k, cos, sin, axis_map)
+    cos_full = torch.cat([cos_half, cos_half], dim=-1)
+    sin_full = torch.cat([sin_half, sin_half], dim=-1)
 
+    cos_ref = cos_full.unsqueeze(1)
+    sin_ref = sin_full.unsqueeze(1)
+
+    ref_q, ref_k = apply_multimodal_rotary_pos_emb(q, k, cos_ref, sin_ref, mrope_section, unsqueeze_dim=1)
+
+    out_q, out_k = mrope_triton(q, k, cos_half, sin_half, axis_map)
     assert torch.allclose(out_q, ref_q, rtol=1e-3, atol=1e-3)
     assert torch.allclose(out_k, ref_k, rtol=1e-3, atol=1e-3)