import torch
from typing import Optional
from vllm.attention import AttentionMetadata
from vllm import _mlu_ops as mlu_ops
from vllm_mlu._mlu_utils import *
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.distributed import get_tensor_model_parallel_rank
from vllm.distributed.parallel_state import get_tensor_model_parallel_group
from vllm.model_executor.models.qwen2 import Qwen2Attention
from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantLinearMethod


def vllm__model_executor__models__qwen2__Qwen2Attention__forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
    residual: Optional[torch.Tensor] = None,
    smooth_quant_scale: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    assert smooth_quant_scale is None
    qkv, _ = self.qkv_proj(hidden_states)
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: pack q & k to fit tmo.apply_rotary
    '''
    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    '''
    =====================================================
    Modify by Context Communication Computation Parallel
    =====================================================
    @brief: call flash_attn_sq_mm_allreduce to finish forward
    '''
    if (attn_metadata.prefill_metadata) and \
       (kv_cache[0].numel() > 0) and \
       (hasattr(self.o_proj, 'quant_method')) and \
       (isinstance(self.o_proj.quant_method, SmoothQuantLinearMethod)) and \
       (self.o_proj.quant_method.quant_config.input_quant_method == "per_token"):
        rank = get_tensor_model_parallel_rank()
        pg = get_tensor_model_parallel_group().device_group
        cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)

        return torch.ops.vllm.context_attn_comm_cmpt_parallel_flash_attention_v2(
                    q, k, v,
                    self.num_heads, self.head_dim, self.num_kv_heads,
                    kv_cache, self.attn.impl.kv_cache_dtype,
                    1.0, 1.0, self.scaling,
                    cncl_comm,
                    self.o_proj.smooth, self.o_proj.qweight,
                    self.o_proj.per_channel_scale.to(torch.float),
                    self.o_proj.quant_method.parallel_num,
                    residual, self.attn.impl.sliding_window, self.attn.impl.alibi_slopes
                )
    '''
    =====================================================
    End of Context Communication Computation Parallel
    =====================================================
    '''
    attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add residual
    '''
    output, _ = self.o_proj(attn_output, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return output


MluHijackObject.undo_hijack(Qwen2Attention, 
                            Qwen2Attention.forward)
MluHijackObject.apply_hijack(Qwen2Attention,
                             Qwen2Attention.forward,
                             vllm__model_executor__models__qwen2__Qwen2Attention__forward)