{"repo": "cambricon/vllm-mlu", "files": {"vllm_mlu/v1/worker/gpu_model_runner.py": {"dc984838c63aef3de46bbadedd82d9848a19c5d2": {"t": 1785674303.8937426, "n": 180895, "h": "eb8acf657e8fae3c", "c": "# SPDX-License-Identifier: Apache-2.0\n# SPDX-FileCopyrightText: Copyright contributors to the vLLM-MLU project\n\n# SPDX-License-Identifier: Apache-2.0\n\nfrom copy import copy\nimport gc\nimport time\nfrom contextlib import contextmanager\nfrom itertools import product\nfrom typing import TYPE_CHECKING, Dict, List, Tuple, cast\nimport re\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom tqdm import tqdm\nimport cnpx\n\nimport vllm.envs as envs\nfrom vllm.attention.layer import Attention\nfrom vllm.compilation.counter import compilation_counter\nfrom vllm.compilation.monitor import set_cudagraph_capturing_enabled\nfrom vllm.config import (\n    CompilationMode,\n    CUDAGraphMode,\n    VllmConfig,\n    get_layers_from_vllm_config,\n)\nfrom vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer\nfrom vllm.distributed.eplb.eplb_state import EplbState\nfrom vllm.distributed.kv_transfer import has_kv_transfer_group, get_kv_transfer_group\nfrom vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase\nfrom vllm.distributed.parallel_state import (\n    get_dcp_group,\n    get_pp_group,\n    get_tp_group,\n    graph_capture,\n    is_global_first_rank,\n    prepare_communication_buffer_for_model,\n    get_tensor_model_parallel_rank,\n)\nfrom vllm.forward_context import BatchDescriptor, set_forward_context\nfrom vllm.logger import init_logger\nfrom vllm.lora.layers import LoRAMapping\nfrom vllm.lora.request import LoRARequest\nfrom vllm.model_executor.layers.attention_layer_base import AttentionLayerBase\nfrom vllm.model_executor.layers.linear import RowParallelLinear\nfrom vllm.model_executor.model_loader import get_model_loader\nfrom vllm.model_executor.models.interfaces import (\n    SupportsMultiModal,\n    is_mixture_of_experts,\n    supports_eagle3,\n    supports_multimodal_pruning,\n)\nfrom vllm.model_executor.models.interfaces_base import VllmModelForPooling\nfrom vllm.multimodal import MULTIMODAL_REGISTRY\nfrom vllm.multimodal.utils import group_mm_kwargs_by_modality\nfrom vllm.sampling_params import SamplingType\nfrom vllm.sequence import IntermediateTensors\nfrom vllm.utils.import_utils import LazyLoader\nfrom vllm.utils.math_utils import cdiv\nfrom vllm.utils.mem_constants import GiB_bytes\nfrom vllm.utils.mem_utils import DeviceMemoryProfiler\nfrom vllm.utils.platform_utils import is_pin_memory_available\nfrom vllm.utils.torch_utils import (\n    get_dtype_size,\n    kv_cache_dtype_str_to_dtype,\n    supports_dynamo,\n    weak_ref_tensor,\n)\nfrom vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder\nfrom vllm.v1.attention.backends.mla.common import MLACommonMetadata\nfrom vllm.v1.attention.backends.utils import (\n    CommonAttentionMetadata,\n    get_dcp_local_seq_lens,\n    split_attn_metadata,\n)\nfrom vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput\nfrom vllm.v1.cudagraph_dispatcher import CudagraphDispatcher\nfrom vllm.v1.kv_cache_interface import (\n    AttentionSpec,\n    EncoderOnlyAttentionSpec,\n    KVCacheConfig,\n    KVCacheSpec,\n    FullAttentionSpec,\n    MambaSpec,\n)\nfrom vllm.v1.outputs import (\n    EMPTY_MODEL_RUNNER_OUTPUT,\n    KVConnectorOutput,\n    LogprobsTensors,\n    ModelRunnerOutput,\n    make_empty_encoder_model_runner_output,\n    LogprobsLists,\n    SamplerOutput,\n)\nfrom vllm.v1.sample.logits_processor import build_logitsprocs\nfrom vllm.v1.sample.metadata import SamplingMetadata\nfrom vllm.v1.sample.rejection_sampler import RejectionSampler\nfrom vllm.v1.sample.sampler import Sampler\nfrom vllm.v1.spec_decode.eagle import EagleProposer\nfrom vllm.v1.spec_decode.medusa import MedusaProposer\nfrom vllm.v1.spec_decode.metadata import SpecDecodeMetadata\nfrom vllm.v1.spec_decode.ngram_proposer import NgramProposer\nfrom vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer\nfrom vllm.v1.utils import record_function_or_nullcontext\nfrom vllm.v1.worker.dp_utils import coordinate_batch_across_dp\nfrom vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch\nfrom vllm.v1.worker.gpu_model_runner import (\n    AsyncGPUModelRunnerOutput,\n    ExecuteModelState,\n    GPUModelRunner,\n    PerLayerAttnMetadata,\n)\nfrom vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper\nfrom vllm.v1.worker.ubatch_utils import (\n    UBatchSlice,\n    UBatchSlices,\n    check_ubatch_thresholds,\n)\nfrom vllm.v1.worker.utils import (\n    AttentionGroup,\n    MultiModalBudget,\n    bind_kv_cache,\n    is_residual_scattered_for_sp,\n    sanity_check_mm_encoder_outputs,\n    scatter_mm_placeholders,\n)\nif TYPE_CHECKING:\n    import xgrammar as xgr\n    import xgrammar.kernels.apply_token_bitmask_inplace_torch_compile as xgr_torch_compile  # noqa: E501\nelse:\n    xgr = LazyLoader(\"xgr\", globals(), \"xgrammar\")\n    xgr_torch_compile = LazyLoader(\n        \"xgr_torch_compile\", globals(),\n        \"xgrammar.kernels.apply_token_bitmask_inplace_torch_compile\")\n\nimport vllm_mlu._mlu_utils as mlu_envs\nfrom vllm_mlu.compilation.mlu_graph import MLUGraphWrapper\nfrom vllm_mlu.distributed.parallel_state import mlu_graph_capture\nfrom vllm_mlu.model_executor.layers.feed_forward import FeedForward\nfrom vllm_mlu.model_executor.layers.rotary_embedding.base import MLURotaryEmbedding\nfrom vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp\nfrom vllm_mlu.v1.kv_cache_interface import MLUMLAAttentionSpec\nfrom vllm_mlu.v1.attention.backends.flash_attn import FlashAttentionMetadata, pad_attn_metadata\nfrom vllm_mlu.v1.attention.backends.utils import (\n    COMMON_METADATA_STR,\n    MLUCommonAttentionMetadata,\n    MLUInferMode,\n    get_common_metadata,\n    unpad_common_attn_metadata)\nfrom vllm_mlu.model_executor.models.sp_utils import set_sp_forward_context\nfrom vllm_mlu.v1.sample.sampler import MluSampler\nfrom vllm_mlu.v1.spec_decode.dp_eagle import DPMluEagleProposer\nfrom vllm_mlu.v1.spec_decode.eagle import MluEagleProposer\nimport vllm_mlu._mlu_utils as mlu_envs\n\nlogger = init_logger(__name__)\n\n\n_NUM_WARMUP_ITERS = 2\n\n\ndef _model_forward_pre_hook(self, args, kwargs):\n    '''\n    This hook function will be called before model.forward\n    '''\n    assert len(args) == 0 and len(kwargs) > 0, \\\n        f\"The pre-forward's expected inputs are not passed by kwargs. \" + \\\n        f\"Expected len(args)=0, len(kwargs)>0, \" + \\\n        f\"now, len(args)={len(args)}, len(kwargs)={len(kwargs)}.\"\n\n    common_metadata: MLUCommonAttentionMetadata = get_common_metadata()\n\n    if common_metadata:\n        # Prepare attributes for all rope in model\n        MLURotaryEmbedding.set_mlu_var_v1(common_metadata=common_metadata)\n\n    if self.config.model_type == \"deepseek_v4\":\n        args, kwargs = self.update_forward_args(args, kwargs)\n\n    return (args, kwargs)\n\n\n# Wrapper for ModelRunnerOutput to support overlapped execution.\nclass AsyncMLUModelRunnerOutput(AsyncGPUModelRunnerOutput):\n    def __init__(\n        self,\n        model_runner_output: ModelRunnerOutput,\n        sampled_token_ids: torch.Tensor,\n        logprobs_tensors: torch.Tensor | None,\n        invalid_req_indices: list[int],\n        async_output_copy_stream: torch.mlu.Stream,\n        vocab_size: int,\n    ):\n        self._model_runner_output = model_runner_output\n        self._invalid_req_indices = invalid_req_indices\n\n        # Event on the copy stream so we can synchronize the non-blocking copy.\n        self.async_copy_ready_event = torch.mlu.Event()\n\n        # Keep a reference to the device tensor to avoid it being\n        # deallocated until we finish copying it to the host.\n        self._sampled_token_ids = sampled_token_ids\n        self.vocab_size = vocab_size\n        self._logprobs_tensors = logprobs_tensors\n\n        # Initiate the copy on a separate stream, but do not synchronize it.\n        default_stream = torch.mlu.current_stream()\n        with torch.mlu.stream(async_output_copy_stream):\n            async_output_copy_stream.wait_stream(default_stream)\n            self.sampled_token_ids_cpu = self._sampled_token_ids.to(\n                \"cpu\", non_blocking=True\n            )\n            self._logprobs_tensors_cpu = (\n                self._logprobs_tensors.to_cpu_nonblocking()\n                if self._logprobs_tensors\n                else None\n            )\n            self.async_copy_ready_event.record()\n\ndef apply_grammar_bitmask(\n    scheduler_output: SchedulerOutput,\n    grammar_output: GrammarOutput,\n    input_batch: InputBatch,\n    logits: torch.Tensor,\n) -> None:\n    \"\"\"\n    Apply grammar bitmask to output logits of the model with xgrammar function.\n\n    Args:\n        scheduler_output (SchedulerOutput): The result of engine scheduling.\n        input_batch (InputBatch): The input of model runner.\n        logits (torch.Tensor): The output logits of model forward.\n    \"\"\"\n    # Serialization of np.ndarray is much more efficient than a tensor,\n    # so we receive it in that format.\n    grammar_bitmask = grammar_output.grammar_bitmask\n\n    # We receive the structured output bitmask from the scheduler,\n    # compacted to contain bitmasks only for structured output requests.\n    # The order of the requests in the bitmask is not guaranteed to be the\n    # same as the order of the requests in the gpu runner's batch. We need\n    # to sort the bitmask to match the order of the requests used here.\n\n    # Get the batch indices of the structured output requests.\n    # Keep track of the number of speculative tokens scheduled for every\n    # request in the batch, as the logit indices are offset by this amount.\n    struct_out_req_batch_indices: dict[str, int] = {}\n    cumulative_offset = 0\n    seq = sorted(input_batch.req_id_to_index.items(), key=lambda x: x[1])\n    for req_id, batch_index in seq:\n        logit_index = batch_index + cumulative_offset\n        cumulative_offset += len(\n            scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])\n        )\n        if req_id in grammar_output.structured_output_request_ids:\n            struct_out_req_batch_indices[req_id] = logit_index\n\n    out_indices = []\n\n    # Reorder the bitmask to match the order of the requests in the batch.\n    sorted_bitmask = np.full(\n        shape=(logits.shape[0], grammar_bitmask.shape[1]),\n        fill_value=-1,\n        dtype=grammar_bitmask.dtype,\n    )\n    cumulative_index = 0\n    for req_id in grammar_output.structured_output_request_ids:\n        num_spec_tokens = len(\n            scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])\n        )\n        if req_id in struct_out_req_batch_indices:\n            logit_index = struct_out_req_batch_indices[req_id]\n            for i in range(1 + num_spec_tokens):\n                sorted_bitmask[logit_index + i] = grammar_bitmask[cumulative_index + i]\n                out_indices.append(logit_index + i)\n        cumulative_index += 1 + num_spec_tokens\n\n    # Copy async to device as tensor.\n    grammar_bitmask = torch.from_numpy(sorted_bitmask).to(\n        logits.device, non_blocking=True\n    )\n\n    # If the length of out indices and the logits have the same shape\n    # we don't need to pass indices to the kernel,\n    # since the bitmask is already aligned with the logits.\n    skip_out_indices = len(out_indices) == logits.shape[0]\n\n    index_tensor = None\n    if not skip_out_indices:\n        # xgrammar expects a python list of indices but it will actually work with\n        # a tensor. If we copy the tensor ourselves here we can do it in a non_blocking\n        # manner and there should be no cpu sync within xgrammar.\n        index_tensor = torch.tensor(\n            out_indices, dtype=torch.int32, device=\"cpu\", pin_memory=True\n        )\n        index_tensor = index_tensor.to(logits.device, non_blocking=True)\n\n    '''\n    =============================\n    Modify by vllm_mlu\n    =============================\n    @brief: remove index_put_ from inductor lowering denylist to\n    avoid torch.compile error when using xgrammar\n    '''\n    from torch_mlu._inductor import remove_from_lowering_denylist\n    remove_from_lowering_denylist([torch.ops.aten.index_put_])\n    '''\n    ==================\n    End of MLU Hijack\n    ==================\n    '''\n\n    xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor)\n\nclass MLUModelRunner(GPUModelRunner):\n\n    def _init_kv_state(\n        self,\n    ):\n        hf_config = self.model_config.hf_config\n        if hf_config.model_type != \"deepseek_v4\":\n            return\n\n        CACHED_STATE_NUM = self.scheduler_config.max_num_seqs\n        hf_config.cached_state_num = CACHED_STATE_NUM\n        self.kv_state_free_slots = set(range(CACHED_STATE_NUM))\n        self.req_id_to_kv_state = dict()\n\n    def _insert_req_id(\n        self,\n        scheduler_output: \"SchedulerOutput\",\n    ):\n        for new_req_data in scheduler_output.scheduled_new_reqs:\n            req_id = new_req_data.req_id\n            assert req_id not in self.req_id_to_kv_state, \\\n                f\"try to insert req_id: {req_id}, which has been stored int kv_state.\"\n            assert self.kv_state_free_slots, \"fail to allocate kv states\"\n            slot = self.kv_state_free_slots.pop()\n            self.req_id_to_kv_state[req_id] = slot\n\n    def _remove_req_id(\n        self,\n        scheduler_output: \"SchedulerOutput\",\n    ):\n        for req_id in scheduler_output.finished_req_ids:\n            slot = self.req_id_to_kv_state.pop(req_id)\n            self.kv_state_free_slots.add(slot)\n\n    def __init__(\n        self,\n        vllm_config: VllmConfig,\n        device: torch.device,\n    ):\n        self.vllm_config = vllm_config\n        self.model_config = vllm_config.model_config\n        self.cache_config = vllm_config.cache_config\n        self.compilation_config = vllm_config.compilation_config\n        self.lora_config = vllm_config.lora_config\n        self.load_config = vllm_config.load_config\n        self.parallel_config = vllm_config.parallel_config\n        self.scheduler_config = vllm_config.scheduler_config\n        self.speculative_config = vllm_config.speculative_config\n        self.observability_config = vllm_config.observability_config\n        self.mlu_config = vllm_config.mlu_config\n\n        from vllm.model_executor.models.utils import set_cpu_offload_max_bytes\n\n        set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3))\n\n        model_config = self.model_config\n        cache_config = self.cache_config\n        scheduler_config = self.scheduler_config\n        parallel_config = self.parallel_config\n        self.device = device\n        self.pin_memory = is_pin_memory_available()\n        self.dtype = self.model_config.dtype\n        self.kv_cache_dtype = kv_cache_dtype_str_to_dtype(\n            cache_config.cache_dtype, self.model_config\n        )\n\n        self.is_pooling_model = model_config.runner_type == \"pooling\"\n        self.enable_prompt_embeds = model_config.enable_prompt_embeds\n        self.is_multimodal_raw_input_only_model = (\n            model_config.is_multimodal_raw_input_only_model\n        )\n        # This will be overridden in load_model()\n        self.is_multimodal_pruning_enabled = False\n        self.max_model_len = model_config.max_model_len\n\n        # Always set to false after the first forward pass\n        self.calculate_kv_scales = self.cache_config.calculate_kv_scales\n        self.dcp_world_size = self.parallel_config.decode_context_parallel_size\n        self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group\n        self.max_num_tokens = scheduler_config.max_num_batched_tokens\n        self.max_num_reqs = scheduler_config.max_num_seqs\n\n        # Broadcast PP output for external_launcher (torchrun)\n        # to make sure we are synced across pp ranks\n        # TODO: Support overlapping mirco-batches\n        # https://github.com/vllm-project/vllm/issues/18019\n        self.broadcast_pp_output = (\n            self.parallel_config.distributed_executor_backend == \"external_launcher\"\n            and len(get_pp_group().ranks) > 0\n        )\n\n        # Model-related.\n        self.num_query_heads = model_config.get_num_attention_heads(parallel_config)\n        self.hidden_size = model_config.get_hidden_size()\n        self.attention_chunk_size = model_config.attention_chunk_size\n        # Only relevant for models using ALiBi (e.g, MPT)\n        self.use_alibi = model_config.uses_alibi\n\n        self.cascade_attn_enabled = not self.model_config.disable_cascade_attn\n\n        # Multi-modal data support\n        self.mm_registry = MULTIMODAL_REGISTRY\n        self.uses_mrope = model_config.uses_mrope\n        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(\n            model_config\n        )\n\n        if self.model_config.is_encoder_decoder:\n            # Maximum length of the encoder input, only for encoder-decoder\n            # models.\n            self.max_encoder_len = scheduler_config.max_num_encoder_input_tokens\n        else:\n            self.max_encoder_len = 0\n\n        # Sampler\n        \"\"\"\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: use tmo topk_topp_sampler to sample.\n        \"\"\"\n        sampler_cls = (MluSampler\n                       if self.model_config.is_deepseek_mla or self.model_config.is_longcat_flash\n                       else Sampler)\n        self.sampler = sampler_cls(logprobs_mode=self.model_config.logprobs_mode)\n        \"\"\"\n        =================\n        End of MLU Hijack\n        =================\n        \"\"\"\n\n        \"\"\"\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: Add an extra field to indicate infer mode.\n        \"\"\"\n        self.mlu_infer_mode = MLUInferMode.PREFILL_ONLY\n        \"\"\"\n        =================\n        End of MLU Hijack\n        =================\n        \"\"\"\n\n        self.eplb_state: EplbState | None = None\n        \"\"\"\n        State of the expert parallelism load balancer.\n\n        Will be lazily initialized when the model is loaded.\n        \"\"\"\n\n        # Lazy initializations\n        # self.model: nn.Module  # Set after load_model\n        # Initialize in initialize_kv_cache\n        self.kv_caches: list[torch.Tensor] = []\n        # indexes: [kv_cache_group_id][attn_group]\n        self.attn_groups: list[list[AttentionGroup]] = []\n        # self.kv_cache_config: KVCacheConfig\n\n        # mm_hash ->  encoder_output\n        self.encoder_cache: dict[str, torch.Tensor] = {}\n\n        self.use_aux_hidden_state_outputs = False\n        # Set up speculative decoding.\n        # NOTE(Jiayi): currently we put the entire draft model on\n        # the last PP rank. This is not ideal if there are many\n        # layers in the draft model.\n        if self.speculative_config and get_pp_group().is_last_rank:\n            self.drafter: (\n                NgramProposer | SuffixDecodingProposer | EagleProposer | MedusaProposer\n            )\n            if self.speculative_config.method == \"ngram\":\n                self.drafter = NgramProposer(self.vllm_config)\n            elif self.speculative_config.method == \"suffix\":\n                self.drafter = SuffixDecodingProposer(self.vllm_config)\n            elif self.speculative_config.use_eagle():\n                '''\n                =============================\n                Modify by vllm_mlu\n                =============================\n                @brief: Use MluEagleProposer instead of EagleProposer\n                '''\n                if vllm_config.mlu_config.enable_custom_data_parallel_opt:\n                    proposer_cls = DPMluEagleProposer\n                else:\n                    proposer_cls = MluEagleProposer\n                self.drafter = proposer_cls(self.vllm_config, self.device, self)\n                self.previous_hidden_states = torch.zeros(\n                    (self.max_num_tokens, self.hidden_size),\n                    dtype=self.dtype,\n                    device=self.device)\n                '''\n                =============================\n                End of MLU Hijack\n                =============================\n                '''\n                if self.speculative_config.method == \"eagle3\":\n                    self.use_aux_hidden_state_outputs = True\n            elif self.speculative_config.method == \"medusa\":\n                self.drafter = MedusaProposer(\n                    vllm_config=self.vllm_config, device=self.device\n                )\n            else:\n                raise ValueError(\n                    \"Unknown speculative decoding method: \"\n                    f\"{self.speculative_config.method}\"\n                )\n            self.rejection_sampler = RejectionSampler(self.sampler)\n\n        self.num_spec_tokens = 0\n        if self.speculative_config:\n            self.num_spec_tokens = self.speculative_config.num_speculative_tokens\n\n        # Request states.\n        self.requests: dict[str, CachedRequestState] = {}\n        self.comm_stream = torch.mlu.Stream()\n\n        # Input Batch\n        # NOTE(Chen): Ideally, we should initialize the input batch inside\n        # `initialize_kv_cache` based on the kv cache config. However, as in\n        # https://github.com/vllm-project/vllm/pull/18298, due to some unknown\n        # reasons, we have to initialize the input batch before `load_model`,\n        # quantization + weight offloading will fail otherwise. As a temporary\n        # solution, we initialize the input batch here, and re-initialize it\n        # in `initialize_kv_cache` if the block_sizes here is different from\n        # the block_sizes in the kv cache config.\n        custom_logitsprocs = model_config.logits_processors\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: Adjust `max_model_len` to expand input_batch.token_ids_cpu_tensor, prevent\n                overflow when the total length (including speculative tokens) exceeds max_model_len.\n        '''\n        max_model_len_revise=max(self.max_model_len, self.max_encoder_len)\n        if self.num_spec_tokens > 1:\n            max_model_len_revise = max_model_len_revise + self.num_spec_tokens - 1\n        self.input_batch = InputBatch(\n            max_num_reqs=self.max_num_reqs,\n            # We need to use the encoder length for encoder-decoer\n            # because of KV cache for cross-attention.\n            max_model_len=max_model_len_revise,\n            max_num_batched_tokens=self.max_num_tokens,\n            device=self.device,\n            pin_memory=self.pin_memory,\n            vocab_size=self.model_config.get_vocab_size(),\n            block_sizes=[self.cache_config.block_size],\n            kernel_block_sizes=[self.cache_config.block_size],\n            is_spec_decode=bool(self.vllm_config.speculative_config),\n            logitsprocs=build_logitsprocs(\n                self.vllm_config,\n                self.device,\n                self.pin_memory,\n                self.is_pooling_model,\n                custom_logitsprocs,\n            ),\n            # We currently don't know whether a particular custom logits processor\n            # uses output token ids so we set this conservatively.\n            logitsprocs_need_output_token_ids=bool(custom_logitsprocs),\n            is_pooling_model=self.is_pooling_model,\n            dcp_kv_cache_interleave_size=self.parallel_config.dcp_kv_cache_interleave_size,\n        )\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n        self.use_async_scheduling = self.scheduler_config.async_scheduling\n        # Separate cuda stream for overlapping transfer of sampled token ids from\n        # GPU to CPU when async scheduling is enabled.\n        self.async_output_copy_stream: torch.mlu.Stream | None = None\n        # cuda event to synchronize use of reused CPU tensors between steps\n        # when async scheduling is enabled.\n        self.prepare_inputs_event: torch.mlu.Event | None = None\n        if self.use_async_scheduling:\n            self.async_output_copy_stream = torch.mlu.Stream()\n            self.prepare_inputs_event = torch.mlu.Event()\n\n        # self.cudagraph_batch_sizes sorts in ascending order.\n        if (\n            self.compilation_config.cudagraph_capture_sizes\n            and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE\n        ):\n            self.cudagraph_batch_sizes = sorted(\n                self.compilation_config.cudagraph_capture_sizes\n            )\n\n        # Cache the device properties.\n        self._init_device_properties()\n\n        # Persistent buffers for CUDA graphs.\n        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: change postions dtype from int64 to int32\n        @brief: add seq_start_loc buffer for chunk fa\n        '''\n        self.positions = self._make_buffer(self.max_num_tokens, dtype=torch.int32)\n        self.seq_start_loc = self._make_buffer(\n            self.max_num_reqs + 1, dtype=torch.int32\n        )\n\n        self.prefill_enable_mlugraph = self.mlu_config.prefill_enable_mlugraph\n        self.prefill_mlugraph_batch_size = self.mlu_config.prefill_mlugraph_batch_size\n        self.prefill_mlugraph_seq_len = self.mlu_config.prefill_mlugraph_seq_len\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: Add kv_state buffer for deepseekv4\n        '''\n        self._init_kv_state()\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n        self.query_start_loc = self._make_buffer(\n            self.max_num_reqs + 1, dtype=torch.int32\n        )\n        self.seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)\n        if self.dcp_world_size > 1:\n            self.dcp_local_seq_lens = self._make_buffer(\n                self.max_num_reqs, dtype=torch.int32\n            )\n        # Because inputs_embeds may be bfloat16 and we don't need a numpy\n        # version of this tensor, avoid a RuntimeError by not creating a\n        # numpy buffer.\n        self.inputs_embeds = self._make_buffer(\n            self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False\n        )\n        self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool)\n        self.discard_request_indices = self._make_buffer(\n            self.max_num_reqs, dtype=torch.int64\n        )\n        self.num_discarded_requests = 0\n\n        self.num_decode_draft_tokens = self._make_buffer(\n            self.max_num_reqs, dtype=torch.int32\n        )\n        self.num_accepted_tokens = self._make_buffer(\n            self.max_num_reqs, dtype=torch.int64\n        )\n\n        # Only relevant for multimodal models\n        if self.supports_mm_inputs:\n            self.is_mm_embed = self._make_buffer(self.max_num_tokens, dtype=torch.bool)\n\n        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)\n        if self.uses_mrope:\n            # NOTE: `mrope_positions` is implemented with one additional dummy\n            # position on purpose to make it non-contiguous so that it can work\n            # with torch compile.\n            # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923\n\n            # NOTE: When M-RoPE is enabled, position ids are 3D regardless of\n            # the modality of inputs. For text-only inputs, each dimension has\n            # identical position IDs, making M-RoPE functionally equivalent to\n            # 1D-RoPE.\n            # See page 5 of https://arxiv.org/abs/2409.12191\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: change postions dtype from int64 to int32\n            '''\n            self.mrope_positions = self._make_buffer(\n                (3, self.max_num_tokens + 1), dtype=torch.int32\n            )\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n\n        # None in the first PP rank. The rest are set after load_model.\n        self.intermediate_tensors: IntermediateTensors | None = None\n\n        # OPTIMIZATION: Cache the tensors rather than creating them every step.\n        # Keep in int64 to avoid overflow with long context\n        self.arange_np = np.arange(\n            max(self.max_num_reqs + 1, self.max_model_len, self.max_num_tokens),\n            dtype=np.int64,\n        )\n\n        # Layer pairings for cross-layer KV sharing.\n        # If an Attention layer `layer_name` is in the keys of this dict, it\n        # means this layer will perform attention using the keys and values\n        # from the KV cache of `shared_kv_cache_layers[layer_name]`.\n        self.shared_kv_cache_layers: dict[str, str] = {}\n        self.kv_sharing_fast_prefill_eligible_layers: set[str] = set()\n\n        self.kv_sharing_fast_prefill_logits_indices = None\n        if self.cache_config.kv_sharing_fast_prefill:\n            self.kv_sharing_fast_prefill_logits_indices = torch.zeros(\n                self.max_num_tokens, dtype=torch.int32, device=self.device\n            )\n\n        self.uniform_decode_query_len = 1 + self.num_spec_tokens\n\n        # Cudagraph dispatcher for runtime cudagraph dispatching.\n        self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)\n\n        self.mm_budget = (\n            MultiModalBudget(\n                self.model_config,\n                self.scheduler_config,\n                self.mm_registry,\n            )\n            if self.supports_mm_inputs\n            else None\n        )\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: enable reorder batch to ensure correct splitting\n            between prefill chunks and decode tokens in chunked prefill mode.\n        '''\n        self.reorder_batch_threshold: int | None = self.uniform_decode_query_len\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        # Attention layers that are only in the KVCacheConfig of the runner\n        # (e.g., KV sharing, encoder-only attention), but not in the\n        # KVCacheConfig of the scheduler.\n        self.runner_only_attn_layers: set[str] = set()\n\n        # Cached outputs.\n        self._draft_token_ids: list[list[int]] | torch.Tensor | None = None\n        self.transfer_event = torch.mlu.Event()\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: change sampled_token_ids dtype from int64 to int32\n        @brief: add draft accepted counter\n        '''\n        self.sampled_token_ids_pinned_cpu = torch.empty(\n            (self.max_num_reqs, 1),\n            dtype=torch.int32,\n            device=\"cpu\",\n            pin_memory=self.pin_memory,\n        )\n\n        # Pre-allocated tensor for copying valid sampled token counts to CPU,\n        # with dedicated stream for overlapping and event for coordination.\n        self.valid_sampled_token_count_event: torch.mlu.Event | None = None\n        self.valid_sampled_token_count_copy_stream: torch.mlu.Stream | None = None\n        if self.use_async_scheduling and self.num_spec_tokens:\n            self.valid_sampled_token_count_event = torch.mlu.Event()\n            self.valid_sampled_token_count_copy_stream = torch.mlu.Stream()\n        self.valid_sampled_token_count_cpu = torch.empty(\n            self.max_num_reqs,\n            dtype=torch.int64,\n            device=\"cpu\",\n            pin_memory=self.pin_memory,\n        )\n        self.total_draft_tokens = 0\n        self.total_accepted_tokens = 0\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        # Ephemeral state transferred between execute_model() and sample_tokens().\n        self.execute_model_state: ExecuteModelState | None = None\n        self.kv_connector_output: KVConnectorOutput | None = None\n        \n        self.execute_cnpx_mark = None\n        self.request_cnpx_mark = {}\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: support qwen3-next\n        '''\n        self.mamba_block_num = 1\n        self.mamba_tensor_size = 0\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n    # Note: used for model runner override.\n    def _init_device_properties(self) -> None:\n        \"\"\"Initialize attributes from torch.cuda.get_device_properties\"\"\"\n        self.device_properties = torch.mlu.get_device_properties(self.device)\n        self.num_sms = self.device_properties.multi_processor_count\n\n    # Note: used for model runner override.\n    def _sync_device(self) -> None:\n        torch.mlu.synchronize()\n\n    \n    def get_accept_rate(self) -> float:\n        if self.total_draft_tokens == 0:\n            return 0.0\n        return self.total_accepted_tokens / self.total_draft_tokens\n\n    def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: only support pad tokens in decode mode.\n        '''\n        if (\n            self.mlu_infer_mode == MLUInferMode.DECODE_ONLY\n            and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE\n            and hasattr(self, \"cudagraph_batch_sizes\")\n            and self.cudagraph_batch_sizes\n            and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]\n        ):\n            # Use CUDA graphs.\n            # Add padding to the batch size.\n            return self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        # Eager mode.\n        # Pad tokens to multiple of tensor_parallel_size when\n        # enabled collective fusion for SP\n        tp_size = self.vllm_config.parallel_config.tensor_parallel_size\n        if (\n            self.compilation_config.pass_config.enable_sequence_parallelism\n            and tp_size > 1\n        ):\n            return round_up(num_scheduled_tokens, tp_size)\n        return num_scheduled_tokens\n\n    def _prepare_inputs(\n        self,\n        scheduler_output: \"SchedulerOutput\",\n        num_scheduled_tokens: np.ndarray,\n        max_num_scheduled_tokens: int,\n    ) -> tuple[\n        torch.Tensor,\n        SpecDecodeMetadata | None,\n        UBatchSlices | None,\n        torch.Tensor | None,\n    ]:\n        \"\"\"\n        :return: tuple[\n            logits_indices, spec_decode_metadata,\n            ubatch_slices, num_tokens_across_dp,\n        ]\n        \"\"\"\n        self._insert_req_id(scheduler_output)\n\n        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens\n        assert total_num_scheduled_tokens > 0\n        num_reqs = self.input_batch.num_reqs\n        assert num_reqs > 0\n\n        # OPTIMIZATION: Start copying the block table first.\n        # This way, we can overlap the copy with the following CPU operations.\n        self.input_batch.block_table.commit_block_table(num_reqs)\n\n        # Get request indices.\n        # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]\n        req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)\n\n        # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]\n        # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]\n        cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens)\n\n        # Get positions.\n        positions_np = self.positions.np[:total_num_scheduled_tokens]\n        np.add(\n            self.input_batch.num_computed_tokens_cpu[req_indices],\n            arange,\n            out=positions_np,\n        )\n\n        # Calculate M-RoPE positions.\n        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)\n        if self.uses_mrope:\n            self._calc_mrope_positions(scheduler_output)\n\n        # Get token indices.\n        # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]\n        # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]\n        # where M is the max_model_len.\n        token_indices = (\n            positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]\n        )\n        token_indices_tensor = torch.from_numpy(token_indices)\n\n        # # Get the number of scheduled tokens for each request.\n        # req_ids = self.input_batch.req_ids\n        # tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]\n        # num_scheduled_tokens = np.array(tokens, dtype=np.int32)\n        # max_num_scheduled_tokens = max(tokens)\n\n\n        # NOTE(woosuk): We use torch.index_select instead of np.take here\n        # because torch.index_select is much faster than np.take for large\n        # tensors.\n        torch.index_select(\n            self.input_batch.token_ids_cpu_tensor.flatten(),\n            0,\n            token_indices_tensor,\n            out=self.input_ids.cpu[:total_num_scheduled_tokens],\n        )\n        if self.enable_prompt_embeds:\n            is_token_ids = self.input_batch.is_token_ids_tensor.flatten()\n            torch.index_select(\n                is_token_ids,\n                0,\n                token_indices_tensor,\n                out=self.is_token_ids.cpu[:total_num_scheduled_tokens],\n            )\n\n        # Because we did not pre-allocate a massive prompt_embeds CPU tensor on\n        # the InputBatch, we need to fill in the prompt embeds into the expected\n        # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor.\n        if self.input_batch.req_prompt_embeds:\n            output_idx = 0\n            for req_idx in range(num_reqs):\n                num_sched = num_scheduled_tokens[req_idx]\n\n                # Skip if this request doesn't have embeddings\n                if req_idx not in self.input_batch.req_prompt_embeds:\n                    output_idx += num_sched\n                    continue\n\n                # Skip if no tokens scheduled\n                if num_sched <= 0:\n                    output_idx += num_sched\n                    continue\n\n                req_embeds = self.input_batch.req_prompt_embeds[req_idx]\n                start_pos = self.input_batch.num_computed_tokens_cpu[req_idx]\n\n                # Skip if trying to read beyond available embeddings\n                if start_pos >= req_embeds.shape[0]:\n                    output_idx += num_sched\n                    continue\n\n                # Copy available embeddings\n                end_pos = start_pos + num_sched\n                actual_end = min(end_pos, req_embeds.shape[0])\n                actual_num_sched = actual_end - start_pos\n\n                if actual_num_sched > 0:\n                    self.inputs_embeds.cpu[\n                        output_idx : output_idx + actual_num_sched\n                    ].copy_(req_embeds[start_pos:actual_end])\n\n                output_idx += num_sched\n\n        self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)\n        self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)\n\n        # Prepare the attention metadata.\n        self.query_start_loc.np[0] = 0\n        self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens\n        # Note: pad query_start_loc to be non-decreasing, as kernels\n        # like FlashAttention requires that\n        self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1])\n        self.query_start_loc.copy_to_gpu()\n        query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]\n\n        num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens\n        num_tokens_padded = self._get_num_input_tokens(num_tokens_unpadded)\n        uniform_decode = (\n            max_num_scheduled_tokens == self.uniform_decode_query_len\n        ) and (total_num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)\n\n        # Disable DP padding when running eager to avoid excessive padding when\n        # running prefills. This lets us set enforce_eager on the prefiller in\n        # a P/D setup and still use CUDA graphs (enabled by this padding) on the\n        # decoder.\n        allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE\n\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: build mlu dp metadata for dp opt.\n        '''\n        ubatch_slices, num_tokens_across_dp = None, None\n        if self.vllm_config.mlu_config.enable_custom_data_parallel_opt:\n            cur_num_reqs = (num_reqs if self.mlu_infer_mode.is_prefill_only\n                            else num_reqs * (1 + self.num_spec_tokens))\n            query_len_per_batch = (self.query_start_loc.np[1:] -\n                                   self.query_start_loc.np[:-1]).tolist()\n            dp_metadata = self._get_data_parallel_metadata(\n                total_num_scheduled_tokens, cur_num_reqs,\n                self.mlu_infer_mode.is_decode_only, query_len_per_batch[:cur_num_reqs]\n            )\n            # replace num_tokens_across_dp with dp_metadata for return\n            num_tokens_across_dp = dp_metadata\n        else:\n            ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(\n                num_tokens_unpadded=num_tokens_unpadded,\n                parallel_config=self.parallel_config,\n                allow_microbatching=True,\n                allow_dp_padding=allow_dp_padding,\n                num_tokens_padded=num_tokens_padded,\n                uniform_decode=uniform_decode,\n                num_scheduled_tokens_per_request=num_scheduled_tokens,\n            )\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        self.seq_lens.np[:num_reqs] = (\n            self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens\n        )\n        # Fill unused with 0 for full cuda graph mode.\n        self.seq_lens.np[num_reqs:].fill(0)\n        self.seq_lens.copy_to_gpu()\n\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: add seq_start_loc for chunk fa.\n        '''\n        self.seq_start_loc.np[0] = 0\n        self.seq_start_loc.np[1:num_reqs + 1] = np.cumsum(self.seq_lens.np[:num_reqs])\n        self.seq_start_loc.np[num_reqs + 1 :].fill(self.seq_start_loc.np[num_reqs])\n        self.seq_start_loc.copy_to_gpu()\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids]\n        num_tokens_np = np.array(num_tokens, dtype=np.int32)\n\n        # Record the index of requests that should not be sampled,\n        # so that we could clear the sampled tokens before returning\n        discard_requests_mask = self.seq_lens.np[:num_reqs] < num_tokens_np\n        discard_request_indices = np.nonzero(discard_requests_mask)[0]\n        self.num_discarded_requests = len(discard_request_indices)\n        self.discard_request_indices.np[: self.num_discarded_requests] = (\n            discard_request_indices\n        )\n\n        self.discard_request_indices.copy_to_gpu(self.num_discarded_requests)\n\n        # Copy the tensors to the GPU.\n        self._prepare_input_ids(\n            scheduler_output,\n            total_num_scheduled_tokens,\n            cu_num_tokens,\n        )\n\n        if self.uses_mrope:\n            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)\n            self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(\n                self.mrope_positions.cpu[:, :total_num_scheduled_tokens],\n                non_blocking=True,\n            )\n        else:\n            # Common case (1D positions)\n            self.positions.copy_to_gpu(total_num_scheduled_tokens)\n\n        use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0\n        if not use_spec_decode:\n            # NOTE(woosuk): Due to chunked prefills, the batch may contain\n            # partial requests. While we should not sample any token\n            # from these partial requests, we do so for simplicity.\n            # We will ignore the sampled tokens from the partial requests.\n            # TODO: Support prompt logprobs.\n            logits_indices = query_start_loc[1:] - 1\n            num_draft_tokens = None\n            spec_decode_metadata = None\n            num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)\n        else:\n            # Get the number of draft tokens for each request.\n            # Iterate over the dictionary rather than all requests since not all\n            # requests have draft tokens.\n            num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)\n            # For chunked prefills, use -1 as mask rather than 0, as guided\n            # decoding may rollback speculative tokens.\n            num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32)\n            for (\n                req_id,\n                draft_token_ids,\n            ) in scheduler_output.scheduled_spec_decode_tokens.items():\n                req_idx = self.input_batch.req_id_to_index[req_id]\n                num_draft_tokens[req_idx] = len(draft_token_ids)\n                num_decode_draft_tokens[req_idx] = (\n                    len(draft_token_ids)\n                    if (\n                        self.input_batch.num_computed_tokens_cpu[req_idx]\n                        >= self.input_batch.num_prompt_tokens[req_idx]\n                    )\n                    else -1\n                )\n            spec_decode_metadata = self._calc_spec_decode_metadata(\n                num_draft_tokens, cu_num_tokens\n            )\n            logits_indices = spec_decode_metadata.logits_indices\n            num_sampled_tokens = num_draft_tokens + 1\n            # For DECODE only cuda graph of some attention backends (e.g., GDN).\n            self.num_decode_draft_tokens.np[:num_reqs] = num_decode_draft_tokens\n            self.num_decode_draft_tokens.np[num_reqs:].fill(-1)\n            self.num_decode_draft_tokens.copy_to_gpu()\n\n        # Hot-Swap lora model\n        if self.lora_config:\n            assert (\n                np.sum(num_sampled_tokens)\n                <= self.vllm_config.scheduler_config.max_num_batched_tokens\n            )\n            self.set_active_loras(\n                self.input_batch, num_scheduled_tokens, num_sampled_tokens\n            )\n\n        return (\n            logits_indices,\n            spec_decode_metadata,\n            ubatch_slices,\n            num_tokens_across_dp,\n        )\n\n    def get_model(self) -> nn.Module:\n        # get raw model out of the cudagraph wrapper.\n        if isinstance(self.model, (MLUGraphWrapper, UBatchWrapper)):\n            return self.model.unwrap()\n        return self.model\n\n    def _execute_mm_encoder(self, scheduler_output: \"SchedulerOutput\"):\n        # Batch the multi-modal inputs using the helper method.\n        mm_kwargs, mm_hashes_pos = self._batch_mm_kwargs_from_scheduler(\n            scheduler_output\n        )\n\n        if not mm_kwargs:\n            return\n\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: v1 offline benchmark\n        '''\n        self.mm_time_markers = []\n\n        if mlu_envs.VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:\n            mm_start = torch.mlu.Event(enable_timing=True)\n            mm_start.record()\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        # Batch mm inputs as much as we can: if a request in the batch has\n        # multiple modalities or a different modality than the previous one,\n        # we process it separately to preserve item order.\n        # FIXME(ywang96): This is a hacky way to deal with multiple modalities\n        # in the same batch while still being able to benefit from batching\n        # multimodal inputs. The proper solution should be reordering the\n        # encoder outputs.\n        model = cast(SupportsMultiModal, self.model)\n        encoder_outputs = []\n        for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(\n            mm_kwargs,\n            device=self.device,\n            pin_memory=self.pin_memory,\n            merge_by_field_config=model.merge_by_field_config,\n            multimodal_cpu_fields=model.multimodal_cpu_fields,\n        ):\n            curr_group_outputs = []\n\n            # EVS-related change.\n            # (ekhvedchenia): Temporary hack to limit peak memory usage when\n            # processing multimodal data. This solves the issue with scheduler\n            # putting too many video samples into a single batch. Scheduler\n            # uses pruned vision tokens count to compare it versus compute\n            # budget which is incorrect (Either input media size or non-pruned\n            # output vision tokens count should be considered)\n            # TODO(ywang96): Fix memory profiling to take EVS into account and\n            # remove this hack.\n            if (\n                self.is_multimodal_pruning_enabled\n                and modality == \"video\"\n                and num_items > 1\n            ):\n                for video_mm_kwargs_item in filter(\n                    lambda item: item.modality == \"video\", mm_kwargs\n                ):\n                    _, _, micro_batch_mm_inputs = next(\n                        group_mm_kwargs_by_modality(\n                            [video_mm_kwargs_item],\n                            device=self.device,\n                            pin_memory=self.pin_memory,\n                            merge_by_field_config=model.merge_by_field_config,\n                            multimodal_cpu_fields=model.multimodal_cpu_fields,\n                        )\n                    )\n\n                    micro_batch_outputs = model.embed_multimodal(\n                        **micro_batch_mm_inputs\n                    )\n\n                    curr_group_outputs.extend(micro_batch_outputs)\n            else:\n                # Run the encoder.\n                # `curr_group_outputs` is either of the following:\n                # 1. A tensor of shape (num_items, feature_size, hidden_size)\n                # in case feature_size is fixed across all multimodal items.\n                # 2. A list or tuple (length: num_items) of tensors,\n                # each of shape (feature_size, hidden_size) in case the feature\n                # size is dynamic depending on the input multimodal items.\n                curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)\n\n            sanity_check_mm_encoder_outputs(\n                curr_group_outputs,\n                expected_num_items=num_items,\n            )\n            encoder_outputs.extend(curr_group_outputs)\n\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: v1 offline benchmark\n        '''\n        if encoder_outputs and mlu_envs.VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:\n            mm_end = torch.mlu.Event(enable_timing=True)\n            mm_end.record()\n            self.mm_time_markers.append([mm_start, mm_end])\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        # Cache the encoder outputs by mm_hash\n        for (mm_hash, pos_info), output in zip(mm_hashes_pos, encoder_outputs):\n            self.encoder_cache[mm_hash] = scatter_mm_placeholders(\n                output,\n                is_embed=pos_info.is_embed,\n            )\n\n    def _update_states(self, scheduler_output: \"SchedulerOutput\") -> None:\n        \"\"\"Update the cached states and the persistent batch with the scheduler\n        output.\n\n        The updated states are used by the `_prepare_inputs` function to create\n        the input GPU tensors for the model.\n\n        The SamplingMetadata is updated and copied to the GPU if there is a\n        new/resumed/paused/finished request in the batch.\n        \"\"\"\n        self._remove_req_id(scheduler_output)\n\n        # Remove finished requests from the cached states.\n        for req_id in scheduler_output.finished_req_ids:\n            self.requests.pop(req_id, None)\n        # Remove the finished requests from the persistent batch.\n        # NOTE(woosuk): There could be an edge case where finished_req_ids and\n        # scheduled_req_ids overlap. This happens when a request is aborted and\n        # then resubmitted with the same ID. In this case, we treat them as two\n        # distinct requests - clearing the cached states for the first request\n        # and handling the second as a new request.\n        for req_id in scheduler_output.finished_req_ids:\n            self.input_batch.remove_request(req_id)\n\n        # Free the cached encoder outputs.\n        for mm_hash in scheduler_output.free_encoder_mm_hashes:\n            self.encoder_cache.pop(mm_hash, None)\n\n        # Remove the unscheduled requests from the persistent batch.\n        # NOTE(woosuk): The unscheduled requests are either preempted requests\n        # or running requests that are not scheduled in this step. We remove\n        # them from the persistent batch but keep their cached states since\n        # they will be scheduled again sometime in the future.\n        scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()\n        cached_req_ids = self.input_batch.req_id_to_index.keys()\n        unscheduled_req_ids = cached_req_ids - scheduled_req_ids\n        # NOTE(woosuk): The persistent batch optimization assumes that\n        # consecutive batches contain mostly the same requests. If batches\n        # have low request overlap (e.g., alternating between two distinct\n        # sets of requests), this optimization becomes very inefficient.\n        for req_id in unscheduled_req_ids:\n            self.input_batch.remove_request(req_id)\n\n        reqs_to_add: list[CachedRequestState] = []\n        # Add new requests to the cached states.\n        for new_req_data in scheduler_output.scheduled_new_reqs:\n            req_id = new_req_data.req_id\n            sampling_params = new_req_data.sampling_params\n            pooling_params = new_req_data.pooling_params\n\n            if (\n                sampling_params\n                and sampling_params.sampling_type == SamplingType.RANDOM_SEED\n            ):\n                generator = torch.Generator(device=self.device)\n                generator.manual_seed(sampling_params.seed)\n            else:\n                generator = None\n\n            if self.is_pooling_model:\n                assert pooling_params is not None\n                task = pooling_params.task\n                assert task is not None, \"You did not set `task` in the API\"\n\n                model = cast(VllmModelForPooling, self.get_model())\n                to_update = model.pooler.get_pooling_updates(task)\n                to_update.apply(pooling_params)\n\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: supoort disagg for mlu.\n            '''\n            req_state = CachedRequestState(\n                req_id=req_id,\n                prompt_token_ids=new_req_data.prompt_token_ids,\n                prompt_embeds=new_req_data.prompt_embeds,\n                mm_features=new_req_data.mm_features,\n                sampling_params=sampling_params,\n                pooling_params=pooling_params,\n                generator=generator,\n                block_ids=new_req_data.block_ids,\n                num_computed_tokens=new_req_data.num_computed_tokens,\n                output_token_ids=getattr(new_req_data, 'new_token_ids', None) or [],\n                lora_request=new_req_data.lora_request,\n            )\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n            self.requests[req_id] = req_state\n\n            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)\n            if self.uses_mrope:\n                self._init_mrope_positions(req_state)\n\n            reqs_to_add.append(req_state)\n\n        # Update the states of the running/resumed requests.\n        is_last_rank = get_pp_group().is_last_rank\n        req_data = scheduler_output.scheduled_cached_reqs\n\n        # Wait until valid_sampled_tokens_count is copied to cpu,\n        # then use it to update actual num_computed_tokens of each request.\n        valid_sampled_token_count = self._get_valid_sampled_token_count()\n\n        for i, req_id in enumerate(req_data.req_ids):\n            req_state = self.requests[req_id]\n            num_computed_tokens = req_data.num_computed_tokens[i]\n            new_block_ids = req_data.new_block_ids[i]\n            resumed_from_preemption = req_id in req_data.resumed_req_ids\n            num_output_tokens = req_data.num_output_tokens[i]\n            req_index = self.input_batch.req_id_to_index.get(req_id)\n\n            # prev_num_draft_len is used in async scheduling mode with\n            # spec decode. it indicates if need to update num_computed_tokens\n            # of the request. for example:\n            # fist step: num_computed_tokens = 0, spec_tokens = [],\n            # prev_num_draft_len = 0.\n            # second step: num_computed_tokens = 100(prompt lenth),\n            # spec_tokens = [a,b], prev_num_draft_len = 0.\n            # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d],\n            # prev_num_draft_len = 2.\n            # num_computed_tokens in first step and second step does't contain\n            # the spec tokens length, but in third step it contains the\n            # spec tokens length. we only need to update num_computed_tokens\n            # when prev_num_draft_len > 0.\n            if req_state.prev_num_draft_len:\n                if req_index is None:\n                    req_state.prev_num_draft_len = 0\n                else:\n                    assert self.input_batch.prev_req_id_to_index is not None\n                    prev_req_index = self.input_batch.prev_req_id_to_index[req_id]\n                    num_accepted = valid_sampled_token_count[prev_req_index] - 1\n                    num_rejected = req_state.prev_num_draft_len - num_accepted\n                    num_computed_tokens -= num_rejected\n                    req_state.output_token_ids.extend([-1] * num_accepted)\n\n            # Update the cached states.\n            req_state.num_computed_tokens = num_computed_tokens\n\n            if not is_last_rank:\n                # When using PP, the scheduler sends the sampled tokens back,\n                # because there's no direct communication between the first-\n                # stage worker and the last-stage worker.\n                new_token_ids = req_data.new_token_ids[i]\n                # Add the sampled token(s) from the previous step (if any).\n                # This doesn't include \"unverified\" tokens like spec tokens.\n                num_new_tokens = (\n                    num_computed_tokens + len(new_token_ids) - req_state.num_tokens\n                )\n                if num_new_tokens == 1:\n                    # Avoid slicing list in most common case.\n                    req_state.output_token_ids.append(new_token_ids[-1])\n                elif num_new_tokens > 0:\n                    req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])\n            elif num_output_tokens < len(req_state.output_token_ids):\n                # Some output tokens were discarded due to a sync-KV-load\n                # failure. Align the cached state.\n                del req_state.output_token_ids[num_output_tokens:]\n                if req_index is not None:\n                    end_idx = (\n                        self.input_batch.num_prompt_tokens[req_index]\n                        + num_output_tokens\n                    )\n                    self.input_batch.num_tokens[req_index] = end_idx\n                    self.input_batch.num_tokens_no_spec[req_index] = end_idx\n\n            # Update the block IDs.\n            if not resumed_from_preemption:\n                if new_block_ids is not None:\n                    # Append the new blocks to the existing block IDs.\n                    for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):\n                        block_ids.extend(new_ids)\n            else:\n                assert req_index is None\n                assert new_block_ids is not None\n                # The request is resumed from preemption.\n                # Replace the existing block IDs with the new ones.\n                req_state.block_ids = new_block_ids\n\n            if req_index is None:\n                # The request is not in the persistent batch.\n                # The request was either preempted and resumed later, or was not\n                # scheduled in the previous step and needs to be added again.\n\n                if self.use_async_scheduling and num_output_tokens > 0:\n                    # We must recover the output token ids for resumed requests in the\n                    # async scheduling case, so that correct input_ids are obtained.\n                    resumed_token_ids = req_data.all_token_ids[req_id]\n                    req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]\n\n                reqs_to_add.append(req_state)\n                continue\n\n            # Update the persistent batch.\n            self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens\n            if new_block_ids is not None:\n                self.input_batch.block_table.append_row(new_block_ids, req_index)\n\n            # For the last rank, we don't need to update the token_ids_cpu\n            # because the sampled tokens are already cached.\n            if not is_last_rank:\n                # Add new_token_ids to token_ids_cpu.\n                start_token_index = num_computed_tokens\n                end_token_index = num_computed_tokens + len(new_token_ids)\n                self.input_batch.token_ids_cpu[\n                    req_index, start_token_index:end_token_index\n                ] = new_token_ids\n                self.input_batch.num_tokens_no_spec[req_index] = end_token_index\n                self.input_batch.num_tokens[req_index] = end_token_index\n\n            # Add spec_token_ids to token_ids_cpu.\n            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(\n                req_id, []\n            )\n            num_spec_tokens = len(spec_token_ids)\n            # For async scheduling, token_ids_cpu assigned from\n            # spec_token_ids are placeholders and will be overwritten in\n            # _prepare_input_ids.\n            if num_spec_tokens:\n                start_index = self.input_batch.num_tokens_no_spec[req_index]\n                end_token_index = start_index + num_spec_tokens\n                self.input_batch.token_ids_cpu[\n                    req_index, start_index:end_token_index\n                ] = spec_token_ids\n                # NOTE(woosuk): `num_tokens` here may include spec tokens.\n                self.input_batch.num_tokens[req_index] += num_spec_tokens\n\n            # When speculative decoding is used with structured output,\n            # the scheduler can drop draft tokens that do not\n            # conform to the schema. This can result in\n            # scheduler_output.scheduled_spec_decode_tokens being empty,\n            # even when speculative decoding is enabled.\n            self.input_batch.spec_token_ids[req_index] = spec_token_ids\n\n            # there are no draft tokens with async scheduling,\n            # we clear the spec_decoding info in scheduler_output and\n            # use normal sampling but rejection_sampling.\n            if self.use_async_scheduling:\n                req_state.prev_num_draft_len = num_spec_tokens\n                if num_spec_tokens and self._draft_token_ids is None:\n                    scheduler_output.total_num_scheduled_tokens -= num_spec_tokens\n                    scheduler_output.num_scheduled_tokens[req_id] -= num_spec_tokens\n                    scheduler_output.scheduled_spec_decode_tokens.pop(req_id, None)\n        # Add the new or resumed requests to the persistent batch.\n        # The smaller empty indices are filled first.\n        for request in reqs_to_add:\n            self.input_batch.add_request(request)\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: cache the unverified spec_decode token\n            '''\n            req_id = request.req_id\n            req_index = self.input_batch.req_id_to_index.get(req_id)\n            assert req_index is not None\n            num_tokens = self.input_batch.num_tokens[req_index]\n            end_token_index = num_tokens\n            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(\n                req_id, []\n            )\n            if spec_token_ids:\n                start_index = end_token_index\n                end_token_index += len(spec_token_ids)\n                self.input_batch.token_ids_cpu[\n                    req_index, start_index:end_token_index] = spec_token_ids\n            # NOTE(woosuk): `num_tokens` here may include spec decode tokens.\n            self.input_batch.num_tokens[req_index] = end_token_index\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n\n        # Condense the batched states if there are gaps left by removed requests\n        self.input_batch.condense()\n        # Allow attention backend to reorder the batch, potentially\n        self._may_reorder_batch(scheduler_output)\n        # Refresh batch metadata with any pending updates.\n        self.input_batch.refresh_metadata()\n\n    @torch.inference_mode()\n    def execute_model(\n        self,\n        scheduler_output: \"SchedulerOutput\",\n        intermediate_tensors: IntermediateTensors | None = None,\n    ) -> ModelRunnerOutput | IntermediateTensors | None:\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: clear time markers before execute model.\n        '''\n        self.time_markers = []\n        self.mm_time_markers = []\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n        if self.execute_model_state is not None:\n            raise RuntimeError(\n                \"State error: sample_tokens() must be called \"\n                \"after execute_model() returns None.\"\n            )   \n        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens\n        with record_function_or_nullcontext(\"gpu_model_runner: preprocess\"):\n            with self.synchronize_input_prep():\n                # Update persistent batch states.\n                self._update_states(scheduler_output)\n\n                if has_ec_transfer() and get_ec_transfer().is_producer:\n                    with self.maybe_get_ec_connector_output(\n                        scheduler_output,\n                        encoder_cache=self.encoder_cache,\n                    ) as ec_connector_output:\n                        self._execute_mm_encoder(scheduler_output)\n                        return make_empty_encoder_model_runner_output(scheduler_output)\n\n                if not num_scheduled_tokens:\n                    if not has_kv_transfer_group():\n                        # Return empty ModelRunnerOutput if no work to do.\n                        return EMPTY_MODEL_RUNNER_OUTPUT\n                    return self.kv_connector_no_forward(\n                        scheduler_output, self.vllm_config\n                    )\n                if self.cache_config.kv_sharing_fast_prefill:\n                    assert not self.input_batch.num_prompt_logprobs, (\n                        \"--kv-sharing-fast-prefill produces incorrect \"\n                        \"logprobs for prompt tokens, tokens, please disable \"\n                        \"it when the requests need prompt logprobs\"\n                    )\n\n                num_reqs = self.input_batch.num_reqs\n                req_ids = self.input_batch.req_ids\n                tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]\n                num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)\n                max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())\n\n                '''\n                =============================\n                Modify by vllm_mlu\n                =============================\n                @brief: add mlu_infer_mode.\n                '''\n                max_computed_tokens = np.max(self.input_batch.num_computed_tokens_cpu[:num_reqs])\n                self.mlu_infer_mode = MLUInferMode.build(\n                    max_query_len=max_num_scheduled_tokens,\n                    max_computed_tokens=max_computed_tokens,\n                    uniform_decode_query_len=self.uniform_decode_query_len,\n                )\n                '''\n                ==================\n                End of MLU Hijack\n                ==================\n                '''\n\n                (\n                    logits_indices,\n                    spec_decode_metadata,\n                    ubatch_slices,\n                    num_tokens_across_dp,\n                ) = self._prepare_inputs(\n                    scheduler_output, num_scheduled_tokens_np, max_num_scheduled_tokens\n                )\n\n                cascade_attn_prefix_lens = None\n                # Disable cascade attention when using microbatching (DBO)\n                if self.cascade_attn_enabled and ubatch_slices is None:\n                    # Pre-compute cascade attention prefix lengths\n                    # NOTE: Must be AFTER _prepare_inputs uses self.input_batch state\n                    cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(\n                        num_scheduled_tokens_np,\n                        scheduler_output.num_common_prefix_blocks,\n                    )\n\n                # TODO(lucas): move cudagraph dispatching here:\n                #   https://github.com/vllm-project/vllm/issues/23789\n\n                total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens\n                use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0\n                attn_metadata, spec_decode_common_attn_metadata = (\n                    self._build_attention_metadata(\n                        total_num_scheduled_tokens=total_num_scheduled_tokens,\n                        max_num_scheduled_tokens=max_num_scheduled_tokens,\n                        num_reqs=num_reqs,\n                        ubatch_slices=ubatch_slices,\n                        logits_indices=logits_indices,\n                        use_spec_decode=use_spec_decode,\n                        scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,\n                        cascade_attn_prefix_lens=cascade_attn_prefix_lens,\n                        mlu_infer_mode=self.mlu_infer_mode,\n                    )\n                )\n\n                dp_rank = self.parallel_config.data_parallel_rank\n                if ubatch_slices:\n                    assert num_tokens_across_dp is not None\n                    num_input_tokens = int(num_tokens_across_dp[dp_rank].item())\n                    self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)\n                elif num_tokens_across_dp is not None:\n                    num_input_tokens = int(num_tokens_across_dp[dp_rank].item())\n                else:\n                    '''\n                    =============================\n                    Modify by vllm_mlu\n                    =============================\n                    pad num_input_tokens after supporting pad decode graph.\n                    '''\n                    max_num_tokens = (\n                        self.scheduler_config.max_num_seqs * self.uniform_decode_query_len\n                    )\n                    capture_already = False\n                    K = 0\n                    if hasattr(self.speculative_config, \"num_speculative_tokens\"):\n                        K = self.speculative_config.num_speculative_tokens\n                    if (hasattr(self, 'cudagraph_batch_sizes') and\n                        self.cudagraph_batch_sizes is not None):\n                        decode_cudagraph_batch_sizes = [\n                            x\n                            for x in self.cudagraph_batch_sizes\n                            if max_num_tokens >= x >= self.uniform_decode_query_len\n                        ]\n                        capture_already = len(decode_cudagraph_batch_sizes) > 0 and \\\n                            num_reqs*(1+K) <= max(decode_cudagraph_batch_sizes)\n                    if self.mlu_infer_mode == MLUInferMode.DECODE_ONLY and not \\\n                        all(x == K + 1 for x in scheduler_output.num_scheduled_tokens.values()):\n                        capture_already = False\n                    if capture_already:\n                        num_input_tokens = self._get_num_input_tokens(\n                            scheduler_output.total_num_scheduled_tokens\n                        )\n                    else:\n                        num_input_tokens = scheduler_output.total_num_scheduled_tokens\n                    '''\n                    ==================\n                    End of MLU Hijack\n                    ==================\n                    '''\n\n                (\n                    input_ids,\n                    inputs_embeds,\n                    positions,\n                    intermediate_tensors,\n                    model_kwargs,\n                    ec_connector_output,\n                ) = self._preprocess(\n                    scheduler_output, num_input_tokens, intermediate_tensors\n                )\n\n                '''\n                =============================\n                Modify by vllm_mlu\n                =============================\n                @breif: add padding for attention metadata in decode graph.\n                '''\n                # padding\n                self._padding_attn_metadata(attn_metadata,\n                    input_ids, inputs_embeds, capture_already,\n                    num_input_tokens, num_scheduled_tokens)\n                '''\n                ==================\n                End of MLU Hijack\n                ==================\n                '''\n\n            uniform_decode = (\n                max_num_scheduled_tokens == self.uniform_decode_query_len\n            ) and (num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)\n            batch_descriptor = BatchDescriptor(\n                num_tokens=num_input_tokens,\n                uniform_decode=uniform_decode,\n                has_lora=len(self.input_batch.lora_id_to_lora_request) > 0,\n            )\n            cudagraph_runtime_mode, batch_descriptor = (\n                self.cudagraph_dispatcher.dispatch(\n                    batch_descriptor,\n                    use_cascade_attn=cascade_attn_prefix_lens is not None,\n                )\n            )\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @breif: add prefill graph & capture already check\n            '''\n            if (self.prefill_enable_mlugraph and\n                attn_metadata.get(COMMON_METADATA_STR) is not None and\n                attn_metadata[COMMON_METADATA_STR].infer_mode == MLUInferMode.PREFILL_ONLY):\n                cudagraph_runtime_mode = CUDAGraphMode.FULL\n            if not capture_already:\n                cudagraph_runtime_mode = CUDAGraphMode.NONE\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n\n        # Set cudagraph mode to none if calc_kv_scales is true.\n        # KV scales calculation involves dynamic operations that are incompatible\n        # with CUDA graph capture.\n        if self.calculate_kv_scales:\n            cudagraph_runtime_mode = CUDAGraphMode.NONE\n            # Mark KV scales as calculated after the first forward pass\n            self.calculate_kv_scales = False\n            \n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: debug disagg cnpx.\n        '''\n        if mlu_envs.VLLM_DISAGG_CNPX_EXECUTE:\n            self.execute_cnpx_mark = cnpx.rangeStart(\"DP_\" + str(self.parallel_config.data_parallel_rank) + \"_TP_\" \\\n                                     + str(get_tensor_model_parallel_rank()) + \"_execute_model\" + \\\n                                     (\"_no_graph\" if cudagraph_runtime_mode == CUDAGraphMode.NONE else \"\"))\n        if mlu_envs.VLLM_DISAGG_CNPX_REQUEST:\n            self.request_cnpx_mark.clear()\n            for req in scheduler_output.scheduled_new_reqs:\n                self.request_cnpx_mark[req.req_id] = cnpx.rangeStart(req.req_id)\n            for req_id in scheduler_output.scheduled_cached_reqs.req_ids:\n                self.request_cnpx_mark[req_id] = cnpx.rangeStart(req_id)\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        if mlu_envs.VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:\n            start = torch.mlu.Event(enable_timing=True)\n            start.record()\n\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @breif: add set_sp_forward_context for sequence parallel.\n        '''\n        # Run the model.\n        # Use persistent buffers for CUDA graphs.\n        with (\n            set_forward_context(\n                attn_metadata,\n                self.vllm_config,\n                num_tokens=num_input_tokens,\n                num_tokens_across_dp=num_tokens_across_dp,\n                cudagraph_runtime_mode=cudagraph_runtime_mode,\n                batch_descriptor=batch_descriptor,\n                ubatch_slices=ubatch_slices,\n            ),\n            set_sp_forward_context(\n                attn_metadata,\n                self.vllm_config,\n                num_input_tokens,\n            ),\n            record_function_or_nullcontext(\"gpu_model_runner: forward\"),\n            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,\n        ):\n            if self.model_config.hf_config.model_type == \"deepseek_v4\":\n                model_kwargs[\"batch_to_kv_state\"] = torch.tensor([\n                    self.req_id_to_kv_state[req_id] for req_id in self.input_batch._req_ids\n                ], dtype=torch.int32, device=input_ids.device)\n            model_output = self._model_forward(\n                input_ids=input_ids,\n                positions=positions,\n                intermediate_tensors=intermediate_tensors,\n                inputs_embeds=inputs_embeds,\n                prefill_enable_mlugraph=self.prefill_enable_mlugraph,\n                **model_kwargs,\n            )\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        with record_function_or_nullcontext(\"gpu_model_runner: postprocess\"):\n            if self.use_aux_hidden_state_outputs:\n                # True when EAGLE 3 is used.\n                hidden_states, aux_hidden_states = model_output\n            else:\n                # Common case.\n                hidden_states = model_output\n                aux_hidden_states = None\n\n            if not self.broadcast_pp_output:\n                # Common case.\n                if not get_pp_group().is_last_rank:\n                    # Return the intermediate tensors.\n                    assert isinstance(hidden_states, IntermediateTensors)\n                    hidden_states.kv_connector_output = kv_connector_output\n                    return hidden_states\n\n                if self.is_pooling_model:\n                    # Return the pooling output.\n                    output = self._pool(\n                        hidden_states, num_scheduled_tokens, num_scheduled_tokens_np\n                    )\n                    output.kv_connector_output = kv_connector_output\n                    '''\n                    =============================\n                    Modify by vllm_mlu\n                    =============================\n                    @breif: add time markers for pooling model\n                    '''\n                    if mlu_envs.VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:\n                        end = torch.mlu.Event(enable_timing=True)\n                        end.record()\n                        self.time_markers.append([start, end])\n                    '''\n                    ==================\n                    End of MLU Hijack\n                    ==================\n                    '''\n                    return output\n\n                sample_hidden_states = hidden_states[logits_indices]\n                logits = self.model.compute_logits(sample_hidden_states)\n            else:\n                # Rare case.\n                assert not self.is_pooling_model\n\n                sample_hidden_states = hidden_states[logits_indices]\n                if not get_pp_group().is_last_rank:\n                    all_gather_tensors = {\n                        \"residual\": not is_residual_scattered_for_sp(\n                            self.vllm_config, num_input_tokens\n                        )\n                    }\n                    get_pp_group().send_tensor_dict(\n                        hidden_states.tensors,\n                        all_gather_group=get_tp_group(),\n                        all_gather_tensors=all_gather_tensors,\n                    )\n                    logits = None\n                else:\n                    logits = self.model.compute_logits(sample_hidden_states)\n\n                model_output_broadcast_data = {}\n                if logits is not None:\n                    model_output_broadcast_data[\"logits\"] = logits.contiguous()\n\n                model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(\n                    model_output_broadcast_data, src=len(get_pp_group().ranks) - 1\n                )\n                assert model_output_broadcast_data is not None\n                logits = model_output_broadcast_data[\"logits\"]\n\n        if mlu_envs.VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:\n            end = torch.mlu.Event(enable_timing=True)\n            end.record()\n            self.time_markers.append([start, end])\n\n        self.execute_model_state = ExecuteModelState(\n            scheduler_output,\n            logits,\n            spec_decode_metadata,\n            spec_decode_common_attn_metadata,\n            hidden_states,\n            sample_hidden_states,\n            aux_hidden_states,\n            kv_connector_output,\n        )\n        return None\n    \n    def response_remote_alloc_once(self) -> None:\n        if has_kv_transfer_group():\n            kv_connector = get_kv_transfer_group()\n            assert isinstance(kv_connector, KVConnectorBase)\n            kv_connector.response_remote_alloc_once()\n\n    @torch.inference_mode\n    def sample_tokens(\n        self, grammar_output: \"GrammarOutput | None\"\n    ) -> ModelRunnerOutput | AsyncMLUModelRunnerOutput | IntermediateTensors:\n        kv_connector_output = self.kv_connector_output\n        self.kv_connector_output = None\n\n        if self.execute_model_state is None:\n            # Nothing to do (PP non-final rank case), output isn't used.\n            if not kv_connector_output:\n                return None  # noqa\n\n            # In case of PP with kv transfer, we need to pass through the\n            # kv_connector_output\n            if kv_connector_output.is_empty():\n                return EMPTY_MODEL_RUNNER_OUTPUT\n\n            output = copy(EMPTY_MODEL_RUNNER_OUTPUT)\n            output.kv_connector_output = kv_connector_output\n            return output\n\n        # Unpack ephemeral state.\n        (\n            scheduler_output,\n            logits,\n            spec_decode_metadata,\n            spec_decode_common_attn_metadata,\n            hidden_states,\n            sample_hidden_states,\n            aux_hidden_states,\n            ec_connector_output,\n        ) = self.execute_model_state\n        # Clear ephemeral state.\n        self.execute_model_state = None\n\n        # Apply structured output bitmasks if present.\n        if grammar_output is not None:\n            apply_grammar_bitmask(\n                scheduler_output, grammar_output, self.input_batch, logits\n            )\n\n        with record_function_or_nullcontext(\"gpu_model_runner: sample\"):\n            sampler_output = self._sample(logits, spec_decode_metadata)\n\n        self.input_batch.prev_sampled_token_ids = None\n\n        def propose_draft_token_ids(\n            sampled_token_ids: torch.Tensor | list[np.ndarray],\n        ) -> None:\n            assert spec_decode_common_attn_metadata is not None\n            with record_function_or_nullcontext(\"gpu_model_runner: draft\"):\n                self._draft_token_ids = self.propose_draft_token_ids(\n                    scheduler_output,\n                    sampled_token_ids,\n                    self.input_batch.sampling_metadata,\n                    hidden_states,\n                    sample_hidden_states,\n                    aux_hidden_states,\n                    spec_decode_metadata,\n                    spec_decode_common_attn_metadata,\n                )\n\n        use_padded_batch_for_eagle = (\n            self.speculative_config\n            and self.speculative_config.use_eagle()\n            and not self.speculative_config.disable_padded_drafter_batch\n        )\n        effective_drafter_max_model_len = self.max_model_len\n        if effective_drafter_max_model_len is None:\n            effective_drafter_max_model_len = self.model_config.max_model_len\n        if (\n            self.speculative_config\n            and self.speculative_config.draft_model_config is not None\n            and self.speculative_config.draft_model_config.max_model_len is not None\n        ):\n            effective_drafter_max_model_len = (\n                self.speculative_config.draft_model_config.max_model_len\n            )\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: Force `input_fits_in_drafter` to be True to ensure that `self.uniform_decode_query_len` tokens are scheduled per batch during model execution.\n                This is required for graph validation and to keep the batch token count consistent with `self.uniform_decode_query_len` immediately after the prefill stage.\n        '''\n        # input_fits_in_drafter = spec_decode_common_attn_metadata and (\n        #     spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens\n        #     <= effective_drafter_max_model_len\n        # )\n        input_fits_in_drafter = True\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n        if use_padded_batch_for_eagle:\n            sampled_token_ids = sampler_output.sampled_token_ids\n            if input_fits_in_drafter:\n                # EAGLE speculative decoding can use the GPU sampled tokens\n                # as inputs, and does not need to wait for bookkeeping to finish.\n                propose_draft_token_ids(sampled_token_ids)\n            elif self.valid_sampled_token_count_event is not None:\n                next_token_ids, valid_sampled_tokens_count = (\n                    self.drafter.prepare_next_token_ids_padded(\n                        spec_decode_common_attn_metadata,\n                        sampled_token_ids,\n                        self.requests,\n                        self.input_batch,\n                        self.discard_request_indices.gpu,\n                        self.num_discarded_requests,\n                    )\n                )\n                self._copy_valid_sampled_token_count(\n                    next_token_ids, valid_sampled_tokens_count\n                )\n\n        with record_function_or_nullcontext(\"gpu_model_runner: bookkeep\"):\n            (\n                num_nans_in_logits,\n                logprobs_lists,\n                valid_sampled_token_ids,\n                prompt_logprobs_dict,\n                req_ids_output_copy,\n                req_id_to_index_output_copy,\n                invalid_req_indices,\n            ) = self._bookkeeping_sync(\n                scheduler_output,\n                sampler_output,\n                logits,\n                hidden_states,\n                scheduler_output.total_num_scheduled_tokens,\n                spec_decode_metadata,\n            )\n\n        if (\n            self.speculative_config\n            and not use_padded_batch_for_eagle\n            and input_fits_in_drafter\n        ):\n            # ngram and other speculative decoding methods use the sampled\n            # tokens on the CPU, so they are run after bookkeeping.\n            propose_draft_token_ids(valid_sampled_token_ids)\n\n        with record_function_or_nullcontext(\"gpu_model_runner: eplb\"):\n            self.eplb_step()\n        with record_function_or_nullcontext(\"gpu_model_runner: ModelRunnerOutput\"):\n            output = ModelRunnerOutput(\n                req_ids=req_ids_output_copy,\n                req_id_to_index=req_id_to_index_output_copy,\n                sampled_token_ids=valid_sampled_token_ids,\n                logprobs=logprobs_lists,\n                prompt_logprobs_dict=prompt_logprobs_dict,\n                pooler_output=[],\n                kv_connector_output=kv_connector_output,\n                ec_connector_output=ec_connector_output\n                if self.supports_mm_inputs\n                else None,\n                num_nans_in_logits=num_nans_in_logits,\n            )\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: supoort disagg for mlu.\n        '''\n        if has_kv_transfer_group():\n            get_kv_transfer_group().wait_for_save()\n            get_kv_transfer_group().clear_connector_metadata()\n        \n        if mlu_envs.VLLM_DISAGG_CNPX_EXECUTE:\n            current_stream = torch.mlu.current_stream()\n            current_stream.synchronize()\n            cnpx.rangeEnd(self.execute_cnpx_mark)\n        if mlu_envs.VLLM_DISAGG_CNPX_REQUEST:\n            current_stream = torch.mlu.current_stream()\n            current_stream.synchronize()\n            for req in scheduler_output.scheduled_new_reqs:\n                cnpx.rangeEnd(self.request_cnpx_mark[req.req_id])\n            for req_id in scheduler_output.scheduled_cached_reqs.req_ids:\n                cnpx.rangeEnd(self.request_cnpx_mark[req_id])\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n        if not self.use_async_scheduling:\n            return output\n        with record_function_or_nullcontext(\n            \"gpu_model_runner: AsyncGPUModelRunnerOutput\"\n        ):\n            async_output = AsyncMLUModelRunnerOutput(\n                model_runner_output=output,\n                sampled_token_ids=sampler_output.sampled_token_ids,\n                logprobs_tensors=sampler_output.logprobs_tensors,\n                invalid_req_indices=invalid_req_indices,\n                async_output_copy_stream=self.async_output_copy_stream,\n                vocab_size=self.input_batch.vocab_size,\n            )\n        with record_function_or_nullcontext(\n            \"gpu_model_runner: set_async_sampled_token_ids\"\n        ):\n            # Save ref of sampled_token_ids CPU tensor if the batch contains\n            # any requests with sampling params that require output ids.\n            self.input_batch.set_async_sampled_token_ids(\n                async_output.sampled_token_ids_cpu,\n                async_output.async_copy_ready_event,\n            )\n\n        return async_output\n\n    def propose_draft_token_ids(\n        self,\n        scheduler_output: \"SchedulerOutput\",\n        sampled_token_ids: torch.Tensor | list[np.ndarray],\n        sampling_metadata: SamplingMetadata,\n        hidden_states: torch.Tensor,\n        sample_hidden_states: torch.Tensor,\n        aux_hidden_states: torch.Tensor | None,\n        spec_decode_metadata: SpecDecodeMetadata | None,\n        common_attn_metadata: MLUCommonAttentionMetadata,\n    ) -> torch.Tensor | list[list[int]]:\n        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: draft model will build new FlashMLAMetadata,\n                so just unpad common_attn_metadata here.\n        '''\n        unpad_common_attn_metadata(\n            common_metadata=common_attn_metadata,\n            num_reqs=self.input_batch.num_reqs,\n            num_scheduled_tokens=num_scheduled_tokens\n        )\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n        if self.speculative_config.method == \"ngram\":\n            assert isinstance(sampled_token_ids, list)\n            assert isinstance(self.drafter, NgramProposer)\n            draft_token_ids = self.drafter.propose(\n                sampled_token_ids,\n                self.input_batch.req_ids,\n                self.input_batch.num_tokens_no_spec,\n                self.input_batch.token_ids_cpu,\n                self.input_batch.spec_decode_unsupported_reqs,\n            )\n        elif self.speculative_config.method == \"suffix\":\n            assert isinstance(sampled_token_ids, list)\n            assert isinstance(self.drafter, SuffixDecodingProposer)\n            draft_token_ids = self.drafter.propose(self.input_batch, sampled_token_ids)\n        elif self.speculative_config.method == \"medusa\":\n            assert isinstance(sampled_token_ids, list)\n            assert isinstance(self.drafter, MedusaProposer)\n\n            if sample_hidden_states.shape[0] == len(sampled_token_ids):\n                # The input to the target model does not include draft tokens.\n                hidden_states = sample_hidden_states\n            else:\n                indices = []\n                offset = 0\n                assert spec_decode_metadata is not None, (\n                    \"No spec decode metadata for medusa\"\n                )\n                for num_draft, tokens in zip(\n                    spec_decode_metadata.num_draft_tokens, sampled_token_ids\n                ):\n                    indices.append(offset + tokens.shape[0] - 1)\n                    offset += num_draft + 1\n                indices = torch.tensor(indices, device=self.device)\n                hidden_states = sample_hidden_states[indices]\n\n            draft_token_ids = self.drafter.propose(\n                target_hidden_states=hidden_states,\n                sampling_metadata=sampling_metadata,\n            )\n        elif self.speculative_config.use_eagle():\n            assert isinstance(self.drafter, EagleProposer)\n\n            if self.speculative_config.disable_padded_drafter_batch:\n                # When padded-batch is disabled, the sampled_token_ids should be\n                # the cpu-side list[list[int]] of valid sampled tokens for each\n                # request, with invalid requests having empty lists.\n                assert isinstance(sampled_token_ids, list), (\n                    \"sampled_token_ids should be a python list when\"\n                    \"padded-batch is disabled.\"\n                )\n                next_token_ids = self.drafter.prepare_next_token_ids_cpu(\n                    sampled_token_ids,\n                    self.requests,\n                    self.input_batch,\n                    scheduler_output.num_scheduled_tokens,\n                )\n            else:\n                # When using padded-batch, the sampled_token_ids should be\n                # the gpu tensor of sampled tokens for each request, of shape\n                # (num_reqs, num_spec_tokens + 1) with rejected tokens having\n                # value -1.\n                assert isinstance(sampled_token_ids, torch.Tensor), (\n                    \"sampled_token_ids should be a torch.Tensor when\"\n                    \"padded-batch is enabled.\"\n                )\n                next_token_ids, valid_sampled_tokens_count = (\n                    self.drafter.prepare_next_token_ids_padded(\n                        common_attn_metadata,\n                        sampled_token_ids,\n                        self.requests,\n                        self.input_batch,\n                        self.discard_request_indices.gpu,\n                        self.num_discarded_requests,\n                    )\n                )\n                self._copy_valid_sampled_token_count(\n                    next_token_ids, valid_sampled_tokens_count\n                )\n\n            if spec_decode_metadata is None:\n                token_indices_to_sample = None\n                # input_ids can be None for multimodal models.\n                target_token_ids = self.input_ids.gpu[:num_scheduled_tokens]\n                target_positions = self._get_positions(num_scheduled_tokens)\n                if self.use_aux_hidden_state_outputs:\n                    assert aux_hidden_states is not None\n                    target_hidden_states = torch.cat(\n                        [h[:num_scheduled_tokens] for h in aux_hidden_states], dim=-1\n                    )\n                else:\n                    target_hidden_states = hidden_states[:num_scheduled_tokens]\n                num_rejected_tokens_gpu = None\n                token_indices = None\n            else:\n                if self.speculative_config.disable_padded_drafter_batch:\n                    token_indices_to_sample = None\n                    common_attn_metadata, token_indices = self.drafter.prepare_inputs(\n                        common_attn_metadata,\n                        sampled_token_ids,\n                        spec_decode_metadata.num_draft_tokens,\n                    )\n                else:\n                    common_attn_metadata, token_indices, token_indices_to_sample, num_rejected_tokens_gpu = (\n                        self.drafter.prepare_inputs_padded(\n                            common_attn_metadata,\n                            spec_decode_metadata,\n                            valid_sampled_tokens_count,\n                        )\n                    )\n\n                target_token_ids = self.input_ids.gpu[token_indices]\n                target_positions = self._get_positions(token_indices)\n                if self.use_aux_hidden_state_outputs:\n                    assert aux_hidden_states is not None\n                    target_hidden_states = torch.cat(\n                        [h[token_indices] for h in aux_hidden_states], dim=-1\n                    )\n                else:\n                    target_hidden_states = hidden_states[token_indices]\n                '''\n                =============================\n                Modify by vllm_mlu\n                =============================\n                @brief: add debug info for draft accepted rate\n                '''\n                if mlu_envs.VLLM_MTP_DEBUG:\n                    batch_total_draft = sum(spec_decode_metadata.num_draft_tokens)\n                    batch_total_rejected = sum(num_rejected_tokens_gpu)\n                    self.total_draft_tokens += batch_total_draft\n                    self.total_accepted_tokens += (\n                        batch_total_draft - batch_total_rejected)\n                    if batch_total_draft > 0:\n                        batch_accept_rate = (\n                            batch_total_draft - batch_total_rejected\n                            ) / batch_total_draft\n                        print(f\"Batch Accept Rate: {batch_accept_rate:.4f}, \"\n                              f\"Total Accept Rate: {self.get_accept_rate():.4f}\")\n                '''\n                ==================\n                End of MLU Hijack\n                ==================\n                '''\n            if self.supports_mm_inputs:\n                mm_embed_inputs = self._gather_mm_embeddings(\n                    scheduler_output,\n                    shift_computed_tokens=1,\n                )\n            else:\n                mm_embed_inputs = None\n\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: keep full scheduled tokens for draft model compute\n            '''\n            target_token_ids = target_token_ids[:num_scheduled_tokens]\n            target_positions = target_positions[:num_scheduled_tokens]\n            target_hidden_states = target_hidden_states[:num_scheduled_tokens]\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n            draft_token_ids = self.drafter.propose(\n                target_token_ids=target_token_ids,\n                target_positions=target_positions,\n                target_hidden_states=target_hidden_states,\n                next_token_ids=next_token_ids,\n                last_token_indices=token_indices_to_sample,\n                sampling_metadata=sampling_metadata,\n                common_attn_metadata=common_attn_metadata,\n                num_rejected_tokens=num_rejected_tokens_gpu,\n                token_indices=token_indices,\n                time_markers=self.time_markers,\n            )\n\n        return draft_token_ids\n\n    def load_model(self, eep_scale_up: bool = False) -> None:\n        \"\"\"\n        Args:\n            eep_scale_up: the model loading is for elastic EP scale up.\n        \"\"\"\n        logger.info_once(\n            \"Starting to load model %s...\",\n            self.model_config.model,\n            scope=\"global\",\n        )\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief:\n        1. Set max_batched_token for SparseMoeMlp when enable avg moe.\n        2. modify rope's max_position_embeddings to max_model_len.\n        Those MUST be set before init model.\n        '''\n        if mlu_envs.VLLM_AVG_MOE_EN:\n            logger.warning(\"Inference with Moe avg dispatch, \"\n                \"it's only for deepseek-v3/r1 model's performance test,\"\n                \" and will result in precision anomalies. Be careful!\")\n            SparseMoeMlp.max_batched_token = max(self.model_config.max_model_len,\n                                                 self.scheduler_config.max_num_batched_tokens)\n        MLURotaryEmbedding.max_seq_len = self.model_config.max_model_len\n        MLURotaryEmbedding.max_model_len = self.model_config.max_model_len\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n        global_expert_loads, old_global_expert_indices_per_model, rank_mapping = (\n            EplbState.get_eep_state(self.parallel_config)\n            if eep_scale_up\n            else (None, None, None)\n        )\n\n        if self.parallel_config.enable_eplb:\n            self.eplb_state = EplbState(self.parallel_config, self.device)\n            eplb_models = 0\n        with DeviceMemoryProfiler() as m:\n            time_before_load = time.perf_counter()\n            model_loader = get_model_loader(self.load_config)\n            self.model = model_loader.load_model(\n                vllm_config=self.vllm_config, model_config=self.model_config\n            )\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: register model pre forward for rope optimization\n            '''\n            self.model.register_forward_pre_hook(_model_forward_pre_hook, with_kwargs=True)\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n            if self.lora_config:\n                self.model = self.load_lora_model(\n                    self.model, self.vllm_config, self.device\n                )\n            if hasattr(self, \"drafter\"):\n                logger.info_once(\"Loading drafter model...\")\n                self.drafter.load_model(self.model)\n                '''\n                =============================\n                Modify by vllm_mlu\n                =============================\n                @brief: Apply forward prehook to draft model.\n                '''\n                self.drafter.model.register_forward_pre_hook(_model_forward_pre_hook, with_kwargs=True)\n                '''\n                ==================\n                End of MLU Hijack\n                ==================\n                '''\n                if (\n                    hasattr(self.drafter, \"model\")\n                    and is_mixture_of_experts(self.drafter.model)\n                    and self.parallel_config.enable_eplb\n                ):\n                    logger.info_once(\n                        \"EPLB is enabled for drafter model %s.\",\n                        self.vllm_config.speculative_config.draft_model_config.model,\n                    )\n\n                    global_expert_load = (\n                        global_expert_loads[eplb_models]\n                        if global_expert_loads\n                        else None\n                    )\n                    old_global_expert_indices = (\n                        old_global_expert_indices_per_model[eplb_models]\n                        if old_global_expert_indices_per_model\n                        else None\n                    )\n                    if self.eplb_state is None:\n                        self.eplb_state = EplbState(self.parallel_config, self.device)\n                    self.eplb_state.add_model(\n                        self.drafter.model,\n                        self.vllm_config.speculative_config.draft_model_config,\n                        global_expert_load,\n                        old_global_expert_indices,\n                        rank_mapping,\n                    )\n                    eplb_models += 1\n\n            if self.use_aux_hidden_state_outputs:\n                if not supports_eagle3(self.get_model()):\n                    raise RuntimeError(\n                        \"Model does not support EAGLE3 interface but \"\n                        \"aux_hidden_state_outputs was requested\"\n                    )\n\n                # Try to get auxiliary layers from speculative config,\n                # otherwise use model's default layers\n                aux_layers = self._get_eagle3_aux_layers_from_config()\n                if aux_layers:\n                    logger.info(\n                        \"Using auxiliary layers from speculative config: %s\",\n                        aux_layers,\n                    )\n                else:\n                    aux_layers = self.model.get_eagle3_aux_hidden_state_layers()\n\n                self.model.set_aux_hidden_state_layers(aux_layers)\n            time_after_load = time.perf_counter()\n        self.model_memory_usage = m.consumed_memory\n        logger.info_once(\n            \"Model loading took %.4f GiB memory and %.6f seconds\",\n            self.model_memory_usage / GiB_bytes,\n            time_after_load - time_before_load,\n            scope=\"local\",\n        )\n        prepare_communication_buffer_for_model(self.model)\n        self.is_multimodal_pruning_enabled = (\n            supports_multimodal_pruning(self.get_model())\n            and self.model_config.multimodal_config.is_multimodal_pruning_enabled()\n        )\n\n        if is_mixture_of_experts(self.model) and self.parallel_config.enable_eplb:\n            logger.info_once(\"EPLB is enabled for model %s.\", self.model_config.model)\n            global_expert_load = (\n                global_expert_loads[eplb_models] if global_expert_loads else None\n            )\n            old_global_expert_indices = (\n                old_global_expert_indices_per_model[eplb_models]\n                if old_global_expert_indices_per_model\n                else None\n            )\n            assert self.eplb_state is not None\n            self.eplb_state.add_model(\n                self.model,\n                self.model_config,\n                global_expert_load,\n                old_global_expert_indices,\n                rank_mapping,\n            )\n\n        if (\n            self.vllm_config.compilation_config.mode\n            == CompilationMode.STOCK_TORCH_COMPILE\n            and supports_dynamo()\n        ):\n            backend = self.vllm_config.compilation_config.init_backend(self.vllm_config)\n            compilation_counter.stock_torch_compile_count += 1\n            self.model.compile(fullgraph=True, backend=backend)\n            return\n        # for other compilation modes, cudagraph behavior is controlled by\n        # CudagraphWraper and CudagraphDispatcher of vllm.\n\n        # wrap the model with full cudagraph wrapper if needed.\n        if (\n            self.compilation_config.cudagraph_mode.has_full_cudagraphs()\n            and not self.parallel_config.enable_dbo\n        ):\n            self.model = MLUGraphWrapper(\n                self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL\n            )\n        elif self.parallel_config.enable_dbo:\n            if self.compilation_config.cudagraph_mode.has_full_cudagraphs():\n                self.model = UBatchWrapper(\n                    self.model, self.vllm_config, CUDAGraphMode.FULL, self.device\n                )\n            else:\n                self.model = UBatchWrapper(\n                    self.model, self.vllm_config, CUDAGraphMode.NONE, self.device\n                )\n\n    def _get_prompt_logprobs_dict(\n        self,\n        hidden_states: torch.Tensor,\n        num_scheduled_tokens: dict[str, int],\n    ) -> dict[str, LogprobsTensors | None]:\n        num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs\n        if not num_prompt_logprobs_dict:\n            return {}\n\n        in_progress_dict = self.input_batch.in_progress_prompt_logprobs_cpu\n        prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}\n\n        # Since prompt logprobs are a rare feature, prioritize simple,\n        # maintainable loop over optimal performance.\n        completed_prefill_reqs = []\n        for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items():\n            num_tokens = num_scheduled_tokens[req_id]\n\n            # Get metadata for this request.\n            request = self.requests[req_id]\n            if request.prompt_token_ids is None:\n                # Prompt logprobs is incompatible with prompt embeddings\n                continue\n\n            num_prompt_tokens = len(request.prompt_token_ids)\n            prompt_token_ids = torch.tensor(request.prompt_token_ids).to(\n                self.device, non_blocking=True\n            )\n\n            # Set up target LogprobsTensors object.\n            logprobs_tensors = in_progress_dict.get(req_id)\n            if not logprobs_tensors:\n                # Create empty logprobs CPU tensors for the entire prompt.\n                # If chunked, we'll copy in slice by slice.\n                logprobs_tensors = LogprobsTensors.empty_cpu(\n                    num_prompt_tokens - 1, num_prompt_logprobs + 1\n                )\n                in_progress_dict[req_id] = logprobs_tensors\n\n            # Determine number of logits to retrieve.\n            start_idx = request.num_computed_tokens\n            start_tok = start_idx + 1\n            num_remaining_tokens = num_prompt_tokens - start_tok\n            if num_tokens <= num_remaining_tokens:\n                # This is a chunk, more tokens remain.\n                # In the == case, there are no more prompt logprobs to produce\n                # but we want to defer returning them to the next step where we\n                # have new generated tokens to return.\n                num_logits = num_tokens\n            else:\n                # This is the last chunk of prompt tokens to return.\n                num_logits = num_remaining_tokens\n                completed_prefill_reqs.append(req_id)\n                prompt_logprobs_dict[req_id] = logprobs_tensors\n\n            if num_logits <= 0:\n                # This can happen for the final chunk if we prefilled exactly\n                # (num_prompt_tokens - 1) tokens for this request in the prior\n                # step. There are no more prompt logprobs to produce.\n                '''\n                =============================\n                Modify by vllm_mlu\n                =============================\n                @brief: remove the prompt_logprobs for final chunk request\n                '''\n                del prompt_logprobs_dict[req_id]\n                '''\n                ==================\n                End of MLU Hijack\n                ==================\n                '''\n                continue\n\n            # Get the logits corresponding to this req's prompt tokens.\n            # If this is a partial request (i.e. chunked prefill),\n            # then there is prompt logprob generated for each index.\n            req_idx = self.input_batch.req_id_to_index[req_id]\n            offset = self.query_start_loc.np[req_idx].item()\n            prompt_hidden_states = hidden_states[offset : offset + num_logits]\n            logits = self.model.compute_logits(prompt_hidden_states)\n\n            # Get the \"target\" tokens for each index. For prompt at index i,\n            # the token at prompt index i+1 is the \"sampled\" token we want\n            # to gather the logprob for.\n            tgt_token_ids = prompt_token_ids[start_tok : start_tok + num_logits]\n\n            # Compute prompt logprobs.\n            logprobs = self.sampler.compute_logprobs(logits)\n            token_ids, logprobs, ranks = self.sampler.gather_logprobs(\n                logprobs, num_prompt_logprobs, tgt_token_ids\n            )\n\n            # Transfer GPU->CPU async.\n            chunk_slice = slice(start_idx, start_idx + num_logits)\n            logprobs_tensors.logprob_token_ids[chunk_slice].copy_(\n                token_ids, non_blocking=True\n            )\n            logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, non_blocking=True)\n            logprobs_tensors.selected_token_ranks[chunk_slice].copy_(\n                ranks, non_blocking=True\n            )\n\n        # Remove requests that have completed prefill from the batch\n        # num_prompt_logprobs_dict.\n        for req_id in completed_prefill_reqs:\n            del num_prompt_logprobs_dict[req_id]\n            del in_progress_dict[req_id]\n\n        # Must synchronize the non-blocking GPU->CPU transfers.\n        if prompt_logprobs_dict:\n            self._sync_device()\n\n        return prompt_logprobs_dict\n\n    def _build_attention_metadata(\n        self,\n        total_num_scheduled_tokens: int,\n        max_num_scheduled_tokens: int,\n        num_reqs: int,\n        ubatch_slices: UBatchSlices | None = None,\n        logits_indices: torch.Tensor | None = None,\n        use_spec_decode: bool = False,\n        for_cudagraph_capture: bool = False,\n        scheduled_encoder_inputs: dict[str, list[int]] | None = None,\n        cascade_attn_prefix_lens: list[list[int]] | None = None,\n        mlu_infer_mode: MLUInferMode | None = None,\n    ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]:\n        \"\"\"\n        :return: tuple[attn_metadata, spec_decode_common_attn_metadata]\n        \"\"\"\n        logits_indices_padded = None\n        num_logits_indices = 0\n        if logits_indices is not None:\n            num_logits_indices = logits_indices.size(0)\n            if self.cache_config.kv_sharing_fast_prefill:\n                logits_indices_padded = self._prepare_kv_sharing_fast_prefill(\n                    logits_indices\n                )\n\n        # update seq_lens of decode reqs under DCP.\n        if self.dcp_world_size > 1:\n            self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens(\n                self.seq_lens.cpu[:num_reqs],\n                self.dcp_world_size,\n                self.dcp_rank,\n                self.parallel_config.dcp_kv_cache_interleave_size,\n            )\n            self.dcp_local_seq_lens.copy_to_gpu(num_reqs)\n\n        attn_metadata: PerLayerAttnMetadata = {}\n        if ubatch_slices is not None:\n            attn_metadata = [dict() for _ in range(len(ubatch_slices))]\n\n        # Used in the below loop\n        query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]\n        query_start_loc_cpu = self.query_start_loc.cpu[: num_reqs + 1]\n        seq_lens = self.seq_lens.gpu[:num_reqs]\n        seq_lens_cpu = self.seq_lens.cpu[:num_reqs]\n        num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[\n            :num_reqs\n        ]\n        dcp_local_seq_lens = (\n            self.dcp_local_seq_lens.gpu[:num_reqs] if self.dcp_world_size > 1 else None\n        )\n        spec_decode_common_attn_metadata = None\n\n        if for_cudagraph_capture:\n            # For some attention backends (e.g. FA) with sliding window models we need\n            # to make sure the backend see a max_seq_len that is larger to the sliding\n            # window size when capturing to make sure the correct kernel is selected.\n            max_seq_len = self.max_model_len\n        else:\n            max_seq_len = self.seq_lens.np[:num_reqs].max().item()\n\n        if use_spec_decode:\n            self.num_accepted_tokens.np[:num_reqs] = (\n                self.input_batch.num_accepted_tokens_cpu[:num_reqs]\n            )\n            self.num_accepted_tokens.np[num_reqs:].fill(1)\n            self.num_accepted_tokens.copy_to_gpu()\n\n        # Prepare the attention metadata for each KV cache group and make layers\n        # in the same group share the same metadata.\n        for kv_cache_gid, kv_cache_group in enumerate(\n            self.kv_cache_config.kv_cache_groups\n        ):\n            encoder_seq_lens = self._get_encoder_seq_lens(\n                scheduled_encoder_inputs or {},\n                kv_cache_group.kv_cache_spec,\n                num_reqs,\n            )\n\n            if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):\n                # Encoder-only layers do not have KV cache, so we need to\n                # create a dummy block table and slot mapping for them.\n                blk_table_tensor = torch.zeros(\n                    (num_reqs, 1),\n                    dtype=torch.int32,\n                    device=self.device,\n                )\n                slot_mapping = torch.zeros(\n                    (total_num_scheduled_tokens,),\n                    dtype=torch.int64,\n                    device=self.device,\n                )\n            else:\n                blk_table = self.input_batch.block_table[kv_cache_gid]\n                blk_table_tensor = blk_table.get_device_tensor(num_reqs)\n                slot_mapping = blk_table.slot_mapping.gpu[:total_num_scheduled_tokens]\n\n                # Fill unused with -1. Needed for reshape_and_cache in full cuda\n                # graph mode.\n                blk_table.slot_mapping.gpu[total_num_scheduled_tokens:].fill_(-1)\n            \"\"\"\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: replace CommonAttentionMetadata with MLUCommonAttentionMetadata\n            \"\"\"\n            common_attn_metadata = MLUCommonAttentionMetadata(\n                    query_start_loc=query_start_loc,\n                    query_start_loc_cpu=query_start_loc_cpu,\n                    seq_lens=seq_lens,\n                    seq_lens_cpu=seq_lens_cpu,\n                    num_computed_tokens_cpu=num_computed_tokens_cpu,\n                    num_reqs=num_reqs,\n                    num_actual_tokens=total_num_scheduled_tokens,\n                    max_query_len=max_num_scheduled_tokens,\n                    max_seq_len=max_seq_len,\n                    block_table_tensor=blk_table_tensor,\n                    slot_mapping=slot_mapping,\n                    causal=True,\n                    dcp_local_seq_lens=dcp_local_seq_lens,\n                    seq_start_loc=self.seq_start_loc.gpu[: num_reqs + 1],\n                    seq_start_loc_cpu=self.seq_start_loc.cpu[: num_reqs + 1],\n                    infer_mode=mlu_infer_mode,\n                    num_prefill_query_tokens=total_num_scheduled_tokens,\n                    num_prefill_kv_tokens=total_num_scheduled_tokens,\n                )\n            \"\"\"\n            =================\n            End of MLU Hijack\n            =================\n            \"\"\"\n            if self.speculative_config and spec_decode_common_attn_metadata is None:\n                if isinstance(self.drafter, EagleProposer):\n                    \"\"\"\n                    =============================\n                    Modify by vllm_mlu\n                    =============================\n                    @brief: replace attn metadata name to prefill_attn name\n                    \"\"\"\n                    attn_layer_name = self.drafter.attn_layer_names[0]\n                    if self.model_config.is_deepseek_mla and attn_layer_name.endswith(\"self_attn.attn\"):\n                        attn_layer_name = attn_layer_name.replace(\n                            \"self_attn.attn\", \"self_attn.mla_attn\")\n                    if attn_layer_name in kv_cache_group.layer_names:\n                        spec_decode_common_attn_metadata = common_attn_metadata\n                    \"\"\"\n                    =================\n                    End of MLU Hijack\n                    =================\n                    \"\"\"\n                else:\n                    spec_decode_common_attn_metadata = common_attn_metadata\n\n            for attn_gid, attn_group in enumerate(self.attn_groups[kv_cache_gid]):\n                cascade_attn_prefix_len = (\n                    cascade_attn_prefix_lens[kv_cache_gid][attn_gid]\n                    if cascade_attn_prefix_lens\n                    else 0\n                )\n                builder = attn_group.get_metadata_builder()\n\n                extra_attn_metadata_args = {}\n                if use_spec_decode and isinstance(builder, GDNAttentionMetadataBuilder):\n                    extra_attn_metadata_args = dict(\n                        num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs],\n                        num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[\n                            :num_reqs\n                        ],\n                    )\n\n                if ubatch_slices is not None:\n                    common_attn_metadata_list = split_attn_metadata(\n                        ubatch_slices, common_attn_metadata\n                    )\n                    for ubid, common_attn_metadata in enumerate(\n                        common_attn_metadata_list\n                    ):\n                        builder = attn_group.get_metadata_builder(ubatch_id=ubid)\n                        if for_cudagraph_capture:\n                            attn_metadata_i = builder.build_for_cudagraph_capture(\n                                common_attn_metadata\n                            )\n                        else:\n                            attn_metadata_i = builder.build(\n                                common_prefix_len=cascade_attn_prefix_len,\n                                common_attn_metadata=common_attn_metadata,\n                            )\n                        for layer_name in kv_cache_group.layer_names:\n                            assert type(attn_metadata) is list\n                            attn_metadata[ubid][layer_name] = attn_metadata_i\n                else:\n                    assert isinstance(attn_metadata, dict)\n                    if for_cudagraph_capture:\n                        attn_metadata_i = builder.build_for_cudagraph_capture(\n                            common_attn_metadata\n                        )\n                    else:\n                        attn_metadata_i = builder.build(\n                            common_prefix_len=cascade_attn_prefix_len,\n                            common_attn_metadata=common_attn_metadata,\n                            **extra_attn_metadata_args,\n                        )\n                    for layer_name in attn_group.layer_names:\n                        attn_metadata[layer_name] = attn_metadata_i\n                \"\"\"\n                =============================\n                Modify by vllm_mlu\n                =============================\n                @brief: bind decode_attn metadata to prefill_attn\n                \"\"\"\n                for layer_name in attn_group.layer_names:\n                    if (\n                        self.model_config.is_deepseek_mla\n                        and layer_name.endswith(\"self_attn.mla_attn\")\n                    ):\n                        prefill_attn_name = layer_name.replace(\n                            \"self_attn.mla_attn\", \"self_attn.attn\"\n                        )\n                        attn_metadata[prefill_attn_name] = attn_metadata[layer_name]\n\n                    # matches self_attn.0.attn or self_attn.1.attn for longcat-flash\n                    if (\n                        self.model_config.is_longcat_flash\n                        and (match := re.match(r\".*self_attn\\.(0|1)\\.mla_attn$\", layer_name))\n                    ):\n                        # Extract the captured digit (0 or 1)\n                        digit = match.group(1)\n                        prefill_attn_name = layer_name.replace(\n                            f\"self_attn.{digit}.mla_attn\",\n                            f\"self_attn.{digit}.attn\"\n                        )\n                        attn_metadata[prefill_attn_name] = attn_metadata_i\n                \"\"\"\n                =================\n                End of MLU Hijack\n                =================\n                \"\"\"\n            \"\"\"\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: Add common_attn_metadata to attn_metadata\n            \"\"\"\n            attn_metadata[COMMON_METADATA_STR] = common_attn_metadata\n            \"\"\"\n            =================\n            End of MLU Hijack\n            =================\n            \"\"\"\n\n        return attn_metadata, spec_decode_common_attn_metadata\n\n    def _padding_attn_metadata(\n        self,\n        attn_metadata: MLACommonMetadata | FlashAttentionMetadata,\n        input_ids: torch.Tensor | None,\n        inputs_embeds: torch.Tensor | None,\n        captured_already: bool,\n        num_input_tokens: int,\n        num_scheduled_tokens: int\n    ) -> None:\n        common_metadata = attn_metadata[COMMON_METADATA_STR]\n        decode_only = common_metadata.is_decode_only\n        if decode_only and captured_already:\n            # If the model is decode only, we can use full graph.\n            # use_full_graph = use_full_graph # and captured_already\n            # Update attn_metadata for full graph.\n            K = 0\n            if (self.speculative_config is not None\n                and self.speculative_config.num_speculative_tokens > 0\n            ):\n                K = self.speculative_config.num_speculative_tokens\n\n            if num_input_tokens != num_scheduled_tokens:\n                for kv_cache_group_id, kv_cache_group_spec in enumerate(\n                        self.kv_cache_config.kv_cache_groups):\n                    block_table = self.input_batch.block_table[kv_cache_group_id]\n                    first_layer_name = kv_cache_group_spec.layer_names[0]\n                    attn_metadata_i = attn_metadata[first_layer_name]\n                    num_reqs = self.input_batch.num_reqs\n                    num_padded_reqs = self.vllm_config.pad_for_cudagraph(num_reqs * (1 + K)) // (1 + K)\n                    pad_attn_metadata(\n                        attn_metadata_i, common_metadata, block_table,\n                        self, num_scheduled_tokens, num_input_tokens,\n                        num_reqs, num_padded_reqs,\n                    )\n\n\n    \"\"\"\n    =============================\n    Modify by vllm_mlu\n    =============================\n    @brief: Add prefill input parameters\n        @parameters: is_capturing_prefill, prefill_batch_size, prefill_seq_len\n    \"\"\"\n    @torch.inference_mode()\n    def _dummy_run(\n        self,\n        num_tokens: int,\n        is_capturing_prefill: bool = False,\n        prefill_batch_size: int = None,\n        prefill_seq_len: int = None,\n        cudagraph_runtime_mode: CUDAGraphMode | None = None,\n        force_attention: bool = False,\n        uniform_decode: bool = False,\n        allow_microbatching: bool = True,\n        skip_eplb: bool = False,\n        is_profile: bool = False,\n        create_mixed_batch: bool = False,\n        remove_lora: bool = True,\n        activate_lora: bool = False,\n    ) -> tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Run a dummy forward pass to warm up/profile run or capture the\n        CUDA graph for the model.\n\n        Args:\n            num_tokens: Number of tokens to run the dummy forward pass.\n            cudagraph_runtime_mode: used to control the behavior.\n                - if not set will determine the cudagraph mode based on using\n                    the self.cudagraph_dispatcher.\n                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run\n                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.\n                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is\n                    needed.\n            force_attention: If True, always create attention metadata. Used to\n                warm up attention backend when mode is NONE.\n            uniform_decode: If True, the batch is a uniform decode batch.\n            skip_eplb: If True, skip EPLB state update.\n            is_profile: If True, this is a profile run.\n            create_mixed_batch: If True, create a mixed batch with both decode\n                (1 token) and prefill (multiple tokens) requests.\n            remove_lora: If False, dummy LoRAs are not destroyed after the run\n            activate_lora: If False, dummy_run is performed without LoRAs.\n        \"\"\"\n        assert (\n            cudagraph_runtime_mode is None\n            or cudagraph_runtime_mode.valid_runtime_modes()\n        )\n\n        # If cudagraph_mode.decode_mode() == FULL and\n        # cudagraph_mode.separate_routine(). This means that we are using\n        # different graphs and/or modes for mixed prefill-decode batches vs.\n        # uniform decode batches. A uniform decode batch means that all\n        # requests have identical query length, except a potential virtual\n        # request (shorter) in the batch account for padding.\n        # Uniform decode batch could either be common pure decode, where\n        # max_query_len == 1, or speculative decode, where\n        # max_query_len == 1 + num_spec_decode_tokens.\n\n        # When setting max_query_len = 1, we switch to and capture the optimized\n        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization\n        # for GQA/MQA.\n        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens\n\n        # Set num_scheduled_tokens based on num_tokens and max_num_seqs\n        # for dummy run with LoRA so that the num_reqs collectively\n        # has num_tokens in total.\n        assert num_tokens <= self.scheduler_config.max_num_batched_tokens\n        max_num_reqs = self.scheduler_config.max_num_seqs\n        if create_mixed_batch:\n            assert not uniform_decode\n            # Create mixed batch:\n            # first half decode tokens, second half one prefill\n            num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2)\n            num_prefill_tokens = num_tokens - num_decode_tokens\n            num_reqs = num_decode_tokens + 1\n\n            # Create decode requests (1 token each) followed by prefill request\n            num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]\n            # Note: Overriding max_query_len to be the prefill tokens\n            max_query_len = num_prefill_tokens\n        elif uniform_decode:\n            assert not create_mixed_batch\n            num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len))\n            num_scheduled_tokens_list = [max_query_len] * num_reqs\n            if num_tokens % max_query_len != 0:\n                num_scheduled_tokens_list[-1] = num_tokens % max_query_len\n        else:\n            num_reqs = min(num_tokens, max_num_reqs)\n            min_tokens_per_req = num_tokens // num_reqs\n            num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs\n            num_scheduled_tokens_list[-1] += num_tokens % num_reqs\n\n        assert sum(num_scheduled_tokens_list) == num_tokens\n        assert len(num_scheduled_tokens_list) == num_reqs\n        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)\n        total_num_scheduled_tokens = int(num_scheduled_tokens.sum())\n        num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)\n\n        # Disable DP padding when running eager\n        allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE\n\n        # We currently only microbatch if the number of tokens is\n        # over a certain threshold.\n        ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(\n            num_tokens_unpadded=total_num_scheduled_tokens,\n            parallel_config=self.vllm_config.parallel_config,\n            allow_microbatching=allow_microbatching,\n            allow_dp_padding=allow_dp_padding,\n            num_tokens_padded=total_num_scheduled_tokens,\n            uniform_decode=uniform_decode,\n            num_scheduled_tokens_per_request=num_scheduled_tokens,\n        )\n        num_tokens_after_padding = num_tokens\n        if num_tokens_across_dp is not None:\n            dp_rank = self.parallel_config.data_parallel_rank\n            num_tokens_after_padding = int(num_tokens_across_dp[dp_rank])\n\n        attn_metadata: PerLayerAttnMetadata | None = None\n\n        # If force_attention is True, we always capture attention. Otherwise,\n        # it only happens for cudagraph_runtime_mode=FULL.\n        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:\n            \"\"\"\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: use prefill_seq_len to build seq_lens\n                when prefill capture\n            \"\"\"\n            if create_mixed_batch:\n                # In the mixed batch mode (used for FI warmup), we use\n                # shorter sequence lengths to run faster.\n                # TODO(luka) better system for describing dummy batches\n                seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]\n            elif is_capturing_prefill:\n                seq_lens = prefill_seq_len\n            else:\n                seq_lens = max_query_len\n            \"\"\"\n            =================\n            End of MLU Hijack\n            =================\n            \"\"\"\n            self.seq_lens.np[:num_reqs] = seq_lens\n            self.seq_lens.np[num_reqs:] = 0\n            self.seq_lens.copy_to_gpu()\n\n            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)\n            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens\n            self.query_start_loc.copy_to_gpu()\n            \"\"\"\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: compute seq_start_loc and mlu_infer_mode.\n            @brief: use prefill_batch_size to build seq_start_loc\n            \"\"\"\n            cu_seqlens_k = np.cumsum(self.seq_lens.np[:num_reqs])\n            self.seq_start_loc.np[0] = 0\n            self.seq_start_loc.np[1 : num_reqs + 1] = cu_seqlens_k\n            self.seq_start_loc.copy_to_gpu()\n\n            max_computed_tokens = np.max(self.input_batch.num_computed_tokens_cpu[:num_reqs])\n            mlu_infer_mode = MLUInferMode.build(\n                max_query_len=max_query_len,\n                max_computed_tokens=max_computed_tokens,\n                uniform_decode_query_len=self.uniform_decode_query_len)\n            if is_capturing_prefill:\n                attn_metadata, _ = self._build_attention_metadata(\n                    total_num_scheduled_tokens=num_tokens,\n                    max_num_scheduled_tokens=max_query_len,\n                    num_reqs=prefill_batch_size,\n                    ubatch_slices=ubatch_slices,\n                    for_cudagraph_capture=True,\n                    mlu_infer_mode=MLUInferMode.PREFILL_ONLY,\n                )\n            else:\n                attn_metadata, _ = self._build_attention_metadata(\n                    total_num_scheduled_tokens=num_tokens,\n                    max_num_scheduled_tokens=max_query_len,\n                    num_reqs=num_reqs,\n                    ubatch_slices=ubatch_slices,\n                    for_cudagraph_capture=True,\n                    mlu_infer_mode=mlu_infer_mode,\n                )\n            \"\"\"\n            =================\n            End of MLU Hijack\n            =================\n            \"\"\"\n\n        with self.maybe_dummy_run_with_lora(\n            self.lora_config,\n            num_scheduled_tokens,\n            num_sampled_tokens,\n            activate_lora,\n            remove_lora,\n        ):\n            # Make sure padding doesn't exceed max_num_tokens\n            assert num_tokens_after_padding <= self.max_num_tokens\n            model_kwargs = self._init_model_kwargs(num_tokens_after_padding)\n            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:\n                input_ids = None\n                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_after_padding]\n                model_kwargs = {\n                    **model_kwargs,\n                    **self._dummy_mm_kwargs(num_reqs),\n                }\n            elif self.enable_prompt_embeds:\n                input_ids = None\n                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_after_padding]\n                model_kwargs = self._init_model_kwargs(num_tokens_after_padding)\n            else:\n                input_ids = self.input_ids.gpu[:num_tokens_after_padding]\n                inputs_embeds = None\n\n            if self.uses_mrope:\n                positions = self.mrope_positions.gpu[:, :num_tokens_after_padding]\n            else:\n                positions = self.positions.gpu[:num_tokens_after_padding]\n\n            if get_pp_group().is_first_rank:\n                intermediate_tensors = None\n            else:\n                if self.intermediate_tensors is None:\n                    self.intermediate_tensors = (\n                        self.model.make_empty_intermediate_tensors(\n                            batch_size=self.max_num_tokens,\n                            dtype=self.model_config.dtype,\n                            device=self.device,\n                        )\n                    )\n\n                intermediate_tensors = self.sync_and_slice_intermediate_tensors(\n                    num_tokens_after_padding, None, False\n                )\n\n            # filter out the valid batch descriptor\n            _cg_mode, batch_descriptor = (\n                self.cudagraph_dispatcher.dispatch(\n                    BatchDescriptor(\n                        num_tokens=num_tokens_after_padding,\n                        uniform_decode=uniform_decode,\n                        has_lora=activate_lora and self.lora_config is not None,\n                    )\n                )\n                if not is_profile\n                else (CUDAGraphMode.NONE, None)\n            )\n            \"\"\"\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: adjust cudagraph mode for\n                prefill graph capture\n            \"\"\"\n            if is_capturing_prefill:\n                _cg_mode = cudagraph_runtime_mode\n            \"\"\"\n            =================\n            End of MLU Hijack\n            =================\n            \"\"\"\n            if cudagraph_runtime_mode is not None:\n                # we allow forcing NONE when the dispatcher disagrees to support\n                # warm ups for cudagraph capture\n                assert (\n                    cudagraph_runtime_mode == CUDAGraphMode.NONE\n                    or cudagraph_runtime_mode == _cg_mode\n                ), (\n                    f\"Cudagraph runtime mode mismatch at dummy_run. \"\n                    f\"Expected {_cg_mode}, but got {cudagraph_runtime_mode}.\"\n                )\n            else:\n                cudagraph_runtime_mode = _cg_mode\n\n            if ubatch_slices is not None:\n                # Adjust values to reflect a single ubatch.\n                # TODO(sage,lucas): this is cruft that should be addressed in\n                #  the padding refactor.\n                num_tokens_after_padding = ubatch_slices[0].num_tokens\n                if num_tokens_across_dp is not None:\n                    num_tokens_across_dp[:] = num_tokens_after_padding\n\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @breif: add set_sp_forward_context for sequence parallel.\n            '''\n            with (\n                self.maybe_randomize_inputs(input_ids),\n                set_forward_context(\n                    attn_metadata,\n                    self.vllm_config,\n                    num_tokens=num_tokens_after_padding,\n                    num_tokens_across_dp=num_tokens_across_dp,\n                    cudagraph_runtime_mode=cudagraph_runtime_mode,\n                    batch_descriptor=batch_descriptor,\n                    ubatch_slices=ubatch_slices,\n                ),\n                set_sp_forward_context(\n                    attn_metadata,\n                    self.vllm_config,\n                    num_tokens_after_padding,\n                ),\n            ):\n                if self.model_config.hf_config.model_type == \"deepseek_v4\":\n                    assert self.kv_state_free_slots, \\\n                        \"At least one slot is needed to run dummy model\"\n                    model_kwargs[\"batch_to_kv_state\"] = torch.tensor([\n                            list(self.kv_state_free_slots)[0]\n                        ] * num_reqs,\n                        dtype=torch.int32,\n                        device=input_ids.device,\n                    )\n                outputs = self.model(\n                    is_capturing_prefill=is_capturing_prefill,\n                    prefill_batch_size=prefill_batch_size,\n                    prefill_seq_len=prefill_seq_len,\n                    input_ids=input_ids,\n                    positions=positions,\n                    intermediate_tensors=intermediate_tensors,\n                    inputs_embeds=inputs_embeds,\n                    **model_kwargs,\n                )\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n            if self.use_aux_hidden_state_outputs:\n                hidden_states, _ = outputs\n            else:\n                hidden_states = outputs\n\n            if self.speculative_config and self.speculative_config.use_eagle():\n                assert isinstance(self.drafter, EagleProposer)\n                use_cudagraphs = (\n                    cudagraph_runtime_mode == CUDAGraphMode.FULL\n                    and not self.speculative_config.enforce_eager\n                )\n\n                # Note(gnovack) - We need to disable cudagraphs for one of the two\n                # lora cases when cudagraph_specialize_lora is enabled. This is a\n                # short term mitigation for issue mentioned in\n                # https://github.com/vllm-project/vllm/issues/28334\n                if self.compilation_config.cudagraph_specialize_lora and activate_lora:\n                    use_cudagraphs = False\n\n                self.drafter.dummy_run(\n                    attn_metadata,\n                    num_tokens,\n                    use_cudagraphs=use_cudagraphs,\n                )\n\n        # This is necessary to avoid blocking DP.\n        # For dummy runs, we typically skip EPLB since we don't have any real\n        # requests to process.\n        # However, in DP settings, there may be cases when some DP ranks do\n        # not have any requests to process, so they're executing dummy batches.\n        # In such cases, we still have to trigger EPLB to make sure\n        # ranks execute the rearrangement in synchronization.\n        if not skip_eplb:\n            self.eplb_step(is_dummy=True, is_profile=is_profile)\n\n        logit_indices = np.cumsum(num_scheduled_tokens) - 1\n        logit_indices_device = torch.from_numpy(logit_indices).to(\n            self.device, non_blocking=True\n        )\n        return hidden_states, hidden_states[logit_indices_device]\n\n    def _capture_cudagraphs(\n        self,\n        is_capturing_prefill: bool = False,\n        prefill_batch_size: int = 0,\n        prefill_seq_len: int = 0,\n        compilation_cases: list[tuple[int, bool]] = []  ,\n        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,\n        uniform_decode: bool = False,\n    ):\n        assert (\n            cudagraph_runtime_mode != CUDAGraphMode.NONE\n            and cudagraph_runtime_mode.valid_runtime_modes()\n        ), f\"Invalid cudagraph runtime mode: {cudagraph_runtime_mode}\"\n\n        # Only rank 0 should print progress bar during capture\n        if is_global_first_rank():\n            compilation_cases = tqdm(\n                compilation_cases,\n                disable=not self.load_config.use_tqdm_on_load,\n                desc=\"Capturing CUDA graphs ({}, {})\".format(\n                    \"prefill\" if is_capturing_prefill else \"decode\",\n                    cudagraph_runtime_mode.name,\n                ),\n            )\n            if (self.speculative_config is not None\n                and self.speculative_config.num_speculative_tokens > 0\n            ):\n                compilation_cases = tqdm(\n                    compilation_cases,\n                    disable=not self.load_config.use_tqdm_on_load,\n                    desc=\"Capturing CUDA draft graphs ({}, {})\".format(\n                        \"decode\",\n                        cudagraph_runtime_mode.name,\n                    ),\n                )\n\n        # We skip EPLB here since we don't want to record dummy metrics\n        for num_tokens, activate_lora in compilation_cases:\n            # We currently only capture ubatched graphs when its a FULL\n            # cudagraph, a uniform decode batch, and the number of tokens\n            # is above the threshold. Otherwise we just capture a non-ubatched\n            # version of the graph\n            allow_microbatching = (\n                self.parallel_config.enable_dbo\n                and cudagraph_runtime_mode == CUDAGraphMode.FULL\n                and uniform_decode\n                and check_ubatch_thresholds(\n                    config=self.vllm_config.parallel_config,\n                    num_tokens=num_tokens,\n                    uniform_decode=uniform_decode,\n                )\n            )\n\n            for _ in range(self.compilation_config.cudagraph_num_of_warmups):\n                # Use CUDAGraphRuntimeStyle.NONE (default) for warmup.\n                # But be careful, warm up with `NONE`is orthogonal to\n                # if we want to warm up attention or not. This is\n                # different from the case where `FULL` implies capture\n                # attention while `PIECEWISE` implies no attention.\n                force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL\n                self._dummy_run(\n                    num_tokens=num_tokens,\n                    is_capturing_prefill=is_capturing_prefill,\n                    prefill_batch_size=prefill_batch_size,\n                    prefill_seq_len=prefill_seq_len,\n                    cudagraph_runtime_mode=CUDAGraphMode.NONE,\n                    force_attention=force_attention,\n                    uniform_decode=uniform_decode,\n                    allow_microbatching=allow_microbatching,\n                    skip_eplb=True,\n                    remove_lora=False,\n                    activate_lora=activate_lora,\n                )\n            self._dummy_run(\n                num_tokens=num_tokens,\n                is_capturing_prefill=is_capturing_prefill,\n                prefill_batch_size=prefill_batch_size,\n                prefill_seq_len=prefill_seq_len,\n                cudagraph_runtime_mode=cudagraph_runtime_mode,\n                uniform_decode=uniform_decode,\n                allow_microbatching=allow_microbatching,\n                skip_eplb=True,\n                remove_lora=False,\n                activate_lora=activate_lora,\n            )\n        self.maybe_remove_all_loras(self.lora_config)\n\n    def capture_model(self) -> int:\n        if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE:\n            logger.warning(\n                \"Skipping CUDA graph capture. To turn on CUDA graph capture, \"\n                \"ensure `cudagraph_mode` was not manually set to `NONE`\"\n            )\n            return 0\n\n        compilation_counter.num_gpu_runner_capture_triggers += 1\n\n        start_time = time.perf_counter()\n\n        @contextmanager\n        def freeze_gc():\n            # Optimize garbage collection during CUDA graph capture.\n            # Clean up, then freeze all remaining objects from being included\n            # in future collections.\n            gc.collect()\n            should_freeze = not envs.VLLM_ENABLE_CUDAGRAPH_GC\n            if should_freeze:\n                gc.freeze()\n            try:\n                yield\n            finally:\n                if should_freeze:\n                    gc.unfreeze()\n                    gc.collect()\n\n        # Trigger CUDA graph capture for specific shapes.\n        # Capture the large shapes first so that the smaller shapes\n        # can reuse the memory pool allocated for the large shapes.\n        set_cudagraph_capturing_enabled(True)\n        with freeze_gc(), mlu_graph_capture(device=self.device):\n            start_free_gpu_memory = torch.mlu.mem_get_info()[0]\n            cudagraph_mode = self.compilation_config.cudagraph_mode\n            assert cudagraph_mode is not None\n\n            if self.lora_config:\n                if self.compilation_config.cudagraph_specialize_lora:\n                    lora_cases = [True, False]\n                else:\n                    lora_cases = [True]\n            else:\n                lora_cases = [False]\n\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: prefill graph capture\n            '''\n            if self.prefill_enable_mlugraph:\n                # capture prefill mlugraph\n                batch_size = self.prefill_mlugraph_batch_size\n                seq_len = self.prefill_mlugraph_seq_len\n                num_tokens = batch_size * seq_len\n                assert num_tokens <= self.scheduler_config.max_num_batched_tokens\n                assert batch_size <= self.scheduler_config.max_num_seqs\n                logger.info(\"Capture prefill mlugraph for batch size \"\n                                f\"{batch_size} and seq len {seq_len}\")\n                prefill_compilation_cases = list(\n                    product([num_tokens], lora_cases)\n                )\n                self._capture_cudagraphs(\n                    is_capturing_prefill=True,\n                    prefill_batch_size=batch_size,\n                    prefill_seq_len=seq_len,\n                    compilation_cases=prefill_compilation_cases,\n                    cudagraph_runtime_mode=CUDAGraphMode.FULL,\n                    uniform_decode=False,\n                )\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n\n            if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:\n                cudagraph_runtime_mode = cudagraph_mode.mixed_mode()\n                # make sure we capture the largest batch size first\n                compilation_cases = list(\n                    product(reversed(self.cudagraph_batch_sizes), lora_cases)\n                )\n                self._capture_cudagraphs(\n                    compilation_cases=compilation_cases,\n                    cudagraph_runtime_mode=cudagraph_runtime_mode,\n                    uniform_decode=False,\n                )\n\n            # Capture full cudagraph for uniform decode batches if we\n            # don't already have full mixed prefill-decode cudagraphs.\n            if (\n                cudagraph_mode.decode_mode() == CUDAGraphMode.FULL\n                and hasattr(self, 'cudagraph_batch_sizes')\n                and cudagraph_mode.separate_routine()\n            ):\n                max_num_tokens = (\n                    self.scheduler_config.max_num_seqs * self.uniform_decode_query_len\n                )\n                decode_cudagraph_batch_sizes = [\n                    x\n                    for x in self.cudagraph_batch_sizes\n                    if max_num_tokens >= x >= self.uniform_decode_query_len\n                ]\n                compilation_cases_decode = list(\n                    product(reversed(decode_cudagraph_batch_sizes), lora_cases)\n                )\n                self._capture_cudagraphs(\n                    compilation_cases=compilation_cases_decode,\n                    cudagraph_runtime_mode=CUDAGraphMode.FULL,\n                    uniform_decode=True,\n                )\n\n            torch.mlu.synchronize()\n            end_free_gpu_memory = torch.mlu.mem_get_info()[0]\n\n        # Disable cudagraph capturing globally, so any unexpected cudagraph\n        # capturing will be detected and raise an error after here.\n        # Note: We don't put it into graph_capture context manager because\n        # we may do lazy capturing in future that still allows capturing\n        # after here.\n        set_cudagraph_capturing_enabled(False)\n\n        end_time = time.perf_counter()\n        elapsed_time = end_time - start_time\n        cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory\n        # This usually takes 5~20 seconds.\n        logger.info_once(\n            \"Graph capturing finished in %.0f secs, took %.2f GiB\",\n            elapsed_time,\n            cuda_graph_size / (1 << 30),\n            scope=\"local\",\n        )\n        return cuda_graph_size\n\n    def _allocate_kv_cache_tensors(\n        self,\n        kv_cache_config: KVCacheConfig\n    ) -> dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:\n        \"\"\"\n        Initializes the KV cache buffer with the correct size. The buffer needs\n        to be reshaped to the desired shape before being used by the models.\n\n        Args:\n            kv_cache_config: The KV cache config\n        Returns:\n            dict[str, torch.Tensor]: A map between layer names to their\n            corresponding memory buffer for KV cache.\n        \"\"\"\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: support qwen3-next, deepseek v3.2 indexer cache and mlu kv8\n        '''\n\n        kv_cache_group = kv_cache_config.kv_cache_groups[0]\n\n        if self.mlu_config.enable_mamba_split_page_size:\n            # hybrid attention, try to find full attention\n            for group in kv_cache_config.kv_cache_groups:\n                if isinstance(group.kv_cache_spec, FullAttentionSpec):\n                    kv_cache_group = group\n                    break\n            self.mamba_block_num = self.mlu_config.mamba_support_max_batch_size\n            self.mamba_tensor_size = (kv_cache_group.kv_cache_spec.page_size_bytes \\\n                * self.mlu_config.mamba_to_attn_block_ratio * self.mamba_block_num)\n            logger.info(f\"one linear attn layer cache tensor size {self.mamba_tensor_size}\")\n        \n        kv_cache_raw_tensors: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}\n\n        for kv_cache_tensor in kv_cache_config.kv_cache_tensors:\n            kv_cache_spec = kv_cache_group.kv_cache_spec\n            assert kv_cache_tensor.size % kv_cache_spec.page_size_bytes == 0\n            num_blocks = kv_cache_tensor.size // kv_cache_spec.page_size_bytes\n            if kv_cache_spec.dtype in [torch.int8, torch.uint8]:\n                # mlu kv8\n                assert isinstance(kv_cache_spec, AttentionSpec)\n                cache_ = torch.zeros(\n                    num_blocks * kv_cache_spec.cache_size_bytes,\n                    dtype=torch.int8, device=self.device,\n                )\n                scale_ = torch.zeros(\n                    num_blocks * kv_cache_spec.scale_size_bytes,\n                    dtype=torch.int8, device=self.device,\n                )\n            else:\n                # not mlu kv8\n                cache_ = torch.zeros(\n                    num_blocks * kv_cache_spec.cache_size_bytes,\n                    dtype=torch.int8,\n                    device=self.device\n                )\n                scale_ = torch.tensor([], dtype=torch.int8, device=self.device)\n\n            if (isinstance(kv_cache_spec, MLUMLAAttentionSpec)\n                and kv_cache_spec.index_n_heads > 0):\n                index_cache_ = torch.zeros((num_blocks *\n                                            kv_cache_spec.index_cache_size_bytes),\n                                            dtype=torch.int8,\n                                            device=self.device)\n            else:\n                index_cache_ = torch.tensor([], dtype=torch.int8, device=self.device)\n\n            for layer_name in kv_cache_tensor.shared_by:\n                '''\n                =============================\n                Modify by vllm_mlu\n                =============================\n                @brief: support qwen3-next\n                '''\n                if self.mlu_config.enable_mamba_split_page_size:\n                    if 'linear_attn' in layer_name:\n                        mamba_tensor = torch.zeros(\n                            self.mamba_tensor_size, dtype=torch.int8, device=self.device\n                        )\n                        kv_cache_raw_tensors[layer_name] = [mamba_tensor, scale_, index_cache_]\n                    else:\n                        kv_cache_raw_tensors[layer_name] = [cache_, scale_, index_cache_]\n                else:\n                    kv_cache_raw_tensors[layer_name] = [cache_, scale_, index_cache_]\n                '''\n                ==================\n                End of MLU Hijack\n                ==================\n                '''\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        layer_names = set()\n        for group in kv_cache_config.kv_cache_groups:\n            for layer_name in group.layer_names:\n                if layer_name in self.runner_only_attn_layers:\n                    continue\n                layer_names.add(layer_name)\n        assert layer_names == set(kv_cache_raw_tensors.keys()), (\n            \"Some layers are not correctly initialized\"\n        )\n        return kv_cache_raw_tensors\n\n    def _reshape_kv_cache_tensors(\n        self,\n        kv_cache_config: KVCacheConfig,\n        kv_cache_raw_tensors: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]],\n        kernel_block_sizes: list[int],\n    ) -> dict[str, torch.Tensor]:\n        \"\"\"\n        Reshape the KV cache tensors to the desired shape and dtype.\n\n        Args:\n            kv_cache_config: The KV cache config\n            kv_cache_raw_tensors: The KV cache buffer of each layer, with\n            correct size but uninitialized shape.\n            kernel_block_sizes: The kernel block sizes for each KV cache group.\n        Returns:\n            Dict[str, torch.Tensor]: A map between layer names to their\n            corresponding memory buffer for KV cache.\n        \"\"\"\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: support mlu kv8 and deepseek v3.2 indexer\n        '''\n        kv_caches: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}\n        has_attn, has_mamba = False, False\n        for group in self._kv_cache_spec_attn_group_iterator():\n            kv_cache_spec = group.kv_cache_spec\n            attn_backend = group.backend\n            if group.kv_cache_group_id == len(kernel_block_sizes):\n                # There may be a last group for layers without kv cache.\n                continue\n            kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]\n            for layer_name in group.layer_names:\n                if layer_name in self.runner_only_attn_layers:\n                    continue\n                raw_tensor = kv_cache_raw_tensors[layer_name]\n                cache_, scale_, index_cache_ = raw_tensor\n                total_numel = cache_.numel() + scale_.numel() + index_cache_.numel()\n                assert total_numel % kv_cache_spec.page_size_bytes == 0\n                num_blocks = total_numel // kv_cache_spec.page_size_bytes\n                if isinstance(kv_cache_spec, AttentionSpec):\n                    has_attn = True\n                    num_blocks_per_kv_block = (\n                        kv_cache_spec.block_size // kernel_block_size\n                    )\n                    kernel_num_blocks = num_blocks * num_blocks_per_kv_block\n\n                    kv_cache_shape = attn_backend.get_kv_cache_shape(\n                        kernel_num_blocks,\n                        kernel_block_size,\n                        kv_cache_spec.num_kv_heads,\n                        kv_cache_spec.head_size,\n                        cache_dtype_str=self.cache_config.cache_dtype,\n                    )\n                    dtype = kv_cache_spec.dtype\n                    try:\n                        kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()\n                        assert len(kv_cache_stride_order) == len(kv_cache_shape)\n                    except (AttributeError, NotImplementedError):\n                        kv_cache_stride_order = tuple(range(len(kv_cache_shape)))\n                    # The allocation respects the backend-defined stride order\n                    # to ensure the semantic remains consistent for each\n                    # backend. We first obtain the generic kv cache shape and\n                    # then permute it according to the stride order which could\n                    # result in a non-contiguous tensor.\n                    kv_cache_shape = tuple(\n                        kv_cache_shape[i] for i in kv_cache_stride_order\n                    )\n                    # Maintain original KV shape view.\n                    inv_order = [\n                        kv_cache_stride_order.index(i)\n                        for i in range(len(kv_cache_stride_order))\n                    ]\n                    cache_ = (\n                        cache_\n                        .view(dtype)\n                        .view(kv_cache_shape)\n                        .permute(*inv_order)\n                    )\n                    # Reshape kv cache scale tensor\n                    if dtype in [torch.int8, torch.uint8]:\n                        kv_cache_scale_shape = attn_backend.get_kv_cache_scale_shape(\n                            kernel_num_blocks,\n                            kernel_block_size,\n                            kv_cache_spec.num_kv_heads,\n                        )\n                        scale_ = (\n                            scale_\n                            .view(torch.float32)\n                            .view(kv_cache_scale_shape)\n                        )\n                    # Reshape index_cache\n                    if (isinstance(kv_cache_spec, MLUMLAAttentionSpec)\n                        and kv_cache_spec.index_n_heads > 0):\n                        index_cache_shape = (\n                            kernel_num_blocks,\n                            kv_cache_spec.index_n_heads,\n                            kernel_block_size,\n                            kv_cache_spec.index_head_dim,\n                        )\n                        index_cache_ = index_cache_.view(dtype).view(index_cache_shape)\n                    kv_caches[layer_name] = [cache_, scale_, index_cache_]\n                elif isinstance(kv_cache_spec, MambaSpec):\n                    has_mamba = True\n                    cache_ = kv_cache_raw_tensors[layer_name]\n                    raw_tensor, scale_, index_cache_ = cache_\n                    state_tensors = []\n                    storage_offset_bytes = 0\n                    for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):\n                        dtype_size = get_dtype_size(dtype)\n                        num_element_per_page = (\n                            kv_cache_spec.page_size_bytes // dtype_size\n                        )\n                        '''\n                        =============================\n                        Modify by vllm_mlu\n                        =============================\n                        @brief: support qwen3-next\n                        '''\n                        if self.mlu_config.enable_mamba_split_page_size:\n                            num_element_per_page *= self.mlu_config.mamba_to_attn_block_ratio\n                            num_blocks = self.mamba_block_num\n                        '''\n                        ==================\n                        End of MLU Hijack\n                        ==================\n                        '''\n                        target_shape = (num_blocks, *shape)\n                        stride = torch.empty(target_shape).stride()\n                        target_stride = (num_element_per_page, *stride[1:])\n                        assert storage_offset_bytes % dtype_size == 0\n                        tensor = torch.as_strided(\n                            raw_tensor.view(dtype),\n                            size=target_shape,\n                            stride=target_stride,\n                            storage_offset=storage_offset_bytes // dtype_size,\n                        )\n                        state_tensors.append(tensor)\n                        storage_offset_bytes += stride[0] * dtype_size\n\n                    kv_caches[layer_name] = state_tensors\n                else:\n                    raise NotImplementedError\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: support qwen3-next\n        '''\n        if has_attn and has_mamba and not self.mlu_config.enable_mamba_split_page_size:\n            self._update_hybrid_attention_mamba_layout(kv_caches)\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        return kv_caches\n\n    def initialize_kv_cache_tensors(\n        self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]\n    ) -> dict[str, torch.Tensor]:\n        \"\"\"\n        Initialize the memory buffer for KV cache.\n\n        Args:\n            kv_cache_config: The KV cache config\n            kernel_block_sizes: The kernel block sizes for each KV cache group.\n\n        Returns:\n            Dict[str, torch.Tensor]: A map between layer names to their\n            corresponding memory buffer for KV cache.\n        \"\"\"\n        # Initialize the memory buffer for KV cache\n        kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)\n        # Change the memory buffer to the desired shape\n        kv_caches = self._reshape_kv_cache_tensors(\n            kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes\n        )\n\n        if self.speculative_config and self.speculative_config.use_eagle():\n            assert isinstance(self.drafter, EagleProposer)\n            # validate all draft model layers belong to the same kv cache\n            # group\n            self.drafter.validate_same_kv_cache_group(kv_cache_config)\n\n        # Set up cross-layer KV cache sharing\n        for layer_name, target_layer_name in self.shared_kv_cache_layers.items():\n            logger.debug(\"%s reuses KV cache of %s\", layer_name, target_layer_name)\n            kv_caches[layer_name] = kv_caches[target_layer_name]\n\n        num_attn_module = (\n            2 if self.model_config.hf_config.model_type == \"longcat_flash\" else 1\n        )\n        bind_kv_cache(\n            kv_caches,\n            self.compilation_config.static_forward_context,\n            self.kv_caches,\n            num_attn_module,\n        )\n\n        '''\n        =============================\n        Modify by vllm_mlu\n        =============================\n        @brief: bind kv cache to deepseek prefill attn\n        '''\n        if self.model_config.is_deepseek_mla:\n            forward_context = self.vllm_config.compilation_config.static_forward_context\n            for layer_name, kv_cache in kv_caches.items():\n                if layer_name.endswith(\"self_attn.mla_attn\"):\n                    layer_name = layer_name.replace(\n                        \"self_attn.mla_attn\", \"self_attn.attn\")\n                forward_context[layer_name].kv_cache = [kv_cache]\n\n        # matches self_attn.0.attn or self_attn.1.attn\n        if self.model_config.is_longcat_flash:\n            forward_context = self.vllm_config.compilation_config.static_forward_context\n            for layer_name, kv_cache in kv_caches.items():\n                if (match := re.match(r\".*self_attn\\.(0|1)\\.mla_attn$\", layer_name)):\n                    digit = match.group(1)  # Extract the captured digit (0 or 1)\n                    layer_name = layer_name.replace(\n                        f\"self_attn.{digit}.mla_attn\",\n                        f\"self_attn.{digit}.attn\"\n                    )\n                forward_context[layer_name].kv_cache = [kv_cache]\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        return kv_caches\n\n    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:\n        \"\"\"\n        Generates the KVCacheSpec by parsing the kv cache format from each\n        Attention module in the static forward context.\n        Returns:\n            KVCacheSpec: A dictionary mapping layer names to their KV cache\n            format. Layers that do not need KV cache are not included.\n        \"\"\"\n\n        # block_size = self.vllm_config.cache_config.block_size\n        # use_mla = self.vllm_config.model_config.use_mla\n        kv_cache_spec: dict[str, KVCacheSpec] = {}\n        attn_layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)\n        for layer_name, attn_module in attn_layers.items():\n            if isinstance(attn_module, Attention) and (\n                kv_tgt_layer := attn_module.kv_sharing_target_layer_name\n            ):\n                # The layer doesn't need its own KV cache and will use that of\n                # the target layer. We skip creating a KVCacheSpec for it, so\n                # that KV cache management logic will act as this layer does\n                # not exist, and doesn't allocate KV cache for the layer. This\n                # enables the memory saving of cross-layer kv sharing, allowing\n                # a given amount of memory to accommodate longer context lengths\n                # or enable more requests to be processed simultaneously.\n                self.shared_kv_cache_layers[layer_name] = kv_tgt_layer\n                continue\n            '''\n            =============================\n            Modify by vllm_mlu\n            =============================\n            @brief: skip deepseek prefill attn init kv_cache\n            '''\n            if (\n                self.model_config.is_deepseek_mla\n                and layer_name.endswith(\"self_attn.attn\")\n            ):\n                continue\n            # matches self_attn.0.attn or self_attn.1.attn\n            if (\n                self.model_config.is_longcat_flash\n                and re.match(r\".*self_attn\\.(0|1)\\.attn$\", layer_name)\n            ):\n                continue\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n            # Skip modules that don't need KV cache (eg encoder-only attention)\n            if spec := attn_module.get_kv_cache_spec(self.vllm_config):\n                kv_cache_spec[layer_name] = spec\n\n        return kv_cache_spec\n\n    def reset_capture_context(self,\n                              prefill_enable_mlugraph: bool,\n                              batch_size: int,\n                              input_len: int):\n        self.graph_runners = {}\n        self.context_graph_runner = None\n        self.graph_memory_pool = None\n\n        # reset prefill mlugraph infos\n        self.prefill_enable_mlugraph = prefill_enable_mlugraph\n        self.prefill_mlugraph_batch_size = batch_size\n        self.prefill_mlugraph_seq_len = input_len\n\n        gc.collect()\n        torch.mlu.empty_cache()\n\n    def _copy_valid_sampled_token_count(\n        self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor\n    ) -> None:\n        if self.valid_sampled_token_count_event is None:\n            return\n\n        '''\n        =============================\n        Modify by vllm_mlu\n        @brief:  replace current stream for MLU device.\n        =======\n        '''\n        default_stream = torch.mlu.current_stream()\n        # Initialize a new stream to overlap the copy operation with\n        # prepare_input of draft model.\n        with torch.mlu.stream(self.valid_sampled_token_count_copy_stream):\n            self.valid_sampled_token_count_copy_stream.wait_stream(default_stream)  # type: ignore\n            counts = valid_sampled_tokens_count\n            counts_cpu = self.valid_sampled_token_count_cpu\n            counts_cpu[: counts.shape[0]].copy_(counts, non_blocking=True)\n            self.valid_sampled_token_count_event.record()\n        '''\n        ==================\n        End of MLU Hijack\n        ==================\n        '''\n\n        self.input_batch.prev_sampled_token_ids = next_token_ids.unsqueeze(1)\n\n\n    def _bookkeeping_sync(\n        self,\n        scheduler_output: \"SchedulerOutput\",\n        sampler_output: SamplerOutput,\n        logits: torch.Tensor | None,\n        hidden_states: torch.Tensor,\n        num_scheduled_tokens: int,\n        spec_decode_metadata: SpecDecodeMetadata | None,\n    ) -> tuple[\n        dict[str, int],\n        LogprobsLists | None,\n        list[np.ndarray],\n        dict[str, LogprobsTensors | None],\n        list[str],\n        dict[str, int],\n        list[int],\n    ]:\n        num_nans_in_logits = {}\n        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:\n            num_nans_in_logits = self._get_nans_in_logits(logits)\n\n        discard_sampled_tokens_req_indices = self.discard_request_indices.np[\n            : self.num_discarded_requests\n        ]\n        for i in discard_sampled_tokens_req_indices:\n            gen = self.input_batch.generators.get(int(i))\n            if gen is not None:\n                gen.set_offset(gen.get_offset() - 4)\n        # Copy some objects so they don't get modified after returning.\n        # This is important when using async scheduling.\n        req_ids_output_copy = self.input_batch.req_ids.copy()\n        req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()\n\n        num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]\n        sampled_token_ids = sampler_output.sampled_token_ids\n        invalid_req_indices = []\n        valid_sampled_token_ids: list[np.ndarray]\n        if not self.use_async_scheduling:\n            # Get the valid generated tokens.\n            max_gen_len = sampled_token_ids.shape[-1]\n            if max_gen_len == 1:\n                # No spec decode tokens.\n                valid_sampled_token_ids = self._to_list(sampled_token_ids)\n            else:\n                # Includes spec decode tokens.\n                valid_sampled_token_ids = self.rejection_sampler.parse_output(\n                    sampled_token_ids,\n                    self.input_batch.vocab_size,\n                )\n            # Mask out the sampled tokens that should not be sampled.\n            for i in discard_sampled_tokens_req_indices:\n                valid_sampled_token_ids[int(i)] = np.array([])\n        else:\n            valid_sampled_token_ids = []\n            invalid_req_indices = discard_sampled_tokens_req_indices.tolist()\n            invalid_req_indices_set = set(invalid_req_indices)\n\n            # Cache the sampled tokens on the GPU and avoid CPU sync.\n            # These will be copied into input_ids in the next step\n            # when preparing inputs.\n            # With spec decoding, this is done in propose_draft_token_ids().\n            if self.input_batch.prev_sampled_token_ids is None:\n                assert sampled_token_ids.shape[-1] == 1\n                self.input_batch.prev_sampled_token_ids = sampled_token_ids\n            self.input_batch.prev_req_id_to_index = {\n                req_id: i\n                for i, req_id in enumerate(self.input_batch.req_ids)\n                if i not in invalid_req_indices_set\n            }\n\n        # Cache the sampled tokens in the model runner, so that the scheduler\n        # doesn't need to send them back.\n        # NOTE(woosuk): As an exception, when using PP, the scheduler sends\n        # the sampled tokens back, because there's no direct communication\n        # between the first-stage worker and the last-stage worker.\n        req_ids = self.input_batch.req_ids\n        logprobs_tensors = sampler_output.logprobs_tensors\n        cu_num_accepted_tokens = (\n            [0] if spec_decode_metadata and logprobs_tensors else None\n        )\n        for req_idx in range(num_sampled_tokens):\n            sampled_ids: np.ndarray | None\n            if self.use_async_scheduling:\n                sampled_ids = (\n                    np.array([-1]) if req_idx not in invalid_req_indices_set else None\n                )\n            else:\n                sampled_ids = valid_sampled_token_ids[req_idx]\n\n            num_sampled_ids: int = (\n                sampled_ids.shape[0] if sampled_ids is not None else 0\n            )\n\n            if cu_num_accepted_tokens is not None:\n                cu_num_accepted_tokens.append(\n                    cu_num_accepted_tokens[-1] + num_sampled_ids\n                )\n\n            if sampled_ids is None or num_sampled_ids == 0:\n                continue\n\n            start_idx = self.input_batch.num_tokens_no_spec[req_idx]\n            end_idx = start_idx + num_sampled_ids\n            '''\n            =============================\n            Modify by vllm_mlu\n            @brief:  end_idx may exceed max_model_len for sepculative tokens in MTP mode.\n            =======\n            '''\n            num_async_sched_tokens = 1 if self.use_async_scheduling else 0\n            max_model_len = self.num_spec_tokens + self.max_model_len + num_async_sched_tokens\n            assert end_idx <= max_model_len, (\n                \"Sampled token IDs exceed the max model length. \"\n                f\"Total number of tokens: {end_idx} > max_model_len: \"\n                f\"{max_model_len}\"\n            )\n            if end_idx > self.max_model_len:\n                end_idx = self.max_model_len\n                sampled_ids = sampled_ids[:end_idx - start_idx]\n            '''\n            ==================\n            End of MLU Hijack\n            ==================\n            '''\n\n            self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids\n            self.input_batch.is_token_ids[req_idx, start_idx:end_idx] = True\n            self.input_batch.num_tokens_no_spec[req_idx] = end_idx\n            self.input_batch.num_tokens[req_idx] = end_idx\n\n            req_id = req_ids[req_idx]\n            req_state = self.requests[req_id]\n            req_state.output_token_ids.extend(sampled_ids)\n\n        logprobs_lists = (\n            logprobs_tensors.tolists(cu_num_accepted_tokens)\n            if not self.use_async_scheduling and logprobs_tensors is not None\n            else None\n        )\n\n        # Compute prompt logprobs if needed.\n        prompt_logprobs_dict = self._get_prompt_logprobs_dict(\n            hidden_states[:num_scheduled_tokens],\n            scheduler_output.num_scheduled_tokens,\n        )\n\n        return (\n            num_nans_in_logits,\n            logprobs_lists,\n            valid_sampled_token_ids,\n            prompt_logprobs_dict,\n            req_ids_output_copy,\n            req_id_to_index_output_copy,\n            invalid_req_indices,\n        )\n\n"}}}}