import torch
from typing import Optional
from vllm import _mlu_ops as mlu_ops
from vllm.distributed import get_tensor_model_parallel_rank
from vllm.distributed.parallel_state import get_tensor_model_parallel_group
from vllm_mlu._mlu_utils import get_is_prompt
from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantLinearMethod
from vllm_mlu.mlu_hijack_utils import MluHijackObject


def vllm_mlu__model_executor__layers__quantization__smoothquant__SmoothQuantLinearMethod__apply(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
    residual: Optional[torch.Tensor] = None
) -> torch.Tensor:
    quant_input = None
    input_scale = None
    if self.quant_config.input_quant_method == "per_token":
        quant_input, input_scale = mlu_ops.per_token_smooth_quantize(x, layer.smooth, None)
    if self.quant_config.input_quant_method == "per_tensor":
        quant_input = x if self.skip_quant_input else mlu_ops.quantize(x, layer.scale_to_int, None)

    '''
    =====================================================
    Modify by Context Communication Computation Parallel
    =====================================================
    @brief: call parallel op
    '''
    if hasattr(self, 'parallel_num') and get_is_prompt():
        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)
        params = [cncl_comm, quant_input, input_scale, layer.qweight, layer.per_channel_scale, 
                self.compute_dtype, bias, residual, 1.0, 1.0, self.parallel_num]
        out = mlu_ops.smooth_quant_matmul_allreduce(*params)
    else:
        out = mlu_ops.smooth_quant_matmul(quant_input, input_scale, layer.qweight,
                                            layer.per_channel_scale, self.compute_dtype, bias, residual)
    '''
    =====================================================
    End of Context Communication Computation Parallel
    =====================================================
    '''
    return out


MluHijackObject.apply_hijack(SmoothQuantLinearMethod,
                             SmoothQuantLinearMethod.apply,
                             vllm_mlu__model_executor__layers__quantization__smoothquant__SmoothQuantLinearMethod__apply)
