[add] mrope triton

Author: sangchengmeng

lightllm/models/qwen2_vl/layer_infer/transformer_layer_infer.py

@@ -7,7 +7,6 @@
 from lightllm.models.qwen2_vl.triton_kernel.mrope import mrope_triton
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 
-
 class Qwen2VLTransformerLayerInfer(LlamaTransformerLayerInfer):
     def __init__(self, layer_num, network_config, mode=[]):
         super().__init__(layer_num, network_config, mode)
@@ -19,10 +18,10 @@ def _get_qkv(self, input, cache_kv, infer_state, layer_weight):
             input, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
         ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
         seq_len, _ = q.shape
-        q = q.view(1, seq_len, -1, self.head_dim_).transpose(1, 2)
-        k = cache_kv[:, : self.tp_k_head_num_, :].view(1, seq_len, -1, self.head_dim_).transpose(1, 2)
+        q = q.view(1, seq_len, -1, self.head_dim_).transpose(1, 2).contiguous()
+        k = cache_kv[:, : self.tp_k_head_num_, :].view(1, seq_len, -1, self.head_dim_).transpose(1, 2).contiguous()
         new_q, new_k = mrope_triton(q, k, infer_state.position_cos, infer_state.position_sin, self.mrope_section)
-        new_q = new_q.transpose(1, 2).reshape(1, seq_len, -1)
+        new_q = new_q.transpose(1, 2).reshape(1, seq_len, -1).contiguous()
         cache_kv[:, : self.tp_k_head_num_, :] = new_k.squeeze(0).permute(1, 0, 2)
 
         return new_q, cache_kv

lightllm/models/qwen2_vl/triton_kernel/mrope.py

@@ -3,7 +3,6 @@
 import triton
 import triton.language as tl
 
-
 @triton.jit
 def mrope_kernel_combined(
     Q_ptr,
@@ -170,6 +169,3 @@ def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim
     print(f"torch {t_ref:.2f} ms/iter")
     print(f"triton {t_tri:.2f} ms/iter")
 
-
-if __name__ == "__main__":
-    test()