import gc
import inspect
import itertools
import time
import weakref
from typing import Dict, List, Optional, Set, Tuple, Union

import numpy as np
import torch
import torch.distributed
import torch.nn as nn

from vllm.attention import AttentionMetadata
from vllm.attention.selector import get_attn_backend
from vllm.attention.backends.utils import CommonAttentionState
from vllm.compilation.compile_context import set_compile_context
from vllm.config import VllmConfig
from vllm.distributed import get_pp_group, get_world_group
from vllm.distributed.parallel_state import graph_capture
from vllm.forward_context import set_forward_context
from vllm.inputs import INPUT_REGISTRY, InputRegistry
from vllm.logger import init_logger
from vllm.lora.layers import LoRAMapping
from vllm.lora.request import LoRARequest
from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
from vllm.model_executor import SamplingMetadataCache
from vllm.model_executor.layers.sampler import SamplerOutput
from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalKwargs,
                             MultiModalRegistry)
from vllm.prompt_adapter.layers import PromptAdapterMapping
from vllm.prompt_adapter.request import PromptAdapterRequest
from vllm.prompt_adapter.worker_manager import (
    LRUCacheWorkerPromptAdapterManager)
from vllm.sampling_params import SamplingParams
from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
from vllm.utils import (GiB_bytes, PyObjectCache,
                        async_tensor_h2d, flatten_2d_lists,
                        is_pin_memory_available)
from vllm.worker.model_runner_base import (ModelRunnerBase,
                                           dump_input_when_exception)
from vllm.model_executor.layers.linear import RowParallelLinear
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
import vllm.worker.model_runner as ori_model_runner
from vllm.worker.model_runner import (
    ModelInputForGPU, ModelInputForGPUWithSamplingMetadata,
    ModelRunnerBase, LORA_WARMUP_RANK,
    _get_max_graph_batch_size, _BATCH_SIZES_TO_CAPTURE
)
from vllm.worker.mlu_model_runner import (ModelInputForMLUBuilder,
                                          MLUModelRunnerBase,
                                          MLUModelRunner, MLUGraphRunner,
                                          mlu_graph_capture)
from vllm_mlu.model_executor.layers.rotary_embedding import MLURotaryEmbedding
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu._mlu_utils import *

logger = init_logger(__name__)


def _model_forward_pre_hook(self, args, kwargs):
    '''
    This hook function will be called before model.forward
    '''
    assert len(args) == 0 and len(kwargs) > 0, \
        f"The pre-forward's expected inputs are not passed by kwargs. " + \
        f"Expected len(args)=0, len(kwargs)>0, " + \
        f"now, len(args)={len(args)}, len(kwargs)={len(kwargs)}."

    input_ids = kwargs['input_ids']
    attn_metadata = kwargs['attn_metadata']
    kv_caches = kwargs['kv_caches']

    # Prepare attributes for all rope in model
    MLURotaryEmbedding.set_mlu_var(input_ids=input_ids,
                                   attn_metadata=attn_metadata)
    RowParallelLinear.is_prompt = True if attn_metadata.prefill_metadata else False
    FeedForward.is_prompt = True if attn_metadata.prefill_metadata else False
    if attn_metadata.prefill_metadata and kv_caches[0][0].numel() > 0:
        set_is_prompt(True)
    else:
        set_is_prompt(False)
    # For unpaged mode to reshape kv to cache
    if not USE_PAGED:
        prefill_meta = attn_metadata.prefill_metadata
        slot_mapping_flat = attn_metadata.slot_mapping.flatten()
        slot_mapping_unpaged = slot_mapping_flat % BlockSizeInfo.BLOCK_SIZE
        batch_ids = slot_mapping_flat[MLURotaryEmbedding.cu_seq_lens[:-1]] // BlockSizeInfo.BLOCK_SIZE
        attn_metadata.cu_seq_lens = MLURotaryEmbedding.cu_seq_lens
        attn_metadata.max_seq_len = prefill_meta.max_prefill_seq_len if prefill_meta else 1
        attn_metadata.slot_mapping_unpaged = None if prefill_meta else slot_mapping_unpaged
        attn_metadata.batch_ids = batch_ids

    return (args, kwargs)


class ModelInputForMLUBuilder_V2(ModelInputForMLUBuilder):

    def build(self) -> ModelInputForGPU:
        """Finalize the builder intermediate data and
        create on-device tensors.
        """
        # Combine and flatten intermediate data.
        input_tokens = []
        for inter_data in self.inter_data_list:
            for cur_input_tokens in inter_data.input_tokens:
                input_tokens.extend(cur_input_tokens)

        if not input_tokens:
            # This may happen when all prefill requests hit
            # prefix caching and there is no decode request.
            return self.model_input_cls()

        mrope_input_positions: Optional[List[List[int]]] = None
        if any(inter_data.mrope_input_positions is not None
                for inter_data in self.inter_data_list):
            mrope_input_positions = [[] for _ in range(3)]
            for idx in range(3):
                for inter_data in self.inter_data_list:
                    msections = inter_data.mrope_input_positions
                    if msections is None:
                        for _seq_input_positions in inter_data.input_positions:
                            mrope_input_positions[idx].extend(
                                _seq_input_positions)
                    else:
                        for _seq_mrope_input_positions in msections:
                            mrope_input_positions[idx].extend(
                                _seq_mrope_input_positions[idx])
            input_positions = None
        else:
            input_positions = []
            for inter_data in self.inter_data_list:
                for cur_input_positions in inter_data.input_positions:
                    input_positions.extend(cur_input_positions)

        seq_lens = []
        query_lens = []
        max_decode_seq_len = 0
        max_encoder_seq_len = 0
        for inter_data in self.inter_data_list:
            seq_lens.extend(inter_data.seq_lens)
            query_lens.extend(inter_data.query_lens)
            if not inter_data.is_prompt:
                max_decode_seq_len = max(max_decode_seq_len,
                                            max(inter_data.seq_lens))
                if self.runner.model_config.is_encoder_decoder:
                    max_encoder_seq_len = max(max_encoder_seq_len,
                                                inter_data.encoder_seq_len)

        # Mapping from request IDs to sequence IDs. Used for Jamba models
        # that manages the cache by itself.
        request_ids_to_seq_ids = {
            data.request_id: data.seq_ids
            for data in self.inter_data_list
        }

        cuda_graph_pad_size = self._get_cuda_graph_pad_size(
            num_seqs=len(seq_lens),
            max_decode_seq_len=max_decode_seq_len,
            max_encoder_seq_len=max_encoder_seq_len)

        batch_size = len(input_tokens)
        if cuda_graph_pad_size != -1:
            # If cuda graph can be used, pad tensors accordingly.
            # See `capture_model` API for more details.
            # vLLM uses cuda graph only for decoding requests.
            batch_size += cuda_graph_pad_size

        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: check mlugraph enable
        ''' 
        # If mlu graph is enabled, but still run with eager mode,
        # print warning message.
        # This behavior can be disabled by export VLLM_GRAPH_DEBUG=0.
        if VLLM_GRAPH_DEBUG \
                and self.decode_only \
                and not self.runner.model_config.enforce_eager \
                and cuda_graph_pad_size == -1 \
                and max_decode_seq_len > self.runner.max_seq_len_to_capture:
            logger.warning(f"Because one of the following conditions is not met, MLU Graphs will not be enabled.\n" +
                        f"1. batch_size({batch_size}) <= max_batch_size_to_capture({self.runner.max_batchsize_to_capture})\n" +
                        f"2. max_seq_len({max_decode_seq_len}) <= max_seq_len_to_capture({self.runner.max_seq_len_to_capture})\n" +
                        f"Use 'export VLLM_GRAPH_DEBUG=false' to disable this warning.")        
        '''
        ==================
        End of MLU Hijack
        ==================
        '''

        # Tokens and positions.
        if cuda_graph_pad_size:
            input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size))
        assert self.runner.device is not None
        input_tokens_tensor = async_tensor_h2d(input_tokens, torch.long,
                                                self.runner.device,
                                                self.runner.pin_memory)
        if mrope_input_positions is not None:
            for idx in range(3):
                mrope_input_positions[idx].extend(
                    itertools.repeat(0, cuda_graph_pad_size))
            input_positions_tensor = async_tensor_h2d(mrope_input_positions,
                                                        torch.int32,
                                                        self.runner.device,
                                                        self.runner.pin_memory)
        else:
            input_positions.extend(itertools.repeat(0, cuda_graph_pad_size))
            input_positions_tensor = async_tensor_h2d(input_positions,
                                                        torch.int32,
                                                        self.runner.device,
                                                        self.runner.pin_memory)
        # Sequence and query lengths.
        if cuda_graph_pad_size:
            seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size))

        # Attention metadata.
        attn_metadata = self.attn_metadata_builder.build(
            seq_lens, query_lens, cuda_graph_pad_size, batch_size)

        # LoRA data.
        lora_requests = set()
        lora_mapping = None
        if self.enable_lora:
            lora_requests = set(r for data in self.inter_data_list
                                for r in data.lora_requests)
            lora_index_mapping = flatten_2d_lists([
                flatten_2d_lists(inter_data.lora_index_mapping)
                for inter_data in self.inter_data_list
            ])
            if cuda_graph_pad_size:
                lora_index_mapping.extend(
                    itertools.repeat(0, cuda_graph_pad_size))
            lora_prompt_mapping = flatten_2d_lists([
                flatten_2d_lists(inter_data.lora_prompt_mapping)
                for inter_data in self.inter_data_list
            ])

            lora_mapping = LoRAMapping(
                **dict(index_mapping=lora_index_mapping,
                        prompt_mapping=lora_prompt_mapping,
                        is_prefill=not self.decode_only))

        # Prompt adapter data.
        prompt_adapter_requests: Set[PromptAdapterRequest] = set()
        prompt_adapter_mapping = None
        if self.enable_prompt_adapter:
            prompt_adapter_requests = set(
                data.prompt_adapter_request for data in self.inter_data_list
                if data.prompt_adapter_request is not None)
            prompt_adapter_index_mapping = flatten_2d_lists([
                inter_data.prompt_adapter_index_mapping
                for inter_data in self.inter_data_list
            ])
            if cuda_graph_pad_size:
                prompt_adapter_index_mapping.extend(
                    itertools.repeat(0, cuda_graph_pad_size))
            prompt_adapter_prompt_mapping = flatten_2d_lists([
                inter_data.prompt_adapter_prompt_mapping
                for inter_data in self.inter_data_list
            ])
            prompt_adapter_mapping = PromptAdapterMapping(
                prompt_adapter_index_mapping,
                prompt_adapter_prompt_mapping,
            )

        # Multi-modal data.
        multi_modal_kwargs_list = [
            data.multi_modal_kwargs for data in self.inter_data_list
            if data.multi_modal_kwargs is not None
        ]
        multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)

        return self.model_input_cls(
            input_tokens=input_tokens_tensor,
            input_positions=input_positions_tensor,
            attn_metadata=attn_metadata,
            seq_lens=seq_lens,
            query_lens=query_lens,
            lora_mapping=lora_mapping,
            lora_requests=lora_requests,
            multi_modal_kwargs=multi_modal_kwargs,
            request_ids_to_seq_ids=request_ids_to_seq_ids,
            finished_requests_ids=self.finished_requests_ids,
            prompt_adapter_mapping=prompt_adapter_mapping,
            prompt_adapter_requests=prompt_adapter_requests)


class MLUModelRunnerBase_V2(MLUModelRunnerBase):

    def __init__(
        self,
        vllm_config: VllmConfig,
        kv_cache_dtype: Optional[str] = "auto",
        is_driver_worker: bool = False,
        return_hidden_states: bool = False,
        input_registry: InputRegistry = INPUT_REGISTRY,
        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
    ):

        ModelRunnerBase.__init__(self, vllm_config)
        model_config = self.model_config
        cache_config = self.cache_config

        self.is_driver_worker = is_driver_worker
        self.return_hidden_states = return_hidden_states

        self.device = self.device_config.device
        self.pin_memory = is_pin_memory_available()

        self.kv_cache_dtype = kv_cache_dtype
        self.sliding_window = model_config.get_sliding_window()
        self.block_size = cache_config.block_size
        BlockSizeInfo.set_block_size(self.block_size)
        self.max_seq_len_to_capture = self.model_config.max_seq_len_to_capture
        self.max_batchsize_to_capture = _get_max_graph_batch_size(
            self.scheduler_config.max_num_seqs)

        self.graph_runners: List[Dict[int, MLUGraphRunner]] = [
            {} for _ in range(self.parallel_config.pipeline_parallel_size)
        ]
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: support context mlugraph
        ''' 
        self.context_graph_runners: List[MLUGraphRunner] = [
            None for _ in range(self.parallel_config.pipeline_parallel_size)
        ]
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        self.graph_memory_pool: Optional[Tuple[
            int, int]] = None  # Set during graph capture.

        self.has_inner_state = model_config.has_inner_state

        # When using CUDA graph, the input block tables must be padded to
        # max_seq_len_to_capture. However, creating the block table in
        # Python can be expensive. To optimize this, we cache the block table
        # in numpy and only copy the actual input content at every iteration.
        # The shape of the cached block table will be
        # (max batch size to capture, max seq len to capture / block size).
        self.graph_block_tables = np.zeros(
            (self.max_batchsize_to_capture, self.get_max_block_per_batch()),
            dtype=np.int32)

        # Attention-free but stateful models like Mamba need a placeholder attn
        # backend, as the attention metadata is needed to manage internal state.
        # However we must bypass attention selection altogether for some models
        # used for speculative decoding to avoid a divide-by-zero in
        # model_config.get_head_size()
        num_attn_heads = self.model_config.get_num_attention_heads(
            self.parallel_config)
        needs_attn_backend = (num_attn_heads != 0
                              or self.model_config.is_attention_free)

        self.attn_backend = get_attn_backend(
            self.model_config.get_head_size(),
            self.model_config.dtype,
            self.kv_cache_dtype,
            self.block_size,
            self.model_config.is_attention_free,
            use_mla=self.model_config.is_deepseek_v2,
        ) if needs_attn_backend else None
        if self.attn_backend:
            self.attn_state = self.attn_backend.get_state_cls()(
                weakref.proxy(self))
        else:
            self.attn_state = CommonAttentionState(weakref.proxy(self))

        # Multi-modal data support
        self.input_registry = input_registry
        self.mm_registry = mm_registry
        self.multi_modal_input_mapper = mm_registry \
            .create_input_mapper(model_config)
        self.mm_registry.init_mm_limits_per_prompt(self.model_config)

        # Lazy initialization
        self.model: nn.Module  # Set after load_model
        # Set after load_model.
        self.lora_manager: Optional[LRUCacheWorkerLoRAManager] = None
        self.prompt_adapter_manager: LRUCacheWorkerPromptAdapterManager = None

        set_cpu_offload_max_bytes(
            int(self.cache_config.cpu_offload_gb * 1024**3))

        # Used to cache python objects
        self.inter_data_cache: Dict[int, PyObjectCache] = {}

        # Using the PythonizationCache in Pipeline-Parallel clobbers the
        # SequenceGroupToSample object. In Pipeline-Parallel, we have
        # more than 1 Scheduler, resulting in a potential back-to-back
        # prepare_model_inputs() call. This clobbers the cached
        # SequenceGroupToSample objects, as we reset the cache during
        # every prepare_model_inputs() call.
        self.sampling_metadata_cache: SamplingMetadataCache = \
              SamplingMetadataCache() \
                if self.parallel_config.pipeline_parallel_size == 1 else None

    def init_llmassistor(self):
        try :
            import llmassistor
        except Exception as e:
            raise ImportError(f"import llmassistor failed: {e}")
        llmdebug = llmassistor.LLMDebug(
            local_rank=get_world_group().local_rank,
        )
        llmdebug.trace(self.model)
        self.llmdebug = llmdebug

    def get_max_seq_len_for_profile_run(self, max_mm_tokens: int) -> int:
        """
        For some multimodal model, for example, llava, the max_mm_tokens(image|video) + 
        max_position_embeddings(text) is smaller than the input len when profile run, which
        causes runtime error for rope.
        So here, we calculate the max seq len for profile run, and compare it with
        max_mm_tokens + max_position_embeddings for a WalkAround.
        """
        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
        max_num_seqs = self.scheduler_config.max_num_seqs
        max_num_seqs = max(1, min(max_num_seqs, max_num_batched_tokens // max_mm_tokens))
        acc_seq_len = 0
        for group_id in range(max_num_seqs):
            seq_len = (max_num_batched_tokens // max_num_seqs +
                        (group_id < max_num_batched_tokens % max_num_seqs))
            acc_seq_len += seq_len
        return acc_seq_len

    def load_model(self) -> None:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: 1. modify rope's max_position_embeddings to max_model_len
                2. register model pre forward for rope optimization 
        '''
        MLURotaryEmbedding.max_seq_len = self.model_config.max_model_len
        if self.model_config.multimodal_config is not None:
            # for multimodal models, the max seq len for language model is
            # max_mm_tokens + max_position_embeddings
            max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(self.model_config)
            profile_run_max_seq_len = self.get_max_seq_len_for_profile_run(max_mm_tokens)
            MLURotaryEmbedding.max_seq_len = max(profile_run_max_seq_len,
                                                 MLURotaryEmbedding.max_seq_len + max_mm_tokens)
        super(MLUModelRunnerBase, self).load_model()
        self.model.register_forward_pre_hook(_model_forward_pre_hook, with_kwargs=True)
        if VLLM_DUMP_OUTPUTS:
            self.init_llmassistor()
        '''
        ==================
        End of MLU Hijack
        ==================
        '''

    def get_max_block_per_batch(self) -> int:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: support unpaged mode 
        '''
        block_size = self.block_size
        if USE_PAGED:
            return (self.max_seq_len_to_capture + block_size - 1) // block_size
        else:
            return 1
        '''
        ==================
        End of MLU Hijack
        ==================
        '''

    @torch.inference_mode()
    def profile_run(self) -> None:
        # Avoid dumping useless profile dummy tensors.
        if VLLM_DUMP_OUTPUTS:
            self.llmdebug.enable = False

        # Enable top-k sampling to reflect the accurate memory usage.
        sampling_params = SamplingParams(top_p=0.99, top_k=self.vocab_size - 1)
        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
        max_num_seqs = self.scheduler_config.max_num_seqs
        # This represents the maximum number of different requests
        # that will have unique loras, an therefore the max amount of memory
        # consumption create dummy lora request copies from the lora request
        # passed in, which contains a lora from the lora warmup path.
        dummy_lora_requests: List[LoRARequest] = []
        dummy_lora_requests_per_seq: List[LoRARequest] = []
        if self.lora_config:
            assert self.lora_manager is not None
            with self.lora_manager.dummy_lora_cache():
                for idx in range(self.lora_config.max_loras):
                    lora_id = idx + 1
                    dummy_lora_request = LoRARequest(
                        lora_name=f"warmup_{lora_id}",
                        lora_int_id=lora_id,
                        lora_path="/not/a/real/path",
                    )
                    self.lora_manager.add_dummy_lora(dummy_lora_request,
                                                        rank=LORA_WARMUP_RANK)
                    dummy_lora_requests.append(dummy_lora_request)
                dummy_lora_requests_per_seq = [
                    dummy_lora_requests[idx % len(dummy_lora_requests)]
                    for idx in range(max_num_seqs)
                ]

        # Profile memory usage with max_num_sequences sequences and the total
        # number of tokens equal to max_num_batched_tokens.
        seqs: List[SequenceGroupMetadata] = []
        # Additional GPU memory may be needed for multi-modal encoding, which
        # needs to be accounted for when calculating the GPU blocks for
        # vLLM blocker manager.
        # To exercise the worst scenario for GPU memory consumption,
        # the number of seqs (batch_size) is chosen to maximize the number
        # of images processed.

        max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(
            self.model_config)
        if max_mm_tokens > 0:
            max_num_seqs_orig = max_num_seqs
            max_num_seqs = min(max_num_seqs,
                                max_num_batched_tokens // max_mm_tokens)
            if max_num_seqs < 1:
                expr = (f"min({max_num_seqs_orig}, "
                        f"{max_num_batched_tokens} // {max_mm_tokens})")
                logger.warning(
                    "Computed max_num_seqs (%s) to be less than 1. "
                    "Setting it to the minimum value of 1.", expr)
                max_num_seqs = 1

        batch_size = 0
        for group_id in range(max_num_seqs):
            seq_len = (max_num_batched_tokens // max_num_seqs +
                        (group_id < max_num_batched_tokens % max_num_seqs))
            batch_size += seq_len

            dummy_data = self.input_registry \
                .dummy_data_for_profiling(self.model_config,
                                            seq_len,
                                            self.mm_registry)

            seq = SequenceGroupMetadata(
                request_id=str(group_id),
                is_prompt=True,
                seq_data={group_id: dummy_data.seq_data},
                sampling_params=sampling_params,
                block_tables=None,
                lora_request=dummy_lora_requests_per_seq[group_id]
                if dummy_lora_requests_per_seq else None,
                multi_modal_data=dummy_data.multi_modal_data,
                multi_modal_placeholders=dummy_data.multi_modal_placeholders,
            )
            seqs.append(seq)

        # Run the model with the dummy inputs.
        num_layers = self.model_config.get_num_layers(self.parallel_config)
        # use an empty tensor instead of `None`` to force Dynamo to pass
        # it by reference, rather by specializing on the value ``None``.
        # the `dtype` argument does not matter, and we use `float32` as
        # a placeholder (it has wide hardware support).
        # it is important to create tensors inside the loop, rather than
        # multiplying the list, to avoid Dynamo from treating them as
        # tensor aliasing.
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: support kv cache int8
        '''
        kv_caches = []
        for _ in range(num_layers):
            kv_cache_ = torch.tensor([], dtype=torch.float32, device=self.device)
            kv_cache_scale_ = torch.tensor([], dtype=torch.float32, device=self.device)
            kv_caches.append([kv_cache_, kv_cache_scale_])
        '''
        ==================
        End of MLU Hijack
        ==================
        '''

        finished_requests_ids = [seq.request_id for seq in seqs]
        model_input = self.prepare_model_input(
            seqs, finished_requests_ids=finished_requests_ids)
        intermediate_tensors = None
        if not get_pp_group().is_first_rank:
            intermediate_tensors = self.model.make_empty_intermediate_tensors(
                batch_size=batch_size,
                dtype=self.model_config.dtype,
                device=self.device)

        graph_batch_size = self.max_batchsize_to_capture
        batch_size_capture_list = [
            bs for bs in _BATCH_SIZES_TO_CAPTURE if bs <= graph_batch_size
        ]
        if self.model_config.enforce_eager:
            batch_size_capture_list = []
        with set_compile_context(batch_size_capture_list):
            self.execute_model(model_input, kv_caches, intermediate_tensors)
        torch.mlu.synchronize()

        if VLLM_DUMP_OUTPUTS:
            self.llmdebug.enable = True

        return

    @torch.inference_mode()
    def capture_model(
        self,
        kv_caches: List[List[torch.Tensor]],
        num_gpu_blocks: int
    ) -> None:
        """Cuda graph capture a model.

        Note that CUDA graph's performance gain is negligible if number
        of batched tokens are larger than 200. And since CUDA graph
        requires fixed sized tensors, supporting large/variable batch
        size requires high GPU memory overhead. Thus, vLLM only captures
        decoding requests. Mixed batch (chunked prefill + decoding) or
        prefill requests are not captured.

        Since it is used for decoding-only, it assumes there's only 1 token
        per sequence in the batch.
        """
        assert not self.model_config.enforce_eager
        logger.info("Capturing the model for MLU graphs. This may lead to "
                    "unexpected consequences if the model is not static. To "
                    "run the model in eager mode, set 'enforce_eager=True' or "
                    "use '--enforce-eager' in the CLI.")
        logger.info("MLU graphs can take additional 1~3 GiB memory per MLU. "
                    "If you are running out of memory, consider decreasing "
                    "`gpu_memory_utilization` or enforcing eager mode. "
                    "You can also reduce the `max_num_seqs` as needed "
                    "to decrease memory usage.")
        start_time = time.perf_counter()
        start_free_gpu_memory = torch.mlu.mem_get_info()[0]

        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: support fewer gpu blocks mlugraph capture 
        '''
        # Prepare dummy inputs. These will be reused for all batch sizes.
        max_batch_size = min(self.max_batchsize_to_capture, num_gpu_blocks)
        input_tokens = torch.zeros(max_batch_size, dtype=torch.long).mlu()
        input_positions = torch.zeros(max_batch_size, dtype=torch.int32).mlu()
        if self.model_config.uses_mrope:
            input_positions = torch.tile(input_positions, (3, 1))
        # Prepare dummy previous_hidden_states only if needed by the model.
        # This is used by draft models such as EAGLE.
        previous_hidden_states = None
        if "previous_hidden_states" in inspect.signature(
                self.model.forward).parameters:
            previous_hidden_states = torch.empty(
                [max_batch_size,
                 self.model_config.get_hidden_size()],
                dtype=self.model_config.dtype,
                device=self.device)

        intermediate_inputs = None
        if not get_pp_group().is_first_rank:
            intermediate_inputs = self.model.make_empty_intermediate_tensors(
                batch_size=max_batch_size,
                dtype=self.model_config.dtype,
                device=self.device)

        graph_batch_size = self.max_batchsize_to_capture
        graph_batch_size = min(graph_batch_size, max_batch_size)
        batch_size_capture_list = [
            bs for bs in ori_model_runner._BATCH_SIZES_TO_CAPTURE if bs <= graph_batch_size
        ]

        with self.attn_state.graph_capture(
                max_batch_size), mlu_graph_capture() as graph_capture_context:
            # NOTE: Capturing the largest batch size first may help reduce the
            # memory usage of CUDA graph.
            for virtual_engine in range(
                    self.parallel_config.pipeline_parallel_size):
                for batch_size in reversed(batch_size_capture_list):
                    attn_metadata = (
                        self.attn_state.graph_capture_get_metadata_for_batch(
                            batch_size,
                            is_encoder_decoder_model=self.model_config.
                            is_encoder_decoder))

                    if self.lora_config:
                        lora_mapping = LoRAMapping(
                            **dict(index_mapping=[0] * batch_size,
                                   prompt_mapping=[0] * batch_size,
                                   is_prefill=False))
                        self.set_active_loras(set(), lora_mapping)

                    if self.prompt_adapter_config:
                        prompt_adapter_mapping = PromptAdapterMapping(
                            [-1] * batch_size,
                            [-1] * batch_size,
                        )
                        self.set_active_prompt_adapters(
                            set(), prompt_adapter_mapping)

                    graph_runner = MLUGraphRunner(
                        self.model, self.attn_backend.get_name(),
                        self.attn_state.graph_clone(batch_size),
                        self.model_config.is_encoder_decoder)

                    capture_inputs = {
                        "input_ids":
                        input_tokens[:batch_size],
                        "positions":
                        input_positions[..., :batch_size],
                        "intermediate_inputs":
                        intermediate_inputs[:batch_size]
                        if intermediate_inputs is not None else None,
                        "kv_caches":
                        kv_caches[virtual_engine],
                        "attn_metadata":
                        attn_metadata,
                        "memory_pool":
                        self.graph_memory_pool,
                        "stream":
                        graph_capture_context.stream
                    }
                    if previous_hidden_states is not None:
                        capture_inputs[
                            "previous_hidden_states"] = previous_hidden_states[:
                                                                               batch_size]

                    if self.has_inner_state:
                        # Only used by Mamba-based models CUDA graph atm (Jamba)
                        capture_inputs.update({
                            "seqlen_agnostic_capture_inputs":
                            self.model.get_seqlen_agnostic_capture_inputs(
                                batch_size)
                        })
                    if self.model_config.is_encoder_decoder:
                        # add the additional inputs to capture for
                        # encoder-decoder models.
                        self._update_inputs_to_capture_for_enc_dec_model(
                            capture_inputs)

                    with set_forward_context(attn_metadata):
                        graph_runner.capture(**capture_inputs)
                    self.graph_memory_pool = graph_runner.graph.pool()
                    self.graph_runners[virtual_engine][batch_size] = (
                        graph_runner)

        end_time = time.perf_counter()
        end_free_gpu_memory = torch.mlu.mem_get_info()[0]
        elapsed_time = end_time - start_time
        mlu_graph_size = start_free_gpu_memory - end_free_gpu_memory
        # This usually takes < 10 seconds.
        logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
                    elapsed_time, mlu_graph_size / GiB_bytes)

        # Set the actually capture batch list to _BATCH_SIZES_TO_CAPTURE
        ori_model_runner._BATCH_SIZES_TO_CAPTURE = batch_size_capture_list
        '''
        ==================
        End of MLU Hijack
        ==================
        '''

    @torch.inference_mode()
    def capture_model_with_context(
        self,
        kv_caches: List[List[torch.Tensor]],
        num_gpu_blocks: int
    ) -> None:
        """Cuda graph capture a model.

        Note that CUDA graph's performance gain is negligible if number
        of batched tokens are larger than 200. And since CUDA graph
        requires fixed sized tensors, supporting large/variable batch
        size requires high GPU memory overhead. Thus, vLLM only captures
        decoding requests. Mixed batch (chunked prefill + decoding) or
        prefill requests are not captured.

        Since it is used for decoding-only, it assumes there's only 1 token
        per sequence in the batch.

        NOTE: Different from capture_model function, we capture both context
        and decoder cuda graph here.
        """

        assert not self.has_inner_state, \
            "Mlugraph does not support mamba-based models(Jamba)"

        enable_context_mlugraph = True
        ctx_graph_bs, ctx_graph_seq_len = self.model_config.get_context_mlugraph_bs_and_seq()
        if ctx_graph_bs > max(ori_model_runner._BATCH_SIZES_TO_CAPTURE):
            enable_context_mlugraph = False
            logger.warning(f"Context mlugraph batch size {ctx_graph_bs} is greater "
                        f"than max batch size to capture {max(ori_model_runner._BATCH_SIZES_TO_CAPTURE)}, "
                        "can not enable context mlugraph.")
        if ctx_graph_bs > num_gpu_blocks:
            enable_context_mlugraph = False
            logger.warning(f"Context mlugraph batch size {ctx_graph_bs} is greater "
                        f"than available gpu blocks {num_gpu_blocks}, can not enable "
                        "context mlugraph.")
        if ctx_graph_seq_len > self.max_seq_len_to_capture:
            enable_context_mlugraph = False
            logger.warning(f"Context mlugraph sequence length {ctx_graph_seq_len} is greater "
                        f"than max sequence length to capture: {self.max_seq_len_to_capture}, "
                        "can not enable context mlugraph.")

        if enable_context_mlugraph:
            logger.info(f"Enable context mlugraph for batch size {ctx_graph_bs} and "
                        f"sequence length {ctx_graph_seq_len}")
        logger.info("Capturing the model for MLU graphs. This may lead to "
                    "unexpected consequences if the model is not static. To "
                    "run the model in eager mode, set 'enforce_eager=True' or "
                    "use '--enforce-eager' in the CLI.")
        logger.info("MLU graphs can take additional 1~3 GiB memory per MLU. "
                    "If you are running out of memory, consider decreasing "
                    "`gpu_memory_utilization` or enforcing eager mode. "
                    "You can also reduce the `max_num_seqs` as needed "
                    "to decrease memory usage.")
        start_time = time.perf_counter()
        start_free_gpu_memory = torch.mlu.mem_get_info()[0]

        # Prepare dummy inputs. These will be reused for all batch sizes.
        max_batch_size = min(self.max_batchsize_to_capture, num_gpu_blocks)
        max_num_tokens = max_batch_size
        if enable_context_mlugraph:
            max_batch_size = max(max_batch_size, ctx_graph_bs)
            max_num_tokens = max(max_num_tokens, ctx_graph_bs * ctx_graph_seq_len)

        input_tokens = torch.zeros(max_num_tokens, dtype=torch.long).mlu()
        input_positions = torch.zeros(max_num_tokens, dtype=torch.int32).mlu()
        if self.model_config.uses_mrope:
            input_positions = torch.tile(input_positions, (3, 1))
        # Prepare dummy previous_hidden_states only if needed by the model.
        # This is used by draft models such as EAGLE.
        previous_hidden_states = None
        if "previous_hidden_states" in inspect.signature(
                self.model.forward).parameters:
            previous_hidden_states = torch.empty(
                [max_batch_size,
                self.model_config.get_hidden_size()],
                dtype=self.model_config.dtype,
                device=self.device)

        intermediate_inputs = None
        if not get_pp_group().is_first_rank:
            intermediate_inputs = self.model.make_empty_intermediate_tensors(
                batch_size=max_num_tokens,
                dtype=self.model_config.dtype,
                device=self.device)

        graph_batch_size = self.max_batchsize_to_capture
        graph_batch_size = min(graph_batch_size, max_batch_size)
        batch_size_capture_list = [
            bs for bs in ori_model_runner._BATCH_SIZES_TO_CAPTURE if bs <= graph_batch_size
        ]

        def capture_mlugraph(is_prefill, batch_size, num_tokens, attn_metadata):
            if self.lora_config:
                lora_mapping = LoRAMapping(
                    **dict(index_mapping=[0] * num_tokens,
                        prompt_mapping=[0] * batch_size,
                        is_prefill=is_prefill))
                self.set_active_loras(set(), lora_mapping)

            if self.prompt_adapter_config:
                prompt_adapter_mapping = PromptAdapterMapping(
                    [-1] * num_tokens,
                    [-1] * batch_size,
                )
                self.set_active_prompt_adapters(
                    set(), prompt_adapter_mapping)

            capture_inputs = {
                "input_ids":
                input_tokens[:num_tokens],
                "positions":
                input_positions[..., :num_tokens],
                "intermediate_inputs":
                intermediate_inputs[:num_tokens]
                if intermediate_inputs is not None else None,
                "kv_caches":
                kv_caches[virtual_engine],
                "attn_metadata":
                attn_metadata,
                "memory_pool":
                self.graph_memory_pool,
                "stream":
                graph_capture_context.stream
            }
            if previous_hidden_states is not None:
                capture_inputs["previous_hidden_states"
                            ] = previous_hidden_states[:batch_size]

            graph_runner = MLUGraphRunner(
                self.model, self.attn_backend.get_name(),
                self.attn_state.graph_clone(batch_size),
                self.model_config.is_encoder_decoder)

            with set_forward_context(attn_metadata):
                graph_runner.capture(**capture_inputs)
            self.graph_memory_pool = graph_runner.graph.pool()
            if is_prefill:
                self.context_graph_runners[virtual_engine] = graph_runner
            else:
                self.graph_runners[virtual_engine][batch_size] = (
                    graph_runner)

        with self.attn_state.graph_capture_with_context(
                    ctx_graph_bs, max_batch_size, max_num_tokens
                ), mlu_graph_capture() as graph_capture_context:
            # NOTE: Capturing the largest batch size first may help reduce the
            # memory usage of CUDA graph.
            # NOTE: If enable context mlugraph, because the output buffer is reused and
            # created when the first mlugraph captured, we must compare max token num for
            # context and decoder to decide which process should be capture first.
            for virtual_engine in range(self.parallel_config.pipeline_parallel_size):
                # capture mlugraph for context when hiddens_states size(bs * seq_len)
                # is bigger than decoder(max bs size in bs capture list)
                if enable_context_mlugraph and (ctx_graph_bs * ctx_graph_seq_len 
                                                >= batch_size_capture_list[-1]):
                    self.attn_state.fill_seq_lens_tensor(ctx_graph_seq_len)
                    context_attn_metadata = (
                        self.attn_state.graph_capture_get_metadata_for_context(
                        ctx_graph_bs, ctx_graph_seq_len))
                    capture_mlugraph(True, ctx_graph_bs, ctx_graph_bs * ctx_graph_seq_len,
                                    context_attn_metadata)

                # capture mlugraph for decoder
                self.attn_state.fill_seq_lens_tensor(1)
                for batch_size in reversed(batch_size_capture_list):
                    decoder_attn_metadata = (
                        self.attn_state.graph_capture_get_metadata_for_batch(batch_size))
                    capture_mlugraph(False, batch_size, batch_size, decoder_attn_metadata)

                # capture mlugraph for context when hiddens_states size(bs * seq_len)
                # is smaller than decoder(max bs size in bs capture list)
                if enable_context_mlugraph and (ctx_graph_bs * ctx_graph_seq_len 
                                                < batch_size_capture_list[-1]):
                    self.attn_state.fill_seq_lens_tensor(ctx_graph_seq_len)
                    context_attn_metadata = (
                        self.attn_state.graph_capture_get_metadata_for_context(
                        ctx_graph_bs, ctx_graph_seq_len))
                    capture_mlugraph(True, ctx_graph_bs, ctx_graph_bs * ctx_graph_seq_len,
                                    context_attn_metadata)

        end_time = time.perf_counter()
        end_free_gpu_memory = torch.mlu.mem_get_info()[0]
        elapsed_time = end_time - start_time
        mlu_graph_size = start_free_gpu_memory - end_free_gpu_memory
        # This usually takes < 10 seconds.
        logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
                    elapsed_time, mlu_graph_size / GiB_bytes)

        # Set the actually capture batch list to _BATCH_SIZES_TO_CAPTURE
        ori_model_runner._BATCH_SIZES_TO_CAPTURE = batch_size_capture_list

    def reset_capture_context(self):
        self.graph_runners = [
            {} for _ in range(self.parallel_config.pipeline_parallel_size)
        ]
        self.context_graph_runners = [
            None for _ in range(self.parallel_config.pipeline_parallel_size)
        ]
        self.graph_memory_pool = None  # Set during graph capture.

        gc.collect()
        torch.mlu.empty_cache()


class MLUModelRunner_V2(MLUModelRunner):

    @torch.inference_mode()
    @dump_input_when_exception(exclude_args=[0], exclude_kwargs=["self"])
    def execute_model(
        self,
        model_input: ModelInputForGPUWithSamplingMetadata,
        kv_caches: List[torch.Tensor],
        intermediate_tensors: Optional[IntermediateTensors] = None,
        num_steps: int = 1,
    ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:
        if num_steps > 1:
            raise ValueError("num_steps > 1 is not supported in ModelRunner")

        if self.lora_config:
            assert model_input.lora_requests is not None
            assert model_input.lora_mapping is not None
            self.set_active_loras(model_input.lora_requests,
                                  model_input.lora_mapping)

        if self.prompt_adapter_config:
            assert model_input.prompt_adapter_requests is not None
            assert model_input.prompt_adapter_mapping is not None
            self.set_active_prompt_adapters(
                model_input.prompt_adapter_requests,
                model_input.prompt_adapter_mapping)

        self.attn_state.begin_forward(model_input)

        # Currently cuda graph is only supported by the decode phase.
        assert model_input.attn_metadata is not None
        prefill_meta = model_input.attn_metadata.prefill_metadata
        decode_meta = model_input.attn_metadata.decode_metadata
        # TODO(andoorve): We can remove this once all
        # virtual engines share the same kv cache.
        virtual_engine = model_input.virtual_engine
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: support context mlugraph
        '''
        use_context_mlugraph = False
        if prefill_meta is None and decode_meta.use_cuda_graph:
            assert model_input.input_tokens is not None
            graph_batch_size = model_input.input_tokens.shape[0]
            model_executable = self.graph_runners[virtual_engine][
                graph_batch_size]
        # The context mlugraph is None when profile run
        elif (decode_meta is None and prefill_meta.use_cuda_graph
              and self.context_graph_runners[virtual_engine] is not None):
            use_context_mlugraph = True
            model_executable = self.context_graph_runners[virtual_engine]
        else:
            model_executable = self.model
        '''
        ==================
        End of MLU Hijack
        ==================
        '''

        multi_modal_kwargs = model_input.multi_modal_kwargs or {}
        seqlen_agnostic_kwargs = {
            "finished_requests_ids": model_input.finished_requests_ids,
            "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
        } if self.has_inner_state else {}
        if (self.observability_config is not None
                and self.observability_config.collect_model_forward_time):
            model_forward_start = torch.mlu.Event(enable_timing=True)
            model_forward_end = torch.mlu.Event(enable_timing=True)
            model_forward_start.record()

        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: add mlu metrics 
        '''
        # Add time markers for model_executable+compute_logits
        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
            use_cuda_graph = ((prefill_meta is None and decode_meta.use_cuda_graph)
                              or use_context_mlugraph)
            # if use_cuda_graph, the start timestamp will be inserted inside MLUGraphRunner.forward()
            if not use_cuda_graph:
                start = torch.mlu.Event(enable_timing=True)
                start.record()

        with set_forward_context(model_input.attn_metadata):
            hidden_or_intermediate_states = model_executable(
                input_ids=model_input.input_tokens,
                positions=model_input.input_positions,
                kv_caches=kv_caches,
                attn_metadata=model_input.attn_metadata,
                intermediate_tensors=intermediate_tensors,
                **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
                                             device=self.device),
                **seqlen_agnostic_kwargs)

        if (self.observability_config is not None
                and self.observability_config.collect_model_forward_time):
            model_forward_end.record()

        # Compute the logits in the last pipeline stage.
        if not get_pp_group().is_last_rank:
            if (self.is_driver_worker
                    and hidden_or_intermediate_states is not None
                    and isinstance(hidden_or_intermediate_states,
                                   IntermediateTensors)
                    and self.observability_config is not None
                    and self.observability_config.collect_model_forward_time):
                model_forward_end.synchronize()
                model_forward_time = model_forward_start.elapsed_time(
                    model_forward_end)
                orig_model_forward_time = 0.0
                if intermediate_tensors is not None:
                    orig_model_forward_time = intermediate_tensors.tensors.get(
                        "model_forward_time", torch.tensor(0.0)).item()
                hidden_or_intermediate_states.tensors["model_forward_time"] = (
                    torch.tensor(model_forward_time + orig_model_forward_time))
            return hidden_or_intermediate_states

        logits = self.model.compute_logits(hidden_or_intermediate_states,
                                           model_input.sampling_metadata)

        # Add time markers for model_executable+compute_logits
        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
            end_marker = torch.mlu.Event(enable_timing=True)
            end_marker.record()
            if use_cuda_graph:
                self.time_markers = (model_executable.start, end_marker)
            else:
                self.time_markers = (start, end_marker)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        if not self.is_driver_worker:
            return []

        if model_input.async_callback is not None:
            model_input.async_callback()

        # Sample the next token.
        output: SamplerOutput = self.model.sample(
            logits=logits,
            sampling_metadata=model_input.sampling_metadata,
        )
        if (self.observability_config is not None
                and self.observability_config.collect_model_forward_time
                and output is not None):
            model_forward_end.synchronize()
            model_forward_time = model_forward_start.elapsed_time(
                model_forward_end)
            orig_model_forward_time = 0.0
            if intermediate_tensors is not None:
                orig_model_forward_time = intermediate_tensors.tensors.get(
                    "model_forward_time", torch.tensor(0.0)).item()
            # If there are multiple workers, we are still tracking the latency
            # from the start time of the driver worker to the end time of the
            # driver worker. The model forward time will then end up covering
            # the communication time as well.
            output.model_forward_time = (orig_model_forward_time +
                                         model_forward_time)

        if self.return_hidden_states:
            # we only need to pass hidden states of most recent token
            assert model_input.sampling_metadata is not None
            indices = model_input.sampling_metadata.selected_token_indices
            if model_input.is_prompt:
                hidden_states = hidden_or_intermediate_states.index_select(
                    0, indices)
                output.prefill_hidden_states = hidden_or_intermediate_states
            elif decode_meta.use_cuda_graph:
                hidden_states = hidden_or_intermediate_states[:len(indices)]
            else:
                hidden_states = hidden_or_intermediate_states

            output.hidden_states = hidden_states

        return [output]


class MLUGraphRunner_V2(MLUGraphRunner):

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
        **kwargs,
    ) -> torch.Tensor:
        # KV caches are fixed tensors, so we don't need to copy them.
        del kv_caches

        # Copy the input tensors to the input buffers.
        self.input_buffers["input_ids"].copy_(input_ids, non_blocking=True)
        self.input_buffers["positions"].copy_(positions, non_blocking=True)

        if self.backend_name != "NO_ATTENTION":
            self.input_buffers["slot_mapping"].copy_(
                attn_metadata.slot_mapping, non_blocking=True)

        self.attn_state.prepare_graph_input_buffers(
            self.input_buffers, attn_metadata, self._is_encoder_decoder_model)

        if "seqlen_agnostic_capture_inputs" in self.input_buffers:
            self.model.copy_inputs_before_cuda_graphs(self.input_buffers,
                                                      **kwargs)

        if "previous_hidden_states" in self.input_buffers:
            self.input_buffers["previous_hidden_states"].copy_(
                kwargs["previous_hidden_states"], non_blocking=True)

        if intermediate_tensors is not None:
            for key in intermediate_tensors.tensors:
                if key != "model_execute_time" and key != "model_forward_time":
                    self.input_buffers[key].copy_(intermediate_tensors[key],
                                                  non_blocking=True)
        if self._is_encoder_decoder_model:
            self.input_buffers["encoder_input_ids"].copy_(
                kwargs['encoder_input_ids'], non_blocking=True)
            self.input_buffers["encoder_positions"].copy_(
                kwargs['encoder_positions'], non_blocking=True)

        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: add mlu metrics
        '''
        # Add time markers for MLUGraph mode
        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
            self.start = torch.mlu.Event(enable_timing=True)
            self.start.record()
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        # Run the graph.
        self.graph.replay()
        # Return the output tensor.
        if get_pp_group().is_last_rank:
            return self.output_buffers["hidden_states"]

        return self.output_buffers


MluHijackObject.apply_hijack(ModelInputForMLUBuilder,
                             ModelInputForMLUBuilder.build,
                             ModelInputForMLUBuilder_V2.build)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             MLUModelRunnerBase.__init__,
                             MLUModelRunnerBase_V2.__init__)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             "init_llmassistor",
                             MLUModelRunnerBase_V2.init_llmassistor)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             "get_max_seq_len_for_profile_run",
                             MLUModelRunnerBase_V2.get_max_seq_len_for_profile_run)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             MLUModelRunnerBase.load_model,
                             MLUModelRunnerBase_V2.load_model)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             MLUModelRunnerBase.get_max_block_per_batch,
                             MLUModelRunnerBase_V2.get_max_block_per_batch)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             MLUModelRunnerBase.profile_run,
                             MLUModelRunnerBase_V2.profile_run)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             MLUModelRunnerBase.capture_model,
                             MLUModelRunnerBase_V2.capture_model)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             "capture_model_with_context",
                             MLUModelRunnerBase_V2.capture_model_with_context)
MluHijackObject.apply_hijack(MLUModelRunnerBase,
                             "reset_capture_context",
                             MLUModelRunnerBase_V2.reset_capture_context)
MluHijackObject.apply_hijack(MLUModelRunner,
                             MLUModelRunner.execute_model,
                             MLUModelRunner_V2.execute_model)
MluHijackObject.apply_hijack(MLUGraphRunner,
                             MLUGraphRunner.forward,
                             MLUGraphRunner_V2.forward)
