diff --git a/pyproject.toml b/pyproject.toml
deleted file mode 100644
index 07616c858..000000000
--- a/pyproject.toml
+++ /dev/null
@@ -1,173 +0,0 @@
-[build-system]
-# Should be mirrored in requirements/build.txt
-requires = [
-    "cmake>=3.26",
-    "ninja",
-    "packaging",
-    "setuptools>=61",
-    "setuptools-scm>=8.0",
-    "torch == 2.6.0",
-    "wheel",
-    "jinja2",
-]
-build-backend = "setuptools.build_meta"
-
-[project]
-name = "vllm"
-authors = [{name = "vLLM Team"}]
-license = { "file"= "LICENSE" }
-readme = "README.md"
-description = "A high-throughput and memory-efficient inference and serving engine for LLMs"
-classifiers = [
-    "Programming Language :: Python :: 3.9",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.11",
-    "Programming Language :: Python :: 3.12",
-    "License :: OSI Approved :: Apache Software License",
-    "Intended Audience :: Developers",
-    "Intended Audience :: Information Technology",
-    "Intended Audience :: Science/Research",
-    "Topic :: Scientific/Engineering :: Artificial Intelligence",
-    "Topic :: Scientific/Engineering :: Information Analysis",
-]
-requires-python = ">=3.9"
-dynamic = [ "version", "dependencies", "optional-dependencies"]
-
-[project.urls]
-Homepage="https://github.com/vllm-project/vllm"
-Documentation="https://vllm.readthedocs.io/en/latest/"
-Slack="http://slack.vllm.ai/"
-
-[project.scripts]
-vllm = "vllm.entrypoints.cli.main:main"
-
-[tool.setuptools_scm]
-# no extra settings needed, presence enables setuptools-scm
-
-[tool.setuptools.packages.find]
-where = ["."]
-exclude = ["benchmarks", "csrc", "docs", "examples", "tests*"]
-namespaces = false
-
-[tool.yapfignore]
-ignore_patterns = [
-    "build/**",
-]
-
-[tool.ruff]
-# Allow lines to be as long as 80.
-line-length = 80
-exclude = [
-    # External file, leaving license intact
-    "examples/other/fp8/quantizer/quantize.py"
-]
-
-[tool.ruff.lint.per-file-ignores]
-"vllm/third_party/**" = ["ALL"]
-"vllm/version.py" = ["F401"]
-"vllm/_version.py" = ["ALL"]
-# Python 3.8 typing. TODO: Remove these excludes after v1.0.0
-"vllm/adapter_commons/**/*.py" = ["UP006", "UP035"]
-"vllm/attention/**/*.py" = ["UP006", "UP035"]
-"vllm/compilation/**/*.py" = ["UP006", "UP035"]
-"vllm/core/**/*.py" = ["UP006", "UP035"]
-"vllm/device_allocator/**/*.py" = ["UP006", "UP035"]
-"vllm/distributed/**/*.py" = ["UP006", "UP035"]
-"vllm/engine/**/*.py" = ["UP006", "UP035"]
-"vllm/executor/**/*.py" = ["UP006", "UP035"]
-"vllm/lora/**/*.py" = ["UP006", "UP035"]
-"vllm/model_executor/**/*.py" = ["UP006", "UP035"]
-"vllm/platforms/**/*.py" = ["UP006", "UP035"]
-"vllm/plugins/**/*.py" = ["UP006", "UP035"]
-"vllm/profiler/**/*.py" = ["UP006", "UP035"]
-"vllm/prompt_adapter/**/*.py" = ["UP006", "UP035"]
-"vllm/spec_decode/**/*.py" = ["UP006", "UP035"]
-"vllm/transformers_utils/**/*.py" = ["UP006", "UP035"]
-"vllm/triton_utils/**/*.py" = ["UP006", "UP035"]
-"vllm/vllm_flash_attn/**/*.py" = ["UP006", "UP035"]
-"vllm/worker/**/*.py" = ["UP006", "UP035"]
-"vllm/utils.py" = ["UP006", "UP035"]
-
-[tool.ruff.lint]
-select = [
-    # pycodestyle
-    "E",
-    # Pyflakes
-    "F",
-    # pyupgrade
-    "UP",
-    # flake8-bugbear
-    "B",
-    # flake8-simplify
-    "SIM",
-    # isort
-    # "I",
-    "G",
-]
-ignore = [
-    # star imports
-    "F405", "F403",
-    # lambda expression assignment
-    "E731",
-    # Loop control variable not used within loop body
-    "B007",
-    # f-string format
-    "UP032",
-    # Can remove once 3.10+ is the minimum Python version
-    "UP007",
-]
-
-[tool.mypy]
-ignore_missing_imports = true
-check_untyped_defs = true
-follow_imports = "silent"
-
-# After fixing type errors resulting from follow_imports: "skip" -> "silent",
-# move the directory here and remove it from tools/mypy.sh
-files = [
-    "vllm/*.py",
-    "vllm/adapter_commons",
-    "vllm/assets",
-    "vllm/entrypoints",
-    "vllm/core",
-    "vllm/inputs",
-    "vllm/logging_utils",
-    "vllm/multimodal",
-    "vllm/platforms",
-    "vllm/transformers_utils",
-    "vllm/triton_utils",
-    "vllm/usage",
-]
-# TODO(woosuk): Include the code from Megatron and HuggingFace.
-exclude = [
-    "vllm/model_executor/parallel_utils/|vllm/model_executor/models/",
-    # Ignore triton kernels in ops.
-    'vllm/attention/ops/.*\.py$'
-]
-
-[tool.codespell]
-ignore-words-list = "dout, te, indicies, subtile, ElementE"
-skip = "tests/models/fixtures/*,tests/prompts/*,benchmarks/sonnet.txt,tests/lora/data/*,build/*,vllm/third_party/*"
-
-[tool.isort]
-use_parentheses = true
-skip_gitignore = true
-
-[tool.pytest.ini_options]
-markers = [
-    "skip_global_cleanup",
-    "core_model: enable this model test in each PR instead of only nightly",
-    "cpu_model: enable this model test in CPU tests",
-    "quant_model: run this model test under Quantized category",
-    "split: run this test as part of a split",
-    "distributed: run this test only in distributed GPU tests",
-    "skip_v1: do not run this test with v1",
-    "optional: optional tests that are automatically skipped, include --optional to run them",
-]
-
-[tool.pymarkdown]
-plugins.md004.style = "sublist" # ul-style
-plugins.md013.enabled = false # line-length
-plugins.md041.enabled = false # first-line-h1
-plugins.md033.enabled = false # inline-html
-plugins.md024.allow_different_nesting = true # no-duplicate-headers
diff --git a/setup.py b/setup.py
index b0cc2f481..2a2728e83 100755
--- a/setup.py
+++ b/setup.py
@@ -13,11 +13,26 @@ from shutil import which
 
 import torch
 from packaging.version import Version, parse
-from setuptools import Extension, setup
+from setuptools import Extension, setup, find_packages
 from setuptools.command.build_ext import build_ext
 from setuptools_scm import get_version
 from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
+from setuptools.command.install import install
 
+class InstallWithMlu(install):
+    def run(self):
+        # 执行原始的install命令安装vllm
+        super().run()
+
+        # 安装vllm_mlu
+        mlu_path = os.path.join(ROOT_DIR, "vllm_mlu")
+        if os.path.exists(mlu_path):
+            print("Installing vllm_mlu after vllm...")
+            subprocess.check_call([
+                sys.executable, "-m", "pip", "install", "."
+            ], cwd=mlu_path)
+        else:
+            print(f"vllm_mlu directory not found at {mlu_path}")
 
 def load_module_from_path(module_name, path):
     spec = importlib.util.spec_from_file_location(module_name, path)
@@ -30,6 +45,8 @@ def load_module_from_path(module_name, path):
 ROOT_DIR = Path(__file__).parent
 logger = logging.getLogger(__name__)
 
+os.environ['VLLM_TARGET_DEVICE'] = 'empty'
+
 # cannot import envs directly because it depends on vllm,
 #  which is not installed yet
 envs = load_module_from_path('envs', os.path.join(ROOT_DIR, 'vllm', 'envs.py'))
@@ -538,7 +555,8 @@ def get_gaudi_sw_version():
 
 
 def get_vllm_version() -> str:
-    version = get_version(write_to="vllm/_version.py")
+    # version = get_version(write_to="vllm/_version.py")
+    version = "0.8.3"
     sep = "+" if "+" not in version else "."  # dev versions might contain +
 
     if _no_device():
@@ -655,11 +673,13 @@ if _is_cuda():
 if _build_custom_ops():
     ext_modules.append(CMakeExtension(name="vllm._C"))
 
+packages = find_packages(include=['vllm', 'vllm.*'])
 package_data = {
     "vllm": [
         "py.typed",
         "model_executor/layers/fused_moe/configs/*.json",
         "model_executor/layers/quantization/utils/configs/*.json",
+        "**/*",
     ]
 }
 
@@ -667,7 +687,7 @@ if _no_device():
     ext_modules = []
 
 if not ext_modules:
-    cmdclass = {}
+    cmdclass = {"install": InstallWithMlu}
 else:
     cmdclass = {
         "build_ext":
@@ -686,6 +706,12 @@ setup(
         "audio": ["librosa", "soundfile"],  # Required for audio processing
         "video": []  # Kept for backwards compatibility
     },
+    entry_points={
+        'console_scripts': [
+            'vllm = vllm.entrypoints.cli.main:main',
+        ],
+    },
     cmdclass=cmdclass,
     package_data=package_data,
+    packages=packages,
 )
diff --git a/tests/spec_decode/test_multi_step_worker.py b/tests/spec_decode/test_multi_step_worker.py
index ca37c9a68..a01af18c6 100644
--- a/tests/spec_decode/test_multi_step_worker.py
+++ b/tests/spec_decode/test_multi_step_worker.py
@@ -749,8 +749,8 @@ def test_draft_proposals_mixed_k():
 
     assert proposals.proposal_lens.shape == torch.Size([batch_size])
     assert proposals.proposal_lens.tolist() == [
-        k for _ in range(expected_num_proposal_seqs - 1)
-    ] + [0 for _ in range(expected_num_no_proposal_seqs)] + [k]
+        k for _ in range(expected_num_proposal_seqs)
+    ] + [0 for _ in range(expected_num_no_proposal_seqs)]
 
 
 @torch.inference_mode()
diff --git a/vllm/compilation/fusion.py b/vllm/compilation/fusion.py
index b46f5f522..7925e8ee4 100644
--- a/vllm/compilation/fusion.py
+++ b/vllm/compilation/fusion.py
@@ -29,8 +29,21 @@ def empty_fp32(*args, **kwargs):
     return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda")
 
 
-RMS_OP = torch.ops._C.rms_norm.default
-RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
+'''
+=============================
+Modify by vllm_mlu
+=============================
+@brief: skip custom op implemention
+'''
+# RMS_OP = torch.ops._C.rms_norm.default
+# RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
+RMS_OP = None
+RMS_ADD_OP = None
+'''
+==================
+End of MLU Hijack
+==================
+'''
 
 
 class QuantKey(NamedTuple):
@@ -57,13 +70,25 @@ kFp8StaticTensorSym = QuantKey(FP8_DTYPE, True, True, True)
 kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, False, True, True)
 kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, False, False, True)
 
+
+'''
+=============================
+Modify by vllm_mlu
+=============================
+@brief: skip fp8 quant op implemention
+'''
 QUANT_OPS: Dict[QuantKey, OpOverload] = {
-    kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default,  # noqa
-    kFp8DynamicTensorSym:
-    torch.ops._C.dynamic_scaled_fp8_quant.default,  # noqa
-    kFp8DynamicTokenSym:
-    torch.ops._C.dynamic_per_token_scaled_fp8_quant.default,  # noqa
+    # kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default,  # noqa
+    # kFp8DynamicTensorSym:
+    # torch.ops._C.dynamic_scaled_fp8_quant.default,  # noqa
+    # kFp8DynamicTokenSym:
+    # torch.ops._C.dynamic_per_token_scaled_fp8_quant.default,  # noqa
 }
+'''
+==================
+End of MLU Hijack
+==================
+'''
 
 
 class FusedRMSQuantKey(NamedTuple):
@@ -80,16 +105,27 @@ class FusedRMSQuantKey(NamedTuple):
                 f"{'' if self.fused_add else 'out'} residual)")
 
 
+'''
+=============================
+Modify by vllm_mlu
+=============================
+@brief: skip fp8 quant op implemention
+'''
 FUSED_OPS: Dict[FusedRMSQuantKey, OpOverload] = {
-    FusedRMSQuantKey(kFp8StaticTensorSym, False):
-    torch.ops._C.rms_norm_static_fp8_quant.default,  # noqa
-    FusedRMSQuantKey(kFp8StaticTensorSym, True):
-    torch.ops._C.fused_add_rms_norm_static_fp8_quant.default,  # noqa
-    FusedRMSQuantKey(kFp8DynamicTokenSym, False):
-    torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa
-    FusedRMSQuantKey(kFp8DynamicTokenSym, True):
-    torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa
+    # FusedRMSQuantKey(kFp8StaticTensorSym, False):
+    # torch.ops._C.rms_norm_static_fp8_quant.default,  # noqa
+    # FusedRMSQuantKey(kFp8StaticTensorSym, True):
+    # torch.ops._C.fused_add_rms_norm_static_fp8_quant.default,  # noqa
+    # FusedRMSQuantKey(kFp8DynamicTokenSym, False):
+    # torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa
+    # FusedRMSQuantKey(kFp8DynamicTokenSym, True):
+    # torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa
 }
+'''
+==================
+End of MLU Hijack
+==================
+'''
 
 
 class QuantMultiOutputMatch(MultiOutputMatch):
diff --git a/vllm/config.py b/vllm/config.py
index bd52fc90b..bc42267ce 100644
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -1542,6 +1542,7 @@ class ParallelConfig:
         if self.data_parallel_size > 1:
             # Data parallel was specified in the engine args.
             self.data_parallel_master_port = get_open_port()
+            self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP
             # TODO multi-node
         else:
             # Otherwise fall back to env vars (e.g. for offline SPMD case).
@@ -1561,7 +1562,7 @@ class ParallelConfig:
         ray_only_devices: list[str] = []
         from vllm.platforms import current_platform
         if (current_platform.device_type in ray_only_devices
-                and self.world_size > 1):
+                and self.world_size_across_dp > 1):
             if self.distributed_executor_backend is None:
                 self.distributed_executor_backend = "ray"
             if self.distributed_executor_backend != "ray":
@@ -1569,18 +1570,20 @@ class ParallelConfig:
                     f"{current_platform.device_type.upper()} backend only "
                     "supports Ray for distributed inference.")
 
-        if self.distributed_executor_backend is None and self.world_size > 1:
+        if self.distributed_executor_backend is None and self.world_size_across_dp > 1:
             # We use multiprocessing by default if world_size fits on the
             # current node and we aren't in a ray placement group.
 
             from vllm.executor import ray_utils
             backend = "mp"
+            if self.data_parallel_size > 1:
+                backend = "ray"
             ray_found = ray_utils.ray_is_available()
             if current_platform.is_neuron():
                 # neuron uses single process to control multiple devices
                 backend = "uni"
             elif (current_platform.is_cuda()
-                  and cuda_device_count_stateless() < self.world_size):
+                  and cuda_device_count_stateless() < self.world_size_across_dp):
                 if not ray_found:
                     raise ValueError("Unable to load Ray which is "
                                      "required for multi-node inference, "
@@ -1600,7 +1603,7 @@ class ParallelConfig:
             logger.info("Defaulting to use %s for distributed inference",
                         backend)
 
-        if self.distributed_executor_backend is None and self.world_size == 1:
+        if self.distributed_executor_backend is None and self.world_size_across_dp == 1:
             self.distributed_executor_backend = "uni"
 
         self._verify_args()
diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
index fa493fefb..47ddef178 100644
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -831,11 +831,12 @@ def init_distributed_environment(
         # adjust the world size to take into account data parallelism
         world_size = parallel_config.world_size_across_dp
         ip = parallel_config.data_parallel_master_ip
-        port = parallel_config.get_next_dp_init_port()
+        # port = parallel_config.get_next_dp_init_port()
+        port = parallel_config.data_parallel_master_port
         distributed_init_method = f"tcp://{ip}:{port}"  # noqa
         logger.info(
-            "Adjusting world_size=%d rank=%d distributed_init_method=%s for DP",
-            world_size, rank, distributed_init_method)
+            "Adjusting world_size=%d rank=%d local_rank=%d distributed_init_method=%s for DP",
+            world_size, rank, local_rank, distributed_init_method)
     if not torch.distributed.is_initialized():
         assert distributed_init_method is not None, (
             "distributed_init_method must be provided when initializing "
@@ -1089,7 +1090,10 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     gc.collect()
     from vllm.platforms import current_platform
     if not current_platform.is_cpu():
-        torch.cuda.empty_cache()
+        if current_platform.is_out_of_tree():
+            current_platform.empty_cache()
+        else:
+            torch.cuda.empty_cache()
     try:
         torch._C._host_emptyCache()
     except AttributeError:
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
index 89c9b6747..3fd7a9798 100644
--- a/vllm/engine/arg_utils.py
+++ b/vllm/engine/arg_utils.py
@@ -211,6 +211,13 @@ class EngineArgs:
     reasoning_parser: Optional[str] = None
     use_tqdm_on_load: bool = True
 
+    # mlu platform params
+    moe_tp_size: int = -1
+    moe_ep_size: int = -1
+    attn_data_parallel_size: int = -1
+    parallelize_shared_expert: bool = False
+    disable_prefill_force_to_use_reduce_scatter: bool = False
+
     def __post_init__(self):
         if not self.tokenizer:
             self.tokenizer = self.model
@@ -361,7 +368,7 @@ class EngineArgs:
         parser.add_argument(
             '--kv-cache-dtype',
             type=str,
-            choices=['auto', 'fp8', 'fp8_e5m2', 'fp8_e4m3'],
+            choices=['auto', 'fp8', 'fp8_e5m2', 'fp8_e4m3', 'int8'],
             default=EngineArgs.kv_cache_dtype,
             help='Data type for kv cache storage. If "auto", will use model '
             'data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. '
@@ -1001,6 +1008,28 @@ class EngineArgs:
             "Note that even if this is set to False, cascade attention will be "
             "only used when the heuristic tells that it's beneficial.")
 
+        # mlu platform params
+        parser.add_argument('--moe-tp-size',
+                            type=int,
+                            default=-1,
+                            help='Number of moe tensor parallel replicas')
+        parser.add_argument('--moe-ep-size',
+                            type=int,
+                            default=-1,
+                            help='Number of moe expert parallel replicas')
+        parser.add_argument('--attn-data-parallel-size',
+                            type=int,
+                            default=-1,
+                            help='Number of attention data parallel replicas')
+        parser.add_argument('--parallelize-shared-expert',
+                            action='store_true',
+                            help='Parallelize shared expert for data parallel')
+        parser.add_argument('--disable-prefill-force-to-use-reduce-scatter',
+                            action='store_true',
+                            help='Prefill force to use reduce scatter can have trade off. '
+                                 'This flag makes prefill can use common all-reduce strategy '
+                                 'for data parallel.')
+
         return parser
 
     @classmethod
diff --git a/vllm/executor/ray_distributed_executor.py b/vllm/executor/ray_distributed_executor.py
index 9b0b98731..f5d554345 100644
--- a/vllm/executor/ray_distributed_executor.py
+++ b/vllm/executor/ray_distributed_executor.py
@@ -318,7 +318,7 @@ class RayDistributedExecutor(DistributedExecutorBase):
         n_ips = len(all_ips)
         n_nodes = len(node_workers)
 
-        if n_nodes != n_ips:
+        if self.parallel_config.data_parallel_size <= 1 and n_nodes != n_ips:
             raise RuntimeError(
                 f"Every node should have a unique IP address. Got {n_nodes}"
                 f" nodes with node ids {list(node_workers.keys())} and "
diff --git a/vllm/executor/ray_utils.py b/vllm/executor/ray_utils.py
index 37cc07bfb..66d0055a3 100644
--- a/vllm/executor/ray_utils.py
+++ b/vllm/executor/ray_utils.py
@@ -89,7 +89,10 @@ try:
             # current device.
             import torch
             if not self.compiled_dag_cuda_device_set:
-                torch.cuda.set_device(self.worker.device)
+                if current_platform.is_out_of_tree():
+                    current_platform.set_device(self.worker.device)
+                else:   
+                    torch.cuda.set_device(self.worker.device)
                 self.compiled_dag_cuda_device_set = True
 
             output = self.worker._execute_model_spmd(execute_model_req,
@@ -112,6 +115,8 @@ try:
                 if current_platform.is_tpu():
                     # Not needed
                     pass
+                elif current_platform.is_out_of_tree():
+                    current_platform.set_device(self.worker.device)
                 else:
                     import torch
                     torch.cuda.set_device(self.worker.device)
@@ -184,7 +189,7 @@ def _verify_bundles(placement_group: "PlacementGroup",
         node_id_to_bundle[node_id].append(bundles[bundle_idx])
     driver_node_id = ray.get_runtime_context().get_node_id()
 
-    if driver_node_id not in node_id_to_bundle:
+    if parallel_config.data_parallel_size <= 1 and (driver_node_id not in node_id_to_bundle):
         raise RuntimeError(
             f"driver node id {driver_node_id} is not included in a placement "
             f"group {placement_group.id}. Node id -> bundles "
@@ -351,8 +356,49 @@ def initialize_ray_cluster(
         # vLLM engine is also a worker to execute model with an accelerator,
         # so it requires to have the device in a current node. Check if
         # the current node has at least one device.
-        current_ip = get_ip()
-        current_node_id = ray.get_runtime_context().get_node_id()
+
+
+        head_ip = get_ip()
+
+        def sort_by_driver_then_worker_ip(ip_and_id):
+            ip = ip_and_id[0]
+            return (0 if ip == head_ip else 1, ip)
+
+        def _get_gpu_mapping():
+            nodes = ray.nodes()
+            node_gpu_mapping = []
+            for node_info in nodes:
+                if node_info.get("alive", False):
+                    node_id = node_info["NodeID"]
+                    ip = node_info["NodeManagerAddress"]
+                    num_gpus = node_info["Resources"].get("GPU", 0)
+                    node_gpu_mapping.append((ip, node_id, int(num_gpus)))
+            node_gpu_mapping = sorted(node_gpu_mapping, key=sort_by_driver_then_worker_ip)
+            return node_gpu_mapping
+
+        def _find_target_gpu(node_gpu_mapping, dp_rank):
+            accumulated_gpus = 0
+            world_size = parallel_config.world_size
+            logger.info(f"for initialize_ray_cluster, world_size = {world_size}")
+            for ip, node_id, num_gpus in node_gpu_mapping:
+                if dp_rank * world_size < accumulated_gpus + num_gpus:
+                    gpu_index = dp_rank * world_size - accumulated_gpus
+                    return (ip, node_id, gpu_index)
+                accumulated_gpus += num_gpus
+            raise ValueError(f"dp_rank {dp_rank} exceeds the num_gpus {accumulated_gpus}")
+
+        node_gpu_mapping = _get_gpu_mapping()
+        current_dp_rank = parallel_config.data_parallel_rank
+        selected_node_ip, selected_node_id, gpu_index = _find_target_gpu(node_gpu_mapping, current_dp_rank)
+        logger.debug(f"selected_node_ip = {selected_node_ip}, selected_node_id = {selected_node_id}, gpu_index = {gpu_index}")
+
+        current_ip = selected_node_ip
+        current_node_id = selected_node_id
+
+        # current_ip = get_ip()
+        # current_node_id = ray.get(get_remote_node_id.remote())
+
+
         current_node_resource = available_resources_per_node()[current_node_id]
         if current_node_resource.get(device_str, 0) < 1:
             raise ValueError(
diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py
index 0e35d8a80..80ac5f42d 100644
--- a/vllm/model_executor/layers/fused_moe/layer.py
+++ b/vllm/model_executor/layers/fused_moe/layer.py
@@ -437,7 +437,7 @@ class FusedMoE(torch.nn.Module):
         # Use expert parallelism instead of tensor parallelism?
         vllm_config = get_current_vllm_config()
         use_ep = (vllm_config.parallel_config.enable_expert_parallel
-                  and self.tp_size > 1)
+                  and self.tp_size * self.dp_size > 1)
 
         # For smuggling this layer into the fused moe custom op
         self.use_direct_call = self.dp_size == 1
diff --git a/vllm/model_executor/models/baichuan.py b/vllm/model_executor/models/baichuan.py
index 6a3112b5f..9777712fb 100644
--- a/vllm/model_executor/models/baichuan.py
+++ b/vllm/model_executor/models/baichuan.py
@@ -392,6 +392,7 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP,
         self.lm_head = ParallelLMHead(config.vocab_size,
                                       config.hidden_size,
                                       quant_config=quant_config)
+        self.lm_head.weight.weight_loader_org = self.lm_head.weight.weight_loader
         self.lm_head.weight.weight_loader = self.lm_head_weight_loader
         if self.config.tie_word_embeddings:
             self.lm_head.weight = self.model.embed_tokens.weight
@@ -448,7 +449,9 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP,
         if is_baichuan2:
             loaded_weight = torch.nn.functional.normalize(loaded_weight)
 
-        default_weight_loader(param, loaded_weight)
+        weight_loader = getattr(param, "weight_loader_org",
+                                        default_weight_loader)
+        weight_loader(param, loaded_weight)
 
 
 class BaichuanForCausalLM(BaiChuanBaseForCausalLM):
diff --git a/vllm/model_executor/models/eagle.py b/vllm/model_executor/models/eagle.py
index 3e4a5040b..8b73868f5 100644
--- a/vllm/model_executor/models/eagle.py
+++ b/vllm/model_executor/models/eagle.py
@@ -35,7 +35,7 @@ class DummyInputLayerNorm(nn.Module):
 
 class DummyOutputNorm(nn.Module):
 
-    def forward(self, x, residual):
+    def forward(self, x, residual=None):
         if residual is None:
             return x
         else:
diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py
index c4d02e5dd..2831a5a12 100644
--- a/vllm/model_executor/models/qwen2.py
+++ b/vllm/model_executor/models/qwen2.py
@@ -263,7 +263,11 @@ class Qwen2DecoderLayer(nn.Module):
     })
 class Qwen2Model(nn.Module):
 
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+    def __init__(self,
+                 *,
+                 vllm_config: VllmConfig,
+                 prefix: str = "",
+                 decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer):
         super().__init__()
 
         config = vllm_config.model_config.hf_config
@@ -297,12 +301,14 @@ class Qwen2Model(nn.Module):
         else:
             self.embed_tokens = PPMissingLayer()
 
+        # Use the provided decoder layer type or default to Qwen2DecoderLayer
+        decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
-            lambda prefix: Qwen2DecoderLayer(config=config,
-                                             cache_config=cache_config,
-                                             quant_config=quant_config,
-                                             prefix=prefix),
+            lambda prefix: decoder_layer_type(config=config,
+                                              cache_config=cache_config,
+                                              quant_config=quant_config,
+                                              prefix=prefix),
             prefix=f"{prefix}.layers",
         )
 
diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py
new file mode 100644
index 000000000..9c14038e6
--- /dev/null
+++ b/vllm/model_executor/models/qwen3.py
@@ -0,0 +1,329 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2024 The Qwen team.
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only Qwen3 model compatible with HuggingFace weights."""
+from typing import Iterable, Optional, Set, Tuple, Union
+
+import torch
+from torch import nn
+from transformers import Qwen3Config
+
+from vllm.attention import Attention, AttentionType
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (QKVParallelLinear,
+                                               RowParallelLinear)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.sequence import IntermediateTensors
+
+from .interfaces import SupportsLoRA, SupportsPP
+from .qwen2 import Qwen2MLP as Qwen3MLP
+from .qwen2 import Qwen2Model
+from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix
+
+logger = init_logger(__name__)
+
+
+class Qwen3Attention(nn.Module):
+
+    def __init__(self,
+                 hidden_size: int,
+                 num_heads: int,
+                 num_kv_heads: int,
+                 max_position: int = 4096 * 32,
+                 head_dim: Optional[int] = None,
+                 rms_norm_eps: float = 1e-06,
+                 qkv_bias: bool = False,
+                 rope_theta: float = 10000,
+                 cache_config: Optional[CacheConfig] = None,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 rope_scaling: Optional[Tuple] = None,
+                 prefix: str = "",
+                 attn_type: str = AttentionType.DECODER) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim or hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=qkv_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=max_position,
+            base=self.rope_theta,
+            rope_scaling=rope_scaling,
+        )
+        self.attn = Attention(self.num_heads,
+                              self.head_dim,
+                              self.scaling,
+                              num_kv_heads=self.num_kv_heads,
+                              cache_config=cache_config,
+                              quant_config=quant_config,
+                              prefix=f"{prefix}.attn",
+                              attn_type=attn_type)
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        # Add qk-norm
+        q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim,
+                           self.head_dim)
+        q_by_head = self.q_norm.forward_native(q_by_head)
+        q = q_by_head.view(q.shape)
+        k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim,
+                           self.head_dim)
+        k_by_head = self.k_norm.forward_native(k_by_head)
+        k = k_by_head.view(k.shape)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Qwen3DecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: Qwen3Config,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        # Requires transformers > 4.32.0
+        rope_theta = getattr(config, "rope_theta", 1000000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+
+        # By default, Qwen3 uses causal attention as it is a decoder-only model.
+        # You can override the HF config with `is_causal=False` to enable
+        # bidirectional attention, which is used in some embedding models
+        # (e.g. Alibaba-NLP/gte-Qwen3-7B-instruct)
+        if getattr(config, "is_causal", True):
+            attn_type = AttentionType.DECODER
+        else:
+            attn_type = AttentionType.ENCODER_ONLY
+
+        self.self_attn = Qwen3Attention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            max_position=config.max_position_embeddings,
+            num_kv_heads=config.num_key_value_heads,
+            rope_theta=rope_theta,
+            rms_norm_eps=config.rms_norm_eps,
+            qkv_bias=getattr(config, 'attention_bias', False),
+            head_dim=getattr(config, 'head_dim', None),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            rope_scaling=rope_scaling,
+            prefix=f"{prefix}.self_attn",
+            attn_type=attn_type,
+        )
+        self.mlp = Qwen3MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size,
+                                       eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                                eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(
+                hidden_states, residual)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+ALL_DECODER_LAYER_TYPES = {
+    "attention": Qwen3DecoderLayer,
+}
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
+        # otherwise (seq_len, ).
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+    })
+class Qwen3Model(Qwen2Model):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config,
+                         prefix=prefix,
+                         decoder_layer_type=Qwen3DecoderLayer)
+
+
+class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": [
+            "gate_proj",
+            "up_proj",
+        ],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        lora_config = vllm_config.lora_config
+
+        self.config = config
+        self.lora_config = lora_config
+
+        self.quant_config = quant_config
+        self.model = Qwen3Model(vllm_config=vllm_config,
+                                prefix=maybe_prefix(prefix, "model"))
+
+        if get_pp_group().is_last_rank:
+            if config.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(config.vocab_size,
+                                              config.hidden_size,
+                                              quant_config=quant_config,
+                                              prefix=maybe_prefix(
+                                                  prefix, "lm_head"))
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.sampler = get_sampler()
+
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.get_input_embeddings(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        hidden_states = self.model(input_ids, positions, intermediate_tensors,
+                                   inputs_embeds)
+        return hidden_states
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[torch.Tensor]:
+        logits = self.logits_processor(self.lm_head, hidden_states,
+                                       sampling_metadata)
+        return logits
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[SamplerOutput]:
+        next_tokens = self.sampler(logits, sampling_metadata)
+        return next_tokens
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["lm_head."]
+                           if self.config.tie_word_embeddings else None),
+        )
+        return loader.load_weights(weights)
diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py
new file mode 100644
index 000000000..f0ef79dfd
--- /dev/null
+++ b/vllm/model_executor/models/qwen3_moe.py
@@ -0,0 +1,538 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2024 The Qwen team.
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only Qwen3MoE model compatible with HuggingFace weights."""
+from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+
+from vllm.attention import Attention
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import (get_pp_group,
+                              get_tensor_model_parallel_world_size,
+                              tensor_model_parallel_all_reduce)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
+                                               QKVParallelLinear,
+                                               ReplicatedLinear,
+                                               RowParallelLinear)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead, VocabParallelEmbedding)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.sequence import IntermediateTensors
+
+from .interfaces import SupportsPP
+from .utils import (AutoWeightsLoader, extract_layer_index,
+                    is_pp_missing_parameter,
+                    make_empty_intermediate_tensors_factory, make_layers,
+                    maybe_prefix)
+
+logger = init_logger(__name__)
+
+
+class Qwen3MoeMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        reduce_results: bool = True,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size, [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj")
+        self.down_proj = RowParallelLinear(intermediate_size,
+                                           hidden_size,
+                                           bias=False,
+                                           quant_config=quant_config,
+                                           reduce_results=reduce_results,
+                                           prefix=f"{prefix}.down_proj")
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. "
+                             "Only silu is supported for now.")
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Qwen3MoeSparseMoeBlock(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+        if self.tp_size > config.num_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than "
+                f"the number of experts {config.num_experts}.")
+
+        self.experts = FusedMoE(num_experts=config.num_experts,
+                                top_k=config.num_experts_per_tok,
+                                hidden_size=config.hidden_size,
+                                intermediate_size=config.moe_intermediate_size,
+                                reduce_results=False,
+                                renormalize=config.norm_topk_prob,
+                                quant_config=quant_config,
+                                prefix=f"{prefix}.experts")
+
+        self.gate = ReplicatedLinear(config.hidden_size,
+                                     config.num_experts,
+                                     bias=False,
+                                     quant_config=None,
+                                     prefix=f"{prefix}.gate")
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # NOTE: hidden_states can have either 1D or 2D shape.
+        orig_shape = hidden_states.shape
+        hidden_dim = hidden_states.shape[-1]
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.experts(hidden_states=hidden_states,
+                                           router_logits=router_logits)
+        final_hidden_states = final_hidden_states
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                final_hidden_states)
+
+        return final_hidden_states.view(orig_shape)
+
+
+class Qwen3MoeAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_theta: float = 10000,
+        rope_scaling: Optional[Dict[str, Any]] = None,
+        max_position_embeddings: int = 8192,
+        head_dim: Optional[int] = None,
+        rms_norm_eps: float = 1e-06,
+        qkv_bias: bool = False,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim or (hidden_size // self.total_num_heads)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+
+        self.qkv_proj = QKVParallelLinear(hidden_size,
+                                          self.head_dim,
+                                          self.total_num_heads,
+                                          self.total_num_kv_heads,
+                                          bias=qkv_bias,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.qkv_proj")
+
+        self.o_proj = RowParallelLinear(self.total_num_heads * self.head_dim,
+                                        hidden_size,
+                                        bias=False,
+                                        quant_config=quant_config,
+                                        prefix=f"{prefix}.o_proj")
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=max_position_embeddings,
+            base=rope_theta,
+            rope_scaling=rope_scaling,
+        )
+        self.attn = Attention(self.num_heads,
+                              self.head_dim,
+                              self.scaling,
+                              num_kv_heads=self.num_kv_heads,
+                              cache_config=cache_config,
+                              quant_config=quant_config,
+                              prefix=f"{prefix}.attn")
+
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        # Add qk-norm
+        q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim,
+                           self.head_dim)
+        q_by_head = self.q_norm.forward_native(q_by_head)
+        q = q_by_head.view(q.shape)
+
+        k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim,
+                           self.head_dim)
+        k_by_head = self.k_norm.forward_native(k_by_head)
+        k = k_by_head.view(k.shape)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Qwen3MoeDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 10000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
+        self.self_attn = Qwen3MoeAttention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            rms_norm_eps=config.rms_norm_eps,
+            qkv_bias=getattr(config, 'attention_bias', False),
+            head_dim=getattr(config, 'head_dim', None),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+
+        # `mlp_only_layers` in the config.
+        layer_idx = extract_layer_index(prefix)
+        mlp_only_layers = ([] if not hasattr(config, "mlp_only_layers") else
+                           config.mlp_only_layers)
+        if (layer_idx not in mlp_only_layers) and (
+                config.num_experts > 0 and
+            (layer_idx + 1) % config.decoder_sparse_step == 0):
+            self.mlp = Qwen3MoeSparseMoeBlock(config=config,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.mlp")
+        else:
+            self.mlp = Qwen3MoeMLP(hidden_size=config.hidden_size,
+                                   intermediate_size=config.intermediate_size,
+                                   hidden_act=config.hidden_act,
+                                   quant_config=quant_config,
+                                   prefix=f"{prefix}.mlp")
+        self.input_layernorm = RMSNorm(config.hidden_size,
+                                       eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                                eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(
+                hidden_states, residual)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+@support_torch_compile
+class Qwen3MoeModel(nn.Module):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.config = config
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=f"{prefix}.embed_tokens")
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: Qwen3MoeDecoderLayer(config=config,
+                                                cache_config=cache_config,
+                                                quant_config=quant_config,
+                                                prefix=prefix),
+            prefix=f"{prefix}.layers",
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.make_empty_intermediate_tensors = (
+            make_empty_intermediate_tensors_factory(
+                ["hidden_states", "residual"], config.hidden_size))
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.get_input_embeddings(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+        for i in range(self.start_layer, self.end_layer):
+            layer = self.layers[i]
+            hidden_states, residual = layer(positions, hidden_states, residual)
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({
+                "hidden_states": hidden_states,
+                "residual": residual
+            })
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        # Params for weights, fp8 weight scales, fp8 activation scales
+        # (param_name, weight_name, expert_id, shard_id)
+        expert_params_mapping = FusedMoE.make_expert_params_mapping(
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.num_experts)
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            for (param_name, weight_name, shard_id) in stacked_params_mapping:
+                # Skip non-stacked layers and experts (experts handled below).
+                if weight_name not in name:
+                    continue
+                # We have mlp.experts[0].gate_proj in the checkpoint.
+                # Since we handle the experts below in expert_params_mapping,
+                # we need to skip here BEFORE we update the name, otherwise
+                # name will be updated to mlp.experts[0].gate_up_proj, which
+                # will then be updated below in expert_params_mapping
+                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
+                if "mlp.experts" in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if ((name.endswith(".bias") or name.endswith("_bias"))
+                        and name not in params_dict):
+                    continue
+                # Skip layers on other devices.
+                if is_pp_missing_parameter(name, self):
+                    continue
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                for mapping in expert_params_mapping:
+                    param_name, weight_name, expert_id, shard_id = mapping
+                    if weight_name not in name:
+                        continue
+                    name = name.replace(weight_name, param_name)
+                    # Skip layers on other devices.
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    # Skip loading extra bias for GPTQ models.
+                    if ((name.endswith(".bias") or name.endswith("_bias"))
+                            and name not in params_dict):
+                        continue
+                    param = params_dict[name]
+                    weight_loader = param.weight_loader
+                    weight_loader(param,
+                                  loaded_weight,
+                                  name,
+                                  shard_id=shard_id,
+                                  expert_id=expert_id)
+                    break
+                else:
+                    # Skip loading extra bias for GPTQ models.
+                    if ((name.endswith(".bias") or name.endswith("_bias"))
+                            and name not in params_dict):
+                        continue
+                    # Skip layers on other devices.
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    # Remapping the name of FP8 kv-scale.
+                    if name.endswith("kv_scale"):
+                        remapped_kv_scale_name = name.replace(
+                            ".kv_scale", ".attn.kv_scale")
+                        if remapped_kv_scale_name not in params_dict:
+                            logger.warning_once(
+                                "Found kv scale in the checkpoint "
+                                f"(e.g. {name}), but not found the expected "
+                                f"name in the model "
+                                f"(e.g. {remapped_kv_scale_name}). "
+                                "kv-scale is not loaded.")
+                            continue
+                        else:
+                            name = remapped_kv_scale_name
+                    param = params_dict[name]
+                    weight_loader = getattr(param, "weight_loader",
+                                            default_weight_loader)
+                    weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class Qwen3MoeForCausalLM(nn.Module, SupportsPP):
+
+    fall_back_to_pt_during_load = False
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+        self.model = Qwen3MoeModel(vllm_config=vllm_config,
+                                   prefix=maybe_prefix(prefix, "model"))
+        self.lm_head = ParallelLMHead(config.vocab_size,
+                                      config.hidden_size,
+                                      quant_config=quant_config)
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.sampler = get_sampler()
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.get_input_embeddings(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        hidden_states = self.model(input_ids, positions, intermediate_tensors,
+                                   inputs_embeds)
+        return hidden_states
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[torch.Tensor]:
+        logits = self.logits_processor(self.lm_head, hidden_states,
+                                       sampling_metadata)
+        return logits
+
+    def sample(
+        self,
+        logits: Optional[torch.Tensor],
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[SamplerOutput]:
+        next_tokens = self.sampler(logits, sampling_metadata)
+        return next_tokens
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["rotary_emb.inv_freq"]),
+        )
+        return loader.load_weights(weights)
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index c0a3c59ba..02be49fd6 100644
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -101,6 +101,8 @@ _TEXT_GENERATION_MODELS = {
     "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
     "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
     "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
+    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
+    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
     "RWForCausalLM": ("falcon", "FalconForCausalLM"),
     "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
     "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
diff --git a/vllm/spec_decode/mqa_scorer.py b/vllm/spec_decode/mqa_scorer.py
index 6275c460e..5d84b3299 100644
--- a/vllm/spec_decode/mqa_scorer.py
+++ b/vllm/spec_decode/mqa_scorer.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+import torch
 
 from vllm.sequence import (ExecuteModelRequest, SequenceData,
                            SequenceGroupMetadata, get_all_seq_ids)
@@ -149,8 +150,15 @@ class MQAScorer(SpeculativeScorer):
 
         hidden_states = None
         if target_sampler_output.hidden_states is not None:
-            hidden_states = target_sampler_output.hidden_states.reshape(
-                bs, (k + 1), -1)
+            hidden_states = target_sampler_output.hidden_states
+            num_tokens = hidden_states.numel() // hidden_states.shape[-1]
+            if num_tokens < bs * (k + 1):
+                pad_size = bs * (k + 1) - num_tokens
+                pad_value = [0] * hidden_states.dim() * 2
+                pad_value[-1] = pad_size
+                hidden_states = torch.nn.functional.pad(
+                    hidden_states, pad_value)
+            hidden_states = hidden_states.reshape(bs, (k + 1), -1)
 
         return SpeculativeScores(probs=all_probs,
                                  token_ids=all_tokens,
diff --git a/vllm/spec_decode/top1_proposer.py b/vllm/spec_decode/top1_proposer.py
index b538923c0..0478df16f 100644
--- a/vllm/spec_decode/top1_proposer.py
+++ b/vllm/spec_decode/top1_proposer.py
@@ -144,7 +144,7 @@ class Top1Proposer(SpeculativeProposer):
             # quota for nonzero_proposal
             new_k = 0
             if (self.max_proposal_len is None
-                    or seq_len + proposal_len < self.max_proposal_len):
+                    or (seq_len < self.max_proposal_len and proposal_len > 0)):
                 new_k = proposal_len
                 nonzero_proposal_len_seqs.append(seq_group_metadata)
                 nonzero_proposal_len_indices.append(i)
@@ -254,12 +254,15 @@ class Top1Proposer(SpeculativeProposer):
             size=(batch_size, *proposal_tokens.shape[1:]),
             fill_value=-1,
         )
-        entire_proposal_tokens[nonzero_proposal_len_indices] = proposal_tokens
+        valid_batch_size = min(proposal_tokens.shape[0], batch_size)
+        valid_proposal_len_indices = nonzero_proposal_len_indices[
+            0: valid_batch_size]
+        entire_proposal_tokens[valid_proposal_len_indices] = proposal_tokens
         entire_proposal_probs = proposal_probs.new_zeros(
             batch_size,
             *proposal_probs.shape[1:],
         )
-        entire_proposal_probs[nonzero_proposal_len_indices] = proposal_probs
+        entire_proposal_probs[valid_proposal_len_indices] = proposal_probs
 
         proposal_tokens, proposal_probs = (
             entire_proposal_tokens,
@@ -269,6 +272,6 @@ class Top1Proposer(SpeculativeProposer):
         proposal_lens_tensor = torch.zeros(batch_size,
                                            dtype=torch.long,
                                            device=self._device)
-        proposal_lens_tensor[nonzero_proposal_len_indices] = proposal_len
+        proposal_lens_tensor[valid_proposal_len_indices] = proposal_len
 
         return proposal_tokens, proposal_probs, proposal_lens_tensor
diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py
index e948e59b8..048caca8d 100644
--- a/vllm/v1/engine/core_client.py
+++ b/vllm/v1/engine/core_client.py
@@ -388,7 +388,7 @@ class MPClient(EngineCoreClient):
         # Paths and sockets for IPC.
         self.output_path = get_open_zmq_ipc_path()
 
-        new_core_engine = lambda index, local_dp_rank=None: CoreEngine(
+        new_core_engine = lambda index, local_dp_rank: CoreEngine(
             vllm_config, executor_class, log_stats, self.ctx, self.output_path,
             index, local_dp_rank)
 
@@ -411,9 +411,8 @@ class MPClient(EngineCoreClient):
 
         # Default case - single core engine.
         dp_rank = vllm_config.parallel_config.data_parallel_rank
-        local_dp_rank = vllm_config.parallel_config.data_parallel_rank_local
         core_engine = new_core_engine(
-            dp_rank, local_dp_rank if local_dp_rank is not None else dp_rank)
+            dp_rank, dp_rank)
         core_engines.append(core_engine)
         self.core_engine = core_engine
 
@@ -765,11 +764,12 @@ class DPAsyncMPClient(AsyncMPClient):
         else:
             # Send request to chosen engine and dp start loop
             # control message to all other engines.
-            self.num_engines_running += len(self.core_engines)
+            engines = [engine for engine in self.core_engines if engine is chosen_engine or not engine.num_reqs_in_flight]
+            self.num_engines_running += len(engines)
             await asyncio.gather(*[
                 engine.send_multipart(msg if engine is
                                       chosen_engine else self.start_dp_msg)
-                for engine in self.core_engines
+                for engine in engines
             ])
 
         self._ensure_output_queue_task()
@@ -792,11 +792,11 @@ class DPAsyncMPClient(AsyncMPClient):
                 # If there are requests in flight here, they must have
                 # been sent after the engines paused. We must make
                 # sure to start the other engines:
-                self.num_engines_running = len(self.core_engines)
+                engines = [engine for engine in self.core_engines if not engine.num_reqs_in_flight]
+                self.num_engines_running += len(engines)
                 coros = [
                     engine.send_multipart(self.start_dp_msg)
-                    for engine in self.core_engines
-                    if not engine.num_reqs_in_flight
+                    for engine in engines
                 ]
                 if coros:
                     await asyncio.gather(*coros)
diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py
index 1d5175eb6..28a36acab 100644
--- a/vllm/v1/executor/multiproc_executor.py
+++ b/vllm/v1/executor/multiproc_executor.py
@@ -46,6 +46,9 @@ class MultiprocExecutor(Executor):
         # The child processes will send SIGUSR1 when unrecoverable
         # errors happen.
         def sigusr1_handler(signum, frame):
+            import traceback
+            logger.fatal("Singal received. Current stack:\n%s",
+                         "".join(traceback.format_stack(frame)))
             logger.fatal(
                 "MulitprocExecutor got fatal signal from worker processes, "
                 "shutting down. See stack trace above for root cause issue.")
diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
index 86e6d9752..94d49eae7 100644
--- a/vllm/worker/model_runner.py
+++ b/vllm/worker/model_runner.py
@@ -74,6 +74,19 @@ TModelInputForGPU = TypeVar('TModelInputForGPU', bound="ModelInputForGPU")
 torch._dynamo.config.cache_size_limit = 128
 torch._dynamo.config.accumulated_cache_size_limit = 128
 
+# --- FLAGSCALE MODIFICATION BEG ---
+# Know more about FlagGems: https://github.com/FlagOpen/FlagGems
+import os
+if os.getenv("USE_FLAGGEMS", "false").lower() in ("1", "true", "yes"):
+    try:
+        import flag_gems
+        flag_gems.enable()
+        logger.info("Successfully enabled flag_gems as default ops implementation.")
+    except ImportError:
+        logger.warning("Failed to import 'flag_gems'. Falling back to default implementation.")
+    except Exception as e:
+        logger.warning(f"Failed to enable 'flag_gems': {e}. Falling back to default implementation.")
+# --- FLAGSCALE MODIFICATION END ---
 
 @dataclass(frozen=True)
 class ModelInputForGPU(ModelRunnerInputBase):
diff --git a/vllm_mlu/requirements.txt b/vllm_mlu/requirements.txt
new file mode 100644
index 000000000..3b3e66923
--- /dev/null
+++ b/vllm_mlu/requirements.txt
@@ -0,0 +1,18 @@
+# Dependencies for Cambricon MLUs
+ray == 2.40.0
+triton >= 3.2.0
+torch == 2.6.0
+torch-mlu >= 1.25.2
+torch_mlu_ops >= 1.4.1
+xformers == 0.0.24
+
+matplotlib # for VLLM_SCHEDULER_PROFILE
+datasets # for benchmark scripts
+
+# Dependencies for tests
+transformers==4.51.0  # Avoid v4.51.0 due to OOM issues in LoRA UT
+pytest-asyncio==0.24.0
+modelscope
+sentence-transformers==3.4.1
+runai-model-streamer==0.11.0
+runai-model-streamer-s3==0.11.0
diff --git a/vllm_mlu/setup.py b/vllm_mlu/setup.py
new file mode 100644
index 000000000..8a3c04d16
--- /dev/null
+++ b/vllm_mlu/setup.py
@@ -0,0 +1,114 @@
+import io
+import logging
+import os
+import re
+import subprocess
+
+from typing import List
+from setuptools import find_namespace_packages, setup
+
+
+ROOT_DIR = os.path.dirname(__file__)
+logger = logging.getLogger(__name__)
+
+
+def get_path(*filepath) -> str:
+    return os.path.join(ROOT_DIR, *filepath)
+
+
+def get_commit_id() -> str:
+    """
+    get the current commit of vllm mlu
+    """
+    git_short_hash = subprocess.run(
+        ['git', 'rev-parse', '--short', 'HEAD'],
+        stdout=subprocess.PIPE,
+        text=True
+    ).stdout.strip()
+    return git_short_hash
+
+
+def get_vllm_version() -> str:
+    """
+    get vllm version
+    """
+    with open(get_path("tools/build.property"), 'r') as file:
+        content = file.read()
+
+    results = re.findall(r'VLLM_VERSION=([\d|\.]+)\+mlu([\d|\.]+)\.pt(\d+)', content)
+
+    assert results, "fail to get vllm, vllm_mlu and pytorch version."
+
+    version = f"{results[-1][0]}+mlu{results[-1][1]}.pt{results[-1][2]}"
+
+    if get_commit_id():
+        version += "." + get_commit_id()
+    return version
+
+
+def read_readme() -> str:
+    """Read the README file if present."""
+    p = get_path("README.md")
+    if os.path.isfile(p):
+        return io.open(get_path("README.md"), "r", encoding="utf-8").read()
+    else:
+        return ""
+
+
+def get_requirements() -> List[str]:
+    """Get Python package dependencies from requirements.txt."""
+
+    def _read_requirements(filename: str) -> List[str]:
+        with open(get_path(filename)) as f:
+            requirements = f.read().strip().split("\n")
+        resolved_requirements = []
+        for line in requirements:
+            if line.startswith("-r "):
+                resolved_requirements += _read_requirements(line.split()[1])
+            elif line.startswith("--"):
+                continue
+            else:
+                resolved_requirements.append(line)
+        return resolved_requirements
+
+    return _read_requirements("requirements.txt")
+
+
+ext_modules = []
+cmdclass = {}
+
+
+setup(
+    name="vllm_mlu",
+    version=get_vllm_version(),
+    author="Cambricon vLLM Team",
+    license="Apache 2.0",
+    description=("A high-throughput and memory-efficient inference and "
+                 "serving engine for LLMs on MLU backend"),
+    long_description=read_readme(),
+    long_description_content_type="text/markdown",
+    url="",
+    project_urls={
+        "Homepage": "https://github.com/vllm-project/vllm",
+        "Documentation": "https://vllm.readthedocs.io/en/latest/",
+    },
+    classifiers=[
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "License :: OSI Approved :: Apache Software License",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+    packages=find_namespace_packages(),
+    include_package_data=True,
+    python_requires=">=3.8",
+    install_requires=get_requirements(),
+    ext_modules = ext_modules,
+    cmdclass=cmdclass,
+    entry_points={
+        'vllm.platform_plugins': ["mlu = vllm_mlu:register_mlu_platform"],
+        'vllm.general_plugins': ["mlu_hijack = vllm_mlu:register_mlu_hijack"]
+    }
+)
diff --git a/vllm_mlu/tools/build.property b/vllm_mlu/tools/build.property
new file mode 100644
index 000000000..347689733
--- /dev/null
+++ b/vllm_mlu/tools/build.property
@@ -0,0 +1,10 @@
+VLLM_VERSION=0.8.3+mlu0.8.0.pt26
+TORCH_MLU_OPS_VERSION=1.5.0+pt26
+CATCH_VERSION=1.26.1+torch2.6.0
+CNCL_VERSION=1.27.2-1
+CNNL_VERSION=2.1.0-1
+CNNLEXTRA_VERSION=2.0.3-1
+CNTOOLKIT_VERSION=4.1.2-1
+MLUOPS_VERSION=1.6.0-1
+TRITON_VERSION=3.2.0+mlu1.5.0
+XFORMERS_VERSION=0.0.24+mlu0.7.0.pt2.6
diff --git a/vllm_mlu/vllm_mlu/__init__.py b/vllm_mlu/vllm_mlu/__init__.py
new file mode 100644
index 000000000..d296626f3
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/__init__.py
@@ -0,0 +1,25 @@
+#
+# Copyright (c) 2025 Cambricon Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-mlu project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+def register_mlu_platform():
+    """Register the MLU platform."""
+    return "vllm_mlu.platforms.mlu.MLUPlatform"
+
+
+def register_mlu_hijack():
+    from vllm_mlu import mlu_hijack
+    return
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/_mlu_ops.py b/vllm_mlu/vllm_mlu/_mlu_ops.py
new file mode 100644
index 000000000..c5f368fb1
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/_mlu_ops.py
@@ -0,0 +1,1102 @@
+from typing import List, Optional, Tuple
+
+import torch
+
+import math
+import triton
+import triton.language as tl
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+try:
+    import torch_mlu_ops as tmo
+except ImportError as e:
+    logger.warning("Failed to import from TMO OPS with %r", e)
+
+@triton.jit
+def _triton_advance_step(input_tokens_ptr,
+                         sampled_token_ids_ptr,
+                         input_positions_ptr,
+                         seq_lens_ptr,
+                         slot_mapping_ptr,
+                         block_tables_ptr,
+                         block_tables_stride,
+                         num_seqs,
+                         num_queries,
+                         block_size,
+                         TILE_SIZE: tl.constexpr,
+):
+    """
+    The triton implementation of advance step.
+    Reference: https://github.com/vllm-project/vllm/blob/v0.6.1/csrc/prepare_inputs/advance_step.cu#L14-L55
+    """
+    # Set meta info.
+    pid = tl.program_id(axis=0)
+    offsets = pid * TILE_SIZE + tl.arange(0, TILE_SIZE)
+    mask = offsets < num_queries
+
+    # Update input_tokens.
+    sampled_token_ids = tl.load(sampled_token_ids_ptr + offsets, mask=mask)
+    tl.store(input_tokens_ptr + offsets, sampled_token_ids, mask=mask)
+
+    seq_lens = tl.load(seq_lens_ptr + offsets, mask=mask)
+    next_seq_lens = seq_lens + 1
+    next_input_pos = next_seq_lens - 1
+
+    # Update seq_lens.
+    tl.store(seq_lens_ptr + offsets, next_seq_lens, mask=mask)
+
+    # Update input_positions.
+    tl.store(input_positions_ptr + offsets, next_input_pos, mask=mask)
+
+    # Calculate slot num.
+    block_index = next_input_pos // block_size
+    block_offset = next_input_pos % block_size
+    block_tables = tl.load(block_tables_ptr + block_tables_stride * offsets + block_index, mask=mask)
+    slot_num = block_tables * block_size + block_offset
+
+    # Update slot_mapping.
+    tl.store(slot_mapping_ptr + offsets, slot_num, mask=mask)
+
+
+
+def rotary_embedding(
+    input: torch.Tensor,
+    sin_cache: torch.Tensor,
+    cos_cache: torch.Tensor,
+    position_ids: Optional[torch.Tensor],
+    cu_seqlens: Optional[torch.Tensor],
+    interleaved: bool,
+    discrete: bool,
+    dynamic_ntk: bool,
+    max_seqlen: int,
+) -> torch.Tensor:
+    return tmo.apply_rotary(
+                input, sin_cache, cos_cache,
+                position_ids, cu_seqlens, interleaved,
+                discrete, dynamic_ntk, max_seqlen)
+
+
+def fused_rms_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    store_output_before_norm: bool,
+    quant_scale: torch.Tensor = None,
+    dynamic_quant: bool = False,
+    out: torch.Tensor = None,
+) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+    return tmo.fused_rms_norm(
+                x, residual, gamma, beta, bias,
+                eps, store_output_before_norm, quant_scale,
+                out, dynamic_quant)
+
+
+def fused_layer_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    gamma: torch.Tensor,
+    beta: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float,
+    store_output_before_norm: bool,
+    quant_scale: torch.Tensor = None,
+    dynamic_quant: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+    return tmo.fused_layer_norm(
+                x, residual, gamma, beta, bias,
+                eps, store_output_before_norm, quant_scale,
+                None, dynamic_quant)
+
+
+def flash_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    cu_seq_lens_q: torch.Tensor,
+    cu_seq_lens_kv: torch.Tensor,
+    alibi_slope: torch.Tensor,
+    attn_bias: torch.Tensor,
+    max_seq_len_q: int,
+    max_seq_len_kv: int,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    compute_dtype: torch.dtype = torch.float,
+    return_lse: bool = False,
+    block_tables: torch.Tensor = None,
+    k_cache_quant_scale: torch.Tensor = None,
+    v_cache_quant_scale: torch.Tensor = None
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    if v is None:
+        v = k
+    return tmo.flash_attention(
+        q, k, v, out,
+        cu_seq_lens_q, cu_seq_lens_kv,
+        alibi_slope, attn_bias,
+        max_seq_len_q, max_seq_len_kv,
+        softmax_scale, is_causal,
+        window_size_left, window_size_right,
+        compute_dtype, return_lse,
+        block_tables, k_cache_quant_scale,
+        v_cache_quant_scale)
+
+
+def split_head_nums(q_head_num, kv_head_num, max_q_head_num):
+    """
+    Split q_head_num such that:
+    1. The maximum value of the split q_head_num does not exceed max_q_head_num.
+    2. kv_head_num is split into the same number of parts as q_head_num.
+    3. Each split q_head_num can be evenly divided by the corresponding kv_head_num.
+    4. If kv_head_num < 1, it is adjusted to 1.
+
+    Parameters:
+    - q_head_num: int, the q_head_num to be split.
+    - kv_head_num: int, the kv_head_num to be split.
+    - max_q_head_num: int, the maximum supported q_head_num after splitting.
+
+    Returns:
+    - q_splits: list, the split q_head_num.
+    - kv_splits: list, the split kv_head_num.
+    """
+    if q_head_num <= 0 or kv_head_num <= 0:
+        return "q_head_num and kv_head_num must be positive integers!"
+
+    q_splits = []
+    kv_splits = []
+
+    # Residual value
+    remaining_q = q_head_num
+    remaining_kv = kv_head_num
+
+    while remaining_q > 0:
+        # Attempt to split q_head_num such that the maximum value does not exceed max_q_head_num.
+        for q_part in range(min(max_q_head_num, remaining_q), 0, -1):
+            # Ensure that q_part can be allocated and the corresponding kv_part is greater than or equal to 1.
+            if remaining_q % q_part == 0:
+                # Ensure that kv_part is greater than or equal to 1.
+                kv_part = max(remaining_kv // (remaining_q // q_part), 1)
+                # Ensure that q_part is divisible by kv_part.
+                if q_part % kv_part == 0:
+                    # Record the split values.
+                    q_splits.append(q_part)
+                    kv_splits.append(kv_part)
+                    remaining_q -= q_part
+                    remaining_kv -= kv_part
+                    break
+        else:
+            err_msg = f"Unable to find split method for q_head_num:{q_head_num}, kv_head_num:{kv_head_num}"
+            raise RuntimeError(err_msg)
+
+    return q_splits, kv_splits
+
+
+def repeat_elements(input_list, n):
+    """
+    Repeat each element in the list n times consecutively.
+
+    Parameters:
+    - input_list: list, the input list.
+    - n: int, the number of times each element should be repeated.
+
+    Returns:
+    - list, a new list containing the repeated elements.
+    """
+    if not isinstance(input_list, list) or not isinstance(n, int) or n < 0:
+        raise ValueError("The input must be a list, and the repetition count n must be an integer greater than or equal to 0.")
+
+    # Repeat each element n times using a list comprehension.
+    return [item for item in input_list for _ in range(n)]
+
+
+def single_query_cached_kv_attn(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    out: torch.Tensor,
+    block_tables: torch.Tensor,
+    context_lens: torch.Tensor,
+    k_cache_quant_scale: Optional[torch.Tensor],
+    v_cache_quant_scale: Optional[torch.Tensor],
+    alibi_slopes: Optional[torch.Tensor],
+    max_contxt_len: int,
+    windows_size_left: int,
+    windows_size_right: int,
+    softmax_scale: float,
+    q_head_dim: Optional[int] = 2,
+    kv_head_dim: Optional[int] = 1,
+    seq_q_dim: Optional[int] = 1,
+    max_seq_q_mul_q_divide_kv: Optional[int] = 48,
+    head_size_v: Optional[int] = -1,
+    compute_dtype: Optional[torch.dtype] = torch.float32,
+) -> None:
+    # FIXME(chenxiaobing): TMO only support windows_size_right = -1 yet.
+    windows_size_right = -1
+
+    # singleQwithkvCache limits seq_q * q_divide_kv <= max_seq_q_mul_q_divide_kv now.
+    # When the limitation is fixed, we should delete the split process.
+    seq_q = q.shape[seq_q_dim]
+    q_head_num = q.shape[q_head_dim]
+    kv_head_num = k_cache.shape[kv_head_dim]
+    q_divide_kv = q_head_num // kv_head_num
+    if seq_q * q_divide_kv <= max_seq_q_mul_q_divide_kv:
+        tmo.single_query_cached_kv_attn(
+            q, k_cache, v_cache, out,
+            block_tables, context_lens,
+            k_cache_quant_scale, v_cache_quant_scale,
+            alibi_slopes, max_contxt_len,
+            windows_size_left, windows_size_right, softmax_scale,
+            head_size_v=head_size_v,
+            compute_dtype=compute_dtype)
+    else:
+        max_q_head_num = max_seq_q_mul_q_divide_kv * kv_head_num // seq_q
+        q_head_num_sizes, kv_head_num_sizes = split_head_nums(q_head_num, kv_head_num, max_q_head_num)
+        parts_num = len(q_head_num_sizes)
+        q_parts = torch.split(q, q_head_num_sizes, dim=q_head_dim)
+        out_parts = torch.split(out, q_head_num_sizes, dim=q_head_dim)
+        alibi_slopes_parts = [None] * parts_num
+        if alibi_slopes:
+            alibi_slopes_parts = torch.split(alibi_slopes, q_head_num_sizes, dim=0)
+
+        kv_parts_num = parts_num
+        if parts_num > kv_head_num:
+            assert parts_num % kv_head_num == 0, f"parts_num:{parts_num} need by divided by kv_head_num:{kv_head_num} when parts_num > kv_head_num"
+            kv_parts_num = kv_head_num
+            kv_head_num_sizes = kv_head_num_sizes[:kv_parts_num]
+
+        if len(kv_head_num_sizes) > 1:
+            k_cache_parts = torch.split(k_cache, kv_head_num_sizes, dim=kv_head_dim)
+            v_cache_parts = torch.split(v_cache, kv_head_num_sizes, dim=kv_head_dim)
+            k_cache_quant_scale_parts = [None] * kv_parts_num
+            v_cache_quant_scale_parts = [None] * kv_parts_num
+            if k_cache_quant_scale:
+                k_cache_quant_scale_dim = 1 if k_cache_quant_scale.dim() == 2 else kv_head_dim
+                k_cache_quant_scale_parts = torch.split(k_cache_quant_scale, kv_head_num_sizes, dim=k_cache_quant_scale_dim)
+            if v_cache_quant_scale:
+                v_cache_quant_scale_dim = 1 if v_cache_quant_scale.dim() == 2 else kv_head_dim
+                v_cache_quant_scale_parts = torch.split(v_cache_quant_scale, kv_head_num_sizes, dim=v_cache_quant_scale_dim)
+        else:
+            k_cache_parts = [k_cache]
+            v_cache_parts = [v_cache]
+            k_cache_quant_scale_parts = [k_cache_quant_scale]
+            v_cache_quant_scale_parts = [v_cache_quant_scale]
+
+        if parts_num > kv_parts_num:
+            repeate_num = parts_num // kv_parts_num
+            k_cache_parts = repeat_elements(k_cache_parts, repeate_num)
+            v_cache_parts = repeat_elements(v_cache_parts, repeate_num)
+            k_cache_quant_scale_parts = repeat_elements(k_cache_quant_scale_parts, repeate_num)
+            v_cache_quant_scale_parts = repeat_elements(v_cache_quant_scale_parts, repeate_num)
+
+        for q_value, k_cache_value, v_cache_value, out_value, k_cache_quant_scale_value, v_cache_quant_scale_value, alibi_slopes_value in zip(
+                q_parts, k_cache_parts, v_cache_parts, out_parts, k_cache_quant_scale_parts, v_cache_quant_scale_parts,
+                alibi_slopes_parts):
+            tmo.single_query_cached_kv_attn(
+                q_value, k_cache_value.contiguous(), v_cache_value.contiguous() if v_cache_value is not None else None,
+                out_value, block_tables, context_lens,
+                k_cache_quant_scale_value, v_cache_quant_scale_value,
+                alibi_slopes_value, max_contxt_len,
+                windows_size_left, windows_size_right, softmax_scale,
+                head_size_v=head_size_v,
+                compute_dtype=compute_dtype)
+
+
+def reshape_paged_cache(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    slot_mapping: torch.Tensor
+) -> None:
+    tmo.reshape_paged_cache(k, v, k_cache, v_cache, slot_mapping)
+
+
+def swap_blocks(
+    dst: torch.Tensor,
+    src: torch.Tensor,
+    block_mapping: torch.Tensor
+) -> None:
+    # FIXME: Remove this conversion after
+    # tmo.swap_blocks support block_mapping tensor.
+    block_mapping = block_mapping.tolist()
+    block_mapping = {src: dst for src, dst in block_mapping}
+    return tmo.swap_blocks(dst, src, block_mapping)
+
+
+def copy_blocks(
+    k_caches: List[torch.Tensor],
+    v_caches: List[torch.Tensor],
+    block_mapping: torch.Tensor
+) -> None:
+    # FIXME: Remove this conversion after
+    # tmo.swap_blocks support block_mapping tensor.
+    block_mapping = block_mapping.tolist()
+    result_dict = {}
+    for row in block_mapping:
+        key = row[0]
+        values = row[1:]
+        if key in result_dict:
+            result_dict[key].extend(values)
+        else:
+            result_dict[key] = values
+    return tmo.copy_blocks(k_caches, v_caches, result_dict)
+
+
+def active(
+    input: torch.Tensor,
+    act_mode: str,
+    is_gated: bool
+) -> torch.Tensor:
+    return tmo.active(input, act_mode, is_gated)
+
+
+def fused_moe(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    bias1: Optional[torch.Tensor],
+    bias2: Optional[torch.Tensor],
+    residual: Optional[torch.Tensor],
+    input_smooth: Optional[torch.Tensor],
+    act_smooth: Optional[torch.Tensor],
+    w1_scale: Optional[torch.Tensor],
+    w2_scale: Optional[torch.Tensor],
+    topk: int,
+    renormalize: bool,
+    gated: bool,
+    act_mode: str,
+    start_expert_id: int = 0,
+    block_n: int = 0,
+    cncl_comm: int = 0
+) -> torch.Tensor:
+    return tmo.fused_moe(
+        hidden_states, gating_output,
+        w1, w2, bias1, bias2, residual,
+        input_smooth, act_smooth,
+        w1_scale, w2_scale, topk,
+        renormalize, gated, act_mode, start_expert_id,
+        block_n, cncl_comm)
+
+
+def matmul(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    c: Optional[torch.Tensor] = None,
+    act_mode: str = 'none',
+    alpha: float = 1.0,
+    beta: float = .0
+) -> torch.Tensor:
+    return tmo.matmul(a, b, bias, c, act_mode, alpha, beta)
+
+
+def weight_only_quant_matmul(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    scale: torch.Tensor,
+    zero: torch.Tensor = None,
+    bias: torch.Tensor = None,
+    c: torch.Tensor = None,
+    act_mode: str = "none",
+    quant_bit_size: int = 8,
+    alpha: float = 1.0,
+    beta: float = 1.0
+) -> torch.Tensor:
+    return tmo.weight_only_quant_matmul(
+                a, b,
+                scale, zero, bias, c,
+                act_mode, quant_bit_size, alpha, beta)
+
+
+def smooth_quant_matmul(
+    a: torch.Tensor,
+    a_scale: torch.Tensor,
+    b: torch.Tensor,
+    b_scale: torch.Tensor,
+    dtype: torch.dtype,
+    bias: torch.Tensor = None,
+    c: torch.Tensor = None,
+    act_mode: str = "none",
+    alpha: float = 1.0,
+    beta: float = 1.0
+) -> torch.Tensor:
+    return tmo.smooth_quant_matmul(
+                a, a_scale,
+                b, b_scale,
+                dtype, bias, c,
+                act_mode, alpha, beta)
+
+
+def per_token_smooth_quantize(x: torch.Tensor,
+                              smooth: torch.Tensor,
+                              zero: torch.Tensor = None,
+                              token_count: torch.Tensor = None,
+                              act_mode: str = "none",
+                              active_coef: float = 1.0,
+                              is_gated: bool = False
+                             ) -> Tuple[torch.Tensor, torch.Tensor]:
+    return tmo.per_token_smooth_quantize(x, smooth, zero, token_count, act_mode, active_coef, is_gated)
+
+
+def quantize(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    zero: torch.Tensor = None
+) -> torch.Tensor:
+    return tmo.quantize(x, scale, zero)
+
+
+def quant_to_paged_cache(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k_cache_quant_scale: torch.Tensor,
+    v_cache_quant_scale: torch.Tensor,
+    slot_mapping: torch.Tensor,
+) -> None:
+    return tmo.quant_to_paged_cache(
+        k, v, k_cache, v_cache, k_cache_quant_scale, v_cache_quant_scale, slot_mapping
+    )
+
+
+def advance_step(num_seqs: int,
+                 num_queries: int,
+                 block_size: int,
+                 input_tokens: torch.Tensor,
+                 sampled_token_ids: torch.Tensor,
+                 input_positions: torch.Tensor,
+                 seq_lens: torch.Tensor,
+                 slot_mapping: torch.Tensor,
+                 block_tables: torch.Tensor,
+                 TILE_SIZE: int = 64) -> None:
+    """
+    Advance a step on MLU for existing inputs for a multi-step runner, which
+    will update input_tokens/seq_lens/input_positions/slot_mapping inplace.
+    """
+    def verify_tensor(
+        name: str,
+        tensor: torch.Tensor,
+        size_0: int,
+        size_1: int,
+        dtype: torch.dtype,
+    ):
+        """
+        Auxiliary function to check whether input is valid.
+        """
+        size_0_cond = (size_0 == -1 or tensor.size(0) == size_0)
+        size_1_cond = (size_1 == -1 or tensor.size(1) == size_1)
+        if not (size_0_cond and size_1_cond and tensor.is_contiguous and tensor.dtype == dtype):
+            raise ValueError(
+                f"The input to advance_step is invalid with tensor name = {name}, "
+                f"shape = {tensor.shape}, "
+                f"is_cont = {tensor.is_contiguous()}, "
+                f"type = {tensor.dtype}, "
+                f"is not as expected: shape[{size_0}, {size_1}], type = {dtype}"
+            )
+
+    verify_tensor("input_tokens", input_tokens, num_seqs, -1, torch.int64)
+    verify_tensor("sampled_token_ids", sampled_token_ids, num_queries, 1, torch.int64)
+    verify_tensor("input_positions", input_positions, num_seqs, -1, torch.int32)
+    verify_tensor("seq_lens", seq_lens, num_seqs, -1, torch.int32)
+    verify_tensor("slot_mapping", slot_mapping, num_seqs, -1, torch.int32)
+    verify_tensor("block_tables", block_tables, num_seqs, -1, torch.int32)
+
+    grid = (math.ceil(num_queries / TILE_SIZE), )
+    _triton_advance_step[grid](input_tokens,
+                               sampled_token_ids,
+                               input_positions,
+                               seq_lens,
+                               slot_mapping,
+                               block_tables,
+                               block_tables.stride(0),
+                               num_seqs,
+                               num_queries,
+                               block_size,
+                               TILE_SIZE)
+
+def preload(
+    weight: torch.Tensor,
+    size: int
+) -> None:
+    """
+    Preload weights of layer.
+
+    Args:
+        weight (torch.Tensor): Weight to preload。
+        size (int): Preload size (byte)。
+
+    Returns:
+        None
+    """
+    return tmo.preload(weight, size)
+
+
+def matmul_allreduce(
+    cncl_comm,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    c: Optional[torch.Tensor] = None,
+    alpha: float = 1.0,
+    beta: float = .0,
+    block_m: int = 0
+) -> torch.Tensor:
+    return tmo.matmul_allreduce(cncl_comm=cncl_comm,
+                                a=a, b=b,
+                                bias=bias, c=c,
+                                alpha=alpha,
+                                beta=beta,
+                                block_m=block_m)
+
+
+def smooth_quant_matmul_allreduce(
+    cncl_comm,
+    a: torch.Tensor,
+    a_scale: torch.Tensor,
+    b: torch.Tensor,
+    b_scale: torch.Tensor,
+    dtype: torch.dtype,
+    bias: torch.Tensor = None,
+    c: torch.Tensor = None,
+    alpha: float = 1.0,
+    beta: float = 1.0,
+    block_m: int = 0):
+    return tmo.smooth_quant_matmul_allreduce(
+                cncl_comm=cncl_comm,
+                a=a, a_scale=a_scale,
+                b=b, b_scale=b_scale,
+                dtype=dtype, bias=bias, c=c,
+                alpha=alpha, beta=beta, block_m=block_m)
+
+
+def quant_matmul_allreduce(
+    cncl_comm,
+    a_tensor: torch.Tensor,
+    a_scale: Optional[torch.Tensor],
+    a_zero: Optional[torch.Tensor],
+    b_tensor: torch.Tensor,
+    b_scale: Optional[torch.Tensor],
+    b_zero: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    c_tensor: Optional[torch.Tensor],
+    c_scale: Optional[torch.Tensor],
+    c_zero: Optional[torch.Tensor],
+    gemm_output_scale: Optional[torch.Tensor],
+    gemm_output_zero: Optional[torch.Tensor],
+    data_type: Optional[str],
+    quant_algo: str,
+    a_quant_layout: str,
+    b_quant_layout: str,
+    quant_bit_size: int = 8,
+    alpha: float = 1.0,
+    beta: float = 1.0,
+    trans_a: bool = False,
+    trans_b: bool = True,
+    block_m: int = 0
+) -> torch.Tensor:
+    return tmo.quant_matmul_allreduce(
+        cncl_comm=cncl_comm, a_tensor=a_tensor, a_scale=a_scale, a_zero=a_zero,
+        b_tensor=b_tensor, b_scale=b_scale, b_zero=b_zero, bias=bias,
+        c_tensor=c_tensor, c_scale=c_scale, c_zero=c_zero,
+        gemm_output_scale=gemm_output_scale, gemm_output_zero=gemm_output_zero,
+        data_type=data_type, quant_algo=quant_algo,
+        a_quant_layout=a_quant_layout, b_quant_layout=b_quant_layout,
+        quant_bit_size=quant_bit_size,
+        alpha=alpha, beta=beta, trans_a=trans_a, trans_b=trans_b, block_m=block_m)
+
+
+def flash_attn_sq_mm_allreduce(
+    cncl_comm: int,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seq_lens_q: Optional[torch.Tensor],
+    cu_seq_lens_kv: Optional[torch.Tensor],
+    alibi_slope: Optional[torch.Tensor],
+    attn_bias: Optional[torch.Tensor],
+    smooth: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    max_seq_len_q: int,
+    max_seq_len_kv: int,
+    softmax_scale: float,
+    is_causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    compute_dtype: torch.dtype = torch.float,
+    block_seq: int = 0) -> torch.Tensor:
+    return tmo.flash_attn_sq_mm_allreduce(cncl_comm, q, k, v,
+                                cu_seq_lens_q, cu_seq_lens_kv, alibi_slope, attn_bias, smooth, weight, weight_scale,
+                                bias, max_seq_len_q, max_seq_len_kv, softmax_scale, is_causal, window_size_left,
+                                window_size_right, compute_dtype, block_seq)
+
+#Moe inner kernels
+def moe_softmax_topk(input: torch.Tensor,
+                     topk: int,
+                     normalize: bool = False,
+                     num_expert_group: int = -1,
+                     topk_group: int = 0,
+                     mask: Optional[torch.Tensor] = None,
+                     normed_by : str = "topk_logit",
+                     route_scale : float = 1.0) -> Tuple[torch.Tensor]:
+    return tmo.moe_softmax_topk(input, topk, normalize, num_expert_group, topk_group, mask, normed_by, route_scale)
+
+def moe_sigmoid_topk(input: torch.Tensor,
+                     topk: int,
+                     normalize: bool = False,
+                     num_expert_group: int = -1,
+                     topk_group: int = 0,
+                     route_scale: float = 1.0,
+                     score_bias: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor]:
+    return tmo.moe_sigmoid_topk(input, topk, normalize, num_expert_group,
+                                topk_group, route_scale = route_scale,
+                                score_bias = score_bias)
+
+def moe_gen_idx(expert_id: torch.Tensor,
+                expert_num: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    return tmo.moe_gen_idx(expert_id, expert_num)
+
+def moe_expand_input(input: torch.Tensor,
+                     gather_idx: torch.Tensor,
+                     cusum_token_count: Optional[torch.Tensor] = None,
+                     start_expert_id: int = 0,
+                     expert_size: int = 0) -> torch.Tensor:
+    return tmo.moe_expand_input(input, gather_idx,
+                                cusum_token_count,
+                                start_expert_id, expert_size)
+
+def moe_active(input: torch.Tensor,
+               act_mode: str,
+               is_gated: bool,
+               output: Optional[torch.Tensor] = None,
+               bias: Optional[torch.Tensor] = None,
+               cusum_token_count: Optional[torch.Tensor] = None,
+               start_expert_id: int = 0,
+               expert_size: int = 0) -> torch.Tensor:
+    return tmo.moe_active(input, act_mode, is_gated, output,
+                          bias, cusum_token_count,
+                          start_expert_id, expert_size)
+
+def group_gemm(a: torch.Tensor,
+               b: torch.Tensor,
+               m_list: torch.Tensor,
+               expand_idx: Optional[torch.Tensor],
+               c: Optional[torch.Tensor],
+               alpha: Optional[torch.Tensor],
+               beta: Optional[torch.Tensor],
+               max_m: int = 0
+               ) -> torch.Tensor:
+    return tmo.group_gemm(a, b, m_list, expand_idx,
+                              c, alpha, beta, max_m)
+
+def smooth_quant_group_gemm(a: torch.Tensor,
+                            b: torch.Tensor,
+                            m_list: torch.Tensor,
+                            expand_idx: Optional[torch.Tensor],
+                            c: Optional[torch.Tensor],
+                            alpha: Optional[torch.Tensor],
+                            beta: Optional[torch.Tensor],
+                            a_scale: torch.Tensor,
+                            b_scale: torch.Tensor,
+                            dtype,
+                            max_m: int = 0,
+                            output: Optional[torch.Tensor] = None,
+                            ) -> torch.Tensor:
+    return tmo.smooth_quant_group_gemm(a, b, m_list, expand_idx, c, alpha, beta,
+                                       a_scale, b_scale, dtype, max_m, d=output)
+
+def moe_combine_result(input: torch.Tensor,
+                       reduce_weight: torch.Tensor,
+                       gather_ids: torch.Tensor,
+                       residual: Optional[torch.Tensor],
+                       cusum_token_count: Optional[torch.Tensor],
+                       start_expert_id: int,
+                       expert_size: int,
+                       bias: Optional[torch.Tensor] = None,
+                       output: Optional[torch.Tensor] = None,
+                       ) -> torch.Tensor:
+    return tmo.moe_combine_result(input, reduce_weight, gather_ids,
+                                  residual, cusum_token_count,
+                                  start_expert_id, expert_size, bias, output=output)
+
+def moe_quantize(x: torch.Tensor,
+                 smooth: torch.Tensor,
+                 zero: Optional[torch.Tensor] = None,
+                 token_count: Optional[torch.Tensor] = None,
+                 gather_index: Optional[torch.Tensor] = None,
+                 gather_index_start_position: Optional[torch.Tensor] = None,
+                 output: Optional[torch.Tensor] = None,
+                 output_scale: Optional[torch.Tensor] = None,
+                 dynamic_quant: bool = True,
+                 act_mode: str = "none",
+                 active_coef: float = 1.0,
+                 is_gated: bool = False,
+                 quant_type: torch.dtype = torch.int8
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+    return tmo.moe_quantize(x, smooth, zero, token_count, gather_index, gather_index_start_position,
+                            output, output_scale, dynamic_quant, act_mode, active_coef, is_gated, quant_type)
+
+
+def dequant_from_paged_cache(key: torch.Tensor,
+                             value: Optional[torch.Tensor],
+                             key_cache: torch.Tensor,
+                             value_cache: Optional[torch.Tensor],
+                             key_cache_quant_scale: torch.Tensor,
+                             value_cache_quant_scale: Optional[torch.Tensor],
+                             context_lengths: torch.Tensor,
+                             max_context_len: int,
+                             context_seq_offset: Optional[torch.Tensor],
+                             block_tables: torch.Tensor,
+                             quant_mode: int = 0,
+                             quant_bit: int = 8) -> None:
+    tmo.dequant_from_paged_cache(
+        key, value, key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale,
+        context_lengths, max_context_len, context_seq_offset, block_tables, quant_mode, quant_bit)
+
+def scaled_quantize(
+    input: torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+    zero: Optional[torch.Tensor] = None,
+    scale_ub: Optional[torch.Tensor] = None,
+    quant_type: torch.dtype = torch.int8,
+    quant_mode: str = "dynamic_per_token"
+) -> Tuple[torch.Tensor]:
+    """
+    Apply quantization to the input tensor input.
+
+    Args:
+        input (torch.Tensor): The tensor to be quantized. Shape is (..., C).
+        scale (Optional[torch.Tensor], optional): The scale multipled to the input tensor.  Shape is (C) or (1).
+        zero (Optional[torch.Tensor], optional):  Not supported, must pass None.
+        scale_ub (Optional[torch.Tensor], optional): The output_scale upper bound.
+            Take effect only if quant_type == torch.float8_e4m3fn and quant_mode == "dynamic_per_token".
+            Value out of range will be truncated.
+        quant_type (optional): Output data type, can be torch.int8, torch.float8_e4m3fn. Defaults to torch.int8.
+        quant_mode (str, optional): quantize mode, which can be:
+          - "dynamic_per_token"
+          - "static_per_tensor"
+          - "static_per_channel"
+
+    Type:
+        input: float, half, bfloat16.
+        scale: float.
+        scale_ub: float.
+
+    Returns:
+        Tuple[torch.Tensor]: Returns (output, output_scale) if quant_mode == "dynamic_per_token",
+        otherwise returns output only.
+    """
+    return tmo.scaled_quantize(input, scale, zero, scale_ub, quant_type=quant_type, quant_mode=quant_mode)
+
+def scaled_matmul(a: torch.Tensor,
+                  b: torch.Tensor,
+                  a_scale: Optional[torch.Tensor],
+                  b_scale: torch.Tensor,
+                  output_dtype: torch.dtype,
+                  bias: torch.Tensor = None,
+                  c: torch.Tensor = None,
+                  act_mode: str = "none",
+                  quant_bit_size: int = 8,
+                  alpha: float = 1.0,
+                  beta: float = 1.0,
+                  use_hp_active: bool = False):
+    """
+    Perform quantized matrix multiplication on tensor a and b.
+
+    Args:
+        a (torch.Tensor): Shape is (M, K).
+        b (torch.Tensor): If quant_bit_size = 8, shape is (N, K).
+                          If quant_bit_size = 4, shape is (N, K//2).
+        a_scale (Optional[torch.Tensor]): Shape can be (M).
+        b_scale (torch.Tensor): If use groupwise quantization, shape must be (N, group_num), data type must be
+            the same as a; otherwise shape must be (N), data type must be float.
+        output_dtype (torch.dtype): Specify the data type of output, must be torch.half or torch.bfloat16.
+        bias (torch.Tensor, optional): Shape is (N).
+        c (torch.Tensor, optional): Shape is (M, N).
+        act_mode (str, optional): Choose the activation algorithm, must be 'silu', 'gelu' or 'none'. If use groupwise
+            quantization, act_mode must be 'none'.
+        quant_bit_size (int, optional): The data format of b. Defaults to 8.
+        alpha (float, optional): coefficient of acted. Defaults to 1.0.
+        beta (float, optional): coefficient of c. Defaults to 1.0.
+        use_hp_active (bool, optional): Describing the algorithm that used in the implementation of the activation function.
+             When the value is true, use the high-precision algorithm, otherwise use the fastest algorithm of activation.
+             Defaults to False.
+
+    Type:
+        a: int8, half, bfloat16, float8_e4m3fn
+        a_scale: float
+        b: int8, float8_e4m3fn
+        b_scale: float
+        bias: same as output
+        c: same as output
+        output: specified by output_dtype
+
+    Returns:
+        A tensor with the shape of (M, N).
+    """
+    return tmo.scaled_matmul(a,
+                             b,
+                             a_scale,
+                             b_scale,
+                             output_dtype,
+                             bias,
+                             c,
+                             act_mode,
+                             quant_bit_size,
+                             alpha,
+                             beta,
+                             use_hp_active)
+
+def fused_mla_kv(kv: torch.Tensor,
+                 sin: torch.Tensor,
+                 cos: torch.Tensor,
+                 position_id: torch.Tensor,
+                 gamma: torch.Tensor,
+                 kv_cache: torch.Tensor,
+                 kv_cache_scale: Optional[torch.Tensor],
+                 slot_mapping: Optional[torch.Tensor],
+                 cache_bs_id: Optional[torch.Tensor] = None,
+                 cache_seq_offset: Optional[torch.Tensor] = None,
+                 is_paged_cache: bool = True,
+                 eps: float = 1e-5,
+                 interleaved: bool = True):
+    return tmo.fused_mla_kv(kv, sin, cos, position_id, gamma, kv_cache, kv_cache_scale, slot_mapping, cache_bs_id,
+                            cache_seq_offset, is_paged_cache, eps, interleaved)
+
+def fused_mla_q(q: torch.Tensor,
+                gamma: torch.Tensor,
+                smooth_quant_scale: torch.Tensor,
+                weight_b: torch.Tensor,
+                weight_b_scale: torch.Tensor,
+                weight_c: torch.Tensor,
+                sin: torch.Tensor,
+                cos: torch.Tensor,
+                position_id: torch.Tensor,
+                output: Optional[torch.Tensor] = None,
+                eps: float = 1e-6,
+                interleaved: bool = True) -> torch.Tensor:
+    return tmo.fused_mla_q(q, gamma, smooth_quant_scale, weight_b, weight_b_scale, weight_c, sin, cos, position_id,
+                           output, eps, interleaved)
+
+
+def gather_cache(
+    kv_cache: List[torch.Tensor],      # [[1, num_blocks, num_kv_heads, block_size, head_size]
+                                       #  [1, num_blocks, num_kv_heads, block_size] if kv_cache_dtype=int8]
+    dst: torch.Tensor,                 # [tot_tokens, entrys...]
+    block_table: torch.Tensor,         # [batch, block_indices]
+    cu_seq_lens: torch.Tensor,         # [batch+1]
+    batch_size: int,
+    seq_starts: torch.Tensor = None,   # Optional: [batch]
+    kv_cache_dtype: str = 'auto',
+) -> None:
+    """
+    Gathers sequences from src_cache into dst based on block_table and cu_seq_lens.
+
+    Args:
+        src_cache: Source KV cache tensor of shape [[1, num_blocks, num_kv_heads, block_size, head_size],
+          [1, num_blocks, num_kv_heads, block_size] if cache_dtype=int8].
+        dst: Destination tensor of shape [tot_tokens, entrys...].
+        block_table: Tensor of shape [batch, block_indices] mapping sequences to blocks.
+        cu_seq_lens: Tensor of shape [batch+1] with cumulative sequence lengths.
+        batch_size: Number of sequences in the batch.
+        seq_starts: Optional tensor of shape [batch] for block index offsets.
+    """
+    assert len(kv_cache) > 0 and kv_cache[0].numel() > 0, "kv cache can't be empty in gather_cache"
+    src_cache = kv_cache[0][0]
+    # Validate inputs
+    assert src_cache.device == dst.device == block_table.device == cu_seq_lens.device, \
+        "All tensors must be on the same device"
+    assert block_table.dtype == torch.int32, "block_table must be int32"
+    assert cu_seq_lens.dtype == torch.int32, "cu_seq_lens must be int32"
+    if kv_cache_dtype != 'int8':
+        assert src_cache.dtype == dst.dtype, "src_cache and dst must have the same dtype when no quantized"
+    if seq_starts is not None:
+        assert seq_starts.dtype == torch.int32, "seq_starts must be int32"
+        assert seq_starts.device == src_cache.device, "seq_starts must be on the same device"
+
+    # Extract dimensions
+    num_blocks, num_kv_heads, block_size, head_size = src_cache.shape
+    # When using MLA during decode it becomes MQA, the num_kv_heads is fixed to 1,
+    # so src_cache can be view to [num_blocks, block_size, head_size]
+    assert num_kv_heads == 1, "mla force num_kv_heads to 1"
+    src_cache = src_cache.view(num_blocks, block_size, -1)
+    entry_shape = src_cache.shape[2:]  # ENTRIES...
+    tot_tokens = cu_seq_lens[-1]
+    assert tot_tokens > 0, "tot_tokens should > 0"
+    assert tot_tokens <= dst.shape[0], "tot_tokens should <= dst.shape[0]"
+    dst_cache = dst[:tot_tokens]
+
+    # Ensure cu_seq_lens matches batch_size
+    assert cu_seq_lens.size(0) == batch_size + 1, "cu_seq_lens must have batch_size + 1 elements"
+
+    # Compute sequence lengths
+    seq_lens = cu_seq_lens[1:] - cu_seq_lens[:-1]  # [BATCH]
+    tot_blocks_per_seq = (seq_lens + block_size - 1) // block_size  # ceil_div
+
+    # Handle seq_starts offset
+    block_offsets = torch.zeros(batch_size, dtype=torch.int32, device=src_cache.device)
+    if seq_starts is not None:
+        block_offsets = seq_starts // block_size
+
+    # Flatten src_cache for easier indexing: [NUM_BLOCKS * BLOCK_SIZE, ENTRIES...]
+    src_flat = src_cache.view(num_blocks * block_size, *entry_shape)
+
+    # Prepare output indices
+    dst_indices = []
+    for bid in range(batch_size):
+        seq_len = seq_lens[bid]
+        if seq_len <= 0:
+            continue
+        seq_start = cu_seq_lens[bid]
+        tot_blocks = tot_blocks_per_seq[bid]
+        offset = block_offsets[bid]
+
+        # Compute block indices for this sequence
+        block_ids = block_table[bid, offset:offset + tot_blocks]
+
+        # Compute token indices within blocks
+        token_indices = torch.arange(seq_len, device=src_cache.device)
+        block_indices = token_indices // block_size
+        within_block = token_indices % block_size
+
+        # Map to src_flat indices
+        src_indices = block_ids[block_indices] * block_size + within_block
+        dst_indices.append(src_indices + seq_start)
+
+    # Concatenate all indices
+    dst_indices = torch.cat(dst_indices)
+
+    # Gather data
+    dst_flat = src_flat[dst_indices]
+    if kv_cache_dtype == 'int8':
+        src_cache_scale = kv_cache[1][0]
+        src_scale_flat = src_cache_scale.view(num_blocks * block_size)
+        dst_scale_flat = src_scale_flat[dst_indices]
+        dst_flat = dst_flat * dst_scale_flat.unsqueeze(-1)
+
+    dst_cache.view(-1, *entry_shape).copy_(dst_flat.view(tot_tokens, *entry_shape))
+
+
+def extract_uncausal_lse(attn_softmax_lse: torch.Tensor, cu_seq_lens: torch.Tensor) -> torch.Tensor:
+    """
+    Extract valid LSE values from attn_softmax_lse using cu_seq_lens and concatenate them to shape (nheads, total_seqlen).
+
+    Args:
+        attn_softmax_lse (torch.Tensor): Shape (batch_size, nheads, seqlen), LSE values from flash_attn_varlen_func.
+        cu_seq_lens (torch.Tensor): Shape (batch_size + 1,), cumulative sequence lengths.
+
+    Returns:
+        torch.Tensor: Shape (nheads, total_seqlen), concatenated valid LSE values for each head.
+    """
+    # Get dimensions
+    batch_size, nheads, seqlen = attn_softmax_lse.shape
+
+    # Compute sequence lengths
+    seq_lens = cu_seq_lens[1:] - cu_seq_lens[:-1]  # Shape: (batch_size,)
+
+    # Validate inputs
+    assert seq_lens.max() <= seqlen, "Sequence length exceeds attn_softmax_lse seqlen dimension."
+    assert len(cu_seq_lens) == batch_size + 1, "cu_seq_lens length must be batch_size + 1."
+
+    # Create a mask for valid positions: (batch_size, seqlen)
+    mask = torch.arange(seqlen, device=attn_softmax_lse.device).unsqueeze(0) < seq_lens.unsqueeze(1)
+
+    # Reshape attn_softmax_lse to (batch_size * seqlen, nheads) for easier indexing
+    lse_flat = attn_softmax_lse.permute(0, 2, 1).reshape(batch_size * seqlen, nheads)
+
+    # Flatten the mask to (batch_size * seqlen,)
+    mask_flat = mask.view(batch_size * seqlen)
+
+    # Select valid LSE values: (total_seqlen, nheads)
+    valid_lse = lse_flat[mask_flat]
+
+    # Reshape to (nheads, total_seqlen)
+    result_lse = valid_lse.t()  # Transpose: (nheads, total_seqlen)
+
+    assert result_lse.shape[-1] == cu_seq_lens[-1], "total lse seq should equal query len"
+
+    return result_lse
+
+
+def merge_attn_states(
+    output: torch.Tensor,
+    prefix_output: torch.Tensor,
+    prefix_lse: torch.Tensor,
+    suffix_output: torch.Tensor,
+    suffix_lse: torch.Tensor,
+    output_lse: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    Merges partial attention states (prefix and suffix) into a single output.
+    Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005.
+
+    Args:
+        output: Output tensor of shape [num_tokens, num_query_heads, head_size].
+        prefix_output: Prefix attention output, same shape as output.
+        prefix_lse: Prefix log-sum-exp, shape [num_query_heads, num_tokens].
+        suffix_output: Suffix attention output, same shape as output.
+        suffix_lse: Suffix log-sum-exp, same shape as prefix_lse.
+        output_lse: Optional output log-sum-exp, same shape as prefix_lse.
+    """
+    # Input validation
+    assert output.shape == prefix_output.shape == suffix_output.shape, \
+        "Output and input tensors must have the same shape"
+    assert prefix_lse.shape == suffix_lse.shape, \
+        "Prefix and suffix LSE tensors must have the same shape"
+    if output_lse is not None:
+        assert output_lse.shape == prefix_lse.shape, \
+            "Output LSE must have the same shape as input LSE tensors"
+
+    # Handle inf values (replace inf with -inf for consistency)
+    p_lse = torch.where(
+        prefix_lse == float('inf'),
+        torch.tensor(float('-inf'), device=prefix_lse.device),
+        prefix_lse
+    )
+    s_lse = torch.where(
+        suffix_lse == float('inf'),
+        torch.tensor(float('-inf'), device=suffix_lse.device),
+        suffix_lse
+    )
+
+    # Compute maximum LSE for numerical stability
+    max_lse = torch.maximum(p_lse, s_lse)  # Shape: [num_query_heads, num_tokens]
+
+    # Normalize LSE terms
+    p_lse = p_lse - max_lse  # Shape: [num_query_heads, num_tokens]
+    s_lse = s_lse - max_lse  # Shape: [num_query_heads, num_tokens]
+
+    # Compute sum of exponentials
+    out_se = torch.exp(p_lse) + torch.exp(s_lse)  # Shape: [num_query_heads, num_tokens]
+
+    # Compute output_lse if provided
+    if output_lse is not None:
+        output_lse.copy_(torch.log(out_se) + max_lse)
+
+    # Compute scaling factors
+    p_scale = torch.exp(p_lse) / out_se  # Shape: [num_query_heads, num_tokens]
+    s_scale = torch.exp(s_lse) / out_se  # Shape: [num_query_heads, num_tokens]
+
+    # Reshape scales for broadcasting
+    p_scale = p_scale.unsqueeze(-1)  # Shape: [num_query_heads, num_tokens, 1]
+    s_scale = s_scale.unsqueeze(-1)  # Shape: [num_query_heads, num_tokens, 1]
+
+    # Transpose outputs to match scaling dimensions
+    prefix_output = prefix_output.permute(1, 0, 2)  # Shape: [num_query_heads, num_tokens, head_size]
+    suffix_output = suffix_output.permute(1, 0, 2)  # Shape: [num_query_heads, num_tokens, head_size]
+
+    # Compute merged output
+    out = prefix_output * p_scale + suffix_output * s_scale  # Shape: [num_query_heads, num_tokens, head_size]
+
+    # Transpose back and store in output
+    output.copy_(out.permute(1, 0, 2))  # Shape: [num_tokens, num_query_heads, head_size]
diff --git a/vllm_mlu/vllm_mlu/_mlu_utils.py b/vllm_mlu/vllm_mlu/_mlu_utils.py
new file mode 100644
index 000000000..7d02d142e
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/_mlu_utils.py
@@ -0,0 +1,86 @@
+import os
+import torch
+
+
+def _check_env(env, default=False):
+    if env in os.environ:
+        return os.environ[env].lower() in ["true", "1"]
+    return default
+
+
+def _check_env_value(env, default=0):
+    if env in os.environ:
+        if not os.environ[env].isdigit():
+            raise ValueError(f"'{env}' should be set with integer")
+        value = int(os.environ[env])
+        return value
+    return default
+
+
+# VLLM_LATENCY_DEBUG: Get more kernel info for benchmark latency.
+VLLM_LATENCY_DEBUG = _check_env("VLLM_LATENCY_DEBUG", default=False)
+
+# VLLM_LATENCY_DEBUG_NO_DEVICE: Get more kernel info(without device) for benchmark latency.
+VLLM_LATENCY_DEBUG_NO_DEVICE = _check_env("VLLM_LATENCY_DEBUG_NO_DEVICE", default=False)
+
+# VLLM_DUMP_TENSORS: Dump each layer outputs when running vLLM inference.
+VLLM_DUMP_OUTPUTS = _check_env("VLLM_DUMP_OUTPUTS", default=False)
+
+# VLLM_DUMP_MLU_INFO: Get device info when running vLLM inference.
+VLLM_DUMP_MLU_INFO = _check_env("VLLM_DUMP_MLU_INFO", default=False)
+
+# VLLM_SCHEDULER_PROFILE: Profiling vLLM scheduler.
+VLLM_SCHEDULER_PROFILE = _check_env("VLLM_SCHEDULER_PROFILE", default=False)
+
+# VLLM_GRAPH_DEBUG: Debug the graph status when running decoder, default value is True.
+# Set to False to disable warning messages.
+VLLM_GRAPH_DEBUG = _check_env("VLLM_GRAPH_DEBUG", default=True)
+
+# CHUNKED_PIPELINE_PARALLEL_EN: use chunked pipeline parallel, default value is False.
+CHUNKED_PIPELINE_PARALLEL_EN = _check_env("CHUNKED_PIPELINE_PARALLEL_EN", default=False)
+
+# CONTEXT_PARALLEL_EN: use context parallel, default value is False.
+CONTEXT_PARALLEL_EN = _check_env("CONTEXT_PARALLEL_EN", default=False)
+
+# EXPERT_PARALLEL_EN: use expert parallel, default value is False.
+EXPERT_PARALLEL_EN = _check_env("EXPERT_PARALLEL_EN", default=False)
+
+# ATTN_DATA_PARALLEL_EN: use attn data parallel, default value is False.
+ATTN_DATA_PARALLEL_EN = _check_env("ATTN_DATA_PARALLEL_EN", default=False)
+
+# VLLM_AVG_MOE_EN: make moe experts workload balance, default value is False.
+VLLM_AVG_MOE_EN = _check_env("VLLM_AVG_MOE_EN", default=False)
+
+# VLLM_LOGITS_USE_ALL_GATHER: use allgather for logits collection, default value is False.
+VLLM_LOGITS_USE_ALL_GATHER = _check_env("VLLM_LOGITS_USE_ALL_GATHER", default=False)
+
+VLLM_LATENCY_DEBUG_EN = (VLLM_LATENCY_DEBUG or VLLM_LATENCY_DEBUG_NO_DEVICE)
+VLLM_LATENCY_DEBUG_WITH_DEVICE_EN = (VLLM_LATENCY_DEBUG and not VLLM_LATENCY_DEBUG_NO_DEVICE)
+VLLM_DUMP_MLU_INFO_EN = (VLLM_LATENCY_DEBUG_WITH_DEVICE_EN and VLLM_DUMP_MLU_INFO)
+CUSTOM_VLLM_HIJACK_EN = (CHUNKED_PIPELINE_PARALLEL_EN or CONTEXT_PARALLEL_EN
+                         or EXPERT_PARALLEL_EN or ATTN_DATA_PARALLEL_EN)
+
+VLLM_PRELOAD_SIZE = _check_env_value("VLLM_PRELOAD_SIZE", default=0)
+
+# ATTN_PARALLEL_NUM & FFN_PARALLEL_NUM: use context comm cmpt parallel.
+ATTN_PARALLEL_NUM = 'ATTN_PARALLEL_NUM'
+FFN_PARALLEL_NUM = 'FFN_PARALLEL_NUM'
+
+
+def check_context_comm_cmpt_parallel():
+    return (ATTN_PARALLEL_NUM in os.environ) or (FFN_PARALLEL_NUM in os.environ)
+
+
+def set_is_prompt(flag):
+    global IS_PROMPT
+    IS_PROMPT=flag
+
+
+def get_is_prompt():
+    return IS_PROMPT
+
+def get_dpsk_mcc_max_parallel_num():
+    return _check_env_value("DPSK_MCC_PARALLEL_NUM", default=1)
+
+def is_dpsk_mcc_enabled():
+    return get_dpsk_mcc_max_parallel_num() > 1
diff --git a/vllm_mlu/vllm_mlu/attention/__init__.py b/vllm_mlu/vllm_mlu/attention/__init__.py
new file mode 100644
index 000000000..8b1378917
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/attention/__init__.py
@@ -0,0 +1 @@
+
diff --git a/vllm_mlu/vllm_mlu/attention/backends/__init__.py b/vllm_mlu/vllm_mlu/attention/backends/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/attention/backends/mlu_flash_attn.py b/vllm_mlu/vllm_mlu/attention/backends/mlu_flash_attn.py
new file mode 100644
index 000000000..3111d59d9
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/attention/backends/mlu_flash_attn.py
@@ -0,0 +1,1571 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from contextlib import contextmanager
+from itertools import accumulate
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
+
+import torch
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm.attention.backends.abstract import (AttentionLayer,
+                                              AttentionMetadata,
+                                              AttentionType)
+from vllm.attention.backends.utils import (
+    PAD_SLOT_ID, get_num_prefill_decode_query_kv_tokens,
+    get_seq_len_block_table_args)
+from vllm.utils import (async_tensor_h2d,
+                        make_tensor_with_pad)
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
+                                              AttentionLayer,
+                                              AttentionMetadata,
+                                              AttentionMetadataBuilder,
+                                              AttentionType)
+from vllm.attention.backends.utils import (
+    PAD_SLOT_ID, CommonAttentionState, compute_slot_mapping,
+    compute_slot_mapping_start_idx, get_num_prefill_decode_query_kv_tokens,
+    get_seq_len_block_table_args, is_all_cross_attn_metadata_set,
+    is_all_encoder_attn_metadata_set, is_block_tables_empty)
+from vllm.multimodal import MultiModalPlaceholderMap
+from vllm.utils import async_tensor_h2d, make_tensor_with_pad, cdiv, round_down
+
+
+if TYPE_CHECKING:
+    from vllm.worker.mlu_model_runner import (ModelInputForMLUBuilder,
+                                              ModelInputForGPUWithSamplingMetadata)
+    from vllm.worker.model_runner_base import ModelRunnerBase
+
+class MLUFlashAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> List[int]:
+        return [32, 64, 80, 96, 128, 160, 192, 224, 256, 512, 576]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> Type["MLUFlashAttentionImpl"]:
+        return MLUFlashAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> Type["MLUFlashAttentionMetadata"]:
+        return MLUFlashAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> Type["MLUFlashAttentionMetadataBuilder"]:
+        return MLUFlashAttentionMetadataBuilder
+
+    @staticmethod
+    def get_state_cls() -> Type["MLUAttentionState"]:
+        return MLUAttentionState
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        return (2, num_blocks, num_kv_heads, block_size, head_size)
+
+    @staticmethod
+    def get_kv_cache_scale_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+    ) -> Tuple[int, ...]:
+        return (2, num_blocks, num_kv_heads, block_size)
+
+    @staticmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: torch.Tensor,
+    ) -> None:
+        src_key_cache = src_kv_cache[0]
+        dst_key_cache = dst_kv_cache[0]
+        mlu_ops.swap_blocks(dst_key_cache, src_key_cache, src_to_dst)
+
+        src_value_cache = src_kv_cache[1]
+        dst_value_cache = dst_kv_cache[1]
+        mlu_ops.swap_blocks(dst_value_cache, src_value_cache, src_to_dst)
+
+    @staticmethod
+    def copy_blocks(
+        kv_caches: List[List[torch.Tensor]],
+        src_to_dists: torch.Tensor,
+    ) -> None:
+        key_caches = [kv_cache[0][0] for kv_cache in kv_caches]
+        value_caches = [kv_cache[0][1] for kv_cache in kv_caches]
+        mlu_ops.copy_blocks(key_caches, value_caches, src_to_dists)
+
+        kv_cache_scales = [kv_cache[1] for kv_cache in kv_caches]
+        if len(kv_cache_scales) > 0 and kv_cache_scales[0].numel() > 0:
+            key_cache_scales = [kv_cache_scale[0] for kv_cache_scale in kv_cache_scales]
+            value_cache_scales = [kv_cache_scale[1] for kv_cache_scale in kv_cache_scales]
+            mlu_ops.copy_blocks(key_cache_scales, value_cache_scales, src_to_dists)
+
+
+class MLUMLAFlashAttentionBackend(MLUFlashAttentionBackend):
+
+    @staticmethod
+    def get_state_cls() -> Type["MLUMLAAttentionState"]:
+        return MLUMLAAttentionState
+
+    @staticmethod
+    def get_builder_cls() -> Type["MLUMLAFlashAttentionMetadataBuilder"]:
+        return MLUMLAFlashAttentionMetadataBuilder
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        return (1, num_blocks, num_kv_heads, block_size, head_size)
+
+    @staticmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: torch.Tensor,
+    ) -> None:
+        src_key_cache = src_kv_cache[0]
+        dst_key_cache = dst_kv_cache[0]
+        mlu_ops.swap_blocks(dst_key_cache, src_key_cache, src_to_dst)
+
+    @staticmethod
+    def get_kv_cache_scale_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+    ) -> Tuple[int, ...]:
+        return (1, num_blocks, num_kv_heads, block_size)
+
+    @staticmethod
+    def copy_blocks(
+        kv_caches: List[List[torch.Tensor]],
+        src_to_dists: torch.Tensor,
+    ) -> None:
+        key_caches = [kv_cache[0][0] for kv_cache in kv_caches]
+        mlu_ops.copy_blocks(key_caches, None, src_to_dists)
+
+        kv_cache_scales = [kv_cache[1] for kv_cache in kv_caches]
+        if len(kv_cache_scales) > 0 and kv_cache_scales[0].numel() > 0:
+            key_cache_scales = [kv_cache_scale[0] for kv_cache_scale in kv_cache_scales]
+            mlu_ops.copy_blocks(key_cache_scales, None, src_to_dists)
+
+
+class MLUAttentionState(CommonAttentionState):
+
+    @contextmanager
+    def graph_capture(self, max_batch_size: int):
+
+        self._is_graph_capturing = True
+
+        self._graph_slot_mapping = torch.full((max_batch_size, ),
+                                              PAD_SLOT_ID,
+                                              dtype=torch.int32,
+                                              device=self.runner.device)
+        self._graph_seq_lens = torch.ones(max_batch_size,
+                                          dtype=torch.int32,
+                                          device=self.runner.device)
+        self._graph_block_tables = torch.from_numpy(
+            self.runner.graph_block_tables).to(device=self.runner.device)
+
+        yield
+
+        self._is_graph_capturing = False
+        del self._graph_slot_mapping
+        del self._graph_seq_lens
+        del self._graph_block_tables
+
+    def get_graph_input_buffers(
+            self,
+            attn_metadata: AttentionMetadata,
+            is_encoder_decoder_model: bool = False) -> Dict[str, Any]:
+        input_buffers = {
+            "slot_mapping": attn_metadata.slot_mapping,
+            "seq_lens_tensor": None,
+            "block_tables": None,
+        }
+        if attn_metadata.num_prefills > 0:
+            input_buffers["seq_lens_tensor"] = attn_metadata.prefill_metadata.seq_lens_tensor
+            input_buffers["block_tables"] = attn_metadata.prefill_metadata.block_tables
+        else:
+            input_buffers["seq_lens_tensor"] = attn_metadata.decode_metadata.seq_lens_tensor
+            input_buffers["block_tables"] = attn_metadata.decode_metadata.block_tables
+        if is_encoder_decoder_model:
+            # The encoder decoder model works only with XFormers and
+            # Flash Attention backend. Assert the same.
+            assert self.runner.attn_backend.get_name() in\
+                ["XFORMERS", "FLASH_ATTN"], \
+                f"Expected attn_backend name to be either 'XFORMERS' or "\
+                f"'FLASH_ATTN', but "\
+                f"got '{self.runner.attn_backend.get_name()}'"
+            self._add_additonal_input_buffers_for_enc_dec_model(
+                attn_metadata=attn_metadata, input_buffers=input_buffers)
+        return input_buffers
+
+    def prepare_graph_input_buffers(
+            self,
+            input_buffers: Dict[str, Any],
+            attn_metadata: AttentionMetadata,
+            is_encoder_decoder_model: bool = False) -> None:
+        metadata = attn_metadata.prefill_metadata if \
+            attn_metadata.num_prefills > 0 else attn_metadata.decode_metadata
+
+        input_buffers["seq_lens_tensor"].copy_(
+            metadata.seq_lens_tensor, non_blocking=True)
+        input_buffers["block_tables"].copy_(
+            metadata.block_tables, non_blocking=True)
+        if is_encoder_decoder_model:
+            # The encoder decoder model works only with XFormers and
+            # Flash Attention backend. Assert the same.
+            assert self.runner.attn_backend.get_name() in\
+                ["XFORMERS", "FLASH_ATTN"], \
+                f"Expected attn_backend name to be either 'XFORMERS' or "\
+                f"'FLASH_ATTN', but "\
+                f"got '{self.runner.attn_backend.get_name()}'"
+            self._prepare_input_buffers_for_enc_dec_model(
+                attn_metadata, input_buffers)
+
+    @contextmanager
+    def graph_capture_with_context(
+        self,
+        ctx_graph_batch_size: int,
+        max_batch_size: int,
+        max_num_tokens: int
+    ):
+        self._is_graph_capturing = True
+        self._graph_slot_mapping = torch.full((max_num_tokens, ),
+                                            PAD_SLOT_ID,
+                                            dtype=torch.int32,
+                                            device=self.runner.device)
+        self._graph_seq_lens = torch.ones(max_batch_size,
+                                        dtype=torch.int32,
+                                        device=self.runner.device)
+        # block tables used for decode mlugraph input buffer
+        self._graph_block_tables = torch.from_numpy(
+            self.runner.graph_block_tables).to(device=self.runner.device)
+        # block tables used for context mlugraph input buffer
+        self._ctx_graph_block_tables = torch.zeros((ctx_graph_batch_size, 0),
+                                                    dtype=self._graph_block_tables.dtype,
+                                                    device=self.runner.device)
+        yield
+        self._is_graph_capturing = False
+        del self._graph_slot_mapping
+        del self._graph_seq_lens
+        del self._graph_block_tables
+        del self._ctx_graph_block_tables
+
+    def fill_seq_lens_tensor(
+        self,
+        seq_len: int
+    ) -> None:
+        self._graph_seq_lens.fill_(seq_len)
+
+    def graph_capture_get_metadata_for_context(
+        self,
+        batch_size: int,
+        seq_len: int,
+        is_encoder_decoder_model: bool = False
+    ):
+        assert self._is_graph_capturing
+
+        query_start_loc = torch.zeros(batch_size + 1,
+                                    dtype=torch.int32,
+                                    device=self.runner.device)
+        seq_start_loc = torch.zeros(batch_size + 1,
+                                    dtype=torch.int32,
+                                    device=self.runner.device)
+        context_lens_tensor = torch.zeros(batch_size,
+                                        dtype=torch.int32,
+                                        device=self.runner.device)
+        torch.cumsum(self._graph_seq_lens[:batch_size],
+                    dim=0,
+                    dtype=query_start_loc.dtype,
+                    out=query_start_loc[1:])
+        torch.cumsum(self._graph_seq_lens[:batch_size],
+                    dim=0,
+                    dtype=seq_start_loc.dtype,
+                    out=seq_start_loc[1:])
+
+        num_tokens = batch_size * seq_len
+        attn_metadata = self.runner.attn_backend.make_metadata(
+            num_prefills=batch_size,
+            num_prefill_tokens=num_tokens,
+            num_decode_tokens=0,
+            slot_mapping=self._graph_slot_mapping[:num_tokens],
+            multi_modal_placeholder_index_maps=None,
+            enable_kv_scales_calculation=True,
+            seq_lens=[seq_len] * batch_size,
+            seq_lens_tensor=self._graph_seq_lens[:batch_size],
+            max_query_len=seq_len,
+            max_decode_query_len=0,
+            max_prefill_seq_len=seq_len,
+            max_decode_seq_len=0,
+            query_start_loc=query_start_loc,
+            seq_start_loc=seq_start_loc,
+            context_lens_tensor=context_lens_tensor,
+            block_tables=self._ctx_graph_block_tables,
+            use_cuda_graph=True,
+            chunked_prefill_enabled=self.runner.scheduler_config.chunked_prefill_enabled,
+        )
+        return attn_metadata
+
+
+@dataclass
+class MLUFlashAttentionMetadata(AttentionMetadata):
+    """Metadata for FlashAttentionBackend.
+
+    NOTE: Any python object stored here is not updated when it is
+    cuda-graph replayed. If you have values that need to be changed
+    dynamically, it should be stored in tensor. The tensor has to be
+    updated from `CUDAGraphRunner.forward` API.
+    """
+    # (batch_size,). The sequence length per sequence. Sequence length means
+    # the computed tokens + new tokens None if it is a decoding.
+    seq_lens: Optional[List[int]]
+    # seq_lens stored as a tensor.
+    seq_lens_tensor: Optional[torch.Tensor]
+
+    # NOTE(sang): Definition of context_len, query_len, and seq_len.
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seq_len ---------------------|
+    #                                   |-- query_len ---|
+
+    # Maximum sequence length among prefill batch. 0 if there are decoding
+    # requests only.
+    max_prefill_seq_len: int
+    # Maximum sequence length among decode batch. 0 if there are prefill
+    # requests only.
+    max_decode_seq_len: int
+    # (batch_size,) A tensor of context lengths (tokens that are computed
+    # so far).
+    context_lens_tensor: Optional[torch.Tensor]
+
+    # (batch_size, max_blocks_per_seq).
+    # Block addresses per sequence. (Seq id -> list of physical block)
+    # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
+    # in the kv cache. Each block can contain up to block_size tokens.
+    # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
+    # captured.
+    block_tables: Optional[torch.Tensor]
+
+    # Whether or not if cuda graph is enabled.
+    # Cuda-graph is currently enabled for decoding only.
+    # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
+
+    use_cuda_graph: bool
+
+    # Maximum query length in the batch.
+    max_query_len: Optional[int] = None
+
+    # Max number of query tokens among request in the batch.
+    max_decode_query_len: Optional[int] = None
+
+    # (batch_size + 1,). The cumulative subquery lengths of the sequences in
+    # the batch, used to index into subquery. E.g., if the subquery length
+    # is [4, 6], it is [0, 4, 10].
+    query_start_loc: Optional[torch.Tensor] = None
+    # (batch_size + 1,). The cumulative sequence lengths of the sequences in
+    # the batch, used to index into sequence. E.g., if the sequence length is
+    # [4, 6], it is [0, 4, 10].
+    seq_start_loc: Optional[torch.Tensor] = None
+
+    _cached_prefill_metadata: Optional["MLUFlashAttentionMetadata"] = None
+    _cached_decode_metadata: Optional["MLUFlashAttentionMetadata"] = None
+
+    # Begin encoder attn & enc/dec cross-attn fields...
+
+    # Encoder sequence lengths representation
+    encoder_seq_lens: Optional[List[int]] = None
+    encoder_seq_lens_tensor: Optional[torch.Tensor] = None
+    # (batch_size + 1,). The cumulative sequence lengths of the sequences in
+    # the batch, used to index into sequence. E.g., if the sequence length is
+    # [4, 6], it is [0, 4, 10].
+    encoder_seq_start_loc: Optional[torch.Tensor] = None
+    # Maximum sequence length among encoder sequences
+    max_encoder_seq_len: Optional[int] = None
+    # Number of tokens input to encoder
+    num_encoder_tokens: Optional[int] = None
+
+    # Cross-attention memory-mapping data structures: slot mapping
+    # and block tables
+    cross_slot_mapping: Optional[torch.Tensor] = None
+    cross_block_tables: Optional[torch.Tensor] = None
+
+    # Chuned prefill enabled
+    chunked_prefill_enabled: Optional[bool] = False
+
+    # FA compute data type
+    compute_dtype: Optional[torch.dtype] = torch.float32
+
+    is_profile_run: bool = False
+
+    # New for MLA (compared to FlashAttention)
+    # For chunked prefill
+    context_chunk_cu_seq_lens: Optional[torch.Tensor] = None
+    context_chunk_starts: Optional[torch.Tensor] = None
+    context_chunk_seq_tot: Optional[List[int]] = None
+    context_chunk_max_seq_lens: Optional[List[int]] = None
+    page_size: Optional[int] = 16
+    # Set by MLUMLAAttentionState in `begin_forward` so it doesn't get broadcasted
+    context_chunk_workspace: Optional[torch.Tensor] = None
+    context_chunk_workspace_size: Optional[int] = 0
+
+    @property
+    def is_all_encoder_attn_metadata_set(self):
+        '''
+        All attention metadata required for encoder attention is set.
+        '''
+        return is_all_encoder_attn_metadata_set(self)
+
+    @property
+    def is_all_cross_attn_metadata_set(self):
+        '''
+        All attention metadata required for enc/dec cross-attention is set.
+
+        Superset of encoder attention required metadata.
+        '''
+        return is_all_cross_attn_metadata_set(self)
+
+    @property
+    def prefill_metadata(self) -> Optional["MLUFlashAttentionMetadata"]:
+        if self.num_prefills == 0:
+            return None
+
+        if self._cached_prefill_metadata is not None:
+            return self._cached_prefill_metadata
+
+        assert ((self.seq_lens is not None)
+                or (self.encoder_seq_lens is not None))
+        assert ((self.seq_lens_tensor is not None)
+                or (self.encoder_seq_lens_tensor is not None))
+
+        # Compute some attn_metadata fields which default to None
+        query_start_loc = (None if self.query_start_loc is None else
+                           self.query_start_loc[:self.num_prefills + 1])
+        slot_mapping = (None if self.slot_mapping is None else
+                        self.slot_mapping[:self.num_prefill_tokens])
+        seq_lens = (None if self.seq_lens is None else
+                    self.seq_lens[:self.num_prefills])
+        seq_lens_tensor = (None if self.seq_lens_tensor is None else
+                           self.seq_lens_tensor[:self.num_prefills])
+        seq_start_loc = (None if self.seq_start_loc is None else
+                         self.seq_start_loc[:self.num_prefills + 1])
+        context_lens_tensor = (None if self.context_lens_tensor is None else
+                               self.context_lens_tensor[:self.num_prefills])
+        block_tables = (None if self.block_tables is None else
+                        self.block_tables[:self.num_prefills])
+
+        self._cached_prefill_metadata = MLUFlashAttentionMetadata(
+            num_prefills=self.num_prefills,
+            num_prefill_tokens=self.num_prefill_tokens,
+            num_decode_tokens=0,
+            slot_mapping=slot_mapping,
+            multi_modal_placeholder_index_maps=self.
+            multi_modal_placeholder_index_maps,
+            enable_kv_scales_calculation=self.enable_kv_scales_calculation,
+            seq_lens=seq_lens,
+            seq_lens_tensor=seq_lens_tensor,
+            max_query_len=self.max_query_len,
+            max_prefill_seq_len=self.max_prefill_seq_len,
+            max_decode_query_len=0,
+            max_decode_seq_len=0,
+            query_start_loc=query_start_loc,
+            seq_start_loc=seq_start_loc,
+            context_lens_tensor=context_lens_tensor,
+            block_tables=block_tables,
+            use_cuda_graph=self.use_cuda_graph,
+            # Begin encoder & cross attn fields below...
+            encoder_seq_lens=self.encoder_seq_lens,
+            encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,
+            encoder_seq_start_loc=self.encoder_seq_start_loc,
+            max_encoder_seq_len=self.max_encoder_seq_len,
+            cross_slot_mapping=self.cross_slot_mapping,
+            cross_block_tables=self.cross_block_tables,
+            chunked_prefill_enabled=self.chunked_prefill_enabled,
+            compute_dtype=self.compute_dtype,
+            is_profile_run=self.is_profile_run,
+            # Chunk prefill specific
+            context_chunk_cu_seq_lens=self.context_chunk_cu_seq_lens,
+            context_chunk_starts=self.context_chunk_starts,
+            context_chunk_seq_tot=self.context_chunk_seq_tot,
+            context_chunk_max_seq_lens=self.context_chunk_max_seq_lens)
+        return self._cached_prefill_metadata
+
+    @property
+    def decode_metadata(self) -> Optional["MLUFlashAttentionMetadata"]:
+        if self.num_decode_tokens == 0:
+            return None
+
+        if self._cached_decode_metadata is not None:
+            return self._cached_decode_metadata
+        assert ((self.seq_lens_tensor is not None)
+                or (self.encoder_seq_lens_tensor is not None))
+
+        # Compute some attn_metadata fields which default to None
+        slot_mapping = (None if self.slot_mapping is None else
+                        self.slot_mapping[self.num_prefill_tokens:])
+        seq_lens_tensor = (None if self.seq_lens_tensor is None else
+                           self.seq_lens_tensor[self.num_prefills:])
+        block_tables = (None if self.block_tables is None else
+                        self.block_tables[self.num_prefills:])
+
+        self._cached_decode_metadata = MLUFlashAttentionMetadata(
+            num_prefills=0,
+            num_prefill_tokens=0,
+            num_decode_tokens=self.num_decode_tokens,
+            slot_mapping=slot_mapping,
+            multi_modal_placeholder_index_maps=None,
+            enable_kv_scales_calculation=True,
+            seq_lens=None,
+            seq_lens_tensor=seq_lens_tensor,
+            max_decode_query_len=self.max_decode_query_len,
+            max_query_len=self.max_query_len,
+            max_prefill_seq_len=0,
+            max_decode_seq_len=self.max_decode_seq_len,
+            # Batch may be composed of prefill|decodes, adjust query start
+            # indices to refer to the start of decodes. E.g.
+            # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].
+            query_start_loc=(self.query_start_loc[self.num_prefills:] -
+                             self.query_start_loc[self.num_prefills])
+            if self.query_start_loc is not None else None,
+            seq_start_loc=self.seq_start_loc[self.num_prefills:]
+            if self.seq_start_loc is not None else None,
+            context_lens_tensor=None,
+            block_tables=block_tables,
+            use_cuda_graph=self.use_cuda_graph,
+            # Begin encoder & cross attn fields below...
+            encoder_seq_lens=self.encoder_seq_lens,
+            encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,
+            encoder_seq_start_loc=self.encoder_seq_start_loc,
+            max_encoder_seq_len=self.max_encoder_seq_len,
+            cross_slot_mapping=self.cross_slot_mapping,
+            cross_block_tables=self.cross_block_tables,
+            chunked_prefill_enabled=self.chunked_prefill_enabled,
+            compute_dtype=self.compute_dtype,
+            is_profile_run=self.is_profile_run)
+        return self._cached_decode_metadata
+
+    def advance_step(self,
+                     model_input: "ModelInputForGPUWithSamplingMetadata",
+                     sampled_token_ids: Optional[torch.Tensor],
+                     block_size: int,
+                     num_seqs: int,
+                     num_queries: int,
+                     turn_prefills_into_decodes: bool = False):
+        """
+        Update metadata in-place to advance one decode step.
+        """
+        # When using cudagraph, the num_seqs is padded to the next captured
+        # batch sized, but num_queries tracks the actual number of requests in
+        # the batch. For --enforce-eager mode, num_seqs == num_queries
+        if num_seqs != num_queries:
+            assert num_seqs > num_queries
+            assert self.use_cuda_graph
+
+        if turn_prefills_into_decodes:
+            # When Mutli-Step is enabled with Chunked-Prefill, prefills and
+            # decodes are scheduled together. In the first step, all the
+            # prefills turn into decodes. This update reflects that
+            # conversion.
+            assert self.num_decode_tokens + self.num_prefills == num_seqs
+            self.num_decode_tokens += self.num_prefills
+            self.num_prefills = 0
+            self.num_prefill_tokens = 0
+            self.max_prefill_seq_len = 0
+            self.max_query_len = 1
+
+            self.slot_mapping = self.slot_mapping[:num_seqs]
+        else:
+            assert self.seq_lens is not None
+            assert self.max_decode_seq_len == max(self.seq_lens)
+
+        assert self.num_prefills == 0
+        assert self.num_prefill_tokens == 0
+        assert self.num_decode_tokens == num_seqs
+        assert self.slot_mapping.shape == (num_seqs, )
+
+        assert self.seq_lens is not None
+        assert len(self.seq_lens) == num_seqs
+        assert self.seq_lens_tensor is not None
+        assert self.seq_lens_tensor.shape == (num_seqs, )
+        assert self.max_query_len == 1
+        assert self.max_prefill_seq_len == 0
+
+        assert self.query_start_loc is not None
+        assert self.query_start_loc.shape == (num_queries + 1, )
+        assert self.seq_start_loc is not None
+        assert self.seq_start_loc.shape == (num_seqs + 1, )
+
+        assert self.context_lens_tensor is not None
+        assert self.context_lens_tensor.shape == (num_queries, )
+
+        assert self.block_tables is not None
+        assert self.block_tables.shape[0] == num_seqs
+
+        # Update query lengths. Note that we update only queries and not seqs,
+        # since tensors may be padded due to captured cuda graph batch size
+        for i in range(num_queries):
+            self.seq_lens[i] += 1
+        self.max_decode_seq_len = max(self.seq_lens)
+
+        mlu_ops.advance_step(num_seqs=num_seqs,
+                             num_queries=num_queries,
+                             block_size=block_size,
+                             input_tokens=model_input.input_tokens,
+                             sampled_token_ids=sampled_token_ids,
+                             input_positions=model_input.input_positions,
+                             seq_lens=self.seq_lens_tensor,
+                             slot_mapping=self.slot_mapping,
+                             block_tables=self.block_tables)
+
+
+class MLUFlashAttentionMetadataBuilder(
+        AttentionMetadataBuilder[MLUFlashAttentionMetadata]):
+
+    def __init__(self, input_builder: "ModelInputForMLUBuilder"):
+        self.input_builder = input_builder
+        self.runner = input_builder.runner
+        self.sliding_window = input_builder.sliding_window
+        self.block_size = input_builder.block_size
+
+    def prepare(self):
+        self.slot_mapping: List[int] = []
+        self.prefill_seq_lens: List[int] = []
+        self.context_lens: List[int] = []
+        self.block_tables: List[List[int]] = []
+        self.curr_seq_lens: List[int] = []
+        self.multimodal_placeholder_maps: Dict[
+            str,
+            MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)
+        self.num_prefills = 0
+        self.num_prefill_tokens = 0
+        self.num_decode_tokens = 0
+        self.has_prefix_cache_hit = False
+
+    def _add_seq_group(
+            self, inter_data: "ModelInputForMLUBuilder.InterDataForSeqGroup",
+            chunked_prefill_enabled: bool, prefix_cache_hit: bool):
+        """Add a sequence group to the metadata. Specifically update/append
+        1. context length.
+        2. block table.
+        3. slot mapping.
+        """
+        is_prompt = inter_data.is_prompt
+        block_tables = inter_data.block_tables
+
+        for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
+             curr_sliding_window_block) in zip(
+                 inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],
+                 inter_data.orig_seq_lens, inter_data.seq_lens,
+                 inter_data.query_lens, inter_data.context_lens,
+                 inter_data.curr_sliding_window_blocks):
+            self.context_lens.append(context_len)
+
+            if is_prompt:
+                mm_maps = inter_data.multi_modal_placeholder_maps
+                if mm_maps:
+                    for modality, placeholders in mm_maps.items():
+                        self.multimodal_placeholder_maps[modality].extend(
+                            placeholders)
+
+                self.num_prefills += 1
+                self.num_prefill_tokens += token_len
+                self.prefill_seq_lens.append(seq_len)
+            else:
+                self.num_decode_tokens += query_len
+                self.curr_seq_lens.append(curr_seq_len)
+
+            # Compute block table.
+            # TODO(sang): Combine chunked prefill and prefix caching by
+            # only allowing multiple of block_size chunk size.
+            # NOTE: This only works for oooooooxxx style attention.
+            block_table = []
+            if prefix_cache_hit:
+                # NOTE(woosuk): For flash-attn, the block table should
+                # include the entries for the incoming prefill tokens.
+                block_table = block_tables[seq_id]
+            elif ((chunked_prefill_enabled or not is_prompt)
+                  and block_tables is not None):
+                if curr_sliding_window_block == 0:
+                    block_table = block_tables[seq_id]
+                else:
+                    block_table = block_tables[seq_id][
+                        -curr_sliding_window_block:]
+            self.block_tables.append(block_table)
+
+            # Compute slot mapping.
+            is_profile_run = is_block_tables_empty(block_tables)
+            start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
+                                                       context_len,
+                                                       self.sliding_window)
+            compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
+                                 seq_len, context_len, start_idx,
+                                 self.block_size, inter_data.block_tables)
+
+    def _get_graph_runner_block_tables(
+            self, num_seqs: int,
+            block_tables: List[List[int]]) -> torch.Tensor:
+        # The shape of graph_block_tables is
+        # [max batch size, max context len // block size].
+        max_batch_size, max_blocks = self.runner.graph_block_tables.shape
+        assert max_batch_size >= num_seqs
+
+        graph_block_tables = self.runner.graph_block_tables[:num_seqs]
+        for i, block_table in enumerate(block_tables):
+            if block_table:
+                num_blocks = len(block_table)
+                if num_blocks <= max_blocks:
+                    graph_block_tables[i, :num_blocks] = block_table
+                else:
+                    # It may be possible to have more blocks allocated due
+                    # to lookahead slots of multi-step, however, they are
+                    # not used anyway, so can be safely ignored.
+                    graph_block_tables[
+                        i, :max_blocks] = block_table[:max_blocks]
+
+        return torch.from_numpy(graph_block_tables).to(
+            device=self.runner.device, non_blocking=True)
+
+    def build(self, seq_lens: List[int], query_lens: List[int],
+              cuda_graph_pad_size: int, batch_size: int):
+        """Build attention metadata with on-device tensors.
+
+        Args:
+            seq_lens: The maybe padded sequence lengths of the input sequences.
+            query_lens: The query lengths of the input sequences.
+            cuda_graph_pad_size: The padding size for cuda graph.
+                                 -1 if cuda graph is not used.
+            batch_size: The maybe padded batch size.
+        """
+        prefix_cache_hit = any([
+            inter_data.prefix_cache_hit
+            for inter_data in self.input_builder.inter_data_list
+        ])
+        for inter_data in self.input_builder.inter_data_list:
+            self._add_seq_group(inter_data,
+                                self.input_builder.chunked_prefill_enabled,
+                                prefix_cache_hit)
+
+        device = self.runner.device
+        use_captured_graph = cuda_graph_pad_size != -1
+
+        max_query_len = max(query_lens)
+        decode_query_lens = query_lens[self.num_prefills:]
+        if len(decode_query_lens) > 0:
+            max_decode_query_len = max(decode_query_lens)
+        else:
+            max_decode_query_len = 1
+        max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
+        max_decode_seq_len = max(self.curr_seq_lens, default=0)
+        num_decode_tokens = self.num_decode_tokens
+        query_start_loc = list(accumulate(query_lens, initial=0))
+        seq_start_loc = list(accumulate(seq_lens, initial=0))
+
+        num_seqs = len(seq_lens)
+        if use_captured_graph:
+            self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)
+            self.block_tables.extend([] * cuda_graph_pad_size)
+            num_decode_tokens = batch_size - self.num_prefill_tokens
+            block_tables = self._get_graph_runner_block_tables(
+                num_seqs, self.block_tables)
+        else:
+            block_tables = make_tensor_with_pad(
+                self.block_tables,
+                pad=0,
+                dtype=torch.int,
+                device=device,
+            )
+        assert max_query_len > 0, ("query_lens: {}".format(query_lens))
+
+        assert device is not None
+        context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,
+                                               device, self.runner.pin_memory)
+        seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,
+                                           self.runner.pin_memory)
+        slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.int32,
+                                               device, self.runner.pin_memory)
+        query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,
+                                                  device,
+                                                  self.runner.pin_memory)
+        seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,
+                                                device, self.runner.pin_memory)
+        placeholder_index_maps = {
+            modality: placeholder_map.index_map()
+            for modality, placeholder_map in
+            self.multimodal_placeholder_maps.items()
+        }
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: Check if we can use context mlugraph for the given input.
+        '''
+        if (self.runner.model_config.use_context_mlugraph()
+                and num_decode_tokens == 0 and self.num_prefills > 0):
+            ctx_graph_bs, ctx_graph_seq_len = (
+                self.runner.model_config.get_context_mlugraph_bs_and_seq())
+            use_captured_graph = len(seq_lens) == ctx_graph_bs and all(
+                seq_len == ctx_graph_seq_len for seq_len in seq_lens)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        return MLUFlashAttentionMetadata(
+            num_prefills=self.num_prefills,
+            slot_mapping=slot_mapping_tensor,
+            num_prefill_tokens=self.num_prefill_tokens,
+            num_decode_tokens=num_decode_tokens,
+            seq_lens=seq_lens,
+            multi_modal_placeholder_index_maps=placeholder_index_maps,
+            enable_kv_scales_calculation=True,
+            seq_lens_tensor=seq_lens_tensor,
+            max_query_len=max_query_len,
+            max_decode_query_len=max_decode_query_len,
+            max_prefill_seq_len=max_prefill_seq_len,
+            max_decode_seq_len=max_decode_seq_len,
+            query_start_loc=query_start_loc_tensor,
+            seq_start_loc=seq_start_loc_tensor,
+            context_lens_tensor=context_lens_tensor,
+            block_tables=block_tables,
+            use_cuda_graph=use_captured_graph,
+            chunked_prefill_enabled=self.input_builder.chunked_prefill_enabled,
+            is_profile_run=self.runner.in_profile_run,
+        )
+
+
+class MLUFlashAttentionImpl(AttentionImpl):
+    """
+    If the input tensors contain prompt tokens, the layout is as follows:
+    |<--------------- num_prefill_tokens ----------------->|	
+    |<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->|
+
+    Otherwise, the layout is as follows:	
+    |<----------------- num_decode_tokens ------------------>|	
+    |<--decode_0-->|..........|<--decode_M-1-->|<--padding-->|
+
+    Generation tokens can contain padding when cuda-graph is used.
+    Currently, prompt tokens don't contain any padding.
+
+    The prompts might have different lengths, while the generation tokens
+    always have length 1.
+
+    If chunked prefill is enabled, prefill tokens and decode tokens can be
+    batched together in a flattened 1D query.
+
+    |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|
+    |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|
+
+    Currently, cuda graph is disabled for chunked prefill, meaning there's no
+    padding between prefill and decode tokens.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: Optional[List[float]],
+        sliding_window: Optional[int],
+        kv_cache_dtype: str,
+        blocksparse_params: Optional[Dict[str, Any]] = None,
+        logits_soft_cap: Optional[float] = None,
+        attn_type: str = AttentionType.DECODER,
+        **extra_impl_args,
+    ) -> None:
+        if blocksparse_params is not None:
+            raise ValueError(
+                "FlashAttention does not support block-sparse attention.")
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        if alibi_slopes is not None:
+            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32).mlu()
+        self.alibi_slopes = alibi_slopes
+        self.sliding_window = ((sliding_window - 1,
+                                -1) if sliding_window is not None else (-1, -1))
+        self.kv_cache_dtype = kv_cache_dtype
+        if logits_soft_cap is None:
+            # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
+            logits_soft_cap = 0
+        self.logits_soft_cap = logits_soft_cap
+
+        assert self.num_heads % self.num_kv_heads == 0
+        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+
+        support_head_sizes = MLUFlashAttentionBackend.get_supported_head_sizes()
+        if head_size not in support_head_sizes:
+            raise ValueError(
+                f"Head size {head_size} is not supported by FlashAttention. "
+                f"Supported head sizes are: {support_head_sizes}.")
+        self.attn_type = attn_type
+        self.use_fused_mla_qkv = extra_impl_args.get("use_fused_mla_qkv", False)
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: List[torch.Tensor],
+        attn_metadata: MLUFlashAttentionMetadata,
+        output: Optional[torch.Tensor] = None,
+        kwargs: Optional[dict[str, Any]] = {},
+    ) -> torch.Tensor:
+        """Forward pass with FlashAttention.
+
+        Args:
+            query: shape = [num_tokens, num_heads, head_size]
+            key: shape = [num_tokens, num_kv_heads, head_size]
+            value: shape = [num_tokens, num_kv_heads, v_head_size]
+            output: shape = [num_tokens, num_heads, head_size]
+            kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
+                NOTE: kv_cache will be an empty tensor with shape [0]
+                for profiling run.
+            attn_metadata: Metadata for attention.
+        NOTE: It in-place updates the output tensor.
+        """
+        # NOTE(woosuk): FlashAttention does not support FP8 KV cache.
+        assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0, (
+            "key/v_scale is not supported in FlashAttention.")
+
+        assert output is not None, "Output tensor must be provided."
+
+        attn_type = self.attn_type
+        if (attn_type == AttentionType.ENCODER
+                and (not attn_metadata.is_all_encoder_attn_metadata_set)):
+            raise AttributeError("Encoder attention requires setting "
+                                 "encoder metadata attributes.")
+        elif (attn_type == AttentionType.ENCODER_DECODER
+              and (not attn_metadata.is_all_cross_attn_metadata_set)):
+            raise AttributeError("Encoder/decoder cross-attention "
+                                 "requires setting cross-attention "
+                                 "metadata attributes.")
+
+        kv_cache_dtype: str = self.kv_cache_dtype
+        softmax_scale: float = self.scale
+        window_size = self.sliding_window
+        alibi_slopes: Optional[torch.Tensor] = self.alibi_slopes
+        logits_soft_cap: Optional[float] = self.logits_soft_cap
+        use_mla: Optional[bool] = self.use_mla
+
+        only_decode = kwargs.get("only_decode", False)
+        only_prefill = kwargs.get("only_prefill", False)
+        prefill_meta = None if only_decode else attn_metadata.prefill_metadata
+        decode_meta = None if only_prefill else attn_metadata.decode_metadata
+
+        (num_prefill_query_tokens, num_prefill_kv_tokens,
+        num_decode_query_tokens) = \
+            get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type)
+
+        if decode_meta:
+            if only_decode:
+                decode_query = query
+                decode_output = output
+            else:
+                decode_query = query[num_prefill_query_tokens:]
+                decode_output = output[num_prefill_query_tokens:]
+            assert decode_query.shape[0] == num_decode_query_tokens
+
+        # QKV for prefill.
+        if not (only_decode or only_prefill):
+            query = query[:num_prefill_query_tokens]
+
+        if prefill_meta:
+            prefill_output = output if only_prefill else output[:num_prefill_query_tokens]
+            assert query.shape[0] == num_prefill_query_tokens
+
+        out_lse = None
+
+        if kv_cache[0].numel() > 0:
+            kv_cache_, kv_cache_scale_ = kv_cache
+            key_cache = kv_cache_[0]
+            value_cache = None if use_mla else kv_cache_[1]
+            key_cache_scale, value_cache_scale = None, None
+            if kv_cache_scale_.numel() > 0:
+                key_cache_scale = kv_cache_scale_[0]
+                value_cache_scale = None if use_mla else kv_cache_scale_[1]
+
+            # We skip updating the KV cache under two conditions:
+            #  a. When the Attention Type is ENCODER. In this phase, we compute
+            #     only the encoder attention without updating the cache.
+            #  b. When both Key and Value are None. This occurs during
+            #     cross-attention computation in the decoding phase, where the
+            #     KV cache is already populated with the cross-attention
+            #     tensor. Thus, we skip cache updates during this time.
+            if (attn_type != AttentionType.ENCODER) and (key is not None) and (
+                    value is not None):
+                if attn_type == AttentionType.ENCODER_DECODER:
+                    # Update cross-attention KV cache (prefill-only)
+                    updated_slot_mapping = attn_metadata.cross_slot_mapping
+                else:
+                    # Update self-attention KV cache (prefill/decode)
+                    updated_slot_mapping = attn_metadata.slot_mapping
+
+                # Reshape the input keys and values and store them in the cache.
+                # If kv_cache is not provided, the new key and value tensors are
+                # not cached. This happens during the initial memory
+                # profiling run.
+                value_to_cache = None if use_mla else value
+                if use_mla and (prefill_meta or self.use_fused_mla_qkv):
+                    # MLA save cache info in models before flashattn
+                    pass
+                else:
+                    if only_decode:
+                        updated_slot_mapping = updated_slot_mapping[num_prefill_query_tokens:]
+                    if only_prefill:
+                        updated_slot_mapping = updated_slot_mapping[:num_prefill_query_tokens]
+                    if kv_cache_dtype == 'int8':
+                        mlu_ops.quant_to_paged_cache(
+                            key,
+                            value_to_cache,
+                            key_cache,
+                            value_cache,
+                            key_cache_scale,
+                            value_cache_scale,
+                            updated_slot_mapping.flatten()
+                        )
+                    else:
+                        mlu_ops.reshape_paged_cache(
+                            key,
+                            value_to_cache,
+                            key_cache,
+                            value_cache,
+                            updated_slot_mapping.flatten()
+                        )
+
+        if prefill_meta:
+            alibi_slopes = None if alibi_slopes is None else \
+                                alibi_slopes.repeat(attn_metadata.num_prefills, 1)
+            # Prompt run.
+            if (kv_cache[0].numel() == 0 or prefill_meta.block_tables is None
+                    or prefill_meta.block_tables.numel() == 0 or use_mla):
+                # normal attention
+                # When block_tables are not filled, it means q and k are the
+                # prompt, and they have the same length.
+                q_seq_start_loc, q_seq_len, k_seq_start_loc, k_seq_len = \
+                    _get_query_key_seq_metadata(prefill_meta, True, attn_type)
+
+                key = key[:num_prefill_kv_tokens]
+                value = value[:num_prefill_kv_tokens]
+                prefill_causal = _get_causal_option(attn_type)
+
+                prefill_causal = kwargs.get("prefill_causal", prefill_causal)
+                q_seq_start_loc = kwargs.get("cu_seq_lens_q", q_seq_start_loc)
+                k_seq_start_loc = kwargs.get("cu_seq_lens_kv", k_seq_start_loc)
+                q_seq_len = kwargs.get("max_seq_len_q", q_seq_len)
+                k_seq_len = kwargs.get("max_seq_len_kv", k_seq_len)
+                return_lse = kwargs.get("return_lse", False)
+
+                attn_output_list = mlu_ops.flash_attention(
+                    query,
+                    key,
+                    value,
+                    prefill_output,
+                    q_seq_start_loc,
+                    k_seq_start_loc,
+                    alibi_slopes,
+                    None,
+                    q_seq_len,
+                    k_seq_len,
+                    softmax_scale,
+                    prefill_causal,
+                    -1 if window_size is None else window_size[0],
+                    -1 if window_size is None else window_size[1],
+                    prefill_meta.compute_dtype,
+                    return_lse
+                )
+                if return_lse:
+                    out_lse = attn_output_list[1]
+            else:
+                # prefix-enabled attention
+                assert attn_type == AttentionType.DECODER, (
+                    "Only decoder-only models support prefix caching")
+                assert prefill_meta.seq_lens is not None
+                max_seq_len = max(prefill_meta.seq_lens)
+
+                if kv_cache_dtype == 'int8' and \
+                        prefill_meta.chunked_prefill_enabled:
+                    _, head_num_kv, _, head_size_qk = key_cache.shape
+                    total_seqlens = prefill_meta.seq_start_loc[-1].item()
+                    key_cache_float = torch.zeros((total_seqlens, head_num_kv, head_size_qk),
+                                                  dtype=query.dtype,
+                                                  device=key_cache.device)
+                    value_cache_float = None
+                    if value_cache is not None:
+                        _, head_num_kv, _, head_size_v = value_cache.shape
+                        value_cache_float = torch.zeros((total_seqlens, head_num_kv, head_size_v),
+                                                        dtype=query.dtype,
+                                                        device=value_cache.device)
+                    mlu_ops.dequant_from_paged_cache(
+                        key=key_cache_float,
+                        value=value_cache_float,
+                        key_cache=key_cache,
+                        value_cache=value_cache,
+                        key_cache_quant_scale=key_cache_scale,
+                        value_cache_quant_scale=value_cache_scale,
+                        context_lengths=prefill_meta.seq_lens_tensor,
+                        max_context_len=prefill_meta.max_prefill_seq_len,
+                        context_seq_offset=None,
+                        block_tables=prefill_meta.block_tables,
+                        quant_mode=1,
+                        quant_bit=8
+                    )
+                    block_tables = None
+                else:
+                    key_cache_float = key_cache
+                    value_cache_float = value_cache
+                    block_tables = prefill_meta.block_tables
+
+                mlu_ops.flash_attention(
+                    query,
+                    key_cache_float,
+                    value_cache_float,
+                    prefill_output,
+                    prefill_meta.query_start_loc,
+                    prefill_meta.seq_start_loc,
+                    alibi_slopes,
+                    None,
+                    prefill_meta.max_query_len,
+                    max_seq_len,
+                    softmax_scale,
+                    True,
+                    -1 if window_size is None else window_size[0],
+                    -1 if window_size is None else window_size[1],
+                    prefill_meta.compute_dtype,
+                    False,
+                    block_tables
+                )
+
+        if decode_meta:
+            # Decoding run.
+            # Use flash_attn_varlen_func kernel for speculative decoding
+            # because different queries might have different lengths.
+            alibi_slopes = None if alibi_slopes is None \
+                            else alibi_slopes.repeat(attn_metadata.num_decode_tokens, 1)
+
+            assert decode_meta.max_decode_query_len is not None
+            # use only for actual varlen decoding
+            if decode_meta.max_decode_query_len > 1:
+                assert attn_type == AttentionType.DECODER, (
+                    "Only decoder-only models support max_decode_query_len > 1"
+                )
+                if use_mla:
+                    op_result = None
+                else:
+                    op_result = decode_output
+                output_tmp = mlu_ops.flash_attention(
+                    decode_query,
+                    key_cache,
+                    value_cache,
+                    op_result,
+                    decode_meta.query_start_loc,
+                    decode_meta.seq_start_loc,
+                    alibi_slopes,
+                    None,
+                    decode_meta.max_decode_query_len,
+                    decode_meta.max_decode_seq_len,
+                    softmax_scale,
+                    True,
+                    -1 if window_size is None else window_size[0],
+                    -1 if window_size is None else window_size[1],
+                    decode_meta.compute_dtype,
+                    False,
+                    decode_meta.block_tables
+                )
+                if use_mla:
+                    decode_output.copy_(output_tmp[..., :decode_output.shape[-1]])
+            else:
+                decode_query = decode_query.view(-1, 1, self.num_heads, self.head_size)
+                output_head_size = value.shape[-1] if use_mla else self.head_size
+                decode_output = decode_output.view(-1, 1, self.num_heads, output_head_size)
+                # Use flash_attn_with_kvcache for normal decoding.
+                (
+                    seq_lens_arg,
+                    max_context_len,
+                    block_tables_arg,
+                ) = get_seq_len_block_table_args(decode_meta, False, attn_type)
+                if use_mla:
+                    value_cache = None
+                    value_cache_scale = None
+                mlu_ops.single_query_cached_kv_attn(
+                    decode_query,
+                    key_cache,
+                    value_cache,
+                    decode_output,
+                    block_tables_arg,
+                    seq_lens_arg,
+                    key_cache_scale,
+                    value_cache_scale,
+                    alibi_slopes,
+                    max_context_len,
+                    -1 if window_size is None else window_size[0],
+                    -1 if window_size is None else window_size[1],
+                    softmax_scale,
+                    head_size_v=value.shape[-1] if use_mla else -1,
+                    compute_dtype=decode_meta.compute_dtype
+                )
+        return output, out_lse
+
+
+def _get_query_key_seq_metadata(
+    attn_metadata,
+    is_prompt: bool,
+    attn_type: str,
+) -> tuple:
+    """
+    Returns sequence metadata for key and query based on the specified 
+    attention type and whether input is a prompt.
+
+    This function computes the starting locations and maximum sequence lengths 
+    for key and query sequences for different attention types.
+
+    Args:
+        attn_metadata: The attention metadata object
+        is_prompt (bool): A flag indicating if the input is a prompt
+        attn_type (AttentionType): The type of attention being used.
+
+    Returns:
+        tuple: A tuple containing four integers:
+            - Starting location for the query sequence.
+            - Maximum sequence length for the query sequence.
+            - Starting location for the key sequence.
+            - Maximum sequence length for the key sequence.
+
+    Raises:
+        AttributeError: If an invalid attention type is provided.
+    """
+    if attn_type == AttentionType.DECODER:
+        # Decoder self-attention
+        # Choose max_seq_len based on whether we are in prompt_run
+        if is_prompt:
+            max_seq_len = attn_metadata.max_prefill_seq_len
+        else:
+            max_seq_len = attn_metadata.max_decode_seq_len
+        return (attn_metadata.query_start_loc, attn_metadata.max_query_len,
+                attn_metadata.seq_start_loc, max_seq_len)
+
+    elif attn_type == AttentionType.ENCODER_DECODER:
+        # This is cross attention between the where the key
+        # is the precomputed encoder attention and query
+        # is the input sequence.
+        # Choose query max length based on whether it is prompt
+        # or not.
+        if is_prompt:
+            max_seq_len = attn_metadata.max_prefill_seq_len
+        else:
+            max_seq_len = attn_metadata.max_decode_seq_len
+        return (attn_metadata.seq_start_loc, max_seq_len,
+                attn_metadata.encoder_seq_start_loc,
+                attn_metadata.max_encoder_seq_len)
+    elif attn_type == AttentionType.ENCODER:
+        # For encoder attention both the query and the key are same i.e the
+        # encoder sequence.
+        return (attn_metadata.encoder_seq_start_loc,
+                attn_metadata.max_encoder_seq_len,
+                attn_metadata.encoder_seq_start_loc,
+                attn_metadata.max_encoder_seq_len)
+    elif attn_type == AttentionType.ENCODER_ONLY:
+        assert is_prompt, "Should not have decode for encoder only model."
+        return (attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len,
+                attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len)
+    else:
+        raise AttributeError(f"Invalid attention type {str(attn_type)}")
+
+
+def _get_causal_option(attn_type: str) -> bool:
+    """
+    Determine whether the given attention type is suitable for causal 
+    attention mechanisms.
+
+    Args:
+        attn_type (AttentionType): The type of attention being evaluated
+
+    Returns:
+        bool: Returns `True` if the attention type is suitable for causal 
+        attention (i.e., not encoder, encoder-only, or encoder-decoder), 
+        otherwise returns `False`.
+    """
+    return not (attn_type == AttentionType.ENCODER
+                or attn_type == AttentionType.ENCODER_ONLY
+                or attn_type == AttentionType.ENCODER_DECODER)
+
+
+class MLUMLAAttentionState(MLUAttentionState):
+    def __init__(self, runner: "ModelRunnerBase"):
+        super().__init__(runner)
+
+        scheduler_config = runner.scheduler_config
+        self.model_config = runner.model_config
+        cache_config = runner.cache_config
+
+        self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
+
+        if self.chunked_prefill_enabled:
+            self.context_chunk_workspace_size = min(
+                # Max sure there is enough for 8 full length request or at least
+                # 4 pages of cache per request
+                max(
+                    8 * self.model_config.max_model_len, 4 *
+                    scheduler_config.max_num_seqs * cache_config.block_size),
+                # For long-context models try not to over-allocate limiting
+                # kv-cache space, limiting it to 64k tokens,
+                # which would result in the workspace being:
+                #   2*(576)*(64*1024) = 144mb
+                # (assuming 576 MLA head dim, and fp16)
+                # which would result in up-projected context being
+                #   2*(192*128)*(64*1024) = 3gb
+                # (assuming 192 QK head dim, 128 heads, and fp16)
+                128 * 1024)
+            assert self.context_chunk_workspace_size >= \
+                scheduler_config.max_num_seqs * cache_config.block_size
+
+    def begin_forward(self, model_input):
+        if self.chunked_prefill_enabled:
+            if not hasattr(self, "context_chunk_workspace"):
+                # not self.runner.device does not return the correct device
+                # for this process, (init_device sets the correct device but
+                # only on the Worker). The only way Ive figured out to get the
+                # correct device is to allocate the workspace on the first call
+                # to begin_forward and use the device of the input tokens
+                assert model_input.input_tokens is not None
+                self.context_chunk_workspace = torch.empty(
+                    (self.context_chunk_workspace_size,
+                     self.model_config.get_head_size()),
+                    dtype=self.model_config.dtype,
+                    device=model_input.input_tokens.device,
+                )
+
+            model_input.attn_metadata.context_chunk_workspace = \
+                self.context_chunk_workspace
+            model_input.attn_metadata.context_chunk_workspace_size = \
+                self.context_chunk_workspace_size
+
+
+class MLUMLAFlashAttentionMetadataBuilder(MLUFlashAttentionMetadataBuilder):
+    def __init__(self, input_builder: "ModelInputForMLUBuilder"):
+        super().__init__(input_builder)
+
+        self.chunked_prefill_enabled = \
+            self.runner.scheduler_config.chunked_prefill_enabled
+
+        if self.chunked_prefill_enabled:
+            attn_state = self.input_builder.runner.attn_state
+            self.context_chunk_workspace_size = \
+                attn_state.context_chunk_workspace_size
+            self.page_size = self.runner.block_size
+
+    def build(self, seq_lens: List[int], query_lens: List[int],
+              cuda_graph_pad_size: int, batch_size: int):
+        """Build attention metadata with on-device tensors.
+
+        Args:
+            seq_lens: The maybe padded sequence lengths of the input sequences.
+            query_lens: The query lengths of the input sequences.
+            cuda_graph_pad_size: The padding size for cuda graph.
+                                 -1 if cuda graph is not used.
+            batch_size: The maybe padded batch size.
+        """
+        prefix_cache_hit = any([
+            inter_data.prefix_cache_hit
+            for inter_data in self.input_builder.inter_data_list
+        ])
+        for inter_data in self.input_builder.inter_data_list:
+            self._add_seq_group(inter_data,
+                                self.input_builder.chunked_prefill_enabled,
+                                prefix_cache_hit)
+
+        device = self.runner.device
+        use_captured_graph = cuda_graph_pad_size != -1
+
+        max_query_len = max(query_lens)
+        decode_query_lens = query_lens[self.num_prefills:]
+        if len(decode_query_lens) > 0:
+            max_decode_query_len = max(decode_query_lens)
+        else:
+            max_decode_query_len = 1
+        max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
+        max_decode_seq_len = max(self.curr_seq_lens, default=0)
+        num_decode_tokens = self.num_decode_tokens
+        query_start_loc = list(accumulate(query_lens, initial=0))
+        seq_start_loc = list(accumulate(seq_lens, initial=0))
+
+        num_seqs = len(seq_lens)
+        if use_captured_graph:
+            self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)
+            self.block_tables.extend([] * cuda_graph_pad_size)
+            num_decode_tokens = batch_size - self.num_prefill_tokens
+            block_tables = self._get_graph_runner_block_tables(
+                num_seqs, self.block_tables)
+        else:
+            block_tables = make_tensor_with_pad(
+                self.block_tables,
+                pad=0,
+                dtype=torch.int,
+                device=device,
+            )
+        assert max_query_len > 0, ("query_lens: {}".format(query_lens))
+
+        assert device is not None
+        context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,
+                                               device, self.runner.pin_memory)
+        seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,
+                                           self.runner.pin_memory)
+        slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.int32,
+                                               device, self.runner.pin_memory)
+        query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,
+                                                  device,
+                                                  self.runner.pin_memory)
+        seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,
+                                                device, self.runner.pin_memory)
+        placeholder_index_maps = {
+            modality: placeholder_map.index_map()
+            for modality, placeholder_map in
+            self.multimodal_placeholder_maps.items()
+        }
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: Check if we can use context mlugraph for the given input.
+        '''
+        if (self.runner.model_config.use_context_mlugraph()
+                and num_decode_tokens == 0 and self.num_prefills > 0):
+            ctx_graph_bs, ctx_graph_seq_len = (
+                self.runner.model_config.get_context_mlugraph_bs_and_seq())
+            use_captured_graph = len(seq_lens) == ctx_graph_bs and all(
+                seq_len == ctx_graph_seq_len for seq_len in seq_lens)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: add mla chunked prefill logic.
+        '''
+        context_chunk_cu_seq_lens = None
+        context_chunk_starts = None
+        context_chunk_seq_tot = None
+        context_chunk_max_seq_lens = None
+        page_size=16
+        if self.chunked_prefill_enabled:
+            (context_chunk_cu_seq_lens, context_chunk_starts,
+             context_chunk_seq_tot,
+             context_chunk_max_seq_lens) = context_chunk_param(
+                 self.chunked_prefill_enabled, self.num_prefills,
+                 context_lens_tensor, self.context_chunk_workspace_size,
+                 self.page_size, device)
+            page_size = self.page_size
+
+        return MLUFlashAttentionMetadata(
+            num_prefills=self.num_prefills,
+            slot_mapping=slot_mapping_tensor,
+            num_prefill_tokens=self.num_prefill_tokens,
+            num_decode_tokens=num_decode_tokens,
+            seq_lens=seq_lens,
+            multi_modal_placeholder_index_maps=placeholder_index_maps,
+            enable_kv_scales_calculation=True,
+            seq_lens_tensor=seq_lens_tensor,
+            max_query_len=max_query_len,
+            max_decode_query_len=max_decode_query_len,
+            max_prefill_seq_len=max_prefill_seq_len,
+            max_decode_seq_len=max_decode_seq_len,
+            query_start_loc=query_start_loc_tensor,
+            seq_start_loc=seq_start_loc_tensor,
+            context_lens_tensor=context_lens_tensor,
+            block_tables=block_tables,
+            use_cuda_graph=use_captured_graph,
+            chunked_prefill_enabled=self.input_builder.chunked_prefill_enabled,
+            is_profile_run=self.runner.in_profile_run,
+            # Chunk prefill specific
+            context_chunk_cu_seq_lens=context_chunk_cu_seq_lens,
+            context_chunk_starts=context_chunk_starts,
+            context_chunk_seq_tot=context_chunk_seq_tot,
+            context_chunk_max_seq_lens=context_chunk_max_seq_lens,
+            page_size=page_size,
+        )
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+
+def context_chunk_param(chunked_prefill_enabled, num_prefills,
+                        context_lens_tensor, context_chunk_workspace_size,
+                        page_size, device):
+    context_chunk_cu_seq_lens = None
+    context_chunk_starts = None
+    context_chunk_seq_tot = None
+    context_chunk_max_seq_lens = None
+
+    if chunked_prefill_enabled \
+        and num_prefills > 0 \
+        and context_lens_tensor is not None \
+        and context_lens_tensor[:num_prefills].max() > 0:
+
+        # NOTE: it is recommend you read the `Chunked Prefill` section in
+        # the comment at the top of the file before trying to understand
+        # the following code
+
+        num_prefills_with_context = \
+            (context_lens_tensor[:num_prefills] > 0).sum().item()
+
+        # currently we allocate an equal amount of workspace for each
+        # prefill in the batch, we could probably use a more advanced
+        # algorithm here and allocate more workspace to prefills with
+        # longer context lengths
+        max_context_chunk = \
+            context_chunk_workspace_size // num_prefills_with_context
+
+        # align max_context_chunk to page_size by rounding down,
+        # currently the `gather_cache` kernel cannot handle
+        # `context_chunk_starts` that are not aligned to page_size
+        max_context_chunk = round_down(max_context_chunk, page_size)
+        assert max_context_chunk > 0
+        num_chunks = cdiv(context_lens_tensor.max(), max_context_chunk)
+
+        # if `max_context_chunk = 256`, `num_chunks = 3`, and
+        #   `num_prefills_with_context = 4`, create a tensor that looks like
+        #  [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
+        context_chunk_starts = \
+            torch.arange(num_chunks, device=device, dtype=torch.int32)\
+            .unsqueeze(1).expand(-1, num_prefills)\
+            * max_context_chunk
+        chunk_ends = torch.min(context_lens_tensor[:num_prefills]\
+            .unsqueeze(0), context_chunk_starts + max_context_chunk)
+        chunk_seq_lens = (chunk_ends - context_chunk_starts).clamp(min=0)
+        _context_chunk_cu_seq_lens = chunk_seq_lens.cumsum(dim=1).to(
+            torch.int32)
+        zero = torch.zeros(num_chunks, dtype=torch.int32, device=device)\
+            .unsqueeze(-1)
+        context_chunk_cu_seq_lens = \
+            torch.cat([zero, _context_chunk_cu_seq_lens], dim=1)
+        context_chunk_max_seq_lens = \
+            chunk_seq_lens.max(dim=1).values.tolist()
+        context_chunk_seq_tot = chunk_seq_lens.sum(dim=1).tolist()
+        assert max(context_chunk_seq_tot) <= \
+            context_chunk_workspace_size
+
+    return context_chunk_cu_seq_lens, context_chunk_starts, context_chunk_seq_tot, context_chunk_max_seq_lens
diff --git a/vllm_mlu/vllm_mlu/attention/layer.py b/vllm_mlu/vllm_mlu/attention/layer.py
new file mode 100644
index 000000000..9221d75ea
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/attention/layer.py
@@ -0,0 +1,277 @@
+from typing import Any, Dict, List, Optional
+
+import torch
+
+import vllm.envs as envs
+from vllm.attention import AttentionType
+from vllm.attention.layer import Attention
+from vllm.attention.selector import backend_name_to_enum, get_attn_backend
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.config import CacheConfig, get_current_vllm_config
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
+from vllm.platforms import current_platform
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+class Attention_MluHjack(Attention):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: Optional[int] = None,
+        alibi_slopes: Optional[List[float]] = None,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        blocksparse_params: Optional[Dict[str, Any]] = None,
+        logits_soft_cap: Optional[float] = None,
+        per_layer_sliding_window: Optional[int] = None,
+        use_mla: bool = False,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+        **extra_impl_args,
+    ) -> None:
+        """
+        The KV cache is stored inside this class and is accessed via
+        `self.kv_cache`.
+        """
+        super(Attention, self).__init__()
+        if per_layer_sliding_window is not None:
+            # per-layer sliding window
+            sliding_window = per_layer_sliding_window
+        elif cache_config is not None:
+            # model-level sliding window
+            sliding_window = cache_config.sliding_window
+        else:
+            sliding_window = None
+
+        if cache_config is not None:
+            kv_cache_dtype = cache_config.cache_dtype
+            block_size = cache_config.block_size
+            is_attention_free = cache_config.is_attention_free
+            calculate_kv_scales = cache_config.calculate_kv_scales
+        else:
+            kv_cache_dtype = "auto"
+            block_size = 16
+            is_attention_free = False
+            calculate_kv_scales = False
+        if num_kv_heads is None:
+            num_kv_heads = num_heads
+
+        # The default k/v_scale is set to 1.0. This is ignored
+        # when kv-cache is not fp8, and should be used with
+        # kv-cache in fp8_e5m2. For kv-cache in fp8_e4m3, we
+        # expect the pre-quantized k/v_scale to be loaded along
+        # with the model weights.
+        self.kv_cache_dtype = kv_cache_dtype
+        self.calculate_kv_scales = calculate_kv_scales
+        self._k_scale = torch.tensor(1.0, dtype=torch.float32)
+        self._v_scale = torch.tensor(1.0, dtype=torch.float32)
+        # FlashAttn doesn't support quantizing the kv-cache only
+        # but requires q to be quantized as well.
+        self._q_scale = torch.tensor(1.0, dtype=torch.float32)
+
+        # We also keep the float32 versions of k/v_scale for attention
+        # backends that don't support tensors (Flashinfer)
+        self._k_scale_float = 1.0
+        self._v_scale_float = 1.0
+
+        self.use_mla = use_mla
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.num_kv_heads = num_kv_heads
+        self.sliding_window = sliding_window
+
+        quant_method = quant_config.get_quant_method(
+            self, prefix=prefix) if quant_config else None
+        if quant_method is not None and not isinstance(
+                quant_method, UnquantizedLinearMethod):
+            assert isinstance(quant_method, BaseKVCacheMethod)
+            # TODO (mgoin): kv cache dtype should be specified in the FP8
+            # checkpoint config and become the "auto" behavior
+            if self.kv_cache_dtype == "fp8_e5m2":
+                raise ValueError("fp8_e5m2 kv-cache is not supported with "
+                                 "fp8 checkpoints.")
+            # If quantization is enabled, we make "k_scale" and "v_scale"
+            # parameters so that it can be loaded from the model checkpoint.
+            # The k/v_scale will then be converted back to native float32
+            # values after weight loading.
+            self.quant_method = quant_method
+            self.quant_method.create_weights(self)
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: insert v_head_dim, cause mlu_flash_attn support qk_v_head_dim not in same
+        '''
+        if use_mla:
+            self.v_head_dim = extra_impl_args['v_head_dim']
+        else:
+            self.v_head_dim = head_size
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        # During model initialization, the default dtype is set as the model
+        # weight and activation dtype.
+        dtype = torch.get_default_dtype()
+        attn_backend = get_attn_backend(head_size,
+                                        dtype,
+                                        kv_cache_dtype,
+                                        block_size,
+                                        is_attention_free,
+                                        blocksparse_params is not None,
+                                        use_mla=use_mla)
+        impl_cls = attn_backend.get_impl_cls()
+        self.impl = impl_cls(num_heads, head_size, scale, num_kv_heads,
+                             alibi_slopes, sliding_window, kv_cache_dtype,
+                             blocksparse_params, logits_soft_cap, attn_type,
+                             **extra_impl_args)
+        self.backend = backend_name_to_enum(attn_backend.get_name())
+        self.dtype = dtype
+
+        # For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
+        # torch.compile works by registering the attention as one giant
+        # opaque custom op. For other platforms, we directly call them
+        # and let torch.compile handle them.
+        self.use_direct_call = not current_platform.is_cuda_alike(
+        ) and not current_platform.is_cpu()
+
+        self.use_output = attn_backend.accept_output_buffer
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+        self.layer_name = prefix
+        self.attn_type = attn_type
+        # use a placeholder kv cache tensor during init, which will be replaced
+        # by bind_kv_cache
+        # this variable will not be accessed if use_direct_call is True
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: support kv8, keep use_mla
+        '''
+        self.kv_cache = [
+            [torch.tensor([]),
+             torch.tensor([])] for _ in range(get_current_vllm_config(
+            ).parallel_config.pipeline_parallel_size)
+        ]
+        self.impl.use_mla = use_mla
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
+        self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
+        self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        # For some alternate attention backends like MLA the attention output
+        # shape does not match the query shape, so we optionally let the model
+        # definition specify the output tensor shape.
+        output_shape: Optional[torch.Size] = None,
+        kwargs: Optional[dict[str, Any]] = {},
+    ) -> torch.Tensor:
+        """
+        The KV cache is stored inside this class and is accessed via
+        `self.kv_cache`.
+
+        Attention metadata (`attn_metadata`) is set using a context manager in
+        the model runner's `execute_model` method. It is accessed via forward
+        context using
+        `vllm.forward_context.get_forward_context().attn_metadata`.
+        """
+        if self.calculate_kv_scales:
+            attn_metadata = get_forward_context().attn_metadata
+            if attn_metadata.enable_kv_scales_calculation:
+                self.calc_kv_scales(query, key, value)
+        if self.use_output:
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: use v_head_dim when is_mla=True
+            '''
+            output_shape = (output_shape
+                            if output_shape is not None else query.shape)
+            if self.use_mla:
+                output_shape = [output_shape[0], self.num_heads * self.v_head_dim]
+            output = torch.empty(output_shape,
+                                 dtype=query.dtype,
+                                 device=query.device)
+            hidden_size = output_shape[-1]
+            # Reshape the query, key, and value tensors.
+            # NOTE(woosuk): We do this outside the custom op to minimize the
+            # CPU overheads from the non-CUDA-graph regions.
+            query = query.view(-1, self.num_heads, self.head_size)
+            output = output.view(-1, self.num_heads, self.v_head_dim)
+            if key is not None:
+                key = key.view(-1, self.num_kv_heads, self.head_size)
+            if value is not None:
+                value = value.view(-1, self.num_kv_heads, self.v_head_dim)
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            output_lse = None
+            if self.use_direct_call:
+                forward_context: ForwardContext = get_forward_context()
+                attn_metadata = forward_context.attn_metadata
+                self_kv_cache = self.kv_cache[forward_context.virtual_engine]
+                attn_output_list = self.impl.forward(self,
+                                                     query,
+                                                     key,
+                                                     value,
+                                                     self_kv_cache,
+                                                     attn_metadata,
+                                                     output=output,
+                                                     kwargs=kwargs)
+                if len(attn_output_list) > 1:
+                    output_lse = attn_output_list[1]
+            else:
+                torch.ops.vllm.unified_attention_with_output(
+                    query, key, value, output, self.layer_name)
+            if output_lse is not None:
+                return output.view(-1, hidden_size), output_lse
+            else:
+                return output.view(-1, hidden_size)
+        else:
+            if self.use_direct_call:
+                forward_context = get_forward_context()
+                attn_metadata = forward_context.attn_metadata
+                self_kv_cache = self.kv_cache[forward_context.virtual_engine]
+                return self.impl.forward(self,
+                                         query,
+                                         key,
+                                         value,
+                                         self_kv_cache,
+                                         attn_metadata,
+                                         output=None,
+                                         kwargs=kwargs)
+            else:
+                return torch.ops.vllm.unified_attention(
+                    query, key, value, self.layer_name)
+
+
+MluHijackObject.apply_hijack(Attention,
+                             Attention.__init__,
+                             Attention_MluHjack.__init__)
+MluHijackObject.apply_hijack(Attention,
+                             Attention.forward,
+                             Attention_MluHjack.forward)
diff --git a/vllm_mlu/vllm_mlu/attention/ops/__init__.py b/vllm_mlu/vllm_mlu/attention/ops/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/attention/ops/prefix_prefill.py b/vllm_mlu/vllm_mlu/attention/ops/prefix_prefill.py
new file mode 100644
index 000000000..95081cd07
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/attention/ops/prefix_prefill.py
@@ -0,0 +1,157 @@
+# The kernels in this file are adapted from LightLLM's context_attention_fwd:
+# https://github.com/ModelTC/lightllm/blob/main/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py
+
+import torch
+import triton
+import triton.language as tl
+
+import vllm.attention.ops.prefix_prefill
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+if triton.__version__ >= "2.1.0":
+
+    @torch.inference_mode()
+    def vllm__attention__ops__prefix_prefill__context_attention_fwd(q,
+                              k,
+                              v,
+                              o,
+                              k_cache,
+                              v_cache,
+                              b_loc,
+                              b_start_loc,
+                              b_seq_len,
+                              b_ctx_len,
+                              max_input_len,
+                              alibi_slopes=None,
+                              sliding_window=None):
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: use to many memory when block is 64
+        '''
+        BLOCK = 16
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        # shape constraints
+        Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
+        assert Lq == Lk and Lk == Lv
+        # round up Lk to a power of 2 - this is required for Triton block size
+        Lk_padded = triton.next_power_of_2(Lk)
+
+        sm_scale = 1.0 / (Lq**0.5)
+        batch, head = b_seq_len.shape[0], q.shape[1]
+        num_queries_per_kv = q.shape[1] // k.shape[1]
+
+        grid = (batch, head, triton.cdiv(max_input_len, BLOCK))  # batch, head,
+
+        num_warps = 8 if Lk <= 64 else 8
+        if alibi_slopes is not None:
+            assert Lk == Lk_padded
+            vllm.attention.ops.prefix_prefill._fwd_kernel_alibi[grid](
+                q,
+                k,
+                v,
+                k_cache,
+                v_cache,
+                b_loc,
+                sm_scale,
+                b_start_loc,
+                b_seq_len,
+                b_ctx_len,
+                alibi_slopes,
+                v_cache.shape[3],
+                8,
+                o,
+                b_loc.stride(0),
+                b_loc.stride(1),
+                q.stride(0),
+                q.stride(1),
+                q.stride(2),
+                k.stride(0),
+                k.stride(1),
+                k.stride(2),
+                v.stride(0),
+                v.stride(1),
+                v.stride(2),
+                o.stride(0),
+                o.stride(1),
+                o.stride(2),
+                k_cache.stride(0),
+                k_cache.stride(1),
+                k_cache.stride(2),
+                k_cache.stride(3),
+                k_cache.stride(
+                    4
+                ),  #[num_blocks, num_kv_heads, head_size/x, block_size, x]
+                v_cache.stride(0),
+                v_cache.stride(1),
+                v_cache.stride(2),
+                v_cache.stride(
+                    3),  #[num_blocks, num_kv_heads, head_size, block_size]
+                num_queries_per_kv=num_queries_per_kv,
+                BLOCK_M=BLOCK,
+                BLOCK_DMODEL=Lk,
+                BLOCK_N=BLOCK,
+                num_warps=num_warps,
+                num_stages=1,
+            )
+            return
+
+        vllm.attention.ops.prefix_prefill._fwd_kernel[grid](
+            q,
+            k,
+            v,
+            k_cache,
+            v_cache,
+            b_loc,
+            sm_scale,
+            b_start_loc,
+            b_seq_len,
+            b_ctx_len,
+            v_cache.shape[3],
+            8,
+            o,
+            b_loc.stride(0),
+            b_loc.stride(1),
+            q.stride(0),
+            q.stride(1),
+            q.stride(2),
+            k.stride(0),
+            k.stride(1),
+            k.stride(2),
+            v.stride(0),
+            v.stride(1),
+            v.stride(2),
+            o.stride(0),
+            o.stride(1),
+            o.stride(2),
+            k_cache.stride(0),
+            k_cache.stride(1),
+            k_cache.stride(2),
+            k_cache.stride(3),
+            k_cache.stride(
+                4),  #[num_blocks, num_kv_heads, head_size/x, block_size, x]
+            v_cache.stride(0),
+            v_cache.stride(1),
+            v_cache.stride(2),
+            v_cache.stride(
+                3),  #[num_blocks, num_kv_heads, head_size, block_size]
+            num_queries_per_kv=num_queries_per_kv,
+            BLOCK_M=BLOCK,
+            BLOCK_DMODEL=Lk,
+            BLOCK_DMODEL_PADDED=Lk_padded,
+            BLOCK_N=BLOCK,
+            SLIDING_WINDOW=sliding_window if sliding_window is not None else 0,
+            num_warps=num_warps,
+            num_stages=1,
+        )
+        return
+
+MluHijackObject.apply_hijack(vllm.attention.ops.prefix_prefill,
+                             vllm.attention.ops.prefix_prefill.context_attention_fwd,
+                             vllm__attention__ops__prefix_prefill__context_attention_fwd)
diff --git a/vllm_mlu/vllm_mlu/attention/ops/triton_flash_attention.py b/vllm_mlu/vllm_mlu/attention/ops/triton_flash_attention.py
new file mode 100644
index 000000000..2084104a0
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/attention/ops/triton_flash_attention.py
@@ -0,0 +1,802 @@
+#!/usr/bin/env python
+"""
+Fused Attention
+===============
+
+This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
+(https://tridao.me/publications/flash2/flash2.pdf)
+Credits: OpenAI kernel team, AMD ML Frameworks Triton team
+
+Features supported:
+
+1) Fwd with causal masking
+2) Any sequence lengths without padding (currently fwd kernel only)
+3) Support for different sequence lengths for q and k
+4) Nested tensor API currently does not support dropout or bias.
+
+Not currently supported:
+
+1) Non power of two head dims
+
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+torch_dtype: tl.constexpr = torch.float16
+
+
+@triton.jit
+def cdiv_fn(x, y):
+    return (x + y - 1) // y
+
+
+@triton.jit
+def max_fn(x, y):
+    return tl.math.max(x, y)
+
+
+@triton.jit
+def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride):
+    ms = tl.arange(0, m)
+    ns = tl.arange(0, n)
+    return philox_offset + ms[:, None] * stride + ns[None, :]
+
+
+@triton.jit
+def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride):
+    rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n,
+                                  stride).to(tl.uint32)
+    # TODO: use tl.randint for better performance
+    return tl.rand(philox_seed, rng_offsets)
+
+
+@triton.jit
+def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride):
+    rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n,
+                             stride)
+    rng_keep = rng_output > dropout_p
+    return rng_keep
+
+
+@triton.jit
+def load_fn(block_ptr, first, second, pad):
+    if first and second:
+        tensor = tl.load(block_ptr, boundary_check=(0, 1), padding_option=pad)
+    elif first:
+        tensor = tl.load(block_ptr, boundary_check=(0, ), padding_option=pad)
+    elif second:
+        tensor = tl.load(block_ptr, boundary_check=(1, ), padding_option=pad)
+    else:
+        tensor = tl.load(block_ptr)
+    return tensor
+
+
+@triton.jit
+def _attn_fwd_inner(
+    acc,
+    l_i,
+    m_i,
+    q,
+    K_block_ptr,
+    V_block_ptr,
+    start_m,
+    actual_seqlen_k,
+    dropout_p,
+    philox_seed,
+    batch_philox_offset,
+    encoded_softmax_block_ptr,
+    block_min,
+    block_max,
+    offs_n_causal,
+    masked_blocks,
+    n_extra_tokens,
+    bias_ptr,
+    IS_CAUSAL: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    OFFS_M: tl.constexpr,
+    OFFS_N: tl.constexpr,
+    PRE_LOAD_V: tl.constexpr,
+    MASK_STEPS: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    RETURN_ENCODED_SOFTMAX: tl.constexpr,
+    PADDED_HEAD: tl.constexpr,
+):
+    # loop over k, v, and update accumulator
+    for start_n in range(block_min, block_max, BLOCK_N):
+        # For padded blocks, we will overrun the tensor size if
+        # we load all BLOCK_N. For others, the blocks are all within range.
+        k = load_fn(
+            K_block_ptr,
+            PADDED_HEAD,
+            MASK_STEPS and (n_extra_tokens != 0),
+            "zero",
+        )
+        if PRE_LOAD_V:
+            v = load_fn(
+                V_block_ptr,
+                MASK_STEPS and (n_extra_tokens != 0),
+                PADDED_HEAD,
+                "zero",
+            )
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        # We start from end of seqlen_k so only the first iteration would need
+        # to be checked for padding if it is not a multiple of block_n
+        # TODO: This can be optimized to only be true for the padded block.
+        if MASK_STEPS:  # noqa: SIM102
+            # If this is the last block / iteration, we want to
+            # mask if the sequence length is not a multiple of block size
+            # a solution is to always do BLOCK_M // BLOCK_N + 1 steps
+            # if not is_modulo_mn. last step might get wasted but that is okay.
+            # check if this masking works for that case.
+            if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0):
+                boundary_m = tl.full([BLOCK_M],
+                                     actual_seqlen_k,
+                                     dtype=tl.int32)
+                size_n = start_n + OFFS_N[None, :]
+                mask = size_n < boundary_m[:, None]
+                qk = tl.where(mask, qk, float("-inf"))
+        if IS_CAUSAL:
+            causal_boundary = start_n + offs_n_causal
+            causal_mask = OFFS_M[:, None] >= causal_boundary[None, :]
+            qk = tl.where(causal_mask, qk, float("-inf"))
+        # -- compute qk ----
+        qk += tl.dot(q, k)
+        if bias_ptr is not None:
+            bias = load_fn(bias_ptr, False, MASK_STEPS
+                           and (n_extra_tokens != 0), "zero")
+            # While bias is added after multiplying qk with sm_scale, our
+            # optimization to use 2^x instead of e^x results in an additional
+            # scale factor of log2(e) which we must also multiply the bias with.
+            qk += bias * 1.44269504089
+        m_ij = tl.maximum(m_i, tl.max(qk, 1))
+        qk = qk - m_ij[:, None]
+        p = tl.math.exp2(qk)
+
+        # CAVEAT: Must update l_ij before applying dropout
+        l_ij = tl.sum(p, 1)
+        if ENABLE_DROPOUT:
+            philox_offset = (batch_philox_offset +
+                             start_m * BLOCK_M * actual_seqlen_k + start_n -
+                             BLOCK_N)
+            keep = dropout_mask(
+                philox_seed,
+                philox_offset,
+                dropout_p,
+                BLOCK_M,
+                BLOCK_N,
+                actual_seqlen_k,
+            )
+            if RETURN_ENCODED_SOFTMAX:
+                tl.store(
+                    encoded_softmax_block_ptr,
+                    tl.where(keep, p,
+                             -p).to(encoded_softmax_block_ptr.type.element_ty),
+                )
+            p = tl.where(keep, p, 0.0)
+        elif RETURN_ENCODED_SOFTMAX:
+            tl.store(
+                encoded_softmax_block_ptr,
+                p.to(encoded_softmax_block_ptr.type.element_ty),
+            )
+        # -- update output accumulator --
+        alpha = tl.math.exp2(m_i - m_ij)
+        acc = acc * alpha[:, None]
+        if not PRE_LOAD_V:
+            v = load_fn(
+                V_block_ptr,
+                MASK_STEPS and (n_extra_tokens != 0),
+                PADDED_HEAD,
+                "zero",
+            )
+        # -- update m_i and l_i
+        l_i = l_i * alpha + l_ij
+        # update m_i and l_i
+        m_i = m_ij
+        acc += tl.dot(p.to(V_block_ptr.type.element_ty), v)
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+        if bias_ptr is not None:
+            bias_ptr = tl.advance(bias_ptr, (0, BLOCK_N))
+        if RETURN_ENCODED_SOFTMAX:
+            encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
+                                                   (0, BLOCK_N))
+    return acc, l_i, m_i
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {
+                "BLOCK_M": 256,
+                "BLOCK_N": 64,
+                "PRE_LOAD_V": False,
+            },
+            num_stages=1,
+            num_warps=8,
+        ),
+        triton.Config(
+            {
+                "BLOCK_M": 128,
+                "BLOCK_N": 128,
+                "PRE_LOAD_V": False,
+            },
+            num_stages=1,
+            num_warps=4,
+        ),
+        triton.Config(
+            {
+                "BLOCK_M": 256,
+                "BLOCK_N": 128,
+                "PRE_LOAD_V": False,
+            },
+            num_stages=1,
+            num_warps=8,
+        ),
+        triton.Config(
+            {
+                "BLOCK_M": 128,
+                "BLOCK_N": 64,
+                "PRE_LOAD_V": True,
+            },
+            num_stages=1,
+            num_warps=4,
+        ),
+        triton.Config(
+            {
+                "BLOCK_M": 128,
+                "BLOCK_N": 64,
+                "PRE_LOAD_V": False,
+            },
+            num_stages=1,
+            num_warps=4,
+        ),
+        triton.Config(
+            {
+                "BLOCK_M": 64,
+                "BLOCK_N": 64,
+                "PRE_LOAD_V": False,
+            },
+            num_stages=1,
+            num_warps=8,
+        ),
+        triton.Config(
+            {
+                "BLOCK_M": 32,
+                "BLOCK_N": 32,
+                "PRE_LOAD_V": False,
+            },
+            num_stages=1,
+            num_warps=8,
+        ),
+        # TODO: This config fails with head_size not pow2 with data mismatches.
+        #    triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16,
+        #                   'PRE_LOAD_V': False}, num_stages=1, num_warps=4),
+        triton.Config(
+            {
+                "BLOCK_M": 16,
+                "BLOCK_N": 16,
+                "PRE_LOAD_V": False,
+            },
+            num_stages=1,
+            num_warps=4,
+        ),
+    ],
+    key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'],
+)
+@triton.jit
+def attn_fwd(
+    Q,
+    K,
+    V,
+    bias,
+    sm_scale,
+    L,
+    Out,
+    stride_qz,
+    stride_qh,
+    stride_qm,
+    stride_qk,
+    stride_kz,
+    stride_kh,
+    stride_kn,
+    stride_kk,
+    stride_vz,
+    stride_vh,
+    stride_vk,
+    stride_vn,
+    stride_oz,
+    stride_oh,
+    stride_om,
+    stride_on,
+    stride_bz,
+    stride_bh,
+    stride_bm,
+    stride_bn,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    dropout_p,
+    philox_seed,
+    philox_offset_base,
+    encoded_softmax,
+    HQ: tl.constexpr,
+    HK: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    MAX_SEQLENS_Q: tl.constexpr,
+    MAX_SEQLENS_K: tl.constexpr,
+    VARLEN: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    PRE_LOAD_V: tl.constexpr,
+    BIAS_TYPE: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    RETURN_ENCODED_SOFTMAX: tl.constexpr,
+):
+    start_m = tl.program_id(0)
+    off_h_q = tl.program_id(1)
+    off_z = tl.program_id(2)
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    if VARLEN:
+        cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
+        cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
+        seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
+        # We have a one-size-fits-all grid in id(0). Some seqlens might be too
+        # small for all start_m so for those we return early.
+        if start_m * BLOCK_M > seqlen_q:
+            return
+        cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
+        cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
+        seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
+    else:
+        cu_seqlens_q_start = 0
+        cu_seqlens_k_start = 0
+        seqlen_q = MAX_SEQLENS_Q
+        seqlen_k = MAX_SEQLENS_K
+
+    # Now we compute whether we need to exit early due to causal masking.
+    # This is because for seqlen_q > seqlen_k, M rows of the attn scores
+    # are completely masked, resulting in 0s written to the output, and
+    # inf written to LSE. We don't need to do any GEMMs in this case.
+    # This block of code determines what N is, and if this WG is operating
+    # on those M rows.
+    n_blocks = cdiv_fn(seqlen_k, BLOCK_N)
+    if IS_CAUSAL:
+        # If seqlen_q == seqlen_k, the attn scores are a square matrix.
+        # If seqlen_q != seqlen_k, attn scores are rectangular which means
+        # the causal mask boundary is bottom right aligned, and ends at either
+        # the top edge (seqlen_q < seqlen_k) or left edge.
+        # This captures the decrease in n_blocks if we have a rectangular attn
+        # matrix
+        n_blocks_seqlen = cdiv_fn(
+            (start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N)
+        # This is what adjusts the block_max for the current WG, only
+        # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks
+        n_blocks = min(n_blocks, n_blocks_seqlen)
+        # If we have no blocks after adjusting for seqlen deltas, this WG is
+        # part of the blocks that are all 0. We exit early.
+        if n_blocks <= 0:
+            o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om +
+                        off_h_q * stride_oh)
+            O_block_ptr = tl.make_block_ptr(
+                base=Out + o_offset,
+                shape=(seqlen_q, BLOCK_DMODEL),
+                strides=(stride_om, stride_on),
+                offsets=(start_m * BLOCK_M, 0),
+                block_shape=(BLOCK_M, BLOCK_DMODEL),
+                order=(1, 0),
+            )
+            acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty)
+            # We still need to write 0s to the result
+            # tl.store(O_block_ptr,
+            # acc.to(Out.type.element_ty), boundary_check=(0,1))
+            # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q
+            #          + offs_m
+            # We store inf to LSE, not -inf because in the bwd pass,
+            # we subtract this
+            # from qk which makes it -inf, such that exp(qk - inf) = 0
+            # for these masked blocks.
+            # l = tl.full([BLOCK_M], value=float("inf"), dtype=tl.float32)
+            # tl.store(l_ptrs, l)
+            # TODO: Should dropout and return encoded softmax be handled here?
+            return
+
+    # If MQA / GQA, set the K and V head offsets appropriately.
+    GROUP_SIZE: tl.constexpr = HQ // HK
+    off_h_k = off_h_q // GROUP_SIZE if GROUP_SIZE != 1 else off_h_q
+
+    n_extra_tokens = 0
+    if seqlen_k < BLOCK_N:
+        n_extra_tokens = BLOCK_N - seqlen_k
+    elif seqlen_k % BLOCK_N:
+        n_extra_tokens = seqlen_k % BLOCK_N
+    padded_head = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL
+
+    # Compute pointers for all the tensors used in this kernel.
+    q_offset = (off_z * stride_qz + off_h_q * stride_qh +
+                cu_seqlens_q_start * stride_qm)
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + q_offset,
+        shape=(seqlen_q, ACTUAL_BLOCK_DMODEL),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    k_offset = (off_z * stride_kz + off_h_k * stride_kh +
+                cu_seqlens_k_start * stride_kn)
+    K_block_ptr = tl.make_block_ptr(
+        base=K + k_offset,
+        shape=(ACTUAL_BLOCK_DMODEL, seqlen_k),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_DMODEL, BLOCK_N),
+        order=(0, 1),
+    )
+    v_offset = (off_z * stride_vz + off_h_k * stride_vh +
+                cu_seqlens_k_start * stride_vk)
+    V_block_ptr = tl.make_block_ptr(
+        base=V + v_offset,
+        shape=(seqlen_k, ACTUAL_BLOCK_DMODEL),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    if BIAS_TYPE != 0:
+        bias_ptr = tl.make_block_ptr(
+            base=bias + off_h_q * stride_bh,
+            shape=(seqlen_q, seqlen_k),
+            strides=(stride_bm, stride_bn),
+            offsets=(start_m * BLOCK_M, 0),
+            block_shape=(BLOCK_M, BLOCK_N),
+            order=(1, 0),
+        )
+    else:
+        bias_ptr = None
+    if ENABLE_DROPOUT:
+        batch_philox_offset = philox_offset_base \
+                              + (off_z * HQ + off_h_q) \
+                              * seqlen_q * seqlen_k
+    else:
+        batch_philox_offset = 0
+    # We can ask to return the dropout mask without actually doing any dropout.
+    # In this case, we return an invalid pointer so indicate the mask is not i
+    # valid.
+    # TODO: Fix encoded softmax. It currently uses just h_q in the base offset.
+    if RETURN_ENCODED_SOFTMAX:
+        encoded_softmax_block_ptr = tl.make_block_ptr(
+            base=encoded_softmax + off_h_q * seqlen_q * seqlen_k,
+            shape=(seqlen_q, seqlen_k),
+            strides=(seqlen_k, 1),
+            offsets=(start_m * BLOCK_M, 0),
+            block_shape=(BLOCK_M, BLOCK_N),
+            order=(1, 0),
+        )
+    else:
+        encoded_softmax_block_ptr = 0
+    # initialize pointer to m and l
+    m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
+    l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    # scale sm_scale by log_2(e) and use 2^x in the loop as we do not
+    # have native e^x support in HW.
+    qk_scale = sm_scale * 1.44269504089
+    # Q is loaded once at the beginning and shared by all N blocks.
+    q = load_fn(Q_block_ptr, True, padded_head, "zero")
+    q = (q * qk_scale).to(Q_block_ptr.type.element_ty)
+
+    # Here we compute how many full and masked blocks we have.
+    padded_block_k = n_extra_tokens != 0
+    is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0)
+    if IS_CAUSAL:
+        # There are always at least BLOCK_M // BLOCK_N masked blocks.
+        # Additionally there might be one more due to dissimilar seqlens.
+        masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn)
+    else:
+        # Padding on Q does not need to be masked in the FA loop.
+        masked_blocks = padded_block_k
+    # if IS_CAUSAL, not is_modulo_mn does not always result in an additional
+    # block. In this case we might exceed n_blocks so pick the min.
+    masked_blocks = min(masked_blocks, n_blocks)
+    n_full_blocks = n_blocks - masked_blocks
+    block_min = 0
+    block_max = n_blocks * BLOCK_N
+    # Compute for full blocks. Here we set causal to false regardless of its
+    # value because there is no masking. Similarly we do not need padding.
+    if n_full_blocks > 0:
+        block_max = (n_blocks - masked_blocks) * BLOCK_N
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            K_block_ptr,
+            V_block_ptr,
+            start_m,
+            seqlen_k,
+            dropout_p,
+            philox_seed,
+            batch_philox_offset,
+            encoded_softmax_block_ptr,
+            # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _
+            block_min,
+            block_max,
+            0,
+            0,
+            0,
+            bias_ptr,
+            # IS_CAUSAL, ....
+            False,
+            BLOCK_M,
+            BLOCK_DMODEL,
+            BLOCK_N,
+            offs_m,
+            offs_n,
+            # _, MASK_STEPS, ...
+            PRE_LOAD_V,
+            False,
+            ENABLE_DROPOUT,
+            RETURN_ENCODED_SOFTMAX,
+            padded_head,
+        )
+        block_min = block_max
+        block_max = n_blocks * BLOCK_N
+
+    tl.debug_barrier()
+    # Remaining blocks, if any, are full / not masked.
+    if masked_blocks > 0:
+        offs_n_causal = offs_n + (seqlen_q - seqlen_k) if IS_CAUSAL else 0
+        K_block_ptr = tl.advance(K_block_ptr, (0, n_full_blocks * BLOCK_N))
+        V_block_ptr = tl.advance(V_block_ptr, (n_full_blocks * BLOCK_N, 0))
+        if bias_ptr is not None:
+            bias_ptr = tl.advance(bias_ptr, (0, n_full_blocks * BLOCK_N))
+        if RETURN_ENCODED_SOFTMAX:
+            encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
+                                                   (0, n_full_blocks))
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            K_block_ptr,
+            V_block_ptr,
+            start_m,
+            seqlen_k,
+            dropout_p,
+            philox_seed,
+            batch_philox_offset,
+            encoded_softmax_block_ptr,
+            block_min,
+            block_max,
+            offs_n_causal,
+            masked_blocks,
+            n_extra_tokens,
+            bias_ptr,
+            IS_CAUSAL,
+            BLOCK_M,
+            BLOCK_DMODEL,
+            BLOCK_N,
+            offs_m,
+            offs_n,
+            # _, MASK_STEPS, ...
+            PRE_LOAD_V,
+            True,
+            ENABLE_DROPOUT,
+            RETURN_ENCODED_SOFTMAX,
+            padded_head,
+        )
+    # epilogue
+    acc = acc / l_i[:, None]
+    if ENABLE_DROPOUT:
+        acc = acc / (1 - dropout_p)
+    # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M,
+    # then we have one block with a row of all NaNs which come from computing
+    # softmax over a row of all -infs (-inf - inf = NaN). We check for that here
+    # and store 0s where there are NaNs as these rows should've been zeroed out.
+    end_m_idx = (start_m + 1) * BLOCK_M
+    start_m_idx = start_m * BLOCK_M
+    causal_start_idx = seqlen_q - seqlen_k
+    acc = acc.to(Out.type.element_ty)
+    if IS_CAUSAL:  # noqa: SIM102
+        if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx:
+            out_mask_boundary = tl.full((BLOCK_DMODEL, ),
+                                        causal_start_idx,
+                                        dtype=tl.int32)
+            mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)
+            out_ptrs_mask = (mask_m_offsets[:, None] >=
+                             out_mask_boundary[None, :])
+            z = 0.0
+            acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty))
+    # write back LSE
+    # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q + offs_m
+    # If seqlen_q not multiple of BLOCK_M, we need to mask out the last
+    # few rows. This is only true for the last M block. For others,
+    # overflow_size will be -ve
+    # overflow_size = end_m_idx - seqlen_q
+    # if overflow_size > 0:
+    #    boundary = tl.full((BLOCK_M,), BLOCK_M - overflow_size, dtype=tl.int32)
+    #    # This is a > check because mask being 0 blocks the store.
+    #    l_ptrs_mask = boundary > tl.arange(0, BLOCK_M)
+    #    tl.store(l_ptrs, m_i + tl.math.log2(l_i), mask=l_ptrs_mask)
+    # else:
+    #    tl.store(l_ptrs, m_i + tl.math.log2(l_i))
+
+    # write back O
+    o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om +
+                off_h_q * stride_oh)
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + o_offset,
+        shape=(seqlen_q, ACTUAL_BLOCK_DMODEL),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    # Need boundary check on this to make sure the padding from the
+    # Q and KV tensors in both dims are not part of what we store back.
+    # TODO: Do the boundary check optionally.
+    tl.store(O_block_ptr, acc, boundary_check=(0, 1))
+
+
+def check_args(
+    q,
+    k,
+    v,
+    o,
+    varlen=True,
+    max_seqlens=None,
+    cu_seqlens_q=None,
+    cu_seqlens_k=None,
+):
+    assert q.dim() == k.dim() and q.dim() == v.dim()
+    if varlen:
+        assert q.dim() == 3
+        total_q, nheads_q, head_size = q.shape
+        total_k, nheads_k, _ = k.shape
+        assert cu_seqlens_q is not None
+        assert cu_seqlens_k is not None
+        assert len(cu_seqlens_q) == len(cu_seqlens_k)
+    else:
+        assert q.dim() == 4
+        batch, nheads_q, seqlen_q, head_size = q.shape
+        _, nheads_k, seqlen_k, _ = k.shape
+        assert max_seqlens > 0
+    assert k.shape == v.shape
+    assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1]
+    # TODO: Change assert if we support qkl f8 and v f16
+    assert q.dtype == k.dtype and q.dtype == v.dtype
+    assert head_size <= 256
+    assert o.shape == q.shape
+    assert (nheads_q % nheads_k) == 0
+
+
+class _attention(torch.autograd.Function):
+
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        o,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlens_q,
+        max_seqlens_k,
+        causal=False,
+        sm_scale=1.0,
+        bias=None,
+    ):
+        if o is None:
+            o = torch.empty_like(q, dtype=v.dtype)
+
+        check_args(
+            q,
+            k,
+            v,
+            o,
+            varlen=True,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+        )
+        if True:  # varlen
+            total_q, nheads_q, head_size = q.shape
+            total_k, nheads_k, _ = k.shape
+            batch = len(cu_seqlens_q) - 1
+            q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
+            k_strides = (0, k.stride(1), k.stride(0), k.stride(2))
+            v_strides = (0, v.stride(1), v.stride(0), v.stride(2))
+            o_strides = (0, o.stride(1), o.stride(0), o.stride(2))
+        else:
+            batch, seqlen_q, nheads_q, head_size = q.shape
+            _, seqlen_k, nheads_k, _ = k.shape
+            q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))
+            k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))
+            v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))
+            o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
+
+        # Get closest power of 2 over or equal to 32.
+        unpadded_head_dims = {32, 64, 128, 256}
+        if head_size not in unpadded_head_dims:
+            padded_d_model = None
+            for i in unpadded_head_dims:
+                if i > head_size:
+                    padded_d_model = i
+                    break
+            assert padded_d_model is not None
+        else:
+            padded_d_model = head_size
+
+        grid = lambda META: (
+            triton.cdiv(max_seqlens_q, META["BLOCK_M"]),
+            nheads_q,
+            batch,
+        )
+
+        encoded_softmax = None
+
+        # Seed the RNG so we get reproducible results for testing.
+        philox_seed = 0x1BF52
+        philox_offset = 0x1D4B42
+
+        if bias is not None:
+            bias_strides = (
+                bias.stride(0),
+                bias.stride(1),
+                bias.stride(2),
+                bias.stride(3),
+            )
+        else:
+            bias_strides = (0, 0, 0, 0)
+
+        attn_fwd[grid](
+            q,
+            k,
+            v,
+            bias,
+            sm_scale,
+            None,
+            o,
+            *q_strides,
+            *k_strides,
+            *v_strides,
+            *o_strides,
+            *bias_strides,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            dropout_p=0.0,
+            philox_seed=philox_seed,
+            philox_offset_base=philox_offset,
+            encoded_softmax=encoded_softmax,
+            HQ=nheads_q,
+            HK=nheads_k,
+            ACTUAL_BLOCK_DMODEL=head_size,
+            MAX_SEQLENS_Q=max_seqlens_q,
+            MAX_SEQLENS_K=max_seqlens_k,
+            IS_CAUSAL=causal,
+            VARLEN=True,
+            BLOCK_DMODEL=padded_d_model,
+            BIAS_TYPE=0 if bias is None else 1,
+            ENABLE_DROPOUT=False,
+            RETURN_ENCODED_SOFTMAX=False,
+        )
+
+        ctx.grid = grid
+        ctx.sm_scale = sm_scale
+        ctx.BLOCK_DMODEL = head_size
+        ctx.causal = causal
+        ctx.dropout_p = 0.0
+        ctx.philox_seed = philox_seed
+        ctx.philox_offset = philox_offset
+        ctx.encoded_softmax = encoded_softmax
+        ctx.return_encoded_softmax = False
+        return o, encoded_softmax
+
+
+triton_attention = _attention.apply
diff --git a/vllm_mlu/vllm_mlu/config.py b/vllm_mlu/vllm_mlu/config.py
new file mode 100644
index 000000000..eac896a78
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/config.py
@@ -0,0 +1,278 @@
+from typing import Tuple
+
+import torch
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+from vllm.config import (ModelConfig, CacheConfig, LoRAConfig,
+                         VllmConfig, CompilationConfig, CompilationLevel,
+                         _DEFAULT_MAX_NUM_BATCHED_TOKENS)
+from vllm.utils import random_uuid
+from vllm_mlu._mlu_utils import is_dpsk_mcc_enabled
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+vllm__config__LoRAConfig__verify_with_model_config_org = LoRAConfig.verify_with_model_config
+
+
+def vllm__config__CacheConfig___verify_cache_dtype(self) -> None:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add kv_cache_dtype int8 support
+    '''
+    if self.cache_dtype == "auto":
+        pass
+    elif self.cache_dtype in ("fp8", "fp8_e4m3", "fp8_e5m2"):
+        logger.info(
+            "Using fp8 data type to store kv cache. It reduces the GPU "
+            "memory footprint and boosts the performance. "
+            "Meanwhile, it may cause accuracy drop without a proper "
+            "scaling factor")
+    elif self.cache_dtype == 'int8':
+        logger.info(
+            "Using int8 data type to store kv cache. It reduces the MLU "
+            "memory footprint and boosts the performance. ")
+    else:
+        raise ValueError(f"Unknown kv cache dtype: {self.cache_dtype}")
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__config__ModelConfig__get_head_size(self) -> int:
+    # TODO remove hard code
+    if self.is_deepseek_mla:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: force to return 576.
+        '''
+        # qk_rope_head_dim = getattr(self.hf_text_config, "qk_rope_head_dim",
+        #                             0)
+        # if self.use_mla:
+        #     return self.hf_text_config.kv_lora_rank + qk_rope_head_dim
+        # else:
+        #     qk_nope_head_dim = getattr(self.hf_text_config,
+        #                                 "qk_nope_head_dim", 0)
+        #     if qk_rope_head_dim and qk_nope_head_dim:
+        #         return qk_rope_head_dim + qk_nope_head_dim
+        return 576
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+    if hasattr(self.hf_text_config,
+                "model_type") and (self.hf_text_config.model_type
+                                    == "zamba2"):
+        return self.hf_text_config.attention_head_dim
+
+    if self.is_attention_free:
+        return 0
+
+    if hasattr(self.hf_text_config, "head_dim"):
+        return self.hf_text_config.head_dim
+    # FIXME(woosuk): This may not be true for all models.
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: adjust num_heads and num_attention_heads.
+    '''
+    if hasattr(self.hf_text_config, "num_heads"):
+        num_attention_heads = self.hf_text_config.num_heads
+    else:
+        num_attention_heads = self.hf_text_config.num_attention_heads
+
+    return (self.hf_text_config.hidden_size // num_attention_heads)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__config__ModelConfig__set_context_mlugraph_info(
+    self,
+    enable_context_mlugraph: bool,
+    batch_size: int,
+    seq_len: int
+) -> None:
+    self.enable_context_mlugraph = enable_context_mlugraph
+    self.context_batch_size_to_capture = batch_size
+    self.context_seq_len_to_capture = seq_len
+
+
+def vllm__config__ModelConfig__use_context_mlugraph(self) -> bool:
+    return hasattr(self, "enable_context_mlugraph") and self.enable_context_mlugraph
+
+
+def vllm__config__ModelConfig__get_context_mlugraph_bs_and_seq(self) -> Tuple[int, int]:
+    return self.context_batch_size_to_capture, self.context_seq_len_to_capture
+
+
+def vllm__config__ModelConfig__is_embedding_task(self) -> bool:
+    return self.runner_type == "pooling"
+
+
+def vllm__config__LoRAConfig__verify_with_model_config(self, model_config: ModelConfig):
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: do not support quantization with lora for now
+    '''
+    if model_config.quantization:
+        raise ValueError("MLU backend does not support quantization with lora for now")
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    vllm__config__LoRAConfig__verify_with_model_config_org(self, model_config)
+
+
+class VllmConfig_Mluhijack(VllmConfig):
+
+    def __post_init__(self):
+        """Verify configs are valid & consistent with each other.
+        """
+        if self.model_config is not None:
+            self.model_config.verify_async_output_proc(self.parallel_config,
+                                                       self.speculative_config,
+                                                       self.device_config)
+            self.model_config.verify_with_parallel_config(self.parallel_config)
+
+        if self.cache_config is not None:
+            self.cache_config.verify_with_parallel_config(self.parallel_config)
+
+        if self.lora_config:
+            self.lora_config.verify_with_cache_config(self.cache_config)
+            self.lora_config.verify_with_model_config(self.model_config)
+            self.lora_config.verify_with_scheduler_config(
+                self.scheduler_config)
+        if self.prompt_adapter_config:
+            self.prompt_adapter_config.verify_with_model_config(
+                self.model_config)
+
+        if self.quant_config is None and \
+            self.model_config is not None and self.load_config is not None:
+            self.quant_config = VllmConfig._get_quantization_config(
+                self.model_config, self.load_config)
+
+        from vllm.platforms import current_platform
+        if self.scheduler_config is not None and \
+            self.model_config is not None and \
+            self.scheduler_config.chunked_prefill_enabled and \
+            self.model_config.dtype == torch.float32 and \
+            current_platform.get_device_capability() == (7, 5):
+            logger.warning_once(
+                "Turing devices tensor cores do not support float32 matmul. "
+                "To workaround this limitation, vLLM will set 'ieee' input "
+                "precision for chunked prefill triton kernels.")
+
+        if self.compilation_config is None:
+            self.compilation_config = CompilationConfig()
+        if envs.VLLM_USE_V1 and self.model_config is not None and \
+            not self.model_config.enforce_eager:
+            # NOTE(woosuk): Currently, we use inductor because the piecewise
+            # CUDA graphs do not work properly with the custom CUDA kernels.
+            # FIXME(woosuk): Disable inductor to reduce the compilation time
+            # and avoid any potential issues with the inductor.
+            # FIXME(rob): Add function to set all of these.
+            self.compilation_config.custom_ops = ["none"]
+            self.compilation_config.use_cudagraph = True
+            self.compilation_config.use_inductor = True
+            self.compilation_config.cudagraph_num_of_warmups = 1
+            self.compilation_config.pass_config.enable_fusion = False
+            self.compilation_config.pass_config.enable_noop = False
+            self.compilation_config.level = CompilationLevel.PIECEWISE
+            self.compilation_config.set_splitting_ops_for_v1()
+
+        self._set_cudagraph_sizes()
+
+        if self.cache_config is not None and \
+            self.cache_config.cpu_offload_gb > 0 and \
+            self.compilation_config.level != CompilationLevel.NO_COMPILATION \
+                and not envs.VLLM_USE_V1:
+            logger.warning(
+                "CPU offload is not supported with `torch.compile` in v0 yet."
+                " Disabling `torch.compile`.")
+            self.compilation_config.level = CompilationLevel.NO_COMPILATION
+
+        if ((not envs.VLLM_USE_V1) and self.lora_config is not None
+                and self.compilation_config.level
+                != CompilationLevel.NO_COMPILATION):
+            logger.warning(
+                "LoRA for V0 is not supported with `torch.compile` yet. "
+                "Disabling `torch.compile`.")
+            self.compilation_config.level = CompilationLevel.NO_COMPILATION
+
+        if self.model_config and self.model_config.use_mla and \
+            not (current_platform.is_cuda() or current_platform.is_rocm()):
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: do not modify max_num_batched_tokens values.
+            @brief: MLU also supports chunked prefill.
+            '''
+            if not current_platform.is_out_of_tree() or is_dpsk_mcc_enabled(
+            ) or (self.speculative_config is not None
+                  and self.speculative_config.num_speculative_tokens
+                  > 0) or self.scheduler_config.num_scheduler_steps > 1:
+                self.scheduler_config.enable_chunked_prefill = False
+                self.scheduler_config.chunked_prefill_enabled = False
+                logger.warning(
+                    "Chunked prefill is disabled on a non-MLU platform, or dpsk mcc or mtp is enabled,"
+                    " or num_scheduler_steps > 1, forcing chunked prefill and prefix caching to be disabled.")
+
+            # self.scheduler_config.max_num_batched_tokens = max(
+            #     self.scheduler_config.max_model_len,
+            #     _DEFAULT_MAX_NUM_BATCHED_TOKENS)
+
+            if self.cache_config is not None:
+                self.cache_config.enable_prefix_caching = False
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+        current_platform.check_and_update_config(self)
+
+        if not self.instance_id:
+            self.instance_id = random_uuid()[:5]
+
+
+MluHijackObject.apply_hijack(ModelConfig,
+                             "set_context_mlugraph_info",
+                             vllm__config__ModelConfig__set_context_mlugraph_info)
+MluHijackObject.apply_hijack(ModelConfig,
+                             "use_context_mlugraph",
+                             vllm__config__ModelConfig__use_context_mlugraph)
+MluHijackObject.apply_hijack(ModelConfig,
+                             "get_context_mlugraph_bs_and_seq",
+                             vllm__config__ModelConfig__get_context_mlugraph_bs_and_seq)
+MluHijackObject.apply_hijack(ModelConfig,
+                             "is_embedding_task",
+                             vllm__config__ModelConfig__is_embedding_task)
+MluHijackObject.apply_hijack(CacheConfig,
+                             CacheConfig._verify_cache_dtype,
+                             vllm__config__CacheConfig___verify_cache_dtype)
+MluHijackObject.apply_hijack(ModelConfig,
+                             ModelConfig.get_head_size,
+                             vllm__config__ModelConfig__get_head_size)
+MluHijackObject.apply_hijack(LoRAConfig,
+                             LoRAConfig.verify_with_model_config,
+                             vllm__config__LoRAConfig__verify_with_model_config)
+MluHijackObject.apply_hijack(VllmConfig,
+                             VllmConfig.__post_init__,
+                             VllmConfig_Mluhijack.__post_init__)
diff --git a/vllm_mlu/vllm_mlu/core/__init__.py b/vllm_mlu/vllm_mlu/core/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/core/scheduler.py b/vllm_mlu/vllm_mlu/core/scheduler.py
new file mode 100644
index 000000000..7f763018a
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/core/scheduler.py
@@ -0,0 +1,328 @@
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from collections import deque
+from typing import Deque, List, Optional, Set, Callable
+from vllm.config import CacheConfig, LoRAConfig, SchedulerConfig
+from vllm.core.scheduler import (SchedulingBudget, SchedulerPrefillOutputs,
+                                 SchedulerRunningOutputs, SchedulerOutputs, Scheduler)
+from vllm.sequence import SequenceGroup
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+vllm__core__scheduler__Scheduler____init____org = Scheduler.__init__
+vllm__core__scheduler__Scheduler___schedule_prefills__org = Scheduler._schedule_prefills
+vllm__core__scheduler__Scheduler___schedule_running__org = Scheduler._schedule_running
+vllm__core__scheduler__Scheduler___schedule__org = Scheduler._schedule
+
+
+def vllm__core__scheduler__Scheduler__init_scheduler_view(self):
+    logger.info(f"vLLM scheduler profiling start...")
+    
+    self.df = pd.DataFrame(
+        data={
+            'waiting': [], 'running': [], 'swapped': [], 'finished': [],
+            'wait_to_run_reqs': [], 'run_to_wait_reqs': [], 'wait_to_run_tokens': [],
+            'batch_utils': [], 'block_utils': [], 'preempt_ratio': []
+        },
+        dtype=np.float32
+    )
+    self.sched_step = 0
+    self.running_seqs = 0
+    self.waiting_seqs = 0
+    self.swapped_seqs = 0
+    self.finished_seqs = 0
+    self.total_seqs = 0
+    self.running_to_waiting_seqs = 0
+    self.waiting_to_running_seqs = 0
+    self.wait_to_run_tokens = 0
+    self.batch_utils = 0
+    self.block_utils = 0
+    self.preempt_ratio = 0
+
+    self.finished_seq_groups = []
+
+
+def summary_finished_seq_groups(seq_groups: List[SequenceGroup]):
+    df = pd.DataFrame(
+        data={
+            'ttft/s': [], 'time_in_queue/s': [], 'context_latency/s': [], 'decoder_latency/s': []
+        },
+        dtype=np.float32
+    )
+    for seq_group in seq_groups:
+        ttft = seq_group.metrics.first_token_time - seq_group.metrics.arrival_time
+        time_in_queue = seq_group.metrics.time_in_queue
+        context_latency = seq_group.metrics.first_token_time - seq_group.metrics.first_scheduled_time
+        decoder_latency = seq_group.metrics.finished_time - seq_group.metrics.first_token_time
+        decoder_token_num = seq_group.get_seqs()[0].get_output_len() - 1
+        per_token_latency = decoder_latency if decoder_token_num == 0 \
+                                            else decoder_latency / decoder_token_num
+        df_ = pd.DataFrame(
+            [[ttft, time_in_queue, context_latency, decoder_latency, per_token_latency, decoder_token_num]],
+            columns=['ttft/s', 'time_in_queue/s', 'context_latency/s', 'decoder_latency/s', 'per_token_latency/s', 'decoder_tokens'],
+            index=[str(seq_group.request_id)]
+        )
+        df = pd.concat([df, df_])
+    sum_, max_, mean_, min_, p99_ = df.sum(), df.max(), df.mean(), df.min(), df.quantile(0.99)
+    df.loc['Sum'] = sum_
+    df.loc['Max'] = max_
+    df.loc['Mean'] = mean_
+    df.loc['Min'] = min_
+    df.loc['P99'] = p99_
+    return df
+
+
+def vllm__core__scheduler__Scheduler__save_scheduler_view(self, scheduler_idx=0):
+    logger.info(f"vLLM scheduler profiling save...")
+    plt.rcParams.update({'font.size': 8})
+    figure = plt.figure(figsize=(6.4, 5.6))
+    gs = figure.add_gridspec(3, hspace=0)
+    axes = gs.subplots(sharex=True, sharey=False)
+    figure.suptitle("Cambricon vLLM Scheduler View")
+    # scheduler queue view
+    self.df.plot(ax=axes[0], y=['waiting', 'running', 'swapped', 'finished'])
+    axes[0].set_xlabel('X-LLMEngineStep', loc='left')
+    axes[0].set_ylabel('Y-ReqNum', loc='top')
+    # utilization
+    self.df.plot(ax=axes[1], y=['batch_utils', 'block_utils', 'preempt_ratio'])
+    axes[1].set_xlabel('X-LLMEngineStep', loc='left')
+    axes[1].set_ylabel('Y-Utilization(%)', loc='top')
+    # token view
+    self.df.plot(ax=axes[2], y=['wait_to_run_tokens'])
+    axes[2].set_xlabel('X-LLMEngineStep', loc='left')
+    axes[2].set_ylabel('Y-TokenNum', loc='top')
+    for ax in axes:
+        ax.label_outer()
+        ax.legend(loc='upper right')
+    figure.tight_layout()
+    figure.savefig(f"vllm_scheduler{scheduler_idx}_view.svg", dpi=300, format='svg')
+    plt.close(figure)
+
+    time_df = summary_finished_seq_groups(self.finished_seq_groups)
+
+    sched_df = self.df.copy(deep=True)
+    max_, mean_, min_ = sched_df.max(), sched_df.mean(), sched_df.min()
+    sched_df.loc["Max"] = max_
+    sched_df.loc["Mean"] = mean_
+    sched_df.loc["Min"] = min_
+    with pd.option_context('display.max_rows', None,
+                           'display.max_columns', None,
+                           'display.max_colwidth', None,
+                           'display.float_format', '{:^6,.2f}'.format,
+                           'expand_frame_repr', False):
+        logger.info(sched_df.loc[["Max", "Mean", "Min"]])
+        logger.info(time_df.loc[["Sum", "Max", "Mean", "Min", "P99"]])
+    sched_df.astype(str).to_csv(f"vllm_scheduler{scheduler_idx}_step_view.csv", mode="w")
+    time_df.astype(str).to_csv(f"vllm_scheduler{scheduler_idx}_reqs_view.csv", mode="w")
+
+
+def vllm__core__scheduler__Scheduler____init__(
+    self,
+    scheduler_config: SchedulerConfig,
+    cache_config: CacheConfig,
+    lora_config: Optional[LoRAConfig],
+    pipeline_parallel_size: int = 1,
+    output_proc_callback: Optional[Callable] = None,
+) -> None:
+    vllm__core__scheduler__Scheduler____init____org(
+        self=self,
+        scheduler_config=scheduler_config,
+        cache_config=cache_config,
+        lora_config=lora_config,
+        pipeline_parallel_size=pipeline_parallel_size,
+        output_proc_callback=output_proc_callback
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add for scheduler profiling
+    '''
+    self.init_scheduler_view()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__core__scheduler__Scheduler___schedule_prefills(
+    self,
+    budget: SchedulingBudget,
+    curr_loras: Optional[Set[int]],
+    enable_chunking: bool = False,
+) -> SchedulerPrefillOutputs:
+    prefills = vllm__core__scheduler__Scheduler___schedule_prefills__org(
+        self=self,
+        budget=budget,
+        curr_loras=curr_loras,
+        enable_chunking=enable_chunking
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add for scheduler profiling
+    '''
+    self.waiting_to_running_seqs = len(prefills.seq_groups)
+    self.wait_to_run_tokens = budget.num_batched_tokens
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return prefills
+
+
+def vllm__core__scheduler__Scheduler___schedule_running(
+    self,
+    budget: SchedulingBudget,
+    curr_loras: Optional[Set[int]],
+    enable_chunking: bool = False,
+) -> SchedulerRunningOutputs:
+    running_scheduled = vllm__core__scheduler__Scheduler___schedule_running__org(
+        self=self,
+        budget=budget,
+        curr_loras=curr_loras,
+        enable_chunking=enable_chunking
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add for scheduler profiling
+    '''
+    self.running_to_waiting_seqs += len(running_scheduled.preempted)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return running_scheduled
+
+
+def vllm__core__scheduler__Scheduler___schedule(self) -> SchedulerOutputs:
+    scheduler_outputs = vllm__core__scheduler__Scheduler___schedule__org(self)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add for scheduler profiling
+    '''
+    self.sched_step += 1
+    self.running_seqs = len(self.running)
+    self.waiting_seqs = len(self.waiting)
+    self.swapped_seqs = len(self.swapped)
+
+    total_seqs_ = self.running_seqs + self.waiting_seqs + self.swapped_seqs + self.finished_seqs
+    if total_seqs_ == 0:
+        return
+
+    if total_seqs_ > self.total_seqs:
+        self.total_seqs = total_seqs_
+
+    self.batch_utils = self.running_seqs / self.scheduler_config.max_num_seqs
+    self.block_utils = (self.block_manager.num_total_gpu_blocks -
+                        self.block_manager.get_num_free_gpu_blocks()) / self.block_manager.num_total_gpu_blocks
+    self.preempt_ratio = self.running_to_waiting_seqs / self.total_seqs
+
+    df_ = pd.DataFrame(
+        [[self.waiting_seqs, self.running_seqs, self.swapped_seqs,
+          self.waiting_to_running_seqs, self.running_to_waiting_seqs, self.wait_to_run_tokens,
+          self.batch_utils, self.block_utils, self.preempt_ratio]],
+        columns=['waiting', 'running', 'swapped',
+                 'wait_to_run_reqs', 'run_to_wait_reqs', 'wait_to_run_tokens',
+                 'batch_utils', 'block_utils', 'preempt_ratio'],
+        index=[str(self.sched_step)])
+    self.df = pd.concat([self.df, df_])
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return scheduler_outputs
+
+
+def vllm__core__scheduler__Scheduler__free_finished_seq_groups(self) -> None:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add for scheduler profiling
+    '''
+    finished_seq_groups_ = []
+    remaining: Deque[SequenceGroup] = deque()
+    for seq_group in self.running:
+        self._free_finished_seq_group(seq_group)
+        if not seq_group.is_finished():
+            remaining.append(seq_group)
+        else:
+            finished_seq_groups_.append(seq_group)
+
+    self.finished_seqs += len(finished_seq_groups_)
+    self.finished_seq_groups += finished_seq_groups_
+
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    self.running = remaining
+
+    # Handle async stopped sequence groups
+    # (ones that reached max model len)
+    if self._async_stopped:
+        for seq_group in self._async_stopped:
+            self._free_seq_group_cross_attn_blocks(seq_group)
+            self._finished_requests_ids.append(seq_group.request_id)
+
+            # Free finished seqs
+            self._free_finished_seqs(seq_group)
+
+        self._async_stopped.clear()
+
+
+def vllm__core__scheduler__Scheduler____del__(self):
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add for scheduler profiling
+    '''
+    self.save_scheduler_view()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(Scheduler,
+                             Scheduler.__init__,
+                             vllm__core__scheduler__Scheduler____init__)
+MluHijackObject.apply_hijack(Scheduler,
+                             Scheduler._schedule_prefills,
+                             vllm__core__scheduler__Scheduler___schedule_prefills)
+MluHijackObject.apply_hijack(Scheduler,
+                             Scheduler._schedule_running,
+                             vllm__core__scheduler__Scheduler___schedule_running)
+MluHijackObject.apply_hijack(Scheduler,
+                             Scheduler._schedule,
+                             vllm__core__scheduler__Scheduler___schedule)
+MluHijackObject.apply_hijack(Scheduler,
+                             Scheduler.free_finished_seq_groups,
+                             vllm__core__scheduler__Scheduler__free_finished_seq_groups)
+MluHijackObject.apply_hijack(Scheduler,
+                             "__del__",
+                             vllm__core__scheduler__Scheduler____del__)
+MluHijackObject.apply_hijack(Scheduler,
+                             "init_scheduler_view",
+                             vllm__core__scheduler__Scheduler__init_scheduler_view)
+MluHijackObject.apply_hijack(Scheduler,
+                             "save_scheduler_view",
+                             vllm__core__scheduler__Scheduler__save_scheduler_view)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/device_info.py b/vllm_mlu/vllm_mlu/device_info.py
new file mode 100644
index 000000000..559785802
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/device_info.py
@@ -0,0 +1,297 @@
+import os
+from vllm.logger import init_logger
+from vllm_mlu.mlu_hijack_utils import TypedDict
+from vllm_mlu._mlu_utils import VLLM_DUMP_MLU_INFO_EN
+from typing import List, Tuple, Dict, Union
+import re
+import subprocess 
+import logging
+from vllm_mlu.dump_info import get_deepseek_v2_flops
+
+logger = init_logger(__name__)
+
+data_type_byte_width_map = {
+                            "int8":1,
+                            "int4":0.5,
+                            "float16":2,
+                            "float32":4,
+                            "int32":4,
+                            "bfloat16":2
+                            }
+
+HFUInfo = {
+    "context_hfu":0.0,
+    "decoder_hfu":0.0
+    }
+
+def str_to_data_type(param):
+    data_type_map = {
+        "float16": "CNNL_DTYPE_HALF",
+        "float32": "CNNL_DTYPE_FLOAT",
+        "int32": "CNNL_DTYPE_INT32",
+        "int8": "CNNL_DTYPE_INT8",
+        "int4": "CNNL_DTYPE_INT8", 
+        "bfloat16": "CNNL_DTYPE_BFLOAT16",
+        "invalid": "CNNL_DTYPE_INVALID",
+        "fp8_e4m3fn": "CNNL_DTYPE_FLOAT8_E4M3FN",
+        "fp8_e5m2": "CNNL_DTYPE_FLOAT8_E5M2",
+    }
+
+    return data_type_map.get(param, "CNNL_DTYPE_INVALID")
+
+def get_deepseek_v2_model_param(bcfg, expert_param, ffn_param):
+    num_attention_heads = bcfg["num_attention_heads"]
+    qk_nope_head_dim = bcfg["qk_nope_head_dim"]
+    qk_rope_head_dim = bcfg["qk_rope_head_dim"]
+    v_head_dim = bcfg["v_head_dim"]
+    q_lora_rank = bcfg["q_lora_rank"]
+    kv_lora_rank = bcfg["kv_lora_rank"]
+    hidden_size = bcfg["hidden_size"]
+    layer_num = bcfg["layer_num"]
+
+    qa_params = hidden_size * q_lora_rank
+    kva_params = hidden_size * (kv_lora_rank + qk_rope_head_dim)
+    qb_params = q_lora_rank * num_attention_heads * (qk_nope_head_dim + qk_rope_head_dim)
+    kvb_params = kv_lora_rank * num_attention_heads * (qk_nope_head_dim + v_head_dim)
+    o_params = v_head_dim * num_attention_heads * hidden_size
+
+    mla_weight_params = (qa_params + kva_params + qb_params + kvb_params + o_params) * layer_num
+    model_params = mla_weight_params + expert_param + ffn_param
+
+    return model_params
+
+def get_model_param(bcfg, once_batch):
+    batch = once_batch
+    hidden_size = bcfg["hidden_size"]
+    ffn_inner_size = bcfg["ffn_inner_size"]
+    moe_inner_size = bcfg["moe_inner_size"]
+    layer_num = bcfg["layer_num"]
+    r = bcfg["head_num"] / bcfg["head_num_kv"]
+    layer_moe_num = bcfg["moe_layer_num"]
+    layer_ffn_num = layer_num - layer_moe_num
+    experts_num = bcfg["experts_num"]
+    topk_num = bcfg["topk_num"]
+    shared_expert_intermediate_size = bcfg["shared_expert_intermediate_size"]
+    expert_num = min(experts_num, batch * topk_num)
+    ffn_num = 3.0 if bcfg["use_gated_ffn"] else 2.0
+
+    # shared_expert_intermediate_size = shared_expert_num * moe_intermediate_size,
+    # when adding a new moe model, must fix it when parse model.json
+    # in moe model may ffn & moe layers appear at the same time, need compute them devided
+    expert_param = layer_moe_num * ffn_num * hidden_size * (expert_num * moe_inner_size + shared_expert_intermediate_size)
+    ffn_param = layer_ffn_num * ffn_num * hidden_size * ffn_inner_size
+
+    if bcfg["model_type"] == "deepseek_v2":
+        # for deepseek_v2, using MLA replace attention，compute mla_weight_param here.
+        return get_deepseek_v2_model_param(bcfg, expert_param, ffn_param)
+
+    cla_coeffient = bcfg["cla_coeffient"]
+    attn_param = (2.0 + 2.0 / r * cla_coeffient) * layer_num * hidden_size * hidden_size
+
+    return expert_param + ffn_param + attn_param
+
+def get_deepseek_v2_kv_cache(bcfg, input_seq_len, output_seq_len, layer_num, kv_cache_byte_width):
+    num_attention_heads = bcfg["head_num"]
+    qk_nope_head_dim = bcfg["qk_nope_head_dim"]
+    qk_rope_head_dim = bcfg["qk_rope_head_dim"]
+    v_head_dim = bcfg["v_head_dim"]
+
+    kv_cache = num_attention_heads * (qk_nope_head_dim + qk_rope_head_dim + v_head_dim) * \
+               (input_seq_len + output_seq_len / 2) * layer_num * kv_cache_byte_width
+    return kv_cache
+
+def get_kv_cache(bcfg, once_batch, input_seq_len_, output_length):
+    batch          = once_batch
+    head_num_kv    = bcfg["head_num_kv"]
+    head_size      = bcfg["head_size"]
+    tp_num         = bcfg["tp_num"]
+    layer_num      = bcfg["layer_num"]
+    input_seq_len  = input_seq_len_
+    output_seq_len = output_length
+    kv_cache_byte_width = data_type_byte_width_map[bcfg["kv_cache_dtype"]]
+    if bcfg["model_type"] == "deepseek_v2":
+        return get_deepseek_v2_kv_cache(bcfg, input_seq_len,
+                                    output_seq_len, layer_num, kv_cache_byte_width)
+    cla_coeffient = bcfg["cla_coeffient"]
+    kv_cache = cla_coeffient * batch * max(head_num_kv, tp_num) * head_size * (input_seq_len + output_seq_len / 2) * 2 * layer_num * kv_cache_byte_width
+    return kv_cache
+    
+    
+def get_decoder_io_efficiency(hfu_model_config, batch_size, input_len, output_len, generate_latency):
+    model_param = get_model_param(hfu_model_config, batch_size)
+    model_byte_width = data_type_byte_width_map[hfu_model_config["filter_data_type"]]
+    param_scale = model_param * model_byte_width
+    kv_cache = get_kv_cache(hfu_model_config, batch_size, input_len, output_len)
+    # '(param_scale + kv_cache)': the number of bytes loaded during the decoder stage
+    # '(param_scale + kv_cache) / unit / unit ': the number of Megabytes loaded during the decoder stage
+    # '(param_scale + kv_cache) / unit / unit / bandwidth * unit': the last 'unit' is used to convert seconds to milliseconds of theoretical loading time
+    unit = 1000.0
+    _, _, _, bandwidth = get_device_attribute()
+    memory_MB = (param_scale + kv_cache) / unit / unit
+    io_efficiency = memory_MB / bandwidth * unit / generate_latency / hfu_model_config["tp_num"]
+
+    return io_efficiency
+
+
+def get_flops(bcfg, once_batch, input_seq_len, output_length, card_num, each_peak_compute, hfu_info):
+    batch = once_batch
+    seq_len = input_seq_len
+    hidden_size = bcfg["hidden_size"]
+    voc_size = bcfg["vocab_size"]
+    ffn_size = bcfg["ffn_inner_size"]
+    moe_size = bcfg["moe_inner_size"]
+    shared_expert_intermediate_size = bcfg["shared_expert_intermediate_size"]
+    layer_num = bcfg["layer_num"]
+    out_seq = output_length
+    seq_len_decode = seq_len + out_seq / 2
+    r = bcfg["head_num"] / bcfg["head_num_kv"]
+    bsh2 = batch * seq_len * hidden_size * hidden_size
+    cla_coeffient = bcfg["cla_coeffient"]
+
+    if bcfg["model_type"] in ['deepseek_v2','deepseek_v3']:
+        context_atn_pre, context_atn_qk, context_atn_qkv, context_atn_post = (
+            get_deepseek_v2_flops(bcfg, batch, seq_len, hidden_size)
+        )
+    else:
+        context_atn_pre = 2 * bsh2 + 4 * bsh2 / r * cla_coeffient
+        context_atn_qk = 2 * batch * seq_len * seq_len * hidden_size
+        context_atn_qkv = 2 * batch * seq_len * seq_len * hidden_size
+        context_atn_post = 2 * batch * seq_len * hidden_size * hidden_size
+    context_lm_head = 2 * batch * seq_len * hidden_size * voc_size
+    context_ffn = 0
+    bh2 = batch * hidden_size * hidden_size
+    decode_atn_pre = 2 * bh2 + 4 * bh2 / r * cla_coeffient
+    decode_atn_qk = 2 * batch * seq_len_decode * hidden_size
+    decode_atn_qkv = 2 * batch * seq_len_decode * hidden_size
+    decode_atn_post = 2 * batch * hidden_size * hidden_size
+    decode_lm_head = 2 * batch * hidden_size * voc_size
+    decode_ffn = 0
+    coeffient = 6 if bcfg["use_gated_ffn"] else 4
+    if bcfg["experts_num"] == 0:
+        context_ffn = coeffient * batch * seq_len * hidden_size * ffn_size
+        decode_ffn = coeffient * batch * hidden_size * ffn_size
+    else:
+        context_ffn = batch * seq_len * hidden_size * (coeffient * (moe_size * bcfg["topk_num"] + shared_expert_intermediate_size) + 2 * bcfg["experts_num"])
+        decode_ffn = batch * hidden_size * (coeffient * (moe_size * bcfg["topk_num"] + shared_expert_intermediate_size) + 2 * bcfg["experts_num"])
+
+    if bcfg["use_causal_mask"]:
+        c = 0.5
+        context_atn_qk *= c
+        context_atn_qkv *= c
+
+    total_peak_compute = each_peak_compute
+    for i in range(len(total_peak_compute)):
+        total_peak_compute[i] = card_num * each_peak_compute[i] + 1e-6
+
+    hfu_info["context_hfu"] = context_lm_head / total_peak_compute[1]
+    hfu_info["decoder_hfu"] = decode_lm_head / total_peak_compute[1]
+    if bcfg["kv_cache_dtype"] != "int8":
+        hfu_info["context_hfu"] += (layer_num * (context_atn_qk + context_atn_qkv) / total_peak_compute[1])
+        hfu_info["decoder_hfu"] += (layer_num * (decode_atn_qk + decode_atn_qkv) / total_peak_compute[1])
+    else:
+        hfu_info["context_hfu"] += (layer_num * (context_atn_qk + context_atn_qkv) / total_peak_compute[0])
+        hfu_info["decoder_hfu"] += (layer_num * (decode_atn_qk + decode_atn_qkv) / total_peak_compute[0])
+
+    if bcfg["smooth_quant_type"] == "invalid":
+        hfu_info["context_hfu"] += (layer_num * (context_atn_pre + context_atn_post + context_ffn) / total_peak_compute[1])
+        hfu_info["decoder_hfu"] += (layer_num * (decode_atn_pre + decode_atn_post + decode_ffn) / total_peak_compute[1])
+    else:
+        hfu_info["context_hfu"] += (layer_num * (context_atn_pre + context_atn_post + context_ffn) / total_peak_compute[2])
+        hfu_info["decoder_hfu"] += (layer_num * (decode_atn_pre + decode_atn_post + decode_ffn) / total_peak_compute[2])
+
+
+def get_flops_inner(hfu_model_config, batch_size, input_seq_len, output_length, tp_num, hfu_info):
+    if VLLM_DUMP_MLU_INFO_EN:
+        each_peak_compute = [1.0, 1.0, 1.0]
+        get_hfu_peak_flops(hfu_model_config, each_peak_compute)
+        tmp = [f"{num:.5e}" for num in each_peak_compute]
+        get_flops(hfu_model_config, batch_size, input_seq_len, output_length, tp_num, each_peak_compute, hfu_info)
+    else:
+        each_peak_compute = [1.0, 1.0, 1.0]
+        get_flops(hfu_model_config, batch_size, input_seq_len, output_length, 1, hfu_info)
+    return hfu_info
+
+def dump_information(dump_info):
+    if VLLM_DUMP_MLU_INFO_EN:
+        millisecond2second_unit = 1000
+        dump_info.hfu_info["context_hfu"] = dump_info.hfu_info["context_hfu"] / (dump_info.context_latency_device / millisecond2second_unit)
+        dump_info.hfu_info["decoder_hfu"] = dump_info.hfu_info["decoder_hfu"] / (dump_info.generate_latency_device / millisecond2second_unit)
+
+        print(f"Context HFU-visible:   {dump_info.hfu_info['context_hfu']:.3%}")
+        print(f"Decoder HFU-visible:   {dump_info.hfu_info['decoder_hfu']:.3%}")
+        print(f"Decoder IO Efficiency: {dump_info.io_efficiency:.3%}")
+
+def get_device_attribute():
+    result = subprocess.run(['cnmon', 'info'], 
+                                stdout=subprocess.PIPE, 
+                                stderr=subprocess.PIPE, 
+                                text=True, 
+                                check=True)
+
+    pattern = r"MLU\s+\d+-\d+"
+    ranges = re.findall(pattern, result.stdout)
+    unique_ranges = list(dict.fromkeys(ranges))
+    cluster_num = len(unique_ranges)
+    first_range = unique_ranges[0]
+    match = re.search(r"(\d+)-(\d+)", first_range)
+    start = int(match.group(1))  
+    end = int(match.group(2))
+    core_num = end - start + 1 
+    match = re.search(r"Default\s+:\s+(\d+)\s+MHz", result.stdout)
+    frequency = int(match.group(1))  
+    match = re.search(r"MEM BandWidth\s+:\s+(\d+)\s+MB/s", result.stdout)
+    if match is None:
+        match = re.search(r"DDR BandWidth\s+:\s+(\d+)\s+MB/s", result.stdout)
+    bandwidth = int(match.group(1))
+    return core_num, cluster_num, frequency, bandwidth
+
+
+def get_hfu_peak_flops(hfu_model_config, each_peak_compute):
+    device_id = 0
+    device_name = os.popen(f"cnmon info -c {device_id} | grep 'Product Name' | cut -d: -f2").read().strip()
+
+    input_data_type = hfu_model_config["data_type"]
+    filter_data_type = hfu_model_config["filter_data_type"]
+        
+    base_value_map: Dict[Tuple[str, str], float] = {
+        ("U5", "CNNL_DTYPE_FLOAT"): int("10000000000", 2),
+        ("U5", "CNNL_DTYPE_HALF"): int("1000000000000", 2),
+        ("U5", "CNNL_DTYPE_INT16"): int("100000000000", 2),
+        ("U5", "CNNL_DTYPE_BFLOAT16"): int("1000000000000", 2),
+        ("U5", "CNNL_DTYPE_INT8"): int("10000000000000", 2),
+        ("0B", "CNNL_DTYPE_INT8"): int("100000000000000", 2),
+        ("0B", "CNNL_DTYPE_INT16"): int("10000000000000", 2),
+        ("0B", "CNNL_DTYPE_HALF"): int("10000000000000", 2),
+        ("0B", "CNNL_DTYPE_BFLOAT16"): int("10000000000000", 2),
+        ("0B", "CNNL_DTYPE_FLOAT"): int("10000000000", 2),
+        ("0B", "CNNL_DTYPE_FLOAT8_E4M3FN"): int("100000000000000", 2),
+        ("0B", "CNNL_DTYPE_FLOAT8_E5M2"): int("100000000000000", 2),
+    }
+    
+    cur_plat = ""
+    if "0B" in device_name or "80-X" in device_name:
+        cur_plat = "0B"
+    elif "U5" in device_name:
+        cur_plat = "U5"
+    else:
+        raise TypeError(f"current platform {device_name} not support HFU info get!")
+        
+
+    if cur_plat:
+        keys = [
+            (cur_plat, "CNNL_DTYPE_INT16"),  # kv8
+            (cur_plat, str_to_data_type(input_data_type)),
+            (cur_plat, str_to_data_type(filter_data_type)),
+        ]
+        freq_unit = 1000.0
+        for i, key in enumerate(keys):
+            if key not in base_value_map:
+                raise TypeError(f"current data type {key} not support! input_data_type: {input_data_type}, filter_data_type: {filter_data_type}")
+            else:
+                each_peak_compute[i] = base_value_map[key]
+                core_num, cluster_num, frequency, _ = get_device_attribute()
+                each_peak_compute[i] *= frequency * freq_unit * freq_unit * cluster_num * core_num 
+    else:
+        raise TypeError(f"current platform {device_name} not support HFU info get!")
+
+    return each_peak_compute
diff --git a/vllm_mlu/vllm_mlu/distributed/__init__.py b/vllm_mlu/vllm_mlu/distributed/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/distributed/device_communicators/__init__.py b/vllm_mlu/vllm_mlu/distributed/device_communicators/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/distributed/device_communicators/mlu_communicator.py b/vllm_mlu/vllm_mlu/distributed/device_communicators/mlu_communicator.py
new file mode 100644
index 000000000..4e6d0dab2
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/distributed/device_communicators/mlu_communicator.py
@@ -0,0 +1,19 @@
+from typing import Optional
+
+import torch
+from torch.distributed import ProcessGroup
+
+from vllm.distributed.device_communicators.base_device_communicator import \
+    DeviceCommunicatorBase
+
+
+class MLUCommunicator(DeviceCommunicatorBase):
+
+    def __init__(self,
+                 cpu_group: ProcessGroup,
+                 device: Optional[torch.device] = None,
+                 device_group: Optional[ProcessGroup] = None,
+                 unique_name: str = ""):
+        super().__init__(cpu_group, device, device_group, unique_name)
+        # init device according to rank
+        self.device = torch.mlu.current_device()
diff --git a/vllm_mlu/vllm_mlu/distributed/kv_transfer/kv_connector/simple_connector.py b/vllm_mlu/vllm_mlu/distributed/kv_transfer/kv_connector/simple_connector.py
new file mode 100644
index 000000000..a42681a12
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/distributed/kv_transfer/kv_connector/simple_connector.py
@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Simple KV Cache Connector for Distributed Machine Learning Inference
+
+The SimpleConnector transfers KV caches between prefill vLLM worker (KV cache
+producer) and decode vLLM worker (KV cache consumer) using PyNcclPipe or
+MooncakePipe.
+
+But the logic can be extended to support other pipe and lookup buffer.
+"""
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
+import torch
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
+from vllm.distributed.kv_transfer.kv_lookup_buffer.simple_buffer import (
+    SimpleBuffer)
+from vllm.distributed.kv_transfer.kv_connector.simple_connector import SimpleConnector
+from vllm.logger import init_logger
+from vllm.sequence import IntermediateTensors
+
+if TYPE_CHECKING:
+    from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
+
+logger = init_logger(__name__)
+
+
+class SimpleConnector_MluHijack(SimpleConnector):
+
+    def send_kv_caches_and_hidden_states(
+        self,
+        model_executable: torch.nn.Module,
+        model_input: "ModelInputForGPUWithSamplingMetadata",
+        kv_caches: List[List[torch.Tensor]],
+        hidden_or_intermediate_states: Union[torch.Tensor,
+                                             IntermediateTensors],
+    ) -> None:
+
+        input_tokens_tensor = model_input.input_tokens
+        seq_lens = model_input.attn_metadata.seq_lens
+        slot_mapping_flat = model_input.attn_metadata.slot_mapping.flatten()
+        num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
+        start_layer = model_executable.model.start_layer
+        end_layer = model_executable.model.end_layer
+
+        model_config = model_executable.model.config
+        num_heads = int(model_config.num_key_value_heads / self.tp_size)
+        hidden_size = model_config.hidden_size
+        num_attention_heads = model_config.num_attention_heads
+
+        # Deepseek's MLA (Multi-head Latent Attention) uses two different
+        # kv_cache shapes based on whether VLLM_MLA_DISABLE is set to 0.
+        # When VLLM_MLA_DISABLE=0 (default), forward absorb is applied,
+        # resulting in a kv_cache shape of [num_blks, blk_size, 1,
+        # kv_lora_rank + qk_rope_head_dim].
+        # When VLLM_MLA_DISABLE=1, standard FA is used instead, leading
+        # to a kv_cache shape of [2, num_blks, blk_size,
+        # num_key_value_heads / tp, qk_nope_head_dim + qk_rope_head_dim].
+        # For more details, see vllm/attention/backends/mla/common.py.
+        if self.is_deepseek_mla and self.use_mla_opt:
+            head_size = model_config.kv_lora_rank + \
+                model_config.qk_rope_head_dim
+            num_heads = 1
+        elif self.is_deepseek_mla and not self.use_mla_opt:
+            head_size = model_config.qk_nope_head_dim + \
+                model_config.qk_rope_head_dim
+        else:
+            head_size = getattr(model_config, "head_dim",
+                                int(hidden_size // num_attention_heads))
+
+        # query_lens contains new KV caches that are added to vLLM.
+        # so we will send them to decode instance
+        # FIXME(Kuntai): This assume that all requests are prefill.
+        for idx, slen in enumerate(seq_lens):
+            start_pos = sum(seq_lens[:idx])
+            end_pos = start_pos + slen
+
+            if start_pos >= num_prefill_tokens:
+                # vllm/worker/model_runner.py::_prepare_model_input_tensors:
+                # - input_tokens[:num_prefill_tokens] contains prefill tokens.
+                # - input_tokens[num_prefill_tokens:] contains decode tokens.
+                logger.warning("You have some decode requests while using "
+                               "SimpleConnector. Their KVCache won't be sent.")
+                break
+
+            current_tokens = input_tokens_tensor[start_pos:end_pos]
+
+            keys, values = [], []
+
+            for layer_id in range(start_layer, end_layer):
+                kv_cache = kv_caches[layer_id - start_layer]
+                kv_scale = kv_cache[1]
+                kv_cache = kv_cache[0]
+
+                if self.is_deepseek_mla and self.use_mla_opt:
+                    key_cache = kv_cache.reshape(-1, num_heads, head_size)
+                    value_cache = kv_cache.reshape(-1, num_heads, head_size)
+                else:
+                    key_cache = kv_cache[0].reshape(-1, num_heads, head_size)
+                    value_cache = kv_cache[1].reshape(-1, num_heads, head_size)
+
+                current_slot_mapping = slot_mapping_flat[start_pos:end_pos]
+
+                keys.append(key_cache[current_slot_mapping].unsqueeze(0))
+                values.append(value_cache[current_slot_mapping].unsqueeze(0))
+
+            keys = torch.cat(keys, dim=0)
+            values = torch.cat(values, dim=0)
+
+            self.insert(current_tokens,
+                        torch.ones_like(current_tokens,
+                                        dtype=bool), keys, values,
+                        hidden_or_intermediate_states[start_pos:end_pos])
+
+        logger.debug("[rank%d]: KV send DONE.", torch.distributed.get_rank())
+
+    def recv_kv_caches_and_hidden_states(
+        self, model_executable: torch.nn.Module,
+        model_input: "ModelInputForGPUWithSamplingMetadata",
+        kv_caches: List[List[torch.Tensor]]
+    ) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,
+               "ModelInputForGPUWithSamplingMetadata"]:
+
+        # When bypass_model_exec is set to False, it means that at least for one
+        # request its corresponding KV cache or hidden state is missing.
+        # In this case we need to do prefilling to recompute missing KV cache
+        # and hidden states.
+        bypass_model_exec = True
+
+        model_config = model_executable.model.config
+
+        input_tokens_tensor = model_input.input_tokens
+        seq_lens = model_input.attn_metadata.seq_lens
+        num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
+        slot_mapping = model_input.attn_metadata.slot_mapping.flatten()
+
+        hidden_or_intermediate_states_for_one_req = []
+
+        input_tokens_list = []
+        num_computed_tokens_list = []
+        start_pos_list = []
+
+        # enumerate different requests
+        # FIXME(Kuntai): This impl assumes that all requests are prefill.
+        for idx, slen in enumerate(seq_lens):
+            start_pos = sum(seq_lens[:idx])
+            end_pos = start_pos + slen
+
+            if start_pos >= num_prefill_tokens:
+                # This can happen during inflight batching. See:
+                # vllm/worker/model_runner.py::_prepare_model_input_tensors:
+                # - input_tokens[:num_prefill_tokens] contains prefill tokens.
+                # - input_tokens[num_prefill_tokens:] contains decode tokens.
+                logger.warning("You should set --enable_chunked_prefill=False "
+                               "and --max_num_batched_tokens "
+                               "should be equal to --max_seq_len_to_capture")
+                bypass_model_exec = False
+                assert start_pos == num_prefill_tokens
+                break
+
+            current_tokens = input_tokens_tensor[start_pos:end_pos]
+            num_tokens = slen
+
+            # collecting data for rebuilding the input
+            input_tokens_list.append(current_tokens)
+            start_pos_list.append(start_pos)
+
+            ret = self.select(current_tokens,
+                              torch.ones_like(current_tokens, dtype=bool))
+            if ret[0] is None:
+                # didn't find any match.
+                bypass_model_exec = False
+                num_computed_tokens_list.append(0)
+                continue
+
+            roi: torch.Tensor = ret[1]
+            keys: torch.Tensor = ret[2]
+            values: torch.Tensor = ret[3]
+            hidden: torch.Tensor = ret[4]
+
+            num_computed_tokens = roi.shape[0]
+            num_computed_tokens_list.append(num_computed_tokens)
+
+            # check if both KV cache and the hidden states are received
+            # If not, need to redo the forwarding to compute missing states
+            if not all([(num_computed_tokens == num_tokens), hidden is not None
+                        ]):
+                bypass_model_exec = False
+
+            # update the end position based on how many tokens are cached.
+            end_pos = start_pos + num_computed_tokens
+
+            # put received KV caches into paged memory
+            for i in range(model_executable.model.start_layer,
+                           model_executable.model.end_layer):
+
+                kv_cache = kv_caches[i - model_executable.model.start_layer]
+                layer = model_executable.model.layers[i]
+
+                if self.is_deepseek_mla and self.use_mla_opt:
+                    layer.self_attn.attn = layer.self_attn.mla_attn
+                    k_c_normed_k_pe = keys[
+                        i - model_executable.model.start_layer].to(
+                            kv_cache.device).squeeze(1)
+                    k_c_normed = k_c_normed_k_pe[:, :model_config.kv_lora_rank]
+                    k_pe = k_c_normed_k_pe[:, model_config.kv_lora_rank:]
+                    ops.concat_and_cache_mla(
+                        k_c_normed,
+                        k_pe,
+                        kv_cache,
+                        slot_mapping[start_pos:end_pos],
+                        layer.self_attn.attn.kv_cache_dtype,
+                        layer.self_attn.attn._k_scale,
+                    )
+                else:
+                    key_cache, value_cache = kv_cache[0][0], kv_cache[0][1]
+                    ops.reshape_and_cache_flash(
+                        keys[i - model_executable.model.start_layer].to(
+                            key_cache.device),
+                        values[i - model_executable.model.start_layer].to(
+                            value_cache.device),
+                        key_cache,
+                        value_cache,
+                        slot_mapping[start_pos:end_pos],
+                        layer.self_attn.attn.kv_cache_dtype,
+                        layer.self_attn.attn._k_scale,
+                        layer.self_attn.attn._v_scale,
+                    )
+
+            hidden_or_intermediate_states_for_one_req.append(hidden)
+
+        if not bypass_model_exec:
+            # Some of the KV cache is not retrieved
+            # Here we will fall back to normal model forwarding
+            # But optionally you can adjust model_input so that you only do
+            # prefilling on those tokens that are missing KV caches.
+            logger.warning(
+                "[rank%d]: Failed to receive all KVs and hidden "
+                "states, redo model forwarding.", torch.distributed.get_rank())
+            hidden_or_intermediate_states = None
+
+        else:
+            logger.debug(
+                "[rank%d]: Successfully received all KVs and hidden "
+                "states, skip model forwarding.", torch.distributed.get_rank())
+            hidden_or_intermediate_states = torch.cat(
+                hidden_or_intermediate_states_for_one_req, dim=0)
+
+        return hidden_or_intermediate_states, bypass_model_exec, model_input
diff --git a/vllm_mlu/vllm_mlu/distributed/parallel_state.py b/vllm_mlu/vllm_mlu/distributed/parallel_state.py
new file mode 100644
index 000000000..230f1bccb
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/distributed/parallel_state.py
@@ -0,0 +1,42 @@
+import contextlib
+import gc
+import torch
+from contextlib import contextmanager
+from typing import Optional
+
+from vllm.distributed import parallel_state
+from vllm.distributed.parallel_state import (GroupCoordinator,
+                                             GraphCaptureContext,
+                                             destroy_model_parallel,
+                                             destroy_distributed_environment)
+from vllm.logger import init_logger
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+@contextmanager
+def vllm__distributed__parallel_state__GroupCoordinator__graph_capture(
+    self,
+    graph_capture_context: Optional[GraphCaptureContext] = None
+):
+    if graph_capture_context is None:
+        stream = torch.mlu.Stream()
+        graph_capture_context = GraphCaptureContext(stream)
+    else:
+        stream = graph_capture_context.stream
+
+    # ensure all initialization operations complete before attempting to
+    # capture the graph on another stream
+    curr_stream = torch.mlu.current_stream()
+    if curr_stream != stream:
+        stream.wait_stream(curr_stream)
+
+    with torch.mlu.stream(stream):
+        yield graph_capture_context
+
+
+MluHijackObject.apply_hijack(GroupCoordinator,
+                             GroupCoordinator.graph_capture,
+                             vllm__distributed__parallel_state__GroupCoordinator__graph_capture)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/dump_info.py b/vllm_mlu/vllm_mlu/dump_info.py
new file mode 100644
index 000000000..242e8189f
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/dump_info.py
@@ -0,0 +1,380 @@
+import os
+from vllm.logger import init_logger
+from vllm_mlu.mlu_hijack_utils import get_is_gated, TypedDict
+import json
+from vllm.transformers_utils.config import get_config
+from vllm.entrypoints.llm import LLM
+from vllm_mlu._mlu_utils import VLLM_DUMP_MLU_INFO_EN
+logger = init_logger(__name__)
+
+def get_deepseek_v2_flops(bcfg, batch, seq_len, hidden_size):
+    qk_nope_head_dim = bcfg["qk_nope_head_dim"]
+    qk_rope_head_dim = bcfg["qk_rope_head_dim"]
+    v_head_dim       = bcfg["v_head_dim"]
+    q_lora_rank      = bcfg["q_lora_rank"]
+    kv_lora_rank     = bcfg["kv_lora_rank"]
+    context_atn_pre = 2 * batch * seq_len * \
+                 (hidden_size * q_lora_rank + \
+                 hidden_size * (kv_lora_rank + qk_rope_head_dim) + \
+                 q_lora_rank * bcfg["head_num"] * (qk_nope_head_dim + qk_rope_head_dim) + \
+                 kv_lora_rank * bcfg["head_num"] * (qk_nope_head_dim + v_head_dim))
+    context_atn_qk = 2 * batch * seq_len * seq_len * bcfg["head_num"] * (qk_nope_head_dim + qk_rope_head_dim)
+    context_atn_qkv = 2 * batch * seq_len * seq_len * bcfg["head_num"] * v_head_dim
+    context_atn_post = 2 * batch * seq_len * bcfg["head_num"] * v_head_dim * hidden_size
+    return context_atn_pre, context_atn_qk, context_atn_qkv, context_atn_post
+
+
+FlopsInfo = {
+    "context_flops":0.0,
+    "decoder_flops":0.0
+    }
+
+
+class LLMDumpInfo:
+    def __init__(self,
+                 tensor_parallel_size=None,
+                 dtype=None, kv_cache_dtype=None,
+                 quantization=None,
+                 model=None, batch_size=None,
+                 input_len=None,
+                 output_len=None,
+                 trust_remote_code=None)->None:
+        self.hfu_info = None
+        self.flops_info = None
+        self.hfu_model_config = TypedDict
+        self.io_efficiency = 0
+        self.context_latency_device = 0
+        self.generate_latency_device = 0
+
+        self.tensor_parallel_size = tensor_parallel_size
+        self.dtype = dtype
+        self.kv_cache_dtype = kv_cache_dtype
+        self.quantization = quantization
+        self.batch_size = batch_size
+        self.input_len = input_len
+        self.output_len = output_len
+        self.model = model
+        self.model_config = None
+        self.trust_remote_code = trust_remote_code
+
+
+    def init_param(self,
+                   tensor_parallel_size=None,
+                   dtype=None,
+                   kv_cache_dtype=None,
+                   quantization=None,
+                   model=None,
+                   batch_size=None,
+                   input_len=None,
+                   output_len=None,
+                   trust_remote_code=None,
+                   context_latency_device=None,
+                   generate_latency_device=None):
+        if tensor_parallel_size != None:
+            self.tensor_parallel_size = tensor_parallel_size
+        if dtype != None:
+            self.dtype = dtype
+        if kv_cache_dtype != None:
+            self.kv_cache_dtype = kv_cache_dtype
+        if quantization != None:
+            self.quantization = quantization
+        if model != None:
+            self.model = model
+        if batch_size != None:
+            self.batch_size = batch_size
+        if input_len != None:
+            self.input_len = input_len
+        if output_len != None:
+            self.output_len = output_len
+        if trust_remote_code != None:
+            self.trust_remote_code = trust_remote_code
+        if context_latency_device != None:
+            self.context_latency_device = context_latency_device
+        if generate_latency_device != None:
+            self.generate_latency_device = generate_latency_device
+
+        # paser the model config
+        if self.model_config == None and self.model != None and self.trust_remote_code != None:
+            self.model_config = get_config(self.model, self.trust_remote_code)
+
+    def initialize_hfu_model_config(self):
+        # prepare input
+        if hasattr(self.model_config, "hidden_size"):
+            self.hfu_model_config["hidden_size"] = self.model_config.hidden_size
+        elif hasattr(self.model_config, "audio_config") and hasattr(self.model_config.audio_config, "d_model"):
+            self.hfu_model_config["hidden_size"] = self.model_config.audio_config.d_model
+        else:
+            logger.error("The model's config.json does not contain hidden_size or audio_config.d_model.")
+        self.hfu_model_config["vocab_size"]   = self.model_config.vocab_size
+        self.hfu_model_config["cla_coeffient"]   = 1.0
+
+        possible_keys_ffn_size = [
+            # chatglm3-6b-32k
+            "ffn_hidden_size",
+            # llama3-8b-hf
+            "intermediate_size",
+        ]
+        possible_kv_heads = [
+            # chatglm3-6b-32k
+            "multi_query_group_num",
+            # llama3-8b-hf
+            "num_key_value_heads",
+            # falcon-180B-chat
+            "num_kv_heads",
+        ]
+        possible_num_attention_heads = [
+            "num_attention_heads",
+            "n_heads",
+        ]
+        moe_size=None
+        ffn_size=None
+        if getattr(self.model_config, "moe_intermediate_size", None):
+            moe_size = getattr(self.model_config, "moe_intermediate_size", None)
+        for key in possible_keys_ffn_size:
+            ffn_size = getattr(self.model_config, key, None)
+            if ffn_size is not None:
+                break
+        if self.model_config.model_type in ['bloom'] and ffn_size is None:
+            ffn_size = self.model_config.hidden_size * 4
+        if self.model_config.model_type in ['qwen']:
+           ffn_size = self.model_config.intermediate_size // 2
+        if ffn_size is None and moe_size is None:
+            logger.warning("The model's config.json does not contain any of the following"
+                        "keys to determine the ffn_size or moe_size: "
+                        f"{possible_keys_ffn_size}. ")
+
+        for key in possible_num_attention_heads:
+            num_attention_heads = getattr(self.model_config, key, None)
+            if num_attention_heads is not None:
+                break
+                
+        if num_attention_heads is None:
+            if hasattr(self.model_config, "audio_config") and hasattr(self.model_config.audio_config, "encoder_attention_heads"):
+                num_attention_heads = self.model_config.audio_config.encoder_attention_heads
+            else:
+                logger.error("The model's config.json does not contain any of the following"
+                            "keys to determine the num_attention_heads: "
+                            f"{possible_num_attention_heads}. ")
+
+        for key in possible_kv_heads:
+            kv_heads = getattr(self.model_config, key, None)
+            if kv_heads is not None:
+                break
+
+        if kv_heads is None:
+            logger.warning("The model's config.json does not contain any of the following"
+                        "keys to determine the kv_heads: "
+                        f"{possible_kv_heads}, use num_attention_heads to replace")
+            kv_heads = num_attention_heads
+        self.hfu_model_config["ffn_inner_size"] =  0 if ffn_size is None else ffn_size
+        self.hfu_model_config["moe_inner_size"] =  0 if moe_size is None else moe_size
+        self.hfu_model_config["moe_layer_num"]  =  0 if moe_size is None else self.model_config.num_hidden_layers
+        if getattr(self.model_config, "num_hidden_layers", None):
+            num_hidden_layers = getattr(self.model_config, "num_hidden_layers", None)
+        elif hasattr(self.model_config, "audio_config") and hasattr(self.model_config.audio_config, "encoder_layers"):
+            num_hidden_layers = self.model_config.audio_config.encoder_layers
+        else:
+            logger.error("The model's config.json does not contain num_hidden_layers or audio_config.encoder_layers.")
+        self.hfu_model_config["layer_num"] = num_hidden_layers
+        self.hfu_model_config["head_num"]       =  num_attention_heads
+        self.hfu_model_config["head_size"]      =  self.hfu_model_config["hidden_size"] / self.hfu_model_config["head_num"]
+        self.hfu_model_config["head_num_kv"]    =  kv_heads
+        self.hfu_model_config["tp_num"]         =  self.tensor_parallel_size
+        if hasattr(self.model_config, "shared_expert_intermediate_size") and self.model_config.shared_expert_intermediate_size is not None:
+            self.hfu_model_config["shared_expert_intermediate_size"] = self.model_config.shared_expert_intermediate_size
+        else:
+            self.hfu_model_config["shared_expert_intermediate_size"] = 0
+        self.hfu_model_config["use_gated_ffn"]   =  get_is_gated()
+        if hasattr(self.model_config, "n_shared_experts") and self.model_config.n_shared_experts is not None:
+            self.hfu_model_config["shared_expert_intermediate_size"] = self.model_config.n_shared_experts * moe_size
+        else:
+            self.hfu_model_config["shared_experts"]  = 0
+        if hasattr(self.model_config, "num_experts") and self.model_config.num_experts is not None:
+            self.hfu_model_config["experts_num"]     =  self.model_config.num_experts
+            if self.model_config.model_type == 'hunyuan':
+                self.hfu_model_config["topk_num"]        =  self.model_config.moe_topk
+            else:
+                self.hfu_model_config["topk_num"]        =  self.model_config.num_experts_per_tok
+        elif hasattr(self.model_config, "num_local_experts"):
+            self.hfu_model_config["experts_num"]     =  self.model_config.num_local_experts
+            if self.model_config.model_type == 'hunyuan':
+                self.hfu_model_config["topk_num"]        =  self.model_config.moe_topk
+            else:
+                self.hfu_model_config["topk_num"]        =  self.model_config.num_experts_per_tok
+        elif hasattr(self.model_config, "n_routed_experts"):
+            self.hfu_model_config["experts_num"]     =  self.model_config.n_routed_experts
+            if self.model_config.model_type == 'hunyuan':
+                self.hfu_model_config["topk_num"]        =  self.model_config.moe_topk
+            else:
+                self.hfu_model_config["topk_num"]        =  self.model_config.num_experts_per_tok
+        else:
+            self.hfu_model_config["experts_num"]     =  0
+        if hasattr(self.model_config, "model_type") and self.model_config.model_type is not None:
+            self.hfu_model_config["model_type"] = self.model_config.model_type
+            # when adding a moe model, need fix moe/ffn info, like
+            # moe_inner_size, ffn_inner_size, moe_layer_num, shared_expert_intermediate_size.
+            # add for mixtral
+            if self.model_config.model_type == "mixtral":
+                self.hfu_model_config["moe_inner_size"] = ffn_size
+                self.hfu_model_config["ffn_inner_size"] =  0
+                self.hfu_model_config["moe_layer_num"]  = self.model_config.num_hidden_layers
+            # add for deepseek-v2
+            if self.model_config.model_type in ["deepseek_v2", "deepseek_v3"]:
+                if hasattr(self.model_config, "first_k_dense_replace") and self.model_config.first_k_dense_replace is not None:
+                    self.hfu_model_config["moe_layer_num"] = self.model_config.num_hidden_layers - self.model_config.first_k_dense_replace
+                if hasattr(self.model_config, "qk_nope_head_dim") and self.model_config.qk_nope_head_dim is not None:
+                    self.hfu_model_config["qk_nope_head_dim"] = self.model_config.qk_nope_head_dim
+                if hasattr(self.model_config, "qk_rope_head_dim") and self.model_config.qk_rope_head_dim is not None:
+                    self.hfu_model_config["qk_rope_head_dim"] = self.model_config.qk_rope_head_dim
+                if hasattr(self.model_config, "v_head_dim") and self.model_config.v_head_dim is not None:
+                    self.hfu_model_config["v_head_dim"] = self.model_config.v_head_dim
+                if hasattr(self.model_config, "q_lora_rank") and self.model_config.q_lora_rank is not None:
+                    self.hfu_model_config["q_lora_rank"] = self.model_config.q_lora_rank
+                else:
+                    self.hfu_model_config["q_lora_rank"] = 0
+                if hasattr(self.model_config, "kv_lora_rank") and self.model_config.kv_lora_rank is not None:
+                    self.hfu_model_config["kv_lora_rank"] = self.model_config.kv_lora_rank
+            # add for Hunyuan
+            if self.model_config.model_type == "hunyuan":
+                self.hfu_model_config["cla_coeffient"] = 0.5 # huanyuan model use CLA2
+                if hasattr(self.model_config, "num_shared_expert") and self.model_config.num_shared_expert is not None:
+                    self.hfu_model_config["shared_expert_intermediate_size"] = self.model_config.num_shared_expert * self.model_config.intermediate_size
+                if not self.hfu_model_config["moe_inner_size"] and self.model_config.intermediate_size is not None:
+                    self.hfu_model_config["moe_inner_size"] = self.model_config.intermediate_size
+                if not self.hfu_model_config["moe_layer_num"] and hasattr(self.model_config, "num_experts"):
+                    self.hfu_model_config["moe_layer_num"] = self.model_config.num_hidden_layers
+
+        self.hfu_model_config["use_causal_mask"]     =  True  # the flash attention is only use causal_mask in vllm
+
+        if self.dtype == "auto":
+            self.hfu_model_config["data_type"] = "float16"
+        else:
+            self.hfu_model_config["data_type"] = self.dtype
+
+        if self.quantization != None:
+            config = None
+            if os.path.exists(self.model + "/quantize_config.json"):
+                with open(self.model + "/quantize_config.json", 'r') as file:
+                    config = json.load(file)
+            elif os.path.exists(self.model + "/config.json"):
+                with open(self.model + "/config.json", 'r') as file:
+                    config = json.load(file).get("quantization_config", None)
+            assert config is not None, "The model's quantize_config.json does not exist."
+
+            if config["quant_mode"] == "SmoothQuant":
+                self.hfu_model_config["smooth_quant_type"] = b"SmoothQuant"
+            else:
+                self.hfu_model_config["smooth_quant_type"] = "invalid"
+            if self.quantization == 'fp8':
+                self.hfu_model_config["filter_data_type"] = f"fp8_{config['fmt']}"
+            else:
+                self.hfu_model_config["filter_data_type"] = ("int" +  str(config['bits']))
+        else:
+            self.hfu_model_config["smooth_quant_type"] = "invalid"
+            self.hfu_model_config["filter_data_type"] = self.hfu_model_config["data_type"]
+
+        if self.kv_cache_dtype == "auto":
+            self.hfu_model_config["kv_cache_dtype"]      =  self.hfu_model_config["data_type"]
+        else:
+            self.hfu_model_config["kv_cache_dtype"]      =  self.kv_cache_dtype
+
+
+    def get_flops(self):
+        self.batch_size = self.batch_size
+        seq_len = self.input_len
+        hidden_size = self.hfu_model_config["hidden_size"]
+        voc_size = self.hfu_model_config["vocab_size"]
+        ffn_size = self.hfu_model_config["ffn_inner_size"]
+        moe_size = self.hfu_model_config["moe_inner_size"]
+        shared_expert_intermediate_size = self.hfu_model_config["shared_expert_intermediate_size"]
+        layer_num = self.hfu_model_config["layer_num"]
+        out_seq = self.output_len
+        seq_len_decode = seq_len + out_seq / 2
+        r = self.hfu_model_config["head_num"] / self.hfu_model_config["head_num_kv"]
+        bsh2 = self.batch_size * seq_len * hidden_size * hidden_size
+        cla_coeffient = self.hfu_model_config["cla_coeffient"]
+        if self.hfu_model_config["model_type"] in ["deepseek_v2", "deepseek_v3"]:
+            context_atn_pre, context_atn_qk, context_atn_qkv, context_atn_post = (
+                get_deepseek_v2_flops(self.hfu_model_config, self.batch_size, seq_len, hidden_size)
+            )
+        else:
+            context_atn_pre = 2 * bsh2 + 4 * bsh2 / r * cla_coeffient
+            context_atn_qk = 2 * self.batch_size * seq_len * seq_len * hidden_size
+            context_atn_qkv = 2 * self.batch_size * seq_len * seq_len * hidden_size
+            context_atn_post = 2 * self.batch_size * seq_len * hidden_size * hidden_size
+        context_lm_head = 2 * self.batch_size * seq_len * hidden_size * voc_size
+        context_ffn = 0
+        bh2 = self.batch_size * hidden_size * hidden_size
+        decode_atn_pre = 2 * bh2 + 4 * bh2 / r * cla_coeffient
+        decode_atn_qk = 2 * self.batch_size * seq_len_decode * hidden_size
+        decode_atn_qkv = 2 * self.batch_size * seq_len_decode * hidden_size
+        decode_atn_post = 2 * self.batch_size * hidden_size * hidden_size
+        decode_lm_head = 2 * self.batch_size * hidden_size * voc_size
+        decode_ffn = 0
+        coeffient = 6 if self.hfu_model_config["use_gated_ffn"] else 4
+        if self.hfu_model_config["experts_num"] == 0:
+            context_ffn = coeffient * self.batch_size * seq_len * hidden_size * ffn_size
+            decode_ffn = coeffient * self.batch_size * hidden_size * ffn_size
+        else:
+            context_ffn = self.batch_size * seq_len * hidden_size * (coeffient * (moe_size * self.hfu_model_config["topk_num"] + shared_expert_intermediate_size) + 2 * self.hfu_model_config["experts_num"])
+            decode_ffn = self.batch_size * hidden_size * (coeffient * (moe_size * self.hfu_model_config["topk_num"] + shared_expert_intermediate_size) + 2 * self.hfu_model_config["experts_num"])
+
+        if self.hfu_model_config["use_causal_mask"]:
+            c = 0.5
+            context_atn_qk *= c
+            context_atn_qkv *= c
+
+        self.flops_info["context_flops"] = context_lm_head
+        self.flops_info["decoder_flops"] = decode_lm_head
+        if self.hfu_model_config["kv_cache_dtype"] != b"int8":
+            self.flops_info["context_flops"] += (layer_num * (context_atn_qk + context_atn_qkv))
+            self.flops_info["decoder_flops"] += (layer_num * (decode_atn_qk + decode_atn_qkv))
+        else:
+            self.flops_info["context_flops"] += (layer_num * (context_atn_qk + context_atn_qkv))
+            self.flops_info["decoder_flops"] += (layer_num * (decode_atn_qk + decode_atn_qkv))
+
+        if self.hfu_model_config["smooth_quant_type"] == b"invalid":
+            self.flops_info["context_flops"] += (layer_num * (context_atn_pre + context_atn_post + context_ffn))
+            self.flops_info["decoder_flops"] += (layer_num * (decode_atn_pre + decode_atn_post + decode_ffn))
+        else:
+            self.flops_info["context_flops"] += (layer_num * (context_atn_pre + context_atn_post + context_ffn))
+            self.flops_info["decoder_flops"] += (layer_num * (decode_atn_pre + decode_atn_post + decode_ffn))
+
+
+    def get_decoder_io_efficiency(self):
+        try:
+            from vllm_mlu.device_info import get_decoder_io_efficiency
+            self.io_efficiency = get_decoder_io_efficiency(self.hfu_model_config, self.batch_size, self.input_len, self.output_len, self.generate_latency_device)
+        except:
+            logger.info("Unsupport io_efficiency get_decoder_io_efficiency function")
+
+    def get_device_output_info(self):
+        self.initialize_hfu_model_config()
+
+        if VLLM_DUMP_MLU_INFO_EN:
+            from vllm_mlu.device_info import get_flops_inner, HFUInfo
+            self.hfu_info = HFUInfo
+            get_flops_inner(self.hfu_model_config, self.batch_size, self.input_len, self.output_len, self.tensor_parallel_size, self.hfu_info)
+            self.get_decoder_io_efficiency()
+        else:
+            self.flops_info = FlopsInfo
+            self.get_flops()
+
+    def has_information_dump(self):
+        if VLLM_DUMP_MLU_INFO_EN:
+            try:
+                import vllm_mlu.device_info
+                return True
+            except:
+                return False
+        return False
+
+    def dump(self):
+        self.get_device_output_info()
+
+    def dump_performance_info(self):
+        try:
+            from vllm_mlu.device_info import dump_information
+            dump_information(LLM.dump_info)
+        except:
+            logger.info("Unsupport dump performance information")
diff --git a/vllm_mlu/vllm_mlu/engine/__init__.py b/vllm_mlu/vllm_mlu/engine/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/engine/arg_utils.py b/vllm_mlu/vllm_mlu/engine/arg_utils.py
new file mode 100644
index 000000000..b0cfdc438
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/engine/arg_utils.py
@@ -0,0 +1,36 @@
+from vllm.config import ModelConfig
+from vllm.engine.arg_utils import EngineArgs
+from vllm.logger import init_logger
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+vllm__engine__arg_utils__EngineArgs__create_model_config_org = EngineArgs.create_model_config
+vllm__engine__arg_utils__EngineArgs__add_cli_args_org = EngineArgs.add_cli_args
+
+
+def vllm__engine__arg_utils__EngineArgs__create_model_config(self) -> ModelConfig:
+    model_config = vllm__engine__arg_utils__EngineArgs__create_model_config_org(self)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: set context mlugraph info for model config
+    '''
+    model_config.set_context_mlugraph_info(
+        getattr(self, "enable_context_mlugraph", False),
+        getattr(self, "context_batch_size_to_capture", None),
+        getattr(self, "context_seq_len_to_capture", None))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return model_config
+
+
+MluHijackObject.apply_hijack(EngineArgs,
+                             EngineArgs.create_model_config,
+                             vllm__engine__arg_utils__EngineArgs__create_model_config)
diff --git a/vllm_mlu/vllm_mlu/engine/async_llm_engine.py b/vllm_mlu/vllm_mlu/engine/async_llm_engine.py
new file mode 100644
index 000000000..33b539878
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/engine/async_llm_engine.py
@@ -0,0 +1,35 @@
+
+from vllm.engine.async_llm_engine import AsyncLLMEngine
+from vllm.logger import init_logger
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+# for client init/reset server scheduler profile data
+async def vllm__engine__async_llm_engine__AsyncLLMEngine__init_scheduler_view(self):
+    for scheduler in self.engine.scheduler:
+        if hasattr(scheduler, "init_scheduler_view"):
+            scheduler.init_scheduler_view()
+        else:
+            logger.warning("Can not find any scheduler view, " + 
+                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")
+
+
+# for client pulling server scheduler profile data
+async def vllm__engine__async_llm_engine__AsyncLLMEngine__save_scheduler_view(self):
+    for idx, scheduler in enumerate(self.engine.scheduler):
+        if hasattr(scheduler, "save_scheduler_view"):
+            scheduler.save_scheduler_view(idx)
+        else:
+            logger.warning("Can not find any scheduler view, " + 
+                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")
+
+
+MluHijackObject.apply_hijack(AsyncLLMEngine,
+                             "init_scheduler_view",
+                             vllm__engine__async_llm_engine__AsyncLLMEngine__init_scheduler_view)
+MluHijackObject.apply_hijack(AsyncLLMEngine,
+                             "save_scheduler_view",
+                             vllm__engine__async_llm_engine__AsyncLLMEngine__save_scheduler_view)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/engine/llm_engine.py b/vllm_mlu/vllm_mlu/engine/llm_engine.py
new file mode 100644
index 000000000..ab753060d
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/engine/llm_engine.py
@@ -0,0 +1,59 @@
+from vllm.engine.llm_engine import LLMEngine
+from vllm.logger import init_logger
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__engine__llm_engine__LLMEngine__get_latency(self):
+    latency = self.model_executor.get_latency()
+    return latency
+
+
+def vllm__engine__llm_engine__LLMEngine__get_memory_usage(self):
+    return self.model_executor.get_memory_usage()
+
+
+def vllm__engine__llm_engine__LLMEngine__get_block_usage(self):
+    assert len(self.scheduler) == 1, f"Only support pipeline_parallel_size=1."
+    num_free_gpu_blocks = self.scheduler[0].block_manager.get_num_free_gpu_blocks()
+    num_free_cpu_blocks = self.scheduler[0].block_manager.get_num_free_cpu_blocks()
+    return (num_free_gpu_blocks, num_free_cpu_blocks)
+
+
+# for client init/reset server scheduler profile data
+def vllm__engine__llm_engine__LLMEngine__init_scheduler_view(self):
+    for scheduler in self.scheduler:
+        if hasattr(scheduler, "init_scheduler_view"):
+            scheduler.init_scheduler_view()
+        else:
+            logger.warning("Can not find any scheduler view, " +
+                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")
+
+
+# for client pulling server scheduler profile data
+def vllm__engine__llm_engine__LLMEngine__save_scheduler_view(self):
+    for idx, scheduler in enumerate(self.scheduler):
+        if hasattr(scheduler, "save_scheduler_view"):
+            scheduler.save_scheduler_view(idx)
+        else:
+            logger.warning("Can not find any scheduler view, " +
+                           "please 'export VLLM_SCHEDULER_PROFILE=true' first.")
+
+
+MluHijackObject.apply_hijack(LLMEngine,
+                             "init_scheduler_view",
+                             vllm__engine__llm_engine__LLMEngine__init_scheduler_view)
+MluHijackObject.apply_hijack(LLMEngine,
+                             "save_scheduler_view",
+                             vllm__engine__llm_engine__LLMEngine__save_scheduler_view)
+MluHijackObject.apply_hijack(LLMEngine,
+                             "get_latency",
+                             vllm__engine__llm_engine__LLMEngine__get_latency)
+MluHijackObject.apply_hijack(LLMEngine,
+                             "get_memory_usage",
+                             vllm__engine__llm_engine__LLMEngine__get_memory_usage)
+MluHijackObject.apply_hijack(LLMEngine,
+                             "get_block_usage",
+                             vllm__engine__llm_engine__LLMEngine__get_block_usage)
diff --git a/vllm_mlu/vllm_mlu/engine/multiprocessing/__init__.py b/vllm_mlu/vllm_mlu/engine/multiprocessing/__init__.py
new file mode 100644
index 000000000..0271c1b84
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/engine/multiprocessing/__init__.py
@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class RPCSchedulerProfileRequest(Enum):
+    INIT_SCHEDULER_VIEW = 1
+    SAVE_SCHEDULER_VIEW = 2
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/engine/multiprocessing/client.py b/vllm_mlu/vllm_mlu/engine/multiprocessing/client.py
new file mode 100644
index 000000000..5957471bd
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/engine/multiprocessing/client.py
@@ -0,0 +1,32 @@
+from vllm.engine.multiprocessing.client import MQLLMEngineClient
+from vllm.logger import init_logger
+
+from vllm_mlu.engine.multiprocessing import RPCSchedulerProfileRequest
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+class MQLLMEngineClient_V2(MQLLMEngineClient):
+
+    async def init_scheduler_view(self):
+        """Send INIT_SCHEDULER_VIEW request to RPC Server."""
+
+        await self._send_one_way_rpc_request(
+            request=RPCSchedulerProfileRequest.INIT_SCHEDULER_VIEW,
+            socket=self.input_socket)
+
+
+    async def save_scheduler_view(self):
+        """Send SAVE_SCHEDULER_VIEW request to RPC Server."""
+
+        await self._send_one_way_rpc_request(
+            request=RPCSchedulerProfileRequest.SAVE_SCHEDULER_VIEW,
+            socket=self.input_socket)
+
+
+MluHijackObject.apply_hijack(MQLLMEngineClient,
+                             "init_scheduler_view",
+                             MQLLMEngineClient_V2.init_scheduler_view)
+MluHijackObject.apply_hijack(MQLLMEngineClient,
+                             "save_scheduler_view",
+                             MQLLMEngineClient_V2.save_scheduler_view)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/engine/multiprocessing/engine.py b/vllm_mlu/vllm_mlu/engine/multiprocessing/engine.py
new file mode 100644
index 000000000..2c3fd2e45
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/engine/multiprocessing/engine.py
@@ -0,0 +1,193 @@
+import pickle
+from typing import Optional
+
+import cloudpickle
+import zmq
+
+from vllm import SamplingParams
+from vllm.engine.llm_engine import LLMEngine
+# yapf conflicts with isort for this block
+# yapf: disable
+from vllm.engine.multiprocessing import (IPC_DATA_EXT, IPC_HEALTH_EXT,
+                                         IPC_INPUT_EXT, IPC_OUTPUT_EXT,
+                                         RPCAbortRequest, RPCProcessRequest,
+                                         RPCUProfileRequest, RPCLoadAdapterRequest,
+                                         RPCResetPrefixCacheRequest,
+                                         RPCSleepRequest,
+                                         RPCWakeUpRequest,
+                                         RPCIsSleepingRequest)
+from vllm.engine.multiprocessing.engine import (MQLLMEngine,
+                                                POLLING_TIMEOUT_MS)
+from vllm.logger import init_logger
+
+from vllm_mlu.engine.multiprocessing import RPCSchedulerProfileRequest
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+vllm__engine__multiprocessing__engine__MQLLMEngine____init____org = MQLLMEngine.__init__
+
+class MQLLMEngine_V2(MQLLMEngine):
+
+    def __init__(self,
+                 ipc_path: str,
+                 use_async_sockets: bool,
+                 *args,
+                 log_requests: bool = True,
+                 **kwargs) -> None:
+        # For MQLLMEngine, we can use cached outputs, since each new request
+        # output is immediately pickled and send over the socket, which frees
+        # the python object to be reused again.
+        kwargs['use_cached_outputs'] = True
+
+        self.engine = LLMEngine(*args, **kwargs)
+        self.log_requests = log_requests
+
+        self.use_async_sockets = use_async_sockets
+        if self.use_async_sockets:
+            self.engine.process_request_outputs_callback = \
+                self._async_socket_engine_callback
+
+        self.ctx = zmq.Context()  # type: ignore[attr-defined]
+
+        # Receive input from the client.
+        self.input_socket = self.ctx.socket(zmq.constants.PULL)
+        self.input_socket.bind(f"{ipc_path}{IPC_INPUT_EXT}")
+
+        # Send output stream back to client.
+        self.output_socket = self.ctx.socket(zmq.constants.PUSH)
+        self.output_socket.bind(f"{ipc_path}{IPC_OUTPUT_EXT}")
+
+        # Send heartbeats back to client.
+        self.heartbeat_socket = self.ctx.socket(zmq.constants.PUSH)
+        self.heartbeat_socket.bind(f"{ipc_path}{IPC_HEALTH_EXT}")
+
+        # IPC path for the data socket.
+        self.data_ipc_path = f"{ipc_path}{IPC_DATA_EXT}"
+
+        # Error state.
+        self._errored_with: Optional[BaseException] = None
+
+        self.collect_scheduler_view = False
+
+    def run_engine_loop(self):
+        """Core busy loop of the LLMEngine."""
+
+        while True:
+            if not self.engine.has_unfinished_requests():
+                # Poll until there is work to do.
+                while self.input_socket.poll(timeout=POLLING_TIMEOUT_MS) == 0:
+                    # When there's no work, check on engine health and send
+                    # health status back to client
+                    self._health_check()
+                    self.engine.do_log_stats()
+                    logger.debug("Waiting for new requests in engine loop.")
+
+            # Handle any input from the client.
+            self.handle_new_input()
+
+            '''
+            =============================
+            Add by vllm_mlu
+            =============================
+            @brief: support scheduler view
+            '''
+            if self.collect_scheduler_view:
+                self.collect_scheduler_view = False
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+
+            # Engine step.
+            request_outputs = self.engine_step()
+
+            # Send request outputs (if async, done in engine_step callback).
+            if not self.use_async_sockets:
+                self._send_outputs(request_outputs)
+
+    def handle_new_input(self):
+        """Handle new input from the socket"""
+        try:
+            while self.input_socket.poll(timeout=0) != 0:
+                frames = self.input_socket.recv_multipart(copy=False)
+                request = pickle.loads(frames[0].buffer)
+
+                '''
+                =============================
+                Add by vllm_mlu
+                =============================
+                @brief: support scheduler view
+                '''
+
+                if isinstance(request, RPCProcessRequest):
+                    if len(frames) > 1:
+                        # Use cloudpickle for logits processors
+                        assert isinstance(request.params, SamplingParams)
+                        lprocs = cloudpickle.loads(frames[1].buffer)
+                        request.params.logits_processors = lprocs
+                    self._handle_process_request(request)
+                elif isinstance(request, RPCAbortRequest):
+                    self._handle_abort_request(request)
+                elif isinstance(request, RPCUProfileRequest):
+                    if request == RPCUProfileRequest.START_PROFILE:
+                        self.start_profile()
+                    else:
+                        self.stop_profile()
+                elif isinstance(request, RPCLoadAdapterRequest):
+                    self._handle_load_adapter_request(request)
+                elif isinstance(request, RPCResetPrefixCacheRequest):
+                    self.reset_prefix_cache()
+                elif isinstance(request, RPCSleepRequest):
+                    self.sleep(request.value)
+                elif isinstance(request, RPCWakeUpRequest):
+                    self.wake_up()
+                elif isinstance(request, RPCIsSleepingRequest):
+                    self._handle_is_sleeping_request(request)
+                elif isinstance(request, RPCSchedulerProfileRequest):
+                    self.collect_scheduler_view = True
+                    if request == RPCSchedulerProfileRequest.INIT_SCHEDULER_VIEW:
+                        self.init_scheduler_view()
+                    elif request == RPCSchedulerProfileRequest.SAVE_SCHEDULER_VIEW:
+                        self.save_scheduler_view()
+                else:
+                    raise ValueError("Unknown RPCRequest Type: "
+                                     f"{type(request)}")
+                '''
+                ==================
+                End of MLU Hijack
+                ==================
+                '''
+
+        except Exception as e:
+            self._set_errored(e)
+            self._send_unhealthy(e)
+            raise e
+
+    def init_scheduler_view(self):
+        """Init scheduler view."""
+        self.engine.init_scheduler_view()
+
+    def save_scheduler_view(self):
+        """Save scheduler view."""
+        self.engine.save_scheduler_view()
+
+
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             MQLLMEngine.__init__,
+                             MQLLMEngine_V2.__init__)
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             MQLLMEngine.run_engine_loop,
+                             MQLLMEngine_V2.run_engine_loop)
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             MQLLMEngine.handle_new_input,
+                             MQLLMEngine_V2.handle_new_input)
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             "init_scheduler_view",
+                             MQLLMEngine_V2.init_scheduler_view)
+MluHijackObject.apply_hijack(MQLLMEngine,
+                             "save_scheduler_view",
+                             MQLLMEngine_V2.save_scheduler_view)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/entrypoints/__init__.py b/vllm_mlu/vllm_mlu/entrypoints/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/entrypoints/llm.py b/vllm_mlu/vllm_mlu/entrypoints/llm.py
new file mode 100644
index 000000000..6a01d4a16
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/entrypoints/llm.py
@@ -0,0 +1,329 @@
+import cloudpickle
+from tqdm import tqdm
+from typing import Optional, Union, Any
+from vllm.entrypoints.llm import LLM
+from vllm.config import CompilationConfig
+from vllm.engine.arg_utils import (EngineArgs, HfOverrides, PoolerConfig,
+                                   TaskOption)
+from vllm.engine.llm_engine import LLMEngine
+from vllm.outputs import PoolingRequestOutput, RequestOutput
+from vllm.usage.usage_lib import UsageContext
+from vllm.utils import Counter, deprecate_args
+from vllm_mlu._mlu_utils import VLLM_LATENCY_DEBUG_EN, VLLM_LATENCY_DEBUG_WITH_DEVICE_EN
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.mlu_metric import LLMMetric
+from vllm_mlu.dump_info import LLMDumpInfo
+from vllm.logger import init_logger
+
+
+logger = init_logger(__name__)
+
+
+@deprecate_args(
+    start_index=2,  # Ignore self and model
+    is_deprecated=lambda: LLM.DEPRECATE_INIT_POSARGS,
+    additional_message=(
+        "All positional arguments other than `model` will be "
+        "replaced with keyword arguments in an upcoming version."),
+)
+def vllm__entrypoints__llm__LLM____init__(
+    self,
+    model: str,
+    tokenizer: Optional[str] = None,
+    tokenizer_mode: str = "auto",
+    skip_tokenizer_init: bool = False,
+    trust_remote_code: bool = False,
+    allowed_local_media_path: str = "",
+    tensor_parallel_size: int = 1,
+    dtype: str = "auto",
+    quantization: Optional[str] = None,
+    revision: Optional[str] = None,
+    tokenizer_revision: Optional[str] = None,
+    seed: Optional[int] = None,
+    gpu_memory_utilization: float = 0.9,
+    swap_space: float = 4,
+    cpu_offload_gb: float = 0,
+    enforce_eager: Optional[bool] = None,
+    max_seq_len_to_capture: int = 8192,
+    disable_custom_all_reduce: bool = False,
+    disable_async_output_proc: bool = False,
+    hf_overrides: Optional[HfOverrides] = None,
+    mm_processor_kwargs: Optional[dict[str, Any]] = None,
+    # After positional args are removed, move this right below `model`
+    task: TaskOption = "auto",
+    override_pooler_config: Optional[PoolerConfig] = None,
+    compilation_config: Optional[Union[int, dict[str, Any]]] = None,
+    **kwargs,
+) -> None:
+    '''
+    LLM constructor.
+
+    Note: if enforce_eager is unset (enforce_eager is None)
+    it defaults to False.
+    '''
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: 1) Initialize LLMDumpInfo
+            2) Initialize context mlugraph params
+    '''
+    LLM.dump_info.init_param(
+        tensor_parallel_size=tensor_parallel_size,
+        dtype=dtype,
+        kv_cache_dtype=kwargs.get('kv_cache_dtype', 'default_value'),
+        quantization=quantization,
+        model=model,
+        trust_remote_code=kwargs.get('trust_remote_code', 'default_value'))
+
+    enable_context_mlugraph = kwargs.pop("enable_context_mlugraph", False)
+    context_batch_size_to_capture = kwargs.pop("context_batch_size_to_capture", None)
+    context_seq_len_to_capture = kwargs.pop("context_seq_len_to_capture", None)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    if "disable_log_stats" not in kwargs:
+        kwargs["disable_log_stats"] = True
+
+    if "worker_cls" in kwargs:
+        worker_cls = kwargs["worker_cls"]
+        # if the worker_cls is not qualified string name,
+        # we serialize it using cloudpickle to avoid pickling issues
+        if isinstance(worker_cls, type):
+            kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
+
+    if compilation_config is not None:
+        if isinstance(compilation_config, (int, dict)):
+            compilation_config_instance = CompilationConfig.from_cli(
+                str(compilation_config))
+        else:
+            compilation_config_instance = compilation_config
+    else:
+        compilation_config_instance = None
+
+    engine_args = EngineArgs(
+        model=model,
+        task=task,
+        tokenizer=tokenizer,
+        tokenizer_mode=tokenizer_mode,
+        skip_tokenizer_init=skip_tokenizer_init,
+        trust_remote_code=trust_remote_code,
+        allowed_local_media_path=allowed_local_media_path,
+        tensor_parallel_size=tensor_parallel_size,
+        dtype=dtype,
+        quantization=quantization,
+        revision=revision,
+        tokenizer_revision=tokenizer_revision,
+        seed=seed,
+        gpu_memory_utilization=gpu_memory_utilization,
+        swap_space=swap_space,
+        cpu_offload_gb=cpu_offload_gb,
+        enforce_eager=enforce_eager,
+        max_seq_len_to_capture=max_seq_len_to_capture,
+        disable_custom_all_reduce=disable_custom_all_reduce,
+        disable_async_output_proc=disable_async_output_proc,
+        hf_overrides=hf_overrides,
+        mm_processor_kwargs=mm_processor_kwargs,
+        override_pooler_config=override_pooler_config,
+        compilation_config=compilation_config_instance,
+        **kwargs,
+    )
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: set context mlugraph params for EngineArgs
+    '''
+    setattr(engine_args, "enable_context_mlugraph", enable_context_mlugraph)
+    setattr(engine_args, "context_batch_size_to_capture", context_batch_size_to_capture)
+    setattr(engine_args, "context_seq_len_to_capture", context_seq_len_to_capture)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    # Create the Engine (autoselects V0 vs V1)
+    self.llm_engine = LLMEngine.from_engine_args(
+        engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
+    self.engine_class = type(self.llm_engine)
+
+    self.request_counter = Counter()
+    self.default_sampling_params: Union[dict[str, Any], None] = None
+
+
+def vllm__entrypoints__llm__LLM__get_metrics(
+    self,
+    metrics_idx_start,
+    only_average,
+    input_len,
+    output_len,
+    tp_nums,
+    quantization,
+    dump_info=None,
+    show_per_iter=False,
+    is_embedding_task=False,
+    mm_kwargs={},
+    total_prefill_steps=1,
+) -> None:
+    '''
+    @brief:该函数用来打印vLLM调用generate接口过程中代码统计的各项性能指标数据
+    @params:
+        metrics_idx_start: 考虑存在调用generate接口为warmup过程的情况，
+        因此设置该参数可忽略统计[0,metrics_idx_start)之间的数据,默认为0,即所有性能数据有效。
+        only_average: True 只打印N次调用generate接口的平均性能 False 打印每次调用generate接口的性能及其均值 若N次性能数据波动较大，需自行排查测试环境是否稳定。
+        其余参数:均为模型配置参数
+    '''
+    if VLLM_LATENCY_DEBUG_EN:
+        self.metric.calc_metric(self.llm_engine.model_config.model,
+                                self.llm_engine.model_config.dtype,
+                                metrics_idx_start, only_average,
+                                input_len, output_len, tp_nums,
+                                quantization, dump_info, show_per_iter,
+                                is_embedding_task, mm_kwargs, total_prefill_steps)
+    else:
+        print("Warnning:please set VLLM_LATENCY_DEBUG=true!")
+
+
+def vllm__entrypoints__llm__LLM___run_engine(
+        self, *, use_tqdm: bool
+) -> list[Union[RequestOutput, PoolingRequestOutput]]:
+    # Initialize tqdm.
+    if use_tqdm:
+        num_requests = self.llm_engine.get_num_unfinished_requests()
+        pbar = tqdm(
+            total=num_requests,
+            desc="Processed prompts",
+            dynamic_ncols=True,
+            postfix=(f"est. speed input: {0:.2f} toks/s, "
+                        f"output: {0:.2f} toks/s"),
+        )
+
+    '''
+    =============================
+    Added by vllm_mlu
+    =============================
+    '''
+    is_latency_debug = VLLM_LATENCY_DEBUG_EN
+    # Record start
+    if is_latency_debug:
+        speculative_config = self.llm_engine.model_executor.speculative_config
+        total_request_num = self.llm_engine.get_num_unfinished_requests()
+        if not self.llm_engine.model_config.is_embedding_task():
+            peak_memory, block_memory, num_total_gpu_blocks, num_total_cpu_blocks = \
+                self.llm_engine.get_memory_usage()
+            self.metric.update_memory_usage(peak_memory, block_memory, num_total_gpu_blocks, num_total_cpu_blocks)
+        e2e_start_time = self.metric.get_mlu_cost_time()
+    '''
+    ==================
+    End of addition
+    ==================
+    '''
+
+    # Run the engine.
+    outputs: list[Union[RequestOutput, PoolingRequestOutput]] = []
+    total_in_toks = 0
+    total_out_toks = 0
+    while self.llm_engine.has_unfinished_requests():
+        '''
+        =============================
+        Added by vllm_mlu
+        =============================
+        '''
+        if is_latency_debug:
+            start_time = self.metric.get_mlu_cost_time()
+        '''
+        ==================
+        End of addition
+        ==================
+        '''
+        step_outputs = self.llm_engine.step()
+        '''
+        =============================
+        Added by vllm_mlu
+        =============================
+        '''
+        if is_latency_debug:
+            end_time = self.metric.get_mlu_cost_time()
+            step_latency = end_time - start_time
+            batch_size = len(step_outputs)
+            if batch_size > 0 and speculative_config and speculative_config.num_speculative_tokens < 1:
+                assert batch_size == total_request_num, \
+                    f"LLM has received {total_request_num} requests, but only processed {batch_size} requests in the current step.\n" + \
+                    f"If you are running benchmark_latency test, please check if the input is correct.\n" + \
+                    f"Otherwise, please set env VLLM_LATENCY_DEBUG=false, then run test again.\n"
+            num_free_gpu_blocks, num_free_cpu_blocks = self.llm_engine.get_block_usage()
+            self.metric.update_step_block_usage(num_free_gpu_blocks, num_free_cpu_blocks)
+            self.metric.update_step_latency(step_latency)
+            if hasattr(self.llm_engine, "stat_loggers"):
+                spec_decode_metrics = self.llm_engine.stat_loggers['logging'].spec_decode_metrics
+                if spec_decode_metrics is not None:
+                    self.metric.update_spec_decode_metrics(spec_decode_metrics)
+            if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+                self.metric.update_step_latency_device(self.llm_engine.get_latency())
+        '''
+        ==================
+        End of addition
+        ==================
+        '''
+        for output in step_outputs:
+            if output.finished:
+                outputs.append(output)
+                if use_tqdm:
+                    if isinstance(output, RequestOutput):
+                        # Calculate tokens only for RequestOutput
+                        n = len(output.outputs)
+                        assert output.prompt_token_ids is not None
+                        total_in_toks += len(output.prompt_token_ids) * n
+                        in_spd = total_in_toks / pbar.format_dict["elapsed"]
+                        total_out_toks += sum(
+                            len(stp.token_ids) for stp in output.outputs)
+                        out_spd = (total_out_toks /
+                                    pbar.format_dict["elapsed"])
+                        pbar.postfix = (
+                            f"est. speed input: {in_spd:.2f} toks/s, "
+                            f"output: {out_spd:.2f} toks/s")
+                        pbar.update(n)
+                    else:
+                        pbar.update(1)
+    '''
+    =============================
+    Added by vllm_mlu
+    =============================
+    '''
+    if is_latency_debug:
+        e2e_end_time = self.metric.get_mlu_cost_time()
+        e2e_latency = e2e_end_time - e2e_start_time
+        self.metric.add_metrics(total_request_num, e2e_latency)
+    '''
+    ==================
+    End of addition
+    ==================
+    '''
+
+    if use_tqdm:
+        pbar.close()
+    # Sort the outputs by request ID.
+    # This is necessary because some requests may be finished earlier than
+    # its previous requests.
+    return sorted(outputs, key=lambda x: int(x.request_id))
+
+
+LLM.metric = LLMMetric()
+
+LLM.dump_info = LLMDumpInfo()
+
+MluHijackObject.apply_hijack(LLM,
+                             LLM.__init__,
+                             vllm__entrypoints__llm__LLM____init__)
+MluHijackObject.apply_hijack(LLM,
+                             "get_metrics",
+                             vllm__entrypoints__llm__LLM__get_metrics)
+MluHijackObject.apply_hijack(LLM,
+                             LLM._run_engine,
+                             vllm__entrypoints__llm__LLM___run_engine)
diff --git a/vllm_mlu/vllm_mlu/entrypoints/openai/__init__.py b/vllm_mlu/vllm_mlu/entrypoints/openai/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/entrypoints/openai/serving_engine.py b/vllm_mlu/vllm_mlu/entrypoints/openai/serving_engine.py
new file mode 100644
index 000000000..28b57f8aa
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/entrypoints/openai/serving_engine.py
@@ -0,0 +1,51 @@
+
+from http import HTTPStatus
+from typing import Optional
+
+from vllm.entrypoints.openai.protocol import ErrorResponse
+from vllm.entrypoints.openai.serving_engine import OpenAIServing, AnyRequest
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+async def vllm__entrypoints__openai__serving_engine__OpenAIServing___check_model(
+    self,
+    request: AnyRequest,
+) -> Optional[ErrorResponse]:
+    if self._is_model_supported(request.model):
+        return None
+    if request.model in [
+            lora.lora_name for lora in self.models.lora_requests
+    ]:
+        return None
+    if request.model in [
+            prompt_adapter.prompt_adapter_name
+            for prompt_adapter in self.models.prompt_adapter_requests
+    ]:
+        return None
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: when client send a request with model=init/save_scheduler_view,
+            scheduler will dump profile data.
+    '''
+    if request.model == "init_scheduler_view":
+        await self.engine_client.init_scheduler_view()
+    if request.model == "save_scheduler_view":
+        await self.engine_client.save_scheduler_view()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return self.create_error_response(
+        message=f"The model `{request.model}` does not exist.",
+        err_type="NotFoundError",
+        status_code=HTTPStatus.NOT_FOUND)
+
+
+MluHijackObject.apply_hijack(OpenAIServing,
+                             OpenAIServing._check_model,
+                             vllm__entrypoints__openai__serving_engine__OpenAIServing___check_model)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/executor/__init__.py b/vllm_mlu/vllm_mlu/executor/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/executor/executor_base.py b/vllm_mlu/vllm_mlu/executor/executor_base.py
new file mode 100644
index 000000000..0cf55e36c
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/executor/executor_base.py
@@ -0,0 +1,40 @@
+from vllm.utils import run_method
+from vllm.executor.executor_base import ExecutorBase
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__executor__executor_base__ExecutorBase__get_latency(self) -> float:
+    '''
+    requires that torch.mlu.synchronize() be executed before this function
+    for getting an accurate reading
+    '''
+    latency = run_method(self.driver_worker,
+                         "get_latency", args=[], kwargs={})
+    return latency
+
+
+def vllm__executor__executor_base__ExecutorBase__recapture_model(
+    self,
+    enable_context_mlugraph,
+    context_batch_size_to_capture,
+    context_seq_len_to_capture
+) -> None:
+    return self.collective_rpc("recapture_model", args=(
+        enable_context_mlugraph, context_batch_size_to_capture, context_seq_len_to_capture))
+
+
+def vllm__executor__executor_base__ExecutorBase__get_memory_usage(self):
+    memory_usage = run_method(self.driver_worker,
+                              "get_memory_usage", args=[], kwargs={})
+    return memory_usage
+
+
+MluHijackObject.apply_hijack(ExecutorBase,
+                             "get_latency",
+                             vllm__executor__executor_base__ExecutorBase__get_latency)
+MluHijackObject.apply_hijack(ExecutorBase,
+                             "recapture_model",
+                             vllm__executor__executor_base__ExecutorBase__recapture_model)
+MluHijackObject.apply_hijack(ExecutorBase,
+                             "get_memory_usage",
+                             vllm__executor__executor_base__ExecutorBase__get_memory_usage)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/executor/multiproc_worker_utils.py b/vllm_mlu/vllm_mlu/executor/multiproc_worker_utils.py
new file mode 100644
index 000000000..1ad09a58a
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/executor/multiproc_worker_utils.py
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import sys
+from multiprocessing import Queue
+from typing import Any, Callable
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.utils import get_mp_context, run_method
+from vllm.executor import multiproc_worker_utils
+from vllm.executor.multiproc_worker_utils import (_TERMINATE,
+                                                  Result, _add_prefix)
+from vllm.platforms import current_platform
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__executor__multiproc_worker_utils___run_worker_process(
+    worker_factory: Callable[[VllmConfig, int], Any],
+    task_queue: Queue,
+    result_queue: Queue,
+    vllm_config: VllmConfig,
+    rank: int,
+) -> None:
+    """Worker process event loop"""
+
+    # Add process-specific prefix to stdout and stderr
+    process_name = get_mp_context().current_process().name
+    pid = os.getpid()
+    _add_prefix(sys.stdout, process_name, pid)
+    _add_prefix(sys.stderr, process_name, pid)
+
+    # Initialize worker
+    worker = worker_factory(vllm_config, rank)
+    del worker_factory
+
+    # Accept tasks from the engine in task_queue
+    # and return task output in result_queue
+    logger.info("Worker ready; awaiting tasks")
+    try:
+        for items in iter(task_queue.get, _TERMINATE):
+            output = None
+            exception = None
+            task_id, method, args, kwargs = items
+            try:
+                output = run_method(worker, method, args, kwargs)
+            except SystemExit:
+                raise
+            except KeyboardInterrupt:
+                break
+            except BaseException as e:
+                logger.exception(
+                    "Exception in worker %s while processing method %s.",
+                    process_name, method)
+                exception = e
+            result_queue.put(
+                Result(task_id=task_id, value=output, exception=exception))
+    except KeyboardInterrupt:
+        pass
+    except Exception:
+        logger.exception("Worker failed")
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: avoid torch gpu migration go into cuda
+    '''
+    if not current_platform.is_out_of_tree():
+        # Flush TunableOp results when TunableOp is enabled and
+        # online (in situ) tuning is enabled.
+        # Offline tuning API (record_untuned_is_enabled()) only
+        # available in PyTorch 2.6 or later.
+        if torch.cuda.is_available():
+            import torch.cuda.tunable as tunable
+            if (tunable.is_enabled() and tunable.tuning_is_enabled()
+                    and not tunable.record_untuned_is_enabled()):
+                tunable.write_file()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    logger.info("Worker exiting")
+
+
+MluHijackObject.apply_hijack(
+    multiproc_worker_utils,
+    multiproc_worker_utils._run_worker_process,
+    vllm__executor__multiproc_worker_utils___run_worker_process
+)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/executor/ray_distributed_executor.py b/vllm_mlu/vllm_mlu/executor/ray_distributed_executor.py
new file mode 100644
index 000000000..08f4e3f18
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/executor/ray_distributed_executor.py
@@ -0,0 +1,349 @@
+import os
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+import vllm.envs as envs
+from vllm.executor.ray_utils import (RayWorkerWrapper, ray)
+from vllm.executor.ray_distributed_executor import (RayWorkerMetaData,
+                                                    RayDistributedExecutor)
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils import (get_distributed_init_method,
+                        get_ip, get_open_port)
+
+if ray is not None:
+    from ray.actor import ActorHandle
+    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+else:
+    ActorHandle = None
+
+if TYPE_CHECKING:
+    from ray.util.placement_group import PlacementGroup
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+class RayDistributedExecutor_MluHijack(RayDistributedExecutor):
+
+    def _configure_ray_workers_use_nsight(self,
+                                          ray_remote_kwargs) -> Dict[str, Any]:
+        # If nsight profiling is enabled, we need to set the profiling
+        # configuration for the ray workers as runtime env.
+        runtime_env = ray_remote_kwargs.setdefault("runtime_env", {})
+        runtime_env.update({
+            # use default cnperf config
+            "nsight": "default"
+        })
+
+        return ray_remote_kwargs
+
+    def _init_workers_ray(self, placement_group: "PlacementGroup",
+                          **ray_remote_kwargs):
+        num_gpus = envs.VLLM_RAY_PER_WORKER_GPUS
+
+        # The driver dummy worker does not actually use any resources.
+        # It holds the resource for the driver worker.
+        self.driver_dummy_worker: Optional[RayWorkerWrapper] = None
+        # The remaining workers are the actual ray actors.
+        self.workers: List[RayWorkerWrapper] = []
+
+        # Used in ray compiled DAG: indexed first by PP rank,
+        # and then TP rank. In other words, the inner list is
+        # the TP group of workers for a PP rank.
+        self.pp_tp_workers: List[List[RayWorkerWrapper]] = []
+
+        if self.parallel_config.ray_workers_use_nsight:
+            ray_remote_kwargs = self._configure_ray_workers_use_nsight(
+                ray_remote_kwargs)
+
+        logger.info("use_ray_spmd_worker: %s", self.use_ray_spmd_worker)
+
+        # Create the workers.
+        bundle_indices: List[int]
+        if envs.VLLM_RAY_BUNDLE_INDICES:
+            # Use the bundle indices specified by the user.
+            bundle_indices = list(
+                map(int, envs.VLLM_RAY_BUNDLE_INDICES.split(",")))
+            assert len(bundle_indices) == self.parallel_config.world_size, \
+            ("VLLM_RAY_BUNDLE_INDICES must have the same size"
+            f" as the world size, but got {bundle_indices=} "
+            f"and {self.parallel_config.world_size=}")
+            assert len(set(bundle_indices)) == len(bundle_indices), \
+            ("VLLM_RAY_BUNDLE_INDICES cannot have duplicate values,"
+            f" but got {bundle_indices=}")
+        else:
+            # use the first N bundles that have GPU resources.
+            bundle_indices = []
+            for bundle_id, bundle in enumerate(placement_group.bundle_specs):
+                if bundle.get(current_platform.ray_device_key, 0):
+                    bundle_indices.append(bundle_id)
+            bundle_indices = bundle_indices[:self.parallel_config.world_size]
+
+        worker_metadata: List[RayWorkerMetaData] = []
+        driver_ip = get_ip()
+        for rank, bundle_id in enumerate(bundle_indices):
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: support ray + cnperf-cli
+            '''
+            if self.parallel_config.ray_workers_use_nsight:
+                ray_remote_kwargs['runtime_env'].update({
+                    "nsight": {
+                        "o": f"cnperf_rank_{rank}",
+                        "force_overwrite": "true"
+                    }
+                })
+                if rank == 0:
+                    ray_remote_kwargs['runtime_env'].update({
+                        "nsight": {}
+                    })
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            scheduling_strategy = PlacementGroupSchedulingStrategy(
+                placement_group=placement_group,
+                placement_group_capture_child_tasks=True,
+                placement_group_bundle_index=bundle_id,
+            )
+
+            if current_platform.ray_device_key == "GPU":
+                # NV+AMD GPUs, and Intel XPUs
+                worker = ray.remote(
+                    num_cpus=0,
+                    num_gpus=num_gpus,
+                    scheduling_strategy=scheduling_strategy,
+                    **ray_remote_kwargs,
+                )(RayWorkerWrapper).remote(vllm_config=self.vllm_config,
+                                           rpc_rank=rank)
+            else:
+                worker = ray.remote(
+                    num_cpus=0,
+                    num_gpus=0,
+                    resources={current_platform.ray_device_key: num_gpus},
+                    scheduling_strategy=scheduling_strategy,
+                    **ray_remote_kwargs,
+                )(RayWorkerWrapper).remote(vllm_config=self.vllm_config,
+                                           rpc_rank=rank)
+            worker_metadata.append(
+                RayWorkerMetaData(worker=worker, created_rank=rank))
+
+        worker_ips = ray.get([
+            each.worker.get_node_ip.remote()  # type: ignore[attr-defined]
+            for each in worker_metadata
+        ])
+
+        for each, ip in zip(worker_metadata, worker_ips):
+            each.ip = ip
+
+        if not self.use_ray_spmd_worker:
+            for i, each in enumerate(worker_metadata):
+                # find and remove the dummy worker from the list
+                worker = each.worker
+                worker_ip = each.ip
+                if self.driver_dummy_worker is None and worker_ip == driver_ip:
+                    # If the worker is on the same node as the driver, we use it
+                    # as the resource holder for the driver process.
+                    self.driver_dummy_worker = worker
+                    self.driver_worker = RayWorkerWrapper(
+                        vllm_config=self.vllm_config, rpc_rank=0)
+                    worker_metadata.pop(i)
+                    break
+
+        logger.debug("workers: %s", worker_metadata)
+        logger.debug("driver_dummy_worker: %s", self.driver_dummy_worker)
+        if not self.use_ray_spmd_worker and self.driver_dummy_worker is None:
+            raise ValueError(
+                "Ray does not allocate any GPUs on the driver node."
+                f"Driver IP: {driver_ip}, worker IPs: {worker_ips}."
+                "Consider adjusting the Ray placement group or running "
+                "the driver on a GPU node.")
+
+        ip_counts: Dict[str, int] = {}
+        for ip in worker_ips:
+            ip_counts[ip] = ip_counts.get(ip, 0) + 1
+
+        def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
+            """
+            Sort the workers based on 3 properties:
+            1. If the worker is on the same node as the driver (vllm engine),
+                it should be placed first.
+            2. Then, if the worker is on a node with fewer workers, it should
+                be placed first.
+            3. Finally, if the work is on a node with smaller IP address, it
+                should be placed first.
+            """
+            ip = item.ip
+            return (0 if ip == driver_ip else 1, ip_counts[ip], ip)
+
+        # After sorting, the workers on the same node will be
+        # close to each other, and the workers on the driver
+        # node will be placed first.
+        sorted_worker_metadata = sorted(worker_metadata,
+                                        key=sort_by_driver_then_worker_ip)
+        start_rank = 0 if self.use_ray_spmd_worker else 1
+        for i, item in enumerate(sorted_worker_metadata):
+            item.adjusted_rank = i + start_rank
+        self.workers = [item.worker for item in sorted_worker_metadata]
+        rerank_mapping = {
+            item.created_rank: item.adjusted_rank
+            for item in sorted_worker_metadata
+        }
+        self._run_workers("adjust_rank", rerank_mapping)
+
+        # Get the set of GPU IDs used on each node.
+        worker_node_and_gpu_ids = []
+        for worker in [self.driver_dummy_worker] + self.workers:
+            if worker is None:
+                # driver_dummy_worker can be None when using ray spmd worker.
+                continue
+            worker_node_and_gpu_ids.append(
+                ray.get(worker.get_node_and_gpu_ids.remote()) \
+            ) # type: ignore
+
+        node_workers = defaultdict(list)  # node id -> list of worker ranks
+        node_gpus = defaultdict(list)  # node id -> list of gpu ids
+
+        for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):
+            node_workers[node_id].append(i)
+            # `gpu_ids` can be a list of strings or integers.
+            # convert them to integers for consistency.
+            # NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs),
+            # string sorting is not sufficient.
+            # see https://github.com/vllm-project/vllm/issues/5590
+            gpu_ids = [int(x) for x in gpu_ids]
+            node_gpus[node_id].extend(gpu_ids)
+        for node_id, gpu_ids in node_gpus.items():
+            node_gpus[node_id] = sorted(gpu_ids)
+
+        all_ips = set(worker_ips + [driver_ip])
+        n_ips = len(all_ips)
+        n_nodes = len(node_workers)
+
+        if n_nodes != n_ips:
+            raise RuntimeError(
+                f"Every node should have a unique IP address. Got {n_nodes}"
+                f" nodes with node ids {list(node_workers.keys())} and "
+                f"{n_ips} unique IP addresses {all_ips}. Please check your"
+                " network configuration. If you set `VLLM_HOST_IP`"
+                " environment variable, make sure it is unique for"
+                " each node.")
+
+        # Set environment variables for the driver and workers.
+        all_args_to_update_environment_variables = [{
+            current_platform.device_control_env_var:
+            ",".join(map(str, node_gpus[node_id])),
+        } for (node_id, _) in worker_node_and_gpu_ids]
+
+        # Environment variables to copy from driver to workers
+        env_vars_to_copy = [
+            v for v in envs.environment_variables
+            if v not in self.WORKER_SPECIFIC_ENV_VARS
+            and v not in self.non_carry_over_env_vars
+        ]
+
+        env_vars_to_copy.extend(current_platform.additional_env_vars)
+
+        # Copy existing env vars to each worker's args
+        for args in all_args_to_update_environment_variables:
+            # TODO: refactor platform-specific env vars
+            for name in env_vars_to_copy:
+                if name in os.environ:
+                    args[name] = os.environ[name]
+
+        logger.info("non_carry_over_env_vars from config: %s",
+                    self.non_carry_over_env_vars)
+        logger.info(
+            "Copying the following environment variables to workers: %s",
+            [v for v in env_vars_to_copy if v in os.environ])
+        logger.info(
+            "If certain env vars should NOT be copied to workers, add them to "
+            "%s file", self.non_carry_over_env_vars_file)
+
+        self._env_vars_for_all_workers = (
+            all_args_to_update_environment_variables)
+
+        self._run_workers("update_environment_variables",
+                          self._get_env_vars_to_be_updated())
+
+        if len(node_gpus) == 1:
+            # in single node case, we don't need to get the IP address.
+            # the loopback address is sufficient
+            # NOTE: a node may have several IP addresses, one for each
+            # network interface. `get_ip()` might return any of them,
+            # while they might not work for communication inside the node
+            # if the network setup is complicated. Using the loopback address
+            # solves this issue, as it always works for communication inside
+            # the node.
+            driver_ip = "127.0.0.1"
+        distributed_init_method = get_distributed_init_method(
+            driver_ip, get_open_port())
+
+        # Initialize the actual workers inside worker wrapper.
+        all_kwargs = []
+        for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids):
+            local_rank = node_workers[node_id].index(rank)
+            kwargs = dict(
+                vllm_config=self.vllm_config,
+                local_rank=local_rank,
+                rank=rank,
+                distributed_init_method=distributed_init_method,
+                is_driver_worker=(not self.parallel_config)
+                or (rank % self.parallel_config.tensor_parallel_size == 0),
+            )
+            all_kwargs.append(kwargs)
+        self._run_workers("init_worker", all_kwargs)
+
+        self._run_workers("init_device")
+        self._run_workers("load_model",
+                          max_concurrent_workers=self.parallel_config.
+                          max_parallel_loading_workers)
+
+        if self.use_ray_spmd_worker:
+            for pp_rank in range(self.parallel_config.pipeline_parallel_size):
+                self.pp_tp_workers.append([])
+                for tp_rank in range(
+                        self.parallel_config.tensor_parallel_size):
+                    # PP=2, TP=4
+                    # pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
+                    rank = (pp_rank * self.parallel_config.tensor_parallel_size
+                            ) + tp_rank
+                    assert len(self.pp_tp_workers[pp_rank]) == tp_rank
+                    assert pp_rank < len(self.pp_tp_workers)
+                    self.pp_tp_workers[pp_rank].append(self.workers[rank])
+
+        # This is the list of workers that are rank 0 of each TP group EXCEPT
+        # global rank 0. These are the workers that will broadcast to the
+        # rest of the workers.
+        self.tp_driver_workers: List[RayWorkerWrapper] = []
+        # This is the list of workers that are not drivers and not the first
+        # worker in a TP group. These are the workers that will be
+        # broadcasted to.
+        self.non_driver_workers: List[RayWorkerWrapper] = []
+
+        # Enforce rank order for correct rank to return final output.
+        for index, worker in enumerate(self.workers):
+            # The driver worker is rank 0 and not in self.workers.
+            rank = index + 1
+            if rank % self.parallel_config.tensor_parallel_size == 0:
+                self.tp_driver_workers.append(worker)
+            else:
+                self.non_driver_workers.append(worker)
+
+
+
+MluHijackObject.apply_hijack(
+    RayDistributedExecutor,
+    RayDistributedExecutor._configure_ray_workers_use_nsight,
+    RayDistributedExecutor_MluHijack._configure_ray_workers_use_nsight
+)
+MluHijackObject.apply_hijack(
+    RayDistributedExecutor,
+    RayDistributedExecutor._init_workers_ray,
+    RayDistributedExecutor_MluHijack._init_workers_ray
+)
diff --git a/vllm_mlu/vllm_mlu/lora/__init__.py b/vllm_mlu/vllm_mlu/lora/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/lora/fully_sharded_layers.py b/vllm_mlu/vllm_mlu/lora/fully_sharded_layers.py
new file mode 100644
index 000000000..1625900f0
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/lora/fully_sharded_layers.py
@@ -0,0 +1,68 @@
+from typing import Optional
+
+import torch
+
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.lora.fully_sharded_layers import RowParallelLinearWithShardedLoRA
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__lora__fully_sharded_layers__RowParallelLinearWithShardedLoRA__apply(
+    self,
+    x: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    residual: Optional[torch.Tensor]
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual and bias in matmul
+    '''
+    output = self.base_layer.quant_method.apply(
+        self.base_layer, x, bias, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    x = x.view(-1, x.shape[-1])
+    output, out_orig_shape = output.view(-1,
+                                         output.shape[-1]), output.shape
+    buffer = torch.zeros(
+        (self.n_slices, x.shape[0], self.lora_a_stacked[0].shape[2]),
+        dtype=torch.float32,
+        device=x.device,
+    )
+
+    self.punica_wrapper.add_shrink(buffer, x, self.lora_a_stacked, 1.0)
+    buffer = tensor_model_parallel_all_reduce(buffer)
+
+    # following S-LoRA, allows the fusing of all_gather and all_reduce
+    # by adding the column partitioned lora output to a slice of output
+    # tensor, which is a partial sum due to row parallel. All that
+    # remains is a standard all_reduce. User should be aware though that
+    # the output is not the same as a normal row_parallel, it should be
+    # reduced before being used
+    # NOTE offset are based on the rank.
+    shard_size = self.lora_b_stacked[0].shape[2]
+    offset_start = self.tp_rank * shard_size
+    self.punica_wrapper.add_expand(
+        output,
+        buffer,
+        self.lora_b_stacked,
+        self.lora_bias_stacked,
+        self.output_slices,
+        offset_start=offset_start,
+        add_input=True,
+    )
+    output = output.view(*out_orig_shape)
+    return output
+
+
+MluHijackObject.apply_hijack(
+    RowParallelLinearWithShardedLoRA,
+    RowParallelLinearWithShardedLoRA.apply,
+    vllm__lora__fully_sharded_layers__RowParallelLinearWithShardedLoRA__apply
+)
diff --git a/vllm_mlu/vllm_mlu/lora/layers.py b/vllm_mlu/vllm_mlu/lora/layers.py
new file mode 100644
index 000000000..acb820316
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/lora/layers.py
@@ -0,0 +1,218 @@
+from typing import List, Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers import PretrainedConfig
+
+from vllm.config import LoRAConfig
+from vllm.distributed import (split_tensor_along_last_dim,
+                              tensor_model_parallel_all_reduce)
+from vllm.lora.layers import (ColumnParallelLinearWithLoRA,
+                              RowParallelLinearWithLoRA,
+                              LinearScalingRotaryEmbeddingWithLoRA)
+
+from vllm_mlu.model_executor.layers.rotary_embedding import (
+    MLURotaryEmbedding, MLULinearScalingRotaryEmbedding)
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+vllm__lora__layers__ColumnParallelLinearWithLoRA__forward_org = ColumnParallelLinearWithLoRA.forward
+
+
+'''
+=============================
+Modify by vllm_mlu
+=============================
+@brief: add smooth_quant_scale parameter.
+'''
+def vllm__lora__layers__ColumnParallelLinearWithLoRA__forward(
+    self,
+    input_,
+    smooth_quant_scale: Optional[torch.Tensor] = None
+) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:
+    assert smooth_quant_scale is None, "LoRA does not support smooth quant yet."
+    return vllm__lora__layers__ColumnParallelLinearWithLoRA__forward_org(self, input_)
+'''
+==================
+End of MLU Hijack
+==================
+'''
+
+
+def vllm__lora__layers__RowParallelLinearWithLoRA__apply(
+    self,
+    x: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    residual: Optional[torch.Tensor]
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual and bias in matmul
+    '''
+    output = self.base_layer.quant_method.apply(
+        self.base_layer, x, bias, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.punica_wrapper.add_lora_linear(output, x, self.lora_a_stacked,
+                                        self.lora_b_stacked,
+                                        self.lora_bias_stacked, 1.0,
+                                        self.output_slices)
+    return output
+
+
+def vllm__lora__layers__RowParallelLinearWithLoRA__forward(
+    self,
+    input_: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None
+) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:
+    # Set up backprop all-reduce.
+    if self.base_layer.input_is_parallel:
+        input_parallel = input_
+    else:
+        # TODO: simplify code below
+        splitted_input = split_tensor_along_last_dim(
+            input_, num_partitions=self.base_layer.tp_size)
+        input_parallel = splitted_input[self.tp_rank].contiguous()
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: 1) apply residual fusion in matmul like RowParallelLinear
+    2) add bias in matmul, not after all reduce
+    '''
+    # Matrix multiply.
+    bias_ = (None if (self.base_layer.tp_rank > 0 or self.base_layer.skip_bias_add) else self.base_layer.bias)
+    residual_ = None if self.base_layer.tp_rank > 0 else residual
+    output_parallel = self.apply(input_parallel, bias_, residual_)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    if self.base_layer.reduce_results and self.base_layer.tp_size > 1:
+        output = tensor_model_parallel_all_reduce(output_parallel)
+    else:
+        output = output_parallel
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: do not add bias after all_reduce
+    '''
+    output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    if not self.base_layer.return_bias:
+        return output
+    return output, output_bias
+
+
+def vllm__lora__layers__LinearScalingRotaryEmbeddingWithLoRA__create_lora_weights(
+    self,
+    max_loras: int,
+    lora_config: LoRAConfig,
+    model_config: Optional[PretrainedConfig] = None,
+) -> None:
+    scaling_factors = (list(lora_config.long_lora_scaling_factors)
+                       if lora_config.long_lora_scaling_factors else [])
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: change LinearScalingRotaryEmbedding to MLULinearScalingRotaryEmbedding
+    '''
+    base_scaling_factor = (self.base_layer.scaling_factor if isinstance(
+        self.base_layer, MLULinearScalingRotaryEmbedding) else 1.0)
+    scaling_factors = sorted(
+        list(set([base_scaling_factor] + scaling_factors)))
+    self.base_layer = MLULinearScalingRotaryEmbedding(
+        self.base_layer.head_size,
+        self.base_layer.rotary_dim,
+        self.base_layer.max_position_embeddings,
+        self.base_layer.base,
+        self.base_layer.is_neox_style,
+        scaling_factors,
+        self.base_layer.dtype,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__lora__layers__LinearScalingRotaryEmbeddingWithLoRA__forward(
+    self,
+    positions: torch.Tensor,
+    qk: torch.Tensor
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: change function prototype to meet forward_mlu in rope
+    '''
+    return self.base_layer(
+        positions,
+        qk,
+        offsets=self.punica_wrapper.long_lora_indices,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+@classmethod
+def vllm__lora__layers__LinearScalingRotaryEmbeddingWithLoRA__can_replace_layer(
+    cls,
+    source_layer: nn.Module,
+    lora_config: LoRAConfig,
+    packed_modules_list: List,
+    model_config: Optional[PretrainedConfig],
+) -> bool:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: change origin rope type to mlu rope
+    '''
+    return (type(source_layer) is MLULinearScalingRotaryEmbedding
+            or type(source_layer) is MLURotaryEmbedding)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(RowParallelLinearWithLoRA,
+                             RowParallelLinearWithLoRA.apply,
+                             vllm__lora__layers__RowParallelLinearWithLoRA__apply)
+MluHijackObject.apply_hijack(ColumnParallelLinearWithLoRA,
+                             ColumnParallelLinearWithLoRA.forward,
+                             vllm__lora__layers__ColumnParallelLinearWithLoRA__forward)
+MluHijackObject.apply_hijack(RowParallelLinearWithLoRA,
+                             RowParallelLinearWithLoRA.forward,
+                             vllm__lora__layers__RowParallelLinearWithLoRA__forward)
+MluHijackObject.apply_hijack(LinearScalingRotaryEmbeddingWithLoRA,
+                             LinearScalingRotaryEmbeddingWithLoRA.create_lora_weights,
+                             vllm__lora__layers__LinearScalingRotaryEmbeddingWithLoRA__create_lora_weights)
+MluHijackObject.apply_hijack(LinearScalingRotaryEmbeddingWithLoRA,
+                             LinearScalingRotaryEmbeddingWithLoRA.forward,
+                             vllm__lora__layers__LinearScalingRotaryEmbeddingWithLoRA__forward)
+MluHijackObject.apply_hijack(LinearScalingRotaryEmbeddingWithLoRA,
+                             LinearScalingRotaryEmbeddingWithLoRA.can_replace_layer,
+                             vllm__lora__layers__LinearScalingRotaryEmbeddingWithLoRA__can_replace_layer)
diff --git a/vllm_mlu/vllm_mlu/lora/ops/__init__.py b/vllm_mlu/vllm_mlu/lora/ops/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/lora/ops/triton_ops/__init__.py b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/__init__.py
new file mode 100644
index 000000000..1d40fde38
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/__init__.py
@@ -0,0 +1,9 @@
+from vllm_mlu.lora.ops.triton_ops.sgmv_expand import sgmv_expand_mlu
+from vllm_mlu.lora.ops.triton_ops.sgmv_expand_slice import sgmv_expand_slice_mlu
+from vllm_mlu.lora.ops.triton_ops.sgmv_shrink import sgmv_shrink_mlu
+
+__all__ = [
+    "sgmv_expand_mlu",
+    "sgmv_expand_slice_mlu",
+    "sgmv_shrink_mlu",
+]
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_expand.py b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_expand.py
new file mode 100644
index 000000000..56d72e9ce
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_expand.py
@@ -0,0 +1,235 @@
+import torch
+import triton
+import triton.language as tl
+
+from vllm_mlu.lora.ops.triton_ops.utils import adjust_kernel_block_size
+
+from vllm.utils import direct_register_custom_op
+
+
+@triton.jit
+def _sgmv_expand_kernel_mlu(
+    input_ptr,
+    lora_ptr,
+    out_ptr,
+    N,
+    K,
+    b_seq_start_loc,
+    seq_lens,
+    lora_indices,
+    xm_stride,
+    xk_stride,  # 1
+    l0_stride,  # hidden_size*max_rank
+    lora_k_stride,
+    lora_n_stride,
+    cm_stride,
+    cn_stride,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    EVEN_K: tl.constexpr,
+    ADD_INPUTS: tl.constexpr,
+    CAST_TYPE: tl.constexpr,
+):
+    """
+    The sgmv's expand triton kernel is based on GroupGEMM.
+    """
+    pid = tl.program_id(axis=0)
+    cur_batch = tl.program_id(axis=1)
+    cta_n_num = tl.cdiv(N, BLOCK_N)
+    pid_m = pid // cta_n_num
+    pid_n = pid % cta_n_num
+    M = tl.load(seq_lens + cur_batch)
+    if pid_m * BLOCK_M > M:
+        return
+    lora_index = tl.load(lora_indices + cur_batch)
+    if lora_index == -1:
+        return
+    cur_seq_start = tl.load(b_seq_start_loc + cur_batch)
+    offset_m = tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
+    offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    offset_k = tl.arange(0, BLOCK_K)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: adjust kernel impl to fit mlu.
+    '''
+    a_ptr = input_ptr + cur_seq_start * xm_stride + offset_m[:, None] * xm_stride + \
+             offset_k[None, :] * xk_stride
+    b_ptr = lora_ptr + l0_stride * lora_index + \
+            offset_k[:, None] * lora_n_stride + offset_n[None, :] * lora_k_stride
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(tl.cdiv(K, BLOCK_K)):
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: adjust kernel impl to fit mlu.
+        '''
+        if EVEN_K:
+            tiled_a = tl.load(a_ptr, mask=offset_m[:, None] < M)
+            tiled_b = tl.load(b_ptr, mask=offset_n[None, :] < N)
+        else:
+            tiled_a = tl.load(a_ptr,
+                              mask=((offset_k[None, :] < K - k * BLOCK_K) & (offset_m[:, None] < M)),
+                              other=0)
+            tiled_b = tl.load(b_ptr,
+                              mask=((offset_k[:, None] < K - k * BLOCK_K) & (offset_n[None, :] < N)),
+                              other=0)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        if CAST_TYPE:
+            tiled_a = tiled_a.to(lora_ptr.dtype.element_ty)
+        accumulator += tl.dot(
+            tiled_a,
+            tiled_b,
+        )
+        a_ptr += BLOCK_K * xk_stride
+        b_ptr += BLOCK_K * lora_n_stride
+    tiled_c = accumulator.to(lora_ptr.dtype.element_ty)
+    offset_cm = cur_seq_start + tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
+    offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    c_ptr = (out_ptr + offset_cm[:, None] * cm_stride +
+             offset_cn[None, :] * cn_stride)
+    M = tl.load(seq_lens + cur_batch)
+    c_mask = (offset_cm[:, None] <
+              (cur_seq_start + M)) & (offset_cn[None, :] < N)
+    if ADD_INPUTS:
+        tiled_out = tl.load(c_ptr, mask=c_mask)
+        tiled_c += tiled_out
+    tl.store(c_ptr, tiled_c, mask=c_mask)
+
+
+@torch.inference_mode()
+def sgmv_expand_mlu(
+    inputs: torch.Tensor,
+    lora_b_weights: torch.Tensor,
+    output_tensor: torch.Tensor,
+    b_seq_start_loc: torch.Tensor,
+    seq_len_tensor: torch.Tensor,
+    lora_indices_tensor: torch.Tensor,
+    batches: int,
+    max_seq_length: int,
+    token_nums: int,
+    add_inputs: bool = False,
+) -> None:
+    """
+    Args:
+        inputs (torch.Tensor): input tensor
+        lora_b_weights (torch.Tensor): lora'a weight
+        output_tensor (torch.Tensor): output tensor
+        b_seq_start_loc (torch.Tensor): (batch_size,). The cumulative
+            sequence lengths of the sequences in the batch, used to index
+            into sequence. E.g., if the sequence length is [4, 6], it is
+            [0, 4, 10].
+        seq_len_tensor (torch.Tensor): (batch_size,). Record the sequence
+            length of the sequences in the batch.
+        lora_indices_tensor (torch.Tensor): (batch_size,). The LoRA index
+            corresponding to each batch. An index of -1 means no lora should be
+            applied.
+        batches (int): batch size
+        max_seq_length (int): The max sequence lengths of the sequences in the
+            batch.
+        token_nums (int): The token numbers in the batch. Used to verify if the
+            token numbers in the inputs matches the one in the metadata.
+        add_inputs (bool, optional): Defaults to False, adds the final lora
+            results to the output.
+    """
+
+    assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
+    assert lora_b_weights.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ]
+    assert inputs.size(0) == token_nums
+    assert inputs.size(1) == lora_b_weights.size(-1)
+    assert b_seq_start_loc.size(0) == batches
+    assert lora_indices_tensor.size(0) == batches
+    assert inputs.is_contiguous()
+    assert output_tensor.is_contiguous()
+
+    if lora_b_weights.ndim == 4:  # shape:(lora_num,1,size,rank)
+        assert lora_b_weights.size(1) == 1
+        lora_b_weights = lora_b_weights.squeeze(dim=1)
+    else:
+        assert lora_b_weights.ndim == 3  # shape:(lora_num,size,rank)
+
+    assert lora_b_weights.is_contiguous()
+
+    # TODO tuning this config
+
+    N, K = lora_b_weights.shape[-2:]  # K= rank,N=hidden_size
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Workaround: Adjust block size to meet mlu restrictions.
+
+    The grid of mlu triton kernel must less than 65536, it will be out of bound when
+    the input seq is very long, and causes runtime error. So we need to adjust the block
+    size to avoid this.
+    '''
+    BLOCK_M, BLOCK_N = adjust_kernel_block_size(max_seq_length, 32, N, 32)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    BLOCK_K = 16
+    EVEN_K = K % BLOCK_K == 0
+    ADD_INPUTS = add_inputs
+    CAST_TYPE = False
+    if inputs.dtype == torch.float32 and lora_b_weights.dtype in [
+            torch.float16,
+            torch.bfloat16,
+    ]:
+        CAST_TYPE = True
+    grid = (
+        triton.cdiv(max_seq_length, BLOCK_M) * triton.cdiv(N, BLOCK_N),
+        batches,
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: call _sgmv_expand_kernel_mlu
+    '''
+    _sgmv_expand_kernel_mlu[grid](
+        inputs,
+        lora_b_weights,
+        output_tensor,
+        N,
+        K,
+        b_seq_start_loc,
+        seq_len_tensor,
+        lora_indices_tensor,
+        inputs.stride(0),
+        inputs.stride(1),
+        lora_b_weights.stride(0),
+        lora_b_weights.stride(1),
+        lora_b_weights.stride(2),
+        output_tensor.stride(0),
+        output_tensor.stride(1),
+        BLOCK_M,
+        BLOCK_N,
+        BLOCK_K,
+        EVEN_K,
+        ADD_INPUTS,
+        CAST_TYPE,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return
diff --git a/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_expand_slice.py b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_expand_slice.py
new file mode 100644
index 000000000..19af45e73
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_expand_slice.py
@@ -0,0 +1,245 @@
+import torch
+import triton
+import triton.language as tl
+
+from vllm_mlu.lora.ops.triton_ops.utils import adjust_kernel_block_size
+
+from vllm.utils import direct_register_custom_op
+
+@triton.jit
+def _sgmv_expand_slice_kernel_mlu(
+    input_ptr,
+    lora_ptr,
+    out_ptr,
+    N,
+    K,
+    b_seq_start_loc,
+    seq_lens,
+    lora_indices,
+    xm_stride,
+    xk_stride,  # 1
+    l0_stride,  # hidden_size*max_rank
+    lora_k_stride,
+    lora_n_stride,
+    cm_stride,
+    cn_stride,
+    slice_offset,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    EVEN_K: tl.constexpr,
+    ADD_INPUTS: tl.constexpr,
+    CAST_TYPE: tl.constexpr,
+):
+    """
+
+    Similar to the 'sgmv_expand' operator, but with an added parameter
+    'slice_offset'. The reason for not reusing the 'sgmv_expand' operator
+    might be that in the future, we could implement a fusion operator to
+    achieve the current functionality instead of having to call it multiple
+    times.
+    """
+    pid = tl.program_id(axis=0)
+    cur_batch = tl.program_id(axis=1)
+    cta_n_num = tl.cdiv(N, BLOCK_N)
+    pid_m = pid // cta_n_num
+    pid_n = pid % cta_n_num
+    M = tl.load(seq_lens + cur_batch)
+    if pid_m * BLOCK_M > M:
+        return
+    lora_index = tl.load(lora_indices + cur_batch)
+    if lora_index == -1:
+        return
+    cur_seq_start = tl.load(b_seq_start_loc + cur_batch)
+    offset_m = tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
+    offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    offset_k = tl.arange(0, BLOCK_K)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: adjust kernel impl to fit mlu.
+    '''
+    a_ptr = input_ptr + cur_seq_start * xm_stride + offset_m[:, None] * xm_stride + \
+            offset_k[None, :] * xk_stride
+    b_ptr = lora_ptr + l0_stride * lora_index + \
+            offset_k[:, None] * lora_n_stride + offset_n[None, :] * lora_k_stride
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(tl.cdiv(K, BLOCK_K)):
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: adjust kernel impl to fit mlu.
+        '''
+        if EVEN_K:
+            tiled_a = tl.load(a_ptr, mask=offset_m[:, None] < M)
+            tiled_b = tl.load(b_ptr, mask=offset_n[None, :] < N)
+        else:
+            tiled_a = tl.load(a_ptr,
+                              mask=((offset_k[None, :] < K - k * BLOCK_K) & (offset_m[:, None] < M)),
+                              other=0)
+            tiled_b = tl.load(b_ptr,
+                              mask=((offset_k[:, None] < K - k * BLOCK_K) & (offset_n[None, :] < N)),
+                              other=0)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        if CAST_TYPE:
+            tiled_a = tiled_a.to(lora_ptr.dtype.element_ty)
+        accumulator += tl.dot(
+            tiled_a,
+            tiled_b,
+        )
+        a_ptr += BLOCK_K * xk_stride
+        b_ptr += BLOCK_K * lora_n_stride
+    tiled_c = accumulator.to(lora_ptr.dtype.element_ty)
+    offset_cm = cur_seq_start + tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
+    offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + slice_offset
+    c_ptr = (out_ptr + offset_cm[:, None] * cm_stride +
+             offset_cn[None, :] * cn_stride)
+    M = tl.load(seq_lens + cur_batch)
+    c_mask = (offset_cm[:, None] < (cur_seq_start + M)) & (offset_cn[None, :] <
+                                                           (slice_offset + N))
+    if ADD_INPUTS:
+        tiled_out = tl.load(c_ptr, mask=c_mask)
+        tiled_c += tiled_out
+    tl.store(c_ptr, tiled_c, mask=c_mask)
+
+
+@torch.inference_mode()
+def sgmv_expand_slice_mlu(
+    inputs: torch.Tensor,
+    lora_b_weights: torch.Tensor,
+    output_tensor: torch.Tensor,
+    b_seq_start_loc: torch.Tensor,
+    seq_len_tensor: torch.Tensor,
+    lora_indices_tensor: torch.Tensor,
+    batches: int,
+    max_seq_length: int,
+    token_nums: int,
+    slice_offset: int,
+    slice_size: int,
+    add_inputs: bool = False,
+) -> None:
+    """_summary_
+
+    Args:
+        inputs (torch.Tensor): input tensor
+        lora_b_weights (torch.Tensor): lora'a weight
+        output_tensor (torch.Tensor): output tensor
+        b_seq_start_loc (torch.Tensor): (batch_size,). The cumulative
+            sequence lengths of the sequences in the batch, used to index
+            into sequence. E.g., if the sequence length is [4, 6], it is
+            [0, 4, 10].
+        seq_len_tensor (torch.Tensor): (batch_size,). Record the sequence
+            length of the sequences in the batch
+        lora_indices_tensor (torch.Tensor): (batch_size,). The LoRA index
+            corresponding to each batch. An index of -1 means no lora should be
+            applied.
+        batches (int): batch size
+        max_seq_length (int): The max sequence lengths of the sequences
+            in the batch
+        token_nums (int): The token numbers in the batch. Used to verify if the
+            token numbers in the inputs matches the one in the metadata.
+        slice_offset (int): output_tensor's offset
+        slice_size (int): current output_tensor's size
+        add_inputs (bool, optional): Defaults to False, adds the final lora
+            results to the output.
+    """
+
+    assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
+    assert lora_b_weights.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ]
+    assert inputs.size(0) == token_nums
+    assert inputs.size(1) == lora_b_weights.size(-1)
+    assert b_seq_start_loc.size(0) == batches
+    assert lora_indices_tensor.size(0) == batches
+    assert slice_size == lora_b_weights.size(-2)
+    assert inputs.is_contiguous()
+    assert output_tensor.is_contiguous()
+
+    if lora_b_weights.ndim == 4:  # shape:(lora_num,1,size,rank)
+        assert lora_b_weights.size(1) == 1
+        lora_b_weights = lora_b_weights.squeeze(dim=1)
+    else:
+        assert lora_b_weights.ndim == 3  # shape:(lora_num,size,rank)
+
+    assert lora_b_weights.is_contiguous()
+
+    # TODO tuning this config
+    N, K = lora_b_weights.shape[-2:]  # K= rank,N=hidden_size
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Workaround: Adjust block size to meet mlu restrictions.
+
+    The grid of mlu triton kernel must less than 65536, it will be out of bound when
+    the input seq is very long, and causes runtime error. So we need to adjust the block
+    size to avoid this.
+    '''
+    BLOCK_M, BLOCK_N = adjust_kernel_block_size(max_seq_length, 32, N, 32)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    BLOCK_K = 16
+    EVEN_K = K % BLOCK_K == 0
+    ADD_INPUTS = add_inputs
+    CAST_TYPE = False
+    if inputs.dtype == torch.float32 and lora_b_weights.dtype in [
+            torch.float16,
+            torch.bfloat16,
+    ]:
+        CAST_TYPE = True
+    grid = (
+        triton.cdiv(max_seq_length, BLOCK_M) * triton.cdiv(N, BLOCK_N),
+        batches,
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: call _sgmv_expand_kernel_mlu
+    '''
+    _sgmv_expand_slice_kernel_mlu[grid](
+        inputs,
+        lora_b_weights,
+        output_tensor,
+        N,
+        K,
+        b_seq_start_loc,
+        seq_len_tensor,
+        lora_indices_tensor,
+        inputs.stride(0),
+        inputs.stride(1),
+        lora_b_weights.stride(0),
+        lora_b_weights.stride(1),
+        lora_b_weights.stride(2),
+        output_tensor.stride(0),
+        output_tensor.stride(1),
+        slice_offset,
+        BLOCK_M,
+        BLOCK_N,
+        BLOCK_K,
+        EVEN_K,
+        ADD_INPUTS,
+        CAST_TYPE,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return
diff --git a/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_shrink.py b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_shrink.py
new file mode 100644
index 000000000..16d0ba914
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/sgmv_shrink.py
@@ -0,0 +1,228 @@
+import torch
+import triton
+import triton.language as tl
+
+from vllm_mlu.lora.ops.triton_ops.utils import adjust_kernel_block_size
+
+from vllm.utils import direct_register_custom_op
+
+
+@triton.jit
+def _sgmv_shrink_kernel_mlu(
+    input_ptr,
+    lora_ptr,
+    out_ptr,
+    N,
+    K,
+    b_seq_start_loc,
+    seq_lens,
+    lora_indices,
+    scaling,
+    xm_stride,  # hidden_size
+    xk_stride,  # 1
+    l0_stride,  # hidden_size*max_rank
+    lora_k_stride,
+    lora_n_stride,
+    cm_stride,
+    cn_stride,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    EVEN_K: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+):
+    """
+    The sgmv's shrink triton kernel is based on GroupGEMM+SPLIT-K.
+    The GEMM of Multi-LoRA can be considered as GroupGEMM. Additionally,
+    introducing SPLIT-K can improve performance
+    """
+    pid = tl.program_id(axis=0)
+    pid_sk = tl.program_id(axis=1)
+    cur_batch = tl.program_id(axis=2)
+    cta_n_num = tl.cdiv(N, BLOCK_N)
+    pid_m = pid // cta_n_num
+    pid_n = pid % cta_n_num
+
+    M = tl.load(seq_lens + cur_batch)
+    if pid_m * BLOCK_M > M:
+        return
+    lora_index = tl.load(lora_indices + cur_batch)
+    if lora_index == -1:
+        return
+    cur_seq_start = tl.load(b_seq_start_loc + cur_batch)
+    offset_m = tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
+    offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: adjust kernel impl to fit mlu.
+    '''
+    a_ptr = input_ptr + cur_seq_start * xm_stride + offset_m[:, None] * xm_stride + \
+            offset_k[None, :] * xk_stride
+    b_ptr = lora_ptr + l0_stride * lora_index + offset_n[None, :] * lora_k_stride + \
+            offset_k[:, None] * lora_n_stride
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: adjust kernel impl to fit mlu.
+        '''
+        if EVEN_K:
+            tiled_a = tl.load(a_ptr, mask=offset_m[:, None] < M)
+            tiled_b = tl.load(b_ptr, mask=offset_n[None, :] < N)
+        else:
+            k_remaining = K - k * (BLOCK_K * SPLIT_K)
+            tiled_a = tl.load(a_ptr,
+                              mask=((offset_k[None, :] < k_remaining) & (offset_m[:, None] < M)),
+                              other=0.0)
+            tiled_b = tl.load(b_ptr,
+                              mask=((offset_k[:, None] < k_remaining) & (offset_n[None, :] < N)),
+                              other=0.0)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        accumulator += tl.dot(tiled_a, tiled_b)
+
+        a_ptr += BLOCK_K * SPLIT_K * xk_stride
+        b_ptr += BLOCK_K * SPLIT_K * lora_n_stride
+    offset_cm = cur_seq_start + tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
+
+    offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    c_ptr = (out_ptr + offset_cm[:, None] * cm_stride +
+             offset_cn[None, :] * cn_stride)
+    c_mask = (offset_cm[:, None] <
+              (cur_seq_start + M)) & (offset_cn[None, :] < N)
+    accumulator *= scaling
+    # handles write-back with reduction-splitting
+    if SPLIT_K == 1:
+        tl.store(c_ptr, accumulator, mask=c_mask)
+    else:
+        tl.atomic_add(c_ptr, accumulator, mask=c_mask)
+
+
+@torch.inference_mode()
+def sgmv_shrink_mlu(
+    inputs: torch.Tensor,
+    lora_a_weights: torch.Tensor,
+    output_tensor: torch.Tensor,
+    b_seq_start_loc: torch.Tensor,
+    seq_len_tensor: torch.Tensor,
+    lora_indices_tensor: torch.Tensor,
+    batches: int,
+    max_seq_length: int,
+    token_nums: int,
+    scaling: float,
+) -> None:
+    """
+    Args:
+        inputs (torch.Tensor): input tensor
+        lora_a_weights (torch.Tensor): lora'a weight
+        output_tensor (torch.Tensor): output tensor
+        b_seq_start_loc (torch.Tensor): (batch_size,). The cumulative
+            sequence lengths of the sequences in the batch, used to index
+            into sequence. E.g., if the sequence length is [4, 6], it is
+            [0, 4].
+        seq_len_tensor (torch.Tensor): (batch_size,). Record the sequence
+            length of the sequences in the batch.
+        lora_indices_tensor (torch.Tensor): (batch_size,). The LoRA index
+            corresponding to each batch. An index of -1 means no lora should be
+            applied.
+        batches (int): batch size
+        max_seq_length (int): The max sequence lengths of the sequences in the
+            batch.
+        token_nums (int): The token numbers in the batch. Used to verify if the
+            token numbers in the inputs matches the one in the metadata.
+        scaling (float): Scaling factor.
+    """
+    assert inputs.dtype == lora_a_weights.dtype
+    assert inputs.dtype in [torch.float16, torch.bfloat16]
+    assert lora_a_weights.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ]
+    assert inputs.size(0) == token_nums
+    assert inputs.size(1) == lora_a_weights.size(-1)
+    assert b_seq_start_loc.size(0) == batches
+    assert lora_indices_tensor.size(0) == batches
+    assert inputs.is_contiguous()
+
+    if lora_a_weights.ndim == 4:  # shape:(lora_num,1,rank, size)
+        assert lora_a_weights.size(1) == 1
+        lora_a_weights = lora_a_weights.squeeze(dim=1)
+    else:
+        assert lora_a_weights.ndim == 3  # shape:(lora_num,rank, size)
+    assert lora_a_weights.is_contiguous()
+    assert output_tensor.is_contiguous()
+    # TODO tuning this config
+    N, K = lora_a_weights.shape[-2:]  # K=hidden_size,N=rank
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Workaround: adjust block size to meet mlu restrictions.
+
+    The grid of mlu triton kernel must less than 65536, it will be out of bound when
+    the input seq is very long, and causes runtime error. So we need to adjust the block
+    size to avoid this.
+    '''
+    BLOCK_M, BLOCK_N = adjust_kernel_block_size(max_seq_length, 32, N, 16)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    BLOCK_K = 32
+    SPLIT_K = 8
+    EVEN_K = K % (BLOCK_K * SPLIT_K) == 0
+    grid = (
+        triton.cdiv(max_seq_length, BLOCK_M) * triton.cdiv(N, BLOCK_N),
+        SPLIT_K,
+        batches,
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: call _sgmv_shrink_kernel_mlu
+    '''
+    _sgmv_shrink_kernel_mlu[grid](
+        inputs,
+        lora_a_weights,
+        output_tensor,
+        N,
+        K,
+        b_seq_start_loc,
+        seq_len_tensor,
+        lora_indices_tensor,
+        scaling,
+        inputs.stride(0),
+        inputs.stride(1),
+        lora_a_weights.stride(0),
+        lora_a_weights.stride(1),
+        lora_a_weights.stride(2),
+        output_tensor.stride(0),
+        output_tensor.stride(1),
+        BLOCK_M,
+        BLOCK_N,
+        BLOCK_K,
+        EVEN_K,
+        SPLIT_K,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return
diff --git a/vllm_mlu/vllm_mlu/lora/ops/triton_ops/utils.py b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/utils.py
new file mode 100644
index 000000000..988139a3d
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/lora/ops/triton_ops/utils.py
@@ -0,0 +1,38 @@
+from typing import Tuple
+from math import ceil
+
+_MLU_MAX_GRID_SIZE = 65536
+
+def adjust_kernel_block_size(
+    m: int,
+    block_m: int,
+    n: int,
+    block_n: int
+) -> Tuple[int, int]:
+    """Adjust block size to meet mlu triton grid restrictions.
+
+    Calculation of the max block size in candidates list:
+
+    LLama3.1-8b-tp1     max n is 14336
+    LLama3.1-70b-tp4    max n is 7168
+    LLama3.1-405b-tp8   max n is 6656
+
+    when n is 14336, the max sequence length of block size 256 can be
+    floor(65536 / ceil(14336 / 256)) * 256 = 299520.
+    """
+    candidates_list = [16, 32, 64, 96, 128, 192, 256]
+    candidates_list_len = len(candidates_list)
+    m_idx = 1
+    n_idx = 0 if block_n == 16 else 1
+    while m_idx < candidates_list_len and n_idx < candidates_list_len:
+        block_m = candidates_list[m_idx]
+        block_n = candidates_list[n_idx]
+        if ceil(m / block_m) * ceil(n / block_n) < _MLU_MAX_GRID_SIZE:
+            break
+        if m_idx < candidates_list_len:
+            m_idx += 1
+        if n_idx < candidates_list_len:
+            n_idx += 1
+    if ceil(m / block_m) * ceil(n / block_n) >= _MLU_MAX_GRID_SIZE:
+        raise ValueError(f"the max seq len {m} is too long for lora triton kernel")
+    return block_m, block_n
diff --git a/vllm_mlu/vllm_mlu/lora/punica_wrapper/__init__.py b/vllm_mlu/vllm_mlu/lora/punica_wrapper/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/lora/punica_wrapper/punica_mlu.py b/vllm_mlu/vllm_mlu/lora/punica_wrapper/punica_mlu.py
new file mode 100644
index 000000000..af88d9aa5
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/lora/punica_wrapper/punica_mlu.py
@@ -0,0 +1,86 @@
+"""
+Based on:
+Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023). 
+Punica: Multi-Tenant LoRA Serving. 
+https://arxiv.org/abs/2310.18547
+"""
+
+from typing import Optional, Tuple, Union, final
+
+import torch
+
+from vllm.triton_utils import HAS_TRITON
+
+if HAS_TRITON:
+    from vllm_mlu.lora.ops.triton_ops import sgmv_expand_mlu
+    from vllm_mlu.lora.ops.triton_ops import sgmv_expand_slice_mlu
+    from vllm_mlu.lora.ops.triton_ops import sgmv_shrink_mlu
+
+from vllm.lora.punica_wrapper.punica_cpu import PunicaWrapperCPU
+
+
+@final
+class PunicaWrapperMLU(PunicaWrapperCPU):
+    """
+    PunicaWrapperMLU is designed to manage and provide metadata for the punica 
+    kernel. The main function is to maintain the state information for 
+    Multi-LoRA, and to provide the interface for the punica triton kernel.
+    """
+
+    def _shrink_prefill(
+        self,
+        y: torch.Tensor,
+        x: torch.Tensor,
+        w_t_all: torch.Tensor,
+        scale: float,
+    ):
+        #No LoRA request, so return directly
+        if self.no_lora:
+            return
+        sgmv_shrink_mlu(
+            x,
+            w_t_all,
+            y,
+            *self.prefill_metadata,
+            scale,
+        )
+
+    def _expand_prefill(
+        self,
+        y: torch.Tensor,
+        x: torch.Tensor,
+        w_t_all: torch.Tensor,
+        add_inputs: bool,
+    ):
+        #No LoRA request, so return directly
+        if self.no_lora:
+            return
+        sgmv_expand_mlu(
+            x,
+            w_t_all,
+            y,
+            *self.prefill_metadata,
+            add_inputs,
+        )
+
+    def _expand_slice_prefill(
+        self,
+        y: torch.Tensor,
+        x: torch.Tensor,
+        w_t_all: torch.Tensor,
+        y_offset: int,
+        y_slice_size: int,
+        add_inputs: bool,
+    ):
+        #No LoRA request, so return directly
+        if self.no_lora:
+            return
+        sgmv_expand_slice_mlu(
+            x,
+            w_t_all,
+            y,
+            *self.prefill_metadata,
+            y_offset,
+            y_slice_size,
+            add_inputs,
+        )
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/common/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/common/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/layers/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/layers/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/layers/feed_forward.py b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/layers/feed_forward.py
new file mode 100644
index 000000000..d33524381
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/layers/feed_forward.py
@@ -0,0 +1,112 @@
+import torch
+from typing import Optional
+
+from vllm.distributed.parallel_state import get_tp_group, get_tensor_model_parallel_group
+from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+from vllm.distributed import get_tensor_model_parallel_rank
+from vllm.lora.layers import BaseLayerWithLoRA
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu._mlu_utils import *
+
+
+def vllm_mlu__model_executor__layers__feed_forward__FeedForward__forward(
+    self, 
+    hidden_states, 
+    residual: Optional[torch.Tensor] = None
+):
+    self.prepare_weight()
+
+    if self.use_bt_ffn is False:
+        return self.forward_naive(hidden_states, residual, None)
+
+    up_proj = getattr(self, self.up_proj_name)
+    down_proj = getattr(self, self.down_proj_name)
+    residual_ = None if self.tp_rank > 0 else residual
+    if (self.quant_config is None and not isinstance(up_proj, BaseLayerWithLoRA)
+            and not isinstance(down_proj, BaseLayerWithLoRA)):
+        # The matmul formula is the following:
+        #   mul_out = alpha * (matmul(input, filter, transpose\_b=True) + bias) + beta * residual
+        #   output = active(mul_out)
+        # Notes: We cannot use the activation function in matmul because it does not support gated operation
+        #  we might support its in tmo matmul in the future
+        fc1 = mlu_ops.matmul(hidden_states.view(-1, self.hidden_size), up_proj.weight, up_proj.bias, 
+                            None, 'none', self.alpha, self.beta)
+        act_out = mlu_ops.active(fc1, self.hidden_act, self.is_gated)
+        beta = 1.0 if residual_ is not None else 0.0
+        '''
+        =======================================
+        Modify by custom vllm_mlu
+        =======================================
+        @brief: call parallel op and abandon original reduce if parallel_num is set
+        '''
+        is_parallel_enable = hasattr(self, 'parallel_num') and get_is_prompt()
+        if is_parallel_enable:
+            rank = get_tensor_model_parallel_rank()
+            pg = get_tensor_model_parallel_group().device_group
+            cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
+            out_ = mlu_ops.matmul_allreduce(cncl_comm, act_out, down_proj.weight, None, residual_,
+                                            self.alpha, beta, self.parallel_num)
+        else:
+            out_ = mlu_ops.matmul(act_out, down_proj.weight, None, residual_, 'none', self.alpha, beta)
+        '''
+        =======================================
+        End of custom MLU Hijack
+        =======================================
+        '''
+        # bias if existed need to add after second matmul according to the original design of vllm
+        '''
+        =============================
+        Modify by custom vllm_mlu
+        =============================
+        @brief: when preload_size is set, call GroupCoordinator.all_reduce() directly and 
+        use async_op to set all_reduce paralleled with preload 
+        '''
+        if self.reduce_results and self.tp_size > 1 and not is_parallel_enable:
+            if hasattr(self, 'preload_size') and self.preload_size > 0 and not self.is_prompt:
+                handle = get_tp_group().all_reduce(out_, async_op=True)
+                _MB = 1 << 20
+                mlu_ops.preload(self.preloaded_weights[0].data, self.preload_size * _MB)
+                preloaded_weights_size = self.preloaded_weights[0].numel() * self.preloaded_weights[0].element_size()
+                if preloaded_weights_size < (self.preload_size * _MB) and len(self.preloaded_weights) > 1:
+                    mlu_ops.preload(self.preloaded_weights[1].data, (self.preload_size * _MB) - preloaded_weights_size)
+                handle.wait()
+                out = out_
+            else:
+                out = tensor_model_parallel_all_reduce(out_)
+        else:
+            out = out_
+        '''
+        =========================
+        End of custom MLU Hijack
+        =========================
+        '''
+        # do the bias add if needed
+        if not self.skip_bias_add:
+            out = out + down_proj.bias if down_proj.bias is not None else out
+        else:
+            return out, down_proj.bias
+    else:
+        fc1, bias = up_proj(hidden_states)
+        if bias is not None:
+            fc1 += bias
+        input_scale= None
+        if (fc1.shape[-1] < 24000 and
+            self.quant_config is not None and self.quant_config.get_name() == "SmoothQuant" and
+            self.quant_config.input_quant_method == "per_token"):
+            down_proj.quant_method.skip_quant_input = True
+            fc1, input_scale = mlu_ops.per_token_smooth_quantize(fc1, down_proj.smooth, None, None, act_mode=self.hidden_act, is_gated=self.is_gated)
+        else:
+            fc1 = mlu_ops.active(fc1, self.hidden_act, self.is_gated)
+        out, bias = down_proj(fc1, residual=residual_, smooth_quant_scale=input_scale)
+
+        if self.skip_bias_add:
+            return out, bias
+    return out
+
+
+MluHijackObject.apply_hijack(FeedForward,
+                             FeedForward.forward,
+                             vllm_mlu__model_executor__layers__feed_forward__FeedForward__forward)
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/layers/linear.py b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/layers/linear.py
new file mode 100644
index 000000000..93ce2116f
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/layers/linear.py
@@ -0,0 +1,123 @@
+from typing import Optional
+import torch
+from vllm.distributed.parallel_state import get_tp_group, get_tensor_model_parallel_group
+from vllm.distributed import get_tensor_model_parallel_rank, split_tensor_along_last_dim
+from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod, RowParallelLinear
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu._mlu_utils import *
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__model_executor__layers__linear__UnquantizedLinearMethod__apply(
+    self,
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    residual: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    beta = 1.0 if residual is not None else 0.0
+    res_shape = x.shape[0:-1] + (layer.weight.shape[0], )
+    '''
+    =====================================================
+    Modify by custom vllm_mlu
+    =====================================================
+    @brief: call parallel op if parallel_num is set
+    '''
+    if hasattr(self, 'parallel_num') and get_is_prompt():
+        rank = get_tensor_model_parallel_rank()
+        pg = get_tensor_model_parallel_group().device_group
+        cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
+        return mlu_ops.matmul_allreduce(cncl_comm, x.view(-1, x.shape[-1]), layer.weight, 
+                                        bias, residual, 1.0, beta, self.parallel_num).view(res_shape)
+    return mlu_ops.matmul(x.reshape(x.numel() // x.shape[-1], x.shape[-1]),
+                          layer.weight, bias, residual, 'none', 1.0, beta).view(res_shape)
+    '''
+    =====================================================
+    End of custom MLU Hijack
+    =====================================================
+    '''
+
+
+def vllm__model_executor__layers__linear__RowParallelLinear__forward(
+    self, 
+    input_, 
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None
+):
+    if self.input_is_parallel:
+        input_parallel = input_
+    else:
+        tp_rank = get_tensor_model_parallel_rank()
+        splitted_input = split_tensor_along_last_dim(
+            input_, num_partitions=self.tp_size)
+        input_parallel = splitted_input[tp_rank].contiguous()
+
+    # Matrix multiply.
+    assert self.quant_method is not None
+    # Only fuse bias add into GEMM for rank 0 (this ensures that
+    # bias will not get added more than once in TP>1 case)
+    bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
+    residual_ = None if self.tp_rank > 0 else residual
+    '''
+    =====================================================
+    Modify by custom vllm_mlu
+    =====================================================
+    @brief: abandon original reduce if parallel_num is set
+    '''
+    is_parallel_enable = hasattr(self.quant_method, 'parallel_num') and get_is_prompt()
+    '''
+    =====================================================
+    End of custom MLU Hijack
+    =====================================================
+    '''
+    kwargs = {'bias': bias_, 'residual': residual_}
+    if smooth_quant_scale is not None:
+        kwargs['input_scale'] = smooth_quant_scale
+    output_parallel = self.quant_method.apply(self,
+                                              input_parallel,
+                                              **kwargs)
+    '''
+    =============================
+    Modify by custom vllm_mlu
+    =============================
+    @brief: when preload_size is set, call GroupCoordinator.all_reduce() directly and
+    use async_op to set all_reduce paralleled with preload
+    '''
+    if self.reduce_results and self.tp_size > 1 and not is_parallel_enable:
+        if hasattr(self, 'preload_size') and self.preload_size > 0 and not self.is_prompt:
+            handle = get_tp_group().all_reduce(output_parallel, async_op=True)
+            _MB = 1 << 20
+            mlu_ops.preload(self.preloaded_weights[0].data, self.preload_size * _MB)
+            preloaded_weights_size = self.preloaded_weights[0].numel() * self.preloaded_weights[0].element_size()
+            if preloaded_weights_size < (self.preload_size * _MB) and len(self.preloaded_weights) > 1:
+                mlu_ops.preload(self.preloaded_weights[1].data, (self.preload_size * _MB) - preloaded_weights_size)
+            handle.wait()
+            output = output_parallel
+        else:
+            output = tensor_model_parallel_all_reduce(output_parallel)
+    else:
+        output = output_parallel
+    '''
+    =========================
+    End of custom MLU Hijack
+    =========================
+    '''
+    output_bias = self.bias if self.skip_bias_add else None
+
+    if not self.return_bias:
+        return output
+    return output, output_bias
+
+
+MluHijackObject.undo_hijack(UnquantizedLinearMethod, 
+                            UnquantizedLinearMethod.apply)
+MluHijackObject.apply_hijack(UnquantizedLinearMethod,
+                             UnquantizedLinearMethod.apply,
+                             vllm__model_executor__layers__linear__UnquantizedLinearMethod__apply)
+MluHijackObject.undo_hijack(RowParallelLinear, 
+                            RowParallelLinear.forward)
+MluHijackObject.apply_hijack(RowParallelLinear,
+                             RowParallelLinear.forward,
+                             vllm__model_executor__layers__linear__RowParallelLinear__forward)
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/model_loader/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/model_loader/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/model_loader/loader.py b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/model_loader/loader.py
new file mode 100644
index 000000000..a3b56b476
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/common/model_executor/model_loader/loader.py
@@ -0,0 +1,145 @@
+import os
+import torch
+from torch import nn
+from typing import Optional
+
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader.loader import DefaultModelLoader
+from vllm.config import VllmConfig, ModelConfig, ParallelConfig
+
+from vllm_mlu._mlu_utils import *
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+
+
+logger = init_logger(__name__)
+
+
+def get_parallel_num(
+    model_config: ModelConfig,
+    parallel_config: ParallelConfig
+):
+    attention_parallel_num = os.environ.get(ATTN_PARALLEL_NUM)
+    ffn_parallel_num = os.environ.get(FFN_PARALLEL_NUM)
+    if attention_parallel_num and attention_parallel_num.isdecimal():
+        attention_parallel_num = int(attention_parallel_num)
+    else:
+        attention_parallel_num = 0
+    if ffn_parallel_num and ffn_parallel_num.isdecimal():
+        ffn_parallel_num = int(ffn_parallel_num)
+    else:
+        ffn_parallel_num = 0
+
+    if parallel_config.tensor_parallel_size == 1:
+        raise ValueError("Can not use context_comm_cmpt_parallel when tp num is 1.")
+    if (attention_parallel_num <= 0 and ffn_parallel_num <= 0):
+        raise ValueError("attention_parallel_num and ffn_parallel_num must be positive integers.")
+
+    hidden_size = model_config.get_hidden_size()
+    ffn_parallel_num = max(ffn_parallel_num, 1)
+    if hidden_size % ffn_parallel_num != 0:
+        raise ValueError(f"Hidden_size: {hidden_size} must be divisible by ffn_parallel_num: {ffn_parallel_num}")
+
+    return attention_parallel_num, ffn_parallel_num
+
+
+def get_attr_by_path(obj, path):
+    # Split the path by dots to get individual attributes
+    attributes = path.split('.')
+    # Iterate through the attributes to access nested members
+    for attr in attributes:
+        if not hasattr(obj, attr):
+            return None
+        obj = getattr(obj, attr)
+    return obj
+
+
+def set_custom_attributes(model, model_config, parallel_config):
+    attn_row_parallel_layers = []
+    attn_weights = []
+    ffn_row_parallel_layers = []
+    ffn_weights = []
+    sparse_moe_mlp_layers = []
+    for module in model.modules():
+        if module.__class__.__name__ == "FeedForward":
+            ffn_weight = []
+            if hasattr(module, "up_proj_name"):
+                up_proj_name = getattr(module, "up_proj_name")
+                up_proj = getattr(module, up_proj_name)
+                if hasattr(up_proj, "weight"):
+                    ffn_weight.append(up_proj.weight)
+            if hasattr(module, "down_proj_name"):
+                down_proj_name = getattr(module, "down_proj_name")
+                down_proj = getattr(module, down_proj_name)
+                if hasattr(down_proj, "weight"):
+                    ffn_weight.append(down_proj.weight)
+            if ffn_weight is not None:
+                ffn_weights.append(ffn_weight)
+            ffn_row_parallel_layers.append(module)
+        for child_module in module.children():
+            if child_module.__class__.__name__ == "Attention":
+                for sibling_module in module.children():
+                    if sibling_module.__class__.__name__ == "QKVParallelLinear":
+                        if hasattr(sibling_module, "weight"):
+                            weight = getattr(sibling_module, "weight")
+                            attn_weights.append([weight])
+                    if sibling_module.__class__.__name__ == "RowParallelLinear":
+                        attn_row_parallel_layers.append(sibling_module)
+        if module.__class__.__name__ == "SparseMoeMlp" or issubclass(module.__class__, SparseMoeMlp):
+            sparse_moe_mlp_layers.append(module)
+
+    if VLLM_PRELOAD_SIZE > 0:
+        if (len(attn_row_parallel_layers) \
+            == len(attn_weights) \
+            == len(ffn_row_parallel_layers) \
+            == len(ffn_weights)) and \
+            len(attn_row_parallel_layers) != 0:
+
+            for i in range(len(attn_row_parallel_layers)):
+                attn_row_parallel_layers[i].preloaded_weights = ffn_weights[i]
+                attn_row_parallel_layers[i].preload_size = VLLM_PRELOAD_SIZE
+                if i < len(attn_row_parallel_layers) - 1:
+                    ffn_row_parallel_layers[i].preloaded_weights = attn_weights[i+1]
+                    ffn_row_parallel_layers[i].preload_size = VLLM_PRELOAD_SIZE
+        else:
+            logger.warning("%s does not support preload weight!", model.__class__.__name__)
+
+    # context compute communication parallel
+    if check_context_comm_cmpt_parallel():
+        attention_parallel_num, ffn_parallel_num = get_parallel_num(model_config, parallel_config)
+        for o_proj in attn_row_parallel_layers:
+            setattr(o_proj.quant_method, 'parallel_num', attention_parallel_num)
+        
+        if len(sparse_moe_mlp_layers) != 0:
+            for sparse_moe_mlp in sparse_moe_mlp_layers:
+                setattr(sparse_moe_mlp, 'parallel_num', ffn_parallel_num)
+        else:
+            for ffn in ffn_row_parallel_layers:
+                setattr(ffn, 'parallel_num', ffn_parallel_num)
+
+
+vllm__model_executor__model_loader__loader__DefaultModelLoader__load_model__org = DefaultModelLoader.load_model
+
+
+def vllm__model_executor__model_loader__loader__DefaultModelLoader__load_model(
+        self, vllm_config: VllmConfig) -> nn.Module:
+    model = vllm__model_executor__model_loader__loader__DefaultModelLoader__load_model__org(
+                self, vllm_config=vllm_config)
+    '''
+    =============================
+    Modify by custom vllm_mlu
+    =============================
+    @brief: According to the layer name in models, set custom optimize attributes.
+    '''
+    set_custom_attributes(model, vllm_config.model_config, vllm_config.parallel_config)
+    '''
+    =========================
+    End of custom MLU Hijack
+    =========================
+    '''
+    return model
+
+
+MluHijackObject.apply_hijack(DefaultModelLoader,
+                            DefaultModelLoader.load_model,
+                            vllm__model_executor__model_loader__loader__DefaultModelLoader__load_model)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/README.md b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/README.md
new file mode 100644
index 000000000..31b78807d
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/README.md
@@ -0,0 +1,17 @@
+### 简介
+
+该劫持代码实现了vllm Context通算并行功能。开启后可在部分数据规模和切分数量上对Context Latency指标有优化效果。目前是可选功能，默认不开启。
+
+### 开启方法
+
+- 设置环境变量ATTN_PARALLEL_NUM和FFN_PARALLEL_NUM为正整数，分别控制attention和ffn部分的通算并行切分数量。两个环境变量相互独立，可以同时开启。例如输入export ATTN_PARALLEL_NUM=2 FFN_PARALLEL_NUM=4，则表示两部分均开启并行，attention数据拆分为2份，ffn数据拆分为4份。
+
+- 需要保证tensor_parallel_size大于1。
+
+- 开启ffn部分的通算并行时，需要保证hidden_size能被FFN_PARALLEL_NUM整除。
+
+### 注意事项
+
+- 开启通算并行功能时，由于算子限制，Mixtral系列模型、Qwen2（包含Qwen1.5和Qwen2.5）系列模型在smoothquant量化下只支持batch_size = 1，且算子默认切分数为4，ATTN_PARALLEL_NUM不生效。
+
+- smoothquant量化下，vllm_mlu ffn部分不调用tmo matmul算子，该部分通算融合不生效。
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/attention/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/attention/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/attention/context_comm_cmpt_parallel_fa.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/attention/context_comm_cmpt_parallel_fa.py
new file mode 100644
index 000000000..7560000d9
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/attention/context_comm_cmpt_parallel_fa.py
@@ -0,0 +1,231 @@
+import torch
+from typing import List, Optional
+
+from vllm.attention.backends.abstract import AttentionType
+from vllm.attention.backends.utils import get_num_prefill_decode_query_kv_tokens
+
+from vllm_mlu._mlu_utils import *
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu.attention.backends.mlu_flash_attn import (MLUFlashAttentionMetadata,
+                                                        _get_query_key_seq_metadata,
+                                                        _get_causal_option)
+
+
+def context_attn_comm_cmpt_parallel_flash_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    num_heads: int,
+    head_size: int,
+    num_kv_heads: int,
+    attn_metadata: MLUFlashAttentionMetadata,
+    kv_cache: List[torch.Tensor],
+    kv_cache_dtype: str,
+    softmax_scale: float,
+    cncl_comm: int,
+    smooth: torch.Tensor,
+    qweight: torch.Tensor,
+    per_channel_scale: torch.Tensor,
+    parallel_num: int,
+    residual: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    attn_type: Optional[AttentionType] = AttentionType.DECODER,
+    use_mla: bool = False,
+) -> torch.Tensor:
+    """Forward pass with FlashAttention.
+
+    Args:
+        query: shape = [num_tokens, num_heads, head_size]
+        key: shape = [num_tokens, num_kv_heads, head_size]
+        value: shape = [num_tokens, num_kv_heads, head_size]
+        output: shape = [num_tokens, num_heads, head_size]
+        kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
+            NOTE: kv_cache will be an empty tensor with shape [0]
+            for profiling run.
+        attn_metadata: Metadata for attention.
+    NOTE: It in-place updates the output tensor.
+    """
+
+    # Reshape the query, key, and value tensors.
+    query = query.view(-1, num_heads, head_size)
+    if (key is not None) and (key is not None):
+        key = key.view(-1, num_kv_heads, head_size)
+        value = value.view(-1, num_kv_heads, head_size)
+
+    if (attn_type == AttentionType.ENCODER
+            and (not attn_metadata.is_all_encoder_attn_metadata_set)):
+        raise AttributeError("Encoder attention requires setting "
+                                "encoder metadata attributes.")
+    elif (attn_type == AttentionType.ENCODER_DECODER
+            and (not attn_metadata.is_all_cross_attn_metadata_set)):
+        raise AttributeError("Encoder/decoder cross-attention "
+                                "requires setting cross-attention "
+                                "metadata attributes.")
+
+    if kv_cache[0].numel() > 0:
+        kv_cache_, kv_cache_scale_ = kv_cache
+        key_cache = kv_cache_[0]
+        value_cache = None if use_mla else kv_cache_[1]
+        key_cache_scale, value_cache_scale = None, None
+        if kv_cache_scale_.numel() > 0:
+            key_cache_scale = kv_cache_scale_[0]
+            value_cache_scale = None if use_mla else kv_cache_scale_[1]
+
+        # We skip updating the KV cache under two conditions:
+        #  a. When the Attention Type is ENCODER. In this phase, we compute
+        #     only the encoder attention without updating the cache.
+        #  b. When both Key and Value are None. This occurs during
+        #     cross-attention computation in the decoding phase, where the
+        #     KV cache is already populated with the cross-attention
+        #     tensor. Thus, we skip cache updates during this time.
+        if (attn_type != AttentionType.ENCODER) and (key is not None) and (
+                value is not None):
+            if attn_type == AttentionType.ENCODER_DECODER:
+                # Update cross-attention KV cache (prefill-only)
+                updated_slot_mapping = attn_metadata.cross_slot_mapping
+            else:
+                # Update self-attention KV cache (prefill/decode)
+                updated_slot_mapping = attn_metadata.slot_mapping
+
+            # Reshape the input keys and values and store them in the cache.
+            # If kv_cache is not provided, the new key and value tensors are
+            # not cached. This happens during the initial memory
+            # profiling run.
+            value_to_cache = None if use_mla else value
+            if use_mla and attn_metadata.prefill_metadata:
+                # MLA save cache info in models before flashattn
+                pass
+            else:
+                if kv_cache_dtype == 'int8':
+                    mlu_ops.quant_to_paged_cache(
+                        key,
+                        value_to_cache,
+                        key_cache,
+                        value_cache,
+                        key_cache_scale,
+                        value_cache_scale,
+                        updated_slot_mapping.flatten()
+                    )
+                else:
+                    mlu_ops.reshape_paged_cache(
+                        key,
+                        value_to_cache,
+                        key_cache,
+                        value_cache,
+                        updated_slot_mapping.flatten()
+                    )
+
+    (num_prefill_query_tokens, num_prefill_kv_tokens,
+    num_decode_query_tokens) = \
+        get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type)
+    decode_query = query[num_prefill_query_tokens:]
+    # QKV for prefill.
+    query = query[:num_prefill_query_tokens]
+    assert query.shape[0] == num_prefill_query_tokens
+    assert decode_query.shape[0] == num_decode_query_tokens
+
+    if prefill_meta := attn_metadata.prefill_metadata:
+        alibi_slopes = None if alibi_slopes is None else \
+                            alibi_slopes.repeat(attn_metadata.num_prefills, 1)
+        # Prompt run.
+        if (kv_cache[0].numel() == 0 or prefill_meta.block_tables is None
+                or prefill_meta.block_tables.numel() == 0):
+            # normal attention
+            # When block_tables are not filled, it means q and k are the
+            # prompt, and they have the same length.
+            q_seq_start_loc, q_seq_len, k_seq_start_loc, k_seq_len = \
+                _get_query_key_seq_metadata(prefill_meta, True, attn_type)
+
+            key = key[:num_prefill_kv_tokens]
+            value = value[:num_prefill_kv_tokens]
+
+            output = mlu_ops.flash_attn_sq_mm_allreduce(
+                cncl_comm, 
+                query,
+                key,
+                value, 
+                q_seq_start_loc,
+                k_seq_start_loc, 
+                alibi_slopes,
+                None, 
+                smooth,
+                qweight, 
+                per_channel_scale,
+                None,
+                q_seq_len,
+                k_seq_len,
+                softmax_scale,
+                _get_causal_option(attn_type), 
+                -1 if window_size is None else window_size[0],
+                -1 if window_size is None else window_size[1],
+                torch.float,
+                parallel_num
+            )
+        else:
+            # prefix-enabled attention
+            assert attn_type == AttentionType.DECODER, (
+                "Only decoder-only models support prefix caching")
+            assert prefill_meta.seq_lens is not None
+            max_seq_len = max(prefill_meta.seq_lens)
+
+            if kv_cache_dtype == 'int8' and \
+                    prefill_meta.chunked_prefill_enabled:
+                _, head_num_kv, _, head_size_qk = key_cache.shape
+                total_seqlens = prefill_meta.seq_start_loc[-1].item()
+                key_cache_float = torch.zeros((total_seqlens, head_num_kv, head_size_qk),
+                                                dtype=query.dtype,
+                                                device=key_cache.device)
+                value_cache_float = None
+                if value_cache is not None:
+                    _, head_num_kv, _, head_size_v = value_cache.shape
+                    value_cache_float = torch.zeros((total_seqlens, head_num_kv, head_size_v),
+                                                    dtype=query.dtype,
+                                                    device=value_cache.device)
+                mlu_ops.dequant_from_paged_cache(
+                    key=key_cache_float,
+                    value=value_cache_float,
+                    key_cache=key_cache,
+                    value_cache=value_cache,
+                    key_cache_quant_scale=key_cache_scale,
+                    value_cache_quant_scale=value_cache_scale,
+                    context_lengths=prefill_meta.seq_lens_tensor,
+                    max_context_len=prefill_meta.max_prefill_seq_len,
+                    context_seq_offset=None,
+                    block_tables=prefill_meta.block_tables,
+                    quant_mode=1,
+                    quant_bit=8
+                )
+                block_tables = None
+            else:
+                key_cache_float = key_cache
+                value_cache_float = value_cache
+                block_tables = prefill_meta.block_tables
+
+            output = mlu_ops.flash_attn_sq_mm_allreduce(
+                cncl_comm, 
+                query,
+                key_cache_float,
+                value_cache_float, 
+                prefill_meta.query_start_loc,
+                prefill_meta.seq_start_loc, 
+                alibi_slopes,
+                None, 
+                smooth,
+                qweight, 
+                per_channel_scale,
+                None,
+                prefill_meta.max_query_len,
+                max_seq_len,
+                softmax_scale,
+                True, 
+                -1 if window_size is None else window_size[0],
+                -1 if window_size is None else window_size[1],
+                torch.float,
+                parallel_num
+            )
+
+    # Add residual.
+    if residual is not None:
+        output = output + residual
+    return output
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/custom_model/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/custom_model/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/custom_model/custom.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/custom_model/custom.py
new file mode 100644
index 000000000..1ae6d7a99
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/custom_model/custom.py
@@ -0,0 +1,65 @@
+import torch
+import torch.nn.functional as F
+from typing import Optional
+
+from vllm.distributed import tensor_model_parallel_all_reduce, get_tensor_model_parallel_rank
+from vllm.distributed.parallel_state import get_tensor_model_parallel_group
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu._mlu_utils import *
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.model_executor.custom_model.custom import CustomMoeBlock
+
+
+def vllm__module_executor__custom_model__CustomMoeBlock__forward(
+    self, 
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    num_tokens, hidden_dim = hidden_states.shape
+    hidden_states = hidden_states.view(-1, hidden_dim)
+    shared_output = None
+    if self.shared_expert is not None:
+        shared_output = self.shared_expert(hidden_states)
+        if self.shared_expert_gate is not None:
+            shared_output = F.sigmoid(
+                self.shared_expert_gate(hidden_states)) * shared_output
+
+    # router_logits: (num_tokens, n_experts)
+    router_logits, _ = self.gate(hidden_states)
+    residual_ = None if self.rank > 0 else residual
+    '''
+    =====================================================
+    Modify by Context Communication Computation Parallel
+    =====================================================
+    @brief: call fused_moe
+    '''
+    params = [hidden_states, router_logits, self.w1, self.w2, None, None, 
+        residual_, self.input_smooth, self.act_smooth, self.w1_scale, self.w2_scale, 
+        self.top_k, self.config.norm_topk_prob, self.config.is_gated, self.config.hidden_act, 0]
+    if hasattr(self, 'parallel_num') and get_is_prompt():
+        rank = get_tensor_model_parallel_rank()
+        pg = get_tensor_model_parallel_group().device_group
+        cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
+        params.extend([self.parallel_num, cncl_comm])
+    final_hidden_states = mlu_ops.fused_moe(*params)
+    '''
+    =====================================================
+    End of Context Communication Computation Parallel
+    =====================================================
+    '''
+
+    if shared_output is not None:
+        final_hidden_states = final_hidden_states + shared_output
+        
+    reduce_results = (self.config.use_parallel_residual == False)
+    if reduce_results:
+        final_hidden_states = tensor_model_parallel_all_reduce(
+            final_hidden_states)
+
+    return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+MluHijackObject.apply_hijack(CustomMoeBlock,
+                             CustomMoeBlock.forward,
+                             vllm__module_executor__custom_model__CustomMoeBlock__forward)
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/quantization/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/quantization/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/quantization/smoothquant.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/quantization/smoothquant.py
new file mode 100644
index 000000000..420b1ad50
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/quantization/smoothquant.py
@@ -0,0 +1,53 @@
+import torch
+from typing import Optional
+
+from vllm.distributed import get_tensor_model_parallel_rank
+from vllm.distributed.parallel_state import get_tensor_model_parallel_group
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu._mlu_utils import get_is_prompt
+from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantLinearMethod
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm_mlu__model_executor__layers__quantization__smoothquant__SmoothQuantLinearMethod__apply(
+    self,
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    residual: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    quant_input = None
+    input_scale = None
+    if self.quant_config.input_quant_method == "per_token":
+        quant_input, input_scale = mlu_ops.per_token_smooth_quantize(x, layer.smooth, None)
+    if self.quant_config.input_quant_method == "per_tensor":
+        quant_input = x if self.skip_quant_input else mlu_ops.quantize(x, layer.scale_to_int, None)
+
+    '''
+    =====================================================
+    Modify by Context Communication Computation Parallel
+    =====================================================
+    @brief: call parallel op
+    '''
+    if hasattr(self, 'parallel_num') and get_is_prompt():
+        rank = get_tensor_model_parallel_rank()
+        pg = get_tensor_model_parallel_group().device_group
+        cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
+        params = [cncl_comm, quant_input, input_scale, layer.qweight, layer.per_channel_scale, 
+                self.compute_dtype, bias, residual, 1.0, 1.0, self.parallel_num]
+        out = mlu_ops.smooth_quant_matmul_allreduce(*params)
+    else:
+        out = mlu_ops.smooth_quant_matmul(quant_input, input_scale, layer.qweight,
+                                            layer.per_channel_scale, self.compute_dtype, bias, residual)
+    '''
+    =====================================================
+    End of Context Communication Computation Parallel
+    =====================================================
+    '''
+    return out
+
+
+MluHijackObject.apply_hijack(SmoothQuantLinearMethod,
+                             SmoothQuantLinearMethod.apply,
+                             vllm_mlu__model_executor__layers__quantization__smoothquant__SmoothQuantLinearMethod__apply)
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/sparse_moe_mlp.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/sparse_moe_mlp.py
new file mode 100644
index 000000000..fc3b68a50
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/layers/sparse_moe_mlp.py
@@ -0,0 +1,91 @@
+"""Inference-only MOE model."""
+import torch
+from torch import nn
+from typing import Optional
+from vllm.distributed import tensor_model_parallel_all_reduce, get_tensor_model_parallel_rank
+from vllm.distributed.parallel_state import get_tensor_model_parallel_group
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu._mlu_utils import *
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+
+
+def vllm_mlu__model_executor__layers__sparse_moe_mlp__SparseMoeMlp__forward(
+    self, 
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    orig_hidden_states_shape = hidden_states.shape
+    hidden_states = hidden_states.view(-1, self.hidden_size)
+    # expert_logits: [num_tokens, self.num_experts_per_rank]
+    expert_logits, _ = self.gate(hidden_states)
+    final_hidden_states = self.forward_experts(hidden_states, expert_logits, residual)
+
+    '''
+    =====================================================
+    Modify by Context Communication Computation Parallel
+    =====================================================
+    @brief: disbale reduce if parallel op used
+    '''
+    is_parallel_enable = hasattr(self, 'parallel_num') and get_is_prompt()
+    if self.tp_size > 1 and not is_parallel_enable:
+        final_hidden_states = tensor_model_parallel_all_reduce(
+            final_hidden_states)
+    '''
+    =====================================================
+    End of Context Communication Computation Parallel
+    =====================================================
+    '''
+
+    output = final_hidden_states.view(orig_hidden_states_shape)
+    return output
+
+
+def vllm_mlu__model_executor__layers__sparse_moe_mlp__SparseMoeMlp__forward_experts(
+    self, 
+    hidden_states, 
+    expert_logits,
+    residual: Optional[torch.Tensor] = None
+):
+    residual_ = None if self.tp_rank > 0 else residual
+    if self.is_use_fused_moe and self.expert_group != 1:
+        final_hidden_states = self.forward_group_experts(hidden_states, expert_logits, residual_)
+    elif self.is_use_fused_moe:
+        self.pack_params()
+        '''
+        =====================================================
+        Modify by Context Communication Computation Parallel
+        =====================================================
+        @brief: call fused_moe all_reduce
+        '''
+        is_parallel_enable = hasattr(self, 'parallel_num') and get_is_prompt()
+        if is_parallel_enable:
+            residual_ = residual
+        params = [hidden_states, expert_logits, self.w13, self.w2, self.b13, self.b2, 
+            residual_, self.a13_scale, self.a2_scale, self.w13_scale, self.w2_scale, 
+            self.top_k, self.renormalize, self.is_gated, self.hidden_act, self.start_expert_id]
+        if is_parallel_enable:
+            rank = get_tensor_model_parallel_rank()
+            pg = get_tensor_model_parallel_group().device_group
+            cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
+            params.extend([self.parallel_num, cncl_comm])
+        final_hidden_states = mlu_ops.fused_moe(*params)
+        '''
+        =====================================================
+        End of Context Communication Computation Parallel
+        =====================================================
+        '''
+    else:
+        final_hidden_states = self.forward_experts_nofused(hidden_states, expert_logits)
+        if residual_ is not None:
+            final_hidden_states = final_hidden_states + residual_
+
+    return final_hidden_states
+
+
+MluHijackObject.apply_hijack(SparseMoeMlp,
+                             SparseMoeMlp.forward,
+                             vllm_mlu__model_executor__layers__sparse_moe_mlp__SparseMoeMlp__forward)
+MluHijackObject.apply_hijack(SparseMoeMlp,
+                             SparseMoeMlp.forward_experts,
+                             vllm_mlu__model_executor__layers__sparse_moe_mlp__SparseMoeMlp__forward_experts)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/mixtral_quant.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/mixtral_quant.py
new file mode 100644
index 000000000..b66bc26be
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/mixtral_quant.py
@@ -0,0 +1,103 @@
+import torch
+from typing import List, Optional
+
+from vllm.distributed import get_tensor_model_parallel_rank
+from vllm.distributed.parallel_state import get_tensor_model_parallel_group
+from vllm.model_executor.models.mixtral_quant import MixtralAttention
+from vllm.forward_context import ForwardContext, get_forward_context
+
+from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantLinearMethod
+from vllm_mlu.mlu_custom.context_comm_cmpt_parallel.attention.context_comm_cmpt_parallel_fa import (
+    context_attn_comm_cmpt_parallel_flash_attention
+)
+from vllm_mlu._mlu_utils import *
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__model_executor__models__mixtral__MixtralAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.qkv_proj(hidden_states)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    ''' 
+    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    '''
+    =====================================================
+    Modify by Context Communication Computation Parallel
+    =====================================================
+    @brief: call flash_attn_sq_mm_allreduce to finish forward
+    '''
+    forward_context: ForwardContext = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
+    self_kv_cache = self.attn.kv_cache[forward_context.virtual_engine]
+    
+    if (attn_metadata.prefill_metadata
+            and not attn_metadata.prefill_metadata.is_profile_run
+            and hasattr(self.o_proj, 'quant_method')
+            and isinstance(self.o_proj.quant_method, SmoothQuantLinearMethod)
+            and self.o_proj.quant_method.quant_config.input_quant_method == "per_token"):
+        rank = get_tensor_model_parallel_rank()
+        pg = get_tensor_model_parallel_group().device_group
+        cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
+
+        return context_attn_comm_cmpt_parallel_flash_attention(
+                    q, k, v,
+                    self.num_heads,
+                    self.head_dim,
+                    self.num_kv_heads,
+                    attn_metadata,
+                    self_kv_cache,
+                    self.attn.impl.kv_cache_dtype,
+                    self.scaling,
+                    cncl_comm,
+                    self.o_proj.smooth,
+                    self.o_proj.qweight, 
+                    self.o_proj.per_channel_scale.to(torch.float),
+                    self.o_proj.quant_method.parallel_num,
+                    residual,
+                    self.attn.impl.sliding_window,
+                    self.attn.impl.alibi_slopes,
+                    self.attn.impl.attn_type,
+                    self.attn.impl.use_mla,
+                )
+    '''
+    =====================================================
+    End of Context Communication Computation Parallel
+    =====================================================
+    '''
+
+    attn_output = self.attn(q, k, v)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    ''' 
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+MluHijackObject.apply_hijack(MixtralAttention,
+                             MixtralAttention.forward,
+                             vllm__model_executor__models__mixtral__MixtralAttention__forward)
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/qwen2.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/qwen2.py
new file mode 100644
index 000000000..7ad628e80
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/qwen2.py
@@ -0,0 +1,106 @@
+import torch
+from typing import Optional
+from vllm.attention import AttentionMetadata
+
+from vllm.distributed import get_tensor_model_parallel_rank
+from vllm.distributed.parallel_state import get_tensor_model_parallel_group
+from vllm.model_executor.models.qwen2 import Qwen2Attention
+from vllm.forward_context import ForwardContext, get_forward_context
+
+from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantLinearMethod
+from vllm_mlu.mlu_custom.context_comm_cmpt_parallel.attention.context_comm_cmpt_parallel_fa import (
+    context_attn_comm_cmpt_parallel_flash_attention
+)
+from vllm_mlu._mlu_utils import *
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__model_executor__models__qwen2__Qwen2Attention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    assert smooth_quant_scale is None
+    qkv, _ = self.qkv_proj(hidden_states)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    '''
+    =====================================================
+    Modify by Context Communication Computation Parallel
+    =====================================================
+    @brief: call flash_attn_sq_mm_allreduce to finish forward
+    '''
+    forward_context: ForwardContext = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
+    self_kv_cache = self.attn.kv_cache[forward_context.virtual_engine]
+
+    if (attn_metadata.prefill_metadata
+            and not attn_metadata.prefill_metadata.is_profile_run
+            and hasattr(self.o_proj, 'quant_method')
+            and isinstance(self.o_proj.quant_method, SmoothQuantLinearMethod)
+            and self.o_proj.quant_method.quant_config.input_quant_method == "per_token"):
+        rank = get_tensor_model_parallel_rank()
+        pg = get_tensor_model_parallel_group().device_group
+        cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
+
+        return context_attn_comm_cmpt_parallel_flash_attention(
+                    q, k, v,
+                    self.num_heads,
+                    self.head_dim,
+                    self.num_kv_heads,
+                    attn_metadata,
+                    self_kv_cache,
+                    self.attn.impl.kv_cache_dtype,
+                    self.scaling,
+                    cncl_comm,
+                    self.o_proj.smooth,
+                    self.o_proj.qweight, 
+                    self.o_proj.per_channel_scale.to(torch.float),
+                    self.o_proj.quant_method.parallel_num,
+                    residual,
+                    self.attn.impl.sliding_window,
+                    self.attn.impl.alibi_slopes,
+                    self.attn.impl.attn_type,
+                    self.attn.impl.use_mla,
+                )
+    '''
+    =====================================================
+    End of Context Communication Computation Parallel
+    =====================================================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+MluHijackObject.undo_hijack(Qwen2Attention, 
+                            Qwen2Attention.forward)
+MluHijackObject.apply_hijack(Qwen2Attention,
+                             Qwen2Attention.forward,
+                             vllm__model_executor__models__qwen2__Qwen2Attention__forward)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/qwen2_moe.py b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/qwen2_moe.py
new file mode 100644
index 000000000..98de9ddf5
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/context_comm_cmpt_parallel/model_executor/models/qwen2_moe.py
@@ -0,0 +1,57 @@
+import torch
+import torch.nn.functional as F
+from typing import Optional
+
+from vllm.distributed import tensor_model_parallel_all_reduce
+
+from vllm_mlu.model_executor.models.qwen2_moe import Qwen2MoeSparseMoeBlock
+from vllm_mlu._mlu_utils import *
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm_mlu__model_executor__models__qwen2_moe__Qwen2MoeSparseMoeBlock__forward(
+    self, 
+    hidden_states: torch.Tensor, 
+    residual: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    num_tokens, hidden_dim = hidden_states.shape
+    hidden_states = hidden_states.view(-1, hidden_dim)
+    shared_output = None
+    if self.shared_expert is not None:
+        shared_output = self.shared_expert(hidden_states)
+        if self.shared_expert_gate is not None:
+            gate_output = self.shared_expert_gate(hidden_states)
+            shared_output = F.sigmoid(gate_output[0]) * shared_output
+
+    # router_logits: (num_tokens, n_experts)
+    router_logits, _ = self.gate(hidden_states)
+    final_hidden_states = self.forward_experts(hidden_states, router_logits, residual)
+
+    '''
+    =====================================================
+    Modify by Context Communication Computation Parallel
+    =====================================================
+    @brief: disbale reduce if parallel op used
+    '''
+    is_parallel_enable = hasattr(self, 'parallel_num') and get_is_prompt()
+    if self.tp_size > 1:
+        if is_parallel_enable:
+            shared_output = tensor_model_parallel_all_reduce(shared_output)
+            if shared_output is not None:
+                final_hidden_states = final_hidden_states + shared_output
+        else:
+            if shared_output is not None:
+                final_hidden_states = final_hidden_states + shared_output
+            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
+    '''
+    =====================================================
+    End of Context Communication Computation Parallel
+    =====================================================
+    '''
+
+    return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+MluHijackObject.apply_hijack(Qwen2MoeSparseMoeBlock,
+                             Qwen2MoeSparseMoeBlock.forward,
+                             vllm_mlu__model_executor__models__qwen2_moe__Qwen2MoeSparseMoeBlock__forward)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/preload/README.md b/vllm_mlu/vllm_mlu/mlu_custom/preload/README.md
new file mode 100644
index 000000000..24308ffab
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/preload/README.md
@@ -0,0 +1,32 @@
+### 简介
+
+该劫持代码实现在vLLM的解码通信过程中预加载下一层的权重，从而减少解码的延迟。
+
+### 支持模型
+
+仅支持以下模型，不支持量化后的模型以及MOE模型。
+- Baichuan
+- Bloom
+- ChatGLM
+- Falcon
+- GPTNeoX
+- Llama
+- Qwen
+- Qwen2
+
+### 支持板卡
+
+300系列不支持，其他系列支持。
+
+### 使用方法
+
+- 设置环境变量export VLLM_PRELOAD_SIZE=<PRELOAD_SIZE>，<PRELOAD_SIZE>表示预加载权重的大小，单位：MB。
+- 参数设置参考：在低带宽资源环境下，对于模型Llama-65B，不同batch_sized和preload_size对应的性能优化收益如下。
+
+| batch\preload  |  8   |  16  |  24  |  32  |  48  |  64  | 
+|:--------------:|:----:|:----:|:----:|:----:|:----:|:----:|
+|      1         | 4.9% | 10.0%| 9.5% | 6.7% |-2.4% | -7.1%|
+|      8         | 3.2% | 6.3% | 8.9% | 11.2%| 6.0% | 1.8% |
+|      16        | 2.3% | 5.1% | 7.5% | 9.2% | 8.3% | 4.3% |
+|      24        | 2.3% | 4.8% | 7.4% | 9.1% | 9.5% | 6.0% |
+|      32        | 2.1% | 4.3% | 7.0% | 8.7% | 10.1%| 8.1% |
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/preload/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/preload/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/preload/distributed/__init__.py b/vllm_mlu/vllm_mlu/mlu_custom/preload/distributed/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/mlu_custom/preload/distributed/parallel_state.py b/vllm_mlu/vllm_mlu/mlu_custom/preload/distributed/parallel_state.py
new file mode 100644
index 000000000..99c513d71
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_custom/preload/distributed/parallel_state.py
@@ -0,0 +1,109 @@
+import torch
+import torch.distributed as dist
+
+from vllm.distributed.parallel_state import GroupCoordinator
+from vllm.distributed.device_communicators.base_device_communicator import (
+    DeviceCommunicatorBase
+)
+
+from vllm_mlu.distributed.device_communicators.mlu_communicator import MLUCommunicator
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__distributed__parallel_state__GroupCoordinator__all_reduce(
+    self,
+    input_: torch.Tensor,
+    async_op: bool = False,
+) -> torch.Tensor:
+    """
+    User-facing all-reduce function before we actually call the
+    all-reduce operation.
+
+    We need this because Dynamo does not support passing an arbitrary
+    object (`self` in this case) to a custom op. We need to pass the
+        group name as a string, and then look up the group coordinator from
+        the group name, dispatch the all-reduce operation to the group
+        coordinator.
+
+    In addition, PyTorch custom ops do not support mutation or returning
+    a new tensor in the same op. So we always make the all-reduce operation
+    out-of-place.
+    """
+    # Bypass the function if we are using only 1 GPU.
+    if self.world_size == 1:
+        return input_
+
+    if self.use_custom_op_call:
+        return torch.ops.vllm.all_reduce(input_,
+                                            group_name=self.unique_name)
+    else:
+        '''
+        =============================
+        Modify by custom vllm_mlu
+        =============================
+        @brief: use async all reduce when preload weights.
+        '''
+        return self._all_reduce_out_place(input_, async_op=async_op)
+        '''
+        ==================
+        End of custom MLU Hijack
+        ==================
+        '''
+
+
+def vllm__distributed__parallel_state__GroupCoordinator___all_reduce_out_place(
+    self,
+    input_: torch.Tensor,
+    async_op: bool = False,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by custom vllm_mlu
+    =============================
+    @brief: use async all reduce when preload weights.
+    '''
+    return self.device_communicator.all_reduce(input_, async_op=async_op)
+    '''
+    ==================
+    End of custom MLU Hijack
+    ==================
+    '''
+
+
+def vllm__distributed__device_communicators__base_device_communicator__DeviceCommunicatorBase__all_reduce(
+    self,
+    input_: torch.Tensor,
+    async_op: bool = False,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by custom vllm_mlu
+    =============================
+    @brief: use async all reduce when preload weights.
+    '''
+    handle = dist.all_reduce(input_, group=self.device_group, async_op=async_op)
+    if async_op:
+        return handle
+    return input_
+    '''
+    ==================
+    End of custom MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(
+    GroupCoordinator,
+    GroupCoordinator.all_reduce,
+    vllm__distributed__parallel_state__GroupCoordinator__all_reduce
+)
+MluHijackObject.apply_hijack(
+    GroupCoordinator,
+    GroupCoordinator._all_reduce_out_place,
+    vllm__distributed__parallel_state__GroupCoordinator___all_reduce_out_place
+)
+MluHijackObject.apply_hijack(
+    DeviceCommunicatorBase,
+    DeviceCommunicatorBase.all_reduce,
+    vllm__distributed__device_communicators__base_device_communicator__DeviceCommunicatorBase__all_reduce
+)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/mlu_hijack.py b/vllm_mlu/vllm_mlu/mlu_hijack.py
new file mode 100644
index 000000000..ccd76aff5
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_hijack.py
@@ -0,0 +1,121 @@
+import logging
+from logging import Logger
+
+from transformers import AutoConfig
+
+from vllm.model_executor.models import ModelRegistry
+
+from vllm_mlu.transformers_utils.configs import CustomConfig
+from vllm_mlu.model_executor.layers.quantization import (
+    register_real_mlu_quantization_methods
+)
+from vllm_mlu._mlu_utils import *
+
+
+def mlu_init_logger(name: str) -> Logger:
+    """Initialize loggers for vllm_mlu module,
+    and keep the configuration consistent with the vllm module"""
+    mlu_logger = logging.getLogger(name)
+    vllm_logger = logging.Logger.manager.loggerDict.get('vllm', None)
+    if vllm_logger:
+        mlu_logger.setLevel(vllm_logger.level)
+        mlu_logger.propagate = vllm_logger.propagate
+        mlu_logger.handlers = vllm_logger.handlers
+    return mlu_logger
+
+
+from vllm import logger
+logger.init_logger = mlu_init_logger
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+logger.info(f"Run vLLM in paged mode, Apply MLU optimization !")
+
+
+import vllm_mlu.config
+import vllm_mlu.utils
+import vllm_mlu.attention.ops.prefix_prefill
+import vllm_mlu.attention.layer
+if VLLM_SCHEDULER_PROFILE:
+    import vllm_mlu.core.scheduler
+    import vllm_mlu.engine.async_llm_engine
+    import vllm_mlu.engine.multiprocessing.client
+    import vllm_mlu.engine.multiprocessing.engine
+    import vllm_mlu.entrypoints.openai.serving_engine
+import vllm_mlu.distributed.parallel_state
+import vllm_mlu.engine.arg_utils
+import vllm_mlu.engine.llm_engine
+import vllm_mlu.entrypoints.llm
+import vllm_mlu.executor.multiproc_worker_utils
+import vllm_mlu.executor.executor_base
+import vllm_mlu.executor.ray_distributed_executor
+import vllm_mlu.lora.fully_sharded_layers
+import vllm_mlu.lora.layers
+import vllm_mlu.model_executor.layers.linear
+import vllm_mlu.model_executor.layers.rotary_embedding
+
+register_real_mlu_quantization_methods()
+
+import vllm_mlu.model_executor.layers.quantization.utils.w8a8_utils
+import vllm_mlu.model_executor.layers.quantization.fp8
+import vllm_mlu.model_executor.layers.activation
+import vllm_mlu.model_executor.layers.layernorm
+import vllm_mlu.model_executor.layers.fused_moe.layer
+import vllm_mlu.model_executor.model_loader.tensorizer
+import vllm_mlu.model_executor.models.deepseek_v2
+import vllm_mlu.model_executor.models.baichuan
+import vllm_mlu.model_executor.models.bert
+import vllm_mlu.model_executor.models.bloom
+import vllm_mlu.model_executor.models.chatglm
+import vllm_mlu.model_executor.models.clip
+import vllm_mlu.model_executor.models.gpt_neox
+import vllm_mlu.model_executor.models.llama
+import vllm_mlu.model_executor.models.mixtral
+import vllm_mlu.model_executor.models.qwen
+import vllm_mlu.model_executor.models.qwen2
+import vllm_mlu.model_executor.models.qwen2_moe
+import vllm_mlu.model_executor.models.qwen3
+import vllm_mlu.model_executor.models.qwen3_moe
+import vllm_mlu.model_executor.models.qwen2_vl
+import vllm_mlu.model_executor.models.qwen2_5_vl
+import vllm_mlu.model_executor.models.falcon
+import vllm_mlu.model_executor.models.internlm2
+import vllm_mlu.model_executor.models.intern_vit
+import vllm_mlu.model_executor.models.hunyuan
+import vllm_mlu.model_executor.models.layer_utils
+import vllm_mlu.model_executor.models.mllama
+import vllm_mlu.model_executor.models.deepseek_mtp
+import vllm_mlu.spec_decode.spec_decode_worker
+import vllm_mlu.spec_decode.multi_step_worker
+import vllm_mlu.spec_decode.ngram_worker
+import vllm_mlu.worker.cache_engine
+
+if VLLM_PRELOAD_SIZE > 0:
+    logger.info("Apply feature -> Preload Weight !")
+    import vllm_mlu.mlu_custom.preload.distributed.parallel_state
+
+    import vllm_mlu.mlu_custom.common.model_executor.layers.feed_forward
+    import vllm_mlu.mlu_custom.common.model_executor.layers.linear
+    import vllm_mlu.mlu_custom.common.model_executor.model_loader.loader
+
+if check_context_comm_cmpt_parallel():
+    logger.info("Apply feature -> Context Communication Computation Parallel !")
+    import vllm_mlu.mlu_custom.context_comm_cmpt_parallel.model_executor.custom_model.custom
+    import vllm_mlu.mlu_custom.context_comm_cmpt_parallel.model_executor.layers.quantization.smoothquant
+    import vllm_mlu.mlu_custom.context_comm_cmpt_parallel.model_executor.layers.sparse_moe_mlp
+    import vllm_mlu.mlu_custom.context_comm_cmpt_parallel.model_executor.models.mixtral_quant
+    import vllm_mlu.mlu_custom.context_comm_cmpt_parallel.model_executor.models.qwen2_moe
+    import vllm_mlu.mlu_custom.context_comm_cmpt_parallel.model_executor.models.mixtral_quant
+    
+    import vllm_mlu.mlu_custom.common.model_executor.layers.feed_forward
+    import vllm_mlu.mlu_custom.common.model_executor.layers.linear
+    import vllm_mlu.mlu_custom.common.model_executor.model_loader.loader
+
+
+AutoConfig.register("custom", CustomConfig)
+ModelRegistry.register_model(
+    model_arch="CustomForCausalLM",
+    model_cls="vllm_mlu.model_executor.custom_model.custom:CustomForCausalLM"
+)
diff --git a/vllm_mlu/vllm_mlu/mlu_hijack_utils.py b/vllm_mlu/vllm_mlu/mlu_hijack_utils.py
new file mode 100644
index 000000000..01717c0bb
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_hijack_utils.py
@@ -0,0 +1,101 @@
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+IS_GATED=False
+
+class MluHijackObject:
+    hijack_objs = []
+
+    @classmethod
+    def apply_hijack(cls, obj, org_func, hijack_func,
+                     verify_orig_func_exists: bool = False):
+        """
+        Optional Args:
+            verify_orig_func_exists (bool): If True, verifies that hijack succeeds
+        """
+        cls.hijack_objs.append((obj, org_func, hijack_func))
+
+        if type(org_func) == str:
+            org_func_name = org_func
+        else:
+            if isinstance(org_func, property):
+                split_name = org_func.fget.__name__.split('__')
+            else:
+                split_name = org_func.__name__.split('__')
+            org_func_name = split_name[-1]
+            if org_func_name == "":
+                assert split_name[-2] != "", f"invalid {org_func.__name__} to apply hijack"
+                org_func_name = split_name[-2] + "__"
+                if len(split_name) >= 3 and split_name[-3] == "":
+                    org_func_name = "__" + org_func_name
+
+        if verify_orig_func_exists and not hasattr(obj, org_func_name):
+            raise AttributeError(f"function {org_func_name} is not part of {obj}")
+
+        setattr(obj, org_func_name, hijack_func)
+
+        if (verify_orig_func_exists and getattr(obj, org_func_name) is not hijack_func):
+            raise AttributeError(
+                f"function {org_func_name} of {obj} failed to be swapped to {hijack_func}")
+
+
+    @classmethod
+    def undo_hijack(cls, obj_ = None, hijack_func_ = None):
+        if obj_ and hijack_func_:
+            for obj, org_func, hijack_func in cls.hijack_objs:
+                if obj_ == obj and hijack_func == hijack_func_:
+                    if type(org_func) == str:
+                        if hasattr(obj, org_func):
+                            delattr(obj, org_func)
+                    else:
+                        org_func_name = org_func.__name__
+                        setattr(obj, org_func_name, org_func)
+            return
+        for obj, org_func, hijack_func in cls.hijack_objs:
+            if type(org_func) == str:
+                if hasattr(obj, org_func):
+                    delattr(obj, org_func)
+            else:
+                org_func_name = org_func.__name__
+                setattr(obj, org_func_name, org_func)
+
+
+TypedDict = {
+    "hidden_size": 0,
+    "vocab_size": 0,
+    "ffn_inner_size": 0,
+    "moe_inner_size": 0,
+    "layer_num": 0,
+    "moe_layer_num": 0,
+    "head_num": 0,
+    "head_size": 0,
+    "head_num_kv": 0,
+    "tp_num": 0,
+    "shared_expert_intermediate_size": 0,
+    "shared_experts": 0,
+    "qk_nope_head_dim": 0,
+    "qk_rope_head_dim": 0,
+    "q_lora_rank": 0.0,
+    "num_attention_heads": 0,
+    "kv_lora_rank": 0,
+    "v_head_dim": 0,
+    "use_gated_ffn": False,
+    "experts_num": 0,
+    "topk_num": 0,
+    "use_causal_mask": False,
+    "cla_coeffient": 0,
+    "kv_cache_dtype": "",
+    "smooth_quant_type": "",
+    "data_type": "",
+    "model_type": "",
+    "filter_data_type": "",
+}
+
+
+def set_is_gated(flag):
+    global IS_GATED
+    IS_GATED=flag
+
+def get_is_gated():
+    return IS_GATED
diff --git a/vllm_mlu/vllm_mlu/mlu_metric.py b/vllm_mlu/vllm_mlu/mlu_metric.py
new file mode 100644
index 000000000..466391a9d
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/mlu_metric.py
@@ -0,0 +1,383 @@
+import torch
+import time
+import statistics
+import pandas as pd
+import numpy as np
+import json
+import os
+from datetime import datetime
+from vllm.logger import init_logger
+from vllm_mlu._mlu_utils import VLLM_LATENCY_DEBUG_WITH_DEVICE_EN
+from vllm.model_executor.layers.quantization import get_quantization_config
+
+logger = init_logger(__name__)
+
+
+class LLMMetric:
+    def __init__(self)->None:
+        self.batch_size_list = []
+        self.context_latency_list = []
+        self.e2e_latency_list = []
+        self.per_token_latency_list = [ [] ]
+        self.per_token_latency_device_list = [ [] ]
+        self.peak_memory = 0
+        self.block_memory = 0
+        self.num_total_gpu_blocks = 0
+        self.num_total_cpu_blocks = 0
+        self.num_free_gpu_blocks_list = [ [] ]
+        self.num_free_cpu_blocks_list = [ [] ]
+        self.num_spec_tokens = 0
+        self.draft_acceptance_rate = 0.0
+
+    def reset_metric(self):
+        self.batch_size_list = []
+        self.context_latency_list = []
+        self.e2e_latency_list = []
+        self.per_token_latency_list = [ [] ]
+        self.per_token_latency_device_list = [ [] ]
+        self.num_free_gpu_blocks_list = [ [] ]
+        self.num_free_cpu_blocks_list = [ [] ]
+        self.num_spec_tokens = 0
+        self.draft_acceptance_rate = 0.0
+
+    def get_mlu_cost_time(self):
+        torch.mlu.synchronize()
+        return time.time()
+
+    def is_prefill_stage(self):
+        return len(self.per_token_latency_list[-1]) == 0
+
+    def update_memory_usage(self, peak_memory, block_memory, num_total_gpu_blocks, num_total_cpu_blocks):
+        self.peak_memory = peak_memory
+        self.block_memory = block_memory
+        self.num_total_gpu_blocks = num_total_gpu_blocks
+        self.num_total_cpu_blocks = num_total_cpu_blocks
+
+    def update_step_block_usage(self, num_free_gpu_blocks, num_free_cpu_blocks):
+        self.num_free_gpu_blocks_list[-1].append(num_free_gpu_blocks)
+        self.num_free_cpu_blocks_list[-1].append(num_free_cpu_blocks)
+
+    def update_step_latency(self, step_latency):
+        self.per_token_latency_list[-1].append(step_latency)
+
+    def update_step_latency_device(self, step_latency):
+        self.per_token_latency_device_list[-1].append(step_latency)
+
+    def update_spec_decode_metrics(self, spec_decode_metrics):
+        self.num_spec_tokens = spec_decode_metrics.num_spec_tokens
+        self.draft_acceptance_rate = spec_decode_metrics.draft_acceptance_rate
+
+    def add_metrics(self, batch_size, e2e_latency)->None:
+        self.batch_size_list.append(batch_size)
+        self.e2e_latency_list.append(e2e_latency)
+        self.per_token_latency_list.append([]) # new iter
+        self.per_token_latency_device_list.append([])
+        self.num_free_gpu_blocks_list.append([])
+        self.num_free_cpu_blocks_list.append([])
+
+    def get_weight_dtype_str(self, model_path, model_dtype, quantization) -> str:
+        # get weight dtype based on quantization config if exists
+        if quantization == 'fp8':
+            return quantization
+        if quantization is not None:
+            quant_method = get_quantization_config(quantization)
+            # combine the model path with the quantization config file name
+            quant_config_paths = quant_method.get_config_filenames()
+            # if there are multiple quantization config files, return the first one existed
+            for quant_config_path in quant_config_paths:
+                quant_config_path = os.path.join(model_path, quant_config_path)
+                # check if the quantization config file exists
+                if not os.path.exists(quant_config_path):
+                    continue
+                with open(quant_config_path, 'r') as f:
+                    quant_config = json.load(f)
+                    quant_config = quant_method.from_config(quant_config)
+                    # for smoothquant and weightonly, return the quantization name with the weight bits
+                    if quantization == "smoothquant" or quantization == ["weightonly"]:
+                        return "{}-int{}".format(quant_config.get_name(), quant_config.weight_bits)
+                    else:
+                        # for other quantization methods, return the quantization name
+                        return quant_config.get_name()
+            # if the quantization config file does not exist, just return the quanization name
+            return quant_config_path.get_name()
+        else:
+            # remove the prefix of model dtype from torch config
+            return str(model_dtype).split(".")[-1]
+
+    def to_csv(self, filename: str, show_per_iter=False) -> None:
+        if show_per_iter:
+            df = pd.DataFrame(self.metrics_data)
+            df = pd.DataFrame([df.iloc[-1]], columns=df.columns)
+            memory_df = pd.DataFrame(self.memory_metrics_data)
+            memory_df = pd.DataFrame([memory_df.iloc[-1]], columns=memory_df.columns)
+        else:
+            df = pd.DataFrame(self.metrics_data)
+            memory_df = pd.DataFrame(self.memory_metrics_data)
+        df_mean = df.mean().round(3)
+        memory_df_mean = memory_df.mean().round(3)
+        header = ["datetime", "model",
+                  "weight dtype", self.batch_size_name,
+                  ]
+        header = header + list(self.mm_kwargs.keys())
+        header = header + ["input len", "output len", "tp",
+            self.context_latency_name, self.per_token_latency_name]
+        data = [datetime.now().strftime("%Y-%m-%d %H:%M:%S"), self.model,
+                self.weight_dtype_str, int(self.metrics_data[self.batch_size_name][0])]
+        data = data + [self.mm_kwargs[k] for k in self.mm_kwargs.keys()]
+        data = data + [self.input_len, self.output_len, self.tp,
+                       df_mean[self.context_latency_name], df_mean[self.per_token_latency_name]]
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+            header += [self.context_latency_device_name, self.per_token_latency_device_name]
+            data += [df_mean[self.context_latency_device_name], df_mean[self.per_token_latency_device_name]]
+        header += [self.e2e_latency_name, self.e2e_throughput_name, self.decoder_throughput_name,
+                   self.k_name, self.acceptance_rate_name, self.decode_times_name,
+                   self.peak_memory_name, self.block_memory_name, self.max_kv_memory_name, self.mean_kv_memory_name,
+                   self.max_kv_usage_name, self.mean_kv_usage_name]
+        data += [
+            df_mean[self.e2e_latency_name], df_mean[self.e2e_throughput_name], df_mean[self.decoder_throughput_name],
+            self.num_spec_tokens, df_mean[self.acceptance_rate_name], df_mean[self.decode_times_name],
+            memory_df_mean[self.peak_memory_name], memory_df_mean[self.block_memory_name],
+            memory_df_mean[self.max_kv_memory_name], memory_df_mean[self.mean_kv_memory_name],
+            memory_df_mean[self.max_kv_usage_name], memory_df_mean[self.mean_kv_usage_name]
+        ]
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN and self.save_hfu_info:
+            header += [self.context_hfu_name, self.decoder_hfu_name, self.decoder_io_efficiency_name]
+            data += [
+                df_mean[self.context_hfu_name], df_mean[self.decoder_hfu_name],
+                df_mean[self.decoder_io_efficiency_name]
+            ]
+        data_dict = dict(zip(header, data))
+        df_csv = pd.DataFrame(data_dict, index=[0])
+        append = False
+        if os.path.isfile(filename):
+            try:
+                df_old = pd.read_csv(filename)
+                append = (df_old.columns.tolist() == header)
+            except Exception as e:
+                logger.info(f"Existing {filename} failed to be read and will be overwritten")
+        if append:
+            df_csv.to_csv(filename, mode='a', header=False, index=False)
+            logger.info(f"Metric appended to existing {filename}")
+        else:
+            df_csv.to_csv(filename, index=False)
+            logger.info(f"Metric written to {filename}")
+
+    def calc_metric(self, model, model_dtype, metrics_idx_start, only_average,
+                    input_len, output_len, tp_nums, quantization, dump_info=None,
+                    show_per_iter=False, is_embedding_task=False, mm_kwargs={},
+                    total_prefill_steps=1
+    ) -> None:
+        keep_digits = 2
+
+        def round_fn(data):
+            return round(data, keep_digits)
+
+        metrics_idx_end = len(self.per_token_latency_list) - 1 # without last []
+        idx_range = range(metrics_idx_start, metrics_idx_end)
+
+        # specify entries to write to csv
+        self.mm_kwargs = mm_kwargs # multimodal args
+        self.batch_size_name = "batch size"
+        self.input_len = input_len
+        self.output_len = output_len
+        self.tp = tp_nums
+        self.model = model
+        self.context_latency_name = "context latency(ms)"
+        self.per_token_latency_name = "per token latency(ms)"
+        self.context_latency_device_name = "context latency device(ms)"
+        self.per_token_latency_device_name = "per token latency device(ms)"
+        self.e2e_latency_name = "e2e latency(ms)"
+        self.e2e_throughput_name = "e2e throughput(tokens/s)"
+        self.decoder_throughput_name = "decoder throughput(tokens/s)"
+        self.k_name = "K"
+        self.acceptance_rate_name = "acceptance rate"
+        self.decode_times_name = "decode times"
+        self.weight_dtype_str = self.get_weight_dtype_str(model, model_dtype, quantization)
+
+        metrics_data = [
+            (
+                self.batch_size_name, [int(self.batch_size_list[i]) for i in idx_range]
+            ),
+            (
+                self.context_latency_name, [round_fn(1000 * sum(self.per_token_latency_list[i][:total_prefill_steps])) for i in idx_range]
+            ),
+            (
+                self.per_token_latency_name, [
+                    0.0 if len(self.per_token_latency_list[i]) <= total_prefill_steps else \
+                    round_fn(statistics.fmean(self.per_token_latency_list[i][total_prefill_steps:]) * 1000) for i in idx_range
+                ]
+            ),
+            (
+                self.e2e_latency_name, [round_fn(1000 * self.e2e_latency_list[i]) for i in idx_range]
+            ),
+            (
+                self.e2e_throughput_name, [
+                    round_fn(((output_len + total_prefill_steps) / self.e2e_latency_list[i]) * self.batch_size_list[i]) \
+                        for i in idx_range
+                ]
+            ),
+            (
+                self.decoder_throughput_name, [
+                    0.0 if len(self.per_token_latency_list[i]) <= total_prefill_steps else \
+                    round_fn((output_len / sum(self.per_token_latency_list[i][total_prefill_steps:])) * self.batch_size_list[i]) \
+                        for i in idx_range
+                ]
+            ),
+            (
+                self.k_name, self.num_spec_tokens
+            ),
+            (
+                self.acceptance_rate_name, float(self.draft_acceptance_rate)
+            ),
+            (
+                self.decode_times_name, [
+                    0 if len(self.per_token_latency_list[i]) <= total_prefill_steps else \
+                    len(self.per_token_latency_list[i][total_prefill_steps:]) for i in idx_range
+                ]
+            ),
+        ]
+
+        insert_latency_device = VLLM_LATENCY_DEBUG_WITH_DEVICE_EN
+        if insert_latency_device:
+            context_latency_device = [round_fn(sum(self.per_token_latency_device_list[i][:total_prefill_steps])) for i in idx_range]
+            per_token_latency_device = [0.0 if len(self.per_token_latency_device_list[i]) <= total_prefill_steps else \
+                                        round_fn(statistics.fmean(self.per_token_latency_device_list[i][total_prefill_steps:])) for i in idx_range]
+            metrics_data.insert(3, (self.context_latency_device_name, context_latency_device))
+            metrics_data.insert(4, (self.per_token_latency_device_name, per_token_latency_device))
+
+        self.metrics_data = dict(metrics_data)
+
+        # Print
+        df = pd.DataFrame(self.metrics_data)
+        if show_per_iter:
+            df = pd.DataFrame([df.iloc[-1]], columns=df.columns)
+        else:
+            df.loc["Average(" + str(metrics_idx_end-metrics_idx_start) + "iters)"] = df.mean().round(keep_digits)
+            if only_average:
+                df = pd.DataFrame([df.iloc[-1]], columns=df.columns)
+
+        df.index.name = 'iter index'
+        df[self.batch_size_name] = df[self.batch_size_name].astype(int)
+        df[self.k_name] = df[self.k_name].astype(int)
+
+        self.peak_memory_name = "profile memory(GB)"
+        self.block_memory_name = "total cache memory(GB)"
+        self.max_kv_memory_name = "max cache used(GB)"
+        self.mean_kv_memory_name = "mean cache used(GB)"
+        self.max_kv_usage_name = "max cache usage(%)"
+        self.mean_kv_usage_name = "mean cache usage(%)"
+        memory_metrics_data = [
+            (
+                self.peak_memory_name, [round_fn(self.peak_memory / 1024 / 1024 / 1024) for i in idx_range]
+            ),
+            (
+                self.block_memory_name, [round_fn(self.block_memory / 1024 / 1024 / 1024) for i in idx_range]
+            ),
+            (
+                self.max_kv_memory_name, [
+                    0.0 if len(self.num_free_gpu_blocks_list[i]) == 1 else \
+                    round_fn((1.0 - min(self.num_free_gpu_blocks_list[i]) / self.num_total_gpu_blocks) \
+                                                     * self.block_memory / 1024 / 1024 / 1024) \
+                        for i in idx_range]
+            ),
+            (
+                self.mean_kv_memory_name, [
+                    0.0 if len(self.num_free_gpu_blocks_list[i]) == 1 else \
+                    round_fn((1.0 - statistics.fmean(self.num_free_gpu_blocks_list[i]) / self.num_total_gpu_blocks) \
+                                                     * self.block_memory / 1024 / 1024 / 1024) \
+                        for i in idx_range]
+            ),
+            (
+                self.max_kv_usage_name, [
+                    0.0 if len(self.num_free_gpu_blocks_list[i]) == 1 else \
+                    round_fn((1.0 - min(self.num_free_gpu_blocks_list[i]) / self.num_total_gpu_blocks) * 100.0) \
+                        for i in idx_range]
+            ),
+            (
+                self.mean_kv_usage_name, [
+                    0.0 if len(self.num_free_gpu_blocks_list[i]) == 1 else \
+                    round_fn((1.0 - statistics.fmean(self.num_free_gpu_blocks_list[i]) / self.num_total_gpu_blocks) * 100.0) \
+                        for i in idx_range]
+            )
+        ]
+
+        self.memory_metrics_data = dict(memory_metrics_data)
+
+        # Print
+        memory_df = pd.DataFrame(self.memory_metrics_data)
+        if show_per_iter:
+            memory_df = pd.DataFrame([memory_df.iloc[-1]], columns=memory_df.columns)
+        else:
+            memory_df.loc["Average(" + str(metrics_idx_end-metrics_idx_start) + "iters)"] = memory_df.mean().round(keep_digits)
+            if only_average:
+                memory_df = pd.DataFrame([memory_df.iloc[-1]], columns=memory_df.columns)
+
+        memory_df.index.name = 'iter index'
+
+        pd.set_option('display.colheader_justify', 'center')
+        pd.set_option('display.max_columns', None)
+        pd.set_option('display.max_rows', None)
+        print("********************************* Test Info****************************")
+        mm_params_text = " ".join(f"{key}:{value}" for key, value in self.mm_kwargs.items())
+        print("Generation Config {} input len:{} output len:{} tp_nums:{} quantization:{}".format(
+            mm_params_text, input_len,output_len,tp_nums,quantization))
+
+        if dump_info and insert_latency_device:
+            dump_info.init_param(batch_size=self.metrics_data['batch size'][0],
+                                 input_len=input_len, output_len=output_len,
+                                 context_latency_device=np.mean(self.metrics_data['context latency device(ms)']),
+                                 generate_latency_device=np.mean(self.metrics_data['per token latency device(ms)']))
+            dump_info.dump()
+
+        print("*************************Performance Info******************************")
+        print(f"Total prefill steps: {total_prefill_steps}")
+        print(df.to_string())
+        if not is_embedding_task:
+            # embedding task does not do profile run, so does not have memory infos
+            print(memory_df.to_string())
+        if insert_latency_device :
+            context_latency = np.mean(self.metrics_data['context latency device(ms)'])
+            generate_latency = np.mean(self.metrics_data['per token latency device(ms)'])
+
+            millisecond2second_unit = 1000
+            if dump_info and dump_info.has_information_dump():
+                dump_info.dump_performance_info()
+            elif dump_info:
+                context_tflops_per_second = 0
+                decoder_tflops_per_second = 0
+                flops2Tflops = 1000 * 1000 * 1000 * 1000
+                context_tflops = dump_info.flops_info["context_flops"] / flops2Tflops
+                decoder_tflops = dump_info.flops_info["decoder_flops"] / flops2Tflops
+                if not context_latency:
+                    logger.warning("context_latency is 0, context_tflops_per_second unable to output correctly")
+                else:
+                    context_tflops_per_second = context_tflops / (context_latency / millisecond2second_unit)
+
+                if not generate_latency:
+                    # embedding task does not have generate process
+                    if not is_embedding_task:
+                        logger.warning("generate_latency is 0, decoder_tflops_per_second unable to output correctly")
+                else:
+                    decoder_tflops_per_second = decoder_tflops / (generate_latency / millisecond2second_unit)
+                print("Context tflops:  {} Tflops".format(context_tflops))
+                print("Generate tflops: {} Tflops".format(decoder_tflops))
+                print("Context tflops_per_second:  {} Tflops/s".format(context_tflops_per_second))
+                print("Generate tflops_per_second: {} Tflops/s".format(decoder_tflops_per_second))
+                df["context_tflops"] = context_tflops
+                df["decoder_tflops"] = decoder_tflops
+                df["context_tflops_per_second"] = context_tflops_per_second
+                df["decoder_tflops_per_second"] = decoder_tflops_per_second
+                if (not context_tflops) or (not decoder_tflops):
+                    logger.warning("the flops is 0, Please check if the model file is correctly parsed!!!!!!!!!")
+        print("***********************************************************************")
+        # collect context_hfu and
+        self.save_hfu_info = False
+        if insert_latency_device and dump_info and dump_info.has_information_dump():
+            self.save_hfu_info = True
+            self.context_hfu_name = "Context HFU"
+            self.decoder_hfu_name = "Decoder  HFU"
+            self.decoder_io_efficiency_name = "Decoder IO Efficiency"
+            self.metrics_data[self.context_hfu_name] = dump_info.hfu_info["context_hfu"] * 100
+            self.metrics_data[self.decoder_hfu_name] = dump_info.hfu_info["decoder_hfu"] * 100
+            self.metrics_data[self.decoder_io_efficiency_name] = dump_info.io_efficiency * 100
+        self.to_csv(os.getenv("OUTPUT_CSV_PATH", "output.csv"), show_per_iter=show_per_iter)
diff --git a/vllm_mlu/vllm_mlu/model_executor/__init__.py b/vllm_mlu/vllm_mlu/model_executor/__init__.py
new file mode 100755
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/model_executor/custom_model/__init__.py b/vllm_mlu/vllm_mlu/model_executor/custom_model/__init__.py
new file mode 100644
index 000000000..6d4e14061
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/custom_model/__init__.py
@@ -0,0 +1 @@
+import vllm_mlu.model_executor.custom_model.custom
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/model_executor/custom_model/custom.py b/vllm_mlu/vllm_mlu/model_executor/custom_model/custom.py
new file mode 100644
index 000000000..ad772484e
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/custom_model/custom.py
@@ -0,0 +1,622 @@
+from collections import namedtuple
+from typing import Any, Dict, Iterable, Union, List, Optional, Tuple
+import math
+import torch
+from torch import nn
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm.config import CacheConfig, VllmConfig
+from vllm_mlu.transformers_utils.configs import CustomConfig
+from vllm.attention import Attention, AttentionMetadata
+from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
+                              get_tensor_model_parallel_world_size,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (QKVParallelLinear,
+                                               RowParallelLinear,
+                                               ReplicatedLinear)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
+from vllm.model_executor.models.interfaces import SupportsPP
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.sequence import IntermediateTensors
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm.model_executor.models.utils import PPMissingLayer, make_layers
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, is_per_tensor_smoothquant,
+    is_per_token_smoothquant, quant_fusion_with_rmsnorm,
+    quant_fusion_with_layernorm)
+
+
+def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
+    closest_power_of_2 = 2**math.floor(math.log2(total_num_heads))
+    base = torch.tensor(
+        2**(-(2**-(math.log2(closest_power_of_2) - 3))),
+        dtype=torch.float32,
+    )
+    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
+    slopes = torch.pow(base, powers)
+
+    if closest_power_of_2 != total_num_heads:
+        extra_base = torch.tensor(
+            2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
+            dtype=torch.float32,
+        )
+        num_remaining_heads = min(closest_power_of_2,
+                                  total_num_heads - closest_power_of_2)
+        extra_powers = torch.arange(start=1,
+                                    end=1 + 2 * num_remaining_heads,
+                                    step=2,
+                                    dtype=torch.int32)
+        slopes = torch.cat(
+            [slopes, torch.pow(extra_base, extra_powers)], dim=0)
+    return slopes
+
+
+class LayerNorm(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.bias = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        x = x.view(-1, self.weight.data.shape[0])
+        if residual is not None:
+            residual = residual.view(-1, self.weight.data.shape[0])
+            return mlu_ops.fused_layer_norm(x, residual, self.weight.data, self.bias.data, None, self.variance_epsilon, True)
+        else:
+            return mlu_ops.fused_layer_norm(x, residual, self.weight.data, self.bias.data, None, self.variance_epsilon, False)
+
+
+_NORM_DICT: Dict[str, nn.Module] = {"rmsnorm": RMSNorm, "layernorm": LayerNorm}
+
+
+class CustomMoeBlock(nn.Module):
+
+    def __init__(
+        self,
+        config: CustomConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+    ):
+        super().__init__()
+        self.config = config
+        self.rank = get_tensor_model_parallel_rank()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.n_routed_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        if self.tp_size > self.n_routed_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than "
+                f"the number of experts {self.n_routed_experts}.")
+
+        self.moe_intermediate_size = self.config.moe_intermediate_size // self.tp_size
+
+        if quant_config is None:
+            self.w1 = nn.Parameter(
+                torch.empty(self.config.num_experts,
+                            2 * self.moe_intermediate_size if self.config.is_gated else self.moe_intermediate_size,
+                            self.config.hidden_size,
+                            dtype=torch.get_default_dtype()), requires_grad=False)
+            self.w2 = nn.Parameter(
+                torch.empty(self.config.num_experts,
+                            self.config.hidden_size,
+                            self.moe_intermediate_size,
+                            dtype=torch.get_default_dtype()), requires_grad=False)
+            self.w1_scale = None
+            self.w2_scale = None
+            self.input_smooth = None
+            self.act_smooth = None
+        else:
+            assert quant_config.weight_bits == 8
+            self.w1 = nn.Parameter(
+                torch.empty(self.config.num_experts,
+                            2 * self.moe_intermediate_size if self.config.is_gated else self.moe_intermediate_size,
+                            self.config.hidden_size,
+                            device="mlu",
+                            dtype=torch.int8), requires_grad=False)
+            self.w2 = nn.Parameter(
+                torch.empty(self.config.num_experts,
+                            self.config.hidden_size,
+                            self.moe_intermediate_size,
+                            device="mlu",
+                            dtype=torch.int8), requires_grad=False)
+            self.w1_scale = nn.Parameter(
+                torch.empty(
+                    self.config.num_experts,
+                    2 * self.moe_intermediate_size if self.config.is_gated else self.moe_intermediate_size,
+                    device="mlu",
+                    dtype=torch.float32), requires_grad=False)
+            self.w2_scale = nn.Parameter(
+                torch.empty(
+                    self.config.num_experts,
+                    self.config.hidden_size,
+                    device="mlu",
+                    dtype=torch.float32), requires_grad=False)
+            self.input_smooth = None
+            self.act_smooth = None
+            if quant_config.quant_mode == "SmoothQuant":
+                self.input_smooth =nn.Parameter(
+                    torch.empty(
+                        self.config.num_experts,
+                        self.config.hidden_size,
+                        device="mlu",
+                        dtype=torch.float32), requires_grad=False)
+                self.act_smooth =nn.Parameter(
+                    torch.empty(
+                        self.config.num_experts,
+                        self.moe_intermediate_size,
+                        device="mlu",
+                        dtype=torch.float32), requires_grad=False)
+
+        self.gate = ReplicatedLinear(config.hidden_size,
+                                     self.n_routed_experts,
+                                     bias=False,
+                                     quant_config=None)
+        if config.shared_expert_intermediate_size > 0:
+            self.shared_expert = FeedForward(hidden_size=config.hidden_size,
+                                             intermediate_size=config.shared_expert_intermediate_size,
+                                             hidden_act = self.config.hidden_act,
+                                             up_proj_name='gate_up_proj',
+                                             is_gated=self.config.is_gated,
+                                             down_proj_name='down_proj',
+                                             bias=False,
+                                             quant_config=quant_config,
+                                             reduce_results=False)
+        else:
+            self.shared_expert = None
+        self.shared_expert_gate = torch.nn.Linear(config.hidden_size,
+                                                  1,
+                                                  bias=False)
+
+    def forward(self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_output = None
+        if self.shared_expert is not None:
+            shared_output = self.shared_expert(hidden_states)
+            if self.shared_expert_gate is not None:
+                shared_output = F.sigmoid(
+                    self.shared_expert_gate(hidden_states)) * shared_output
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        residual_ = None if self.rank > 0 else residual
+        final_hidden_states = mlu_ops.fused_moe(hidden_states,
+                                               router_logits,
+                                               self.w1,
+                                               self.w2,
+                                               None,
+                                               None,
+                                               residual_,
+                                               self.input_smooth,
+                                               self.act_smooth,
+                                               self.w1_scale,
+                                               self.w2_scale,
+                                               self.top_k,
+                                               self.config.norm_topk_prob,
+                                               self.config.is_gated,
+                                               self.config.hidden_act)
+        if shared_output is not None:
+            final_hidden_states = final_hidden_states + shared_output
+
+        reduce_results = (self.config.use_parallel_residual == False)
+        if reduce_results:
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                final_hidden_states)
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+class CustomAttention(nn.Module):
+
+    def __init__(
+        self,
+        config: CustomConfig,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = ""
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False)
+        self.hidden_size = config.hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = config.num_attention_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads)
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = self.hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.kv_scale = 1.0
+        self.scaling = self.head_dim**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            self.hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=attention_bias,
+            quant_config=quant_config,
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            self.hidden_size,
+            bias=attention_bias,
+            quant_config=quant_config,
+            skip_bias_add=(self.config.use_parallel_residual and attention_bias),
+            reduce_results = (self.config.use_parallel_residual == False),
+        )
+
+        self.alibi_slopes = None
+        self.rotary_emb = None
+        if self.config.position_embedding_type == "ALIBI":
+            tp_rank = get_tensor_model_parallel_rank()
+            head_start = tp_rank * self.num_heads
+            head_end = (tp_rank + 1) * self.num_heads
+            alibi_slopes = _get_alibi_slopes(self.total_num_heads)
+            self.alibi_slopes = alibi_slopes[head_start:head_end].tolist()
+        else:
+            rope_theta = getattr(config, "rope_theta", 10000)
+            rope_scaling = getattr(config, "rope_scaling", None)
+            if rope_scaling is not None and getattr(
+                    config, "original_max_position_embeddings", None):
+                rope_scaling["original_max_position_embeddings"] = (
+                    config.original_max_position_embeddings)
+            max_position_embeddings = getattr(config, "max_sequence_length", 8192)
+            is_neox_style = getattr(config, "is_neox_style", False)
+            self.rotary_emb = get_rope(
+                self.head_dim,
+                rotary_dim=self.head_dim,
+                max_position=max_position_embeddings,
+                base=rope_theta,
+                is_neox_style=is_neox_style,
+                rope_scaling=rope_scaling,
+            )
+
+        self.attn = Attention(self.num_heads,
+                              self.head_dim,
+                              self.scaling,
+                              num_kv_heads=self.num_kv_heads,
+                              alibi_slopes=self.alibi_slopes,
+                              cache_config=cache_config,
+                              prefix=f"{prefix}.attn")
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+        smooth_quant_scale: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        if self.rotary_emb:
+            qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+            self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+        attn_output = self.attn(q, k, v)
+        output, bias = self.o_proj(attn_output, residual)
+        if self.o_proj.skip_bias_add and get_tensor_model_parallel_rank() == 0:
+            output += bias
+        return output
+
+
+class CustomDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: CustomConfig,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = ""
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.self_attn = CustomAttention(
+            config=config,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn"
+        )
+
+        mlp_bias = getattr(config, "mlp_bias", False) or getattr(config, "bias", False)
+        is_gated = getattr(config, "is_gated", False)
+
+        if config.num_experts is not None:
+            self.mlp = CustomMoeBlock(config=config,
+                                    quant_config=quant_config)
+        else:
+            self.mlp = FeedForward(hidden_size=config.hidden_size,
+                                   intermediate_size=config.intermediate_size,
+                                   hidden_act=self.config.hidden_act,
+                                   up_proj_name='up_proj',
+                                   is_gated=is_gated,
+                                   down_proj_name='down_proj',
+                                   bias=mlp_bias,
+                                   quant_config=quant_config,
+                                   skip_bias_add=(self.config.use_parallel_residual and mlp_bias),
+                                   reduce_results = (self.config.use_parallel_residual == False))
+
+        self.input_layernorm = _NORM_DICT[self.config.norm_type](config.hidden_size, eps=config.norm_eps)
+        self.post_attention_layernorm = _NORM_DICT[self.config.norm_type](config.hidden_size, eps=config.norm_eps)
+
+        # perf per-tensor sq cases by fusing quantization in layernorm
+        self.is_per_tesnor_sq_perf_cases = (is_per_tensor_smoothquant(quant_config) and
+                                            not self.config.apply_residual_connection_post_layernorm)
+        self.is_per_token_sq_perf_cases = (is_per_token_smoothquant(quant_config) and
+                                            not self.config.apply_residual_connection_post_layernorm)
+        if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+            self.self_attn.qkv_proj.quant_method.skip_quant_input = True
+            self.quant_fusion_attn_layernorm = None
+            self.is_moe = config.num_experts is not None
+            self.use_rmsnorm = self.config.norm_type == "rmsnorm"
+            if not self.is_moe:
+                self.mlp.up_proj.quant_method.skip_quant_input = True
+                self.quant_fusion_mlp_layernorm = None
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.config.use_parallel_residual:
+            # x = x + attn(ln1(x)) + mlp(ln2(x))
+            layernorm_output = self.input_layernorm(hidden_states)
+            attention_output = self.self_attn(
+                positions=positions,
+                hidden_states=layernorm_output,
+            )
+
+            layernorm_output = self.post_attention_layernorm(hidden_states)
+            if self.mlp.skip_bias_add:
+                mlp_output, mlp_bias = self.mlp(layernorm_output)
+                if get_tensor_model_parallel_rank() == 0:
+                    mlp_output += mlp_bias
+            else:
+                mlp_output = self.mlp(layernorm_output)
+
+            if get_tensor_model_parallel_rank() == 0:
+                hidden_states = mlp_output + attention_output + hidden_states
+            else:
+                hidden_states = mlp_output + attention_output
+
+            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+            return hidden_states
+        else:
+            # rmsnorm use fused_rms_norm to get better performance
+            # if apply_residual_connection_post_layernorm:
+            #     x = ln1(x) + attn(ln1(x))
+            #     x = ln2(x) + mlp(ln2(x))
+            # else:
+            #     x = x + attn(ln1(x))
+            #     x = x + mlp(ln2(x))
+            attn_layernorm = self.input_layernorm
+            mlp_layernorm = self.post_attention_layernorm
+            if self.is_per_tesnor_sq_perf_cases:
+                quant_fusion_func = (quant_fusion_with_rmsnorm if
+                                     self.use_rmsnorm else quant_fusion_with_layernorm)
+                if self.quant_fusion_attn_layernorm is None:
+                    self.quant_fusion_attn_layernorm = quant_fusion_func(
+                        self.input_layernorm, self.self_attn.qkv_proj.scale_to_int)
+                attn_layernorm = self.quant_fusion_attn_layernorm
+                if not self.is_moe:
+                    if self.quant_fusion_mlp_layernorm is None:
+                        self.quant_fusion_mlp_layernorm = quant_fusion_func(
+                            self.post_attention_layernorm, self.mlp.up_proj.scale_to_int)
+                    mlp_layernorm = self.quant_fusion_mlp_layernorm
+            elif self.is_per_token_sq_perf_cases:
+                quant_fusion_func = (quant_fusion_with_rmsnorm if
+                                     self.use_rmsnorm else quant_fusion_with_layernorm)
+                if self.quant_fusion_attn_layernorm is None:
+                    self.quant_fusion_attn_layernorm = quant_fusion_func(
+                        self.input_layernorm, self.self_attn.qkv_proj.smooth, dynamic_quant=True)
+                attn_layernorm = self.quant_fusion_attn_layernorm
+                if not self.is_moe:
+                    if self.quant_fusion_mlp_layernorm is None:
+                        self.quant_fusion_mlp_layernorm = quant_fusion_func(
+                            self.post_attention_layernorm, self.mlp.up_proj.smooth, dynamic_quant=True)
+                    mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+            post_norm_fuse_en=(self.is_per_token_sq_perf_cases and not self.is_moe)
+            return decoder_layer_forward_base(positions=positions,
+                                      hidden_states=hidden_states,
+                                      input_layernorm=attn_layernorm,
+                                      self_attn=self.self_attn,
+                                      post_layernorm=mlp_layernorm,
+                                      mlp=self.mlp,
+                                      apply_residual_connection_post_layernorm=self.config.apply_residual_connection_post_layernorm,
+                                      input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+                                      post_norm_fuse_en=post_norm_fuse_en)
+
+
+class CustomModel(nn.Module):
+
+    def __init__(
+        self,
+        config: CustomConfig,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        if get_pp_group().is_first_rank or (config.tie_word_embeddings
+                                            and get_pp_group().is_last_rank):
+            embed_layer = VocabParallelEmbedding if self.config.use_parallel_embedding else nn.Embedding
+            self.embed_tokens = embed_layer(
+                config.vocab_size,
+                config.hidden_size,
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: CustomDecoderLayer(config=config,
+                                             cache_config=cache_config,
+                                             quant_config=quant_config,
+                                             prefix=prefix),
+            prefix="custom_model")
+
+        if get_pp_group().is_last_rank:
+            self.norm = _NORM_DICT[self.config.norm_type](config.hidden_size, eps=config.norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor],
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_tokens(input_ids)
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+
+        for i in range(self.start_layer, self.end_layer):
+            layer = self.layers[i]
+            hidden_states = layer(
+                positions,
+                hidden_states,
+            )
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({
+                "hidden_states": hidden_states
+            })
+
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+
+class CustomForCausalLM(nn.Module, SupportsPP):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        self.config = vllm_config.model_config.hf_text_config
+        self.quant_config = vllm_config.quant_config
+        self.cache_config = vllm_config.cache_config
+        self._verify_params()
+        self.model = CustomModel(self.config, self.cache_config, self.quant_config)
+
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(self.config.vocab_size, self.config.hidden_size)
+            self.logits_processor = LogitsProcessor(self.config.vocab_size)
+            self.sampler = Sampler()
+        else:
+            self.lm_head = PPMissingLayer()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        hidden_states = self.model(input_ids, positions, intermediate_tensors)
+        return hidden_states
+
+    def compute_logits(self, hidden_states: torch.Tensor,
+                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
+        logits = self.logits_processor(self.lm_head, hidden_states,
+                                       sampling_metadata)
+        return logits
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[SamplerOutput]:
+        next_tokens = self.sampler(logits, sampling_metadata)
+        return next_tokens
+    
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        pass
+
+    def make_empty_intermediate_tensors(
+            self, batch_size: int, dtype: torch.dtype,
+            device: torch.device) -> IntermediateTensors:
+        return IntermediateTensors({
+            "hidden_states":
+            torch.zeros((batch_size, self.config.hidden_size),
+                        dtype=dtype,
+                        device=device),
+        })
+
+    def _verify_params(self) -> None:
+        if (self.config.max_sequence_length) is None or \
+           (self.config.num_hidden_layers) is None or \
+           (self.config.hidden_size) is None or \
+           (self.config.vocab_size) is None or \
+           (self.config.num_attention_heads) is None:
+            raise ValueError(
+                "max_sequence_length, num_hidden_layers, hidden_size, vocab_size, "
+                "num_attention_heads, must be vaild int values")
+
+        if self.config.hidden_act not in ["silu", "gelu"]:
+            raise ValueError(
+                "CustomConfig hidden_act must be one of [silu, gelu]. Got "
+                f"{self.config.hidden_act}.")
+
+        if self.config.position_embedding_type not in ["ALIBI", "ROPE"]:
+            raise ValueError(
+                "position_embedding_type must be one of [ALIBI, ROPE]. Got "
+                f"{self.config.position_embedding_type}.")
+
+        if self.config.num_experts is not None:
+            if self.config.num_experts_per_tok is None:
+                raise ValueError(
+                    "num_experts_per_tok must be a valid int value when num_experts is not None")
+            if self.config.moe_intermediate_size is None:
+                raise ValueError(
+                    "moe_intermediate_size must be a valid int value when num_experts is not None")
+            if self.config.shared_expert_intermediate_size is None:
+                raise ValueError(
+                    "shared_expert_intermediate_size must be a valid int value when num_experts is not None")
+            if self.config.norm_topk_prob is None:
+                raise ValueError(
+                    "norm_topk_prob must be a valid bool value when num_experts is not None")
+            if self.config.mlp_bias is True:
+                raise ValueError(
+                    "mlp_bias must be False when num_experts is not None")
+            if self.quant_config is not None and self.quant_config.get_name() != "SmoothQuant":
+                raise ValueError(
+                    "moe only support smoothquant now")
+        else:
+            if self.config.intermediate_size is None:
+                raise ValueError(
+                    "intermediate_size must be a valid int value when num_experts is None")
+
+        if self.config.norm_type not in ["rmsnorm", "layernorm"]:
+            raise ValueError(
+                "norm_type must be one of [rmsnorm, layernorm]. Got "
+                f"{self.config.norm_type}.")
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/__init__.py b/vllm_mlu/vllm_mlu/model_executor/layers/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/activation.py b/vllm_mlu/vllm_mlu/model_executor/layers/activation.py
new file mode 100644
index 000000000..571e4648e
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/activation.py
@@ -0,0 +1,22 @@
+import torch
+from vllm.model_executor.layers.activation import QuickGELU
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu import _mlu_ops as mlu_ops
+
+def vllm__model_executor__activation__QuickGELU__forward_oot(self, x: torch.Tensor) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: implement forward_oot
+    '''
+    return mlu_ops.active(x, 'quick_gelu', False)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+MluHijackObject.apply_hijack(QuickGELU,
+                             QuickGELU.forward_oot,
+                             vllm__model_executor__activation__QuickGELU__forward_oot)
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/feed_forward.py b/vllm_mlu/vllm_mlu/model_executor/layers/feed_forward.py
new file mode 100755
index 000000000..03ac5a0a3
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/feed_forward.py
@@ -0,0 +1,189 @@
+import torch
+import torch.nn.functional as F
+from typing import Optional
+
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+from vllm.logger import init_logger
+from vllm.lora.layers import BaseLayerWithLoRA
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    ColumnParallelLinear,
+    RowParallelLinear
+)
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu.mlu_hijack_utils import set_is_gated
+from vllm.distributed import get_tensor_model_parallel_rank
+
+logger = init_logger(__name__)
+
+class FeedForward(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        up_proj_name: str,
+        is_gated: bool,
+        down_proj_name: str,
+        bias: bool,
+        quant_config: Optional[QuantizationConfig] = None,
+        skip_bias_add: bool = False,
+        reduce_results: bool = True,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.is_gated = is_gated
+        self.bias = bias
+        self.up_proj_name = up_proj_name
+        self.down_proj_name = down_proj_name
+        self.quant_config = quant_config
+        self.is_initialized = False
+        self.skip_bias_add = skip_bias_add
+        self.reduce_results = reduce_results
+        self.use_bt_ffn = True
+        self.tp_size = get_tensor_model_parallel_world_size()
+        set_is_gated(self.is_gated)
+        self.tp_rank = get_tensor_model_parallel_rank()
+
+        # up_proj with gate or not
+        if self.is_gated:
+            up_proj = MergedColumnParallelLinear(hidden_size,
+                                                 [intermediate_size] * 2,
+                                                 bias=bias,
+                                                 quant_config=quant_config,
+                                                 prefix=f"{prefix}.{up_proj_name}")
+        else:
+            up_proj = ColumnParallelLinear(hidden_size,
+                                           intermediate_size,
+                                           bias=bias,
+                                           skip_bias_add=skip_bias_add,
+                                           quant_config=quant_config,
+                                           prefix=f"{prefix}.{up_proj_name}")
+        self.register_module(up_proj_name, up_proj)
+
+        # down_proj
+        down_proj = RowParallelLinear(intermediate_size,
+                                      hidden_size,
+                                      bias=bias,
+                                      skip_bias_add=skip_bias_add,
+                                      reduce_results=reduce_results,
+                                      quant_config=quant_config,
+                                      prefix=f"{prefix}.{down_proj_name}")
+        self.register_module(down_proj_name, down_proj)
+
+    def prepare_weight(self):
+        if not self.is_initialized:
+            # alpha and beta are 1.0 and 0.0 respectively due to the fact that we don't need residual for now
+            self.alpha = 1.0
+            self.beta = 0.0
+            # place it here to avoid the overhead of calling it in the forward pass
+            self.is_initialized = True
+
+    def _forward(self, hidden_states):
+        self.prepare_weight()
+        up_proj = getattr(self, self.up_proj_name)
+        down_proj = getattr(self, self.down_proj_name)
+        act_dict = {
+            "relu": F.relu,
+            "gelu": F.gelu,
+            "silu": F.silu,
+        }
+        fc1 = F.linear(hidden_states, up_proj.weight, bias=up_proj.bias)
+        if self.is_gated:
+            d = fc1.shape[-1] // 2
+            fc1 = act_dict[self.hidden_act](fc1[..., :d]) * fc1[..., d:]
+        else:
+            fc1 = act_dict[self.hidden_act](fc1)
+        fc2 = F.linear(fc1, down_proj.weight, bias=None)
+        fc2 = tensor_model_parallel_all_reduce(fc2)
+        if not self.skip_bias_add:
+            fc2 = fc2 + down_proj.bias if down_proj.bias is not None else fc2
+        return fc2
+
+    def forward_naive(
+        self,
+        hidden_states,
+        residual: Optional[torch.Tensor] = None,
+        smooth_quant_scale: Optional[torch.Tensor] = None
+    ):
+        '''
+        used by quant_tools
+        '''
+        assert self.quant_config is None, "ffn naive forward dosen't support quantization"
+        assert smooth_quant_scale is None, "ffn naive forward dosen't support smooth_quant_scale"
+
+        up_proj = getattr(self, self.up_proj_name)
+        down_proj = getattr(self, self.down_proj_name)
+        residual_ = None if self.tp_rank > 0 else residual
+        fc1, bias = up_proj(hidden_states)
+        if bias is not None:
+            fc1 += bias
+        fc1 = mlu_ops.active(fc1, self.hidden_act, self.is_gated)
+        out, bias = down_proj(fc1, residual=residual_)
+
+        if self.skip_bias_add:
+            return out, bias
+        return out
+
+    def forward(
+        self,
+        hidden_states,
+        residual: Optional[torch.Tensor] = None,
+        smooth_quant_scale: Optional[torch.Tensor] = None
+    ):
+        self.prepare_weight()
+
+        if self.use_bt_ffn is False:
+            return self.forward_naive(hidden_states, residual, None)
+
+        up_proj = getattr(self, self.up_proj_name)
+        down_proj = getattr(self, self.down_proj_name)
+        residual_ = None if self.tp_rank > 0 else residual
+        if (self.quant_config is None and not isinstance(up_proj, BaseLayerWithLoRA)
+                and not isinstance(down_proj, BaseLayerWithLoRA)):
+            # The matmul formula is the following:
+            #   mul_out = alpha * (matmul(input, filter, transpose\_b=True) + bias) + beta * residual
+            #   output = active(mul_out)
+            # Notes: We cannot use the activation function in matmul because it does not support gated operation
+            #  we might support its in tmo matmul in the future
+            fc1 = mlu_ops.matmul(hidden_states.view(-1, self.hidden_size), up_proj.weight, up_proj.bias,
+                                None, 'none', self.alpha, self.beta)
+            act_out = mlu_ops.active(fc1, self.hidden_act, self.is_gated)
+            beta = 0.0
+            if residual_ is not None:
+                beta = 1.0
+                residual_ = residual_.view(-1, residual_.shape[-1])
+            out_ = mlu_ops.matmul(act_out, down_proj.weight, None, residual_, 'none', self.alpha, beta)
+            # bias if existed need to add after second matmul according to the original design of vllm
+            if self.reduce_results:
+                out = tensor_model_parallel_all_reduce(out_)
+            else:
+                out = out_
+            # do the bias add if needed
+            if not self.skip_bias_add:
+                out = out + down_proj.bias if down_proj.bias is not None else out
+            else:
+                return out, down_proj.bias
+        else:
+            fc1, bias = up_proj(hidden_states, smooth_quant_scale=smooth_quant_scale)
+            if bias is not None:
+                fc1 += bias
+            input_scale= None
+            if (fc1.shape[-1] < 24000 and
+                self.quant_config is not None and self.quant_config.get_name() == "SmoothQuant" and
+                self.quant_config.input_quant_method == "per_token"):
+                down_proj.quant_method.skip_quant_input = True
+                fc1, input_scale = mlu_ops.per_token_smooth_quantize(fc1, down_proj.smooth, None, None, act_mode=self.hidden_act, is_gated=self.is_gated)
+            else:
+                fc1 = mlu_ops.active(fc1, self.hidden_act, self.is_gated)
+            out, bias = down_proj(fc1, residual=residual_, smooth_quant_scale=input_scale)
+
+            if self.skip_bias_add:
+                return out, bias
+        return out
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/__init__.py b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/fused_moe.py b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/fused_moe.py
new file mode 100644
index 000000000..4ca40494c
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/fused_moe.py
@@ -0,0 +1,635 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused MoE kernel."""
+import functools
+import json
+import os
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.logger import init_logger
+from vllm.utils import direct_register_custom_op
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+    write_zeros_to_output, get_default_config)
+
+from vllm_mlu.model_executor.layers.fused_moe.moe_align_block_size import (
+    moe_align_block_size)
+from vllm_mlu.model_executor.layers.fused_moe.utils import _fp8_quantize
+import vllm_mlu._mlu_ops as mlu_ops
+
+logger = init_logger(__name__)
+
+@triton.jit
+def fused_moe_kernel(
+        # Pointers to matrices
+        a_ptr,
+        b_ptr,
+        c_ptr,
+        a_scale_ptr,
+        b_scale_ptr,
+        topk_weights_ptr,
+        sorted_token_ids_ptr,
+        expert_ids_ptr,
+        num_tokens_post_padded_ptr,
+        # Matrix dimensions
+        N,
+        K,
+        EM,
+        num_valid_tokens,
+        # The stride variables represent how much to increase the ptr by when
+        # moving by 1 element in a particular dimension. E.g. `stride_am` is
+        # how much to increase `a_ptr` by to get the element one row down
+        # (A has M rows).
+        stride_am,
+        stride_ak,
+        stride_be,
+        stride_bk,
+        stride_bn,
+        stride_cm,
+        stride_cn,
+        stride_asm,
+        stride_ask,
+        stride_bse,
+        stride_bsk,
+        stride_bsn,
+        # Block size for block-wise quantization
+        group_n: tl.constexpr,
+        group_k: tl.constexpr,
+        # Meta-parameters
+        BLOCK_SIZE_M: tl.constexpr,
+        BLOCK_SIZE_N: tl.constexpr,
+        BLOCK_SIZE_K: tl.constexpr,
+        GROUP_SIZE_M: tl.constexpr,
+        MUL_ROUTED_WEIGHT: tl.constexpr,
+        top_k: tl.constexpr,
+        compute_type: tl.constexpr,
+        use_fp8_w8a8: tl.constexpr,
+        use_int8_w8a16: tl.constexpr):
+    """
+    Implements the fused computation for a Mixture of Experts (MOE) using
+    token and expert matrices.
+
+    Key Parameters:
+    - A: The input tensor representing tokens with shape (*, K), where '*' can
+        be any shape representing batches and K is the feature dimension of
+        each token.
+    - B: The stacked MOE weight tensor with shape (E, N, K), where E is
+        the number of experts, K is the input feature dimension, and N is
+        the output feature dimension.
+    - C: The output cache tensor with shape (M, topk, N), where M is the
+        total number of tokens post padding, topk is the number of times
+        each token is repeated, and N is the output feature dimension.
+    - sorted_token_ids: A tensor containing the sorted indices of tokens,
+        repeated topk times and arranged by the expert index they are
+        assigned to.
+    - expert_ids: A tensor containing the indices of the expert for each
+        block. It determines which expert matrix from B should be used for
+        each block in A.
+    This kernel performs the multiplication of a token by its corresponding
+    expert matrix as determined by `expert_ids`. The sorting of
+    `sorted_token_ids` by expert index and padding ensures divisibility by
+    BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix
+    multiplication across different blocks processed by the same expert.
+    """
+    # -----------------------------------------------------------
+    # Map program ids `pid` to the block of C it should compute.
+    # This is done in a grouped ordering to promote L2 data reuse.
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Split the program ID into two dimensions (pid_0 and pid_1)
+    '''
+    pid_0 = tl.program_id(axis=0)
+    pid_1 = tl.program_id(axis=1)
+    pid = pid_1 * tl.num_programs(axis=0) + pid_0
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # ----------------------------------------------------------
+    # Create pointers for the first blocks of A and B.
+    # We will advance this pointer as we move in the K direction
+    # and accumulate
+    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
+    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(
+        tl.int64)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
+    token_mask = offs_token < num_valid_tokens
+
+    off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+    if off_experts == -1:
+        # -----------------------------------------------------------
+        # Write back zeros to the output when the expert is not
+        # in the current expert parallel rank.
+        write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N,
+                              offs_token, token_mask, BLOCK_SIZE_M,
+                              BLOCK_SIZE_N, compute_type)
+        return
+
+    offs_bn = (pid_n * BLOCK_SIZE_N +
+               tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
+                      offs_k[None, :] * stride_ak)
+
+    b_ptrs = b_ptr + off_experts * stride_be + (offs_k[:, None] * stride_bk +
+                                                offs_bn[None, :] * stride_bn)
+    if use_int8_w8a16:
+        b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + offs_bn[
+            None, :] * stride_bsn
+        b_scale = tl.load(b_scale_ptrs)
+
+    if use_fp8_w8a8:
+        if group_k > 0 and group_n > 0:
+            a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm
+            offs_bsn = offs_bn // group_n
+            b_scale_ptrs = (b_scale_ptr + off_experts * stride_bse +
+                            offs_bsn * stride_bsn)
+        else:
+            a_scale = tl.load(a_scale_ptr)
+            b_scale = tl.load(b_scale_ptr + off_experts)
+
+    # -----------------------------------------------------------
+    # Iterate to compute a block of the C matrix.
+    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
+    # of fp32 values for higher accuracy.
+    # `accumulator` will be converted back to fp16 after the loop.
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # Load the next block of A and B, generate a mask by checking the
+        # K dimension.
+        a = tl.load(a_ptrs,
+                    mask=token_mask[:, None] &
+                    (offs_k[None, :] < K - k * BLOCK_SIZE_K),
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        # We accumulate along the K dimension.
+        if use_int8_w8a16:
+            accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
+        elif use_fp8_w8a8:
+            if group_k > 0 and group_n > 0:
+                k_start = k * BLOCK_SIZE_K
+                offs_ks = k_start // group_k
+                a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask,
+                                  mask=token_mask,
+                                  other=0.0)
+                b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
+
+                accumulator += tl.dot(a, b) * a_scale[:,
+                                                      None] * b_scale[None, :]
+            else:
+                accumulator = tl.dot(a, b, acc=accumulator)
+        else:
+            accumulator += tl.dot(a, b)
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token,
+                             mask=token_mask,
+                             other=0)
+        accumulator = accumulator * moe_weight[:, None]
+    if use_int8_w8a16:
+        accumulator = (accumulator * b_scale).to(compute_type)
+    elif use_fp8_w8a8:
+        if group_k > 0 and group_n > 0:
+            accumulator = accumulator.to(compute_type)
+        else:
+            accumulator = (accumulator * a_scale * b_scale).to(compute_type)
+    else:
+        accumulator = accumulator.to(compute_type)
+    # -----------------------------------------------------------
+    # Write back the block of the output
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[
+        None, :]
+    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, accumulator, mask=c_mask)
+
+
+def invoke_fused_moe_kernel(A: torch.Tensor,
+                            B: torch.Tensor,
+                            C: torch.Tensor,
+                            A_scale: Optional[torch.Tensor],
+                            B_scale: Optional[torch.Tensor],
+                            B_zp: Optional[torch.Tensor],
+                            topk_weights: Optional[torch.Tensor],
+                            sorted_token_ids: torch.Tensor,
+                            expert_ids: torch.Tensor,
+                            num_tokens_post_padded: torch.Tensor,
+                            mul_routed_weight: bool,
+                            top_k: int,
+                            config: Dict[str, Any],
+                            compute_type: tl.dtype,
+                            use_fp8_w8a8: bool,
+                            use_int8_w8a16: bool,
+                            use_int4_w4a16: bool,
+                            block_shape: Optional[List[int]] = None) -> None:
+    assert topk_weights is not None or not mul_routed_weight
+    assert topk_weights is None or topk_weights.stride(1) == 1
+    assert sorted_token_ids.stride(0) == 1
+
+    if use_fp8_w8a8:
+        assert B_scale is not None
+        assert (block_shape is None or triton.cdiv(B.shape[-2], block_shape[0])
+                == B_scale.shape[-2])
+        assert (block_shape is None or triton.cdiv(B.shape[-1], block_shape[1])
+                == B_scale.shape[-1])
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Only support fp8_w8a8
+    '''
+    assert use_fp8_w8a8, "use_fp8_w8a8 must be True"
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    M = A.shape[0]
+    num_tokens = M * top_k
+
+    EM = sorted_token_ids.shape[0]
+    if A.shape[0] < config["BLOCK_SIZE_M"]:
+        # optimize for small batch_size.
+        # We assume that top_ids of each token is unique, so
+        # so num_valid_experts <= batch_size <= BLOCK_SIZE_M,
+        # and we can skip some invalid blocks.
+        EM = min(sorted_token_ids.shape[0],
+                 A.shape[0] * top_k * config['BLOCK_SIZE_M'])
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Split the program ID into two dimensions (pid_0, pid_1)
+    '''
+    grid = lambda META: (triton.cdiv(EM, META['BLOCK_SIZE_M']), triton.cdiv(
+        B.shape[1], META['BLOCK_SIZE_N']), )
+
+    assert not (use_int8_w8a16 or use_int4_w4a16)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    if use_fp8_w8a8:
+        config = config.copy()
+        BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K")
+        if block_shape is not None:
+            BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0],
+                                                 block_shape[1]))
+        fused_moe_kernel[grid](
+            A,
+            B,
+            C,
+            A_scale,
+            B_scale,
+            topk_weights,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            B.shape[1],
+            B.shape[2],
+            EM,
+            num_tokens,
+            A.stride(0),
+            A.stride(1),
+            B.stride(0),
+            B.stride(2),
+            B.stride(1),
+            C.stride(1),
+            C.stride(2),
+            A_scale.stride(0)
+            if A_scale is not None and A_scale.ndim == 2 else 0,
+            A_scale.stride(1)
+            if A_scale is not None and A_scale.ndim == 2 else 0,
+            B_scale.stride(0)
+            if B_scale is not None and B_scale.ndim >= 2 else 0,
+            B_scale.stride(2)
+            if B_scale is not None and B_scale.ndim == 3 else 0,
+            B_scale.stride(1)
+            if B_scale is not None and B_scale.ndim >= 2 else 0,
+            0 if block_shape is None else block_shape[0],
+            0 if block_shape is None else block_shape[1],
+            MUL_ROUTED_WEIGHT=mul_routed_weight,
+            top_k=top_k,
+            compute_type=compute_type,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            BLOCK_SIZE_K=BLOCK_SIZE_K,
+            **config,
+        )
+
+
+def inplace_fused_experts(hidden_states: torch.Tensor,
+                          w1: torch.Tensor,
+                          w2: torch.Tensor,
+                          topk_weights: torch.Tensor,
+                          topk_ids: torch.Tensor,
+                          activation: str = "silu",
+                          apply_router_weight_on_input: bool = False,
+                          use_fp8_w8a8: bool = False,
+                          use_int8_w8a16: bool = False,
+                          use_int4_w4a16: bool = False,
+                          global_num_experts: int = -1,
+                          expert_map: Optional[torch.Tensor] = None,
+                          w1_scale: Optional[torch.Tensor] = None,
+                          w2_scale: Optional[torch.Tensor] = None,
+                          w1_zp: Optional[torch.Tensor] = None,
+                          w2_zp: Optional[torch.Tensor] = None,
+                          a1_scale: Optional[torch.Tensor] = None,
+                          a2_scale: Optional[torch.Tensor] = None,
+                          block_shape: Optional[List[int]] = None) -> None:
+    fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
+                       activation, apply_router_weight_on_input, use_fp8_w8a8,
+                       use_int8_w8a16, use_int4_w4a16, global_num_experts,
+                       expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
+                       a2_scale, block_shape)
+
+
+def inplace_fused_experts_fake(
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: str = "silu",
+        apply_router_weight_on_input: bool = False,
+        use_fp8_w8a8: bool = False,
+        use_int8_w8a16: bool = False,
+        use_int4_w4a16: bool = False,
+        global_num_experts: int = -1,
+        expert_map: Optional[torch.Tensor] = None,
+        w1_scale: Optional[torch.Tensor] = None,
+        w2_scale: Optional[torch.Tensor] = None,
+        w1_zp: Optional[torch.Tensor] = None,
+        w2_zp: Optional[torch.Tensor] = None,
+        a1_scale: Optional[torch.Tensor] = None,
+        a2_scale: Optional[torch.Tensor] = None,
+        block_shape: Optional[List[int]] = None) -> None:
+    pass
+
+
+direct_register_custom_op(
+    op_name="inplace_fused_experts_mlu",
+    op_func=inplace_fused_experts,
+    mutates_args=["hidden_states"],
+    fake_impl=inplace_fused_experts_fake,
+    dispatch_key="PrivateUse1",
+)
+
+def fused_experts(hidden_states: torch.Tensor,
+                  w1: torch.Tensor,
+                  w2: torch.Tensor,
+                  topk_weights: torch.Tensor,
+                  topk_ids: torch.Tensor,
+                  inplace: bool = False,
+                  activation: str = "silu",
+                  apply_router_weight_on_input: bool = False,
+                  use_fp8_w8a8: bool = False,
+                  use_int8_w8a16: bool = False,
+                  use_int4_w4a16: bool = False,
+                  global_num_experts: int = -1,
+                  expert_map: Optional[torch.Tensor] = None,
+                  w1_scale: Optional[torch.Tensor] = None,
+                  w2_scale: Optional[torch.Tensor] = None,
+                  w1_zp: Optional[torch.Tensor] = None,
+                  w2_zp: Optional[torch.Tensor] = None,
+                  a1_scale: Optional[torch.Tensor] = None,
+                  a2_scale: Optional[torch.Tensor] = None,
+                  block_shape: Optional[List[int]] = None,
+                  allow_deep_gemm: bool = False) -> torch.Tensor:
+    torch.ops.vllm.inplace_fused_experts_mlu(
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            global_num_experts=global_num_experts,
+            expert_map=expert_map,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape)
+    return hidden_states
+
+
+def fused_experts_impl(hidden_states: torch.Tensor,
+                       w1: torch.Tensor,
+                       w2: torch.Tensor,
+                       topk_weights: torch.Tensor,
+                       topk_ids: torch.Tensor,
+                       inplace: bool = False,
+                       activation: str = "silu",
+                       apply_router_weight_on_input: bool = False,
+                       use_fp8_w8a8: bool = False,
+                       use_int8_w8a16: bool = False,
+                       use_int4_w4a16: bool = False,
+                       global_num_experts: int = -1,
+                       expert_map: Optional[torch.Tensor] = None,
+                       w1_scale: Optional[torch.Tensor] = None,
+                       w2_scale: Optional[torch.Tensor] = None,
+                       w1_zp: Optional[torch.Tensor] = None,
+                       w2_zp: Optional[torch.Tensor] = None,
+                       a1_scale: Optional[torch.Tensor] = None,
+                       a2_scale: Optional[torch.Tensor] = None,
+                       block_shape: Optional[List[int]] = None):
+    # Check constraints.
+    if use_int4_w4a16:
+        assert hidden_states.shape[1] // 2 == w1.shape[
+            2], "Hidden size mismatch"
+    else:
+        assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
+
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
+    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+    assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
+    assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
+    assert hidden_states.dtype in [
+        torch.float32, torch.float16, torch.bfloat16
+    ]
+
+    num_tokens, _ = hidden_states.shape
+    E, N, _ = w1.shape
+    K = w2.shape[1]
+    if global_num_experts == -1:
+        global_num_experts = E
+    top_k_num = topk_ids.shape[1]
+    # We execute the fused_moe kernel in chunks to circumvent this issue:
+    # https://github.com/vllm-project/vllm/issues/5938
+    CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
+    M = min(num_tokens, CHUNK_SIZE)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Only use the default config
+    '''
+    config = get_default_config(M, E, N, w1.shape[2], topk_ids.shape[1],
+                                hidden_states.dtype, False, block_shape)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    # We can reuse the memory between these because by the time we need
+    # cache3, we're done with cache1
+    cache13 = torch.empty(M * top_k_num * max(N, K),
+                          device=hidden_states.device,
+                          dtype=hidden_states.dtype)
+    intermediate_cache1 = cache13[:M * top_k_num * N].view(M, top_k_num, N)
+    intermediate_cache3 = cache13[:M * top_k_num * K].view(M, top_k_num, K)
+
+    if hidden_states.dtype == torch.bfloat16:
+        compute_type = tl.bfloat16
+    elif hidden_states.dtype == torch.float16:
+        compute_type = tl.float16
+    elif hidden_states.dtype == torch.float32:
+        compute_type = tl.float32
+    else:
+        raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
+
+    if inplace:
+        out_hidden_states = hidden_states
+    else:
+        out_hidden_states = torch.empty_like(hidden_states)
+
+    for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+        begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
+                                          min((chunk + 1) * CHUNK_SIZE,
+                                              num_tokens))
+        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+        tokens_in_chunk, _ = curr_hidden_states.shape
+
+        if tokens_in_chunk == 0:
+            break
+
+        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+        a1q_scale: Optional[torch.Tensor] = None
+
+        if use_fp8_w8a8:
+            qcurr_hidden_states, a1q_scale = _fp8_quantize(
+                curr_hidden_states, a1_scale, block_shape)
+        sorted_token_ids, expert_ids, num_tokens_post_padded = (
+            moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'],
+                                 global_num_experts, expert_map))
+
+        invoke_fused_moe_kernel(qcurr_hidden_states,
+                                w1,
+                                intermediate_cache1,
+                                a1q_scale,
+                                w1_scale,
+                                w1_zp,
+                                curr_topk_weights,
+                                sorted_token_ids,
+                                expert_ids,
+                                num_tokens_post_padded,
+                                apply_router_weight_on_input,
+                                top_k_num,
+                                config,
+                                compute_type=compute_type,
+                                use_fp8_w8a8=use_fp8_w8a8,
+                                use_int8_w8a16=use_int8_w8a16,
+                                use_int4_w4a16=use_int4_w4a16,
+                                block_shape=block_shape)
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: Activate by mlu_ops
+        '''
+        intermediate_cache2 = mlu_ops.active(intermediate_cache1.view(-1, N),
+                                              act_mode=activation,
+                                              is_gated=True)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        a2q_scale: Optional[torch.Tensor] = None
+
+        if use_fp8_w8a8:
+            qintermediate_cache2, a2q_scale = _fp8_quantize(
+                intermediate_cache2, a2_scale, block_shape)
+        invoke_fused_moe_kernel(qintermediate_cache2,
+                                w2,
+                                intermediate_cache3 if topk_ids.shape[1] > 1 else out_hidden_states[begin_chunk_idx:end_chunk_idx],
+                                a2q_scale,
+                                w2_scale,
+                                w2_zp,
+                                curr_topk_weights,
+                                sorted_token_ids,
+                                expert_ids,
+                                num_tokens_post_padded,
+                                not apply_router_weight_on_input,
+                                1,
+                                config,
+                                compute_type=compute_type,
+                                use_fp8_w8a8=use_fp8_w8a8,
+                                use_int8_w8a16=use_int8_w8a16,
+                                use_int4_w4a16=use_int4_w4a16,
+                                block_shape=block_shape)
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: replace moe_sum with torch.sum
+        Reference Links: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py#L1513
+        '''
+        if topk_ids.shape[1] == 2:
+            torch.add(
+                intermediate_cache3[:, 0],
+                intermediate_cache3[:, 1],
+                out=out_hidden_states[begin_chunk_idx:end_chunk_idx],
+            ).squeeze(dim=1)
+        elif topk_ids.shape[1] > 2:
+            torch.sum(
+                intermediate_cache3.view(*intermediate_cache3.shape),
+                dim=1,
+                out=out_hidden_states[begin_chunk_idx:end_chunk_idx],
+            )
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    return out_hidden_states
+
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/layer.py b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/layer.py
new file mode 100644
index 000000000..28012543f
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/layer.py
@@ -0,0 +1,52 @@
+from typing import Optional, Callable
+
+import torch
+
+from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__model_executor__layers__fused_moe__layer__UnquantizedFusedMoEMethod__forward_oot(
+    self,
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    top_k: int,
+    renormalize: bool,
+    use_grouped_topk: bool = False,
+    topk_group: Optional[int] = None,
+    num_expert_group: Optional[int] = None,
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
+    custom_routing_function: Optional[Callable] = None,
+    scoring_func: str = "softmax",
+    e_score_correction_bias: Optional[torch.Tensor] = None,
+    apply_router_weight_on_input: bool = False,
+    activation: str = "silu",
+) -> torch.Tensor:
+    from vllm_mlu import _mlu_ops as mlu_ops
+
+    assert use_grouped_topk is False and num_expert_group is None and topk_group is None, \
+        f"Following params: use_grouped_topk, num_expert_group, topk_group are not support yet."
+    return mlu_ops.fused_moe(
+        x,
+        router_logits,
+        layer.w13_weight, layer.w2_weight,
+        None, None, # bias1, bias2
+        None, # residual
+        None, # input_smooth
+        None, # act_smooth
+        None, None, # w1_scale, w2_scale
+        top_k,
+        renormalize,
+        True, # gated
+        activation
+    )
+
+
+MluHijackObject.apply_hijack(
+    UnquantizedFusedMoEMethod,
+    UnquantizedFusedMoEMethod.forward_oot,
+    vllm__model_executor__layers__fused_moe__layer__UnquantizedFusedMoEMethod__forward_oot
+)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/moe_align_block_size.py b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/moe_align_block_size.py
new file mode 100644
index 000000000..ba34586a2
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/moe_align_block_size.py
@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.utils import round_up
+from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
+    moe_align_block_size_triton)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+    expert_map: Optional[torch.Tensor] = None,
+    pad_sorted_ids: bool = False
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Aligns the token distribution across experts to be compatible with block
+    size for matrix multiplication.
+
+    Parameters:
+    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
+        top-k expert indices for each token.
+    - block_size: The block size used in block matrix multiplication.
+    - num_experts: The total number of experts.
+    - expert_map: A tensor of shape [num_experts] that maps the expert index
+        from the global space to the local index space of the current
+        expert parallel shard. If the expert is not in the current expert
+        parallel shard, the mapping is set to -1.
+    - pad_sorted_ids: A flag indicating whether the sorted_token_ids length
+      should be padded to a multiple of block_size,
+
+    Returns:
+    - sorted_token_ids: A tensor containing the sorted token indices according
+        to their allocated expert.
+    - expert_ids: A tensor indicating the assigned expert index for each block.
+    - num_tokens_post_padded: The total number of tokens after padding,
+        ensuring divisibility by block_size.
+
+    This function pads the number of tokens that each expert needs to process
+    so that it is divisible by block_size.
+    Padding ensures that during block matrix multiplication, the dimensions
+    align correctly.
+
+    Example:
+    Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
+    block_size = 4, and num_experts = 4:
+    - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
+        with each expert needing to process 3 tokens.
+    - As block_size is 4, we pad 1 token for each expert.
+    - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
+    - Then append padding tokens [12, 12, 12, 12] for each block.
+    - After sorting by expert index, we obtain token_ids
+        [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
+        Tokens 12 are non-existent (padding) and are ignored in
+        the subsequent matrix multiplication.
+    - The padding ensures that the total number of tokens is now divisible
+        by block_size for proper block matrix operations.
+    """
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    if pad_sorted_ids:
+        max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
+    sorted_ids = torch.empty((max_num_tokens_padded, ),
+                             dtype=torch.int32,
+                             device=topk_ids.device)
+    sorted_ids.fill_(topk_ids.numel())
+    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
+    # Expert ids must be zeroed out to prevent index out of bounds error while
+    # mapping global expert ids to local expert ids in expert parallelism.
+    expert_ids = torch.zeros((max_num_m_blocks, ),
+                             dtype=torch.int32,
+                             device=topk_ids.device)
+    num_tokens_post_pad = torch.empty((1),
+                                      dtype=torch.int32,
+                                      device=topk_ids.device)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Only use triton to implement moe_align_block_size
+    '''
+    moe_align_block_size_triton(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_ids,
+        expert_ids,
+        num_tokens_post_pad,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    if expert_map is not None:
+        expert_ids = expert_map[expert_ids]
+
+    return sorted_ids, expert_ids, num_tokens_post_pad
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/utils.py b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/utils.py
new file mode 100644
index 000000000..5634c5dab
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/fused_moe/utils.py
@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+from math import prod
+from typing import List, Optional, Tuple
+
+import torch
+
+from vllm_mlu.model_executor.layers.quantization.utils.fp8_utils import (
+    per_token_group_quant_fp8)
+from vllm.utils import cdiv
+
+
+def _fp8_quantize(
+    A: torch.Tensor,
+    A_scale: Optional[torch.Tensor],
+    block_shape: List[int],
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Perform fp8 quantization on the inputs.  If a block_shape
+    is provided, the output will be blocked.
+    """
+    assert block_shape is not None
+    assert len(block_shape) == 2
+    _, block_k = block_shape[0], block_shape[1]
+    A, A_scale = per_token_group_quant_fp8(A, block_k)
+    assert cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
+    return A, A_scale
+
+
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/layernorm.py b/vllm_mlu/vllm_mlu/model_executor/layers/layernorm.py
new file mode 100644
index 000000000..c6fc579d7
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/layernorm.py
@@ -0,0 +1,35 @@
+from typing import Optional, Tuple, Union
+
+import torch
+
+from vllm.model_executor.layers.layernorm import RMSNorm
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__model_executor__layers__layernorm__RMSNorm__forward_oot(
+    self,
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+) -> Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]:
+    from vllm_mlu import _mlu_ops as mlu_ops
+
+    org_shape = x.shape
+    x = x.view(-1, self.weight.data.shape[0])
+    if out is not None:
+        out = out.view(-1, self.weight.data.shape[0])
+    if residual is not None:
+        residual = residual.view(-1, self.weight.data.shape[0])
+        x = mlu_ops.fused_rms_norm(x, residual, self.weight.data, None, None, self.variance_epsilon, True, out = out)
+    else:
+        x = mlu_ops.fused_rms_norm(x, residual, self.weight.data, None, None, self.variance_epsilon, False, out = out)
+    if out is None and isinstance(x, torch.Tensor):
+        x = x.view(org_shape)
+    return x
+
+MluHijackObject.apply_hijack(
+    RMSNorm,
+    RMSNorm.forward_oot,
+    vllm__model_executor__layers__layernorm__RMSNorm__forward_oot
+)
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/linear.py b/vllm_mlu/vllm_mlu/model_executor/layers/linear.py
new file mode 100644
index 000000000..12b4d4b87
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/linear.py
@@ -0,0 +1,128 @@
+import torch
+from typing import Optional
+from torch.nn.parameter import Parameter
+from vllm.distributed import (get_tensor_model_parallel_rank,
+                              split_tensor_along_last_dim,
+                              tensor_model_parallel_all_gather,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.linear import (
+    UnquantizedLinearMethod,
+    ColumnParallelLinear,
+    RowParallelLinear,
+    WEIGHT_LOADER_V2_SUPPORTED
+)
+from vllm.logger import init_logger
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu import _mlu_ops as mlu_ops
+
+logger = init_logger(__name__)
+
+
+WEIGHT_LOADER_V2_SUPPORTED.extend([
+    "GPTQMluLinearMethod",
+    "AWQMluLinearMethod"
+])
+
+
+def vllm__module_executor__layers__linear__UnquantizedLinearMethod__apply(
+    self,
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    residual: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    beta = 0.0
+    if residual is not None:
+         beta = 1.0
+         residual = residual.view(-1, residual.shape[-1])
+    res_shape = x.shape[0:-1] + (layer.weight.shape[0], )
+    return mlu_ops.matmul(x.reshape(x.numel() // x.shape[-1], x.shape[-1]),
+                          layer.weight,
+                          bias, residual, 'none', 1.0, beta).view(res_shape)
+
+
+def vllm__module_executor__layers__linear__RowParallelLinear__forward(
+    self,
+    input_,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None
+) -> tuple[torch.Tensor, Optional[Parameter]]:
+    if self.input_is_parallel:
+        input_parallel = input_
+    else:
+        tp_rank = get_tensor_model_parallel_rank()
+        splitted_input = split_tensor_along_last_dim(
+            input_, num_partitions=self.tp_size)
+        input_parallel = splitted_input[tp_rank].contiguous()
+
+    # Matrix multiply.
+    assert self.quant_method is not None
+    # Only fuse bias add into GEMM for rank 0 (this ensures that
+    # bias will not get added more than once in TP>1 case)
+    bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
+    residual_ = None if self.tp_rank > 0 else residual
+    kwargs = {'bias': bias_, 'residual': residual_}
+    if smooth_quant_scale is not None:
+        kwargs['input_scale'] = smooth_quant_scale
+    output_parallel = self.quant_method.apply(self,
+                                              input_parallel,
+                                              **kwargs)
+    if self.reduce_results and self.tp_size > 1:
+        output = tensor_model_parallel_all_reduce(output_parallel)
+    else:
+        output = output_parallel
+
+    output_bias = self.bias if self.skip_bias_add else None
+
+    if not self.return_bias:
+        return output
+    return output, output_bias
+
+
+def vllm__module_executor__layers__linear__ColumnParallelLinear__forward(
+    self,
+    input_,
+    smooth_quant_scale: Optional[torch.Tensor] = None
+):
+    bias = self.bias if not self.skip_bias_add else None
+
+    # Matrix multiply.
+    assert self.quant_method is not None
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Add input_scale parameter.
+    '''
+    kwargs = {'bias': bias}
+    if smooth_quant_scale is not None:
+        kwargs['input_scale'] = smooth_quant_scale
+    output_parallel = self.quant_method.apply(self,
+                                              input_,
+                                              **kwargs)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    if self.gather_output:
+        # All-gather across the partitions.
+        output = tensor_model_parallel_all_gather(output_parallel)
+    else:
+        output = output_parallel
+    output_bias = self.bias if self.skip_bias_add else None
+    if not self.return_bias:
+        return output
+    return output, output_bias
+
+
+MluHijackObject.apply_hijack(UnquantizedLinearMethod,
+                             UnquantizedLinearMethod.apply,
+                             vllm__module_executor__layers__linear__UnquantizedLinearMethod__apply)
+MluHijackObject.apply_hijack(RowParallelLinear,
+                             RowParallelLinear.forward,
+                             vllm__module_executor__layers__linear__RowParallelLinear__forward)
+MluHijackObject.apply_hijack(ColumnParallelLinear,
+                             ColumnParallelLinear.forward,
+                             vllm__module_executor__layers__linear__ColumnParallelLinear__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/__init__.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/__init__.py
new file mode 100644
index 000000000..23157ba2d
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/__init__.py
@@ -0,0 +1,34 @@
+from vllm.model_executor.layers.quantization import (
+    QUANTIZATION_METHODS, register_quantization_config
+)
+
+MLU_QUANTIZATION_METHODS= [
+    "smoothquant",
+    "weightonly",
+    "awq_mlu",
+    "gptq_mlu",
+]
+
+
+def register_fake_mlu_quantization_methods():
+    for quant_method in MLU_QUANTIZATION_METHODS:
+        if quant_method not in QUANTIZATION_METHODS:
+            QUANTIZATION_METHODS.append(quant_method)
+
+
+def remove_fake_mlu_quantization_methods():
+    for quant_method in MLU_QUANTIZATION_METHODS:
+        if quant_method in QUANTIZATION_METHODS:
+            QUANTIZATION_METHODS.remove(quant_method)
+
+
+def register_real_mlu_quantization_methods():
+    remove_fake_mlu_quantization_methods()
+    from vllm_mlu.model_executor.layers.quantization.weightonly import WeightOnlyConfig
+    from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantConfig
+    from vllm_mlu.model_executor.layers.quantization.awq_mlu import AWQMluConfig
+    from vllm_mlu.model_executor.layers.quantization.gptq_mlu import GPTQMluConfig
+    register_quantization_config("weightonly")(WeightOnlyConfig)
+    register_quantization_config("smoothquant")(SmoothQuantConfig)
+    register_quantization_config("awq_mlu")(AWQMluConfig)
+    register_quantization_config("gptq_mlu")(GPTQMluConfig)
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/awq_mlu.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/awq_mlu.py
new file mode 100644
index 000000000..39ab6a5f3
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/awq_mlu.py
@@ -0,0 +1,415 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
+from vllm.model_executor.layers.quantization import register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+from vllm.model_executor.parameter import (GroupQuantScaleParameter,
+                                           PackedvLLMParameter)
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.scalar_type import ScalarType, scalar_types
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+from vllm_mlu import _mlu_ops as mlu_ops
+
+logger = init_logger(__name__)
+
+MLU_SUPPORTED_GROUP_SIZES = [64, 128, 256, 512]
+
+# We only support gptq and awq over 300 serials and only support int4 and int8 precision
+def query_mlu_supported_quant_types(has_zp: bool,
+                                       device_capability: Optional[int] = None
+                                       ):
+    if device_capability is None:
+        major, minor = current_platform.get_device_capability()
+        device_capability = major * 10 + minor
+
+    if device_capability < 50:
+        return []
+
+    if has_zp:
+        # AWQ style, unsigned + zero-point
+        return [scalar_types.uint4, scalar_types.uint8]
+    else:
+        # GPTQ style, unsigned + symmetric bias
+        return [scalar_types.uint4b8, scalar_types.uint8b128]
+
+
+def check_mlu_supported(
+        quant_type: ScalarType,
+        group_size: Optional[int],
+        has_zp: bool,
+        device_capability: Optional[int] = None) -> Tuple[bool, Optional[str]]:
+
+    if device_capability is None:
+        major, minor = get_device_capability()
+        device_capability = major * 10 + minor
+
+    supported_types = query_mlu_supported_quant_types(
+        has_zp, device_capability)
+
+    if quant_type not in supported_types:
+        return (False, f"Mlu does not support weight_bits = {quant_type}. "
+                f"Only types = {supported_types} "
+                f"are supported (for group_size = {group_size}, "
+                f"device_capability = {device_capability}, zp = {has_zp}).")
+    if (group_size is None or group_size not in MLU_SUPPORTED_GROUP_SIZES):
+        return (False, f"Mlu does not support group_size = {group_size}. "
+                f"Only group_sizes = {MLU_SUPPORTED_GROUP_SIZES} "
+                "are supported.")
+
+    return True
+
+
+# @register_quantization_config("awq_mlu")
+class AWQMluConfig(QuantizationConfig):
+    """Config class for AWQMlu.
+
+    Reference: https://arxiv.org/abs/2306.00978
+    """
+
+    # num_bits -> type
+    TYPE_MAP = {
+        4: {
+            False: scalar_types.uint4b8,
+            True: scalar_types.uint4,
+        },
+        8: {
+            False: scalar_types.uint8b128,
+            True: scalar_types.uint8,
+        }
+    }
+
+    VERSION = ["gemm"]
+
+    def __init__(
+        self,
+        weight_bits: int,
+        group_size: int,
+        zero_point: bool,
+        lm_head_quantized: bool,
+        version: str = "gemm",
+    ) -> None:
+        self.weight_bits = weight_bits
+        self.group_size = group_size
+        self.zero_point = zero_point
+        self.lm_head_quantized = lm_head_quantized
+        self.pack_factor = 32 // self.weight_bits
+        self.version = version
+        self.support_scale_zeros = False
+
+        if self.weight_bits not in [4, 8]:
+            raise ValueError(
+                "Currently, only 4/8-bit weight quantization is supported for "
+                f"AWQMlu, but got {self.weight_bits} bits.")
+        if self.version not  in self.VERSION:
+            raise ValueError(
+                "Currently, only gemm, gemv version is supported for "
+                f"AWQMlu, but got verion:{self.version}.")
+
+        if self.version in ["gemm"]:
+            self.order_map = {4: [0, 2, 4, 6, 1, 3, 5, 7], 8: [0, 2, 1, 3]}
+            self.reverse_order_map = {4 : [0, 4, 1, 5, 2, 6, 3, 7], 8: [0, 2, 1, 3]}
+        else:
+            self.order_map = {4: [0, 1, 2, 3, 4, 5, 6, 7], 8: [0, 1, 2, 3]}
+            self.reverse_order_map = {4: [0, 1, 2, 3, 4, 5, 6, 7], 8: [0, 1, 2, 3]}
+
+    def __repr__(self) -> str:
+        return (f"AWQMluConfig(weight_bits={self.weight_bits}, "
+                f"group_size={self.group_size}, "
+                f"zero_point={self.zero_point}), "
+                f"lm_head_quantized={self.lm_head_quantized})")
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "awq_mlu"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.half, torch.bfloat16, torch.float32]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 50
+
+    @staticmethod
+    def get_config_filenames() -> List[str]:
+        return ["quant_config.json", "quantize_config.json"]
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "AWQMluConfig":
+        weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])
+        group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])
+        zero_point = cls.get_from_keys(config, ["zero_point"])
+        lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
+                                                 default=False)
+        version = cls.get_from_keys_or(config, ["version"],
+                                                 default="gemm")
+        return cls(weight_bits, group_size, zero_point, lm_head_quantized, version)
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional["AWQMluLinearMethod"]:
+        if (isinstance(layer, LinearBase) or
+            (isinstance(layer, ParallelLMHead) and self.lm_head_quantized)):
+            return AWQMluLinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return ["gelu", "gelu_fast", "gelu_new", "gelu_pytorch_tanh"]
+
+    @classmethod
+    def override_quantization_method(cls, hf_quant_cfg,
+                                     user_quant) -> Optional[str]:
+        can_convert = cls.is_awq_mlu_compatible(hf_quant_cfg)
+        is_valid_user_quant = (user_quant is None or user_quant == "awq"
+                               or user_quant == "awq_mlu")
+
+        if can_convert and is_valid_user_quant:
+            msg = ("The model is convertible to {} during runtime."
+                   " Using {} kernel.".format(cls.get_name(), cls.get_name()))
+            logger.info(msg)
+            return cls.get_name()
+
+        if can_convert and user_quant == "awq":
+            logger.info("Detected that the model can run with awq_mlu"
+                        ", however you specified quantization=awq explicitly,"
+                        " so forcing awq. Use quantization=awq_mlu for"
+                        " faster inference")
+        return None
+
+    @classmethod
+    def is_awq_mlu_compatible(cls, quant_config: Dict[str, Any]):
+        # Extract data from quant config.
+        quant_method = quant_config.get("quant_method", "").lower()
+        num_bits = quant_config.get("bits", None)
+        group_size = quant_config.get("group_size", None)
+        has_zp = quant_config.get("zero_point", None)
+        version = quant_config.get("version", "gemm")
+
+        if quant_method != "awq":
+            return False
+
+        # If we cannot find the info needed in the config, cannot convert.
+        if (num_bits is None or group_size is None or has_zp is None):
+            return False
+
+        if num_bits not in cls.TYPE_MAP:
+            return False
+
+        if version not in cls.VERSION:
+            return False
+
+        return check_mlu_supported(quant_type=cls.TYPE_MAP[num_bits][has_zp],
+                                      group_size=group_size,
+                                      has_zp=has_zp)
+
+class AWQMluLinearMethod(LinearMethodBase):
+    """Linear method for AWQMlu.
+
+    Args:
+        quant_config: The AWQMlu quantization config.
+    """
+
+    def __init__(self, quant_config: AWQMluConfig):
+        self.quant_config = quant_config
+
+    def create_weights(self, layer: torch.nn.Module,
+                       input_size_per_partition: int,
+                       output_partition_sizes: List[int], input_size: int,
+                       output_size: int, params_dtype: torch.dtype,
+                       **extra_weight_attrs):
+        if input_size_per_partition % self.quant_config.group_size != 0:
+            raise ValueError(
+                "The input size is not aligned with the quantized "
+                "weight shape. This can be caused by too large "
+                "tensor parallel size.")
+
+        output_size_per_partition = sum(output_partition_sizes)
+        if output_size_per_partition % self.quant_config.pack_factor != 0:
+            raise ValueError(
+                "The output size is not aligned with the quantized "
+                "weight shape. This can be caused by too large "
+                "tensor parallel size.")
+
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        qweight = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition,
+                output_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader)
+
+        qzeros = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition // self.quant_config.group_size,
+                output_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader)
+
+        scales = GroupQuantScaleParameter(data=torch.empty(
+            input_size_per_partition // self.quant_config.group_size,
+            output_size_per_partition,
+            dtype=params_dtype,
+        ),
+                                          input_dim=0,
+                                          output_dim=1,
+                                          weight_loader=weight_loader)
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("qzeros", qzeros)
+        layer.register_parameter("scales", scales)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        packed_qweight, scale_zeros = self.extract_autoawq(layer)
+        if self.quant_config.zero_point and (not self.quant_config.support_scale_zeros):
+            layer.qweight = torch.nn.Parameter(packed_qweight.contiguous(), requires_grad=False)
+            layer.qzeros = None
+            layer.scales = None
+        else:
+            layer.qweight = torch.nn.Parameter(packed_qweight.contiguous(), requires_grad=False)
+            if scale_zeros is not None:
+                layer.qzeros = torch.nn.Parameter(scale_zeros.contiguous(), requires_grad=False)
+            else:
+                layer.qzeros = None
+            layer.scales = torch.nn.Parameter(layer.scales.data.transpose(0, 1).contiguous(), requires_grad=False)
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None,
+              residual: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.quant_config.zero_point and not self.quant_config.support_scale_zeros:
+            output = mlu_ops.matmul(x, layer.qweight, bias)
+            if residual is not None:
+                output = output + residual
+        else:
+            output = mlu_ops.weight_only_quant_matmul(x,
+                                                  layer.qweight,
+                                                  layer.scales,
+                                                  layer.qzeros,
+                                                  bias,
+                                                  residual,
+                                                  "none",
+                                                  self.quant_config.weight_bits)
+
+        return output
+
+    def extract_autoawq(self, layer: torch.nn.Module):
+        qweight = layer.qweight.data
+        qzeros = layer.qzeros.data
+        scales = layer.scales.data
+        bits = self.quant_config.weight_bits
+        group_size = self.quant_config.group_size
+
+        # Unpack the qweight and qzeros tensors
+        iweight, izeros = self.unpack_awq_int32_into_int8(qweight, qzeros, bits)
+        # Reverse the order of the iweight and izeros tensors
+        iweight, izeros = self.reverse_awq_order(iweight, izeros, bits)
+    
+        # overflow checks
+        iweight = torch.bitwise_and(iweight, (2**bits) - 1)
+        if izeros is not None:
+            izeros = torch.bitwise_and(izeros, (2**bits) - 1)
+
+        if self.quant_config.zero_point and (not self.quant_config.support_scale_zeros):
+            scales = scales.repeat_interleave(group_size, dim=0)
+            if izeros is not None:
+                izeros = izeros.repeat_interleave(group_size, dim=0)
+                fweight = (iweight - izeros) * scales
+            else:
+                fweight = iweight * scales
+            # transpose [ci, co] -> [co, ci]
+            fweight = fweight.transpose(0, 1)
+            
+            return fweight, None
+
+        if self.quant_config.zero_point and self.quant_config.support_scale_zeros and izeros is not None:
+            scale_zeros = izeros.to(scales.dtype) * -1 * scales
+            # transpose [ci, co] -> [co, ci]
+            scale_zeros = scale_zeros.transpose(0, 1)
+        else:
+            scale_zeros = None
+
+        # transpose [ci, co] -> [co, ci]
+        iweight = iweight.to(torch.int8).transpose(0, 1)
+
+        if bits == 4:
+            higher_bit_tensor = iweight[:, 1::2]
+            lower_bit_tensor = iweight[:, 0::2]
+            packed_qweight = self.combine_low_bits(higher_bit_tensor, lower_bit_tensor)
+        else:
+            packed_qweight = iweight
+
+        return packed_qweight, scale_zeros
+
+    def unpack_awq_int32_into_int8(self, qweight: torch.Tensor, qzeros: torch.Tensor, bits: int):
+        shifts = torch.arange(0, 32, bits, device=qweight.device)
+        dtype = torch.int16 if bits == 8 else torch.int8
+        # unpacking columnwise
+        iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to(dtype)
+        iweights = iweights.view(iweights.shape[0], -1)
+        if not self.quant_config.zero_point or self.quant_config.support_scale_zeros:
+            iweights = torch.bitwise_and(iweights - 2**(bits - 1), (2 ** bits) - 1)
+
+        # unpacking columnwise
+        if qzeros is not None:
+            izeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to(dtype)
+            izeros = izeros.view(izeros.shape[0], -1)
+            if not self.quant_config.zero_point:
+                izeros = torch.bitwise_and(izeros - 2**(bits - 1), (2 ** bits) - 1)
+        else:
+            izeros = None
+
+        return iweights, izeros
+
+    def reverse_awq_order(self, iweights: torch.Tensor, izeros: torch.Tensor, bits: int):
+        reverse_order_tensor = torch.arange(iweights.shape[-1], dtype=torch.int32, device=iweights.device)
+        reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
+        reverse_order_tensor = reverse_order_tensor[:, self.quant_config.reverse_order_map[bits]]
+        reverse_order_tensor = reverse_order_tensor.view(-1)
+
+        rweights = iweights[:, reverse_order_tensor]
+        if izeros is not None:
+            rzeros = izeros[:, reverse_order_tensor]
+
+        return rweights, rzeros
+
+    def combine_low_bits(self, tensor_a, tensor_b):
+        """
+        Combine the lower 4 bits of two int8 tensors into a new int8 tensor.
+
+        Args:
+        tensor_a (torch.Tensor): First tensor of type int8.
+        tensor_b (torch.Tensor): Second tensor of type int8.
+
+        Returns:
+        torch.Tensor: New tensor of type int8, combining lower 4 bits of tensor_a and tensor_b.
+        """
+        # 确保输入是 int8 类型
+        if tensor_a.dtype != torch.int8 or tensor_b.dtype != torch.int8:
+            raise ValueError("Both tensors must be of int8 type.")
+
+        # 提取每个 tensor 的低4位
+        low_bits_a = torch.bitwise_and(tensor_a, 0x0F)  # 保留 tensor_a 的低4位
+        low_bits_b = torch.bitwise_and(tensor_b, 0x0F)  # 保留 tensor_b 的低4位
+
+        # 将 tensor_a 的低4位左移4位
+        shifted_low_bits_a = low_bits_a << 4
+
+        # 组合两个 tensor 的低4位
+        combined = torch.bitwise_or(shifted_low_bits_a, low_bits_b)
+
+        return combined
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/fp8.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/fp8.py
new file mode 100644
index 000000000..3ff56f3af
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/fp8.py
@@ -0,0 +1,561 @@
+import torch
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+from typing import Any, Dict, List, Optional, Callable
+from vllm import envs
+from vllm import _custom_ops as ops
+from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.fp8 import (
+    ACTIVATION_SCHEMES, Fp8Config, Fp8LinearMethod, Fp8MoEMethod)
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
+    apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin)
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    convert_to_channelwise, cutlass_block_fp8_supported, cutlass_fp8_supported,
+    normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale,
+    maybe_create_device_identity, Fp8LinearOp)
+from vllm.model_executor.parameter import (
+    BlockQuantScaleParameter, ChannelQuantScaleParameter,
+    ModelWeightParameter, PerTensorScaleParameter)
+from vllm.platforms import current_platform
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+import vllm_mlu._mlu_ops as mlu_ops
+
+logger = init_logger(__name__)
+
+
+def vllm__model_executor__layers__quantization__fp8__Fp8Config__get_min_capability(cls) -> int:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: For Cambricon MLU, the minimum capability is 60.
+    '''
+    return 60
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__model_executor__layers__quantization__fp8__Fp8Config____init__(
+    self,
+    is_checkpoint_fp8_serialized: bool = False,
+    activation_scheme: str = "dynamic",
+    ignored_layers: Optional[List[str]] = None,
+    weight_block_size: Optional[List[int]] = None,
+    activation_quant_method: Optional[str] = None,
+    weight_quant_method: Optional[str] = None,
+) -> None:
+    self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
+    if is_checkpoint_fp8_serialized:
+        logger.warning("Detected fp8 checkpoint. Please note that the "
+                        "format is experimental and subject to change.")
+    if activation_scheme not in ACTIVATION_SCHEMES:
+        raise ValueError(
+            f"Unsupported activation scheme {activation_scheme}")
+    self.activation_scheme = activation_scheme
+    self.ignored_layers = ignored_layers or []
+    if weight_block_size is not None:
+        if not is_checkpoint_fp8_serialized:
+            raise ValueError(
+                "The block-wise quantization only supports fp8-serialized "
+                "checkpoint for now.")
+        if len(weight_block_size) != 2:
+            raise ValueError(
+                "The quantization block size of weight must have 2 "
+                f"dimensions, but got {len(weight_block_size)} dimensions")
+        if activation_scheme != "dynamic":
+            raise ValueError("The block-wise quantization only supports "
+                                "dynamic activation scheme for now, but got "
+                                f"{activation_scheme} activation scheme.")
+    self.weight_block_size = weight_block_size
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Add class members activation_quant_method and weight_quant_method to
+    indicate the granularity of quantization.
+    '''
+    self.activation_quant_method = activation_quant_method
+    self.weight_quant_method = weight_quant_method
+
+    assert (self.weight_block_size or \
+        self.activation_quant_method == "per_token" and self.weight_quant_method == "per_channel"
+        and self.activation_scheme == "dynamic"), "Only support block-wise quantization, or "\
+            "input dynamic per-token output per-channel quantization yet."
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+@classmethod
+def vllm__model_executor__layers__quantization__fp8__Fp8Config__from_config(
+    cls, config: Dict[str, Any]
+) -> "Fp8Config":
+    quant_method = cls.get_from_keys(config, ["quant_method"])
+    is_checkpoint_fp8_serialized = ("fp8" in quant_method)
+    activation_scheme = cls.get_from_keys(config, ["activation_scheme"])
+    ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)
+    weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"],
+                                                None)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Add config members activation_quant_method and weight_quant_method to
+    indicate the granularity of quantization.
+    '''
+    activation_quant_method = cls.get_from_keys_or(config, ["activation_quant_method"], 'per_token')
+    weight_quant_method = cls.get_from_keys_or(config, ["weight_quant_method"], None)
+    return cls(is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
+                activation_scheme=activation_scheme,
+                ignored_layers=ignored_layers,
+                weight_block_size=weight_block_size,
+                activation_quant_method=activation_quant_method,
+                weight_quant_method=weight_quant_method)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__model_executor__layers__quantization__fp8__Fp8LinearMethod__create_weights(
+    self,
+    layer: torch.nn.Module,
+    input_size_per_partition: int,
+    output_partition_sizes: List[int],
+    input_size: int,
+    output_size: int,
+    params_dtype: torch.dtype,
+    **extra_weight_attrs,
+):
+    maybe_create_device_identity()
+
+    output_size_per_partition = sum(output_partition_sizes)
+    weight_loader = extra_weight_attrs.get("weight_loader")
+
+    if self.block_quant:
+        tp_size = get_tensor_model_parallel_world_size()
+        assert self.quant_config.weight_block_size is not None
+        block_n, block_k = (
+            self.quant_config.weight_block_size[0],
+            self.quant_config.weight_block_size[1],
+        )
+        # Required by row parallel
+        if (tp_size > 1
+                and input_size // input_size_per_partition == tp_size
+                and input_size_per_partition % block_k != 0):
+            raise ValueError(
+                f"Weight input_size_per_partition = "
+                f"{input_size_per_partition} is not divisible by "
+                f"weight quantization block_k = {block_k}.")
+        # Required by column parallel or enabling merged weights
+        if (tp_size > 1 and output_size // output_size_per_partition
+                == tp_size) or len(output_partition_sizes) > 1:
+            for output_partition_size in output_partition_sizes:
+                if output_partition_size % block_n != 0:
+                    raise ValueError(
+                        f"Weight output_partition_size = "
+                        f"{output_partition_size} is not divisible by "
+                        f"weight quantization block_n = {block_n}.")
+
+    layer.logical_widths = output_partition_sizes
+
+    layer.input_size_per_partition = input_size_per_partition
+    layer.output_size_per_partition = output_size_per_partition
+    layer.orig_dtype = params_dtype
+
+    # WEIGHT
+    weight_dtype = (torch.float8_e4m3fn
+                    if self.quant_config.is_checkpoint_fp8_serialized else
+                    params_dtype)
+
+    weight = ModelWeightParameter(data=torch.empty(
+        output_size_per_partition,
+        input_size_per_partition,
+        dtype=weight_dtype),
+                                    input_dim=1,
+                                    output_dim=0,
+                                    weight_loader=weight_loader)
+    layer.register_parameter("weight", weight)
+
+    # If checkpoint is serialized fp8, load them.
+    # Otherwise, wait until process_weights_after_loading.
+    if self.quant_config.is_checkpoint_fp8_serialized:
+        # WEIGHT SCALE
+        if not self.block_quant:
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: Support weight per channel quantization.
+            '''
+            if self.weight_per_channel:
+                scale = ChannelQuantScaleParameter(
+                    data=torch.empty(sum(output_partition_sizes), dtype=torch.float32),
+                    output_dim=0,
+                    weight_loader=weight_loader,
+                )
+            else:
+                scale = PerTensorScaleParameter(
+                    data=torch.empty(len(output_partition_sizes),
+                                        dtype=torch.float32),
+                    weight_loader=weight_loader,
+                )
+            scale[:] = torch.finfo(torch.float32).min
+            layer.register_parameter("weight_scale", scale)
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+        else:
+            assert self.quant_config.activation_scheme == "dynamic"
+            scale = BlockQuantScaleParameter(
+                data=torch.empty(
+                    (output_size_per_partition + block_n - 1) // block_n,
+                    (input_size_per_partition + block_k - 1) // block_k,
+                    dtype=torch.float32,
+                ),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=weight_loader,
+            )
+            scale[:] = torch.finfo(torch.float32).min
+            # The weight_scale_inv name is intentional for deepseekv3
+            layer.register_parameter("weight_scale_inv", scale)
+
+        # INPUT ACTIVATION SCALE
+        if self.quant_config.activation_scheme == "static":
+            scale = PerTensorScaleParameter(data=torch.empty(
+                len(output_partition_sizes), dtype=torch.float32),
+                                            weight_loader=weight_loader)
+
+            scale[:] = torch.finfo(torch.float32).min
+            layer.register_parameter("input_scale", scale)
+        else:
+            layer.register_parameter("input_scale", None)
+
+
+def vllm__model_executor__layers__quantization__fp8__Fp8LinearMethod____init__(
+    self,
+    quant_config: Fp8Config
+):
+    self.quant_config = quant_config
+    self.cutlass_fp8_supported = cutlass_fp8_supported()
+    self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
+
+    # For GPUs that lack FP8 hardware support, we can leverage the Marlin
+    # kernel for fast weight-only FP8 quantization
+    self.use_marlin = (not current_platform.has_device_capability(89)
+                        or envs.VLLM_TEST_FORCE_FP8_MARLIN)
+    # Disable marlin for rocm
+    if current_platform.is_rocm():
+        self.use_marlin = False
+
+    self.block_quant = self.quant_config.weight_block_size is not None
+    if self.block_quant:
+        # Marlin doesn't support block-wise fp8
+        self.use_marlin = False
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Add config members activation_quant_method and weight_quant_method to
+    indicate the granularity of quantization.
+    '''
+    self.weight_per_channel = (self.quant_config.weight_quant_method == 'per_channel')
+    self.activation_per_token = (self.quant_config.activation_quant_method == 'per_token')
+    if self.weight_per_channel and  self.activation_per_token:
+        self.use_marlin = False
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.fp8_linear = Fp8LinearOp(
+        # Default to using per_token quantization if cutlass is supported
+        use_per_token_if_dynamic=cutlass_fp8_supported())
+
+
+def vllm__model_executor__layers__quantization__fp8__Fp8LinearMethod__process_weights_after_loading(
+    self,
+    layer: Module,
+) -> None:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: For dynamic activation and channel-wise weight quantization,
+    additional processing is not needed.
+    '''
+    if self.quant_config.is_checkpoint_fp8_serialized and \
+        self.weight_per_channel and \
+        self.quant_config.activation_scheme == "dynamic":
+        return
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    # TODO(rob): refactor block quant into separate class.
+    if self.block_quant:
+        assert self.quant_config.activation_scheme == "dynamic"
+        if current_platform.is_fp8_fnuz():
+            weight, weight_scale_inv, _ = \
+                normalize_e4m3fn_to_e4m3fnuz(
+                    weight=layer.weight,
+                    weight_scale=layer.weight_scale_inv)
+        else:
+            weight = layer.weight.data
+            weight_scale_inv = layer.weight_scale_inv.data
+
+        weight = self._maybe_pad_weight(weight)
+
+        # Torch.compile cannot use Parameter subclasses.
+        layer.weight = Parameter(weight, requires_grad=False)
+        layer.weight_scale_inv = Parameter(weight_scale_inv,
+                                            requires_grad=False)
+        return
+
+    # If checkpoint not serialized fp8, quantize the weights.
+    if not self.quant_config.is_checkpoint_fp8_serialized:
+        qweight, weight_scale = ops.scaled_fp8_quant(layer.weight,
+                                                        scale=None)
+
+        # If using marlin (w8a16), kernel uses channelwise weights,
+        # so extend the weight scales to be channelwise.
+        if self.use_marlin:
+            assert weight_scale.numel() == 1
+            weight_scale = convert_to_channelwise(
+                weight_scale.expand(len(layer.logical_widths)),
+                layer.logical_widths)
+
+        # Update the layer with the new values.
+        layer.weight = Parameter(qweight.t(), requires_grad=False)
+        layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+        layer.input_scale = None
+
+    # If checkpoint is fp8, handle that there are N scales for N
+    # shards in a fused module
+    else:
+        layer.weight_scale = torch.nn.Parameter(layer.weight_scale.data,
+                                                requires_grad=False)
+        if self.quant_config.activation_scheme == "static":
+            layer.input_scale = torch.nn.Parameter(layer.input_scale.data,
+                                                    requires_grad=False)
+        # If using marlin (w8a16), kernel uses channelwise weights,
+        # so extend the weight scales to be channelwise.
+        if self.use_marlin:
+            weight = layer.weight
+            weight_scale = convert_to_channelwise(layer.weight_scale,
+                                                    layer.logical_widths)
+
+        # If using w8a8, torch._scaled_mm needs per tensor, so
+        # requantize the logical shards as a single weight.
+        else:
+            # Dequant -> Quant with max scale so we can run per tensor.
+            weight = layer.weight
+            weight_scale = layer.weight_scale
+
+            if current_platform.is_fp8_fnuz():
+                weight, weight_scale, input_scale = \
+                    normalize_e4m3fn_to_e4m3fnuz(
+                        weight=weight,
+                        weight_scale=weight_scale,
+                        input_scale=layer.input_scale)
+                if input_scale is not None:
+                    layer.input_scale = Parameter(input_scale,
+                                                    requires_grad=False)
+
+            weight_scale, weight = requantize_with_max_scale(
+                weight=weight,
+                weight_scale=weight_scale,
+                logical_widths=layer.logical_widths,
+            )
+
+        weight = self._maybe_pad_weight(weight)
+        # Update layer with new values.
+        layer.weight = Parameter(weight.t(), requires_grad=False)
+        layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+        if self.quant_config.activation_scheme == "static":
+            layer.input_scale = Parameter(layer.input_scale.max(),
+                                            requires_grad=False)
+
+    if self.use_marlin:
+        prepare_fp8_layer_for_marlin(layer)
+        # Activations not quantized for marlin.
+        del layer.input_scale
+
+def vllm__model_executor__layers__quantization__fp8__Fp8LinearMethod__apply(
+    self,
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    residual: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    assert residual is None, "residual is not supported yet."
+
+    if self.use_marlin:
+        return apply_fp8_marlin_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            workspace=layer.workspace,
+            size_n=layer.output_size_per_partition,
+            size_k=layer.input_size_per_partition,
+            bias=bias)
+
+    # Note: lazy import to avoid triton import error.
+    from vllm_mlu.model_executor.layers.quantization.utils.fp8_utils import (
+        apply_w8a8_block_fp8_linear)
+    if self.block_quant:
+        assert self.quant_config.weight_block_size is not None
+        return apply_w8a8_block_fp8_linear(
+            input=x,
+            weight=layer.weight,
+            block_size=self.quant_config.weight_block_size,
+            weight_scale=layer.weight_scale_inv,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
+        )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Use activation per token quantization based on quantization config.
+    '''
+    return self.fp8_linear.apply(
+        input=x,
+        weight=layer.weight,
+        weight_scale=layer.weight_scale,
+        input_scale=layer.input_scale,
+        bias=bias,
+        use_per_token_if_dynamic=self.activation_per_token,
+        weight_per_channel=self.weight_per_channel,
+        activation_per_token=self.activation_per_token)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+def vllm__model_executor__layers__quantization__fp8__Fp8MoEMethod__apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+        top_k: int,
+        renormalize: bool,
+        use_grouped_topk: bool = False,
+        topk_group: Optional[int] = None,
+        num_expert_group: Optional[int] = None,
+        global_num_experts: int = -1,
+        expert_map: Optional[torch.Tensor] = None,
+        custom_routing_function: Optional[Callable] = None,
+        scoring_func: str = "softmax",
+        e_score_correction_bias: Optional[torch.Tensor] = None,
+        apply_router_weight_on_input: bool = False,
+        activation: str = "silu",
+    ) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Use moe_softmax_topk and moe_sigmoid_topk of mlu_ops to implement FusedMoE.select_experts
+    '''
+    from vllm_mlu.model_executor.layers.fused_moe.fused_moe import fused_experts
+    routed_scaling_factor = 1.
+    if scoring_func == "softmax":
+        topk_weights, topk_ids = mlu_ops.moe_softmax_topk(router_logits,
+                                                            top_k,
+                                                            renormalize,
+                                                            num_expert_group,
+                                                            topk_group,
+                                                            route_scale=routed_scaling_factor)
+    elif scoring_func == "sigmoid":
+        topk_weights, topk_ids = mlu_ops.moe_sigmoid_topk(router_logits, top_k, renormalize,
+                                                            num_expert_group, topk_group,
+                                                            routed_scaling_factor,
+                                                            e_score_correction_bias)
+    else:
+        raise ValueError(f"Unsupported scoring function: {scoring_func}")
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    return fused_experts(
+        x,
+        layer.w13_weight,
+        layer.w2_weight,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=True,
+        global_num_experts=global_num_experts,
+        apply_router_weight_on_input=apply_router_weight_on_input,
+        expert_map=expert_map,
+        w1_scale=(layer.w13_weight_scale_inv
+                    if self.block_quant else layer.w13_weight_scale),
+        w2_scale=(layer.w2_weight_scale_inv
+                    if self.block_quant else layer.w2_weight_scale),
+        a1_scale=layer.w13_input_scale,
+        a2_scale=layer.w2_input_scale,
+        block_shape=self.quant_config.weight_block_size,
+        allow_deep_gemm=self.allow_deep_gemm,
+    )
+
+
+
+MluHijackObject.apply_hijack(
+    Fp8LinearMethod,
+    Fp8LinearMethod.apply,
+    vllm__model_executor__layers__quantization__fp8__Fp8LinearMethod__apply
+)
+MluHijackObject.apply_hijack(
+    Fp8Config,
+    Fp8Config.get_min_capability,
+    vllm__model_executor__layers__quantization__fp8__Fp8Config__get_min_capability
+)
+MluHijackObject.apply_hijack(
+    Fp8Config,
+    Fp8Config.__init__,
+    vllm__model_executor__layers__quantization__fp8__Fp8Config____init__
+)
+MluHijackObject.apply_hijack(
+    Fp8Config,
+    Fp8Config.from_config,
+    vllm__model_executor__layers__quantization__fp8__Fp8Config__from_config
+)
+MluHijackObject.apply_hijack(
+    Fp8LinearMethod,
+    Fp8LinearMethod.create_weights,
+    vllm__model_executor__layers__quantization__fp8__Fp8LinearMethod__create_weights
+)
+MluHijackObject.apply_hijack(
+    Fp8LinearMethod,
+    Fp8LinearMethod.__init__,
+    vllm__model_executor__layers__quantization__fp8__Fp8LinearMethod____init__
+)
+MluHijackObject.apply_hijack(
+    Fp8LinearMethod,
+    Fp8LinearMethod.process_weights_after_loading,
+    vllm__model_executor__layers__quantization__fp8__Fp8LinearMethod__process_weights_after_loading
+)
+MluHijackObject.apply_hijack(
+    Fp8MoEMethod,
+    Fp8MoEMethod.apply,
+    vllm__model_executor__layers__quantization__fp8__Fp8MoEMethod__apply
+)
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/gptq_mlu.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/gptq_mlu.py
new file mode 100644
index 000000000..10e3fdc8f
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/gptq_mlu.py
@@ -0,0 +1,439 @@
+from fractions import Fraction
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
+from vllm.model_executor.layers.quantization import register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
+                                           GroupQuantScaleParameter,
+                                           PackedColumnParameter,
+                                           PackedvLLMParameter,
+                                           RowvLLMParameter)
+from vllm.scalar_type import ScalarType, scalar_types
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+from vllm_mlu import _mlu_ops as mlu_ops
+
+logger = init_logger(__name__)
+
+MLU_SUPPORTED_GROUP_SIZES = [64, 128, 256, 512]
+
+# We only support gptq and awq over 300 serials and only support int4 and int8 precision
+def query_mlu_supported_quant_types(has_zp: bool,
+                                       device_capability: Optional[int] = None
+                                       ):
+    if device_capability is None:
+        major, minor = current_platform.get_device_capability()
+        device_capability = major * 10 + minor
+
+    if device_capability < 50:
+        return []
+
+    if has_zp:
+        # AWQ style, unsigned + zero-point
+        return [scalar_types.uint4, scalar_types.uint8]
+    else:
+        # GPTQ style, unsigned + symmetric bias
+        return [scalar_types.uint4b8, scalar_types.uint8b128]
+
+
+def check_mlu_supported(
+        quant_type: ScalarType,
+        group_size: Optional[int],
+        has_zp: bool,
+        device_capability: Optional[int] = None) -> Tuple[bool, Optional[str]]:
+
+    if device_capability is None:
+        major, minor = get_device_capability()
+        device_capability = major * 10 + minor
+
+    supported_types = query_mlu_supported_quant_types(
+        has_zp, device_capability)
+
+    if quant_type not in supported_types:
+        return (False, f"Mlu does not support weight_bits = {quant_type}. "
+                f"Only types = {supported_types} "
+                f"are supported (for group_size = {group_size}, "
+                f"device_capability = {device_capability}, zp = {has_zp}).")
+    if (group_size is None or group_size not in MLU_SUPPORTED_GROUP_SIZES):
+        return (False, f"Mlu does not support group_size = {group_size}. "
+                f"Only group_sizes = {MLU_SUPPORTED_GROUP_SIZES} "
+                "are supported.")
+
+    return True
+
+
+# @register_quantization_config("gptq_mlu")
+class GPTQMluConfig(QuantizationConfig):
+    """Config class for GPTQMlu.
+
+    Reference: https://arxiv.org/abs/2210.17323
+    """
+
+    # (num_bits, is_sym) -> quant_type
+    TYPE_MAP = {
+        (4, True): scalar_types.uint4b8,
+        (8, True): scalar_types.uint8b128,
+    }
+
+    def __init__(
+        self,
+        weight_bits: int,
+        group_size: int,
+        desc_act: bool,
+        is_sym: bool,
+        lm_head_quantized: bool,
+    ) -> None:
+        self.weight_bits = weight_bits
+        self.group_size = group_size
+        self.desc_act = desc_act
+        self.is_sym = is_sym
+        self.lm_head_quantized = lm_head_quantized
+        self.pack_factor = Fraction(32, self.weight_bits)
+        self.support_scale_zeros = False
+
+        if self.weight_bits not in [4, 8]:
+            raise ValueError(
+                "Currently, only 4/8-bit weight quantization is "
+                f"supported for GPTQMlu, but got {self.weight_bits} bits.")
+
+    def __repr__(self) -> str:
+        return (f"GPTQMluConfig(weight_bits={self.weight_bits}, "
+                f"group_size={self.group_size}, "
+                f"desc_act={self.desc_act}),"
+                f"lm_head_quantized={self.lm_head_quantized}")
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "gptq_mlu"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.half, torch.bfloat16, torch.float32]
+
+    @classmethod
+    # Need to figure it out
+    def get_min_capability(cls) -> int:
+        return 50
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        return ["quant_config.json", "quantize_config.json"]
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "GPTQMluConfig":
+        weight_bits = cls.get_from_keys(config, ["bits"])
+        group_size = cls.get_from_keys(config, ["group_size"])
+        desc_act = cls.get_from_keys(config, ["desc_act"])
+        is_sym = cls.get_from_keys(config, ["sym"])
+        lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
+                                                 default=False)
+        return cls(weight_bits, group_size, desc_act, is_sym, lm_head_quantized)
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional["GPTQMluLinearMethod"]:
+        if (isinstance(layer, LinearBase) or
+            (isinstance(layer, ParallelLMHead) and self.lm_head_quantized)):
+            return GPTQMluLinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return ["gelu", "gelu_fast", "gelu_new", "gelu_pytorch_tanh"]
+
+    @classmethod
+    def is_gptq_mlu_compatible(cls, quant_config: Dict[str, Any]):
+        # Extract data from quant config.
+        quant_method = quant_config.get("quant_method", "").lower()
+        num_bits = quant_config.get("bits", None)
+        group_size = quant_config.get("group_size", None)
+        sym = quant_config.get("sym", None)
+        desc_act = quant_config.get("desc_act", None)
+
+        if quant_method != "gptq":
+            return False
+
+        # If we cannot find the info needed in the config, cannot convert.
+        if (num_bits is None or group_size is None or sym is None
+                or desc_act is None):
+            return False
+
+        if (num_bits, sym) not in cls.TYPE_MAP:
+            return False
+
+        return check_mlu_supported(quant_type=cls.TYPE_MAP[(num_bits, sym)],
+                                      group_size=group_size, has_zp=False)
+
+    @classmethod
+    def override_quantization_method(cls, hf_quant_cfg,
+                                     user_quant) -> Optional[str]:
+        can_convert = cls.is_gptq_mlu_compatible(hf_quant_cfg)
+
+        is_valid_user_quant = (user_quant is None or user_quant == "gptq"
+                               or user_quant == "gptq_mlu")
+
+        if can_convert and is_valid_user_quant:
+            msg = ("The model is convertible to {} during runtime."
+                   " Using {} kernel.".format(cls.get_name(), cls.get_name()))
+            logger.info(msg)
+            return cls.get_name()
+
+        return None
+
+class GPTQMluLinearMethod(LinearMethodBase):
+    """Linear method for GPTQMlu.
+
+    Args:
+        quant_config: The GPTQMlu quantization config.
+    """
+
+    def __init__(self, quant_config: GPTQMluConfig):
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        del output_size  # Unused.
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        if input_size_per_partition % self.quant_config.group_size != 0:
+            raise ValueError(
+                "The input size is not aligned with the quantized "
+                "weight shape. This can be caused by too large "
+                "tensor parallel size.")
+        output_size_per_partition = sum(output_partition_sizes)
+        if (output_size_per_partition % self.quant_config.pack_factor.numerator
+                != 0):
+            raise ValueError(
+                "The output size is not aligned with the quantized "
+                "weight shape. This can be caused by too large "
+                "tensor parallel size.")
+
+        if self.quant_config.group_size != -1:
+            group_size = self.quant_config.group_size
+        else:
+            group_size = input_size
+
+        scale_and_zero_size = input_size // group_size
+        scale_and_zero_input_dim = None
+        if (input_size != input_size_per_partition) and (self.quant_config.group_size !=
+                                                         -1) and (not self.quant_config.desc_act):
+            scale_and_zero_size = input_size_per_partition // group_size
+            scale_and_zero_input_dim = 0
+
+        qweight = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition // self.quant_config.pack_factor,
+                output_size_per_partition,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=0,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader)
+
+        g_idx = RowvLLMParameter(data=torch.tensor(
+            [
+                i // self.quant_config.group_size
+                for i in range(input_size_per_partition)
+            ],
+            dtype=torch.int32,
+        ),
+                                 input_dim=0,
+                                 weight_loader=weight_loader)
+        qzeros_args = {
+            "data":
+            torch.empty(
+                scale_and_zero_size,
+                output_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            "weight_loader":
+            weight_loader
+        }
+        weight_scale_args = {
+            "data":
+            torch.empty(
+                scale_and_zero_size,
+                output_size_per_partition,
+                dtype=params_dtype,
+            ),
+            "weight_loader":
+            weight_loader
+        }
+        if scale_and_zero_input_dim is None:
+            scales = ChannelQuantScaleParameter(output_dim=1,
+                                                **weight_scale_args)
+            qzeros = PackedColumnParameter(
+                output_dim=1,
+                packed_dim=1,
+                packed_factor=self.quant_config.pack_factor,
+                **qzeros_args)
+
+        else:
+            scales = GroupQuantScaleParameter(output_dim=1,
+                                              input_dim=0,
+                                              **weight_scale_args)
+            qzeros = PackedvLLMParameter(
+                input_dim=0,
+                output_dim=1,
+                packed_dim=1,
+                packed_factor=self.quant_config.pack_factor,
+                **qzeros_args)
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("g_idx", g_idx)
+        layer.register_parameter("qzeros", qzeros)
+        layer.register_parameter("scales", scales)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        self.device = layer.qweight.data.device
+        if self.quant_config.desc_act:
+            g_idx_list = layer.g_idx.data.tolist()
+            g_idx_unique = list(dict.fromkeys(g_idx_list))
+            g_idx = torch.tensor(g_idx_unique, dtype=layer.g_idx.data.dtype, device=self.device)
+            scales = layer.scales.data[g_idx]
+        else:
+            scales = layer.scales.data
+
+        packed_qweight, scale_zeros = self.extract_autogptq(layer, scales)
+        if (not self.quant_config.is_sym) and (not self.quant_config.support_scale_zeros):
+            layer.qweight = torch.nn.Parameter(packed_qweight.contiguous(), requires_grad=False)
+            layer.qzeros = None
+            layer.scales = None
+        else:
+            layer.qweight = torch.nn.Parameter(packed_qweight.contiguous(), requires_grad=False)
+            if scale_zeros is not None:
+                layer.qzeros = torch.nn.Parameter(scale_zeros.contiguous(), requires_grad=False)
+            else:
+                layer.qzeros = None
+            layer.scales = torch.nn.Parameter(scales.transpose(0, 1).contiguous(), requires_grad=False)
+
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None,
+              residual: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if (not self.quant_config.is_sym) and (not self.quant_config.support_scale_zeros):
+            output = mlu_ops.matmul(x, layer.qweight, bias)
+            if residual is not None:
+                output = output + residual
+        else:
+            output = mlu_ops.weight_only_quant_matmul(x,
+                                                  layer.qweight,
+                                                  layer.scales,
+                                                  layer.qzeros,
+                                                  bias,
+                                                  residual,
+                                                  "none",
+                                                  self.quant_config.weight_bits)
+
+        return output
+
+
+    def extract_autogptq(self, layer: torch.nn.Module, scales: torch.Tensor):
+        bits = self.quant_config.weight_bits
+        group_size = self.quant_config.group_size
+        # Unpack the qweight and qzeros tensors
+        iweight = self.unpack_gptq_qweight_int32_into_int8(layer.qweight.data, bits)
+        izeros = self.unpack_gptq_qzeros_int32_into_int8(layer.qzeros.data, bits)
+
+        # overflow checks
+        iweight = torch.bitwise_and(iweight, (2**bits) - 1)
+        if izeros is not None:
+            izeros = torch.bitwise_and(izeros, (2**bits) - 1)
+
+        if not self.quant_config.is_sym and (not self.quant_config.support_scale_zeros):
+            scales = scales.repeat_interleave(group_size, dim=0)
+            if izeros is not None:
+                izeros = izeros.repeat_interleave(group_size, dim=0)
+                fweight = (iweight - izeros) * scales
+            else:
+                fweight = iweight * scales
+            # transpose [ci, co] -> [co, ci]
+            fweight = fweight.transpose(0, 1)
+            
+            return fweight, None
+
+        if not self.quant_config.is_sym and self.quant_config.support_scale_zeros and izeros is not None:
+            scale_zeros = izeros.to(scales.dtype) * -1 * scales
+            # transpose [ci, co] -> [co, ci]
+            scale_zeros = scale_zeros.transpose(0, 1)
+        else:
+            scale_zeros = None
+
+        # transpose [ci, co] -> [co, ci]
+        iweight = iweight.to(torch.int8).transpose(0, 1)
+
+        if bits == 4:
+            higher_bit_tensor = iweight[:, 1::2]
+            lower_bit_tensor = iweight[:, 0::2]
+            packed_qweight = self.combine_low_bits(higher_bit_tensor, lower_bit_tensor)
+        else:
+            packed_qweight = iweight
+
+        return packed_qweight, scale_zeros
+
+    def unpack_gptq_qweight_int32_into_int8(self, qweight: torch.Tensor, bits: int):
+        shifts = torch.arange(0, 32, bits, device=qweight.device)
+        dtype = torch.int16 if bits == 8 else torch.int8
+        # unpacking columnwise
+        iweight = torch.bitwise_right_shift(qweight[:, None, :], shifts[None, :, None]).to(dtype)
+        iweight = iweight.view(-1, iweight.shape[-1])
+        # minus 2**(bit-1)
+        if self.quant_config.is_sym or self.quant_config.support_scale_zeros:
+            iweight = torch.bitwise_and(iweight - 2**(bits - 1), (2 ** bits) - 1)
+
+        return iweight
+
+    def unpack_gptq_qzeros_int32_into_int8(self, qzeros: torch.Tensor, bits: int):
+        shifts = torch.arange(0, 32, bits, device=qzeros.device)
+        dtype = torch.int16 if bits == 8 else torch.int8
+        # unpacking columnwise
+        izeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to(dtype)
+        izeros = izeros.view(izeros.shape[0], -1)
+        izeros = izeros + 1
+        # minus 2**(bit-1)
+        if self.quant_config.is_sym:
+            izeros = torch.bitwise_and(izeros - 2**(bits - 1), (2 ** bits) - 1)
+
+        return izeros
+
+    def combine_low_bits(self, tensor_a, tensor_b):
+        """
+        Combine the lower 4 bits of two int8 tensors into a new int8 tensor.
+
+        Args:
+        tensor_a (torch.Tensor): First tensor of type int8.
+        tensor_b (torch.Tensor): Second tensor of type int8.
+
+        Returns:
+        torch.Tensor: New tensor of type int8, combining lower 4 bits of tensor_a and tensor_b.
+        """
+        # 确保输入是 int8 类型
+        if tensor_a.dtype != torch.int8 or tensor_b.dtype != torch.int8:
+            raise ValueError("Both tensors must be of int8 type.")
+
+        # 提取每个 tensor 的低4位
+        low_bits_a = torch.bitwise_and(tensor_a, 0x0F)  # 保留 tensor_a 的低4位
+        low_bits_b = torch.bitwise_and(tensor_b, 0x0F)  # 保留 tensor_b 的低4位
+
+        # 将 tensor_a 的低4位左移4位
+        shifted_low_bits_a = low_bits_a << 4
+
+        # 组合两个 tensor 的低4位
+        combined = torch.bitwise_or(shifted_low_bits_a, low_bits_b)
+
+        return combined
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/smoothquant.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/smoothquant.py
new file mode 100755
index 000000000..dd250fdda
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/smoothquant.py
@@ -0,0 +1,239 @@
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from vllm.model_executor.layers.linear import (LinearMethodBase, LinearBase,
+                                               set_weight_attrs)
+from vllm.model_executor.layers.quantization import register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
+                                           GroupQuantScaleParameter,
+                                           ModelWeightParameter,
+                                           RowvLLMParameter)
+from vllm_mlu import _mlu_ops as mlu_ops
+
+
+# @register_quantization_config("smoothquant")
+class SmoothQuantConfig(QuantizationConfig):
+    """Config class for SmoothQuant.
+    """
+
+    def __init__(
+        self,
+        weight_bits: int,
+        quant_mode: str,  # smoothquant
+        input_quant_method: str, # per token/per tensor
+        group_size: int,
+    ) -> None:
+        self.weight_bits = weight_bits
+        self.quant_mode = quant_mode
+        self.input_quant_method = input_quant_method
+        self.group_size = group_size
+
+        if quant_mode == "SmoothQuant" and (self.input_quant_method != "per_token" and self.input_quant_method != "per_tensor"):
+            raise ValueError(
+                "Currently, only per_token or per_tensor input quantization is supported for "
+                f"SmoothQuant, but got {self.input_quant_method}.")
+
+        self.pack_factor = 8 // self.weight_bits
+
+    def __repr__(self) -> str:
+        return (f"SmoothQuantConfig(weight_bits={self.weight_bits}, "
+                f"input_quant_method={self.input_quant_method}, "
+                f"quant_mode={self.quant_mode}, "
+                f"group_size={self.group_size})")
+
+    @classmethod
+    def get_name(self) -> str:
+        return "SmoothQuant"
+
+    @classmethod
+    def get_supported_act_dtypes(self) -> List[torch.dtype]:
+        return [torch.half, torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(self) -> int:
+        return 30
+
+    @staticmethod
+    def get_config_filenames() -> List[str]:
+        return ["quantize_config.json"]
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "SmoothQuantConfig":
+        weight_bits = cls.get_from_keys(config, ["bits"])
+        input_quant_method = cls.get_from_keys(config, ["input_quant_method"])
+        quant_mode = cls.get_from_keys(config, ["quant_mode"])
+        group_size = cls.get_from_keys_or(config, ["group_size"], 1)
+        return cls(weight_bits, quant_mode, input_quant_method, group_size)
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional["SmoothQuantLinearMethod"]:
+        if isinstance(layer, LinearBase):
+            return SmoothQuantLinearMethod(self, prefix)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return ["gelu", "gelu_fast", "gelu_new", "gelu_pytorch_tanh"]
+
+
+class SmoothQuantLinearMethod(LinearMethodBase):
+    """Linear method for SmoothQuant.
+
+    Args:
+        quant_config: The SmoothQuant quantization config.
+    """
+
+    def __init__(self, quant_config: SmoothQuantConfig, prefix: str):
+        self.quant_config = quant_config
+        # for per-tensor case, we can skip quant input for the first attn|ffn linear
+        #   and fusion this step in layernorm to get better performance
+        self.skip_quant_input = False
+        self.compute_dtype = torch.get_default_dtype()
+        self.is_expert = 'expert' in prefix
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        output_size_per_partition = sum(output_partition_sizes)
+        if (output_size_per_partition % self.quant_config.pack_factor != 0):
+            raise ValueError(
+                "The output size is not aligned with the quantized "
+                "weight shape. This can be caused by too large "
+                "tensor parallel size.")
+
+        pack_factor, group_num = 1, 1
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        if self.is_expert and self.quant_config.group_size > 1:
+            if input_size_per_partition % self.quant_config.group_size != 0:
+                raise ValueError(
+                    f"The input size {input_size_per_partition} is not aligned with the quantized "
+                    f"weight shape. This can be caused by too large "
+                    f"tensor parallel size. group_size: {self.quant_config.group_size}.")
+
+            pack_factor = self.quant_config.pack_factor
+            group_num = (input_size + self.quant_config.group_size - 1) // self.quant_config.group_size
+            if input_size_per_partition != input_size:
+                group_num = (input_size_per_partition + self.quant_config.group_size - 1) // self.quant_config.group_size
+
+        qweight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // pack_factor,
+                device="mlu",
+                dtype=torch.int8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        if group_num == 1:
+            per_channel_scale = ChannelQuantScaleParameter(
+                data=torch.empty(
+                    output_size_per_partition,
+                    device="mlu",
+                    dtype=torch.float32,
+                ),
+                output_dim=0,
+                weight_loader=weight_loader,
+            )
+        else:
+            per_channel_scale = GroupQuantScaleParameter(
+                data=torch.empty(
+                    output_size_per_partition,
+                    group_num,
+                    device="mlu",
+                    dtype=torch.float32,
+                ),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=weight_loader,
+            )
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("per_channel_scale", per_channel_scale)
+
+        if self.quant_config.input_quant_method == "per_token":
+            smooth = RowvLLMParameter(
+                data=torch.empty(
+                    input_size_per_partition,
+                    device="mlu",
+                    dtype=torch.float32,
+                ),
+                input_dim=0,
+                weight_loader=weight_loader,
+            )
+            set_weight_attrs(smooth, {
+                "ignore_warning": True,
+            })
+            layer.register_parameter("smooth", smooth)
+        if self.quant_config.input_quant_method == "per_tensor":
+            scale_to_int = RowvLLMParameter(
+                data=torch.empty(
+                    input_size_per_partition,
+                    device="mlu",
+                    dtype=torch.float32,
+                ),
+                input_dim=0,
+                weight_loader=weight_loader,
+            )
+            set_weight_attrs(scale_to_int, {
+                "ignore_warning": True,
+            })
+            layer.register_parameter("scale_to_int", scale_to_int)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if self.quant_config.input_quant_method == "per_token" and layer.smooth.dtype != torch.float:
+            layer.smooth = Parameter(layer.smooth.to(torch.float), requires_grad=False)
+        if self.quant_config.input_quant_method == "per_tensor" and layer.scale_to_int.dtype != torch.float:
+            layer.scale_to_int = Parameter(layer.scale_to_int.to(torch.float), requires_grad=False)
+        if layer.per_channel_scale.dtype != torch.float:
+            layer.per_channel_scale = Parameter(layer.per_channel_scale.to(torch.float), requires_grad=False)
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None,
+              residual: Optional[torch.Tensor] = None,
+              input_scale: Optional[torch.Tensor] = None,
+              use_tp_weight : bool = False
+        ) -> torch.Tensor:
+        layer_smooth = layer.smooth
+        layer_qweight = layer.qweight
+        layer_per_channel_scale = layer.per_channel_scale
+        if use_tp_weight:
+            if hasattr(layer, 'tp_smooth'):
+                layer_smooth = layer.tp_smooth
+            if hasattr(layer, 'tp_qweight'):
+                layer_qweight = layer.tp_qweight
+            if hasattr(layer, 'tp_per_channel_scale'):
+                layer_per_channel_scale = layer.tp_per_channel_scale
+
+        quant_input = None
+        if self.skip_quant_input:
+            quant_input = x
+        elif self.quant_config.input_quant_method == "per_token":
+            quant_input, input_scale = mlu_ops.per_token_smooth_quantize(x, layer_smooth, None)
+        elif self.quant_config.input_quant_method == "per_tensor":
+            quant_input = mlu_ops.quantize(x, layer.scale_to_int, None)
+        else:
+            raise ValueError(
+                "Currently, only per_token or per_tensor input quantization is supported for "
+                f"SmoothQuant, but got {self.input_quant_method}.")
+        quant_input_shape = quant_input.shape
+        if len(quant_input_shape) > 2:
+            quant_input = quant_input.view(-1, quant_input_shape[-1])
+            input_scale = input_scale.view(-1)
+        out = mlu_ops.smooth_quant_matmul(quant_input, input_scale, layer_qweight,
+                                          layer_per_channel_scale, self.compute_dtype, bias, residual)
+        if len(quant_input_shape) > 2:
+            out = out.view(*quant_input_shape[:-1], out.shape[-1])
+        return out
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/utils/__init__.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/utils/fp8_utils.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/utils/fp8_utils.py
new file mode 100644
index 000000000..fe87bef16
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/utils/fp8_utils.py
@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapted from https://github.com/sgl-project/sglang/pull/2575
+import functools
+import json
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    CUTLASS_BLOCK_FP8_SUPPORTED)
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    _w8a8_block_fp8_matmul,
+    _per_token_group_quant_fp8_colmajor)
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+def apply_w8a8_block_fp8_linear(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    block_size: List[int],
+    weight_scale: torch.Tensor,
+    input_scale: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+    cutlass_block_fp8_supported: bool = CUTLASS_BLOCK_FP8_SUPPORTED,
+) -> torch.Tensor:
+    assert input_scale is None
+    # View input as 2D matrix for fp8 methods
+    input_2d = input.view(-1, input.shape[-1])
+    output_shape = [*input.shape[:-1], weight.shape[0]]
+
+    shape_supported_by_cutlass = (weight.shape[0] % 128 == 0
+                                  and weight.shape[1] % 128 == 0)
+    if current_platform.is_rocm():
+        # TODO this is never used, as cutlass_block_fp8_supported is False
+        scale_a_shape = ((input_2d.shape[-1] // block_size[1], ) +
+                         input_2d.shape[:-1])[::-1]
+        scale_b_shape = (weight_scale.view(-1, 1)
+                         if weight_scale.dim() <= 1 else weight_scale.T).shape
+        ar, ac = scale_a_shape
+        br, bc = scale_b_shape
+        if (ac > 1 or bc > 1 or ar not in (1, input_2d.shape[0])
+                or br not in (1, weight.shape[0])):
+            shape_supported_by_cutlass = False
+    if cutlass_block_fp8_supported and shape_supported_by_cutlass:
+        q_input, x_scale = per_token_group_quant_fp8(input_2d,
+                                                     block_size[1],
+                                                     column_major_scales=True)
+        output = ops.cutlass_scaled_mm(q_input,
+                                       weight.T,
+                                       out_dtype=input.dtype,
+                                       scale_a=x_scale,
+                                       scale_b=weight_scale.T)
+    else:
+        q_input, x_scale = per_token_group_quant_fp8(input_2d,
+                                                     block_size[1],
+                                                     column_major_scales=False)
+        output = w8a8_block_fp8_matmul(q_input,
+                                       weight,
+                                       x_scale,
+                                       weight_scale,
+                                       block_size,
+                                       output_dtype=input.dtype)
+    if bias is not None:
+        output = output + bias
+    return output.to(dtype=input.dtype).view(*output_shape)
+
+@triton.jit
+def _per_token_group_quant_fp8(
+    # Pointers to inputs and output
+    y_ptr,
+    y_q_ptr,
+    y_s_ptr,
+    group_size: tl.constexpr,
+    # Avoid to divide zero
+    eps,
+    # Information for float8
+    fp8_min,
+    fp8_max,
+    # Meta-parameters
+    GROUP_BLOCK_NUM,
+):
+    """A Triton-accelerated function to perform per-token-group
+    quantization on a tensor.
+    This function converts the tensor values into float8 values.
+    """
+    # Map the program id to the row of X and Y it should compute.
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: split for limit the memory usage(65536)
+    '''
+    g_id = tl.program_id(0)
+    y_ptr += g_id * group_size * GROUP_BLOCK_NUM
+    y_q_ptr += g_id * group_size * GROUP_BLOCK_NUM
+    y_s_ptr += g_id * GROUP_BLOCK_NUM
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    cols = tl.arange(0, group_size)  # N <= BLOCK
+    for group_id in range(GROUP_BLOCK_NUM):
+        cols_group = cols + group_id * group_size
+        y = tl.load(y_ptr + cols_group).to(tl.float32)
+        # Quant
+        _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
+        y_s = _absmax / fp8_max
+        y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
+
+        tl.store(y_q_ptr + cols_group, y_q)
+        tl.store(y_s_ptr + group_id, y_s)
+
+
+def per_token_group_quant_fp8(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    dtype: Optional[torch.dtype] = None,
+    column_major_scales: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Function to perform per-token-group quantization on an input tensor `x`.
+    It converts the tensor values into signed float8 values and returns the
+    quantized tensor along with the scaling factor used for quantization.
+    Args:
+        x: The input tensor with ndim >= 2.
+        group_size: The group size used for quantization.
+        eps: The minimum to avoid dividing zero.
+        dtype: The dype of output tensor. Note that only `torch.float8_e4m3fn`
+        is supported for now.
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the
+        scaling factor for quantization.
+    """
+    dtype = current_platform.fp8_dtype() if dtype is None else dtype
+    assert (x.shape[-1] % group_size == 0), (
+        f"the last dimension of `x` {x.shape[-1]} must be divisible "
+        f"by `group_size` {group_size}")
+    assert x.stride(-1) == 1, "`x` groups must be contiguous"
+
+    finfo = torch.finfo(dtype)
+    fp8_min = finfo.min
+    fp8_max = finfo.max
+
+    x_q = torch.empty_like(x, device=x.device, dtype=dtype)
+    M = x.numel() // group_size
+    N = group_size
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: split for limit the memory usage(65536)
+    '''
+    group_per_block = 1
+    while M >= 65536:
+        group_per_block *= 2
+        M = x.numel() // (group_size * group_per_block)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    if column_major_scales:
+        shape = (x.shape[-1] // group_size, ) + x.shape[:-1]
+        x_s = torch.empty(shape, device=x.device,
+                          dtype=torch.float32).permute(-1, -2)
+    else:
+        shape = x.shape[:-1] + (x.shape[-1] // group_size, )
+        x_s = torch.empty(shape, device=x.device, dtype=torch.float32)
+
+    BLOCK = triton.next_power_of_2(N)
+    # heuristics for number of warps
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: set num_warps to 1 for triton-mlu
+    '''
+    num_warps = 1
+    num_stages = 1
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    if column_major_scales:
+        _per_token_group_quant_fp8_colmajor[(M, )](
+            x,
+            x_q,
+            x_s,
+            group_size,
+            x.shape[1],
+            x.stride(0),
+            x_s.stride(1),
+            eps,
+            fp8_min=fp8_min,
+            fp8_max=fp8_max,
+            BLOCK=BLOCK,
+            num_warps=num_warps,
+            num_stages=num_stages,
+        )
+    else:
+        _per_token_group_quant_fp8[(M, )](
+            x,
+            x_q,
+            x_s,
+            group_size,
+            eps,
+            fp8_min=fp8_min,
+            fp8_max=fp8_max,
+            GROUP_BLOCK_NUM=group_per_block,
+            num_warps=num_warps,
+            num_stages=num_stages,
+        )
+
+    return x_q, x_s
+
+
+def w8a8_block_fp8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor,
+    Bs: torch.Tensor,
+    block_size: List[int],
+    output_dtype: torch.dtype = torch.float16,
+) -> torch.Tensor:
+    """This function performs matrix multiplication with block-wise
+    quantization.
+    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
+    The output is returned in the specified `output_dtype`.
+    Args:
+        A: The input tensor, e.g., activation.
+        B: The input tensor, e.g., weight.
+        As: The per-token-group quantization scale for `A`.
+        Bs: The per-block quantization scale for `B`.
+        block_size: The block size for per-block quantization. It should
+        be 2-dim, e.g., [128, 128].
+        output_dytpe: The dtype of the returned tensor.
+    Returns:
+        torch.Tensor: The result of matmul.
+    """
+    assert len(block_size) == 2
+    block_n, block_k = block_size[0], block_size[1]
+
+    assert A.shape[-1] == B.shape[-1]
+    assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
+    assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
+    M = A.numel() // A.shape[-1]
+
+    assert B.ndim == 2 and Bs.ndim == 2
+    N, K = B.shape
+    assert triton.cdiv(N, block_n) == Bs.shape[0]
+    assert triton.cdiv(K, block_k) == Bs.shape[1]
+
+    C_shape = A.shape[:-1] + (N, )
+    C = A.new_empty(C_shape, dtype=output_dtype)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: using default config for triton-mlu
+    '''
+    # Default config
+    # Block-wise quant: BLOCK_SIZE_N must be divisible by block_size[0]
+    # BLOCK_SIZE_K must be divisible by block_size[1]
+    config = {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": block_size[0],
+        "BLOCK_SIZE_K": block_size[1],
+        "GROUP_SIZE_M": 32,
+        "num_warps": 1,
+        "num_stages": 1,
+    }
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    def grid(META):
+        return (triton.cdiv(M, META["BLOCK_SIZE_M"]) *
+                triton.cdiv(N, META["BLOCK_SIZE_N"]), )
+
+    _w8a8_block_fp8_matmul[grid](
+        A,
+        B,
+        C,
+        As,
+        Bs,
+        M,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(-2),
+        A.stride(-1),
+        B.stride(1),
+        B.stride(0),
+        C.stride(-2),
+        C.stride(-1),
+        As.stride(-2),
+        As.stride(-1),
+        Bs.stride(1),
+        Bs.stride(0),
+        **config,
+    )
+
+    return C
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/utils/w8a8_utils.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/utils/w8a8_utils.py
new file mode 100644
index 000000000..80a786543
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/utils/w8a8_utils.py
@@ -0,0 +1,180 @@
+from typing import Optional
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    Fp8LinearOp, USE_ROWWISE_TORCH_SCALED_MM, TORCH_DEVICE_IDENTITY)
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__model_executor__layers__quantization__utils__w8a8_util__Fp8LinearOp__apply(
+    self,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_scale: Optional[torch.Tensor] = None,
+    input_scale_ub: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+    # TODO(luka) remove this parameter in favor of __init__
+    use_per_token_if_dynamic: Optional[bool] = None,
+    weight_per_channel: bool = True,
+    activation_per_token: bool = True,
+) -> torch.Tensor:
+    # ops.scaled_fp8_quant supports both dynamic and static quant.
+    #   If dynamic, layer.input_scale is None and x_scale computed from x.
+    #   If static, layer.input_scale is scalar and x_scale is input_scale.
+
+    # View input as 2D matrix for fp8 methods
+    input_2d = input.view(-1, input.shape[-1])
+    output_shape = [*input.shape[:-1], weight.shape[1]]
+
+    # TODO(luka) this is here because currently MLA only decides this
+    #  during the forward method instead of in __init__.
+    if use_per_token_if_dynamic is None:
+        use_per_token_if_dynamic = self.use_per_token_if_dynamic
+
+    # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
+    if self.cutlass_fp8_supported:
+        qinput, x_scale = ops.scaled_fp8_quant(
+            input_2d,
+            input_scale,
+            scale_ub=input_scale_ub,
+            use_per_token_if_dynamic=use_per_token_if_dynamic)
+
+        # Fused GEMM_DQ
+        output = ops.cutlass_scaled_mm(qinput,
+                                        weight,
+                                        out_dtype=input.dtype,
+                                        scale_a=x_scale,
+                                        scale_b=weight_scale,
+                                        bias=bias)
+        return output.view(*output_shape)
+
+    elif weight_per_channel and activation_per_token:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: Add support for activation-per-token weight-per-channel quantization.
+        '''
+        q_input, x_scale = mlu_ops.scaled_quantize(
+            input_2d,# x
+            None, # scale
+            None, # zero
+            None, # scale_ub
+            quant_type=torch.float8_e4m3fn,
+            quant_mode='dynamic_per_token'
+        )
+        output_shape = [*input.shape[:-1], weight.shape[0]]
+        output = mlu_ops.scaled_matmul(
+            q_input, # a
+            weight, # b
+            x_scale, # a_scale
+            weight_scale, # b_scale
+            input.dtype, # output_dtype
+            bias, # bias
+        )
+        return output.view(output_shape)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+    # torch.scaled_mm supports per tensor weights + activations only
+    # so fallback to naive if per channel or per token
+    else:
+        # Maybe apply padding to output, see comment in __init__
+        qinput, x_scale = ops.scaled_fp8_quant(
+            input_2d,
+            input_scale,
+            num_token_padding=self.output_padding,
+            use_per_token_if_dynamic=use_per_token_if_dynamic)
+
+        per_tensor_weights = (weight_scale.numel() == 1)
+        per_tensor_activations = (x_scale.numel() == 1)
+
+        if per_tensor_weights and per_tensor_activations:
+            # Fused GEMM_DQ
+            output = torch._scaled_mm(qinput,
+                                        weight,
+                                        out_dtype=input.dtype,
+                                        scale_a=x_scale,
+                                        scale_b=weight_scale,
+                                        bias=bias)
+            # A fix for discrepancy in scaled_mm which returns tuple
+            # for torch < 2.5 and a single value in torch >= 2.5
+            if type(output) is tuple and len(output) == 2:
+                output = output[0]
+
+            return torch.narrow(output, 0, 0,
+                                input_2d.shape[0]).view(*output_shape)
+
+        elif (use_per_token_if_dynamic and not per_tensor_weights
+                and not per_tensor_activations
+                and USE_ROWWISE_TORCH_SCALED_MM):
+            # For now validated on ROCm platform
+            # fp8 rowwise scaling in torch._scaled_mm is introduced in
+            # https://github.com/pytorch/pytorch/pull/144432 using hipBLASLt
+            # and ROCm 6.3, which only exists in torch 2.7 and above.
+            # For CUDA platform please validate if the
+            # torch._scaled_mm support rowwise scaled GEMM
+            # Fused GEMM_DQ Rowwise GEMM
+            output = torch._scaled_mm(qinput,
+                                        weight,
+                                        out_dtype=input.dtype,
+                                        scale_a=x_scale,
+                                        scale_b=weight_scale.t(),
+                                        bias=bias)
+
+            output = torch.narrow(output, 0, 0, input_2d.shape[0])
+            output = output.view(*output_shape)
+            return output
+
+        else:
+            # Fallback for channelwise case, where we use unfused DQ
+            # due to limitations with scaled_mm
+
+            # Symmetric quantized GEMM by definition computes the following:
+            #   C = (s_x * X) (s_w * W) + bias
+            # This is equivalent to dequantizing the weights and activations
+            # before applying a GEMM.
+            #
+            # In order to compute quantized operands, a quantized kernel
+            # will rewrite the above like so:
+            #   C = s_w * s_x * (X * W) + bias
+            #
+            # For the scaled_mm fallback case, we break this down, since it
+            # does not support s_w being a vector.
+
+            # GEMM
+            # This computes C = (X * W).
+            # Output in fp32 to allow subsequent ops to happen in-place
+            output = torch._scaled_mm(qinput,
+                                        weight,
+                                        scale_a=TORCH_DEVICE_IDENTITY,
+                                        scale_b=TORCH_DEVICE_IDENTITY,
+                                        out_dtype=torch.float32)
+            # A fix for discrepancy in scaled_mm which returns tuple
+            # for torch < 2.5 and a single value in torch >= 2.5
+            if type(output) is tuple and len(output) == 2:
+                output = output[0]
+            # Unpad (undo num_token_padding)
+            output = torch.narrow(output, 0, 0, input_2d.shape[0])
+            x_scale = torch.narrow(x_scale, 0, 0, input_2d.shape[0])
+
+            # DQ
+            # C = sw * sx * (X * W) + bias
+            output = output * x_scale * weight_scale.t()
+            if bias is not None:
+                output = output + bias
+            return output.to(dtype=input.dtype).view(*output_shape)
+
+
+MluHijackObject.apply_hijack(
+    Fp8LinearOp,
+    Fp8LinearOp.apply,
+    vllm__model_executor__layers__quantization__utils__w8a8_util__Fp8LinearOp__apply
+)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/quantization/weightonly.py b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/weightonly.py
new file mode 100755
index 000000000..ea1074eb8
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/quantization/weightonly.py
@@ -0,0 +1,149 @@
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+from vllm.model_executor.layers.linear import (LinearMethodBase, LinearBase,
+                                               set_weight_attrs)
+from vllm.model_executor.layers.quantization import register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+from vllm_mlu import _mlu_ops as mlu_ops
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+# @register_quantization_config("weightonly")
+class WeightOnlyConfig(QuantizationConfig):
+    """Config class for WeightOnly.
+    """
+
+    def __init__(
+        self,
+        weight_bits: int,
+        quant_mode: str,  # weight_only
+    ) -> None:
+        self.weight_bits = weight_bits
+        self.quant_mode = quant_mode
+
+        if quant_mode == "WeightOnly" and (self.weight_bits != 8 and self.weight_bits != 4):
+            raise ValueError(
+                "Currently, only 8/4-bit weight quantization is supported for "
+                f"weight_only, but got {self.weight_bits} bits.")
+        self.pack_factor = 8 // self.weight_bits
+
+    def __repr__(self) -> str:
+        return (f"WeightOnlyConfig(weight_bits={self.weight_bits}, "
+                f"quant_mode={self.quant_mode})")
+
+    def get_name(self) -> str:
+        return "WeightOnly"
+
+    def get_supported_act_dtypes(self) -> List[torch.dtype]:
+        return [torch.half, torch.bfloat16]
+
+    def get_min_capability(self) -> int:
+        return 30
+
+    @staticmethod
+    def get_config_filenames() -> List[str]:
+        return ["quantize_config.json"]
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "WeightOnlyConfig":
+        weight_bits = cls.get_from_keys(config, ["bits"])
+        try:
+            quant_mode = cls.get_from_keys(config, ["quant_mode"])
+        except Exception:
+            quant_mode = "WeightOnly"
+        return cls(weight_bits, quant_mode)
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional["WeightOnlyLinearMethod"]:
+        if isinstance(layer, LinearBase):
+            return WeightOnlyLinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return ["gelu", "gelu_fast", "gelu_new", "gelu_pytorch_tanh"]
+
+
+class WeightOnlyLinearMethod(LinearMethodBase):
+    """Linear method for WeightOnly.
+
+    Args:
+        quant_config: The WeightOnly quantization config.
+    """
+
+    def __init__(self, quant_config: WeightOnlyConfig):
+        self.quant_config = quant_config
+
+    def create_weights(
+        self, 
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int], 
+        input_size: int,
+        output_size: int, 
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ) -> Dict[str, Any]:
+        output_size_per_partition = sum(output_partition_sizes)
+        if self.quant_config.quant_mode == "WeightOnly":
+            scale_and_zero_input_dim = None
+            if output_size != output_size_per_partition:
+                scale_and_zero_input_dim = 0
+            qweight = Parameter(
+                torch.empty(
+                    output_size_per_partition,
+                    input_size_per_partition // self.quant_config.pack_factor,
+                    device="mlu",
+                    dtype=torch.int8,
+                ),
+                requires_grad=False,
+            )
+            set_weight_attrs(qweight, {
+                "input_dim": 1,
+                "output_dim": 0,
+            })
+            scales = Parameter(
+                torch.empty(
+                    output_size_per_partition,
+                    device="mlu",
+                    dtype=params_dtype,
+                ),
+                requires_grad=False,
+            )
+            set_weight_attrs(scales, {
+                "input_dim": scale_and_zero_input_dim,
+                "output_dim": 0,
+            })
+            layer.register_parameter("qweight", qweight)
+            set_weight_attrs(qweight, extra_weight_attrs)
+            layer.register_parameter("scales", scales)
+            set_weight_attrs(scales, extra_weight_attrs)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if layer.scales.dtype != torch.float:
+            layer.scales = Parameter(layer.scales.to(torch.float), requires_grad=False)
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None,
+              residual: Optional[torch.Tensor] = None) -> torch.Tensor:
+        x_shape = x.shape
+        if len(x_shape) > 2:
+            x = x.view(-1, x_shape[-1])
+        out = mlu_ops.weight_only_quant_matmul(x,
+                                              layer.qweight,
+                                              layer.scales,
+                                              None,
+                                              bias,
+                                              residual,
+                                              "none",
+                                              self.quant_config.weight_bits)
+        if len(x_shape) > 2:
+            out = out.view(*x_shape[:-1], out.shape[-1])
+        return out
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/rotary_embedding.py b/vllm_mlu/vllm_mlu/model_executor/layers/rotary_embedding.py
new file mode 100644
index 000000000..938551ddb
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/rotary_embedding.py
@@ -0,0 +1,738 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import math
+import torch
+
+from vllm.attention import AttentionMetadata
+from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.rotary_embedding import (
+    RotaryEmbedding, MRotaryEmbedding,
+    LinearScalingRotaryEmbedding, DeepseekScalingRotaryEmbedding,
+    DynamicNTKScalingRotaryEmbedding, YaRNScalingRotaryEmbedding,
+    Phi3LongRoPEScaledRotaryEmbedding, Llama4VisionRotaryEmbedding,
+    _yarn_find_correction_range, _ROPE_DICT,
+    yarn_get_mscale, _yarn_linear_ramp_mask)
+from vllm.model_executor.layers import rotary_embedding
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+def get_long_max_model_max_position_emb(max_position_embeddings, scaling_factor):
+    if MLURotaryEmbedding.max_seq_len != None and \
+            MLURotaryEmbedding.max_seq_len > max_position_embeddings * scaling_factor:
+        logger.warning(f"User-specified max_model_len ({MLURotaryEmbedding.max_seq_len}) is different with " +
+                        f"max_position_embedding ({max_position_embeddings}) * scaling_factor ({scaling_factor}) " +
+                        "from model's config.json, This may lead to incorrect model outputs or MLU errors. " +
+                        f"Make sure the value is correct and within the model context size. " +
+                        f"Set max_position_embedding={MLURotaryEmbedding.max_seq_len}.")
+        return math.ceil(MLURotaryEmbedding.max_seq_len / scaling_factor)
+    return max_position_embeddings
+
+
+@CustomOp.register("rotary_embedding_mlu")
+class MLURotaryEmbedding(RotaryEmbedding, CustomOp):
+
+    cu_seq_lens : torch.Tensor = None
+    max_seq_len : int = None
+    max_model_len : int = None
+    is_prompt : bool = False
+    is_chunked : bool = False
+    set_cos_sin : bool = False
+    cos_ : torch.Tensor = None
+    sin_ : torch.Tensor = None
+    positions_: torch.Tensor = None
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+    ) -> None:
+        CustomOp.__init__(self)
+        self.head_size = head_size
+        self.rotary_dim = rotary_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        self.is_neox_style = is_neox_style
+        self.dtype = dtype
+
+        if MLURotaryEmbedding.max_seq_len != None \
+                and self.max_position_embeddings < MLURotaryEmbedding.max_seq_len and \
+                not isinstance(self, (YaRNScalingRotaryEmbedding, DeepseekScalingRotaryEmbedding)):
+            logger.warning(f"User-specified max_model_len ({MLURotaryEmbedding.max_seq_len}) is different with " +
+                            f"max_position_embedding ({max_position_embeddings}) from model's config.json, " +
+                            f"This may lead to incorrect model outputs or MLU errors. " +
+                            f"Make sure the value is correct and within the model context size. " +
+                            f"Set max_position_embedding={MLURotaryEmbedding.max_seq_len}.")
+            self.max_position_embeddings = MLURotaryEmbedding.max_seq_len
+        cache = self._compute_cos_sin_cache()
+        if isinstance(self, MLULinearScalingRotaryEmbedding):
+            logger.debug(f"Using mlu defining _compute_cos_sin_cache due to the special tensor composition")
+        elif is_neox_style:
+            cache_pos = cache.shape[0]
+            cache = cache.reshape(cache_pos, 2, -1)
+            cache = torch.tile(cache, (1, 1, 2)).reshape(cache_pos, -1)
+        else:
+            cache = cache.repeat_interleave(2, dim=-1)
+
+        cache = cache.to(dtype)
+        self.cos_sin_cache: torch.Tensor
+        self.register_buffer("cos_sin_cache", cache, persistent=False)
+
+
+    @classmethod
+    def set_mlu_var(
+        cls,
+        input_ids: torch.Tensor,
+        attn_metadata: AttentionMetadata
+    ) -> None:
+        cls.unset_mlu_var()
+        is_chunked = False
+        is_prompt = False
+        prefill_metadata = attn_metadata.prefill_metadata
+        decode_metadata = attn_metadata.decode_metadata
+        if prefill_metadata:
+            cu_seq_lens = prefill_metadata.query_start_loc
+            rope_max_seq_len = prefill_metadata.max_query_len
+            is_prompt = True
+            # Workaround: mlugraph does not support torch.ne|eq|equal .etc for now,
+            # because context mlugraph always uses in benchmark latency, and in this
+            # case, query_start_loc always equals to seq_start_loc, so we can set
+            # is_chunked to False directly.
+            if prefill_metadata.use_cuda_graph:
+                is_chunked = False
+            elif decode_metadata or \
+                max(prefill_metadata.seq_lens) != prefill_metadata.max_query_len:
+                is_chunked = True
+        if decode_metadata:
+            if prefill_metadata:
+                cu_seq_lens = attn_metadata.query_start_loc
+                rope_max_seq_len = min(cls.max_model_len, max(rope_max_seq_len,
+                    attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens))
+            else:
+                # input_ids is pack mode, and the decode_seq_len = 1
+                cu_seq_lens = torch.arange(0, input_ids.shape[0] + 1, 1, dtype=torch.int32, device="mlu")
+                rope_max_seq_len = 1
+        cls.cu_seq_lens = cu_seq_lens
+        cls.max_seq_len = rope_max_seq_len
+        cls.is_prompt = is_prompt
+        cls.is_chunked = is_chunked
+
+        if prefill_metadata:
+            cls.prefill_max_seq_len = prefill_metadata.max_query_len
+            cls.prefill_cu_seq_lens = prefill_metadata.query_start_loc
+
+        if decode_metadata:
+            cls.decode_max_seq_len = decode_metadata.max_decode_query_len
+            cls.decode_cu_seq_lens = decode_metadata.query_start_loc
+            if cls.decode_cu_seq_lens is None:
+                cls.decode_cu_seq_lens = cu_seq_lens
+
+    @classmethod
+    def unset_mlu_var(cls):
+        cls.cu_seq_lens = None
+        cls.max_seq_len = None
+        cls.is_prompt = False
+        cls.is_chunked = False
+        cls.set_cos_sin = False
+        cls.cos_ = None
+        cls.sin_ = None
+        cls.positions_ = None
+
+    def _get_cos_sin(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        cos, sin = self.cos_sin_cache.chunk(2, dim=-1)
+        sin = sin.view(-1, self.rotary_dim)
+        cos = cos.view(-1, self.rotary_dim)
+        return cos, sin
+
+    def _get_positions_with_offsets_mlu(
+        self,
+        positions: torch.Tensor,
+        offsets: torch.Tensor
+    ) -> torch.Tensor:
+        if offsets.numel() != positions.numel():
+            raise Exception("rope offsets numel mismatch with positions, "
+                            f"positions: {positions.numel()}, offsets: {offsets.numel()}")
+        return (positions + offsets).to(torch.int32)
+
+    def forward_oot(
+        self,
+        positions: torch.Tensor,
+        x: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        from vllm_mlu import _mlu_ops as mlu_ops
+
+        # ops.rotary_embedding()/batched_rotary_embedding()
+        # are in-place operations that update the query and key tensors.
+        if MLURotaryEmbedding.set_cos_sin == False:
+            MLURotaryEmbedding.cos_, MLURotaryEmbedding.sin_ = self._get_cos_sin()
+            MLURotaryEmbedding.set_cos_sin = True
+        interleaved = True
+        if self.is_neox_style:
+            interleaved = False
+
+        if offsets is not None:
+            if MLURotaryEmbedding.positions_ is None:
+                MLURotaryEmbedding.positions_ = (
+                    self._get_positions_with_offsets_mlu(positions, offsets))
+            position_ids = MLURotaryEmbedding.positions_
+            discrete = True
+        elif MLURotaryEmbedding.is_chunked or not MLURotaryEmbedding.is_prompt:
+            position_ids = positions
+            discrete = True
+        else:
+            position_ids = None
+            discrete = False
+
+        x = mlu_ops.rotary_embedding(
+            x,
+            MLURotaryEmbedding.sin_,
+            MLURotaryEmbedding.cos_,
+            position_ids,
+            MLURotaryEmbedding.cu_seq_lens,
+            interleaved,
+            discrete,
+            False,
+            MLURotaryEmbedding.max_seq_len
+        )
+        return x
+
+
+class MLULinearScalingRotaryEmbedding(MLURotaryEmbedding, LinearScalingRotaryEmbedding):
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        scaling_factors: Union[List[float], float],
+        dtype: torch.dtype,
+    ) -> None:
+        if isinstance(scaling_factors, float):
+            scaling_factors = [scaling_factors]
+        self.scaling_factors: List[float] = scaling_factors  # noqa
+        MLURotaryEmbedding.__init__(self, head_size, rotary_dim,
+                                    max_position_embeddings, base,
+                                    is_neox_style, dtype)
+        # Lazy initialized.
+        self._scaling_factor_to_offset: Dict[float, int]
+
+    def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
+        """Compute the inverse frequency."""
+        half_dim = self.rotary_dim // 2
+        if self.is_neox_style:
+            inv_freq = 1.0 / (base ** ((torch.arange(
+                0, self.rotary_dim, 1, dtype=torch.float32, device="mlu") % half_dim) * 2 / self.rotary_dim)
+            )
+        else:
+            inv_freq = 1.0 / (
+                    base
+                    ** ( torch.arange(0, self.rotary_dim, 1, device="mlu", dtype=torch.float32) // 2 * 2
+                        / self.rotary_dim
+                    )
+            )
+        return inv_freq
+
+    def _compute_cos_sin_cache(self) -> torch.Tensor:
+        inv_freq = self._compute_inv_freq(self.base)
+        cache_list: List[torch.Tensor] = []
+        # offsets to the next cache in a tensor.
+        # Each offset corresponds to the same index in scaling_factors.
+        offsets: List[int] = []
+        for scaling_factor in self.scaling_factors:
+            # NOTE(woosuk): self.max_position_embeddings is the original
+            # maximum length before applying the rope scaling.
+            # Thus, the maximum length after applying the rope scaling is
+            # self.max_position_embeddings * self.scaling_factor.
+            max_len = self.max_position_embeddings * scaling_factor
+            t = torch.arange(max_len, dtype=torch.float, device="mlu")
+            t = t / scaling_factor
+
+            freqs = torch.einsum("i,j -> ij", t, inv_freq)
+            cos = freqs.cos()
+            sin = freqs.sin()
+            cache = torch.cat((cos, sin), dim=-1)
+            if not cache_list:
+                offset = 0
+            else:
+                last_offset = offsets[-1]
+                next_max_len = cache_list[-1].shape[0]
+                offset = last_offset + next_max_len
+            offsets.append(offset)
+            cache_list.append(cache)
+        self._scaling_factor_to_offset = {
+            float(scaling_factor): offsets[i]
+            for i, scaling_factor in enumerate(self.scaling_factors)
+        }
+        assert len(self.scaling_factors) == len(offsets)
+        return torch.cat(cache_list, dim=0)
+
+
+class MLUDynamicNTKScalingRotaryEmbedding(MLURotaryEmbedding, DynamicNTKScalingRotaryEmbedding):
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        scaling_factor: float,
+        dtype: torch.dtype,
+    ) -> None:
+        self.scaling_factor = scaling_factor
+        MLURotaryEmbedding.__init__(self, head_size, rotary_dim,
+                                    max_position_embeddings, base,
+                                    is_neox_style, dtype)
+
+
+class MLUDeepseekScalingRotaryEmbedding(MLURotaryEmbedding, DeepseekScalingRotaryEmbedding):
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        scaling_factor: float,
+        dtype: torch.dtype,
+        *,
+        extrapolation_factor: float = 1,
+        attn_factor: float = 1,
+        beta_fast: int = 32,
+        beta_slow: int = 1,
+        mscale: float = 1,
+        mscale_all_dim: float = 0,
+    ) -> None:
+        self.scaling_factor = scaling_factor
+        self.extrapolation_factor = extrapolation_factor
+        self.attn_factor = attn_factor
+        self.beta_fast = beta_fast
+        self.beta_slow = beta_slow
+        # Get n-d magnitude scaling corrected for interpolation.
+        self.mscale = float(
+            yarn_get_mscale(self.scaling_factor, float(mscale)) /
+            yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) *
+            attn_factor)
+        MLURotaryEmbedding.__init__(self, head_size, rotary_dim,
+                                    max_position_embeddings, base,
+                                    is_neox_style, dtype)
+
+    def forward_mlu_rot(self, input, position_ids, interleaved, discrete, cu_seq_lens, max_seq_len):
+        """only one input rotary implementation"""
+        from vllm_mlu import _mlu_ops as mlu_ops
+        if input is None:
+            return None
+        if self.rotary_dim < self.head_size:
+            input_pass = input[..., self.rotary_dim:]
+        input_rot = input[..., :self.rotary_dim]
+        input_rot = mlu_ops.rotary_embedding(
+            input_rot,
+            MLURotaryEmbedding.sin_,
+            MLURotaryEmbedding.cos_,
+            position_ids,
+            cu_seq_lens,
+            interleaved,
+            discrete,
+            False,
+            max_seq_len
+        )
+
+        if self.rotary_dim < self.head_size:
+            input = torch.cat((input_rot, input_pass), dim=-1)
+        else:
+            input = input_rot
+
+        return input
+
+    def get_param(self, positions):
+        if MLURotaryEmbedding.set_cos_sin == False:
+            MLURotaryEmbedding.cos_, MLURotaryEmbedding.sin_ = self._get_cos_sin()
+            MLURotaryEmbedding.set_cos_sin = True
+        interleaved = True
+        if self.is_neox_style:
+            interleaved = False
+        if MLURotaryEmbedding.is_chunked or not MLURotaryEmbedding.is_prompt:
+            position_ids = positions
+            discrete = True
+        else:
+            position_ids = None
+            discrete = False
+
+        return position_ids, interleaved, discrete
+
+    def forward_oot(
+        self,
+        positions: torch.Tensor,
+        query: Optional[torch.Tensor] = None,
+        key: Optional[torch.Tensor] = None,
+        offsets: Optional[torch.Tensor] = None,
+        only_prefill: Optional[bool] = False,
+        only_decode: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """PyTorch-native implementation equivalent to forward()."""
+        position_ids, interleaved, discrete = self.get_param(positions)
+        if only_prefill:
+            cu_seq_lens = MLURotaryEmbedding.prefill_cu_seq_lens
+            max_seq_len = MLURotaryEmbedding.prefill_max_seq_len
+        elif only_decode:
+            cu_seq_lens = MLURotaryEmbedding.decode_cu_seq_lens
+            max_seq_len = MLURotaryEmbedding.decode_max_seq_len
+        else:
+            cu_seq_lens = MLURotaryEmbedding.cu_seq_lens
+            max_seq_len = MLURotaryEmbedding.max_seq_len
+
+        query = self.forward_mlu_rot(query, position_ids, interleaved, discrete, cu_seq_lens, max_seq_len)
+        key = self.forward_mlu_rot(key, position_ids, interleaved, discrete, cu_seq_lens, max_seq_len)
+
+        return query, key
+
+    def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: change device cuda to mlu
+        '''
+        pos_freqs = self.base**(torch.arange(
+            0, self.rotary_dim, 2, dtype=torch.float, device="mlu") /
+                                self.rotary_dim)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        inv_freq_extrapolation = 1.0 / pos_freqs
+        inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
+
+        low, high = _yarn_find_correction_range(self.beta_fast, self.beta_slow,
+                                                self.rotary_dim, self.base,
+                                                self.max_position_embeddings)
+        # Get n-d rotational scaling corrected for extrapolation
+        inv_freq_mask = (1 - _yarn_linear_ramp_mask(
+            low, high, self.rotary_dim // 2,
+            dtype=torch.float)) * self.extrapolation_factor
+        inv_freq = inv_freq_interpolation * (
+            1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
+        return inv_freq
+
+    def _compute_cos_sin_cache(self) -> torch.Tensor:
+        inv_freq = self._compute_inv_freq(self.scaling_factor)
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: change device cuda to mlu
+        '''
+        t = torch.arange(self.max_position_embeddings * self.scaling_factor,
+                         device="mlu",
+                         dtype=torch.float32)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        freqs = torch.einsum("i,j -> ij", t, inv_freq)
+        cos = (freqs.cos() * self.mscale)
+        sin = (freqs.sin() * self.mscale)
+        cache = torch.cat((cos, sin), dim=-1)
+        return cache
+
+    forward = MLURotaryEmbedding.forward
+
+
+class MLULlama3RotaryEmbedding(MLURotaryEmbedding):
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+        scaling_factor: float,
+        low_freq_factor: float,
+        high_freq_factor: float,
+        orig_max_position: int,
+    ) -> None:
+        self.scaling_factor = scaling_factor
+        self.low_freq_factor = low_freq_factor
+        self.high_freq_factor = high_freq_factor
+        self.orig_max_position = orig_max_position
+        super().__init__(head_size, rotary_dim, max_position_embeddings, base,
+                         is_neox_style, dtype)
+
+
+class DynamicNTKAlphaRotaryEmbedding(RotaryEmbedding):
+    """RotaryEmbedding extended with Dynamic NTK scaling.
+    Credits to the Reddit users /u/bloc97 and /u/emozilla
+    """
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        scaling_alpha: float,
+        dtype: torch.dtype,
+    ) -> None:
+        self.scaling_alpha = scaling_alpha
+        super().__init__(head_size, rotary_dim, max_position_embeddings, base,
+                         is_neox_style, dtype)
+
+    def _compute_cos_sin_cache(self) -> torch.Tensor:
+        max_len = self.max_position_embeddings
+        base = self.base * self.scaling_alpha ** (self.rotary_dim / (self.rotary_dim - 2))
+
+        inv_freq = self._compute_inv_freq(base)
+        t = torch.arange(max_len, dtype=torch.float)
+
+        freqs = torch.einsum("i,j -> ij", t, inv_freq)
+        cos = freqs.cos()
+        sin = freqs.sin()
+        cache = torch.cat((cos, sin), dim=-1)
+        return cache
+
+
+class MLUDynamicNTKAlphaRotaryEmbedding(MLURotaryEmbedding, DynamicNTKAlphaRotaryEmbedding):
+    """RotaryEmbedding extended with Dynamic NTK scaling.
+    Credits to the Reddit users /u/bloc97 and /u/emozilla
+    """
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        scaling_alpha: float,
+        dtype: torch.dtype,
+    ) -> None:
+        self.scaling_alpha = scaling_alpha
+        MLURotaryEmbedding.__init__(
+            self, head_size, rotary_dim, max_position_embeddings, base,
+            is_neox_style, dtype)
+
+
+class MLUMRotaryEmbedding(MLURotaryEmbedding, MRotaryEmbedding):
+    """Rotary Embedding with Multimodal Sections."""
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+        mrope_section: Optional[List[int]] = None,
+    ) -> None:
+        # In Qwen2.5-VL, the maximum index value is related to the duration of
+        # the input video. We enlarge max_position_embeddings to 4 times to get
+        # a larger the cos and sin cache.
+        self.cache_max_position_num = max_position_embeddings * 4
+        MLURotaryEmbedding.__init__(
+            self, head_size, rotary_dim, self.cache_max_position_num, base,
+            is_neox_style, dtype)
+
+        self.mrope_section = mrope_section
+        if self.mrope_section:
+            assert sum(self.mrope_section) == rotary_dim // 2
+
+    def forward_oot(
+        self,
+        positions: torch.Tensor,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert positions.ndim == 1 or positions.ndim == 2
+        num_tokens = positions.shape[-1]
+
+        cos_sin = self.cos_sin_cache[positions]
+        cos, sin = cos_sin.chunk(2, dim=-1)
+        if positions.ndim == 2:
+            assert self.mrope_section
+            num_section = len(self.mrope_section)
+            mrope_section = self.mrope_section * 2
+            cos = torch.cat([
+                m[i % num_section]
+                for i, m in enumerate(cos.split(mrope_section, dim=-1))
+            ],
+                            dim=-1)
+            sin = torch.cat([
+                m[i % num_section]
+                for i, m in enumerate(sin.split(mrope_section, dim=-1))
+            ],
+                            dim=-1)
+        from vllm_mlu import _mlu_ops as mlu_ops
+        interleaved = True
+        if self.is_neox_style:
+            interleaved = False
+        position_ids = None
+        discrete = False
+        # mlu_ops.rotary_embedding() is a in-place operation that update the query and key tensors.
+        x = mlu_ops.rotary_embedding(x,
+                                     sin,
+                                     cos,
+                                     position_ids,
+                                     MLURotaryEmbedding.cu_seq_lens,
+                                     interleaved,
+                                     discrete,
+                                     False,
+                                     MLURotaryEmbedding.max_seq_len)
+        return x
+
+    forward = MLURotaryEmbedding.forward
+
+
+def vllm__model_executor__layers__rotary_embedding__get_rope(
+    head_size: int,
+    rotary_dim: int,
+    max_position: int,
+    base: int,
+    is_neox_style: bool = True,
+    rope_scaling: Optional[Dict[str, Any]] = None,
+    dtype: Optional[torch.dtype] = None,
+    partial_rotary_factor: float = 1.0,
+) -> RotaryEmbedding:
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+    if rope_scaling is not None:
+        # Transforms every value that is a list into a tuple for caching calls
+        rope_scaling_tuple = {
+            k: tuple(v) if isinstance(v, list) else v
+            for k, v in rope_scaling.items()
+        }
+        rope_scaling_args = tuple(rope_scaling_tuple.items())
+    else:
+        rope_scaling_args = None
+    if partial_rotary_factor < 1.0:
+        rotary_dim = int(rotary_dim * partial_rotary_factor)
+    key = (head_size, rotary_dim, max_position, base, is_neox_style,
+           rope_scaling_args, dtype)
+    if key in _ROPE_DICT:
+        return _ROPE_DICT[key]
+
+    if rope_scaling is None:
+        rotary_emb = MLURotaryEmbedding(head_size, rotary_dim, max_position, base,
+                                        is_neox_style, dtype)
+    else:
+        scaling_type = rope_scaling["rope_type"]
+
+        if scaling_type == "llama3":
+            scaling_factor = rope_scaling["factor"]
+            low_freq_factor = rope_scaling["low_freq_factor"]
+            high_freq_factor = rope_scaling["high_freq_factor"]
+            original_max_position = rope_scaling[
+                "original_max_position_embeddings"]
+            rotary_emb = MLULlama3RotaryEmbedding(head_size, rotary_dim,
+                                                  max_position, base,
+                                                  is_neox_style, dtype,
+                                                  scaling_factor, low_freq_factor,
+                                                  high_freq_factor,
+                                                  original_max_position)
+        elif scaling_type == "mllama4":
+            rotary_emb = Llama4VisionRotaryEmbedding(head_size, rotary_dim,
+                                                     max_position, base,
+                                                     is_neox_style, dtype)
+        elif scaling_type == "default":
+            if "mrope_section" in rope_scaling:
+                rotary_emb = MLUMRotaryEmbedding(
+                    head_size,
+                    rotary_dim,
+                    max_position,
+                    base,
+                    is_neox_style,
+                    dtype,
+                    mrope_section=rope_scaling["mrope_section"],
+                )
+            else:
+                rotary_emb = MLURotaryEmbedding(
+                    head_size,
+                    rotary_dim,
+                    max_position,
+                    base,
+                    is_neox_style,
+                    dtype,
+                )
+        elif scaling_type == "linear":
+            scaling_factor = rope_scaling["factor"]
+            rotary_emb = MLULinearScalingRotaryEmbedding(head_size, rotary_dim,
+                                                         max_position, base,
+                                                         is_neox_style,
+                                                         scaling_factor, dtype)
+        elif scaling_type == "dynamic":
+            if "alpha" in rope_scaling:
+                rotary_emb = MLUDynamicNTKAlphaRotaryEmbedding(
+                    head_size, rotary_dim, max_position, base, is_neox_style,
+                    rope_scaling["alpha"], dtype)
+            else:
+                scaling_factor = rope_scaling["factor"]
+                rotary_emb = MLUDynamicNTKScalingRotaryEmbedding(
+                    head_size, rotary_dim, max_position, base, is_neox_style,
+                    scaling_factor, dtype)
+        elif scaling_type == "yarn":
+            scaling_factor = rope_scaling["factor"]
+            original_max_position = rope_scaling[
+                "original_max_position_embeddings"]
+            extra_kwargs = {
+                k: v
+                for k, v in rope_scaling.items()
+                if k in ("extrapolation_factor", "attn_factor", "beta_fast",
+                         "beta_slow")
+            }
+            original_max_position = get_long_max_model_max_position_emb(original_max_position, scaling_factor)
+            rotary_emb = YaRNScalingRotaryEmbedding(head_size, rotary_dim,
+                                                    original_max_position,
+                                                    base, is_neox_style,
+                                                    scaling_factor, dtype,
+                                                    **extra_kwargs)
+        elif scaling_type == "deepseek_yarn":
+            scaling_factor = rope_scaling["factor"]
+            original_max_position = rope_scaling[
+                "original_max_position_embeddings"]
+            # assert max_position == original_max_position * scaling_factor
+            extra_kwargs = {
+                k: v
+                for k, v in rope_scaling.items()
+                if k in ("extrapolation_factor", "attn_factor", "beta_fast",
+                         "beta_slow", "mscale", "mscale_all_dim")
+            }
+            original_max_position = get_long_max_model_max_position_emb(original_max_position, scaling_factor)
+            rotary_emb = MLUDeepseekScalingRotaryEmbedding(
+                head_size, rotary_dim, original_max_position, base,
+                is_neox_style, scaling_factor, dtype, **extra_kwargs)
+        elif scaling_type == "longrope":
+            short_factor = rope_scaling["short_factor"]
+            long_factor = rope_scaling["long_factor"]
+            original_max_position = rope_scaling[
+                "original_max_position_embeddings"]
+            extra_kwargs = {
+                k: v
+                for k, v in rope_scaling.items()
+                if k in ("short_mscale", "long_mscale")
+            }
+            rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(
+                head_size, rotary_dim, max_position, original_max_position,
+                base, is_neox_style, dtype, short_factor, long_factor,
+                **extra_kwargs)
+        else:
+            raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+    _ROPE_DICT[key] = rotary_emb
+    return rotary_emb
+
+
+MluHijackObject.apply_hijack(rotary_embedding,
+                             rotary_embedding.get_rope,
+                             vllm__model_executor__layers__rotary_embedding__get_rope)
diff --git a/vllm_mlu/vllm_mlu/model_executor/layers/sparse_moe_mlp.py b/vllm_mlu/vllm_mlu/model_executor/layers/sparse_moe_mlp.py
new file mode 100644
index 000000000..92af55b49
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/layers/sparse_moe_mlp.py
@@ -0,0 +1,613 @@
+"""Inference-only MOE model."""
+from typing import Optional
+
+import torch
+from torch import nn
+
+from vllm.distributed import (get_tensor_model_parallel_rank,
+                              get_tensor_model_parallel_world_size,
+                              get_tensor_model_parallel_group,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk
+
+from vllm_mlu._mlu_utils import *
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.model_executor.layers.quantization.weightonly import WeightOnlyConfig
+from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantConfig
+
+class SparseMoeMlp(nn.Module):
+    """
+    Tensor Parallel evenly splits each expert's weight and distributes them to different ranks,
+    which means each rank holds partial weight of all experts.
+    While Expert Parallel evenly distributes some of the experts' full weight to different ranks,
+    which means each rank holds part of the experts' full weight.
+
+    As a result, each rank in the Tensor Parallel group receives all tokens' hidden states for all experts,
+    then computes using the partial weights, while for Expert Parallel, each rank only receives
+    part of tokens' hidden states for experts on this rank, then computes using the full weights.
+
+    When both Tensor Parallel and Expert Parallel are enabled, each rank handles
+    a portion of the expert weights matrices (as in EP mode) and these weights are further sliced
+    across ranks (as in TP mode). This hybrid approach aims to balance the workload more evenly across ranks,
+    enhancing efficiency and reducing the likelihood of bottlenecks associated with EP mode alone.
+    """
+    reduce_weight : torch.Tensor = None
+    expert_id     : torch.Tensor = None
+    is_expert_avg : bool = False
+    max_batched_token : int = 2048
+
+    def __init__(
+        self,
+        num_experts: int,
+        top_k: int,
+        hidden_size: int,
+        intermediate_size: int,
+        up_proj_name: str,
+        is_gated: bool,
+        down_proj_name: str,
+        has_bias: bool,
+        skip_bias_add: bool = False,
+        renormalize:bool = False,
+        hidden_act: str = "silu",
+        params_dtype: Optional[torch.dtype] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        is_use_fused_moe: bool = False,
+        expert_group: Optional[int] = 1,
+        topk_group: Optional[int] = 1,
+        scoring_func: str = "softmax",
+        topk_method: str = "",
+        routed_scaling_factor: float = 1.0,
+    ):
+        super().__init__()
+        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_group = get_tensor_model_parallel_group()
+        self.num_total_experts = num_experts
+        self.top_k = top_k
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.up_proj_name = up_proj_name
+        self.is_gated = is_gated
+        self.down_proj_name = down_proj_name
+        self.has_bias = has_bias
+        self.renormalize = renormalize
+        self.hidden_act = hidden_act
+        self.quant_config = quant_config
+        self.is_use_fused_moe = is_use_fused_moe
+        self.expert_group = expert_group
+        self.topk_group = topk_group
+        self.scoring_func = scoring_func
+        self.routed_scaling_factor = routed_scaling_factor
+        # fused_moe doesn't support weightonly quantization
+        if isinstance(quant_config, WeightOnlyConfig):
+            self.is_use_fused_moe = False
+
+        if params_dtype is None:
+            params_dtype = torch.get_default_dtype()
+        self.params_dtype = params_dtype
+
+        if VLLM_AVG_MOE_EN and not SparseMoeMlp.is_expert_avg:
+            n_tokens = SparseMoeMlp.max_batched_token
+            val = 1.0 / float(num_experts)
+            SparseMoeMlp.reduce_weight    = torch.full((n_tokens, top_k), val, device="mlu", dtype=torch.float32)
+            #avg 0,20,40,60,80...120,140,  1,21,...121,141, 2...142,  ......  19,...159,  0,20,......
+            import math
+            batch_table = math.ceil(n_tokens * top_k / num_experts) * num_experts
+            hi_val = batch_table // num_experts
+            table = (torch.arange(hi_val * num_experts, device="mlu", dtype=torch.int32) % num_experts).view(hi_val, expert_group, num_experts // expert_group).transpose(1, 2)
+            SparseMoeMlp.expert_id        = table.flatten()[:n_tokens*top_k].view(n_tokens, top_k)
+            SparseMoeMlp.is_expert_avg = True
+
+        # NOTE: The bias for fc2 is only applied on tp_rank 0. If we added it on all nodes the allreduce() would
+        # contain multiple copies of the bias. The bias on other node will be ignored, and may be set to nullptr
+        self.skip_bias_add = True if self.tp_rank > 0 else False
+
+        assert self.intermediate_size % self.tp_size == 0, (
+            f"need intermediate_size:{self.intermediate_size} % tp_size:{self.tp_size} == 0")
+
+        self.num_experts_per_rank = self.num_total_experts
+
+        self.start_expert_id = 0
+        self.end_expert_id = self.start_expert_id + self.num_experts_per_rank
+
+        # Gate always runs at half / full precision for now.
+        self.gate = ReplicatedLinear(self.hidden_size,
+                                     self.num_total_experts,
+                                     bias=False,
+                                     params_dtype=self.params_dtype,
+                                     quant_config=None)
+        if topk_method == "noaux_tc":
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.empty(self.num_total_experts, device="mlu"))
+        else:
+            self.gate.e_score_correction_bias = None
+        self.is_fp8_block_wise = isinstance(self.quant_config, Fp8Config) and self.quant_config.weight_block_size is not None
+        if self.is_fp8_block_wise:
+            self.experts = FusedMoE(
+                num_experts=self.num_experts_per_rank,
+                top_k=self.top_k,
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size,
+                reduce_results=False,
+                renormalize=self.renormalize,
+                quant_config=self.quant_config,
+                use_grouped_topk=True,
+                num_expert_group=self.expert_group,
+                topk_group=self.topk_group,
+                prefix=f"experts",
+                scoring_func=self.scoring_func,
+                e_score_correction_bias=self.gate.e_score_correction_bias)
+        else:
+            self.experts = nn.ModuleList([
+                FeedForward(hidden_size=self.hidden_size,
+                            intermediate_size=self.intermediate_size,
+                        hidden_act=self.hidden_act,
+                        up_proj_name=self.up_proj_name,
+                        is_gated=self.is_gated,
+                        down_proj_name=self.down_proj_name,
+                        bias=self.has_bias,
+                        quant_config=self.quant_config,
+                        skip_bias_add=self.skip_bias_add,
+                        reduce_results=False,
+                        prefix=f"experts.{idx}") for idx in range(self.num_experts_per_rank)
+        ])
+
+        self.init_pack_param()
+
+
+    def init_pack_param(self):
+        self.w13 = None
+        self.w2 = None
+        self.b13 = None
+        self.b2 = None
+        self.w13_scale = None
+        self.w2_scale = None
+        self.a13_scale = None
+        self.a2_scale = None
+        self.pack_params_done = False
+
+
+    def map_param_data(self, param_list, is_use_first_data=False):
+        if len(param_list) == 0:
+            return None
+
+        if is_use_first_data or len(param_list) == 1:
+            first_data = param_list[0].data
+            for param in param_list[1: -1]:
+                param.data = first_data
+            if is_use_first_data:
+                out_param = first_data.view_as(param_list[0])
+            else:
+                out_param = first_data.view(len(param_list), *first_data.shape)
+        else:
+            packed_param = torch._utils._flatten_dense_tensors(param_list)
+            data_list = torch._utils._unflatten_dense_tensors(packed_param, param_list)
+            for data, param in zip(data_list, param_list):
+                param.data = data
+            out_param = packed_param.view(len(param_list), *data_list[0].shape)
+
+        torch.mlu.empty_cache()
+
+        return out_param
+
+
+    def pack_unquantized_params(self, w13, w2, b13, b2):
+        for expert in self.experts:
+            up_proj = getattr(expert, self.up_proj_name)
+            down_proj = getattr(expert, self.down_proj_name)
+            w13.append(up_proj.weight)
+            w2.append(down_proj.weight)
+            if self.has_bias:
+                b13.append(up_proj.bias)
+                b2.append(down_proj.bias)
+
+
+    def pack_smoothquant_params(self, w13, w2, b13, b2, w13_scale, w2_scale, a13_scale, a2_scale):
+        for expert in self.experts:
+            up_proj = getattr(expert, self.up_proj_name)
+            down_proj = getattr(expert, self.down_proj_name)
+            w13.append(up_proj.qweight)
+            w2.append(down_proj.qweight)
+            if self.has_bias:
+                b13.append(up_proj.bias)
+                b2.append(down_proj.bias)
+            w13_scale.append(up_proj.per_channel_scale)
+            w2_scale.append(down_proj.per_channel_scale)
+            if self.quant_config.input_quant_method == "per_token":
+                a13_scale.append(up_proj.smooth)
+                a2_scale.append(down_proj.smooth)
+            else:
+                a13_scale.append(up_proj.scale_to_int)
+                a2_scale.append(down_proj.scale_to_int)
+
+
+    def pack_weightonly_params(self, w13, w2, b13, b2, w13_scale, w2_scale):
+        for expert in self.experts:
+            up_proj = getattr(expert, self.up_proj_name)
+            down_proj = getattr(expert, self.down_proj_name)
+            w13.append(up_proj.qweight)
+            w2.append(down_proj.qweight)
+            if self.has_bias:
+                b13.append(up_proj.bias)
+                b2.append(down_proj.bias)
+            w13_scale.append(up_proj.scales)
+            w2_scale.append(down_proj.scales)
+
+    def pack_fp8_params_without_activation_scheme(self, w13, w2, b13, b2, w13_scale, w2_scale):
+        for expert in self.experts:
+            up_proj = getattr(expert, self.up_proj_name)
+            down_proj = getattr(expert, self.down_proj_name)
+            w13.append(up_proj.weight)
+            w2.append(down_proj.weight)
+            if self.has_bias:
+                b13.append(up_proj.bias)
+                b2.append(down_proj.bias)
+            w13_scale.append(up_proj.weight_scale)
+            w2_scale.append(down_proj.weight_scale)
+
+
+    def pack_params(self):
+        if self.pack_params_done or self.is_fp8_block_wise:
+            return
+
+        w13 = []
+        w2 = []
+        b13 = []
+        b2 = []
+        w13_scale = []
+        w2_scale = []
+        a13_scale = []
+        a2_scale = []
+
+        if self.quant_config is None:
+            self.pack_unquantized_params(w13, w2, b13, b2)
+        elif isinstance(self.quant_config, SmoothQuantConfig):
+            self.pack_smoothquant_params(w13, w2, b13, b2, w13_scale, w2_scale, a13_scale, a2_scale)
+        elif isinstance(self.quant_config, WeightOnlyConfig):
+            self.pack_weightonly_params(w13, w2, b13, b2, w13_scale, w2_scale)
+        elif isinstance(self.quant_config, Fp8Config) and self.quant_config.activation_scheme == 'dynamic':
+            self.pack_fp8_params_without_activation_scheme(w13, w2, b13, b2, w13_scale, w2_scale)
+        else:
+            raise ValueError(f'Unsupported quantization:{self.quant_config}')
+
+        # pack weight
+        self.w13 = self.map_param_data(w13)
+        self.w2 = self.map_param_data(w2)
+
+        # pack bias
+        if self.has_bias:
+            self.b13 = self.map_param_data(b13)
+            # NOTE: The bias for fc2 is only applied on tp_rank 0. If we added it on all nodes the allreduce() would
+            # contain multiple copies of the bias. The bias on other node will be ignored, and may be set to nullptr
+            if self.skip_bias_add is False:
+                self.b2 = self.map_param_data(b2)
+
+
+        # pack weight scale
+        if len(w13_scale) > 0:
+            self.w13_scale = self.map_param_data(w13_scale)
+        if len(w2_scale) > 0:
+            self.w2_scale = self.map_param_data(w2_scale)
+
+        # pack activate scale
+        if len(a13_scale) > 0:
+            self.a13_scale = self.map_param_data(a13_scale)
+        if len(a2_scale) > 0:
+            self.a2_scale = self.map_param_data(a2_scale)
+            
+        if isinstance(self.quant_config, SmoothQuantConfig) and self.quant_config.group_size > 1 and self.is_use_fused_moe:
+            assert self.w13_scale is not None and self.w2_scale is not None, "w13_scale and w2_scale must be not None"
+            self.w13_scale = self.w13_scale.permute(2, 0, 1).contiguous()
+            self.w2_scale = self.w2_scale.permute(2, 0, 1).contiguous()
+
+        # pack smooth variables for moe_quantize if fp8
+        # FIXME: replace smooth to None after tmo supports.
+        if isinstance(self.quant_config, Fp8Config):
+            expert_size = self.w13.shape[0]
+            fp8_smooth_2_hidden_size = self.w13.shape[1] // 2 if self.is_gated else self.w13.shape[1]
+            self.fp8_smooth_1 = torch.ones([expert_size, self.hidden_size], device=self.w13.device, dtype=torch.float32)
+            self.fp8_smooth_2 = torch.ones([expert_size, fp8_smooth_2_hidden_size], device=self.w13.device, dtype=torch.float32)
+        
+        self.pack_params_done = True
+
+
+    def forward(self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None) -> torch.Tensor:
+        orig_hidden_states_shape = hidden_states.shape
+        hidden_states = hidden_states.view(-1, self.hidden_size)
+        # expert_logits: [num_tokens, self.num_experts_per_rank]
+        expert_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.forward_experts(hidden_states, expert_logits, residual)
+
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                final_hidden_states)
+
+        output = final_hidden_states.view(orig_hidden_states_shape)
+        return output
+
+
+    def forward_experts(self, hidden_states, expert_logits, residual: Optional[torch.Tensor] = None,
+                        shared_output: Optional[torch.Tensor] = None):
+        assert not (residual is not None and shared_output is not None)
+        residual_ = None if self.tp_rank > 0 else residual
+
+        # change only for deepseek_model without residual_
+        if shared_output is not None:
+            residual_ = shared_output
+        if self.is_use_fused_moe and self.expert_group != 1:
+            final_hidden_states = self.forward_group_experts(hidden_states, expert_logits, residual_)
+        elif self.is_use_fused_moe:
+            self.pack_params()
+            final_hidden_states = mlu_ops.fused_moe(hidden_states=hidden_states,
+                                                    gating_output=expert_logits,
+                                                    w1=self.w13,
+                                                    w2=self.w2,
+                                                    bias1=self.b13,
+                                                    bias2=self.b2,
+                                                    residual=residual_,
+                                                    input_smooth=self.a13_scale,
+                                                    act_smooth=self.a2_scale,
+                                                    w1_scale=self.w13_scale,
+                                                    w2_scale=self.w2_scale,
+                                                    topk=self.top_k,
+                                                    renormalize=self.renormalize,
+                                                    gated=self.is_gated,
+                                                    act_mode=self.hidden_act,
+                                                    start_expert_id=self.start_expert_id)
+        else:
+            final_hidden_states = self.forward_experts_nofused(hidden_states, expert_logits)
+            if residual_ is not None:
+                final_hidden_states = final_hidden_states + residual_
+        return final_hidden_states
+
+
+    def forward_experts_nofused(self, hidden_states, expert_logits):
+        hidden_states_shape = hidden_states.shape
+        if self.scoring_func == "softmax":
+            topk_values, topk_indices = self.topk_softmax(expert_logits)
+        elif self.scoring_func == "sigmoid":
+            gating_output = expert_logits.to(torch.float32)
+            gating_output = gating_output.view(-1, gating_output.size(-1))
+            topk_values, topk_indices = grouped_topk(hidden_states, gating_output, self.top_k, self.renormalize,
+                                                    self.expert_group, self.topk_group, self.scoring_func,
+                                                    self.gate.e_score_correction_bias)
+            topk_values = topk_values.to(hidden_states.dtype)
+            topk_indices = topk_indices.to(torch.int64)
+        expand_gather_idx, scatter_idx, expand_token_count, cusum_token_count = self.generate_gather_idx(
+            topk_indices)
+        # no expert is routed, then expand_gather_idx, expand_scatter_idx has no item,
+        # expand_token_count and expand_cusum_token_count has item but the value is all zero
+        # so this rank should only return final_hidden_states with zero value
+        if expand_gather_idx.numel() == 0:
+            final_hidden_states = torch.zeros_like(hidden_states,
+                                                   dtype=hidden_states.dtype,
+                                                   device=hidden_states.device)
+            return final_hidden_states
+
+        expand_hidden_states = self.expand_input(hidden_states, expand_gather_idx)
+
+        expand_output_list = []
+        expand_cusum_token_count = cusum_token_count[self.start_expert_id:self.end_expert_id +
+                                                     1] - cusum_token_count[self.start_expert_id]
+        for expert_idx, num_tokens_per_expert in enumerate(expand_token_count):
+            if num_tokens_per_expert > 0:
+                expert_hidden_states = expand_hidden_states[
+                    expand_cusum_token_count[expert_idx]:expand_cusum_token_count[expert_idx + 1]]
+                expert_output = self.experts[expert_idx](expert_hidden_states)
+                expert_output = expert_output[0] if isinstance(expert_output, (tuple, list)) else expert_output
+                expand_output_list.append(expert_output)
+        expand_output = torch.cat(expand_output_list, dim=0)
+        final_hidden_states = self.combine_moe(expand_output, scatter_idx, cusum_token_count, hidden_states_shape,
+                                               topk_values)
+
+        return final_hidden_states
+
+    def forward_group_experts(self, hidden_states, expert_logits, residual_):
+        if self.is_fp8_block_wise:
+            output = self.experts(hidden_states=hidden_states,
+                                router_logits=expert_logits) * self.routed_scaling_factor
+            if residual_ is not None:
+                output = output + residual_
+            return output
+        is_fp8_quant = isinstance(self.quant_config, Fp8Config)
+        ori_input_shape = hidden_states.shape
+        dtype = hidden_states.dtype
+        self.pack_params()
+        gating_output=expert_logits.to(torch.float32)
+        w1=self.w13
+        w2=self.w2
+        bias1=self.b13
+        bias2=self.b2
+        input_smooth=self.a13_scale
+        act_smooth=self.a2_scale
+        w1_scale=self.w13_scale
+        w2_scale=self.w2_scale
+        topk=self.top_k
+        renormalized=self.renormalize
+        gated=self.is_gated
+        act_mode=self.hidden_act
+        quant_input=None
+
+        start_expert_id=self.start_expert_id
+        expert_num = gating_output.size(-1)
+        expert_size = w1.size(0)
+        max_m = hidden_states.shape[0]
+        hidden_states = hidden_states.view(-1, hidden_states.size(-1))
+        gating_output = gating_output.view(-1, gating_output.size(-1))
+        residual_ = residual_.view(-1, residual_.size(-1)) if residual_ is not None else None
+        # Check smooth quant parameters.
+        per_token_sq = False
+        if not is_fp8_quant:
+            check_list = [input_smooth, act_smooth, w1_scale, w2_scale]
+            if all(x is not None for x in check_list):
+                per_token_sq = True
+
+            if not (all(x is None for x in check_list) or all(x is not None for x in check_list)):
+                raise ValueError("input_smooth, act_smooth, w1_scale and w2_scale must be present "
+                                "and absent at the same time.")
+
+        # softmax_topk
+        if self.scoring_func == "softmax":
+            reduce_weight, expert_id = mlu_ops.moe_softmax_topk(gating_output, topk, renormalized, self.expert_group,
+                                                                self.topk_group, route_scale=self.routed_scaling_factor)
+        elif self.scoring_func == "sigmoid":
+            reduce_weight, expert_id = mlu_ops.moe_sigmoid_topk(gating_output, topk, renormalized,
+                                                                self.expert_group, self.topk_group,
+                                                                self.routed_scaling_factor,
+                                                                self.gate.e_score_correction_bias)
+        else:
+            raise ValueError(f"Unsupported scoring function: {self.scoring_func}")
+
+        if VLLM_AVG_MOE_EN:
+            n_tokens = hidden_states.shape[0]
+            reduce_weight = SparseMoeMlp.reduce_weight[:n_tokens]
+            expert_id     = SparseMoeMlp.expert_id[:n_tokens]
+        # gen_idx
+        expand_idx, combine_idx, token_count, cusum_token_count = mlu_ops.moe_gen_idx(expert_id, expert_num)
+        # check quant
+        if is_fp8_quant and self.quant_config.activation_quant_method == 'per_token':
+            quant_input, input_scale = mlu_ops.moe_quantize(
+                hidden_states,
+                self.fp8_smooth_1,
+                zero=None,
+                token_count=token_count[start_expert_id:start_expert_id+expert_size],
+                gather_index=expand_idx,
+                gather_index_start_position=cusum_token_count[start_expert_id].unsqueeze(0),
+                output=None,
+                output_scale=None,
+                dynamic_quant=True,
+                quant_type=torch.float8_e4m3fn
+            )
+        elif per_token_sq:
+            quant_input, input_scale = mlu_ops.moe_quantize(hidden_states,
+                input_smooth, None, token_count[start_expert_id:start_expert_id+expert_size], expand_idx,
+                cusum_token_count[start_expert_id].unsqueeze(0))
+        else:
+            expand_hidden_states = mlu_ops.moe_expand_input(hidden_states, expand_idx,
+                    cusum_token_count, start_expert_id, expert_size)
+
+        if (is_fp8_quant and self.quant_config.activation_quant_method == 'per_token') or per_token_sq:
+            gemm_out = mlu_ops.smooth_quant_group_gemm(quant_input, w1,
+                                                        token_count[start_expert_id:start_expert_id+expert_size],
+                                                        None, None, None, None,
+                                                        input_scale, w1_scale, dtype, max_m)
+        else:
+            gemm_out = mlu_ops.group_gemm(expand_hidden_states, w1,
+                                       token_count[start_expert_id:start_expert_id+expert_size],
+                                       None, None, None, None, max_m)
+        # add_bias_active
+        if is_fp8_quant and self.quant_config.activation_quant_method == 'per_token':
+            act_out = mlu_ops.moe_active(gemm_out, act_mode, gated, gemm_out[:,:gemm_out.shape[-1]//2], bias=bias1, cusum_token_count=cusum_token_count, start_expert_id=start_expert_id, expert_size=expert_size)
+            quant_input, input_scale = mlu_ops.moe_quantize(
+                act_out,
+                self.fp8_smooth_2,
+                zero=None,
+                token_count=token_count[start_expert_id:start_expert_id+expert_size],
+                gather_index=None,
+                gather_index_start_position=None,
+                output=quant_input[:,:act_out.shape[-1]],
+                output_scale=None,
+                dynamic_quant=True,
+                quant_type=torch.float8_e4m3fn
+            )
+        elif per_token_sq:
+            quant_input = quant_input[:, :gemm_out.shape[-1] // 2]
+            input_scale = input_scale[:gemm_out.shape[0]]
+            quant_input, input_scale = mlu_ops.moe_quantize(gemm_out, act_smooth, None,
+                                                            token_count[start_expert_id:start_expert_id+expert_size],
+                                                            output=quant_input,
+                                                            output_scale=input_scale,
+                                                            act_mode=act_mode,
+                                                            is_gated=self.is_gated)
+    
+        if (is_fp8_quant and self.quant_config.activation_quant_method == 'per_token') or per_token_sq:
+            # Remove the reference to gemm_out tensor. 
+            # If that was the only reference, the tensor’s memory becomes eligible for deallocation
+            # So that we can reuse this memory for the new allocation of next gemm operation
+            del gemm_out
+            gemm_out = mlu_ops.smooth_quant_group_gemm(quant_input, w2,
+                                                    token_count[start_expert_id:start_expert_id+expert_size],
+                                                    None, None, None, None, input_scale, w2_scale, dtype, max_m)
+        else:
+            act_out = mlu_ops.moe_active(gemm_out, act_mode, gated, gemm_out[:,:gemm_out.shape[-1]//2], bias1, cusum_token_count, start_expert_id, expert_size)
+            gemm_out = mlu_ops.group_gemm(act_out, w2,
+                                       token_count[start_expert_id:start_expert_id+expert_size],
+                                       None, None, None, None, max_m)
+
+        # we reuse the memory of hidden_states to store the output
+        output = mlu_ops.moe_combine_result(gemm_out, reduce_weight, combine_idx,
+                                        residual_, cusum_token_count, start_expert_id,
+                                        expert_size, bias2, output=hidden_states)
+        return output.view(ori_input_shape)
+
+
+    def topk_softmax(self, expert_logits):
+        # expert_logits: [num_tokens, self.num_experts_per_rank]
+        # topk_values: [num_tokens, self.top_k]
+        # topk_indices: [num_tokens, self.top_k]
+        if self.renormalize:
+            topk_values, topk_indices = torch.topk(expert_logits, self.top_k, dim=-1)
+            topk_values = torch.softmax(topk_values, -1)
+        else:
+            router_probs = torch.softmax(expert_logits, -1)
+            topk_values, topk_indices = torch.topk(router_probs, self.top_k, dim=-1)
+
+        return topk_values, topk_indices
+
+
+    def generate_gather_idx(self, topk_indices):
+        device = topk_indices.device
+        # gather_expand_idx: [num_tokens * self.top_k]
+        sorted_expert_id, indices = topk_indices.flatten().sort()
+        gather_idx = indices // self.top_k
+
+        seqs = torch.arange(indices.numel(), dtype=indices.dtype, device=indices.device)
+        scatter_idx=torch.zeros((indices.numel(),), dtype=seqs.dtype, device=seqs.device).scatter(0, indices, seqs)
+
+        # token_count: [self.num_experts_per_rank]
+        partial_token_index, partial_token_count = sorted_expert_id.unique(sorted=True, return_counts=True)
+        zero_token_count = torch.zeros(self.num_total_experts, dtype=partial_token_count.dtype, device=device)
+        token_count = zero_token_count.scatter(dim=0, index=partial_token_index, src=partial_token_count)
+        # cusum_token_count: [self.num_experts_per_rank + 1]
+        cusum_token_count = torch.cat(
+            [torch.tensor([0], dtype=token_count.dtype, device=device),
+             token_count.cumsum(dim=0)])
+
+        num_tokens_before_expert = cusum_token_count[self.start_expert_id]
+        num_tokens_including_expert = cusum_token_count[self.end_expert_id]
+
+        expand_gather_idx = gather_idx[num_tokens_before_expert:num_tokens_including_expert]
+        expand_token_count = token_count[self.start_expert_id:self.end_expert_id]
+
+        return expand_gather_idx, scatter_idx, expand_token_count, cusum_token_count
+
+
+    def expand_input(self, hidden_states, expand_gather_idx):
+        expand_hidden_states = hidden_states[expand_gather_idx]
+        return expand_hidden_states
+
+
+    def combine_moe(self, expand_output, scatter_idx, cusum_token_count, hidden_states_shape, topk_values):
+        num_tokens, hidden_size = hidden_states_shape
+        num_tokens_before_expert = cusum_token_count[self.start_expert_id]
+        num_tokens_after_expert = cusum_token_count[-1] - cusum_token_count[self.end_expert_id]
+
+        expand_output_before_expert = torch.zeros((num_tokens_before_expert, hidden_size),
+                                                   dtype=expand_output.dtype,
+                                                   device=expand_output.device)
+        expand_output_after_expert = torch.zeros((num_tokens_after_expert, hidden_size),
+                                                   dtype=expand_output.dtype,
+                                                   device=expand_output.device)
+        unscatted_output = torch.cat([expand_output_before_expert, expand_output, expand_output_after_expert], dim=0)
+        scatter_output = unscatted_output[scatter_idx]
+        hidden_states_weight = topk_values.flatten().unsqueeze(-1)
+        weighted_hidden_states = scatter_output * hidden_states_weight
+        unreduced_hidden_states = weighted_hidden_states.view(num_tokens, self.top_k, hidden_size)
+        final_hidden_states = unreduced_hidden_states.sum(dim=1)
+
+        return final_hidden_states
diff --git a/vllm_mlu/vllm_mlu/model_executor/model_loader/__init__.py b/vllm_mlu/vllm_mlu/model_executor/model_loader/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/model_executor/model_loader/tensorizer.py b/vllm_mlu/vllm_mlu/model_executor/model_loader/tensorizer.py
new file mode 100644
index 000000000..fd47b4946
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/model_loader/tensorizer.py
@@ -0,0 +1,71 @@
+import time
+import torch
+
+from vllm.model_executor.model_loader.tensorizer import (TensorizerAgent,
+                                                         TensorDeserializer,
+                                                         get_mem_usage,
+                                                         _read_stream,
+                                                         convert_bytes)
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.logger import init_logger
+logger = init_logger(__name__)
+
+
+def vllm__model_executor__model_loader__tensorizer__TensorizerAgent__deserialize(self):
+    """
+    Deserialize the model using the TensorDeserializer. This method is
+    specifically for vLLM models using tensorizer's plaid_mode.
+
+    The deserializer makes use of tensorizer_args.stream_params
+    to configure the behavior of the stream when loading tensors from a
+    serialized model. The deserializer_params are used to configure the
+    behavior of the TensorDeserializer when loading tensors themselves.
+    Documentation on these params can be found in TensorizerArgs
+
+    Returns:
+        nn.Module: The deserialized model.
+    """
+    before_mem = get_mem_usage()
+    start = time.perf_counter()
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use mlu device
+    '''
+    with _read_stream(
+            self.tensorizer_config.tensorizer_uri,
+            **self.tensorizer_args.stream_params
+    ) as stream, TensorDeserializer(
+            stream,
+            dtype=self.tensorizer_config.dtype,
+            device=f'mlu:{torch.mlu.current_device()}',
+            **self.tensorizer_args.deserializer_params) as deserializer:
+        deserializer.load_into_module(self.model)
+        end = time.perf_counter()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    total_bytes_str = convert_bytes(deserializer.total_tensor_bytes)
+    duration = end - start
+    per_second = convert_bytes(deserializer.total_tensor_bytes / duration)
+    after_mem = get_mem_usage()
+    deserializer.close()
+    logger.info("Deserialized %s in %0.2fs, %s/s", total_bytes_str,
+                end - start, per_second)
+    logger.info("Memory usage before: %s", before_mem)
+    logger.info("Memory usage after: %s", after_mem)
+
+    self._check_tensors_on_meta_device()
+    self._resize_lora_embeddings()
+    del self.model.vllm_tensorized_marker
+    return self.model.eval()
+
+
+MluHijackObject.apply_hijack(
+    TensorizerAgent,
+    TensorizerAgent.deserialize,
+    vllm__model_executor__model_loader__tensorizer__TensorizerAgent__deserialize
+)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/__init__.py b/vllm_mlu/vllm_mlu/model_executor/models/__init__.py
new file mode 100755
index 000000000..8b1378917
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/__init__.py
@@ -0,0 +1 @@
+
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/baichuan.py b/vllm_mlu/vllm_mlu/model_executor/models/baichuan.py
new file mode 100644
index 000000000..ba4bafef0
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/baichuan.py
@@ -0,0 +1,313 @@
+import torch
+from typing import List, Optional, Union
+from transformers import PretrainedConfig
+from vllm.attention import Attention, AttentionMetadata
+from vllm.config import CacheConfig
+from vllm.distributed import (get_tensor_model_parallel_rank,
+                              get_tensor_model_parallel_world_size)
+from vllm.sequence import IntermediateTensors
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (QKVParallelLinear,
+                                               RowParallelLinear)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.models.baichuan import (
+    _get_alibi_slopes, BaiChuanAttention,
+    BaiChuanDecoderLayer, BaiChuanModel)
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__baichuan__BaiChuanAttention__init__(
+    self,
+    hidden_size: int,
+    num_heads: int,
+    position_embedding: str,
+    rope_theta: float = 10000,
+    max_position_embeddings: int = 8192,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(BaiChuanAttention, self).__init__()
+    self.hidden_size = hidden_size
+    tensor_model_parallel_world_size = get_tensor_model_parallel_world_size(
+    )
+    self.total_num_heads = num_heads
+    assert self.total_num_heads % tensor_model_parallel_world_size == 0
+    self.num_heads = (self.total_num_heads //
+                        tensor_model_parallel_world_size)
+    self.head_dim = hidden_size // self.total_num_heads
+    self.postion_embedding = position_embedding
+    self.rope_theta = rope_theta
+    self.max_position_embeddings = max_position_embeddings
+
+    # pylint: disable=invalid-name
+    self.W_pack = QKVParallelLinear(
+        hidden_size,
+        self.head_dim,
+        self.total_num_heads,
+        self.total_num_heads,
+        bias=False,
+        quant_config=quant_config,
+    )
+    self.o_proj = RowParallelLinear(
+        self.total_num_heads * self.head_dim,
+        hidden_size,
+        bias=False,
+        quant_config=quant_config,
+    )
+    # Create the alibi slopes and slice them.
+    if self.postion_embedding == "ALIBI":
+        tp_rank = get_tensor_model_parallel_rank()
+        head_start = tp_rank * self.num_heads
+        head_end = (tp_rank + 1) * self.num_heads
+        alibi_slopes = _get_alibi_slopes(self.total_num_heads)
+        alibi_slopes = alibi_slopes[head_start:head_end].tolist()
+
+        scaling = self.head_dim**-0.5
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: add cache_config to support kv8
+        '''
+        self.attn = Attention(self.num_heads,
+                              self.head_dim,
+                              scaling,
+                              alibi_slopes=alibi_slopes,
+                              cache_config=cache_config,
+                              quant_config=quant_config,
+                              prefix=f"{prefix}.attn")
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    else:
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=self.max_position_embeddings,
+            base=self.rope_theta,
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attn = Attention(self.num_heads,
+                              self.head_dim,
+                              self.scaling,
+                              cache_config=cache_config,
+                              quant_config=quant_config,
+                              prefix=f"{prefix}.attn")
+
+
+def vllm__module_executor__models__baichuan__BaiChuanAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.W_pack(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.chunk(chunks=3, dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk, _ = qkv.split([self.num_heads * self.head_dim * 2, self.num_heads * self.head_dim], dim=-1)
+    if self.postion_embedding != "ALIBI":
+        self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__baichuan__BaiChuanDecoderLayer__init__(
+    self,
+    config: PretrainedConfig,
+    position_embedding: str,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = ""
+):
+    super(BaiChuanDecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    rope_theta = getattr(config, "rope_theta", 10000)
+    max_position_embeddings = getattr(config, "max_position_embeddings",
+                                        8192)
+    self.self_attn = BaiChuanAttention(
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        position_embedding=position_embedding,
+        rope_theta=rope_theta,
+        max_position_embeddings=max_position_embeddings,
+        cache_config=cache_config,
+        quant_config=quant_config,
+        prefix=f"{prefix}.self_attn",
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.mlp = FeedForward(
+        hidden_size=config.hidden_size,
+        intermediate_size=config.intermediate_size,
+        hidden_act='silu',
+        up_proj_name='gate_up_proj',
+        is_gated=True,
+        down_proj_name='down_proj',
+        bias=False,
+        quant_config=quant_config
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                    eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf per-tensor sq cases if suitable
+    '''
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attn.W_pack.quant_method.skip_quant_input = True
+        self.mlp.gate_up_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+        self.quant_fusion_mlp_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__baichuan__BaiChuanDecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    mlp_layernorm = self.post_attention_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attn.W_pack.smooth
+                mlp_quant_scale = self.mlp.gate_up_proj.smooth
+            else:
+                attn_quant_scale = self.self_attn.W_pack.scale_to_int
+                mlp_quant_scale = self.mlp.gate_up_proj.scale_to_int
+
+        if self.quant_fusion_attn_layernorm is None:
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
+                self.post_attention_layernorm, mlp_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+    return decoder_layer_forward_base(
+        positions=positions,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.self_attn,
+        post_layernorm=mlp_layernorm,
+        mlp=self.mlp,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+        post_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__baichuan__BaiChuanModel__forward(
+    self,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors],
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids=input_ids,
+        positions=positions,
+        intermediate_tensors=intermediate_tensors,
+        inputs_embeds=inputs_embeds,
+        layers=self.layers,
+        start_layer=self.start_layer,
+        end_layer=self.end_layer,
+        get_input_embeddings=self.embed_tokens,
+        norm=self.norm
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(BaiChuanAttention,
+                             BaiChuanAttention.__init__,
+                             vllm__module_executor__models__baichuan__BaiChuanAttention__init__)
+MluHijackObject.apply_hijack(BaiChuanAttention,
+                             BaiChuanAttention.forward,
+                             vllm__module_executor__models__baichuan__BaiChuanAttention__forward)
+MluHijackObject.apply_hijack(BaiChuanDecoderLayer,
+                             BaiChuanDecoderLayer.__init__,
+                             vllm__module_executor__models__baichuan__BaiChuanDecoderLayer__init__)
+MluHijackObject.apply_hijack(BaiChuanDecoderLayer,
+                             BaiChuanDecoderLayer.forward,
+                             vllm__module_executor__models__baichuan__BaiChuanDecoderLayer__forward)
+MluHijackObject.apply_hijack(BaiChuanModel,
+                             BaiChuanModel.forward,
+                             vllm__module_executor__models__baichuan__BaiChuanModel__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/bert.py b/vllm_mlu/vllm_mlu/model_executor/models/bert.py
new file mode 100644
index 000000000..b73b850a9
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/bert.py
@@ -0,0 +1,54 @@
+import torch
+
+from vllm.model_executor.models.bert import BertSelfOutput, BertOutput
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__module_executor__models__bert__BertSelfOutput__forward(
+    self,
+    hidden_states: torch.Tensor,
+    input_tensor: torch.Tensor
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: apply residual fusion
+    '''
+    hidden_states, _ = self.dense(hidden_states, input_tensor)
+    hidden_states = self.LayerNorm(hidden_states)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return hidden_states
+
+
+def vllm__module_executor__models__bert__BertOutput__forward(
+    self,
+    hidden_states: torch.Tensor,
+    input_tensor: torch.Tensor
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: apply residual fusion
+    '''
+    hidden_states, _ = self.dense(hidden_states, input_tensor)
+    hidden_states = self.LayerNorm(hidden_states)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return hidden_states
+
+
+MluHijackObject.apply_hijack(BertSelfOutput,
+                             BertSelfOutput.forward,
+                             vllm__module_executor__models__bert__BertSelfOutput__forward)
+MluHijackObject.apply_hijack(BertOutput,
+                             BertOutput.forward,
+                             vllm__module_executor__models__bert__BertOutput__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/bloom.py b/vllm_mlu/vllm_mlu/model_executor/models/bloom.py
new file mode 100644
index 000000000..1d2229240
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/bloom.py
@@ -0,0 +1,169 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from transformers import BloomConfig
+
+from vllm.config import CacheConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+from vllm.model_executor.models.bloom import BloomAttention, BloomBlock
+from vllm.attention import AttentionMetadata
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base,
+    is_per_tensor_smoothquant,
+    is_per_token_smoothquant,
+    quant_fusion_with_layernorm
+)
+
+logger = init_logger(__name__)
+
+def vllm__module_executor__models__bloom__BloomAttention__forward(
+    self,
+    position_ids: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    del position_ids  # Unused.
+    qkv, _ = self.query_key_value(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.chunk(chunks=3, dim=-1)
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    output, _ = self.dense(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+def vllm__module_executor__models__bloom__BloomBlock__init__(
+    self,
+    config: BloomConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(BloomBlock, self).__init__()
+    hidden_size = config.hidden_size
+
+    self.input_layernorm = nn.LayerNorm(hidden_size,
+                                        eps=config.layer_norm_epsilon)
+    self.self_attention = BloomAttention(config,
+                                         cache_config,
+                                         quant_config,
+                                         prefix=f"{prefix}.self_attention")
+    self.post_attention_layernorm = nn.LayerNorm(
+            hidden_size, eps=config.layer_norm_epsilon)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.mlp = FeedForward(hidden_size=hidden_size,
+                           intermediate_size=hidden_size * 4,
+                           hidden_act='gelu',
+                           up_proj_name="dense_h_to_4h",
+                           is_gated=False,
+                           down_proj_name="dense_4h_to_h",
+                           bias=True,
+                           quant_config=quant_config)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.apply_residual_connection_post_layernorm = (
+            config.apply_residual_connection_post_layernorm)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf sq cases if suitable
+    '''
+    self.is_per_tesnor_sq_perf_cases = (is_per_tensor_smoothquant(quant_config) and
+        not self.apply_residual_connection_post_layernorm)
+    self.is_per_token_sq_perf_cases = (is_per_token_smoothquant(quant_config) and
+        not self.apply_residual_connection_post_layernorm)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attention.query_key_value.quant_method.skip_quant_input = True
+        self.mlp.dense_h_to_4h.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+        self.quant_fusion_mlp_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__bloom__BloomBlock__forward(
+    self,
+    position_ids: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    mlp_layernorm = self.post_attention_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attention.query_key_value.smooth
+                mlp_quant_scale = self.mlp.dense_h_to_4h.smooth
+            else:
+                attn_quant_scale = self.self_attention.query_key_value.scale_to_int
+                mlp_quant_scale = self.mlp.dense_h_to_4h.scale_to_int
+
+            self.quant_fusion_attn_layernorm = quant_fusion_with_layernorm(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+            self.quant_fusion_mlp_layernorm = quant_fusion_with_layernorm(
+                self.post_attention_layernorm, mlp_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+    return decoder_layer_forward_base(
+        positions=position_ids,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.self_attention,
+        post_layernorm=mlp_layernorm,
+        mlp=self.mlp,
+        apply_residual_connection_post_layernorm=self.apply_residual_connection_post_layernorm,
+        position_name='position_ids',
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+        post_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+MluHijackObject.apply_hijack(BloomAttention,
+                             BloomAttention.forward,
+                             vllm__module_executor__models__bloom__BloomAttention__forward)
+MluHijackObject.apply_hijack(BloomBlock,
+                             BloomBlock.__init__,
+                             vllm__module_executor__models__bloom__BloomBlock__init__)
+MluHijackObject.apply_hijack(BloomBlock,
+                             BloomBlock.forward,
+                             vllm__module_executor__models__bloom__BloomBlock__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/chatglm.py b/vllm_mlu/vllm_mlu/model_executor/models/chatglm.py
new file mode 100644
index 000000000..87d5a21a3
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/chatglm.py
@@ -0,0 +1,190 @@
+import torch
+
+from torch.nn import LayerNorm
+from typing import Optional
+from vllm.attention import AttentionMetadata
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm.logger import init_logger
+
+from vllm.transformers_utils.configs import ChatGLMConfig
+from vllm.model_executor.models.chatglm import GLMAttention, GLMBlock
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base,
+    is_per_tensor_smoothquant,
+    is_per_token_smoothquant,
+    quant_fusion_with_layernorm,
+    quant_fusion_with_rmsnorm
+)
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__chatglm__GLMAttention__forward(
+    self,
+    hidden_states: torch.Tensor,
+    position_ids: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.query_key_value(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+    self.rotary_emb(position_ids, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    context_layer = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    attn_output, _ = self.dense(context_layer, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return attn_output
+
+
+def vllm__module_executor__models__chatglm__GLMBlock__init__(
+    self,
+    config: ChatGLMConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(GLMBlock, self).__init__()
+    self.apply_residual_connection_post_layernorm = (
+        config.apply_residual_connection_post_layernorm)
+
+    self.fp32_residual_connection = config.fp32_residual_connection
+
+    layer_norm_func = RMSNorm if config.rmsnorm else LayerNorm
+    # Layernorm on the input data.
+    self.input_layernorm = layer_norm_func(config.hidden_size,
+                                            eps=config.layernorm_epsilon)
+
+    # Self attention.
+    self.self_attention = GLMAttention(config,
+                                       cache_config,
+                                       quant_config,
+                                       prefix=f"{prefix}.self_attention")
+    self.hidden_dropout = config.hidden_dropout
+
+    # Layernorm on the attention output
+    self.post_attention_layernorm = layer_norm_func(
+        config.hidden_size, eps=config.layernorm_epsilon)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: 1) use FeedForward instead of MLP
+            2) prepare to perf per-tensor sq cases if suitable
+    '''
+    # MLP
+    self.mlp = FeedForward(
+        hidden_size=config.hidden_size,
+        intermediate_size=config.ffn_hidden_size,
+        hidden_act='silu',
+        up_proj_name='dense_h_to_4h',
+        is_gated=True,
+        down_proj_name='dense_4h_to_h',
+        bias=config.add_bias_linear,
+        quant_config=quant_config
+    )
+
+    self.is_per_tesnor_sq_perf_cases = (is_per_tensor_smoothquant(quant_config) and
+        not self.apply_residual_connection_post_layernorm)
+    self.is_per_token_sq_perf_cases = (is_per_token_smoothquant(quant_config) and
+        not self.apply_residual_connection_post_layernorm)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attention.query_key_value.quant_method.skip_quant_input = True
+        self.mlp.dense_h_to_4h.quant_method.skip_quant_input = True
+        self.use_rmsnorm = config.rmsnorm
+        self.quant_fusion_attn_layernorm = None
+        self.quant_fusion_mlp_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__chatglm__GLMBlock__forward(
+    self,
+    hidden_states: torch.Tensor,
+    position_ids: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    mlp_layernorm = self.post_attention_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            quant_fusion_func = (quant_fusion_with_rmsnorm if
+                                 self.use_rmsnorm else quant_fusion_with_layernorm)
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attention.query_key_value.smooth
+                mlp_quant_scale = self.mlp.dense_h_to_4h.smooth
+            else:
+                attn_quant_scale = self.self_attention.query_key_value.scale_to_int
+                mlp_quant_scale = self.mlp.dense_h_to_4h.scale_to_int
+
+            self.quant_fusion_attn_layernorm = quant_fusion_func(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+            self.quant_fusion_mlp_layernorm = quant_fusion_func(
+                self.post_attention_layernorm, mlp_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+    return decoder_layer_forward_base(
+        positions=position_ids,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.self_attention,
+        post_layernorm=mlp_layernorm,
+        mlp=self.mlp,
+        apply_residual_connection_post_layernorm=self.apply_residual_connection_post_layernorm,
+        position_name='position_ids',
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+        post_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(GLMAttention,
+                             GLMAttention.forward,
+                             vllm__module_executor__models__chatglm__GLMAttention__forward)
+MluHijackObject.apply_hijack(GLMBlock,
+                             GLMBlock.__init__,
+                             vllm__module_executor__models__chatglm__GLMBlock__init__)
+MluHijackObject.apply_hijack(GLMBlock,
+                             GLMBlock.forward,
+                             vllm__module_executor__models__chatglm__GLMBlock__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/clip.py b/vllm_mlu/vllm_mlu/model_executor/models/clip.py
new file mode 100644
index 000000000..4c0a2b955
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/clip.py
@@ -0,0 +1,231 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from transformers import CLIPVisionConfig
+from transformers.models.clip.modeling_clip import CLIPSdpaAttention
+
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.models.clip import (CLIPVisionModel,
+                                             CLIPVisionTransformer,
+                                             CLIPEncoderLayer,
+                                             CLIPAttention)
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__clip__CLIPAttention__forward(
+    self,
+    hidden_states: torch.Tensor,
+    residual: torch.Tensor
+):
+    """Input shape: Batch x Time x Channel"""
+    bsz, tgt_len, _ = hidden_states.size()
+
+    qkv_states, _ = self.qkv_proj(hidden_states)
+    query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
+
+    query_states = query_states.view(bsz, tgt_len,
+                                     self.num_heads_per_partition,
+                                     self.head_dim)
+    key_states = key_states.view(bsz, tgt_len,
+                                 self.num_heads_per_partition,
+                                 self.head_dim)
+    value_states = value_states.view(bsz, tgt_len,
+                                     self.num_heads_per_partition,
+                                     self.head_dim)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf attn using tmo flash attn
+    '''
+    if self.dropout is None or self.dropout == 0.0:
+        # Always true for inference
+        from vllm_mlu import _mlu_ops as mlu_ops
+
+        out = mlu_ops.flash_attention(
+            query_states,
+            key_states,
+            value_states,
+            out=None,
+            cu_seq_lens_q=None,
+            cu_seq_lens_kv=None,
+            alibi_slope=None,
+            attn_bias=None,
+            max_seq_len_q=tgt_len,
+            max_seq_len_kv=tgt_len,
+            softmax_scale=self.scale,
+            is_causal=False
+        )
+    else:
+        from xformers import ops as xops
+
+        out = xops.memory_efficient_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            p=self.dropout,
+            scale=self.scale
+        )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    out = out.view(bsz, tgt_len, -1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    ''' 
+    attn_output, _ = self.out_proj(out, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return attn_output, None
+
+
+def vllm__module_executor__models__clip__CLIPEncoderLayer____init__(
+    self,
+    config: CLIPVisionConfig,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(CLIPEncoderLayer, self).__init__()
+
+    num_heads = config.num_attention_heads
+    tp_size = get_tensor_model_parallel_world_size()
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf attn using tmo flash attn, do not check xformers
+    '''
+    if num_heads % tp_size == 0:
+        self.self_attn = CLIPAttention(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.use_parallel_attn = True
+    else:
+        logger.warning("Use CLIPSdpaAttention for clip model, this can be slow.")
+        self.self_attn = CLIPSdpaAttention(config)
+        self.use_parallel_attn = False
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.layer_norm1 = nn.LayerNorm(config.hidden_size,
+                                    eps=config.layer_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.mlp = FeedForward(hidden_size=config.hidden_size,
+                           intermediate_size=config.intermediate_size,
+                           is_gated=False,
+                           bias=True,
+                           hidden_act=config.hidden_act,
+                           quant_config=quant_config,
+                           up_proj_name='fc1',
+                           down_proj_name='fc2',
+                           prefix=f"{prefix}.mlp")
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.layer_norm2 = nn.LayerNorm(config.hidden_size,
+                                    eps=config.layer_norm_eps)
+
+
+def vllm__module_executor__models__clip__CLIPEncoderLayer__forward(
+    self,
+    hidden_states: torch.Tensor
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: apply residual fusion
+    '''
+    residual = hidden_states
+    if self.use_parallel_attn:
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, _ = self.self_attn(hidden_states, residual)
+    else:
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, _ = self.self_attn(hidden_states)
+        hidden_states = residual + hidden_states
+
+    residual = hidden_states
+    hidden_states = self.layer_norm2(hidden_states)
+    bsz, tgt_len, _ = hidden_states.size()
+    hidden_states = self.mlp(hidden_states, residual)
+    hidden_states = hidden_states.view(bsz, tgt_len, -1)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return hidden_states
+
+
+def vllm__module_executor__models__clip__CLIPVisionModel____init__(
+    self,
+    config: CLIPVisionConfig,
+    quant_config: Optional[QuantizationConfig] = None,
+    *,
+    num_hidden_layers_override: Optional[int] = None,
+    require_post_norm: Optional[bool] = None,
+    prefix: str = "",
+) -> None:
+    super(CLIPVisionModel, self).__init__()
+
+    tp_size = get_tensor_model_parallel_world_size()
+    num_heads = config.num_attention_heads
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf attn using tmo flash attn, do not check xformers
+    '''
+    self.shard_weight = num_heads % tp_size == 0
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.vision_model = CLIPVisionTransformer(
+        config=config,
+        quant_config=quant_config,
+        num_hidden_layers_override=num_hidden_layers_override,
+        require_post_norm=require_post_norm,
+        prefix=f"{prefix}.vision_model",
+    )
+
+
+MluHijackObject.apply_hijack(CLIPAttention,
+                             CLIPAttention.forward,
+                             vllm__module_executor__models__clip__CLIPAttention__forward)
+MluHijackObject.apply_hijack(CLIPEncoderLayer,
+                             CLIPEncoderLayer.__init__,
+                             vllm__module_executor__models__clip__CLIPEncoderLayer____init__)
+MluHijackObject.apply_hijack(CLIPEncoderLayer,
+                             CLIPEncoderLayer.forward,
+                             vllm__module_executor__models__clip__CLIPEncoderLayer__forward)
+MluHijackObject.apply_hijack(CLIPVisionModel,
+                             CLIPVisionModel.__init__,
+                             vllm__module_executor__models__clip__CLIPVisionModel____init__)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/deepseek_mtp.py b/vllm_mlu/vllm_mlu/model_executor/models/deepseek_mtp.py
new file mode 100644
index 000000000..12dec4ac6
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/deepseek_mtp.py
@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Iterable, Set, Tuple
+
+import torch
+
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name, DeepseekV2MLAAttention
+from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+def vllm__module_executor__models__deepseek_mtp__DeepSeekMTP__load_weights(
+        self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
+    stacked_params_mapping = [
+        ("gate_up_proj", "gate_proj", 0),
+        ("gate_up_proj", "up_proj", 1),
+    ]
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: delete expert_params_mapping for useless
+    '''
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    params_dict = dict(self.named_parameters())
+    loaded_params: Set[str] = set()
+    for name, loaded_weight in weights:
+        if "rotary_emb.inv_freq" in name:
+            continue
+        spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
+        if spec_layer is None:
+            continue
+        name = self._rewrite_spec_layer_name(spec_layer, name)
+        for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            # Skip non-stacked layers and experts (experts handled below).
+            if weight_name not in name:
+                continue
+            # We have mlp.experts[0].gate_proj in the checkpoint.
+            # Since we handle the experts below in expert_params_mapping,
+            # we need to skip here BEFORE we update the name, otherwise
+            # name will be updated to mlp.experts[0].gate_up_proj, which
+            # will then be updated below in expert_params_mapping
+            # for mlp.experts[0].gate_gate_up_proj, which breaks load.
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
+            '''
+            name = name.replace(weight_name, param_name)
+            if (("mlp.experts." in name or "mlp.shared_experts." in name or "mlp.shared_expert_gate." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            # Skip loading extra bias for GPTQ models.
+            if name.endswith(".bias") and name not in params_dict:
+                continue
+
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            weight_loader(param, loaded_weight, shard_id)
+            break
+        else:
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: delete expert_params_mapping weight load
+            '''
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            # Skip loading extra bias for GPTQ models.
+            if name.endswith(".bias") and name not in params_dict:
+                continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition
+            '''
+            if (("mlp.experts." in name or "mlp.shared_experts." in name or "mlp.shared_expert_gate." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            weight_loader(param, loaded_weight)
+
+        loaded_params.add(name)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack_params
+    '''
+    for name, m in self.model.named_modules():
+        if isinstance(m, SparseMoeMlp) or isinstance(m, DeepseekV2MLAAttention):
+            m.pack_params()
+        if (isinstance(m, DeepseekV2MLAAttention)
+                 and m.use_fused_mla_qkv
+                 and self.config.q_lora_rank is not None):
+            m.weight_c = m.w_kc.view(m.num_local_heads, m.kv_lora_rank, -1).contiguous()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    return loaded_params
+
+
+MluHijackObject.apply_hijack(DeepSeekMTP,
+                             DeepSeekMTP.load_weights,
+                             vllm__module_executor__models__deepseek_mtp__DeepSeekMTP__load_weights)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/deepseek_v2.py b/vllm_mlu/vllm_mlu/model_executor/models/deepseek_v2.py
new file mode 100644
index 000000000..570817aa8
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/deepseek_v2.py
@@ -0,0 +1,996 @@
+import types
+import torch
+from torch import nn
+from typing import Any, Dict, Iterable, Optional, Tuple, Set
+from transformers import PretrainedConfig
+
+from vllm.attention.backends.abstract import AttentionType
+from vllm.attention import Attention, AttentionMetadata
+from vllm.attention.backends.utils import get_num_prefill_decode_query_kv_tokens
+from vllm.config import CacheConfig, ModelConfig
+from vllm.forward_context import (ForwardContext,
+                                  get_forward_context)
+from vllm.distributed import (get_tensor_model_parallel_world_size,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
+                                               ReplicatedLinear,
+                                               RowParallelLinear)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader, maybe_remap_kv_scale_name)
+from vllm.model_executor.models.utils import is_pp_missing_parameter
+from vllm.model_executor.models.deepseek_v2 import (
+    DeepseekV2Attention, DeepseekV2MLAAttention, DeepseekV2DecoderLayer,
+    DeepseekV2ForCausalLM, yarn_get_mscale, get_spec_layer_idx_from_weight_name)
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+from vllm.model_executor.layers.quantization.utils.fp8_utils import scaled_dequantize
+
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+from vllm_mlu.model_executor.models.layer_utils import quant_fusion_with_rmsnorm, is_per_token_smoothquant
+from vllm_mlu import _mlu_ops as mlu_ops
+
+
+class DeepseekV2MoE(SparseMoeMlp):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__(num_experts=config.n_routed_experts,
+                         top_k=config.num_experts_per_tok,
+                         hidden_size=config.hidden_size,
+                         intermediate_size=config.moe_intermediate_size,
+                         up_proj_name="gate_up_proj",
+                         is_gated=True,
+                         down_proj_name="down_proj",
+                         has_bias=False,
+                         skip_bias_add=False,
+                         renormalize=config.norm_topk_prob,
+                         hidden_act=config.hidden_act,
+                         params_dtype=None,
+                         quant_config=quant_config,
+                         is_use_fused_moe=True,
+                         expert_group=config.n_group,
+                         topk_group=config.topk_group,
+                         scoring_func=config.scoring_func,
+                         topk_method=config.topk_method,
+                         routed_scaling_factor=config.routed_scaling_factor)
+        self.config = config
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.n_shared_experts = config.n_shared_experts
+        self.routed_scaling_factor = config.routed_scaling_factor
+        if self.tp_size > config.n_routed_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than "
+                f"the number of experts {config.n_routed_experts}.")
+
+        if config.hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {config.hidden_act}. "
+                             "Only silu is supported for now.")
+
+        self.gate = ReplicatedLinear(config.hidden_size,
+                                     config.n_routed_experts,
+                                     bias=False,
+                                     quant_config=None,
+                                     prefix=f"{prefix}.gate")
+        if config.topk_method == "noaux_tc":
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.empty(config.n_routed_experts))
+        else:
+            self.gate.e_score_correction_bias = None
+        if config.n_shared_experts is not None:
+            intermediate_size = (config.moe_intermediate_size *
+                                 config.n_shared_experts)
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: replace MLP with FeedForward.
+            '''
+            self.shared_experts = FeedForward(hidden_size=config.hidden_size,
+                                             intermediate_size=intermediate_size,
+                                             hidden_act=config.hidden_act,
+                                             up_proj_name='gate_up_proj',
+                                             is_gated=True,
+                                             down_proj_name='down_proj',
+                                             bias=False,
+                                             quant_config=quant_config,
+                                             reduce_results=False)
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+        if self.is_fp8_block_wise:
+            self.experts.e_score_correction_bias = self.gate.e_score_correction_bias
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        if self.n_shared_experts is not None:
+            shared_output = self.shared_experts(hidden_states)
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: replace experts() with forward_experts, which defined by SparseMoeMlp.
+        '''
+        final_hidden_states = self.forward_experts(
+            hidden_states, router_logits, shared_output = shared_output)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                final_hidden_states)
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+def _compute_prefill_context(
+    self,
+    q: torch.Tensor,
+    kv_c_and_k_pe_cache: torch.Tensor,
+    attn_metadata: AttentionMetadata,
+):
+    prefill_metadata = attn_metadata.prefill_metadata
+    assert prefill_metadata is not None
+    assert prefill_metadata.context_chunk_seq_tot is not None
+    assert prefill_metadata.context_chunk_cu_seq_lens is not None
+    assert prefill_metadata.context_chunk_starts is not None
+    assert prefill_metadata.context_chunk_max_seq_lens is not None
+    assert prefill_metadata.context_lens_tensor is not None
+
+    output = None
+    iters = len(prefill_metadata.context_chunk_seq_tot)
+
+    # Fetch from attn_metadata directly, since it late bound by
+    # MLAAttentionState, grabbing it directly `attn_metadata` can avoid
+    # any weirdness around prefill_metadata caching
+    assert attn_metadata.context_chunk_workspace is not None
+    workspace = attn_metadata.context_chunk_workspace
+
+    for i in range(iters):
+        toks = prefill_metadata.context_chunk_seq_tot[i]
+
+        mlu_ops.gather_cache(
+            kv_cache=kv_c_and_k_pe_cache,
+            dst=workspace,
+            block_table=prefill_metadata.block_tables,
+            cu_seq_lens=prefill_metadata.context_chunk_cu_seq_lens[i],
+            batch_size=prefill_metadata.num_prefills,
+            seq_starts=prefill_metadata.context_chunk_starts[i],
+            kv_cache_dtype=self.attn.kv_cache_dtype
+        )
+
+        kv_c_normed = workspace[:toks][..., :self.kv_lora_rank]
+        k_pe = workspace[:toks][..., self.kv_lora_rank:].unsqueeze(1)
+
+        kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
+            -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim],
+                                  dim=-1)
+
+        k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
+
+        q = q.reshape(-1, self.num_local_heads * self.qk_head_dim)
+        k = k.reshape(-1, self.num_local_heads * self.qk_head_dim)
+        v = v.reshape(-1, self.num_local_heads * self.v_head_dim)
+
+        prefill_kwargs = {
+            "only_prefill": True,
+            "prefill_causal": False,
+            "cu_seq_lens_q": prefill_metadata.query_start_loc,
+            "cu_seq_lens_kv": prefill_metadata.context_chunk_cu_seq_lens[i],
+            "max_seq_len_q": prefill_metadata.max_query_len,
+            "max_seq_len_kv": prefill_metadata.context_chunk_max_seq_lens[i],
+            "return_lse": True,
+        }
+        attn_output, attn_softmax_lse = self.attn(q, k, v, kwargs=prefill_kwargs)
+        attn_output = attn_output.reshape(-1, self.num_local_heads, self.v_head_dim)
+        attn_softmax_lse = mlu_ops.extract_uncausal_lse(
+            attn_softmax_lse, prefill_metadata.query_start_loc)
+        if output is None:
+            output = attn_output
+            output_lse = attn_softmax_lse
+        else:
+            output_tmp = torch.empty_like(output)
+            output_lse_tmp = torch.empty_like(output_lse)
+            mlu_ops.merge_attn_states(
+                output=output_tmp,
+                output_lse=output_lse_tmp,
+                prefix_output=output,
+                prefix_lse=output_lse,
+                suffix_output=attn_output,
+                suffix_lse=attn_softmax_lse,
+            )
+            output_lse = output_lse_tmp
+
+    return output, output_lse
+
+
+def forward_prefill(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata: AttentionMetadata,
+) -> torch.Tensor:
+    if self.q_lora_rank is not None:
+        q = self.q_a_proj(hidden_states)[0]
+        q_scale = None
+        if hasattr(self.q_b_proj.quant_method,"quant_config") and \
+                is_per_token_smoothquant(self.q_b_proj.quant_method.quant_config):
+            self.q_b_proj.quant_method.skip_quant_input = True
+            quant_scale = self.q_b_proj.smooth
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                        self.q_a_layernorm, quant_scale, dynamic_quant=True)
+            q, q_scale = self.quant_fusion_attn_layernorm(q)
+            q = self.q_b_proj(q, q_scale)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+        else:
+            q = self.q_a_layernorm(q)
+            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads,
+                                         self.qk_head_dim)
+    else:
+        q = self.q_proj(hidden_states)[0].view(-1, self.num_local_heads,
+                                               self.qk_head_dim)
+    q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim],
+                           dim=-1)
+    latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+    kv_a, _ = latent_cache.split(
+        [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+    latent_cache = latent_cache.unsqueeze(1)
+    kv_a = self.kv_a_layernorm(kv_a)
+    kv = self.kv_b_proj(kv_a)[0]
+    kv = kv.view(-1, self.num_local_heads,
+                 self.qk_nope_head_dim + self.v_head_dim)
+    k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+    k_pe = latent_cache[:, :, self.kv_lora_rank:]
+    q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe, only_prefill=True)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: MLA save cache before flashattn
+    '''
+    if len(kv_cache) != 0  and kv_cache[0].numel() > 0:
+        key_cache = kv_cache[0][0]
+        key_value = torch.concat((kv_a.unsqueeze(1), k_pe), dim=-1)
+        updated_slot_mapping = attn_metadata.slot_mapping[:key_value.size(0)]
+        if self.attn.kv_cache_dtype == 'int8':
+            key_cache_scale = kv_cache[1][0]
+            mlu_ops.quant_to_paged_cache(key_value,
+                                            None,
+                                            key_cache,
+                                            None,
+                                            key_cache_scale,
+                                            None,
+                                            updated_slot_mapping.flatten())
+        else:
+            mlu_ops.reshape_paged_cache(key_value,
+                                        None,
+                                        key_cache,
+                                        None,
+                                        updated_slot_mapping.flatten())
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    k = torch.empty_like(q)
+    k[..., :self.qk_nope_head_dim] = k_nope
+    k[..., self.qk_nope_head_dim:] = k_pe
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: mlu attention not pad but qk_head_dim 192 v_head_dim 128.
+    '''
+    q = q.reshape(-1, self.num_local_heads * self.qk_head_dim)
+    k = k.reshape(-1, self.num_local_heads * self.qk_head_dim)
+    v = v.reshape(-1, self.num_local_heads * self.v_head_dim)
+
+    cu_seq_lens_q = attn_metadata.prefill_metadata.query_start_loc
+    max_seq_len_q = attn_metadata.prefill_metadata.max_query_len
+    prefill_kwargs = {
+        "only_prefill": True,
+        "prefill_causal": True,
+        "cu_seq_lens_q": cu_seq_lens_q,
+        "cu_seq_lens_kv": cu_seq_lens_q,
+        "max_seq_len_q": max_seq_len_q,
+        "max_seq_len_kv": max_seq_len_q,
+        "return_lse": attn_metadata.prefill_metadata.context_chunk_seq_tot is not None,
+    }
+    attn_output = self.attn(q, k, v, kwargs=prefill_kwargs)
+    # chunked prefill
+    if attn_metadata.prefill_metadata.context_chunk_seq_tot is not None:
+        suffix_output, suffix_lse = attn_output
+        suffix_output = suffix_output.reshape(-1, self.num_local_heads, self.v_head_dim)
+        suffix_lse = mlu_ops.extract_uncausal_lse(suffix_lse, cu_seq_lens_q)
+        context_output, context_lse = self._compute_prefill_context(q, kv_cache, attn_metadata)
+        attn_output = torch.empty_like(suffix_output)
+        mlu_ops.merge_attn_states(
+            output=attn_output,
+            prefix_output=context_output,
+            prefix_lse=context_lse,
+            suffix_output=suffix_output,
+            suffix_lse=suffix_lse,
+        )
+    attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    output, _ = self.o_proj(attn_output)
+    return output
+
+
+def forward_decoder(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata: AttentionMetadata,
+) -> torch.Tensor:
+    q_len = hidden_states.shape[0]
+    q_input = hidden_states.new_empty(
+        q_len, self.num_local_heads, self.kv_lora_rank + self.qk_rope_head_dim
+    )
+    if self.q_lora_rank is not None:
+        q = self.q_a_proj(hidden_states)[0]
+        q_scale = None
+        if hasattr(self.q_b_proj.quant_method,"quant_config") and \
+                is_per_token_smoothquant(self.q_b_proj.quant_method.quant_config):
+            self.q_b_proj.quant_method.skip_quant_input = True
+            quant_scale = self.q_b_proj.smooth
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                        self.q_a_layernorm, quant_scale, dynamic_quant=True)
+            q, q_scale = self.quant_fusion_attn_layernorm(q)
+            q = self.q_b_proj(q, q_scale)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+        else:
+            q = self.q_a_layernorm(q)
+            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+    else:
+        q = self.q_proj(hidden_states)[0].view(
+            -1, self.num_local_heads, self.qk_head_dim
+        )
+    q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+    torch.bmm(q_nope.transpose(0, 1), self.w_kc, out=q_input[..., : self.kv_lora_rank].transpose(0, 1))
+
+    latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+    v_input = latent_cache[..., : self.kv_lora_rank]
+    k_input = latent_cache
+    self.kv_a_layernorm(v_input, out=k_input[..., : self.kv_lora_rank])
+
+    k_input = k_input.unsqueeze(1)
+    k_pe = k_input[..., self.kv_lora_rank :]
+    q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe, only_decode=True)
+    q_input[..., self.kv_lora_rank :] = q_pe
+
+    q_input = q_input.reshape(q_input.shape[0], -1)
+    k_input = k_input.reshape(k_input.shape[0], -1)
+    v_input = v_input.reshape(v_input.shape[0], -1)
+    decode_kwargs = {"only_decode":True}
+    attn_output = self.attn_decoder(q_input, k_input, v_input, kwargs=decode_kwargs)
+    attn_output = attn_output.reshape(-1, self.num_local_heads,
+                                      self.kv_lora_rank)
+    attn_bmm_output = torch.empty(
+        q_len, self.num_local_heads, self.v_head_dim, device=attn_output.device, dtype=attn_output.dtype)
+    torch.bmm(attn_output.transpose(0, 1), self.w_vc.transpose(1, 2), out=attn_bmm_output.transpose(0, 1))
+    attn_output = attn_bmm_output.flatten(1, 2)
+    output, _ = self.o_proj(attn_output)
+    return output
+
+
+def forward_decoder_fused_mla_q(self, q, output, position):
+    if q.dim() == 2:
+        batch, input_size = q.shape
+        seq = 1
+    else:
+        batch, seq, input_size = q.shape
+    q = q.view(batch, seq, -1)
+
+    if output is not None:
+        output = output.view(batch, seq, output.shape[1], output.shape[2])
+
+    gamma = self.q_a_layernorm.weight.data
+    smooth_quant_scale = None
+    weight_b_scale = None
+    if hasattr(self.q_b_proj.quant_method,"quant_config") and \
+            is_per_token_smoothquant(self.q_b_proj.quant_method.quant_config):
+        smooth_quant_scale = self.q_b_proj.smooth
+        weight_b_scale = self.q_b_proj.per_channel_scale
+        weight_b = self.q_b_proj.qweight.data
+    else:
+        weight_b = self.q_b_proj.qweight.data
+
+    position_id, interleaved, _ = self.rotary_emb.get_param(position)
+    sin = self.rotary_emb.sin_
+    cos = self.rotary_emb.cos_
+    eps = self.q_a_layernorm.variance_epsilon
+
+    return mlu_ops.fused_mla_q(q, gamma, smooth_quant_scale, weight_b, weight_b_scale, self.weight_c, sin, cos,
+                               position_id, output, eps, interleaved)
+
+
+def forward_decoder_fused_mla_kv(self, kv_latent_cache, position, kv_cache, attn_metadata):
+    head_num = 1
+    if kv_latent_cache.dim() == 2:
+        batch, head_size = kv_latent_cache.shape
+        seq = 1
+    else:
+        batch, seq, head_size = kv_latent_cache.shape
+
+    kv = kv_latent_cache.view(batch, seq, head_num, head_size)
+
+    position_id, interleaved, _ = self.rotary_emb.get_param(position)
+    sin = self.rotary_emb.sin_
+    cos = self.rotary_emb.cos_
+
+    gamma = self.kv_a_layernorm.weight.data
+    key_cache = None
+    key_cache_scale = None
+    if kv_cache[0].numel() > 0:
+        kv_cache_, kv_cache_scale_ = kv_cache
+        key_cache = kv_cache_[0]
+        if kv_cache_scale_.numel() > 0:
+            key_cache_scale = kv_cache_scale_[0]
+
+    is_paged_cache = True
+    slot_mapping = None
+    cache_bs_id = None
+    cache_seq_offset = None
+    slot_mapping = attn_metadata.slot_mapping
+    slot_mapping = slot_mapping[attn_metadata.num_prefill_tokens:].view(batch, seq)
+    eps = self.kv_a_layernorm.variance_epsilon
+    return mlu_ops.fused_mla_kv(kv, sin, cos, position_id, gamma, key_cache, key_cache_scale, slot_mapping, cache_bs_id,
+                                cache_seq_offset, is_paged_cache, eps, interleaved)
+
+
+def forward_decoder_fused(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata: AttentionMetadata,
+) -> torch.Tensor:
+    q_len = hidden_states.shape[0]
+    q_input = hidden_states.new_empty(
+        q_len, self.num_local_heads, self.kv_lora_rank + self.qk_rope_head_dim
+    )
+    is_fused_mla_q = False
+    if self.q_lora_rank is not None:
+        if self.use_fused_qkv_a:
+            assert self.pack_params_done, "q_a_proj and kv_a_proj weights hasn't merged"
+            qkv_latent_cache = self.qkv_a_proj(hidden_states)[0]
+            q = qkv_latent_cache[..., : self.q_lora_rank]
+            latent_cache = qkv_latent_cache[..., self.q_lora_rank:]
+        else:
+            q = self.q_a_proj(hidden_states)[0]
+            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+
+        if hasattr(self.q_b_proj.quant_method,"quant_config") and \
+                is_per_token_smoothquant(self.q_b_proj.quant_method.quant_config):
+            forward_decoder_fused_mla_q(self, q, q_input, positions)
+            is_fused_mla_q = True
+        else:
+            q = self.q_a_layernorm(q)
+            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+    else:
+        q = self.q_proj(hidden_states)[0].view(
+            -1, self.num_local_heads, self.qk_head_dim
+        )
+        latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+
+    if is_fused_mla_q is False:
+        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+        torch.bmm(q_nope.transpose(0, 1), self.w_kc, out=q_input[..., :self.kv_lora_rank].transpose(0, 1))
+        q_pe, _ = self.rotary_emb(positions, q_pe, None, only_decode=True)
+        q_input[..., self.kv_lora_rank:] = q_pe
+
+    forward_decoder_fused_mla_kv(self, latent_cache, positions, kv_cache, attn_metadata)
+
+    k_input = latent_cache
+    k_input = k_input.unsqueeze(1)
+
+    v_input = latent_cache[..., : self.kv_lora_rank]
+
+    q_input = q_input.reshape(q_input.shape[0], -1)
+    k_input = k_input.reshape(k_input.shape[0], -1)
+    v_input = v_input.reshape(v_input.shape[0], -1)
+    decode_kwargs = {"only_decode":True}
+    attn_output = self.attn_decoder(q_input, k_input, v_input, kwargs=decode_kwargs)
+    attn_output = attn_output.reshape(-1, self.num_local_heads,
+                                      self.kv_lora_rank)
+    attn_bmm_output = torch.empty(
+        q_len, self.num_local_heads, self.v_head_dim, device=attn_output.device, dtype=attn_output.dtype)
+    torch.bmm(attn_output.transpose(0, 1), self.w_vc.transpose(1, 2), out=attn_bmm_output.transpose(0, 1))
+    attn_output = attn_bmm_output.flatten(1, 2)
+    output, _ = self.o_proj(attn_output)
+    return output
+
+
+def vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__pack_params(self):
+    if self.pack_params_done:
+        return
+    if self.use_fused_qkv_a:
+        self.q_a_proj.weight.data = self.qkv_a_proj.weight.data[:self.q_lora_rank, ...]
+        self.kv_a_proj_with_mqa.weight.data = self.qkv_a_proj.weight.data[self.q_lora_rank:, ...]
+    if self.is_fp8_block_wise:
+        kv_b_proj_dequant_weight = scaled_dequantize(self.kv_b_proj.weight, self.kv_b_proj.weight_scale_inv, self.kv_b_proj.quant_method.quant_config.weight_block_size, self.kv_b_proj.params_dtype)
+        w_kc, w_vc = kv_b_proj_dequant_weight.unflatten(0, (-1, self.qk_nope_head_dim + self.v_head_dim)).split([self.qk_nope_head_dim, self.v_head_dim], dim=1)
+        self.w_kc = w_kc
+        self.w_vc = w_vc
+    
+    self.pack_params_done = True
+
+
+def vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__init(
+    self,
+    config: PretrainedConfig,
+    hidden_size: int,
+    num_heads: int,
+    qk_nope_head_dim: int,
+    qk_rope_head_dim: int,
+    v_head_dim: int,
+    q_lora_rank: int,
+    kv_lora_rank: int,
+    rope_theta: float = 10000,
+    rope_scaling: Optional[Dict[str, Any]] = None,
+    max_position_embeddings: int = 8192,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(DeepseekV2MLAAttention, self).__init__()
+    self.hidden_size = hidden_size
+    self.qk_nope_head_dim = qk_nope_head_dim
+    self.qk_rope_head_dim = qk_rope_head_dim
+    self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+    self.v_head_dim = v_head_dim
+    self.q_lora_rank = q_lora_rank
+    self.kv_lora_rank = kv_lora_rank
+    self.num_heads = num_heads
+    tp_size = get_tensor_model_parallel_world_size()
+    assert num_heads % tp_size == 0
+    self.num_local_heads = num_heads // tp_size
+    self.scaling = self.qk_head_dim**-0.5
+    self.rope_theta = rope_theta
+    self.max_position_embeddings = max_position_embeddings
+    self.use_fused_mla_qkv = False
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: skip q_a_proj, kv_a_proj_with_mqa, kv_b_proj weight quant,
+            split kv_b_proj weight if not is_fp8_block_wise
+    '''
+    self.is_fp8_block_wise = isinstance(quant_config, Fp8Config) and quant_config.weight_block_size is not None
+
+    if self.q_lora_rank is not None:
+        self.q_a_proj = ReplicatedLinear(self.hidden_size,
+                                         self.q_lora_rank,
+                                         bias=False,
+                                         quant_config=quant_config if self.is_fp8_block_wise else None,
+                                         prefix=f"{prefix}.q_a_proj")
+        self.q_a_layernorm = RMSNorm(self.q_lora_rank,
+                                     eps=config.rms_norm_eps)
+        self.q_b_proj = ColumnParallelLinear(q_lora_rank,
+                                             self.num_heads *
+                                             self.qk_head_dim,
+                                             bias=False,
+                                             quant_config=quant_config,
+                                             prefix=f"{prefix}.q_b_proj")
+    else:
+        self.q_proj = ColumnParallelLinear(self.hidden_size,
+                                           self.num_heads *
+                                           self.qk_head_dim,
+                                           bias=False,
+                                           quant_config=quant_config,
+                                           prefix=f"{prefix}.q_proj")
+
+    self.kv_a_proj_with_mqa = ReplicatedLinear(
+        self.hidden_size,
+        self.kv_lora_rank + self.qk_rope_head_dim,
+        bias=False,
+        quant_config=quant_config if self.is_fp8_block_wise else None,
+        prefix=f"{prefix}.kv_a_proj_with_mqa")
+    self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,
+                                  eps=config.rms_norm_eps)
+    self.kv_b_proj = ColumnParallelLinear(
+        self.kv_lora_rank,
+        self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+        bias=False,
+        quant_config=quant_config if self.is_fp8_block_wise else None,
+        prefix=f"{prefix}.kv_b_proj")
+    kv_b_proj_weight = self.kv_b_proj.weight
+    w_kc, w_vc = kv_b_proj_weight.unflatten(
+        0, (-1, self.qk_nope_head_dim + self.v_head_dim)
+        ).split([self.qk_nope_head_dim, self.v_head_dim], dim=1)
+    self.w_kc = w_kc
+    self.w_vc = w_vc
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    # O projection.
+    self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
+                                    self.hidden_size,
+                                    bias=False,
+                                    quant_config=quant_config,
+                                    prefix=f"{prefix}.o_proj")
+    if rope_scaling:
+        rope_scaling["rope_type"] = 'deepseek_yarn'
+        self.use_normal_rope = False
+    else:
+        self.use_normal_rope = True
+    self.rotary_emb = get_rope(qk_rope_head_dim,
+                               rotary_dim=qk_rope_head_dim,
+                               max_position=max_position_embeddings,
+                               base=rope_theta,
+                               rope_scaling=rope_scaling,
+                               is_neox_style=False)
+
+    if rope_scaling:
+        mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
+        scaling_factor = rope_scaling["factor"]
+        mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
+        self.scaling = self.scaling * mscale * mscale
+
+    # self.attn = Attention(self.num_heads,
+    #                       self.qk_head_dim,
+    #                       self.scaling,
+    #                       num_kv_heads=self.num_heads)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add fused mla_q/kv and fused_qkv_a.
+    @brief: mlu attention support head_size 192.
+    '''
+    self.use_fused_mla_qkv = cache_config.cache_dtype != 'int8' and not self.is_fp8_block_wise
+    self.use_fused_qkv_a = False and self.use_fused_mla_qkv and self.q_lora_rank is not None and quant_config is not None
+
+    self.attn = Attention(self.num_local_heads,
+                          self.qk_nope_head_dim + self.qk_rope_head_dim,
+                          self.scaling,
+                          num_kv_heads=self.num_local_heads,
+                          cache_config=cache_config,
+                          quant_config=quant_config,
+                          prefix=f"{prefix}.attn",
+                          use_mla=True,
+                          v_head_dim=self.v_head_dim)
+    self.attn_decoder = Attention(self.num_local_heads,
+                          self.kv_lora_rank + self.qk_rope_head_dim,
+                          self.scaling,
+                          num_kv_heads=1,
+                          cache_config=cache_config,
+                          quant_config=quant_config,
+                          prefix=f"{prefix}.mla_attn",
+                          use_mla=True,
+                          v_head_dim=self.kv_lora_rank,
+                          use_fused_mla_qkv=self.use_fused_mla_qkv)
+    self.forward_prefill = types.MethodType(forward_prefill, self)
+    self.forward_decoder = types.MethodType(forward_decoder, self)
+    self._compute_prefill_context = types.MethodType(_compute_prefill_context, self)
+    if self.use_fused_mla_qkv:
+        self.forward_decoder_fused = types.MethodType(forward_decoder_fused, self)
+        self.weight_c = self.w_kc.transpose(-2, -1).contiguous()
+        if self.use_fused_qkv_a:
+            # self.qkv_a_proj is the fusion of self.q_a_proj and self.kv_a_proj_with_mqa
+            self.qkv_a_proj = ReplicatedLinear(self.hidden_size,
+                                               self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim,
+                                               bias=False,
+                                               quant_config=None,
+                                               prefix=f"{prefix}.qkv_a_proj")
+    self.pack_params_done = False
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use normal computation for prefill and use weight absorption for extend/decode.
+            pack_params() is called for dummy model
+    '''
+    forward_context: ForwardContext = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
+
+    if attn_metadata.is_profile_run and attn_metadata.chunked_prefill_enabled and \
+        attn_metadata.context_chunk_workspace is not None:
+        # During the profile run try to simulate to worse case output size
+        # for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context`
+        # since this can be large
+        _ = torch.empty(
+            (attn_metadata.context_chunk_workspace.shape[0],
+             self.num_heads, self.qk_nope_head_dim + self.v_head_dim),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+    num_prefill_tokens = attn_metadata.num_prefill_tokens
+    kv_cache = self.attn.kv_cache[forward_context.virtual_engine]
+    self.pack_params()
+    output = torch.empty_like(hidden_states)
+    if attn_metadata.prefill_metadata:
+        attn_metadata.prefill_metadata.compute_dtype = torch.float16
+        prefill_positions = positions[:num_prefill_tokens, ...]
+        prefill_hidden_states = hidden_states[:num_prefill_tokens, ...]
+        output[:num_prefill_tokens] = self.forward_prefill(prefill_positions, prefill_hidden_states, kv_cache,
+                                                           attn_metadata)
+    if attn_metadata.decode_metadata:
+        attn_metadata.decode_metadata.compute_dtype = torch.float16
+        forward_decoder_func = self.forward_decoder_fused if self.use_fused_mla_qkv else self.forward_decoder
+        decode_positions = positions[num_prefill_tokens:, ...]
+        decode_hidden_states = hidden_states[num_prefill_tokens:, ...]
+        output[num_prefill_tokens:] = forward_decoder_func(decode_positions, decode_hidden_states, kv_cache,
+                                                           attn_metadata)
+    return output
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__deepseek_v2__DeepseekV2DecoderLayer__init__(
+    self,
+    config: PretrainedConfig,
+    prefix: str,
+    model_config: ModelConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+) -> None:
+    super(DeepseekV2DecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    rope_theta = getattr(config, "rope_theta", 10000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+    max_position_embeddings = getattr(config, "max_position_embeddings",
+                                      8192)
+    # DecoderLayers are created with `make_layers` which passes the prefix
+    # with the layer's index.
+    layer_idx = int(prefix.split(sep='.')[-1])
+    if model_config.use_mla:
+        attn_cls = DeepseekV2MLAAttention
+    else:
+        attn_cls = DeepseekV2Attention
+    self.self_attn = attn_cls(
+        config=config,
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        qk_nope_head_dim=config.qk_nope_head_dim,
+        qk_rope_head_dim=config.qk_rope_head_dim,
+        v_head_dim=config.v_head_dim,
+        q_lora_rank=config.q_lora_rank
+        if hasattr(config, "q_lora_rank") else None,
+        kv_lora_rank=config.kv_lora_rank,
+        rope_theta=rope_theta,
+        rope_scaling=rope_scaling,
+        max_position_embeddings=max_position_embeddings,
+        cache_config=cache_config,
+        quant_config=quant_config,
+        prefix=f"{prefix}.self_attn",
+    )
+
+    if (config.n_routed_experts is not None
+            and layer_idx >= config.first_k_dense_replace
+            and layer_idx % config.moe_layer_freq == 0):
+        self.mlp = DeepseekV2MoE(
+            config=config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+    else:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: replace MLP with FeedForward.
+        '''
+        self.mlp = FeedForward(hidden_size=config.hidden_size,
+                               intermediate_size=config.intermediate_size,
+                               hidden_act=config.hidden_act,
+                               up_proj_name='gate_up_proj',
+                               is_gated=True,
+                               down_proj_name='down_proj',
+                               bias=False,
+                               quant_config=quant_config,
+                               prefix=f"{prefix}.mlp")
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                   eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    self.routed_scaling_factor = config.routed_scaling_factor
+
+
+def vllm__module_executor__models__deepseek_v2__DeepseekV2ForCausalLM__load_weights(
+    self,
+    weights: Iterable[Tuple[str, torch.Tensor]]
+) -> Set[str]:
+
+    stacked_params_mapping = [
+        # (param_name, shard_name, shard_id)
+        ("gate_up_proj", "gate_proj", 0),
+        ("gate_up_proj", "up_proj", 1),
+    ]
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: expert_params_mapping for fp8 block_wise
+    '''
+    is_fp8_block_wise = isinstance(self.quant_config, Fp8Config) and self.quant_config.weight_block_size is not None
+    expert_params_mapping = FusedMoE.make_expert_params_mapping(
+        ckpt_gate_proj_name="gate_proj",
+        ckpt_down_proj_name="down_proj",
+        ckpt_up_proj_name="up_proj",
+        num_experts=self.config.n_routed_experts) if is_fp8_block_wise else {}
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    params_dict = dict(self.named_parameters())
+    loaded_params: Set[str] = set()
+    for name, loaded_weight in weights:
+        import re
+        pattern = r'layers\.([0-9]*)\.'
+        match = re.search(pattern, name)
+        if match:
+            layer_id = int(match.group(1))
+            if layer_id >= self.config.num_hidden_layers:
+                continue
+        if "rotary_emb.inv_freq" in name:
+            continue
+
+        spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
+        if spec_layer is not None:
+            continue  # skip spec decode layers for main model
+
+        for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            # Skip non-stacked layers and experts (experts handled below).
+            if weight_name not in name:
+                continue
+            # We have mlp.experts[0].gate_proj in the checkpoint.
+            # Since we handle the experts below in expert_params_mapping,
+            # we need to skip here BEFORE we update the name, otherwise
+            # name will be updated to mlp.experts[0].gate_up_proj, which
+            # will then be updated below in expert_params_mapping
+            # for mlp.experts[0].gate_gate_up_proj, which breaks load.
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
+            '''
+            if is_fp8_block_wise and ("mlp.experts." in name) and name not in params_dict:
+                continue
+            name = name.replace(weight_name, param_name)
+            if (("mlp.experts." in name or "mlp.shared_experts." in name or "mlp.shared_expert_gate." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            # Skip loading extra bias for GPTQ models.
+            if name.endswith(".bias") and name not in params_dict:
+                continue
+
+            if is_pp_missing_parameter(name, self):
+                continue
+
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            weight_loader(param, loaded_weight, shard_id)
+            break
+        else:
+            for mapping in expert_params_mapping:
+                param_name, weight_name, expert_id, shard_id = mapping
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param,
+                              loaded_weight,
+                              name,
+                              shard_id=shard_id,
+                              expert_id=expert_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if (("mlp.experts." in name or "mlp.shared_experts." in name or "mlp.shared_expert_gate." in name)
+                        and name not in params_dict):
+                    continue
+
+                # Remapping the name of FP8 kv-scale.
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+        loaded_params.add(name)
+    
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack params
+    '''
+    for name, m in self.model.named_modules():
+        if isinstance(m, SparseMoeMlp) or isinstance(m, DeepseekV2MLAAttention):
+            m.pack_params()
+        if (isinstance(m, DeepseekV2MLAAttention)
+                 and m.use_fused_mla_qkv
+                 and self.config.q_lora_rank is not None):
+            m.weight_c = m.w_kc.transpose(-2, -1).contiguous()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    return loaded_params
+
+
+MluHijackObject.apply_hijack(DeepseekV2MLAAttention,
+                             DeepseekV2MLAAttention.forward,
+                             vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__forward)
+MluHijackObject.apply_hijack(DeepseekV2MLAAttention,
+                             DeepseekV2MLAAttention.__init__,
+                             vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__init)
+MluHijackObject.apply_hijack(DeepseekV2MLAAttention,
+                             "pack_params",
+                             vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__pack_params)
+MluHijackObject.apply_hijack(DeepseekV2DecoderLayer,
+                             DeepseekV2DecoderLayer.__init__,
+                             vllm__module_executor__models__deepseek_v2__DeepseekV2DecoderLayer__init__)
+MluHijackObject.apply_hijack(DeepseekV2ForCausalLM,
+                             DeepseekV2ForCausalLM.load_weights,
+                             vllm__module_executor__models__deepseek_v2__DeepseekV2ForCausalLM__load_weights)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/falcon.py b/vllm_mlu/vllm_mlu/model_executor/models/falcon.py
new file mode 100755
index 000000000..3b7cd5b6e
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/falcon.py
@@ -0,0 +1,253 @@
+import math
+import torch
+from torch.nn import LayerNorm
+
+from typing import Optional, Union
+from vllm.attention import Attention
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+from transformers import FalconConfig as HF_FalconConfig
+from vllm.transformers_utils.configs import RWConfig
+FalconConfig = Union[HF_FalconConfig, RWConfig]
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.linear import (QKVParallelLinear,
+                                               RowParallelLinear)
+from vllm.distributed import (get_tensor_model_parallel_rank,
+                              get_tensor_model_parallel_world_size)
+
+from vllm.model_executor.models.falcon import (FalconAttention,
+                                               FalconDecoderLayer,
+                                               _get_alibi_slopes)
+
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__falcon__FalconAttention____init__(
+    self,
+    config: FalconConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(FalconAttention, self).__init__()
+
+    self.hidden_size = config.hidden_size
+    tp_size = get_tensor_model_parallel_world_size()
+
+    self.total_num_heads = config.num_attention_heads
+    assert self.total_num_heads % tp_size == 0
+    self.num_heads = self.total_num_heads // tp_size
+    self.head_dim = self.hidden_size // self.total_num_heads
+    assert self.head_dim * self.total_num_heads == self.hidden_size
+
+    self.new_decoder_architecture = config.new_decoder_architecture
+    self.multi_query = config.multi_query
+
+    if self.new_decoder_architecture:
+        self.total_num_kv_heads = config.num_kv_heads
+    elif self.multi_query:
+        self.total_num_kv_heads = 1
+    else:
+        self.total_num_kv_heads = self.total_num_heads
+    if self.total_num_kv_heads >= tp_size:
+        # Number of KV heads is greater than TP size, so we partition
+        # the KV heads across multiple tensor parallel GPUs.
+        assert self.total_num_kv_heads % tp_size == 0
+    else:
+        # Number of KV heads is less than TP size, so we replicate
+        # the KV heads across multiple tensor parallel GPUs.
+        assert tp_size % self.total_num_kv_heads == 0
+    self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+    self.query_key_value = QKVParallelLinear(
+        self.hidden_size,
+        self.head_dim,
+        self.total_num_heads,
+        self.total_num_kv_heads,
+        bias=config.bias,
+        skip_bias_add=True,
+        quant_config=quant_config,
+    )
+    self.q_size = self.num_heads * self.head_dim
+    self.kv_size = self.num_kv_heads * self.head_dim
+
+    # Layer-wise attention scaling
+    self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
+    self.reduce_row_parallel_results = not (config.new_decoder_architecture
+                                            or config.parallel_attn)
+    self.dense = RowParallelLinear(
+        self.hidden_size,
+        self.hidden_size,
+        bias=config.bias,
+        skip_bias_add=True,
+        quant_config=quant_config,
+        reduce_results=self.reduce_row_parallel_results)
+
+    self.use_rotary = config.rotary
+    self.use_alibi = config.alibi
+    assert not (self.use_rotary and self.use_alibi), (
+        "Rotary and alibi are mutually exclusive.")
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: set cache_config for rotary & alibi
+    ''' 
+    if self.use_rotary:
+        rope_theta = getattr(config, "rope_theta", 10000)
+        max_position_embeddings = getattr(config,
+                                            "max_position_embeddings", 8192)
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=max_position_embeddings,
+            base=rope_theta,
+        )
+        self.attn = Attention(self.num_heads,
+                                self.head_dim,
+                                self.inv_norm_factor,
+                                num_kv_heads=self.num_kv_heads,
+                                cache_config=cache_config,
+                                quant_config=quant_config,
+                                prefix=f"{prefix}.attn")
+    elif self.use_alibi:
+        tp_rank = get_tensor_model_parallel_rank()
+        head_start = tp_rank * self.num_heads
+        head_end = (tp_rank + 1) * self.num_heads
+        alibi_slopes = (_get_alibi_slopes(self.total_num_heads) *
+                        self.inv_norm_factor)
+        alibi_slopes = alibi_slopes[head_start:head_end].tolist()
+        self.attn = Attention(self.num_heads,
+                                self.head_dim,
+                                self.inv_norm_factor,
+                                num_kv_heads=self.num_kv_heads,
+                                alibi_slopes=alibi_slopes,
+                                cache_config=cache_config,
+                                quant_config=quant_config,
+                                prefix=f"{prefix}.attn")
+    else:
+        self.attn = Attention(self.num_heads,
+                                self.head_dim,
+                                scale=self.inv_norm_factor,
+                                num_kv_heads=self.num_kv_heads,
+                                cache_config=cache_config,
+                                quant_config=quant_config,
+                                prefix=f"{prefix}.attn")
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__falcon__FalconAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    qkv, bias = self.query_key_value(hidden_states)
+    if bias is not None:
+        qkv += bias
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    ''' 
+    if self.use_rotary:
+        qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+        self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    attn_output, bias = self.dense(attn_output)
+    return attn_output, bias
+
+
+def vllm__module_executor__models__falcon__FalconDecoderLayer____init__(
+    self,
+    config: FalconConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(FalconDecoderLayer, self).__init__()
+    hidden_size = config.hidden_size
+    self.num_heads = config.num_attention_heads
+    self.self_attention = FalconAttention(
+        config,
+        cache_config,
+        quant_config,
+        prefix=f"{prefix}.self_attention")
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.reduce_row_parallel_results = not (config.new_decoder_architecture
+                                            or config.parallel_attn)
+    self.mlp = FeedForward(hidden_size=hidden_size,
+                            intermediate_size=hidden_size * 4,
+                            hidden_act='gelu',
+                            up_proj_name='dense_h_to_4h',
+                            is_gated=False,
+                            down_proj_name='dense_4h_to_h',
+                            bias=config.bias,
+                            quant_config=quant_config,
+                            skip_bias_add=True,
+                            reduce_results=self.reduce_row_parallel_results)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.config = config
+
+    if (not hasattr(config, "num_ln_in_parallel_attn")):
+        config.num_ln_in_parallel_attn = None
+
+    if (config.num_ln_in_parallel_attn is None
+            and config.new_decoder_architecture):
+        config.num_ln_in_parallel_attn = 2
+
+    if not config.parallel_attn:
+        self.post_attention_layernorm = LayerNorm(
+            hidden_size, eps=config.layer_norm_epsilon)
+        self.input_layernorm = LayerNorm(hidden_size,
+                                            eps=config.layer_norm_epsilon)
+    else:
+        if config.num_ln_in_parallel_attn == 2:
+            # The layer norm before self-attention
+            self.ln_attn = LayerNorm(hidden_size,
+                                        eps=config.layer_norm_epsilon)
+            # The layer norm before the MLP
+            self.ln_mlp = LayerNorm(hidden_size,
+                                    eps=config.layer_norm_epsilon)
+        else:
+            self.input_layernorm = LayerNorm(hidden_size,
+                                                eps=config.layer_norm_epsilon)
+
+    self.reduce_row_parallel_results = not (config.new_decoder_architecture
+                                            or config.parallel_attn)
+
+
+MluHijackObject.apply_hijack(FalconAttention,
+                             FalconAttention.__init__,
+                             vllm__module_executor__models__falcon__FalconAttention____init__)
+MluHijackObject.apply_hijack(FalconAttention,
+                             FalconAttention.forward,
+                             vllm__module_executor__models__falcon__FalconAttention__forward)
+MluHijackObject.apply_hijack(FalconDecoderLayer,
+                             FalconDecoderLayer.__init__,
+                             vllm__module_executor__models__falcon__FalconDecoderLayer____init__)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/gpt_neox.py b/vllm_mlu/vllm_mlu/model_executor/models/gpt_neox.py
new file mode 100644
index 000000000..bfd4b3a15
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/gpt_neox.py
@@ -0,0 +1,239 @@
+import torch
+from torch import nn
+from typing import Optional
+from transformers import GPTNeoXConfig
+
+from vllm.attention import Attention
+from vllm.config import CacheConfig
+from vllm.distributed import (get_tensor_model_parallel_world_size,
+                              get_tensor_model_parallel_rank,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.linear import (QKVParallelLinear,
+                                               RowParallelLinear)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.models.gpt_neox import GPTNeoXAttention, GPTNeoXLayer
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__gpt_neox__GPTNeoXAttention__init__(
+    self,
+    config: GPTNeoXConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(GPTNeoXAttention, self).__init__()
+    self.total_num_heads = config.num_attention_heads
+    self.hidden_size = config.hidden_size
+    self.head_size = self.hidden_size // self.total_num_heads
+    self.bias = getattr(config, "attention_bias", True)
+
+    tensor_model_parallel_world_size = (
+        get_tensor_model_parallel_world_size())
+    assert self.total_num_heads % tensor_model_parallel_world_size == 0
+    self.num_heads = (self.total_num_heads //
+                      tensor_model_parallel_world_size)
+
+    self.query_key_value = QKVParallelLinear(
+        config.hidden_size,
+        self.head_size,
+        self.total_num_heads,
+        bias=self.bias,
+        quant_config=quant_config,
+    )
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: do not do allreduce in linear and skip bias add when use_parallel_residual
+    '''
+    if config.use_parallel_residual:
+        reduce_results = False
+        skip_bias_add = True
+    else:
+        reduce_results = True
+        skip_bias_add = False
+
+    self.dense = RowParallelLinear(
+        config.hidden_size,
+        config.hidden_size,
+        bias=self.bias,
+        quant_config=quant_config,
+        reduce_results=reduce_results,
+        skip_bias_add=skip_bias_add,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    scaling = self.head_size**-0.5
+    rotary_dim = int(self.head_size * config.rotary_pct)
+    assert rotary_dim % 2 == 0
+    rope_theta = getattr(config, "rope_theta", 10000)
+    max_position_embeddings = getattr(config, "max_position_embeddings",
+                                      8192)
+    self.rotary_emb = get_rope(
+        self.head_size,
+        rotary_dim=rotary_dim,
+        max_position=max_position_embeddings,
+        base=rope_theta,
+    )
+    self.attn = Attention(self.num_heads,
+                          self.head_size,
+                          scaling,
+                          cache_config=cache_config,
+                          quant_config=quant_config,
+                          prefix=f"{prefix}.attn")
+
+
+def vllm__module_executor__models__gpt_neox__GPTNeoXAttention__forward(
+    self,
+    position_ids: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    qkv, _ = self.query_key_value(hidden_states)
+    q, k, v = qkv.chunk(chunks=3, dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    ''' 
+    qk, _ = qkv.split([self.num_heads * self.head_size * 2, self.num_heads * self.head_size], dim=-1)
+    self.rotary_emb(position_ids, qk.view(-1, self.num_heads + self.num_heads, self.head_size))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add bias for rank 0 when use_parallel_residual
+    ''' 
+    output, bias = self.dense(attn_output)
+    if self.dense.skip_bias_add and get_tensor_model_parallel_rank() == 0:
+        output += bias
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__gpt_neox__GPTNeoXLayer__init__(
+    self,
+    config: GPTNeoXConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(GPTNeoXLayer, self).__init__()
+    self.use_parallel_residual = config.use_parallel_residual
+    self.input_layernorm = nn.LayerNorm(config.hidden_size,
+                                        eps=config.layer_norm_eps)
+    self.post_attention_layernorm = nn.LayerNorm(config.hidden_size,
+                                                 eps=config.layer_norm_eps)
+    self.attention = GPTNeoXAttention(config,
+                                      cache_config,
+                                      quant_config,
+                                      prefix=f"{prefix}.attention")
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: 1) use FeedForward instead of MLP
+            2) do not do allreduce in row linear and skip bias add in it
+    '''
+    if self.use_parallel_residual:
+        reduce_results = False
+        skip_bias_add = True
+    else:
+        reduce_results = True
+        skip_bias_add = False
+
+    self.mlp = FeedForward(hidden_size=config.hidden_size,
+                           intermediate_size=config.intermediate_size,
+                           hidden_act='gelu',
+                           up_proj_name='dense_h_to_4h',
+                           is_gated=False,
+                           down_proj_name='dense_4h_to_h',
+                           bias=True,
+                           quant_config=quant_config,
+                           skip_bias_add=skip_bias_add,
+                           reduce_results=reduce_results)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__gpt_neox__GPTNeoXLayer__forward(
+    self,
+    position_ids: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    attn_input = self.input_layernorm(hidden_states)
+    attn_output = self.attention(
+        position_ids=position_ids,
+        hidden_states=attn_input,
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: only do one allreduce when use_parallel_residual
+    '''
+    if self.use_parallel_residual:
+        # pseudocode:
+        # x = x + attn(ln1(x)) + mlp(ln2(x))
+        mlp_input = self.post_attention_layernorm(hidden_states)
+        mlp_output, mlp_bias = self.mlp(mlp_input)
+        if get_tensor_model_parallel_rank() == 0:
+            mlp_output += mlp_bias
+            hidden_states = mlp_output + attn_output + hidden_states
+        else:
+            hidden_states = mlp_output + attn_output
+        hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+    else:
+        # pseudocode:
+        # x = x + attn(ln1(x))
+        # x = x + mlp(ln2(x))
+        attn_output = attn_output + hidden_states
+        mlp_input = self.post_attention_layernorm(attn_output)
+        mlp_output = self.mlp(mlp_input)
+        hidden_states = mlp_output + attn_output
+    return hidden_states
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(GPTNeoXAttention,
+                             GPTNeoXAttention.__init__,
+                             vllm__module_executor__models__gpt_neox__GPTNeoXAttention__init__)
+MluHijackObject.apply_hijack(GPTNeoXAttention,
+                             GPTNeoXAttention.forward,
+                             vllm__module_executor__models__gpt_neox__GPTNeoXAttention__forward)
+MluHijackObject.apply_hijack(GPTNeoXLayer,
+                             GPTNeoXLayer.__init__,
+                             vllm__module_executor__models__gpt_neox__GPTNeoXLayer__init__)
+MluHijackObject.apply_hijack(GPTNeoXLayer,
+                             GPTNeoXLayer.forward,
+                             vllm__module_executor__models__gpt_neox__GPTNeoXLayer__forward)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/hunyuan.py b/vllm_mlu/vllm_mlu/model_executor/models/hunyuan.py
new file mode 100755
index 000000000..8d00373ce
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/hunyuan.py
@@ -0,0 +1,508 @@
+import torch
+import re
+
+from typing import Optional, Tuple, Iterable, Union
+from transformers import PretrainedConfig
+
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.model_executor.models.utils import is_pp_missing_parameter
+from vllm.sequence import IntermediateTensors
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader, maybe_remap_kv_scale_name)
+
+from vllm_mlu.model_executor.models.mlu_models.hunyuan import (
+    HunYuanAttention, HunYuanDecoderLayer, HunYuanForCausalLM, HunYuanModel)
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.model_executor.models.layer_utils import (
+    hunyuan_decoder_layer_forward_base, hunyuan_decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant, quant_fusion_with_rmsnorm)
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+class HunYuanSparseMoeBlock(SparseMoeMlp):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+    ):
+        super().__init__(num_experts=config.num_experts,
+                         top_k=config.moe_topk,
+                         hidden_size=config.hidden_size,
+                         intermediate_size=config.intermediate_size,
+                         up_proj_name="gate_up_proj",
+                         is_gated=True,
+                         down_proj_name="down_proj",
+                         has_bias=False,
+                         skip_bias_add=False,
+                         renormalize=True if config.moe_topk>1 else False,
+                         hidden_act=config.hidden_act,
+                         params_dtype=None,
+                         quant_config=quant_config,
+                         is_use_fused_moe=True)
+        self.config = config
+        self.shared_mlp = None
+        if config.use_mixed_mlp_moe > 0:
+            self.shared_mlp = FeedForward(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size * config.num_shared_expert,
+                hidden_act=config.hidden_act,
+                up_proj_name='gate_up_proj',
+                is_gated=True,
+                down_proj_name='down_proj',
+                bias=False,
+                quant_config=quant_config,
+                reduce_results=False
+            )
+
+
+    def forward(self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_output = None
+        if self.shared_mlp is not None:
+            shared_output = self.shared_mlp(hidden_states)
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.forward_experts(hidden_states, router_logits, residual)
+
+        if shared_output is not None:
+            final_hidden_states = final_hidden_states + shared_output
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+def vllm__module_executor__models__hunyuan__HunYuanAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    kv_states: Optional[Tuple[torch.Tensor]] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if self.attention_type == "self":
+        qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: pack q & k to fit tmo.apply_rotary
+        '''
+        qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+        self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        ori_k = k
+        if self.use_qk_norm:
+            q = self.query_layernorm(
+                    q.reshape(-1, self.num_heads, self.head_dim).contiguous()
+                ).reshape(-1, self.num_heads*self.head_dim)
+            k = self.key_layernorm(
+                    k.reshape(-1, self.num_kv_heads, self.head_dim).contiguous()
+                ).reshape(-1, self.num_kv_heads*self.head_dim)
+    elif self.attention_type == "cross":
+        assert kv_states is not None
+        ori_k, v = kv_states # use last layer kv,
+        k = ori_k
+        q, _ = self.q_proj(hidden_states, smooth_quant_scale)
+        k_tmp = torch.empty_like(k) # Todo: reduant rotary embedding
+        qk_temp = torch.cat((q, k_tmp), dim=-1)
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: pack q & k to fit tmo.apply_rotary
+        '''
+        self.rotary_emb(positions, qk_temp.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        if self.use_qk_norm:
+            q = self.query_layernorm(
+                    q.view(-1, self.num_heads, self.head_dim).contiguous()
+                ).reshape(-1, self.num_heads*self.head_dim)
+            k = self.key_layernorm(
+                    k.view(-1, self.num_kv_heads, self.head_dim).contiguous()
+                ).reshape(-1, self.num_kv_heads*self.head_dim)
+    else:
+        raise RuntimeError("Not support attnention type")
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output, (ori_k, v)
+
+
+def vllm__module_executor__models__hunyuan__HunYuanDecoderLayer____init__(
+    self,
+    config: PretrainedConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+    layer_id: int = -1,
+) -> None:
+    super(HunYuanDecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    rope_theta = getattr(config, "rope_theta", 10000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+    if rope_scaling is not None and getattr(
+                config, "original_max_position_embeddings", None):
+            rope_scaling["original_max_position_embeddings"] = (
+                config.original_max_position_embeddings)
+    max_position_embeddings = getattr(config, "max_position_embeddings",
+                                        8192)
+    # Support abacusai/Smaug-72B-v0.1 with attention_bias
+    # Support internlm/internlm-7b with bias
+    attention_bias = getattr(config, "attention_bias", False) or getattr(
+        config, "bias", False)
+    cla_factor = getattr(config, "cla_share_factor", 1)
+    attention_type = "cross" \
+        if layer_id >= 0 and layer_id % cla_factor != 0 else "self"
+    self.self_attn = HunYuanAttention(
+        config=config,
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=getattr(config, "num_key_value_heads",
+                            config.num_attention_heads),
+        rope_theta=rope_theta,
+        rope_scaling=rope_scaling,
+        max_position_embeddings=max_position_embeddings,
+        quant_config=quant_config,
+        bias=attention_bias,
+        cache_config=cache_config,
+        prefix=f"{prefix}.self_attn",
+        attention_type=attention_type,
+    )
+
+    if getattr(config, "num_experts", None):
+        self.mlp = HunYuanSparseMoeBlock(config=config,
+                                        quant_config=quant_config)
+    else:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: use FeedForward instead of MLP
+        '''
+        self.mlp = FeedForward(hidden_size=config.hidden_size,
+                               intermediate_size=config.intermediate_size,
+                               hidden_act=config.hidden_act,
+                               up_proj_name='gate_up_proj',
+                               is_gated=True,
+                               down_proj_name='down_proj',
+                               bias=getattr(config, "mlp_bias", False),
+                               prefix=f"{prefix}.mlp",
+                               quant_config=quant_config)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                    eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf per-tensor sq cases if suitable. For moe
+        model, we only do quant fusion in attn block.
+    '''
+    self.is_per_tensor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tensor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.self_attn.attention_type == "self":
+            self.self_attn.qkv_proj.quant_method.skip_quant_input = True
+        if self.self_attn.attention_type == "cross":
+            self.self_attn.q_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__hunyuan__HunYuanForCausalLM__load_weights(
+    self,
+    weights: Iterable[Tuple[str, torch.Tensor]]
+):
+    cla_factor = getattr(self.config, "cla_share_factor", 1)
+    start_expert_id = 0
+    stacked_params_mapping = [
+        # (param_name, shard_name, shard_id)
+        (".qkv_proj", ".q_proj", "q"),
+        (".qkv_proj", ".k_proj", "k"),
+        (".qkv_proj", ".v_proj", "v"),
+        (".gate_up_proj", ".gate_proj", 0),
+        (".gate_up_proj", ".up_proj", 1),
+    ]
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: delete expert_params_mapping for useless
+    '''
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    params_dict = dict(self.named_parameters())
+    for name, loaded_weight in weights:
+        if "rotary_emb.inv_freq" in name:
+            continue
+        if ("rotary_emb.cos_cached" in name
+                or "rotary_emb.sin_cached" in name):
+            # Models trained using ColossalAI may include these tensors in
+            # the checkpoint. Skip them.
+            continue
+        # With tie_word_embeddings, we can skip lm_head.weight
+        # The weight might appear unnecessarily in the files if the model is
+        # processed with quantization, LoRA, fine-tuning, etc.
+        if self.config.tie_word_embeddings and "lm_head.weight" in name:
+            continue
+        if (self.quant_config is not None and
+                (scale_name := self.quant_config.get_cache_scale(name))):
+            # Loading kv cache scales for compressed-tensors quantization
+            param = params_dict[scale_name]
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            loaded_weight = loaded_weight[0]
+            weight_loader(param, loaded_weight)
+            continue
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: replace expert_id in weight to named_expert_id in params_dict
+        '''
+        if start_expert_id > 0 and "mlp.experts." in name:
+            expert_str = re.search(r'experts\.\d+', name).group(0)
+            expert_id=int(expert_str.split(".")[1])
+            named_expert_id = expert_id - start_expert_id
+            old_expert_name = f"experts.{expert_id}"
+            new_expert_name = f"experts.{named_expert_id}"
+            name = name.replace(old_expert_name, new_expert_name)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            if weight_name not in name:
+                continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: delete if "mlp.experts" in name: continue condition
+            '''
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            # cross layer only have q_proj, skip qkv pack
+            if weight_name == ".q_proj":
+                match = re.search(r'layers\.\d+', name)
+                if match:
+                    layer_id = int(match.group(0).split('.')[-1])
+                    if cla_factor > 1 and layer_id % cla_factor != 0:
+                        continue
+            name = name.replace(weight_name, param_name)
+            # Skip loading extra bias for GPTQ models.
+            if (name.endswith(".bias") and name not in params_dict):
+                continue
+
+            if is_pp_missing_parameter(name, self):
+                continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
+            '''
+            # Skip experts that are not assigned to this worker.
+            if (("mlp.experts." in name or "mlp.shared_mlp." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            weight_loader(param, loaded_weight, shard_id)
+            break
+        else:
+            if (name.endswith(".bias") and name not in params_dict):
+                continue
+            # Remapping the name of FP8 kv-scale.
+            name = maybe_remap_kv_scale_name(name, params_dict)
+            if name is None:
+                continue
+
+            if is_pp_missing_parameter(name, self):
+                continue
+
+            if "mlp.gate.wg." in name:
+                name = name.replace("wg.", "")
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition
+            '''
+            # Skip experts that are not assigned to this worker.
+            if (("mlp.experts." in name or "mlp.shared_mlp." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            weight_loader(param, loaded_weight)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack params
+    '''
+    for name, m in self.model.named_modules():
+        if isinstance(m, SparseMoeMlp):
+            m.pack_params()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__hunyuan__HunYuanDecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_states: Optional[Tuple[torch.Tensor]] = None,
+    residual: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    if self.is_per_tensor_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.self_attn.attention_type == "self":
+                self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                    self.input_layernorm, self.self_attn.qkv_proj.scale_to_int)
+            if self.self_attn.attention_type == "cross":
+                self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                    self.input_layernorm, self.self_attn.q_proj.scale_to_int)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+    elif self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.self_attn.attention_type == "self":
+                self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                    self.input_layernorm, self.self_attn.qkv_proj.smooth, dynamic_quant=True)
+            if self.self_attn.attention_type == "cross":
+                self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                    self.input_layernorm, self.self_attn.q_proj.smooth, dynamic_quant=True)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+    return hunyuan_decoder_layer_forward_base(
+        positions=positions,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.self_attn,
+        post_layernorm=self.post_attention_layernorm,
+        mlp=self.mlp,
+        kv_states=kv_states,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__hunyuan__HunYuanModel__forward(
+    self,
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors],
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return hunyuan_decoder_model_forward_base_pp(
+        config=self.config,
+        input_ids=input_ids,
+        positions=positions,
+        intermediate_tensors=intermediate_tensors,
+        layers=self.layers,
+        start_layer=self.start_layer,
+        end_layer=self.end_layer,
+        get_input_embeddings=self.get_input_embeddings,
+        norm=self.norm,
+        inputs_embeds=inputs_embeds
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(HunYuanAttention,
+                             HunYuanAttention.forward,
+                             vllm__module_executor__models__hunyuan__HunYuanAttention__forward)
+MluHijackObject.apply_hijack(HunYuanDecoderLayer,
+                             HunYuanDecoderLayer.__init__,
+                             vllm__module_executor__models__hunyuan__HunYuanDecoderLayer____init__)
+MluHijackObject.apply_hijack(HunYuanForCausalLM,
+                             HunYuanForCausalLM.load_weights,
+                             vllm__module_executor__models__hunyuan__HunYuanForCausalLM__load_weights)
+MluHijackObject.apply_hijack(HunYuanDecoderLayer,
+                             HunYuanDecoderLayer.forward,
+                             vllm__module_executor__models__hunyuan__HunYuanDecoderLayer__forward)
+MluHijackObject.apply_hijack(HunYuanModel,
+                             HunYuanModel.forward,
+                             vllm__module_executor__models__hunyuan__HunYuanModel__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/intern_vit.py b/vllm_mlu/vllm_mlu/model_executor/models/intern_vit.py
new file mode 100644
index 000000000..c28ecb720
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/intern_vit.py
@@ -0,0 +1,142 @@
+import torch
+import torch.nn as nn
+
+from typing import Optional
+from transformers import PretrainedConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.intern_vit import (InternVisionEncoderLayer,
+                                                   InternSdpaAttention, NORM2FN)
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.model_executor.models.layer_utils import intern_vision_encoder_layer_base
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__intern_vit__InternVisionEncoderLayer____init__(
+    self,
+    config: PretrainedConfig,
+    quant_config: Optional[QuantizationConfig] = None,
+    *,
+    num_dummy_heads: int = 0,
+    prefix: str = "",
+) -> None:
+    super(InternVisionEncoderLayer, self).__init__()
+    self.embed_dim = config.hidden_size
+    self.intermediate_size = config.intermediate_size
+    self.norm_type = config.norm_type
+
+    self.attn = self._init_attn(config,
+                                quant_config,
+                                num_dummy_heads=num_dummy_heads,
+                                prefix=f"{prefix}.attn")
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.mlp = FeedForward(hidden_size=config.hidden_size,
+                           intermediate_size=config.intermediate_size,
+                           hidden_act=config.hidden_act,
+                           up_proj_name='fc1',
+                           is_gated=False,
+                           down_proj_name='fc2',
+                           bias=True,
+                           quant_config=quant_config,
+                           prefix=f'{prefix}.mlp')
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.norm1 = NORM2FN[self.norm_type](self.embed_dim,
+                                         eps=config.layer_norm_eps)
+    self.norm2 = NORM2FN[self.norm_type](self.embed_dim,
+                                         eps=config.layer_norm_eps)
+
+    self.ls1 = nn.Parameter(config.initializer_factor *
+                            torch.ones(self.embed_dim))
+    self.ls2 = nn.Parameter(config.initializer_factor *
+                                torch.ones(self.embed_dim))
+
+
+def vllm__module_executor__models__intern_vit__InternVisionEncoderLayer__forward(
+    self,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+
+    return intern_vision_encoder_layer_base(
+        hidden_states=hidden_states,
+        self_attn=self.attn,
+        norm1=self.norm1,
+        norm2=self.norm2,
+        ls1=self.ls1,
+        ls2=self.ls2,
+        mlp=self.mlp
+    )
+
+
+def vllm__module_executor__models__intern_vit__InternSdpaAttention__forward(
+    self,
+    x: torch.Tensor
+) -> torch.Tensor:
+    B, N, C = x.shape
+    qkv = self.qkv(x)
+    q, k, v = qkv.chunk(3, dim=-1)
+
+    q = q.view(B, N, self.num_heads, self.head_dim)
+    k = k.view(B, N, self.num_heads, self.head_dim)
+    v = v.view(B, N, self.num_heads, self.head_dim)
+
+    if self.qk_normalization:
+        B_, N_, H_, D_ = q.shape
+        q = self.q_norm.forward_native(q.flatten(-2,
+                                                 -1)).view(B_, N_, H_, D_)
+        k = self.k_norm.forward_native(k.flatten(-2,
+                                                 -1)).view(B_, N_, H_, D_)
+    
+    batch, seq_len_q, q_head_num, head_size = q.shape
+    seq_len_k = k.shape[1]
+    softmax_scale = self.head_dim ** -0.5
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: replace SDPA with flash attn.
+    '''
+    x = mlu_ops.flash_attention(
+        q, k, v,
+        None, # out
+        None, # cu_seq_lens_q
+        None, # cu_seq_lens_kv
+        None, # alibi_slop
+        None, # attn_bias
+        seq_len_q, # max_seq_len_q
+        seq_len_k, # max_seq_len_kv
+        softmax_scale, # softmax_scale
+        False
+    )
+    x = x.reshape(x.shape[0],x.shape[1], -1)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    x = self.proj(x)
+    return x
+
+
+MluHijackObject.apply_hijack(InternVisionEncoderLayer,
+                             InternVisionEncoderLayer.__init__,
+                             vllm__module_executor__models__intern_vit__InternVisionEncoderLayer____init__)
+MluHijackObject.apply_hijack(InternVisionEncoderLayer,
+                             InternVisionEncoderLayer.forward,
+                             vllm__module_executor__models__intern_vit__InternVisionEncoderLayer__forward)
+MluHijackObject.apply_hijack(InternSdpaAttention,
+                             InternSdpaAttention.forward,
+                             vllm__module_executor__models__intern_vit__InternSdpaAttention__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/internlm2.py b/vllm_mlu/vllm_mlu/model_executor/models/internlm2.py
new file mode 100644
index 000000000..2149f6ada
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/internlm2.py
@@ -0,0 +1,299 @@
+import torch
+
+from typing import Optional, Tuple, Iterable, Union, Set
+from transformers import PretrainedConfig
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.models.internlm2 import (
+    InternLM2Attention, InternLMDecoderLayer, InternLM2ForCausalLM, InternLM2Model)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.utils import (is_pp_missing_parameter)
+from vllm.sequence import IntermediateTensors
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__internlm2__InternLM2Attention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.wqkv(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk, v = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.wo(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__internlm2__InternLMDecoderLayer____init__(
+    self,
+    config: PretrainedConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(InternLMDecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    rope_theta = getattr(config, "rope_theta", 10000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+    max_position_embeddings = getattr(config, "max_position_embeddings",
+                                        8192)
+    self.attention = InternLM2Attention(
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=config.num_key_value_heads,
+        rope_theta=rope_theta,
+        rope_scaling=rope_scaling,
+        max_position_embeddings=max_position_embeddings,
+        cache_config=cache_config,
+        quant_config=quant_config,
+        prefix=f"{prefix}.attention",
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.feed_forward = FeedForward(
+        hidden_size=self.hidden_size,
+        intermediate_size=config.intermediate_size,
+        hidden_act=config.hidden_act,
+        up_proj_name='gate_up_proj',
+        is_gated=True,
+        down_proj_name='w2',
+        bias=False,
+        quant_config=quant_config,
+        prefix=f"{prefix}.feed_forward",
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.attention_norm = RMSNorm(config.hidden_size,
+                                  eps=config.rms_norm_eps)
+    self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf per-tensor sq cases if suitable
+    '''
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.attention.wqkv.quant_method.skip_quant_input = True
+        self.feed_forward.gate_up_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+        self.quant_fusion_mlp_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__internlm2__InternLM2ForCausalLM__load_weights(
+    self,
+    weights: Iterable[Tuple[str, torch.Tensor]]
+) -> Set[str]:
+    stacked_params_mapping = [
+        # (param_name, shard_name, shard_id)
+        ("gate_up_proj", "w1", 0),
+        ("gate_up_proj", "w3", 1),
+    ]
+    params_dict = dict(self.named_parameters())
+    loaded_params: Set[str] = set()
+    for name, loaded_weight in weights:
+        if "rotary_emb.inv_freq" in name:
+            continue
+        for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            if weight_name not in name:
+                continue
+            name = name.replace(weight_name, param_name)
+            # Skip loading extra bias for GPTQ models.
+            if name.endswith(".bias") and name not in params_dict:
+                continue
+            if is_pp_missing_parameter(name, self):
+                continue
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            weight_loader(param, loaded_weight, shard_id)
+            break
+        else:
+            # Skip loading extra bias for GPTQ models.
+            if name.endswith(".bias") and name not in params_dict:
+                continue
+            if is_pp_missing_parameter(name, self):
+                continue
+            param = params_dict[name]
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: support load quant weights and params
+            '''
+            if "wqkv" in name and 'smooth' not in name and 'scale_to_int' not in name:
+                config = self.config
+                kv_groups = (config.num_attention_heads //
+                                config.num_key_value_heads)
+                head_dim = config.hidden_size // config.num_attention_heads
+                if 'weight' in name:
+                    loaded_weight = loaded_weight.view(-1, 2 + kv_groups,
+                                                        head_dim,
+                                                        loaded_weight.shape[-1])
+                    wq, wk, wv = torch.split(loaded_weight, [kv_groups, 1, 1],
+                                                dim=1)
+                    wq = wq.reshape(-1, wq.shape[-1])
+                    wk = wk.reshape(-1, wk.shape[-1])
+                    wv = wv.reshape(-1, wv.shape[-1])
+                elif 'scale' in name:
+                    loaded_weight = loaded_weight.view(-1, 2 + kv_groups, head_dim)
+                    wq, wk, wv = torch.split(loaded_weight, [kv_groups, 1, 1],
+                                                dim=1)
+                    wq = wq.reshape(-1)
+                    wk = wk.reshape(-1)
+                    wv = wv.reshape(-1)
+                else:
+                    logger.error(f"unsupport internlm2 quant param: {name}, shape: {loaded_weight.shape}")
+                weight_loader = param.weight_loader
+                weight_loader(param, wq, 'q')
+                weight_loader(param, wk, 'k')
+                weight_loader(param, wv, 'v')
+            else:
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+        loaded_params.add(name)
+    return loaded_params
+
+
+def vllm__module_executor__models__internlm2__InternLMDecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.attention_norm
+    mlp_layernorm = self.ffn_norm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.attention.wqkv.smooth
+                mlp_quant_scale = self.feed_forward.gate_up_proj.smooth
+            else:
+                attn_quant_scale = self.attention.wqkv.scale_to_int
+                mlp_quant_scale = self.feed_forward.gate_up_proj.scale_to_int
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.attention_norm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
+                self.ffn_norm, mlp_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+    return decoder_layer_forward_base(
+        positions, hidden_states,
+        attn_layernorm,
+        self.attention,
+        mlp_layernorm,
+        self.feed_forward,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+        post_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+def vllm__module_executor__models__internlm2__InternLM2Model__forward(
+    self,
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors],
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids, positions, intermediate_tensors,
+        self.layers, self.start_layer, self.end_layer,
+        self.tok_embeddings,
+        self.norm,
+        inputs_embeds
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(InternLM2Attention,
+                             InternLM2Attention.forward,
+                             vllm__module_executor__models__internlm2__InternLM2Attention__forward)
+MluHijackObject.apply_hijack(InternLMDecoderLayer,
+                             InternLMDecoderLayer.__init__,
+                             vllm__module_executor__models__internlm2__InternLMDecoderLayer____init__)
+MluHijackObject.apply_hijack(InternLM2ForCausalLM,
+                             InternLM2ForCausalLM.load_weights,
+                             vllm__module_executor__models__internlm2__InternLM2ForCausalLM__load_weights)
+MluHijackObject.apply_hijack(InternLMDecoderLayer,
+                             InternLMDecoderLayer.forward,
+                             vllm__module_executor__models__internlm2__InternLMDecoderLayer__forward)
+MluHijackObject.apply_hijack(InternLM2Model,
+                             InternLM2Model.forward,
+                             vllm__module_executor__models__internlm2__InternLM2Model__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/layer_utils.py b/vllm_mlu/vllm_mlu/model_executor/models/layer_utils.py
new file mode 100755
index 000000000..8ffe42758
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/layer_utils.py
@@ -0,0 +1,281 @@
+import torch
+from typing import Callable, Optional, List, Union, Tuple
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm.attention import AttentionMetadata
+from vllm.sequence import IntermediateTensors
+from vllm.distributed import get_pp_group
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from transformers import PretrainedConfig
+from vllm_mlu._mlu_utils import check_context_comm_cmpt_parallel
+
+
+def intern_vision_encoder_layer_base(
+    hidden_states: torch.Tensor,
+    self_attn: Callable,
+    norm1: Callable,
+    norm2: Callable,
+    ls1: Callable,
+    ls2: Callable,
+    mlp: Callable
+) -> torch.Tensor:
+    hidden_states = hidden_states + self_attn(
+                        norm1(hidden_states).reshape(hidden_states.shape)
+                    ) * ls1
+    hidden_states = hidden_states + mlp(
+                        norm2(hidden_states).reshape(hidden_states.shape)
+                    ).reshape(hidden_states.shape) * ls2
+    return hidden_states
+
+
+def hunyuan_decoder_layer_forward_base(
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    input_layernorm: Callable,
+    self_attn: Callable,
+    post_layernorm: Callable,
+    mlp: Callable,
+    kv_states: Optional[Tuple[torch.Tensor]] = None,
+    apply_residual_connection_post_layernorm: bool = False,
+    position_name: str = 'positions',
+    input_norm_fuse_en: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    smooth_quant_scale = None
+    if input_norm_fuse_en:
+        layernorm_output, smooth_quant_scale = input_layernorm(hidden_states)
+    else:
+        layernorm_output = input_layernorm(hidden_states)
+    if apply_residual_connection_post_layernorm:
+        residual = layernorm_output
+    else:
+        residual = hidden_states
+    
+    # Self Attention
+    attention_output, ori_kv_states = self_attn(
+        **{position_name: positions},
+        hidden_states=layernorm_output,
+        residual=residual,
+        kv_states=kv_states,
+        smooth_quant_scale=smooth_quant_scale,
+    )
+    
+    layernorm_output = post_layernorm(attention_output)
+    if apply_residual_connection_post_layernorm:
+        residual = layernorm_output
+    else:
+        residual = attention_output
+
+    # Fully Connected      
+    hidden_states = mlp(layernorm_output, residual)
+    return hidden_states, ori_kv_states
+
+
+def decoder_layer_forward_base(
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    input_layernorm: Callable,
+    self_attn: Callable,
+    post_layernorm: Callable,
+    mlp: Callable,
+    apply_residual_connection_post_layernorm: bool = False,
+    position_name: str = 'positions',
+    input_norm_fuse_en: bool = False,
+    post_norm_fuse_en: bool = False,
+) -> torch.Tensor:
+    smooth_quant_scale = None
+    if input_norm_fuse_en:
+        layernorm_output, smooth_quant_scale = input_layernorm(hidden_states)
+    else:
+        layernorm_output = input_layernorm(hidden_states)
+
+    if apply_residual_connection_post_layernorm:
+        residual = layernorm_output
+    else:
+        residual = hidden_states
+    
+    # Self Attention
+    attention_output = self_attn(
+        **{position_name: positions},
+        hidden_states=layernorm_output,
+        residual=residual,
+        smooth_quant_scale=smooth_quant_scale,
+    )
+
+    smooth_quant_scale = None
+    if post_norm_fuse_en:
+        layernorm_output, smooth_quant_scale = post_layernorm(attention_output)
+    else:
+        layernorm_output = post_layernorm(attention_output)
+
+    if apply_residual_connection_post_layernorm:
+        residual = layernorm_output
+    else:
+        residual = attention_output
+
+    # Fully Connected
+    kwargs = dict()
+    if post_norm_fuse_en:
+        kwargs['smooth_quant_scale'] = smooth_quant_scale
+    hidden_states = mlp(layernorm_output, residual, **kwargs)
+    return hidden_states
+
+
+def decoder_model_forward_base(
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    layers: torch.nn.ModuleList,
+    get_input_embeddings: Callable,
+    norm: Callable
+) -> torch.Tensor:
+    hidden_states = get_input_embeddings(input_ids)
+    for i in range(len(layers)):
+        layer = layers[i]
+        hidden_states = layer(
+            positions,
+            hidden_states,
+        )
+    hidden_states = norm(hidden_states)
+    return hidden_states
+
+
+def hunyuan_decoder_model_forward_base_pp(
+    config: PretrainedConfig,
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors],
+    layers: torch.nn.ModuleList,
+    start_layer: int,
+    end_layer: int,
+    get_input_embeddings: Callable,
+    norm: Callable,
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    if get_pp_group().is_first_rank:
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = get_input_embeddings(input_ids)
+    else:
+        assert intermediate_tensors is not None
+        hidden_states = intermediate_tensors["hidden_states"]
+
+    cla_factor = getattr(config, "cla_share_factor", 1)
+    prev_kv_states = None
+    for i in range(start_layer, end_layer):
+        layer = layers[i]
+        hidden_states, kv_states = layer(
+            positions,
+            hidden_states,
+            prev_kv_states,
+        )
+        if (i - start_layer) % cla_factor == 0:
+            prev_kv_states = kv_states
+        else:
+            prev_kv_states = None
+
+    if not get_pp_group().is_last_rank:
+        return IntermediateTensors({
+            "hidden_states": hidden_states,
+        })
+
+    hidden_states = norm(hidden_states)
+    return hidden_states
+
+
+def decoder_model_forward_base_pp(
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors],
+    layers: torch.nn.ModuleList,
+    start_layer: int,
+    end_layer: int,
+    get_input_embeddings: Callable,
+    norm: Callable,
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    if get_pp_group().is_first_rank:
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = get_input_embeddings(input_ids)
+    else:
+        assert intermediate_tensors is not None
+        hidden_states = intermediate_tensors["hidden_states"]
+
+    for i in range(start_layer, end_layer):
+        layer = layers[i]
+        hidden_states = layer(
+            positions,
+            hidden_states,
+        )
+
+    if not get_pp_group().is_last_rank:
+        return IntermediateTensors({
+            "hidden_states": hidden_states,
+        })
+
+    hidden_states = norm(hidden_states)
+    return hidden_states
+
+
+def is_smoothquant(quant_config: QuantizationConfig) -> bool:
+    return (quant_config is not None
+            and quant_config.get_name() == "SmoothQuant")
+
+
+def is_per_tensor_smoothquant(quant_config: QuantizationConfig) -> bool:
+    return (is_smoothquant(quant_config)
+            and quant_config.input_quant_method == "per_tensor")
+
+
+def is_per_token_smoothquant(quant_config: QuantizationConfig) -> bool:
+    if check_context_comm_cmpt_parallel():
+        return False
+    return (is_smoothquant(quant_config)
+            and quant_config.input_quant_method == "per_token")
+
+
+def quant_fusion_with_layernorm(
+    op: torch.nn.LayerNorm,
+    quant_scale: torch.Tensor,
+    dynamic_quant: bool = False,
+) -> Callable:
+    bias = None
+    if op.bias is not None:
+        bias = op.bias.data
+
+    def func(x: torch.Tensor) -> torch.Tensor:
+        return mlu_ops.fused_layer_norm(
+            x,
+            None,
+            op.weight.data,
+            bias,
+            None,
+            op.eps,
+            False,
+            quant_scale,
+            dynamic_quant)
+
+    return func
+
+
+def quant_fusion_with_rmsnorm(
+    op: RMSNorm,
+    quant_scale: torch.Tensor,
+    dynamic_quant: bool = False,
+) -> Callable:
+
+    def func(x: torch.Tensor) -> torch.Tensor:
+        return mlu_ops.fused_rms_norm(
+            x,
+            None,
+            op.weight.data,
+            None,
+            None,
+            op.variance_epsilon,
+            False,
+            quant_scale,
+            dynamic_quant)
+
+    return func
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/llama.py b/vllm_mlu/vllm_mlu/model_executor/models/llama.py
new file mode 100644
index 000000000..fb58701cc
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/llama.py
@@ -0,0 +1,281 @@
+import torch
+
+from typing import Dict, Optional, Union, Any
+from transformers import LlamaConfig
+
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.models.llama import LlamaAttention, LlamaDecoderLayer, LlamaModel
+from vllm.sequence import IntermediateTensors
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+vllm__module_executor__models__llama__LlamaAttention__init__org = LlamaAttention.__init__
+
+
+def vllm__module_executor__models__llama__LlamaAttention____init__(
+    self,
+    config: LlamaConfig,
+    hidden_size: int,
+    num_heads: int,
+    num_kv_heads: int,
+    rope_theta: float = 10000,
+    rope_scaling: Optional[Dict[str, Any]] = None,
+    max_position_embeddings: int = 8192,
+    quant_config: Optional[QuantizationConfig] = None,
+    bias: bool = False,
+    bias_o_proj: bool = False,
+    cache_config: Optional[CacheConfig] = None,
+    prefix: str = "",
+) -> None:
+    vllm__module_executor__models__llama__LlamaAttention__init__org(
+        self,
+        config,
+        hidden_size,
+        num_heads,
+        num_kv_heads,
+        rope_theta,
+        rope_scaling,
+        max_position_embeddings,
+        quant_config,
+        bias,
+        bias_o_proj,
+        cache_config,
+        prefix)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add rope_scaling params
+    ''' 
+    self.rope_scaling = rope_scaling
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__llama__LlamaAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    if self.rope_scaling is not None and self.rope_scaling["rope_type"] == "longrope":
+        q, k = self.rotary_emb(positions, q, k)
+    else:
+        qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+        self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__llama__LlamaDecoderLayer____init__(
+    self,
+    config: LlamaConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(LlamaDecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    rope_theta = getattr(config, "rope_theta", 10000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+    if rope_scaling is not None and getattr(
+            config, "original_max_position_embeddings", None):
+        rope_scaling["original_max_position_embeddings"] = (
+            config.original_max_position_embeddings)
+    max_position_embeddings = getattr(config, "max_position_embeddings",
+                                        8192)
+    # Support abacusai/Smaug-72B-v0.1 with attention_bias
+    # Support internlm/internlm-7b with bias
+    attention_bias = getattr(config, "attention_bias", False) or getattr(
+        config, "bias", False)
+    bias_o_proj = attention_bias
+    # support internlm/internlm3-8b with qkv_bias
+    if hasattr(config, 'qkv_bias'):
+        attention_bias = config.qkv_bias
+
+    self.self_attn = LlamaAttention(
+        config=config,
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=getattr(config, "num_key_value_heads",
+                                config.num_attention_heads),
+        rope_theta=rope_theta,
+        rope_scaling=rope_scaling,
+        max_position_embeddings=max_position_embeddings,
+        quant_config=quant_config,
+        bias=attention_bias,
+        bias_o_proj=bias_o_proj,
+        cache_config=cache_config,
+        prefix=f"{prefix}.self_attn",
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.mlp = FeedForward(hidden_size=config.hidden_size,
+                            intermediate_size=config.intermediate_size,
+                            hidden_act='silu',
+                            up_proj_name='gate_up_proj',
+                            is_gated=True,
+                            down_proj_name='down_proj',
+                            bias=getattr(config, "mlp_bias", False),
+                            quant_config=quant_config,
+                            prefix=f"{prefix}.mlp")
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                    eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf sq cases if suitable
+    '''
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attn.qkv_proj.quant_method.skip_quant_input = True
+        self.mlp.gate_up_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+        self.quant_fusion_mlp_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+def vllm__module_executor__models__llama__LlamaDecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    mlp_layernorm = self.post_attention_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attn.qkv_proj.smooth
+                mlp_quant_scale = self.mlp.gate_up_proj.smooth
+            else:
+                attn_quant_scale = self.self_attn.qkv_proj.scale_to_int
+                mlp_quant_scale = self.mlp.gate_up_proj.scale_to_int
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
+                self.post_attention_layernorm, mlp_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+    return decoder_layer_forward_base(
+        positions, hidden_states,
+        attn_layernorm,
+        self.self_attn,
+        mlp_layernorm,
+        self.mlp,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+        post_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+def vllm__module_executor__models__llama__LlamaModel__forward(
+    self,
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors],
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids, positions, intermediate_tensors,
+        self.layers, self.start_layer, self.end_layer,
+        self.get_input_embeddings,
+        self.norm,
+        inputs_embeds
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+MluHijackObject.apply_hijack(LlamaAttention,
+                             LlamaAttention.__init__,
+                             vllm__module_executor__models__llama__LlamaAttention____init__)
+MluHijackObject.apply_hijack(LlamaAttention,
+                             LlamaAttention.forward,
+                             vllm__module_executor__models__llama__LlamaAttention__forward)
+MluHijackObject.apply_hijack(LlamaDecoderLayer,
+                             LlamaDecoderLayer.__init__,
+                             vllm__module_executor__models__llama__LlamaDecoderLayer____init__)
+MluHijackObject.apply_hijack(LlamaDecoderLayer,
+                             LlamaDecoderLayer.forward,
+                             vllm__module_executor__models__llama__LlamaDecoderLayer__forward)
+MluHijackObject.apply_hijack(LlamaModel,
+                             LlamaModel.forward,
+                             vllm__module_executor__models__llama__LlamaModel__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/mixtral.py b/vllm_mlu/vllm_mlu/model_executor/models/mixtral.py
new file mode 100644
index 000000000..83665708a
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/mixtral.py
@@ -0,0 +1,366 @@
+import torch
+import re
+
+from typing import Optional, Tuple, Union, Iterable, Set
+from transformers import MixtralConfig
+
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.models.mixtral import (MixtralAttention, MixtralDecoderLayer,
+                                                MixtralForCausalLM, MixtralModel)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.model_loader.weight_utils import (default_weight_loader,
+                                                           maybe_remap_kv_scale_name)
+from vllm.model_executor.models.utils import is_pp_missing_parameter
+from vllm.sequence import IntermediateTensors
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.mlu_hijack_utils import MluHijackObject, set_is_gated
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+
+
+def vllm__module_executor__models__mixtral__MixtralAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__mixtral__MixtralDecoderLayer____init__(
+    self,
+    config: MixtralConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(MixtralDecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    # Requires transformers > 4.32.0
+    rope_theta = getattr(config, "rope_theta", 10000)
+    self.self_attn = MixtralAttention(
+        config=config,
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        max_position=config.max_position_embeddings,
+        num_kv_heads=config.num_key_value_heads,
+        rope_theta=rope_theta,
+        cache_config=cache_config,
+        quant_config=quant_config,
+        prefix=f"{prefix}.self_attn")
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: replace MixtralMoE to SparseMoeMlp
+    '''
+    self.block_sparse_moe = SparseMoeMlp(
+        num_experts=config.num_local_experts,
+        top_k=config.num_experts_per_tok,
+        hidden_size=config.hidden_size,
+        intermediate_size=config.intermediate_size,
+        up_proj_name="w13",
+        is_gated=True,
+        down_proj_name="w2",
+        has_bias=False,
+        skip_bias_add=False,
+        renormalize=True,
+        hidden_act=config.hidden_act,
+        params_dtype=None,
+        quant_config=quant_config,
+        is_use_fused_moe=True
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                   eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf per-tensor sq cases if suitable. MoE gate linear always runs
+        in half/full precision for now, so we only do quant fusion in attn block.
+    '''
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attn.qkv_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__mixtral__MixtralForCausalLM__load_weights(
+    self,
+    weights: Iterable[Tuple[str, torch.Tensor]]
+):
+    start_expert_id = 0
+    stacked_params_mapping = [
+        # (param_name, shard_name, shard_id)
+        ("qkv_proj", "q_proj", "q"),
+        ("qkv_proj", "k_proj", "k"),
+        ("qkv_proj", "v_proj", "v"),
+        ("w13", "w1", 0),
+        ("w13", "w3", 1),
+        ]
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: delete expert_params_mapping for useless
+    '''
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    params_dict = dict(self.named_parameters())
+    loaded_params: Set[str] = set()
+    for name, loaded_weight in weights:
+        if "rotary_emb.inv_freq" in name:
+            continue
+
+        if (self.quant_config is not None and
+            (scale_name := self.quant_config.get_cache_scale(name))):
+            # Loading kv cache quantization scales
+            param = params_dict[scale_name]
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
+                                loaded_weight[0])
+            weight_loader(param, loaded_weight)
+            loaded_params.add(scale_name)
+            continue
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: replace expert_id in weight to named_expert_id in params_dict
+        '''
+        if start_expert_id > 0 and "block_sparse_moe.experts." in name:
+            expert_str = re.search(r'experts\.\d+', name).group(0)
+            expert_id=int(expert_str.split(".")[1])
+            named_expert_id = expert_id - start_expert_id
+            old_expert_name = f"experts.{expert_id}"
+            new_expert_name = f"experts.{named_expert_id}"
+            name = name.replace(old_expert_name, new_expert_name)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            if weight_name not in name:
+                continue
+            name = name.replace(weight_name, param_name)
+            # Skip loading extra bias for GPTQ models.
+            if ((name.endswith(".bias") or name.endswith("_bias"))
+                        and name not in params_dict):
+                continue
+            if is_pp_missing_parameter(name, self):
+                continue
+            if name.endswith("scale"):
+                # Remapping the name of FP8 kv-scale.
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition
+            '''
+            # Skip experts that are not assigned to this worker.
+            if (("block_sparse_moe.experts." in name) and (name not in params_dict)):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            weight_loader(param, loaded_weight, shard_id)
+            break
+        else:
+            # Skip loading extra bias for GPTQ models.
+            if ((name.endswith(".bias") or name.endswith("_bias"))
+                            and name not in params_dict):
+                continue
+            if is_pp_missing_parameter(name, self):
+                continue
+            # Remapping the name of FP8 kv-scale.
+            name = maybe_remap_kv_scale_name(name, params_dict)
+            if name is None:
+                continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition
+            '''
+            # Skip experts that are not assigned to this worker.
+            if (("block_sparse_moe.experts." in name) and (name not in params_dict)):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            weight_loader(param, loaded_weight)
+        loaded_params.add(name)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack params
+    '''
+    for name, m in self.model.named_modules():
+        if isinstance(m, SparseMoeMlp):
+            m.pack_params()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return loaded_params
+
+
+def vllm__module_executor__models__mixtral__MixtralDecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attn.qkv_proj.smooth
+            else:
+                attn_quant_scale = self.self_attn.qkv_proj.scale_to_int
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+
+    return decoder_layer_forward_base(
+        positions, hidden_states,
+        attn_layernorm,
+        self.self_attn,
+        self.post_attention_layernorm,
+        self.block_sparse_moe,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__mixtral__MixtralModel__forward(
+    self,
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors],
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids, positions, intermediate_tensors,
+        self.layers, self.start_layer, self.end_layer,
+        self.embed_tokens,
+        self.norm,
+        inputs_embeds
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+set_is_gated(True)
+MluHijackObject.apply_hijack(MixtralAttention,
+                             MixtralAttention.forward,
+                             vllm__module_executor__models__mixtral__MixtralAttention__forward)
+MluHijackObject.apply_hijack(MixtralDecoderLayer,
+                             MixtralDecoderLayer.__init__,
+                             vllm__module_executor__models__mixtral__MixtralDecoderLayer____init__)
+MluHijackObject.apply_hijack(MixtralForCausalLM,
+                             MixtralForCausalLM.load_weights,
+                             vllm__module_executor__models__mixtral__MixtralForCausalLM__load_weights)
+MluHijackObject.apply_hijack(MixtralDecoderLayer,
+                             MixtralDecoderLayer.forward,
+                             vllm__module_executor__models__mixtral__MixtralDecoderLayer__forward)
+MluHijackObject.apply_hijack(MixtralModel,
+                             MixtralModel.forward,
+                             vllm__module_executor__models__mixtral__MixtralModel__forward)
+
+
+from vllm.model_executor.models.registry import ModelRegistry, _LazyRegisteredModel
+
+ModelRegistry.models['QuantMixtralForCausalLM'] = _LazyRegisteredModel(
+    module_name=f"vllm.model_executor.models.mixtral",
+    class_name="MixtralForCausalLM",
+)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/mllama.py b/vllm_mlu/vllm_mlu/model_executor/models/mllama.py
new file mode 100644
index 000000000..7cfeb8ecb
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/mllama.py
@@ -0,0 +1,218 @@
+# Copyright 2024 the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Mllama model."""
+from typing import (List, Optional, Tuple)
+
+import torch
+import torch.utils.checkpoint
+
+from vllm_mlu import _mlu_ops as mlu_ops
+from vllm.attention.selector import _Backend
+from vllm.attention import AttentionMetadata
+from vllm.attention.ops.paged_attn import PagedAttention
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.models.mllama import (MllamaTextCrossAttention,
+                                               MllamaTextModel,
+                                               MllamaVisionSdpaAttention)
+from vllm.logger import init_logger
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__model_executor__models__mllama__MllamaTextModel__forward(
+    self,
+    input_ids: torch.LongTensor,
+    positions: Optional[torch.LongTensor],
+    cross_attention_states: Optional[torch.LongTensor],
+    cross_attention_mask: Optional[torch.LongTensor],
+    kv_range_for_decode: Optional[List[Tuple[int, int]]],
+    full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor,
+                                                  torch.Tensor]],
+    skip_cross_attention: bool,
+) -> torch.Tensor:
+    inputs_embeds = self.embed_tokens(input_ids)
+    hidden_states = inputs_embeds
+
+    for idx, decoder_layer in enumerate(self.layers):
+        if idx in self.cross_attention_layers:
+            if not skip_cross_attention:
+                hidden_states = decoder_layer(
+                    hidden_states=hidden_states,
+                    cross_attention_states=cross_attention_states,
+                    cross_attention_mask=cross_attention_mask,
+                    kv_range_for_decode=kv_range_for_decode,
+                    full_text_row_masked_out_mask=
+                    full_text_row_masked_out_mask,
+                )
+        else:
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: fuse residual into decoder layer.
+            '''
+            hidden_states = decoder_layer(
+                positions=positions,
+                hidden_states=hidden_states,
+            )
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+    hidden_states = self.norm(hidden_states)
+    return hidden_states
+
+
+def vllm__model_executor__models__mllama__MllamaVisionSdpaAttention__forward(
+    self,
+    hidden_state: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.qkv_proj(hidden_state)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    q = q.view(q.shape[0], q.shape[1], self.num_local_heads,
+               self.head_dim)
+    k = k.view(k.shape[0], k.shape[1], self.num_local_heads,
+               self.head_dim)
+    v = v.view(v.shape[0], v.shape[1], self.num_local_heads,
+               self.head_dim)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: replace SDPA with flash attn.
+    '''
+    batch, seq_len_q, q_head_num, head_size = q.shape
+    seq_len_k = k.shape[1]
+    softmax_scale = head_size ** -0.5
+    attention_mask = attention_mask.repeat(1, q_head_num, 1, 1)
+    attn_output = mlu_ops.flash_attention(
+        q, k, v,
+        None, # out
+        None, # cu_seq_lens_q
+        None, # cu_seq_lens_kv
+        None, # alibi_slop
+        attention_mask, # attn_bias
+        seq_len_q, # max_seq_len_q
+        seq_len_k, # max_seq_len_kv
+        softmax_scale, # softmax_scale
+        False, # is_casual
+    )
+
+    attn_output = attn_output.reshape(attn_output.shape[0],
+                                      attn_output.shape[1], -1).contiguous()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    output, _ = self.o_proj(attn_output)
+    return output
+
+
+def vllm__model_executor__models__mllama__MllamaTextCrossAttention___attention_with_mask(
+    self,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    attention_mask: torch.Tensor,
+    kv_range_for_decode: List[Tuple[int, int]],
+) -> torch.Tensor:
+    kv_cache = self.attn.kv_cache[self.pipeline_parallel_rank]
+    attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
+    # Skip writing kv-cache for the initial profiling run.
+    # TODO (NickLucche) replace with custom attn bias and use standard attn
+    if len(kv_cache[0].shape) > 1:
+        i = torch.ones(1, dtype=torch.float32)
+        if self.attn.backend in (_Backend.FLASH_ATTN,
+                                    _Backend.FLASH_ATTN_VLLM_V1):
+            cached_k = torch.cat([k[s:e] for s, e in kv_range_for_decode])
+            cached_v = torch.cat([v[s:e] for s, e in kv_range_for_decode])
+            mlu_ops.reshape_paged_cache(
+                cached_k,
+                cached_v,
+                kv_cache[0][0],
+                kv_cache[0][1],
+                attn_metadata.cross_slot_mapping,
+            )
+        elif self.attn.backend in (_Backend.XFORMERS, _Backend.ROCM_FLASH,
+                                    _Backend.TORCH_SDPA):
+            key_cache, value_cache = PagedAttention.split_kv_cache(
+                kv_cache, self.num_local_key_value_heads, self.head_dim)
+            cached_k = torch.cat([k[s:e] for s, e in kv_range_for_decode])
+            cached_v = torch.cat([v[s:e] for s, e in kv_range_for_decode])
+            PagedAttention.write_to_paged_cache(
+                cached_k, cached_v, key_cache, value_cache,
+                attn_metadata.cross_slot_mapping, "auto", i, i)
+        else:
+            raise ValueError(
+                f"Unsupported Attention backend {self.attn.backend} "
+                "enum found. Expected the Attention backend to be "
+                "FLASH_ATTN, FLASH_ATTN_VLLM_V1, "
+                "XFORMERS or TORCH_SDPA.")
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: replace SDPA with flash attn.
+    '''
+    # We have to call torch.sdpa for prefill when using a
+    # custom cross-attention mask. Because the mask is not a
+    # standard causal mask, neither a block diagonal mask which
+    # can be optimized by xformers.BlockDiagonalMask.
+    # The mask is specially calculated for supporting multi
+    # images and interleaved images.
+    seq_len_q, q_head_num, head_size = q.shape
+    softmax_scale = head_size ** -0.5
+    cu_seq_lens_q = attn_metadata.seq_start_loc
+    cu_seq_lens_kv = attn_metadata.encoder_seq_start_loc
+
+    max_seq_len_q = attn_metadata.max_prefill_seq_len
+    max_seq_len_kv = attn_metadata.max_encoder_seq_len
+    attn_output = mlu_ops.flash_attention(
+        q, k, v,
+        None, # out
+        cu_seq_lens_q,
+        cu_seq_lens_kv,
+        None, # alibi_slope
+        None, # attn_bias
+        max_seq_len_q,
+        max_seq_len_kv,
+        softmax_scale,
+        False, # is_causal
+    )
+
+    output = attn_output.reshape(seq_len_q, self.num_local_heads * self.head_dim)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+MluHijackObject.apply_hijack(MllamaTextCrossAttention,
+                             MllamaTextCrossAttention._attention_with_mask,
+                             vllm__model_executor__models__mllama__MllamaTextCrossAttention___attention_with_mask)
+MluHijackObject.apply_hijack(MllamaTextModel,
+                             MllamaTextModel.forward,
+                             vllm__model_executor__models__mllama__MllamaTextModel__forward)
+MluHijackObject.apply_hijack(MllamaVisionSdpaAttention,
+                             MllamaVisionSdpaAttention.forward,
+                             vllm__model_executor__models__mllama__MllamaVisionSdpaAttention__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/mlu_models/__init__.py b/vllm_mlu/vllm_mlu/model_executor/models/mlu_models/__init__.py
new file mode 100644
index 000000000..e2723c1ac
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/mlu_models/__init__.py
@@ -0,0 +1,12 @@
+from vllm.model_executor.models.registry import ModelRegistry
+
+
+_MLU_VLLM_MODELS = {
+    "HunYuanForCausalLM": ("hunyuan", "HunYuanForCausalLM"),
+}
+
+for model_arch, (mod_relname, cls_name) in _MLU_VLLM_MODELS.items():
+    ModelRegistry.register_model(
+        model_arch=model_arch,
+        model_cls=f"vllm_mlu.model_executor.models.mlu_models.{mod_relname}:{cls_name}"
+    )
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/mlu_models/hunyuan.py b/vllm_mlu/vllm_mlu/model_executor/models/mlu_models/hunyuan.py
new file mode 100644
index 000000000..a8f4efe30
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/mlu_models/hunyuan.py
@@ -0,0 +1,670 @@
+# coding=utf-8
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved.
+#
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent/Tencent-Hunyuan-Large/blob/main/License.docx
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only HunYuan model compatible with HuggingFace weights."""
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+
+import re
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+
+from vllm.attention import Attention, AttentionMetadata
+from vllm.config import CacheConfig, LoRAConfig, VllmConfig
+from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
+                              get_tensor_model_parallel_world_size,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
+                                               QKVParallelLinear,
+                                               ReplicatedLinear,
+                                               ColumnParallelLinear,
+                                               RowParallelLinear)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader, maybe_remap_kv_scale_name)
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.sequence import IntermediateTensors
+
+from vllm.model_executor.models.interfaces import SupportsLoRA
+from vllm.model_executor.models.utils import (
+    PPMissingLayer, is_pp_missing_parameter, make_layers, maybe_prefix)
+
+
+class HunYuanMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        bias: bool = False,
+        prefix: str = "",
+        reduce_results: bool = True,
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj")
+        self.down_proj = RowParallelLinear(input_size=intermediate_size,
+                                           output_size=hidden_size,
+                                           bias=bias,
+                                           quant_config=quant_config,
+                                           prefix=f"{prefix}.down_proj",
+                                           reduce_results=reduce_results)
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. "
+                             "Only silu is supported for now.")
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class HunYuanSparseMoeBlock(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+        if self.tp_size > config.num_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than "
+                f"the number of experts {config.num_experts}.")
+
+        self.experts = FusedMoE(num_experts=config.num_experts,
+                                top_k=config.moe_topk,
+                                hidden_size=config.hidden_size,
+                                intermediate_size=config.intermediate_size,
+                                reduce_results=False,
+                                renormalize=True if config.moe_topk>1 else False,
+                                quant_config=quant_config)
+
+        self.gate = ReplicatedLinear(config.hidden_size,
+                                     config.num_experts,
+                                     bias=False,
+                                     quant_config=None)
+        if config.use_mixed_mlp_moe > 0:
+            self.shared_mlp = HunYuanMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size * config.num_shared_expert,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                reduce_results=False,
+            )
+        else:
+            self.shared_mlp = None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # NOTE: hidden_states can have either 1D or 2D shape.
+        orig_shape = hidden_states.shape
+        hidden_dim = hidden_states.shape[-1]
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_output = None
+        if self.shared_mlp is not None:
+            shared_output = self.shared_mlp(hidden_states)
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.experts(hidden_states=hidden_states,
+                                           router_logits=router_logits)
+        if shared_output is not None:
+            final_hidden_states = final_hidden_states + shared_output
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                final_hidden_states)
+
+        return final_hidden_states.view(orig_shape)
+
+
+class HunYuanAttention(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_theta: float = 10000,
+        rope_scaling: Optional[Dict[str, Any]] = None,
+        max_position_embeddings: int = 8192,
+        quant_config: Optional[QuantizationConfig] = None,
+        bias: bool = False,
+        cache_config: Optional[CacheConfig] = None,
+        prefix: str = "",
+        attention_type: str = "self",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
+        self.head_dim = getattr(config, "head_dim",
+                                self.hidden_size // self.total_num_heads)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.use_qk_norm = config.use_qk_norm
+        self.attention_type = attention_type
+
+        if attention_type == "self":
+            self.qkv_proj = QKVParallelLinear(
+                hidden_size=hidden_size,
+                head_size=self.head_dim,
+                total_num_heads=self.total_num_heads,
+                total_num_kv_heads=self.total_num_kv_heads,
+                bias=bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.qkv_proj",
+            )
+        elif attention_type == "cross":
+            self.q_proj = ColumnParallelLinear(
+                hidden_size,
+                hidden_size,
+                bias=bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_proj",
+            )
+        else:
+            raise RuntimeError("Not support attnention type")
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        is_neox_style = True
+        if quant_config is not None and quant_config.get_name() == "gguf":
+            is_neox_style = False
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=max_position_embeddings,
+            base=rope_theta,
+            rope_scaling=rope_scaling,
+            is_neox_style=is_neox_style,
+        )
+        self.attn = Attention(self.num_heads,
+                              self.head_dim,
+                              self.scaling,
+                              num_kv_heads=self.num_kv_heads,
+                              cache_config=cache_config,
+                              quant_config=quant_config,
+                              prefix=f"{prefix}.attn")
+
+        if self.use_qk_norm:
+            self.query_layernorm = RMSNorm(self.head_dim,
+                                           eps=config.rms_norm_eps)
+            self.key_layernorm = RMSNorm(self.head_dim,
+                                         eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_states: Optional[Tuple[torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        if self.attention_type == "self":
+            qkv, _ = self.qkv_proj(hidden_states)
+            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+            q, k = self.rotary_emb(positions, q, k)
+            ori_k = k
+            if self.use_qk_norm:
+                q = self.query_layernorm(q.view(-1, self.num_heads, self.head_dim).contiguous())
+                k = self.key_layernorm(k.view(-1, self.num_kv_heads, self.head_dim).contiguous())
+        elif self.attention_type == "cross":
+            assert kv_states is not None
+            ori_k, v = kv_states # use last layer kv,
+            k = ori_k
+            q, _ = self.q_proj(hidden_states)
+            k_tmp = torch.empty_like(k) # Todo: reduant rotary embedding
+            q, _ = self.rotary_emb(positions, q, k_tmp)
+            if self.use_qk_norm:
+                q = self.query_layernorm(q.view(-1, self.num_heads, self.head_dim).contiguous())
+                k = self.key_layernorm(k.view(-1, self.num_kv_heads, self.head_dim).contiguous())
+        else:
+            raise RuntimeError("Not support attnention type")
+
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output, (ori_k, v)
+
+
+class HunYuanDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+        layer_id: int = -1,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 10000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        if rope_scaling is not None and getattr(
+                config, "original_max_position_embeddings", None):
+            rope_scaling["original_max_position_embeddings"] = (
+                config.original_max_position_embeddings)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
+        # Support abacusai/Smaug-72B-v0.1 with attention_bias
+        # Support internlm/internlm-7b with bias
+        attention_bias = getattr(config, "attention_bias", False) or getattr(
+            config, "bias", False)
+        cla_factor = getattr(config, "cla_share_factor", 1)
+        attention_type = "cross" \
+            if layer_id >= 0 and layer_id % cla_factor != 0 else "self"
+        self.self_attn = HunYuanAttention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=getattr(config, "num_key_value_heads",
+                                 config.num_attention_heads),
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            quant_config=quant_config,
+            bias=attention_bias,
+            cache_config=cache_config,
+            prefix=f"{prefix}.self_attn",
+            attention_type=attention_type,
+        )
+        if getattr(config, "num_experts", None):
+            self.mlp = HunYuanSparseMoeBlock(config=config,
+                                             quant_config=quant_config)
+        else:
+            self.mlp = HunYuanMLP(
+                hidden_size=self.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                bias=getattr(config, "mlp_bias", False),
+                prefix=f"{prefix}.mlp",
+            )
+        self.input_layernorm = RMSNorm(config.hidden_size,
+                                       eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                                eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        kv_states: Optional[Tuple[torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(
+                hidden_states, residual)
+        hidden_states, ori_kv_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            kv_states=kv_states,
+        )
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual, ori_kv_states
+
+
+class HunYuanModel(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        lora_config: Optional[LoRAConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.padding_idx = config.pad_token_id
+        lora_vocab = (lora_config.lora_extra_vocab_size *
+                      (lora_config.max_loras or 1)) if lora_config else 0
+        self.vocab_size = config.vocab_size + lora_vocab
+        self.org_vocab_size = config.vocab_size
+        if get_pp_group().is_first_rank or (config.tie_word_embeddings
+                                            and get_pp_group().is_last_rank):
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                config.hidden_size,
+                org_num_embeddings=config.vocab_size,
+                quant_config=quant_config,
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: HunYuanDecoderLayer(config=config,
+                                               layer_id=int(
+                                                    prefix.split(".")[-1]),
+                                               cache_config=cache_config,
+                                               quant_config=quant_config,
+                                               prefix=prefix),
+            prefix=f"{prefix}.layers")
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor],
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors],
+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.get_input_embeddings(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        cla_factor = getattr(self.config, "cla_share_factor", 1)
+        prev_kv_states = None
+        for i in range(self.start_layer, self.end_layer):
+            layer = self.layers[i]
+            hidden_states, residual, kv_states = layer(
+                positions,
+                hidden_states,
+                residual,
+                prev_kv_states,
+            )
+
+            if (i - self.start_layer) % cla_factor == 0:
+                prev_kv_states = kv_states
+            else:
+                prev_kv_states = None
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({
+                "hidden_states": hidden_states,
+                "residual": residual
+            })
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+
+class HunYuanForCausalLM(nn.Module, SupportsLoRA):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": [
+            "gate_proj",
+            "up_proj",
+        ],
+    }
+
+    # LoRA specific attributes
+    supported_lora_modules = [
+        "qkv_proj", "o_proj", "gate_up_proj", "down_proj", "embed_tokens",
+        "lm_head"
+    ]
+    embedding_modules = {
+        "embed_tokens": "input_embeddings",
+        "lm_head": "output_embeddings",
+    }
+    embedding_padding_modules = ["lm_head"]
+    bitsandbytes_stacked_params_mapping = {
+        # shard_name, weight_name, index
+        "q_proj": ("qkv_proj", 0),
+        "k_proj": ("qkv_proj", 1),
+        "v_proj": ("qkv_proj", 2),
+        "gate_proj": ("gate_up_proj", 0),
+        "up_proj": ("gate_up_proj", 1),
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        lora_config = vllm_config.lora_config
+        self.config = config
+        self.lora_config = lora_config
+
+        self.quant_config = quant_config
+        self.model = HunYuanModel(config,
+                                  cache_config,
+                                  quant_config,
+                                  lora_config=lora_config,
+                                  prefix=maybe_prefix(prefix, "model"))
+        if get_pp_group().is_last_rank:
+            self.unpadded_vocab_size = config.vocab_size
+            if lora_config:
+                self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
+            self.lm_head = ParallelLMHead(
+                self.unpadded_vocab_size,
+                config.hidden_size,
+                org_num_embeddings=config.vocab_size,
+                padding_size=DEFAULT_VOCAB_PADDING_SIZE
+                # We need bigger padding if using lora for kernel
+                # compatibility
+                if not lora_config else lora_config.lora_vocab_padding_size,
+                quant_config=quant_config,
+            )
+            if config.tie_word_embeddings:
+                self.lm_head.weight = self.model.embed_tokens.weight
+
+            logit_scale = getattr(config, "logit_scale", 1.0)
+            self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
+                                                    config.vocab_size,
+                                                    logit_scale)
+            self.sampler = Sampler()
+        else:
+            self.lm_head = PPMissingLayer()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        model_output = self.model(input_ids, positions, intermediate_tensors)
+        return model_output
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[torch.Tensor]:
+        logits = self.logits_processor(self.lm_head, hidden_states,
+                                       sampling_metadata)
+        return logits
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[SamplerOutput]:
+        next_tokens = self.sampler(logits, sampling_metadata)
+        return next_tokens
+
+    def make_empty_intermediate_tensors(
+            self, batch_size: int, dtype: torch.dtype,
+            device: torch.device) -> IntermediateTensors:
+        return IntermediateTensors({
+            "hidden_states":
+            torch.zeros((batch_size, self.config.hidden_size),
+                        dtype=dtype,
+                        device=device),
+            "residual":
+            torch.zeros((batch_size, self.config.hidden_size),
+                        dtype=dtype,
+                        device=device),
+        })
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        cla_factor = getattr(self.config, "cla_share_factor", 1)
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+
+        if getattr(self.config, "num_experts", None):
+            # Params for weights, fp8 weight scales, fp8 activation scales
+            # (param_name, weight_name, expert_id, shard_id)
+            expert_params_mapping = FusedMoE.make_expert_params_mapping(
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=self.config.num_experts)
+        else:
+            expert_params_mapping = {}
+
+        params_dict = dict(self.named_parameters())
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if ("rotary_emb.cos_cached" in name
+                    or "rotary_emb.sin_cached" in name):
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+            # With tie_word_embeddings, we can skip lm_head.weight
+            # The weight might appear unnecessarily in the files if the model is
+            # processed with quantization, LoRA, fine-tuning, etc.
+            if self.config.tie_word_embeddings and "lm_head.weight" in name:
+                continue
+            if (self.quant_config is not None and
+                (scale_name := self.quant_config.get_cache_scale(name))):
+                # Loading kv cache scales for compressed-tensors quantization
+                param = params_dict[scale_name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                loaded_weight = loaded_weight[0]
+                weight_loader(param, loaded_weight)
+                continue
+            for (param_name, weight_name, shard_id) in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                if "mlp.experts" in name:
+                    continue
+                # cross layer only have q_proj, skip qkv pack
+                if weight_name == ".q_proj":
+                    match = re.search(r'layers\.\d+', name)
+                    if match:
+                        layer_id = int(match.group(0).split('.')[-1])
+                        if cla_factor > 1 and layer_id % cla_factor != 0:
+                            continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                for mapping in expert_params_mapping:
+                    param_name, weight_name, expert_id, shard_id = mapping
+                    if weight_name not in name:
+                        continue
+                    name = name.replace(weight_name, param_name)
+                    # Skip layers on other devices.
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    param = params_dict[name]
+                    weight_loader = param.weight_loader
+                    weight_loader(param,
+                                  loaded_weight,
+                                  name,
+                                  shard_id=shard_id,
+                                  expert_id=expert_id)
+                    break
+                else:
+                    # Remapping the name of FP8 kv-scale.
+                    name = maybe_remap_kv_scale_name(name, params_dict)
+                    if name is None:
+                        continue
+
+                    if is_pp_missing_parameter(name, self):
+                        continue
+
+                    if "mlp.gate.wg." in name:
+                        name = name.replace("wg.", "")
+
+                    param = params_dict[name]
+                    weight_loader = getattr(param, "weight_loader",
+                                            default_weight_loader)
+                    weight_loader(param, loaded_weight)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/qwen.py b/vllm_mlu/vllm_mlu/model_executor/models/qwen.py
new file mode 100644
index 000000000..4f26ba5ca
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/qwen.py
@@ -0,0 +1,199 @@
+import torch
+
+from typing import Optional, Union
+from transformers import PretrainedConfig
+
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.models.qwen import QWenAttention, QWenBlock, QWenModel
+from vllm.sequence import IntermediateTensors
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__qwen__QwenAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.c_attn(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.chunk(chunks=3, dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk, _ = qkv.split([self.head_dim * self.num_heads * 2, self.head_dim * self.num_heads], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.c_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__qwen__QWenBlock__init__(
+    self,
+    config: PretrainedConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+):
+    super(QWenBlock, self).__init__()
+    self.ln_1 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+
+    rope_theta = getattr(config, "rope_theta", 10000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+    self.attn = QWenAttention(config.hidden_size,
+                              config.num_attention_heads,
+                              config.max_position_embeddings,
+                              rope_theta=rope_theta,
+                              rope_scaling=rope_scaling,
+                              cache_config=cache_config,
+                              quant_config=quant_config,
+                              prefix=f"{prefix}.attn")
+
+    self.ln_2 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: 1) use FeedForward instead of MLP
+            2) prepare to perf per-tensor sq cases if suitable
+    '''
+    self.mlp = FeedForward(hidden_size=config.hidden_size,
+                            intermediate_size=config.intermediate_size // 2,
+                            hidden_act='silu',
+                            up_proj_name='gate_up_proj',
+                            is_gated=True,
+                            down_proj_name='c_proj',
+                            bias=False,
+                            quant_config=quant_config)
+
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.attn.c_attn.quant_method.skip_quant_input = True
+        self.mlp.gate_up_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+        self.quant_fusion_mlp_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen__QWenBlock__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.ln_1
+    mlp_layernorm = self.ln_2
+    if self.is_per_tesnor_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.ln_1, self.attn.c_attn.scale_to_int)
+            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
+                self.ln_2, self.mlp.gate_up_proj.scale_to_int)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+    elif self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.ln_1, self.attn.c_attn.smooth, dynamic_quant=True)
+            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
+                self.ln_2, self.mlp.gate_up_proj.smooth, dynamic_quant=True)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+    return decoder_layer_forward_base(
+        positions=positions,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.attn,
+        post_layernorm=mlp_layernorm,
+        mlp=self.mlp,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+        post_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen__QWenModel__forward(
+    self,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors],
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids, positions, intermediate_tensors,
+        self.h, self.start_layer, self.end_layer,
+        self.get_input_embeddings,
+        self.ln_f,
+        inputs_embeds
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(QWenAttention,
+                             QWenAttention.forward,
+                             vllm__module_executor__models__qwen__QwenAttention__forward)
+MluHijackObject.apply_hijack(QWenBlock,
+                             QWenBlock.__init__,
+                             vllm__module_executor__models__qwen__QWenBlock__init__)
+MluHijackObject.apply_hijack(QWenBlock,
+                             QWenBlock.forward,
+                             vllm__module_executor__models__qwen__QWenBlock__forward)
+MluHijackObject.apply_hijack(QWenModel,
+                             QWenModel.forward,
+                             vllm__module_executor__models__qwen__QWenModel__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/qwen2.py b/vllm_mlu/vllm_mlu/model_executor/models/qwen2.py
new file mode 100644
index 000000000..9f5d4853c
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/qwen2.py
@@ -0,0 +1,230 @@
+import torch
+
+from typing import Optional, Union
+from transformers import Qwen2Config
+
+from vllm.attention import AttentionType
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.models.qwen2 import Qwen2Attention, Qwen2DecoderLayer, Qwen2Model
+from vllm.sequence import IntermediateTensors
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__qwen2__Qwen2Attention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__qwen2__Qwen2DecoderLayer____init__(
+    self,
+    config: Qwen2Config,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(Qwen2DecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    # Requires transformers > 4.32.0
+    rope_theta = getattr(config, "rope_theta", 1000000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+
+    # By default, Qwen2 uses causal attention as it is a decoder-only model.
+    # You can override the HF config with `is_causal=False` to enable
+    # bidirectional attention, which is used in some embedding models
+    # (e.g. Alibaba-NLP/gte-Qwen2-7B-instruct)
+    if getattr(config, "is_causal", True):
+        attn_type = AttentionType.DECODER
+    else:
+        attn_type = AttentionType.ENCODER_ONLY
+
+    self.self_attn = Qwen2Attention(
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        max_position=config.max_position_embeddings,
+        num_kv_heads=config.num_key_value_heads,
+        rope_theta=rope_theta,
+        cache_config=cache_config,
+        quant_config=quant_config,
+        rope_scaling=rope_scaling,
+        prefix=f"{prefix}.self_attn",
+        attn_type=attn_type,
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.mlp = FeedForward(hidden_size=config.hidden_size,
+                            intermediate_size=config.intermediate_size,
+                            hidden_act='silu',
+                            up_proj_name='gate_up_proj',
+                            is_gated=True,
+                            down_proj_name='down_proj',
+                            bias=False,
+                            quant_config=quant_config,
+                            prefix=f"{prefix}.mlp")
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                    eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf per-tensor sq cases if suitable
+    '''
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attn.qkv_proj.quant_method.skip_quant_input = True
+        self.mlp.gate_up_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+        self.quant_fusion_mlp_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen2__Qwen2DecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    mlp_layernorm = self.post_attention_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attn.qkv_proj.smooth
+                mlp_quant_scale = self.mlp.gate_up_proj.smooth
+            else:
+                attn_quant_scale = self.self_attn.qkv_proj.scale_to_int
+                mlp_quant_scale = self.mlp.gate_up_proj.scale_to_int
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
+                self.post_attention_layernorm, mlp_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+    return decoder_layer_forward_base(
+        positions=positions,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.self_attn,
+        post_layernorm=mlp_layernorm,
+        mlp=self.mlp,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+        post_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen2__Qwen2Model__forward(
+    self,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors] = None,
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids=input_ids,
+        positions=positions,
+        intermediate_tensors=intermediate_tensors,
+        layers=self.layers,
+        start_layer=self.start_layer,
+        end_layer=self.end_layer,
+        get_input_embeddings=self.embed_tokens,
+        norm=self.norm,
+        inputs_embeds=inputs_embeds
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(Qwen2Attention,
+                             Qwen2Attention.forward,
+                             vllm__module_executor__models__qwen2__Qwen2Attention__forward)
+MluHijackObject.apply_hijack(Qwen2DecoderLayer,
+                             Qwen2DecoderLayer.__init__,
+                             vllm__module_executor__models__qwen2__Qwen2DecoderLayer____init__)
+MluHijackObject.apply_hijack(Qwen2DecoderLayer,
+                             Qwen2DecoderLayer.forward,
+                             vllm__module_executor__models__qwen2__Qwen2DecoderLayer__forward)
+MluHijackObject.apply_hijack(Qwen2Model,
+                             Qwen2Model.forward,
+                             vllm__module_executor__models__qwen2__Qwen2Model__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/qwen2_5_vl.py b/vllm_mlu/vllm_mlu/model_executor/models/qwen2_5_vl.py
new file mode 100644
index 000000000..5fa98254c
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/qwen2_5_vl.py
@@ -0,0 +1,276 @@
+from typing import Optional
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from vllm.logger import init_logger
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig
+from vllm.model_executor.models.qwen2_5_vl import (
+    Qwen2_5_VisionTransformer, Qwen2_5_VisionAttention)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.attention.selector import _Backend
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionTransformer__init_org = Qwen2_5_VisionTransformer.__init__
+
+def vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionTransformer____init__(
+    self,
+    vision_config: Qwen2_5_VLVisionConfig,
+    norm_eps: float = 1e-6,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = ""
+):
+    vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionTransformer__init_org(
+        self, vision_config, norm_eps, quant_config, prefix)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use mlu_ops.flash_atten for better performance
+    '''
+    self.attn_backend = _Backend.FLASH_ATTN
+
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+def vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionTransformer__forward(
+    self,
+    x: torch.Tensor,
+    grid_thw: torch.Tensor
+):
+    # patchify
+    hidden_states = x.to(device=self.device, dtype=self.dtype)
+    hidden_states = self.patch_embed(hidden_states)
+
+    # compute position embedding
+    rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+    # windows attention
+    window_index, cu_window_seqlens = self.get_window_index(grid_thw)
+    cu_window_seqlens = torch.tensor(
+        cu_window_seqlens,
+        device=hidden_states.device,
+        dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32)
+    cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
+    seq_len, _ = hidden_states.size()
+    hidden_states = hidden_states.reshape(
+        seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+    hidden_states = hidden_states[window_index, :, :]
+    hidden_states = hidden_states.reshape(seq_len, -1)
+    rotary_pos_emb = rotary_pos_emb.reshape(
+        seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+    rotary_pos_emb = rotary_pos_emb[window_index, :, :]
+    rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: compute cos and sin once for rope
+    '''
+    # compute cos sin for apply_rope
+    cos = rotary_pos_emb.cos()
+    sin = rotary_pos_emb.sin()
+    cos = repeat(cos, "... d -> ... (2 d)")
+    sin = repeat(sin, "... d -> ... (2 d)")
+    rotary_pos_emb.cos = cos
+    rotary_pos_emb.sin = sin
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    # compute cu_seqlens
+    cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2],
+                                         grid_thw[:, 0]).cumsum(
+                                             dim=0, dtype=torch.int32)
+    cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)
+
+    # transformers
+    hidden_states = hidden_states.unsqueeze(1)
+
+    # pre-compute seqlens for window/full attn to reduce cuMemcpy operations
+    max_seqlen_full, seqlens_full = self.compute_attn_mask_seqlen(
+        cu_seqlens)
+    max_seqlen_window, seqlens_window = self.compute_attn_mask_seqlen(
+        cu_window_seqlens)
+    for layer_num, blk in enumerate(self.blocks):
+        if layer_num in self.fullatt_block_indexes:
+            cu_seqlens_now = cu_seqlens
+            max_seqlen_now = max_seqlen_full
+            seqlens_now = seqlens_full
+        else:
+            cu_seqlens_now = cu_window_seqlens
+            max_seqlen_now = max_seqlen_window
+            seqlens_now = seqlens_window
+
+        hidden_states = blk(
+            hidden_states,
+            cu_seqlens=cu_seqlens_now,
+            rotary_pos_emb=rotary_pos_emb,
+            max_seqlen=max_seqlen_now,
+            seqlens=seqlens_now,
+        )
+
+    # For Qwen2.5-VL-3B, float16 will overflow at last block
+    # for long visual tokens sequences.
+    if hidden_states.dtype == torch.float16:
+        hidden_states = cast_overflow_tensors(hidden_states)
+
+    # adapter
+    hidden_states = self.merger(hidden_states)
+    reverse_indices = torch.argsort(window_index)
+    hidden_states = hidden_states[reverse_indices, :]
+    return hidden_states
+
+def vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionAttention__forward(
+    self,
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    rotary_pos_emb: torch.Tensor,
+    max_seqlen: Optional[int] = None,  # Only used for Flash Attention
+    seqlens: Optional[list[int]] = None,  # Only used for xFormers
+):
+    # [s, b, c] --> [s, b, 3 * head * head_dim]
+    x, _ = self.qkv(x)
+
+    # [s, b, 3 * head * head_dim] -> 3 * [s, b, head, head_dim]
+    q, k, v = self.split_qkv(x)
+    batch_size = q.shape[1]
+
+    q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()
+               for x in (q, k, v))
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: apply mlu_ops.apply_rotary
+    '''
+    head_dim = q.shape[-1]
+    if rotary_pos_emb is not None:
+        sin = rotary_pos_emb.sin
+        cos = rotary_pos_emb.cos
+        from vllm_mlu import _mlu_ops as mlu_ops
+        q = q.float()
+        q = mlu_ops.rotary_embedding(
+            q, sin, cos, None, None, False, False, False, q.shape[1]
+        )
+        k = k.float()
+        k = mlu_ops.rotary_embedding(
+            k, sin, cos, None, None, False, False, False, k.shape[1]
+        )
+        q = q.type_as(v)
+        k = k.type_as(v)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    if self.attn_backend == _Backend.FLASH_ATTN:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: apply mlu_ops.flash_attention
+        '''
+        q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
+        output = mlu_ops.flash_attention(q,
+                                         k,
+                                         v,
+                                         out=None,
+                                         cu_seq_lens_q=cu_seqlens,
+                                         cu_seq_lens_kv=cu_seqlens,
+                                         max_seq_len_q=max_seqlen,
+                                         max_seq_len_kv=max_seqlen,
+                                         alibi_slope=None,
+                                         attn_bias=None,
+                                         softmax_scale=head_dim ** -0.5,
+                                         is_causal=False)
+        context_layer = rearrange(output,
+                                  "(b s) ... -> b s ...",
+                                  b=batch_size)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    elif self.attn_backend == _Backend.TORCH_SDPA:
+        # Execute attention entry by entry for speed & less VRAM.
+        outputs = []
+        for i in range(1, len(cu_seqlens)):
+            start_idx = cu_seqlens[i - 1]
+            end_idx = cu_seqlens[i]
+            q_i = q[:, start_idx:end_idx]
+            k_i = k[:, start_idx:end_idx]
+            v_i = v[:, start_idx:end_idx]
+            q_i, k_i, v_i = (rearrange(x, "b s h d -> b h s d")
+                             for x in [q_i, k_i, v_i])
+            output_i = F.scaled_dot_product_attention(q_i,
+                                                      k_i,
+                                                      v_i,
+                                                      dropout_p=0.0)
+            output_i = rearrange(output_i, "b h s d -> b s h d ")
+            outputs.append(output_i)
+        context_layer = torch.cat(outputs, dim=1)
+    elif self.attn_backend == _Backend.XFORMERS:
+        from xformers import ops as xops
+        from xformers.ops.fmha.attn_bias import BlockDiagonalMask
+
+        seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+        attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,
+                                                   kv_seqlen=None)
+
+        context_layer = xops.memory_efficient_attention_forward(
+            q, k, v, attn_bias=attn_bias, p=0, scale=None)
+    context_layer = rearrange(context_layer,
+                              "b s h d -> s b (h d)").contiguous()
+
+    output, _ = self.proj(context_layer)
+    return output
+
+
+vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionAttention__init_org = Qwen2_5_VisionAttention.__init__
+
+def vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionAttention____init__(
+    self,
+    embed_dim: Optional[int] = None,
+    num_heads: Optional[int] = None,
+    projection_size: Optional[int] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionAttention__init_org(
+            self, embed_dim, num_heads, projection_size, quant_config, prefix)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use mlu_ops.flash_atten for better performance
+    '''
+    self.attn_backend = _Backend.FLASH_ATTN
+
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+MluHijackObject.apply_hijack(Qwen2_5_VisionTransformer,
+                             Qwen2_5_VisionTransformer.__init__,
+                             vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionTransformer____init__)
+MluHijackObject.apply_hijack(Qwen2_5_VisionTransformer,
+                             Qwen2_5_VisionTransformer.forward,
+                             vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionTransformer__forward)
+MluHijackObject.apply_hijack(Qwen2_5_VisionAttention,
+                             Qwen2_5_VisionAttention.forward,
+                             vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionAttention__forward)
+MluHijackObject.apply_hijack(Qwen2_5_VisionAttention,
+                             Qwen2_5_VisionAttention.__init__,
+                             vllm__module_executor__models__qwen2_5_vl__Qwen2_5_VisionAttention____init__)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/qwen2_moe.py b/vllm_mlu/vllm_mlu/model_executor/models/qwen2_moe.py
new file mode 100644
index 000000000..4d75039c0
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/qwen2_moe.py
@@ -0,0 +1,459 @@
+import torch
+import re
+import torch.nn.functional as F
+
+from typing import Optional, Tuple, Iterable, Set
+from transformers import PretrainedConfig
+
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.models.qwen2_moe import (
+    Qwen2MoeAttention, Qwen2MoeDecoderLayer, Qwen2MoeForCausalLM, Qwen2MoeModel)
+from vllm.model_executor.models.utils import (extract_layer_index,
+                                              is_pp_missing_parameter)
+from vllm.sequence import IntermediateTensors
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+class Qwen2MoeSparseMoeBlock(SparseMoeMlp):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+    ):
+        super().__init__(num_experts=config.num_experts,
+                         top_k=config.num_experts_per_tok,
+                         hidden_size=config.hidden_size,
+                         intermediate_size=config.moe_intermediate_size,
+                         up_proj_name="gate_up_proj",
+                         is_gated=True,
+                         down_proj_name="down_proj",
+                         has_bias=False,
+                         skip_bias_add=False,
+                         renormalize=config.norm_topk_prob,
+                         hidden_act=config.hidden_act,
+                         params_dtype=None,
+                         quant_config=quant_config,
+                         is_use_fused_moe=True)
+        self.config = config
+        self.shared_expert = None
+        self.shared_expert_gate = None
+        if config.shared_expert_intermediate_size > 0:
+            self.shared_expert = FeedForward(hidden_size=config.hidden_size,
+                                             intermediate_size=config.shared_expert_intermediate_size,
+                                             hidden_act=config.hidden_act,
+                                             up_proj_name='gate_up_proj',
+                                             is_gated=True,
+                                             down_proj_name='down_proj',
+                                             bias=False,
+                                             quant_config=quant_config,
+                                             reduce_results=False)
+            self.shared_expert_gate = ReplicatedLinear(config.hidden_size,
+                                                       1,
+                                                       bias=False,
+                                                       params_dtype=self.params_dtype,
+                                                       quant_config=None)
+
+
+    def forward(self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_output = None
+        if self.shared_expert is not None:
+            shared_output = self.shared_expert(hidden_states)
+            if self.shared_expert_gate is not None:
+                gate_output = self.shared_expert_gate(hidden_states)
+                shared_output = F.sigmoid(gate_output[0]) * shared_output
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.forward_experts(hidden_states, router_logits, residual)
+
+        if shared_output is not None:
+            final_hidden_states = final_hidden_states + shared_output
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+
+def vllm__module_executor__models__qwen2moe__Qwen2MoeAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__qwen2moe__Qwen2MoeDecoderLayer____init__(
+    self,
+    config: PretrainedConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(Qwen2MoeDecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    rope_theta = getattr(config, "rope_theta", 10000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+    max_position_embeddings = getattr(config, "max_position_embeddings",
+                                        8192)
+    self.self_attn = Qwen2MoeAttention(
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=config.num_key_value_heads,
+        rope_theta=rope_theta,
+        rope_scaling=rope_scaling,
+        max_position_embeddings=max_position_embeddings,
+        cache_config=cache_config,
+        quant_config=quant_config,
+        prefix=f"{prefix}.self_attn",
+    )
+
+    # Note: Qwen/Qwen2-57B-A14B-Instruct does not have
+    # `mlp_only_layers` in the config.
+    layer_idx = extract_layer_index(prefix)
+    mlp_only_layers = ([] if not hasattr(config, "mlp_only_layers") else
+                        config.mlp_only_layers)
+    if (layer_idx not in mlp_only_layers) and (
+            config.num_experts > 0 and
+        (layer_idx + 1) % config.decoder_sparse_step == 0):
+        self.mlp = Qwen2MoeSparseMoeBlock(config=config,
+                                          quant_config=quant_config)
+    else:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: use FeedForward instead of MLP
+        '''
+        self.mlp = FeedForward(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            up_proj_name='gate_up_proj',
+            is_gated=True,
+            down_proj_name='down_proj',
+            bias=False,
+            quant_config=quant_config,
+        )
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                    eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf per-tensor sq cases if suitable. For moe
+        model, we only do quant fusion in attn block.
+    '''
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attn.qkv_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen2moe__Qwen2MoeForCausalLM__load_weights(
+    self,
+    weights: Iterable[Tuple[str, torch.Tensor]]
+):
+    start_expert_id = 0
+    stacked_params_mapping = [
+        # (param_name, shard_name, shard_id)
+        ("qkv_proj", "q_proj", "q"),
+        ("qkv_proj", "k_proj", "k"),
+        ("qkv_proj", "v_proj", "v"),
+        ("gate_up_proj", "gate_proj", 0),
+        ("gate_up_proj", "up_proj", 1),
+    ]
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: delete expert_params_mapping for useless
+    '''
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    params_dict = dict(self.named_parameters())
+    loaded_params: Set[str] = set()
+    for name, loaded_weight in weights:
+        if "rotary_emb.inv_freq" in name:
+            continue
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: replace expert_id in weight to named_expert_id in params_dict
+        '''
+        if start_expert_id > 0 and "mlp.experts." in name:
+            expert_str = re.search(r'experts\.\d+', name).group(0)
+            expert_id=int(expert_str.split(".")[1])
+            named_expert_id = expert_id - start_expert_id
+            old_expert_name = f"experts.{expert_id}"
+            new_expert_name = f"experts.{named_expert_id}"
+            name = name.replace(old_expert_name, new_expert_name)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            if weight_name not in name:
+                continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: delete if "mlp.experts" in name: continue condition
+            '''
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            name = name.replace(weight_name, param_name)
+            # Skip loading extra bias for GPTQ models.
+            if ((name.endswith(".bias") or name.endswith("_bias"))
+                        and name not in params_dict):
+                continue
+            # Skip layers on other devices.
+            if is_pp_missing_parameter(name, self):
+                continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
+            '''
+            # Skip experts that are not assigned to this worker.
+            if (("mlp.experts." in name or "mlp.shared_expert." in name or "mlp.shared_expert_gate." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            weight_loader(param, loaded_weight, shard_id)
+            break
+        else:
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: delete for mapping in expert_params_mapping condition
+            '''
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            # Skip loading extra bias for GPTQ models.
+            if ((name.endswith(".bias") or name.endswith("_bias"))
+                            and name not in params_dict):
+                continue
+            # Skip layers on other devices.
+            if is_pp_missing_parameter(name, self):
+                continue
+            # Remapping the name of FP8 kv-scale.
+            if name.endswith("kv_scale"):
+                remapped_kv_scale_name = name.replace(
+                    ".kv_scale", ".attn.kv_scale")
+                if remapped_kv_scale_name not in params_dict:
+                    logger.warning_once(
+                        "Found kv scale in the checkpoint "
+                        f"(e.g. {name}), but not found the expected "
+                        f"name in the model "
+                        f"(e.g. {remapped_kv_scale_name}). "
+                        "kv-scale is not loaded.")
+                    continue
+                else:
+                    name = remapped_kv_scale_name
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition
+            '''
+            # Skip experts that are not assigned to this worker.
+            if (("mlp.experts." in name or "mlp.shared_expert." in name or "mlp.shared_expert_gate." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            weight_loader(param, loaded_weight)
+        loaded_params.add(name)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack params
+    '''
+    for name, m in self.model.named_modules():
+        if isinstance(m, SparseMoeMlp):
+            m.pack_params()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return loaded_params
+
+
+def vllm__module_executor__models__qwen2moe__Qwen2MoeDecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attn.qkv_proj.smooth
+            else:
+                attn_quant_scale = self.self_attn.qkv_proj.scale_to_int
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+
+    return decoder_layer_forward_base(
+        positions=positions,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.self_attn,
+        post_layernorm=self.post_attention_layernorm,
+        mlp=self.mlp,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen2moe__Qwen2MoeModel__forward(
+    self,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors] = None,
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids=input_ids,
+        positions=positions,
+        intermediate_tensors=intermediate_tensors,
+        layers=self.layers,
+        start_layer=self.start_layer,
+        end_layer=self.end_layer,
+        get_input_embeddings=self.embed_tokens,
+        norm=self.norm,
+        inputs_embeds=inputs_embeds,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(Qwen2MoeAttention,
+                             Qwen2MoeAttention.forward,
+                             vllm__module_executor__models__qwen2moe__Qwen2MoeAttention__forward)
+MluHijackObject.apply_hijack(Qwen2MoeDecoderLayer,
+                             Qwen2MoeDecoderLayer.__init__,
+                             vllm__module_executor__models__qwen2moe__Qwen2MoeDecoderLayer____init__)
+MluHijackObject.apply_hijack(Qwen2MoeForCausalLM,
+                             Qwen2MoeForCausalLM.load_weights,
+                             vllm__module_executor__models__qwen2moe__Qwen2MoeForCausalLM__load_weights)
+MluHijackObject.apply_hijack(Qwen2MoeDecoderLayer,
+                             Qwen2MoeDecoderLayer.forward,
+                             vllm__module_executor__models__qwen2moe__Qwen2MoeDecoderLayer__forward)
+MluHijackObject.apply_hijack(Qwen2MoeModel,
+                             Qwen2MoeModel.forward,
+                             vllm__module_executor__models__qwen2moe__Qwen2MoeModel__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/qwen2_vl.py b/vllm_mlu/vllm_mlu/model_executor/models/qwen2_vl.py
new file mode 100644
index 000000000..0d9f1c8e7
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/qwen2_vl.py
@@ -0,0 +1,299 @@
+from typing import Optional
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from transformers.models.qwen2_vl.configuration_qwen2_vl import (
+    Qwen2VLVisionConfig)
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import QuickGELU
+from vllm.model_executor.models.qwen2_vl import (
+    Qwen2VisionMLP, Qwen2VisionTransformer,
+    Qwen2VisionAttention, Qwen2VLForConditionalGeneration)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.attention.selector import _Backend
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+vllm__module_executor__models__qwen2_vl__Qwen2VisionTransformer__init__org = Qwen2VisionTransformer.__init__
+
+def vllm__module_executor__models__qwen2_vl__Qwen2VisionTransformer__init__(
+    self,
+    vision_config: Qwen2VLVisionConfig,
+    norm_eps: float = 1e-6,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    vllm__module_executor__models__qwen2_vl__Qwen2VisionTransformer__init__org(
+        self, vision_config, norm_eps, quant_config, prefix
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use mlu_ops.flash_atten for better performance
+    '''
+    self.attn_backend = _Backend.FLASH_ATTN
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen2_vl__Qwen2VisionTransformer__forward(
+    self,
+    x: torch.Tensor,
+    grid_thw: torch.Tensor
+):
+    # patchify
+    x = x.to(device=self.device, dtype=self.dtype)
+    x = self.patch_embed(x)
+
+    # compute position embedding
+    rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    # compute cos sin for apply_rope
+    cos = rotary_pos_emb.cos()
+    sin = rotary_pos_emb.sin()
+    cos = repeat(cos, "... d -> ... (2 d)")
+    sin = repeat(sin, "... d -> ... (2 d)")
+    rotary_pos_emb.cos = cos
+    rotary_pos_emb.sin = sin
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    # compute cu_seqlens
+    cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2],
+                                         grid_thw[:, 0]).cumsum(
+                                             dim=0, dtype=torch.int32)
+    cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)
+
+    # transformers
+    x = x.unsqueeze(1)
+
+    # pre-compute seqlens for attn mask to reduce cuMemcpy operations
+    max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
+    for blk in self.blocks:
+        x = blk(
+            x,
+            cu_seqlens=cu_seqlens,
+            rotary_pos_emb=rotary_pos_emb,
+            max_seqlen=max_seqlen,
+            seqlens=seqlens,
+        )
+
+    # adapter
+    x = self.merger(x)
+
+    return x
+
+
+def vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention__forward(
+    self,
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    rotary_pos_emb: torch.Tensor,
+    max_seqlen: Optional[int] = None,  # Only used for Flash Attention
+    seqlens: Optional[list[int]] = None,  # Only used for xFormers
+) -> torch.Tensor:
+    # [s, b, c] --> [s, b, 3 * head * head_dim]
+    x, _ = self.qkv(x)
+
+    # [s, b, 3 * head * head_dim] -> 3 * [s, b, head, head_dim]
+    q, k, v = self.split_qkv(x)
+    batch_size = q.shape[1]
+
+    q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous()
+               for x in (q, k, v))
+
+    if rotary_pos_emb is not None:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: apply mlu_ops.apply_rotary
+        '''
+        head_dim = q.shape[-1]
+        sin = rotary_pos_emb.sin
+        cos = rotary_pos_emb.cos
+        from vllm_mlu import _mlu_ops as mlu_ops
+        q = q.float()
+        q = mlu_ops.rotary_embedding(
+            q, sin, cos, None, None, False, False, False, q.shape[1]
+        )
+        k = k.float()
+        k = mlu_ops.rotary_embedding(
+            k, sin, cos, None, None, False, False, False, k.shape[1]
+        )
+        q = q.type_as(v)
+        k = k.type_as(v)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+    if self.attn_backend == _Backend.FLASH_ATTN:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: apply mlu_ops.flash_attention
+        '''
+        q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
+        # max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+        output = mlu_ops.flash_attention(q,
+                                         k,
+                                         v,
+                                         out=None,
+                                         cu_seq_lens_q=cu_seqlens,
+                                         cu_seq_lens_kv=cu_seqlens,
+                                         max_seq_len_q=max_seqlen,
+                                         max_seq_len_kv=max_seqlen,
+                                         alibi_slope=None,
+                                         attn_bias=None,
+                                         softmax_scale=head_dim ** -0.5,
+                                         is_causal=False)
+        context_layer = rearrange(output,
+                                  "(b s) ... -> b s ...",
+                                  b=batch_size)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    elif self.attn_backend == _Backend.TORCH_SDPA:
+        # Execute attention entry by entry for speed & less VRAM.
+        outputs = []
+        for i in range(1, len(cu_seqlens)):
+            start_idx = cu_seqlens[i - 1]
+            end_idx = cu_seqlens[i]
+            q_i = q[:, start_idx:end_idx]
+            k_i = k[:, start_idx:end_idx]
+            v_i = v[:, start_idx:end_idx]
+            q_i, k_i, v_i = (rearrange(x, "b s h d -> b h s d")
+                                for x in [q_i, k_i, v_i])
+            output_i = F.scaled_dot_product_attention(q_i,
+                                                        k_i,
+                                                        v_i,
+                                                        dropout_p=0.0)
+            output_i = rearrange(output_i, "b h s d -> b s h d ")
+            outputs.append(output_i)
+        context_layer = torch.cat(outputs, dim=1)
+    elif self.attn_backend == _Backend.XFORMERS:
+        from xformers import ops as xops
+        from xformers.ops.fmha.attn_bias import BlockDiagonalMask
+
+        attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,
+                                                    kv_seqlen=None,
+                                                    device=q.device)
+
+        context_layer = xops.memory_efficient_attention_forward(
+            q, k, v, attn_bias=attn_bias, p=0, scale=None)
+    context_layer = rearrange(context_layer,
+                              "b s h d -> s b (h d)").contiguous()
+
+    output, _ = self.proj(context_layer)
+    return output
+
+
+vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention__init_org = Qwen2VisionAttention.__init__
+
+def vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention____init__(
+    self,
+    embed_dim: int,
+    num_heads: int,
+    projection_size: int,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention__init_org(
+            self, embed_dim, num_heads, projection_size, quant_config, prefix)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use mlu_ops.flash_atten for better performance
+    '''
+    self.attn_backend = _Backend.FLASH_ATTN
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+def vllm__module_executor__models__qwen2_vl___maybe_ignore_quant_config(
+    self,
+    quant_config: QuantizationConfig
+):
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: quantization for vit not yet supported
+    '''
+    if quant_config is not None:
+        logger.warning("Quantization for VisionTransformer not yet supported.")
+        return None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm_module_executor__models__qwen2_vl__Qwen2VisionMLP__forward(
+    self,
+    x: torch.Tensor
+) -> torch.Tensor:
+    x_parallel, _ = self.fc1(x)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: better acc than mlu_ops.active for half precision
+    '''
+    if x_parallel.dtype == torch.half and isinstance(self.act, QuickGELU):
+        x_parallel = self.act.forward_native(x_parallel)
+    else:
+        x_parallel = self.act(x_parallel)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    x, _ = self.fc2(x_parallel)
+    return x
+
+
+MluHijackObject.apply_hijack(Qwen2VisionTransformer,
+                             Qwen2VisionTransformer.__init__,
+                             vllm__module_executor__models__qwen2_vl__Qwen2VisionTransformer__init__)
+MluHijackObject.apply_hijack(Qwen2VisionTransformer,
+                             Qwen2VisionTransformer.forward,
+                             vllm__module_executor__models__qwen2_vl__Qwen2VisionTransformer__forward)
+MluHijackObject.apply_hijack(Qwen2VisionAttention,
+                             Qwen2VisionAttention.forward,
+                             vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention__forward)
+MluHijackObject.apply_hijack(Qwen2VisionAttention,
+                             Qwen2VisionAttention.__init__,
+                             vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention____init__)
+MluHijackObject.apply_hijack(Qwen2VLForConditionalGeneration,
+                             Qwen2VLForConditionalGeneration._maybe_ignore_quant_config,
+                             vllm__module_executor__models__qwen2_vl___maybe_ignore_quant_config)
+MluHijackObject.apply_hijack(Qwen2VisionMLP,
+                             Qwen2VisionMLP.forward,
+                             vllm_module_executor__models__qwen2_vl__Qwen2VisionMLP__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/qwen3.py b/vllm_mlu/vllm_mlu/model_executor/models/qwen3.py
new file mode 100644
index 000000000..e0938e754
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/qwen3.py
@@ -0,0 +1,245 @@
+import torch
+
+from typing import Optional, Union
+from transformers import Qwen3Config
+
+from vllm.attention import AttentionType
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.models.qwen3 import Qwen3Attention, Qwen3DecoderLayer, Qwen3Model
+from vllm.sequence import IntermediateTensors
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__module_executor__models__qwen3__Qwen3Attention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    # qwen3
+    # Add qk-norm
+    q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim,
+                        self.head_dim)
+    q_by_head = self.q_norm.forward_native(q_by_head)
+    q = q_by_head.view(q.shape)
+    k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim,
+                        self.head_dim)
+    k_by_head = self.k_norm.forward_native(k_by_head)
+    k = k_by_head.view(k.shape)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk = torch.cat([q, k], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    q, k = qk.split([self.q_size, self.kv_size], dim=-1)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+
+def vllm__module_executor__models__qwen3__Qwen3DecoderLayer____init__(
+    self,
+    config: Qwen3Config,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(Qwen3DecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    # Requires transformers > 4.32.0
+    rope_theta = getattr(config, "rope_theta", 1000000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+
+    # By default, Qwen3 uses causal attention as it is a decoder-only model.
+    # You can override the HF config with `is_causal=False` to enable
+    # bidirectional attention, which is used in some embedding models
+    # (e.g. Alibaba-NLP/gte-Qwen3-7B-instruct)
+    if getattr(config, "is_causal", True):
+        attn_type = AttentionType.DECODER
+    else:
+        attn_type = AttentionType.ENCODER_ONLY
+
+    self.self_attn = Qwen3Attention(
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        max_position=config.max_position_embeddings,
+        num_kv_heads=config.num_key_value_heads,
+        rope_theta=rope_theta,
+        # qwen3
+        rms_norm_eps=config.rms_norm_eps,
+        qkv_bias=getattr(config, 'attention_bias', False),
+        head_dim=getattr(config, 'head_dim', None),
+        cache_config=cache_config,
+        quant_config=quant_config,
+        rope_scaling=rope_scaling,
+        prefix=f"{prefix}.self_attn",
+        attn_type=attn_type,
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use FeedForward instead of MLP
+    '''
+    self.mlp = FeedForward(hidden_size=config.hidden_size,
+                            intermediate_size=config.intermediate_size,
+                            hidden_act='silu',
+                            up_proj_name='gate_up_proj',
+                            is_gated=True,
+                            down_proj_name='down_proj',
+                            bias=False,
+                            quant_config=quant_config,
+                            prefix=f"{prefix}.mlp")
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                    eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf per-tensor sq cases if suitable
+    '''
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attn.qkv_proj.quant_method.skip_quant_input = True
+        self.mlp.gate_up_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+        self.quant_fusion_mlp_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen3__Qwen3DecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    mlp_layernorm = self.post_attention_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attn.qkv_proj.smooth
+                mlp_quant_scale = self.mlp.gate_up_proj.smooth
+            else:
+                attn_quant_scale = self.self_attn.qkv_proj.scale_to_int
+                mlp_quant_scale = self.mlp.gate_up_proj.scale_to_int
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
+                self.post_attention_layernorm, mlp_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+        mlp_layernorm = self.quant_fusion_mlp_layernorm
+
+    return decoder_layer_forward_base(
+        positions=positions,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.self_attn,
+        post_layernorm=mlp_layernorm,
+        mlp=self.mlp,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases,
+        post_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen3__Qwen3Model__forward(
+    self,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors] = None,
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, IntermediateTensors]:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids=input_ids,
+        positions=positions,
+        intermediate_tensors=intermediate_tensors,
+        layers=self.layers,
+        start_layer=self.start_layer,
+        end_layer=self.end_layer,
+        get_input_embeddings=self.embed_tokens,
+        norm=self.norm,
+        inputs_embeds=inputs_embeds
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(Qwen3Attention,
+                             Qwen3Attention.forward,
+                             vllm__module_executor__models__qwen3__Qwen3Attention__forward)
+MluHijackObject.apply_hijack(Qwen3DecoderLayer,
+                             Qwen3DecoderLayer.__init__,
+                             vllm__module_executor__models__qwen3__Qwen3DecoderLayer____init__)
+MluHijackObject.apply_hijack(Qwen3DecoderLayer,
+                             Qwen3DecoderLayer.forward,
+                             vllm__module_executor__models__qwen3__Qwen3DecoderLayer__forward)
+MluHijackObject.apply_hijack(Qwen3Model,
+                             Qwen3Model.forward,
+                             vllm__module_executor__models__qwen3__Qwen3Model__forward)
diff --git a/vllm_mlu/vllm_mlu/model_executor/models/qwen3_moe.py b/vllm_mlu/vllm_mlu/model_executor/models/qwen3_moe.py
new file mode 100644
index 000000000..e92dbf35b
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/model_executor/models/qwen3_moe.py
@@ -0,0 +1,444 @@
+import torch
+import re
+import torch.nn.functional as F
+
+from typing import Optional, Tuple, Iterable, Set
+from transformers import PretrainedConfig
+
+from vllm.config import CacheConfig
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.models.qwen3_moe import (
+    Qwen3MoeAttention, Qwen3MoeDecoderLayer, Qwen3MoeForCausalLM, Qwen3MoeModel)
+from vllm.model_executor.models.utils import (extract_layer_index,
+                                              is_pp_missing_parameter)
+from vllm.sequence import IntermediateTensors
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.models.layer_utils import (
+    decoder_layer_forward_base, decoder_model_forward_base_pp,
+    is_per_tensor_smoothquant, is_per_token_smoothquant,
+    quant_fusion_with_rmsnorm)
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+class Qwen3MoeSparseMoeBlock(SparseMoeMlp):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__(num_experts=config.num_experts,
+                         top_k=config.num_experts_per_tok,
+                         hidden_size=config.hidden_size,
+                         intermediate_size=config.moe_intermediate_size,
+                         up_proj_name="gate_up_proj",
+                         is_gated=True,
+                         down_proj_name="down_proj",
+                         has_bias=False,
+                         skip_bias_add=False,
+                         renormalize=config.norm_topk_prob,
+                         hidden_act=config.hidden_act,
+                         params_dtype=None,
+                         quant_config=quant_config,
+                         is_use_fused_moe=True)
+
+    def forward(self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.forward_experts(hidden_states, router_logits, residual)
+
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+def vllm__module_executor__models__qwen3moe__Qwen3MoeAttention__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+    smooth_quant_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
+    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+    # Add qk-norm
+    q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim,
+                       self.head_dim)
+    q_by_head = self.q_norm.forward_native(q_by_head)
+    q = q_by_head.view(q.shape)
+
+    k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim,
+                       self.head_dim)
+    k_by_head = self.k_norm.forward_native(k_by_head)
+    k = k_by_head.view(k.shape)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack q & k to fit tmo.apply_rotary
+    '''
+    qk = torch.cat([q, k], dim=-1)
+    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
+    q, k = qk.split([self.q_size, self.kv_size], dim=-1)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    attn_output = self.attn(q, k, v)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add residual
+    '''
+    output, _ = self.o_proj(attn_output, residual)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return output
+
+def vllm__module_executor__models__qwen3moe__Qwen3MoeDecoderLayer____init__(
+    self,
+    config: PretrainedConfig,
+    cache_config: Optional[CacheConfig] = None,
+    quant_config: Optional[QuantizationConfig] = None,
+    prefix: str = "",
+) -> None:
+    super(Qwen3MoeDecoderLayer, self).__init__()
+    self.hidden_size = config.hidden_size
+    rope_theta = getattr(config, "rope_theta", 10000)
+    rope_scaling = getattr(config, "rope_scaling", None)
+    max_position_embeddings = getattr(config, "max_position_embeddings",
+                                      8192)
+    self.self_attn = Qwen3MoeAttention(
+        hidden_size=self.hidden_size,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=config.num_key_value_heads,
+        rope_theta=rope_theta,
+        rope_scaling=rope_scaling,
+        max_position_embeddings=max_position_embeddings,
+        rms_norm_eps=config.rms_norm_eps,
+        qkv_bias=getattr(config, 'attention_bias', False),
+        head_dim=getattr(config, 'head_dim', None),
+        cache_config=cache_config,
+        quant_config=quant_config,
+        prefix=f"{prefix}.self_attn",
+    )
+
+    # `mlp_only_layers` in the config.
+    layer_idx = extract_layer_index(prefix)
+    mlp_only_layers = ([] if not hasattr(config, "mlp_only_layers") else
+                       config.mlp_only_layers)
+    if (layer_idx not in mlp_only_layers) and (
+            config.num_experts > 0 and
+        (layer_idx + 1) % config.decoder_sparse_step == 0):
+        self.mlp = Qwen3MoeSparseMoeBlock(config=config,
+                                          quant_config=quant_config,
+                                          prefix=f"{prefix}.mlp")
+    else:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: use FeedForward instead of MLP
+        '''
+        self.mlp = FeedForward(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            up_proj_name='gate_up_proj',
+            is_gated=True,
+            down_proj_name='down_proj',
+            bias=False,
+            quant_config=quant_config,
+        )
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    self.input_layernorm = RMSNorm(config.hidden_size,
+                                   eps=config.rms_norm_eps)
+    self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                            eps=config.rms_norm_eps)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: prepare to perf per-tensor sq cases if suitable. For moe
+        model, we only do quant fusion in attn block.
+    '''
+    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
+    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        self.self_attn.qkv_proj.quant_method.skip_quant_input = True
+        self.quant_fusion_attn_layernorm = None
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen3moe__Qwen3MoeForCausalLM__load_weights(
+    self,
+    weights: Iterable[Tuple[str, torch.Tensor]]
+):
+    start_expert_id = 0
+    stacked_params_mapping = [
+        # (param_name, shard_name, shard_id)
+        ("qkv_proj", "q_proj", "q"),
+        ("qkv_proj", "k_proj", "k"),
+        ("qkv_proj", "v_proj", "v"),
+        ("gate_up_proj", "gate_proj", 0),
+        ("gate_up_proj", "up_proj", 1),
+    ]
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: delete expert_params_mapping for useless
+    '''
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    params_dict = dict(self.named_parameters())
+    loaded_params: Set[str] = set()
+    for name, loaded_weight in weights:
+        if "rotary_emb.inv_freq" in name:
+            continue
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: replace expert_id in weight to named_expert_id in params_dict
+        '''
+        if start_expert_id > 0 and "mlp.experts." in name:
+            expert_str = re.search(r'experts\.\d+', name).group(0)
+            expert_id = int(expert_str.split(".")[1])
+            named_expert_id = expert_id - start_expert_id
+            old_expert_name = f"experts.{expert_id}"
+            new_expert_name = f"experts.{named_expert_id}"
+            name = name.replace(old_expert_name, new_expert_name)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            if weight_name not in name:
+                continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: delete if "mlp.experts" in name: continue condition
+            '''
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            name = name.replace(weight_name, param_name)
+            # Skip loading extra bias for GPTQ models.
+            if ((name.endswith(".bias") or name.endswith("_bias"))
+                        and name not in params_dict):
+                continue
+            # Skip layers on other devices.
+            if is_pp_missing_parameter(name, self):
+                continue
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
+            '''
+            # Skip experts that are not assigned to this worker.
+            if (("mlp.experts." in name or "mlp.shared_expert." in name or "mlp.shared_expert_gate." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            weight_loader(param, loaded_weight, shard_id)
+            break
+        else:
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: delete for mapping in expert_params_mapping condition
+            '''
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            # Skip loading extra bias for GPTQ models.
+            if ((name.endswith(".bias") or name.endswith("_bias"))
+                            and name not in params_dict):
+                continue
+            # Skip layers on other devices.
+            if is_pp_missing_parameter(name, self):
+                continue
+            # Remapping the name of FP8 kv-scale.
+            if name.endswith("kv_scale"):
+                remapped_kv_scale_name = name.replace(
+                    ".kv_scale", ".attn.kv_scale")
+                if remapped_kv_scale_name not in params_dict:
+                    logger.warning_once(
+                        "Found kv scale in the checkpoint "
+                        f"(e.g. {name}), but not found the expected "
+                        f"name in the model "
+                        f"(e.g. {remapped_kv_scale_name}). "
+                        "kv-scale is not loaded.")
+                    continue
+                else:
+                    name = remapped_kv_scale_name
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add expert skiped condition
+            '''
+            # Skip experts that are not assigned to this worker.
+            if (("mlp.experts." in name or "mlp.shared_expert." in name or "mlp.shared_expert_gate." in name)
+                    and name not in params_dict):
+                continue
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            weight_loader(param, loaded_weight)
+        loaded_params.add(name)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: pack params
+    '''
+    for name, m in self.model.named_modules():
+        if isinstance(m, SparseMoeMlp):
+            m.pack_params()
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return loaded_params
+
+
+def vllm__module_executor__models__qwen3moe__Qwen3MoeDecoderLayer__forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: perf model by:
+    1) add residual in matmul;
+    2) fuse quantization in layernorm in per-tensor sq case;
+    '''
+    attn_layernorm = self.input_layernorm
+    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
+        if self.quant_fusion_attn_layernorm is None:
+            if self.is_per_token_sq_perf_cases:
+                attn_quant_scale = self.self_attn.qkv_proj.smooth
+            else:
+                attn_quant_scale = self.self_attn.qkv_proj.scale_to_int
+            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
+                self.input_layernorm, attn_quant_scale,
+                dynamic_quant=self.is_per_token_sq_perf_cases)
+        attn_layernorm = self.quant_fusion_attn_layernorm
+
+    return decoder_layer_forward_base(
+        positions=positions,
+        hidden_states=hidden_states,
+        input_layernorm=attn_layernorm,
+        self_attn=self.self_attn,
+        post_layernorm=self.post_attention_layernorm,
+        mlp=self.mlp,
+        input_norm_fuse_en=self.is_per_token_sq_perf_cases
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__module_executor__models__qwen3moe__Qwen3MoeModel__forward(
+    self,
+    input_ids: torch.Tensor,
+    positions: torch.Tensor,
+    intermediate_tensors: Optional[IntermediateTensors] = None,
+    inputs_embeds: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    '''
+    return decoder_model_forward_base_pp(
+        input_ids=input_ids,
+        positions=positions,
+        intermediate_tensors=intermediate_tensors,
+        layers=self.layers,
+        start_layer=self.start_layer,
+        end_layer=self.end_layer,
+        get_input_embeddings=self.embed_tokens,
+        norm=self.norm,
+        inputs_embeds=inputs_embeds,
+    )
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+MluHijackObject.apply_hijack(Qwen3MoeAttention,
+                             Qwen3MoeAttention.forward,
+                             vllm__module_executor__models__qwen3moe__Qwen3MoeAttention__forward)
+MluHijackObject.apply_hijack(Qwen3MoeDecoderLayer,
+                             Qwen3MoeDecoderLayer.__init__,
+                             vllm__module_executor__models__qwen3moe__Qwen3MoeDecoderLayer____init__)
+MluHijackObject.apply_hijack(Qwen3MoeForCausalLM,
+                             Qwen3MoeForCausalLM.load_weights,
+                             vllm__module_executor__models__qwen3moe__Qwen3MoeForCausalLM__load_weights)
+MluHijackObject.apply_hijack(Qwen3MoeDecoderLayer,
+                             Qwen3MoeDecoderLayer.forward,
+                             vllm__module_executor__models__qwen3moe__Qwen3MoeDecoderLayer__forward)
+MluHijackObject.apply_hijack(Qwen3MoeModel,
+                             Qwen3MoeModel.forward,
+                             vllm__module_executor__models__qwen3moe__Qwen3MoeModel__forward)
diff --git a/vllm_mlu/vllm_mlu/platforms/__init__.py b/vllm_mlu/vllm_mlu/platforms/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/platforms/mlu.py b/vllm_mlu/vllm_mlu/platforms/mlu.py
new file mode 100644
index 000000000..31938da64
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/platforms/mlu.py
@@ -0,0 +1,175 @@
+from functools import lru_cache
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import os
+import torch
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+from vllm.platforms.interface import (DeviceCapability, Platform,
+                                      PlatformEnum)
+
+from vllm_mlu.model_executor.layers.quantization import (
+    register_fake_mlu_quantization_methods
+)
+
+if TYPE_CHECKING:
+    from vllm.config import ModelConfig, VllmConfig
+    from vllm.utils import FlexibleArgumentParser
+else:
+    ModelConfig = None
+    VllmConfig = None
+    FlexibleArgumentParser = None
+
+from vllm_mlu._mlu_utils import VLLM_LOGITS_USE_ALL_GATHER
+
+logger = init_logger(__name__)
+
+
+envs.environment_variables.update({
+    "MLU_VISIBLE_DEVICES":
+    lambda: os.environ.get("MLU_VISIBLE_DEVICES", None)
+})
+
+
+register_fake_mlu_quantization_methods()
+
+
+class MLUPlatform(Platform):
+    _enum = PlatformEnum.OOT
+    device_name: str = "mlu"
+    device_type: str = "mlu"
+    dispatch_key: str = "MLU"
+    ray_device_key: str = "GPU"
+    device_control_env_var: str = "MLU_VISIBLE_DEVICES"
+    simple_compile_backend: str = "eager"  # Disable torch.compile()
+
+    supported_quantization: list[str] = ["weightonly", "smoothquant",
+                                         "awq_mlu", "gptq_mlu", "fp8"]
+    additional_env_vars: list[str] = ["VLLM_LATENCY_DEBUG",
+                                      "VLLM_LATENCY_DEBUG_NO_DEVICE",
+                                      "MLU_GRAPH_CAPTURE_LIST",
+                                      "DPSK_MCC_PARALLEL_NUM",
+                                      "VLLM_LOGITS_USE_ALL_GATHER"]
+
+    @classmethod
+    def pre_register_and_update(cls,
+                                parser: Optional[FlexibleArgumentParser] = None
+                                ) -> None:
+        pass
+
+    @classmethod
+    def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
+                             kv_cache_dtype, block_size, use_v1, use_mla) -> str:
+        if use_v1:
+            logger.warning(
+                f"AttentionBackend v1 is not supported on MLU now, forcing AttentionBackend to v0.")
+        if use_mla:
+            logger.info(f"Select MLUMLAFlashAttentionBackend.")
+            return "vllm_mlu.attention.backends.mlu_flash_attn.MLUMLAFlashAttentionBackend"
+        logger.info(f"Select MLUFlashAttentionBackend.")
+        return "vllm_mlu.attention.backends.mlu_flash_attn.MLUFlashAttentionBackend"
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
+        major, minor = torch.mlu.get_device_capability(device_id)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return torch.mlu.get_device_name(device_id)
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        device_props = torch.mlu.get_device_properties(device_id)
+        return device_props.total_memory
+
+    @classmethod
+    def set_device(cls, device: torch.device):
+        torch.mlu.set_device(device)
+
+    @classmethod
+    def empty_cache(cls):
+        torch.mlu.empty_cache()
+
+    @classmethod
+    def synchronize(cls):
+        torch.mlu.synchronize()
+
+    @classmethod
+    def mem_get_info(cls) -> Tuple[int, int]:
+        return torch.mlu.mem_get_info()
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
+        return True
+
+    @classmethod
+    def inference_mode(cls):
+        return torch.no_grad()
+
+    @classmethod
+    def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        from vllm.config import CompilationLevel
+        compilation_config = vllm_config.compilation_config
+        if compilation_config.level != CompilationLevel.NO_COMPILATION:
+            logger.warning(
+                "Compilation level %s is not supported on MLU now, forcing compilation level to NO_COMPILATION",
+                compilation_config.level)
+            compilation_config.level = CompilationLevel.NO_COMPILATION
+
+        parallel_config = vllm_config.parallel_config
+        scheduler_config = vllm_config.scheduler_config
+
+        if parallel_config.worker_cls == "auto":
+            if scheduler_config.is_multi_step:
+                if envs.VLLM_USE_V1:
+                    raise NotImplementedError(f"vllm v1 is not supported on mlu platform.")
+                else:
+                    parallel_config.worker_cls = \
+                        "vllm_mlu.worker.multi_step_mlu_worker.MLUMultiStepWorker"
+            elif vllm_config.speculative_config:
+                parallel_config.worker_cls = \
+                        "vllm.spec_decode.spec_decode_worker.create_spec_worker"
+                parallel_config.sd_worker_cls = \
+                        "vllm_mlu.worker.mlu_worker.MLUWorker"
+            else:
+                if envs.VLLM_USE_V1:
+                    raise NotImplementedError(f"vllm v1 is not supported on mlu platform.")
+                else:
+                    parallel_config.worker_cls = "vllm_mlu.worker.mlu_worker.MLUWorker"
+
+        cache_config = vllm_config.cache_config
+        if cache_config and cache_config.block_size is None:
+            cache_config.block_size = 16
+
+    @classmethod
+    def get_current_memory_usage(cls,
+                                 device: Optional[torch.types.Device] = None
+                                 ) -> float:
+        torch.mlu.reset_peak_memory_stats(device)
+        return torch.mlu.max_memory_allocated(device)
+
+    @classmethod
+    def get_punica_wrapper(cls) -> str:
+        return "vllm_mlu.lora.punica_wrapper.punica_mlu.PunicaWrapperMLU"
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "vllm_mlu.distributed.device_communicators.mlu_communicator.MLUCommunicator"
+
+    @classmethod
+    def use_all_gather(cls) -> bool:
+        from vllm.config import get_current_vllm_config
+        parallel_config = get_current_vllm_config().parallel_config
+        return (VLLM_LOGITS_USE_ALL_GATHER or parallel_config.distributed_executor_backend
+                == "external_launcher")
+
+    @classmethod
+    def supports_v1(cls, model_config: ModelConfig) -> bool:
+        """Returns whether the current platform can support v1 for the supplied
+        model configuration.
+        """
+        return False
diff --git a/vllm_mlu/vllm_mlu/spec_decode/__init__.py b/vllm_mlu/vllm_mlu/spec_decode/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/spec_decode/mlu_draft_model_runner.py b/vllm_mlu/vllm_mlu/spec_decode/mlu_draft_model_runner.py
new file mode 100644
index 000000000..565db35b2
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/spec_decode/mlu_draft_model_runner.py
@@ -0,0 +1,398 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List, Optional
+
+import torch
+
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.sampler import SamplerOutput
+from vllm.logger import init_logger
+from vllm.multimodal import MultiModalKwargs
+from vllm.sequence import ExecuteModelRequest, IntermediateTensors
+from vllm.worker.model_runner_base import (ModelRunnerBase,
+                                           ModelRunnerInputBase,
+                                           ModelRunnerWrapperBase)
+
+from vllm_mlu.attention.backends.mlu_flash_attn import MLUFlashAttentionMetadata
+from vllm_mlu._mlu_utils import VLLM_LATENCY_DEBUG_WITH_DEVICE_EN
+
+logger = init_logger(__name__)
+
+# A flag to enable debug prints for the updated input tensors
+# before each step.
+debug_advance_input = False
+# A flag to allow GPU advance step for draft model runner.
+# Set to False for debugging.
+allow_gpu_advance_step = True
+
+
+class MLUTP1DraftModelRunner(ModelRunnerWrapperBase):
+    """Specialized model runner for speculative decoding draft model.
+    Since the draft model always execute k forward passes consecutively to
+    generate k speculative tokens in a single speculative decoding step,
+    we could get rid of most CPU-GPU synchronization and data transfer
+    overheads by keeping model input and output tensors on GPU all the time.
+
+    TODOs:
+    1. Currently supports only flash-attn, add support for other attn_backends.
+    2. Support TP > 1 (this requires some designs because we do not expect
+       any broadcasting inside execute_model).
+    """
+
+    def __init__(self, model_runner: ModelRunnerBase):
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief:
+        1) annotate if hasattr(...)
+        2) add self.time_markers
+        '''
+        # if hasattr(
+        #         model_runner,
+        #         "return_hidden_states") and model_runner.return_hidden_states:
+        #     raise ValueError(
+        #         "return_hidden_states is not supported for MLUTP1DraftModelRunner."
+        #     )
+        super().__init__(model_runner)
+
+        self.time_markers = []
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        self.indices_of_seq_with_bonus_tokens = None
+
+    def _update_sampling_metadata(self, sampling_metadata, num_seqs,
+                                  num_queries):
+
+        assert sampling_metadata.num_prompts == 0
+        assert len(sampling_metadata.seq_groups) == num_queries
+        assert sampling_metadata.selected_token_indices.shape == (
+            num_queries, )
+        # assert sampling_metadata.categorized_sample_indices == TODO: Add if needed # noqa: E501
+
+        # Verify that all sequences are decodes
+        for i in range(num_queries):
+            seq_group = sampling_metadata.seq_groups[i]
+
+            assert seq_group.is_prompt is False  # No prompt
+            assert seq_group.prompt_logprob_indices == []  # No prompt
+            assert seq_group.sample_indices == [i]  # Simple
+
+    def _gpu_advance_step(self, model_input: ModelRunnerInputBase,
+                          last_output: SamplerOutput) -> ModelRunnerInputBase:
+        # Currently, we expect "decode mode" only
+        assert not model_input.is_prompt
+
+        # Get num_seqs
+        num_seqs = len(model_input.seq_lens)
+        num_queries = len(model_input.query_lens)
+
+        # Get output tokens GPU tensor
+        sampled_token_ids = last_output.sampled_token_ids
+        assert sampled_token_ids is not None
+
+        # Update attn_metadata
+        attn_metadata = model_input.attn_metadata
+        assert isinstance(attn_metadata, MLUFlashAttentionMetadata)
+
+        attn_metadata.advance_step(model_input, sampled_token_ids,
+                                   self.block_size, num_seqs, num_queries)
+
+        # Update sampling_metadata
+        sampling_metadata = model_input.sampling_metadata
+        self._update_sampling_metadata(sampling_metadata, num_seqs,
+                                       num_queries)
+
+        # Create new input
+        new_model_input = self._model_input_cls(
+            input_tokens=model_input.input_tokens,
+            input_positions=model_input.input_positions,
+            attn_metadata=attn_metadata,
+            seq_lens=attn_metadata.seq_lens,
+            query_lens=model_input.query_lens,
+            lora_mapping=model_input.lora_mapping,
+            lora_requests=model_input.lora_requests,
+            multi_modal_kwargs=model_input.multi_modal_kwargs,
+            sampling_metadata=model_input.sampling_metadata,
+            is_prompt=False,
+        )
+
+        # Ensure we skip CPU samples
+        assert new_model_input.sampling_metadata.skip_sampler_cpu_output is True
+        # We can reuse sampling tensors since every decode iteration is the same
+        new_model_input.sampling_metadata.reuse_sampling_tensors = True
+
+        if debug_advance_input:
+            logger.debug("NEW INPUT: ")
+            logger.debug("  input_tokens = %s", new_model_input.input_tokens)
+            logger.debug("  input_positions = %s",
+                         new_model_input.input_positions)
+            logger.debug("  seq_lens = %d", new_model_input.seq_lens)
+            logger.debug("  query_lens = %d", new_model_input.query_lens)
+            logger.debug("  attn_metadata:")
+            logger.debug("    seq_lens_tensor: %s",
+                         attn_metadata.seq_lens_tensor)
+            logger.debug("    slot_mapping: %s", attn_metadata.slot_mapping)
+            logger.debug("    block_tables: %s", attn_metadata.block_tables)
+
+        return new_model_input
+
+    def supports_gpu_multi_step(self, execute_model_req: ExecuteModelRequest):
+        """Determines if draft_model_runner GPU multi-step can be used.
+        Currently required conditions are:
+            1. Only decodes
+            2. Only flash-attn
+            3. No LORA
+            4. No prompt_adapter_config
+        """
+        if not allow_gpu_advance_step:
+            return False
+
+        # We allow multi-step GPU only in decode mode
+        for seq_group in execute_model_req.seq_group_metadata_list:
+            if seq_group.is_prompt:
+                return False
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: remove TRITON_MLA
+        '''
+        # TODO: Add support for other attn backends
+        if self.attn_backend.get_name() not in ("FLASH_ATTN", ):
+            return False
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        # TODO: Add support for LORA
+        if self.lora_config:
+            return False
+
+        # TODO: Add soft-tuning prompt adapter support
+        return not self.prompt_adapter_config
+
+    def set_indices_of_seq_with_bonus_tokens(self,
+                                             indices_of_seq_with_bonus_tokens):
+        self.indices_of_seq_with_bonus_tokens = indices_of_seq_with_bonus_tokens
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        model_input: ModelRunnerInputBase,
+        kv_caches: List[torch.Tensor],
+        previous_hidden_states: Optional[torch.Tensor] = None,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        num_steps: int = 1,
+        **kwargs,
+    ) -> Optional[List[SamplerOutput]]:
+        """Executes num_steps forward passes with advacement of input tensors
+        on the GPU. Look at supports_gpu_multi_step(..) for pre-conditions.
+
+        Optimizations used:
+            1. Input tensors are updated on the GPU directly
+            2. Skips GPU=>CPU serialization of sampler outputs (we don't need
+                them since we do batch expansion later that uses GPU outputs)
+            3. Reuses sampling tensors (since we run only decodes and they have
+                a repeating sampling logic)
+        """
+
+        # When num_steps == 1, we execute the fallback here for the GPU
+        # advance_step, which runs prepare_inputs on CPU and for each spec
+        # iteration invokes this function only once
+        # (Look at multi-step-worker code)
+        is_fallback = num_steps == 1
+        if not is_fallback:
+            # Since we do not broadcast data inside execute_model anymore,
+            # we need to figure out the best way to support TP > 1 in this
+            # case, because we will at least need to broadcast the sampled
+            # tokens to all workers.
+            if not self.is_driver_worker:
+                raise ValueError("MLUTP1DraftModelRunner only supports TP=1.")
+
+            # Sanity
+            if self.lora_config is not None:
+                raise ValueError("MLUTP1DraftModelRunner has no support for LORA")
+            if self.prompt_adapter_config is not None:
+                raise ValueError("MLUTP1DraftModelRunner has no support for "
+                                 "prompt_adapter_config")
+            if model_input.multi_modal_kwargs:
+                raise ValueError(
+                    "MLUTP1DraftModelRunner has no support for multi_modal_kwargs"
+                )
+        else:
+            if self.lora_config:
+                assert model_input.lora_requests is not None
+                assert model_input.lora_mapping is not None
+                self.set_active_loras(model_input.lora_requests,
+                                      model_input.lora_mapping)
+
+            if self.prompt_adapter_config:
+                assert model_input.prompt_adapter_requests is not None
+                assert model_input.prompt_adapter_mapping is not None
+                self.set_active_prompt_adapters(
+                    model_input.prompt_adapter_requests,
+                    model_input.prompt_adapter_mapping)
+
+            self.attn_state.begin_forward(model_input)
+
+        # Detect exec mode
+        assert model_input.attn_metadata is not None
+        use_cuda_graph = False
+        if model_input.attn_metadata.num_prefills > 0:
+            # In this case, execute_model(..) was called directly
+            if num_steps > 1:
+                raise ValueError(
+                    "execute_model(..) of draft_model_runner can be called "
+                    "directly only with a single-step prefill")
+        else:
+            # We can skip CPU samples for spec token generation.
+            # (We do allow CPU samples for num_steps == 1 to support the
+            # fallback case, where supports_gpu_multi_step(..) does not pass)
+            model_input.sampling_metadata.skip_sampler_cpu_output = (
+                not is_fallback)
+
+            # Attn attr defines if we use cuda graphs
+            use_cuda_graph = model_input.attn_metadata.use_cuda_graph
+
+        # Get model
+        if use_cuda_graph:
+            graph_batch_size = model_input.input_tokens.shape[0]
+            model_executable = (self.graph_runners[model_input.virtual_engine]
+                                [graph_batch_size])
+
+            if previous_hidden_states is not None:
+                hidden_states = torch.cat([
+                    previous_hidden_states,
+                    torch.empty([
+                        graph_batch_size - previous_hidden_states.shape[0],
+                        *previous_hidden_states.shape[1:]
+                    ],
+                                dtype=previous_hidden_states.dtype,
+                                device=previous_hidden_states.device)
+                ])
+            else:
+                hidden_states = None
+        else:
+            model_executable = self.model
+            hidden_states = previous_hidden_states
+
+        outputs: List[SamplerOutput] = []
+        for step in range(num_steps):
+            multi_modal_kwargs = model_input.multi_modal_kwargs or {}
+
+            model_execute_kwargs = {"previous_hidden_states": hidden_states} \
+                if previous_hidden_states is not None else {}
+
+            compute_logits_kwargs = {}
+            # Run model
+            if hasattr(self.model.config, "num_nextn_predict_layers"):
+                # for DeepSeek MTP only to use the corresponding layer for
+                # each step
+                spec_step_idx = kwargs.get("spec_step_idx", step)
+                model_execute_kwargs["spec_step_idx"] = spec_step_idx
+                compute_logits_kwargs["spec_step_idx"] = spec_step_idx
+
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add mlu metrics
+            '''
+            # Add time markers for model_executable+compute_logits
+            if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+                self.time_markers = []
+                prefill_meta = model_input.attn_metadata.prefill_metadata
+                decode_meta = model_input.attn_metadata.decode_metadata
+                use_cuda_graph = ((prefill_meta is None and decode_meta.use_cuda_graph))
+                # if use_cuda_graph, the start timestamp will be inserted inside MLUGraphRunner.forward()
+                if not use_cuda_graph:
+                    start = torch.mlu.Event(enable_timing=True)
+                    start.record()
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            with set_forward_context(model_input.attn_metadata,
+                                     self.vllm_config):
+                hidden_states = model_executable(
+                    input_ids=model_input.input_tokens,
+                    positions=model_input.input_positions,
+                    intermediate_tensors=intermediate_tensors,
+                    **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
+                                                 device=self.device),
+                    **model_execute_kwargs,
+                )
+
+            # Compute the logits.
+            logits = self.model.compute_logits(hidden_states,
+                                               model_input.sampling_metadata,
+                                               **compute_logits_kwargs)
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: add mlu metrics
+            '''
+            # Add time markers for model_executable+compute_logits
+            if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+                end_marker = torch.mlu.Event(enable_timing=True)
+                end_marker.record()
+                if use_cuda_graph:
+                    self.time_markers.append((model_executable.start, end_marker))
+                else:
+                    self.time_markers.append((start, end_marker))
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+
+            if not self.is_driver_worker:
+                return []
+            # Sample the next token.
+            output = self.model.sample(
+                logits=logits,
+                sampling_metadata=model_input.sampling_metadata,
+            )
+            outputs.append(output)
+
+            if self.return_hidden_states and is_fallback:
+                if use_cuda_graph:
+                    indices = model_input.sampling_metadata\
+                      .selected_token_indices
+                    output.hidden_states = hidden_states[:len(indices)]
+                else:
+                    output.hidden_states = hidden_states
+
+            if model_input.attn_metadata.num_prefills == 0 \
+                and self.indices_of_seq_with_bonus_tokens is not None:
+                assert output.sampled_token_ids is not None
+                # output.sampled_token_ids should be of shape (num_seqs, 1)
+                nums_seqs, num_tokens_per_seq = output.sampled_token_ids.shape
+                assert num_tokens_per_seq == 1
+                count = 0
+                for i in range(nums_seqs):
+                    bonus_seq_idx = self.indices_of_seq_with_bonus_tokens[
+                        count]
+                    if i != bonus_seq_idx:
+                        # The following might cause a cpu->gpu sync
+                        # However, the performance impact is negligible as we
+                        # benchmarked on H100.
+                        output.sampled_token_ids[
+                            i, :] = model_input.input_tokens[bonus_seq_idx]
+                    else:
+                        count += 1
+
+            # Prepare inputs for the next step
+            if step != num_steps - 1:
+                model_input = self._gpu_advance_step(model_input, outputs[-1])
+
+        return outputs
diff --git a/vllm_mlu/vllm_mlu/spec_decode/mlu_metrics.py b/vllm_mlu/vllm_mlu/spec_decode/mlu_metrics.py
new file mode 100644
index 000000000..9cb76176c
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/spec_decode/mlu_metrics.py
@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import time
+from typing import Optional, Union
+
+import torch
+
+from vllm.model_executor.layers.spec_decode_base_sampler import (
+    SpecDecodeBaseSampler)
+from vllm.utils import is_pin_memory_available
+
+from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics, Timer
+
+
+class MLUAsyncMetricsCollector:
+    """Class which copies rejection/typical-acceptance sampler metrics
+    from the device to CPU on a non-default Torch stream.
+    """
+
+    def __init__(self,
+                 spec_decode_sampler: SpecDecodeBaseSampler,
+                 timer: Optional[Timer] = None,
+                 collect_interval_s: float = 5.0):
+        self.spec_decode_sampler = spec_decode_sampler
+        self._timer = time.time if timer is None else timer
+
+        self._rank: Optional[int] = None
+
+        # We don't have a device set yet.
+        self._copy_stream: Optional[torch.mlu.Stream] = None
+
+        self._in_flight_copy: Optional[torch.mlu.Event] = None
+
+        pin_memory = is_pin_memory_available()
+        self._aggregate_num_accepted_tokens = torch.tensor(
+            0, dtype=torch.long, device="cpu", pin_memory=pin_memory)
+        self._aggregate_num_emitted_tokens = torch.tensor(
+            0, dtype=torch.long, device="cpu", pin_memory=pin_memory)
+        self._aggregate_num_draft_tokens = 0
+
+        self._rejsample_metrics_collect_interval_s = collect_interval_s
+        self._last_metrics_collect_time = self._timer()
+
+    def init_gpu_tensors(self, rank: int) -> None:
+        self._rank = rank
+        self._copy_stream = torch.mlu.Stream()
+
+    def init_tensors(self,
+                     rank: int,
+                     device_type: Union[torch.device, str] = 'mlu') -> None:
+        self._rank = rank
+        if isinstance(device_type, torch.device):
+            device_type = device_type.type
+        if device_type == 'mlu':
+            self._copy_stream = torch.mlu.Stream()
+
+    def maybe_collect_rejsample_metrics(
+            self, k: int) -> Optional[SpecDecodeWorkerMetrics]:
+        # currently using mlu.Event, skip for any non_mlu_alike platform
+        from vllm.platforms import current_platform
+        if not current_platform.is_out_of_tree():
+            return None
+
+        # If a copy was initiated in the previous call, collect and return.
+        if self._in_flight_copy is not None:
+            ready_event = self._in_flight_copy
+            self._in_flight_copy = None
+            return self._collect_rejsample_metrics(k, ready_event)
+
+        # Otherwise, check if we should start a new copy.
+        if self._should_collect_rejsample_metrics(self._timer()):
+            assert self._in_flight_copy is None
+            self._in_flight_copy = self._copy_rejsample_metrics_async()
+
+        return None
+
+    def _should_collect_rejsample_metrics(self, now: float) -> bool:
+        """Return whether or not this iteration should print sampling
+        metrics.
+        """
+        if self._rank != 0:
+            return False
+
+        return now - self._last_metrics_collect_time >= self._rejsample_metrics_collect_interval_s  # noqa: E501
+
+    def _copy_rejsample_metrics_async(self) -> torch.mlu.Event:
+        """Copy rejection/typical-acceptance sampling metrics
+        (number of accepted tokens, etc) to CPU asynchronously.
+
+        Returns a mlu event recording when the copy is complete.
+        """
+        assert self._copy_stream is not None
+        self._copy_stream.wait_stream(torch.mlu.current_stream())
+
+        with torch.mlu.stream(self._copy_stream):
+            self._aggregate_num_accepted_tokens.copy_(
+                self.spec_decode_sampler.num_accepted_tokens,
+                non_blocking=True)
+            self._aggregate_num_emitted_tokens.copy_(
+                self.spec_decode_sampler.num_emitted_tokens, non_blocking=True)
+            # Number of draft tokens is calculated on CPU, so no copy is
+            # required.
+            self._aggregate_num_draft_tokens = (
+                self.spec_decode_sampler.num_draft_tokens)
+
+        aggregate_metrics_ready = torch.mlu.Event()
+        aggregate_metrics_ready.record(self._copy_stream)
+
+        return aggregate_metrics_ready
+
+    def _collect_rejsample_metrics(
+            self, k: int,
+            ready_event: torch.mlu.Event) -> SpecDecodeWorkerMetrics:
+        """Create metrics object from statistics copied asynchronously.
+
+        Args:
+            k: int. The number of speculative tokens; used to determine system
+                efficiency.
+            ready_event: torch.mlu.Event. The mlu event recording when the
+                async GPU->CPU copy is complete.
+        """
+
+        ready_event.synchronize()
+
+        # update time of last collection
+        self._last_metrics_collect_time = self._timer()
+
+        accepted_tokens = self._aggregate_num_accepted_tokens.item()
+        emitted_tokens = self._aggregate_num_emitted_tokens.item()
+        draft_tokens = self._aggregate_num_draft_tokens
+
+        max_num_emitted_tokens = self.get_max_num_emitted_tokens(
+            draft_tokens, k)
+
+        if draft_tokens > 0:
+            draft_acceptance_rate = accepted_tokens / draft_tokens
+        else:
+            draft_acceptance_rate = float("nan")
+
+        if max_num_emitted_tokens > 0:
+            system_efficiency = emitted_tokens / max_num_emitted_tokens
+        else:
+            system_efficiency = float("nan")
+
+        return SpecDecodeWorkerMetrics(
+            num_spec_tokens=k,
+            draft_acceptance_rate=draft_acceptance_rate,
+            system_efficiency=system_efficiency,
+            accepted_tokens=accepted_tokens,
+            draft_tokens=draft_tokens,
+            emitted_tokens=emitted_tokens,
+        )
+
+    @staticmethod
+    def get_max_num_emitted_tokens(draft_tokens: int, k: int) -> int:
+        """Calculate the number of emitted tokens, assuming all tokens are
+        accepted.
+
+        This is equal to the number of sequences that have been speculated on,
+        times (speculation len + 1). The +1 comes from the bonus token.
+        """
+        # Determine the number of sequences that have been speculated on. Since
+        # the batch size can be variable, we divide by k.
+        assert draft_tokens % k == 0
+        total_num_spec_seqs = draft_tokens // k
+
+        # A single sequence may emit k accepted tokens and one bonus token in
+        # the best case.
+        num_emitted_per_seq_if_all_accepted = k + 1
+
+        # The max num of emitted tokens is the number of speculated sequences
+        # times the max emitted per seq.
+        return total_num_spec_seqs * num_emitted_per_seq_if_all_accepted
diff --git a/vllm_mlu/vllm_mlu/spec_decode/multi_step_worker.py b/vllm_mlu/vllm_mlu/spec_decode/multi_step_worker.py
new file mode 100644
index 000000000..5f2997010
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/spec_decode/multi_step_worker.py
@@ -0,0 +1,82 @@
+from typing import List, Set, Tuple
+
+import torch
+
+from vllm.model_executor.layers.sampler import SamplerOutput
+from vllm.platforms import current_platform
+from vllm.sequence import ExecuteModelRequest
+from vllm.spec_decode.multi_step_worker import MultiStepWorker
+
+from vllm_mlu.spec_decode.mlu_draft_model_runner import MLUTP1DraftModelRunner
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+@torch.inference_mode()
+def vllm__sepc_decode__multi_step_worker__MultiStepWorker__sampler_output(
+    self,
+    execute_model_req: ExecuteModelRequest,
+    sample_len: int,
+    seq_ids_with_bonus_token_in_last_step: Set[int],
+) -> Tuple[List[SamplerOutput], bool]:
+    """Run the model forward pass sample_len times. Returns the list of
+    sampler output, one per model forward pass, along with indicator of
+    whether torch tensor in sampler output need to be transposed in latter
+    sampler_output_to_torch logic.
+
+    For multi step worker, this indicator shall be True.
+    """
+    self._raise_if_unsupported(execute_model_req)
+    # Expand the batch for sequences with a bonus token.
+    # Perform a forward pass on the expanded batch and filter the
+    # response to retain only the original sequences' responses.
+    expanded_request, indices_of_seq_with_bonus_tokens =\
+        self._expand_execute_model_request(
+            execute_model_req, seq_ids_with_bonus_token_in_last_step)
+
+    # Run model sample_len times.
+    model_outputs: List[SamplerOutput] = []
+    if current_platform.is_out_of_tree() and isinstance(
+            self.model_runner, MLUTP1DraftModelRunner
+    ) and self.model_runner.supports_gpu_multi_step(expanded_request):
+        # Here we run the draft_model_runner with multi-step prepare
+        # on the GPU directly
+        expanded_request.num_steps = sample_len
+        self.model_runner.set_indices_of_seq_with_bonus_tokens(
+            indices_of_seq_with_bonus_tokens)
+        model_outputs = self.execute_model(
+            execute_model_req=expanded_request)
+    else:
+        # Here we run multi-step directly, with every step prepared
+        # on the CPU.
+        # TODO: Remove this branch once DraftModelRunner supports TP>1
+        # and other restrictions that are part of DraftModelRunner's
+        # supports_gpu_multi_step(..)
+        if expanded_request.previous_hidden_states is not None:
+            self.worker.model_runner.return_hidden_states = True
+        for _ in range(sample_len):
+            model_output: List[SamplerOutput] = self.worker.execute_model(
+                execute_model_req=expanded_request)
+            assert (len(model_output) == 1
+                    ), "composing multistep workers not supported"
+            model_output = model_output[0]
+            self._maybe_update_previous_hidden_states(
+                model_output, expanded_request)
+
+            self._append_new_tokens(
+                model_output, expanded_request.seq_group_metadata_list,
+                indices_of_seq_with_bonus_tokens)
+            model_outputs.append(model_output)
+
+    # move indices to device to avoid stream sync
+    indices_of_seq_with_bonus_tokens = torch.tensor(
+        indices_of_seq_with_bonus_tokens, device=self.device)
+    filtered_model_outputs = self._filter_model_output(
+        model_outputs, indices_of_seq_with_bonus_tokens)
+    return filtered_model_outputs, True
+
+
+MluHijackObject.apply_hijack(
+    MultiStepWorker,
+    MultiStepWorker.sampler_output,
+    vllm__sepc_decode__multi_step_worker__MultiStepWorker__sampler_output
+)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/spec_decode/ngram_worker.py b/vllm_mlu/vllm_mlu/spec_decode/ngram_worker.py
new file mode 100644
index 000000000..f3981aab9
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/spec_decode/ngram_worker.py
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+from vllm.config import VllmConfig
+from vllm.spec_decode.ngram_worker import NGramWorker
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+
+NGramWorker__init__org = NGramWorker.__init__
+
+
+class NGramWorker_MluHijack(NGramWorker):
+    """NGramWorker provides a light drafter without need for model.
+
+    Current NGramWorker only implements prompt lookup decoding,
+    and in future we may also do RAG type drafter and other scenarios
+    which don't rely on LLM model to give proposals.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        local_rank: int,
+        device_type: str = "mlu",
+        **kwargs,
+    ):
+        NGramWorker__init__org(self, vllm_config,
+                               local_rank, device_type, **kwargs)
+
+
+MluHijackObject.apply_hijack(NGramWorker,
+                             NGramWorker.__init__,
+                             NGramWorker_MluHijack.__init__)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/spec_decode/spec_decode_worker.py b/vllm_mlu/vllm_mlu/spec_decode/spec_decode_worker.py
new file mode 100644
index 000000000..3a14c9bba
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/spec_decode/spec_decode_worker.py
@@ -0,0 +1,397 @@
+from collections import defaultdict
+from typing import Any, Dict, Optional, Set, Tuple
+import torch
+
+from vllm.logger import init_logger
+from vllm.config import ParallelConfig
+from vllm.model_executor.layers.rejection_sampler import RejectionSampler
+from vllm.model_executor.layers.spec_decode_base_sampler import (
+    SpecDecodeBaseSampler)
+from vllm.model_executor.layers.typical_acceptance_sampler import (
+    TypicalAcceptanceSampler)
+from vllm.platforms import current_platform
+from vllm.sequence import HiddenStates
+from vllm.spec_decode.interfaces import SpeculativeScorer
+from vllm.spec_decode.medusa_worker import MedusaWorker
+from vllm.spec_decode.mlp_speculator_worker import MLPSpeculatorWorker
+from vllm.spec_decode.multi_step_worker import MultiStepWorker
+from vllm.spec_decode.ngram_worker import NGramWorker
+from vllm.spec_decode.proposer_worker_base import ProposerWorkerBase
+from vllm.spec_decode.smaller_tp_proposer_worker import SmallerTpProposerWorker
+from vllm.worker.cache_engine import CacheEngine
+from vllm.worker.worker_base import WorkerBase
+from vllm.spec_decode.spec_decode_worker import SpecDecodeWorker, split_num_cache_blocks_evenly
+
+from vllm_mlu.spec_decode.mlu_metrics import MLUAsyncMetricsCollector
+from vllm_mlu.spec_decode.mlu_draft_model_runner import MLUTP1DraftModelRunner
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+from vllm_mlu.worker.mlu_worker import MLUWorker
+
+logger = init_logger(__name__)
+
+
+@classmethod
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__create_worker(
+    cls,
+    scorer_worker: WorkerBase,
+    draft_worker_kwargs: Dict[str, Any],
+    disable_mqa_scorer: bool,
+    disable_by_batch_size: Optional[int],
+    draft_token_acceptance_method: str,
+    typical_acceptance_sampler_posterior_threshold: float,
+    typical_acceptance_sampler_posterior_alpha: float,
+    disable_logprobs: bool,
+    disable_log_stats: bool,
+    num_speculative_tokens: int,
+) -> "SpecDecodeWorker":
+
+    allow_zero_draft_token_step = True
+    enable_lm_head_weight_load = False
+    num_spec_prefill_steps = 1
+    ngram_prompt_lookup_max = (
+        draft_worker_kwargs.pop("ngram_prompt_lookup_max"))
+    ngram_prompt_lookup_min = (
+        draft_worker_kwargs.pop("ngram_prompt_lookup_min"))
+    draft_model_config = draft_worker_kwargs["vllm_config"].model_config
+    draft_parallel_config: ParallelConfig = draft_worker_kwargs[
+        'vllm_config'].parallel_config
+    if ngram_prompt_lookup_max > 0:
+        draft_worker_kwargs[
+            "device_type"] = scorer_worker.device_config.device.type
+        proposer_worker = NGramWorker(**draft_worker_kwargs)
+        proposer_worker.set_ngram_window_size(ngram_prompt_lookup_min,
+                                                ngram_prompt_lookup_max)
+    else:
+        draft_tp = draft_parallel_config.tensor_parallel_size
+        target_tp = scorer_worker.parallel_config.tensor_parallel_size
+
+        if draft_model_config.hf_config.model_type == "mlp_speculator":
+            proposer_worker = MLPSpeculatorWorker(**draft_worker_kwargs)
+        elif draft_model_config.hf_config.model_type == "medusa":
+            proposer_worker = MedusaWorker(**draft_worker_kwargs)
+        else:
+            if draft_tp == 1:
+                if current_platform.is_out_of_tree():
+                    draft_worker_kwargs[
+                        "model_runner_cls"] = MLUTP1DraftModelRunner
+            else:
+                if draft_model_config.hf_config.model_type == "eagle":
+                    raise NotImplementedError(
+                        f"{draft_model_config.hf_config.model_type} "
+                        "does not support TP > 1 yet")
+
+                allow_zero_draft_token_step = False
+
+            # Load lm_head weight for eagle in init_device
+            if draft_model_config.hf_config.model_type == "eagle":
+                enable_lm_head_weight_load = True
+
+            proposer_worker = MultiStepWorker(**draft_worker_kwargs)
+            if draft_model_config.hf_config.model_type == "deepseek_mtp":
+                num_spec_prefill_steps = \
+                    draft_model_config.hf_config.n_predict
+
+        proposer_worker = SmallerTpProposerWorker.maybe_wrap_worker(
+            proposer_worker, draft_tp, target_tp)
+
+    logger.info("Configuring SpecDecodeWorker with proposer=%s",
+                type(proposer_worker))
+
+    spec_decode_sampler: SpecDecodeBaseSampler = None
+    if draft_token_acceptance_method == "rejection_sampler":
+        spec_decode_sampler = RejectionSampler()
+    elif draft_token_acceptance_method == "typical_acceptance_sampler":
+        spec_decode_sampler = TypicalAcceptanceSampler(
+            posterior_threshold=\
+                typical_acceptance_sampler_posterior_threshold,
+            posterior_alpha=typical_acceptance_sampler_posterior_alpha,
+        )
+    logger.info(
+        "[Speculative Decoding] Configuring"
+        " SpecDecodeWorker with sampler=%s", type(spec_decode_sampler))
+
+    if not disable_mqa_scorer:
+        if scorer_worker.model_runner.attn_backend.get_name(
+        ) != "FLASH_ATTN":
+            disable_mqa_scorer = True
+            logger.info(
+                "[Speculative Decoding] Disabling MQA scorer as the "
+                "MQA is only available with flash attn backend.")
+
+        if draft_model_config and \
+            draft_model_config.max_model_len < \
+                scorer_worker.model_config.max_model_len:
+            disable_mqa_scorer = True
+            logger.info(
+                "[Speculative Decoding] Disabling MQA scorer as the "
+                "draft model max_model_len is smaller than the target "
+                "model max_model_len.")
+
+        if not scorer_worker.model_runner.model_config.enforce_eager:
+            disable_mqa_scorer = True
+            logger.info(
+                "[Speculative Decoding] Disabling MQA scorer as the "
+                "target model is not running in eager mode.")
+
+    return SpecDecodeWorker(
+        proposer_worker,
+        scorer_worker,
+        disable_mqa_scorer=disable_mqa_scorer,
+        disable_logprobs=disable_logprobs,
+        disable_log_stats=disable_log_stats,
+        disable_by_batch_size=disable_by_batch_size,
+        spec_decode_sampler=spec_decode_sampler,
+        allow_zero_draft_token_step=allow_zero_draft_token_step,
+        enable_lm_head_weight_load=enable_lm_head_weight_load,
+        num_spec_prefill_steps=num_spec_prefill_steps)
+
+
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker____init__(
+    self,
+    proposer_worker: ProposerWorkerBase,
+    scorer_worker: WorkerBase,
+    spec_decode_sampler: SpecDecodeBaseSampler,
+    disable_mqa_scorer: bool = False,
+    disable_logprobs: bool = False,
+    disable_log_stats: bool = False,
+    metrics_collector: Optional[MLUAsyncMetricsCollector] = None,
+    disable_by_batch_size: Optional[int] = None,
+    allow_zero_draft_token_step: Optional[bool] = True,
+    enable_lm_head_weight_load: Optional[bool] = False,
+    num_spec_prefill_steps: int = 1,
+):
+    """
+    Create a SpecDecodeWorker.
+
+    Args:
+        proposer_worker: A worker that can produce speculative tokens for
+            sequences.
+        scorer_worker: A worker that produces probabilities of speculative
+            tokens according to some base model. Typically a vanilla vLLM
+            Worker.
+        spec_decode_sampler: A Torch module used to perform acceptance
+            sampling of the draft tokens in the verification step of
+            speculative decoding. Currently we support two different 
+            types of sampler namely RejectionSampler and
+            TypicalAcceptanceSampler. 'spec_decode_sampler' is either an
+            instance of RejectionSampler or TypicalAcceptanceSampler.
+        disable_mqa_scorer: If set to True, disable the MQA scorer and use
+            the BatchExpansionTop1Scorer instead.
+        disable_logprobs: If set to True, token log probabilities will
+            not be output in both the draft worker and the target worker.
+            If set to False, log probabilities will be output by both.
+        disable_log_stats: If set to True, disable periodic printing of
+            speculative stage times.
+        disable_by_batch_size: If the batch size is larger than this,
+            disable speculative decoding for new incoming requests.
+        metrics_collector: Helper class for collecting metrics; can be set
+            for testing purposes.
+        allow_zero_draft_token_step: whether to allow a step where the draft
+            model generates no draft token; should disallow when the tp of
+            draft model is larger than 1 (TODO: #5814)
+        enable_lm_head_weight_load: whether to load lm_head weight for
+            draft models like eagle.
+        num_spec_prefill_steps: number of speculative prefill steps to run
+            before the speculative decoding starts. This is only used when
+            the draft model is a deepseek_mtp model that requires prefill
+            kv cache separately for each MTP layer.
+    """
+    self.proposer_worker = proposer_worker
+    self.scorer_worker = scorer_worker
+    scorer_runner = getattr(self.scorer_worker, "model_runner", None)
+    self.generators = scorer_runner.get_generators(
+    ) if scorer_runner else None
+    self.disable_by_batch_size = disable_by_batch_size or float("inf")
+    self.spec_decode_sampler = spec_decode_sampler
+    self._allow_zero_draft_token_step = allow_zero_draft_token_step
+    self._enable_lm_head_weight_load = enable_lm_head_weight_load
+    self._metrics = MLUAsyncMetricsCollector(
+        self.spec_decode_sampler
+    ) if metrics_collector is None else metrics_collector
+    # Tracks the sequence IDs that received a bonus token ID in
+    # their last forward pass. Needed only if KV cache is being
+    # used for token generation such as in the case of MultiStepWorker.
+    self._seq_with_bonus_token_in_last_step: Set[int] = set()
+    # Tracks the currently active request ids and the sequence IDs
+    # corresponding to them
+    self._request_id_seq_id_mapping: Dict[str, Set[int]] = defaultdict(set)
+    # Tracks if the proposer worker uses the KV cache or not.
+
+    self.probs_dtype = self.spec_decode_sampler.probs_dtype
+    self.token_id_dtype = self.spec_decode_sampler.token_id_dtype
+    # Lazy initialization.
+    self.scorer: SpeculativeScorer
+    self.disable_mqa_scorer = disable_mqa_scorer
+
+    # Hidden states from target model to pass to proposer
+    # in the subsequent step.
+    self.previous_hidden_states: Optional[HiddenStates] = None
+    self._disable_logprobs = disable_logprobs
+    self._disable_log_stats = disable_log_stats
+    self._num_spec_prefill_steps = num_spec_prefill_steps
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: asign self.proposer_worker.woker with _worker
+    '''
+    if hasattr(self.proposer_worker, "_worker") and not hasattr(self.proposer_worker, "worker"):
+        self.proposer_worker.worker = self.proposer_worker._worker
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+@torch.inference_mode()
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__determine_num_available_blocks(self) -> Tuple[int, int]:
+    """Determine the number of cache blocks to use.
+
+    This is done by profiling the scorer model (which is typically the
+    larger of the two). Then the total memory which would be used by the
+    scorer cache is divided evenly between the proposer and scorer model KV,
+    such that the number of blocks is equal in both KV caches.
+    """
+    num_gpu_blocks, num_cpu_blocks = (
+        self.scorer_worker.determine_num_available_blocks())
+
+    scorer_cache_block_size_bytes = (
+        self.scorer_worker.get_cache_block_size_bytes())
+    proposer_cache_block_size_bytes = (
+        self.proposer_worker.get_cache_block_size_bytes())
+
+    new_num_gpu_blocks = split_num_cache_blocks_evenly(
+        scorer_cache_block_size_bytes, proposer_cache_block_size_bytes,
+        num_gpu_blocks)
+
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: Fix: tmo.reshape_paged_cache limit: RuntimeError: The addressing range of kv_cache cannot exceed 4G
+    @brief: Fix: cnnlScaledDotProductAttn_v5 check failed: cnnlGetTensorElementNum(key_desc) * kv_cache_dbyte <= INT32_MAX
+    @brief: record init memory usage
+    '''
+    if hasattr(self.proposer_worker, "worker") and hasattr(self.proposer_worker.worker.model_config.hf_text_config,
+                                                           "num_attention_heads"):
+        max_num_gpu_blocks = CacheEngine.get_max_num_gpu_blocks(self.proposer_worker.worker.cache_config,
+                                                                self.proposer_worker.worker.model_config,
+                                                                self.proposer_worker.worker.parallel_config)
+        if new_num_gpu_blocks > max_num_gpu_blocks:
+            logger.warning(f"current cache block num {new_num_gpu_blocks} is greater than tmo op limit, "
+                           f"force reduce cache block num to {max_num_gpu_blocks}.")
+        new_num_gpu_blocks = min(new_num_gpu_blocks, max_num_gpu_blocks)
+    # Record memory usage
+    if isinstance(self.scorer_worker, MLUWorker):
+        self.peak_memory = self.scorer_worker.peak_memory
+        self.block_memory = self.scorer_worker.block_memory
+    else:
+        self.peak_memory = 1
+        self.block_memory = 1
+    self.num_gpu_blocks = new_num_gpu_blocks
+    self.num_cpu_blocks = num_cpu_blocks
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+    return new_num_gpu_blocks, num_cpu_blocks
+
+
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__get_latency(self):
+    scorer_latency = self.scorer_worker.get_latency()
+    proposer_latency = 0
+    if not isinstance(self.proposer_worker, NGramWorker):
+        proposer_latency = self.proposer_worker.worker.get_latency()
+
+    return scorer_latency + proposer_latency
+
+
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__get_memory_usage(self):
+    return (self.peak_memory, self.block_memory, self.num_gpu_blocks, self.num_cpu_blocks)
+
+
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__recapture_model(
+    self,
+    enable_context_mlugraph,
+    context_batch_size_to_capture,
+    context_seq_len_to_capture
+) -> None:
+    self.scorer_worker.recapture_model(enable_context_mlugraph, context_batch_size_to_capture,
+                                       context_seq_len_to_capture)
+    if not isinstance(self.proposer_worker, NGramWorker):
+        self.proposer_worker.worker.recapture_model(enable_context_mlugraph, context_batch_size_to_capture,
+                                                    context_seq_len_to_capture)
+
+
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__setup_smooth_hook(self, is_save_input_id: bool = False):
+    self.scorer_worker.setup_smooth_hook(is_save_input_id)
+    if not isinstance(self.proposer_worker, NGramWorker):
+        self.proposer_worker.worker.setup_smooth_hook(is_save_input_id)
+
+
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__remove_hooks(self):
+    self.scorer_worker.remove_hooks()
+    if not isinstance(self.proposer_worker, NGramWorker):
+        self.proposer_worker.worker.remove_hooks()
+
+
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__get_act_range(self):
+    scorer_act_range = self.scorer_worker.get_act_range()
+    proposer_act_range = {}
+    if not isinstance(self.proposer_worker, NGramWorker):
+        proposer_act_range = self.proposer_worker.worker.get_act_range()
+
+    act_range = scorer_act_range | proposer_act_range
+
+    return act_range
+
+
+def vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__get_named_parameters(self):
+    scorer_named_parameters = self.scorer_worker.get_named_parameters()
+    proposer_named_parameters = {}
+    if not isinstance(self.proposer_worker, NGramWorker):
+        proposer_named_parameters = self.proposer_worker.worker.get_named_parameters()
+
+    named_parameters = scorer_named_parameters | proposer_named_parameters
+
+    return named_parameters
+
+
+MluHijackObject.apply_hijack(
+    SpecDecodeWorker,
+    SpecDecodeWorker.create_worker,
+    vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__create_worker
+)
+MluHijackObject.apply_hijack(
+    SpecDecodeWorker,
+    SpecDecodeWorker.__init__,
+    vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker____init__
+)
+MluHijackObject.apply_hijack(SpecDecodeWorker,
+                             SpecDecodeWorker.determine_num_available_blocks,
+                             vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__determine_num_available_blocks)
+MluHijackObject.apply_hijack(SpecDecodeWorker,
+                             "get_latency",
+                             vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__get_latency)
+MluHijackObject.apply_hijack(SpecDecodeWorker,
+                             "get_memory_usage",
+                             vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__get_memory_usage)
+MluHijackObject.apply_hijack(SpecDecodeWorker,
+                             "recapture_model",
+                             vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__recapture_model)
+MluHijackObject.apply_hijack(SpecDecodeWorker,
+                             "setup_smooth_hook",
+                             vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__setup_smooth_hook)
+MluHijackObject.apply_hijack(SpecDecodeWorker,
+                             "remove_hooks",
+                             vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__remove_hooks)
+MluHijackObject.apply_hijack(SpecDecodeWorker,
+                             "get_act_range",
+                             vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__get_act_range)
+MluHijackObject.apply_hijack(SpecDecodeWorker,
+                             "get_named_parameters",
+                             vllm__sepc_decode__spec_decode_worker__SpecDecodeWorker__get_named_parameters)
diff --git a/vllm_mlu/vllm_mlu/transformers_utils/__init__.py b/vllm_mlu/vllm_mlu/transformers_utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/transformers_utils/configs/__init__.py b/vllm_mlu/vllm_mlu/transformers_utils/configs/__init__.py
new file mode 100644
index 000000000..78db7dd3b
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/transformers_utils/configs/__init__.py
@@ -0,0 +1,5 @@
+from vllm_mlu.transformers_utils.configs.custom import CustomConfig
+
+__all__ = [
+    "CustomConfig"
+]
diff --git a/vllm_mlu/vllm_mlu/transformers_utils/configs/custom.py b/vllm_mlu/vllm_mlu/transformers_utils/configs/custom.py
new file mode 100644
index 000000000..3977eb503
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/transformers_utils/configs/custom.py
@@ -0,0 +1,60 @@
+from transformers import PretrainedConfig
+
+class CustomConfig(PretrainedConfig):
+    
+    model_type = "custom"
+
+    def __init__(self,
+                 max_sequence_length=None,
+                 num_hidden_layers=None,
+                 hidden_size=None,
+                 use_parallel_embedding=False,
+                 vocab_size=None,
+                 position_embedding_type="ROPE",
+                 is_neox_style=True,
+                 num_attention_heads=None,
+                 num_key_value_heads=None,
+                 attention_bias=False,
+                 intermediate_size=None,
+                 hidden_act="silu",
+                 is_gated=False,
+                 num_experts=None,
+                 num_experts_per_tok=None,
+                 moe_intermediate_size=None,
+                 shared_expert_intermediate_size=None,
+                 norm_topk_prob=None,
+                 mlp_bias=False,
+                 norm_type="rmsnorm",
+                 norm_eps=1e-05, 
+                 apply_residual_connection_post_layernorm=False,
+                 use_parallel_residual=False,
+                 **kwargs):
+        self.max_sequence_length = max_sequence_length
+        self.num_hidden_layers = num_hidden_layers
+        self.hidden_size = hidden_size
+        self.use_parallel_embedding = use_parallel_embedding
+        self.vocab_size = vocab_size
+        self.position_embedding_type = position_embedding_type  # ALIBI, ROPE
+        self.is_neox_style = is_neox_style  # True: fold_rotary; False: cross_rotary
+        self.num_attention_heads = num_attention_heads
+        if num_key_value_heads is None:
+            self.num_key_value_heads = num_attention_heads
+        else:
+            self.num_key_value_heads = num_key_value_heads
+        self.attention_bias = attention_bias
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act  # silu, gelu
+        self.is_gated = is_gated
+        self.num_experts = num_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.moe_intermediate_size = moe_intermediate_size
+        self.shared_expert_intermediate_size = shared_expert_intermediate_size
+        self.norm_topk_prob = norm_topk_prob
+        self.mlp_bias = mlp_bias
+        self.norm_type = norm_type  # rmsnorm， layernorm
+        self.norm_eps = norm_eps
+        self.apply_residual_connection_post_layernorm = (
+            apply_residual_connection_post_layernorm)
+        self.use_parallel_residual = use_parallel_residual
+        
+        super().__init__(**kwargs)
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/utils.py b/vllm_mlu/vllm_mlu/utils.py
new file mode 100644
index 000000000..09b425e4a
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/utils.py
@@ -0,0 +1,186 @@
+import os
+import torch
+from torch.library import Library
+from functools import lru_cache
+from typing import Optional, Union, Callable, List
+
+import vllm.envs as envs
+from vllm import utils
+from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, is_in_doc_build,
+                        supports_custom_op, vllm_lib,
+                        cuda_is_initialized, is_in_ray_actor)
+from vllm.logger import init_logger
+
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+STR_DTYPE_TO_TORCH_DTYPE["int8"] = torch.int8
+
+
+@lru_cache(maxsize=8)
+def _mlu_device_count_stateless(
+    mlu_visible_devices: Optional[str] = None) -> int:
+
+    if mlu_visible_devices is None:
+        return torch.mlu.device_count()
+    if mlu_visible_devices == "":
+        return 0
+    if "," not in mlu_visible_devices:
+        return 1
+    return len(mlu_visible_devices.split(","))
+
+
+def mlu_device_count_stateless() -> int:
+    """Get number of MLU devices, caching based on the value of
+    MLU_VISIBLE_DEVICES at the time of call.
+    
+    This should be used instead of torch.cuda.device_count()
+    unless MLU_VISIBLE_DEVICES has already been set to the desired
+    value."""
+
+    # This can be removed and simply replaced with torch.cuda.get_device_count
+    # after https://github.com/pytorch/pytorch/pull/122815 is released.
+    return _mlu_device_count_stateless(envs.MLU_VISIBLE_DEVICES)
+
+
+def vllm__utils___maybe_force_spawn():
+    """Check if we need to force the use of the `spawn` multiprocessing start
+    method.
+    """
+    if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == "spawn":
+        return
+
+    reason = None
+    if cuda_is_initialized():
+        reason = "CUDA is initialized"
+    elif is_in_ray_actor():
+        # even if we choose to spawn, we need to pass the ray address
+        # to the subprocess so that it knows how to connect to the ray cluster.
+        # env vars are inherited by subprocesses, even if we use spawn.
+        import ray
+        os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
+        reason = "In a Ray actor and can only be spawned"
+
+    if reason is None:
+        reason = "Force use 'spawn' for MLU platform"
+
+    if reason is not None:
+        logger.warning(
+            "We must use the `spawn` multiprocessing start method. "
+            "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
+            "See https://docs.vllm.ai/en/latest/getting_started/"
+            "troubleshooting.html#python-multiprocessing "
+            "for more information. Reason: %s", reason)
+        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+
+def vllm__utils__get_kv_cache_torch_dtype(
+        cache_dtype: Optional[Union[str, torch.dtype]],
+        model_dtype: Optional[Union[str, torch.dtype]] = None) -> torch.dtype:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: use STR_DTYPE_TO_TORCH_DTYPE to get torch_dtype
+    '''  
+    if isinstance(cache_dtype, str):
+        if cache_dtype == "auto":
+            if isinstance(model_dtype, str):
+                torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
+            elif isinstance(model_dtype, torch.dtype):
+                torch_dtype = model_dtype
+            else:
+                raise ValueError(f"Invalid model dtype: {model_dtype}")
+        elif cache_dtype in ["half", "bfloat16", "float"]:
+            torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
+        elif cache_dtype == "fp8":
+            torch_dtype = torch.uint8
+        elif cache_dtype == 'int8':
+            torch_dtype = torch.int8
+        else:
+            raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
+    elif isinstance(cache_dtype, torch.dtype):
+        torch_dtype = cache_dtype
+    else:
+        raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
+    return torch_dtype
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+'''
+=============================
+Modify by vllm_mlu
+=============================
+@brief: change dispatch_key default value from 'CUDA' to 'MLU'
+'''
+'''
+==================
+End of MLU Hijack
+==================
+'''
+def vllm__utils__direct_register_custom_op(
+    op_name: str,
+    op_func: Callable,
+    mutates_args: List[str],
+    fake_impl: Optional[Callable] = None,
+    target_lib: Optional[Library] = None,
+    dispatch_key: str = "MLU",
+):
+    """
+    `torch.library.custom_op` can have significant overhead because it
+    needs to consider complicated dispatching logic. This function
+    directly registers a custom op and dispatches it to the CUDA backend.
+    See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
+    for more details.
+
+    By default, the custom op is registered to the vLLM library. If you
+    want to register it to a different library, you can pass the library
+    object to the `target_lib` argument.
+
+    IMPORTANT: the lifetime of the operator is tied to the lifetime of the
+    library object. If you want to bind the operator to a different library,
+    make sure the library object is alive when the operator is used.
+    """
+    if is_in_doc_build():
+        return
+
+    if not supports_custom_op():
+        from vllm.platforms import current_platform
+        assert not current_platform.is_cuda_alike(), (
+            "cuda platform needs torch>=2.4 to support custom op, "
+            "chances are you are using an old version of pytorch "
+            "or a custom build of pytorch. It is recommended to "
+            "use vLLM in a fresh new environment and let it install "
+            "the required dependencies.")
+        return
+
+    import torch.library
+    if hasattr(torch.library, "infer_schema"):
+        schema_str = torch.library.infer_schema(op_func,
+                                                mutates_args=mutates_args)
+    else:
+        # for pytorch 2.4
+        import torch._custom_op.impl
+        schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
+    my_lib = target_lib or vllm_lib
+    my_lib.define(op_name + schema_str)
+    my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
+    if fake_impl is not None:
+        my_lib._register_fake(op_name, fake_impl)
+
+
+
+MluHijackObject.apply_hijack(utils,
+                             utils.direct_register_custom_op,
+                             vllm__utils__direct_register_custom_op)
+MluHijackObject.apply_hijack(utils,
+                             utils.get_kv_cache_torch_dtype,
+                             vllm__utils__get_kv_cache_torch_dtype)
+MluHijackObject.apply_hijack(utils,
+                             utils._maybe_force_spawn,
+                             vllm__utils___maybe_force_spawn)
diff --git a/vllm_mlu/vllm_mlu/worker/__init__.py b/vllm_mlu/vllm_mlu/worker/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_mlu/vllm_mlu/worker/cache_engine.py b/vllm_mlu/vllm_mlu/worker/cache_engine.py
new file mode 100644
index 000000000..7217d12eb
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/worker/cache_engine.py
@@ -0,0 +1,180 @@
+"""CacheEngine class for managing the KV cache."""
+from typing import List
+
+import torch
+
+from vllm.config import CacheConfig, ModelConfig, ParallelConfig
+from vllm.logger import init_logger
+from vllm.utils import is_pin_memory_available, get_dtype_size, LayerBlockType, STR_DTYPE_TO_TORCH_DTYPE
+from vllm.worker.cache_engine import CacheEngine
+from vllm_mlu.mlu_hijack_utils import MluHijackObject
+
+logger = init_logger(__name__)
+
+
+def vllm__worker__cache_engine__CacheEngine___allocate_kv_cache(
+    self,
+    num_blocks: int,
+    device: str,
+) -> List[List[torch.Tensor]]:
+    """Allocates KV cache on the specified device."""
+    kv_cache_shape = self.attn_backend.get_kv_cache_shape(
+        num_blocks, self.block_size, self.num_kv_heads, self.head_size)
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: add kv_cache_scale for int8 support
+    ''' 
+    kv_cache_scales_shape = self.attn_backend.get_kv_cache_scale_shape(
+        num_blocks, self.block_size, self.num_kv_heads)
+    pin_memory = is_pin_memory_available() if device == "cpu" else False
+    kv_cache: List[List[torch.Tensor]] = []
+    for _ in range(self.num_attention_layers):
+        # null block in CpuGpuBlockAllocator requires at least that
+        # block to be zeroed-out.
+        # We zero-out everything for simplicity.
+        kv_cache_ = torch.zeros(kv_cache_shape,
+                                dtype=self.dtype,
+                                pin_memory=pin_memory,
+                                device=device)
+        if self.dtype == torch.int8:
+            kv_cache_scale_ = torch.zeros(kv_cache_scales_shape,
+                                          dtype=torch.float32,
+                                          pin_memory=pin_memory,
+                                          device=device)
+        else:
+            kv_cache_scale_ = torch.tensor([],
+                                           dtype=torch.float32,
+                                           device=device)
+        kv_cache.append([kv_cache_, kv_cache_scale_])
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+    return kv_cache
+
+
+def vllm__worker__cache_engine__CacheEngine__swap_in(self, src_to_dst: torch.Tensor) -> None:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: swap kv_cache_scale for int8 support
+    ''' 
+    for i in range(self.num_attention_layers):
+        # swap kv_cache
+        self.attn_backend.swap_blocks(self.cpu_cache[i][0], self.gpu_cache[i][0],
+                                        src_to_dst)
+        if self.dtype == torch.int8:
+            # swap kv_cache_scale
+            self.attn_backend.swap_blocks(self.cpu_cache[i][1], self.gpu_cache[i][1],
+                                            src_to_dst)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+def vllm__worker__cache_engine__CacheEngine__swap_out(self, src_to_dst: torch.Tensor) -> None:
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: swap kv_cache_scale for int8 support
+    ''' 
+    for i in range(self.num_attention_layers):
+        # swap kv_cache
+        self.attn_backend.swap_blocks(self.gpu_cache[i][0], self.cpu_cache[i][0],
+                                        src_to_dst)
+        if self.dtype == torch.int8:
+            # swap kv_cache_scale
+            self.attn_backend.swap_blocks(self.gpu_cache[i][1], self.cpu_cache[i][1],
+                                            src_to_dst)
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+vllm__worker__cache_engine__CacheEngine__get_cache_block_size__org = CacheEngine.get_cache_block_size
+
+@staticmethod
+def vllm__worker__cache_engine__CacheEngine__get_cache_block_size(
+    cache_config: CacheConfig,
+    model_config: ModelConfig,
+    parallel_config: ParallelConfig,
+) -> int:
+    kv_cache_total_size = vllm__worker__cache_engine__CacheEngine__get_cache_block_size__org(
+        cache_config=cache_config,
+        model_config=model_config,
+        parallel_config=parallel_config
+    )
+    '''
+    =============================
+    Modify by vllm_mlu
+    =============================
+    @brief: compute kv_cache_scale total size
+    ''' 
+    num_heads = model_config.get_num_kv_heads(parallel_config)
+    num_attention_layers = model_config.get_num_layers_by_block_type(
+        parallel_config, LayerBlockType.attention)
+
+    kv_cache_scale_total_size = 0
+    if cache_config.cache_dtype == 'int8':
+        key_cache_scale_block = cache_config.block_size * num_heads
+        value_cache_scale_block = key_cache_scale_block if not model_config.use_mla else 0
+        scale_total = num_attention_layers * (key_cache_scale_block + value_cache_scale_block)
+        dtype_size = get_dtype_size(torch.float32)
+        kv_cache_scale_total_size = dtype_size * scale_total
+
+    return kv_cache_total_size + kv_cache_scale_total_size
+    '''
+    ==================
+    End of MLU Hijack
+    ==================
+    '''
+
+
+@staticmethod
+def vllm__worker__cache_engine__CacheEngine__get_max_num_gpu_blocks(
+    cache_config: CacheConfig,
+    model_config: ModelConfig,
+    parallel_config: ParallelConfig,
+) -> int:
+    '''
+    Fix: tmo.reshape_paged_cache limit: RuntimeError: The addressing range of kv_cache cannot exceed 4G
+    Fix: cnnlScaledDotProductAttn_v5 check failed: cnnlGetTensorElementNum(key_desc) * kv_cache_dbyte <= INT32_MAX
+    '''
+    head_size = model_config.get_head_size()
+    num_heads = model_config.get_num_kv_heads(parallel_config)
+    if cache_config.cache_dtype == "auto":
+        dtype = model_config.dtype
+    else:
+        dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
+    dtype_size = get_dtype_size(dtype)
+    max_kv_cache_size = min(4 * 1024 * 1024 * 1024, torch.iinfo(torch.int).max)
+
+    max_num_gpu_blocks = max_kv_cache_size // (cache_config.block_size * num_heads * head_size * dtype_size)
+
+    return max_num_gpu_blocks
+
+
+MluHijackObject.apply_hijack(CacheEngine,
+                             CacheEngine._allocate_kv_cache,
+                             vllm__worker__cache_engine__CacheEngine___allocate_kv_cache)
+MluHijackObject.apply_hijack(CacheEngine,
+                             CacheEngine.swap_in,
+                             vllm__worker__cache_engine__CacheEngine__swap_in)
+MluHijackObject.apply_hijack(CacheEngine,
+                             CacheEngine.swap_out,
+                             vllm__worker__cache_engine__CacheEngine__swap_out)
+MluHijackObject.apply_hijack(CacheEngine,
+                             CacheEngine.get_cache_block_size,
+                             vllm__worker__cache_engine__CacheEngine__get_cache_block_size)
+MluHijackObject.apply_hijack(CacheEngine,
+                             "get_max_num_gpu_blocks",
+                             vllm__worker__cache_engine__CacheEngine__get_max_num_gpu_blocks)
diff --git a/vllm_mlu/vllm_mlu/worker/mlu_enc_dec_model_runner.py b/vllm_mlu/vllm_mlu/worker/mlu_enc_dec_model_runner.py
new file mode 100644
index 000000000..e6200b8b8
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/worker/mlu_enc_dec_model_runner.py
@@ -0,0 +1,380 @@
+import itertools
+from typing import List, Optional, Tuple, Type
+
+import torch
+import torch.distributed
+
+from vllm.attention.backends.abstract import AttentionMetadata
+from vllm.attention.backends.utils import PAD_SLOT_ID
+from vllm.logger import init_logger
+from vllm.sampling_params import SamplingParams
+from vllm.model_executor.layers.sampler import SamplerOutput
+from vllm.sequence import (IntermediateTensors, PoolerOutput,
+                           SequenceGroupMetadata)
+from vllm.utils import make_tensor_with_pad
+from vllm.worker.enc_dec_model_runner import EncoderDecoderModelInput, EncoderDecoderModelRunner
+from vllm.forward_context import set_forward_context
+from vllm.multimodal import MultiModalKwargs
+
+from vllm_mlu.worker.mlu_model_runner import (MLUModelRunnerBase,
+                                              ModelInputForMLUBuilder)
+from vllm_mlu._mlu_utils import VLLM_LATENCY_DEBUG_WITH_DEVICE_EN
+
+logger = init_logger(__name__)
+
+
+class MLUEncoderDecoderModelRunner(EncoderDecoderModelRunner,
+                                   MLUModelRunnerBase[EncoderDecoderModelInput]):
+    _model_input_cls: Type[EncoderDecoderModelInput] = (
+        EncoderDecoderModelInput)
+    _builder_cls: Type[ModelInputForMLUBuilder] = (ModelInputForMLUBuilder)
+
+    @torch.inference_mode()
+    def profile_run(self) -> None:
+        # Enable top-k sampling to reflect the accurate memory usage.
+        sampling_params = SamplingParams(top_p=0.99, top_k=self.vocab_size - 1)
+        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
+        max_num_seqs = self.scheduler_config.max_num_seqs
+
+        # Profile memory usage with max_num_sequences sequences and the total
+        # number of tokens equal to max_num_batched_tokens.
+        seqs: List[SequenceGroupMetadata] = []
+
+        max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(
+            self.model_config)
+        if max_mm_tokens > 0:
+            logger.info("Starting profile run for multi-modal models.")
+
+        batch_size = 0
+        for group_id in range(max_num_seqs):
+            seq_len = (max_num_batched_tokens // max_num_seqs +
+                       (group_id < max_num_batched_tokens % max_num_seqs))
+            batch_size += seq_len
+
+            decoder_dummy_data = self.input_registry \
+                .dummy_data_for_profiling(self.model_config,
+                                          seq_len,
+                                          self.mm_registry,
+                                          is_encoder_data=False)
+            encoder_dummy_data = self.input_registry \
+                .dummy_data_for_profiling(self.model_config,
+                                          seq_len,
+                                          self.mm_registry,
+                                          is_encoder_data=True)
+
+            # Having more tokens is over-conservative but otherwise fine
+            assert len(
+                decoder_dummy_data.seq_data.prompt_token_ids
+            ) >= seq_len, (
+                f"Expected at least {seq_len} dummy tokens for profiling, "
+                f"but got: {len(decoder_dummy_data.seq_data.prompt_token_ids)}"
+            )
+
+            assert decoder_dummy_data.multi_modal_data is None or \
+            encoder_dummy_data.multi_modal_data is None, (
+                "Multi-modal data can't be provided in both encoder and decoder"
+            )
+
+            seq = SequenceGroupMetadata(
+                request_id=str(group_id),
+                is_prompt=True,
+                seq_data={group_id: decoder_dummy_data.seq_data},
+                sampling_params=sampling_params,
+                block_tables=None,
+                encoder_seq_data=encoder_dummy_data.seq_data,
+                cross_block_table=None,
+                multi_modal_data=decoder_dummy_data.multi_modal_data
+                or encoder_dummy_data.multi_modal_data,
+                multi_modal_placeholders=decoder_dummy_data.
+                multi_modal_placeholders
+                or encoder_dummy_data.multi_modal_placeholders)
+            seqs.append(seq)
+
+        finished_requests_ids = [seq.request_id for seq in seqs]
+        model_input = self.prepare_model_input(
+            seqs, finished_requests_ids=finished_requests_ids)
+        intermediate_tensors = None
+        self.execute_model(model_input, None, intermediate_tensors)
+        torch.mlu.synchronize()
+        return
+
+    def _prepare_encoder_model_input_tensors(
+        self,
+        seq_group_metadata_list: List[SequenceGroupMetadata],
+        model_input: EncoderDecoderModelInput,
+    ) -> Tuple[AttentionMetadata, Optional[torch.Tensor],
+               Optional[torch.Tensor]]:
+        """Helper method to prepare the encoder- and cross-attn-related
+        model inputs based on a given sequence group. These additional inputs
+        are used to augment an already-computed `EncoderDecoderModelInput`
+        data structure which already has decoder-related model inputs
+        populated.
+
+        Sets the following attn_metadata fields:
+        * `num_encoder_tokens`
+        * `encoder_seq_lens`
+        * `encoder_seq_lens_tensor`
+        * `max_encoder_seq_len`
+        * `cross_slot_mapping`
+        * `cross_block_tables`
+
+        Constructs a new model inputs data structure, based on
+        (1) the existing fields in the `model_inputs` argument,
+        and (2) the following additional fields which are
+        computed (or in the case of `attn_metadata`, updated) 
+        by this function:
+        * attn_metadata
+        * encoder_input_tokens
+        * encoder_input_positions
+
+        Arguments:
+
+        * seq_group_metadata_list: list of sequence groups for which to
+                                   compute inputs
+        * model_inputs: model inputs data structure with decoder-oriented
+                        fields already computed.
+
+        Return:
+
+        * Updated model inputs data structure
+        """
+
+        if len(seq_group_metadata_list) == 0:
+            return (model_input.attn_metadata, None, None)
+
+        # Since we are not supporting chunked prefill either the entire
+        # batch is prefill or it is decode
+        is_prompt = seq_group_metadata_list[0].is_prompt
+
+        # Build encoder inputs
+        encoder_seq_lens: List[int] = []
+        if is_prompt:
+            # Prefill phase.
+            cross_block_tables = self._empty_int32_tensor().view(
+                len(seq_group_metadata_list), -1)
+
+            # Extract input tokens/positions, cross-attention slot-mapping,
+            # & seq len from each sequence group metadata
+            (
+                encoder_input_tokens,
+                encoder_input_positions,
+                cross_slot_mapping,
+            ) = (
+                [],
+                [],
+                [],
+            )
+            for seq_group_metadata in seq_group_metadata_list:
+                # Build seq lens
+                seq_len = seq_group_metadata.encoder_seq_data.get_len()
+                token_ids = seq_group_metadata.encoder_seq_data.get_token_ids()
+                encoder_seq_lens.append(seq_len)
+
+                # Build slot mapping
+                is_profile_run = (seq_group_metadata.block_tables is None)
+                if is_profile_run:
+                    # During memory profiling, the block tables are not
+                    # initialized yet. In this case, we just use a dummy
+                    # slot mapping.
+                    # In embeddings, the block tables are {seq_id: None}.
+                    cross_slot_mapping.extend([PAD_SLOT_ID] * seq_len)
+                else:
+                    for i in range(0, seq_len):
+                        block_number = seq_group_metadata.cross_block_table[
+                            i // self.block_size]
+                        block_offset = i % self.block_size
+                        slot = block_number * self.block_size + block_offset
+                        cross_slot_mapping.append(slot)
+
+                # Build encoder input tokens
+                encoder_input_tokens.extend(token_ids)
+                encoder_input_positions.extend(list(range(0, seq_len)))
+
+            # Convert tokens/positions & cross-attention
+            # slot-mapping to encoder input tensors
+            encoder_input_tokens_tensor = self._list_to_long_tensor(
+                encoder_input_tokens)
+            encoder_input_positions_tensor = self._list_to_long_tensor(
+                encoder_input_positions)
+            cross_slot_mapping_tensor = self._list_to_int32_tensor(
+                cross_slot_mapping)
+
+        else:
+            # Decode phase.
+            encoder_input_tokens_tensor = self._empty_long_tensor()
+            encoder_input_positions_tensor = self._empty_long_tensor()
+            cross_slot_mapping_tensor = self._empty_int32_tensor()
+            # Extract cross-attention block tables &
+            # seq len from each sequence group metadata.
+            # Cross-attention block tables are empty
+            # during vLLM memory profiling.
+            cross_block_tables = []
+            for seq_group_metadata in seq_group_metadata_list:
+                for _ in range(len(seq_group_metadata.seq_data)):
+                    encoder_seq_lens.append(
+                        seq_group_metadata.encoder_seq_data.get_len())
+                    cross_block_table = seq_group_metadata.cross_block_table
+                    cross_block_tables.append([] if (
+                        cross_block_table is None) else cross_block_table)
+
+            if (model_input.attn_metadata is not None
+                    and model_input.attn_metadata.use_cuda_graph):
+                # We will be using CUDA graph replay for this decode.
+                max_len_of_block_table = self.get_max_block_per_batch()
+                batch_size = len(encoder_seq_lens)
+                graph_batch_size = self.vllm_config.pad_for_cudagraph(
+                    batch_size)
+                assert graph_batch_size >= batch_size
+                cuda_graph_pad_size = graph_batch_size - batch_size
+                # extend the cross_block_tables and encoder_seq_lens to match
+                # the graph_batch_size.
+                cross_block_tables.extend([[]
+                                           for _ in range(cuda_graph_pad_size)
+                                           ])
+                encoder_seq_lens.extend(
+                    itertools.repeat(1, cuda_graph_pad_size))
+
+            else:
+                max_len_of_block_table = max(
+                    len(block_table) for block_table in cross_block_tables)
+
+            cross_block_tables = make_tensor_with_pad(
+                cross_block_tables,
+                max_len=max_len_of_block_table,
+                pad=0,
+                dtype=torch.int32,
+                device=self.device,
+            )
+
+        # Compute encoder sequence lengths & encoder
+        # sequence starting offset tensors
+        max_encoder_seq_len = max(encoder_seq_lens, default=0)
+        encoder_seq_lens_tensor = self._list_to_int32_tensor(encoder_seq_lens)
+        encoder_seq_start_loc = torch.zeros(encoder_seq_lens_tensor.shape[0] +
+                                            1,
+                                            dtype=torch.int32,
+                                            device=self.device)
+        torch.cumsum(encoder_seq_lens_tensor,
+                     dim=0,
+                     dtype=encoder_seq_start_loc.dtype,
+                     out=encoder_seq_start_loc[1:])
+
+        # Update attention metadata with encoder-oriented attributes
+        attn_metadata = model_input.attn_metadata
+        assert attn_metadata is not None
+        (
+            attn_metadata.num_encoder_tokens,
+            attn_metadata.encoder_seq_lens,
+            attn_metadata.encoder_seq_lens_tensor,
+            attn_metadata.max_encoder_seq_len,
+            attn_metadata.encoder_seq_start_loc,
+            attn_metadata.cross_slot_mapping,
+            attn_metadata.cross_block_tables,
+        ) = (
+            sum(encoder_seq_lens),
+            encoder_seq_lens,
+            encoder_seq_lens_tensor,
+            max_encoder_seq_len,
+            encoder_seq_start_loc,
+            cross_slot_mapping_tensor,
+            cross_block_tables,
+        )
+
+        return (attn_metadata, encoder_input_tokens_tensor,
+                encoder_input_positions_tensor)
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        model_input: EncoderDecoderModelInput,
+        kv_caches: List[torch.Tensor],
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        num_steps: int = 1,
+    ) -> Optional[List[PoolerOutput]]:
+        if num_steps > 1:
+            raise ValueError("num_steps > 1 is not supported in "
+                             "EncoderDecoderModelRunner")
+    
+        if (model_input.attn_metadata is not None
+                and model_input.attn_metadata.prefill_metadata is None
+                and model_input.attn_metadata.decode_metadata.use_cuda_graph):
+            assert model_input.input_tokens is not None
+            graph_batch_size = model_input.input_tokens.shape[0]
+            model_executable = self.graph_runners[
+                model_input.virtual_engine][graph_batch_size]
+        else:
+            model_executable = self.model
+    
+        seqlen_agnostic_kwargs = {
+            "finished_requests_ids": model_input.finished_requests_ids,
+            "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
+        } if self.has_inner_state else {}
+    
+        multi_modal_kwargs = model_input.multi_modal_kwargs or {}
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: add mlu metrics 
+        '''
+        # Add time markers for model_executable+compute_logits
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+            prefill_meta = model_input.attn_metadata.prefill_metadata
+            decode_meta = model_input.attn_metadata.decode_metadata
+            use_cuda_graph = ((prefill_meta is None and decode_meta.use_cuda_graph))
+            # if use_cuda_graph, the start timestamp will be inserted inside MLUGraphRunner.forward()
+            if not use_cuda_graph:
+                start = torch.mlu.Event(enable_timing=True)
+                start.record()
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        with set_forward_context(model_input.attn_metadata, self.vllm_config,
+                                 model_input.virtual_engine):
+            hidden_or_intermediate_states = model_executable(
+                input_ids=model_input.input_tokens,
+                positions=model_input.input_positions,
+                encoder_input_ids=model_input.encoder_input_tokens,
+                encoder_positions=model_input.encoder_input_positions,
+                intermediate_tensors=intermediate_tensors,
+                **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
+                                             device=self.device),
+                **seqlen_agnostic_kwargs)
+
+        logits = self.model.compute_logits(hidden_or_intermediate_states,
+                                           model_input.sampling_metadata)
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: add mlu metrics 
+        '''
+        # Add time markers for model_executable+compute_logits
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+            end_marker = torch.mlu.Event(enable_timing=True)
+            end_marker.record()
+            if use_cuda_graph:
+                self.time_markers = (model_executable.start, end_marker)
+            else:
+                self.time_markers = (start, end_marker)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+    
+        if not self.is_driver_worker:
+            return []
+    
+        if model_input.async_callback is not None:
+            model_input.async_callback()
+    
+        # Sample the next token.
+        output: SamplerOutput = self.model.sample(
+            logits=logits,
+            sampling_metadata=model_input.sampling_metadata,
+        )
+    
+        return [output]
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/worker/mlu_model_runner.py b/vllm_mlu/vllm_mlu/worker/mlu_model_runner.py
new file mode 100644
index 000000000..13dbbda35
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/worker/mlu_model_runner.py
@@ -0,0 +1,1413 @@
+import gc
+import inspect
+import itertools
+import time
+import weakref
+import os
+from typing import Dict, List, Optional, Set, Tuple, Union, Type
+
+import numpy as np
+import torch
+import torch.distributed
+import torch.nn as nn
+from contextlib import contextmanager
+from dataclasses import dataclass
+from tqdm import tqdm
+
+from vllm.attention import AttentionMetadata
+from vllm.attention.selector import get_attn_backend
+from vllm.attention.backends.utils import CommonAttentionState
+from vllm.config import VllmConfig
+from vllm.distributed import (get_pp_group, get_world_group,
+                              get_kv_transfer_group)
+from vllm.distributed.parallel_state import (get_tensor_model_parallel_rank,
+                                             get_tp_group)
+from vllm.forward_context import get_forward_context, set_forward_context
+from vllm.inputs import INPUT_REGISTRY, InputRegistry
+from vllm.logger import init_logger
+from vllm.lora.layers import LoRAMapping
+from vllm.lora.request import LoRARequest
+from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
+from vllm.model_executor import SamplingMetadataCache
+from vllm.model_executor.layers.sampler import SamplerOutput
+from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
+from vllm.multimodal import (MULTIMODAL_REGISTRY, MultiModalKwargs,
+                             MultiModalRegistry)
+from vllm.prompt_adapter.layers import PromptAdapterMapping
+from vllm.prompt_adapter.request import PromptAdapterRequest
+from vllm.prompt_adapter.worker_manager import (
+    LRUCacheWorkerPromptAdapterManager)
+from vllm.sampling_params import SamplingParams
+from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
+from vllm.utils import (GiB_bytes, PyObjectCache,
+                        async_tensor_h2d, flatten_2d_lists,
+                        is_pin_memory_available)
+from vllm.worker.model_runner_base import ModelRunnerBase
+from vllm.model_executor.layers.linear import RowParallelLinear
+from vllm.worker.model_runner import (
+    ModelInputForGPU, ModelInputForGPUWithSamplingMetadata,
+    ModelRunnerBase, LORA_WARMUP_RANK,
+)
+from vllm.worker.model_runner import (
+    TModelInputForGPU, ModelInputForGPU,
+    ModelInputForGPUWithSamplingMetadata,
+    ModelInputForGPUBuilder, GPUModelRunnerBase,
+    ModelRunner, CUDAGraphRunner,
+    LORA_WARMUP_RANK, _NUM_WARMUP_ITERS
+)
+
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.model_executor.layers.rotary_embedding import MLURotaryEmbedding
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+from vllm_mlu._mlu_utils import *
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class MLUGraphCaptureContext:
+    stream: torch.mlu.Stream
+
+
+@contextmanager
+def mlu_graph_capture(device: torch.device):
+    """
+    `graph_capture` is a context manager which should surround the code that
+    is capturing the CUDA graph. Its main purpose is to ensure that the
+    some operations will be run after the graph is captured, before the graph
+    is replayed. It returns a `GraphCaptureContext` object which contains the
+    necessary data for the graph capture. Currently, it only contains the
+    stream that the graph capture is running on. This stream is set to the
+    current CUDA stream when the context manager is entered and reset to the
+    default stream when the context manager is exited. This is to ensure that
+    the graph capture is running on a separate stream from the default stream,
+    in order to explicitly distinguish the kernels to capture
+    from other kernels possibly launched on background in the default stream.
+    """
+    context = MLUGraphCaptureContext(torch.mlu.Stream(device=device))
+    with get_tp_group().graph_capture(context), get_pp_group().graph_capture(
+            context):
+        yield context
+
+
+def _model_forward_pre_hook(self, args, kwargs):
+    '''
+    This hook function will be called before model.forward
+    '''
+    assert len(args) == 0 and len(kwargs) > 0, \
+        f"The pre-forward's expected inputs are not passed by kwargs. " + \
+        f"Expected len(args)=0, len(kwargs)>0, " + \
+        f"now, len(args)={len(args)}, len(kwargs)={len(kwargs)}."
+
+    input_ids = kwargs['input_ids']
+
+    attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
+
+    # model inputs can be None for data parallel
+    if attn_metadata is None:
+        return (args, kwargs)
+
+    # Prepare attributes for all rope in model
+    MLURotaryEmbedding.set_mlu_var(input_ids=input_ids,
+                                   attn_metadata=attn_metadata)
+    RowParallelLinear.is_prompt = True if attn_metadata.prefill_metadata else False
+    FeedForward.is_prompt = True if attn_metadata.prefill_metadata else False
+    if (attn_metadata.prefill_metadata
+            and not attn_metadata.prefill_metadata.is_profile_run):
+        set_is_prompt(True)
+    else:
+        set_is_prompt(False)
+
+    return (args, kwargs)
+
+
+class ModelInputForMLUBuilder(ModelInputForGPUBuilder):
+
+    def build(self) -> ModelInputForGPU:
+        """Finalize the builder intermediate data and
+        create on-device tensors.
+        """
+        # Combine and flatten intermediate data.
+        input_tokens = []
+        token_types = []
+        for inter_data in self.inter_data_list:
+            for cur_input_tokens in inter_data.input_tokens:
+                input_tokens.extend(cur_input_tokens)
+            for cur_token_types in inter_data.token_types:
+                token_types.extend(cur_token_types)
+
+        if not input_tokens:
+            # This may happen when all prefill requests hit
+            # prefix caching and there is no decode request.
+            return self.model_input_cls()
+
+        mrope_input_positions: Optional[List[List[int]]] = None
+        if any(inter_data.mrope_input_positions is not None
+               for inter_data in self.inter_data_list):
+            mrope_input_positions = [[] for _ in range(3)]
+            for idx in range(3):
+                for inter_data in self.inter_data_list:
+                    msections = inter_data.mrope_input_positions
+                    if msections is None:
+                        for _seq_input_positions in inter_data.input_positions:
+                            mrope_input_positions[idx].extend(
+                                _seq_input_positions)
+                    else:
+                        for _seq_mrope_input_positions in msections:
+                            mrope_input_positions[idx].extend(
+                                _seq_mrope_input_positions[idx])
+            input_positions = None
+        else:
+            input_positions = []
+            for inter_data in self.inter_data_list:
+                for cur_input_positions in inter_data.input_positions:
+                    input_positions.extend(cur_input_positions)
+
+        seq_lens = []
+        query_lens = []
+        max_decode_seq_len = 0
+        max_encoder_seq_len = 0
+        for inter_data in self.inter_data_list:
+            seq_lens.extend(inter_data.seq_lens)
+            query_lens.extend(inter_data.query_lens)
+            if not inter_data.is_prompt:
+                max_decode_seq_len = max(max_decode_seq_len,
+                                         max(inter_data.seq_lens))
+                if self.runner.model_config.is_encoder_decoder:
+                    max_encoder_seq_len = max(max_encoder_seq_len,
+                                              inter_data.encoder_seq_len)
+
+        # Mapping from request IDs to sequence IDs. Used for Jamba models
+        # that manages the cache by itself.
+        request_ids_to_seq_ids = {
+            data.request_id: data.seq_ids
+            for data in self.inter_data_list
+        }
+
+        cuda_graph_pad_size = self._get_cuda_graph_pad_size(
+            num_seqs=len(seq_lens),
+            max_decode_seq_len=max_decode_seq_len,
+            max_encoder_seq_len=max_encoder_seq_len)
+
+        batch_size = len(input_tokens)
+        if cuda_graph_pad_size != -1:
+            # If cuda graph can be used, pad tensors accordingly.
+            # See `capture_model` API for more details.
+            # vLLM uses cuda graph only for decoding requests.
+            batch_size += cuda_graph_pad_size
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: check mlugraph enable
+        '''
+        # If mlu graph is enabled, but still run with eager mode,
+        # print warning message.
+        # This behavior can be disabled by export VLLM_GRAPH_DEBUG=0.
+        if VLLM_GRAPH_DEBUG \
+                and self.decode_only \
+                and not self.runner.model_config.enforce_eager \
+                and cuda_graph_pad_size == -1 \
+                and max_decode_seq_len > self.runner.max_seq_len_to_capture:
+            logger.warning(f"Because one of the following conditions is not met, MLU Graphs will not be enabled.\n" +
+                        f"1. batch_size({batch_size}) <= max_batch_size_to_capture({self.runner.max_batchsize_to_capture})\n" +
+                        f"2. max_seq_len({max_decode_seq_len}) <= max_seq_len_to_capture({self.runner.max_seq_len_to_capture})\n" +
+                        f"Use 'export VLLM_GRAPH_DEBUG=false' to disable this warning.")
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        # Tokens and positions.
+        if cuda_graph_pad_size:
+            input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size))
+        assert self.runner.device is not None
+        input_tokens_tensor = async_tensor_h2d(input_tokens, torch.long,
+                                               self.runner.device,
+                                               self.runner.pin_memory)
+
+        token_types_tensor = async_tensor_h2d(token_types, torch.long,
+                                               self.runner.device,
+                                               self.runner.pin_memory) \
+                                                if token_types else None
+
+        if mrope_input_positions is not None:
+            for idx in range(3):
+                mrope_input_positions[idx].extend(
+                    itertools.repeat(0, cuda_graph_pad_size))
+            input_positions_tensor = async_tensor_h2d(mrope_input_positions,
+                                                      torch.int32,
+                                                      self.runner.device,
+                                                      self.runner.pin_memory)
+        else:
+            input_positions.extend(itertools.repeat(0, cuda_graph_pad_size))
+            input_positions_tensor = async_tensor_h2d(input_positions,
+                                                      torch.int32,
+                                                      self.runner.device,
+                                                      self.runner.pin_memory)
+        # Sequence and query lengths.
+        if cuda_graph_pad_size:
+            seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size))
+
+        # Attention metadata.
+        attn_metadata = self.attn_metadata_builder.build(
+            seq_lens, query_lens, cuda_graph_pad_size, batch_size)
+
+        # LoRA data.
+        lora_requests = set()
+        lora_mapping = None
+        if self.enable_lora:
+            lora_requests = set(r for data in self.inter_data_list
+                                for r in data.lora_requests)
+            lora_index_mapping = flatten_2d_lists([
+                flatten_2d_lists(inter_data.lora_index_mapping)
+                for inter_data in self.inter_data_list
+            ])
+            if cuda_graph_pad_size:
+                lora_index_mapping.extend(
+                    itertools.repeat(0, cuda_graph_pad_size))
+            lora_prompt_mapping = flatten_2d_lists([
+                flatten_2d_lists(inter_data.lora_prompt_mapping)
+                for inter_data in self.inter_data_list
+            ])
+
+            lora_mapping = LoRAMapping(
+                **dict(index_mapping=lora_index_mapping,
+                       prompt_mapping=lora_prompt_mapping,
+                       is_prefill=not self.decode_only))
+
+        # Prompt adapter data.
+        prompt_adapter_requests: Set[PromptAdapterRequest] = set()
+        prompt_adapter_mapping = None
+        if self.enable_prompt_adapter:
+            prompt_adapter_requests = set(
+                data.prompt_adapter_request for data in self.inter_data_list
+                if data.prompt_adapter_request is not None)
+            prompt_adapter_index_mapping = flatten_2d_lists([
+                inter_data.prompt_adapter_index_mapping
+                for inter_data in self.inter_data_list
+            ])
+            if cuda_graph_pad_size:
+                prompt_adapter_index_mapping.extend(
+                    itertools.repeat(0, cuda_graph_pad_size))
+            prompt_adapter_prompt_mapping = flatten_2d_lists([
+                inter_data.prompt_adapter_prompt_mapping
+                for inter_data in self.inter_data_list
+            ])
+            prompt_adapter_mapping = PromptAdapterMapping(
+                prompt_adapter_index_mapping,
+                prompt_adapter_prompt_mapping,
+            )
+
+        # Multi-modal data.
+        multi_modal_kwargs_list = [
+            data.multi_modal_kwargs for data in self.inter_data_list
+            if data.multi_modal_kwargs is not None
+        ]
+        multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
+
+        return self.model_input_cls(
+            input_tokens=input_tokens_tensor,
+            input_positions=input_positions_tensor,
+            token_types=token_types_tensor,
+            attn_metadata=attn_metadata,
+            seq_lens=seq_lens,
+            query_lens=query_lens,
+            lora_mapping=lora_mapping,
+            lora_requests=lora_requests,
+            multi_modal_kwargs=multi_modal_kwargs,
+            request_ids_to_seq_ids=request_ids_to_seq_ids,
+            finished_requests_ids=self.finished_requests_ids,
+            prompt_adapter_mapping=prompt_adapter_mapping,
+            prompt_adapter_requests=prompt_adapter_requests)
+
+
+class MLUModelRunnerBase(GPUModelRunnerBase[TModelInputForGPU]):
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_dtype: Optional[str] = "auto",
+        is_driver_worker: bool = False,
+        return_hidden_states: bool = False,
+        input_registry: InputRegistry = INPUT_REGISTRY,
+        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
+    ):
+
+        ModelRunnerBase.__init__(self, vllm_config)
+        model_config = self.model_config
+        cache_config = self.cache_config
+
+        self.is_driver_worker = is_driver_worker
+        self.return_hidden_states = return_hidden_states
+
+        self.device = self.device_config.device
+        self.pin_memory = is_pin_memory_available()
+
+        self.kv_cache_dtype = kv_cache_dtype
+        self.sliding_window = model_config.get_sliding_window()
+        self.block_size = cache_config.block_size
+        self.max_seq_len_to_capture = self.model_config.max_seq_len_to_capture
+        self.max_batchsize_to_capture = \
+            self.vllm_config.compilation_config.max_capture_size
+
+        self.graph_runners: List[Dict[int, MLUGraphRunner]] = [
+            {} for _ in range(self.parallel_config.pipeline_parallel_size)
+        ]
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief:
+        1) support context mlugraph
+        2) init self.time_markers which avoids to visit time_marks before calling execute_model
+        '''
+        self.context_graph_runners: List[MLUGraphRunner] = [
+            None for _ in range(self.parallel_config.pipeline_parallel_size)
+        ]
+
+        self.time_markers = []
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        self.graph_memory_pool: Optional[Tuple[
+            int, int]] = None  # Set during graph capture.
+
+        self.has_inner_state = model_config.has_inner_state
+
+        self.in_profile_run = False
+
+        # When using CUDA graph, the input block tables must be padded to
+        # max_seq_len_to_capture. However, creating the block table in
+        # Python can be expensive. To optimize this, we cache the block table
+        # in numpy and only copy the actual input content at every iteration.
+        # The shape of the cached block table will be
+        # (max batch size to capture, max seq len to capture / block size).
+        self.graph_block_tables = np.zeros(
+            (self.max_batchsize_to_capture, self.get_max_block_per_batch()),
+            dtype=np.int32)
+
+        # Attention-free but stateful models like Mamba need a placeholder attn
+        # backend, as the attention metadata is needed to manage internal state.
+        # However we must bypass attention selection altogether for some models
+        # used for speculative decoding to avoid a divide-by-zero in
+        # model_config.get_head_size()
+        num_attn_heads = self.model_config.get_num_attention_heads(
+            self.parallel_config)
+        needs_attn_backend = (num_attn_heads != 0
+                              or self.model_config.is_attention_free)
+
+        self.attn_backend = get_attn_backend(
+            self.model_config.get_head_size(),
+            self.model_config.dtype,
+            self.kv_cache_dtype,
+            self.block_size,
+            self.model_config.is_attention_free,
+            use_mla=self.model_config.use_mla,
+        ) if needs_attn_backend else None
+        if self.attn_backend:
+            self.attn_state = self.attn_backend.get_state_cls()(
+                weakref.proxy(self))
+        else:
+            self.attn_state = CommonAttentionState(weakref.proxy(self))
+
+        # Multi-modal data support
+        self.input_registry = input_registry
+        self.mm_registry = mm_registry
+        self.multi_modal_input_mapper = mm_registry \
+            .create_input_mapper(model_config)
+        self.mm_registry.init_mm_limits_per_prompt(self.model_config)
+
+        # Lazy initialization
+        self.model: nn.Module  # Set after load_model
+        # Set after load_model.
+        self.lora_manager: Optional[LRUCacheWorkerLoRAManager] = None
+        self.prompt_adapter_manager: LRUCacheWorkerPromptAdapterManager = None
+
+        set_cpu_offload_max_bytes(
+            int(self.cache_config.cpu_offload_gb * 1024**3))
+
+        # Used to cache python objects
+        self.inter_data_cache: Dict[int, PyObjectCache] = {}
+
+        # Using the PythonizationCache in Pipeline-Parallel clobbers the
+        # SequenceGroupToSample object. In Pipeline-Parallel, we have
+        # more than 1 Scheduler, resulting in a potential back-to-back
+        # prepare_model_inputs() call. This clobbers the cached
+        # SequenceGroupToSample objects, as we reset the cache during
+        # every prepare_model_inputs() call.
+        self.sampling_metadata_cache: SamplingMetadataCache = \
+              SamplingMetadataCache() \
+                if self.parallel_config.pipeline_parallel_size == 1 else None
+
+        if hasattr(self, "_builder_cls"):
+            # multi-step model runner does not have `_builder_cls`
+            self.builder = self._builder_cls(weakref.proxy(self))
+
+    def init_llmassistor(self):
+        try :
+            import llmassistor
+        except Exception as e:
+            raise ImportError(f"import llmassistor failed: {e}")
+        llmdebug = llmassistor.LLMDebug(
+            local_rank=get_world_group().local_rank,
+        )
+        llmdebug.trace(self.model)
+        self.llmdebug = llmdebug
+
+    def get_max_seq_len_for_profile_run(self, max_mm_tokens: int) -> int:
+        """
+        For some multimodal model, for example, llava, the max_mm_tokens(image|video) +
+        max_position_embeddings(text) is smaller than the input len when profile run, which
+        causes runtime error for rope.
+        So here, we calculate the max seq len for profile run, and compare it with
+        max_mm_tokens + max_position_embeddings for a WalkAround.
+        """
+        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
+        max_num_seqs = self.scheduler_config.max_num_seqs
+        max_num_seqs = max(1, min(max_num_seqs, max_num_batched_tokens // max_mm_tokens))
+        acc_seq_len = 0
+        for group_id in range(max_num_seqs):
+            seq_len = (max_num_batched_tokens // max_num_seqs +
+                        (group_id < max_num_batched_tokens % max_num_seqs))
+            acc_seq_len += seq_len
+        return acc_seq_len
+
+    def load_model(self) -> None:
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: 1. modify rope's max_position_embeddings to max_model_len
+                2. register model pre forward for rope optimization
+        '''
+        MLURotaryEmbedding.max_seq_len = self.model_config.max_model_len
+        if VLLM_AVG_MOE_EN:
+            logger.warning("Inference with Moe avg dispatch, "
+                "it's only for deepseek-v3/r1 model's performance test,"
+                " and will result in precision anomalies. Be careful!")
+            SparseMoeMlp.max_batched_token = max(self.model_config.max_model_len,
+                                                 self.scheduler_config.max_num_batched_tokens)
+        MLURotaryEmbedding.max_model_len = self.model_config.max_model_len
+        if self.model_config.multimodal_config is not None:
+            # for multimodal models, the max seq len for language model is
+            # max_mm_tokens + max_position_embeddings
+            max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(self.model_config)
+            profile_run_max_seq_len = self.get_max_seq_len_for_profile_run(max_mm_tokens)
+            MLURotaryEmbedding.max_seq_len = max(profile_run_max_seq_len,
+                                                 MLURotaryEmbedding.max_seq_len + max_mm_tokens)
+        super(MLUModelRunnerBase, self).load_model()
+        self.model.register_forward_pre_hook(_model_forward_pre_hook, with_kwargs=True)
+        if VLLM_DUMP_OUTPUTS:
+            self.init_llmassistor()
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+    def _dummy_run(self,
+                   max_num_batched_tokens: int,
+                   max_num_seqs: int = 1) -> None:
+        with self.set_in_profile_run():
+            # Avoid dumping useless profile dummy tensors.
+            if VLLM_DUMP_OUTPUTS:
+                self.llmdebug.enable = False
+
+            # Enable top-k sampling to reflect the accurate memory usage.
+            sampling_params = \
+                SamplingParams(top_p=0.99, top_k=self.vocab_size - 1)
+
+            # This represents the maximum number of different requests
+            # that will have unique loras, and therefore the max amount of
+            # memory consumption. Create dummy lora request copies from the
+            # lora request passed in, which contains a lora from the lora
+            # warmup path.
+            dummy_lora_requests: List[LoRARequest] = []
+            dummy_lora_requests_per_seq: List[LoRARequest] = []
+            if self.lora_config:
+                dummy_lora_requests = self._add_dummy_loras(
+                    self.lora_config.max_loras)
+                assert len(dummy_lora_requests) == self.lora_config.max_loras
+                dummy_lora_requests_per_seq = [
+                    dummy_lora_requests[idx % len(dummy_lora_requests)]
+                    for idx in range(max_num_seqs)
+                ]
+
+            # Profile memory usage with max_num_sequences sequences and the
+            # total number of tokens equal to max_num_batched_tokens.
+            seqs: List[SequenceGroupMetadata] = []
+            # Additional GPU memory may be needed for multi-modal encoding,
+            # which needs to be accounted for when calculating the GPU blocks
+            # for vLLM blocker manager.
+            # To exercise the worst scenario for GPU memory consumption,
+            # the number of seqs (batch_size) is chosen to maximize the number
+            # of images processed.
+
+            max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(
+                self.model_config)
+            if max_mm_tokens > 0:
+                max_num_seqs_orig = max_num_seqs
+                max_num_seqs = min(max_num_seqs,
+                                   max_num_batched_tokens // max_mm_tokens)
+                if max_num_seqs < 1:
+                    expr = (f"min({max_num_seqs_orig}, "
+                            f"{max_num_batched_tokens} // {max_mm_tokens})")
+                    logger.warning(
+                        "Computed max_num_seqs (%s) to be less than 1. "
+                        "Setting it to the minimum value of 1.", expr)
+                    max_num_seqs = 1
+
+            batch_size = 0
+            for group_id in range(max_num_seqs):
+                seq_len = (max_num_batched_tokens // max_num_seqs +
+                           (group_id < max_num_batched_tokens % max_num_seqs))
+                batch_size += seq_len
+
+                dummy_data = self.input_registry \
+                    .dummy_data_for_profiling(self.model_config,
+                                            seq_len,
+                                            self.mm_registry)
+
+                seq = SequenceGroupMetadata(
+                    request_id=str(group_id),
+                    is_prompt=True,
+                    seq_data={group_id: dummy_data.seq_data},
+                    sampling_params=sampling_params,
+                    block_tables=None,
+                    lora_request=dummy_lora_requests_per_seq[group_id]
+                    if dummy_lora_requests_per_seq else None,
+                    multi_modal_data=dummy_data.multi_modal_data,
+                    multi_modal_placeholders=dummy_data.
+                    multi_modal_placeholders,
+                )
+                seqs.append(seq)
+
+            # Run the model with the dummy inputs.
+            num_layers = self.model_config.get_num_layers(self.parallel_config)
+            # use an empty tensor instead of `None`` to force Dynamo to pass
+            # it by reference, rather by specializing on the value ``None``.
+            # the `dtype` argument does not matter, and we use `float32` as
+            # a placeholder (it has wide hardware support).
+            # it is important to create tensors inside the loop, rather than
+            # multiplying the list, to avoid Dynamo from treating them as
+            # tensor aliasing.
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: support kv cache int8
+            '''
+            kv_caches = []
+            for _ in range(num_layers):
+                kv_cache_ = torch.tensor([], dtype=torch.float32, device=self.device)
+                kv_cache_scale_ = torch.tensor([], dtype=torch.float32, device=self.device)
+                kv_caches.append([kv_cache_, kv_cache_scale_])
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+            finished_requests_ids = [seq.request_id for seq in seqs]
+            model_input = self.prepare_model_input(
+                seqs, finished_requests_ids=finished_requests_ids)
+            intermediate_tensors = None
+            if not get_pp_group().is_first_rank:
+                intermediate_tensors = \
+                    self.model.make_empty_intermediate_tensors(
+                    batch_size=batch_size,
+                    dtype=self.model_config.dtype,
+                    device=self.device)
+
+            # Disable KV Scale Calculation for dummy data during profile run
+            if model_input.attn_metadata is not None:
+                model_input.attn_metadata.enable_kv_scales_calculation = False
+
+            self.execute_model(model_input, kv_caches, intermediate_tensors)
+            torch.mlu.synchronize()
+            if self.lora_config:
+                self._remove_dummy_loras()
+
+            if VLLM_DUMP_OUTPUTS:
+                self.llmdebug.enable = True
+
+            return
+
+
+    @torch.inference_mode()
+    def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> None:
+        """Cuda graph capture a model.
+
+        Note that CUDA graph's performance gain is negligible if number
+        of batched tokens are larger than 200. And since CUDA graph
+        requires fixed sized tensors, supporting large/variable batch
+        size requires high GPU memory overhead. Thus, vLLM only captures
+        decoding requests. Mixed batch (chunked prefill + decoding) or
+        prefill requests are not captured.
+
+        Since it is used for decoding-only, it assumes there's only 1 token
+        per sequence in the batch.
+        """
+        assert not self.model_config.enforce_eager
+        logger.info("Capturing mlugraphs for decoding. This may lead to "
+                    "unexpected consequences if the model is not static. To "
+                    "run the model in eager mode, set 'enforce_eager=True' or "
+                    "use '--enforce-eager' in the CLI. "
+                    "If out-of-memory error occurs during mlugraph capture,"
+                    " consider decreasing `gpu_memory_utilization` or "
+                    "switching to eager mode. You can also reduce the "
+                    "`max_num_seqs` as needed to decrease memory usage.")
+        start_time = time.perf_counter()
+        start_free_gpu_memory = torch.mlu.mem_get_info()[0]
+
+        # Prepare dummy inputs. These will be reused for all batch sizes.
+        max_batch_size = self.max_batchsize_to_capture
+        input_tokens = torch.zeros(max_batch_size,
+                                   dtype=torch.long,
+                                   device=self.device)
+        input_positions = torch.zeros(max_batch_size,
+                                      dtype=torch.int32,
+                                      device=self.device)
+        if self.model_config.uses_mrope:
+            input_positions = torch.tile(input_positions,
+                                         (3, 1)).mlu(device=self.device)
+        # Prepare dummy previous_hidden_states only if needed by the model.
+        # This is used by draft models such as EAGLE.
+        previous_hidden_states = None
+        if "previous_hidden_states" in inspect.signature(
+                self.model.forward).parameters:
+            previous_hidden_states = torch.empty(
+                [max_batch_size,
+                 self.model_config.get_hidden_size()],
+                dtype=self.model_config.dtype,
+                device=self.device)
+
+        intermediate_inputs = None
+        if not get_pp_group().is_first_rank:
+            intermediate_inputs = self.model.make_empty_intermediate_tensors(
+                batch_size=max_batch_size,
+                dtype=self.model_config.dtype,
+                device=self.device)
+
+        dummy_lora_id: Optional[int] = None
+        dummy_lora_request: LoRARequest = []
+        if self.lora_config:
+            # The goal is to capture the LoRA kernels in cuda graphs.
+            # for this purpose, as single dummy lora is sufficient.
+            dummy_lora_requests = self._add_dummy_loras(num_loras=1)
+            assert len(dummy_lora_requests) == 1
+            dummy_lora_request = dummy_lora_requests[0]
+            dummy_lora_id = dummy_lora_request.lora_int_id
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: capture MLUGraph by given batch list
+        '''
+        # Update cuda graph capture sizes if environ exists.
+        mlu_graph_capture_list = os.getenv("MLU_GRAPH_CAPTURE_LIST", None)
+        if mlu_graph_capture_list:
+            self.vllm_config.compilation_config.cudagraph_capture_sizes = \
+                [int(x) for x in mlu_graph_capture_list.split(",")]
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        with self.attn_state.graph_capture(max_batch_size), mlu_graph_capture(
+                self.device) as graph_capture_context:
+            # NOTE: Capturing the largest batch size first may help reduce the
+            # memory usage of MLU graph.
+            for virtual_engine in range(
+                    self.parallel_config.pipeline_parallel_size):
+                # Only rank 0 should print progress bar during capture
+                cudagraph_capture_sizes = (tqdm(
+                    self.vllm_config.compilation_config.
+                    cudagraph_capture_sizes,
+                    desc="Capturing MLU graph shapes",
+                ) if get_tensor_model_parallel_rank() == 0 else
+                                           self.vllm_config.compilation_config.
+                                           cudagraph_capture_sizes)
+                for batch_size in cudagraph_capture_sizes:
+                    attn_metadata = (
+                        self.attn_state.graph_capture_get_metadata_for_batch(
+                            batch_size,
+                            is_encoder_decoder_model=self.model_config.
+                            is_encoder_decoder))
+                    # Disable KV Scale Calculation for graph capture
+                    attn_metadata.enable_kv_scales_calculation = False
+                    if self.lora_config:
+                        lora_mapping = LoRAMapping(
+                            **dict(index_mapping=[dummy_lora_id] * batch_size,
+                                   prompt_mapping=[dummy_lora_id] * batch_size,
+                                   is_prefill=False))
+                        self.set_active_loras(set([dummy_lora_request]),
+                                              lora_mapping)
+
+                    if self.prompt_adapter_config:
+                        prompt_adapter_mapping = PromptAdapterMapping(
+                            [-1] * batch_size,
+                            [-1] * batch_size,
+                        )
+                        self.set_active_prompt_adapters(
+                            set(), prompt_adapter_mapping)
+                    graph_runner = MLUGraphRunner(
+                        self.model, self.attn_backend.get_name(),
+                        self.attn_state.graph_clone(batch_size),
+                        self.model_config.is_encoder_decoder)
+
+                    capture_inputs = {
+                        "input_ids":
+                        input_tokens[:batch_size],
+                        "positions":
+                        input_positions[..., :batch_size],
+                        "intermediate_inputs":
+                        intermediate_inputs[:batch_size]
+                        if intermediate_inputs is not None else None,
+                        "kv_caches":
+                        kv_caches[virtual_engine],
+                        "attn_metadata":
+                        attn_metadata,
+                        "memory_pool":
+                        self.graph_memory_pool,
+                        "stream":
+                        graph_capture_context.stream
+                    }
+                    if previous_hidden_states is not None:
+                        capture_inputs[
+                            "previous_hidden_states"] = previous_hidden_states[:
+                                                                               batch_size]
+
+                    if self.has_inner_state:
+                        # Only used by Mamba-based models CUDA graph atm (Jamba)
+                        capture_inputs.update({
+                            "seqlen_agnostic_capture_inputs":
+                            self.model.get_seqlen_agnostic_capture_inputs(
+                                batch_size)
+                        })
+                    if self.model_config.is_encoder_decoder:
+                        # add the additional inputs to capture for
+                        # encoder-decoder models.
+                        self._update_inputs_to_capture_for_enc_dec_model(
+                            capture_inputs)
+
+                    with set_forward_context(attn_metadata, self.vllm_config,
+                                             virtual_engine):
+                        graph_runner.capture(**capture_inputs)
+                    self.graph_memory_pool = graph_runner.graph.pool()
+                    self.graph_runners[virtual_engine][batch_size] = (
+                        graph_runner)
+
+        if self.lora_config:
+            self._remove_dummy_loras()
+
+        end_time = time.perf_counter()
+        end_free_gpu_memory = torch.mlu.mem_get_info()[0]
+        elapsed_time = end_time - start_time
+        cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
+        # This usually takes < 10 seconds.
+        logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
+                    elapsed_time, cuda_graph_size / GiB_bytes)
+
+    @torch.inference_mode()
+    def capture_model_with_context(self, kv_caches: List[List[torch.Tensor]]) -> None:
+        """Cuda graph capture a model.
+
+        Note that CUDA graph's performance gain is negligible if number
+        of batched tokens are larger than 200. And since CUDA graph
+        requires fixed sized tensors, supporting large/variable batch
+        size requires high GPU memory overhead. Thus, vLLM only captures
+        decoding requests. Mixed batch (chunked prefill + decoding) or
+        prefill requests are not captured.
+
+        Since it is used for decoding-only, it assumes there's only 1 token
+        per sequence in the batch.
+
+        NOTE: Different from capture_model function, we capture both context
+        and decoder cuda graph here.
+        """
+
+        assert not self.has_inner_state, \
+            "Mlugraph does not support mamba-based models(Jamba)"
+
+        enable_context_mlugraph = True
+        ctx_graph_bs, ctx_graph_seq_len = self.model_config.get_context_mlugraph_bs_and_seq()
+        if ctx_graph_bs > self.max_batchsize_to_capture:
+            enable_context_mlugraph = False
+            logger.warning(f"Context mlugraph batch size {ctx_graph_bs} is greater "
+                        f"than max batch size to capture {self.max_batchsize_to_capture}, "
+                        "can not enable context mlugraph.")
+        if ctx_graph_seq_len > self.max_seq_len_to_capture:
+            enable_context_mlugraph = False
+            logger.warning(f"Context mlugraph sequence length {ctx_graph_seq_len} is greater "
+                        f"than max sequence length to capture: {self.max_seq_len_to_capture}, "
+                        "can not enable context mlugraph.")
+        if ctx_graph_seq_len > self.scheduler_config.max_num_batched_tokens:
+            logger.warning(f"Context mlugraph sequence length {ctx_graph_seq_len} is greater "
+                        f"than max num batched tokens to capture: {self.scheduler_config.max_num_batched_tokens}, "
+                        "can not enable context mlugraph.")
+            enable_context_mlugraph = False
+
+        if enable_context_mlugraph:
+            logger.info(f"Enable context mlugraph for batch size {ctx_graph_bs} and "
+                        f"sequence length {ctx_graph_seq_len}")
+
+        logger.info("Capturing mlugraphs for decoding. This may lead to "
+                    "unexpected consequences if the model is not static. To "
+                    "run the model in eager mode, set 'enforce_eager=True' or "
+                    "use '--enforce-eager' in the CLI. "
+                    "If out-of-memory error occurs during mlugraph capture,"
+                    " consider decreasing `gpu_memory_utilization` or "
+                    "switching to eager mode. You can also reduce the "
+                    "`max_num_seqs` as needed to decrease memory usage.")
+        start_time = time.perf_counter()
+        start_free_gpu_memory = torch.mlu.mem_get_info()[0]
+
+        # Prepare dummy inputs. These will be reused for all batch sizes.
+        max_batch_size = self.max_batchsize_to_capture
+        max_num_tokens = max_batch_size
+        if enable_context_mlugraph:
+            max_batch_size = max(max_batch_size, ctx_graph_bs)
+            max_num_tokens = max(max_num_tokens, ctx_graph_bs * ctx_graph_seq_len)
+
+        input_tokens = torch.zeros(max_num_tokens,
+                                   dtype=torch.long,
+                                   device=self.device)
+        input_positions = torch.zeros(max_num_tokens,
+                                      dtype=torch.int32,
+                                      device=self.device)
+        if self.model_config.uses_mrope:
+            input_positions = torch.tile(input_positions,
+                                         (3, 1)).mlu(device=self.device)
+        # Prepare dummy previous_hidden_states only if needed by the model.
+        # This is used by draft models such as EAGLE.
+        previous_hidden_states = None
+        if "previous_hidden_states" in inspect.signature(
+                self.model.forward).parameters:
+            previous_hidden_states = torch.empty(
+                [max_batch_size,
+                self.model_config.get_hidden_size()],
+                dtype=self.model_config.dtype,
+                device=self.device)
+
+        intermediate_inputs = None
+        if not get_pp_group().is_first_rank:
+            intermediate_inputs = self.model.make_empty_intermediate_tensors(
+                batch_size=max_num_tokens,
+                dtype=self.model_config.dtype,
+                device=self.device)
+
+        dummy_lora_id: Optional[int] = None
+        dummy_lora_request: LoRARequest = []
+        if self.lora_config:
+            # The goal is to capture the LoRA kernels in cuda graphs.
+            # for this purpose, as single dummy lora is sufficient.
+            dummy_lora_requests = self._add_dummy_loras(num_loras=1)
+            assert len(dummy_lora_requests) == 1
+            dummy_lora_request = dummy_lora_requests[0]
+            dummy_lora_id = dummy_lora_request.lora_int_id
+
+        # Update cuda graph capture sizes if environ exists.
+        mlu_graph_capture_list = os.getenv("MLU_GRAPH_CAPTURE_LIST", None)
+        if mlu_graph_capture_list:
+            self.vllm_config.compilation_config.cudagraph_capture_sizes = \
+                [int(x) for x in mlu_graph_capture_list.split(",")]
+
+        def capture_mlugraph(is_prefill, batch_size, num_tokens, attn_metadata):
+            if self.lora_config:
+                lora_mapping = LoRAMapping(
+                    **dict(index_mapping=[dummy_lora_id] * num_tokens,
+                            prompt_mapping=[dummy_lora_id] * batch_size,
+                            is_prefill=is_prefill))
+                self.set_active_loras(set([dummy_lora_request]),
+                                      lora_mapping)
+
+            if self.prompt_adapter_config:
+                prompt_adapter_mapping = PromptAdapterMapping(
+                    [-1] * num_tokens,
+                    [-1] * batch_size,
+                )
+                self.set_active_prompt_adapters(
+                    set(), prompt_adapter_mapping)
+
+            capture_inputs = {
+                "input_ids":
+                input_tokens[:num_tokens],
+                "positions":
+                input_positions[..., :num_tokens],
+                "intermediate_inputs":
+                intermediate_inputs[:num_tokens]
+                if intermediate_inputs is not None else None,
+                "kv_caches":
+                kv_caches[virtual_engine],
+                "attn_metadata":
+                attn_metadata,
+                "memory_pool":
+                self.graph_memory_pool,
+                "stream":
+                graph_capture_context.stream
+            }
+            if previous_hidden_states is not None:
+                capture_inputs["previous_hidden_states"
+                            ] = previous_hidden_states[:batch_size]
+
+            graph_runner = MLUGraphRunner(
+                self.model, self.attn_backend.get_name(),
+                self.attn_state.graph_clone(batch_size),
+                self.model_config.is_encoder_decoder)
+
+            with set_forward_context(attn_metadata, self.vllm_config,
+                                     virtual_engine):
+                graph_runner.capture(**capture_inputs)
+            self.graph_memory_pool = graph_runner.graph.pool()
+            if is_prefill:
+                self.context_graph_runners[virtual_engine] = graph_runner
+            else:
+                self.graph_runners[virtual_engine][batch_size] = (
+                    graph_runner)
+
+        with self.attn_state.graph_capture_with_context(
+                    ctx_graph_bs, max_batch_size, max_num_tokens
+                ), mlu_graph_capture(self.device) as graph_capture_context:
+            # NOTE: Capturing the largest batch size first may help reduce the
+            # memory usage of CUDA graph.
+            # NOTE: If enable context mlugraph, because the output buffer is reused and
+            # created when the first mlugraph captured, we must compare max token num for
+            # context and decoder to decide which process should be capture first.
+            for virtual_engine in range(self.parallel_config.pipeline_parallel_size):
+                # capture mlugraph for context when hiddens_states size(bs * seq_len)
+                # is bigger than decoder(max bs size in bs capture list)
+                if enable_context_mlugraph and (ctx_graph_bs * ctx_graph_seq_len
+                                                >= self.max_batchsize_to_capture):
+                    self.attn_state.fill_seq_lens_tensor(ctx_graph_seq_len)
+                    context_attn_metadata = (
+                        self.attn_state.graph_capture_get_metadata_for_context(
+                        ctx_graph_bs, ctx_graph_seq_len))
+                    capture_mlugraph(True, ctx_graph_bs, ctx_graph_bs * ctx_graph_seq_len,
+                                    context_attn_metadata)
+
+                # capture mlugraph for decoder
+                self.attn_state.fill_seq_lens_tensor(1)
+                # Only rank 0 should print progress bar during capture
+                cudagraph_capture_sizes = (tqdm(
+                    self.vllm_config.compilation_config.
+                    cudagraph_capture_sizes,
+                    desc="Capturing MLU graph shapes",
+                ) if get_tensor_model_parallel_rank() == 0 else
+                                           self.vllm_config.compilation_config.
+                                           cudagraph_capture_sizes)
+                for batch_size in cudagraph_capture_sizes:
+                    decoder_attn_metadata = (
+                        self.attn_state.graph_capture_get_metadata_for_batch(batch_size))
+                    capture_mlugraph(False, batch_size, batch_size, decoder_attn_metadata)
+
+                # capture mlugraph for context when hiddens_states size(bs * seq_len)
+                # is smaller than decoder(max bs size in bs capture list)
+                if enable_context_mlugraph and (ctx_graph_bs * ctx_graph_seq_len
+                                                < self.max_batchsize_to_capture):
+                    self.attn_state.fill_seq_lens_tensor(ctx_graph_seq_len)
+                    context_attn_metadata = (
+                        self.attn_state.graph_capture_get_metadata_for_context(
+                        ctx_graph_bs, ctx_graph_seq_len))
+                    capture_mlugraph(True, ctx_graph_bs, ctx_graph_bs * ctx_graph_seq_len,
+                                    context_attn_metadata)
+
+        if self.lora_config:
+            self._remove_dummy_loras()
+
+        end_time = time.perf_counter()
+        end_free_gpu_memory = torch.mlu.mem_get_info()[0]
+        elapsed_time = end_time - start_time
+        cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
+        # This usually takes < 10 seconds.
+        logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
+                    elapsed_time, cuda_graph_size / GiB_bytes)
+
+    def reset_capture_context(self):
+        self.graph_runners = [
+            {} for _ in range(self.parallel_config.pipeline_parallel_size)
+        ]
+        self.context_graph_runners = [
+            None for _ in range(self.parallel_config.pipeline_parallel_size)
+        ]
+        self.graph_memory_pool = None  # Set during graph capture.
+
+        gc.collect()
+        torch.mlu.empty_cache()
+
+
+class MLUModelRunner(MLUModelRunnerBase, ModelRunner):
+    """
+    MLU model runner with sampling step.
+    """
+    _builder_cls: Type[ModelInputForMLUBuilder] = ModelInputForMLUBuilder
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        model_input: ModelInputForGPUWithSamplingMetadata,
+        kv_caches: List[torch.Tensor],
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        num_steps: int = 1,
+        **kwargs,
+    ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:
+        if num_steps > 1:
+            raise ValueError("num_steps > 1 is not supported in ModelRunner")
+
+        if self.lora_config:
+            assert model_input.lora_requests is not None
+            assert model_input.lora_mapping is not None
+            self.set_active_loras(model_input.lora_requests,
+                                  model_input.lora_mapping)
+
+        if self.prompt_adapter_config:
+            assert model_input.prompt_adapter_requests is not None
+            assert model_input.prompt_adapter_mapping is not None
+            self.set_active_prompt_adapters(
+                model_input.prompt_adapter_requests,
+                model_input.prompt_adapter_mapping)
+
+        self.attn_state.begin_forward(model_input)
+
+        # Currently cuda graph is only supported by the decode phase.
+        assert model_input.attn_metadata is not None
+        prefill_meta = model_input.attn_metadata.prefill_metadata
+        decode_meta = model_input.attn_metadata.decode_metadata
+        # TODO(andoorve): We can remove this once all
+        # virtual engines share the same kv cache.
+        virtual_engine = model_input.virtual_engine
+        previous_hidden_states = kwargs.get("previous_hidden_states")
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: support context mlugraph
+        '''
+        use_context_mlugraph = False
+        if prefill_meta is None and decode_meta.use_cuda_graph:
+            assert model_input.input_tokens is not None
+            graph_batch_size = model_input.input_tokens.shape[0]
+            model_executable = self.graph_runners[virtual_engine][
+                graph_batch_size]
+            if previous_hidden_states is not None:
+                previous_hidden_states = torch.cat([
+                    previous_hidden_states,
+                    torch.empty([
+                        graph_batch_size - previous_hidden_states.shape[0],
+                        *previous_hidden_states.shape[1:]
+                    ],
+                                dtype=previous_hidden_states.dtype,
+                                device=previous_hidden_states.device)
+                ])
+        # The context mlugraph is None when profile run
+        elif (decode_meta is None and prefill_meta.use_cuda_graph
+              and self.context_graph_runners[virtual_engine] is not None):
+            use_context_mlugraph = True
+            model_executable = self.context_graph_runners[virtual_engine]
+        else:
+            model_executable = self.model
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        # Receive KV cache in distributed KV cache transfer setting
+        # In disagg prefill setting, it will also recv hidden states and bypass
+        # model forwarding
+        # In KV cache database setting, it will change the model input so that
+        # we can skip prefilling on tokens that successfully received KV caches
+        # NOTE: The receive operation is blocking
+        bypass_model_exec = False
+        if self.need_recv_kv(model_input, kv_caches):
+            hidden_or_intermediate_states, bypass_model_exec, model_input = \
+                get_kv_transfer_group().recv_kv_caches_and_hidden_states(
+                    # model is used to know which layer the current worker
+                    # is working on, so that we can receive KV for only those
+                    # layers.
+                    model_executable,
+                    model_input,
+                    kv_caches=kv_caches
+                )
+
+        multi_modal_kwargs = model_input.multi_modal_kwargs or {}
+        seqlen_agnostic_kwargs = {
+            "finished_requests_ids": model_input.finished_requests_ids,
+            "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
+        } if self.has_inner_state else {}
+        model_kwargs = {}
+        if previous_hidden_states is not None:
+            model_kwargs["previous_hidden_states"] = previous_hidden_states
+        if (self.observability_config is not None
+                and self.observability_config.collect_model_forward_time):
+            model_forward_start = torch.mlu.Event(enable_timing=True)
+            model_forward_end = torch.mlu.Event(enable_timing=True)
+            model_forward_start.record()
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: add mlu metrics
+        '''
+        # Add time markers for model_executable+compute_logits
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+            use_cuda_graph = ((prefill_meta is None and decode_meta.use_cuda_graph)
+                              or use_context_mlugraph)
+            # if use_cuda_graph, the start timestamp will be inserted inside MLUGraphRunner.forward()
+            if not use_cuda_graph:
+                start = torch.mlu.Event(enable_timing=True)
+                start.record()
+
+        if not bypass_model_exec:
+            with set_forward_context(model_input.attn_metadata,
+                                     self.vllm_config, virtual_engine):
+                hidden_or_intermediate_states = model_executable(
+                    input_ids=model_input.input_tokens,
+                    positions=model_input.input_positions,
+                    intermediate_tensors=intermediate_tensors,
+                    **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
+                                                 device=self.device),
+                    **seqlen_agnostic_kwargs,
+                    **model_kwargs,
+                )
+
+        if (self.observability_config is not None
+                and self.observability_config.collect_model_forward_time):
+            model_forward_end.record()
+
+        # Sending KV cache in distributed KV cache transfer setting
+        # NOTE: the send operation is non-blocking
+        if self.need_send_kv(model_input, kv_caches):
+            get_kv_transfer_group().send_kv_caches_and_hidden_states(
+                # model_executable is used to know which layer the current
+                # worker is working on, so that we can send KV for only those
+                # layers.
+                model_executable,
+                model_input,
+                kv_caches,
+                hidden_or_intermediate_states,
+            )
+
+        # Compute the logits in the last pipeline stage.
+        if not get_pp_group().is_last_rank:
+            if (self.is_driver_worker
+                    and hidden_or_intermediate_states is not None
+                    and isinstance(hidden_or_intermediate_states,
+                                   IntermediateTensors)
+                    and self.observability_config is not None
+                    and self.observability_config.collect_model_forward_time):
+                model_forward_end.synchronize()
+                model_forward_time = model_forward_start.elapsed_time(
+                    model_forward_end)
+                orig_model_forward_time = 0.0
+                if intermediate_tensors is not None:
+                    orig_model_forward_time = intermediate_tensors.tensors.get(
+                        "model_forward_time", torch.tensor(0.0)).item()
+                hidden_or_intermediate_states.tensors["model_forward_time"] = (
+                    torch.tensor(model_forward_time + orig_model_forward_time))
+            return hidden_or_intermediate_states
+
+        logits = self.model.compute_logits(hidden_or_intermediate_states,
+                                           model_input.sampling_metadata)
+
+        # Add time markers for model_executable+compute_logits
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+            end_marker = torch.mlu.Event(enable_timing=True)
+            end_marker.record()
+            if use_cuda_graph:
+                self.time_markers = (model_executable.start, end_marker)
+            else:
+                self.time_markers = (start, end_marker)
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        if not self.is_driver_worker:
+            return []
+
+        if model_input.async_callback is not None:
+            model_input.async_callback()
+
+        # Sample the next token.
+        output: SamplerOutput = self.model.sample(
+            logits=logits,
+            sampling_metadata=model_input.sampling_metadata,
+        )
+        if (self.observability_config is not None
+                and self.observability_config.collect_model_forward_time
+                and output is not None):
+            model_forward_end.synchronize()
+            model_forward_time = model_forward_start.elapsed_time(
+                model_forward_end)
+            orig_model_forward_time = 0.0
+            if intermediate_tensors is not None:
+                orig_model_forward_time = intermediate_tensors.tensors.get(
+                    "model_forward_time", torch.tensor(0.0)).item()
+            # If there are multiple workers, we are still tracking the latency
+            # from the start time of the driver worker to the end time of the
+            # driver worker. The model forward time will then end up covering
+            # the communication time as well.
+            output.model_forward_time = (orig_model_forward_time +
+                                         model_forward_time)
+
+        if self.return_hidden_states:
+            # we only need to pass hidden states of most recent token
+            assert model_input.sampling_metadata is not None
+            indices = model_input.sampling_metadata.selected_token_indices
+            if model_input.is_prompt:
+                hidden_states = hidden_or_intermediate_states.index_select(
+                    0, indices)
+                output.prefill_hidden_states = hidden_or_intermediate_states
+            elif decode_meta.use_cuda_graph:
+                hidden_states = hidden_or_intermediate_states[:len(indices)]
+            else:
+                hidden_states = hidden_or_intermediate_states
+
+            output.hidden_states = hidden_states
+
+        return [output]
+
+
+class MLUGraphRunner(CUDAGraphRunner):
+
+    def capture(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_inputs: Optional[IntermediateTensors],
+        kv_caches: List[torch.Tensor],
+        attn_metadata: AttentionMetadata,
+        memory_pool: Optional[Tuple[int, int]],
+        stream: torch.mlu.Stream,
+        **kwargs,
+    ):
+        assert self._graph is None
+        # Run the model a few times without capturing the graph.
+        # This is to make sure that the captured graph does not include the
+        # kernel launches for initial benchmarking (e.g., Triton autotune).
+        # Note one iteration is not enough for torch.compile
+        for _ in range(_NUM_WARMUP_ITERS):
+            self.model(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_inputs,
+                **kwargs,
+            )
+        # Wait for the warm up operations to finish before proceeding with
+        # Graph Capture.
+        torch.mlu.synchronize()
+        # Capture the graph.
+        self._graph = torch.mlu.MLUGraph()
+        with torch.mlu.graph(self._graph, pool=memory_pool, stream=stream):
+            output_hidden_or_intermediate_states = self.model(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_inputs,
+                **kwargs,
+            )
+
+            hidden_or_intermediate_states = (
+                output_hidden_or_intermediate_states
+            )
+
+            del output_hidden_or_intermediate_states
+            # make sure `output_hidden_or_intermediate_states` is deleted
+            # in the graph's memory pool
+            gc.collect()
+        torch.mlu.synchronize()
+
+        # Save the input and output buffers.
+        self.input_buffers = {
+            "input_ids":
+            input_ids,
+            "positions":
+            positions,
+            "kv_caches":
+            kv_caches,
+            **self.attn_state.get_graph_input_buffers(
+                attn_metadata, self._is_encoder_decoder_model),
+            **kwargs,
+        }
+        if intermediate_inputs is not None:
+            self.input_buffers.update(intermediate_inputs.tensors)
+        if get_pp_group().is_last_rank:
+            self.output_buffers = {
+                "hidden_states": hidden_or_intermediate_states
+            }
+        else:
+            self.output_buffers = hidden_or_intermediate_states
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors],
+        **kwargs,
+    ) -> torch.Tensor:
+        attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
+
+        # Copy the input tensors to the input buffers.
+        self.input_buffers["input_ids"].copy_(input_ids, non_blocking=True)
+        if positions is not None:
+            # in some case like MLA, it will reuse positions in metadata
+            # but truncate them to the original size
+            # so the shape is not padded, we need to copy partial only
+            self.input_buffers["positions"][:positions.shape[0]].copy_(
+                positions, non_blocking=True)
+
+        if self.backend_name != "NO_ATTENTION":
+            self.input_buffers["slot_mapping"].copy_(
+                attn_metadata.slot_mapping, non_blocking=True)
+
+        self.attn_state.prepare_graph_input_buffers(
+            self.input_buffers, attn_metadata, self._is_encoder_decoder_model)
+
+        if "seqlen_agnostic_capture_inputs" in self.input_buffers:
+            self.model.copy_inputs_before_cuda_graphs(self.input_buffers,
+                                                      **kwargs)
+
+        if "previous_hidden_states" in self.input_buffers:
+            self.input_buffers["previous_hidden_states"].copy_(
+                kwargs["previous_hidden_states"], non_blocking=True)
+
+        if intermediate_tensors is not None:
+            for key in intermediate_tensors.tensors:
+                if key != "model_execute_time" and key != "model_forward_time":
+                    self.input_buffers[key].copy_(intermediate_tensors[key],
+                                                  non_blocking=True)
+        if self._is_encoder_decoder_model:
+            self.input_buffers["encoder_input_ids"].copy_(
+                kwargs['encoder_input_ids'], non_blocking=True)
+            self.input_buffers["encoder_positions"].copy_(
+                kwargs['encoder_positions'], non_blocking=True)
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: add mlu metrics
+        '''
+        # Add time markers for MLUGraph mode
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+            self.start = torch.mlu.Event(enable_timing=True)
+            self.start.record()
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+        # Run the graph.
+        self.graph.replay()
+        # Return the output tensor.
+        if get_pp_group().is_last_rank:
+            return self.output_buffers["hidden_states"]
+
+        return self.output_buffers
diff --git a/vllm_mlu/vllm_mlu/worker/mlu_pooling_model_runner.py b/vllm_mlu/vllm_mlu/worker/mlu_pooling_model_runner.py
new file mode 100644
index 000000000..065db23c6
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/worker/mlu_pooling_model_runner.py
@@ -0,0 +1,162 @@
+from typing import List, Optional, Type, Union
+
+import torch
+
+from vllm.distributed import get_pp_group
+from vllm.forward_context import set_forward_context
+from vllm.logger import init_logger
+from vllm.multimodal import MultiModalKwargs
+from vllm.sequence import IntermediateTensors, PoolerOutput
+from vllm.worker.pooling_model_runner import (ModelInputForGPUWithPoolingMetadata,
+                                              PoolingModelRunner)
+
+from vllm_mlu.worker.mlu_model_runner import (MLUModelRunnerBase,
+                                              ModelInputForMLUBuilder)
+from vllm_mlu._mlu_utils import VLLM_LATENCY_DEBUG_WITH_DEVICE_EN
+
+logger = init_logger(__name__)
+
+
+class MLUPoolingModelRunner(PoolingModelRunner,
+                            MLUModelRunnerBase[ModelInputForGPUWithPoolingMetadata]):
+    _model_input_cls: Type[ModelInputForGPUWithPoolingMetadata] = (
+        ModelInputForGPUWithPoolingMetadata)
+    _builder_cls: Type[ModelInputForMLUBuilder] = ModelInputForMLUBuilder
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        model_input: ModelInputForGPUWithPoolingMetadata,
+        kv_caches: List[torch.Tensor],
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        num_steps: int = 1,
+    ) -> Optional[Union[List[PoolerOutput], IntermediateTensors]]:
+        if num_steps > 1:
+            raise ValueError(
+                "PoolingModelRunner does not support multi-step execution.")
+
+        if self.lora_config:
+            assert model_input.lora_requests is not None
+            assert model_input.lora_mapping is not None
+            self.set_active_loras(model_input.lora_requests,
+                                  model_input.lora_mapping)
+
+        if self.prompt_adapter_config:
+            assert model_input.prompt_adapter_requests is not None
+            assert model_input.prompt_adapter_mapping is not None
+            self.set_active_prompt_adapters(
+                model_input.prompt_adapter_requests,
+                model_input.prompt_adapter_mapping)
+
+        # Currently cuda graph is only supported by the decode phase.
+        assert model_input.attn_metadata is not None
+        prefill_meta = model_input.attn_metadata.prefill_metadata
+        decode_meta = model_input.attn_metadata.decode_metadata
+        virtual_engine = model_input.virtual_engine
+        # Pooling models are (ab-)used also to integrate non text models that
+        # are not autoregressive (PrithviGeosaptialMAE).
+        # These model might not use attention and do not really have a prefill
+        # and decode phase. The model input is processed in one shot and both
+        # decode_metadata and prefill_metadata would be None for such models.
+        # See the PlaceholderAttentionMetadata class.
+        # TODO: Figure out if cuda_graph is of any use for these models and
+        #  explore how to leverage it.
+        if (prefill_meta is None and decode_meta is not None
+                and decode_meta.use_cuda_graph):
+            assert model_input.input_tokens is not None
+            graph_batch_size = model_input.input_tokens.shape[0]
+            model_executable = self.graph_runners[virtual_engine][
+                graph_batch_size]
+        else:
+            model_executable = self.model
+
+        multi_modal_kwargs = model_input.multi_modal_kwargs or {}
+        seqlen_agnostic_kwargs = {
+            "finished_requests_ids": model_input.finished_requests_ids,
+            "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
+        } if self.has_inner_state else {}
+        if (self.observability_config is not None
+                and self.observability_config.collect_model_forward_time):
+            model_forward_start = torch.mlu.Event(enable_timing=True)
+            model_forward_end = torch.mlu.Event(enable_timing=True)
+            model_forward_start.record()
+
+        cross_enc_kwargs = {}
+        if model_input.token_types is not None:
+            cross_enc_kwargs["token_type_ids"] = model_input.token_types
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: add mlu metrics
+        '''
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+            start_marker = torch.mlu.Event(enable_timing=True)
+            start_marker.record()
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        with set_forward_context(model_input.attn_metadata, self.vllm_config,
+                                 virtual_engine):
+            hidden_or_intermediate_states = model_executable(
+                input_ids=model_input.input_tokens,
+                positions=model_input.input_positions,
+                intermediate_tensors=intermediate_tensors,
+                **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
+                                             device=self.device),
+                **cross_enc_kwargs,
+                **seqlen_agnostic_kwargs)
+
+        if (self.observability_config is not None
+                and self.observability_config.collect_model_forward_time):
+            model_forward_end.record()
+
+        # Only perform pooling in the last pipeline stage.
+        if not get_pp_group().is_last_rank:
+            if (self.is_driver_worker
+                    and hidden_or_intermediate_states is not None
+                    and isinstance(hidden_or_intermediate_states,
+                                   IntermediateTensors)
+                    and self.observability_config is not None
+                    and self.observability_config.collect_model_forward_time):
+                model_forward_end.synchronize()
+                model_forward_time = model_forward_start.elapsed_time(
+                    model_forward_end)
+                orig_model_forward_time = 0.0
+                if intermediate_tensors is not None:
+                    orig_model_forward_time = intermediate_tensors.tensors.get(
+                        "model_forward_time", torch.tensor(0.0)).item()
+                hidden_or_intermediate_states.tensors["model_forward_time"] = (
+                    torch.tensor(model_forward_time + orig_model_forward_time))
+            return hidden_or_intermediate_states
+
+        # Only perform pooling in the driver worker.
+        if not self.is_driver_worker:
+            return []
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: add mlu metrics 
+        '''
+        embedding_result = [
+            self.model.pooler(hidden_states=hidden_or_intermediate_states,
+                              pooling_metadata=model_input.pooling_metadata)
+        ]
+
+        if VLLM_LATENCY_DEBUG_WITH_DEVICE_EN:
+            end_marker = torch.mlu.Event(enable_timing=True)
+            end_marker.record()
+            self.time_markers = (start_marker, end_marker)
+
+        return embedding_result
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/worker/mlu_worker.py b/vllm_mlu/vllm_mlu/worker/mlu_worker.py
new file mode 100644
index 000000000..24d22b7ce
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/worker/mlu_worker.py
@@ -0,0 +1,565 @@
+"""A MLU worker class."""
+import gc
+import os
+import contextlib
+import time
+import functools
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Type, Generator
+from dataclasses import dataclass, field
+
+import torch
+import torch.distributed
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.distributed import (ensure_kv_transfer_initialized,
+                              ensure_model_parallel_initialized,
+                              init_distributed_environment,
+                              set_custom_all_reduce)
+from vllm.model_executor import set_random_seed
+from vllm.model_executor.layers.vocab_parallel_embedding import (VocabParallelEmbedding,
+                                                                 ParallelLMHead)
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear)
+from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLAAttention
+from vllm.platforms import current_platform
+from vllm.sequence import SequenceGroupMetadata
+from vllm.utils import GiB_bytes
+from vllm.worker.cache_engine import CacheEngine
+from vllm.worker.worker_base import WorkerBase
+from vllm.worker.worker import Worker
+from vllm.logger import init_logger
+
+from vllm_mlu.model_executor.layers.feed_forward import FeedForward
+from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
+from vllm_mlu.worker.mlu_pooling_model_runner import MLUPoolingModelRunner
+from vllm_mlu.worker.mlu_enc_dec_model_runner import MLUEncoderDecoderModelRunner
+from vllm_mlu.worker.mlu_model_runner import MLUModelRunnerBase, MLUModelRunner
+
+logger = init_logger(__name__)
+
+
+def default_act_range_value():
+    return {
+        "x": None,
+        "split": None,
+        "is_linear": False,
+        "is_qkv": False,
+        "q_proj_size": 0,
+        "num_kv_head_replicas": 1,
+        "is_merge": False,
+        "input_id": []
+    }
+
+
+@dataclass
+class MLUMemorySnapshot:
+    """Memory snapshot."""
+    torch_peak: int = 0
+    mlu_memory: int = 0
+    torch_memory: int = 0
+    non_torch_memory: int = 0
+    timestamp: float = 0.0
+    auto_measure: bool = True
+
+    def __post_init__(self):
+        if self.auto_measure:
+            self.measure()
+
+    def measure(self):
+        # we measure the torch peak memory usage via allocated_bytes,
+        # rather than `torch.mlu.memory_reserved()` .
+        # After `torch.mlu.reset_peak_memory_stats()`,
+        # `torch.mlu.memory_reserved()` will keep growing, and only shrink
+        # when we call `torch.mlu.empty_cache()` or OOM happens.
+        self.torch_peak = torch.mlu.memory_stats().get(
+            "allocated_bytes.all.peak", 0)
+
+        self.mlu_memory = torch.mlu.mem_get_info(
+        )[1] - torch.mlu.mem_get_info()[0]
+
+        # torch.mlu.memory_reserved() is how many bytes
+        # PyTorch gets from mlu (by calling mluMalloc, etc.)
+        # this is used to measure the non-torch memory usage
+        self.torch_memory = torch.mlu.memory_reserved()
+
+        self.non_torch_memory = self.mlu_memory - self.torch_memory
+        self.timestamp = time.time()
+
+    def __sub__(self, other: "MLUMemorySnapshot") -> "MLUMemorySnapshot":
+        return MLUMemorySnapshot(
+            torch_peak=self.torch_peak - other.torch_peak,
+            mlu_memory=self.mlu_memory - other.mlu_memory,
+            torch_memory=self.torch_memory - other.torch_memory,
+            non_torch_memory=self.non_torch_memory - other.non_torch_memory,
+            timestamp=self.timestamp - other.timestamp,
+            auto_measure=False,
+        )
+
+
+@dataclass
+class MLUMemoryProfilingResult:
+    """Memory profiling result. All numbers are in bytes.
+    """
+    non_kv_cache_memory: int = 0
+    torch_peak_increase: int = 0
+    non_torch_increase: int = 0
+    weights_memory: float = 0
+    before_create: MLUMemorySnapshot = field(default_factory=MLUMemorySnapshot)
+    before_profile: MLUMemorySnapshot = field(default_factory=MLUMemorySnapshot)
+    after_profile: MLUMemorySnapshot = field(default_factory=MLUMemorySnapshot)
+    profile_time: float = 0.0
+
+
+@contextlib.contextmanager
+def mlu_memory_profiling(
+        baseline_snapshot: MLUMemorySnapshot,
+        weights_memory: int) -> Generator[MLUMemoryProfilingResult, None, None]:
+    gc.collect()
+    torch.mlu.empty_cache()
+    torch.mlu.reset_peak_memory_stats()
+
+    result = MLUMemoryProfilingResult()
+
+    result.before_create = baseline_snapshot
+    # the part of memory used for holding the model weights
+    result.weights_memory = weights_memory
+
+    result.before_profile.measure()
+
+    yield result
+
+    gc.collect()
+    torch.mlu.empty_cache()
+
+    result.after_profile.measure()
+
+    diff_profile = result.after_profile - result.before_profile
+    diff_from_create = result.after_profile - result.before_create
+    result.torch_peak_increase = diff_profile.torch_peak
+    result.non_torch_increase = diff_from_create.non_torch_memory
+    result.profile_time = diff_profile.timestamp
+    result.non_kv_cache_memory = result.non_torch_increase + result.torch_peak_increase + result.weights_memory  # noqa
+
+
+class MLUWorker(Worker):
+    """A worker class that executes (a partition of) the model on a GPU.
+
+    Each worker is associated with a single GPU. The worker is responsible for
+    maintaining the KV cache and executing the model on the GPU. In case of
+    distributed inference, each worker is assigned a partition of the model.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        local_rank: int,
+        rank: int,
+        distributed_init_method: str,
+        is_driver_worker: bool = False,
+        model_runner_cls: Optional[Type[MLUModelRunnerBase]] = None,
+    ) -> None:
+        WorkerBase.__init__(self, vllm_config)
+        self.parallel_config.rank = rank
+        self.local_rank = local_rank
+        self.rank = rank
+        self.distributed_init_method = distributed_init_method
+        self.is_driver_worker = is_driver_worker
+        if self.model_config.trust_remote_code:
+            # note: lazy import to avoid importing torch before initializing
+            from vllm.utils import init_cached_hf_modules
+            init_cached_hf_modules()
+
+        # Return hidden states from target model if the draft model is an
+        # mlp_speculator
+        speculative_config = self.speculative_config
+        model_config = self.model_config
+        speculative_args = {} if speculative_config is None \
+            or (speculative_config.draft_model_config.hf_config.model_type ==
+                model_config.hf_config.model_type) \
+            or (speculative_config.draft_model_config.hf_config.model_type
+                not in ("medusa", "mlp_speculator", "eagle", "deepseek_mtp")) \
+                    else {"return_hidden_states": True}
+
+        ModelRunnerClass: Type[MLUModelRunnerBase] = MLUModelRunner
+        if model_config.runner_type == "pooling":
+            ModelRunnerClass = MLUPoolingModelRunner
+        elif self.model_config.is_encoder_decoder:
+            ModelRunnerClass = MLUEncoderDecoderModelRunner
+        self.model_runner: MLUModelRunnerBase = ModelRunnerClass(
+            vllm_config=self.vllm_config,
+            kv_cache_dtype=self.cache_config.cache_dtype,
+            is_driver_worker=is_driver_worker,
+            **speculative_args,
+        )
+        if model_runner_cls is not None:
+            self.model_runner = model_runner_cls(self.model_runner)
+
+        # Uninitialized cache engine. Will be initialized by
+        # initialize_cache.
+        self.cache_engine: List[CacheEngine]
+        # Initialize gpu_cache as pooling models don't initialize kv_caches
+        self.gpu_cache: Optional[List[List[torch.Tensor]]] = None
+        self._seq_group_metadata_cache: Dict[str, SequenceGroupMetadata] = {}
+
+        # Torch profiler. Enabled and configured through env vars:
+        # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
+        if envs.VLLM_TORCH_PROFILER_DIR:
+            torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
+            logger.info("Profiling enabled. Traces will be saved to: %s",
+                        torch_profiler_trace_dir)
+            self.profiler = torch.profiler.profile(
+                activities=[
+                    torch.profiler.ProfilerActivity.CPU,
+                    torch.profiler.ProfilerActivity.MLU,
+                ],
+                with_stack=True,
+                on_trace_ready=torch.profiler.tensorboard_trace_handler(
+                    torch_profiler_trace_dir, use_gzip=True))
+        else:
+            self.profiler = None
+
+    def init_device(self) -> None:
+        if self.device_config.device.type == "mlu":
+            # torch.distributed.all_reduce does not free the input tensor until
+            # the synchronization point. This causes the memory usage to grow
+            # as the number of all_reduce calls increases. This env var disables
+            # this behavior.
+            # Related issue:
+            # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
+            os.environ["TORCH_CNCL_AVOID_RECORD_STREAMS"] = "1"
+
+            # This env var set by Ray causes exceptions with graph building.
+            os.environ.pop("CNCL_ASYNC_ERROR_HANDLING", None)
+            self.device = torch.device(f"mlu:{self.local_rank}")
+            torch.mlu.set_device(self.device)
+
+            _check_if_gpu_supports_dtype(self.model_config.dtype)
+            gc.collect()
+            torch.mlu.empty_cache()
+            torch.mlu.reset_peak_memory_stats()
+            self.baseline_snapshot = MLUMemorySnapshot()
+        else:
+            raise RuntimeError(
+                f"Not support device type: {self.device_config.device}")
+        # Initialize the distributed environment.
+        init_worker_distributed_environment(self.vllm_config, self.rank,
+                                            self.distributed_init_method,
+                                            self.local_rank)
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+    @torch.inference_mode()
+    def determine_num_available_blocks(self) -> Tuple[int, int]:
+        """Profiles the peak memory usage of the model to determine how many
+        KV blocks may be allocated without OOMs.
+
+        The engine will first conduct a profiling of the existing memory usage.
+        Then, it calculate the maximum possible number of GPU and CPU blocks
+        that can be allocated with the remaining free memory.
+
+        .. tip::
+            You may limit the usage of GPU memory
+            by adjusting the `gpu_memory_utilization` parameter.
+        """
+        # Profile the memory usage of the model and get the maximum number of
+        # cache blocks that can be allocated with the remaining free memory.
+        torch.mlu.empty_cache()
+        torch.mlu.reset_peak_memory_stats()
+
+        free_memory_pre_profile, total_gpu_memory = torch.mlu.mem_get_info()
+
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model.
+        with mlu_memory_profiling(
+                self.baseline_snapshot,
+                weights_memory=self.model_runner.model_memory_usage) as result:
+            self.model_runner.profile_run()
+
+        self._assert_memory_footprint_increased_during_profiling()
+
+        memory_for_current_instance = total_gpu_memory * \
+            self.cache_config.gpu_memory_utilization
+        available_kv_cache_memory = (memory_for_current_instance -
+                                     result.non_kv_cache_memory)
+
+        # Calculate the number of blocks that can be allocated with the
+        # profiled peak memory.
+        cache_block_size = self.get_cache_block_size_bytes()
+        if cache_block_size == 0:
+            num_gpu_blocks = 0
+            num_cpu_blocks = 0
+        else:
+            num_gpu_blocks = int(available_kv_cache_memory // cache_block_size)
+            num_cpu_blocks = int(self.cache_config.swap_space_bytes //
+                                 cache_block_size)
+        num_gpu_blocks = max(num_gpu_blocks, 0)
+        num_cpu_blocks = max(num_cpu_blocks, 0)
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: Fix: tmo.reshape_paged_cache limit: RuntimeError: The addressing range of kv_cache cannot exceed 4G
+        @brief: Fix: cnnlScaledDotProductAttn_v5 check failed: cnnlGetTensorElementNum(key_desc) * kv_cache_dbyte <= INT32_MAX
+        '''
+        max_num_gpu_blocks = CacheEngine.get_max_num_gpu_blocks(self.cache_config, self.model_config,
+                                                                self.parallel_config)
+
+        if num_gpu_blocks > max_num_gpu_blocks:
+            logger.warning(f"current cache block num {num_gpu_blocks} is greater than tmo op limit, "
+                           f"force reduce cache block num to {max_num_gpu_blocks}.")
+
+        num_gpu_blocks = min(num_gpu_blocks, max_num_gpu_blocks)
+        available_kv_cache_memory = num_gpu_blocks * cache_block_size
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        msg = (f"Memory profiling takes {result.profile_time:.2f} seconds\n"
+               "the current vLLM instance can use "
+               "total_gpu_memory "
+               f"({(total_gpu_memory / GiB_bytes):.2f}GiB)"
+               " x gpu_memory_utilization "
+               f"({self.cache_config.gpu_memory_utilization:.2f})"
+               f" = {(memory_for_current_instance / GiB_bytes):.2f}GiB\n"
+               "model weights take "
+               f"{(result.weights_memory / GiB_bytes):.2f}GiB;"
+               " non_torch_memory takes "
+               f"{(result.non_torch_increase / GiB_bytes):.2f}GiB;"
+               " PyTorch activation peak memory takes "
+               f"{(result.torch_peak_increase / GiB_bytes):.2f}GiB;"
+               " the rest of the memory reserved for KV Cache is "
+               f"{(available_kv_cache_memory / GiB_bytes):.2f}GiB.")
+
+        logger.info(msg)
+        # Final cleanup
+        gc.collect()
+
+        '''
+        =============================
+        Modify by vllm_mlu
+        =============================
+        @brief: record init memory usage
+        ''' 
+        # Record memory usage
+        self.peak_memory = result.torch_peak_increase
+        self.block_memory = available_kv_cache_memory
+        self.num_gpu_blocks = num_gpu_blocks
+        self.num_cpu_blocks = num_cpu_blocks
+        '''
+        ==================
+        End of MLU Hijack
+        ==================
+        '''
+
+        return num_gpu_blocks, num_cpu_blocks
+
+    def _assert_memory_footprint_increased_during_profiling(self):
+        # NOTE(woosuk): Here we assume that the other processes using the same
+        # MLU did not change their memory usage during the profiling.
+        free_gpu_memory, total = torch.mlu.mem_get_info()
+        mlu_memory = total - free_gpu_memory
+        assert self.baseline_snapshot.mlu_memory < mlu_memory, (
+            "Error in memory profiling. "
+            f"Initial used memory {self.baseline_snapshot.mlu_memory}, "
+            f"currently used memory {mlu_memory}. "
+            f"This happens when the MLU memory was "
+            "not properly cleaned up before initializing the vLLM instance.")
+
+    def _warm_up_model(self) -> None:
+        # warm up sizes that are not in cudagraph capture sizes,
+        # but users still want to compile for better performance,
+        # e.g. for the max-num-batched token size in chunked prefill.
+        warmup_sizes = self.vllm_config.compilation_config.compile_sizes.copy()
+        if not self.model_config.enforce_eager:
+            warmup_sizes = [
+                x for x in warmup_sizes if x not in
+                self.vllm_config.compilation_config.cudagraph_capture_sizes
+            ]
+        for size in sorted(warmup_sizes, reverse=True):
+            logger.info("Compile and warming up model for size %d", size)
+            self.model_runner._dummy_run(size)
+        if not self.model_config.enforce_eager:
+            '''
+            =============================
+            Modify by vllm_mlu
+            =============================
+            @brief: support context mlugraph
+            ''' 
+            if self.model_config.use_context_mlugraph():
+                # Capture MLUGraph both prefill and decode
+                self.model_runner.capture_model_with_context(self.gpu_cache)
+            else:
+                # Capture MLUGraph only decode
+                self.model_runner.capture_model(self.gpu_cache)
+            '''
+            ==================
+            End of MLU Hijack
+            ==================
+            '''
+        # Reset the seed to ensure that the random state is not affected by
+        # the model initialization and profiling.
+        set_random_seed(self.model_config.seed)
+
+    def get_latency(self):
+        time_markers = self.model_runner.time_markers
+        total_latency = 0
+        if not isinstance(time_markers, list):
+            time_markers = [time_markers]
+        for time_marker in time_markers:
+            start, end = time_marker
+            latency = start.elapsed_time(end)
+            total_latency += latency
+        return total_latency
+
+    def get_memory_usage(self):
+        return (self.peak_memory, self.block_memory,
+                self.num_gpu_blocks, self.num_cpu_blocks)
+
+    def recapture_model(
+        self,
+        enable_context_mlugraph,
+        context_batch_size_to_capture,
+        context_seq_len_to_capture
+    ) -> None:
+        # Reset history capture context
+        self.model_runner.reset_capture_context()
+
+        # Re-capture context and decoder mlugraph
+        self.model_runner.model_config.set_context_mlugraph_info(
+            enable_context_mlugraph, context_batch_size_to_capture, context_seq_len_to_capture)
+        self._warm_up_model()
+
+    def stat_tensor(self, name, tensor, act_range, key, dim):
+        logger.debug(f"name:{name}, key:{key}, dim:{dim}, tensor.shape:{tensor.shape}")
+        hidden_dim = tensor.shape[-1]
+        tensor = tensor.view(-1, hidden_dim).abs()
+        comming_max = torch.max(tensor, dim=dim)[0].float()
+
+        if act_range[name][key] is None:
+            act_range[name][key] = comming_max
+        else:
+            act_range[name][key] = torch.max(act_range[name][key], comming_max)
+
+    def stat_input_hook(self, m, x, y, name, act_range, is_linear, is_save_input_id):
+        if isinstance(x, tuple):
+            x = x[0]
+        if isinstance(y, tuple):
+            y = y[0]
+        logger.debug(f"name:{name}, x.shape:{x.shape}, y.shape:{y.shape}, m.weight.shape:{m.weight.shape}")
+        if is_linear:
+            self.stat_tensor(name, x, act_range, "x", 0)
+            if act_range[name]["is_qkv"] and is_save_input_id and ".0." in name:
+                x_cpu = x.clone().to("cpu")
+                act_range[name]["input_id"].append(x_cpu)
+
+    def setup_smooth_hook(self, is_save_input_id: bool = False):
+        model = self.model_runner.model
+        self.act_range = defaultdict(default_act_range_value)
+        self.hooks = []
+        linear_class_list = (ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear)
+        other_class_list = (VocabParallelEmbedding, ParallelLMHead)
+        class_list = linear_class_list + other_class_list
+        row_class_list = (RowParallelLinear)
+
+        for name, m in model.named_modules():
+            if isinstance(m, FeedForward):
+                m.use_bt_ffn = False
+            if isinstance(m, SparseMoeMlp):
+                m.is_use_fused_moe = False
+            if isinstance(m, DeepseekV2MLAAttention):
+                m.use_fused_mla_qkv = False
+
+            if isinstance(m, class_list):
+                is_linear = True if isinstance(m, linear_class_list) else False
+                split_type = "row" if isinstance(m, row_class_list) else "col"
+                self.act_range[name]["split"] = split_type
+                self.act_range[name]["is_linear"] = is_linear
+                if isinstance(m, QKVParallelLinear):
+                    self.act_range[name]["is_qkv"] = True
+                    self.act_range[name]["q_proj_size"] = m.num_heads * m.head_size
+                    self.act_range[name]["num_kv_head_replicas"] = m.num_kv_head_replicas
+                self.act_range[name]["is_merge"] = isinstance(m, MergedColumnParallelLinear)
+
+                logger.info(f"rank:{self.rank}, add hook to {name}, is_linear:{is_linear}, split_type:{split_type}")
+                self.hooks.append(m.register_forward_hook(functools.partial(self.stat_input_hook,
+                                                                            name=name, act_range=self.act_range,
+                                                                            is_linear=is_linear,
+                                                                            is_save_input_id=is_save_input_id)))
+
+    def remove_hooks(self):
+        for h in self.hooks:
+            h.remove()
+
+    def get_act_range(self):
+        act_range = defaultdict(default_act_range_value)
+        for layer_name, layer_range in self.act_range.items():
+            for tensor_key, tensor_value in layer_range.items():
+                if isinstance(tensor_value, torch.Tensor):
+                    act_range[layer_name][tensor_key] = tensor_value.to("cpu")
+                elif tensor_key == "input_id" and isinstance(tensor_value, list):
+                    input_id_len = len(tensor_value)
+                    for i in range(input_id_len):
+                        if isinstance(tensor_value[i], torch.Tensor):
+                            act_range[layer_name][tensor_key].append(tensor_value[i].to("cpu"))
+                        else:
+                            act_range[layer_name][tensor_key].append(tensor_value[i])
+                else:
+                    act_range[layer_name][tensor_key] = tensor_value
+
+        return act_range
+
+    @torch.no_grad()
+    def get_named_parameters(self):
+        name_parameters = {}
+        for name, param in self.model_runner.model.named_parameters():
+            name_parameters[name] = param.to("cpu")
+
+        return name_parameters
+
+
+def init_worker_distributed_environment(
+    vllm_config: VllmConfig,
+    rank: int,
+    distributed_init_method: Optional[str] = None,
+    local_rank: int = -1,
+) -> None:
+    """Initialize the distributed environment."""
+    parallel_config = vllm_config.parallel_config
+    set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
+
+    init_distributed_environment(parallel_config.world_size, rank,
+                                 distributed_init_method, local_rank,
+                                 backend='cncl')
+
+    ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
+                                      parallel_config.pipeline_parallel_size)
+
+    ensure_kv_transfer_initialized(vllm_config)
+
+
+def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
+    # Check if the GPU supports the dtype.
+    if torch_dtype == torch.bfloat16:  # noqa: SIM102
+        if not current_platform.has_device_capability(50):
+            capability = current_platform.get_device_capability()
+            gpu_name = current_platform.get_device_name()
+
+            if capability is None:
+                compute_str = "does not have a compute capability"
+            else:
+                version_str = capability.as_version_str()
+                compute_str = f"has compute capability {version_str}"
+
+            raise ValueError(
+                "Bfloat16 is only supported on MLUs with compute capability "
+                f"of at least 5xx. Your {gpu_name} MLU {compute_str}. "
+                "You can use float16 instead by explicitly setting the"
+                "`dtype` flag in CLI, for example: --dtype=half.")
diff --git a/vllm_mlu/vllm_mlu/worker/multi_step_mlu_model_runner.py b/vllm_mlu/vllm_mlu/worker/multi_step_mlu_model_runner.py
new file mode 100644
index 000000000..a92730794
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/worker/multi_step_mlu_model_runner.py
@@ -0,0 +1,471 @@
+import dataclasses
+import functools
+from dataclasses import dataclass, field
+from typing import (Any, Callable, Dict, List, Optional, Union)
+
+import torch
+
+from vllm.distributed import get_pp_group
+from vllm.logger import init_logger
+from vllm.model_executor.layers.sampler import (SamplerOutput,
+                                                SamplingMetadata)
+from vllm.sequence import (IntermediateTensors, SequenceGroupMetadata)
+from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
+from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
+from vllm.worker.multi_step_model_runner import (
+    ModelOutput,
+    StatefulModelInput,
+    PythonizationCache,
+    _pythonize_sampler_output,
+    MULTI_STEP_ATTENTION_BACKENDS,
+    MULTI_STEP_CHUNKED_PREFILL_ATTENTION_BACKENDS,
+    _get_supported_attention_backends
+)
+
+from vllm_mlu.worker.mlu_model_runner import MLUModelRunnerBase
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class MLUModelOutput(ModelOutput):
+
+    def _pythonize_sampler_output(self, input_metadata: "MLUStatefulModelInput",
+                                  copy_stream: torch.mlu.Stream,
+                                  pinned_sampled_token_buffer: torch.Tensor,
+                                  blocking: bool) -> bool:
+        """
+        If blocking is set, will block until the forward pass for the output is
+        ready and pythonize the output. Upon completing Pythonization, erases
+        self.logprobs (note that a non-blocking call that is performed when
+        the sampler output is not yet ready, will not erase self.logprobs.)
+        """
+        assert self.sampled_token_ids is not None
+        if not blocking and not self.sampler_output_ready_event.query():
+            return False
+
+        if blocking:
+            self.sampler_output_ready_event.synchronize()
+        with torch.mlu.stream(copy_stream):
+            _pythonize_sampler_output(input_metadata, self.sampler_output,
+                                      pinned_sampled_token_buffer,
+                                      self.sampled_token_ids, self.logprobs,
+                                      self.pythonization_cache)
+
+        # Erase the logprobs GPU-side tensor.
+        # Note that although _pythonize_sampler_output() runs in its
+        # own CUDA stream, nonetheless _pythonize_sampler_output()
+        # cannot return until Pythonization is complete; therefore
+        # we know that by the time the CPU reaches this point,
+        # `self.logprobs` is no longer needed.
+        self.logprobs = None
+        return True
+
+
+@dataclass(frozen=False)
+class MLUStatefulModelInput(StatefulModelInput):
+    # ping-pong data structures for multi-step to wait on the previous step
+    step_cuda_events: List[torch.mlu.Event] = field(
+        default_factory=lambda: [torch.mlu.Event(blocking=True)] * 2)
+
+    def record_step_event(self, current_stream: torch.mlu.Stream):
+        # record the event for the current step so that the next step can sync
+        # on it. We modulo by 2 to keep the events in a circular buffer and
+        # support any attn backends that may be supported in the future. ie
+        # Flashinfer would want two DecodeWrappers to overlap the CPU and GPU.
+        self.step_cuda_events[self.current_step & 1] = \
+            torch.mlu.Event(blocking=True)
+        self.step_cuda_events[self.current_step & 1].record(current_stream)
+
+    def add_sampler_output(self,
+                           sampler_output: SamplerOutput,
+                           sampled_token_ids: Optional[torch.Tensor] = None):
+        self.cached_outputs.append(
+            MLUModelOutput(sampler_output=sampler_output,
+                           sampler_output_ready_event=None,
+                           sampled_token_ids=sampled_token_ids,
+                           pythonized=False))
+
+
+# MutableModelInputForGPUWithMultiStepMetadata is not subclass of
+# ModelInputForGPU but it wraps the actual input dataclass and adds multi-step
+# metadata
+# mypy: disable-error-code=type-var
+class MLUMultiStepModelRunner(MLUModelRunnerBase[MLUStatefulModelInput]):
+    # mypy: enable-error-code=type-var
+
+    def __init__(self, base_model_runner: MLUModelRunnerBase, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Check attention backend support.
+        supported_attention_backends: List[str] = \
+            _get_supported_attention_backends(
+                self.scheduler_config.chunked_prefill_enabled)
+        if self.attn_backend.get_name() not in supported_attention_backends:
+            ms_config_str: str = "Multi-Step + Chunked-Prefill" \
+                if self.scheduler_config.chunked_prefill_enabled \
+                      else "Multi-Step"
+            raise ValueError(
+                f"{ms_config_str} not supported for attention backend: "
+                f"{self.attn_backend.get_name()}. Set VLLM_ATTENTION_BACKEND "
+                f"to a value from {supported_attention_backends}.")
+
+        # uses the base model runner to execute the model and wraps it with
+        # multi-step logic
+        self._base_model_runner: MLUModelRunnerBase = base_model_runner
+
+        self.is_multi_step = self.scheduler_config.is_multi_step
+        self.pinned_sampled_token_ids: Optional[torch.Tensor] = None
+
+        # Using the PythonizationCache in Pipeline-Parallel clobbers the
+        # SequenceOutput and CompletionSequenceGroupOutput object.
+        # When cache-reset happens at the last step of a multi-step
+        # execution, there may be other on-going single-step/multi-step
+        # executions. The current caching implementation does not check
+        # for this.
+        self.pythonization_cache = PythonizationCache() \
+            if self.parallel_config.pipeline_parallel_size == 1 else None
+
+    @functools.cached_property
+    def _copy_stream(self):
+        # used to copy tensors from GPU to CPU asynchronously
+        return torch.mlu.Stream()
+
+    def make_model_input_from_broadcasted_tensor_dict(
+            self, tensor_dict: Dict[str, Any]) -> MLUStatefulModelInput:
+        model_input = (MLUStatefulModelInput.from_broadcasted_tensor_dict(
+            tensor_dict,
+            attn_backend=self.attn_backend,
+        ))
+        return model_input
+
+    def prepare_model_input(
+        self,
+        seq_group_metadata_list: List[SequenceGroupMetadata],
+        virtual_engine: int = 0,
+        finished_requests_ids: Optional[List[str]] = None
+    ) -> MLUStatefulModelInput:
+        frozen_model_input: ModelInputForGPUWithSamplingMetadata = \
+              self._base_model_runner.prepare_model_input(
+                    seq_group_metadata_list,
+                    virtual_engine,
+                    finished_requests_ids)
+
+        assert frozen_model_input.query_lens is not None
+        assert frozen_model_input.seq_lens is not None
+        assert frozen_model_input.attn_metadata is not None
+        num_queries = len(frozen_model_input.query_lens)
+        num_seqs = len(frozen_model_input.seq_lens)
+        num_single_step_prefills = frozen_model_input.attn_metadata.num_prefills
+
+        model_input = MLUStatefulModelInput(
+            frozen_model_input=frozen_model_input,
+            num_seqs=num_seqs,
+            num_queries=num_queries,
+            num_single_step_prefills=num_single_step_prefills)
+
+        return model_input
+
+    def _async_process_outputs(self, model_input: MLUStatefulModelInput,
+                               output_proc_callback: Callable):
+        # Proceed with pythonization and output_proc in order.
+        # Stop on the first one that fails to pythonize
+        output_proc_callback()
+
+        cont = True
+        for step_num, model_output in enumerate(model_input.cached_outputs):
+            if not model_output.pythonized:
+                model_output.maybe_pythonize(model_input, self._copy_stream,
+                                             self.pinned_sampled_token_ids)
+                if model_output.pythonized:
+                    ctx = output_proc_callback.keywords["ctx"]
+                    ctx.append_output(
+                        outputs=[model_output.sampler_output],
+                        seq_group_metadata_list=ctx.seq_group_metadata_list,
+                        scheduler_outputs=ctx.scheduler_outputs,
+                        is_async=False,
+                        is_last_step=False,
+                        is_first_step_output=step_num == 0)
+
+                    output_proc_callback()
+                else:
+                    cont = False
+
+            if not cont:
+                break
+
+    def _final_process_outputs(
+            self, model_input: MLUStatefulModelInput,
+            output_proc_callback: Optional[Callable]) -> List[SamplerOutput]:
+        assert model_input.frozen_model_input is not None
+
+        has_async_callback = output_proc_callback is not None
+
+        outputs = []
+        for step_num, output in enumerate(model_input.cached_outputs):
+            is_last_step = step_num == len(model_input.cached_outputs) - 1
+
+            # For non-async case:
+            #   -- We simply add the outputs
+            # For async case:
+            #   -- Invoke callback, pythonize, add to callback queue and repeat
+            #   -- For last output, just add to callback queue
+            if has_async_callback:
+                assert output_proc_callback is not None
+
+                # Invoke callback before pythonize (to overlap with GPU)
+                output_proc_callback()
+
+                # Pythonize
+                if not output.pythonized:
+                    output.pythonize(model_input, self._copy_stream,
+                                     self.pinned_sampled_token_ids)
+
+                    # For non last step, add to callback queue to chain
+                    # callbacks=>pythonize pairs (for GPU overlap)
+                    if not is_last_step:
+                        ctx = output_proc_callback.keywords[  # type: ignore
+                            "ctx"]  # type: ignore
+                        ctx.append_output(
+                            outputs=[output.sampler_output],
+                            seq_group_metadata_list=ctx.
+                            seq_group_metadata_list,
+                            scheduler_outputs=ctx.scheduler_outputs,
+                            is_async=False,
+                            is_last_step=False,
+                            is_first_step_output=step_num == 0)
+                    else:
+                        outputs.append(output.sampler_output)
+            else:
+                output.pythonize(model_input, self._copy_stream,
+                                 self.pinned_sampled_token_ids)
+                outputs.append(output.sampler_output)
+
+        return outputs
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        model_input: MLUStatefulModelInput,
+        kv_caches: List[torch.Tensor],
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        num_steps: int = 1,
+    ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:
+        """ 
+        Execute the model for a single step and update multi-step
+        metadata
+        """
+        assert num_steps == 1, "MLUMultiStepModelRunner only supports num_steps=1"
+        frozen_model_input = model_input.frozen_model_input
+        assert frozen_model_input is not None
+
+        # path for warm up runs
+        if not model_input.is_multi_step:
+            return self._base_model_runner.execute_model(
+                frozen_model_input, None, intermediate_tensors, num_steps)
+
+        # make sure we skip the sampler on the lask rank and only pythonize
+        # if CPU is ahead.
+        if self.is_driver_worker and get_pp_group().is_last_rank:
+            if self.pinned_sampled_token_ids is None:
+                self.pinned_sampled_token_ids = torch.zeros(
+                    (self.scheduler_config.max_num_seqs, 1),
+                    dtype=torch.long,
+                    device="cpu",
+                    pin_memory=True)
+
+            self._base_model_runner.model.sampler.include_gpu_probs_tensor = (
+                True)
+            if frozen_model_input.sampling_metadata:
+                frozen_model_input.sampling_metadata.skip_sampler_cpu_output = (
+                    True)
+
+        # some pre-execute model logic for multi-step:
+        #   - if it's the first step, we need to reset the sampling tensors
+        #   - if it's not the first step, we need to advance the step using the
+        #   appended sampler output from last iteration
+        #   - also maybe pythonize if CPU is ahead of GPU
+
+        stream = torch.mlu.current_stream()
+        if not model_input.is_first_multi_step:
+            # Explicitly block on the previous step's forward to make sure we
+            # don't clobber any GPU tensors still in use.
+            # This is not needed for flashattn backend, but for other attn
+            # backends such as flashinfer that performs extra CPU operations on
+            # input metadata we may need to synchronize any CPU operations that
+            # might clobber enqueued forwards. (prevents CPU from running too
+            # far ahead if needed)
+            model_input.wait_previous_step()
+            model_input = self._advance_step(
+                model_input, model_input.cached_outputs[-1].sampler_output)
+
+            # frozen_model_input may have been updated
+            frozen_model_input = model_input.frozen_model_input
+            assert frozen_model_input is not None
+
+        if model_input.base_output_proc_callback is None:
+            assert frozen_model_input is not None
+            model_input.base_output_proc_callback = \
+                        frozen_model_input.async_callback
+
+        if frozen_model_input.async_callback is not None:
+            assert model_input.base_output_proc_callback is not None
+            async_callback = functools.partial(
+                self._async_process_outputs,
+                model_input=model_input,
+                output_proc_callback=model_input.base_output_proc_callback)
+
+            model_input.frozen_model_input = dataclasses.replace(  # type: ignore
+                model_input.frozen_model_input,
+                async_callback=async_callback)
+            # Update the local instance
+            frozen_model_input = model_input.frozen_model_input
+            assert frozen_model_input is not None
+
+        # Execute the model
+        output = self._base_model_runner.execute_model(frozen_model_input,
+                                                       None,
+                                                       intermediate_tensors,
+                                                       num_steps=1)
+
+        # record the event for the current step so that the next step can sync
+        model_input.record_step_event(stream)
+
+        if get_pp_group().is_last_rank and self.is_driver_worker:
+            assert isinstance(output, list)
+            assert len(
+                output
+            ) == 1, "MultiStepModelRunner requires single-step base_models"
+
+            # event for the pythonization so that we only pythonize if the
+            # tensors are ready. May be able to be combined with the step event
+            output_ready_event = torch.mlu.Event()
+            output_ready_event.record(stream)
+            if self.parallel_config.pipeline_parallel_size > 1:
+                output[0].sampled_token_ids_cpu = output[
+                    0].sampled_token_ids.cpu()
+            model_input.cached_outputs.append(
+                MLUModelOutput(output[0], output_ready_event,
+                               output[0].sampled_token_ids, False,
+                               output[0].logprobs, self.pythonization_cache))
+
+            # These GPU tensors are not required by multi-step;
+            # erase them to ensure they are not pythonized or
+            # transferred to CPU
+            output[0].sampled_token_ids = None
+            output[0].sampled_token_probs = None
+            output[0].logprobs = None
+
+            # Pythonize the output if CPU is ahead and the previous step is
+            # ready.
+            if frozen_model_input.async_callback is None:
+                for model_output in model_input.cached_outputs:
+                    model_output.maybe_pythonize(model_input,
+                                                 self._copy_stream,
+                                                 self.pinned_sampled_token_ids)
+
+        model_input.current_step += 1
+
+        if not get_pp_group().is_last_rank:
+            # Should be IntermediateTensors
+            assert isinstance(output, IntermediateTensors)
+            return output
+        if not self.is_driver_worker:
+            return []
+
+        # Pythonize the output and block if needed since it is the last step
+        if model_input.is_last_step:
+            outputs = self._final_process_outputs(
+                model_input, model_input.base_output_proc_callback)
+            if self.pythonization_cache:
+                self.pythonization_cache.reset()
+            return outputs
+
+        # should be [SamplerOutput]
+        return output
+
+    def _update_sampling_metadata(self, sampling_metadata: SamplingMetadata,
+                                  num_seqs: Optional[int], num_queries: int):
+
+        assert sampling_metadata.num_prompts == 0
+        assert len(sampling_metadata.seq_groups) == num_queries
+        assert sampling_metadata.selected_token_indices.shape == (
+            num_queries, )
+        # assert sampling_metadata.categorized_sample_indices == TODO: Add if needed # noqa: E501
+
+        # Verify that all sequences are decodes
+        for i in range(num_queries):
+            seq_group = sampling_metadata.seq_groups[i]
+
+            assert seq_group.is_prompt is False  # No prompt
+            assert seq_group.prompt_logprob_indices == []  # No prompt
+            assert seq_group.sample_indices == [i]  # Simple
+            assert seq_group.seq_len is None  # Decode
+            assert seq_group.query_len is None  # Decode
+
+    def _advance_step(self, model_input: MLUStatefulModelInput,
+                      out: SamplerOutput) -> MLUStatefulModelInput:
+
+        model_input.maybe_advance_frozen_model_input(self.device,
+                                                     self.pin_memory)
+        frozen_model_input = model_input.frozen_model_input
+        assert frozen_model_input is not None
+        assert frozen_model_input.input_tokens is not None
+        assert frozen_model_input.input_tokens.shape[0] == model_input.num_seqs
+        assert frozen_model_input.attn_metadata is not None
+
+        sampled_token_ids = model_input.cached_outputs[-1].sampled_token_ids
+        num_seqs = model_input.num_seqs
+        num_queries = model_input.num_queries
+        frozen_model_input = model_input.frozen_model_input
+        assert frozen_model_input is not None
+        attn_metadata = frozen_model_input.attn_metadata
+        assert attn_metadata is not None
+
+        turn_prefills_into_decodes: bool = model_input.current_step == 1 and \
+                                    model_input.num_single_step_prefills != 0
+        attn_metadata.advance_step(
+            frozen_model_input,
+            sampled_token_ids,
+            self.block_size,
+            num_seqs,
+            num_queries,
+            turn_prefills_into_decodes=turn_prefills_into_decodes)
+
+        return model_input
+
+    def load_model(self) -> None:
+        self._base_model_runner.load_model()
+        self.model_memory_usage = self._base_model_runner.model_memory_usage
+
+    def save_sharded_state(
+        self,
+        path: str,
+        pattern: Optional[str] = None,
+        max_size: Optional[int] = None,
+    ) -> None:
+        return self._base_model_runner.save_sharded_state(
+            path, pattern, max_size)
+
+    def save_tensorized_model(self,
+                              tensorizer_config: TensorizerConfig) -> None:
+        return self._base_model_runner.save_tensorized_model(tensorizer_config)
+
+    def profile_run(self) -> None:
+        return self._base_model_runner.profile_run()
+
+    def remove_all_loras(self):
+        return self._base_model_runner.remove_all_loras()
+
+    def capture_model(self, kv_caches: List[List]) -> None:
+        return self._base_model_runner.capture_model(kv_caches)
+
+    @property
+    def vocab_size(self) -> int:
+        return self._base_model_runner.vocab_size
+
+    def capture_model_with_context(self, kv_caches: List[List[torch.Tensor]]) -> None:
+        return self._base_model_runner.capture_model_with_context(kv_caches)
+
+    def reset_capture_context(self):
+        return self._base_model_runner.reset_capture_context()
\ No newline at end of file
diff --git a/vllm_mlu/vllm_mlu/worker/multi_step_mlu_worker.py b/vllm_mlu/vllm_mlu/worker/multi_step_mlu_worker.py
new file mode 100644
index 000000000..a72e5f4b3
--- /dev/null
+++ b/vllm_mlu/vllm_mlu/worker/multi_step_mlu_worker.py
@@ -0,0 +1,207 @@
+import dataclasses
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from vllm.distributed import broadcast_tensor_dict, get_pp_group
+from vllm.model_executor.layers.sampler import SamplerOutput
+from vllm.sequence import ExecuteModelRequest
+from vllm.worker.model_runner_base import BroadcastableModelInput
+from vllm.worker.worker import WorkerInput
+
+from vllm_mlu.worker.mlu_worker import MLUWorker
+from vllm_mlu.worker.multi_step_mlu_model_runner import (
+    MLUMultiStepModelRunner, MLUStatefulModelInput)
+
+
+@dataclass
+class MultiStepState:
+    worker_input: WorkerInput
+    model_input: MLUStatefulModelInput
+
+
+class MLUMultiStepWorker(MLUWorker):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        base_model_runner = self.model_runner
+        # for multi-step model, wrap the model runner with MLUMultiStepModelRunner
+        self.model_runner = MLUMultiStepModelRunner(
+            base_model_runner,
+            vllm_config=base_model_runner.vllm_config,
+            kv_cache_dtype=self.cache_config.cache_dtype,
+            is_driver_worker=base_model_runner.is_driver_worker,
+        )
+
+        pipeline_parallel_size = self.parallel_config.pipeline_parallel_size
+        self.multi_step_states: List[
+            Optional[MultiStepState]] = [None] * pipeline_parallel_size
+        self.temp_output = None
+
+    def _get_driver_input_and_broadcast(
+        self, execute_model_req: ExecuteModelRequest
+    ) -> Tuple[BroadcastableModelInput, WorkerInput, Dict[str, torch.Tensor]]:
+        """
+        Get the driver input and broadcast it to other workers.
+        """
+        assert self.is_driver_worker
+        virtual_engine = execute_model_req.virtual_engine
+        is_first_multi_step = execute_model_req.is_first_multi_step
+        if is_first_multi_step:
+            # on first step we prepare the worker input and model input normally
+            worker_input: WorkerInput = self.prepare_worker_input(
+                execute_model_req=execute_model_req)
+            model_input: MLUStatefulModelInput = (
+                self.model_runner.prepare_model_input(
+                    execute_model_req.seq_group_metadata_list,
+                    execute_model_req.virtual_engine,
+                    execute_model_req.finished_requests_ids))
+
+            if execute_model_req.async_callback:
+                model_input.frozen_model_input = dataclasses.replace(  # type: ignore
+                    model_input.frozen_model_input,
+                    async_callback=execute_model_req.async_callback)
+        else:
+            # on subsequent steps we reuse the worker input and model input
+            multi_step_state = self.multi_step_states[virtual_engine]
+            worker_input = multi_step_state.worker_input
+            model_input = multi_step_state.model_input
+            frozen_model_input = model_input.frozen_model_input
+            assert frozen_model_input is not None
+            assert frozen_model_input.attn_metadata is not None
+            # clear the cached metadata so that it can be recomputed on
+            # the workers.
+            frozen_model_input.attn_metadata._cached_prefill_metadata = None
+            frozen_model_input.attn_metadata._cached_decode_metadata = None
+
+        model_input.is_first_multi_step = is_first_multi_step
+        model_input.is_last_step = execute_model_req.is_last_step
+
+        if not is_first_multi_step:
+            # we broadcast the last sampled token ids to all TP workers so they
+            # can update their model input metadata in-place.
+            self._prepare_last_sampled_token_ids_for_tp_workers(
+                execute_model_req=execute_model_req, model_input=model_input)
+
+        if self.do_metadata_broadcast:
+            broadcast_data = worker_input.as_broadcastable_tensor_dict()
+            broadcast_data.update(model_input.as_broadcastable_tensor_dict())
+            broadcast_tensor_dict(broadcast_data, src=0)
+
+        # Retuning empty dict here to keep this compatible with
+        # `LocalOrDistributedWorkerBase._get_driver_input_and_broadcast`
+        return model_input, worker_input, {}
+
+    def _prepare_last_sampled_token_ids_for_tp_workers(
+        self,
+        execute_model_req: ExecuteModelRequest,
+        model_input: MLUStatefulModelInput,
+    ) -> None:
+        """ 
+        Prepare the last sampled token ids for TP workers. If it's the last 
+        PP rank, then the last sampled token ids are already in the model_input.
+        If it is NOT the last PP rank, then we need to get the last sampled
+        token that is cached in the execute_model_req.
+        """
+        if get_pp_group().is_last_rank:
+            assert model_input.cached_outputs[
+                -1].sampler_output.sampled_token_ids is None
+            assert model_input.cached_outputs[-1].sampled_token_ids is not None
+            model_input.last_sampled_token_ids = model_input.cached_outputs[
+                -1].sampled_token_ids
+            # free sampled token ids from the previous step if it has been
+            # pythonized. Cannot free the last sampled token ids because
+            # we need it for GPU advance_step.
+            for output in model_input.cached_outputs[:-1]:
+                if output.pythonized:
+                    output.sampled_token_ids = None
+        else:
+            # otherwise we need to get the cached sampled token ids from the
+            # execute_model_req
+            assert execute_model_req.last_sampled_token_ids is not None
+            model_input.last_sampled_token_ids = (
+                execute_model_req.last_sampled_token_ids.mlu())
+            model_input.add_sampler_output(
+                SamplerOutput(outputs=[], sampled_token_ids=None),
+                model_input.last_sampled_token_ids)
+
+            # free sampled token ids from the previous step.
+            # TODO(will) we could reuse the sampled token ids tensor from
+            # the previous step instead.
+            for output in model_input.cached_outputs[:-1]:
+                output.sampled_token_ids = None
+            assert model_input.cached_outputs[-1].sampled_token_ids is not None
+
+    def prepare_input(
+        self,
+        execute_model_req: Optional[ExecuteModelRequest] = None,
+    ) -> Optional[Tuple[MLUStatefulModelInput, WorkerInput, Dict[str,
+                                                              torch.Tensor]]]:
+        """
+        Depending on the current state of the request and multi step worker,
+        this method may skip the normal _prepare_model_input and
+        _prepare_worker_input methods and instead used cached values.
+        """
+        if self.is_driver_worker:
+            if execute_model_req is None:
+                if self.do_metadata_broadcast:
+                    # This signals that there's no more requests to process for
+                    # now. All workers are running infinite loop with
+                    # broadcast_tensor_dict, and it stops the loop when the
+                    # driver broadcasts an empty input. Send an empty input to
+                    # notify all other workers to stop their execution loop.
+                    broadcast_tensor_dict({}, src=0)
+                return None
+
+            virtual_engine = execute_model_req.virtual_engine
+            (model_input, worker_input,
+             kwargs) = self._get_driver_input_and_broadcast(execute_model_req)
+            assert isinstance(model_input, MLUStatefulModelInput)
+            if execute_model_req.is_first_multi_step:
+                # cache the worker input and model input for the next steps
+                self.multi_step_states[virtual_engine] = MultiStepState(
+                    worker_input=worker_input, model_input=model_input)
+        # if TP workers
+        else:
+            broadcast_data = self._get_worker_input_from_broadcast()
+            # if the driver has sent an empty input, we should stop the worker
+            # loop
+            if broadcast_data is None:
+                return None
+            model_input, worker_input, kwargs = broadcast_data
+            assert isinstance(model_input, MLUStatefulModelInput)
+            virtual_engine = worker_input.virtual_engine
+            if model_input.is_first_multi_step:
+                pass
+                # TODO(will) Can cache the worker input and model input for the
+                # next steps. See below for details
+            else:
+                # TODO(will) possible to also cache and reuse the cached worker
+                # input and model input. The idea is essentially the delta
+                # optimization for model_inputs. Where the TP workers can cache
+                # the model input states and we only broadcast the delta need
+                # for the next step (sampled_token_ids from the previous step)
+
+                assert isinstance(model_input, MLUStatefulModelInput)
+                # we need to update the last sampled token ids in the model
+                # input for the workers so that they can run inplace
+                # advance_step
+                model_input.add_sampler_output(
+                    SamplerOutput(outputs=[], sampled_token_ids=None),
+                    model_input.last_sampled_token_ids)
+
+        assert model_input is not None
+        assert worker_input is not None
+        return model_input, worker_input, kwargs
+
+    def get_latency(self):
+        time_markers = self.model_runner._base_model_runner.time_markers
+        total_latency = 0
+        if not isinstance(time_markers, list):
+            time_markers = [time_markers]
+        for time_marker in time_markers:
+            start, end = time_marker
+            latency = start.elapsed_time(end)
+            total_latency += latency
+        return total_latency
\ No newline at end of file
