diff --git a/.buildkite/check-wheel-size.py b/.buildkite/check-wheel-size.py
index 68aff793a..a378bc6ba 100644
--- a/.buildkite/check-wheel-size.py
+++ b/.buildkite/check-wheel-size.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import os
 import sys
@@ -9,12 +8,12 @@ import zipfile
 # Note that we have 400 MiB quota, please use it wisely.
 # See https://github.com/pypi/support/issues/3792 .
 # Please also sync the value with the one in Dockerfile.
-VLLM_MAX_SIZE_MB = int(os.environ.get("VLLM_MAX_SIZE_MB", 400))
+VLLM_MAX_SIZE_MB = int(os.environ.get('VLLM_MAX_SIZE_MB', 400))
 
 
 def print_top_10_largest_files(zip_file):
     """Print the top 10 largest files in the given zip file."""
-    with zipfile.ZipFile(zip_file, "r") as z:
+    with zipfile.ZipFile(zip_file, 'r') as z:
         file_sizes = [(f, z.getinfo(f).file_size) for f in z.namelist()]
         file_sizes.sort(key=lambda x: x[1], reverse=True)
         for f, size in file_sizes[:10]:
@@ -29,18 +28,14 @@ def check_wheel_size(directory):
                 wheel_path = os.path.join(root, file_name)
                 wheel_size_mb = os.path.getsize(wheel_path) / (1024 * 1024)
                 if wheel_size_mb > VLLM_MAX_SIZE_MB:
-                    print(
-                        f"Not allowed: Wheel {wheel_path} is larger "
-                        f"({wheel_size_mb:.2f} MB) than the limit "
-                        f"({VLLM_MAX_SIZE_MB} MB)."
-                    )
+                    print(f"Not allowed: Wheel {wheel_path} is larger "
+                          f"({wheel_size_mb:.2f} MB) than the limit "
+                          f"({VLLM_MAX_SIZE_MB} MB).")
                     print_top_10_largest_files(wheel_path)
                     return 1
                 else:
-                    print(
-                        f"Wheel {wheel_path} is within the allowed size "
-                        f"({wheel_size_mb:.2f} MB)."
-                    )
+                    print(f"Wheel {wheel_path} is within the allowed size "
+                          f"({wheel_size_mb:.2f} MB).")
     return 0
 
 
@@ -50,4 +45,4 @@ if __name__ == "__main__":
         sys.exit(1)
 
     directory = sys.argv[1]
-    sys.exit(check_wheel_size(directory))
+    sys.exit(check_wheel_size(directory))
\ No newline at end of file
diff --git a/.buildkite/generate_index.py b/.buildkite/generate_index.py
index 7045d8810..36e1b6c01 100644
--- a/.buildkite/generate_index.py
+++ b/.buildkite/generate_index.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
 import os
@@ -23,5 +22,5 @@ with open("index.html", "w") as f:
     print(f"Generated index.html for {args.wheel}")
     # cloudfront requires escaping the '+' character
     f.write(
-        template.format(wheel=filename, wheel_html_escaped=filename.replace("+", "%2B"))
-    )
+        template.format(wheel=filename,
+                        wheel_html_escaped=filename.replace("+", "%2B")))
diff --git a/.buildkite/lm-eval-harness/configs/Meta-Llama-3.2-1B-Instruct-FP8-compressed-tensors.yaml b/.buildkite/lm-eval-harness/configs/Meta-Llama-3.2-1B-Instruct-FP8-compressed-tensors.yaml
deleted file mode 100644
index cca58097e..000000000
--- a/.buildkite/lm-eval-harness/configs/Meta-Llama-3.2-1B-Instruct-FP8-compressed-tensors.yaml
+++ /dev/null
@@ -1,11 +0,0 @@
-# bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m RedHatAI/Llama-3.2-1B-Instruct-FP8 -b "auto" -l 1319 -f 5 -t 1
-model_name: "RedHatAI/Llama-3.2-1B-Instruct-FP8"
-tasks:
-- name: "gsm8k"
-  metrics:
-  - name: "exact_match,strict-match"
-    value: 0.335
-  - name: "exact_match,flexible-extract"
-    value: 0.323
-limit: 1319
-num_fewshot: 5
diff --git a/.buildkite/lm-eval-harness/configs/Qwen2.5-1.5B-Instruct.yaml b/.buildkite/lm-eval-harness/configs/Qwen2.5-1.5B-Instruct.yaml
deleted file mode 100644
index 54579a63a..000000000
--- a/.buildkite/lm-eval-harness/configs/Qwen2.5-1.5B-Instruct.yaml
+++ /dev/null
@@ -1,11 +0,0 @@
-# bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m Qwen/Qwen2.5-1.5B-Instruct -b auto -l 1319 -f 5 -t 1
-model_name: "Qwen/Qwen2.5-1.5B-Instruct"
-tasks:
-- name: "gsm8k"
-  metrics:
-  - name: "exact_match,strict-match"
-    value: 0.54
-  - name: "exact_match,flexible-extract"
-    value: 0.59
-limit: 1319
-num_fewshot: 5
diff --git a/.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml b/.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
deleted file mode 100644
index a2f235f48..000000000
--- a/.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
+++ /dev/null
@@ -1,11 +0,0 @@
-# bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic -b auto -l 1319 -f 5 -t 1
-model_name: "RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic"
-tasks:
-- name: "gsm8k"
-  metrics:
-  - name: "exact_match,strict-match"
-    value: 0.47
-  - name: "exact_match,flexible-extract"
-    value: 0.64
-limit: 1319
-num_fewshot: 5
diff --git a/.buildkite/lm-eval-harness/configs/models-large.txt b/.buildkite/lm-eval-harness/configs/models-large.txt
index 27a1a9a82..37eeac85c 100644
--- a/.buildkite/lm-eval-harness/configs/models-large.txt
+++ b/.buildkite/lm-eval-harness/configs/models-large.txt
@@ -3,4 +3,3 @@ Meta-Llama-3-70B-Instruct.yaml
 Mixtral-8x7B-Instruct-v0.1.yaml
 Qwen2-57B-A14-Instruct.yaml
 DeepSeek-V2-Lite-Chat.yaml
-Meta-Llama-3-8B-QQQ.yaml
diff --git a/.buildkite/lm-eval-harness/configs/models-small.txt b/.buildkite/lm-eval-harness/configs/models-small.txt
index 36e054387..254d01edf 100644
--- a/.buildkite/lm-eval-harness/configs/models-small.txt
+++ b/.buildkite/lm-eval-harness/configs/models-small.txt
@@ -1,6 +1,10 @@
-Qwen2.5-1.5B-Instruct.yaml
+Meta-Llama-3-8B-Instruct.yaml
+Meta-Llama-3-8B-Instruct-FP8-compressed-tensors.yaml
 Meta-Llama-3.2-1B-Instruct-INT8-compressed-tensors.yaml
 Meta-Llama-3-8B-Instruct-INT8-compressed-tensors-asym.yaml
 Meta-Llama-3-8B-Instruct-nonuniform-compressed-tensors.yaml
-Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
+Meta-Llama-3-8B-Instruct-Channelwise-compressed-tensors.yaml
 Qwen1.5-MoE-W4A16-compressed-tensors.yaml
+Qwen2-1.5B-Instruct-INT8-compressed-tensors.yaml
+Qwen2-1.5B-Instruct-FP8W8.yaml
+Meta-Llama-3-8B-QQQ.yaml
diff --git a/.buildkite/lm-eval-harness/conftest.py b/.buildkite/lm-eval-harness/conftest.py
deleted file mode 100644
index c0d60dd53..000000000
--- a/.buildkite/lm-eval-harness/conftest.py
+++ /dev/null
@@ -1,44 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from pathlib import Path
-
-import pytest
-
-
-def pytest_addoption(parser):
-    parser.addoption(
-        "--config-list-file",
-        action="store",
-        help="Path to the file listing model config YAMLs (one per line)",
-    )
-    parser.addoption(
-        "--tp-size",
-        action="store",
-        default="1",
-        help="Tensor parallel size to use for evaluation",
-    )
-
-
-@pytest.fixture(scope="session")
-def config_list_file(pytestconfig, config_dir):
-    rel_path = pytestconfig.getoption("--config-list-file")
-    return config_dir / rel_path
-
-
-@pytest.fixture(scope="session")
-def tp_size(pytestconfig):
-    return pytestconfig.getoption("--tp-size")
-
-
-def pytest_generate_tests(metafunc):
-    if "config_filename" in metafunc.fixturenames:
-        rel_path = metafunc.config.getoption("--config-list-file")
-        config_list_file = Path(rel_path).resolve()
-        config_dir = config_list_file.parent
-        with open(config_list_file, encoding="utf-8") as f:
-            configs = [
-                config_dir / line.strip()
-                for line in f
-                if line.strip() and not line.startswith("#")
-            ]
-        metafunc.parametrize("config_filename", configs)
diff --git a/.buildkite/lm-eval-harness/run-tests.sh b/.buildkite/lm-eval-harness/run-tests.sh
new file mode 100644
index 000000000..26f33b744
--- /dev/null
+++ b/.buildkite/lm-eval-harness/run-tests.sh
@@ -0,0 +1,59 @@
+#!/bin/bash
+
+usage() {
+    echo``
+    echo "Runs lm eval harness on GSM8k using vllm and compares to "
+    echo "precomputed baseline (measured by HF transformers.)"
+    echo
+    echo "usage: ${0} <options>"
+    echo
+    echo "  -c    - path to the test data config (e.g. configs/small-models.txt)"
+    echo "  -t    - tensor parallel size"
+    echo
+}
+
+SUCCESS=0
+
+while getopts "c:t:" OPT; do
+  case ${OPT} in
+    c ) 
+        CONFIG="$OPTARG"
+        ;;
+    t )
+        TP_SIZE="$OPTARG"
+        ;;
+    \? )
+        usage
+        exit 1
+        ;;
+  esac
+done
+
+# Parse list of configs.
+IFS=$'\n' read -d '' -r -a MODEL_CONFIGS < "$CONFIG"
+
+for MODEL_CONFIG in "${MODEL_CONFIGS[@]}"
+do
+    LOCAL_SUCCESS=0
+    
+    echo "=== RUNNING MODEL: $MODEL_CONFIG WITH TP SIZE: $TP_SIZE==="
+
+    export LM_EVAL_TEST_DATA_FILE=$PWD/configs/${MODEL_CONFIG}
+    export LM_EVAL_TP_SIZE=$TP_SIZE
+    pytest -s test_lm_eval_correctness.py || LOCAL_SUCCESS=$?
+
+    if [[ $LOCAL_SUCCESS == 0 ]]; then
+        echo "=== PASSED MODEL: ${MODEL_CONFIG} ==="
+    else
+        echo "=== FAILED MODEL: ${MODEL_CONFIG} ==="
+    fi
+
+    SUCCESS=$((SUCCESS + LOCAL_SUCCESS))
+
+done
+
+if [ "${SUCCESS}" -eq "0" ]; then
+    exit 0
+else
+    exit 1
+fi
diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
index 930adfaf3..6015a83e8 100644
--- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
+++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py
@@ -1,55 +1,69 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
 LM eval harness on model to compare vs HF baseline computed offline.
 Configs are found in configs/$MODEL.yaml
 
-pytest -s -v test_lm_eval_correctness.py \
-    --config-list-file=configs/models-small.txt \
-    --tp-size=1
+* export LM_EVAL_TEST_DATA_FILE=configs/Meta-Llama-3-70B-Instruct.yaml
+* export LM_EVAL_TP_SIZE=4 
+* pytest -s test_lm_eval_correctness.py
 """
 
+import os
+from pathlib import Path
+
 import lm_eval
-import numpy as np
+import numpy
+import pytest
 import yaml
 
 RTOL = 0.08
+TEST_DATA_FILE = os.environ.get(
+    "LM_EVAL_TEST_DATA_FILE",
+    ".buildkite/lm-eval-harness/configs/Meta-Llama-3-8B-Instruct.yaml")
+
+TP_SIZE = os.environ.get("LM_EVAL_TP_SIZE", 1)
+
 
+def launch_lm_eval(eval_config):
+    trust_remote_code = eval_config.get('trust_remote_code', False)
+
+    model_args = f"pretrained={eval_config['model_name']}," \
+                 f"tensor_parallel_size={TP_SIZE}," \
+                 f"add_bos_token=true," \
+                 f"trust_remote_code={trust_remote_code}"
 
-def launch_lm_eval(eval_config, tp_size):
-    trust_remote_code = eval_config.get("trust_remote_code", False)
-    model_args = (
-        f"pretrained={eval_config['model_name']},"
-        f"tensor_parallel_size={tp_size},"
-        f"enforce_eager=true,"
-        f"add_bos_token=true,"
-        f"trust_remote_code={trust_remote_code}"
-    )
     results = lm_eval.simple_evaluate(
         model="vllm",
         model_args=model_args,
         tasks=[task["name"] for task in eval_config["tasks"]],
         num_fewshot=eval_config["num_fewshot"],
         limit=eval_config["limit"],
-        batch_size="auto",
-    )
+        batch_size="auto")
+
     return results
 
 
-def test_lm_eval_correctness_param(config_filename, tp_size):
-    eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
+def test_lm_eval_correctness():
+    eval_config = yaml.safe_load(
+        Path(TEST_DATA_FILE).read_text(encoding="utf-8"))
+
+    if eval_config[
+            "model_name"] == "nm-testing/Meta-Llama-3-70B-Instruct-FBGEMM-nonuniform":  #noqa: E501
+        pytest.skip("FBGEMM is currently failing on main.")
 
-    results = launch_lm_eval(eval_config, tp_size)
+    # Launch eval requests.
+    results = launch_lm_eval(eval_config)
 
+    # Confirm scores match ground truth.
     success = True
     for task in eval_config["tasks"]:
         for metric in task["metrics"]:
             ground_truth = metric["value"]
             measured_value = results["results"][task["name"]][metric["name"]]
-            print(
-                f"{task['name']} | {metric['name']}: "
-                f"ground_truth={ground_truth} | measured={measured_value}"
-            )
-            success = success and np.isclose(ground_truth, measured_value, rtol=RTOL)
+            print(f'{task["name"]} | {metric["name"]}: '
+                  f'ground_truth={ground_truth} | measured={measured_value}')
+            success = success and numpy.isclose(
+                ground_truth, measured_value, rtol=RTOL)
 
+    # Assert at the end, print all scores even on failure for debugging.
     assert success
diff --git a/.buildkite/nightly-benchmarks/README.md b/.buildkite/nightly-benchmarks/README.md
index 72c52d5bb..d3f5fc5cd 100644
--- a/.buildkite/nightly-benchmarks/README.md
+++ b/.buildkite/nightly-benchmarks/README.md
@@ -113,7 +113,7 @@ WARNING: The benchmarking script will save json results by itself, so please do
 
 ### Visualizing the results
 
-The `convert-results-json-to-markdown.py` helps you put the benchmarking results inside a markdown table, by formatting [descriptions.md](performance-benchmarks-descriptions.md) with real benchmarking results.
+The `convert-results-json-to-markdown.py` helps you put the benchmarking results inside a markdown table, by formatting [descriptions.md](tests/descriptions.md) with real benchmarking results.
 You can find the result presented as a table inside the `buildkite/performance-benchmark` job page.
 If you do not see the table, please wait till the benchmark finish running.
 The json version of the table (together with the json version of the benchmark) will be also attached to the markdown file.
diff --git a/.buildkite/nightly-benchmarks/scripts/convert-results-json-to-markdown.py b/.buildkite/nightly-benchmarks/scripts/convert-results-json-to-markdown.py
index a4f1638c1..1030ec24e 100644
--- a/.buildkite/nightly-benchmarks/scripts/convert-results-json-to-markdown.py
+++ b/.buildkite/nightly-benchmarks/scripts/convert-results-json-to-markdown.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import json
 import os
@@ -66,18 +65,18 @@ def read_markdown(file):
 
 
 def results_to_json(latency, throughput, serving):
-    return json.dumps(
-        {
-            "latency": latency.to_dict(),
-            "throughput": throughput.to_dict(),
-            "serving": serving.to_dict(),
-        }
-    )
+    return json.dumps({
+        'latency': latency.to_dict(),
+        'throughput': throughput.to_dict(),
+        'serving': serving.to_dict()
+    })
 
 
 if __name__ == "__main__":
+
     # collect results
     for test_file in results_folder.glob("*.json"):
+
         with open(test_file) as f:
             raw_result = json.loads(f.read())
 
@@ -121,8 +120,7 @@ if __name__ == "__main__":
             for perc in [10, 25, 50, 75, 90, 99]:
                 # Multiply 1000 to convert the time unit from s to ms
                 raw_result.update(
-                    {f"P{perc}": 1000 * raw_result["percentiles"][str(perc)]}
-                )
+                    {f"P{perc}": 1000 * raw_result["percentiles"][str(perc)]})
             raw_result["avg_latency"] = raw_result["avg_latency"] * 1000
 
             # add the result to raw_result
@@ -155,27 +153,26 @@ if __name__ == "__main__":
     serving_results = pd.DataFrame.from_dict(serving_results)
     throughput_results = pd.DataFrame.from_dict(throughput_results)
 
-    raw_results_json = results_to_json(
-        latency_results, throughput_results, serving_results
-    )
+    raw_results_json = results_to_json(latency_results, throughput_results,
+                                       serving_results)
 
     # remapping the key, for visualization purpose
     if not latency_results.empty:
-        latency_results = latency_results[list(latency_column_mapping.keys())].rename(
-            columns=latency_column_mapping
-        )
+        latency_results = latency_results[list(
+            latency_column_mapping.keys())].rename(
+                columns=latency_column_mapping)
     if not serving_results.empty:
-        serving_results = serving_results[list(serving_column_mapping.keys())].rename(
-            columns=serving_column_mapping
-        )
+        serving_results = serving_results[list(
+            serving_column_mapping.keys())].rename(
+                columns=serving_column_mapping)
     if not throughput_results.empty:
-        throughput_results = throughput_results[
-            list(throughput_results_column_mapping.keys())
-        ].rename(columns=throughput_results_column_mapping)
+        throughput_results = throughput_results[list(
+            throughput_results_column_mapping.keys())].rename(
+                columns=throughput_results_column_mapping)
 
-    processed_results_json = results_to_json(
-        latency_results, throughput_results, serving_results
-    )
+    processed_results_json = results_to_json(latency_results,
+                                             throughput_results,
+                                             serving_results)
 
     for df in [latency_results, serving_results, throughput_results]:
         if df.empty:
@@ -187,39 +184,38 @@ if __name__ == "__main__":
         # The GPUs sometimes come in format of "GPUTYPE\nGPUTYPE\n...",
         # we want to turn it into "8xGPUTYPE"
         df["GPU"] = df["GPU"].apply(
-            lambda x: f"{len(x.split('\n'))}x{x.split('\n')[0]}"
-        )
+            lambda x: f"{len(x.split('\n'))}x{x.split('\n')[0]}")
 
     # get markdown tables
-    latency_md_table = tabulate(
-        latency_results, headers="keys", tablefmt="pipe", showindex=False
-    )
-    serving_md_table = tabulate(
-        serving_results, headers="keys", tablefmt="pipe", showindex=False
-    )
-    throughput_md_table = tabulate(
-        throughput_results, headers="keys", tablefmt="pipe", showindex=False
-    )
+    latency_md_table = tabulate(latency_results,
+                                headers='keys',
+                                tablefmt='pipe',
+                                showindex=False)
+    serving_md_table = tabulate(serving_results,
+                                headers='keys',
+                                tablefmt='pipe',
+                                showindex=False)
+    throughput_md_table = tabulate(throughput_results,
+                                   headers='keys',
+                                   tablefmt='pipe',
+                                   showindex=False)
 
     # document the result
     with open(results_folder / "benchmark_results.md", "w") as f:
-        results = read_markdown(
-            "../.buildkite/nightly-benchmarks/"
-            + "performance-benchmarks-descriptions.md"
-        )
+
+        results = read_markdown("../.buildkite/nightly-benchmarks/" +
+                                "performance-benchmarks-descriptions.md")
         results = results.format(
             latency_tests_markdown_table=latency_md_table,
             throughput_tests_markdown_table=throughput_md_table,
             serving_tests_markdown_table=serving_md_table,
-            benchmarking_results_in_json_string=processed_results_json,
-        )
+            benchmarking_results_in_json_string=processed_results_json)
         f.write(results)
 
     # document benchmarking results in json
     with open(results_folder / "benchmark_results.json", "w") as f:
-        results = (
-            latency_results.to_dict(orient="records")
-            + throughput_results.to_dict(orient="records")
-            + serving_results.to_dict(orient="records")
-        )
+
+        results = latency_results.to_dict(
+            orient='records') + throughput_results.to_dict(
+                orient='records') + serving_results.to_dict(orient='records')
         f.write(json.dumps(results))
diff --git a/.buildkite/nightly-benchmarks/scripts/download-tokenizer.py b/.buildkite/nightly-benchmarks/scripts/download-tokenizer.py
index 8532ff7ef..5e17b79d2 100644
--- a/.buildkite/nightly-benchmarks/scripts/download-tokenizer.py
+++ b/.buildkite/nightly-benchmarks/scripts/download-tokenizer.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
 
@@ -15,12 +14,15 @@ def main(model, cachedir):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
-        description="Download and save Hugging Face tokenizer"
-    )
-    parser.add_argument("--model", type=str, required=True, help="Name of the model")
-    parser.add_argument(
-        "--cachedir", type=str, required=True, help="Directory to save the tokenizer"
-    )
+        description="Download and save Hugging Face tokenizer")
+    parser.add_argument("--model",
+                        type=str,
+                        required=True,
+                        help="Name of the model")
+    parser.add_argument("--cachedir",
+                        type=str,
+                        required=True,
+                        help="Directory to save the tokenizer")
 
     args = parser.parse_args()
     main(args.model, args.cachedir)
diff --git a/.buildkite/nightly-benchmarks/scripts/generate-nightly-markdown.py b/.buildkite/nightly-benchmarks/scripts/generate-nightly-markdown.py
index 053fd52c3..0ff95a091 100644
--- a/.buildkite/nightly-benchmarks/scripts/generate-nightly-markdown.py
+++ b/.buildkite/nightly-benchmarks/scripts/generate-nightly-markdown.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
 import json
@@ -12,33 +11,33 @@ from tabulate import tabulate
 
 def parse_arguments():
     parser = argparse.ArgumentParser(
-        description="Parse command line arguments for summary-nightly-results script."
-    )
-    parser.add_argument(
-        "--results-folder",
-        type=str,
-        required=True,
-        help="The folder where the results are stored.",
-    )
-    parser.add_argument(
-        "--description", type=str, required=True, help="Description of the results."
-    )
+        description=
+        'Parse command line arguments for summary-nightly-results script.')
+    parser.add_argument('--results-folder',
+                        type=str,
+                        required=True,
+                        help='The folder where the results are stored.')
+    parser.add_argument('--description',
+                        type=str,
+                        required=True,
+                        help='Description of the results.')
 
     args = parser.parse_args()
     return args
 
 
 def get_perf(df, method, model, metric):
+
     means = []
 
     for qps in [2, 4, 8, 16, "inf"]:
-        target = df["Test name"].str.contains(model)
-        target = target & df["Engine"].str.contains(method)
-        target = target & df["Test name"].str.contains("qps_" + str(qps))
+        target = df['Test name'].str.contains(model)
+        target = target & df['Engine'].str.contains(method)
+        target = target & df['Test name'].str.contains("qps_" + str(qps))
         filtered_df = df[target]
 
         if filtered_df.empty:
-            means.append(0.0)
+            means.append(0.)
         else:
             means.append(filtered_df[metric].values[0])
 
@@ -46,6 +45,7 @@ def get_perf(df, method, model, metric):
 
 
 def get_perf_w_std(df, method, model, metric):
+
     if metric in ["TTFT", "ITL"]:
         mean = get_perf(df, method, model, "Mean " + metric + " (ms)")
         mean = mean.tolist()
@@ -60,8 +60,7 @@ def get_perf_w_std(df, method, model, metric):
     else:
         assert metric == "Tput"
         mean = get_perf(df, method, model, "Input Tput (tok/s)") + get_perf(
-            df, method, model, "Output Tput (tok/s)"
-        )
+            df, method, model, "Output Tput (tok/s)")
         mean = mean.tolist()
         std = None
 
@@ -81,17 +80,18 @@ def main(args):
     # generate markdown table
     df = pd.DataFrame.from_dict(results)
 
-    md_table = tabulate(df, headers="keys", tablefmt="pipe", showindex=False)
+    md_table = tabulate(df, headers='keys', tablefmt='pipe', showindex=False)
 
     with open(args.description) as f:
         description = f.read()
 
-    description = description.format(nightly_results_benchmarking_table=md_table)
+    description = description.format(
+        nightly_results_benchmarking_table=md_table)
 
     with open("nightly_results.md", "w") as f:
         f.write(description)
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     args = parse_arguments()
     main(args)
diff --git a/.buildkite/nightly-benchmarks/scripts/get-lmdeploy-modelname.py b/.buildkite/nightly-benchmarks/scripts/get-lmdeploy-modelname.py
index ddea1d2b1..e5f179a0f 100644
--- a/.buildkite/nightly-benchmarks/scripts/get-lmdeploy-modelname.py
+++ b/.buildkite/nightly-benchmarks/scripts/get-lmdeploy-modelname.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from lmdeploy.serve.openai.api_client import APIClient
 
diff --git a/.buildkite/nightly-benchmarks/scripts/summary-nightly-results.py b/.buildkite/nightly-benchmarks/scripts/summary-nightly-results.py
index fb3b9d5e3..62ee5e10b 100644
--- a/.buildkite/nightly-benchmarks/scripts/summary-nightly-results.py
+++ b/.buildkite/nightly-benchmarks/scripts/summary-nightly-results.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import datetime
 import json
@@ -35,8 +34,10 @@ serving_column_mapping = {
 }
 
 if __name__ == "__main__":
+
     # collect results
     for test_file in results_folder.glob("*.json"):
+
         with open(test_file) as f:
             raw_result = json.loads(f.read())
 
@@ -55,16 +56,17 @@ if __name__ == "__main__":
     serving_results = pd.DataFrame.from_dict(serving_results)
 
     if not serving_results.empty:
-        serving_results = serving_results[list(serving_column_mapping.keys())].rename(
-            columns=serving_column_mapping
-        )
+        serving_results = serving_results[list(
+            serving_column_mapping.keys())].rename(
+                columns=serving_column_mapping)
 
-    serving_md_table_with_headers = tabulate(
-        serving_results, headers="keys", tablefmt="pipe", showindex=False
-    )
+    serving_md_table_with_headers = tabulate(serving_results,
+                                             headers='keys',
+                                             tablefmt='pipe',
+                                             showindex=False)
     # remove the first line of header
-    serving_md_table_lines = serving_md_table_with_headers.split("\n")
-    serving_md_table_without_header = "\n".join(serving_md_table_lines[2:])
+    serving_md_table_lines = serving_md_table_with_headers.split('\n')
+    serving_md_table_without_header = '\n'.join(serving_md_table_lines[2:])
 
     prefix = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
     prefix = prefix + "_" + os.environ.get("CURRENT_LLM_SERVING_ENGINE")
@@ -74,9 +76,10 @@ if __name__ == "__main__":
         # document results with header.
         # for those who wants to reproduce our benchmark.
         f.write(serving_md_table_with_headers)
-        f.write("\n")
+        f.write('\n')
 
     # document benchmarking results in json
     with open(results_folder / f"{prefix}_nightly_results.json", "w") as f:
-        results = serving_results.to_dict(orient="records")
+
+        results = serving_results.to_dict(orient='records')
         f.write(json.dumps(results))
diff --git a/.buildkite/pyproject.toml b/.buildkite/pyproject.toml
deleted file mode 100644
index d5cad1c73..000000000
--- a/.buildkite/pyproject.toml
+++ /dev/null
@@ -1,46 +0,0 @@
-# This local pyproject file is part of the migration from yapf to ruff format.
-# It uses the same core rules as the main pyproject.toml file, but with the
-# following differences:
-# - ruff line length is overridden to 88
-# - deprecated typing ignores (UP006, UP035) have been removed
-
-[tool.ruff]
-line-length = 88
-
-[tool.ruff.lint.per-file-ignores]
-"vllm/third_party/**" = ["ALL"]
-"vllm/version.py" = ["F401"]
-"vllm/_version.py" = ["ALL"]
-
-[tool.ruff.lint]
-select = [
-    # pycodestyle
-    "E",
-    # Pyflakes
-    "F",
-    # pyupgrade
-    "UP",
-    # flake8-bugbear
-    "B",
-    # flake8-simplify
-    "SIM",
-    # isort
-    "I",
-    # flake8-logging-format
-    "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.ruff.format]
-docstring-code-format = true
diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml
index 16b5ad029..a21a657c4 100644
--- a/.buildkite/release-pipeline.yaml
+++ b/.buildkite/release-pipeline.yaml
@@ -1,22 +1,20 @@
 steps:
-  - label: "Build wheel - CUDA 12.8"
-    id: build-wheel-cuda-12-8
+  - label: "Build wheel - CUDA 12.4"
     agents:
       queue: cpu_queue_postmerge
     commands:
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.4.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
       - "mkdir artifacts"
       - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
       - "bash .buildkite/scripts/upload-wheels.sh"
     env:
       DOCKER_BUILDKIT: "1"
 
-  - label: "Build wheel - CUDA 12.6"
-    id: build-wheel-cuda-12-6
+  - label: "Build wheel - CUDA 12.1"
     agents:
       queue: cpu_queue_postmerge
     commands:
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.6.3 --build-arg torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0+PTX' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.1.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
       - "mkdir artifacts"
       - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
       - "bash .buildkite/scripts/upload-wheels.sh"
@@ -30,11 +28,10 @@ steps:
 
   - label: "Build wheel - CUDA 11.8"
     # depends_on: block-build-cu118-wheel
-    id: build-wheel-cuda-11-8
     agents:
       queue: cpu_queue_postmerge
     commands:
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=11.8.0 --build-arg torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0+PTX' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=11.8.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
       - "mkdir artifacts"
       - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
       - "bash .buildkite/scripts/upload-wheels.sh"
@@ -47,49 +44,33 @@ steps:
 
   - label: "Build release image"
     depends_on: block-release-image-build
-    id: build-release-image
     agents:
       queue: cpu_queue_postmerge
     commands:
       - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT --target vllm-openai --progress plain -f docker/Dockerfile ."
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.4.0 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT --target vllm-openai --progress plain -f docker/Dockerfile ."
       - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
 
-  - label: "Annotate release workflow"
-    depends_on:
-      - build-release-image
-      - build-wheel-cuda-12-8
-      - build-wheel-cuda-12-6
-      - build-wheel-cuda-11-8
-    id: annotate-release-workflow
-    agents:
-      queue: cpu_queue_postmerge
-    commands:
-      - "bash .buildkite/scripts/annotate-release.sh"
-
   - label: "Build and publish TPU release image"
     depends_on: ~
     if: build.env("NIGHTLY") == "1"
     agents:
       queue: tpu_queue_postmerge
     commands:
-      - "yes | docker system prune -a"
-      - "git fetch --all"
       - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --tag vllm/vllm-tpu:nightly --tag vllm/vllm-tpu:$BUILDKITE_COMMIT --progress plain -f docker/Dockerfile.tpu ."
       - "docker push vllm/vllm-tpu:nightly"
       - "docker push vllm/vllm-tpu:$BUILDKITE_COMMIT"
     plugins:
       - docker-login#v3.0.0:
-          username: vllmbot
+          username: vllm
           password-env: DOCKERHUB_TOKEN
     env:
       DOCKER_BUILDKIT: "1"
 
   - input: "Provide Release version here"
-    id: input-release-version
     fields:
       - text: "What is the release version?"
-        key: release-version
+        key: "release-version"
 
   - block: "Build CPU release image"
     key: block-cpu-release-image-build
diff --git a/.buildkite/scripts/annotate-release.sh b/.buildkite/scripts/annotate-release.sh
deleted file mode 100755
index 94e0ac239..000000000
--- a/.buildkite/scripts/annotate-release.sh
+++ /dev/null
@@ -1,31 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-# Get release version and strip leading 'v' if present
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version | sed 's/^v//')
-
-if [ -z "$RELEASE_VERSION" ]; then
-  echo "Error: RELEASE_VERSION is empty. 'release-version' metadata might not be set or is invalid."
-  exit 1
-fi
-
-buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
-To download the wheel:
-\`\`\`
-aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu126/vllm-${RELEASE_VERSION}+cu126-cp38-abi3-manylinux1_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu118/vllm-${RELEASE_VERSION}+cu118-cp38-abi3-manylinux1_x86_64.whl . 
-\`\`\`
-
-To download and upload the image:
-
-\`\`\`
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT} vllm/vllm-openai
-docker tag vllm/vllm-openai vllm/vllm-openai:latest
-docker tag vllm/vllm-openai vllm/vllm-openai:v${RELEASE_VERSION}
-docker push vllm/vllm-openai:latest
-docker push vllm/vllm-openai:v${RELEASE_VERSION}
-\`\`\`
-EOF 
\ No newline at end of file
diff --git a/.buildkite/scripts/ci-clean-log.sh b/.buildkite/scripts/ci-clean-log.sh
deleted file mode 100644
index 69d8a3a28..000000000
--- a/.buildkite/scripts/ci-clean-log.sh
+++ /dev/null
@@ -1,17 +0,0 @@
-#!/bin/bash
-# Usage: ./ci_clean_log.sh ci.log
-# This script strips timestamps and color codes from CI log files.
-
-# Check if argument is given
-if [ $# -lt 1 ]; then
-    echo "Usage: $0 ci.log"
-    exit 1
-fi
-
-INPUT_FILE="$1"
-
-# Strip timestamps
-sed -i 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE"
-
-# Strip colorization
-sed -i -r 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE"
diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh
index 6e9af1e72..368f30434 100755
--- a/.buildkite/scripts/hardware_ci/run-amd-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh
@@ -3,9 +3,6 @@
 # This script runs test inside the corresponding ROCm docker container.
 set -o pipefail
 
-# Export Python path
-export PYTHONPATH=".."
-
 # Print ROCm version
 echo "--- Confirming Clean Initial State"
 while true; do
@@ -77,73 +74,38 @@ HF_MOUNT="/root/.cache/huggingface"
 
 commands=$@
 echo "Commands:$commands"
-
-if [[ $commands == *"pytest -v -s basic_correctness/test_basic_correctness.py"* ]]; then
-  commands=${commands//"pytest -v -s basic_correctness/test_basic_correctness.py"/"VLLM_USE_TRITON_FLASH_ATTN=0 pytest -v -s basic_correctness/test_basic_correctness.py"}
-fi
-
-if [[ $commands == *"pytest -v -s models/test_registry.py"* ]]; then
-  commands=${commands//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
-fi
-
-if [[ $commands == *"VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2'"* ]]; then
-  commands=${commands//"VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2'"/"VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2 and not BambaForCausalLM and not Gemma2ForCausalLM and not Grok1ModelForCausalLM and not Zamba2ForCausalLM and not Gemma2Model and not GritLM'"}
-fi
-
-if [[ $commands == *"pytest -v -s compile/test_basic_correctness.py"* ]]; then
-  commands=${commands//"pytest -v -s compile/test_basic_correctness.py"/"VLLM_USE_TRITON_FLASH_ATTN=0 pytest -v -s compile/test_basic_correctness.py"}
-fi
-
-if [[ $commands == *"pytest -v -s lora"* ]]; then
-  commands=${commands//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}
-fi
-
 #ignore certain kernels tests
-if [[ $commands == *" kernels/core"* ]]; then
+if [[ $commands == *" kernels "* ]]; then
   commands="${commands} \
-  --ignore=kernels/core/test_fused_quant_layernorm.py \
-  --ignore=kernels/core/test_permute_cols.py"
-fi
-
-if [[ $commands == *" kernels/attention"* ]]; then
-  commands="${commands} \
-  --ignore=kernels/attention/stest_attention_selector.py \
-  --ignore=kernels/attention/test_blocksparse_attention.py \
-  --ignore=kernels/attention/test_encoder_decoder_attn.py \
-  --ignore=kernels/attention/test_attention_selector.py \
-  --ignore=kernels/attention/test_flash_attn.py \
-  --ignore=kernels/attention/test_flashinfer.py \
-  --ignore=kernels/attention/test_prefix_prefill.py \
-  --ignore=kernels/attention/test_cascade_flash_attn.py \
-  --ignore=kernels/attention/test_mha_attn.py \
-  --ignore=kernels/attention/test_lightning_attn.py \
-  --ignore=kernels/attention/test_attention.py"
-fi
-
-if [[ $commands == *" kernels/quantization"* ]]; then
-  commands="${commands} \
-  --ignore=kernels/quantization/test_int8_quant.py \
-  --ignore=kernels/quantization/test_aqlm.py \
-  --ignore=kernels/quantization/test_machete_mm.py \
-  --ignore=kernels/quantization/test_block_fp8.py \
-  --ignore=kernels/quantization/test_block_int8.py \
-  --ignore=kernels/quantization/test_marlin_gemm.py \
-  --ignore=kernels/quantization/test_cutlass_scaled_mm.py \
-  --ignore=kernels/quantization/test_int8_kernel.py"
-fi
-
-if [[ $commands == *" kernels/mamba"* ]]; then
-  commands="${commands} \
-  --ignore=kernels/mamba/test_mamba_mixer2.py \
-  --ignore=kernels/mamba/test_causal_conv1d.py \
-  --ignore=kernels/mamba/test_mamba_ssm_ssd.py"
-fi
-
-if [[ $commands == *" kernels/moe"* ]]; then
-  commands="${commands} \
-  --ignore=kernels/moe/test_moe.py \
-  --ignore=kernels/moe/test_cutlass_moe.py \
-  --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
+  --ignore=kernels/test_attention_selector.py \
+  --ignore=kernels/test_blocksparse_attention.py \
+  --ignore=kernels/test_causal_conv1d.py \
+  --ignore=kernels/test_cutlass.py \
+  --ignore=kernels/test_encoder_decoder_attn.py \
+  --ignore=kernels/test_flash_attn.py \
+  --ignore=kernels/test_flashinfer.py \
+  --ignore=kernels/test_int8_quant.py \
+  --ignore=kernels/test_machete_gemm.py \
+  --ignore=kernels/test_mamba_ssm.py \
+  --ignore=kernels/test_marlin_gemm.py \
+  --ignore=kernels/test_moe.py \
+  --ignore=kernels/test_prefix_prefill.py \
+  --ignore=kernels/test_rand.py \
+  --ignore=kernels/test_sampler.py \
+  --ignore=kernels/test_cascade_flash_attn.py \
+  --ignore=kernels/test_mamba_mixer2.py \
+  --ignore=kernels/test_aqlm.py \
+  --ignore=kernels/test_machete_mm.py \
+  --ignore=kernels/test_mha_attn.py \
+  --ignore=kernels/test_block_fp8.py \
+  --ignore=kernels/test_cutlass_moe.py \
+  --ignore=kernels/test_mamba_ssm_ssd.py \
+  --ignore=kernels/test_attention.py \
+  --ignore=kernels/test_block_int8.py \
+  --ignore=kernels/test_fused_quant_layernorm.py \
+  --ignore=kernels/test_int8_kernel.py \
+  --ignore=kernels/test_triton_moe_ptpc_fp8.py \
+  --ignore=kernels/test_permute_cols.py"
 fi
 
 #ignore certain Entrypoints/openai tests
@@ -185,8 +147,6 @@ fi
 
 
 PARALLEL_JOB_COUNT=8
-MYPYTHONPATH=".."
-
 # check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs. 
 if [[ $commands == *"--shard-id="* ]]; then
   # assign job count as the number of shards used   
@@ -207,7 +167,6 @@ if [[ $commands == *"--shard-id="* ]]; then
         -e AWS_SECRET_ACCESS_KEY \
         -v "${HF_CACHE}:${HF_MOUNT}" \
         -e "HF_HOME=${HF_MOUNT}" \
-        -e "PYTHONPATH=${MYPYTHONPATH}" \
         --name "${container_name}_${GPU}" \
         "${image_name}" \
         /bin/bash -c "${commands_gpu}" \
@@ -238,7 +197,6 @@ else
           -e AWS_SECRET_ACCESS_KEY \
           -v "${HF_CACHE}:${HF_MOUNT}" \
           -e "HF_HOME=${HF_MOUNT}" \
-          -e "PYTHONPATH=${MYPYTHONPATH}" \
           --name "${container_name}" \
           "${image_name}" \
           /bin/bash -c "${commands}"
diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh
index 36bcb015d..5d863dd82 100755
--- a/.buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh
+++ b/.buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh
@@ -7,7 +7,6 @@ set -ex
 # Setup cleanup
 remove_docker_container() {
   if [[ -n "$container_id" ]]; then
-      podman stop --all -t0
       podman rm -f "$container_id" || true
   fi
   podman system prune -f
@@ -33,12 +32,9 @@ function cpu_tests() {
     set -e
     pip install pytest pytest-asyncio einops peft Pillow soundfile transformers_stream_generator matplotlib
     pip install sentence-transformers datamodel_code_generator
-    pytest -v -s tests/models/language/generation/test_bart.py -m cpu_model
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-openai-community/gpt2]
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-facebook/opt-125m]
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-google/gemma-1.1-2b-it]
-    pytest -v -s tests/models/language/pooling/test_classification.py::test_models[float-jason9693/Qwen2.5-1.5B-apeach]
-    pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model"
+    pytest -v -s tests/models/embedding/language/test_cls_models.py::test_classification_models[float-jason9693/Qwen2.5-1.5B-apeach]
+    pytest -v -s tests/models/embedding/language/test_embedding.py::test_models[half-BAAI/bge-base-en-v1.5]
+    pytest -v -s tests/models/encoder_decoder/language -m cpu_model"
 }
 
 # All of CPU tests are expected to be finished less than 40 mins.
diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh
index bbcde4009..40f3df960 100644
--- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh
@@ -6,70 +6,72 @@ set -ex
 
 # allow to bind to different cores
 CORE_RANGE=${CORE_RANGE:-48-95}
-OMP_CORE_RANGE=${OMP_CORE_RANGE:-48-95}
 NUMA_NODE=${NUMA_NODE:-1}
 
-export CMAKE_BUILD_PARALLEL_LEVEL=32
-
 # Setup cleanup
 remove_docker_container() { 
     set -e; 
-    docker rm -f cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"-avx2 || true; 
+    docker rm -f cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2-"$NUMA_NODE" || true; 
+    docker image rm cpu-test-"$BUILDKITE_BUILD_NUMBER" cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2 || true; 
 }
 trap remove_docker_container EXIT
 remove_docker_container
 
 # Try building the docker image
-numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .
-numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" --tag cpu-test-"$NUMA_NODE"-avx2 --target vllm-test -f docker/Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --tag cpu-test-"$BUILDKITE_BUILD_NUMBER" --target vllm-test -f docker/Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" --tag cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2 --target vllm-test -f docker/Dockerfile.cpu .
 
 # Run the image, setting --shm-size=4g for tensor parallel.
-docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --env VLLM_CPU_OMP_THREADS_BIND="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"
-docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --env VLLM_CPU_OMP_THREADS_BIND="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE"-avx2 cpu-test-"$NUMA_NODE"-avx2
+docker run -itd --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --cpuset-cpus="$CORE_RANGE"  \
+ --cpuset-mems="$NUMA_NODE" --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --shm-size=4g --name cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" cpu-test-"$BUILDKITE_BUILD_NUMBER"
+docker run -itd --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --cpuset-cpus="$CORE_RANGE" \
+ --cpuset-mems="$NUMA_NODE" --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --shm-size=4g --name cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2-"$NUMA_NODE" cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2
 
 function cpu_tests() {
   set -e
   export NUMA_NODE=$2
+  export BUILDKITE_BUILD_NUMBER=$3
 
   # offline inference
-  docker exec cpu-test-"$NUMA_NODE"-avx2 bash -c "
+  docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2-"$NUMA_NODE" bash -c "
     set -e
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"
 
   # Run basic model test
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
+  docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
     set -e
-    pytest -v -s tests/kernels/attention/test_cache.py -m cpu_model
-    pytest -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
-    pytest -v -s tests/models/language/generation -m cpu_model
-    pytest -v -s tests/models/language/pooling -m cpu_model
-    pytest -v -s tests/models/multimodal/generation \
-                --ignore=tests/models/multimodal/generation/test_mllama.py \
-                --ignore=tests/models/multimodal/generation/test_pixtral.py \
-                -m cpu_model"
+    pytest -v -s tests/kernels/test_cache.py -m cpu_model
+    pytest -v -s tests/kernels/test_mla_decode_cpu.py -m cpu_model
+    pytest -v -s tests/models/decoder_only/language -m cpu_model
+    pytest -v -s tests/models/embedding/language -m cpu_model
+    pytest -v -s tests/models/encoder_decoder/language -m cpu_model
+    pytest -v -s tests/models/decoder_only/audio_language -m cpu_model
+    pytest -v -s tests/models/decoder_only/vision_language -m cpu_model"
 
   # Run compressed-tensor test
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
+  docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
     set -e
     pytest -s -v \
     tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_static_setup \
     tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_dynamic_per_token"
 
   # Run AWQ test
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
+  docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
     set -e
-    VLLM_USE_V1=0 pytest -s -v \
+    pytest -s -v \
     tests/quantization/test_ipex_quant.py"
 
   # Run chunked-prefill and prefix-cache test
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
+  docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
     set -e
     pytest -s -v -k cpu_model \
     tests/basic_correctness/test_chunked_prefill.py"  
 
   # online serving
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
+  docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
     set -e
+    export VLLM_CPU_KVCACHE_SPACE=10 
+    export VLLM_CPU_OMP_THREADS_BIND=$1
     python3 -m vllm.entrypoints.openai.api_server --model facebook/opt-125m --dtype half & 
     timeout 600 bash -c 'until curl localhost:8000/v1/models; do sleep 1; done' || exit 1
     python3 benchmarks/benchmark_serving.py \
@@ -81,7 +83,7 @@ function cpu_tests() {
       --tokenizer facebook/opt-125m"
 
   # Run multi-lora tests
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
+  docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
     set -e
     pytest -s -v \
     tests/lora/test_qwen2vl.py"
@@ -89,4 +91,4 @@ function cpu_tests() {
 
 # All of CPU tests are expected to be finished less than 40 mins.
 export -f cpu_tests
-timeout 1h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
+timeout 40m bash -c "cpu_tests $CORE_RANGE $NUMA_NODE $BUILDKITE_BUILD_NUMBER"
diff --git a/.buildkite/scripts/hardware_ci/run-hpu-test.sh b/.buildkite/scripts/hardware_ci/run-hpu-test.sh
index 5efac3ddf..95b6ac37f 100644
--- a/.buildkite/scripts/hardware_ci/run-hpu-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-hpu-test.sh
@@ -10,17 +10,15 @@ docker build -t hpu-test-env -f docker/Dockerfile.hpu .
 # Setup cleanup
 # certain versions of HPU software stack have a bug that can
 # override the exit code of the script, so we need to use
-# separate remove_docker_containers and remove_docker_containers_and_exit
+# separate remove_docker_container and remove_docker_container_and_exit
 # functions, while other platforms only need one remove_docker_container
 # function.
 EXITCODE=1
-remove_docker_containers() { docker rm -f hpu-test || true; docker rm -f hpu-test-tp2 || true; }
-remove_docker_containers_and_exit() { remove_docker_containers; exit $EXITCODE; }
-trap remove_docker_containers_and_exit EXIT
-remove_docker_containers
+remove_docker_container() { docker rm -f hpu-test || true; }
+remove_docker_container_and_exit() { remove_docker_container; exit $EXITCODE; }
+trap remove_docker_container_and_exit EXIT
+remove_docker_container
 
 # Run the image and launch offline inference
 docker run --runtime=habana --name=hpu-test --network=host -e HABANA_VISIBLE_DEVICES=all -e VLLM_SKIP_WARMUP=true --entrypoint="" hpu-test-env python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m
-docker run --runtime=habana --name=hpu-test-tp2 --network=host -e HABANA_VISIBLE_DEVICES=all -e VLLM_SKIP_WARMUP=true --entrypoint="" hpu-test-env python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --tensor-parallel-size 2
-
 EXITCODE=$?
diff --git a/.buildkite/scripts/hardware_ci/run-neuron-test.sh b/.buildkite/scripts/hardware_ci/run-neuron-test.sh
index 3d294ea5f..ec6a080eb 100644
--- a/.buildkite/scripts/hardware_ci/run-neuron-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-neuron-test.sh
@@ -11,14 +11,13 @@ container_name="neuron_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
 HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
-HF_TOKEN=$(aws secretsmanager get-secret-value  --secret-id "ci/vllm-neuron/hf-token" --region us-west-2 --query 'SecretString' --output text | jq -r .VLLM_NEURON_CI_HF_TOKEN)
 
 NEURON_COMPILE_CACHE_URL="$(realpath ~)/neuron_compile_cache"
 mkdir -p "${NEURON_COMPILE_CACHE_URL}"
 NEURON_COMPILE_CACHE_MOUNT="/root/.cache/neuron_compile_cache"
 
 # Try building the docker image
-aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws
+aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-west-2.amazonaws.com
 
 # prune old image and containers to save disk space, and only once a day
 # by using a timestamp file in tmp.
@@ -48,16 +47,8 @@ trap remove_docker_container EXIT
 docker run --rm -it --device=/dev/neuron0 --network bridge \
        -v "${HF_CACHE}:${HF_MOUNT}" \
        -e "HF_HOME=${HF_MOUNT}" \
-       -e "HF_TOKEN=${HF_TOKEN}" \
        -v "${NEURON_COMPILE_CACHE_URL}:${NEURON_COMPILE_CACHE_MOUNT}" \
        -e "NEURON_COMPILE_CACHE_URL=${NEURON_COMPILE_CACHE_MOUNT}" \
        --name "${container_name}" \
        ${image_name} \
-       /bin/bash -c "
-            python3 /workspace/vllm/examples/offline_inference/neuron.py;
-            python3 -m pytest /workspace/vllm/tests/neuron/1_core/ -v --capture=tee-sys;
-            for f in /workspace/vllm/tests/neuron/2_core/*.py; do
-                echo 'Running test file: '$f;
-                python3 -m pytest \$f -v --capture=tee-sys;
-            done
-       "
\ No newline at end of file
+       /bin/bash -c "python3 /workspace/vllm/examples/offline_inference/neuron.py && python3 -m pytest /workspace/vllm/tests/neuron/1_core/ -v --capture=tee-sys && python3 -m pytest /workspace/vllm/tests/neuron/2_core/ -v --capture=tee-sys"
diff --git a/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh b/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
index a2a5c2a02..21982b01b 100755
--- a/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
@@ -1,185 +1,54 @@
 #!/bin/bash
 
-set -xu
-
-
-remove_docker_container() { 
-    docker rm -f tpu-test || true; 
-    docker rm -f vllm-tpu || true;
-}
-
-trap remove_docker_container EXIT
-
-# Remove the container that might not be cleaned up in the previous run.
-remove_docker_container
+set -xue
 
 # Build the docker image.
 docker build -f docker/Dockerfile.tpu -t vllm-tpu .
 
 # Set up cleanup.
-cleanup_docker() {
-  # Get Docker's root directory
-  docker_root=$(docker info -f '{{.DockerRootDir}}')
-  if [ -z "$docker_root" ]; then
-    echo "Failed to determine Docker root directory."
-    exit 1
-  fi
-  echo "Docker root directory: $docker_root"
-  # Check disk usage of the filesystem where Docker's root directory is located
-  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
-  # Define the threshold
-  threshold=70
-  if [ "$disk_usage" -gt "$threshold" ]; then
-    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
-    # Remove dangling images (those that are not tagged and not used by any container)
-    docker image prune -f
-    # Remove unused volumes / force the system prune for old images as well.
-    docker volume prune -f && docker system prune --force --filter "until=72h" --all
-    echo "Docker images and volumes cleanup completed."
-  else
-    echo "Disk usage is below $threshold%. No cleanup needed."
-  fi
-}
-cleanup_docker
+remove_docker_container() { docker rm -f tpu-test || true; }
+trap remove_docker_container EXIT
+# Remove the container that might not be cleaned up in the previous run.
+remove_docker_container
 
 # For HF_TOKEN.
 source /etc/environment
-
+# Run a simple end-to-end example.
 docker run --privileged --net host --shm-size=16G -it \
     -e "HF_TOKEN=$HF_TOKEN" --name tpu-test \
-    vllm-tpu /bin/bash -c '
-set -e # Exit immediately if a command exits with a non-zero status.
-set -u # Treat unset variables as an error.
-
-echo "--- Starting script inside Docker container ---"
-
-# Create results directory
-RESULTS_DIR=$(mktemp -d)
-# If mktemp fails, set -e will cause the script to exit.
-echo "Results will be stored in: $RESULTS_DIR"
-
-# Install dependencies
-echo "--- Installing Python dependencies ---"
-python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
-    && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
-    && python3 -m pip install --progress-bar off lm_eval[api]==0.4.4
-echo "--- Python dependencies installed ---"
-export VLLM_USE_V1=1
-export VLLM_XLA_CHECK_RECOMPILATION=1
-export VLLM_XLA_CACHE_PATH=
-echo "Using VLLM V1"
-
-echo "--- Hardware Information ---"
-tpu-info
-echo "--- Starting Tests ---"
-set +e
-overall_script_exit_code=0
-
-# --- Test Definitions ---
-# If a test fails, this function will print logs and will not cause the main script to exit.
-run_test() {
-    local test_num=$1
-    local test_name=$2
-    local test_command=$3
-    local log_file="$RESULTS_DIR/test_${test_num}.log"
-    local actual_exit_code
-
-    echo "--- TEST_$test_num: Running $test_name ---"
-    
-    # Execute the test command.
-    eval "$test_command" > >(tee -a "$log_file") 2> >(tee -a "$log_file" >&2)
-    actual_exit_code=$?
-
-    echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" # This goes to main log
-    echo "TEST_${test_num}_COMMAND_EXIT_CODE: $actual_exit_code" >> "$log_file" # Also to per-test log
-
-    if [ "$actual_exit_code" -ne 0 ]; then
-        echo "TEST_$test_num ($test_name) FAILED with exit code $actual_exit_code." >&2
-        echo "--- Log for failed TEST_$test_num ($test_name) ---" >&2
-        if [ -f "$log_file" ]; then
-            cat "$log_file" >&2
-        else
-            echo "Log file $log_file not found for TEST_$test_num ($test_name)." >&2
-        fi
-        echo "--- End of log for TEST_$test_num ($test_name) ---" >&2
-        return "$actual_exit_code" # Return the failure code
-    else
-        echo "TEST_$test_num ($test_name) PASSED."
-        return 0 # Return success
-    fi
-}
-
-# Helper function to call run_test and update the overall script exit code
-run_and_track_test() {
-    local test_num_arg="$1"
-    local test_name_arg="$2"
-    local test_command_arg="$3"
-
-    # Run the test
-    run_test "$test_num_arg" "$test_name_arg" "$test_command_arg"
-    local test_specific_exit_code=$?
-
-    # If the test failed, set the overall script exit code to 1
-    if [ "$test_specific_exit_code" -ne 0 ]; then
-        # No need for extra echo here, run_test already logged the failure.
-        overall_script_exit_code=1
-    fi
-}
-
-# --- Actual Test Execution ---
-run_and_track_test 0 "test_perf.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_perf.py"
-run_and_track_test 1 "test_compilation.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_compilation.py"
-run_and_track_test 2 "test_basic.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_basic.py"
-run_and_track_test 3 "test_accuracy.py::test_lm_eval_accuracy_v1_engine" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine"
-run_and_track_test 4 "test_quantization_accuracy.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py"
-run_and_track_test 5 "examples/offline_inference/tpu.py" \
-    "python3 /workspace/vllm/examples/offline_inference/tpu.py"
-run_and_track_test 6 "test_tpu_model_runner.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py"
-run_and_track_test 7 "test_sampler.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py"
-run_and_track_test 8 "test_topk_topp_sampler.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_topk_topp_sampler.py"
-run_and_track_test 9 "test_multimodal.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_multimodal.py"
-run_and_track_test 10 "test_pallas.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py"
-run_and_track_test 11 "test_struct_output_generate.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\""
-run_and_track_test 12 "test_moe_pallas.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py"
-run_and_track_test 13 "test_lora.py" \
-    "VLLM_XLA_CHECK_RECOMPILATION=0 python3 -m pytest -s -v /workspace/vllm/tests/tpu/lora/test_lora.py"
-run_and_track_test 14 "test_tpu_qkv_linear.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_qkv_linear.py"
-run_and_track_test 15 "test_spmd_model_weight_loading.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_spmd_model_weight_loading.py"
-
-# After all tests have been attempted, exit with the overall status.
-if [ "$overall_script_exit_code" -ne 0 ]; then
-    echo "--- One or more tests FAILED. Overall script exiting with failure code 1. ---"
-else
-    echo "--- All tests have completed and PASSED. Overall script exiting with success code 0. ---"
-fi
-exit "$overall_script_exit_code"
-' # IMPORTANT: This is the closing single quote for the bash -c "..." command. Ensure it is present and correct.
+    vllm-tpu /bin/bash -c "python3 -m pip install git+https://github.com/thuml/depyf.git \
+    && python3 -m pip install pytest pytest-asyncio tpu-info \
+    && python3 -m pip install lm_eval[api]==0.4.4 \
+    && export VLLM_XLA_CACHE_PATH= \
+    && export VLLM_USE_V1=1 \
+    && export VLLM_XLA_CHECK_RECOMPILATION=1 \
+    && echo HARDWARE \
+    && tpu-info \
+    && echo TEST_0 \
+    && pytest -v -s /workspace/vllm/tests/v1/tpu/test_perf.py \
+    && echo TEST_1 \
+    && pytest -v -s /workspace/vllm/tests/tpu/test_compilation.py \
+    && echo TEST_2 \
+    && pytest -v -s /workspace/vllm/tests/v1/tpu/test_basic.py \
+    && echo TEST_3 \
+    && pytest -v -s /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine \
+    && echo TEST_4 \
+    && pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py \
+    && echo TEST_5 \
+    && python3 /workspace/vllm/examples/offline_inference/tpu.py \
+    && echo TEST_6 \
+    && pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py \
+    && echo TEST_7 \
+    && pytest -s -v /workspace/vllm/tests/v1/tpu/test_sampler.py \
+    && echo TEST_8 \
+    && pytest -s -v /workspace/vllm/tests/v1/tpu/test_topk_topp_sampler.py \
+    && echo TEST_9 \
+    && pytest -s -v /workspace/vllm/tests/v1/tpu/test_multimodal.py \
+    && echo TEST_10 \
+    && pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py \
+    && echo TEST_11 \
+    && pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py" \
 
-# Capture the exit code of the docker run command
-DOCKER_RUN_EXIT_CODE=$?
 
-# The trap will run for cleanup.
-# Exit the main script with the Docker run command's exit code.
-if [ "$DOCKER_RUN_EXIT_CODE" -ne 0 ]; then
-    echo "Docker run command failed with exit code $DOCKER_RUN_EXIT_CODE."
-    exit "$DOCKER_RUN_EXIT_CODE"
-else
-    echo "Docker run command completed successfully."
-    exit 0
-fi
 # TODO: This test fails because it uses RANDOM_SEED sampling
-# pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \
+# && VLLM_USE_V1=1 pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \
diff --git a/.buildkite/scripts/rerun-test.sh b/.buildkite/scripts/rerun-test.sh
deleted file mode 100644
index d79c0d5f3..000000000
--- a/.buildkite/scripts/rerun-test.sh
+++ /dev/null
@@ -1,18 +0,0 @@
-#!/bin/bash
-
-# Usage: ./rerun_test.sh path/to/test.py::test_name
-
-# Check if argument is given
-if [ $# -lt 1 ]; then
-    echo "Usage: $0 path/to/test.py::test_name"
-    echo "Example: $0 tests/v1/engine/test_engine_core_client.py::test_kv_cache_events[True-tcp]"
-    exit 1
-fi
-
-TEST=$1
-COUNT=1
-
-while pytest -sv "$TEST"; do
-    COUNT=$((COUNT + 1))
-    echo "RUN NUMBER ${COUNT}"
-done
diff --git a/.buildkite/scripts/tpu/cleanup_docker.sh b/.buildkite/scripts/tpu/cleanup_docker.sh
deleted file mode 100755
index 209d9c434..000000000
--- a/.buildkite/scripts/tpu/cleanup_docker.sh
+++ /dev/null
@@ -1,24 +0,0 @@
-#!/bin/bash
-
-set -euo pipefail
-
-docker_root=$(docker info -f '{{.DockerRootDir}}')
-if [ -z "$docker_root" ]; then
-  echo "Failed to determine Docker root directory."
-  exit 1
-fi
-echo "Docker root directory: $docker_root"
-# Check disk usage of the filesystem where Docker's root directory is located
-disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
-# Define the threshold
-threshold=70
-if [ "$disk_usage" -gt "$threshold" ]; then
-  echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
-  # Remove dangling images (those that are not tagged and not used by any container)
-  docker image prune -f
-  # Remove unused volumes / force the system prune for old images as well.
-  docker volume prune -f && docker system prune --force --filter "until=72h" --all
-  echo "Docker images and volumes cleanup completed."
-else
-  echo "Disk usage is below $threshold%. No cleanup needed."
-fi
diff --git a/.buildkite/scripts/tpu/config_v6e_1.env b/.buildkite/scripts/tpu/config_v6e_1.env
deleted file mode 100644
index 441758647..000000000
--- a/.buildkite/scripts/tpu/config_v6e_1.env
+++ /dev/null
@@ -1,14 +0,0 @@
-# Environment config
-TEST_NAME=llama8b
-CONTAINER_NAME=vllm-tpu
-
-# vllm config
-MODEL=meta-llama/Llama-3.1-8B-Instruct
-MAX_NUM_SEQS=512
-MAX_NUM_BATCHED_TOKENS=512
-TENSOR_PARALLEL_SIZE=1
-MAX_MODEL_LEN=2048
-DOWNLOAD_DIR=/mnt/disks/persist
-EXPECTED_THROUGHPUT=8.0
-INPUT_LEN=1800
-OUTPUT_LEN=128
diff --git a/.buildkite/scripts/tpu/docker_run_bm.sh b/.buildkite/scripts/tpu/docker_run_bm.sh
deleted file mode 100755
index 6705da03e..000000000
--- a/.buildkite/scripts/tpu/docker_run_bm.sh
+++ /dev/null
@@ -1,102 +0,0 @@
-#!/bin/bash
-
-if [ ! -f "$1" ]; then
-  echo "Error: The env file '$1' does not exist."
-  exit 1  # Exit the script with a non-zero status to indicate an error
-fi
-
-ENV_FILE=$1
-
-# For testing on local vm, use `set -a` to export all variables
-source /etc/environment
-source $ENV_FILE
-
-remove_docker_container() { 
-    docker rm -f tpu-test || true; 
-    docker rm -f vllm-tpu || true;
-    docker rm -f $CONTAINER_NAME || true;
-}
-
-trap remove_docker_container EXIT
-
-# Remove the container that might not be cleaned up in the previous run.
-remove_docker_container
-
-# Build docker image.
-# TODO: build the image outside the script and share the image with other
-# tpu test if building time is too long.
-DOCKER_BUILDKIT=1 docker build \
-  --build-arg max_jobs=16 \
-  --build-arg USE_SCCACHE=1 \
-  --build-arg GIT_REPO_CHECK=0 \
-  --tag vllm/vllm-tpu-bm \
-  --progress plain -f docker/Dockerfile.tpu .
-
-LOG_ROOT=$(mktemp -d)
-# If mktemp fails, set -e will cause the script to exit.
-echo "Results will be stored in: $LOG_ROOT"
-
-if [ -z "$HF_TOKEN" ]; then
-  echo "Error: HF_TOKEN is not set or is empty."  
-  exit 1
-fi
-
-# Make sure mounted disk or dir exists
-if [ ! -d "$DOWNLOAD_DIR" ]; then
-    echo "Error: Folder $DOWNLOAD_DIR does not exist. This is useually a mounted drive. If no mounted drive, just create a folder."
-    exit 1
-fi
-
-echo "Run model $MODEL"
-echo
-
-echo "starting docker...$CONTAINER_NAME"
-echo    
-docker run \
- -v $DOWNLOAD_DIR:$DOWNLOAD_DIR \
- --env-file $ENV_FILE \
- -e HF_TOKEN="$HF_TOKEN" \
- -e TARGET_COMMIT=$BUILDKITE_COMMIT \
- -e MODEL=$MODEL \
- -e WORKSPACE=/workspace \
- --name $CONTAINER_NAME \
- -d \
- --privileged \
- --network host \
- -v /dev/shm:/dev/shm \
- vllm/vllm-tpu-bm tail -f /dev/null
-
-echo "run script..."
-echo
-docker exec "$CONTAINER_NAME" /bin/bash -c ".buildkite/scripts/hardware_ci/run_bm.sh"
-
-echo "copy result back..."
-VLLM_LOG="$LOG_ROOT/$TEST_NAME"_vllm_log.txt
-BM_LOG="$LOG_ROOT/$TEST_NAME"_bm_log.txt
-docker cp "$CONTAINER_NAME:/workspace/vllm_log.txt" "$VLLM_LOG" 
-docker cp "$CONTAINER_NAME:/workspace/bm_log.txt" "$BM_LOG"
-
-throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g')
-echo "throughput for $TEST_NAME at $BUILDKITE_COMMIT: $throughput"
-
-if [ "$BUILDKITE" = "true" ]; then
-  echo "Running inside Buildkite"
-  buildkite-agent artifact upload "$VLLM_LOG" 
-  buildkite-agent artifact upload "$BM_LOG"
-else
-  echo "Not running inside Buildkite"
-fi
-
-#
-# compare the throughput with EXPECTED_THROUGHPUT 
-# and assert meeting the expectation
-# 
-if [[ -z "$throughput" || ! "$throughput" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
-  echo "Failed to get the throughput"
-  exit 1
-fi
-
-if (( $(echo "$throughput < $EXPECTED_THROUGHPUT" | bc -l) )); then
-  echo "Error: throughput($throughput) is less than expected($EXPECTED_THROUGHPUT)"
-  exit 1
-fi
diff --git a/.buildkite/scripts/tpu/run_bm.sh b/.buildkite/scripts/tpu/run_bm.sh
deleted file mode 100755
index 877669cd9..000000000
--- a/.buildkite/scripts/tpu/run_bm.sh
+++ /dev/null
@@ -1,94 +0,0 @@
-#!/bin/bash
-
-set -euo pipefail
-
-VLLM_LOG="$WORKSPACE/vllm_log.txt"
-BM_LOG="$WORKSPACE/bm_log.txt"
-
-if [ -n "$TARGET_COMMIT" ]; then
-  head_hash=$(git rev-parse HEAD)
-  if [ "$TARGET_COMMIT" != "$head_hash" ]; then
-    echo "Error: target commit $TARGET_COMMIT does not match HEAD: $head_hash"
-    exit 1
-  fi
-fi
-
-echo "model: $MODEL"
-echo
-
-#
-# create a log folder
-#
-mkdir "$WORKSPACE/log"
-
-# TODO: Move to image building.
-pip install pandas
-pip install datasets
-
-#
-# create sonnet_4x
-#
-echo "Create sonnet_4x.txt"
-echo "" > benchmarks/sonnet_4x.txt
-for _ in {1..4}
- do
-  cat benchmarks/sonnet.txt >> benchmarks/sonnet_4x.txt
-done
-
-#
-# start vllm service in backend
-#
-echo "lanching vllm..."
-echo "logging to $VLLM_LOG"
-echo
-
-VLLM_USE_V1=1 vllm serve $MODEL \
- --seed 42 \
- --disable-log-requests \
- --max-num-seqs $MAX_NUM_SEQS \
- --max-num-batched-tokens $MAX_NUM_BATCHED_TOKENS \
- --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
- --no-enable-prefix-caching \
- --download_dir $DOWNLOAD_DIR \
- --max-model-len $MAX_MODEL_LEN > "$VLLM_LOG" 2>&1 &
-
-
-echo "wait for 20 minutes.."
-echo
-# sleep 1200
-# wait for 10 minutes...
-for i in {1..120}; do
-    # TODO: detect other type of errors.
-    if grep -Fq "raise RuntimeError" "$VLLM_LOG"; then
-        echo "Detected RuntimeError, exiting."
-        exit 1
-    elif grep -Fq "Application startup complete" "$VLLM_LOG"; then
-        echo "Application started"
-        break
-    else
-        echo "wait for 10 seconds..."
-        sleep 10
-    fi
-done
-
-#
-# run test
-#
-echo "run benchmark test..."
-echo "logging to $BM_LOG"
-echo
-python benchmarks/benchmark_serving.py \
-    --backend vllm \
-    --model $MODEL  \
-    --dataset-name sonnet \
-    --dataset-path benchmarks/sonnet_4x.txt \
-    --sonnet-input-len $INPUT_LEN \
-    --sonnet-output-len $OUTPUT_LEN \
-    --ignore-eos > "$BM_LOG"
-
-echo "completed..."
-echo
-
-throughput=$(grep "Request throughput (req/s):" "$BM_LOG" | sed 's/[^0-9.]//g')
-echo "throughput: $throughput"
-echo
diff --git a/.buildkite/scripts/upload-wheels.sh b/.buildkite/scripts/upload-wheels.sh
index 037897e53..a681f8927 100644
--- a/.buildkite/scripts/upload-wheels.sh
+++ b/.buildkite/scripts/upload-wheels.sh
@@ -50,11 +50,11 @@ aws s3 cp "$normal_wheel" "s3://vllm-wheels/$BUILDKITE_COMMIT/"
 if [[ $normal_wheel == *"cu118"* ]]; then
     # if $normal_wheel matches cu118, do not upload the index.html
     echo "Skipping index files for cu118 wheels"
-elif [[ $normal_wheel == *"cu126"* ]]; then
-    # if $normal_wheel matches cu126, do not upload the index.html
-    echo "Skipping index files for cu126 wheels"
+elif [[ $normal_wheel == *"cu121"* ]]; then
+    # if $normal_wheel matches cu121, do not upload the index.html
+    echo "Skipping index files for cu121 wheels"
 else
-    # only upload index.html for cu128 wheels (default wheels)
+    # only upload index.html for cu124 wheels (default wheels)
     aws s3 cp index.html "s3://vllm-wheels/$BUILDKITE_COMMIT/vllm/index.html"
     aws s3 cp "s3://vllm-wheels/nightly/index.html" "s3://vllm-wheels/$BUILDKITE_COMMIT/index.html"
 fi
@@ -66,13 +66,12 @@ aws s3 cp "$normal_wheel" "s3://vllm-wheels/nightly/"
 if [[ $normal_wheel == *"cu118"* ]]; then
     # if $normal_wheel matches cu118, do not upload the index.html
     echo "Skipping index files for cu118 wheels"
-elif [[ $normal_wheel == *"cu126"* ]]; then
-    # if $normal_wheel matches cu126, do not upload the index.html
-    echo "Skipping index files for cu126 wheels"
+elif [[ $normal_wheel == *"cu121"* ]]; then
+    # if $normal_wheel matches cu121, do not upload the index.html
+    echo "Skipping index files for cu121 wheels"
 else
-    # only upload index.html for cu128 wheels (default wheels)
+    # only upload index.html for cu124 wheels (default wheels)
     aws s3 cp index.html "s3://vllm-wheels/nightly/vllm/index.html"
 fi
 
-aws s3 cp "$wheel" "s3://vllm-wheels/$version/"
-aws s3 cp index.html "s3://vllm-wheels/$version/vllm/index.html"
+aws s3 cp "$wheel" "s3://vllm-wheels/$version/"
\ No newline at end of file
diff --git a/.buildkite/test-pipeline.yaml b/.buildkite/test-pipeline.yaml
index b739851cb..20d858cb1 100644
--- a/.buildkite/test-pipeline.yaml
+++ b/.buildkite/test-pipeline.yaml
@@ -32,17 +32,16 @@ steps:
 ##### fast check tests  #####
 
 - label: Documentation Build # 2min
-  mirror_hardwares: [amdexperimental]
-  working_dir: "/vllm-workspace/test_docs"
+  working_dir: "/vllm-workspace/test_docs/docs"
   fast_check: true
   no_gpu: True
   commands:
-  - pip install -r ../requirements/docs.txt
-  # TODO: add `--strict` once warnings in docstrings are fixed
-  - mkdocs build
+  - pip install -r ../../requirements/docs.txt
+  - SPHINXOPTS=\"-W\" make html
+  # Check API reference (if it fails, you may have missing mock imports)
+  - grep \"sig sig-object py\" build/html/api/inference_params.html
 
 - label: Async Engine, Inputs, Utils, Worker Test # 24min
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - vllm/
   - tests/mq_llm_engine
@@ -58,13 +57,11 @@ steps:
   - pytest -v -s async_engine # AsyncLLMEngine
   - NUM_SCHEDULER_STEPS=4 pytest -v -s async_engine/test_async_llm_engine.py
   - pytest -v -s test_inputs.py
-  - pytest -v -s test_outputs.py
   - pytest -v -s multimodal
   - pytest -v -s test_utils.py # Utils
   - pytest -v -s worker # Worker
 
 - label: Python-only Installation Test
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - tests/standalone_tests/python_only_compile.sh
   - setup.py
@@ -72,7 +69,7 @@ steps:
   - bash standalone_tests/python_only_compile.sh
 
 - label: Basic Correctness Test # 30min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  #mirror_hardwares: [amd]
   fast_check: true
   torch_nightly: true
   source_file_dependencies:
@@ -89,7 +86,6 @@ steps:
   - VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT=1 pytest -v -s basic_correctness/test_preemption.py
 
 - label: Chunked Prefill Test
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - vllm/
   - tests/basic_correctness/test_chunked_prefill
@@ -98,7 +94,7 @@ steps:
   - VLLM_ATTENTION_BACKEND=FLASH_ATTN pytest -v -s basic_correctness/test_chunked_prefill.py
 
 - label: Core Test # 10min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amd]
   fast_check: true
   source_file_dependencies:
   - vllm/core
@@ -108,10 +104,10 @@ steps:
   - pytest -v -s core
 
 - label: Entrypoints Test # 40min
-  mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/tests"
   fast_check: true
   torch_nightly: true
+  #mirror_hardwares: [amd]
   source_file_dependencies:
   - vllm/
   - tests/entrypoints/llm
@@ -125,12 +121,11 @@ steps:
   - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
   - pytest -v -s entrypoints/llm/test_generate_multiple_loras.py # it needs a clean process
   - VLLM_USE_V1=0 pytest -v -s entrypoints/llm/test_guided_generate.py # it needs a clean process
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_oot_registration.py  --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/test_openai_schema.py
   - pytest -v -s entrypoints/test_chat_utils.py
   - VLLM_USE_V1=0 pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
 
 - label: Distributed Tests (4 GPUs) # 10min
-  mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 4
   source_file_dependencies:
@@ -138,38 +133,32 @@ steps:
   - vllm/core/
   - tests/distributed/test_utils
   - tests/distributed/test_pynccl
-  - tests/distributed/test_events
   - tests/spec_decode/e2e/test_integration_dist_tp4
   - tests/compile/test_basic_correctness
   - examples/offline_inference/rlhf.py
   - examples/offline_inference/rlhf_colocate.py
   - tests/examples/offline_inference/data_parallel.py
   - tests/v1/test_async_llm_dp.py
-  - tests/v1/engine/test_engine_core_client.py
   commands:
   # test with tp=2 and external_dp=2
   - VLLM_USE_V1=0 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
   - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
-  # test with tp=2 and pp=2
-  - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
   # test with internal dp
   - python3 ../examples/offline_inference/data_parallel.py
   - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
-  - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
   - pytest -v -s distributed/test_utils.py
   - pytest -v -s compile/test_basic_correctness.py
   - pytest -v -s distributed/test_pynccl.py
-  - pytest -v -s distributed/test_events.py
   - pytest -v -s spec_decode/e2e/test_integration_dist_tp4.py
   # TODO: create a dedicated test section for multi-GPU example tests
   # when we have multiple distributed example tests
   - pushd ../examples/offline_inference
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
+  - python3 rlhf.py
+  - RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
   - popd
 
 - label: Metrics, Tracing Test # 10min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amd]
   num_gpus: 2
   source_file_dependencies:
   - vllm/
@@ -183,7 +172,7 @@ steps:
 #####  1 GPU test  #####
 
 - label: Regression Test # 5min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  #mirror_hardwares: [amd]
   source_file_dependencies:
   - vllm/
   - tests/test_regression
@@ -193,7 +182,7 @@ steps:
   working_dir: "/vllm-workspace/tests" # optional
 
 - label: Engine Test # 10min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amd]
   source_file_dependencies:
   - vllm/
   - tests/engine
@@ -201,14 +190,13 @@ steps:
   - tests/test_sequence
   - tests/test_config
   - tests/test_logger
-  - tests/test_vllm_port
   commands:
-  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py
   # OOM in the CI unless we run this separately
   - pytest -v -s tokenization
 
 - label: V1 Test
-  mirror_hardwares: [amdexperimental]
+  #mirror_hardwares: [amd]
   source_file_dependencies:
     - vllm/
     - tests/v1
@@ -221,11 +209,10 @@ steps:
     - pytest -v -s v1/worker
     - pytest -v -s v1/structured_output
     - pytest -v -s v1/spec_decode
-    - pytest -v -s v1/kv_connector/unit
     - pytest -v -s v1/test_serial_utils.py
+    - pytest -v -s v1/test_stats.py
     - pytest -v -s v1/test_utils.py
     - pytest -v -s v1/test_oracle.py
-    - pytest -v -s v1/test_metrics_reader.py
     # TODO: accuracy does not match, whether setting
     # VLLM_USE_FLASHINFER_SAMPLER or not on H100.
     - pytest -v -s v1/e2e
@@ -234,8 +221,8 @@ steps:
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
 
 - label: Examples Test # 25min
-  mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/examples"
+  #mirror_hardwares: [amd]
   source_file_dependencies:
   - vllm/entrypoints
   - examples/
@@ -250,7 +237,7 @@ steps:
     - python3 offline_inference/vision_language.py --seed 0
     - python3 offline_inference/vision_language_embedding.py --seed 0
     - python3 offline_inference/vision_language_multi_image.py --seed 0
-    - VLLM_USE_V1=0 python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - VLLM_USE_V1=0 python3 other/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 other/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
     - python3 offline_inference/encoder_decoder.py
     - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     - python3 offline_inference/basic/classify.py
@@ -259,7 +246,7 @@ steps:
     - VLLM_USE_V1=0 python3 offline_inference/profiling.py --model facebook/opt-125m run_num_steps --num-steps 2
 
 - label: Prefix Caching Test # 9min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amd]
   source_file_dependencies:
   - vllm/
   - tests/prefix_caching
@@ -267,7 +254,6 @@ steps:
     - pytest -v -s prefix_caching
 
 - label: Samplers Test # 36min
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - vllm/model_executor/layers
   - vllm/sampling_metadata.py
@@ -277,8 +263,18 @@ steps:
     - pytest -v -s samplers
     - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
 
+- label: LogitsProcessor Test # 5min
+  mirror_hardwares: [amd]
+  source_file_dependencies:
+  - vllm/model_executor/layers
+  - vllm/model_executor/guided_decoding
+  - tests/test_logits_processor
+  - tests/model_executor/test_guided_processors
+  commands:
+    - pytest -v -s test_logits_processor.py
+    - pytest -v -s model_executor/test_guided_processors.py
+
 - label: Speculative decoding tests # 40min
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - vllm/spec_decode
   - tests/spec_decode
@@ -289,7 +285,7 @@ steps:
     - pytest -v -s spec_decode/e2e/test_eagle_correctness.py
 
 - label: LoRA Test %N # 15min each
-  mirror_hardwares: [amdexperimental, amdproduction]
+  #mirror_hardwares: [amd]
   source_file_dependencies:
   - vllm/lora
   - tests/lora
@@ -297,21 +293,15 @@ steps:
   parallelism: 4
 
 - label: PyTorch Compilation Unit Tests
-  mirror_hardwares: [amdexperimental, amdproduction]
-  torch_nightly: true
   source_file_dependencies:
     - vllm/
     - tests/compile
   commands:
     - pytest -v -s compile/test_pass_manager.py
     - pytest -v -s compile/test_fusion.py
-    - pytest -v -s compile/test_silu_mul_quant_fusion.py
     - pytest -v -s compile/test_sequence_parallelism.py
-    - pytest -v -s compile/test_async_tp.py
 
 - label: PyTorch Fullgraph Smoke Test # 9min
-  mirror_hardwares: [amdexperimental, amdproduction]
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/compile
@@ -320,11 +310,8 @@ steps:
   # these tests need to be separated, cannot combine
   - pytest -v -s compile/piecewise/test_simple.py
   - pytest -v -s compile/piecewise/test_toy_llama.py
-  - pytest -v -s compile/piecewise/test_full_cudagraph.py
 
 - label: PyTorch Fullgraph Test # 18min
-  mirror_hardwares: [amdexperimental, amdproduction]
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/compile
@@ -332,7 +319,6 @@ steps:
   - pytest -v -s compile/test_full_graph.py
 
 - label: Kernels Core Operation Test
-  mirror_hardwares: [amdexperimental, amdproduction]
   source_file_dependencies:
   - csrc/
   - tests/kernels/core
@@ -340,7 +326,6 @@ steps:
     - pytest -v -s kernels/core
 
 - label: Kernels Attention Test %N
-  mirror_hardwares: [amdexperimental, amdproduction]
   source_file_dependencies:
   - csrc/attention/
   - vllm/attention
@@ -351,7 +336,6 @@ steps:
   parallelism: 2
 
 - label: Kernels Quantization Test %N
-  mirror_hardwares: [amdexperimental, amdproduction]
   source_file_dependencies:
   - csrc/quantization/
   - vllm/model_executor/layers/quantization
@@ -361,7 +345,6 @@ steps:
   parallelism: 2
 
 - label: Kernels MoE Test
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - csrc/moe/
   - tests/kernels/moe
@@ -370,7 +353,6 @@ steps:
     - pytest -v -s kernels/moe
 
 - label: Kernels Mamba Test
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - csrc/mamba/
   - tests/kernels/mamba
@@ -378,69 +360,48 @@ steps:
     - pytest -v -s kernels/mamba
 
 - label: Tensorizer Test # 11min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  # mirror_hardwares: [amd]
   soft_fail: true
   source_file_dependencies:
   - vllm/model_executor/model_loader
   - tests/tensorizer_loader
-  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
   commands:
     - apt-get update && apt-get install -y curl libsodium23
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s tensorizer_loader
-    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
-
-- label: Model Executor Test
-  mirror_hardwares: [amdexperimental, amdproduction]
-  soft_fail: true
-  source_file_dependencies:
-  - vllm/model_executor
-  - tests/model_executor
-  commands:
-    - apt-get update && apt-get install -y curl libsodium23
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -v -s model_executor
 
 - label: Benchmarks # 9min
-  mirror_hardwares: [amdexperimental, amdproduction]
   working_dir: "/vllm-workspace/.buildkite"
+  mirror_hardwares: [amd]
   source_file_dependencies:
   - benchmarks/
   commands:
   - bash scripts/run-benchmarks.sh
 
 - label: Benchmarks CLI Test # 10min
-  mirror_hardwares: [amdexperimental, amdproduction]
   source_file_dependencies:
   - vllm/
   - tests/benchmarks/
   commands:
   - pytest -v -s benchmarks/
 
-- label: Quantization Test
-  mirror_hardwares: [amdexperimental]
+- label: Quantization Test # 33min
   source_file_dependencies:
   - csrc/
   - vllm/model_executor/layers/quantization
   - tests/quantization
-  commands:
-  # temporary install here since we need nightly, will move to requirements/test.in
-  # after torchao 0.12 release
-  - pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126
-  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization
+  command: VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization
 
 - label: LM Eval Small Models # 53min
-  mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
   source_file_dependencies:
   - csrc/
   - vllm/model_executor/layers/quantization
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small.txt --tp-size=1
+  - bash ./run-tests.sh -c configs/models-small.txt -t 1
 
 - label: OpenAI API correctness
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - csrc/
   - vllm/entrypoints/openai/
@@ -449,7 +410,6 @@ steps:
   - pytest -s entrypoints/openai/correctness/
 
 - label: Encoder Decoder tests # 5min
-  mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - vllm/
   - tests/encoder_decoder
@@ -457,8 +417,8 @@ steps:
     - pytest -v -s encoder_decoder
 
 - label: OpenAI-Compatible Tool Use # 20 min
-  mirror_hardwares: [amdexperimental]
   fast_check: false
+  #mirror_hardwares: [ amd ]
   source_file_dependencies:
     - vllm/
     - tests/tool_use
@@ -470,104 +430,92 @@ steps:
 #####  models test  #####
 
 - label: Basic Models Test # 24min
-  mirror_hardwares: [amdexperimental, amdproduction]
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models
   commands:
     - pytest -v -s models/test_transformers.py
     - pytest -v -s models/test_registry.py
-    - pytest -v -s models/test_utils.py
-    - pytest -v -s models/test_vision.py
-    - pytest -v -s models/test_initialization.py
+    # V1 Test: https://github.com/vllm-project/vllm/issues/14531
+    - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2'
+    - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'llama4'
+    - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'plamo2'
 
-- label: Language Models Test (Standard)
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
+- label: Language Models Test (Standard) # 32min
+  #mirror_hardwares: [amd]
   source_file_dependencies:
   - vllm/
-  - tests/models/language
+  - tests/models/decoder_only/language
+  - tests/models/embedding/language
+  - tests/models/encoder_decoder/language
   commands:
     # Install causal-conv1d for plamo2 models here, as it is not compatible with pip-compile.
-    - pip install 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.0.post8'
-    - pip freeze | grep -E 'torch'
-    - pytest -v -s models/language -m core_model
+    - pip install causal-conv1d
+    - pytest -v -s models/decoder_only/language -m 'core_model or quant_model'
+    - pytest -v -s models/embedding/language -m core_model
 
-- label: Language Models Test (Extended Generation) # 1hr20min
-  mirror_hardwares: [amdexperimental]
+- label: Language Models Test (Extended) # 1h10min
   optional: true
   source_file_dependencies:
   - vllm/
-  - tests/models/language/generation
+  - tests/models/decoder_only/language
+  - tests/models/embedding/language
+  - tests/models/encoder_decoder/language
   commands:
     # Install causal-conv1d for plamo2 models here, as it is not compatible with pip-compile.
-    - pip install 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.0.post8'
-    - pytest -v -s models/language/generation -m 'not core_model'
-
-- label: Language Models Test (Extended Pooling)  # 36min
-  mirror_hardwares: [amdexperimental]
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/pooling
-  commands:
-    - pytest -v -s models/language/pooling -m 'not core_model'
+    - pip install causal-conv1d
+    - pytest -v -s models/decoder_only/language -m 'not core_model and not quant_model'
+    - pytest -v -s models/embedding/language -m 'not core_model'
 
-- label: Multi-Modal Models Test (Standard)
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
+- label: Multi-Modal Models Test (Standard) # 40min
+  #mirror_hardwares: [amd]
   source_file_dependencies:
   - vllm/
-  - tests/models/multimodal
+  - tests/models/decoder_only/audio_language
+  - tests/models/decoder_only/vision_language
+  - tests/models/embedding/vision_language
+  - tests/models/encoder_decoder/audio_language
+  - tests/models/encoder_decoder/vision_language
   commands:
     - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pip freeze | grep -E 'torch'
-    - pytest -v -s models/multimodal/processing
-    - pytest -v -s --ignore models/multimodal/generation/test_whisper.py models/multimodal -m core_model
-    - cd .. && pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
-
-- label: Multi-Modal Models Test (Extended) 1
-  mirror_hardwares: [amdexperimental]
+    - pytest -v -s models/multimodal
+    - pytest -v -s models/decoder_only/audio_language -m 'core_model or quant_model'
+    - pytest -v -s models/decoder_only/vision_language -m 'core_model or quant_model'
+    - pytest -v -s models/embedding/vision_language -m core_model
+    - pytest -v -s models/encoder_decoder/audio_language -m core_model
+    - pytest -v -s models/encoder_decoder/language -m core_model
+    - pytest -v -s models/encoder_decoder/vision_language -m core_model
+    - pytest -v -s models/decoder_only/vision_language/test_interleaved.py
+
+- label: Multi-Modal Models Test (Extended) 1 # 48m
   optional: true
   source_file_dependencies:
   - vllm/
-  - tests/models/multimodal
+  - tests/models/decoder_only/audio_language
+  - tests/models/decoder_only/vision_language
+  - tests/models/embedding/vision_language
+  - tests/models/encoder_decoder/vision_language
   commands:
     - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/processing models/multimodal -m 'not core_model'
-
-- label: Multi-Modal Models Test (Extended) 2
-  mirror_hardwares: [amdexperimental]
+    - pytest -v -s models/decoder_only/audio_language -m 'not core_model and not quant_model'
+    - pytest -v -s models/decoder_only/vision_language/test_models.py -m 'split(group=0) and not core_model and not quant_model'
+    - pytest -v -s --ignore models/decoder_only/vision_language/test_models.py models/decoder_only/vision_language -m 'not core_model and not quant_model'
+    - pytest -v -s models/embedding/vision_language -m 'not core_model'
+    - pytest -v -s models/encoder_decoder/language -m 'not core_model'
+    - pytest -v -s models/encoder_decoder/vision_language -m 'not core_model'
+
+- label: Multi-Modal Models Test (Extended) 2 # 38m
   optional: true
   source_file_dependencies:
   - vllm/
-  - tests/models/multimodal
+  - tests/models/decoder_only/vision_language
   commands:
     - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
-
-- label: Multi-Modal Models Test (Extended) 3
-  mirror_hardwares: [amdexperimental, amdproduction]
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
-
-- label: Quantized Models Test
-  mirror_hardwares: [amdexperimental, amdproduction]
-  source_file_dependencies:
-  - vllm/model_executor/layers/quantization
-  - tests/models/quantization
-  commands:
-    - pytest -v -s models/quantization
+    - pytest -v -s models/decoder_only/vision_language/test_models.py -m 'split(group=1) and not core_model and not quant_model'
 
 # This test is used only in PR development phase to test individual models and should never run on main
 - label: Custom Models Test
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amd]
   optional: true
   commands:
     - echo 'Testing custom models...'
@@ -579,7 +527,7 @@ steps:
 #####  multi gpus test  #####
 
 - label: Distributed Comm Ops Test # 7min
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amd]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
   source_file_dependencies:
@@ -590,7 +538,6 @@ steps:
   - pytest -v -s distributed/test_shm_broadcast.py
 
 - label: 2 Node Tests (4 GPUs in total) # 16min
-  mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
   num_nodes: 2
@@ -609,7 +556,7 @@ steps:
     - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
 
 - label: Distributed Tests (2 GPUs) # 40min
-  mirror_hardwares: [amdexperimental]
+  #mirror_hardwares: [amd]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
   source_file_dependencies:
@@ -624,11 +571,9 @@ steps:
   - vllm/worker/model_runner.py
   - entrypoints/llm/test_collective_rpc.py
   - tests/v1/test_async_llm_dp.py
-  - tests/v1/entrypoints/openai/test_multi_api_servers.py
   - vllm/v1/engine/
   commands:
   - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
-  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
   - pytest -v -s entrypoints/llm/test_collective_rpc.py
   - pytest -v -s ./compile/test_basic_correctness.py
   - pytest -v -s ./compile/test_wrapper.py
@@ -636,8 +581,9 @@ steps:
   - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
   # Avoid importing model tests that cause CUDA reinitialization error
   - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/language -v -s -m 'distributed(num_gpus=2)'
-  - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)'
+  - pytest models/encoder_decoder/language/test_bart.py -v -s -m 'distributed(num_gpus=2)'
+  - pytest models/encoder_decoder/vision_language/test_broadcast.py -v -s -m 'distributed(num_gpus=2)'
+  - pytest models/decoder_only/vision_language/test_models.py -v -s -m 'distributed(num_gpus=2)'
   # test sequence parallel
   - pytest -v -s distributed/test_sequence_parallel.py
   # this test fails consistently.
@@ -648,14 +594,13 @@ steps:
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
 
 - label: Plugin Tests (2 GPUs) # 40min
-  mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
   source_file_dependencies:
   - vllm/plugins/
   - tests/plugins/
   commands:
-  # begin platform plugin and general plugin tests, all the code in-between runs on dummy platform
+  # begin platform plugin tests, all the code in-between runs on dummy platform
   - pip install -e ./plugins/vllm_add_dummy_platform
   - pytest -v -s plugins_tests/test_platform_plugins.py
   - pip uninstall vllm_add_dummy_platform -y
@@ -666,10 +611,8 @@ steps:
   - pytest -v -s distributed/test_distributed_oot.py
   - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
   - pytest -v -s models/test_oot_registration.py # it needs a clean process
-  - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
 
 - label: Multi-step Tests (4 GPUs) # 36min
-  mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 4
   source_file_dependencies:
@@ -690,7 +633,6 @@ steps:
   - pytest -v -s multi_step/test_correctness_llm.py
 
 - label: Pipeline Parallelism Test # 45min
-  mirror_hardwares: [amdexperimental, amdproduction]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 4
   source_file_dependencies:
@@ -704,7 +646,6 @@ steps:
   - pytest -v -s distributed/test_pipeline_parallel.py
 
 - label: LoRA TP Test (Distributed)
-  mirror_hardwares: [amdexperimental, amdproduction]
   num_gpus: 4
   source_file_dependencies:
   - vllm/lora
@@ -720,7 +661,6 @@ steps:
 
 
 - label: Weight Loading Multiple GPU Test  # 33min
-  mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
   source_file_dependencies:
@@ -730,7 +670,6 @@ steps:
     - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
 
 - label: Weight Loading Multiple GPU Test - Large Models # optional
-  mirror_hardwares: [amdexperimental] 
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
   gpu: a100
@@ -769,4 +708,4 @@ steps:
   - vllm/model_executor/layers/quantization
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
+  - bash ./run-tests.sh -c configs/models-large.txt -t 4
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index e98ccd035..76aa5f7a3 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -10,17 +10,14 @@
 /vllm/worker/worker.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill
 /vllm/model_executor/layers/sampler.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth
-/vllm/model_executor/guided_decoding @mgoin @russellb @aarnphm
+/vllm/model_executor/guided_decoding @mgoin @russellb
 /vllm/multimodal @DarkLight1337 @ywang96
 /vllm/vllm_flash_attn @LucasWilkinson
-/vllm/lora @jeejeelee
-/vllm/reasoning @aarnphm
-/vllm/entrypoints @aarnphm
 CMakeLists.txt @tlrmchlsmth
 
 # vLLM V1
 /vllm/v1 @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @comaniac @alexm-redhat
-/vllm/v1/structured_output @mgoin @russellb @aarnphm
+/vllm/v1/structured_output @mgoin @russellb
 
 # Test ownership
 /.buildkite/lm-eval-harness @mgoin @simon-mo
@@ -29,8 +26,8 @@ CMakeLists.txt @tlrmchlsmth
 /tests/distributed/test_multi_node_assignment.py @youkaichao
 /tests/distributed/test_pipeline_parallel.py @youkaichao
 /tests/distributed/test_same_node.py @youkaichao
-/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @simon-mo @aarnphm
-/tests/entrypoints/llm/test_guided_generate.py @mgoin @russellb @aarnphm
+/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @simon-mo
+/tests/entrypoints/llm/test_guided_generate.py @mgoin @russellb
 /tests/kernels @tlrmchlsmth @WoosukKwon
 /tests/model_executor/test_guided_processors.py @mgoin @russellb
 /tests/models @DarkLight1337 @ywang96
@@ -40,11 +37,6 @@ CMakeLists.txt @tlrmchlsmth
 /tests/quantization @mgoin @robertgshaw2-redhat
 /tests/spec_decode @njhill @LiuXiaoxuanPKU
 /tests/test_inputs.py @DarkLight1337 @ywang96
-/tests/v1/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
-/tests/v1/structured_output @mgoin @russellb @aarnphm
+/tests/v1/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb
+/tests/v1/structured_output @mgoin @russellb
 /tests/weight_loading @mgoin @youkaichao
-/tests/lora @jeejeelee
-
-# Docs
-/docs @hmellor
-mkdocs.yaml @hmellor
diff --git a/.github/ISSUE_TEMPLATE/400-bug-report.yml b/.github/ISSUE_TEMPLATE/400-bug-report.yml
index 8c5c28cd7..b96ab4074 100644
--- a/.github/ISSUE_TEMPLATE/400-bug-report.yml
+++ b/.github/ISSUE_TEMPLATE/400-bug-report.yml
@@ -8,16 +8,6 @@ body:
   attributes:
     value: >
       #### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm/issues?q=is%3Aissue+sort%3Acreated-desc+).
-- type: markdown
-  attributes:
-    value: |
-      ⚠️ **SECURITY WARNING:** Please review any text you paste to ensure it does not contain sensitive information such as:
-      - API tokens or keys (e.g., Hugging Face tokens, OpenAI API keys)
-      - Passwords or authentication credentials
-      - Private URLs or endpoints
-      - Personal or confidential data
-      
-      Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
 - type: textarea
   attributes:
     label: Your current environment
@@ -31,12 +21,12 @@ body:
       It is suggested to download and execute the latest script, as vllm might frequently update the diagnosis information needed for accurately and quickly responding to issues.
     value: |
       <details>
-      <summary>The output of <code>python collect_env.py</code></summary>
+      <summary>The output of `python collect_env.py`</summary>
 
       ```text
       Your output of `python collect_env.py` here
       ```
-
+      
       </details>
   validations:
     required: true
@@ -85,20 +75,20 @@ body:
       ```
 
       ```
-      The error message you got, with the full traceback and the error logs with [dump_input.py:##] if present.
+      The error message you got, with the full traceback.
       ```
   validations:
     required: true
 - type: markdown
   attributes:
-    value: |
-      ⚠️ Please separate bugs of `transformers` implementation or usage from bugs of `vllm`. If you think anything is wrong with the model's output:
+    value: >
+      ⚠️ Please separate bugs of `transformers` implementation or usage from bugs of `vllm`. If you think anything is wrong with the models' output:
 
       - Try the counterpart of `transformers` first. If the error appears, please go to [their issues](https://github.com/huggingface/transformers/issues?q=is%3Aissue+is%3Aopen+sort%3Aupdated-desc).
 
       - If the error only appears in vllm, please provide the detailed script of how you run `transformers` and `vllm`, also highlight the difference and what you expect.
 
-      Thanks for reporting 🙏!
+      Thanks for contributing 🎉!
 - type: checkboxes
   id: askllm
   attributes:
diff --git a/.github/ISSUE_TEMPLATE/450-ci-failure.yml b/.github/ISSUE_TEMPLATE/450-ci-failure.yml
deleted file mode 100644
index 7af0e0673..000000000
--- a/.github/ISSUE_TEMPLATE/450-ci-failure.yml
+++ /dev/null
@@ -1,69 +0,0 @@
-name: 🧪 CI failure report
-description: Report a failing test.
-title: "[CI Failure]: "
-labels: ["ci-failure"]
-
-body:
-- type: markdown
-  attributes:
-    value: >
-      #### Include the name of the failing Buildkite step and test file in the title.
-- type: input
-  attributes:
-    label: Name of failing test
-    description: |
-      Paste in the fully-qualified name of the failing test from the logs.
-    placeholder: |
-      `path/to/test_file.py::test_name[params]`
-  validations:
-    required: true
-- type: checkboxes
-  attributes:
-    label: Basic information
-    description: Select all items that apply to the failing test.
-    options:
-      - label: Flaky test
-      - label: Can reproduce locally
-      - label: Caused by external libraries (e.g. bug in `transformers`)
-- type: textarea
-  attributes:
-    label: 🧪 Describe the failing test
-    description: |
-      Please provide a clear and concise description of the failing test.
-    placeholder: |
-      A clear and concise description of the failing test.
-  
-      ```
-      The error message you got, with the full traceback and the error logs with [dump_input.py:##] if present.
-      ```
-  validations:
-    required: true
-- type: textarea
-  attributes:
-    label: 📝 History of failing test
-    description: |
-      Since when did the test start to fail?
-      You can look up its history via [Buildkite Test Suites](https://buildkite.com/organizations/vllm/analytics/suites/ci-1/tests?branch=main).
-
-      If you have time, identify the PR that caused the test to fail on main. You can do so via the following methods:
-
-      - Use Buildkite Test Suites to find the PR where the test failure first occurred, and reproduce the failure locally.
-
-      - Run [`git bisect`](https://git-scm.com/docs/git-bisect) locally.
-
-      - Manually unblock Buildkite steps for suspected PRs on main and check the results. (authorized users only)
-    placeholder: |
-      Approximate timeline and/or problematic PRs
-
-      A link to the Buildkite analytics of the failing test (if available)
-  validations:
-    required: true
-- type: textarea
-  attributes:
-    label: CC List.
-    description: >
-      The list of people you want to CC. Usually, this includes those who worked on the PR that failed the test.
-- type: markdown
-  attributes:
-    value: >
-      Thanks for reporting 🙏!
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index 017ec7ca8..7042e81a8 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -1,18 +1,6 @@
-## Essential Elements of an Effective PR Description Checklist
-- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
-- [ ] The test plan, such as providing test command.
-- [ ] The test results, such as pasting the results comparison before and after, or e2e results
-- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
+FILL IN THE PR DESCRIPTION HERE
 
-PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS ABOVE HAVE BEEN CONSIDERED.
-
-## Purpose
-
-## Test Plan
-
-## Test Result
-
-## (Optional) Documentation Update
+FIX #xxxx (*link existing issues this PR will resolve*)
 
 <!--- pyml disable-next-line no-emphasis-as-heading -->
-**BEFORE SUBMITTING, PLEASE READ <https://docs.vllm.ai/en/latest/contributing>** (anything written below this line will be removed by GitHub Actions)
+**BEFORE SUBMITTING, PLEASE READ <https://docs.vllm.ai/en/latest/contributing/overview.html>** (anything written below this line will be removed by GitHub Actions)
diff --git a/.github/mergify.yml b/.github/mergify.yml
index 5692bb5d3..15fa3660a 100644
--- a/.github/mergify.yml
+++ b/.github/mergify.yml
@@ -36,20 +36,6 @@ pull_request_rules:
       add:
         - frontend
 
-- name: label-llama
-  description: Automatically apply llama label
-  conditions:
-    - or:
-      - files~=^examples/.*llama.*\.py
-      - files~=^tests/.*llama.*\.py
-      - files~=^vllm/entrypoints/openai/tool_parsers/llama.*\.py
-      - files~=^vllm/model_executor/models/.*llama.*\.py
-      - files~=^vllm/transformers_utils/configs/.*llama.*\.py
-  actions:
-    label:
-      add:
-        - llama
-
 - name: label-multi-modality
   description: Automatically apply multi-modality label
   conditions:
@@ -72,7 +58,7 @@ pull_request_rules:
       - files~=^benchmarks/structured_schemas/
       - files=benchmarks/benchmark_serving_structured_output.py
       - files=benchmarks/run_structured_output_benchmark.sh
-      - files=docs/features/structured_outputs.md
+      - files=docs/source/features/structured_outputs.md
       - files=examples/offline_inference/structured_outputs.py
       - files=examples/online_serving/openai_chat_completion_structured_outputs.py
       - files=examples/online_serving/openai_chat_completion_structured_outputs_with_reasoning.py
@@ -149,7 +135,9 @@ pull_request_rules:
       - files~=^tests/entrypoints/openai/tool_parsers/
       - files=tests/entrypoints/openai/test_chat_with_tool_reasoning.py
       - files~=^vllm/entrypoints/openai/tool_parsers/
-      - files=docs/features/tool_calling.md
+      - files=docs/source/features/tool_calling.md
+      - files=docs/source/getting_started/examples/openai_chat_completion_client_with_tools.md
+      - files=docs/source/getting_started/examples/chat_with_tools.md
       - files~=^examples/tool_chat_*
       - files=examples/offline_inference/chat_with_tools.py
       - files=examples/online_serving/openai_chat_completion_client_with_tools_required.py
@@ -175,17 +163,6 @@ pull_request_rules:
 
        https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
 
-- name: assign reviewer for tensorizer changes
-  conditions:
-      - files~=^vllm/model_executor/model_loader/tensorizer.py
-      - files~=^vllm/model_executor/model_loader/tensorizer_loader.py
-      - files~=^tests/entrypoints/openai/test_tensorizer_entrypoint.py
-      - files~=^tests/tensorizer_loader/
-  actions:
-    assign:
-      users:
-        - "sangstar"
-
 - name: remove 'needs-rebase' label when conflict is resolved
   conditions:
       - -conflict
diff --git a/.github/scripts/cleanup_pr_body.sh b/.github/scripts/cleanup_pr_body.sh
index 8d65936fb..3246c6f9b 100755
--- a/.github/scripts/cleanup_pr_body.sh
+++ b/.github/scripts/cleanup_pr_body.sh
@@ -26,7 +26,7 @@ sed -i '/\*\*BEFORE SUBMITTING, PLEASE READ.*\*\*/,$d' "${NEW}"
 
 # Remove HTML <details> section that includes <summary> text of "PR Checklist (Click to Expand)"
 python3 - <<EOF
-import regex as re
+import re
 
 with open("${NEW}", "r") as file:
     content = file.read()
diff --git a/.github/workflows/add_label_automerge.yml b/.github/workflows/add_label_automerge.yml
index 315042fbf..c9d6d4259 100644
--- a/.github/workflows/add_label_automerge.yml
+++ b/.github/workflows/add_label_automerge.yml
@@ -1,6 +1,4 @@
 name: Add label on auto-merge enabled
-permissions:
-    pull-requests: write
 on:
     pull_request_target:
         types:
diff --git a/.github/workflows/cleanup_pr_body.yml b/.github/workflows/cleanup_pr_body.yml
index d5c6b8d43..50fea0c43 100644
--- a/.github/workflows/cleanup_pr_body.yml
+++ b/.github/workflows/cleanup_pr_body.yml
@@ -20,12 +20,7 @@ jobs:
         with:
           python-version: '3.12'
 
-      - name: Install Python dependencies
-        run: |
-          python3 -m pip install --upgrade pip
-          python3 -m pip install regex
-
       - name: Update PR description
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: bash .github/scripts/cleanup_pr_body.sh "${{ github.event.number }}"
+        run: .github/scripts/cleanup_pr_body.sh "${{ github.event.number }}"
diff --git a/.github/workflows/lint-and-deploy.yaml b/.github/workflows/lint-and-deploy.yaml
index 64011922a..7b1d9f699 100644
--- a/.github/workflows/lint-and-deploy.yaml
+++ b/.github/workflows/lint-and-deploy.yaml
@@ -2,9 +2,6 @@ name: Lint and Deploy Charts
 
 on: pull_request
 
-permissions:
-  contents: read
-
 jobs:
   lint-and-deploy:
     runs-on: ubuntu-latest
@@ -69,7 +66,7 @@ jobs:
           export AWS_SECRET_ACCESS_KEY=minioadmin
           sleep 30 && kubectl -n ns-vllm logs -f "$(kubectl -n ns-vllm get pods | awk '/deployment/ {print $1;exit}')" &
           helm install --wait --wait-for-jobs --timeout 5m0s --debug --create-namespace --namespace=ns-vllm test-vllm examples/online_serving/chart-helm -f examples/online_serving/chart-helm/values.yaml --set secrets.s3endpoint=http://minio:9000 --set secrets.s3bucketname=testbucket --set secrets.s3accesskeyid=$AWS_ACCESS_KEY_ID --set secrets.s3accesskey=$AWS_SECRET_ACCESS_KEY --set resources.requests.cpu=1 --set resources.requests.memory=4Gi --set resources.limits.cpu=2 --set resources.limits.memory=5Gi --set image.env[0].name=VLLM_CPU_KVCACHE_SPACE --set image.env[1].name=VLLM_LOGGING_LEVEL --set-string image.env[0].value="1" --set-string image.env[1].value="DEBUG" --set-string extraInit.s3modelpath="opt-125m/" --set-string 'resources.limits.nvidia\.com/gpu=0' --set-string 'resources.requests.nvidia\.com/gpu=0' --set-string image.repository="vllm-cpu-env"
-
+    
       - name: curl test
         run: |
           kubectl -n ns-vllm port-forward service/test-vllm-service 8001:80 &
@@ -82,4 +79,4 @@ jobs:
                           "max_tokens": 7,
                           "temperature": 0
                   }'):$CODE"
-          echo "$CODE"
+          echo "$CODE"
\ No newline at end of file
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index 8e694d181..6ab63a402 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -5,9 +5,6 @@ on:
   push:
     branches: [main]
 
-permissions:
-  contents: read
-
 jobs:
   pre-commit:
     runs-on: ubuntu-latest
diff --git a/.github/workflows/reminder_comment.yml b/.github/workflows/reminder_comment.yml
index 16ae1aadb..27318c2fd 100644
--- a/.github/workflows/reminder_comment.yml
+++ b/.github/workflows/reminder_comment.yml
@@ -1,6 +1,4 @@
 name: PR Reminder Comment Bot
-permissions:
-  pull-requests: write
 on:
   pull_request_target:
     types: [opened]
diff --git a/.gitignore b/.gitignore
index e49d1d6ba..728213ceb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -77,6 +77,10 @@ instance/
 # Scrapy stuff:
 .scrapy
 
+# Sphinx documentation
+docs/_build/
+docs/source/getting_started/examples/
+
 # PyBuilder
 .pybuilder/
 target/
@@ -146,7 +150,6 @@ venv.bak/
 
 # mkdocs documentation
 /site
-docs/examples
 
 # mypy
 .mypy_cache/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index a105b0e14..f76b24c02 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -11,47 +11,42 @@ repos:
   hooks:
   - id: yapf
     args: [--in-place, --verbose]
-    # Keep the same list from yapfignore here to avoid yapf failing without any inputs
-    exclude: '(.buildkite|benchmarks|build|examples)/.*'
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.11.7
+  rev: v0.9.3
   hooks:
   - id: ruff
     args: [--output-format, github, --fix]
-  - id: ruff-format
-    files: ^(.buildkite|benchmarks|examples)/.*
 - repo: https://github.com/codespell-project/codespell
-  rev: v2.4.1
+  rev: v2.4.0
   hooks:
   - id: codespell
     additional_dependencies: ['tomli']
     args: ['--toml', 'pyproject.toml']
 - repo: https://github.com/PyCQA/isort
-  rev: 6.0.1
+  rev: 0a0b7a830386ba6a31c2ec8316849ae4d1b8240d # 6.0.0
   hooks:
   - id: isort
 - repo: https://github.com/pre-commit/mirrors-clang-format
-  rev: v20.1.3
+  rev: v19.1.7
   hooks:
   - id: clang-format
     exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
     types_or: [c++, cuda]
     args: [--style=file, --verbose]
 - repo: https://github.com/jackdewinter/pymarkdown
-  rev: v0.9.29
+  rev: v0.9.27
   hooks:
   - id: pymarkdown
-    exclude: '.*\.inc\.md'
     args: [fix]
 - repo: https://github.com/rhysd/actionlint
   rev: v1.7.7
   hooks:
   - id: actionlint
 - repo: https://github.com/astral-sh/uv-pre-commit
-  rev: 0.6.17
+  rev: 0.6.2
   hooks:
     - id: pip-compile
-      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu128]
+      args: [requirements/test.in, -o, requirements/test.txt]
       files: ^requirements/test\.(in|txt)$
 - repo: local
   hooks:
@@ -60,7 +55,7 @@ repos:
     entry: tools/mypy.sh 0 "local"
     language: python
     types: [python]
-    additional_dependencies: &mypy_deps [mypy==1.11.1, types-cachetools, types-setuptools, types-PyYAML, types-requests, pydantic]
+    additional_dependencies: &mypy_deps [mypy==1.11.1, types-cachetools, types-setuptools, types-PyYAML, types-requests]
     stages: [pre-commit] # Don't run in CI
   - id: mypy-3.9 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.9
@@ -106,8 +101,8 @@ repos:
     args:
       - -c
       - |
-        if ! grep -q "^Signed-off-by: $(git config user.name) <$(git config user.email)>" "$(git rev-parse --git-path COMMIT_EDITMSG)"; then
-          printf "\nSigned-off-by: $(git config user.name) <$(git config user.email)>\n" >> "$(git rev-parse --git-path COMMIT_EDITMSG)"
+        if ! grep -q "^Signed-off-by: $(git config user.name) <$(git config user.email)>" .git/COMMIT_EDITMSG; then
+          printf "\nSigned-off-by: $(git config user.name) <$(git config user.email)>\n" >> .git/COMMIT_EDITMSG
         fi
     language: system
     verbose: true
@@ -130,21 +125,8 @@ repos:
     name: Update Dockerfile dependency graph
     entry: tools/update-dockerfile-graph.sh
     language: script
-  - id: enforce-import-regex-instead-of-re
-    name: Enforce import regex as re
-    entry: python tools/enforce_regex_import.py
-    language: python
-    types: [python]
-    pass_filenames: false
-    additional_dependencies: [regex]
-  # forbid directly import triton
-  - id: forbid-direct-triton-import
-    name: "Forbid direct 'import triton'"
-    entry: python tools/check_triton_import.py
-    language: python
-    types: [python]
+    files: ^docker/Dockerfile$
     pass_filenames: false
-    additional_dependencies: [regex]
   # Keep `suggestion` last
   - id: suggestion
     name: Suggestion
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index 98c3be25f..2781ec223 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -8,8 +8,12 @@ build:
   tools:
     python: "3.12"
 
-mkdocs:
-  configuration: mkdocs.yaml
+sphinx:
+  configuration: docs/source/conf.py
+  fail_on_warning: true
+
+# If using Sphinx, optionally build your docs in additional formats such as PDF
+formats: []
 
 # Optionally declare the Python requirements required to build your docs
 python:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index bd389823f..f199d66eb 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -18,6 +18,32 @@ set(VLLM_TARGET_DEVICE "cuda" CACHE STRING "Target device backend for vLLM")
 message(STATUS "Build type: ${CMAKE_BUILD_TYPE}")
 message(STATUS "Target device: ${VLLM_TARGET_DEVICE}")
 
+option(USE_MACA "Enable MACA Support" ON)
+option(ENABLE_BLAS_API "Enable blas api for fused moe" ON)
+
+if (USE_MACA)
+  message(STATUS "###########################")
+  message(STATUS "Detected Using MACA")
+  message(STATUS "###########################")
+
+  add_compile_definitions(
+    MACA_VERSION_MAJOR=${MACA_VERSION_MAJOR}
+    MACA_VERSION_MINOR=${MACA_VERSION_MINOR}
+    MACA_VERSION_PATCH=${MACA_VERSION_PATCH}
+    MACA_VERSION_BUILD=${MACA_VERSION_BUILD}
+  )
+  message(STATUS "MACA version: ${MACA_VERSION_MAJOR}.${MACA_VERSION_MINOR}.${MACA_VERSION_PATCH}.${MACA_VERSION_BUILD}")
+
+  set(MACA_PATH "$ENV{MACA_PATH}")
+  add_compile_definitions(USE_MACA)
+
+  if (MACA_PATH AND EXISTS ${MACA_PATH})
+    message(STATUS "MACA found at ${MACA_PATH}")
+  else()
+    message(FATAL_ERROR "MACA not found or invalid path, please check your MACA_PATH")
+  endif()
+endif()
+
 include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake)
 
 # Suppress potential warnings about unused manually-specified variables
@@ -45,8 +71,13 @@ set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1
 # requirements.txt files and should be kept consistent.  The ROCm torch
 # versions are derived from docker/Dockerfile.rocm
 #
+if (USE_MACA)
+set(TORCH_SUPPORTED_VERSION_CUDA "2.6.0")
+set(TORCH_SUPPORTED_VERSION_ROCM "2.6.0")
+else()
 set(TORCH_SUPPORTED_VERSION_CUDA "2.7.0")
 set(TORCH_SUPPORTED_VERSION_ROCM "2.7.0")
+endif()
 
 #
 # Try to find python package with an executable that exactly matches
@@ -162,7 +193,7 @@ endif()
 # `VLLM_GPU_LANG`.
 # The final set of arches is stored in `VLLM_GPU_FLAGS`.
 #
-get_torch_gpu_compiler_flags(VLLM_GPU_FLAGS ${VLLM_GPU_LANG})
+get_torch_gpu_compiler_flags(VLLM_GPU_FLAGS ${VLLM_GPU_LANG} USE_MACA)
 
 #
 # Set nvcc parallelism.
@@ -255,6 +286,13 @@ set(VLLM_EXT_SRC
   "csrc/custom_all_reduce.cu"
   "csrc/torch_bindings.cpp")
 
+# support opt of gptq-marlin
+set_source_files_properties(
+  "csrc/quantization/gptq/q_gemm.cu"
+  PROPERTIES
+  COMPILE_FLAGS "-mllvm -metaxgpu-igroup=true -mllvm -metaxgpu-igroup-config=B8 -maxmregcount=128"
+)
+
 if(VLLM_GPU_LANG STREQUAL "CUDA")
   SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
@@ -266,6 +304,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
   endif()
 
+  # Substitue CUTLASS with MATLASS 
+  if (USE_MACA)
+    message(WARNING "Use MACA, Overwrite VLLM_CUTLASS_SRC_DIR.")
+    set(VLLM_CUTLASS_SRC_DIR $ENV{MACA_PATH}/include)
+  endif()
+
   if(VLLM_CUTLASS_SRC_DIR)
     if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
       get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
@@ -286,16 +330,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
         GIT_SHALLOW TRUE
     )
   endif()
-  FetchContent_MakeAvailable(cutlass)
+  #FetchContent_MakeAvailable(cutlass)
 
   list(APPEND VLLM_EXT_SRC
-    "csrc/quantization/aqlm/gemm_kernels.cu"
+    # "csrc/quantization/aqlm/gemm_kernels.cu"
     "csrc/quantization/awq/gemm_kernels.cu"
     "csrc/permute_cols.cu"
     "csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu"
     "csrc/quantization/fp4/nvfp4_quant_entry.cu"
     "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
-    "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
+    # "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
     "csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
     "csrc/cutlass_extensions/common.cpp"
     "csrc/attention/mla/cutlass_mla_entry.cu")
@@ -304,12 +348,26 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     SRCS "${VLLM_EXT_SRC}"
     CUDA_ARCHS "${CUDA_ARCHS}")
 
+  # support opt of gptq-marlin
+  set_source_files_properties(
+    "csrc/quantization/awq/gemm_kernels.cu"
+    PROPERTIES
+    COMPILE_FLAGS " -mllvm -metaxgpu-igroup=true -mllvm -metaxgpu-igroup-config=B8 -maxmregcount=128"
+  )
+
+  # support opt of cutlass w8a8 scale mm
+  set_source_files_properties(
+    "csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu"
+    PROPERTIES
+    COMPILE_FLAGS " -mllvm -metaxgpu-igroup=true -mllvm -misched-postra=true"
+  )  
+
   # Only build Marlin kernels if we are building for at least some compatible archs.
   # Keep building Marlin for 9.0 as there are some group sizes and shapes that
   # are not supported by Machete yet.
   # 9.0 for latest bf16 atomicAdd PTX
   cuda_archs_loose_intersection(MARLIN_ARCHS "8.0;8.7;9.0+PTX" "${CUDA_ARCHS}")
-  if (MARLIN_ARCHS)
+  if (MARLIN_ARCHS AND NOT USE_MACA)
 
     #
     # For the Marlin kernels we automatically generate sources for various
@@ -368,13 +426,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_EXT_SRC "${MARLIN_SRCS}")
     message(STATUS "Building Marlin kernels for archs: ${MARLIN_ARCHS}")
   else()
+    message(WARNING "Use MACA, Skip MARLIN_ARCHS.")
     message(STATUS "Not building Marlin kernels as no compatible archs found"
                    " in CUDA target architectures")
   endif()
 
   # Only build AllSpark kernels if we are building for at least some compatible archs.
   cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
-  if (ALLSPARK_ARCHS)
+  if (ALLSPARK_ARCHS AND NOT USE_MACA)
     set(ALLSPARK_SRCS
        "csrc/quantization/gptq_allspark/allspark_repack.cu"
        "csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
@@ -384,6 +443,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_EXT_SRC "${ALLSPARK_SRCS}")
     message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
   else()
+    message(WARNING "Use MACA, Skip ALLSPARK_ARCHS.")
     message(STATUS "Not building AllSpark kernels as no compatible archs found"
                    " in CUDA target architectures")
   endif()
@@ -393,7 +453,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # The cutlass_scaled_mm kernels for Hopper (c3x, i.e. CUTLASS 3.x) require
   # CUDA 12.0 or later
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.0 AND SCALED_MM_ARCHS)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.0 AND SCALED_MM_ARCHS AND NOT USE_MACA)
     set(SRCS
        "csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm90.cu"
        "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8.cu"
@@ -409,6 +469,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
     message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
   else()
+    message(WARNING "Use MACA, Skip SCALED_MM_ARCHS.")
     if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.0 AND SCALED_MM_ARCHS)
       message(STATUS "Not building scaled_mm_c3x_sm90 as CUDA Compiler version is "
                      "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
@@ -464,6 +525,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
     list(APPEND VLLM_EXT_SRC "${SRCS}")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
+    
+    # support opt of cutlass w8a8 scale mm
+    set_source_files_properties(
+      "csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu"
+      PROPERTIES
+      COMPILE_FLAGS " -mllvm -metaxgpu-igroup=true -mllvm -misched-postra=true"
+    )
     message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
   else()
     if (SCALED_MM_3X_ARCHS)
@@ -502,7 +570,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 
   # FP4 Archs and flags
   cuda_archs_loose_intersection(FP4_ARCHS "10.0a" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND FP4_ARCHS)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND FP4_ARCHS AND NOT USE_MACA)
     set(SRCS
       "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
       "csrc/quantization/fp4/nvfp4_experts_quant.cu"
@@ -515,6 +583,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4=1")
     message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
   else()
+    message(WARNING "Use MACA, Skip FP4_ARCHS.")
     message(STATUS "Not building NVFP4 as no compatible archs were found.")
     # clear FP4_ARCHS
     set(FP4_ARCHS)
@@ -522,7 +591,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 
   # CUTLASS MLA Archs and flags
   cuda_archs_loose_intersection(MLA_ARCHS "10.0a" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND MLA_ARCHS)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND MLA_ARCHS AND NOT USE_MACA)
     set(SRCS
       "csrc/attention/mla/cutlass_mla_kernels.cu")
     set_gencode_flags_for_srcs(
@@ -535,6 +604,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
     message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
   else()
+    message(WARNING "Use MACA, Skip MLA_ARCHS.")
     message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
     # clear MLA_ARCHS
     set(MLA_ARCHS)
@@ -666,6 +736,13 @@ set(VLLM_MOE_EXT_SRC
   "csrc/moe/moe_align_sum_kernels.cu"
   "csrc/moe/topk_softmax_kernels.cu")
 
+if (USE_MACA AND ENABLE_BLAS_API)
+  list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_ops.cpp")
+
+  set(MCBLAS_INCLUDE_DIR $ENV{MACA_PATH}/include/mcblas)
+  set(MCBLAS_LIB $ENV{MACA_PATH}/lib/libmcblas.so)
+endif()
+
 if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_wna16.cu")
 endif()
@@ -685,7 +762,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_MOE_EXT_SRC "${VLLM_MOE_WNA16_SRC}")
   # 9.0 for latest bf16 atomicAdd PTX
   cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0;8.7;9.0+PTX" "${CUDA_ARCHS}")
-  if (MARLIN_MOE_ARCHS)
+  if (MARLIN_MOE_ARCHS AND NOT USE_MACA)
 
     #
     # For the Marlin MOE kernels we automatically generate sources for various
@@ -733,6 +810,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 
     message(STATUS "Building Marlin MOE kernels for archs: ${MARLIN_MOE_ARCHS}")
   else()
+    message(WARNING "Use MACA, Skip MARLIN_MOE_ARCHS.")
     message(STATUS "Not building Marlin MOE kernels as no compatible archs found"
                    " in CUDA target architectures")
   endif()
@@ -750,6 +828,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}")
 endif()
 message(STATUS "Enabling moe extension.")
+
+set(BLAS_API_ARGS "")
+if(USE_MACA AND ENABLE_BLAS_API)
+  message(STATUS "Blas API for fused moe enabled")
+  list(APPEND BLAS_API_ARGS 
+      INCLUDE_DIRECTORIES ${MCBLAS_INCLUDE_DIR}
+      LIBRARIES ${MCBLAS_LIB})
+endif()
+
 define_gpu_extension_target(
   _moe_C
   DESTINATION vllm
@@ -759,6 +846,7 @@ define_gpu_extension_target(
   ARCHITECTURES ${VLLM_GPU_ARCHES}
   INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
   INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
+  ${BLAS_API_ARGS}
   USE_SABI 3
   WITH_SOABI)
 
@@ -783,7 +871,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
 endif()
 
 # For CUDA we also build and ship some external projects.
-if (VLLM_GPU_LANG STREQUAL "CUDA")
+if (VLLM_GPU_LANG STREQUAL "CUDA" AND NOT USE_MACA)
     include(cmake/external_projects/flashmla.cmake)
 
     # vllm-flash-attn should be last as it overwrites some CMake functions
diff --git a/cmake/utils.cmake b/cmake/utils.cmake
index 6d90555f2..79a0105e5 100644
--- a/cmake/utils.cmake
+++ b/cmake/utils.cmake
@@ -89,7 +89,7 @@ endfunction()
 #
 # Get additional GPU compiler flags from torch.
 #
-function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG)
+function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG USE_MACA)
   if (${GPU_LANG} STREQUAL "CUDA")
     #
     # Get common NVCC flags from torch.
@@ -98,6 +98,18 @@ function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG)
       "from torch.utils.cpp_extension import COMMON_NVCC_FLAGS; print(';'.join(COMMON_NVCC_FLAGS))"
       "Failed to determine torch nvcc compiler flags")
 
+    if (USE_MACA)
+      message(WARNING "Use MACA, Overwrite GPU_FLAGS.")
+      set(GPU_FLAGS 
+        "-D__CUDA_NO_HALF_OPERATORS__"
+        "-D__CUDA_NO_HALF_CONVERSIONS__"
+        "-D__CUDA_NO_HALF2_OPERATORS__"
+        "--expt-relaxed-constexpr")
+      list(APPEND GPU_FLAGS 
+        "-mllvm" 
+        "-metaxgpu-GridDim-UseLdu")
+    endif()
+
     if (CUDA_VERSION VERSION_GREATER_EQUAL 11.8)
       list(APPEND GPU_FLAGS "-DENABLE_FP8")
     endif()
@@ -105,8 +117,11 @@ function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG)
       list(REMOVE_ITEM GPU_FLAGS
         "-D__CUDA_NO_HALF_OPERATORS__"
         "-D__CUDA_NO_HALF_CONVERSIONS__"
-        "-D__CUDA_NO_BFLOAT16_CONVERSIONS__"
         "-D__CUDA_NO_HALF2_OPERATORS__")
+      if (not USE_MACA)
+        list(REMOVE_ITEM GPU_FLAGS
+        "-D__CUDA_NO_BFLOAT16_CONVERSIONS__")
+      endif()
     endif()
 
   elseif(${GPU_LANG} STREQUAL "HIP")
@@ -148,6 +163,13 @@ macro(clear_cuda_arches CUDA_ARCH_FLAGS)
     string(REGEX MATCHALL "-gencode arch=[^ ]+" CUDA_ARCH_FLAGS
       ${CMAKE_CUDA_FLAGS})
 
+    # change the config's value of metaxgpu-disable-bsm-offset
+    string(REPLACE "-metaxgpu-disable-bsm-offset=1" "-metaxgpu-disable-bsm-offset=0"
+            CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS})
+
+    # opt of cutlass w8a8
+    string(APPEND CMAKE_CUDA_FLAGS " -mllvm -structurizecfg-skip-uniform-regions=true")
+
     # Remove all `-gencode` flags from `CMAKE_CUDA_FLAGS` since they will be modified
     # and passed back via the `CUDA_ARCHITECTURES` property.
     string(REGEX REPLACE "-gencode arch=[^ ]+ *" "" CMAKE_CUDA_FLAGS
diff --git a/csrc/activation_kernels.cu b/csrc/activation_kernels.cu
index 55e659679..d39c87209 100644
--- a/csrc/activation_kernels.cu
+++ b/csrc/activation_kernels.cu
@@ -1,3 +1,4 @@
+// 2025 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
 #include <ATen/cuda/CUDAContext.h>
 #include <torch/all.h>
 #include <c10/cuda/CUDAGuard.h>
@@ -31,10 +32,120 @@ __global__ void act_and_mul_kernel(
   }
 }
 
+template<typename scalar_t, typename VT, int N, scalar_t (*ACT_FN)(const scalar_t&),
+        bool act_first>
+__global__ void act_and_mul_kernel_bd_opt(
+  scalar_t* __restrict__ out,               // [..., d]
+  const scalar_t* __restrict__ input,       // [..., 2, d]
+  const int d,
+  const int blockDim_x
+){
+    const int64_t token_idx = blockIdx.y;
+    int x_offset = (blockIdx.x * blockDim_x + threadIdx.x) * N;
+    if(x_offset >= d) return;
+    int64_t offset0 = token_idx * d;
+    int64_t offset1 = offset0 << 1;
+    const scalar_t* ptr_input = input + offset1;
+    const scalar_t* ptr_input0 = ptr_input + x_offset;
+    const scalar_t* ptr_input1 = ptr_input0 + d;
+    scalar_t* ptr_output = out + offset0 + x_offset;
+    VT vsrc0 = *(VT*)(ptr_input0);
+    VT vsrc1 = *(VT*)(ptr_input1);
+    VT vdst;
+    scalar_t* ptr_src0 = (scalar_t*)&vsrc0;
+    scalar_t* ptr_src1 = (scalar_t*)&vsrc1;
+    scalar_t* ptr_dst = (scalar_t*)&vdst;
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+        ptr_dst[i] = compute<scalar_t, ACT_FN, act_first>(ptr_src0[i], ptr_src1[i]);
+    }
+    *(VT*)(ptr_output) = vdst;
+}
+
+template<typename scalar_t, typename VT, int N, scalar_t (*ACT_FN)(const scalar_t&),
+        bool act_first>
+__global__ void act_and_mul_kernel_sd_opt(
+  scalar_t* __restrict__ out,               // [..., d]
+  const scalar_t* __restrict__ input,       // [..., 2, d]
+  const int d,
+  const int blockDim_x,
+  const int gridDim_x,
+  const int token_per_block,
+  const int max_token_num) {
+    __shared__ int8_t sm_buffer[16384];
+    int token_offset = blockIdx.x * token_per_block;
+    int out_offset = token_offset * d;
+    int in_offset = out_offset << 1;
+    int num_token = min(max_token_num - token_offset, token_per_block);
+    if(num_token <= 0) return;
+    const scalar_t *ptr_block_input = input + in_offset;
+    scalar_t* ptr_block_output = out + out_offset;
+    int output_size = num_token * d;
+    int input_size = output_size << 1;
+    
+    scalar_t* ptr_sm_buffer = (scalar_t*)sm_buffer;
+    int stride = blockDim_x * N;
+    for(int i = threadIdx.x*N; i < input_size; i += stride) {
+        *(VT*)(ptr_sm_buffer + i) = *(VT*)(ptr_block_input + i);
+    }
+    __syncthreads();
+    for(int i = threadIdx.x; i < output_size; i += blockDim_x) {
+      int token_id = i / d;
+      int x_offset = i % d;
+      scalar_t *ptr_input0 = ptr_sm_buffer + token_id * d * 2 + x_offset;
+      scalar_t *ptr_input1 = ptr_input0 + d;
+      *(ptr_block_output + i) = compute<scalar_t, ACT_FN, act_first>(*ptr_input0, ptr_input1[0]);
+    }
+}
+
+template<typename scalar_t, typename VT, int N, scalar_t (*ACT_FN)(const scalar_t&),
+        bool act_first>
+__global__ void act_and_mul_kernel_sd_fast_opt(
+  scalar_t* __restrict__ out,               // [..., d]
+  const scalar_t* __restrict__ input,       // [..., 2, d]
+  const int d,
+  const int blockDim_x,
+  const int gridDim_x,
+  const int token_per_block,
+  const int max_token_num) {
+    __shared__ int8_t sm_buffer[16384];
+    int token_offset = blockIdx.x * token_per_block;
+    int out_offset = token_offset * d;
+    int in_offset = out_offset << 1;
+    int num_token = min(max_token_num - token_offset, token_per_block);
+    if(num_token <= 0) return;
+    const scalar_t *ptr_block_input = input + in_offset;
+    scalar_t* ptr_block_output = out + out_offset;
+    int output_size = num_token * d;
+    int input_size = output_size << 1;
+    
+    scalar_t* ptr_sm_buffer = (scalar_t*)sm_buffer;
+    int stride = blockDim_x * N;
+    for(int i = threadIdx.x*N; i < input_size; i += stride) {
+        *(VT*)(ptr_sm_buffer + i) = *(VT*)(ptr_block_input + i);
+    }
+    __syncthreads();
+    for(int i = threadIdx.x*N; i < output_size; i += stride) {
+      int token_id = i / d;
+      int x_offset = i % d;
+      scalar_t *ptr_input0 = ptr_sm_buffer + token_id * d * 2 + x_offset;
+      scalar_t *ptr_input1 = ptr_input0 + d;
+      VT vdst;
+      scalar_t *ptr_dst = (scalar_t*)&vdst;
+      #pragma unroll N
+      for(int j = 0; j < N; j++) {
+        ptr_dst[j] = compute<scalar_t, ACT_FN, act_first>(ptr_input0[j], ptr_input1[j]);
+      }
+      *(VT*)(ptr_block_output + i) = vdst;
+    }
+}
+
 template <typename T>
 __device__ __forceinline__ T silu_kernel(const T& x) {
   // x * sigmoid(x)
-  return (T)(((float)x) / (1.0f + expf((float)-x)));
+  // return (T)(((float)x) / (1.0f + expf((float)-x)));
+  float x_f = (float)x;
+  return (T) ((x_f) / (1.0f + __builtin_expf(-x_f)));
 }
 
 template <typename T>
@@ -68,6 +179,54 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
 #define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST)                 \
   int d = input.size(-1) / 2;                                            \
   int64_t num_tokens = input.numel() / input.size(-1);                   \
+  int n = 16 / input.element_size();                                                          \
+  if(((d&(n - 1)) == 0) && d >= 512 * n) {\
+    int blocksize = 512;                                                                  \
+    dim3 gridsize((d + 512*n - 1) / (512*n), num_tokens,1);                               \
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));                     \
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                         \
+    VLLM_DISPATCH_FLOATING_TYPES(                                                         \
+      input.scalar_type(),                                                                \
+      "act_and_mul_kernel_bd_opt",                                                               \
+      [&] {                                                                               \
+        vllm::act_and_mul_kernel_bd_opt<scalar_t, float4, 16 / sizeof(scalar_t), KERNEL<scalar_t>, ACT_FIRST><<<gridsize, blocksize, 0, stream>>>(   \
+          out.data_ptr<scalar_t>(),                                                         \
+          input.data_ptr<scalar_t>(),                                                       \
+          d, blocksize);                                                                    \
+      });                                                                                   \
+  } else if(d < 512 && (d & (n - 1)) == 0) {                                                \
+        int block_token = 16384 / input.element_size() / 2 / d;                             \
+        block_token = block_token / n * n;                                                  \
+        int blocksize = 512;                                                                \
+        int gridsize = (num_tokens + block_token - 1) / block_token;                        \
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(input));                   \
+        const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                       \
+        VLLM_DISPATCH_FLOATING_TYPES(                                                       \
+        input.scalar_type(),                                                                \
+        "act_and_mul_kernel_sd_fast_opt",                                                   \
+        [&] {                                                                               \
+        vllm::act_and_mul_kernel_sd_fast_opt<scalar_t, float4, 16 / sizeof(scalar_t), KERNEL<scalar_t>, ACT_FIRST><<<gridsize, blocksize, 0, stream>>>(   \
+        out.data_ptr<scalar_t>(),                                                       \
+        input.data_ptr<scalar_t>(),                                                     \
+        d, blocksize,gridsize,block_token,num_tokens);                                  \
+        });                                                                                 \
+  } else if(d < 512) {                                                                      \
+        int block_token = 16384 / input.element_size() / 2 / d;                             \
+        block_token = block_token / n * n;                                                  \
+        int blocksize = 512;                                                                \
+        int gridsize = (num_tokens + block_token - 1) / block_token;                        \
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(input));                   \
+        const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                       \
+        VLLM_DISPATCH_FLOATING_TYPES(                                                       \
+        input.scalar_type(),                                                                \
+        "act_and_mul_kernel_sd_opt",                                                        \
+        [&] {                                                                               \
+          vllm::act_and_mul_kernel_sd_opt<scalar_t, float4, 16 / sizeof(scalar_t), KERNEL<scalar_t>, ACT_FIRST><<<gridsize, blocksize, 0, stream>>>(   \
+            out.data_ptr<scalar_t>(),                                                       \
+            input.data_ptr<scalar_t>(),                                                     \
+            d, blocksize,gridsize,block_token,num_tokens);                                  \
+        });                                                                                 \
+  } else {                                                                                  \
   dim3 grid(num_tokens);                                                 \
   dim3 block(std::min(d, 1024));                                         \
   if (num_tokens == 0) {                                                 \
@@ -76,11 +235,15 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
   const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
   VLLM_DISPATCH_FLOATING_TYPES(                                          \
-      input.scalar_type(), "act_and_mul_kernel", [&] {                   \
-        vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST>  \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });
+      input.scalar_type(),                                                                  \
+      "act_and_mul_kernel",                                                                 \
+      [&] {                                                                                 \
+        vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST><<<grid, block, 0, stream>>>(   \
+          out.data_ptr<scalar_t>(),                                                         \
+          input.data_ptr<scalar_t>(),                                                       \
+          d);                                                                               \
+      });                                                                                   \
+  }
 
 void silu_and_mul(torch::Tensor& out,    // [..., d]
                   torch::Tensor& input)  // [..., 2 * d]
diff --git a/csrc/attention/dtype_float16.cuh b/csrc/attention/dtype_float16.cuh
index 3a1815f0e..05ab38d71 100644
--- a/csrc/attention/dtype_float16.cuh
+++ b/csrc/attention/dtype_float16.cuh
@@ -22,6 +22,7 @@
 
 #include "attention_generic.cuh"
 #include "dtype_float32.cuh"
+#include "cuda_fp16.h"
 
 #ifdef USE_ROCM
   #include <hip/hip_fp16.h>
@@ -69,6 +70,12 @@ struct FloatVec<uint4> {
 
 // Utility functions for type conversions.
 inline __device__ uint32_t h0_h0(uint16_t a) {
+#ifdef USE_MACA
+  uint32_t b;
+  b = a;
+  b = b << 16 | b;
+  return b;
+#else
 #ifndef USE_ROCM
   uint32_t b;
   asm volatile("mov.b32 %0, {%1, %1};" : "=r"(b) : "h"(a));
@@ -82,19 +89,35 @@ inline __device__ uint32_t h0_h0(uint16_t a) {
   tmp.u16[1] = a;
   return tmp.u32;
 #endif
+#endif  // USE_MACA
 }
 
 inline __device__ float half_to_float(uint16_t h) {
   float f;
+#ifdef USE_MACA
+  f = __half2float(*(__half*)&h);
+#else
 #ifndef USE_ROCM
   asm volatile("cvt.f32.f16 %0, %1;\n" : "=f"(f) : "h"(h));
 #else
   asm volatile("v_cvt_f32_f16 %0, %1;" : "=v"(f) : "v"(h));
 #endif
+#endif // USE_MACA
   return f;
 }
 
 inline __device__ float2 half2_to_float2(uint32_t v) {
+#ifdef USE_MACA
+  uint16_t lo, hi;
+  union {
+    uint32_t u32;
+    uint16_t u16[2];
+  } tmp;
+  tmp.u32 = v;
+  lo = tmp.u16[0];
+  hi = tmp.u16[1];
+  return make_float2(half_to_float(lo), half_to_float(hi));
+#else
 #ifndef USE_ROCM
   uint16_t lo, hi;
   asm volatile("mov.b32 {%0, %1}, %2;\n" : "=h"(lo), "=h"(hi) : "r"(v));
@@ -110,6 +133,7 @@ inline __device__ float2 half2_to_float2(uint32_t v) {
   ret.y = half_to_float(tmp.u16[1]);
   return ret;
 #endif
+#endif // USE_MACA
 }
 
 inline __device__ uint16_t float_to_half(float f) {
@@ -117,11 +141,16 @@ inline __device__ uint16_t float_to_half(float f) {
     uint32_t u32;
     uint16_t u16[2];
   } tmp;
+#ifdef USE_MACA
+  __half __tmp = __float2half(f);
+  tmp.u16[0] = *(uint16_t*)&__tmp;
+#else
 #ifndef USE_ROCM
   asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f));
 #else
   asm volatile("v_cvt_f16_f32 %0, %1;\n" : "=v"(tmp.u32) : "v"(f));
 #endif
+#endif // USE MACA
   return tmp.u16[0];
 }
 
@@ -130,6 +159,15 @@ inline __device__ uint32_t float2_to_half2(float2 f) {
     uint32_t u32;
     uint16_t u16[2];
   } tmp;
+#ifdef USE_MACA
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  __half2 __tmp = __half2(__float2half(f.x), __float2half(f.y));
+  tmp.u32 = *(uint32_t*)&__tmp;
+  #else
+  tmp.u16[0] = float_to_half(f.x);
+  tmp.u16[1] = float_to_half(f.y);
+  #endif
+#else
 #ifndef USE_ROCM
   #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
   asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n"
@@ -143,27 +181,42 @@ inline __device__ uint32_t float2_to_half2(float2 f) {
   tmp.u16[0] = float_to_half(f.x);
   tmp.u16[1] = float_to_half(f.y);
 #endif
+#endif // USE_MACA
   return tmp.u32;
 }
 
 // Vector addition.
 inline __device__ uint16_t add(uint16_t a, uint16_t b) {
   uint16_t c;
+#ifdef USE_MACA
+  unsigned short __a=(a);
+  unsigned short __b=(b);
+  __half __d=__hadd(*(__half*)&__a,*(__half*)&__b);
+  (c)=*(unsigned short*)&__d;
+#else
 #ifndef USE_ROCM
   asm volatile("add.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b));
 #else
   asm volatile("v_add_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b));
 #endif
+#endif // USE_MACA
   return c;
 }
 
 inline __device__ uint32_t add(uint32_t a, uint32_t b) {
   uint32_t c;
+#ifdef USE_MACA
+  unsigned int __a=(a);
+  unsigned int __b=(b);
+  __half2 __d=__hadd2(*(__half2*)&__a,*(__half2*)&__b);
+  (c)=*(unsigned int*)&__d;
+#else
 #ifndef USE_ROCM
   asm volatile("add.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b));
 #else
   asm volatile("v_pk_add_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b));
 #endif
+#endif // USE_MACA
   return c;
 }
 
@@ -208,22 +261,36 @@ inline __device__ Float8_ add(uint4 a, Float8_ fb) {
 template <>
 inline __device__ uint16_t mul(uint16_t a, uint16_t b) {
   uint16_t c;
+#ifdef USE_MACA
+  unsigned short __a=(a);
+  unsigned short __b=(b);
+  __half __d=__hmul(*(__half*)&__a,*(__half*)&__b);
+  (c)=*(unsigned short*)&__d;
+#else
 #ifndef USE_ROCM
   asm volatile("mul.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b));
 #else
   asm volatile("v_mul_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b));
 #endif
+#endif // USE_MACA
   return c;
 }
 
 template <>
 inline __device__ uint32_t mul(uint32_t a, uint32_t b) {
   uint32_t c;
+#ifdef USE_MACA
+  unsigned int __a=(a);
+  unsigned int __b=(b);
+  __half2 __d=__hmul2(*(__half2*)&__a,*(__half2*)&__b);
+  (c)=*(unsigned int*)&__d;
+#else
 #ifndef USE_ROCM
   asm volatile("mul.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b));
 #else
   asm volatile("v_pk_mul_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b));
 #endif
+#endif //USE_MACA
   return c;
 }
 
@@ -330,6 +397,13 @@ inline __device__ Float8_ mul(uint16_t a, uint4 b) {
 // Vector fused multiply-add.
 inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) {
   uint32_t d;
+#ifdef USE_MACA
+  unsigned int __a=(a);
+  unsigned int __b=(b);
+  unsigned int __c=(c);
+  __half2 __d=__hfma2(*(__half2*)&__a,*(__half2*)&__b,*(__half2*)&__c);
+  (d)=*(unsigned int*)&__d;
+#else
 #ifndef USE_ROCM
   asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
                : "=r"(d)
@@ -339,6 +413,7 @@ inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) {
                : "=v"(d)
                : "v"(a), "v"(b), "v"(c));
 #endif
+#endif // USE_MACA
   return d;
 }
 
diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp
index fab6ca36d..41015355d 100644
--- a/csrc/cumem_allocator.cpp
+++ b/csrc/cumem_allocator.cpp
@@ -3,8 +3,6 @@
 // need to be unsigned long long
 #include <iostream>
 
-extern "C" {
-
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
 
@@ -12,6 +10,8 @@ extern "C" {
 #include <cuda_runtime_api.h>
 #include <cuda.h>
 
+extern "C" {
+
 char error_msg[10240];  // 10KB buffer to store error messages
 CUresult no_error = CUresult(0);
 CUresult error_code = no_error;  // store error code
diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh
index 44709b459..18c38855f 100644
--- a/csrc/custom_all_reduce.cuh
+++ b/csrc/custom_all_reduce.cuh
@@ -56,9 +56,15 @@ struct Signal {
   alignas(128) FlagType _flag[kMaxBlocks];  // incremental flags for each rank
 };
 
+#ifdef USE_MACA
+struct __align__(16) RankData { 
+  const void* ptrs[8]; 
+};
+#else
 struct __align__(16) RankData {
   const void* ptrs[8];
 };
+#endif // USE_MACA
 
 struct __align__(16) RankSignals {
   Signal* signals[8];
@@ -155,6 +161,7 @@ DINLINE O downcast(array_t<float, O::size> val) {
 #if !defined(USE_ROCM)
 
 static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
+  #ifndef USE_MACA
   #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
   asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
                "l"(flag_addr));
@@ -162,10 +169,12 @@ static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
   asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
                "l"(flag_addr));
   #endif
+  #endif // USE_MACA
 }
 
 static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
   FlagType flag;
+  #ifndef USE_MACA
   #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
   asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
                : "=r"(flag)
@@ -175,18 +184,23 @@ static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
                : "=r"(flag)
                : "l"(flag_addr));
   #endif
+  #endif // USE_MACA
   return flag;
 }
 
 static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
+#ifndef USE_MACA
   asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
+#endif // USE_MACA
 }
 
 static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
   FlagType flag;
+#ifndef USE_MACA
   asm volatile("ld.volatile.global.u32 %0, [%1];"
                : "=r"(flag)
                : "l"(flag_addr));
+#endif // USE_MACA
   return flag;
 }
 
@@ -365,7 +379,9 @@ __global__ void __launch_bounds__(512, 1)
 
 using IPC_KEY = std::array<uint8_t, sizeof(cudaIpcMemHandle_t)>;
 static_assert(sizeof(IPC_KEY) == sizeof(cudaIpcMemHandle_t));
+#ifndef USE_MACA
 static_assert(alignof(IPC_KEY) == alignof(cudaIpcMemHandle_t));
+#endif // USE_MACA
 
 class CustomAllreduce {
  public:
diff --git a/csrc/cutlass_extensions/common.hpp b/csrc/cutlass_extensions/common.hpp
index 195872e8e..858cecd2c 100644
--- a/csrc/cutlass_extensions/common.hpp
+++ b/csrc/cutlass_extensions/common.hpp
@@ -1,6 +1,10 @@
 #pragma once
 
+#ifndef USE_MACA
 #include "cutlass/cutlass.h"
+#else
+#include "mctlass/mctlass.h"
+#endif // USE_MACA
 #include <climits>
 #include "cuda_runtime.h"
 #include <iostream>
@@ -8,12 +12,21 @@
 /**
  * Helper function for checking CUTLASS errors
  */
+#ifndef USE_MACA
 #define CUTLASS_CHECK(status)                       \
   {                                                 \
     cutlass::Status error = status;                 \
     TORCH_CHECK(error == cutlass::Status::kSuccess, \
                 cutlassGetStatusString(error));     \
   }
+#else
+#define CUTLASS_CHECK(status)                       \
+  {                                                 \
+    mctlass::Status error = status;                 \
+    TORCH_CHECK(error == mctlass::Status::kSuccess, \
+                mctlassGetStatusString(error));     \
+  }
+#endif // USE_MACA
 
 inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) {
   int max_shared_mem_per_block_opt_in = 0;
@@ -34,7 +47,11 @@ int32_t get_sm_version_num();
 template <typename Kernel>
 struct enable_sm90_or_later : Kernel {
   template <typename... Args>
+#ifndef USE_MACA
   CUTLASS_DEVICE void operator()(Args&&... args) {
+#else
+  MCTLASS_DEVICE void operator()(Args&&... args) {
+#endif // USE_MACA
 #if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 900
     Kernel::operator()(std::forward<Args>(args)...);
 #endif
@@ -44,7 +61,11 @@ struct enable_sm90_or_later : Kernel {
 template <typename Kernel>
 struct enable_sm90_only : Kernel {
   template <typename... Args>
+#ifndef USE_MACA
   CUTLASS_DEVICE void operator()(Args&&... args) {
+#else
+  MCTLASS_DEVICE void operator()(Args&&... args) {
+#endif // USE_MACA
 #if defined __CUDA_ARCH__ && __CUDA_ARCH__ == 900
     Kernel::operator()(std::forward<Args>(args)...);
 #endif
@@ -54,7 +75,11 @@ struct enable_sm90_only : Kernel {
 template <typename Kernel>
 struct enable_sm100_only : Kernel {
   template <typename... Args>
+#ifndef USE_MACA
   CUTLASS_DEVICE void operator()(Args&&... args) {
+#else
+  MCTLASS_DEVICE void operator()(Args&&... args) {
+#endif // USE_MACA
 #if defined __CUDA_ARCH__ && __CUDA_ARCH__ == 1000
     Kernel::operator()(std::forward<Args>(args)...);
 #endif
diff --git a/csrc/mamba/causal_conv1d/causal_conv1d.cu b/csrc/mamba/causal_conv1d/causal_conv1d.cu
index f62d08c17..b1f15930e 100644
--- a/csrc/mamba/causal_conv1d/causal_conv1d.cu
+++ b/csrc/mamba/causal_conv1d/causal_conv1d.cu
@@ -187,7 +187,7 @@ void causal_conv1d_fwd(const at::Tensor &x, const at::Tensor &weight,
 
     // Otherwise the kernel will be launched from cuda:0 device
     // Cast to char to avoid compiler warning about narrowing
-    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
+    at::cuda::CUDAGuard device_guard{static_cast<c10::DeviceIndex>(x.get_device())};
     auto stream = at::cuda::getCurrentCUDAStream().stream();
     DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_fwd", [&] {
             causal_conv1d_fwd_cuda<input_t, weight_t>(params, stream);
@@ -280,7 +280,7 @@ void causal_conv1d_update(const at::Tensor &x,
 
     // Otherwise the kernel will be launched from cuda:0 device
     // Cast to char to avoid compiler warning about narrowing
-    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
+    at::cuda::CUDAGuard device_guard{static_cast<c10::DeviceIndex>(x.get_device())};
     auto stream = at::cuda::getCurrentCUDAStream().stream();
     DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_update", [&] {
             causal_conv1d_update_cuda<input_t, weight_t>(params, stream);
diff --git a/csrc/mamba/mamba_ssm/selective_scan_fwd.cu b/csrc/mamba/mamba_ssm/selective_scan_fwd.cu
index 0c9df925b..093f7f284 100644
--- a/csrc/mamba/mamba_ssm/selective_scan_fwd.cu
+++ b/csrc/mamba/mamba_ssm/selective_scan_fwd.cu
@@ -331,7 +331,17 @@ void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
 
 template<typename input_t, typename weight_t>
 void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
-
+    #ifdef USE_MACA
+        if (params.seqlen <= 256) {
+            selective_scan_fwd_launch<64, 4, input_t, weight_t>(params, stream);
+        } else if (params.seqlen <= 512) {
+            selective_scan_fwd_launch<64, 8, input_t, weight_t>(params, stream);
+        } else if (params.seqlen <= 1024) {
+            selective_scan_fwd_launch<64, 16, input_t, weight_t>(params, stream);
+        } else {
+            selective_scan_fwd_launch<128, 16, input_t, weight_t>(params, stream);
+        }
+    #else
     #ifndef USE_ROCM
         if (params.seqlen <= 128) {           
             selective_scan_fwd_launch<32, 4, input_t, weight_t>(params, stream);
@@ -355,6 +365,7 @@ void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
             selective_scan_fwd_launch<128, 16, input_t, weight_t>(params, stream);
         }
     #endif
+    #endif // USE_MACA
 }
 
 template void selective_scan_fwd_cuda<at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
@@ -649,7 +660,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
     
     // Otherwise the kernel will be launched from cuda:0 device
     // Cast to char to avoid compiler warning about narrowing
-    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    at::cuda::CUDAGuard device_guard{static_cast<c10::DeviceIndex>(u.get_device())};
     auto stream = at::cuda::getCurrentCUDAStream().stream();
     DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_fwd", [&] {
         selective_scan_fwd_cuda<input_t, weight_t>(params, stream);
diff --git a/csrc/moe/moe_align_sum_kernels.cu b/csrc/moe/moe_align_sum_kernels.cu
index 6b6a9d04a..80fb62430 100644
--- a/csrc/moe/moe_align_sum_kernels.cu
+++ b/csrc/moe/moe_align_sum_kernels.cu
@@ -197,6 +197,240 @@ __global__ void moe_align_block_size_global_mem_kernel(
   }
 }
 
+
+
+__device__ __forceinline__ int32_t ScanWarp2(int32_t val, int32_t mask = 8) {
+  int32_t lane = threadIdx.x & 31;
+  int32_t tmp = __shfl_up_sync(0xffffffff, val, 1, mask);
+  if (lane >= 1) {
+    val += tmp;
+  }
+  tmp = __shfl_up_sync(0xffffffff, val, 2, mask);
+  if (lane >= 2) {
+    val += tmp;
+  }
+  tmp = __shfl_up_sync(0xffffffff, val, 4, mask);
+  if (lane >= 4) {
+    val += tmp;
+  }
+  return val;
+}
+
+__device__ __forceinline__ int32_t ScanWarp(int32_t val) {
+  int32_t lane = threadIdx.x & 31;
+  int32_t tmp = __shfl_up_sync(0xffffffff, val, 1);
+  if (lane >= 1) {
+    val += tmp;
+  }
+  tmp = __shfl_up_sync(0xffffffff, val, 2);
+  if (lane >= 2) {
+    val += tmp;
+  }
+  tmp = __shfl_up_sync(0xffffffff, val, 4);
+  if (lane >= 4) {
+    val += tmp;
+  }
+  tmp = __shfl_up_sync(0xffffffff, val, 8);
+  if (lane >= 8) {
+    val += tmp;
+  }
+  tmp = __shfl_up_sync(0xffffffff, val, 16);
+  if (lane >= 16) {
+    val += tmp;
+  }
+  return val;
+}
+
+
+__device__ inline void copy(const void* local, void* data)
+{
+    const float4* in = static_cast<const float4*>(local);
+    float4* out = static_cast<float4*>(data);
+    *out = *in;
+}
+
+template <typename scalar_t, int VPT>
+__global__ void opt_sgl_moe_align_block_size_kernel_stage_2(size_t numel, int32_t* local_offsets, scalar_t* __restrict__ topk_ids, int32_t* sorted_token_ids){
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx * VPT >= numel) return;
+  int32_t expert_local[VPT];
+  copy(topk_ids + idx * VPT, expert_local);
+  for(int i = 0; i < VPT && i + idx * VPT < numel; ++i) {
+    int32_t expert_id = expert_local[i];
+    int32_t rank_post_pad = atomicAdd(&local_offsets[expert_id], 1);
+    sorted_token_ids[rank_post_pad] = i + idx * VPT;
+  }
+}
+
+template <typename scalar_t>
+__global__ void opt_sgl_moe_align_block_size_kernel_stage_1(
+    scalar_t* __restrict__ topk_ids, int32_t* sorted_token_ids,
+    int32_t* expert_ids, int32_t* total_tokens_post_pad, int32_t num_experts,
+    int32_t block_size, size_t numel, int32_t* offsets) {
+  __shared__ int32_t shared_counts[32][8];
+  __shared__ int32_t local_offsets[256];
+  __shared__ int32_t cum_test[257];
+  __shared__ int32_t blocksum[32];
+
+  const int warp_id = threadIdx.x / 32;
+  const int lane_id = threadIdx.x % 32;
+  const int experts_per_warp = 8;
+  const int my_expert_start = warp_id * experts_per_warp;
+
+  int32_t idx = threadIdx.x;
+  if(idx < num_experts){
+    shared_counts[idx / 8][idx % 8] = 0;
+    cum_test[0] = 0;
+  }
+  __syncthreads();
+
+  const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
+  const size_t start_idx = threadIdx.x * tokens_per_thread;
+
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int expert_id = topk_ids[i];
+    int warp_idx = expert_id / experts_per_warp;
+    int expert_offset = expert_id % experts_per_warp;
+    atomicAdd(&shared_counts[warp_idx][expert_offset], 1);
+  }
+  __syncthreads();
+
+  int val = 0;
+  if (threadIdx.x < 256) {
+    int32_t final_val = 0;
+    int row = idx / 8;
+    int line = idx % 8;
+    val = shared_counts[row][line];
+    val = CEILDIV(val, block_size) * block_size;
+  }
+  __syncthreads();
+
+  if(idx < 256) {
+    int tmp = 0;
+    val = ScanWarp(val);
+    if(lane_id == 31) {
+      blocksum[warp_id] = val;
+    }
+  }
+  __syncthreads();
+
+  if(warp_id == 0 && lane_id < 8) {
+    int res = blocksum[lane_id];
+    blocksum[lane_id] = ScanWarp2(res);
+  }
+  __syncthreads();
+
+  if(threadIdx.x < 256 && warp_id > 0){
+    val += blocksum[warp_id - 1];
+  }
+  __syncthreads();
+
+  if(idx < 256){
+    cum_test[idx + 1] = val;
+  }
+  __syncthreads();
+
+  if (threadIdx.x < num_experts) {
+    for (int i = cum_test[threadIdx.x]; i < cum_test[threadIdx.x + 1];
+         i += block_size) {
+      expert_ids[i / block_size] = threadIdx.x;
+    }
+    local_offsets[threadIdx.x] = cum_test[threadIdx.x];
+    offsets[threadIdx.x] = cum_test[threadIdx.x];
+    if(threadIdx.x == 0){
+      *total_tokens_post_pad = cum_test[num_experts];
+    }
+  }
+}
+
+
+template <typename scalar_t>
+__global__ void opt_sgl_moe_align_block_size_kernel(
+    scalar_t* __restrict__ topk_ids, int32_t* sorted_token_ids,
+    int32_t* expert_ids, int32_t* total_tokens_post_pad, int32_t num_experts,
+    int32_t block_size, size_t numel, int32_t* cumsum) {
+  __shared__ int32_t shared_counts[32][8];
+  __shared__ int32_t local_offsets[256];
+  __shared__ int32_t cum_test[257];
+  __shared__ int32_t blocksum[32];
+
+  const int warp_id = threadIdx.x / 32;
+  const int lane_id = threadIdx.x % 32;
+  const int experts_per_warp = 8;
+  const int my_expert_start = warp_id * experts_per_warp;
+
+  int32_t idx = threadIdx.x;
+  if(idx < num_experts){
+    shared_counts[idx / 8][idx % 8] = 0;
+    cum_test[0] = 0;
+  }
+  __syncthreads();
+
+  const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
+  const size_t start_idx = threadIdx.x * tokens_per_thread;
+
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int expert_id = topk_ids[i];
+    int warp_idx = expert_id / experts_per_warp;
+    int expert_offset = expert_id % experts_per_warp;
+    atomicAdd(&shared_counts[warp_idx][expert_offset], 1);
+  }
+  __syncthreads();
+
+  int val = 0;
+  if (threadIdx.x < 256) {
+    int32_t final_val = 0;
+    int row = idx / 8;
+    int line = idx % 8;
+    val = shared_counts[row][line];
+    val = CEILDIV(val, block_size) * block_size;
+  }
+  __syncthreads();
+
+  if(idx < 256) {
+    int tmp = 0;
+    val = ScanWarp(val);
+    if(lane_id == 31) {
+      blocksum[warp_id] = val;
+    }
+  }
+  __syncthreads();
+
+  if(warp_id == 0 && lane_id < 8) {
+    int res = blocksum[lane_id];
+    blocksum[lane_id] = ScanWarp2(res);
+  }
+  __syncthreads();
+
+  if(threadIdx.x < 256 && warp_id > 0){
+    val += blocksum[warp_id - 1];
+  }
+  __syncthreads();
+
+  if(idx < 256){
+    cum_test[idx + 1] = val;
+  }
+  __syncthreads();
+
+  if (threadIdx.x < num_experts) {
+    for (int i = cum_test[threadIdx.x]; i < cum_test[threadIdx.x + 1];
+         i += block_size) {
+      expert_ids[i / block_size] = threadIdx.x;
+    }
+    local_offsets[threadIdx.x] = cum_test[threadIdx.x];
+    if(threadIdx.x == 0){
+      *total_tokens_post_pad = cum_test[num_experts];
+    }
+  }
+  __syncthreads();
+
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int32_t expert_id = topk_ids[i];
+    int32_t rank_post_pad = atomicAdd(&local_offsets[expert_id], 1);
+    sorted_token_ids[rank_post_pad] = i;
+  }
+}
+
 // taken from
 // https://github.com/sgl-project/sglang/commit/cdae77b03dfc6fec3863630550b45bbfc789f957
 template <typename scalar_t>
@@ -290,6 +524,81 @@ __global__ void moe_sum_kernel(
   }
 }
 
+template<typename scalar_t, typename token_cnts_t, int32_t NUM_EXPS, int32_t BLOCK_DIM>
+__global__ void moe_align_block_size_kernel_opt(scalar_t* __restrict__ topk_ids,
+                            int32_t* sorted_token_ids,
+                            int32_t* expert_ids,
+                            int32_t* total_tokens_post_pad,
+                            int32_t num_experts,
+                            int32_t block_size, size_t numel) {
+  __shared__ int32_t shared_counts[32][4];
+  __shared__ int32_t cumsum[NUM_EXPS + 1];
+  __shared__ int32_t blocksum[BLOCK_DIM / WARP_SIZE + 1];
+
+  const size_t tokens_per_thread = CEILDIV(numel, BLOCK_DIM);
+  const size_t start_idx = threadIdx.x * tokens_per_thread;
+  static constexpr int32_t experts_per_warp = 4;
+
+  const int32_t warp_id = threadIdx.x / 32;
+  const int32_t lane_id = threadIdx.x % 32;
+  const int32_t my_expert_start = warp_id * experts_per_warp;
+
+  if (threadIdx.x < NUM_EXPS) {
+    shared_counts[threadIdx.x / experts_per_warp][threadIdx.x % experts_per_warp] = 0;
+    cumsum[0] = 0;
+    blocksum[0] = 0;
+  }
+  __syncthreads();
+
+  for (int32_t i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int32_t expert_id = topk_ids[i];
+    int32_t warp_idx = expert_id / experts_per_warp;
+    int32_t expert_offset = expert_id % experts_per_warp;
+    atomicAdd(&shared_counts[warp_idx][expert_offset], 1);
+  }
+  __syncthreads();
+
+  int32_t val = 0;
+  if (threadIdx.x < NUM_EXPS) {
+    int32_t row = threadIdx.x / experts_per_warp;
+    int32_t line = threadIdx.x % experts_per_warp;
+    val = CEILDIV(shared_counts[row][line], block_size) * block_size;
+    val = ScanWarp(val);
+    if(lane_id == 31) {
+      blocksum[warp_id + 1] = val;
+    }
+  }
+  __syncthreads();
+
+  if(threadIdx.x < BLOCK_DIM / WARP_SIZE) {
+    int32_t res = blocksum[lane_id + 1];
+    blocksum[lane_id + 1] = ScanWarp2(res, BLOCK_DIM / WARP_SIZE);
+  }
+  __syncthreads();
+
+  if (threadIdx.x < NUM_EXPS) {
+    val += blocksum[warp_id];
+    cumsum[threadIdx.x + 1] = val;
+  }
+  __syncthreads();
+
+  if (threadIdx.x < NUM_EXPS) {
+    for (int32_t i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1]; i += block_size) {
+      expert_ids[i / block_size] = threadIdx.x;
+    }
+    if (threadIdx.x == 0) {
+      *total_tokens_post_pad = cumsum[NUM_EXPS];
+    }
+  }
+  __syncthreads();
+
+  for (int32_t i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int32_t expert_id = topk_ids[i];
+    int32_t rank_post_pad = atomicAdd(&cumsum[expert_id], 1);
+    sorted_token_ids[rank_post_pad] = i;
+  }
+}
+
 }  // namespace moe
 }  // namespace vllm
 
@@ -351,7 +660,24 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
               cumsum_buffer.data_ptr<int32_t>());
         });
   } else if (use_i16) {
-    VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
+    if (num_experts == 128){
+      VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
+        topk_ids.scalar_type(), "moe_align_block_size_kernel_opt", [&] {
+          // set dynamic shared mem
+          constexpr int32_t tids = 512;
+          auto kernel =
+              vllm::moe::moe_align_block_size_kernel_opt<scalar_t, uint16_t, 128, tids>;
+          // AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
+          //     (void*)kernel, shared_mem_i16));
+          kernel<<<1, tids, 0, stream>>>(
+              topk_ids.data_ptr<scalar_t>(),
+              sorted_token_ids.data_ptr<int32_t>(),
+              experts_ids.data_ptr<int32_t>(),
+              num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
+              topk_ids.numel());
+        });
+    } else {
+      VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
         topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
           // set dynamic shared mem
           auto kernel =
@@ -365,6 +691,7 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
               num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
               topk_ids.numel());
         });
+    }
   } else {
     VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
         topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
@@ -396,27 +723,44 @@ void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
         // calc needed amount of shared mem for `cumsum` tensors
         auto options_int =
             torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
-        torch::Tensor cumsum_buffer =
+        torch::Tensor offsets_buffer =
             torch::zeros({num_experts + 1}, options_int);
 
-        auto align_kernel =
-            vllm::moe::sgl_moe_align_block_size_kernel<scalar_t>;
-        align_kernel<<<1, 1024, 0, stream>>>(
-            topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>(),
-            experts_ids.data_ptr<int32_t>(),
-            num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
-            topk_ids.numel(), cumsum_buffer.data_ptr<int32_t>());
-
-        const int block_threads = 256;
-        const int num_blocks =
-            (topk_ids.numel() + block_threads - 1) / block_threads;
-        const int max_blocks = 65535;
-        const int actual_blocks = std::min(num_blocks, max_blocks);
-        auto sort_kernel = vllm::moe::sgl_moe_token_sort_kernel<scalar_t>;
-        sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(
-            topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>(),
-            cumsum_buffer.data_ptr<int32_t>(), topk_ids.numel());
-      });
+        if(topk_ids.numel() <= 4096) {
+          auto align_kernel =
+              vllm::moe::opt_sgl_moe_align_block_size_kernel<scalar_t>;
+          align_kernel<<<1, 1024, 0, stream>>>(
+              topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>(),
+              experts_ids.data_ptr<int32_t>(),
+              num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
+                topk_ids.numel(), offsets_buffer.data_ptr<int32_t>());
+        } else {
+           auto kernel = vllm::moe::opt_sgl_moe_align_block_size_kernel_stage_1<scalar_t>;
+          kernel<<<1, 1024, 0, stream>>>(
+              topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>(),
+              experts_ids.data_ptr<int32_t>(),
+              num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
+              topk_ids.numel(), offsets_buffer.data_ptr<int32_t>());
+
+          auto kernel_2 = vllm::moe::opt_sgl_moe_align_block_size_kernel_stage_2<scalar_t, 16 / sizeof(scalar_t)>;
+          int block_size = 256;
+          int grid_size = (topk_ids.numel() + block_size - 1) / block_size;
+          kernel_2<<<grid_size, block_size, 0, stream>>>(
+              topk_ids.numel(), offsets_buffer.data_ptr<int32_t>(), 
+              topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>());
+        }
+
+      //   const int block_threads = 256;
+      //   const int num_blocks =
+      //       (topk_ids.numel() + block_threads - 1) / block_threads;
+      //   const int max_blocks = 65535;
+      //   const int actual_blocks = std::min(num_blocks, max_blocks);
+      //   auto sort_kernel = vllm::moe::sgl_moe_token_sort_kernel<scalar_t>;
+      //   sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(
+      //       topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>(),
+      //       offsets_buffer.data_ptr<int32_t>(), topk_ids.numel());
+    }
+);
 }
 
 void moe_sum(torch::Tensor& input,   // [num_tokens, topk, hidden_size]
diff --git a/csrc/moe/moe_ops.cpp b/csrc/moe/moe_ops.cpp
new file mode 100644
index 000000000..a12aeff32
--- /dev/null
+++ b/csrc/moe/moe_ops.cpp
@@ -0,0 +1,64 @@
+#include "moe_ops.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <mcblas.h>
+#include <maca_fp16.h>
+
+mcblasStatus_t mcblasFusedMoe(mcStream_t stream,
+                              const void *a_ptr,
+                              const void *b_ptr,
+                              void *c_ptr,
+                              const int *sorted_token_ids_ptr,
+                              const int *expert_ids_ptr,
+                              const int *num_tokens_post_padded,
+                              int N,
+                              int K,
+                              int num_valid_tokens,
+                              int sorted_token_ids_len,
+                              int stride_am,
+                              int stride_ak,
+                              int stride_be,
+                              int stride_bk,
+                              int stride_bn,
+                              int stride_cm,
+                              int stride_cn,
+                              int top_k,
+                              bool mul_routed_weight,
+                              const float *topk_weights_ptr,
+                              macaDataType compute_type,
+                              int tileConfig = 0);
+
+void fused_moe_kernel(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& C,
+                    const torch::Tensor& topk_weights, const torch::Tensor& topk_ids,
+                    const torch::Tensor& sorted_token_ids, const torch::Tensor& expert_ids,
+                    const torch::Tensor& num_tokens_post_padded, bool mul_routed_weight, int64_t top_k, int64_t tileConfig) {
+    
+    assert(topk_weights.stride(1) == 1);
+    assert(sorted_token_ids.stride(0) == 1);
+
+    auto stream = at::cuda::getCurrentCUDAStream();
+    macaDataType compute_type = (A.dtype() == at::ScalarType::BFloat16) ? MACA_R_16BF : MACA_R_16F;
+    mcblasFusedMoe(stream,
+                  A.data_ptr(),
+                  B.data_ptr(),
+                  C.data_ptr(),
+                  sorted_token_ids.data_ptr<int>(),
+                  expert_ids.data_ptr<int>(),
+                  num_tokens_post_padded.data_ptr<int>(),
+                  B.size(1),
+                  B.size(2),
+                  topk_ids.numel(),
+                  sorted_token_ids.size(0),
+                  A.stride(0),
+                  A.stride(1),
+                  B.stride(0),
+                  B.stride(2),
+                  B.stride(1),
+                  C.stride(1),
+                  C.stride(2),
+                  static_cast<int>(top_k),
+                  mul_routed_weight,
+                  topk_weights.data_ptr<float>(),
+                  compute_type,
+                  static_cast<int>(tileConfig));
+}
\ No newline at end of file
diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h
index c4faef731..9dc55541a 100644
--- a/csrc/moe/moe_ops.h
+++ b/csrc/moe/moe_ops.h
@@ -30,6 +30,15 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
                              int64_t BLOCK_SIZE_K, int64_t bit);
 #endif
 
+#ifdef USE_MACA
+
+void fused_moe_kernel(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& C,
+                    const torch::Tensor& topk_weights, const torch::Tensor& topk_ids,
+                    const torch::Tensor& sorted_token_ids, const torch::Tensor& expert_ids,
+                    const torch::Tensor& num_tokens_post_padded, bool mul_routed_weight, int64_t top_k, int64_t tileConfig);
+
+#endif
+
 bool moe_permute_unpermute_supported();
 
 void shuffle_rows(const torch::Tensor& input_tensor,
diff --git a/csrc/moe/moe_permute_unpermute_op.cu b/csrc/moe/moe_permute_unpermute_op.cu
index 68f429fac..86aaef251 100644
--- a/csrc/moe/moe_permute_unpermute_op.cu
+++ b/csrc/moe/moe_permute_unpermute_op.cu
@@ -1,12 +1,13 @@
 #include <c10/core/ScalarType.h>
 #include <torch/all.h>
 #include <ATen/cuda/CUDAContext.h>
-#include "permute_unpermute_kernels/moe_permute_unpermute_kernel.h"
-#include "permute_unpermute_kernels/dispatch.h"
 #include "core/registration.h"
 
+#ifndef USE_MACA
 // moe_permute kernels require at least CUDA 12.0
 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
+#include "permute_unpermute_kernels/moe_permute_unpermute_kernel.h"
+#include "permute_unpermute_kernels/dispatch.h"
 
 void moe_permute(
     const torch::Tensor& input,                      // [n_token, hidden]
@@ -216,12 +217,50 @@ void moe_unpermute(const torch::Tensor& input,
 
 #endif
 
+#else
+void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_weights,
+                 torch::Tensor& topk_ids,
+                 const torch::Tensor& token_expert_indicies,
+                 const std::optional<torch::Tensor>& expert_map,
+                 int64_t n_expert, int64_t n_local_expert, int64_t topk,
+                 const std::optional<int64_t>& align_block_size,
+                 torch::Tensor& permuted_input,
+                 torch::Tensor& expert_first_token_offset,
+                 torch::Tensor& src_row_id2dst_row_id_map,
+                 torch::Tensor& m_indices) {
+  TORCH_CHECK(false, "moe_unpermute is not supported on MACA");
+}
+
+void moe_unpermute(
+    const torch::Tensor& permuted_hidden_states,     // [n_token * topk, hidden]
+    const torch::Tensor& topk_weights,               //[n_token, topk]
+    const torch::Tensor& topk_ids,                   // [n_token, topk]
+    const torch::Tensor& src_row_id2dst_row_id_map,  // [n_token, topk]
+    const torch::Tensor& expert_first_token_offset,  // [n_local_expert+1]
+    int64_t n_expert, int64_t n_local_expert, int64_t topk,
+    torch::Tensor& hidden_states  // [n_token, hidden]
+) {
+  TORCH_CHECK(false, "moe_unpermute is not supported on MACA");
+}
+
+void shuffle_rows(const torch::Tensor& input_tensor,
+                  const torch::Tensor& dst2src_map,
+                  torch::Tensor& output_tensor) {
+  TORCH_CHECK(false, "shuffle_rows is not supported on MACA");
+}
+
+#endif // USE_MACA
+
 bool moe_permute_unpermute_supported() {
+#ifndef USE_MACA
 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
   return true;
 #else
   return false;
 #endif
+#else
+  return false;
+#endif
 }
 
 TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
diff --git a/csrc/moe/moe_wna16.cu b/csrc/moe/moe_wna16.cu
index 7b6a111c0..f1badfe23 100644
--- a/csrc/moe/moe_wna16.cu
+++ b/csrc/moe/moe_wna16.cu
@@ -85,9 +85,15 @@ __global__ void moe_wna16_gemm_kernel(
     if (threadIdx.x >= BLOCK_SIZE_N || offset_n >= size_n) return;
 
     float res[64];  // assume BLOCK_SIZE_M <= 64
+#ifndef USE_MACA
     scalar_t2 res2;
     scalar_t2 scale_f2;
     scalar_t2 qzero_f2;
+#else
+    scalar_t2 res2{};
+    scalar_t2 scale_f2{};
+    scalar_t2 qzero_f2{};
+#endif
 
     // note that (size_n * size_k * expert_id) may greater than 2 ** 31
     constexpr int8_t pack_factor = 32 / bit;
@@ -190,7 +196,9 @@ __global__ void moe_wna16_gemm_kernel(
       dequant<scalar_t2, bit>(expert_qweight_tmp[tmp_k % 4], weight_half2);
 
       for (int m = 0; m < num_valid_tokens; m++) {
+#ifndef USE_MACA
         res2 = {};
+#endif
 
 #pragma unroll
         for (int i = 0; i < 16 / bit; i++) {
diff --git a/csrc/moe/moe_wna16_utils.h b/csrc/moe/moe_wna16_utils.h
index 8ef03f0e6..1b74926de 100644
--- a/csrc/moe/moe_wna16_utils.h
+++ b/csrc/moe/moe_wna16_utils.h
@@ -81,18 +81,22 @@ class ScalarType<nv_bfloat16> {
 template <int lut>
 __device__ inline int lop3(int a, int b, int c) {
   int res;
+#ifndef USE_MACA
   asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
                : "=r"(res)
                : "r"(a), "r"(b), "r"(c), "n"(lut));
+#endif
   return res;
 }
 
 template <int start_byte, int mask>
 __device__ inline uint32_t prmt(uint32_t a) {
   uint32_t res;
+#ifndef USE_MACA
   asm volatile("prmt.b32 %0, %1, %2, %3;\n"
                : "=r"(res)
                : "r"(a), "n"(start_byte), "n"(mask));
+#endif
   return res;
 }
 
diff --git a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
index de2c15388..9b3b055f4 100644
--- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
@@ -1,8 +1,7 @@
-
-#include "moe_permute_unpermute_kernel.h"
-
+#ifndef USE_MACA
 // moe_permute kernels require at least CUDA 12.0
 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
+#include "moe_permute_unpermute_kernel.h"
 
 // CubKeyValueSorter definition begin
 CubKeyValueSorter::CubKeyValueSorter()
@@ -229,3 +228,4 @@ void getMIndices(int64_t* expert_first_token_offset,
 }
 
 #endif
+#endif // USE_MACA
diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu
index 10be47966..6f378e8be 100644
--- a/csrc/moe/topk_softmax_kernels.cu
+++ b/csrc/moe/topk_softmax_kernels.cu
@@ -390,6 +390,129 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
     }
 }
 
+__device__ __forceinline__ float get_weight(const int64_t& v) {
+    const float* idx_and_weight = (const float*)&v;
+    return idx_and_weight[0];
+}
+
+template<uint32_t MASK=0xffffffff>
+__device__ __forceinline__ void warpSortDescending(float (&idx_and_weight)[2], int tid) {
+
+    int64_t val = *(int64_t*)idx_and_weight;
+    for (int width = 2; width <= WARP_SIZE; width <<=1) {
+        for (int step = width >> 1; step > 0; step >>=1) {
+            const bool is_not_final_phase = (width != WARP_SIZE);
+            const uint32_t bitmask = (tid & width);
+            const bool direction = is_not_final_phase & (bitmask == 0);
+            int64_t other_temp_val = __shfl_xor_sync(MASK, val, step);
+            int other_tid = tid ^ step;
+
+            float current_weight_bits = get_weight(val);
+            float other_weight_bits = get_weight(other_temp_val);
+            int current_index = val >> 32;
+            int other_index = other_temp_val >> 32;
+
+            bool weight_gt = other_weight_bits > current_weight_bits;
+            bool weight_eq = other_weight_bits == current_weight_bits;
+            bool index_lt = other_index < current_index;
+            bool cond = (tid < other_tid) ^ direction;
+
+            bool swap = (cond & (weight_gt | (weight_eq & index_lt))) |
+                        (!cond & ((other_weight_bits < current_weight_bits) | (weight_eq & (other_index > current_index))));
+
+            val = swap ? other_temp_val : val;
+        }
+    }
+    *(int64_t*)idx_and_weight = val;
+}
+
+template <typename IndType, int VPT = 4, int NUM_EXPERTS = 128, int WARPS_PER_CTA = 8, int BYTES_PER_LDG = 4, int TOPK = 8, int ROWS_PER_CTA = 2, int WARPS_PER_ROW = 4>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
+    void topkGatingSoftmaxDecode(const float* input, const bool* finished, float* output, const int num_rows, IndType* indices,
+        int* source_rows, const int k, const int start_expert, const int end_expert)
+{
+    static int THREADS_PER_ROW = WARPS_PER_ROW * WARP_SIZE;
+    const int thread_row = blockIdx.x * ROWS_PER_CTA + threadIdx.x / 128;
+    if (thread_row >= num_rows)
+    {
+        return;
+    }
+    const int warp_id = (threadIdx.x / WARP_SIZE);
+    const int tid_in_row = threadIdx.x % THREADS_PER_ROW;
+    const int row_in_block = threadIdx.x / THREADS_PER_ROW;
+    float idx_and_weight[2];
+    float idx_and_weight_2[2];
+    const bool row_is_active = finished ? !finished[thread_row] : true;
+    __shared__ float shared_experts[64][2];
+    __shared__ float max_val[8];
+    __shared__ float sum_val[8];
+    __shared__ float max_val_final[ROWS_PER_CTA];
+    __shared__ float rep[ROWS_PER_CTA];
+
+    // We finally start setting up the read pointers for each thread. First, each thread jumps to the start of the
+    // row it will read.
+    const float* thread_row_ptr = input + thread_row * NUM_EXPERTS;
+
+    float row_val = input[thread_row * NUM_EXPERTS + tid_in_row];
+    float thread_max = row_val;
+    float sum_data = expf(row_val);
+    for (int step = 1; step < WARP_SIZE; step <<= 1)
+    {
+        thread_max = max(thread_max, __shfl_xor_sync(0x0000000f, thread_max, step));
+        sum_data += __shfl_xor_sync(0x0000000f, sum_data, step);
+    }
+    if(tid_in_row % WARP_SIZE == 0) {
+        max_val[warp_id] = thread_max;
+        sum_val[warp_id] = sum_data;
+    }
+    __syncthreads();
+    if(tid_in_row < 4) {
+        float mx_v = max_val[tid_in_row + row_in_block * 4];
+        float sm_v = sum_val[tid_in_row + row_in_block * 4];
+        for(int step = 1; step < TOPK; step <<= 1 ) {
+            mx_v = max(mx_v, __shfl_xor_sync(0x0000000f, mx_v, step));
+            sm_v += __shfl_xor_sync(0x0000000f, sm_v, step);
+        }
+        if(tid_in_row == 0) {
+            max_val_final[row_in_block] = mx_v;
+            rep[row_in_block] = 1.0f / (sm_v * expf(-mx_v));
+        }
+    }
+    __syncthreads();
+    row_val = expf(row_val - max_val_final[row_in_block]) * rep[row_in_block];
+    
+    int64_t expert_id_opt = static_cast<int64_t>(tid_in_row);
+    idx_and_weight[0] = row_val;
+    idx_and_weight[1] = 0.0f;
+    *((int64_t*)idx_and_weight) |= (expert_id_opt << 32);
+    warpSortDescending<0xffffffff>(idx_and_weight, threadIdx.x);
+    if (tid_in_row % WARP_SIZE < TOPK) {
+        *((int64_t*)(shared_experts) + warp_id * TOPK + tid_in_row % WARP_SIZE) = *(int64_t*)idx_and_weight;
+    }
+    __syncthreads();
+    if(tid_in_row < WARP_SIZE){
+        *(int64_t*)idx_and_weight_2 = *((int64_t*)shared_experts + (warp_id / WARPS_PER_ROW) * WARP_SIZE + tid_in_row);
+    }
+    warpSortDescending<0xffffffff>(idx_and_weight_2, threadIdx.x);
+    if (tid_in_row < TOPK) {
+        int64_t res = *(int64_t*)idx_and_weight_2;
+        float max_val_ordered = get_weight(res);
+        int expert_id_ordered = (res >> 32);
+        max_val_ordered = max_val_ordered;//expf(max_val_ordered - thread_max_opt[threadIdx.y]) * reciprocal_row_sum;
+
+        // Add a guard to ignore experts not included by this node
+        const bool node_uses_expert = expert_id_ordered >= start_expert && expert_id_ordered < end_expert;
+        const bool should_process_row = row_is_active && node_uses_expert;
+
+        // The lead thread from each sub-group will write out the final results to global memory. (This will be a
+        // single) thread per row of the input/output matrices.
+        const int idx = TOPK * thread_row + tid_in_row;
+        output[idx] = max_val_ordered;
+        indices[idx] = should_process_row ? (expert_id_ordered - start_expert) : NUM_EXPERTS;
+        source_rows[idx] = tid_in_row * num_rows + thread_row;
+    }
+}
+
 namespace detail
 {
 // Constructs some constants needed to partition the work across threads at compile time.
@@ -423,12 +546,50 @@ void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, f
         input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert);
 }
 
+template <int EXPERTS, int WARPS_PER_TB, typename IndType>
+void topkDecodeGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, IndType* indices,
+    int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, cudaStream_t stream) {
+    if (k == 8 and num_rows < 1024){
+        static constexpr int MAX_BYTES_PER_LDG = 16;
+        static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(float) * EXPERTS);
+        using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG>;
+        static constexpr int VPT = Constants::VPT;
+        static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+
+        const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+        const int num_blocks = num_rows + 1 / 2;
+
+        dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
+        topkGatingSoftmaxDecode<<<num_blocks, 256, 0, stream>>>(
+            input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert);
+    } else {
+        static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+
+        static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(float) * EXPERTS);
+        using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG>;
+        static constexpr int VPT = Constants::VPT;
+        static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+        const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+        const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+
+        dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
+        topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG><<<num_blocks, block_dim, 0, stream>>>(
+            input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert);
+    }
+}
+
 #define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB)                       \
     topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB>(         \
         gating_output, nullptr, topk_weights, topk_indicies,            \
         token_expert_indices, num_tokens, topk, 0, num_experts,         \
         stream);
 
+#define LAUNCH_SOFTMAX_OPT(NUM_EXPERTS, WARPS_PER_TB)                   \
+    topkDecodeGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB>(  \
+        gating_output, nullptr, topk_weights, topk_indicies,            \
+        token_expert_indices, num_tokens, topk, 0, num_experts,         \
+        stream);
+
 template <typename IndType>
 void topkGatingSoftmaxKernelLauncher(
     const float* gating_output,
@@ -464,7 +625,7 @@ void topkGatingSoftmaxKernelLauncher(
             LAUNCH_SOFTMAX(64, WARPS_PER_TB);
             break;
         case 128:
-            LAUNCH_SOFTMAX(128, WARPS_PER_TB);
+            LAUNCH_SOFTMAX_OPT(128, WARPS_PER_TB);
             break;
         case 256:
             LAUNCH_SOFTMAX(256, WARPS_PER_TB);
diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp
index a74eb3720..ca37d3a0d 100644
--- a/csrc/moe/torch_bindings.cpp
+++ b/csrc/moe/torch_bindings.cpp
@@ -88,6 +88,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
   m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows);
 
 #endif
+
+#ifdef USE_MACA
+// Fused moe in mcblas
+  m.def(
+      "fused_moe_kernel(Tensor! A, Tensor! B, Tensor! C,"
+      "Tensor! topk_weights, Tensor! topk_ids,"
+      "Tensor! sorted_token_ids, Tensor! expert_ids,"
+      "Tensor! num_tokens_post_padded, bool mul_routed_weight, int top_k, int tileConfig) -> ()");
+  m.impl("fused_moe_kernel", torch::kCUDA, &fused_moe_kernel);
+#endif
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
diff --git a/csrc/ops.h b/csrc/ops.h
index f02f5083a..a202a046c 100644
--- a/csrc/ops.h
+++ b/csrc/ops.h
@@ -182,13 +182,16 @@ torch::Tensor aqlm_dequant(
 
 torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
                        torch::Tensor _scaling_factors, torch::Tensor _zeros,
-                       int64_t split_k_iters);
+                       int64_t split_k_iters, torch::Tensor _temp_space, 
+                       bool dtype_bf16);
 
 torch::Tensor awq_dequantize(torch::Tensor _kernel,
                              torch::Tensor _scaling_factors,
                              torch::Tensor _zeros, int64_t split_k_iters,
                              int64_t thx, int64_t thy);
 
+torch::Tensor awq_to_gptq_4bit(torch::Tensor qweight);
+
 torch::Tensor permute_cols(torch::Tensor const& A, torch::Tensor const& perm);
 #endif
 
@@ -296,11 +299,21 @@ void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
 void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                                torch::Tensor& scales,
                                std::optional<torch::Tensor> const& azp);
+                               
+void fused_silu_mul_dq_mask_quant_pack(
+    torch::Tensor& out,          
+    torch::Tensor const& input, 
+    torch::Tensor const &mask);
+
+void dynamic_scaled_int8_mask_quant( torch::Tensor& out,  torch::Tensor const& input, torch::Tensor const &mask, 
+    torch::Tensor& scales, c10::optional<torch::Tensor> const& azp);
 
 torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
                         torch::Tensor b_gptq_qzeros,
                         torch::Tensor b_gptq_scales, torch::Tensor b_g_idx,
-                        bool use_exllama, int64_t bit);
+                        bool use_exllama, int64_t bit, int64_t group_size, 
+                        torch::Tensor perm_space, torch::Tensor temp_space,
+			            bool dtype_bf16);
 
 void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit);
 
diff --git a/csrc/quantization/awq/Hgemm_nn_128x32x128_8m1n8k_awq.hpp b/csrc/quantization/awq/Hgemm_nn_128x32x128_8m1n8k_awq.hpp
new file mode 100644
index 000000000..614a6251e
--- /dev/null
+++ b/csrc/quantization/awq/Hgemm_nn_128x32x128_8m1n8k_awq.hpp
@@ -0,0 +1,473 @@
+// 2024 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
+#pragma once
+#include "../gptq/Hgemm_common.cuh"
+#include "awq_4bits.cuh"
+#include "maca_fp16.h"
+
+#define input_type __half
+#define output_type __half
+#define scalar_type float
+#define acc_type float
+
+#define SEL0 0x01000504
+#define SEL1 0x03020706
+
+#define SWIZZLE_INDEX(index, phase) (index ^ phase)
+#define ONE_DIM_INDEX(x, y, lda) (y * lda + x)
+#define CLAMP(value, bound, align) (value < bound ? value : bound - align)
+
+#define LDG_B                                                       \
+    {                                                               \
+        ldg_b128_bsm_async(b_sts, dB[0], pred_n0, true);            \
+        ldg_b128_bsm_async(b_sts + 128 * 16, dB[1], pred_n1, true); \
+    }
+
+#define LDG_B_HEAD                                                            \
+    {                                                                         \
+        bool pred_k = rowB_swizzle < ktail;                                   \
+        ldg_b128_bsm_async(b_sts, dB[0], pred_n0 &&pred_k, true);             \
+        ldg_b128_bsm_async(b_sts + 128 * 16, dB[1], pred_n1 && pred_k, true); \
+    }
+
+#define MMA_ABC1                                                             \
+    {                                                                        \
+        for (int index_n = 0; index_n < 2; ++index_n)                        \
+        {                                                                    \
+            for (int index_m = 0; index_m < 8; ++index_m)                    \
+            {                                                                \
+                mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 0, 8)], \
+                                cast_b64(rgB)[ONE_DIM_INDEX(0, index_n, 2)], \
+                                rgC[ONE_DIM_INDEX(index_m, index_n, 8)]);    \
+            }                                                                \
+        }                                                                    \
+    }
+
+#define MMA_ABC2                                                             \
+    {                                                                        \
+        for (int index_n = 0; index_n < 2; ++index_n)                        \
+        {                                                                    \
+            for (int index_m = 0; index_m < 8; ++index_m)                    \
+            {                                                                \
+                mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 1, 8)], \
+                                cast_b64(rgB)[ONE_DIM_INDEX(1, index_n, 2)], \
+                                rgC[ONE_DIM_INDEX(index_m, index_n, 8)]);    \
+            }                                                                \
+        }                                                                    \
+    }
+
+#define LDS_B                                          \
+    {                                                  \
+        cast_b128(rgB)[0] = cast_b128(b_lds)[0];       \
+        cast_b128(rgB)[1] = cast_b128(b_lds)[16 * 16]; \
+    }
+
+#define PERM_C(C2perm)                             \
+    {                                              \
+        Float4VecType C_tmp[8];                    \
+        float *ptri = (float *)C2perm;             \
+        float *ptro = (float *)C_tmp;              \
+        for (int j = 0; j < 4; ++j)                \
+        {                                          \
+            for (int i = 0; i < 8; ++i)            \
+            {                                      \
+                ptro[j * 8 + i] = ptri[j + i * 4]; \
+            }                                      \
+        }                                          \
+        for (int i = 0; i < 8; ++i)                \
+        {                                          \
+            C2perm[i] = C_tmp[i];                  \
+        }                                          \
+    }
+
+#define STS_C(phase)                                            \
+    {                                                           \
+        cast_b128(c_sts[0])[0] = cast_b128(rgC)[0 + phase * 8]; \
+        cast_b128(c_sts[1])[0] = cast_b128(rgC)[1 + phase * 8]; \
+        cast_b128(c_sts[2])[0] = cast_b128(rgC)[2 + phase * 8]; \
+        cast_b128(c_sts[3])[0] = cast_b128(rgC)[3 + phase * 8]; \
+        cast_b128(c_sts[4])[0] = cast_b128(rgC)[4 + phase * 8]; \
+        cast_b128(c_sts[5])[0] = cast_b128(rgC)[5 + phase * 8]; \
+        cast_b128(c_sts[6])[0] = cast_b128(rgC)[6 + phase * 8]; \
+        cast_b128(c_sts[7])[0] = cast_b128(rgC)[7 + phase * 8]; \
+    }
+
+#define REDUCE_C(phase)                                                   \
+    {                                                                     \
+        cast_b128(rgC)[0 + phase * 8] = cast_b128(c_lds[0])[0];           \
+        cast_b128(rgC)[1 + phase * 8] = cast_b128(c_lds[1])[0];           \
+        cast_b128(rgC)[2 + phase * 8] = cast_b128(c_lds[0])[1 * 16 * 32]; \
+        cast_b128(rgC)[3 + phase * 8] = cast_b128(c_lds[1])[1 * 16 * 32]; \
+        cast_b128(rgC)[4 + phase * 8] = cast_b128(c_lds[0])[2 * 16 * 32]; \
+        cast_b128(rgC)[5 + phase * 8] = cast_b128(c_lds[1])[2 * 16 * 32]; \
+        cast_b128(rgC)[6 + phase * 8] = cast_b128(c_lds[0])[3 * 16 * 32]; \
+        cast_b128(rgC)[7 + phase * 8] = cast_b128(c_lds[1])[3 * 16 * 32]; \
+        float *reduc_c = reinterpret_cast<float *>(rgC);                  \
+        for (int loop = 0; loop < 8; ++loop)                              \
+        {                                                                 \
+            float acc = 0;                                                \
+            acc += reduc_c[loop + phase * 32 + 0 * 8];                    \
+            acc += reduc_c[loop + phase * 32 + 1 * 8];                    \
+            acc += reduc_c[loop + phase * 32 + 2 * 8];                    \
+            acc += reduc_c[loop + phase * 32 + 3 * 8];                    \
+            reduc_c[loop + phase * 8] = acc;                              \
+        }                                                                 \
+    }
+
+template <Operation_t TransA,
+          Operation_t TransB,
+          bool IsBetaZero,
+          int tileM,
+          int tileN,
+          int tileK,
+          bool splitk = false,
+          bool Swap = false>
+__global__ __launch_bounds__(256) void Hgemm_nn_128x32x128_8m1n8k_awq_4bit(int m,
+                                                                           int n,
+                                                                           int k,
+                                                                           const scalar_type alpha,
+                                                                           const scalar_type beta,
+                                                                           const quant_packed_type *dA_input,
+                                                                           int lda,
+                                                                           const input_type *dB_input,
+                                                                           int ldb,
+                                                                           output_type *dC_input,
+                                                                           output_type *dC_output,
+                                                                           int ldc,
+                                                                           quant_packed_type *d_zeros,
+                                                                           input_type *d_scales,
+                                                                           int splitk_iters = 1,
+                                                                           acc_type * d_acc_tmp=nullptr)
+{
+    __shared__ uint8_t smem_base[0x8000];
+    int bidx = Swap ? blockIdx.y : blockIdx.x;
+    int bidy = Swap ? blockIdx.x : blockIdx.y;
+    int bidz = blockIdx.z;
+    int tid = threadIdx.x;
+
+    uint64_t arowstride = TransA == OP_N ? 1 : lda;
+    uint64_t acolstride = TransA == OP_N ? lda : 1;
+    uint64_t browstride = TransB == OP_N ? 1 : ldb;
+    uint64_t bcolstride = TransB == OP_N ? ldb : 1;
+
+    const int malign = 8;
+    const int nalign = 1;
+    const int kalign = 8;
+    int align_m = (m + malign - 1) / malign * malign;
+    int align_n = (n + nalign - 1) / nalign * nalign;
+    int align_k = (k + kalign - 1) / kalign * kalign;
+
+    int k_num = (align_k + splitk_iters - 1) / splitk_iters;
+    k_num = (k_num + tileK - 1) / tileK * tileK; // k_num should be aligned to tileK
+    int k_begin = bidz * k_num;
+    int k_end = (bidz + 1) * k_num;
+    k_end = k_end > align_k ? align_k : k_end;
+    k_num = k_end - k_begin;
+    int ktail = k_num % tileK;
+    int kloop = k_begin;
+    if (k_begin > align_k)
+    {
+        return;
+    }    
+
+
+    int slot = tid / 64;
+    int lane = tid & 63;
+    int m64d16 = lane / 16;
+    int m64m16 = lane % 16;
+    output_type *c_base_i = dC_input;
+    output_type *c_base_o = dC_output;
+
+
+    quant_packed_type rga[8];
+    b128VecType rgb[2];
+
+    b64VecType rgA[16], rgB[4], rgCi[4];
+    Float4VecType rgC[16];
+    uint32_t rgZeros[1];
+    b128VecType rgScales[1];
+
+    quant_packed_type *dA[1];
+    input_type *dB[2];
+
+    // ldg A/B head
+
+    int rowA = m64m16 * 8;
+    int colA = m64d16 * 8 + slot * 32;
+    uint current_m = bidx * tileM + rowA;
+    bool pred_m = current_m < align_m;
+    dA[0] = (quant_packed_type *)dA_input + (uint64_t)(current_m / PACK_RATIO) * (uint64_t)(arowstride) +
+            (uint64_t)(colA + k_begin) * (uint64_t)(acolstride / PACK_RATIO);
+
+    int colB = m64d16 + slot * 4;
+    int rowB = m64m16 * 8;
+    int rowB_swizzle = SWIZZLE_INDEX(rowB, (colB % 4 * 8));
+    int current_n = bidy * tileN + colB;
+    bool pred_n0 = current_n < align_n;
+    dB[0] = (input_type *)dB_input + (uint64_t)(rowB_swizzle + k_begin) * (uint64_t)(browstride) +
+            (uint64_t)(current_n) * (uint64_t)bcolstride;
+
+    current_n += 16;
+    bool pred_n1 = current_n < align_n;
+    dB[1] = (input_type *)dB_input + (uint64_t)(rowB_swizzle + k_begin) * (uint64_t)(browstride) +
+            (uint64_t)(current_n) * (uint64_t)bcolstride;
+
+    // ldgB to BSM need swizzle
+    input_type *b_sts = (input_type *)smem_base;
+    b_sts += ONE_DIM_INDEX(rowB, colB, 128);
+
+    // lds need swizzle
+    input_type *b_lds = (input_type *)smem_base;
+    int colB_lds = m64m16;
+    int rowB_lds = m64d16 * 8 + slot * 32;
+    b_lds += ONE_DIM_INDEX(SWIZZLE_INDEX(rowB_lds, (colB_lds % 4 * 8)), colB_lds, 128);
+
+    // ldg zeros and scales
+    quant_packed_type *ldg_zeros_offset = d_zeros + (lda / PACK_RATIO) * (kloop / tileK) + bidx * (tileM / PACK_RATIO) + tid;
+    input_type *ldg_scales_offset = d_scales + lda * (kloop / tileK) + bidx * tileM + tid * 2;
+    quant_packed_type *sts_zeros_offset = (quant_packed_type *)(smem_base + tileK * tileN * sizeof(input_type)) + tid;
+    input_type *sts_scales_offset = (input_type *)(smem_base + tileK * tileN * sizeof(input_type) + tileM / PACK_RATIO * sizeof(quant_packed_type)) + tid * 2;
+    quant_packed_type *lds_zeros_offset = (quant_packed_type *)(smem_base + tileK * tileN * sizeof(input_type)) + m64m16;
+    input_type *lds_scales_offset = (input_type *)(smem_base + tileK * tileN * sizeof(input_type) + tileM / PACK_RATIO * sizeof(quant_packed_type)) + m64m16 * 8;
+
+    // clear C registers
+    rgC[0] = 0;
+    rgC[1] = 0;
+    rgC[2] = 0;
+    rgC[3] = 0;
+    rgC[4] = 0;
+    rgC[5] = 0;
+    rgC[6] = 0;
+    rgC[7] = 0;
+    rgC[8] = 0;
+    rgC[9] = 0;
+    rgC[10] = 0;
+    rgC[11] = 0;
+    rgC[12] = 0;
+    rgC[13] = 0;
+    rgC[14] = 0;
+    rgC[15] = 0;
+
+    int aincr = tileK * acolstride / PACK_RATIO;
+    int bincr = tileK * browstride;
+
+    if (k_num > ktail)
+    {
+        LDG_B;
+        LDG_ZEROS;
+        LDG_SCALES;
+        LDG_A;
+
+        dA[0] += aincr;
+        dB[0] += bincr;
+        dB[1] += bincr;
+
+        // lds_A, need swizzle
+        arrive_gvmcnt(8);
+        barrier();
+        LDS_B;
+        LDS_ZEROS;
+        LDS_SCALES;
+        arrive_bsmcnt(0);
+        barrier();
+
+        ldg_zeros_offset += lda / PACK_RATIO;
+        ldg_scales_offset += lda;
+
+// KLOOP
+#pragma unroll 1
+        for (kloop = tileK; kloop + tileK <= k_num; kloop += tileK)
+        {
+            // load next B
+            LDG_B;
+            LDG_ZEROS;
+            LDG_SCALES;
+
+            // MMA_ABC1;
+            arrive_gvmcnt(4);
+            PERM_A;
+            LDG_A;
+            MMA;
+
+            // sts && lds
+            arrive_gvmcnt(8);
+            barrier();
+            LDS_B;
+            LDS_ZEROS;
+            LDS_SCALES;
+            arrive_bsmcnt(0);
+            barrier();
+
+            dA[0] += aincr;
+            dB[0] += bincr;
+            dB[1] += bincr;
+            ldg_zeros_offset += lda / PACK_RATIO;
+            ldg_scales_offset += lda;
+        }
+
+        arrive_gvmcnt(0);
+        PERM_A;
+        MMA;
+    }
+
+    // final tail kloop
+    if (ktail > 0)
+    {
+        LDG_B_HEAD;
+        LDG_ZEROS;
+        LDG_SCALES;
+        LDG_A_HEAD;
+        arrive_gvmcnt(8);
+        barrier();
+        LDS_B;
+        LDS_ZEROS;
+        LDS_SCALES;
+        arrive_bsmcnt(0);
+        barrier();
+        arrive_gvmcnt(0);
+        PERM_A;
+        MMA;
+    }
+
+    // store C registers into BSM & do reduction
+    int colC = m64m16 + slot * 16;
+    int rowC = m64d16 * 32;
+    scalar_type *c_sts[8];
+    c_sts[0] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC, (colC % 4 * 4)), colC, 128);
+    c_sts[1] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 4, (colC % 4 * 4)), colC, 128);
+    c_sts[2] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 8, (colC % 4 * 4)), colC, 128);
+    c_sts[3] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 12, (colC % 4 * 4)), colC, 128);
+    c_sts[4] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 16, (colC % 4 * 4)), colC, 128);
+    c_sts[5] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 20, (colC % 4 * 4)), colC, 128);
+    c_sts[6] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 24, (colC % 4 * 4)), colC, 128);
+    c_sts[7] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 28, (colC % 4 * 4)), colC, 128);
+
+    colC = m64d16 + slot * 4;
+    rowC = m64m16 * 4;
+    scalar_type *c_lds[2];
+    c_lds[0] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC, (colC % 4 * 4)), colC, 128);
+    c_lds[1] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 64, (colC % 4 * 4)), colC, 128);
+
+    PERM_C(rgC);
+    STS_C(0);
+    arrive_bsmcnt(0);
+    barrier();
+    REDUCE_C(0);
+
+    arrive_bsmcnt(0);
+    barrier();
+
+    PERM_C((rgC + 8))
+    STS_C(1);
+    arrive_bsmcnt(0);
+    barrier();
+    REDUCE_C(1);
+
+    arrive_bsmcnt(0);
+
+    // read C_input if beta!=0
+    colC = m64d16 + slot * 4 + bidy * tileN;
+    rowC = m64m16 * 4 + bidx * tileM;
+    int colC1 = colC + 16;
+    int rowC1 = rowC + 64;
+
+
+    if constexpr (!IsBetaZero && !splitk)
+    {
+        ldg_b64_reg_async(cast_b64(rgCi)[0],
+                          c_base_i + (uint64_t)rowC + (uint64_t)(colC) * (uint64_t)(ldc),
+                          rowC < align_m && colC < align_n, true);
+        ldg_b64_reg_async(cast_b64(rgCi)[1], c_base_i + rowC1 + (uint64_t)(colC) * (uint64_t)(ldc),
+                          rowC1 < align_m && colC < align_n, true);
+        ldg_b64_reg_async(cast_b64(rgCi)[2], c_base_i + rowC + (uint64_t)(colC1) * (uint64_t)(ldc),
+                          rowC < align_m && colC1 < align_n, true);
+        ldg_b64_reg_async(cast_b64(rgCi)[3], c_base_i + rowC1 + (uint64_t)(colC1) * (uint64_t)(ldc),
+                          rowC1 < align_m && colC1 < align_n, true);
+        arrive_gvmcnt(0);
+    }
+
+    if constexpr (IsBetaZero && !splitk)
+    {
+        for (int i = 0; i < 4; ++i) rgCi[i] = static_cast<b64VecType>(0);
+    }
+
+
+    if constexpr (!splitk)
+    {
+        output_type *ptrCi = reinterpret_cast<output_type *>(rgCi);
+        b64VecType *ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC) + (uint64_t)(colC) * (uint64_t)(ldc));
+        if (rowC < align_m && colC < align_n)
+        {
+            output_type result[4];
+            result[0] = static_cast<output_type>(alpha * rgC[0][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            result[1] = static_cast<output_type>(alpha * rgC[0][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            result[2] = static_cast<output_type>(alpha * rgC[0][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            result[3] = static_cast<output_type>(alpha * rgC[0][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC1) + (uint64_t)(colC) * (uint64_t)(ldc));
+        if (rowC1 < align_m && colC < align_n)
+        {
+            output_type result[4];
+            result[0] = static_cast<output_type>(alpha * rgC[1][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            result[1] = static_cast<output_type>(alpha * rgC[1][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            result[2] = static_cast<output_type>(alpha * rgC[1][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            result[3] = static_cast<output_type>(alpha * rgC[1][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC) + (uint64_t)(colC1) * (uint64_t)(ldc));
+        if (rowC < align_m && colC1 < align_n)
+        {
+            output_type result[4];
+            result[0] = static_cast<output_type>(alpha * rgC[2][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            result[1] = static_cast<output_type>(alpha * rgC[2][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            result[2] = static_cast<output_type>(alpha * rgC[2][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            result[3] = static_cast<output_type>(alpha * rgC[2][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC1) + (uint64_t)(colC1) * (uint64_t)(ldc));
+        if (rowC1 < align_m && colC1 < align_n)
+        {
+            output_type result[4];
+            result[0] = static_cast<output_type>(alpha * rgC[3][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            result[1] = static_cast<output_type>(alpha * rgC[3][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            result[2] = static_cast<output_type>(alpha * rgC[3][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            result[3] = static_cast<output_type>(alpha * rgC[3][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            ptrO[0] = cast_b64(result)[0];
+        }
+    }
+    else
+    {
+        acc_type *ptrAcc = d_acc_tmp + (uint64_t)(rowC) + (uint64_t)(colC) * (uint64_t)(ldc);
+        if (rowC < align_m && colC < align_n)
+        {
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC1) + (uint64_t)(colC) * (uint64_t)(ldc);
+        if (rowC1 < align_m && colC < align_n)
+        {
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC) + (uint64_t)(colC1) * (uint64_t)(ldc);
+        if (rowC < align_m && colC1 < align_n)
+        {
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC1) + (uint64_t)(colC1) * (uint64_t)(ldc);
+        if (rowC1 < align_m && colC1 < align_n)
+        {
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][3]));
+        }
+    }
+}
\ No newline at end of file
diff --git a/csrc/quantization/awq/awq_4bits.cuh b/csrc/quantization/awq/awq_4bits.cuh
new file mode 100644
index 000000000..b5012b9b0
--- /dev/null
+++ b/csrc/quantization/awq/awq_4bits.cuh
@@ -0,0 +1,125 @@
+#pragma once
+
+#include "../gptq/Hgemm_common.cuh"
+#include "dequant.cuh"
+#define quant_packed_type uint32_t
+
+#define QBITS 4
+#define PACK_RATIO (32 / QBITS)
+
+#define LDG_A                                                                         \
+    {                                                                                 \
+        ldg_b32_reg_async(rga[0], dA[0] + acolstride * 0 / PACK_RATIO, pred_m, true); \
+        ldg_b32_reg_async(rga[1], dA[0] + acolstride * 1 / PACK_RATIO, pred_m, true); \
+        ldg_b32_reg_async(rga[2], dA[0] + acolstride * 2 / PACK_RATIO, pred_m, true); \
+        ldg_b32_reg_async(rga[3], dA[0] + acolstride * 3 / PACK_RATIO, pred_m, true); \
+        ldg_b32_reg_async(rga[4], dA[0] + acolstride * 4 / PACK_RATIO, pred_m, true); \
+        ldg_b32_reg_async(rga[5], dA[0] + acolstride * 5 / PACK_RATIO, pred_m, true); \
+        ldg_b32_reg_async(rga[6], dA[0] + acolstride * 6 / PACK_RATIO, pred_m, true); \
+        ldg_b32_reg_async(rga[7], dA[0] + acolstride * 7 / PACK_RATIO, pred_m, true); \
+    }
+
+#define LDG_A_HEAD                                                                    \
+    {                                                                                 \
+        bool predk = colA < ktail;                                                    \
+        ldg_b32_reg_async(rga[0], dA[0] + acolstride * 0 / PACK_RATIO, pred_m && predk, true); \
+        ldg_b32_reg_async(rga[1], dA[0] + acolstride * 1 / PACK_RATIO, pred_m && predk, true); \
+        ldg_b32_reg_async(rga[2], dA[0] + acolstride * 2 / PACK_RATIO, pred_m && predk, true); \
+        ldg_b32_reg_async(rga[3], dA[0] + acolstride * 3 / PACK_RATIO, pred_m && predk, true); \
+        ldg_b32_reg_async(rga[4], dA[0] + acolstride * 4 / PACK_RATIO, pred_m && predk, true); \
+        ldg_b32_reg_async(rga[5], dA[0] + acolstride * 5 / PACK_RATIO, pred_m && predk, true); \
+        ldg_b32_reg_async(rga[6], dA[0] + acolstride * 6 / PACK_RATIO, pred_m && predk, true); \
+        ldg_b32_reg_async(rga[7], dA[0] + acolstride * 7 / PACK_RATIO, pred_m && predk, true); \
+    }
+
+// 每个block中用tileM/PACK_RATIO个线程去加载tileM个zeros，zeros为int4类型，每个32bits寄存器可以保存8个zeros
+#define LDG_ZEROS                                                                                                                              \
+    {                                                                                                                                          \
+        ldg_b32_bsm_async(sts_zeros_offset, ldg_zeros_offset, tid < tileM / PACK_RATIO && (bidx * tileM + tid * PACK_RATIO < align_m), false); \
+    }
+
+// 每个block中用tileM/(sizeof(uint32_t)/sizeof(__half))个线程去加载tileM个scales，scale为fp16类型，每个32bits寄存器可以保存2个scales
+#define LDG_SCALES                                                                                                      \
+    {                                                                                                                   \
+        ldg_b32_bsm_async(sts_scales_offset, ldg_scales_offset, tid < 64 && (bidx * tileM + tid * 2 < align_m), false); \
+    }
+
+#define LDS_ZEROS                         \
+    {                                     \
+        rgZeros[0] = lds_zeros_offset[0]; \
+    }
+
+#define LDS_SCALES                                     \
+    {                                                  \
+        rgScales[0] = cast_b128(lds_scales_offset)[0]; \
+    }
+
+// zeros中从低字节到高字节分别对应第 0 2 4 6 1 3 5 7 列， 所以index_shfl = index % 2 * 4 + index / 2
+//    提取出int4的zeros后转成int32，再把weight（矩阵A）还原成fp16，Weight_Q=Scale*(Weight_4Bit-ZeroPoint)
+#define PERM_ELEM(index)                                                                                                          \
+    {                                                                                                                             \
+        __half_raw elem;                                                                                                          \
+        if constexpr (index & 0x1)                                                                                                \
+        {                                                                                                                         \
+            elem.x = cast_b32(rgScales)[index / 2] >> 16;                                                                         \
+        }                                                                                                                         \
+        else                                                                                                                      \
+        {                                                                                                                         \
+            elem.x = cast_b32(rgScales)[index / 2] & 0xffff;                                                                      \
+        }                                                                                                                         \
+        __half scale = __half(elem);                                                                                              \
+        constexpr int index_shfl = index % 2 * 4 + index / 2;                                                                     \
+        uint32_t zero = __builtin_mxc_ubfe(rgZeros[0], QBITS * index_shfl, QBITS);                                                \
+        cast_half(rgA)[index * 4 + 0] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[0], QBITS * index_shfl, QBITS) -  \
+                                                                         zero) , scale);                                                                                    \
+        cast_half(rgA)[index * 4 + 1] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[1], QBITS * index_shfl, QBITS) -  \
+                                                                         zero) , scale);                                                                                    \
+        cast_half(rgA)[index * 4 + 2] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[2], QBITS * index_shfl, QBITS) -  \
+                                                                         zero) , scale);                                                                                    \
+        cast_half(rgA)[index * 4 + 3] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[3], QBITS * index_shfl, QBITS) -  \
+                                                                         zero) , scale);                                                                                    \
+        cast_half(rgA)[index * 4 + 32] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[4], QBITS * index_shfl, QBITS) - \
+                                                                          zero) , scale);                                                                                   \
+        cast_half(rgA)[index * 4 + 33] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[5], QBITS * index_shfl, QBITS) - \
+                                                                          zero) , scale);                                                                                   \
+        cast_half(rgA)[index * 4 + 34] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[6], QBITS * index_shfl, QBITS) - \
+                                                                          zero) , scale);                                                                                   \
+        cast_half(rgA)[index * 4 + 35] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[7], QBITS * index_shfl, QBITS) - \
+                                                                          zero) , scale);                                                                                   \
+    }
+
+#define PERM_A       \
+    {                \
+        PERM_ELEM(0) \
+        PERM_ELEM(1) \
+        PERM_ELEM(2) \
+        PERM_ELEM(3) \
+        PERM_ELEM(4) \
+        PERM_ELEM(5) \
+        PERM_ELEM(6) \
+        PERM_ELEM(7) \
+    }
+
+#define MMA_ELEM(index_m)                                        \
+    mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 0, 8)], \
+                    cast_b64(rgB)[ONE_DIM_INDEX(0, 0, 2)],       \
+                    rgC[ONE_DIM_INDEX(index_m, 0, 8)]);          \
+    mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 0, 8)], \
+                    cast_b64(rgB)[ONE_DIM_INDEX(0, 1, 2)],       \
+                    rgC[ONE_DIM_INDEX(index_m, 1, 8)]);          \
+    mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 1, 8)], \
+                    cast_b64(rgB)[ONE_DIM_INDEX(1, 0, 2)],       \
+                    rgC[ONE_DIM_INDEX(index_m, 0, 8)]);          \
+    mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 1, 8)], \
+                    cast_b64(rgB)[ONE_DIM_INDEX(1, 1, 2)],       \
+                    rgC[ONE_DIM_INDEX(index_m, 1, 8)]);
+
+#define MMA     \
+    MMA_ELEM(0) \
+    MMA_ELEM(1) \
+    MMA_ELEM(2) \
+    MMA_ELEM(3) \
+    MMA_ELEM(4) \
+    MMA_ELEM(5) \
+    MMA_ELEM(6) \
+    MMA_ELEM(7)
diff --git a/csrc/quantization/awq/dequant.cuh b/csrc/quantization/awq/dequant.cuh
new file mode 100644
index 000000000..7d6db705c
--- /dev/null
+++ b/csrc/quantization/awq/dequant.cuh
@@ -0,0 +1,10 @@
+#pragma once
+
+#define BFE(source, bit_start, num_bits) __builtin_mxc_ubfe(source, bit_start, num_bits);
+
+template <typename outputT, typename inputT, int qbits>
+__device__ __forceinline__ outputT extract(inputT *input, size_t loc)
+{
+    constexpr int widthPerElem = sizeof(inputT) * 8 / qbits;
+    return (outputT)BFE(input[loc / widthPerElem], loc % widthPerElem, qbits);
+}
\ No newline at end of file
diff --git a/csrc/quantization/awq/dequantize.cuh b/csrc/quantization/awq/dequantize.cuh
index 5fa4b5f64..0f245d470 100644
--- a/csrc/quantization/awq/dequantize.cuh
+++ b/csrc/quantization/awq/dequantize.cuh
@@ -14,6 +14,7 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
 namespace vllm {
 namespace awq {
 
+template<typename VT>
 __device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
   assert(false);
@@ -39,6 +40,7 @@ __device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) {
   // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW
   // dependency if we issue immediately before required.
   const uint32_t top_i4s = i4s >> 8;
+#ifndef USE_MACA
   // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400
   asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
                : "=r"(h[0])
@@ -59,6 +61,63 @@ __device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) {
                : "=r"(h[3])
                : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM),
                  "n"(immLut));
+#else
+      // >>>> PTX2CPP Success <<<<
+{
+(h[0])=0;
+if((immLut)&0x01)(h[0])|=~(i4s)&~(BOTTOM_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x02)(h[0])|=~(i4s)&~(BOTTOM_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x04)(h[0])|=~(i4s)& (BOTTOM_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x08)(h[0])|=~(i4s)& (BOTTOM_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x10)(h[0])|= (i4s)&~(BOTTOM_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x20)(h[0])|= (i4s)&~(BOTTOM_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x40)(h[0])|= (i4s)& (BOTTOM_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x80)(h[0])|= (i4s)& (BOTTOM_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+}
+
+    // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400
+
+// >>>> PTX2CPP Success <<<<
+{
+(h[1])=0;
+if((immLut)&0x01)(h[1])|=~(i4s)&~(TOP_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x02)(h[1])|=~(i4s)&~(TOP_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x04)(h[1])|=~(i4s)& (TOP_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x08)(h[1])|=~(i4s)& (TOP_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x10)(h[1])|= (i4s)&~(TOP_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x20)(h[1])|= (i4s)&~(TOP_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x40)(h[1])|= (i4s)& (TOP_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x80)(h[1])|= (i4s)& (TOP_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+}
+// >>>> PTX2CPP Success <<<<
+{
+(h[2])=0;
+if((immLut)&0x01)(h[2])|=~(top_i4s)&~(BOTTOM_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x02)(h[2])|=~(top_i4s)&~(BOTTOM_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x04)(h[2])|=~(top_i4s)& (BOTTOM_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x08)(h[2])|=~(top_i4s)& (BOTTOM_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x10)(h[2])|= (top_i4s)&~(BOTTOM_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x20)(h[2])|= (top_i4s)&~(BOTTOM_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x40)(h[2])|= (top_i4s)& (BOTTOM_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x80)(h[2])|= (top_i4s)& (BOTTOM_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+}
+    // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400
+
+// >>>> PTX2CPP Success <<<<
+{
+(h[3])=0;
+if((immLut)&0x01)(h[3])|=~(top_i4s)&~(TOP_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x02)(h[3])|=~(top_i4s)&~(TOP_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x04)(h[3])|=~(top_i4s)& (TOP_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x08)(h[3])|=~(top_i4s)& (TOP_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x10)(h[3])|= (top_i4s)&~(TOP_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x20)(h[3])|= (top_i4s)&~(TOP_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x40)(h[3])|= (top_i4s)& (TOP_MASK)&~(I4s_TO_F16s_MAGIC_NUM);
+if((immLut)&0x80)(h[3])|= (top_i4s)& (TOP_MASK)& (I4s_TO_F16s_MAGIC_NUM);
+}
+
+
+#endif //USE_MACA
 
   // I use inline PTX below because I am not sure if the compiler will emit
   // float2half instructions if I use the half2 ctor. In this case, I chose
@@ -77,6 +136,7 @@ __device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) {
 
   // Finally, we construct the output numbers.
   // Convert elt_01
+#ifndef USE_MACA
   asm volatile("sub.f16x2 %0, %1, %2;\n"
                : "=r"(h[0])
                : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM));
@@ -92,7 +152,54 @@ __device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) {
   asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
                : "=r"(h[3])
                : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64));
+#else
+    // >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(h[0]);
+unsigned int __b=(FP16_TOP_MAGIC_NUM);
+VT __d=__hsub2(*(VT*)&__a,*(VT*)&__b);
+(h[0])=*(unsigned int*)&__d;
+}
+}
+
+    // Convert elt_23
+
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(h[1]);
+unsigned int __b=(ONE_SIXTEENTH);
+unsigned int __c=(NEG_64);
+VT __d=__hfma2(*(VT*)&__a,*(VT*)&__b,*(VT*)&__c);
+(h[1])=*(unsigned int*)&__d;
+}
+}
+    // Convert elt_45
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(h[2]);
+unsigned int __b=(FP16_TOP_MAGIC_NUM);
+VT __d=__hsub2(*(VT*)&__a,*(VT*)&__b);
+(h[2])=*(unsigned int*)&__d;
+}
+}
+
+    // Convert elt_67
+
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(h[3]);
+unsigned int __b=(ONE_SIXTEENTH);
+unsigned int __c=(NEG_64);
+VT __d=__hfma2(*(VT*)&__a,*(VT*)&__b,*(VT*)&__c);
+(h[3])=*(unsigned int*)&__d;
+}
+}
 
+#endif // USE_MACA
   return result;
 #endif
   __builtin_unreachable();  // Suppress missing return statement warning
diff --git a/csrc/quantization/awq/gemm_kernels.cu b/csrc/quantization/awq/gemm_kernels.cu
index 53c47679c..81e79e938 100644
--- a/csrc/quantization/awq/gemm_kernels.cu
+++ b/csrc/quantization/awq/gemm_kernels.cu
@@ -1,3 +1,4 @@
+// 2025 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
 /*
 Adapted from https://github.com/mit-han-lab/llm-awq
 @article{lin2023awq,
@@ -14,9 +15,56 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
 
 #include <cuda_fp16.h>
 
+#include "../gptq/hgemm_gptq.h"
+#include "../gptq/scalar_type.hpp"
+
+//#include "hgemv_nn_splitk_awq.hpp"
+//#include "hgemv_selector.hpp"
+//#include "Hgemm_nn_128x32x128_8m1n8k_awq.hpp"
+
 namespace vllm {
 namespace awq {
+#define input_type __half
+#define output_type __half
+#define quant_packed_type uint32_t
+#define QUANT_GROUP 128
+
+struct DivModFast {
+    DivModFast(int d = 1)
+    {
+        d_ = (d == 0) ? 1 : d;
+        for (l_ = 0;; ++l_) {
+            if ((1U << l_) >= d_)
+                break;
+        }
+        uint64_t one = 1;
+        uint64_t m   = ((one << 32) * ((one << l_) - d_)) / d_ + 1;
+        m_           = static_cast<uint32_t>(m);
+    }
+
+    __device__ __inline__ int div(int idx) const
+    {
+        uint32_t tm = __umulhi(m_, idx); // get high 32-bit of the product
+        return (tm + idx) >> l_;
+    }
+
+    __device__ __inline__ int mod(int idx) const
+    {
+        return idx - d_ * div(idx);
+    }
+
+    __device__ __inline__ void divmod(int idx, int &quo, int &rem)
+    {
+        quo = div(idx);
+        rem = idx - quo * d_;
+    }
+    
+    uint32_t d_; // divisor
+    uint32_t l_; // ceil(log2(d_))
+    uint32_t m_; // m' in the papaer
+};
 
+#if 0
 template <int N>
 __global__ void __launch_bounds__(64)
     gemm_forward_4bit_cuda_m16nXk32(int G, int split_k_iters,
@@ -136,6 +184,7 @@ __global__ void __launch_bounds__(64)
       // - zero and * scale
       // TODO (Haotian): can save 4 assembly instructions if sormulate as deq =
       // q * scale - zero * scale.
+#ifndef USE_MACA
       asm volatile("sub.f16x2 %0, %1, %2;\n"
                    : "=r"(B_loaded_fp16.x)
                    : "r"(B_loaded_fp16.x), "r"(B_loaded_zero.x));
@@ -160,6 +209,7 @@ __global__ void __launch_bounds__(64)
       asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
                    : "=r"(B_loaded_fp16.w)
                    : "r"(B_loaded_fp16.w), "r"(B_loaded_scale.w), "r"(ZERO));
+#endif
       /*
       if (ax0_ax1_fused_0 == 0 && blockIdx_z == 0 && blockIdx_y == 0 && k_0_0 ==
       0 && threadIdx.x == 17 && threadIdx.y == 0){ printf("[x] %X %X %X %X\n",
@@ -176,6 +226,7 @@ __global__ void __launch_bounds__(64)
     for (int k_0_1 = 0; k_0_1 < 2; ++k_0_1) {
       {
         unsigned int addr;
+#ifndef USE_MACA
         __asm__ __volatile__(
             "{ .reg .u64 addr; cvta.to.shared.u64 addr, %1; cvt.u32.u64 %0, "
             "addr; }\n"
@@ -192,11 +243,13 @@ __global__ void __launch_bounds__(64)
               "=r"(((unsigned*)(A_shared_warp + 0))[2]),
               "=r"(((unsigned*)(A_shared_warp + 0))[3])
             : "r"(addr));
+#endif
       }
 
       for (int ax1_0 = 0; ax1_0 < N / 32; ++ax1_0) {
         {
           unsigned int addr;
+#ifndef USE_MACA
           __asm__ __volatile__(
               "{ .reg .u64 addr; cvta.to.shared.u64 addr, %1; cvt.u32.u64 %0, "
               "addr; }\n"
@@ -214,9 +267,11 @@ __global__ void __launch_bounds__(64)
                 "=r"(((unsigned*)(B_shared_warp + (ax1_0 * 8)))[2]),
                 "=r"(((unsigned*)(B_shared_warp + (ax1_0 * 8)))[3])
               : "r"(addr));
+#endif
         }
       }
       for (int j_0_4 = 0; j_0_4 < N / 32; ++j_0_4) {
+#ifndef USE_MACA
   #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
         {
           __asm__ __volatile__(
@@ -329,12 +384,13 @@ __global__ void __launch_bounds__(64)
         }
 
   #endif
+#endif
       }
     }
   }
 
   // TODO: Shang: Hoist loop invariance.
-  for (int ax1_0_1 = 0; ax1_0_1 < (N / 32); ++ax1_0_1) {
+  for (int ax1_0_1 = 0; ax1_0_1 < 4; ++ax1_0_1) {
     for (int local_id = 0; local_id < 8; ++local_id) {
       int row_offset = (((int)blockIdx_y) / j_factors1) * 16 +
                        ((int)threadIdx.x) / 4 + (local_id % 4) / 2 * 8;
@@ -346,20 +402,22 @@ __global__ void __launch_bounds__(64)
   }
 #endif
 }
+#endif
 
+template<typename T, typename VT>
 __global__ void __launch_bounds__(64)
-    dequantize_weights(int* __restrict__ B, half* __restrict__ scaling_factors,
-                       int* __restrict__ zeros, half* __restrict__ C, int G) {
+    dequantize_weights(int* __restrict__ B, T* __restrict__ scaling_factors,
+                       int* __restrict__ zeros, T* __restrict__ C, int G) {
   static constexpr uint32_t ZERO = 0x0;
-  half B_shared[32 * (128 + 8)];
-
-  half* B_shared_ptr2 = B_shared;
+  T B_shared[8];
+  T B_loaded_scale[8];
+  T* B_shared_ptr2 = B_shared;
 
   int N = blockDim.x * gridDim.x;  // 2
   int col = (blockIdx.x * blockDim.x + threadIdx.x);
   int row = blockIdx.y * blockDim.y + threadIdx.y;
   int index1 = 8 * col + 8 * row * N;
-  half* C_ptr2 = C + index1;
+  T* C_ptr2 = C + index1;
 
   int index2 = col + row * N;
   int* B_ptr2 = B + index2;
@@ -367,14 +425,18 @@ __global__ void __launch_bounds__(64)
   int index3 = col + (int)(row / G) * N;
   int* zeros_ptr2 = zeros + index3;
   int index4 = 8 * col + (int)(row / G) * N * 8;
-  half* scaling_factors_ptr2 = scaling_factors + index4;
+  T* scaling_factors_ptr2 = scaling_factors + index4;
 
   uint32_t zeros_loaded = *(uint32_t*)(zeros_ptr2);
-  uint4 B_loaded_zero = dequantize_s4_to_fp16x2(zeros_loaded);
-  uint4 B_loaded_scale = *(uint4*)(scaling_factors_ptr2);
-
+  *(uint4*)B_loaded_scale = *(uint4*)(scaling_factors_ptr2);
   uint32_t B_loaded = *(uint32_t*)B_ptr2;
-  uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2(B_loaded);
+#if 1
+  //zero //4bit scale 8bit b 4bit
+  hgemm_marlin_gptq::awq_dequant_4bits<T>(B_loaded,B_shared, B_loaded_scale, zeros_loaded);
+#else
+  uint4 B_loaded_zero = dequantize_s4_to_fp16x2<VT>(zeros_loaded);
+  uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2<VT>(B_loaded);
+#ifndef USE_MACA
   asm volatile("sub.f16x2 %0, %1, %2;\n"
                : "=r"(B_loaded_fp16.x)
                : "r"(B_loaded_fp16.x), "r"(B_loaded_zero.x));
@@ -399,17 +461,360 @@ __global__ void __launch_bounds__(64)
   asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
                : "=r"(B_loaded_fp16.w)
                : "r"(B_loaded_fp16.w), "r"(B_loaded_scale.w), "r"(ZERO));
+#else
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(B_loaded_fp16.x);
+unsigned int __b=(B_loaded_zero.x);
+VT __d=__hsub2(*(VT*)&__a,*(VT*)&__b);
+(B_loaded_fp16.x)=*(unsigned int*)&__d;
+}
+}
+
+
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(B_loaded_fp16.x);
+unsigned int __b=(B_loaded_scale.x);
+unsigned int __c=(ZERO);
+VT __d=__hfma2(*(VT*)&__a,*(VT*)&__b,*(VT*)&__c);
+(B_loaded_fp16.x)=*(unsigned int*)&__d;
+}
+}
+
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(B_loaded_fp16.y);
+unsigned int __b=(B_loaded_zero.y);
+VT __d=__hsub2(*(VT*)&__a,*(VT*)&__b);
+(B_loaded_fp16.y)=*(unsigned int*)&__d;
+}
+}
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(B_loaded_fp16.y);
+unsigned int __b=(B_loaded_scale.y);
+unsigned int __c=(ZERO);
+VT __d=__hfma2(*(VT*)&__a,*(VT*)&__b,*(VT*)&__c);
+(B_loaded_fp16.y)=*(unsigned int*)&__d;
+}
+}
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(B_loaded_fp16.z);
+unsigned int __b=(B_loaded_zero.z);
+VT __d=__hsub2(*(VT*)&__a,*(VT*)&__b);
+(B_loaded_fp16.z)=*(unsigned int*)&__d;
+}
+}
+
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(B_loaded_fp16.z);
+unsigned int __b=(B_loaded_scale.z);
+unsigned int __c=(ZERO);
+VT __d=__hfma2(*(VT*)&__a,*(VT*)&__b,*(VT*)&__c);
+(B_loaded_fp16.z)=*(unsigned int*)&__d;
+}
+}
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(B_loaded_fp16.w);
+unsigned int __b=(B_loaded_zero.w);
+VT __d=__hsub2(*(VT*)&__a,*(VT*)&__b);
+(B_loaded_fp16.w)=*(unsigned int*)&__d;
+}
+}
 
+
+// >>>> PTX2CPP Success <<<<
+{
+{
+unsigned int __a=(B_loaded_fp16.w);
+unsigned int __b=(B_loaded_scale.w);
+unsigned int __c=(ZERO);
+VT __d=__hfma2(*(VT*)&__a,*(VT*)&__b,*(VT*)&__c);
+(B_loaded_fp16.w)=*(unsigned int*)&__d;
+}
+}
+
+#endif
   *(uint4*)B_shared_ptr2 = B_loaded_fp16;
+#endif
 
   for (int i = 0; i < 8; ++i) {
     *(C_ptr2 + i) = B_shared[i];
   }
 }
 
+template<typename T>
+__global__ void dequantize_weights_opt(int* __restrict__ B, T* __restrict__ scaling_factors,
+                       int* __restrict__ zeros, T* __restrict__ C, int G, int length, int blocksize, int num_elems, DivModFast length_fast) {
+    constexpr int N = 8;
+    T B_loaded_scale[8];
+    T B_shared[8];
+    int tid = blockIdx.x * blocksize + threadIdx.x;
+    if(tid >= num_elems) return;
+    // int row = tid / length;
+    // int col = tid % length;
+    int row, col;
+    length_fast.divmod(tid, row, col);
+    int group_row = row / G;
+    int group_offset = group_row * length + col;
+    int offset = row * length + col;
+    uint32_t* ptr_zeros = (uint32_t*)(zeros + group_offset);
+    uint32_t* ptr_B = (uint32_t*)(B + offset);
+    T* ptr_scale = scaling_factors + group_offset * N;
+    T* ptr_C = C + offset * N;
+    uint32_t zeros_loaded = *(uint32_t*)ptr_zeros;
+    uint32_t B_loaded = *(uint32_t*)ptr_B;
+    *(uint4*)(B_loaded_scale) = *(uint4*)(ptr_scale);
+    hgemm_marlin_gptq::awq_dequant_4bits<T>(B_loaded,B_shared, B_loaded_scale, zeros_loaded);
+    *(float4*)(ptr_C) = *(float4*)(B_shared);
+}
+
+template <int BLOCK_SIZE>
+__global__ void awq_to_gptq_4bit(uint32_t *output, const uint32_t *input, int k, int n) {
+    constexpr int COMPACT_FACTOR = 8;
+    constexpr int QBIT = 4;
+    int tid = threadIdx.x;
+    int tile_idx = blockIdx.x * BLOCK_SIZE + tid;
+    int N_COMPACT = (n + COMPACT_FACTOR - 1) / COMPACT_FACTOR;
+    int K_COMPACT = (k + COMPACT_FACTOR - 1) / COMPACT_FACTOR;
+    int tile_n_idx = tile_idx / K_COMPACT;
+    int tile_k_idx = tile_idx % K_COMPACT;
+
+    uint32_t awq_data[COMPACT_FACTOR];
+    uint32_t temp_data[COMPACT_FACTOR];
+    uint32_t gptq_data[COMPACT_FACTOR];
+
+    int gptq_shift[COMPACT_FACTOR] = {0, 4, 1, 5, 2, 6, 3, 7};
+    int awq_shift[COMPACT_FACTOR] = {0, 4, 1, 5, 2, 6, 3, 7};
+
+    // load k8xn8
+    #pragma unroll
+    for (int i = 0; i < COMPACT_FACTOR; i++) {
+        int gvm_addr_offset = (tile_k_idx * COMPACT_FACTOR + i) * N_COMPACT + tile_n_idx;
+        int pred_k = tile_k_idx * COMPACT_FACTOR + i < k;
+        int pred_n = tile_n_idx * COMPACT_FACTOR < n;
+        if (pred_k && pred_n) {
+            awq_data[i] = *(input + gvm_addr_offset);
+        }
+    }
+
+    // decompress awq_data and recompress to gptq_data
+    #pragma unroll
+    for (int i = 0; i < COMPACT_FACTOR; i++) {
+        #pragma unroll
+        for(int j = 0; j < COMPACT_FACTOR; j++) {
+            temp_data[j] = ((awq_data[j] >> (awq_shift[i] * QBIT)) & 0xf);
+        }
+        #pragma unroll
+        for(int j = 0; j < COMPACT_FACTOR; j++) {
+            gptq_data[i] &= (~(0xf << (gptq_shift[j] * QBIT)));
+            gptq_data[i] |= temp_data[j] << (gptq_shift[j] * QBIT);
+
+        }
+    }
+
+    // store k8xn8
+    #pragma unroll
+    for (int i = 0; i < COMPACT_FACTOR; i++) {
+        int gvm_addr_offset = tile_k_idx * n + tile_n_idx * COMPACT_FACTOR + i;
+        int pred_k = tile_k_idx * COMPACT_FACTOR < k;
+        int pred_n = tile_n_idx * COMPACT_FACTOR + i < n;
+        if (pred_k && pred_n) {
+            *(output + gvm_addr_offset) = gptq_data[i];
+        } else {
+            *(output + gvm_addr_offset) = 0x00000000;
+        }
+    }
+}
+
+
+template <typename input_tp, const vllm::ScalarTypeId w_type_id, typename output_tp, typename quant_packed_tp>
+bool launch_gemm_gptq(int m,
+                      int n,
+                      int k,
+                      int quant_group,
+                      const input_tp *dA,
+                      int lda,
+                      const quant_packed_tp *dB,
+                      int ldb,
+                      output_tp *dC,
+		      float *dC_temp,
+                      int ldc,
+                      quant_packed_tp *d_zeros,
+                      input_tp *d_scales,
+                      const cudaStream_t stream,
+                      int chunks = 1) {
+    using namespace hgemm_marlin_gptq;
+    if(n % 16 != 0) {
+        printf("n %% 16 != 0, n = %d\n", n);
+        return false;
+    }
+    if(k % 32 != 0) {
+        printf("k %% 32 != 0, k = %d\n", k);
+        return false;
+    }
+    //const vllm::ScalarTypeId w_type_id = vllm::kU4B8.id();
+    const int THREADS = 256;
+    int BLOCKS_M = div_ceil(m, SLICE_M);
+    if(BLOCKS_M >= MAX_BLOCKS_M && BLOCKS_M % MAX_BLOCKS_M != 0) {
+        printf("Error: input m is error, m = %d, blocks_m = %d\n", m, BLOCKS_M);
+        return false;
+    }
+    if (BLOCKS_M > MAX_BLOCKS_M) BLOCKS_M = MAX_BLOCKS_M;
+    int BLOCKS_N = 8;
+    //It is better let TILE_K = quant_group
+    //But if quant_group is too large, a quant_group can be divided into two parts
+    int BLOCKS_K = quant_group / SLICE_K;
+    if (quant_group > 128) BLOCKS_K = 128 / SLICE_K;
+    //if (BLOCKS_M == 1 || BLOCKS_M == 2) {
+    //    BLOCKS_N = 16;
+    //}
+    const bool HAS_ACT_ORDER = false;
+    const bool HAS_ZP = (w_type_id == vllm::kU4.id()) || (w_type_id == vllm::kU8.id());
+    int *g_idx = nullptr;
+    bool HAS_NK_PRED = true;
+    bool HAS_M_PRED = true;
+    if (n % TILE_N == 0 && k % TILE_K == 0) {
+        HAS_NK_PRED = false;
+    }
+    if (m % TILE_M == 0) {
+        HAS_M_PRED = false;
+    }
+
+#define LAUNCH_AWQ(threads, bm, bn, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    else if (THREADS == threads && BLOCKS_M == bm && BLOCKS_N == bn \
+        && BLOCKS_K == bk  && HAS_ACT_ORDER == has_act_order \
+        && HAS_ZP == has_zp \
+        && HAS_M_PRED == has_m_pred && HAS_NK_PRED == has_nk_pred) { \
+            launch_gemm_gptq_kernel<input_tp, w_type_id, \
+                    threads, bm, bn, bk, has_act_order, has_zp, has_m_pred, has_nk_pred>( \
+                    (const PackTypeInt4*)dA, \
+                    (const PackTypeInt4*)dB, \
+                    (PackTypeInt4*)dC, \
+                    (PackTypeInt4*)dC_temp, \
+                    (const PackTypeInt4*)d_scales, \
+                    (const PackTypeInt4*)d_zeros, \
+                    nullptr, m, n, k, quant_group, chunks,\
+                    stream); \
+    }
+
+#define LAUNCH_AWQ_K(bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ(256, 1, 8, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ(256, 2, 8, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ(256, 3, 8, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ(256, 4, 8, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ(256, 1, 16, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ(256, 2, 16, bk, has_act_order, has_zp, has_nk_pred, has_m_pred)
+
+
+#define LAUNCH_AWQ_ZP(has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ_K(1, false, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ_K(2, false, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ_K(4, false, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ_K(8, false, has_zp, has_nk_pred, has_m_pred)
+
+#define LAUNCH_AWQ_PRED(has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ_ZP(false, has_nk_pred, has_m_pred) \
+    LAUNCH_AWQ_ZP(true, has_nk_pred, has_m_pred)
+
+    if (false) {
+
+    }
+    LAUNCH_AWQ_PRED(true, true)
+    LAUNCH_AWQ_PRED(true, false)
+    LAUNCH_AWQ_PRED(false, true)
+    LAUNCH_AWQ_PRED(false, false)
+    else {
+        printf("BLOCKS_M=%d, BLOCKS_N=%d, BLOCKS_k=%d, THREADS=%d, HAS_ACT_ORDER=%d, HAS_ZP=%d, quant_group=%d, HAS_M_PRED=%d, HAS_NK_PRED=%d is not supported\n",
+        BLOCKS_M, BLOCKS_N, BLOCKS_K, THREADS, HAS_ACT_ORDER, HAS_ZP, quant_group, HAS_M_PRED, HAS_NK_PRED);
+        return false;
+    }
+
+    return true;
+}
+
+
+template <typename input_tp, const vllm::ScalarTypeId w_type_id, typename output_tp, typename quant_packed_tp>
+bool launch_gemm(int quant_group,
+                int m,
+                int n,
+                int k,
+                const input_tp *dA,
+                int lda,
+                const quant_packed_tp *dB,
+                int ldb,
+                output_tp *dC,
+		            float* dC_temp,
+                int ldc,
+                quant_packed_tp *d_zeros,
+                input_tp *d_scales,
+		            const cudaStream_t stream) {
+    using namespace hgemm_marlin_gptq;
+    //constexpr int max_blocks_m = 4;
+    int total_m_blocks = div_ceil(m, SLICE_M);
+    int chunks = total_m_blocks / MAX_BLOCKS_M;
+    int rest_blocks_m = total_m_blocks % MAX_BLOCKS_M;
+    // printf("m=%d,n=%d,k=%d,lda=%d,ldb=%d,ldc=%d,total_m_blocks=%d,chunks=%d,rest_blocks_m=%d\n",
+    //     m, n, k, lda, ldb, ldc, total_m_blocks, chunks, rest_blocks_m
+    // );
+    //const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    // const int quant_group = 128;
+    bool ret = true;
+    if (chunks > 0) {
+        int real_m = m > chunks * MAX_BLOCKS_M * SLICE_M ? chunks * MAX_BLOCKS_M * SLICE_M : m;
+        ret = launch_gemm_gptq<input_tp, w_type_id, output_tp, quant_packed_tp>(real_m, n, k, quant_group, dA, lda, dB, ldb, dC, dC_temp, ldc, d_zeros, d_scales, stream, chunks);
+    }
+    if (rest_blocks_m > 0) {
+        int m_offset = chunks * MAX_BLOCKS_M * SLICE_M;
+        ret = ret && launch_gemm_gptq<input_tp, w_type_id, output_tp, quant_packed_tp>(m - m_offset, n, k, quant_group, dA + lda * m_offset, lda, dB, ldb, dC + ldc * m_offset, dC_temp + ldc * m_offset, ldc, d_zeros, d_scales, stream, 1);
+    }
+
+    return ret;
+}
+
 }  // namespace awq
 }  // namespace vllm
 
+torch::Tensor awq_to_gptq_4bit(torch::Tensor qweight) {
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  const uint32_t* qweight_ptr = reinterpret_cast<const uint32_t*>(qweight.data_ptr<int>());
+
+  int num_in_channels = qweight.size(0);
+  int num_out_channels = qweight.size(1) * 8;
+
+  int compact_n = (num_out_channels + hgemm_marlin_gptq::PACK_RATIO_4BITS - 1) / hgemm_marlin_gptq::PACK_RATIO_4BITS;
+  int compact_output_k = (num_in_channels + hgemm_marlin_gptq::PACK_RATIO_4BITS - 1) / hgemm_marlin_gptq::PACK_RATIO_4BITS;;
+
+  int block_size = 256;
+  int tile_all_num = compact_n * compact_output_k;
+  int grid_size = (tile_all_num + 255) / 256;
+
+  auto options = torch::TensorOptions()
+                     .dtype(qweight.dtype())
+                     .device(qweight.device());
+
+  torch::Tensor out = torch::zeros({num_out_channels,  compact_output_k}, options);
+  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr<int>());
+
+  vllm::awq::awq_to_gptq_4bit<256><<<grid_size, block_size, 0, stream>>>((uint32_t*)out_ptr, (const uint32_t*)qweight_ptr, num_in_channels, num_out_channels);
+
+  return out;
+}
+
 torch::Tensor awq_dequantize(torch::Tensor _kernel,
                              torch::Tensor _scaling_factors,
                              torch::Tensor _zeros, int64_t split_k_iters,
@@ -419,24 +824,6 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
   int out_c = qout_c * 8;
   int G = in_c / _scaling_factors.size(0);
 
-  int x_thread = thx;
-  int y_thread = thy;
-
-  int x_blocks = 1;
-  int y_blocks = 1;
-  if (thx == 0) {
-    x_thread = qout_c;
-  }
-  if (thy == 0) {
-    y_thread = in_c;
-  }
-  if (thx == 0 && thy == 0) {
-    x_thread = 8;
-    y_thread = 8;
-    x_blocks = (int)(qout_c / 8);
-    y_blocks = (int)(in_c / 8);
-  }
-
   const at::cuda::OptionalCUDAGuard device_guard(device_of(_scaling_factors));
 
   auto options = torch::TensorOptions()
@@ -444,19 +831,48 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
                      .device(_scaling_factors.device());
   at::Tensor _de_kernel = torch::empty({in_c, out_c}, options);
 
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
-  auto de_kernel = reinterpret_cast<half*>(_de_kernel.data_ptr<at::Half>());
-  auto scaling_factors =
-      reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
   auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
-
+#if 1
+  int blocksize = 512;
+  int num_elems = in_c * qout_c;
+  int gridsize = (num_elems + blocksize - 1) / blocksize;
+  if(_scaling_factors.dtype() == at::ScalarType::Half) {
+    auto de_kernel = reinterpret_cast<half*>(_de_kernel.data_ptr<at::Half>());
+    auto scaling_factors =
+        reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
+    vllm::awq::dequantize_weights_opt<__half><<<gridsize, blocksize, 0 , stream>>>(kernel, scaling_factors, zeros, de_kernel, G, qout_c, blocksize, num_elems, vllm::awq::DivModFast(qout_c));
+  } else if(_scaling_factors.dtype() == at::ScalarType::BFloat16) {
+    auto de_kernel = reinterpret_cast<maca_bfloat16*>(_de_kernel.data_ptr<at::BFloat16>());
+    auto scaling_factors =
+        reinterpret_cast<maca_bfloat16*>(_scaling_factors.data_ptr<at::BFloat16>());
+    vllm::awq::dequantize_weights_opt<maca_bfloat16><<<gridsize, blocksize, 0, stream>>>(kernel, scaling_factors, zeros, de_kernel, G, qout_c, blocksize, num_elems, vllm::awq::DivModFast(qout_c));
+  } else {
+    printf("not support this type\n");
+    assert(0);
+  }
+#else
   dim3 num_blocks(x_blocks, y_blocks);
   dim3 threads_per_block(x_thread, y_thread);
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  vllm::awq::dequantize_weights<<<num_blocks, threads_per_block, 0, stream>>>(
-      kernel, scaling_factors, zeros, de_kernel, G);
-
+  if(_scaling_factors.dtype() == at::ScalarType::Half) {
+    auto de_kernel = reinterpret_cast<half*>(_de_kernel.data_ptr<at::Half>());
+    auto scaling_factors =
+        reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
+    vllm::awq::dequantize_weights<__half,__half2><<<num_blocks, threads_per_block, 0, stream>>>(
+        kernel, scaling_factors, zeros, de_kernel, G);
+  } else if(_scaling_factors.dtype() == at::ScalarType::BFloat16) {
+    auto de_kernel = reinterpret_cast<maca_bfloat16*>(_de_kernel.data_ptr<at::BFloat16>());
+    auto scaling_factors =
+        reinterpret_cast<maca_bfloat16*>(_scaling_factors.data_ptr<at::BFloat16>());
+    vllm::awq::dequantize_weights<maca_bfloat16,maca_bfloat162><<<num_blocks, threads_per_block, 0, stream>>>(
+        kernel, scaling_factors, zeros, de_kernel, G);
+  } else {
+    printf("not support this type\n");
+    assert(0);
+  }
+#endif
   return _de_kernel;
 }
 
@@ -468,59 +884,49 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
 
 torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
                        torch::Tensor _scaling_factors, torch::Tensor _zeros,
-                       int64_t split_k_iters) {
+                       int64_t split_k_iters,
+                       torch::Tensor _temp_space,
+                       bool dtype_bf16) {
   int num_in_feats = _in_feats.size(0);
   int num_in_channels = _in_feats.size(1);
   const at::cuda::OptionalCUDAGuard device_guard(device_of(_in_feats));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   auto options = torch::TensorOptions()
                      .dtype(_in_feats.dtype())
                      .device(_in_feats.device());
+                     
+  // int num_out_channels = _kernel.size(1) * 8;
+  int num_out_channels = _kernel.size(0); 
   at::Tensor _out_feats =
-      torch::empty({split_k_iters, num_in_feats, _kernel.size(1) * 8}, options);
-  int num_out_feats = _out_feats.size(-2);
-  int num_out_channels = _out_feats.size(-1);
+      torch::zeros({num_in_feats, num_out_channels}, options);
 
-  auto in_feats = reinterpret_cast<half*>(_in_feats.data_ptr<at::Half>());
+  //auto in_feats = reinterpret_cast<half*>(_in_feats.data_ptr<at::Half>());
   auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
-  auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
-  auto scaling_factors =
-      reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
+  //auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
+  //auto scaling_factors =
+  //    reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
   auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
+  auto temp_space = reinterpret_cast<float*>(_temp_space.data_ptr<float>());
   int group_size = num_in_channels / _scaling_factors.size(0);
 
-  if (num_out_channels % 64 != 0)
-    throw std::invalid_argument("OC is not multiple of cta_N = 64");
-  if (num_out_channels % 8 != 0)
-    throw std::invalid_argument("OC is not multiple of pack_num = 8");
-  if (group_size % 32 != 0)
-    throw std::invalid_argument("Group size should be a multiple of 32");
-  if (num_out_channels % group_size != 0)
-    throw std::invalid_argument("OC is not multiple of Group size");
-
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  if (num_out_channels % 128 == 0) {
-    int j_factors1 = num_out_channels / 128 / 1;
-    dim3 num_blocks((num_out_feats + 16 - 1) / 16 * j_factors1 * split_k_iters);
-    // threadIdx.x: 32
-    // threadIdx.y: i_factors[2] * j_factors[2]
-    dim3 threads_per_block(32, 2);
-    vllm::awq::gemm_forward_4bit_cuda_m16nXk32<128>
-        <<<num_blocks, threads_per_block, 0, stream>>>(
-            group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
-            num_in_feats, num_in_channels, num_out_channels, out_feats);
-  } else if (num_out_channels % 64 == 0) {
-    int j_factors1 = num_out_channels / 64 / 1;
-    dim3 num_blocks(1 * (num_out_feats + 16 - 1) / 16 * j_factors1 *
-                    split_k_iters);
-
-    // threadIdx.x: 32
-    // threadIdx.y: i_factors[2] * j_factors[2]
-    dim3 threads_per_block(32, 2);
-    vllm::awq::gemm_forward_4bit_cuda_m16nXk32<64>
-        <<<num_blocks, threads_per_block, 0, stream>>>(
-            group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
-            num_in_feats, num_in_channels, num_out_channels, out_feats);
+  int lda = num_in_channels;
+  int ldb = num_out_channels;
+  int ldc = num_out_channels;
+
+  if (dtype_bf16) {
+	  using scalar_t = __maca_bfloat16;
+	  vllm::awq::launch_gemm<scalar_t, vllm::kU4.id(), scalar_t, quant_packed_type>(group_size, num_in_feats, num_out_channels, num_in_channels,
+                                     (const scalar_t*)_in_feats.data_ptr(), lda, (const uint32_t*)kernel, ldb, (scalar_t*)_out_feats.data_ptr(), temp_space, ldc,
+                                     (uint32_t*)zeros, (scalar_t*)_scaling_factors.data_ptr(), stream);
+  } else {
+	  auto in_feats = reinterpret_cast<half*>(_in_feats.data_ptr<at::Half>());
+	  auto scaling_factors = reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
+	  auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
+	  vllm::awq::launch_gemm<input_type, vllm::kU4.id(), output_type, quant_packed_type>(group_size, num_in_feats, num_out_channels, num_in_channels,
+                                     (const half*)in_feats, lda, (const uint32_t*)kernel, ldb, (half*)out_feats, nullptr, ldc,
+                                     (uint32_t*)zeros, (half*)scaling_factors, stream);
   }
-  return _out_feats.sum(0);
+
+  return _out_feats;
 }
diff --git a/csrc/quantization/awq/hgemv_nn_splitk_awq.hpp b/csrc/quantization/awq/hgemv_nn_splitk_awq.hpp
new file mode 100644
index 000000000..39570b9da
--- /dev/null
+++ b/csrc/quantization/awq/hgemv_nn_splitk_awq.hpp
@@ -0,0 +1,520 @@
+/**
+ * @file hgemv_nn_splitk.hpp
+ * @author Jiawen Yang (jiawen.yang@metax-tech.com)
+ * @brief *
+ * @version 0.1
+ * @date 2024-03-05
+ * @copyright Copyright (c) 2024
+ *
+ *   fp16 gemv kernel template for some gemv cases
+ *   Note:
+ *    1. BlockDimX * BlockDimY = 64, and BlockDimX should be 8/16/32/64
+ *    2. LoopNum % 2 == 0, so Load_B can use ldg_b32 or ldg_b64
+ *    3. m % (BlockDimX * 8) == 0
+ *    4. k % (ThreadBlock / BlockDimX * LoopNum * SplitKNum) = 0
+ *
+ *    A load layout:
+ *
+ *       **************************** Wave_0 ******************* | Wave_1  ...
+ *       ********* Repeat LoopNum *********                      |
+ *       tid_0(ldg_b128)   tid_0 ... tid_0 | tid_(BlockDimX) ... |
+ *       tid_1                             |                     |
+ *       tid_2                             |                     |
+ *       ……                                |                     |
+ *       tid_(BlockDimX-1)                 |                     |
+ *
+ */
+#pragma once
+
+#include <mc_runtime.h>
+#include <maca_fp16.h>
+#include "../gptq/Hgemm_common.cuh"
+#include "dequant.cuh"
+#define quant_packed_type uint32_t
+typedef __NATIVE_VECTOR__(2, float) v2f;
+
+template<int N, int m_per_thread>
+__device__ __forceinline__ void dequant_fma_awq_int4(
+                const quant_packed_type& a,
+                const v2f (&scale)[m_per_thread/2],
+                const v2f (&zero)[m_per_thread/2],
+                const v2f (&b)[N],
+                v2f (&out)[N][m_per_thread/2]) {
+    uint32_t p0 = a & 0x0f0f0f0f;
+    float o1,o3;
+    asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o1):"r"(p0));
+    asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o3):"r"(p0));
+    v2f a0 = {o1, o3};
+    a0 = __builtin_mxc_pk_fma_f32(a0, scale[0], zero[0]);
+
+    #pragma unroll N
+    for (int y = 0; y < N; y++) {
+        out[y][0] = __builtin_mxc_pk_fma_f32(a0, b[y], out[y][0]);
+    }
+
+    asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o1):"r"(p0));
+    asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o3):"r"(p0));
+    a0 = {o1, o3};
+    a0 = __builtin_mxc_pk_fma_f32(a0, scale[2], zero[2]);
+
+    #pragma unroll N
+    for (int y = 0; y < N; y++) {
+        out[y][2] = __builtin_mxc_pk_fma_f32(a0, b[y], out[y][2]);
+    }
+
+    uint32_t p1 = (a >> 4) & 0x0f0f0f0f;
+    asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o1):"r"(p1));
+    asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o3):"r"(p1));
+    a0 = {o1, o3};
+    a0 = __builtin_mxc_pk_fma_f32(a0, scale[1], zero[1]);
+
+    #pragma unroll N
+    for (int y = 0; y < N; y++) {
+        out[y][1] = __builtin_mxc_pk_fma_f32(a0, b[y], out[y][1]);
+    }
+
+    asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o1):"r"(p1));
+    asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o3):"r"(p1));
+    a0 = {o1, o3};
+    a0 = __builtin_mxc_pk_fma_f32(a0, scale[3], zero[3]);
+
+    #pragma unroll N
+    for (int y = 0; y < N; y++) {
+        out[y][3] = __builtin_mxc_pk_fma_f32(a0, b[y], out[y][3]);
+    }
+};
+
+template <int ThreadBlock, int BlockDimX, int BATCH>
+__global__ __launch_bounds__(256) void hgemv_nn_splitk_awq(const half* __restrict__ srcB,
+                                                        const quant_packed_type* __restrict__ srcA,
+                                                        const quant_packed_type* __restrict__ zeros,
+                                                        const half* __restrict__ scales,
+                                                        half *dst,
+                                                        int m,
+                                                        int n,
+                                                        int k,
+                                                        int srcAStride,
+                                                        int dstStride,
+                                                        int k_div_sk,
+                                                        const int* __restrict__ b_perm = nullptr) {
+    constexpr int QBITS = 4;
+    constexpr int PACK_RATIO = (32 / QBITS);
+    constexpr int QUANT_GROUP = 128;
+    constexpr int LoopNum = QUANT_GROUP / (ThreadBlock / BlockDimX);
+    constexpr int N = BATCH;
+    const int k_stride = k;
+    const int k_block = (k_div_sk + QUANT_GROUP - 1) / QUANT_GROUP * QUANT_GROUP;
+    const int splitKOffset = blockIdx.y * k_block;
+    if (splitKOffset + k_block > k) k = k - splitKOffset;
+    else k = k_block;
+    srcA += splitKOffset * srcAStride / PACK_RATIO;
+    //srcB += splitKOffset;
+
+    constexpr int quant_groups = 1;
+    constexpr int m_per_thread = PACK_RATIO;
+    constexpr int thread_groups = ThreadBlock / BlockDimX;
+    constexpr int group_elements = BlockDimX * m_per_thread;
+    constexpr int reduce_size = ThreadBlock * m_per_thread;
+    constexpr int data_cache_size = LoopNum * thread_groups * N + group_elements + group_elements / 2;
+
+    float *bsm_b_ptr, *bsm_zeros_ptr, *smem;
+    half *bsm_scales_ptr;
+    if constexpr(reduce_size > data_cache_size) {
+        __shared__ float bsm_ptr[reduce_size];
+        bsm_b_ptr = bsm_ptr;
+        bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+        bsm_scales_ptr = (half*)(bsm_zeros_ptr + group_elements);
+        smem = bsm_ptr;
+    } else {
+        __shared__ float bsm_ptr[data_cache_size];
+        bsm_b_ptr = bsm_ptr;
+        bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+        bsm_scales_ptr = (half*)(bsm_zeros_ptr + group_elements);
+        smem = bsm_ptr;
+    }
+
+    const int zeros_stride = srcAStride / PACK_RATIO;
+    const int scales_stride = srcAStride;
+    zeros += splitKOffset * zeros_stride / QUANT_GROUP;
+    scales += splitKOffset * scales_stride / QUANT_GROUP;
+
+    dst += group_elements * blockIdx.x;
+    const int m_offset_a = group_elements * blockIdx.x / PACK_RATIO;
+    const int m_offset_zeros = group_elements / PACK_RATIO * blockIdx.x;
+    const int m_offset_scales = group_elements * blockIdx.x;
+
+    //store splited fma results
+    v2f c_splited[N][m_per_thread/2];
+    for (int i = 0; i < N; i++) {
+        for (int j = 0; j < m_per_thread/2; j++) c_splited[i][j] = {0, 0};
+    }
+
+    int tid = threadIdx.x;
+    int tidCol = tid / BlockDimX;
+    int tidRow = tid % BlockDimX;
+    int this_group_elements = (blockIdx.x + 1) * group_elements <= m ? group_elements : m - blockIdx.x * group_elements;
+    int m_index = tidRow * m_per_thread;
+    for (int i = 0; i < k; i += LoopNum * thread_groups) {
+        int quant_group = i / QUANT_GROUP;
+        constexpr int loading_pack = 2;
+        int loading_count = this_group_elements / loading_pack;
+        //Load needed zeros, scales
+        const int shuffle_rank[8] = {0, 2, 4, 6, 1, 3, 5, 7};
+        for (int x = tid; x < loading_count; x += ThreadBlock) {
+            uint8_t temp_zeros = 0;
+            temp_zeros = ((uint8_t*)(zeros + quant_group * zeros_stride + m_offset_zeros))[x];
+            half temp_scales[2];
+            int packed_group_offset = (x >> 2) << 3;
+            int packed_index = (x << 1) % 8;
+            int dest_offset_0 = packed_group_offset + shuffle_rank[packed_index];
+            int dest_offset_1 = packed_group_offset + shuffle_rank[packed_index + 1];
+            //*((float*)temp_scales) = ((float*)(scales + quant_group * scales_stride + m_offset_scales))[x];
+            temp_scales[0] = scales[quant_group * scales_stride + m_offset_scales + dest_offset_0];
+            temp_scales[1] = scales[quant_group * scales_stride + m_offset_scales + dest_offset_1];
+            uint32_t z = temp_zeros;
+            uint32_t z0 = __builtin_mxc_ubfe(z, 0, QBITS);
+            uint32_t z1 = __builtin_mxc_ubfe(z, 4, QBITS);
+            float s1 = (float)(temp_scales[0]);
+            float s2 = (float)(temp_scales[1]);
+            //Store to shared memory
+            bsm_zeros_ptr[dest_offset_0] = (float)z0 * s1 * -1.0f;
+            bsm_zeros_ptr[dest_offset_1] = (float)z1 * s2 * -1.0f;
+            bsm_scales_ptr[dest_offset_0] = temp_scales[0];
+            bsm_scales_ptr[dest_offset_1] = temp_scales[1];
+        }
+
+        int loop_index = 0;
+
+        //Load B and transform to float
+        if (b_perm != nullptr) {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[b_perm[splitKOffset + i + tidCol * LoopNum + x] + y * k_stride];
+                }
+            }
+        } else {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[splitKOffset + i + tidCol * LoopNum + x + y * k_stride];
+                }
+            }
+        }
+
+        __syncthreads();
+
+        //Load zero and scale from bsm
+        if (m_index < this_group_elements) {
+            v2f local_scales[m_per_thread/2];
+            for (int c = 0; c < m_per_thread / 2; c++) {
+                float s0 = (float)bsm_scales_ptr[m_index + c*2];
+                float s1 = (float)bsm_scales_ptr[m_index + c*2+1];
+                local_scales[c] = {s0, s1};
+            }
+            v2f local_zeros[m_per_thread/2];
+            for (int c = 0; c < m_per_thread/2; c++) local_zeros[c] = {bsm_zeros_ptr[m_index + c*2],bsm_zeros_ptr[m_index + c*2+1]};
+
+    #define DEQUANT_FMA(a, b) \
+            dequant_fma_awq_int4<N, m_per_thread>(a, local_scales, local_zeros, b, c_splited);
+
+            quant_packed_type A[4];
+            const int packed_a_stride = srcAStride / PACK_RATIO;
+            int src_a_offset = (loop_index + i + tidCol * LoopNum) * srcAStride / PACK_RATIO + m_offset_a + m_per_thread * tidRow / PACK_RATIO;
+            A[0] = srcA[src_a_offset];
+            src_a_offset += packed_a_stride;
+            A[1] = srcA[src_a_offset];
+
+            v2f local_b[4][N];
+            //#pragma unroll LoopNum / 4 - 1
+            for (; loop_index < LoopNum - 4; loop_index += 4) {
+                //Load A
+                src_a_offset += packed_a_stride;
+                A[2] = srcA[src_a_offset];
+                src_a_offset += packed_a_stride;
+                A[3] = srcA[src_a_offset];
+
+                for (int y = 0; y < N; y++) {
+                    float s[4];
+                    *(float4*)s = *(float4*)(bsm_b_ptr+tidCol*LoopNum+loop_index+y*thread_groups*LoopNum);
+                    local_b[0][y] = {s[0], s[0]};
+                    local_b[1][y] = {s[1], s[1]};
+                    local_b[2][y] = {s[2], s[2]};
+                    local_b[3][y] = {s[3], s[3]};
+                }
+                DEQUANT_FMA(A[0], local_b[0])
+                DEQUANT_FMA(A[1], local_b[1])
+                src_a_offset += packed_a_stride;
+                A[0] = srcA[src_a_offset];
+                src_a_offset += packed_a_stride;
+                A[1] = srcA[src_a_offset];
+                DEQUANT_FMA(A[2], local_b[2])
+                DEQUANT_FMA(A[3], local_b[3])
+            }
+            src_a_offset += packed_a_stride;
+            A[2] = srcA[src_a_offset];
+            src_a_offset += packed_a_stride;
+            A[3] = srcA[src_a_offset];
+            for (int y = 0; y < N; y++) {
+                float s[4];
+                *(float4*)s = *(float4*)(bsm_b_ptr+tidCol*LoopNum+loop_index+y*thread_groups*LoopNum);
+                local_b[0][y] = {s[0], s[0]};
+                local_b[1][y] = {s[1], s[1]};
+                local_b[2][y] = {s[2], s[2]};
+                local_b[3][y] = {s[3], s[3]};
+            }
+            DEQUANT_FMA(A[0], local_b[0])
+            DEQUANT_FMA(A[1], local_b[1])
+            DEQUANT_FMA(A[2], local_b[2])
+            DEQUANT_FMA(A[3], local_b[3])
+        }
+        __syncthreads();
+    }
+#undef DEQUANT_FMA
+    #pragma unroll N
+    for (int y = 0; y < N; y++) {
+        if (m_index < this_group_elements) {
+            for (int i = 0; i < m_per_thread/2; i++) {
+                smem[tidCol + (tidRow * m_per_thread + i*2) * ThreadBlock / BlockDimX] = c_splited[y][i].x;
+                smem[tidCol + (tidRow * m_per_thread + i*2+1) * ThreadBlock / BlockDimX] = c_splited[y][i].y;
+            }
+        }
+        __syncthreads();
+        constexpr int stride = ThreadBlock / BlockDimX;
+        int data_size = ThreadBlock * m_per_thread;
+        #pragma unroll
+        for (int i = ThreadBlock / BlockDimX / 2; i > 0; i /= 2) {
+            for (int j = tid; j < data_size / 2; j += ThreadBlock) {
+                int reduce_group = j / i;
+                int reduce_index = j % i;
+                smem[reduce_index + reduce_group * stride] += smem[reduce_index + reduce_group * stride + i];
+            }
+            __syncthreads();
+            data_size /= 2;
+        }
+        for (int i = tid; i < this_group_elements; i += ThreadBlock) {
+            atomicAdd(dst + i + y * dstStride, (half)smem[i*stride]);
+        }
+        if constexpr(N > 1) {
+            if (y + 1 < N) {
+                __syncthreads();
+            }
+        }
+    }
+}
+
+template <int BX, int BATCH>
+__global__ __launch_bounds__(256) void hgemv_nn_splitk_awq_kb128(const half* __restrict__ srcB,
+                                                        const quant_packed_type* __restrict__ srcA,
+                                                        const quant_packed_type* __restrict__ zeros,
+                                                        const half* __restrict__ scales,
+                                                        half *dst,
+                                                        int m,
+                                                        int n,
+                                                        int k,
+                                                        int srcAStride,
+                                                        int dstStride,
+                                                        int k_div_sk,
+                                                        const int* __restrict__ b_perm = nullptr) {
+    constexpr int QBITS = 4;
+    constexpr int PACK_RATIO = (32 / QBITS);
+    constexpr int QUANT_GROUP = 128;
+    constexpr int ThreadBlock = 256;
+    constexpr int BlockDimX = BX;
+    constexpr int LoopNum = QUANT_GROUP / (ThreadBlock / BlockDimX);
+    constexpr int N = BATCH;
+    const int k_stride = k;
+    const int k_block = 128;
+    const int splitKOffset = blockIdx.y * k_block;
+
+    k = k_block;
+    srcA += splitKOffset * srcAStride / PACK_RATIO;
+    //srcB += splitKOffset;
+
+    constexpr int quant_groups = 1;
+    constexpr int m_per_thread = PACK_RATIO;
+    constexpr int thread_groups = ThreadBlock / BlockDimX;
+    constexpr int group_elements = BlockDimX * m_per_thread;
+    constexpr int reduce_size = ThreadBlock * m_per_thread;
+    constexpr int data_cache_size = LoopNum * thread_groups * N + group_elements + group_elements / 2;
+    float *bsm_b_ptr, *bsm_zeros_ptr, *smem;
+    half *bsm_scales_ptr;
+    if constexpr(reduce_size > data_cache_size) {
+        __shared__ float bsm_ptr[reduce_size];
+        bsm_b_ptr = bsm_ptr;
+        bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+        bsm_scales_ptr = (half*)(bsm_zeros_ptr + group_elements);
+        smem = bsm_ptr;
+    } else {
+        __shared__ float bsm_ptr[data_cache_size];
+        bsm_b_ptr = bsm_ptr;
+        bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+        bsm_scales_ptr = (half*)(bsm_zeros_ptr + group_elements);
+        smem = bsm_ptr;
+    }
+
+    const int zeros_stride = srcAStride / PACK_RATIO;
+    const int scales_stride = srcAStride;
+    zeros += splitKOffset * zeros_stride / QUANT_GROUP;
+    scales += splitKOffset * scales_stride / QUANT_GROUP;
+
+    dst += group_elements * blockIdx.x;
+    const int m_offset_a = group_elements * blockIdx.x / PACK_RATIO;
+    const int m_offset_zeros = group_elements / PACK_RATIO * blockIdx.x;
+    const int m_offset_scales = group_elements * blockIdx.x;
+
+    //store splited fma results
+    v2f c_splited[N][m_per_thread/2];
+    for (int i = 0; i < N; i++) {
+        for (int j = 0; j < m_per_thread/2; j++) c_splited[i][j] = {0, 0};
+    }
+
+    int tid = threadIdx.x;
+    int tidCol = tid / BlockDimX;
+    int tidRow = tid % BlockDimX;
+    int this_group_elements = (blockIdx.x + 1) * group_elements <= m ? group_elements : m - blockIdx.x * group_elements;
+
+    constexpr int quant_group = 0;
+    constexpr int loading_pack = 2;
+    int loading_count = this_group_elements / loading_pack;
+    //Load needed zeros, scales
+    const int shuffle_rank[8] = {0, 2, 4, 6, 1, 3, 5, 7};
+    for (int x = tid; x < loading_count; x += ThreadBlock) {
+        uint8_t temp_zeros = 0;
+        temp_zeros = ((uint8_t*)(zeros + quant_group * zeros_stride + m_offset_zeros))[x];
+        half temp_scales[2];
+        int packed_group_offset = (x >> 2) << 3;
+        int packed_index = (x << 1) % 8;
+        int dest_offset_0 = packed_group_offset + shuffle_rank[packed_index];
+        int dest_offset_1 = packed_group_offset + shuffle_rank[packed_index + 1];
+        //*((float*)temp_scales) = ((float*)(scales + quant_group * scales_stride + m_offset_scales))[x];
+        temp_scales[0] = scales[quant_group * scales_stride + m_offset_scales + dest_offset_0];
+        temp_scales[1] = scales[quant_group * scales_stride + m_offset_scales + dest_offset_1];
+        uint32_t z = temp_zeros;
+        uint32_t z0 = __builtin_mxc_ubfe(z, 0, QBITS);
+        uint32_t z1 = __builtin_mxc_ubfe(z, 4, QBITS);
+        float s1 = (float)(temp_scales[0]);
+        float s2 = (float)(temp_scales[1]);
+        //Store to shared memory
+        bsm_zeros_ptr[dest_offset_0] = (float)z0 * s1 * -1.0f;
+        bsm_zeros_ptr[dest_offset_1] = (float)z1 * s2 * -1.0f;
+        bsm_scales_ptr[dest_offset_0] = temp_scales[0];
+        bsm_scales_ptr[dest_offset_1] = temp_scales[1];
+    }
+
+    int loop_index = 0;
+
+    //Load B and transform to float
+    if (b_perm != nullptr) {
+        for (int y = 0; y < N; y++) {
+            for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[b_perm[splitKOffset + tidCol * LoopNum + x] + y * k_stride];
+            }
+        }
+    } else {
+        for (int y = 0; y < N; y++) {
+            for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[splitKOffset + tidCol * LoopNum + x + y * k_stride];
+            }
+        }
+    }
+
+    __syncthreads();
+
+    //Load zero and scale from bsm
+    int m_index = tidRow * m_per_thread;
+    if (m_index < this_group_elements) {
+        v2f local_scales[m_per_thread/2];
+        for (int c = 0; c < m_per_thread / 2; c++) {
+            float s0 = (float)bsm_scales_ptr[m_index + c*2];
+            float s1 = (float)bsm_scales_ptr[m_index + c*2+1];
+            local_scales[c] = {s0, s1};
+        }
+        v2f local_zeros[m_per_thread/2];
+        for (int c = 0; c < m_per_thread/2; c++) local_zeros[c] = {bsm_zeros_ptr[m_index + c*2],bsm_zeros_ptr[m_index + c*2+1]};
+
+#define DEQUANT_FMA(a, b) \
+        dequant_fma_awq_int4<N, m_per_thread>(a, local_scales, local_zeros, b, c_splited);
+
+        quant_packed_type A[4];
+        const int packed_a_stride = srcAStride / PACK_RATIO;
+        int src_a_offset = (loop_index + tidCol * LoopNum) * srcAStride / PACK_RATIO + m_offset_a + m_per_thread * tidRow / PACK_RATIO;
+        A[0] = srcA[src_a_offset];
+        src_a_offset += packed_a_stride;
+        A[1] = srcA[src_a_offset];
+
+        v2f local_b[4][N];
+        #pragma unroll LoopNum / 4 - 1
+        for (; loop_index < LoopNum - 4; loop_index += 4) {
+            //Load A
+            src_a_offset += packed_a_stride;
+            A[2] = srcA[src_a_offset];
+            src_a_offset += packed_a_stride;
+            A[3] = srcA[src_a_offset];
+
+            for (int y = 0; y < N; y++) {
+                float s[4];
+                *(float4*)s = *(float4*)(bsm_b_ptr+tidCol*LoopNum+loop_index+y*thread_groups*LoopNum);
+                local_b[0][y] = {s[0], s[0]};
+                local_b[1][y] = {s[1], s[1]};
+                local_b[2][y] = {s[2], s[2]};
+                local_b[3][y] = {s[3], s[3]};
+            }
+            DEQUANT_FMA(A[0], local_b[0])
+            DEQUANT_FMA(A[1], local_b[1])
+            src_a_offset += packed_a_stride;
+            A[0] = srcA[src_a_offset];
+            src_a_offset += packed_a_stride;
+            A[1] = srcA[src_a_offset];
+            DEQUANT_FMA(A[2], local_b[2])
+            DEQUANT_FMA(A[3], local_b[3])
+        }
+        src_a_offset += packed_a_stride;
+        A[2] = srcA[src_a_offset];
+        src_a_offset += packed_a_stride;
+        A[3] = srcA[src_a_offset];
+        for (int y = 0; y < N; y++) {
+            float s[4];
+            *(float4*)s = *(float4*)(bsm_b_ptr+tidCol*LoopNum+loop_index+y*thread_groups*LoopNum);
+            local_b[0][y] = {s[0], s[0]};
+            local_b[1][y] = {s[1], s[1]};
+            local_b[2][y] = {s[2], s[2]};
+            local_b[3][y] = {s[3], s[3]};
+        }
+        DEQUANT_FMA(A[0], local_b[0])
+        DEQUANT_FMA(A[1], local_b[1])
+        DEQUANT_FMA(A[2], local_b[2])
+        DEQUANT_FMA(A[3], local_b[3])
+    }
+    __syncthreads();
+
+#undef DEQUANT_FMA
+    #pragma unroll N
+    for (int y = 0; y < N; y++) {
+        if (m_index < this_group_elements) {
+            for (int i = 0; i < m_per_thread/2; i++) {
+                smem[tidCol + (tidRow * m_per_thread + i*2) * ThreadBlock / BlockDimX] = c_splited[y][i].x;
+                smem[tidCol + (tidRow * m_per_thread + i*2+1) * ThreadBlock / BlockDimX] = c_splited[y][i].y;
+            }
+        }
+        __syncthreads();
+        constexpr int stride = ThreadBlock / BlockDimX;
+        int data_size = ThreadBlock * m_per_thread;
+        #pragma unroll
+        for (int i = ThreadBlock / BlockDimX / 2; i > 0; i /= 2) {
+            for (int j = tid; j < data_size / 2; j += ThreadBlock) {
+                int reduce_group = j / i;
+                int reduce_index = j % i;
+                smem[reduce_index + reduce_group * stride] += smem[reduce_index + reduce_group * stride + i];
+            }
+            __syncthreads();
+            data_size /= 2;
+        }
+        for (int i = tid; i < this_group_elements; i += ThreadBlock) {
+            atomicAdd(dst + i + y * dstStride, (half)smem[i*stride]);
+        }
+        if constexpr(N > 1) {
+            if (y + 1 < N) {
+                __syncthreads();
+            }
+        }
+    }
+}
diff --git a/csrc/quantization/awq/hgemv_selector.hpp b/csrc/quantization/awq/hgemv_selector.hpp
new file mode 100644
index 000000000..b9d12a4e1
--- /dev/null
+++ b/csrc/quantization/awq/hgemv_selector.hpp
@@ -0,0 +1,287 @@
+#include <memory>
+#include <vector>
+#include <algorithm>
+#include "mc_runtime.h"
+#include "maca_fp16.h"
+
+struct KernelEventRecorder {
+    mcEvent_t _start;
+    mcEvent_t _stop;
+    float _eventMs = -1.f;
+
+    KernelEventRecorder() {
+        mcEventCreate(&_start);
+        mcEventCreate(&_stop);
+    }
+
+    ~KernelEventRecorder() {
+        mcEventDestroy(_start);
+        mcEventDestroy(_stop);
+    }
+
+    void start() {
+        mcEventRecord(_start, NULL);
+    }
+
+    float stop() {
+        mcEventRecord(_stop, NULL);
+        mcEventSynchronize(_stop);
+        mcEventElapsedTime(&_eventMs, _start, _stop);
+        return _eventMs;
+    }
+};
+namespace hgemv_selector {
+template<int quant_group=128, int pack_ratio=8, int m_per_thread=4>
+struct GemvParamAutoSelector {
+    int m;
+    int k;
+    int best_block_x = 0;
+    int best_split_k = 0;
+
+    std::pair<int,int> block_x_range;
+    std::pair<int,int> split_k_range;
+    bool _valid = false;
+
+private:
+    std::vector<std::pair<int, int>> param_candidates;
+    int warmup_iters = 0;
+    int current_block_x = 0;
+    int current_split_k = 0;
+    int current_perf_iter = 0;
+    std::vector<float> perf_times;
+    float kernel_best_time_ms_ave = 99999999.0f;
+    float best_band_width;
+    float data_size_gb;
+    std::shared_ptr<KernelEventRecorder> _r;
+    bool _selected = false;
+    const static int MAX_PERF_COUNT = 20;
+
+public:
+    GemvParamAutoSelector(int _m, int _n, int _k, bool allow_imcomplete_bx = false) : m(_m), k(_k)
+    {
+        if (!prepare_parameters(_m, _n, _k, block_x_range.second, split_k_range.second, allow_imcomplete_bx)) return;
+        block_x_range.first = block_x_range.second;
+        split_k_range.first = split_k_range.second;
+        for (int i = 0; i < 3; i++) if (block_x_range.first / 2 >= 16) block_x_range.first /= 2;
+        for (int i = 0; i < 3; i++) if (split_k_range.first % 2 == 0 && split_k_range.first / 2 > 0) split_k_range.first /= 2;
+        for (int i = block_x_range.first; i <= block_x_range.second; i *= 2) {
+            for (int j = split_k_range.first; j <= split_k_range.second; j *= 2) {
+                param_candidates.emplace_back(i, j);
+            }
+        }
+        if (split_k_range.second * quant_group != k) {
+            int max_split_k = k / quant_group;
+            if (max_split_k < 256) {
+                for (int i = block_x_range.first; i <= block_x_range.second; i *= 2) {
+                    param_candidates.emplace_back(i, max_split_k);
+                }
+            }
+        }
+
+        current_block_x = block_x_range.second;
+        current_split_k = split_k_range.second;
+        uint64_t data_size_in_bytes = _n * _k * sizeof(half);
+        data_size_in_bytes += (uint64_t)_m * _k / pack_ratio * sizeof(uint32_t);
+        data_size_in_bytes += (uint64_t)_m * _k / quant_group * sizeof(half);
+        data_size_in_bytes += (uint64_t)_m * _k / quant_group / pack_ratio * sizeof(uint32_t);
+        data_size_in_bytes += (uint64_t)_n * _m * sizeof(half);
+        data_size_gb = float(data_size_in_bytes) / 1024 / 1024 / 1024;
+        warmup_iters = 4;
+        _valid = true;
+    }
+
+    void select_in_warmup(const std::function<bool(int,int)> &f) {
+        if (!valid()) {
+            printf("Cannot run this selector!\n");
+            return;
+        }
+        if (_selected) {
+            f(best_block_x, best_split_k);
+            return;
+        };
+        //Warmup
+        for (auto iter = param_candidates.begin(); iter != param_candidates.end(); iter++) {
+            for (int i = 0; i < 5; i++) {
+                auto &p = *iter;
+                f(p.first, p.second);
+            }
+        }
+        _r.reset(new KernelEventRecorder());
+        kernel_best_time_ms_ave = 9999999.0f;
+        mcDeviceSynchronize();
+
+        for (auto iter = param_candidates.begin(); iter != param_candidates.end(); iter++) {
+            auto &p = *iter;
+            auto &bx = p.first;
+            auto &sk = p.second;
+            mcDeviceSynchronize();
+            _r->start();
+            bool launched = false;
+            for (int i = 0; i < MAX_PERF_COUNT; i++) {
+                launched = f(bx, sk);
+            }
+            auto ms = _r->stop();
+            if (launched && ms / MAX_PERF_COUNT < kernel_best_time_ms_ave) {
+                best_block_x = bx;
+                best_split_k = sk;
+                kernel_best_time_ms_ave = ms / MAX_PERF_COUNT;
+            }
+        }
+
+        _r.reset();
+        _selected = true;
+        warmup_iters = 0;
+    }
+
+    void run(const std::function<bool(int,int)> &f) {
+        if (!valid()) {
+            printf("Cannot run this selector!\n");
+            return;
+        }
+        if (warmup_iters > 0) {
+            f(current_block_x, current_split_k);
+            //printf("Calling warmup current_block_x = %d, current_split_k = %d\n", current_block_x, current_split_k);
+            if (current_block_x > block_x_range.first) current_block_x /= 2;
+            else if (current_split_k > split_k_range.first) current_split_k /= 2;
+            if (current_block_x == block_x_range.first && current_split_k >= split_k_range.first) {
+                current_block_x = block_x_range.second;
+                current_split_k = split_k_range.second;
+            }
+            warmup_iters--;
+            mcDeviceSynchronize();
+            return;
+        }
+
+        if (_selected) {
+            f(best_block_x, best_split_k);
+        } else {
+            if (!_r) {
+                _r = std::shared_ptr<KernelEventRecorder>(new KernelEventRecorder());
+                current_block_x = block_x_range.second;
+                current_split_k = split_k_range.second;
+                current_perf_iter = MAX_PERF_COUNT;
+            }
+            _r->start();
+            auto launched = f(current_block_x, current_split_k);
+            auto ms = _r->stop();
+            //printf("Launched with bx=%d,sk=%d,iter=%d,ms=%fms\n", current_block_x, current_split_k, current_perf_iter, ms);
+            if (!launched) {
+                printf("ERROR: Cannot run gemv with bx = %d, sk = %d\n", current_block_x, current_split_k);
+                return;
+            }
+            if (current_perf_iter-- > 0) {
+                perf_times.emplace_back(ms);
+                return;
+            }
+
+            std::sort(perf_times.begin(), perf_times.end());
+            float total_tm = 0;
+            for (size_t i = 1; i < perf_times.size() - 1; i++) {
+                total_tm += perf_times[i];
+            }
+
+            ms = total_tm /= (MAX_PERF_COUNT - 2);
+            perf_times.clear();
+            current_perf_iter = MAX_PERF_COUNT;
+            //printf("get ave time %fms\n", ms);
+
+            if (ms < kernel_best_time_ms_ave) {
+                best_block_x = current_block_x;
+                best_split_k = current_split_k;
+                kernel_best_time_ms_ave = ms;
+            }
+
+            if (current_split_k > split_k_range.first) {
+                current_split_k /= 2;
+            } else if (current_block_x > block_x_range.first){
+                current_split_k = split_k_range.second;
+                current_block_x /= 2;
+            } else if (current_block_x == block_x_range.first && current_split_k == split_k_range.first) {
+                _selected = true;
+                _r.reset();
+            } else {
+                printf("Error: all parameters are tried! current_block_x = %d, current_split_k = %d, block_x_range = %d,%d, split_k_range = %d,%d, best_block_x = %d, best_split_k = %d\n",
+                    current_block_x, current_split_k,
+                    block_x_range.first, block_x_range.second,
+                    split_k_range.first, split_k_range.second,
+                    best_block_x, best_split_k
+                );
+            }
+        }
+    }
+
+    bool valid() const { return _valid; }
+    bool selected() const {return _selected; }
+
+    float gemv_ave_time_us_cost() {
+        return kernel_best_time_ms_ave;
+    }
+
+    float gemv_bandwidth() {
+        return data_size_gb / (kernel_best_time_ms_ave / 1000.0f);
+    }
+
+    static bool prepare_parameters(int m, int n, int k, int &bx, int &sk, bool allow_imcomplete_bx = false) {
+        if (n > 4) return false;
+        if (k % quant_group != 0) return false;
+        if (m < 16 * m_per_thread) return false;
+        int max_split_k = k / quant_group;
+        int proper_splitk = 1;
+        while (proper_splitk*2 <= 256 && proper_splitk*2 <= max_split_k) proper_splitk *= 2;
+
+        int thread_block = 256;  //A thread block of 512 will wait for more warps when doing syncthreads
+
+        int proper_bx = 16;
+        if (m % (proper_bx * m_per_thread) != 0) return false;
+        //proper_bx * m_per_threads corresponds to group_elements, this value should not be larger than 2048 as limit of shared memory
+        while (proper_bx * 2 * m_per_thread && (m % (proper_bx*2*m_per_thread) == 0) && proper_bx * 2 <= 256 && proper_bx * 2 * m_per_thread <= 2048) proper_bx*=2;
+        if (allow_imcomplete_bx) {
+            int may_proper_bx = proper_bx * 2;
+            if (may_proper_bx <= 256 && may_proper_bx * m_per_thread <= 2048 && (m % (may_proper_bx*m_per_thread)) > may_proper_bx * m_per_thread / 4) {
+                proper_bx = may_proper_bx;
+            }
+        }
+
+        bx = proper_bx;
+        sk = proper_splitk;
+        //printf("m=%d,n=%d,k=%d,get bx=%d,sk=%d,allow_imcomplete_bx=%d\n", m, n, k, bx, sk, allow_imcomplete_bx);
+        return true;
+    }
+};
+
+template<int quant_group=128, int pack_ratio=8, int m_per_thread=4>
+class GemvSelectorHolder {
+private:
+    std::vector<GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread>> _selectors;
+    static std::shared_ptr<GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>> _holder;
+    static GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread> _invalid_selector;
+
+public:
+    static GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread>& selector(int m, int n, int k, bool allow_imcomplete_bx = false) {
+        if (!GemvSelectorHolder::_holder) {
+            GemvSelectorHolder::_holder.reset(new GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>());
+        }
+        int bx, sk;
+        if (!GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread>::prepare_parameters(m, n, k, bx, sk, allow_imcomplete_bx)) {
+            return _invalid_selector;
+        }
+        auto iter = std::find_if(_holder->_selectors.begin(), _holder->_selectors.end(),
+            [&](const GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread>& p) -> bool {
+                return p.m == m && p.k == k;
+            });
+        if (iter != _holder->_selectors.end()) return *iter;
+        GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread> sl(m, n, k, allow_imcomplete_bx);
+        if (!sl.valid()) {
+            return _invalid_selector;
+        }
+        _holder->_selectors.emplace_back(sl);
+        return _holder->_selectors.back();
+    }
+};
+
+template<int quant_group, int pack_ratio, int m_per_thread>
+std::shared_ptr<GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>> GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>::_holder;
+
+template<int quant_group, int pack_ratio, int m_per_thread>
+GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread> GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>::_invalid_selector(0,8,0);
+}
\ No newline at end of file
diff --git a/csrc/quantization/compressed_tensors/int8_quant_kernels.cu b/csrc/quantization/compressed_tensors/int8_quant_kernels.cu
index bf46cce60..5fdd5037f 100644
--- a/csrc/quantization/compressed_tensors/int8_quant_kernels.cu
+++ b/csrc/quantization/compressed_tensors/int8_quant_kernels.cu
@@ -12,7 +12,7 @@
   #include <hipcub/hipcub.hpp>
 #endif
 
-static inline __device__ int8_t float_to_int8_rn(float x) {
+static __forceinline__ __device__ int8_t float_to_int8_rn(float x) {
 #ifdef USE_ROCM
   static constexpr auto i8_min =
       static_cast<float>(std::numeric_limits<int8_t>::min());
@@ -36,8 +36,13 @@ static inline __device__ int8_t float_to_int8_rn(float x) {
   return static_cast<int8_t>(dst);
 #else
   // CUDA path
-  uint32_t dst;
-  asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
+  //uint32_t dst;
+  //asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
+  //return reinterpret_cast<const int8_t&>(dst);
+  int32_t dst;
+  dst = __float2int_rn(x);
+  dst = min(dst, 127);
+  dst = max(dst, -127);
   return reinterpret_cast<const int8_t&>(dst);
 #endif
 }
@@ -71,9 +76,13 @@ static inline __device__ int32_t float_to_int32_rn(float x) {
   return static_cast<int32_t>(dst);
 #else
   // CUDA path
-  uint32_t dst;
-  asm volatile("cvt.rni.sat.s32.f32 %0, %1;" : "=r"(dst) : "f"(x));
-  return reinterpret_cast<const int32_t&>(dst);
+  static constexpr auto i32_min = std::numeric_limits<int32_t>::min();
+  static constexpr auto i32_min_f = static_cast<float>(i32_min);
+  static constexpr auto i32_max = std::numeric_limits<int32_t>::max();
+  static constexpr auto i32_max_f = static_cast<float>(i32_max);
+  x = min(x, i32_max_f);
+  x = max(x, i32_min_f);
+  return __float2int_rn(x);
 #endif
 }
 
@@ -95,9 +104,14 @@ static inline __device__ int8_t int32_to_int8(int32_t x) {
   return static_cast<int8_t>(dst);
 #else
   // CUDA path
-  uint32_t dst;
-  asm volatile("cvt.sat.s8.s32 %0, %1;" : "=r"(dst) : "r"(x));
-  return reinterpret_cast<const int8_t&>(dst);
+  static constexpr auto i8_min =
+      static_cast<int32_t>(std::numeric_limits<int8_t>::min());
+  static constexpr auto i8_max =
+      static_cast<int32_t>(std::numeric_limits<int8_t>::max());
+
+  // saturate
+  int32_t dst = std::clamp(x, i8_min, i8_max);
+  return static_cast<int8_t>(dst);
 #endif
 }
 
@@ -141,21 +155,23 @@ __global__ void static_scaled_int8_azp_quant_kernel(
   }
 }
 
-template <typename scalar_t, typename scale_type>
+template <typename scalar_t, typename scale_type, bool WITHMASK>
 __global__ void dynamic_scaled_int8_quant_kernel(
     scalar_t const* __restrict__ input, int8_t* __restrict__ out,
-    scale_type* scale, const int hidden_size) {
+    scale_type* scale, const int hidden_size, const int num_tokens, int *mask_buffer = NULL) {
+  if constexpr(WITHMASK) {
+    __shared__ int sm_max_token;
+    if(threadIdx.x == 0) sm_max_token = mask_buffer[blockIdx.y]; 
+    __syncthreads();
+    if(blockIdx.x >= sm_max_token) return;
+  }
   int const tid = threadIdx.x;
-  int64_t const token_idx = blockIdx.x;
+  int64_t const token_idx = blockIdx.y * num_tokens + blockIdx.x;
   float absmax_val = 0.0f;
   float const zero = 0.0f;
 
-  // Must be performed using 64-bit math to avoid integer overflow.
-  out += token_idx * hidden_size;
-  input += token_idx * hidden_size;
-
   for (int i = tid; i < hidden_size; i += blockDim.x) {
-    float val = static_cast<float>(input[i]);
+    float val = static_cast<float>(input[token_idx * hidden_size + i]);
     val = val > zero ? val : -val;
     absmax_val = val > absmax_val ? val : absmax_val;
   }
@@ -167,31 +183,330 @@ __global__ void dynamic_scaled_int8_quant_kernel(
   __shared__ float block_absmax_val;
   if (tid == 0) {
     block_absmax_val = block_absmax_val_maybe;
-    scale[token_idx] = block_absmax_val / 127.0f;
+    scale[token_idx] = block_absmax_val * 0.0078740157;
   }
   __syncthreads();
 
-  float const tmp_scale = 127.0f / block_absmax_val;
+  float const tmp_scale = 127.0f *__builtin_mxc_rcpf(block_absmax_val);
   for (int i = tid; i < hidden_size; i += blockDim.x) {
-    out[i] = float_to_int8_rn(static_cast<float>(input[i]) * tmp_scale);
+    out[token_idx * hidden_size + i] = float_to_int8_rn(
+        static_cast<float>(input[token_idx * hidden_size + i]) * tmp_scale);
   }
 }
 
-template <typename scalar_t, typename scale_type, typename azp_type>
-__global__ void dynamic_scaled_int8_azp_quant_kernel(
+template <typename scalar_t, typename scale_type, typename VT, typename VT1, int NUM_THREADS, bool WITHMASK>
+__global__ void dynamic_scaled_int8_quant_kernel_sreg_opt(
     scalar_t const* __restrict__ input, int8_t* __restrict__ out,
-    scale_type* scale, azp_type* azp, const int hidden_size) {
-  int64_t const token_idx = blockIdx.x;
+    scale_type* scale, const int hidden_size, int num_tokens, int* mask_buffer=NULL) {
+  if constexpr(WITHMASK) {
+    __shared__ int sm_max_token;
+    if(threadIdx.x == 0) sm_max_token = mask_buffer[blockIdx.y]; 
+    __syncthreads();
+    if(blockIdx.x >= sm_max_token) return;
+  }
+  int const tid = threadIdx.x;
+  int64_t const token_idx = blockIdx.y * num_tokens + blockIdx.x;
+  float absmax_val = 0.0f;
+  float const zero = 0.0f;
+  constexpr int N = sizeof(VT) / sizeof(scalar_t);
+  float reg_src0[N];
+  scalar_t const* ptr_input = input + token_idx * hidden_size;
+  int reg_length = NUM_THREADS * N;
+  int length = min(hidden_size, reg_length);
+  int index = tid * N;
+  if(index < length) {
+    VT reg_src;
+    reg_src = *(VT*)(ptr_input + index);
+    scalar_t* ptr_reg_src = (scalar_t*)&reg_src;
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+      reg_src0[i] = (float)ptr_reg_src[i];
+    }
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+      float val = abs(reg_src0[i]);
+      absmax_val = max(absmax_val, val);
+    }
+  }
 
-  // Must be performed using 64-bit math to avoid integer overflow.
-  out += token_idx * hidden_size;
-  input += token_idx * hidden_size;
+  using BlockReduce = cub::BlockReduce<float, NUM_THREADS>;
+  __shared__ typename BlockReduce::TempStorage reduceStorage;
+  float const block_absmax_val_maybe =
+      BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, NUM_THREADS);
+  __shared__ float block_absmax_val;
+  if (tid == 0) {
+    block_absmax_val = block_absmax_val_maybe;
+    scale[token_idx] = static_cast<scale_type>(block_absmax_val_maybe * 0.0078740157);
+  }
+  __syncthreads();
+  float const tmp_scale = 127.0f * __builtin_mxc_rcpf(block_absmax_val);
+  int8_t* ptr_output = out + token_idx * hidden_size;
+  if(index < length) {
+    VT1 vdst;
+    int8_t* ptr_reg = (int8_t*)&vdst;
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+      ptr_reg[i] = float_to_int8_rn(reg_src0[i] * tmp_scale);
+    }
+    *(VT1*)(ptr_output + index) = vdst;
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename VT, typename VT1, bool WITHMASK>
+__global__ void dynamic_scaled_int8_quant_kernel_reg_opt(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, int blockDim_x, int num_tokens, int* mask_buffer=NULL) {
+  if constexpr(WITHMASK) {
+    __shared__ int sm_max_token;
+    if(threadIdx.x == 0) sm_max_token = mask_buffer[blockIdx.y]; 
+    __syncthreads();
+    if(blockIdx.x >= sm_max_token) return;
+  }
+  int const tid = threadIdx.x;
+  int64_t const token_idx = blockIdx.y * num_tokens + blockIdx.x;
+  float absmax_val = 0.0f;
+  float const zero = 0.0f;
+  constexpr int N = sizeof(VT) / sizeof(scalar_t);
+  float reg_src0[N];
+  float reg_src1[N];
+  scalar_t const* ptr_input = input + token_idx * hidden_size;
+  int reg_length = 2 * blockDim_x * N;
+  int length = min(hidden_size, reg_length);
+  int index = 2 * tid * N;
+  if(index < length) {
+    VT reg_src = *(VT*)(ptr_input + index);
+    scalar_t* ptr_reg_src = (scalar_t*)&reg_src;
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+      reg_src0[i] = (float)ptr_reg_src[i];
+    }
+    reg_src = *(VT*)(ptr_input + index + N);
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+      reg_src1[i] = (float)ptr_reg_src[i];
+    }
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+      float val = abs(reg_src0[i]);
+      absmax_val =  max(val, absmax_val);
+      val = abs(reg_src1[i]);
+      absmax_val = max(val, absmax_val);
+    }
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 512>;
+  __shared__ typename BlockReduce::TempStorage reduceStorage;
+  float const block_absmax_val_maybe =
+      BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim_x);
+  __shared__ float block_absmax_val;
+  if (tid == 0) {
+    block_absmax_val = (block_absmax_val_maybe);
+    scale[token_idx] = static_cast<scale_type>(block_absmax_val * 0.0078740157);
+  }
+  __syncthreads();
+  float const tmp_scale = 127.0f * __builtin_mxc_rcpf(block_absmax_val);
+  int8_t* ptr_output = out + token_idx * hidden_size;
+  if(index < length) {
+    VT1 vdst;
+    int8_t* ptr_reg = (int8_t*)&vdst;
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+      ptr_reg[i] = float_to_int8_rn(
+             reg_src0[i] * tmp_scale);
+    }
+    ptr_reg = ptr_reg + N;
+    #pragma unroll N
+    for(int i = 0; i < N; i++) {
+      ptr_reg[i] = float_to_int8_rn(
+            reg_src1[i] * tmp_scale);
+    }
+    *(VT1*)(ptr_output + index) = vdst;
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename VT, typename VT1, int NUM_REG, int NUM_THREADS, bool WITHMASK>
+__global__ __launch_bounds__(1024) void dynamic_scaled_int8_quant_kernel_lh_opt(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, int num_tokens, int* mask_buffer=NULL) {
+  if constexpr(WITHMASK) {
+    __shared__ int sm_max_token;
+    if(threadIdx.x == 0) sm_max_token = mask_buffer[blockIdx.y]; 
+    __syncthreads();
+    if(blockIdx.x >= sm_max_token) return;
+  }
+  int const tid = threadIdx.x;
+  int64_t const token_idx = blockIdx.y * num_tokens + blockIdx.x;
+  float absmax_val = 0.0f;
+  float const zero = 0.0f;
+  constexpr int N = sizeof(VT) / sizeof(scalar_t);
+  int stride = NUM_THREADS * N;
+  float reg_src[NUM_REG][N];
+  scalar_t const* ptr_input = input + token_idx * hidden_size;
+  for(int i = tid * N, k = 0; i < hidden_size; i += stride, k++) {
+    VT vsrc = *(VT*)(ptr_input + i);
+    scalar_t *ptr_src = (scalar_t*)&vsrc;
+    #pragma unroll N
+    for(int j = 0; j < N; j++) {
+        float val = static_cast<float>(ptr_src[j]);
+        reg_src[k][j] = val;
+        val = abs(val);
+        absmax_val = max(val, absmax_val);
+    }
+  }
+  using BlockReduce = cub::BlockReduce<float, NUM_THREADS>;
+  __shared__ typename BlockReduce::TempStorage reduceStorage;
+  float const block_absmax_val_maybe =
+      BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, NUM_THREADS);
+  __shared__ float block_absmax_val;
+  if (tid == 0) {
+    block_absmax_val = block_absmax_val_maybe;
+    scale[token_idx] = block_absmax_val * 0.0078740157;
+  }
+  
+  __syncthreads();
+
+  float const tmp_scale = 127.0f * __builtin_mxc_rcpf(block_absmax_val);
+  int8_t* ptr_output = out + token_idx * hidden_size;
+  for(int i = tid * N, k = 0; i < hidden_size; i += stride, k++) {
+    VT1 vdst;
+    int8_t* ptr_reg = (int8_t*)&vdst;
+    #pragma unroll N
+    for(int j = 0; j < N; j++) {
+        ptr_reg[j] = float_to_int8_rn(
+            reg_src[k][j] * tmp_scale);
+    }
+    *(VT1*)(ptr_output + i) = vdst;
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename VT, typename VT1, bool WITHMASK>
+__global__ void dynamic_scaled_int8_quant_kernel_sm_opt(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, int blockDim_x, int num_tokens, int* mask_buffer=NULL) {
+  if constexpr(WITHMASK) {
+    __shared__ int sm_max_token;
+    if(threadIdx.x == 0) sm_max_token = mask_buffer[blockIdx.y]; 
+    __syncthreads();
+    if(blockIdx.x >= sm_max_token) return;
+  }
+  int const tid = threadIdx.x;
+  int64_t const token_idx = blockIdx.y * num_tokens + blockIdx.x;
+  float absmax_val = 0.0f;
+  float const zero = 0.0f;
+  constexpr int N = sizeof(VT) / sizeof(scalar_t);
+  int stride = blockDim_x * N;
+  __shared__ float sm_buffer[8064];
+  scalar_t const* ptr_input = input + token_idx * hidden_size;
+  for(int i = tid * N; i < hidden_size; i += stride) {
+    VT vsrc = *(VT*)(ptr_input + i);
+    scalar_t *ptr_src = (scalar_t*)&vsrc;
+    float* ptr_sm_buffer = sm_buffer + i;
+    #pragma unroll N
+    for(int j = 0; j < N; j++) {
+        float val = static_cast<float>(ptr_src[j]);
+        ptr_sm_buffer[j] = val;
+        val = abs(val);
+        absmax_val = max(val, absmax_val);
+    }
+  }
+  using BlockReduce = cub::BlockReduce<float, 512>;
+  __shared__ typename BlockReduce::TempStorage reduceStorage;
+  float const block_absmax_val_maybe =
+      BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
+  __shared__ float block_absmax_val;
+  if (tid == 0) {
+    block_absmax_val = block_absmax_val_maybe;
+    scale[token_idx] = block_absmax_val * 0.0078740157;
+  }
+  
+  __syncthreads();
 
+  float const tmp_scale = 127.0f *__builtin_mxc_rcpf(block_absmax_val);
+  int8_t* ptr_output = out + token_idx * hidden_size;
+  for(int i = tid * N; i < hidden_size; i += stride) {
+    VT1 vdst;
+    int8_t* ptr_reg = (int8_t*)&vdst;
+    float* ptr_sm_buffer = sm_buffer + i;
+    #pragma unroll N
+    for(int j = 0; j < N; j++) {
+        ptr_reg[j] = float_to_int8_rn(
+            ptr_sm_buffer[j] * tmp_scale);
+    }
+    *(VT1*)(ptr_output + i) = vdst;
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename VT, typename VT1, bool WITHMASK>
+__launch_bounds__(1024) __global__ void dynamic_scaled_int8_quant_kernel_opt(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, const int blockDim_x, const int num_tokens, int* mask_buffer=NULL) {
+  if constexpr(WITHMASK) {
+    __shared__ int sm_max_token;
+    if(threadIdx.x == 0) sm_max_token = mask_buffer[blockIdx.y]; 
+    __syncthreads();
+    if(blockIdx.x >= sm_max_token) return;
+  }
+  constexpr int N = sizeof(VT) / sizeof(scalar_t);
+  int const tid = threadIdx.x * N;
+  int64_t const token_idx = blockIdx.y * num_tokens + blockIdx.x;
+  float absmax_val = 0.0f;
+  int stride = blockDim_x * N;
+  const scalar_t * ptr_input = input + token_idx * hidden_size;
+
+  for (int i = tid ; i < hidden_size; i += stride) {
+    VT vsrc = *(VT*)(ptr_input + i);
+    scalar_t *ptr_src = (scalar_t*)&vsrc;
+    #pragma unroll N
+    for(int j = 0; j < N; j++) {
+        float val = static_cast<float>(ptr_src[j]);
+        val = abs(val);
+        absmax_val = max(val, absmax_val);
+    }
+  }
+
+    using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStorage;
+  float const block_absmax_val_maybe =
+      BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
+  __shared__ float block_absmax_val;
+  if (tid == 0) {
+    block_absmax_val = block_absmax_val_maybe;
+    scale[token_idx] = block_absmax_val * 0.0078740157;
+  }
+  __syncthreads();
+
+  float const tmp_scale = 127.0f *__builtin_mxc_rcpf(block_absmax_val);
+  int8_t* ptr_output = out + token_idx * hidden_size;
+  for (int i = tid; i < hidden_size; i += stride) {
+    VT vsrc = *(VT*)(ptr_input + i);
+    VT1 vdst;
+    scalar_t *ptr_src = (scalar_t*)&vsrc;
+    int8_t* ptr_dst = (int8_t*)&vdst;
+    #pragma unroll N
+    for(int j = 0; j < N; ++j) {
+        ptr_dst[j] = float_to_int8_rn(
+        static_cast<float>(ptr_src[j]) * tmp_scale);
+    }
+    *(VT1*)(ptr_output + i) = vdst;
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename azp_type, bool WITHMASK>
+__global__ void dynamic_scaled_int8_azp_quant_kernel(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, azp_type* azp, const int hidden_size, const int num_tokens, int* mask_buffer=NULL) {
+  if constexpr(WITHMASK) {
+    __shared__ int sm_max_token;
+    if(threadIdx.x == 0) sm_max_token = mask_buffer[blockIdx.y]; 
+    __syncthreads();
+    if(blockIdx.x >= sm_max_token) return;
+  }  
+  int64_t const token_idx = blockIdx.y * num_tokens + blockIdx.x;
   // Scan for the min and max value for this token
   float max_val = std::numeric_limits<float>::min();
   float min_val = std::numeric_limits<float>::max();
   for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
-    auto val = static_cast<float>(input[i]);
+    auto val = static_cast<float>(input[token_idx * hidden_size + i]);
     max_val = std::max(max_val, val);
     min_val = std::min(min_val, val);
   }
@@ -226,11 +541,830 @@ __global__ void dynamic_scaled_int8_azp_quant_kernel(
 
   // Quantize the values
   for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
-    auto const val = static_cast<float>(input[i]);
+    auto const val = static_cast<float>(input[token_idx * hidden_size + i]);
     auto const quant_val =
         int32_to_int8(float_to_int32_rn(val / scale_val) + azp_val);
-    out[i] = quant_val;
+    out[token_idx * hidden_size + i] = quant_val;
+  }
+}
+
+template <typename scalar_t, typename scale_type>
+__global__ void dynamic_scaled_int8_quant_mask_kernel(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, const int num_tokens, const int mask_size, const int grid_size, int * mask) {
+  int const tid = threadIdx.x;
+  __shared__ int sm_mask[16];
+  __shared__ int sm_stride[16];
+  if(tid < mask_size) {
+    sm_mask[tid] = mask[tid];
+  }
+  __syncthreads();
+  if(tid < mask_size) {
+    int tmp = 0;
+    for(int i = 0; i < tid; i++) {
+      tmp += sm_mask[i];
+    }
+    sm_stride[tid] = tmp;
+  }
+  __syncthreads();
+  int total_tokens = sm_stride[mask_size - 1] + sm_mask[mask_size - 1];
+  for(int idx = blockIdx.x; idx < total_tokens; idx += grid_size) {
+    int token_id = mask_size - 1;
+    while(idx < sm_stride[token_id]) {
+      token_id--;
+    }
+    int64_t const token_idx = token_id * num_tokens + idx - sm_stride[token_id];
+    float absmax_val = 0.0f;
+    float const zero = 0.0f;
+
+    for (int i = tid; i < hidden_size; i += blockDim.x) {
+      float val = static_cast<float>(input[token_idx * hidden_size + i]);
+      val = val > zero ? val : -val;
+      absmax_val = val > absmax_val ? val : absmax_val;
+    }
+
+    using BlockReduce = cub::BlockReduce<float, 1024>;
+    __shared__ typename BlockReduce::TempStorage reduceStorage;
+    float const block_absmax_val_maybe =
+        BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
+    __shared__ float block_absmax_val;
+    if (tid == 0) {
+      block_absmax_val = block_absmax_val_maybe;
+      scale[token_idx] = block_absmax_val * 0.0078740157;
+    }
+    __syncthreads();
+
+    float const tmp_scale = 127.0f *__builtin_mxc_rcpf(block_absmax_val);
+    for (int i = tid; i < hidden_size; i += blockDim.x) {
+      out[token_idx * hidden_size + i] = float_to_int8_rn(
+          static_cast<float>(input[token_idx * hidden_size + i]) * tmp_scale);
+    }
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename VT, typename VT1>
+__global__ void dynamic_scaled_int8_quant_mask_kernel_sreg_opt(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, int blockDim_x, int num_tokens, int mask_size, int grid_size, int* mask=NULL) {
+  int const tid = threadIdx.x;
+  __shared__ int sm_mask[16];
+  __shared__ int sm_stride[16];
+  if(tid < mask_size) {
+    sm_mask[tid] = mask[tid];
+  }
+  __syncthreads();
+  if(tid < mask_size) {
+    int tmp = 0;
+    for(int i = 0; i < tid; i++) {
+      tmp += sm_mask[i];
+    }
+    sm_stride[tid] = tmp;
+  }
+  __syncthreads();
+  int total_tokens = sm_stride[mask_size - 1] + sm_mask[mask_size - 1];
+  for(int idx = blockIdx.x; idx < total_tokens; idx += grid_size) {
+      int token_id = mask_size - 1;
+      while(idx < sm_stride[token_id]) {
+        token_id--;
+      }
+      int64_t const token_idx = token_id * num_tokens + idx - sm_stride[token_id];
+      scalar_t absmax_val = static_cast<scalar_t>(0.0f);
+      float const zero = 0.0f;
+      constexpr int N = sizeof(VT) / sizeof(scalar_t);
+      scalar_t reg_src0[N];
+      scalar_t const* ptr_input = input + token_idx * hidden_size;
+      int reg_length = blockDim_x * N;
+      int length = min(hidden_size, reg_length);
+      int index = tid * N;
+      if(index < length) {
+        *(VT*)reg_src0 = *(VT*)(ptr_input + index);
+        #pragma unroll N
+        for(int i = 0; i < N; i++) {
+          scalar_t val = abs(reg_src0[i]);
+          absmax_val = max(absmax_val, val);
+        }
+      }
+
+      using BlockReduce = cub::BlockReduce<scalar_t, 512>;
+      __shared__ typename BlockReduce::TempStorage reduceStorage;
+      float const block_absmax_val_maybe =
+          BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim_x);
+      __shared__ scale_type block_absmax_val;
+      if (tid == 0) {
+        block_absmax_val = static_cast<scale_type>(block_absmax_val_maybe);
+        scale[token_idx] = static_cast<scale_type>(block_absmax_val * 0.0078740157);
+      }
+      __syncthreads();
+      float const tmp_scale = 127.0f *__builtin_mxc_rcpf(block_absmax_val);
+      int8_t* ptr_output = out + token_idx * hidden_size;
+      if(index < length) {
+        VT1 vdst;
+        int8_t* ptr_reg = (int8_t*)&vdst;
+        #pragma unroll N
+        for(int i = 0; i < N; i++) {
+          ptr_reg[i] = float_to_int8_rn(
+                static_cast<float>(reg_src0[i]) * tmp_scale);
+        }
+        *(VT1*)(ptr_output + index) = vdst;
+      }
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename VT, typename VT1>
+__global__ void dynamic_scaled_int8_quant_mask_kernel_reg_opt(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, int blockDim_x, int num_tokens, int mask_size, int grid_size, int* mask=NULL) {
+  int const tid = threadIdx.x;
+  __shared__ int sm_mask[16];
+  __shared__ int sm_stride[16];
+  if(tid < mask_size) {
+    sm_mask[tid] = mask[tid];
   }
+  __syncthreads();
+  if(tid < mask_size) {
+    int tmp = 0;
+    for(int i = 0; i < tid; i++) {
+      tmp += sm_mask[i];
+    }
+    sm_stride[tid] = tmp;
+  }
+  __syncthreads();
+  int total_tokens = sm_stride[mask_size - 1] + sm_mask[mask_size - 1];
+  for(int idx = blockIdx.x; idx < total_tokens; idx += grid_size) 
+  {
+    int token_id = mask_size - 1;
+    while(idx < sm_stride[token_id]) {
+      token_id--;
+    }
+    int64_t const token_idx = token_id * num_tokens + idx - sm_stride[token_id];
+    scalar_t absmax_val = static_cast<scalar_t>(0.0f);
+    float const zero = 0.0f;
+    constexpr int N = sizeof(VT) / sizeof(scalar_t);
+    scalar_t reg_src0[N];
+    scalar_t reg_src1[N];
+    scalar_t const* ptr_input = input + token_idx * hidden_size;
+    int reg_length = 2 * blockDim_x * N;
+    int length = min(hidden_size, reg_length);
+    int index = 2 * tid * N;
+    if(index < length) {
+      *(VT*)reg_src0 = *(VT*)(ptr_input + index);
+      *(VT*)reg_src1 = *(VT*)(ptr_input + index + N);
+      #pragma unroll N
+      for(int i = 0; i < N; i++) {
+        scalar_t val = abs(reg_src0[i]);
+        absmax_val =  max(val, absmax_val);
+        val = abs(reg_src1[i]);
+        absmax_val = max(val, absmax_val);
+      }
+    }
+
+    using BlockReduce = cub::BlockReduce<scalar_t, 512>;
+    __shared__ typename BlockReduce::TempStorage reduceStorage;
+    scalar_t const block_absmax_val_maybe =
+        BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim_x);
+    __shared__ scale_type block_absmax_val;
+    if (tid == 0) {
+      block_absmax_val = static_cast<scale_type>(block_absmax_val_maybe);
+      scale[token_idx] = block_absmax_val * 0.0078740157;
+    }
+    __syncthreads();
+    float const tmp_scale = 127.0f *__builtin_mxc_rcpf(block_absmax_val);
+    int8_t* ptr_output = out + token_idx * hidden_size;
+    if(index < length) {
+      VT1 vdst;
+      int8_t* ptr_reg = (int8_t*)&vdst;
+      constexpr int ON = 2 * N;
+      #pragma unroll N
+      for(int i = 0; i < N; i++) {
+        ptr_reg[i] = float_to_int8_rn(
+              static_cast<float>(reg_src0[i]) * tmp_scale);
+      }
+      ptr_reg = ptr_reg + N;
+      #pragma unroll N
+      for(int i = 0; i < N; i++) {
+        ptr_reg[i] = float_to_int8_rn(
+              static_cast<float>(reg_src1[i]) * tmp_scale);
+      }
+      *(VT1*)(ptr_output + index) = vdst;
+    }
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename VT, typename VT1>
+__global__ void dynamic_scaled_int8_quant_mask_kernel_sm_opt(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, int blockDim_x, int num_tokens, int mask_size, int grid_size, int* mask=NULL) {
+  int const tid = threadIdx.x;
+  __shared__ int sm_mask[16];
+  __shared__ int sm_stride[16];
+  if(tid < mask_size) {
+    sm_mask[tid] = mask[tid];
+  }
+  __syncthreads();
+  if(tid < mask_size) {
+    int tmp = 0;
+    for(int i = 0; i < tid; i++) {
+      tmp += sm_mask[i];
+    }
+    sm_stride[tid] = tmp;
+  }
+  __syncthreads();
+  int total_tokens = sm_stride[mask_size - 1] + sm_mask[mask_size - 1];
+  for(int idx = blockIdx.x; idx < total_tokens; idx += grid_size) {
+    int token_id = mask_size - 1;
+    while(idx < sm_stride[token_id]) {
+      token_id--;
+    }
+    int64_t const token_idx = token_id * num_tokens + idx - sm_stride[token_id];
+    float absmax_val = 0.0f;
+    float const zero = 0.0f;
+    constexpr int N = sizeof(VT) / sizeof(scalar_t);
+    int stride = blockDim_x * N;
+    __shared__ float sm_buffer[8064];
+    scalar_t const* ptr_input = input + token_idx * hidden_size;
+    for(int i = tid * N; i < hidden_size; i += stride) {
+      VT vsrc = *(VT*)(ptr_input + i);
+      scalar_t *ptr_src = (scalar_t*)&vsrc;
+      float* ptr_sm_buffer = sm_buffer + i;
+      #pragma unroll N
+      for(int j = 0; j < N; j++) {
+          float val = static_cast<float>(ptr_src[j]);
+          ptr_sm_buffer[j] = val;
+          val = val > zero ? val : -val;
+          absmax_val = val > absmax_val ? val : absmax_val;
+      }
+    }
+    using BlockReduce = cub::BlockReduce<float, 512>;
+    __shared__ typename BlockReduce::TempStorage reduceStorage;
+    float const block_absmax_val_maybe =
+        BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
+    __shared__ float block_absmax_val;
+    if (tid == 0) {
+      block_absmax_val = block_absmax_val_maybe;
+      scale[token_idx] = block_absmax_val * 0.0078740157;
+    }
+    
+    __syncthreads();
+
+    float const tmp_scale = 127.0f *__builtin_mxc_rcpf(block_absmax_val);
+    int8_t* ptr_output = out + token_idx * hidden_size;
+    for(int i = tid * N; i < hidden_size; i += stride) {
+      VT1 vdst;
+      int8_t* ptr_reg = (int8_t*)&vdst;
+      float* ptr_sm_buffer = sm_buffer + i;
+      #pragma unroll N
+      for(int j = 0; j < N; j++) {
+          ptr_reg[j] = float_to_int8_rn(
+              ptr_sm_buffer[j] * tmp_scale);
+      }
+      *(VT1*)(ptr_output + i) = vdst;
+    }
+  }
+}
+
+template <typename scalar_t, typename scale_type, typename VT, typename VT1>
+__launch_bounds__(1024) __global__ void dynamic_scaled_int8_quant_mask_kernel_opt(
+    scalar_t const* __restrict__ input, int8_t* __restrict__ out,
+    scale_type* scale, const int hidden_size, const int blockDim_x, const int num_tokens, int mask_size, int grid_size, int* mask=NULL) {
+  constexpr int N = sizeof(VT) / sizeof(scalar_t);
+  int const tid = threadIdx.x * N;
+  __shared__ int sm_mask[16];
+  __shared__ int sm_stride[16];
+  if(threadIdx.x < mask_size) {
+    sm_mask[threadIdx.x] = mask[threadIdx.x];
+  }
+  __syncthreads();
+  if(threadIdx.x < mask_size) {
+    int tmp = 0;
+    for(int i = 0; i < threadIdx.x; i++) {
+      tmp += sm_mask[i];
+    }
+    sm_stride[threadIdx.x] = tmp;
+  }
+  __syncthreads();
+  int total_tokens = sm_stride[mask_size - 1] + sm_mask[mask_size - 1];
+  
+  for(int idx = blockIdx.x; idx < total_tokens; idx += grid_size) {
+    int token_id = mask_size - 1;
+    while(idx < sm_stride[token_id]) {
+      token_id--;
+    }
+    int64_t const token_idx = token_id * num_tokens + idx - sm_stride[token_id];
+    float absmax_val = 0.0f;
+    int stride = blockDim_x * N;
+    const scalar_t * ptr_input = input + token_idx * hidden_size;
+
+    for (int i = tid ; i < hidden_size; i += stride) {
+      VT vsrc = *(VT*)(ptr_input + i);
+      scalar_t *ptr_src = (scalar_t*)&vsrc;
+      #pragma unroll N
+      for(int j = 0; j < N; j++) {
+          float val = static_cast<float>(ptr_src[j]);
+          val = val > 0 ? val : -val;
+          absmax_val = val > absmax_val ? val : absmax_val;
+      }
+    }
+
+      using BlockReduce = cub::BlockReduce<float, 1024>;
+    __shared__ typename BlockReduce::TempStorage reduceStorage;
+    float const block_absmax_val_maybe =
+        BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
+    __shared__ float block_absmax_val;
+    if (tid == 0) {
+      block_absmax_val = block_absmax_val_maybe;
+      scale[token_idx] = block_absmax_val * 0.0078740157;
+    }
+    __syncthreads();
+
+    float const tmp_scale = 127.0f *__builtin_mxc_rcpf(block_absmax_val);
+    int8_t* ptr_output = out + token_idx * hidden_size;
+    for (int i = tid; i < hidden_size; i += stride) {
+      VT vsrc = *(VT*)(ptr_input + i);
+      VT1 vdst;
+      scalar_t *ptr_src = (scalar_t*)&vsrc;
+      int8_t* ptr_dst = (int8_t*)&vdst;
+      #pragma unroll N
+      for(int j = 0; j < N; ++j) {
+          ptr_dst[j] = float_to_int8_rn(
+          static_cast<float>(ptr_src[j]) * tmp_scale);
+      }
+      *(VT1*)(ptr_output + i) = vdst;
+    }
+  }
+}
+
+template<typename T, typename T1, typename VT, typename VT1, int NUM_VT> 
+__global__ void silu_and_mul_mask_quant_pack(T* input, T* output,T1* mask, int mask_size, int64_t grid_size, int64_t num_tokens, int64_t hidden_size, int64_t out_stirde, int blockDim_x)
+{
+    constexpr int N = sizeof(VT) / sizeof(T);
+    int const tid = threadIdx.x;
+    __shared__ T1 sm_mask[16];
+    __shared__ T1 sm_stride[16];
+    if(tid < mask_size) {
+        sm_mask[tid] = mask[tid];
+    }
+
+    int64_t hidden_size2 = hidden_size << 1;
+    __syncthreads();
+    if(tid < mask_size) {
+        T1 tmp = 0;
+        for(int i = 0; i < tid; i++) {
+        tmp += sm_mask[i];
+        }
+        sm_stride[tid] = tmp;
+    }
+    int stride = blockDim_x * N;
+    __syncthreads();
+    int64_t total_tokens = sm_stride[mask_size - 1] + sm_mask[mask_size - 1];
+
+    for(int64_t idx = blockIdx.x; idx < total_tokens; idx += grid_size) {
+        float reg_i[NUM_VT][N];
+        int64_t token_id = mask_size - 1;
+        while(idx < sm_stride[token_id]) {
+            token_id--;
+        }
+        int64_t const token_idx = token_id * num_tokens + idx - sm_stride[token_id];
+        const T* ptr_input0 = input + token_idx * hidden_size2;
+        const T* ptr_input1 = ptr_input0 + hidden_size;
+        float absmax_val = 0.0f;
+        for(int i = tid*N, j = 0; i < hidden_size; i += stride, j++) {
+            VT vsrc0, vsrc1;
+            vsrc0 = *(VT*)(ptr_input0 + i);
+            vsrc1 = *(VT*)(ptr_input1 + i);
+            T* ptr_local0 = (T*)&vsrc0;
+            T* ptr_local1 = (T*)&vsrc1;
+            #pragma unroll N
+            for(int k = 0; k < N; k++) {
+                float val0 = static_cast<float>(ptr_local0[k]);
+                float val1 = static_cast<float>(ptr_local1[k]);
+                float sigmoid = val0 * __builtin_mxc_rcpf(1.0f + __builtin_expf(-val0));
+                float gate_up = val1 * sigmoid;
+                reg_i[j][k] = gate_up;
+                absmax_val = max(absmax_val, abs(gate_up));
+            }
+        }
+
+        using BlockReduce = cub::BlockReduce<float, 512>;
+        __shared__ typename BlockReduce::TempStorage reduceStorage;
+        float const block_absmax_val_maybe =
+            BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim_x);
+        __shared__ float block_absmax_val;
+        int8_t* ptr_output = (int8_t*)(output + token_idx * out_stirde);
+        float* ptr_scale = (float*)(ptr_output + hidden_size);
+        if (tid == 0) {
+            block_absmax_val = block_absmax_val_maybe;
+            ptr_scale[0] = block_absmax_val * 0.0078740157;
+        }
+        __syncthreads();
+        float const tmp_scale = 127.0f * __builtin_mxc_rcpf(block_absmax_val);
+        for (int i = tid*N, k = 0; i < hidden_size; i += stride, k++) {
+            VT1 vdst;
+            int8_t* ptr_dst = (int8_t*)&vdst;
+            #pragma unroll N
+            for(int j = 0; j < N; ++j) {
+                ptr_dst[j] = float_to_int8_rn(reg_i[k][j] * tmp_scale);
+            }
+            *(VT1*)(ptr_output + i) = vdst;
+        }
+    }
+}
+
+template<typename T, typename T1, typename VT, typename VT1, int NUM_VT, int NUM_THREADS> 
+__global__ void silu_and_mul_mask_quant_pack_1mask(T* input, T* output,T1* mask, int grid_size, int64_t num_tokens, int64_t hidden_size, int64_t out_stirde)
+{
+    constexpr int N = sizeof(VT) / sizeof(T);
+    int const tid = threadIdx.x;
+    T1 mask_stride;
+    mask_stride = mask[0];
+    int64_t hidden_size2 = hidden_size << 1;
+    int stride = NUM_THREADS * N;
+    int total_tokens = mask_stride;
+
+    for(int idx = blockIdx.x; idx < total_tokens; idx += grid_size) {
+        float reg_i[NUM_VT][N];
+        int64_t const token_idx = idx;
+        const T* ptr_input0 = input + token_idx * hidden_size2;
+        const T* ptr_input1 = ptr_input0 + hidden_size;
+        float absmax_val = 0.0f;
+        for(int i = tid*N, j = 0; i < hidden_size; i += stride, j++) {
+            VT vsrc0, vsrc1;
+            vsrc0 = *(VT*)(ptr_input0 + i);
+            vsrc1 = *(VT*)(ptr_input1 + i);
+            T* ptr_local0 = (T*)&vsrc0;
+            T* ptr_local1 = (T*)&vsrc1;
+            #pragma unroll N
+            for(int k = 0; k < N; k++) {
+                float val0 = static_cast<float>(ptr_local0[k]);
+                float val1 = static_cast<float>(ptr_local1[k]);
+                float sigmoid = val0 * __builtin_mxc_rcpf(1.0f + __builtin_expf(-val0));
+                float gate_up = val1 * sigmoid;
+                reg_i[j][k] = gate_up;
+                absmax_val = max(absmax_val, abs(gate_up));
+            }
+        }
+
+        using BlockReduce = cub::BlockReduce<float, NUM_THREADS>;
+        __shared__ typename BlockReduce::TempStorage reduceStorage;
+        float const block_absmax_val_maybe =
+            BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, NUM_THREADS);
+        __shared__ float block_absmax_val;
+        int8_t* ptr_output = (int8_t*)(output + token_idx * out_stirde);
+        float* ptr_scale = (float*)(ptr_output + hidden_size);
+        if (tid == 0) {
+            block_absmax_val = block_absmax_val_maybe;
+            ptr_scale[0] = block_absmax_val * 0.0078740157;
+        }
+        __syncthreads();
+        float const tmp_scale = 127.0f * __builtin_mxc_rcpf(block_absmax_val);
+        for (int i = tid*N, k = 0; i < hidden_size; i += stride, k++) {
+            VT1 vdst;
+            int8_t* ptr_dst = (int8_t*)&vdst;
+            #pragma unroll N
+            for(int j = 0; j < N; ++j) {
+                ptr_dst[j] = float_to_int8_rn(reg_i[k][j] * tmp_scale);
+            }
+            *(VT1*)(ptr_output + i) = vdst;
+        }
+    }
+}
+
+template<typename T, typename T1, typename VT, typename VT1, typename VMASK_TYPE, int NUM_VT, int NUM_THREADS> 
+__global__ void silu_and_mul_mask_quant_pack_2mask(T* input, T* output,T1* mask, int grid_size, int64_t num_tokens, int64_t hidden_size, int64_t out_stirde)
+{
+    constexpr int N = sizeof(VT) / sizeof(T);
+    int const tid = threadIdx.x;
+    VMASK_TYPE vmask_reg = *(VMASK_TYPE*)mask;
+    T1 mask_stride[2];
+    T1* ptr_mask = (T1*)&vmask_reg;
+    mask_stride[0] = 0;
+    mask_stride[1] = ptr_mask[0];
+    int64_t hidden_size2 = hidden_size << 1;
+    int stride = NUM_THREADS * N;
+    int total_tokens = ptr_mask[0] + ptr_mask[1];
+
+    for(int idx = blockIdx.x; idx < total_tokens; idx += grid_size) {
+        float reg_i[NUM_VT][N];
+        int64_t token_id = idx < mask_stride[1] ? 0 : 1;;
+        int64_t const token_idx = token_id * num_tokens + idx - mask_stride[token_id];
+        const T* ptr_input0 = input + token_idx * hidden_size2;
+        const T* ptr_input1 = ptr_input0 + hidden_size;
+        float absmax_val = 0.0f;
+        for(int i = tid*N, j = 0; i < hidden_size; i += stride, j++) {
+            VT vsrc0, vsrc1;
+            vsrc0 = *(VT*)(ptr_input0 + i);
+            vsrc1 = *(VT*)(ptr_input1 + i);
+            T* ptr_local0 = (T*)&vsrc0;
+            T* ptr_local1 = (T*)&vsrc1;
+            #pragma unroll N
+            for(int k = 0; k < N; k++) {
+                float val0 = static_cast<float>(ptr_local0[k]);
+                float val1 = static_cast<float>(ptr_local1[k]);
+                float sigmoid = val0 * __builtin_mxc_rcpf(1.0f + __builtin_expf(-val0));
+                float gate_up = val1 * sigmoid;
+                reg_i[j][k] = gate_up;
+                absmax_val = max(absmax_val, abs(gate_up));
+            }
+        }
+
+        using BlockReduce = cub::BlockReduce<float, NUM_THREADS>;
+        __shared__ typename BlockReduce::TempStorage reduceStorage;
+        float const block_absmax_val_maybe =
+            BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, NUM_THREADS);
+        __shared__ float block_absmax_val;
+        int8_t* ptr_output = (int8_t*)(output + token_idx * out_stirde);
+        float* ptr_scale = (float*)(ptr_output + hidden_size);
+        if (tid == 0) {
+            block_absmax_val = block_absmax_val_maybe;
+            ptr_scale[0] = block_absmax_val * 0.0078740157;
+        }
+        __syncthreads();
+        float const tmp_scale = 127.0f * __builtin_mxc_rcpf(block_absmax_val);
+        for (int i = tid*N, k = 0; i < hidden_size; i += stride, k++) {
+            VT1 vdst;
+            int8_t* ptr_dst = (int8_t*)&vdst;
+            #pragma unroll N
+            for(int j = 0; j < N; ++j) {
+                ptr_dst[j] = float_to_int8_rn(reg_i[k][j] * tmp_scale);
+            }
+            *(VT1*)(ptr_output + i) = vdst;
+        }
+    }
+}
+
+template<typename T, typename VT, typename VT1, int NUM_VT> 
+__global__ void silu_and_mul_quant(T* input, int8_t* output, float* scale, int64_t hidden_size, int blockDim_x)
+{
+    constexpr int N = sizeof(VT) / sizeof(T);
+    int const tid = threadIdx.x;
+    int stride = blockDim_x * N;
+    int64_t const token_idx = blockIdx.x;
+    int64_t hidden_size2 = hidden_size << 1;
+    const T* ptr_input0 = input + token_idx * hidden_size2;
+    const T* ptr_input1 = ptr_input0 + hidden_size;
+    float absmax_val = 0.0f;
+    float reg_i[NUM_VT][N];
+    for(int i = tid*N, j = 0; i < hidden_size; i += stride, j++) {
+        VT vsrc0, vsrc1;
+        vsrc0 = *(VT*)(ptr_input0 + i);
+        vsrc1 = *(VT*)(ptr_input1 + i);
+        T* ptr_local0 = (T*)&vsrc0;
+        T* ptr_local1 = (T*)&vsrc1;
+        #pragma unroll N
+        for(int k = 0; k < N; k++) {
+            float val0 = static_cast<float>(ptr_local0[k]);
+            float val1 = static_cast<float>(ptr_local1[k]);
+            float sigmoid = val0 * __builtin_mxc_rcpf(1.0f + __builtin_expf(-val0));
+            float gate_up = val1 * sigmoid;
+            reg_i[j][k] = gate_up;
+            absmax_val = max(absmax_val, abs(gate_up));
+        }
+    }
+
+    using BlockReduce = cub::BlockReduce<float, 512>;
+    __shared__ typename BlockReduce::TempStorage reduceStorage;
+    float const block_absmax_val_maybe =
+        BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim_x);
+    __shared__ float block_absmax_val;
+    int8_t* ptr_output = (int8_t*)(output + token_idx * hidden_size);
+    if (tid == 0) {
+        block_absmax_val = block_absmax_val_maybe;
+        scale[token_idx] = block_absmax_val * 0.0078740157;
+    }
+    __syncthreads();
+    float const tmp_scale = 127.0f * __builtin_mxc_rcpf(block_absmax_val);
+    for (int i = tid*N, k = 0; i < hidden_size; i += stride, k++) {
+        VT1 vdst;
+        int8_t* ptr_dst = (int8_t*)&vdst;
+        #pragma unroll N
+        for(int j = 0; j < N; ++j) {
+            ptr_dst[j] = float_to_int8_rn(reg_i[k][j] * tmp_scale);
+        }
+        *(VT1*)(ptr_output + i) = vdst;
+    }
+}
+
+template<typename T, typename VT, typename VT1, int NUM_VT> 
+__global__ void silu_and_mul_sm_quant(T* input, int8_t* output, float* scale, int64_t hidden_size, int blockDim_x)
+{
+    constexpr int N = sizeof(VT) / sizeof(T);
+    int const tid = threadIdx.x;
+    int stride = blockDim_x * N;
+    int64_t const token_idx = blockIdx.x;
+    int64_t hidden_size2 = hidden_size << 1;
+    const T* ptr_input0 = input + token_idx * hidden_size2;
+    const T* ptr_input1 = ptr_input0 + hidden_size;
+    float absmax_val = 0.0f;
+    float reg_i[4][N];
+    __shared__ float sm_gate[4096];
+    int hidden_size1 = stride * 4;
+    for(int i = tid*N, j = 0; i < hidden_size1; i += stride, j++) {
+        VT vsrc0, vsrc1;
+        vsrc0 = *(VT*)(ptr_input0 + i);
+        vsrc1 = *(VT*)(ptr_input1 + i);
+        T* ptr_local0 = (T*)&vsrc0;
+        T* ptr_local1 = (T*)&vsrc1;
+        #pragma unroll N
+        for(int k = 0; k < N; k++) {
+            float val0 = static_cast<float>(ptr_local0[k]);
+            float val1 = static_cast<float>(ptr_local1[k]);
+            float sigmoid = val0 * __builtin_mxc_rcpf(1.0f + __builtin_expf(-val0));
+            float gate_up = val1 * sigmoid;
+            reg_i[j][k] = gate_up;
+            absmax_val = max(absmax_val, abs(gate_up));
+        }
+    }
+    const T* ptr_input2 = ptr_input0 + hidden_size1;
+    const T* ptr_input3 = ptr_input1 + hidden_size1;
+    int remain_hidden_size = hidden_size - hidden_size1;
+    for(int i = tid*N; i < remain_hidden_size; i += stride) {
+        VT vsrc0, vsrc1;
+        vsrc0 = *(VT*)(ptr_input2 + i);
+        vsrc1 = *(VT*)(ptr_input3 + i);
+        T* ptr_local0 = (T*)&vsrc0;
+        T* ptr_local1 = (T*)&vsrc1;
+        float* ptr_sm = sm_gate + i;
+        #pragma unroll N
+        for(int k = 0; k < N; k++) {
+            float val0 = static_cast<float>(ptr_local0[k]);
+            float val1 = static_cast<float>(ptr_local1[k]);
+            float sigmoid = val0 * __builtin_mxc_rcpf(1.0f + __builtin_expf(-val0));
+            float gate_up = val1 * sigmoid;
+            ptr_sm[k] = gate_up;
+            absmax_val = max(absmax_val, abs(gate_up));
+        }
+    }
+
+    using BlockReduce = cub::BlockReduce<float, 512>;
+    __shared__ typename BlockReduce::TempStorage reduceStorage;
+    float const block_absmax_val_maybe =
+        BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim_x);
+    __shared__ float block_absmax_val;
+    int8_t* ptr_output = (int8_t*)(output + token_idx * hidden_size);
+    if (tid == 0) {
+        block_absmax_val = block_absmax_val_maybe;
+        scale[token_idx] = block_absmax_val * 0.0078740157;
+    }
+    __syncthreads();
+    float const tmp_scale = 127.0f * __builtin_mxc_rcpf(block_absmax_val);
+    for (int i = tid*N, k = 0; i < hidden_size1; i += stride, k++) {
+        VT1 vdst;
+        int8_t* ptr_dst = (int8_t*)&vdst;
+        #pragma unroll N
+        for(int j = 0; j < N; ++j) {
+            ptr_dst[j] = float_to_int8_rn(reg_i[k][j] * tmp_scale);
+        }
+        *(VT1*)(ptr_output + i) = vdst;
+    }
+
+    ptr_output = ptr_output + hidden_size1;
+    for(int i = tid*N; i < remain_hidden_size; i += stride) {
+        VT1 vdst;
+        int8_t* ptr_dst = (int8_t*)&vdst;
+        float* ptr_sm = sm_gate + i;
+        #pragma unroll N
+        for(int j = 0; j < N; ++j) {
+            ptr_dst[j] = float_to_int8_rn(ptr_sm[j] * tmp_scale);
+        }
+        *(VT1*)(ptr_output + i) = vdst;
+    }
+}
+
+template<typename T, typename T1>
+void launch_silu_mul_quant_pack(T* input, T* output, T1* mask, int64_t num_tokens, int64_t hidden_size, int64_t out_stride, int64_t mask_size,cudaStream_t stream) {
+    int dev = 0;
+    cudaGetDevice(&dev);
+    int sm_count = 0;
+    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
+    int gridsize = sm_count;
+    int64_t inner_hidden_size = hidden_size / 2;
+    int blocksize = 512;
+    int N = sizeof(float4) / sizeof(T);
+    if(mask_size == 1 && N == 8&&(inner_hidden_size & (N - 1)) == 0 && (out_stride & (N -1)) == 0) {
+        int base = blocksize * N;
+        if(inner_hidden_size <= 64*N) {
+          gridsize = gridsize * 8;
+          silu_and_mul_mask_quant_pack_1mask<T, T1, float4, float2, 1, 64><<<gridsize, 64,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+        } else if(inner_hidden_size <= 128*N) {
+          gridsize = gridsize * 4;
+          silu_and_mul_mask_quant_pack_1mask<T, T1, float4, float2, 1, 128><<<gridsize, 128,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+        } else if(inner_hidden_size <= 256*N) {
+          gridsize = gridsize * 2;
+          silu_and_mul_mask_quant_pack_1mask<T, T1, float4, float2, 1, 256><<<gridsize, 256,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+        } else if(inner_hidden_size <= base) {
+            silu_and_mul_mask_quant_pack_1mask<T, T1, float4, float2, 1, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+        } else if(inner_hidden_size <= base*2) {
+            silu_and_mul_mask_quant_pack_1mask<T, T1, float4, float2, 2, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+        } else if(inner_hidden_size <= base * 3) {
+            silu_and_mul_mask_quant_pack_1mask<T, T1, float4, float2, 3, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+        } else if(inner_hidden_size <= base * 4) {
+            silu_and_mul_mask_quant_pack_1mask<T, T1, float4, float2, 4, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+        } else {
+            assert(0);
+        }
+    } else if(mask_size == 2 && (N == 8&&(inner_hidden_size & (N - 1)) == 0 && (out_stride & (N -1)) == 0)) {
+        int base = blocksize * N;
+        if(sizeof(T1) == 4) {
+          if(inner_hidden_size <= 64 * N) {
+            gridsize = gridsize * 8;
+            silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float2, 1, 64><<<gridsize, 64,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= 128 * N){
+            gridsize = gridsize * 4;
+            silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float2, 1, 128><<<gridsize, 128,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= 256 * N) {
+            gridsize = gridsize * 2;
+            silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float2, 1, 256><<<gridsize, 256,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= base) {
+              silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float2, 1, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= base*2) {
+              silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float2, 2, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= base * 3) {
+              silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float2, 3, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= base * 4) {
+              silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float2, 4, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else {
+              assert(0);
+          }
+        } else if(sizeof(T1) == 8) {
+          if(inner_hidden_size <= 64 * N) {
+            gridsize = gridsize * 8;
+            silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float4, 1, 64><<<gridsize, 64,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= 128 * N) {
+            gridsize = gridsize * 4;
+            silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float4, 1, 128><<<gridsize, 128,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= 256 * N) {
+            gridsize = gridsize * 2;
+            silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float4, 1, 256><<<gridsize, 256,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= base) {
+              silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float4, 1, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= base*2) {
+              silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float4, 2, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= base * 3) {
+              silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float4, 3,512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else if(inner_hidden_size <= base * 4) {
+              silu_and_mul_mask_quant_pack_2mask<T, T1, float4, float2, float4, 4, 512><<<gridsize, 512,0,stream>>>(input, output, mask, gridsize, num_tokens, inner_hidden_size, out_stride);
+          } else {
+              assert(0);
+          }
+        }
+        
+    } else if(N == 8&&(inner_hidden_size & (N - 1)) == 0 && (out_stride & (N -1)) == 0) {
+        int base = blocksize * N;
+        if(inner_hidden_size <= base) {
+            silu_and_mul_mask_quant_pack<T, T1, float4, float2, 1><<<gridsize, blocksize,0,stream>>>(input, output, mask, mask_size, gridsize, num_tokens, inner_hidden_size, out_stride, blocksize);
+        } else if(inner_hidden_size <= base*2) {
+            silu_and_mul_mask_quant_pack<T, T1, float4, float2, 2><<<gridsize, blocksize,0,stream>>>(input, output, mask, mask_size, gridsize, num_tokens, inner_hidden_size, out_stride, blocksize);
+        } else if(inner_hidden_size <= base * 3) {
+            silu_and_mul_mask_quant_pack<T, T1, float4, float2, 3><<<gridsize, blocksize,0,stream>>>(input, output, mask, mask_size, gridsize, num_tokens, inner_hidden_size, out_stride, blocksize);
+        } else if(inner_hidden_size <= base * 4) {
+            silu_and_mul_mask_quant_pack<T, T1, float4, float2, 4><<<gridsize, blocksize,0,stream>>>(input, output, mask, mask_size, gridsize, num_tokens, inner_hidden_size, out_stride, blocksize);
+        } else {
+            assert(0);
+        }
+    } else {
+        assert(0);
+    }
+}
+
+template<typename T>
+void launch_silu_mul_quan_no_mask(T* input, int8_t* output, float* scale, int64_t num_tokens, int64_t hidden_size,cudaStream_t stream) {
+    int64_t inner_hidden_size = hidden_size / 2;
+    int blocksize = 512;
+    int N = sizeof(float4) / sizeof(T);
+    if(N == 8&&(inner_hidden_size & (N - 1)) == 0) {
+        int base = blocksize * N;
+        if(inner_hidden_size <= base) {
+            silu_and_mul_quant<T, float4, float2, 1><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else if(inner_hidden_size <= base*2) {
+            silu_and_mul_quant<T, float4, float2, 2><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else if(inner_hidden_size <= base * 3) {
+            silu_and_mul_quant<T, float4, float2, 3><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else if(inner_hidden_size <= base * 4) {
+            silu_and_mul_quant<T, float4, float2, 4><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else if(inner_hidden_size <= base*4 + 4096) {
+            silu_and_mul_sm_quant<T, float4, float2, 4><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else {
+            printf("silu_and_mul_quant not support\n");
+            assert(0);
+        }
+    } else if(N == 4 && (inner_hidden_size & (N - 1)) == 0) {
+        int base = blocksize * N;
+        if(inner_hidden_size <= base) {
+            silu_and_mul_quant<T, float4, float, 1><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else if(inner_hidden_size <= base*2) {
+            silu_and_mul_quant<T, float4, float, 2><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else if(inner_hidden_size <= base * 3) {
+            silu_and_mul_quant<T, float4, float, 3><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else if(inner_hidden_size <= base * 4) {
+            silu_and_mul_quant<T, float4, float, 4><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else if(inner_hidden_size <= base * 8) {
+            silu_and_mul_quant<T, float4, float, 8><<<num_tokens, blocksize,0,stream>>>(input, output, scale, inner_hidden_size, blocksize);
+        } else {
+            printf("silu_and_mul_quant not support\n");
+            assert(0);
+        }
+    } else {
+        assert(0);
+    }
 }
 
 }  // namespace vllm
@@ -238,7 +1372,7 @@ __global__ void dynamic_scaled_int8_azp_quant_kernel(
 void static_scaled_int8_quant(torch::Tensor& out,          // [..., hidden_size]
                               torch::Tensor const& input,  // [..., hidden_size]
                               torch::Tensor const& scale,
-                              std::optional<torch::Tensor> const& azp) {
+                              c10::optional<torch::Tensor> const& azp) {
   TORCH_CHECK(input.is_contiguous());
   TORCH_CHECK(out.is_contiguous());
   TORCH_CHECK(scale.numel() == 1);
@@ -269,7 +1403,7 @@ void static_scaled_int8_quant(torch::Tensor& out,          // [..., hidden_size]
 void dynamic_scaled_int8_quant(
     torch::Tensor& out,          // [..., hidden_size]
     torch::Tensor const& input,  // [..., hidden_size]
-    torch::Tensor& scales, std::optional<torch::Tensor> const& azp) {
+    torch::Tensor& scales, c10::optional<torch::Tensor> const& azp) {
   TORCH_CHECK(input.is_contiguous());
   TORCH_CHECK(out.is_contiguous());
   TORCH_CHECK(scales.is_contiguous());
@@ -277,22 +1411,184 @@ void dynamic_scaled_int8_quant(
 
   int const hidden_size = input.size(-1);
   int const num_tokens = input.numel() / hidden_size;
-  dim3 const grid(num_tokens);
+  dim3 const grid(num_tokens,1,1);
   dim3 const block(std::min(hidden_size, 1024));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   VLLM_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "dynamic_scaled_int8_quant_kernel", [&] {
         if (!azp) {
-          vllm::dynamic_scaled_int8_quant_kernel<scalar_t, float>
+          int n = 16 / sizeof(scalar_t);
+          if(hidden_size <= 4096 && ((hidden_size & (n - 1)) == 0) && n == 8) {
+            if(hidden_size > 256*n) {
+              vllm::dynamic_scaled_int8_quant_kernel_sreg_opt<scalar_t, float, float4, float2, 512, false><<<grid, 512, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens);
+            } else if(hidden_size > 128*n) {
+              vllm::dynamic_scaled_int8_quant_kernel_sreg_opt<scalar_t, float, float4, float2, 256, false><<<grid, 256, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens);
+            } else if(hidden_size > 64 * n) {
+              vllm::dynamic_scaled_int8_quant_kernel_sreg_opt<scalar_t, float, float4, float2, 128, false><<<grid, 128, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens);
+            } else {
+              vllm::dynamic_scaled_int8_quant_kernel_sreg_opt<scalar_t, float, float4, float2, 64, false><<<grid, 64, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens);
+            }
+          } else if(hidden_size > 4096 &&hidden_size <= 8192 && ((hidden_size & (2*n - 1)) == 0) && n == 8) {
+            int blocksize = 512;
+            vllm::dynamic_scaled_int8_quant_kernel_reg_opt<scalar_t, float, float4, float4, false><<<grid, blocksize, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize, num_tokens);
+          } else if(hidden_size <= 8064 && (hidden_size & (n - 1)) == 0 && n == 8) {
+            int blocksize = 512;
+            vllm::dynamic_scaled_int8_quant_kernel_sm_opt<scalar_t, float, float4, float2,false><<<grid, blocksize, 0, stream>>>(
+              input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize,num_tokens);
+          } else if(hidden_size >= 16384 && hidden_size <= 18432 && (hidden_size & (n - 1)) == 0 && n == 8) {
+            vllm::dynamic_scaled_int8_quant_kernel_lh_opt<scalar_t, float, float4, float2, 3,1024, false><<<grid, 1024, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens);
+          } else if (hidden_size > 8064 && ((hidden_size & (n - 1)) == 0 && n == 8)) {
+            int blocksize = 1024;
+            vllm::dynamic_scaled_int8_quant_kernel_opt<scalar_t, float,float4,float2, false>
+                    <<<grid, blocksize, 0, stream>>>(
+                        input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize, num_tokens);
+          } else {
+              vllm::dynamic_scaled_int8_quant_kernel<scalar_t, float, false>
+                  <<<grid, block, 0, stream>>>(
+                      input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
+                      scales.data_ptr<float>(), hidden_size, num_tokens);
+          }
+        } else {
+          vllm::dynamic_scaled_int8_azp_quant_kernel<scalar_t, float, int32_t, false>
               <<<grid, block, 0, stream>>>(
                   input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
-                  scales.data_ptr<float>(), hidden_size);
+                  scales.data_ptr<float>(), azp->data_ptr<int32_t>(),
+                  hidden_size, num_tokens);
+        }
+      });
+}
+
+void dynamic_scaled_int8_mask_quant(
+    torch::Tensor& out,          // [..., hidden_size]
+    torch::Tensor const& input,  // [..., hidden_size]
+    torch::Tensor const &mask,
+    torch::Tensor& scales, c10::optional<torch::Tensor> const& azp) {
+  TORCH_CHECK(input.is_contiguous());
+  TORCH_CHECK(out.is_contiguous());
+  TORCH_CHECK(scales.is_contiguous());
+  TORCH_CHECK(mask.is_contiguous());
+  TORCH_CHECK(!azp || azp->is_contiguous());
+
+  int const hidden_size = input.size(-1);
+  int const num_tokens = input.numel() / hidden_size;
+  int const mask_size = mask.numel();
+  int const num_tokens_batch = num_tokens / mask_size;
+  dim3 const grid(num_tokens_batch, mask_size, 1);
+  dim3 const block(std::min(hidden_size, 1024));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "dynamic_scaled_int8_quant_kernel_mask", [&] {
+        if (!azp) {
+          int n = 16 / sizeof(scalar_t);
+          if(mask_size < 16) {
+            int dev = 0;
+            cudaGetDevice(&dev);
+            int sm_count = 0;
+            cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
+            int gridsize = sm_count;
+            if(hidden_size <= 4096 && ((hidden_size & (n - 1)) == 0) && n == 8) {
+              int blocksize = 512;
+              vllm::dynamic_scaled_int8_quant_mask_kernel_sreg_opt<scalar_t, float, float4, float2><<<gridsize, blocksize, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize,num_tokens_batch,mask_size,gridsize,mask.data_ptr<int>());
+            } else if(hidden_size > 4096 &&hidden_size <= 8192 && ((hidden_size & (2*n - 1)) == 0) && n == 8) {
+                int blocksize = 512;
+              vllm::dynamic_scaled_int8_quant_mask_kernel_reg_opt<scalar_t, float, float4, float4><<<gridsize, blocksize, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize,num_tokens_batch,mask_size, gridsize,mask.data_ptr<int>());
+            } else if(hidden_size <= 8064 && (hidden_size & (n - 1)) == 0 && n == 8) {
+                int blocksize = 512;
+                vllm::dynamic_scaled_int8_quant_mask_kernel_sm_opt<scalar_t, float, float4, float2><<<gridsize, blocksize, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize,num_tokens_batch,mask_size, gridsize, mask.data_ptr<int>());
+            } else if(hidden_size > 8064 && ((hidden_size & (n - 1)) == 0 && n == 8)){
+                int blocksize = 1024;
+                vllm::dynamic_scaled_int8_quant_mask_kernel_opt<scalar_t, float,float4,float2>
+                        <<<gridsize, blocksize, 0, stream>>>(
+                            input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize,num_tokens_batch,mask_size, gridsize, mask.data_ptr<int>());
+            } else {
+              vllm::dynamic_scaled_int8_quant_mask_kernel<scalar_t, float>
+                    <<<gridsize, block, 0, stream>>>(
+                        input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
+                        scales.data_ptr<float>(), hidden_size, num_tokens_batch, mask_size, gridsize, mask.data_ptr<int>());
+            }
+          } else {
+            if(hidden_size <= 4096 && ((hidden_size & (n - 1)) == 0) && n == 8) {
+              if(hidden_size > 256*n) {
+                vllm::dynamic_scaled_int8_quant_kernel_sreg_opt<scalar_t, float, float4, float2,512,true><<<grid, 512, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens_batch, mask.data_ptr<int>());
+              } else if(hidden_size > 128*n) {
+                vllm::dynamic_scaled_int8_quant_kernel_sreg_opt<scalar_t, float, float4, float2,256,true><<<grid, 256, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens_batch, mask.data_ptr<int>());
+              } else if(hidden_size > 64*n) {
+                vllm::dynamic_scaled_int8_quant_kernel_sreg_opt<scalar_t, float, float4, float2,128,true><<<grid, 128, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens_batch, mask.data_ptr<int>());
+              } else {
+                vllm::dynamic_scaled_int8_quant_kernel_sreg_opt<scalar_t, float, float4, float2,64,true><<<grid, 64, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size, num_tokens_batch, mask.data_ptr<int>());
+              }
+            } else if(hidden_size > 4096 &&hidden_size <= 8192 && ((hidden_size & (2*n - 1)) == 0) && n == 8) {
+              int blocksize = 512;
+              vllm::dynamic_scaled_int8_quant_kernel_reg_opt<scalar_t, float, float4, float4, true><<<grid, blocksize, 0, stream>>>(input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize,num_tokens_batch,mask.data_ptr<int>());
+            } else if(hidden_size <= 8064 && (hidden_size & (n - 1)) == 0 && n == 8) {
+              int blocksize = 512;
+              vllm::dynamic_scaled_int8_quant_kernel_sm_opt<scalar_t, float, float4, float2, true><<<grid, blocksize, 0, stream>>>(
+                input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize, num_tokens_batch,mask.data_ptr<int>());
+            } else if (hidden_size > 8064 && ((hidden_size & (n - 1)) == 0 && n == 8)) {
+              int blocksize = 1024;
+              vllm::dynamic_scaled_int8_quant_kernel_opt<scalar_t, float,float4,float2, true>
+                      <<<grid, blocksize, 0, stream>>>(
+                          input.data_ptr<scalar_t>(),out.data_ptr<int8_t>(),scales.data_ptr<float>(),hidden_size,blocksize, num_tokens_batch, mask.data_ptr<int>());
+            } else {
+                vllm::dynamic_scaled_int8_quant_kernel<scalar_t, float,true>
+                    <<<grid, block, 0, stream>>>(
+                        input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
+                        scales.data_ptr<float>(), hidden_size, num_tokens_batch, mask.data_ptr<int>());
+            }
+          }
         } else {
-          vllm::dynamic_scaled_int8_azp_quant_kernel<scalar_t, float, int32_t>
+          vllm::dynamic_scaled_int8_azp_quant_kernel<scalar_t, float, int32_t,true>
               <<<grid, block, 0, stream>>>(
                   input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
                   scales.data_ptr<float>(), azp->data_ptr<int32_t>(),
-                  hidden_size);
+                  hidden_size, num_tokens_batch, mask.data_ptr<int>());
         }
       });
 }
+
+
+void fused_silu_mul_dq_mask_quant_pack(
+    torch::Tensor& out,          
+    torch::Tensor const& input, 
+    torch::Tensor const &mask)
+{
+  TORCH_CHECK(input.is_contiguous());
+  TORCH_CHECK(out.is_contiguous());
+  TORCH_CHECK(mask.is_contiguous());
+  int64_t const hidden_size = input.size(-1);
+  int64_t const num_tokens = input.numel() / hidden_size;
+  int64_t const mask_size = mask.numel();
+  int64_t const num_tokens_batch = num_tokens / mask_size;
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  int64_t out_stride = ((hidden_size/4 + 2) + 255)/ 256 * 256;
+  switch(mask.element_size()) {
+    case 8:
+      VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "launch_silu_mul_quant_pack", [&] {
+        vllm::launch_silu_mul_quant_pack<scalar_t, int64_t>(input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), mask.data_ptr<int64_t>(), num_tokens_batch, hidden_size, out_stride, mask_size, stream);
+      });
+    break;
+    case 4:
+       VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "launch_silu_mul_quant_pack", [&] {
+        vllm::launch_silu_mul_quant_pack<scalar_t, int32_t>(input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), mask.data_ptr<int32_t>(), num_tokens_batch, hidden_size, out_stride, mask_size, stream);
+      });
+      break;
+    default:
+    return;
+  }
+}
+
+void fused_silu_mul_dq_quant_interface(
+    torch::Tensor& out,
+    torch::Tensor& scale,   
+    torch::Tensor const& input)
+{
+  TORCH_CHECK(input.is_contiguous());
+  TORCH_CHECK(scale.is_contiguous());
+  TORCH_CHECK(out.is_contiguous());
+  int64_t const hidden_size = input.size(-1);
+  int64_t const num_tokens = input.numel() / hidden_size;
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "launch_silu_mul_quan_no_mask", [&] {
+    vllm::launch_silu_mul_quan_no_mask<scalar_t>(input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(), scale.data_ptr<float>(), num_tokens, hidden_size, stream);
+  });
+}
\ No newline at end of file
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
index 865fef5ae..95065bcef 100644
--- a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
+++ b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
@@ -1,14 +1,21 @@
 #include <stddef.h>
 #include <torch/all.h>
+#ifdef USE_MACA
+#include "mctlass/mctlass.h"
+#include "mctlass/epilogue/thread/scale_type.h"
+#else
 #include "cutlass/cutlass.h"
+#endif
 
 #include "scaled_mm_c2x.cuh"
 #include "scaled_mm_c2x_sm75_dispatch.cuh"
+#ifndef USE_MACA
 #include "scaled_mm_c2x_sm80_dispatch.cuh"
 #include "scaled_mm_c2x_sm89_fp8_dispatch.cuh"
 #include "scaled_mm_c2x_sm89_int8_dispatch.cuh"
 
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp"
+#endif
 
 using namespace vllm;
 
@@ -22,6 +29,7 @@ template <template <typename, typename> typename Epilogue,
 void cutlass_scaled_mm_sm75_epilogue(torch::Tensor& out, torch::Tensor const& a,
                                      torch::Tensor const& b,
                                      EpilogueArgs&&... epilogue_args) {
+#ifndef USE_MACA
   TORCH_CHECK(a.dtype() == torch::kInt8);
   TORCH_CHECK(b.dtype() == torch::kInt8);
 
@@ -33,6 +41,7 @@ void cutlass_scaled_mm_sm75_epilogue(torch::Tensor& out, torch::Tensor const& a,
     return cutlass_gemm_sm75_dispatch<int8_t, cutlass::half_t, Epilogue>(
         out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
   }
+#endif // USE_MACA
 }
 
 void cutlass_scaled_mm_sm75(torch::Tensor& out, torch::Tensor const& a,
@@ -40,6 +49,7 @@ void cutlass_scaled_mm_sm75(torch::Tensor& out, torch::Tensor const& a,
                             torch::Tensor const& a_scales,
                             torch::Tensor const& b_scales,
                             std::optional<torch::Tensor> const& bias) {
+#ifndef USE_MACA
   TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
   TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
   if (bias) {
@@ -51,6 +61,195 @@ void cutlass_scaled_mm_sm75(torch::Tensor& out, torch::Tensor const& a,
     return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogue>(
         out, a, b, a_scales, b_scales);
   }
+#else
+  int32_t m = a.size(0);
+  int32_t n = b.size(1);
+  int32_t k = a.size(1);
+  int32_t batch_count = 1;
+  if (a.dim() == 3 && b.dim() == 3) {
+      // a.size = [batch_size, M, K], b.size = [batch_size, K, N]
+      m = a.size(1);
+      n = b.size(2);
+      k = a.size(2);
+      batch_count = a.size(0);
+  }
+
+  using ArchTag = mctlass::arch::Sm80;
+  using ElementA = int8_t;
+  using ElementB = int8_t;
+  using ElementC = mctlass::half_t;
+  using ElementCompute = float;
+  using LayoutA = mctlass::layout::RowMajor;
+  //using LayoutB = mctlass::layout::RowMajor;
+  using LayoutB = mctlass::layout::ColumnMajor;
+  using LayoutC = mctlass::layout::RowMajor;
+
+  if (out.dtype() == torch::kBFloat16)
+  {
+    auto a_ptr = static_cast<ElementA const*>(a.data_ptr());
+    auto b_ptr = static_cast<ElementB const*>(b.data_ptr());
+    auto c_ptr = static_cast<maca_bfloat16*>(out.data_ptr());
+    auto scale_a = a_scales.data_ptr<float>();
+    auto scale_b = b_scales.data_ptr<float>();
+    auto stream = at::cuda::getCurrentCUDAStream(a.get_device());
+    if (bias) {
+      mctlass::epilogue::thread::ScaleType::ScaleBiasKind const scale_type =
+      mctlass::epilogue::thread::ScaleType::ScaleAvBvBias;
+      using mctlassGemmScaleOp = mctlassGemmScale<
+        ElementA,
+        LayoutA,
+        ElementB,
+        LayoutB,
+        maca_bfloat16,
+        LayoutC,
+        ElementCompute,
+        ArchTag,
+        scale_type
+      >;
+      maca_bfloat16 *bias_t;
+      bias_t = static_cast<maca_bfloat16 *>(bias.value().data_ptr());
+      mctlassGemmScaleOp mctlass_op;
+      mctlass::gemm::GemmCoord problem_size(m, n, k);
+      typename mctlassGemmScaleOp::Arguments arguments{
+          mctlass::gemm::GemmUniversalMode::kGemm,
+          problem_size,
+          1,//batch_count
+          {scale_a, scale_b, bias_t},
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          c_ptr,
+          m * k,
+          n * k,
+          m * n,
+          m * n,
+          k,
+          n,
+          n,
+          n
+      };
+      mctlass_op(arguments, NULL, stream);
+    }
+    else{
+      mctlass::epilogue::thread::ScaleType::ScaleBiasKind const scale_type =
+      mctlass::epilogue::thread::ScaleType::ScaleAvBv;
+      using mctlassGemmScaleOp = mctlassGemmScale<
+        ElementA,
+        LayoutA,
+        ElementB,
+        LayoutB,
+        maca_bfloat16,
+        LayoutC,
+        ElementCompute,
+        ArchTag,
+        scale_type
+      >;
+      mctlassGemmScaleOp mctlass_op;
+      mctlass::gemm::GemmCoord problem_size(m, n, k);
+      typename mctlassGemmScaleOp::Arguments arguments{
+          mctlass::gemm::GemmUniversalMode::kGemm,
+          problem_size,
+          1,//batch_count
+          {scale_a, scale_b, nullptr},
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          c_ptr,
+          m * k,
+          n * k,
+          m * n,
+          m * n,
+          k,
+          n,
+          n,
+          n
+      };
+      mctlass_op(arguments, NULL, stream);
+    }
+  }
+  else{
+    auto a_ptr = static_cast<ElementA const*>(a.data_ptr());
+    auto b_ptr = static_cast<ElementB const*>(b.data_ptr());
+    auto c_ptr = static_cast<ElementC*>(out.data_ptr());
+    auto scale_a = a_scales.data_ptr<float>();
+    auto scale_b = b_scales.data_ptr<float>();
+    auto stream = at::cuda::getCurrentCUDAStream(a.get_device());
+    if (bias) {
+      mctlass::epilogue::thread::ScaleType::ScaleBiasKind const scale_type =
+      mctlass::epilogue::thread::ScaleType::ScaleAvBvBias;
+      using mctlassGemmScaleOp = mctlassGemmScale<
+        ElementA,
+        LayoutA,
+        ElementB,
+        LayoutB,
+        ElementC,
+        LayoutC,
+        ElementCompute,
+        ArchTag,
+        scale_type
+      >;
+      ElementC *bias_t;
+      bias_t = static_cast<ElementC *>(bias.value().data_ptr());
+      mctlassGemmScaleOp mctlass_op;
+      mctlass::gemm::GemmCoord problem_size(m, n, k);
+      typename mctlassGemmScaleOp::Arguments arguments{
+          mctlass::gemm::GemmUniversalMode::kGemm,
+          problem_size,
+          1,//batch_count
+          {scale_a, scale_b, bias_t},
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          c_ptr,
+          m * k,
+          n * k,
+          m * n,
+          m * n,
+          k,
+          n,
+          n,
+          n
+      };
+      mctlass_op(arguments, NULL, stream);
+    }
+    else{
+      mctlass::epilogue::thread::ScaleType::ScaleBiasKind const scale_type =
+      mctlass::epilogue::thread::ScaleType::ScaleAvBv;
+      using mctlassGemmScaleOp = mctlassGemmScale<
+        ElementA,
+        LayoutA,
+        ElementB,
+        LayoutB,
+        ElementC,
+        LayoutC,
+        ElementCompute,
+        ArchTag,
+        scale_type
+      >;
+      mctlassGemmScaleOp mctlass_op;
+      mctlass::gemm::GemmCoord problem_size(m, n, k);
+      typename mctlassGemmScaleOp::Arguments arguments{
+          mctlass::gemm::GemmUniversalMode::kGemm,
+          problem_size,
+          1,//batch_count
+          {scale_a, scale_b, nullptr},
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          c_ptr,
+          m * k,
+          n * k,
+          m * n,
+          m * n,
+          k,
+          n,
+          n,
+          n
+      };
+      mctlass_op(arguments, NULL, stream);
+    }
+  }
+#endif // USE_MACA
 }
 
 void cutlass_scaled_mm_azp_sm75(torch::Tensor& out, torch::Tensor const& a,
@@ -60,6 +259,7 @@ void cutlass_scaled_mm_azp_sm75(torch::Tensor& out, torch::Tensor const& a,
                                 torch::Tensor const& azp_adj,
                                 std::optional<torch::Tensor> const& azp,
                                 std::optional<torch::Tensor> const& bias) {
+#ifndef USE_MACA
   TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
   TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
 
@@ -70,8 +270,178 @@ void cutlass_scaled_mm_azp_sm75(torch::Tensor& out, torch::Tensor const& a,
     return cutlass_scaled_mm_sm75_epilogue<c2x::ScaledEpilogueBiasAzp>(
         out, a, b, a_scales, b_scales, azp_adj, bias);
   }
+#elif (MACA_VERSION_MAJOR * 100 + MACA_VERSION_MINOR) >= 231 // MACA version >= 2.31.0.x
+  int32_t m = a.size(0);
+  int32_t n = b.size(1);
+  int32_t k = a.size(1);
+  int32_t batchsize = 1;
+  if (a.dim() == 3 && b.dim() == 3) {
+    // a.size = [batch_size, M, K], b.size = [batch_size, K, N]
+    m = a.size(1);
+    n = b.size(2);
+    k = a.size(2);
+    batchsize = a.size(0);
+  }
+
+  using ArchTag = mctlass::arch::Sm80;
+  using ElementA = int8_t;
+  using ElementB = int8_t;
+
+  using ElementCompute = float;
+  using ElementAccumulator = int32_t;
+
+  using LayoutA = mctlass::layout::RowMajor;
+  using LayoutB = mctlass::layout::ColumnMajor;
+  using LayoutC = mctlass::layout::RowMajor;
+
+  auto stream = at::cuda::getCurrentCUDAStream(a.get_device());
+
+  if (out.dtype() == torch::kBFloat16) {
+    using ElementC = maca_bfloat16;
+    using ElementOutput = ElementC;
+
+    auto a_ptr = static_cast<ElementA const*>(a.data_ptr());
+    auto b_ptr = static_cast<ElementB const*>(b.data_ptr());
+    auto c_ptr = static_cast<ElementC*>(out.data_ptr());
+    auto scale_a = a_scales.data_ptr<float>();
+    auto scale_b = b_scales.data_ptr<float>();
+
+    ElementAccumulator* azp_ptr = NULL;
+    auto azp_adj_ptr = azp_adj.data_ptr<ElementAccumulator>();
+    ElementOutput* bias_t = static_cast<ElementOutput*>(bias.value().data_ptr());
+
+    if (azp) {
+      azp_ptr = static_cast<ElementAccumulator*>(azp.value().data_ptr());
+      mctlass::epilogue::thread::ScaleType::ScaleBiasKind const scale_type =
+          mctlass::epilogue::thread::ScaleType::ScaleAvBvBiasAzpPerTorken;
+      using mctlassGemmScaleOp =
+          mctlassGemmScale<ElementA, LayoutA, ElementB, LayoutB, ElementC,
+                           LayoutC, ElementCompute, ArchTag, scale_type>;
+      mctlassGemmScaleOp mctlass_op;
+      mctlass::gemm::GemmCoord problem_size(m, n, k);
+      typename mctlassGemmScaleOp::Arguments arguments{
+          mctlass::gemm::GemmUniversalMode::kGemm,
+          problem_size,
+          batchsize,
+          {scale_a, scale_b, bias_t, azp_adj_ptr, azp_ptr},
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          c_ptr,
+          m * k,
+          n * k,
+          m * n,
+          m * n,
+          k,
+          n,
+          n,
+          n
+      };
+      mctlass_op(arguments, NULL, stream);
+    } else {
+      mctlass::epilogue::thread::ScaleType::ScaleBiasKind const scale_type =
+          mctlass::epilogue::thread::ScaleType::ScaleAsBvBiasAzp;
+      using mctlassGemmScaleOp =
+          mctlassGemmScale<ElementA, LayoutA, ElementB, LayoutB, ElementC,
+                           LayoutC, ElementCompute, ArchTag, scale_type>;
+      mctlassGemmScaleOp mctlass_op;
+      mctlass::gemm::GemmCoord problem_size(m, n, k);
+      typename mctlassGemmScaleOp::Arguments arguments{
+          mctlass::gemm::GemmUniversalMode::kGemm,
+          problem_size,
+          batchsize,
+          {scale_a, scale_b, bias_t, azp_adj_ptr, azp_ptr},
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          c_ptr,
+          m * k,
+          n * k,
+          m * n,
+          m * n,
+          k,
+          n,
+          n,
+          n
+      };
+      mctlass_op(arguments, NULL, stream);
+    }
+  } else {
+    using ElementC = mctlass::half_t;
+    using ElementOutput = ElementC;
+
+    auto a_ptr = static_cast<ElementA const*>(a.data_ptr());
+    auto b_ptr = static_cast<ElementB const*>(b.data_ptr());
+    auto c_ptr = static_cast<ElementC*>(out.data_ptr());
+    auto scale_a = a_scales.data_ptr<float>();
+    auto scale_b = b_scales.data_ptr<float>();
+
+    ElementAccumulator* azp_ptr = nullptr;
+    auto azp_adj_ptr = azp_adj.data_ptr<ElementAccumulator>();
+    ElementOutput* bias_t = static_cast<ElementOutput*>(bias.value().data_ptr());
+
+    if (azp) {
+      azp_ptr = static_cast<ElementAccumulator*>(azp.value().data_ptr());
+
+      mctlass::epilogue::thread::ScaleType::ScaleBiasKind const scale_type =
+          mctlass::epilogue::thread::ScaleType::ScaleAvBvBiasAzpPerTorken;
+      using mctlassGemmScaleOp =
+          mctlassGemmScale<ElementA, LayoutA, ElementB, LayoutB, ElementC,
+                           LayoutC, ElementCompute, ArchTag, scale_type>;
+      mctlassGemmScaleOp mctlass_op;
+      mctlass::gemm::GemmCoord problem_size(m, n, k);
+      typename mctlassGemmScaleOp::Arguments arguments{
+          mctlass::gemm::GemmUniversalMode::kGemm,
+          problem_size,
+          batchsize,
+          {scale_a, scale_b, bias_t, azp_adj_ptr, azp_ptr},
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          c_ptr,
+          m * k,
+          n * k,
+          m * n,
+          m * n,
+          k,
+          n,
+          n,
+          n
+      };
+      mctlass_op(arguments, NULL, stream);
+    } else {
+      mctlass::epilogue::thread::ScaleType::ScaleBiasKind const scale_type =
+          mctlass::epilogue::thread::ScaleType::ScaleAsBvBiasAzp;
+      using mctlassGemmScaleOp =
+          mctlassGemmScale<ElementA, LayoutA, ElementB, LayoutB, ElementC,
+                           LayoutC, ElementCompute, ArchTag, scale_type>;
+      mctlassGemmScaleOp mctlass_op;
+      mctlass::gemm::GemmCoord problem_size(m, n, k);
+      typename mctlassGemmScaleOp::Arguments arguments{
+          mctlass::gemm::GemmUniversalMode::kGemm,
+          problem_size,
+          batchsize,
+          {scale_a, scale_b, bias_t, azp_adj_ptr, azp_ptr},
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          c_ptr,
+          m * k,
+          n * k,
+          m * n,
+          m * n,
+          k,
+          n,
+          n,
+          n
+      };
+      mctlass_op(arguments, NULL, stream);
+    }
+  }
+#endif //USE_MACA
 }
 
+#ifndef USE_MACA
 template <template <typename, typename> typename Epilogue,
           typename... EpilogueArgs>
 void cutlass_scaled_mm_sm80_epilogue(torch::Tensor& out, torch::Tensor const& a,
@@ -197,3 +567,4 @@ void cutlass_scaled_mm_azp_sm89(torch::Tensor& out, torch::Tensor const& a,
         out, a, b, a_scales, b_scales, azp_adj, bias);
   }
 }
+#endif // USE_MACA
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
index ce7cf2f35..096c90749 100644
--- a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
+++ b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
@@ -6,6 +6,7 @@
 
 // clang-format will break include orders
 // clang-format off
+#ifndef USE_MACA
 #include "cute/tensor.hpp"
 #include "cute/atom/mma_atom.hpp"
 #include "cutlass/numeric_types.h"
@@ -21,6 +22,14 @@
 #include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
 #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
 
+#else
+#include "mctlass/mctlass_ex.h"
+#include "mctlass/half.h"
+#include "mctlass/layout/matrix.h"
+#include "mctlass/epilogue/thread/scale_type.h"
+#include "mctlass/util/command_line.h"
+#endif // USE_MACA
+
 #include "core/math.hpp"
 #include "cutlass_extensions/common.hpp"
 // clang-format on
@@ -42,6 +51,7 @@ namespace vllm {
 // reduce the size of the compiled binary.
 // __CUDA_ARCH__ is not defined in host code, so this lets us smuggle the ifdef
 // into code that will be executed on the device where it is defined.
+#ifndef USE_MACA
 template <typename Kernel>
 struct enable_sm75_to_sm80 : Kernel {
   template <typename... Args>
@@ -51,7 +61,9 @@ struct enable_sm75_to_sm80 : Kernel {
 #endif
   }
 };
+#endif // USE_MACA
 
+#ifndef USE_MACA
 template <typename Kernel>
 struct enable_sm80_to_sm89 : Kernel {
   template <typename... Args>
@@ -128,11 +140,13 @@ struct cutlass_2x_gemm {
 
   using Op = cutlass::gemm::device::GemmUniversalAdapter<KernelType>;
 };
+#endif // USE_MACA
 
 template <typename Gemm, typename... EpilogueArgs>
 inline void cutlass_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
                                 torch::Tensor const& b,
                                 EpilogueArgs&&... epilogue_params) {
+#ifndef USE_MACA
   using ElementAB = typename Gemm::ElementAB;
   using ElementD = typename Gemm::ElementD;
 
@@ -193,6 +207,7 @@ inline void cutlass_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
   CUTLASS_CHECK(gemm_op.can_implement(args));
   cutlass::Status status = gemm_op(args, workspace.data_ptr(), stream);
   CUTLASS_CHECK(status);
+#endif // USE_MACA
 }
 
 template <typename Gemm, typename FallbackGemm, typename... EpilogueArgs>
@@ -200,6 +215,7 @@ inline void fallback_cutlass_gemm_caller(torch::Tensor& out,
                                          torch::Tensor const& a,
                                          torch::Tensor const& b,
                                          EpilogueArgs&&... args) {
+#ifndef USE_MACA
   // In some cases, the GPU isn't able to accommodate the
   // shared memory requirements of the Gemm. In such cases, use
   // the FallbackGemm instead.
@@ -220,6 +236,7 @@ inline void fallback_cutlass_gemm_caller(torch::Tensor& out,
     return cutlass_gemm_caller<FallbackGemm>(
         out, a, b, std::forward<EpilogueArgs>(args)...);
   }
+#endif // USE_MACA
 }
 
 }  // namespace vllm
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm75_dispatch.cuh b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm75_dispatch.cuh
index a562fd896..c7dec5784 100644
--- a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm75_dispatch.cuh
+++ b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm75_dispatch.cuh
@@ -9,6 +9,7 @@
 
 namespace vllm {
 
+#ifndef USE_MACA
 template <typename InType, typename OutType,
           template <typename, typename> typename Epilogue>
 struct sm75_config_default {
@@ -66,6 +67,7 @@ struct sm75_config_M32 {
       cutlass_2x_gemm<cutlass::arch::Sm75, enable_sm75_to_sm80, InType, OutType,
                       Epilogue, TileShape, WarpShape, InstructionShape, 2>;
 };
+#endif // USE_MACA
 
 template <typename InType, typename OutType,
           template <typename, typename> typename Epilogue,
@@ -74,6 +76,7 @@ inline void cutlass_gemm_sm75_dispatch(torch::Tensor& out,
                                        torch::Tensor const& a,
                                        torch::Tensor const& b,
                                        EpilogueArgs&&... args) {
+#ifndef USE_MACA
   static_assert(std::is_same<InType, int8_t>());
   TORCH_CHECK(a.dtype() == torch::kInt8);
   TORCH_CHECK(b.dtype() == torch::kInt8);
@@ -118,6 +121,7 @@ inline void cutlass_gemm_sm75_dispatch(torch::Tensor& out,
     return fallback_cutlass_gemm_caller<Cutlass2xGemmDefault, FallbackGemm>(
         out, a, b, std::forward<EpilogueArgs>(args)...);
   }
+#endif
 }
 
 }  // namespace vllm
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu b/csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
index 348525810..cc3b9f1bb 100644
--- a/csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+++ b/csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
@@ -149,6 +149,9 @@ void cutlass_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
                        torch::Tensor const& b, torch::Tensor const& a_scales,
                        torch::Tensor const& b_scales,
                        std::optional<torch::Tensor> const& bias) {
+#ifdef USE_MACA
+  cutlass_scaled_mm_sm75(c, a, b, a_scales, b_scales, bias);
+#else
   // Checks for conformality
   TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
   TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
@@ -209,6 +212,7 @@ void cutlass_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
       "No compiled cutlass_scaled_mm for a compute capability less than "
       "CUDA device capability: ",
       version_num);
+#endif // USE_MACA
 }
 
 void cutlass_moe_mm(
@@ -317,6 +321,29 @@ void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,
 
   at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
 
+#ifdef USE_MACA
+
+  if (!bias) {
+    // mctlass not support None bias
+
+    int32_t n = b.size(1);
+    int32_t batchsize = 1;
+    if (a.dim() == 3 && b.dim() == 3) {
+      // a.size = [batch_size, M, K], b.size = [batch_size, K, N]
+      n = b.size(2);
+      batchsize = a.size(0);
+    }
+    auto options = torch::TensorOptions()
+                     .dtype(c.dtype())
+                     .device(a.device());
+    torch::Tensor zero_bias = torch::zeros({batchsize,  n}, options);
+    cutlass_scaled_mm_azp_sm75(c, a, b, a_scales, b_scales, azp_adj, azp, zero_bias);
+  } else {
+    cutlass_scaled_mm_azp_sm75(c, a, b, a_scales, b_scales, azp_adj, azp, bias);
+  }
+
+#else
+
   int32_t version_num = get_sm_version_num();
 
 #if defined ENABLE_SCALED_MM_SM90 && ENABLE_SCALED_MM_SM90
@@ -350,4 +377,6 @@ void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,
       "No compiled cutlass_scaled_mm_azp for a compute capability less than "
       "CUDA device capability: ",
       version_num);
+
+#endif // USE_MACA
 }
diff --git a/csrc/quantization/fused_kernels/quant_conversions.cuh b/csrc/quantization/fused_kernels/quant_conversions.cuh
index 4e6118e52..435582408 100644
--- a/csrc/quantization/fused_kernels/quant_conversions.cuh
+++ b/csrc/quantization/fused_kernels/quant_conversions.cuh
@@ -32,7 +32,12 @@ static __device__ __forceinline__ int8_t float_to_int8_rn(float const x) {
 #else
   // CUDA path
   uint32_t dst;
+#ifndef USE_MACA
   asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
+#else
+  dst = (int32_t)(x > 0? x + 0.5: x - 0.5);
+  dst = (char)(dst > 127 ? 127 : (dst < -128 ? -128 : dst));
+#endif // USE_MACA
   return reinterpret_cast<const int8_t&>(dst);
 #endif
 }
diff --git a/csrc/quantization/gguf/ggml-common.h b/csrc/quantization/gguf/ggml-common.h
index 6bef5db3c..8dbab74a3 100644
--- a/csrc/quantization/gguf/ggml-common.h
+++ b/csrc/quantization/gguf/ggml-common.h
@@ -1147,4 +1147,12 @@ static __device__ __forceinline__ uint32_t __vsub4(const uint32_t a, const uint3
            (static_cast<uint8_t>(((a & 0x0000ff00) >>  8) - ((b & 0x0000ff00) >>  8)) <<  8) +
            (static_cast<uint8_t>(((a & 0x000000ff) >>  0) - ((b & 0x000000ff) >>  0)) <<  0);
 }
+#elif defined(USE_MACA)
+typedef int8_t int8x4_t __attribute__((ext_vector_type(4)));
+static __device__ __forceinline__ int __dp4a(const int a, const int b, int c) {
+    const int8x4_t va = reinterpret_cast<const int8x4_t&>(a);
+    const int8x4_t vb = reinterpret_cast<const int8x4_t&>(b);
+    c += va[0] * vb[0] + va[1] * vb[1] + va[2] * vb[2] + va[3] * vb[3];
+    return c;
+}
 #endif // defined(USE_ROCM)
diff --git a/csrc/quantization/gptq/Hgemm_common.cuh b/csrc/quantization/gptq/Hgemm_common.cuh
new file mode 100644
index 000000000..22723dd75
--- /dev/null
+++ b/csrc/quantization/gptq/Hgemm_common.cuh
@@ -0,0 +1,92 @@
+#pragma once
+#include "maca_fp16.h"
+
+using b32VecType = uint32_t;
+using b64VecType = __NATIVE_VECTOR__(2, uint32_t);
+using b128VecType = __NATIVE_VECTOR__(4, uint32_t);
+using b128VecType_i = __NATIVE_VECTOR__(4, int32_t);
+using Float4VecType = __NATIVE_VECTOR__(4, float);
+typedef enum
+{
+    OP_N = 0,
+    OP_T = 1
+} Operation_t;
+
+#define cast_half(ptr) reinterpret_cast<__half *>(ptr)
+#define cast_b16(ptr) reinterpret_cast<uint16_t *>(ptr)
+#define cast_b32(ptr) reinterpret_cast<uint32_t *>(ptr)
+#define cast_b64(ptr) reinterpret_cast<b64VecType *>(ptr)
+#define cast_u64(ptr) reinterpret_cast<uint64_t *>(ptr)
+#define cast_b128(ptr) reinterpret_cast<b128VecType *>(ptr)
+#define cast_b128_i(ptr) reinterpret_cast<b128VecType_i *>(ptr)
+
+#define ldg_b32_reg_noasync(dst, base, pred, ret0_en)                                                \
+    dst = __builtin_mxc_ldg_b32_predicator(cast_b32(base), 0, ret0_en, true, false, false, pred, 1, \
+                                           MACA_ICMP_EQ)[0];
+#define ldg_b64_reg_noasync(dst, base, pred, ret0_en)                                                \
+    dst = __builtin_mxc_ldg_b64_predicator(cast_b64(base), 0, ret0_en, true, false, false, pred, 1, \
+                                           MACA_ICMP_EQ);
+#define ldg_b128_reg_noasync(dst, base, pred, ret0_en)                                                \
+    dst = __builtin_mxc_ldg_b128_predicator(cast_b128(base), 0, ret0_en, true, false, false, pred, 1, \
+                                           MACA_ICMP_EQ);
+
+
+#define ldg_b32_reg_async(dst, base, pred, ret0_en)                                                \
+    dst = __builtin_mxc_ldg_b32_predicator(cast_b32(base), 0, ret0_en, true, false, true, pred, 1, \
+                                           MACA_ICMP_EQ)[0];
+#define ldg_b64_reg_async(dst, base, pred, ret0_en)                                                \
+    dst = __builtin_mxc_ldg_b64_predicator(cast_b64(base), 0, ret0_en, true, false, true, pred, 1, \
+                                           MACA_ICMP_EQ);
+#define ldg_b128_reg_async(dst, base, pred, ret0_en)                                                 \
+    dst = __builtin_mxc_ldg_b128_predicator(cast_b128(base), 0, ret0_en, true, false, true, pred, 1, \
+                                            MACA_ICMP_EQ);
+#define ldg_b64_v4h_reg_async(dst, base, pred, ret0_en)                                                \
+    dst = __builtin_mxc_ldg_b64_predicator(cast_b64(base), 0, ret0_en, true, false, true, pred, 1, \
+                                           MACA_ICMP_EQ);
+
+#define ldg_b32_bsm_noasync(saddr, base, pred, ret0_en)                                                    \
+    __builtin_mxc_ldg_b32_bsm_predicator(cast_b32(saddr), cast_b32(base), 0, ret0_en, true, false, false, \
+                                         pred, 1, MACA_ICMP_EQ);
+#define ldg_b64_bsm_noasync(saddr, base, pred, ret0_en)                                                    \
+    __builtin_mxc_ldg_b64_bsm_predicator(cast_b64(saddr), cast_b64(base), 0, ret0_en, true, false, false, \
+                                         pred, 1, MACA_ICMP_EQ);
+#define ldg_b128_bsm_noasync(saddr, base, pred, ret0_en)                                                \
+    __builtin_mxc_ldg_b128_bsm_predicator(cast_b128(saddr), cast_b128(base), 0, ret0_en, true, false, \
+                                          false, pred, 1, MACA_ICMP_EQ);
+
+
+#define ldg_b32_bsm_async(saddr, base, pred, ret0_en)                                                    \
+    __builtin_mxc_ldg_b32_bsm_predicator(cast_b32(saddr), cast_b32(base), 0, ret0_en, true, false, true, \
+                                         pred, 1, MACA_ICMP_EQ);
+#define ldg_b64_bsm_async(saddr, base, pred, ret0_en)                                                    \
+    __builtin_mxc_ldg_b64_bsm_predicator(cast_b64(saddr), cast_b64(base), 0, ret0_en, true, false, true, \
+                                         pred, 1, MACA_ICMP_EQ);
+#define ldg_b128_bsm_async(saddr, base, pred, ret0_en)                                                \
+    __builtin_mxc_ldg_b128_bsm_predicator(cast_b128(saddr), cast_b128(base), 0, ret0_en, true, false, \
+                                          true, pred, 1, MACA_ICMP_EQ);
+
+#define stg_b32_async(data, base, pred)                                                   \
+    __builtin_mxc_stg_b32_predicator(cast_b32(base), 0, data, true, false, true, pred, 1, \
+                                     MACA_ICMP_EQ);
+#define stg_b64_async(data, base, pred)                                                   \
+    __builtin_mxc_stg_b64_predicator(cast_b64(base), 0, data, true, false, true, pred, 1, \
+                                     MACA_ICMP_EQ);
+#define stg_b128_async(data, base, pred)                                                    \
+    __builtin_mxc_stg_b128_predicator(cast_b128(base), 0, data, true, false, true, pred, 1, \
+                                      MACA_ICMP_EQ);
+
+#define perm_b32(dst, reg1, reg2, selector) dst = __builtin_mxc_byte_perm(reg1, reg2, selector)
+#define mma_16x16x16f16(a_reg, b_reg, c_reg) \
+    c_reg = __builtin_mxc_mma_16x16x16f16(a_reg, b_reg, c_reg)
+
+#define mma_16x16x16bf16(a_reg, b_reg, c_reg) \
+    c_reg = __builtin_mxc_mma_16x16x16bf16(a_reg, b_reg, c_reg)
+
+#define FENCE__ asm volatile(";")
+#define arrive_gvmcnt(num) __builtin_mxc_arrive(64 + num)
+#define arrive_bsmcnt(num) __builtin_mxc_arrive(4096 + 128 * num)
+#define arrive_gvm_bsmcnt(gvm, bsm) __builtin_mxc_arrive(4096 | (128 * bsm) | 64 | gvm)
+#define barrier __builtin_mxc_barrier_inst
+#define barrier_all __builtin_mxc_barrier_ex(0)
+#define barrier_bsm __builtin_mxc_barrier_ex(1)
+#define barrier_inst __builtin_mxc_barrier_ex(2)
diff --git a/csrc/quantization/gptq/Hgemm_nn_128x32x128_8m1n8k_gptq-4bits.hpp b/csrc/quantization/gptq/Hgemm_nn_128x32x128_8m1n8k_gptq-4bits.hpp
new file mode 100644
index 000000000..5a97af2ad
--- /dev/null
+++ b/csrc/quantization/gptq/Hgemm_nn_128x32x128_8m1n8k_gptq-4bits.hpp
@@ -0,0 +1,426 @@
+#pragma once
+#include "Hgemm_common.cuh"
+#include "gptq.cuh"
+#include "maca_fp16.h"
+
+template <Operation_t TransA,
+          Operation_t TransB,
+          bool IsBetaZero,
+          int tileM,
+          int tileN,
+          int tileK,
+          bool doShuffle = true,
+          bool splitk = false,
+          bool Swap = false>
+__global__ __launch_bounds__(256) void Hgemm_nn_128x32x128_8m1n8k_gptq_4bit(int m,
+                                                                            int n,
+                                                                            int k,
+                                                                            const scalar_type alpha,    // alpha is 1.0f for gptq
+                                                                            const scalar_type beta,     // beta is 0.0f for gptq
+                                                                            const quant_packed_type *dA_input,
+                                                                            int lda,
+                                                                            const input_type *dB_input,
+                                                                            int ldb,
+                                                                            output_type *dC_input,
+                                                                            output_type *dC_output,
+                                                                            int ldc,
+                                                                            quant_packed_type *d_zeros,
+                                                                            input_type *d_scales,
+                                                                            int splitk_iters = 1,
+                                                                            acc_type * d_acc_tmp=nullptr)
+{
+    __shared__ uint8_t smem_base[0x8000];
+    int bidx = Swap ? blockIdx.y : blockIdx.x;
+    int bidy = Swap ? blockIdx.x : blockIdx.y;
+    int bidz = blockIdx.z;
+    int tid = threadIdx.x;
+
+    // this kernel only support NN trans mode
+    uint64_t arowstride = 1;
+    uint64_t acolstride = lda;
+    uint64_t browstride = 1;
+    uint64_t bcolstride = ldb;
+
+    const int malign = 8;
+    const int nalign = 1;
+    const int kalign = 8;
+    int align_m = (m + malign - 1) / malign * malign;
+    int align_n = (n + nalign - 1) / nalign * nalign;
+    int align_k = (k + kalign - 1) / kalign * kalign;
+
+    int k_num = (align_k + splitk_iters -1) / splitk_iters;
+    k_num = (k_num + tileK - 1) / tileK * tileK; // k_num should be aligned to tileK
+    int k_begin = bidz * k_num;
+    int k_end = (bidz + 1) * k_num;
+    k_end = k_end > align_k ? align_k : k_end;
+    k_num = k_end - k_begin;
+    int ktail = k_num % tileK;
+    int kloop = k_begin;
+
+    // if k_begin > align_k(we force k_num to be aligned to 8, so it is possible), return immediately
+    if (k_begin >= align_k) 
+    {
+        return;
+    }
+
+    int slot = tid / 64;
+    int lane = tid & 63;
+    int m64d16 = lane / 16;
+    int m64m16 = lane % 16;
+    output_type *c_base_i = dC_input;
+    output_type *c_base_o = dC_output;
+
+
+    quant_packed_type rga[8], rgzeros[1];
+    b128VecType rgb[2];
+    input_type rgscales[2];
+
+    b64VecType rgA[16], rgB[4], rgCi[4];
+    Float4VecType rgC[16];
+    uint32_t rgZeros[1];
+    b64VecType rgScales[2];
+
+    quant_packed_type *dA[2];
+    input_type *dB[2];
+
+    // ldg A/B head
+    bool predm[2], predn[2];
+    int rowA = m64m16 * 4;
+    int colA = (m64d16 * 8 + slot * 32) / PACK_RATIO_4BITS;
+    int current_m = bidx * tileM + rowA;
+    predm[0] = current_m < align_m;
+    dA[0] = (quant_packed_type *)dA_input + (uint64_t)(current_m) * (uint64_t)(arowstride) +
+            (uint64_t)(colA + k_begin / PACK_RATIO_4BITS) * (uint64_t)(acolstride);
+    current_m += 64;
+    predm[1] = current_m < align_m;
+    // dA[1] = (quant_packed_type *)dA_input + (uint64_t)(current_m) * (uint64_t)(arowstride) +
+    //         (uint64_t)(colA + k_begin / PACK_RATIO_4BITS) * (uint64_t)(acolstride);
+    dA[1] = dA[0] + 64 * (uint64_t)(arowstride);
+
+    int colB = m64d16 + slot * 4;
+    int rowB = m64m16 * 8;
+    int rowB_swizzle = SWIZZLE_INDEX(rowB, (colB % 4 * 8));
+    int current_n = bidy * tileN + colB;
+    predn[0] = current_n < align_n;
+    dB[0] = (input_type *)dB_input + (uint64_t)(rowB_swizzle + k_begin) * (uint64_t)(browstride) +
+            (uint64_t)(current_n) * (uint64_t)bcolstride;
+
+    current_n += 16;
+    predn[1] = current_n < align_n;
+    // dB[1] = (input_type *)dB_input + (uint64_t)(rowB_swizzle + k_begin) * (uint64_t)(browstride) +
+    //         (uint64_t)(current_n) * (uint64_t)bcolstride;
+    dB[1] = dB[0] + 16 * (uint64_t)bcolstride;
+
+    // ldgB to BSM need swizzle
+    input_type *b_sts = (input_type *)smem_base;
+    b_sts += ONE_DIM_INDEX(rowB, colB, 128);
+
+    input_type *b_lds = (input_type *)smem_base;
+    int colB_lds = m64m16;
+    int rowB_lds = m64d16 * 8 + slot * 32;
+    b_lds += ONE_DIM_INDEX(SWIZZLE_INDEX(rowB_lds, (colB_lds % 4 * 8)), colB_lds, 128);
+
+    // ldg zeros and scales
+    quant_packed_type *ldg_zeros_offset = d_zeros + (lda / PACK_RATIO_4BITS) * (kloop / tileK) + bidx * (tileM / PACK_RATIO_4BITS) + tid;
+    input_type *ldg_scales_offset = d_scales + lda * (kloop / tileK) + bidx * tileM + tid * 2;
+    quant_packed_type *sts_zeros_offset = (quant_packed_type *)(smem_base + tileK * tileN * sizeof(input_type)) + tid;
+    input_type *sts_scales_offset = (input_type *)(smem_base + tileK * tileN * sizeof(input_type) + tileM / PACK_RATIO_4BITS * sizeof(quant_packed_type)) + tid * 2;
+    uint16_t *lds_zeros_offset = (uint16_t *)(smem_base + tileK * tileN * sizeof(input_type)) + m64m16;
+    input_type *lds_scales_offset = (input_type *)(smem_base + tileK * tileN * sizeof(input_type) + tileM / PACK_RATIO_4BITS * sizeof(quant_packed_type)) + m64m16 * 4;
+
+    // clear C registers
+    rgC[0] = 0;
+    rgC[1] = 0;
+    rgC[2] = 0;
+    rgC[3] = 0;
+    rgC[4] = 0;
+    rgC[5] = 0;
+    rgC[6] = 0;
+    rgC[7] = 0;
+    rgC[8] = 0;
+    rgC[9] = 0;
+    rgC[10] = 0;
+    rgC[11] = 0;
+    rgC[12] = 0;
+    rgC[13] = 0;
+    rgC[14] = 0;
+    rgC[15] = 0;
+
+    int aincr = tileK * acolstride / PACK_RATIO_4BITS;
+    int bincr = tileK * browstride;
+
+    if (k_num > ktail)
+    {
+        LDG_B;      // 2
+        LDG_ZEROS_4BITS;  // 3
+        LDG_SCALES_4BITS; // 4
+        LDG_A1_4BITS;     // 5
+        LDG_A2_4BITS;     // 6
+
+        dA[0] += aincr;
+        dA[1] += aincr;
+        dB[0] += bincr;
+        dB[1] += bincr;
+
+        // lds_A, need swizzle
+        arrive_gvmcnt(2);
+        barrier();
+
+        LDS_B;
+        LDS_ZEROS_4BITS;
+        LDS_SCALES_4BITS;
+        arrive_bsmcnt(0);
+        barrier();
+        ldg_zeros_offset += lda / PACK_RATIO_4BITS;
+        ldg_scales_offset += lda;
+
+// KLOOP
+#pragma unroll 1
+        for (kloop = tileK; kloop + tileK <= k_num; kloop += tileK)
+        {
+            // load next B
+            LDG_B;      // 4
+            LDG_ZEROS_4BITS;  // 5
+            LDG_SCALES_4BITS; // 6
+
+            arrive_gvmcnt(5);
+            PERM_A1_4BITS;
+            LDG_A1_4BITS;
+            MMA1;
+
+            arrive_gvmcnt(5);
+            PERM_A2_4BITS;
+            LDG_A2_4BITS; // 6
+            MMA2;
+
+            // sts && lds
+            arrive_gvmcnt(2);
+            barrier();
+            LDS_B;
+            LDS_ZEROS_4BITS;
+            LDS_SCALES_4BITS;
+            arrive_bsmcnt(0);
+            barrier();
+
+            dA[0] += aincr;
+            dA[1] += aincr;
+            dB[0] += bincr;
+            dB[1] += bincr;
+            ldg_zeros_offset += lda / PACK_RATIO_4BITS;
+            ldg_scales_offset += lda;
+        }
+
+        arrive_gvmcnt(1);
+        PERM_A1_4BITS;
+        MMA1;
+
+        arrive_gvmcnt(0);
+        PERM_A2_4BITS;
+        MMA2;
+    }
+
+    // final tail kloop
+    if (ktail > 0)
+    {
+        LDG_B_HEAD;
+        LDG_ZEROS_4BITS;
+        LDG_SCALES_4BITS;
+        LDG_A1_HEAD_4BITS;
+        LDG_A2_HEAD_4BITS;
+        arrive_gvmcnt(2);
+        barrier();
+        LDS_B;
+        LDS_ZEROS_4BITS;
+        LDS_SCALES_4BITS;
+        arrive_bsmcnt(0);
+        barrier();
+
+        arrive_gvmcnt(1);
+        PERM_A1_4BITS;
+        MMA1;
+        arrive_gvmcnt(0);
+        PERM_A2_4BITS;
+        MMA2;
+    }
+
+    // store C registers into BSM & do reduction
+    int colC = m64m16 + slot * 16;
+    int rowC = m64d16 * 16;
+    scalar_type *c_sts[8];
+    c_sts[0] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC, (colC % 4 * 4)), colC, 128);
+    c_sts[1] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 4, (colC % 4 * 4)), colC, 128);
+    c_sts[2] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 8, (colC % 4 * 4)), colC, 128);
+    c_sts[3] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 12, (colC % 4 * 4)), colC, 128);
+    c_sts[4] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 64, (colC % 4 * 4)), colC, 128);
+    c_sts[5] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 68, (colC % 4 * 4)), colC, 128);
+    c_sts[6] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 72, (colC % 4 * 4)), colC, 128);
+    c_sts[7] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 76, (colC % 4 * 4)), colC, 128);
+
+    colC = m64d16 + slot * 4;
+    rowC = m64m16 * 4;
+    scalar_type *c_lds[2];
+    c_lds[0] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC, (colC % 4 * 4)), colC, 128);
+    c_lds[1] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 64, (colC % 4 * 4)), colC, 128);
+
+    PERM_C(rgC);
+    STS_C(0);
+    arrive_bsmcnt(0);
+    barrier();
+    REDUCE_C(0);
+
+    arrive_bsmcnt(0);
+    barrier();
+
+    PERM_C((rgC + 8));
+    STS_C(1);
+    arrive_bsmcnt(0);
+    barrier();
+    REDUCE_C(1);
+
+    arrive_bsmcnt(0);
+    barrier();
+
+    // read C_input if beta!=0
+    colC = m64d16 + slot * 4 + bidy * tileN;
+    rowC = m64m16 * 4 + bidx * tileM;
+    int colC1 = colC + 16;
+    int rowC1 = rowC + 64;
+
+
+    if constexpr (!IsBetaZero && !splitk)
+    {
+        ldg_b64_v4h_reg_async(cast_b64(rgCi)[0],
+                              c_base_i + (uint64_t)rowC + (uint64_t)(colC) * (uint64_t)(ldc),
+                              rowC < align_m && colC < align_n, true);
+        ldg_b64_v4h_reg_async(cast_b64(rgCi)[1], c_base_i + rowC1 + (uint64_t)(colC) * (uint64_t)(ldc),
+                              rowC1 < align_m && colC < align_n, true);
+        ldg_b64_v4h_reg_async(cast_b64(rgCi)[2], c_base_i + rowC + (uint64_t)(colC1) * (uint64_t)(ldc),
+                              rowC < align_m && colC1 < align_n, true);
+        ldg_b64_v4h_reg_async(cast_b64(rgCi)[3], c_base_i + rowC1 + (uint64_t)(colC1) * (uint64_t)(ldc),
+                              rowC1 < align_m && colC1 < align_n, true);
+        arrive_gvmcnt(0);
+    }
+    
+    if constexpr (IsBetaZero && !splitk)
+    {
+        for (int i = 0; i < 4; ++i) rgCi[i] = static_cast<b64VecType>(0);
+    }
+
+    if constexpr (!splitk)
+    {
+        output_type *ptrCi = reinterpret_cast<output_type *>(rgCi);
+        b64VecType *ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC) + (uint64_t)(colC) * (uint64_t)(ldc));
+        output_type result[4];
+        if (rowC < align_m && colC < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // result[0] = static_cast<output_type>(alpha * rgC[0][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            // result[1] = static_cast<output_type>(alpha * rgC[0][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            // result[2] = static_cast<output_type>(alpha * rgC[0][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            // result[3] = static_cast<output_type>(alpha * rgC[0][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            result[0] = static_cast<output_type>(rgC[0][0]);
+            result[1] = static_cast<output_type>(rgC[0][1]);
+            result[2] = static_cast<output_type>(rgC[0][2]);
+            result[3] = static_cast<output_type>(rgC[0][3]);
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC1) + (uint64_t)(colC) * (uint64_t)(ldc));
+        if (rowC1 < align_m && colC < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // result[0] = static_cast<output_type>(alpha * rgC[1][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            // result[1] = static_cast<output_type>(alpha * rgC[1][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            // result[2] = static_cast<output_type>(alpha * rgC[1][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            // result[3] = static_cast<output_type>(alpha * rgC[1][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            result[0] = static_cast<output_type>(rgC[1][0]);
+            result[1] = static_cast<output_type>(rgC[1][1]);
+            result[2] = static_cast<output_type>(rgC[1][2]);
+            result[3] = static_cast<output_type>(rgC[1][3]);
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC) + (uint64_t)(colC1) * (uint64_t)(ldc));
+        if (rowC < align_m && colC1 < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // result[0] = static_cast<output_type>(alpha * rgC[2][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            // result[1] = static_cast<output_type>(alpha * rgC[2][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            // result[2] = static_cast<output_type>(alpha * rgC[2][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            // result[3] = static_cast<output_type>(alpha * rgC[2][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            result[0] = static_cast<output_type>(rgC[2][0]);
+            result[1] = static_cast<output_type>(rgC[2][1]);
+            result[2] = static_cast<output_type>(rgC[2][2]);
+            result[3] = static_cast<output_type>(rgC[2][3]);
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC1) + (uint64_t)(colC1) * (uint64_t)(ldc));
+        if (rowC1 < align_m && colC1 < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // result[0] = static_cast<output_type>(alpha * rgC[3][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            // result[1] = static_cast<output_type>(alpha * rgC[3][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            // result[2] = static_cast<output_type>(alpha * rgC[3][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            // result[3] = static_cast<output_type>(alpha * rgC[3][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            result[0] = static_cast<output_type>(rgC[3][0]);
+            result[1] = static_cast<output_type>(rgC[3][1]);
+            result[2] = static_cast<output_type>(rgC[3][2]);
+            result[3] = static_cast<output_type>(rgC[3][3]);
+            ptrO[0] = cast_b64(result)[0];
+        }
+    }
+    else
+    {
+        // for acc_type is float
+        acc_type *ptrAcc = d_acc_tmp + (uint64_t)(rowC) + (uint64_t)(colC) * (uint64_t)(ldc);
+        if (rowC < align_m && colC < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][0]));
+            // atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][1]));
+            // atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][2]));
+            // atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][3]));
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(rgC[0][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(rgC[0][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(rgC[0][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(rgC[0][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC1) + (uint64_t)(colC) * (uint64_t)(ldc);
+        if (rowC1 < align_m && colC < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][0]));
+            // atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][1]));
+            // atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][2]));
+            // atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][3]));
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(rgC[1][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(rgC[1][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(rgC[1][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(rgC[1][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC) + (uint64_t)(colC1) * (uint64_t)(ldc);
+        if (rowC < align_m && colC1 < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][0]));
+            // atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][1]));
+            // atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][2]));
+            // atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][3]));
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(rgC[2][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(rgC[2][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(rgC[2][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(rgC[2][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC1) + (uint64_t)(colC1) * (uint64_t)(ldc);
+        if (rowC1 < align_m && colC1 < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][0]));
+            // atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][1]));
+            // atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][2]));
+            // atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][3]));
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(rgC[3][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(rgC[3][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(rgC[3][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(rgC[3][3]));
+        }
+    }
+}
diff --git a/csrc/quantization/gptq/Hgemm_nn_128x32x128_8m1n8k_gptq-8bits.hpp b/csrc/quantization/gptq/Hgemm_nn_128x32x128_8m1n8k_gptq-8bits.hpp
new file mode 100644
index 000000000..2cd807f57
--- /dev/null
+++ b/csrc/quantization/gptq/Hgemm_nn_128x32x128_8m1n8k_gptq-8bits.hpp
@@ -0,0 +1,440 @@
+#pragma once
+#include "Hgemm_common.cuh"
+#include "gptq.cuh"
+#include "maca_fp16.h"
+
+template <Operation_t TransA,
+          Operation_t TransB,
+          bool IsBetaZero,
+          int tileM,
+          int tileN,
+          int tileK,
+          bool doShuffle = true,
+          bool splitk = false,
+          bool Swap = false>
+__global__ __launch_bounds__(256) void Hgemm_nn_128x32x128_8m1n8k_gptq_8bit(int m,
+                                                                            int n,
+                                                                            int k,
+                                                                            const scalar_type alpha,    // alpha is 1.0f for gptq
+                                                                            const scalar_type beta,     // beta is 0.0f for gptq
+                                                                            const quant_packed_type *dA_input,
+                                                                            int lda,
+                                                                            const input_type *dB_input,
+                                                                            int ldb,
+                                                                            output_type *dC_input,
+                                                                            output_type *dC_output,
+                                                                            int ldc,
+                                                                            quant_packed_type *d_zeros,
+                                                                            input_type *d_scales,
+                                                                            int splitk_iters = 1,
+                                                                            acc_type * d_acc_tmp=nullptr,
+                                                                            half* dequanted = nullptr)
+{
+    __shared__ uint8_t smem_base[0x8000];
+    int bidx = Swap ? blockIdx.y : blockIdx.x;
+    int bidy = Swap ? blockIdx.x : blockIdx.y;
+    int bidz = blockIdx.z;
+    int tid = threadIdx.x;
+
+    // this kernel only support NN trans mode
+    uint64_t arowstride = 1;
+    uint64_t acolstride = lda;
+    uint64_t browstride = 1;
+    uint64_t bcolstride = ldb;
+
+    const int malign = 8;
+    const int nalign = 1;
+    const int kalign = 8;
+    int align_m = (m + malign - 1) / malign * malign;
+    int align_n = (n + nalign - 1) / nalign * nalign;
+    int align_k = (k + kalign - 1) / kalign * kalign;
+
+    int k_num = (align_k + splitk_iters -1) / splitk_iters;
+    k_num = (k_num + tileK - 1) / tileK * tileK; // k_num should be aligned to tileK
+    int k_begin = bidz * k_num;
+    int k_end = (bidz + 1) * k_num;
+    k_end = k_end > align_k ? align_k : k_end;
+    k_num = k_end - k_begin;
+    int ktail = k_num % tileK;
+    int kloop = k_begin;
+
+    // if k_begin > align_k(we force k_num to be aligned to 8, so it is possible), return immediately
+    if (k_begin >= align_k) 
+    {
+        return;
+    }
+
+    int slot = tid / 64;
+    int lane = tid & 63;
+    int m64d16 = lane / 16;
+    int m64m16 = lane % 16;
+    output_type *c_base_i = dC_input;
+    output_type *c_base_o = dC_output;
+
+
+    quant_packed_type rga[16], rgzeros[2];
+    b128VecType rgb[2];
+    input_type rgscales[2];
+
+    b64VecType rgA[16], rgB[4], rgCi[4];
+    Float4VecType rgC[16];
+    uint32_t rgZeros[2];
+    b64VecType rgScales[2];
+
+    quant_packed_type *dA[4];
+    input_type *dB[2];
+
+    // ldg A/B head
+    bool predm[2], predn[2];
+    int rowA = m64m16 * 4;
+    int colA = (m64d16 * 8 + slot * 32) / PACK_RATIO_8BITS;
+    int current_m = bidx * tileM + rowA;
+    predm[0] = current_m < align_m;
+    dA[0] = (quant_packed_type *)dA_input + (uint64_t)(current_m) * (uint64_t)(arowstride) +
+            (uint64_t)(colA + k_begin / PACK_RATIO_8BITS) * (uint64_t)(acolstride);
+    dA[2] = dA[0] + align_m;
+    current_m += 64;
+    predm[1] = current_m < align_m;
+    // dA[1] = (quant_packed_type *)dA_input + (uint64_t)(current_m) * (uint64_t)(arowstride) +
+    //         (uint64_t)(colA + k_begin / PACK_RATIO_8BITS) * (uint64_t)(acolstride);
+    dA[1] = dA[0] + 64 * (uint64_t)(arowstride);
+    dA[3] = dA[1] + align_m;
+
+    int colB = m64d16 + slot * 4;
+    int rowB = m64m16 * 8;
+    int rowB_swizzle = SWIZZLE_INDEX(rowB, (colB % 4 * 8));
+    int current_n = bidy * tileN + colB;
+    predn[0] = current_n < align_n;
+    dB[0] = (input_type *)dB_input + (uint64_t)(rowB_swizzle + k_begin) * (uint64_t)(browstride) +
+            (uint64_t)(current_n) * (uint64_t)bcolstride;
+
+    current_n += 16;
+    predn[1] = current_n < align_n;
+    // dB[1] = (input_type *)dB_input + (uint64_t)(rowB_swizzle + k_begin) * (uint64_t)(browstride) +
+    //         (uint64_t)(current_n) * (uint64_t)bcolstride;
+    dB[1] = dB[0] + 16 * (uint64_t)bcolstride;
+
+    // ldgB to BSM need swizzle
+    input_type *b_sts = (input_type *)smem_base;
+    b_sts += ONE_DIM_INDEX(rowB, colB, 128);
+
+    input_type *b_lds = (input_type *)smem_base;
+    int colB_lds = m64m16;
+    int rowB_lds = m64d16 * 8 + slot * 32;
+    b_lds += ONE_DIM_INDEX(SWIZZLE_INDEX(rowB_lds, (colB_lds % 4 * 8)), colB_lds, 128);
+
+    // ldg zeros and scales
+    quant_packed_type *ldg_zeros_offset = d_zeros + (lda / PACK_RATIO_8BITS) * (kloop / tileK) + bidx * (tileM / PACK_RATIO_8BITS) + tid;
+    input_type *ldg_scales_offset = d_scales + lda * (kloop / tileK) + bidx * tileM + tid * 2;
+    quant_packed_type *sts_zeros_offset = (quant_packed_type *)(smem_base + tileK * tileN * sizeof(input_type)) + tid;
+    input_type *sts_scales_offset = (input_type *)(smem_base + tileK * tileN * sizeof(input_type) + tileM / PACK_RATIO_8BITS * sizeof(quant_packed_type)) + tid * 2;
+    quant_packed_type *lds_zeros_offset = (quant_packed_type *)(smem_base + tileK * tileN * sizeof(input_type)) + m64m16;
+    input_type *lds_scales_offset = (input_type *)(smem_base + tileK * tileN * sizeof(input_type) + tileM / PACK_RATIO_8BITS * sizeof(quant_packed_type)) + m64m16 * 4;
+
+    // clear C registers
+    rgC[0] = 0;
+    rgC[1] = 0;
+    rgC[2] = 0;
+    rgC[3] = 0;
+    rgC[4] = 0;
+    rgC[5] = 0;
+    rgC[6] = 0;
+    rgC[7] = 0;
+    rgC[8] = 0;
+    rgC[9] = 0;
+    rgC[10] = 0;
+    rgC[11] = 0;
+    rgC[12] = 0;
+    rgC[13] = 0;
+    rgC[14] = 0;
+    rgC[15] = 0;
+
+    int aincr = tileK * acolstride / PACK_RATIO_8BITS;
+    int bincr = tileK * browstride;
+
+    if (k_num > ktail)
+    {
+        LDG_B;      // 2
+        LDG_ZEROS_8BITS;  // 3
+        LDG_SCALES_8BITS; // 4
+        LDG_A1_8BITS;     // 5
+        LDG_A3_8BITS;     // 6
+        LDG_A2_8BITS;     // 7
+        LDG_A4_8BITS;     // 8
+
+        dA[0] += aincr;
+        dA[1] += aincr;
+        dA[2] += aincr;
+        dA[3] += aincr;
+        dB[0] += bincr;
+        dB[1] += bincr;
+
+        // lds_A, need swizzle
+        arrive_gvmcnt(4);
+        barrier();
+
+        LDS_B;
+        LDS_ZEROS_8BITS;
+        LDS_SCALES_8BITS;
+        arrive_bsmcnt(0);
+        barrier();
+        ldg_zeros_offset += lda / PACK_RATIO_8BITS;
+        ldg_scales_offset += lda;
+        constexpr bool loading_filter = false;//blockIdx.x == 0 && threadIdx.x == 0;
+
+// KLOOP
+#pragma unroll 1
+        for (kloop = tileK; kloop + tileK <= k_num; kloop += tileK)
+        {
+            // load next B
+            LDG_B;      // 6
+            LDG_ZEROS_8BITS;  // 7
+            LDG_SCALES_8BITS; // 8
+
+            arrive_gvmcnt(6);
+            PERM_A1A3_8BITS;
+            LDG_A1_8BITS;
+            LDG_A3_8BITS;
+            MMA1;
+
+            arrive_gvmcnt(6);
+            PERM_A2A4_8BITS;
+            LDG_A2_8BITS; // 6
+            LDG_A4_8BITS;
+            MMA2;
+
+            // sts && lds
+            arrive_gvmcnt(4);
+            barrier();
+            LDS_B;
+            LDS_ZEROS_8BITS;
+            LDS_SCALES_8BITS;
+            arrive_bsmcnt(0);
+            barrier();
+
+            dA[0] += aincr;
+            dA[1] += aincr;
+            dA[2] += aincr;
+            dA[3] += aincr;
+            dB[0] += bincr;
+            dB[1] += bincr;
+            ldg_zeros_offset += lda / PACK_RATIO_8BITS;
+            ldg_scales_offset += lda;
+        }
+
+        arrive_gvmcnt(2);
+        PERM_A1A3_8BITS;
+        MMA1;
+
+        arrive_gvmcnt(0);
+        PERM_A2A4_8BITS;
+        MMA2;
+    }
+
+    // final tail kloop
+    if (ktail > 0)
+    {
+        LDG_B_HEAD;
+        LDG_ZEROS_8BITS;
+        LDG_SCALES_8BITS;
+        LDG_A1_HEAD_8BITS;
+        LDG_A3_HEAD_8BITS;
+        LDG_A2_HEAD_8BITS;
+        LDG_A4_HEAD_8BITS;
+        arrive_gvmcnt(4);
+        barrier();
+        LDS_B;
+        LDS_ZEROS_8BITS;
+        LDS_SCALES_8BITS;
+        arrive_bsmcnt(0);
+        barrier();
+
+        arrive_gvmcnt(2);
+        PERM_A1A3_8BITS;
+        MMA1;
+        arrive_gvmcnt(0);
+        PERM_A2A4_8BITS;
+        MMA2;
+    }
+
+    // store C registers into BSM & do reduction
+    int colC = m64m16 + slot * 16;
+    int rowC = m64d16 * 16;
+    scalar_type *c_sts[8];
+    c_sts[0] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC, (colC % 4 * 4)), colC, 128);
+    c_sts[1] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 4, (colC % 4 * 4)), colC, 128);
+    c_sts[2] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 8, (colC % 4 * 4)), colC, 128);
+    c_sts[3] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 12, (colC % 4 * 4)), colC, 128);
+    c_sts[4] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 64, (colC % 4 * 4)), colC, 128);
+    c_sts[5] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 68, (colC % 4 * 4)), colC, 128);
+    c_sts[6] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 72, (colC % 4 * 4)), colC, 128);
+    c_sts[7] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 76, (colC % 4 * 4)), colC, 128);
+
+    colC = m64d16 + slot * 4;
+    rowC = m64m16 * 4;
+    scalar_type *c_lds[2];
+    c_lds[0] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC, (colC % 4 * 4)), colC, 128);
+    c_lds[1] = (scalar_type *)smem_base + ONE_DIM_INDEX(SWIZZLE_INDEX(rowC + 64, (colC % 4 * 4)), colC, 128);
+
+    PERM_C(rgC);
+    STS_C(0);
+    arrive_bsmcnt(0);
+    barrier();
+    REDUCE_C(0);
+
+    arrive_bsmcnt(0);
+    barrier();
+
+    PERM_C((rgC + 8));
+    STS_C(1);
+    arrive_bsmcnt(0);
+    barrier();
+    REDUCE_C(1);
+
+    arrive_bsmcnt(0);
+    barrier();
+
+    // read C_input if beta!=0
+    colC = m64d16 + slot * 4 + bidy * tileN;
+    rowC = m64m16 * 4 + bidx * tileM;
+    int colC1 = colC + 16;
+    int rowC1 = rowC + 64;
+
+
+    if constexpr (!IsBetaZero && !splitk)
+    {
+        ldg_b64_v4h_reg_async(cast_b64(rgCi)[0],
+                              c_base_i + (uint64_t)rowC + (uint64_t)(colC) * (uint64_t)(ldc),
+                              rowC < align_m && colC < align_n, true);
+        ldg_b64_v4h_reg_async(cast_b64(rgCi)[1], c_base_i + rowC1 + (uint64_t)(colC) * (uint64_t)(ldc),
+                              rowC1 < align_m && colC < align_n, true);
+        ldg_b64_v4h_reg_async(cast_b64(rgCi)[2], c_base_i + rowC + (uint64_t)(colC1) * (uint64_t)(ldc),
+                              rowC < align_m && colC1 < align_n, true);
+        ldg_b64_v4h_reg_async(cast_b64(rgCi)[3], c_base_i + rowC1 + (uint64_t)(colC1) * (uint64_t)(ldc),
+                              rowC1 < align_m && colC1 < align_n, true);
+        arrive_gvmcnt(0);
+    }
+    
+    if constexpr (IsBetaZero && !splitk)
+    {
+        for (int i = 0; i < 4; ++i) rgCi[i] = static_cast<b64VecType>(0);
+    }
+
+    if constexpr (!splitk)
+    {
+        output_type *ptrCi = reinterpret_cast<output_type *>(rgCi);
+        b64VecType *ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC) + (uint64_t)(colC) * (uint64_t)(ldc));
+        output_type result[4];
+        if (rowC < align_m && colC < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // result[0] = static_cast<output_type>(alpha * rgC[0][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            // result[1] = static_cast<output_type>(alpha * rgC[0][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            // result[2] = static_cast<output_type>(alpha * rgC[0][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            // result[3] = static_cast<output_type>(alpha * rgC[0][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            result[0] = static_cast<output_type>(rgC[0][0]);
+            result[1] = static_cast<output_type>(rgC[0][1]);
+            result[2] = static_cast<output_type>(rgC[0][2]);
+            result[3] = static_cast<output_type>(rgC[0][3]);
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC1) + (uint64_t)(colC) * (uint64_t)(ldc));
+        if (rowC1 < align_m && colC < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // result[0] = static_cast<output_type>(alpha * rgC[1][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            // result[1] = static_cast<output_type>(alpha * rgC[1][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            // result[2] = static_cast<output_type>(alpha * rgC[1][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            // result[3] = static_cast<output_type>(alpha * rgC[1][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            result[0] = static_cast<output_type>(rgC[1][0]);
+            result[1] = static_cast<output_type>(rgC[1][1]);
+            result[2] = static_cast<output_type>(rgC[1][2]);
+            result[3] = static_cast<output_type>(rgC[1][3]);
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC) + (uint64_t)(colC1) * (uint64_t)(ldc));
+        if (rowC < align_m && colC1 < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // result[0] = static_cast<output_type>(alpha * rgC[2][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            // result[1] = static_cast<output_type>(alpha * rgC[2][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            // result[2] = static_cast<output_type>(alpha * rgC[2][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            // result[3] = static_cast<output_type>(alpha * rgC[2][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            result[0] = static_cast<output_type>(rgC[2][0]);
+            result[1] = static_cast<output_type>(rgC[2][1]);
+            result[2] = static_cast<output_type>(rgC[2][2]);
+            result[3] = static_cast<output_type>(rgC[2][3]);
+            ptrO[0] = cast_b64(result)[0];
+        }
+        ptrO = reinterpret_cast<b64VecType *>(c_base_o + (uint64_t)(rowC1) + (uint64_t)(colC1) * (uint64_t)(ldc));
+        if (rowC1 < align_m && colC1 < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // result[0] = static_cast<output_type>(alpha * rgC[3][0] + beta * static_cast<scalar_type>(ptrCi[0]));
+            // result[1] = static_cast<output_type>(alpha * rgC[3][1] + beta * static_cast<scalar_type>(ptrCi[1]));
+            // result[2] = static_cast<output_type>(alpha * rgC[3][2] + beta * static_cast<scalar_type>(ptrCi[2]));
+            // result[3] = static_cast<output_type>(alpha * rgC[3][3] + beta * static_cast<scalar_type>(ptrCi[3]));
+            result[0] = static_cast<output_type>(rgC[3][0]);
+            result[1] = static_cast<output_type>(rgC[3][1]);
+            result[2] = static_cast<output_type>(rgC[3][2]);
+            result[3] = static_cast<output_type>(rgC[3][3]);
+            ptrO[0] = cast_b64(result)[0];
+        }
+    }
+    else
+    {
+        // for acc_type is float
+        acc_type *ptrAcc = d_acc_tmp + (uint64_t)(rowC) + (uint64_t)(colC) * (uint64_t)(ldc);
+        if (rowC < align_m && colC < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][0]));
+            // atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][1]));
+            // atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][2]));
+            // atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[0][3]));
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(rgC[0][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(rgC[0][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(rgC[0][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(rgC[0][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC1) + (uint64_t)(colC) * (uint64_t)(ldc);
+        if (rowC1 < align_m && colC < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][0]));
+            // atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][1]));
+            // atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][2]));
+            // atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[1][3]));
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(rgC[1][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(rgC[1][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(rgC[1][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(rgC[1][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC) + (uint64_t)(colC1) * (uint64_t)(ldc);
+        if (rowC < align_m && colC1 < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][0]));
+            // atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][1]));
+            // atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][2]));
+            // atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[2][3]));
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(rgC[2][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(rgC[2][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(rgC[2][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(rgC[2][3]));
+        }
+        ptrAcc = d_acc_tmp + (uint64_t)(rowC1) + (uint64_t)(colC1) * (uint64_t)(ldc);
+        if (rowC1 < align_m && colC1 < align_n)
+        {
+            // for gptq, alpha is 1.0f, beta is 0.0f
+            // atomicAdd(&ptrAcc[0], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][0]));
+            // atomicAdd(&ptrAcc[1], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][1]));
+            // atomicAdd(&ptrAcc[2], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][2]));
+            // atomicAdd(&ptrAcc[3], static_cast<acc_type>(alpha) * static_cast<acc_type>(rgC[3][3]));
+            atomicAdd(&ptrAcc[0], static_cast<acc_type>(rgC[3][0]));
+            atomicAdd(&ptrAcc[1], static_cast<acc_type>(rgC[3][1]));
+            atomicAdd(&ptrAcc[2], static_cast<acc_type>(rgC[3][2]));
+            atomicAdd(&ptrAcc[3], static_cast<acc_type>(rgC[3][3]));
+        }
+    }
+}
diff --git a/csrc/quantization/gptq/dequant.cuh b/csrc/quantization/gptq/dequant.cuh
new file mode 100644
index 000000000..b192a7804
--- /dev/null
+++ b/csrc/quantization/gptq/dequant.cuh
@@ -0,0 +1,11 @@
+// 2024 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
+#pragma once
+
+#define BFE(source, bit_start, num_bits) __builtin_mxc_ubfe(source, bit_start, num_bits);
+
+template <typename outputT, typename inputT, int qbits>
+__device__ __forceinline__ outputT extract(inputT *input, size_t loc)
+{
+    constexpr int widthPerElem = sizeof(inputT) * 8 / qbits;
+    return (outputT)BFE(input[loc / widthPerElem], loc % widthPerElem, qbits);
+}
\ No newline at end of file
diff --git a/csrc/quantization/gptq/gptq.cuh b/csrc/quantization/gptq/gptq.cuh
new file mode 100644
index 000000000..c91bf4168
--- /dev/null
+++ b/csrc/quantization/gptq/gptq.cuh
@@ -0,0 +1,365 @@
+#pragma once
+
+#include "Hgemm_common.cuh"
+#include "dequant.cuh"
+#define quant_packed_type uint32_t
+
+#define Q4BITS 4
+#define PACK_RATIO_4BITS (32 / Q4BITS)
+
+#define Q8BITS 8
+#define PACK_RATIO_8BITS (32 / Q8BITS)
+
+#define input_type __half
+#define output_type __half
+#define scalar_type float
+//#define acc_type float
+#define acc_type __half
+
+#define SEL0 0x01000504
+#define SEL1 0x03020706
+
+#define SWIZZLE_INDEX(index, phase) (index ^ phase)
+#define ONE_DIM_INDEX(x, y, lda) (y * lda + x)
+#define CLAMP(value, bound, align) (value < bound ? value : bound - align)
+
+#define LDG_A1                                                        \
+    {                                                                 \
+        ldg_b128_reg_async(cast_b128(rga)[0], dA[0], predm[0], true); \
+    }
+
+#define LDG_A2                                                        \
+    {                                                                 \
+        ldg_b128_reg_async(cast_b128(rga)[1], dA[1], predm[1], true); \
+    }
+
+#define LDG_A3                                                        \
+    {                                                                 \
+        ldg_b128_reg_async(cast_b128(rga)[2], dA[2], predm[0], true); \
+    }
+
+#define LDG_A4                                                        \
+    {                                                                 \
+        ldg_b128_reg_async(cast_b128(rga)[3], dA[3], predm[1], true); \
+    }
+
+#define LDG_A1_4BITS LDG_A1
+#define LDG_A2_4BITS LDG_A2
+#define LDG_A1_8BITS LDG_A1
+#define LDG_A2_8BITS LDG_A2
+#define LDG_A3_8BITS LDG_A3
+#define LDG_A4_8BITS LDG_A4
+
+#define LDG_A1_HEAD_4BITS                                                      \
+    {                                                                          \
+        bool predk = colA < ktail / PACK_RATIO_4BITS;                          \
+        ldg_b128_reg_async(cast_b128(rga)[0], dA[0], predm[0] && predk, true); \
+    }
+
+#define LDG_A2_HEAD_4BITS                                                      \
+    {                                                                          \
+        bool predk = colA < ktail / PACK_RATIO_4BITS;                          \
+        ldg_b128_reg_async(cast_b128(rga)[1], dA[1], predm[1] && predk, true); \
+    }
+
+#define LDG_A1_HEAD_8BITS                                                      \
+    {                                                                          \
+        bool predk = colA < ktail / PACK_RATIO_8BITS;                          \
+        ldg_b128_reg_async(cast_b128(rga)[0], dA[0], predm[0] && predk, true); \
+    }
+
+#define LDG_A2_HEAD_8BITS                                                      \
+    {                                                                          \
+        bool predk = colA < ktail / PACK_RATIO_8BITS;                          \
+        ldg_b128_reg_async(cast_b128(rga)[1], dA[1], predm[1] && predk, true); \
+    }
+
+#define LDG_A3_HEAD_8BITS                                                      \
+    {                                                                          \
+        bool predk = colA < ktail / PACK_RATIO_8BITS;                          \
+        ldg_b128_reg_async(cast_b128(rga)[2], dA[2], predm[0] && predk, true); \
+    }
+
+#define LDG_A4_HEAD_8BITS                                                      \
+    {                                                                          \
+        bool predk = colA < ktail / PACK_RATIO_8BITS;                          \
+        ldg_b128_reg_async(cast_b128(rga)[3], dA[3], predm[1] && predk, true); \
+    }
+
+#define LDG_ZEROS_4BITS                                                                                                                        \
+    {                                                                                                                                          \
+        ldg_b32_bsm_async(sts_zeros_offset, ldg_zeros_offset, tid < tileM / PACK_RATIO_4BITS && (bidx * tileM + tid * PACK_RATIO_4BITS < align_m), false); \
+    }
+
+#define LDG_ZEROS_8BITS                                                                                                                        \
+    {                                                                                                                                          \
+        ldg_b32_bsm_async(sts_zeros_offset, ldg_zeros_offset, tid < tileM / PACK_RATIO_8BITS && (bidx * tileM + tid * PACK_RATIO_8BITS < align_m), false); \
+    }
+
+#define LDG_SCALES                                                                                                      \
+    {                                                                                                                   \
+        ldg_b32_bsm_async(sts_scales_offset, ldg_scales_offset, tid < 64 && (bidx * tileM + tid * 2 < align_m), false); \
+    }
+
+#define LDG_SCALES_4BITS LDG_SCALES
+#define LDG_SCALES_8BITS LDG_SCALES
+
+#define LDS_ZEROS_4BITS                              \
+    {                                                \
+        cast_b16(rgZeros)[0] = lds_zeros_offset[0];  \
+        cast_b16(rgZeros)[1] = lds_zeros_offset[16]; \
+    }
+
+#define LDS_ZEROS_8BITS                              \
+    {                                                \
+        cast_b32(rgZeros)[0] = lds_zeros_offset[0];  \
+        cast_b32(rgZeros)[1] = lds_zeros_offset[16]; \
+    }
+
+#define LDS_SCALES                                               \
+    {                                                            \
+        cast_b64(rgScales)[0] = cast_b64(lds_scales_offset)[0];  \
+        cast_b64(rgScales)[1] = cast_b64(lds_scales_offset)[16]; \
+    }
+
+#define LDS_SCALES_4BITS LDS_SCALES
+#define LDS_SCALES_8BITS LDS_SCALES
+
+#define PERM_ELEM_4BITS(index)                                                                                                          \
+    {                                                                                                                                   \
+        __half_raw elem;                                                                                                                \
+        if constexpr (index & 0x1)                                                                                                      \
+        {                                                                                                                               \
+            elem.x = rgScales[index / 4][index % 4 / 2] >> 16;                                                                          \
+        }                                                                                                                               \
+        else                                                                                                                            \
+        {                                                                                                                               \
+            elem.x = rgScales[index / 4][index % 4 / 2] & 0xffff;                                                                       \
+        }                                                                                                                               \
+        __half scale = __half(elem);                                                                                                    \
+        uint32_t zero = __builtin_mxc_ubfe(rgZeros[0], Q4BITS * index, Q4BITS) + 1;                                                     \
+        if constexpr (doShuffle)                                                                                                        \
+        {                                                                                                                               \
+            cast_half(rgA)[index * 4 + 0] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 0, Q4BITS) -       \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 1] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 4, Q4BITS) -       \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 2] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 1, Q4BITS) -       \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 3] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 5, Q4BITS) -       \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 32] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 2, Q4BITS) -      \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 33] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 6, Q4BITS) -      \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 34] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 3, Q4BITS) -      \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 35] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 7, Q4BITS) -      \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+        }                                                                                                                               \
+        else                                                                                                                            \
+        {                                                                                                                               \
+            cast_half(rgA)[index * 4 + 0] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 0, Q4BITS) -       \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 1] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 1, Q4BITS) -       \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 2] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 2, Q4BITS) -       \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 3] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 3, Q4BITS) -       \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 32] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 4, Q4BITS) -      \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 33] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 5, Q4BITS) -      \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 34] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 6, Q4BITS) -      \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+            cast_half(rgA)[index * 4 + 35] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q4BITS * 7, Q4BITS) -      \
+                                                                            zero) ,                                                     \
+                                            scale);                                                                                      \
+        }                                                                                                                               \
+    }
+
+#define PERM_A1_4BITS       \
+    {                       \
+        PERM_ELEM_4BITS(0)  \
+        PERM_ELEM_4BITS(1)  \
+        PERM_ELEM_4BITS(2)  \
+        PERM_ELEM_4BITS(3)  \
+    }
+
+#define PERM_A2_4BITS       \
+    {                       \
+        PERM_ELEM_4BITS(4)  \
+        PERM_ELEM_4BITS(5)  \
+        PERM_ELEM_4BITS(6)  \
+        PERM_ELEM_4BITS(7)  \
+    }
+
+#define PERM_ELEM_8BITS(index)                                                                                                          \
+    {                                                                                                                                   \
+        __half_raw elem;                                                                                                                \
+        if constexpr (index & 0x1)                                                                                                      \
+        {                                                                                                                               \
+            elem.x = rgScales[index / 4][index % 4 / 2] >> 16;                                                                          \
+        }                                                                                                                               \
+        else                                                                                                                            \
+        {                                                                                                                               \
+            elem.x = rgScales[index / 4][index % 4 / 2] & 0xffff;                                                                       \
+        }                                                                                                                               \
+        __half scale = __half(elem);                                                                                                    \
+        uint32_t zero = __builtin_mxc_ubfe(rgZeros[index/4], Q8BITS * index, Q8BITS) + 1;                                               \
+                                                                                                                                        \
+        cast_half(rgA)[index * 4 + 0] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q8BITS * 0, Q8BITS) -           \
+                                                                        zero) ,                                                         \
+                                        scale);                                                                                          \
+        cast_half(rgA)[index * 4 + 1] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q8BITS * 1, Q8BITS) -           \
+                                                                        zero) ,                                                         \
+                                        scale);                                                                                          \
+        cast_half(rgA)[index * 4 + 2] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q8BITS * 2, Q8BITS) -           \
+                                                                        zero) ,                                                         \
+                                        scale);                                                                                          \
+        cast_half(rgA)[index * 4 + 3] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index], Q8BITS * 3, Q8BITS) -           \
+                                                                        zero) ,                                                         \
+                                        scale);                                                                                          \
+        cast_half(rgA)[index * 4 + 32] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index+8], Q8BITS * 0, Q8BITS) -        \
+                                                                        zero) ,                                                         \
+                                        scale);                                                                                          \
+        cast_half(rgA)[index * 4 + 33] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index+8], Q8BITS * 1, Q8BITS) -        \
+                                                                        zero) ,                                                         \
+                                        scale);                                                                                          \
+        cast_half(rgA)[index * 4 + 34] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index+8], Q8BITS * 2, Q8BITS) -        \
+                                                                        zero) ,                                                         \
+                                        scale);                                                                                          \
+        cast_half(rgA)[index * 4 + 35] = __hmul((__half)__builtin_mxc_i16_to_f16(__builtin_mxc_ubfe(rga[index+8], Q8BITS * 3, Q8BITS) -        \
+                                                                        zero) ,                                                         \
+                                        scale);                                                                                          \
+    }
+
+#define PERM_A1A3_8BITS     \
+    {                       \
+        PERM_ELEM_8BITS(0)  \
+        PERM_ELEM_8BITS(1)  \
+        PERM_ELEM_8BITS(2)  \
+        PERM_ELEM_8BITS(3)  \
+    }
+
+#define PERM_A2A4_8BITS     \
+    {                       \
+        PERM_ELEM_8BITS(4)  \
+        PERM_ELEM_8BITS(5)  \
+        PERM_ELEM_8BITS(6)  \
+        PERM_ELEM_8BITS(7)  \
+    }
+
+#define MMA_ELEM(index_m)                                        \
+    mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 0, 8)], \
+                    cast_b64(rgB)[ONE_DIM_INDEX(0, 0, 2)],       \
+                    rgC[ONE_DIM_INDEX(index_m, 0, 8)]);          \
+    mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 0, 8)], \
+                    cast_b64(rgB)[ONE_DIM_INDEX(0, 1, 2)],       \
+                    rgC[ONE_DIM_INDEX(index_m, 1, 8)]);          \
+    mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 1, 8)], \
+                    cast_b64(rgB)[ONE_DIM_INDEX(1, 0, 2)],       \
+                    rgC[ONE_DIM_INDEX(index_m, 0, 8)]);          \
+    mma_16x16x16f16(cast_b64(rgA)[ONE_DIM_INDEX(index_m, 1, 8)], \
+                    cast_b64(rgB)[ONE_DIM_INDEX(1, 1, 2)],       \
+                    rgC[ONE_DIM_INDEX(index_m, 1, 8)]);
+
+#define MMA1    \
+    MMA_ELEM(0) \
+    MMA_ELEM(1) \
+    MMA_ELEM(2) \
+    MMA_ELEM(3)
+
+#define MMA2    \
+    MMA_ELEM(4) \
+    MMA_ELEM(5) \
+    MMA_ELEM(6) \
+    MMA_ELEM(7)
+
+#define LDG_B                                                        \
+    {                                                                \
+        ldg_b128_bsm_async(b_sts, dB[0], predn[0], true);            \
+        ldg_b128_bsm_async(b_sts + 128 * 16, dB[1], predn[1], true); \
+    }
+
+#define LDG_B_HEAD                                                            \
+    {                                                                         \
+        bool predk = rowB_swizzle < ktail;                                    \
+        ldg_b128_bsm_async(b_sts, dB[0], predn[0] && predk, true);            \
+        ldg_b128_bsm_async(b_sts + 128 * 16, dB[1], predn[1] && predk, true); \
+    }
+
+#define LDS_B                                          \
+    {                                                  \
+        cast_b128(rgB)[0] = cast_b128(b_lds)[0];       \
+        cast_b128(rgB)[1] = cast_b128(b_lds)[16 * 16]; \
+    }
+
+#define PERM_C(C2perm)                                       \
+    {                                                        \
+        Float4VecType C_tmp[8];                              \
+        float *ptri = (float *)C2perm;                       \
+        float *ptro = (float *)C_tmp;                        \
+        for (int j = 0; j < 4; ++j)                          \
+        {                                                    \
+            for (int i = 0; i < 4; ++i)                      \
+            {                                                \
+                ptro[j * 4 + i] = ptri[j + i * 4];           \
+                ptro[j * 4 + i + 16] = ptri[j + i * 4 + 16]; \
+            }                                                \
+        }                                                    \
+        for (int i = 0; i < 8; ++i)                          \
+        {                                                    \
+            C2perm[i] = C_tmp[i];                            \
+        }                                                    \
+    }
+
+#define STS_C(phase)                                            \
+    {                                                           \
+        cast_b128(c_sts[0])[0] = cast_b128(rgC)[0 + phase * 8]; \
+        cast_b128(c_sts[1])[0] = cast_b128(rgC)[1 + phase * 8]; \
+        cast_b128(c_sts[2])[0] = cast_b128(rgC)[2 + phase * 8]; \
+        cast_b128(c_sts[3])[0] = cast_b128(rgC)[3 + phase * 8]; \
+        cast_b128(c_sts[4])[0] = cast_b128(rgC)[4 + phase * 8]; \
+        cast_b128(c_sts[5])[0] = cast_b128(rgC)[5 + phase * 8]; \
+        cast_b128(c_sts[6])[0] = cast_b128(rgC)[6 + phase * 8]; \
+        cast_b128(c_sts[7])[0] = cast_b128(rgC)[7 + phase * 8]; \
+    }
+
+#define REDUCE_C(phase)                                                   \
+    {                                                                     \
+        cast_b128(rgC)[0 + phase * 8] = cast_b128(c_lds[0])[0];           \
+        cast_b128(rgC)[1 + phase * 8] = cast_b128(c_lds[1])[0];           \
+        cast_b128(rgC)[2 + phase * 8] = cast_b128(c_lds[0])[1 * 16 * 32]; \
+        cast_b128(rgC)[3 + phase * 8] = cast_b128(c_lds[1])[1 * 16 * 32]; \
+        cast_b128(rgC)[4 + phase * 8] = cast_b128(c_lds[0])[2 * 16 * 32]; \
+        cast_b128(rgC)[5 + phase * 8] = cast_b128(c_lds[1])[2 * 16 * 32]; \
+        cast_b128(rgC)[6 + phase * 8] = cast_b128(c_lds[0])[3 * 16 * 32]; \
+        cast_b128(rgC)[7 + phase * 8] = cast_b128(c_lds[1])[3 * 16 * 32]; \
+        float *reduc_c = reinterpret_cast<float *>(rgC);                  \
+        for (int loop = 0; loop < 8; ++loop)                              \
+        {                                                                 \
+            float acc = 0;                                                \
+            acc += reduc_c[loop + phase * 32 + 0 * 8];                    \
+            acc += reduc_c[loop + phase * 32 + 1 * 8];                    \
+            acc += reduc_c[loop + phase * 32 + 2 * 8];                    \
+            acc += reduc_c[loop + phase * 32 + 3 * 8];                    \
+            reduc_c[loop + phase * 8] = acc;                              \
+        }                                                                 \
+    }
diff --git a/csrc/quantization/gptq/hgemm_gptq.h b/csrc/quantization/gptq/hgemm_gptq.h
new file mode 100644
index 000000000..383b42ab8
--- /dev/null
+++ b/csrc/quantization/gptq/hgemm_gptq.h
@@ -0,0 +1,2031 @@
+// 2025 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
+/*
+hgemm gptq 4bits
+
+hgemm inputs:
+half/tf16 a[m*k];
+uint32_t b[k*n/pack_ratio]; //pack_ratio = sizeof(uint32_t) / 4;
+uint32_t zp[k*n] with kU4 or zp[0] with kU4B8;
+half scales[m*k/quant_group];
+
+MMA accepts a m16xk16 b k16xn16
+A tile should be m16n64k128 or m32n64k128
+When m>=32, we should load a when compute, implement this stragedge later?
+
+We will make a block size 256, which means 4 waves, we need all these 4 waves execute on one PEU
+So the maximum shared memory used by a block use not larger than 16K, and MTRegisters should be
+less than 128.
+
+A tile will be divied into serveral sub_tiles, a minimum sub_tile is m16n16k32, so we can have
+m_blocks*k_blocks when loading a and k_blocks*n_blocks when loading b.
+
+We will have n*k/TILE_N/TILE_K tiles in b, and these TILES are diveded into ITERS where ITERS*PEUS>=TILES
+for a proper shape.
+Of course some small shapes will nerver have chance to use all the PEUS
+
+Parallism is not suitable for C500 as parallism will dequant b more than once, that is not acceptable
+*/
+#include <iostream>
+#include <algorithm>
+
+#include <mc_runtime.h>
+#include <maca_fp16.h>
+#include <maca_bfloat16.h>
+
+#include "Hgemm_common.cuh"
+#include "scalar_type.hpp"
+
+using cudaStream_t = mcStream_t;
+
+#define WAVES_PER_BLOCK (THREADS/WAVE)
+#define TILE_M (BLOCKS_M*SLICE_M)
+#define TILE_N (BLOCKS_N*SLICE_N)
+#define TILE_K (BLOCKS_K*SLICE_K)
+#define N_ITERS (TILE_N / (WAVES_PER_BLOCK*SLOT))
+#define LOADING_A_LOOP SLICE_K * TILE_M / (sizeof(PackType) / sizeof(scalar_t)) / THREADS
+#define AS_PTR_B128(x) ((PackTypeInt4*)x)
+#define AS_PTR_B64(x) ((PackTypeInt2*)x)
+#define AS_PTR_B32(x) ((float*)x)
+#define AS_PTR_B16(x) ((half*)x)
+#define AS_PTR_B8(x) ((uint8_t*)x)
+
+#define BF16_HIGH_PRECISION
+
+#define div_ceil(x, y) (x + y - 1) / (y)
+
+//Although quant_group can be any positive value, but it is not a good idea
+//to set quant_group to values that cannot fit the SLICE_K as we will get a
+//very low performance, and we are not ready to support these values
+//Here we annouce that we support quant_group = 32, 64, 128, but actually
+//quant_group = 2^n where n >= 5 is also supported, for very large k.
+static int get_power2(uint32_t v) {
+    uint32_t power = 0;
+    uint32_t mask = 0x00000001;
+    while (power < 32) {
+        if ((v & mask) > 0) break;
+        power++;
+        mask <<= 1;
+    }
+    if ((1 << power) != v) return -1;
+    return static_cast<int>(power);
+}
+
+
+
+namespace hgemm_marlin_gptq {
+
+constexpr static int clean_kernel_thread_num = 512;
+constexpr static int clean_kernel_pack_num = 4;
+constexpr static int reduce_kernel_thread_num = 512;
+constexpr static int reduce_kernel_pack_num = 4;
+
+//#define DEBUG
+using PackTypeInt4 = b128VecType;
+using PackTypeInt2 = b64VecType;
+using PackType = uint32_t;
+
+template<class scalar_t>
+__device__ __forceinline__ void mma_16x16x16(PackTypeInt2& a, PackTypeInt2& b, PackTypeInt4& c) {
+}
+
+template<>
+__device__ __forceinline__ void mma_16x16x16<half>(PackTypeInt2& a, PackTypeInt2& b, PackTypeInt4& c) {
+    mma_16x16x16f16(a, b, c);
+}
+
+template<>
+__device__ __forceinline__ void mma_16x16x16<__maca_bfloat16>(PackTypeInt2& a, PackTypeInt2& b, PackTypeInt4& c) {
+    mma_16x16x16bf16(a, b, c);
+}
+
+#if 0
+#ifdef BF16_HIGH_PRECISION
+__global__ void vectorized_elementwise_fp32tobf16(float* input, __maca_bfloat16* output, int N) {
+    uint64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
+    if (tid < N) {
+        // printf("tid = %d, input = %f, output = %f\n", tid, input[tid], (float)(__maca_bfloat16)input[tid]);
+        *(__maca_bfloat16*)(output+tid) = (__maca_bfloat16)input[tid];
+    }
+}
+#else
+__global__ void vectorized_elementwise_fp16tobf16(__maca_bfloat16* input, int N) {
+    uint64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
+    if (tid < N) {
+        input[tid] = (float)(*(half*)(input+tid));
+    }
+}
+#endif
+#endif
+
+template<
+    const int THREADS, // number of threads in a threadblock
+    const int PACK_NUM
+    >
+__global__ void clean_zero(float* data, size_t num_elem) {
+    int bidx = blockIdx.x;
+    int tidx = threadIdx.x;
+    int idx = (bidx * THREADS + tidx) * PACK_NUM;
+    float zeros[4] = {0.0, 0.0, 0.0, 0.0};
+    if (idx < num_elem) {
+        *((b128VecType*)(&data[idx])) = *((b128VecType*)zeros);
+    }
+}
+
+template<
+    const int THREADS, // number of threads in a threadblock
+    const int PACK_NUM,
+    const bool USE_C = false
+    >
+__global__ void all_reduce(float* in_data, maca_bfloat16* out_data, size_t num_elem) {
+    int bidx = blockIdx.x;
+    int tidx = threadIdx.x;
+    int idx = (bidx * THREADS + tidx) * PACK_NUM;
+
+    if constexpr(PACK_NUM == 4) {
+        if constexpr(USE_C == true) {
+            float temp_in_fp32[PACK_NUM];
+            float temp_out_fp32[PACK_NUM];
+            maca_bfloat16 temp_out_bf16[PACK_NUM];
+
+            bool pred = idx < num_elem;
+            ldg_b64_reg_noasync(*((b64VecType*)temp_out_bf16), ((b64VecType*)(out_data + idx)), pred, true);
+            ldg_b128_reg_noasync(*((b128VecType*)temp_in_fp32), ((b128VecType*)(in_data + idx)), pred, true);
+
+            #pragma unroll
+            for(int i = 0; i < PACK_NUM; i++) {
+                temp_out_fp32[i] = __bfloat162float(temp_out_bf16[i]);
+                temp_out_fp32[i] += temp_in_fp32[i];
+                temp_out_bf16[i] = __float2bfloat16(temp_out_fp32[i]);
+            }
+
+            if (pred) {
+                *((b64VecType*)(out_data + idx)) = *((b64VecType*)temp_out_bf16);
+            }
+        } else {
+            float temp_in_fp32[PACK_NUM];
+            maca_bfloat16 temp_out_bf16[PACK_NUM];
+
+            bool pred = idx < num_elem;
+            ldg_b128_reg_noasync(*((b128VecType*)temp_in_fp32), ((b128VecType*)(in_data + idx)), pred, true);
+
+            #pragma unroll
+            for(int i = 0; i < PACK_NUM; i++) {
+                temp_out_bf16[i] = __float2bfloat16(temp_in_fp32[i]);
+            }
+
+            if (pred) {
+                *((b64VecType*)(out_data + idx)) = *((b64VecType*)temp_out_bf16);
+            }
+        }
+    }
+}
+
+typedef __NATIVE_VECTOR__(2, float) v2f;
+using PackTypeFloat2 = v2f;
+constexpr static int Q4BITS = 4;
+constexpr static int Q8BITS = 8;
+constexpr static int PACK_RATIO_4BITS = sizeof(PackType) * 8 / Q4BITS;
+constexpr static int PACK_RATIO_8BITS = sizeof(PackType) * 8 / Q8BITS;
+constexpr static int SLICE_M = 16;
+constexpr static int SLICE_N = 16;
+constexpr static int SLICE_K = 32;
+constexpr static int PAD_SLICE_K = 40;
+constexpr static int SLOT    = 16;
+constexpr static int WAVE    = 64;
+constexpr static int WAVE_SLOTS = 4;
+constexpr static int PEUS = 13*8*4; //For C500, There are 8 DPC and each DPC have 13 APs, each AP have 4 PEUs
+constexpr static int MAX_BLOCKS_M = 4;
+constexpr static uint32_t seil = 0x03020706u;
+
+__device__ __forceinline__ void f32x2_cvt_bf16x2(uint32_t& dst, float src[2]) {
+    uint32_t tmp[2];
+    tmp[0] = __builtin_mxc_ubfe(*(reinterpret_cast<uint32_t*>(src)), 16, 1);
+    tmp[0] = tmp[0] + *reinterpret_cast<uint32_t*>(src);
+    tmp[0] = (uint32_t)0x7fff + tmp[0];
+    tmp[1] = __builtin_mxc_ubfe(*(reinterpret_cast<uint32_t*>(src + 1)), 16, 1);
+    tmp[1] = tmp[1] + *(reinterpret_cast<uint32_t*>(src + 1));
+    tmp[1] = (uint32_t)0x7fff + tmp[1];
+    dst = __builtin_mxc_byte_perm(tmp[0], tmp[1], seil);
+}
+
+
+// #define CVT_B0TOF32(q, out) asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(out):"r"(q));
+// #define CVT_B1TOF32(q, out) asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(out):"r"(q));
+// #define CVT_B2TOF32(q, out) asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(out):"r"(q));
+// #define CVT_B3TOF32(q, out) asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(out):"r"(q));
+
+#define CVT_B0TOF32(q, out) out = __builtin_mxc_b0_cast_to_f32(q);
+#define CVT_B1TOF32(q, out) out = __builtin_mxc_b1_cast_to_f32(q);
+#define CVT_B2TOF32(q, out) out = __builtin_mxc_b2_cast_to_f32(q);
+#define CVT_B3TOF32(q, out) out = __builtin_mxc_b3_cast_to_f32(q);
+
+//FIXME: We'd rather a quant group will not divided into serveral blocks
+template<int BLOCKS_M, int BLOCKS_N, int BLOCKS_K>
+struct TileManager {
+    int tile_start_row;
+    int tile_start_col;
+    int tiles_k;
+    int my_iters = 0;
+    bool global_pred = true;
+
+    __device__ __forceinline__ void init(int m, int n, int k, int bidx, int iters) {
+        //Calculate tile start row and cols so we can calculate the offset address of a b and c
+        int tile_idx = iters * bidx;
+        int tiles_n = div_ceil(n, TILE_N);
+        tiles_k = div_ceil(k, TILE_K);
+        //if (tile_idx >= tiles_n*tiles_k) return false;
+        global_pred = tile_idx < tiles_n * tiles_k;
+        int tile_col = tile_idx / tiles_k;
+        int tile_row = tile_idx - tile_col * tiles_k;
+        tile_start_col = tile_col;
+        tile_start_row = tile_row;
+        my_iters = tile_idx + iters >= tiles_n*tiles_k ? tiles_n*tiles_k - tile_idx : iters;
+        my_iters = global_pred ? my_iters : 0;
+    }
+
+    __device__ __forceinline__ void next_tile() {
+        tile_start_col = tile_start_row + 1 == tiles_k ? tile_start_col + 1 : tile_start_col;
+        tile_start_row = tile_start_row + 1 == tiles_k ? 0 : tile_start_row + 1;
+        --my_iters;
+        global_pred = my_iters > 0;
+    }
+
+    __device__ __host__ __forceinline__ bool need_save_data() {
+        if (global_pred && my_iters == 1) return true;
+        if (global_pred && tile_start_row + 1 == tiles_k) return true;
+        return false;
+    }
+
+    //support for preloading next tile in current tile calculation
+    //The point is when all quanted values are dequanted and all a are stored to bsm already
+    //Then the registers for scales, zeros, and a_caches are free to use, bsm for scales are already
+    //free to use.
+    //The main problem we will face is that no more registers can be used to cache tile information
+    //when we want to preload next tile data
+    bool flag_save_data = false;
+    int tile_start_col_cache; //Kept for saving result n calculation
+
+    __device__ __forceinline__ void next_tile_pre() {
+        flag_save_data = need_save_data();
+        tile_start_col_cache = tile_start_col;
+        next_tile();
+    }
+
+    __device__ __forceinline__ bool need_save_data_pre() {
+        return flag_save_data;
+    }
+};
+
+struct ThreadView {
+    int tid;
+    int wave_idx;
+    int wave_tid;
+    int slot_idx;
+    int slot_tid;
+
+    __device__ __forceinline__ void init() {
+        tid = threadIdx.x;
+        wave_idx = tid / WAVE;
+        wave_tid = tid % WAVE;
+        slot_idx = wave_tid / SLOT;
+        slot_tid = wave_tid % SLOT;
+    }
+};
+
+/*
+Currently, we develop a version that TILE size is m64n128k64
+So every TILE we load 64x64x2=8192bytes A, 128x64/8x4=4096bytes B, 128x4=512 scales, 128x4=512 zeros if needed
+
+If more batches is needed, I think it is better to run dequant again. We do not count on situation that m>128 now.
+
+Memory View:
+There are 1/2/4 waves according to BLOCKS_N, wave_count = BLOCKS_N, mostly, BLOCKS_N=4, and THREADS=256
+Each wave will process a m16n16k32 sub tile with two mma instructions, we should re-order data in bsm so
+we can load data from shared memory only with tid or wave_tid, where we can reduce the instructions of 
+calculate address offsets of each data:
+
+BSM A order(assume that BLOCKS_M=4):
+|---------------------------------8x4x16 half value ------------------------------------------------------------------------|
+|-use by wave  slotidx 0 slottid 0------|-use by wave  slotidx 0 slottid 1------|...|-use by wave  slotidx 0 slottid 15-----|
+|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|...|-8 half--|-8 half--|-8 half--|-8 half--|
+|m00k00~07|m16k00~07|m32k00~07|m48k00~07|m01k00~07|m17k00~07|m33k00~07|m49k00~07|...|m15k00~07|m31k00~07|m47k00~07|m63k00~07|
+
+|-use by wave  slotidx 1 slottid 0------|-use by wave  slotidx 1 slottid 1------|...|-use by wave  slotidx 1 slottid 15-----|
+|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|...|-8 half--|-8 half--|-8 half--|-8 half--|
+|m00k08~15|m16k08~15|m32k08~15|m48k08~15|m01k08~15|m17k08~15|m33k08~15|m49k08~15|...|m15k08~15|m31k08~15|m47k08~15|m63k08~15|
+
+...
+
+|-use by wave  slotidx 3 slottid 0------|-use by wave  slotidx 3 slottid 1------|...|-use by wave  slotidx 3 slottid 15-----|
+|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|...|-8 half--|-8 half--|-8 half--|-8 half--|
+|m00k24~31|m16k24~31|m32k24~31|m48k24~31|m01k24~31|m17k24~31|m33k24~31|m49k24~31|...|m15k24~31|m31k24~31|m47k24~31|m63k24~31|
+
+---- next k ----
+|---------------------------------8x4x16 half value ------------------------------------------------------------------------|
+|-use by wave  slotidx 0 slottid 0------|-use by wave  slotidx 0 slottid 1------|...|-use by wave  slotidx 0 slottid 15-----|
+|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|...|-8 half--|-8 half--|-8 half--|-8 half--|
+|m00k32~39|m16k32~39|m32k32~39|m48k32~39|m01k32~39|m17k32~39|m33k32~39|m49k32~39|...|m15k32~39|m31k32~39|m47k32~39|m63k32~39|
+
+...
+
+|---------------------------------8x4x16 half value ------------------------------------------------------------------------|
+|-use by wave  slotidx 3 slottid 0------|-use by wave  slotidx 3 slottid 1------|...|-use by wave  slotidx 3 slottid 15-----|
+|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|-8 half--|...|-8 half--|-8 half--|-8 half--|-8 half--|
+|m00k56~63|m16k56~63|m32k56~63|m48k56~63|m01k56~63|m17k56~63|m33k56~63|m49k56~63|...|m15k56~63|m31k56~63|m47k56~63|m63k56~63|
+
+When loading a from bsm, each thread only read 4x8 half value with offset wave_tid:
+b128 out[4];
+b128 *bsm_a_ptr;
+out[0] = bsm_a_ptr[wave_tid*BLOCKS_M];
+out[1] = bsm_a_ptr[wave_tid*BLOCKS_M+1];
+out[2] = bsm_a_ptr[wave_tid*BLOCKS_M+2];
+out[3] = bsm_a_ptr[wave_tid*BLOCKS_M+3];
+
+
+BSM B order: here k means the quanted k, which means each k represents 8 dequanted k data
+//Loop = BLOCKS_N*SLICE_N/16/WAVES
+Loop 0:
+|---------16 uint32_t loaded by wave0-|
+k00n00 k00n01 k00n02 k00n03 ... k00n15
+k01n00 k01n01 k01n02 k01n03 ... k01n15
+k02n00 k02n01 k02n02 k02n03 ... k02n15
+k03n00 k03n01 k03n02 k03n03 ... k03n15
+|---------16 uint32_t loaded by wave1-|
+k00n16 k00n17 k00n18 k00n19 ... k00n31
+...
+k03n16 k03n17 k03n18 k03n19 ... k03n31
+|---------16 uint32_t loaded by wave2-|
+k00n32 k00n33 k00n34 k00n35 ... k00n47
+...
+k03n32 k03n33 k03n34 k03n35 ... k03n47
+|---------16 uint32_t loaded by wave3-|
+k00n48 k00n49 k00n50 k00n51 ... k00n63
+...
+k03n48 k03n49 k03n50 k03n51 ... k03n63
+*/
+template<class scalar_t>
+__device__ __forceinline__ void awq_dequant_4bits(const uint32_t& p, scalar_t (&out)[8], scalar_t (&scale)[8], const uint32_t& scale_zero) {
+    if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+        v2f v2z,v2scale, v2neg;
+        v2neg[0] = -1.0f; v2neg[1] = -1.0f;
+        v2z[0] = 0.0f; v2z[1] = 0.0f;
+        v2f v2zero;
+        float tmp[2];
+        int p0 = p & 0x0f0f0f0f;
+        int z0 = scale_zero & 0x0f0f0f0f;
+        CVT_B0TOF32(p0, tmp[0]);
+        CVT_B2TOF32(p0, tmp[1]);
+        CVT_B0TOF32(z0, v2zero[0]);
+        CVT_B2TOF32(z0, v2zero[1]);
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(v2zero, v2neg, *((v2f*)tmp));
+        v2scale[0] = (float)scale[0]; v2scale[1] = (float)scale[1];
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(*((v2f*)tmp), v2scale, v2z);
+        f32x2_cvt_bf16x2(*((uint32_t*)out), tmp);
+
+        CVT_B1TOF32(p0, tmp[0]);
+        CVT_B3TOF32(p0, tmp[1]);
+        CVT_B1TOF32(z0, v2zero[0]);
+        CVT_B3TOF32(z0, v2zero[1]);
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(v2zero, v2neg, *((v2f*)tmp));
+        v2scale[0] = (float)scale[4]; v2scale[1] = (float)scale[5];
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(*((v2f*)tmp), v2scale, v2z);
+        f32x2_cvt_bf16x2(*((uint32_t*)(out + 4)), tmp);
+
+        p0 = (p >> 4) & 0x0f0f0f0f;
+        z0 = (scale_zero >> 4) & 0x0f0f0f0f;
+        CVT_B0TOF32(p0, tmp[0]);
+        CVT_B2TOF32(p0, tmp[1]);
+        CVT_B0TOF32(z0, v2zero[0]);
+        CVT_B2TOF32(z0, v2zero[1]);
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(v2zero, v2neg, *((v2f*)tmp));
+        v2scale[0] = (float)scale[2]; v2scale[1] = (float)scale[3];
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(*((v2f*)tmp), v2scale, v2z);
+        f32x2_cvt_bf16x2(*((uint32_t*)(out + 2)), tmp);
+
+        CVT_B1TOF32(p0, tmp[0]);
+        CVT_B3TOF32(p0, tmp[1]);
+        CVT_B1TOF32(z0, v2zero[0]);
+        CVT_B3TOF32(z0, v2zero[1]);
+
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(v2zero, v2neg, *((v2f*)tmp));
+        v2scale[0] = (float)scale[6]; v2scale[1] = (float)scale[7];
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(*((v2f*)tmp), v2scale, v2z);
+        f32x2_cvt_bf16x2(*((uint32_t*)(out + 6)), tmp);
+    } else {
+        v2f a0, v2z, v2scale, v2zero, v2neg;
+        v2z[0] = 0; v2z[1] = 0;
+        v2neg[0] = -1.0f; v2neg[1] = -1.0f;
+        int p0 = p & 0x0f0f0f0f;
+        int z0 = scale_zero & 0x0f0f0f0f;
+        CVT_B0TOF32(p0, a0.x);
+        CVT_B2TOF32(p0, a0.y);
+        CVT_B0TOF32(z0, v2zero[0]);
+        CVT_B2TOF32(z0, v2zero[1]);
+        a0 = __builtin_mxc_pk_fma_f32(v2zero, v2neg, a0);
+        v2scale[0] = (float)scale[0]; v2scale[1] = (float)scale[1];
+        a0 = __builtin_mxc_pk_fma_f32(a0, v2scale, v2z);
+        out[0] = (scalar_t)a0.x;
+        out[1] = (scalar_t)a0.y;
+
+        CVT_B1TOF32(p0, a0.x);
+        CVT_B3TOF32(p0, a0.y);
+        CVT_B1TOF32(z0, v2zero[0]);
+        CVT_B3TOF32(z0, v2zero[1]);
+        a0 = __builtin_mxc_pk_fma_f32(v2zero, v2neg, a0);
+        v2scale[0] = (float)scale[4]; v2scale[1] = (float)scale[5];
+        a0 = __builtin_mxc_pk_fma_f32(a0, v2scale, v2z);
+        out[4] = (scalar_t)a0.x;
+        out[5] = (scalar_t)a0.y;
+
+        p0 = (p >> 4) & 0x0f0f0f0f;
+        z0 = (scale_zero >> 4) & 0x0f0f0f0f;
+        CVT_B0TOF32(p0, a0.x);
+        CVT_B2TOF32(p0, a0.y);
+        CVT_B0TOF32(z0, v2zero[0]);
+        CVT_B2TOF32(z0, v2zero[1]);
+        a0 = __builtin_mxc_pk_fma_f32(v2zero, v2neg, a0);
+        v2scale[0] = (float)scale[2]; v2scale[1] = (float)scale[3];
+        a0 = __builtin_mxc_pk_fma_f32(a0, v2scale, v2z);
+        out[2] = (scalar_t)a0.x;
+        out[3] = (scalar_t)a0.y;
+
+        CVT_B1TOF32(p0, a0.x);
+        CVT_B3TOF32(p0, a0.y);
+        CVT_B1TOF32(z0, v2zero[0]);
+        CVT_B3TOF32(z0, v2zero[1]);
+        a0 = __builtin_mxc_pk_fma_f32(v2zero, v2neg, a0);
+        v2scale[0] = (float)scale[6]; v2scale[1] = (float)scale[7];
+        a0 = __builtin_mxc_pk_fma_f32(a0, v2scale, v2z);
+        out[6] = (scalar_t)a0.x;
+        out[7] = (scalar_t)a0.y;
+    }
+}
+
+//deqaunt a uint32_t
+template<class scalar_t>
+__device__ __forceinline__ void dequant_gptq_4bits(const PackType& p, scalar_t (&out)[8], const v2f& scale, const v2f& scale_zero) {
+    if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+        v2f a0;
+        float tmp[2];
+        int p0 = p & 0x0f0f0f0f;
+
+        // CVT_B0TOF32(p0, a0.x);
+        // CVT_B2TOF32(p0, a0.y);
+        // a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+        // out[0] = __float2bfloat16(a0.x);
+        // out[1] = __float2bfloat16(a0.y);
+
+        CVT_B0TOF32(p0, tmp[0]);
+        CVT_B2TOF32(p0, tmp[1]);
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(*((v2f*)tmp), scale, scale_zero);
+        f32x2_cvt_bf16x2(*((uint32_t*)out), tmp);
+
+        // CVT_B1TOF32(p0, a0.x);
+        // CVT_B3TOF32(p0, a0.y);
+        // a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+        // out[4] = __float2bfloat16(a0.x);
+        // out[5] = __float2bfloat16(a0.y);
+
+        CVT_B1TOF32(p0, tmp[0]);
+        CVT_B3TOF32(p0, tmp[1]);
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(*((v2f*)tmp), scale, scale_zero);
+        f32x2_cvt_bf16x2(*((uint32_t*)(out + 4)), tmp);
+
+        p0 = (p >> 4) & 0x0f0f0f0f;
+
+        // CVT_B0TOF32(p0, a0.x);
+        // CVT_B2TOF32(p0, a0.y);
+        // a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+        // out[2] = __float2bfloat16(a0.x);
+        // out[3] = __float2bfloat16(a0.y);
+
+        CVT_B0TOF32(p0, tmp[0]);
+        CVT_B2TOF32(p0, tmp[1]);
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(*((v2f*)tmp), scale, scale_zero);
+        f32x2_cvt_bf16x2(*((uint32_t*)(out + 2)), tmp);
+
+        // CVT_B1TOF32(p0, a0.x);
+        // CVT_B3TOF32(p0, a0.y);
+        // a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+        // out[6] = __float2bfloat16(a0.x);
+        // out[7] = __float2bfloat16(a0.y);
+
+        CVT_B1TOF32(p0, tmp[0]);
+        CVT_B3TOF32(p0, tmp[1]);
+        *((v2f*)tmp) = __builtin_mxc_pk_fma_f32(*((v2f*)tmp), scale, scale_zero);
+        f32x2_cvt_bf16x2(*((uint32_t*)(out + 6)), tmp);
+    } else {
+        v2f a0;
+        int p0 = p & 0x0f0f0f0f;
+        CVT_B0TOF32(p0, a0.x);
+        CVT_B2TOF32(p0, a0.y);
+        a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+        out[0] = (scalar_t)a0.x;
+        out[1] = (scalar_t)a0.y;
+
+        CVT_B1TOF32(p0, a0.x);
+        CVT_B3TOF32(p0, a0.y);
+        a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+        out[4] = (scalar_t)a0.x;
+        out[5] = (scalar_t)a0.y;
+
+        p0 = (p >> 4) & 0x0f0f0f0f;
+        CVT_B0TOF32(p0, a0.x);
+        CVT_B2TOF32(p0, a0.y);
+        a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+        out[2] = (scalar_t)a0.x;
+        out[3] = (scalar_t)a0.y;
+
+        CVT_B1TOF32(p0, a0.x);
+        CVT_B3TOF32(p0, a0.y);
+        a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+        out[6] = (scalar_t)a0.x;
+        out[7] = (scalar_t)a0.y;
+    }
+}
+
+template<class scalar_t>
+__device__ __forceinline__ void dequant_gptq_8bits(const PackType& p, scalar_t (&out)[4], const v2f& scale, const v2f& scale_zero) {
+    v2f a0;
+    CVT_B0TOF32(p, a0.x);
+    CVT_B1TOF32(p, a0.y);
+    a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+    out[0] = (scalar_t)a0.x;
+    out[1] = (scalar_t)a0.y;
+
+    CVT_B2TOF32(p, a0.x);
+    CVT_B3TOF32(p, a0.y);
+    a0 = __builtin_mxc_pk_fma_f32(a0, scale, scale_zero);
+    out[2] = (scalar_t)a0.x;
+    out[3] = (scalar_t)a0.y;
+}
+
+// decompress zero
+__device__ __forceinline__ void decompress_zero_4bits(const PackType& zp, float (&out)[8]) {
+    v2f a0;
+    int p0 = zp & 0x0f0f0f0f;
+    CVT_B0TOF32(p0, a0.x);
+    CVT_B2TOF32(p0, a0.y);
+    out[0] = -a0.x;
+    out[1] = -a0.y;
+
+    CVT_B1TOF32(p0, a0.x);
+    CVT_B3TOF32(p0, a0.y);
+    out[4] = -a0.x;
+    out[5] = -a0.y;
+
+    p0 = (zp >> 4) & 0x0f0f0f0f;
+    CVT_B0TOF32(p0, a0.x);
+    CVT_B2TOF32(p0, a0.y);
+    out[2] = -a0.x;
+    out[3] = -a0.y;
+
+    CVT_B1TOF32(p0, a0.x);
+    CVT_B3TOF32(p0, a0.y);
+    out[6] = -a0.x;
+    out[7] = -a0.y;
+}
+
+namespace __hgemm_singular_blocks_k {
+template<typename scalar_t, const vllm::ScalarTypeId w_type_id, int THREADS, int BLOCKS_M, int BLOCKS_N, int BLOCKS_K, bool HAS_ACT, bool HAS_ZP, bool HAS_M_PRED, bool HAS_NK_PRED>
+struct LoadingManager {
+    constexpr static int FragACount = 2;
+    using FragA = PackType;
+    constexpr static int FragBCount = 1;
+    using FragB = PackType;
+    constexpr static int FragCCount = 4;
+    #ifdef BF16_HIGH_PRECISION
+    using FragC = scalar_t;
+    #else
+    //Directly use half as the final atomic type:
+    //1. Half precision and data range satisfies need of deepseek gemm
+    //2. C500 has no atomic instructions for bfloat16, we cannot atomic a bfloat16 memory
+    //3. The perfect precision type of atomic should be fp32, but the cost is too high to allocate a temp memory for float atomic
+    using atomic_type = half;
+    using FragC = atomic_type;
+    #endif
+    const FragA* A;
+    const FragA* A_loading;
+    const FragB* B;
+    const FragB* B_loading;
+    FragC* C;
+    float* C_temp;
+    using FragScaleLoading = half2;
+    using FragZeroLoading = uint32_t;
+    const FragScaleLoading* scales;
+    const FragScaleLoading* scales_loading;
+    const FragZeroLoading* zeros;
+    const FragZeroLoading* zeros_loading;
+    
+    int m;
+    int n;
+    int k;
+    int quant_group_power2;
+    uint8_t* smem_base;
+    int bidx;
+
+    PackTypeInt4* bsm_a_ptr;
+    scalar_t* bsm_scales_ptr;
+    float* bsm_zeros_ptr;
+    float* remaining_bsm_ptr;
+
+    PackTypeInt2 local_a[BLOCKS_M][2];
+    PackType local_b[N_ITERS];
+    PackType local_b_cache[N_ITERS];
+    scalar_t local_dequanted_b_8bits[N_ITERS][2][PACK_RATIO_8BITS];
+    scalar_t local_dequanted_b[N_ITERS][PACK_RATIO_4BITS];
+    v2f local_scales[N_ITERS];
+    v2f local_zeros[N_ITERS];
+    FragScaleLoading temp_scales;
+    PackType temp_zeros;
+    float output[BLOCKS_M][N_ITERS][4];
+    FragA temp_a[LOADING_A_LOOP];
+
+    TileManager<BLOCKS_M, BLOCKS_N, BLOCKS_K> tile_manager;
+    ThreadView tv;
+
+    __device__ __forceinline__ void set_address(const PackTypeInt4* a,
+        const PackTypeInt4* b,
+        PackTypeInt4* c,
+        PackTypeInt4* c_temp,
+        const PackTypeInt4* scale_ptr,
+        const PackTypeInt4* zp_ptr = nullptr) {
+            A = (const FragA*)a;
+            B = (const FragB*)b;
+            C = (FragC*)c;
+            C_temp = (float*)c_temp;
+            scales = (const FragScaleLoading*)scale_ptr;
+            if constexpr(w_type_id == vllm::kU4.id()) {
+                zeros = (const FragZeroLoading*)zp_ptr;
+            }
+    }
+
+    __device__ __forceinline__ bool debug() {
+        #ifdef DEBUG
+        bool do_print = tv.wave_idx == 1 && tv.slot_idx == 0 && tv.slot_tid == 0;
+        return do_print;
+        #else
+        return false;
+        #endif
+    }
+
+    __device__ __forceinline__ void next_k() {
+        //Update only bsm_a_ptr
+        bsm_a_ptr = (PackTypeInt4*)smem_base;
+        bsm_a_ptr += tv.slot_tid * (PAD_SLICE_K / (sizeof(PackTypeInt4) / sizeof(scalar_t))) + tv.slot_idx;
+    }
+
+    __device__ __forceinline__ void next_k_pre() {
+        A_loading += SLICE_K / FragACount;
+        //B_loading += SLICE_K / PACK_RATIO_4BITS * n;
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            B_loading += SLICE_K / PACK_RATIO_4BITS * n;
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            B_loading += SLICE_K / PACK_RATIO_8BITS * n;
+        }
+    }
+
+    __device__ __forceinline__ void ldg_a(int k_idx) {
+        //32x64/2/256 = 16 / 4 = 4
+        int t = tv.tid;
+        int k_broad = tile_manager.tile_start_row * TILE_K + k_idx * SLICE_K;
+        #pragma unroll LOADING_A_LOOP
+        for (int i = 0; i < LOADING_A_LOOP; i++)  {
+            int reading_m = t / (SLICE_K / FragACount);
+            int reading_k = t % (SLICE_K / FragACount);
+            int gvm_offset = reading_m * k / FragACount + reading_k;
+            FragA* gvm_addr = (FragA*)A_loading + gvm_offset;
+            //FIXME: we cannot do slice k pad as ldg_b32_bsm_async seems does not support padding
+            if constexpr(HAS_M_PRED && HAS_NK_PRED) {
+                bool pred = reading_m < m;
+                bool pred_k = k_broad + reading_k * FragACount < k;
+                pred = pred && pred_k && tile_manager.global_pred;
+                ldg_b32_reg_noasync(temp_a[i], gvm_addr, pred, true);
+            } else if constexpr(HAS_M_PRED) {
+                bool pred = reading_m < m && tile_manager.global_pred;
+                ldg_b32_reg_noasync(temp_a[i], gvm_addr, pred, true);
+            } else if constexpr(HAS_NK_PRED) {
+                bool pred_k = k_broad + reading_k * FragACount < k && tile_manager.global_pred;
+                ldg_b32_reg_noasync(temp_a[i], gvm_addr, pred_k, true);
+            } else {
+                ldg_b32_reg_noasync(temp_a[i], gvm_addr, tile_manager.global_pred, true);
+            }
+            t += THREADS;
+        }
+    }
+
+    __device__ __forceinline__ void sts_a() {
+        FragA* to_bsm_a_ptr = (FragA*)smem_base;
+        int t = tv.tid;
+        #pragma unroll LOADING_A_LOOP
+        for (int i = 0; i < LOADING_A_LOOP; i++)  {
+            int reading_m = t / (SLICE_K / FragACount);
+            int reading_k = t % (SLICE_K / FragACount);
+            int bsm_offset = reading_m * (PAD_SLICE_K / FragACount) + reading_k;
+            *(to_bsm_a_ptr + bsm_offset) = temp_a[i];
+            t += THREADS;
+        }
+    }
+
+    __device__ __forceinline__ void lds_a(int midx) {
+        *((PackTypeInt4*)local_a[midx]) = *bsm_a_ptr;
+        bsm_a_ptr += SLOT * (PAD_SLICE_K / (sizeof(PackTypeInt4) / sizeof(scalar_t)));
+    }
+
+    //TODO: implement when N_ITERS==1 or N_ITERS==3
+    __device__ __forceinline__ void ldg_b(int k_idx, int korder = 0) {
+        if constexpr(HAS_NK_PRED) {
+            bool pred_k = tile_manager.tile_start_row * TILE_K + k_idx * SLICE_K + tv.slot_idx * PACK_RATIO_4BITS + korder < k;
+            bool pred_n = tile_manager.tile_start_col * TILE_N + (tv.wave_idx * SLOT + tv.slot_tid)*N_ITERS < n;
+            bool pred = pred_n && pred_k && tile_manager.global_pred;
+            FragB* addr =  (FragB*)B_loading + korder * n;
+            if constexpr(N_ITERS == 2) {
+                ldg_b64_reg_noasync(*((PackTypeInt2*)local_b_cache), ((PackTypeInt2*)addr), pred, true);
+            } else if constexpr(N_ITERS == 4) {
+                ldg_b128_reg_noasync(*((PackTypeInt4*)local_b_cache), ((PackTypeInt4*)addr), pred, true);
+            }
+        } else {
+            FragB* addr =  (FragB*)B_loading + korder * n;
+            if constexpr(N_ITERS == 2) {
+                ldg_b64_reg_noasync(*((PackTypeInt2*)local_b_cache), ((PackTypeInt2*)addr), tile_manager.global_pred, true);
+            } else if constexpr(N_ITERS == 4) {
+                ldg_b128_reg_noasync(*((PackTypeInt4*)local_b_cache), ((PackTypeInt4*)addr), tile_manager.global_pred, true);
+            }
+        }
+    }
+
+    __device__ __forceinline__ void swap_b_cache(int i) {
+        local_b[i] = local_b_cache[i];
+    }
+
+    __device__ __forceinline__ void ldg_scales() {
+        bool pred = tv.tid < TILE_N / (sizeof(FragScaleLoading) / sizeof(scalar_t)) && tile_manager.global_pred;
+        if constexpr(HAS_NK_PRED) {
+            pred = pred && tv.tid < (n - tile_manager.tile_start_col * TILE_N) / (sizeof(FragScaleLoading) / sizeof(scalar_t));
+        }
+        //FragScaleLoading *scale_bsm = (FragScaleLoading*)(smem_base + 0x2000) + tv.tid;
+        FragScaleLoading *gvm_addr = (FragScaleLoading*)scales_loading + tv.tid;
+        //ldg_b32_bsm_async(scale_bsm, gvm_addr, pred, false);
+        ldg_b32_reg_noasync(*((PackType*)&temp_scales), gvm_addr, pred, true);
+    }
+
+    __device__ __forceinline__ void ldg_zp() {
+        if constexpr(w_type_id == vllm::kU4.id()) {
+            bool pred = (tv.tid < TILE_N / PACK_RATIO_4BITS) && tile_manager.global_pred;
+            if constexpr(HAS_NK_PRED) {
+                pred = pred && tv.tid < ((n - tile_manager.tile_start_col * TILE_N) / PACK_RATIO_4BITS);
+            }
+            FragZeroLoading *gvm_addr = (FragZeroLoading*)zeros_loading + tv.tid;
+            ldg_b32_reg_noasync(*((PackType*)&temp_zeros), gvm_addr, pred, true);
+        }
+    }
+
+    __device__ __forceinline__ void sts_scales() {
+        FragScaleLoading *scale_bsm = (FragScaleLoading*)(smem_base + 0x2000) + tv.tid;
+        *scale_bsm = temp_scales;
+    }
+
+    __device__ __forceinline__ void sts_zeros() {
+        if constexpr(w_type_id == vllm::kU4.id()) {
+            bool pred = (tv.tid < TILE_N / PACK_RATIO_4BITS) && tile_manager.global_pred;
+            if (pred) {
+                float temp[PACK_RATIO_4BITS];
+                decompress_zero_4bits(temp_zeros, temp);
+                float *scale_bsm = (float*)(smem_base + 0x3000) + tv.tid * PACK_RATIO_4BITS;
+                for (int i = 0; i < PACK_RATIO_4BITS; i++) {
+                    *(scale_bsm + i) = temp[i];
+                }
+            }
+        }
+    }
+
+    //TODO: implement when N_ITERS==1 or N_ITERS==3
+    __device__ __forceinline__ void lds_scales() {
+        if constexpr(N_ITERS==2) {
+            *((half2*)local_dequanted_b[0]) = *((half2*)bsm_scales_ptr);
+        } else if constexpr(N_ITERS==4) {
+            *((PackTypeInt2*)local_dequanted_b[0]) = *((PackTypeInt2*)bsm_scales_ptr);
+        }
+    }
+
+    __device__ __forceinline__ void pack_scales() {
+        if constexpr(w_type_id == vllm::kU4B8.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                float s = local_dequanted_b[0][i];
+                float z = -8 * s;
+                local_scales[i] = {s, s};
+                local_zeros[i] = {z, z};
+            }
+        } else if constexpr(w_type_id == vllm::kU4.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                float s;
+                if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+                    s = __bfloat162float(local_dequanted_b[0][i]);
+                } else {
+                    s = local_dequanted_b[0][i];
+                }
+                float z = *(bsm_zeros_ptr + i);
+                z = z * s;
+                local_scales[i] = {s, s};
+                local_zeros[i] = {z, z};
+            }
+        } else if constexpr(w_type_id == vllm::kU8B128.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                float s = local_dequanted_b[0][i];
+                float z = -128 * s;
+                local_scales[i] = {s, s};
+                local_zeros[i] = {z, z};
+            }
+        } else if constexpr(w_type_id == vllm::kU8.id()) {
+            // should apply zeros
+        }
+    }
+
+    __device__ __forceinline__ void dequant(int kdx, int korder = 0) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            dequant_gptq_4bits<scalar_t>(local_b[kdx], local_dequanted_b[kdx], local_scales[kdx], local_zeros[kdx]);
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            dequant_gptq_8bits<scalar_t>(local_b[kdx], local_dequanted_b_8bits[kdx][korder], local_scales[kdx], local_zeros[kdx]);
+        }
+    }
+
+    __device__ __forceinline__ void matmul(int mdx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                mma_16x16x16<scalar_t>(local_a[mdx][0], *((PackTypeInt2*)local_dequanted_b[i]), *((PackTypeInt4*)output[mdx][i]));
+                mma_16x16x16<scalar_t>(local_a[mdx][1], *((PackTypeInt2*)local_dequanted_b[i] + 1), *((PackTypeInt4*)output[mdx][i]));
+	    }
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                mma_16x16x16<scalar_t>(local_a[mdx][0], *((PackTypeInt2*)local_dequanted_b_8bits[i][0]), *((PackTypeInt4*)output[mdx][i]));
+                mma_16x16x16<scalar_t>(local_a[mdx][1], *((PackTypeInt2*)local_dequanted_b_8bits[i][1]), *((PackTypeInt4*)output[mdx][i]));
+            }
+        }
+    }
+
+    __device__ __forceinline__ void clear_c() {
+        #pragma unroll
+        for (int miter = 0; miter < BLOCKS_M; miter++) {
+            #pragma unroll
+            for (int niter = 0; niter < N_ITERS; niter++) {
+                #pragma unroll
+                for (int miter2 = 0; miter2 < 4; miter2++) {
+                    output[miter][niter][miter2] = 0;
+                }
+            }
+        }
+    }
+
+    //functions for preloading next tile data
+    __device__ __forceinline__ void init_address_pre(int _m, int _n, int _k, int _quant_group_power2, int _bidx, int _iters, uint8_t *_smem_base) {
+        tv.init();
+        m = _m;
+        n = _n;
+        k = _k;
+        quant_group_power2 = _quant_group_power2;
+        bidx = _bidx;
+        smem_base = _smem_base;
+        tile_manager.init(m, n, k, bidx, _iters);
+        init_tile_pre(tile_manager.tile_start_col, tile_manager.tile_start_row);
+    }
+
+    __device__ __forceinline__ void init_tile_pre(int col, int row) {
+        //Initialize start slice address and set them to A_loading and B_loading
+        int offset_n = col * TILE_N;
+        int offset_k = row * TILE_K;
+        //A_loading address will always be valid
+        A_loading = A + offset_k / (FragACount);
+        //B_loading = B + (offset_k / PACK_RATIO_4BITS * n + offset_n) / (FragBCount) + tv.slot_idx * n + (tv.wave_idx * SLOT + tv.slot_tid)*N_ITERS;
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            B_loading = B + (offset_k / PACK_RATIO_4BITS * n + offset_n) / (FragBCount) + tv.slot_idx * n + (tv.wave_idx * SLOT + tv.slot_tid)*N_ITERS;
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            B_loading = B + (offset_k / PACK_RATIO_8BITS * n + offset_n) / (FragBCount) + tv.slot_idx * n * 2 + (tv.wave_idx * SLOT + tv.slot_tid)*N_ITERS;
+        }
+        scales_loading = scales + ((offset_k >> quant_group_power2) * n + offset_n) / (sizeof(FragScaleLoading)/sizeof(scalar_t));
+        if constexpr(w_type_id == vllm::kU4.id()) {
+            zeros_loading = zeros + ((offset_k >> quant_group_power2) * n + offset_n) / PACK_RATIO_4BITS;
+        }
+    }
+
+    __device__ __forceinline__ void next_tile_pre() {
+        tile_manager.next_tile_pre();
+        init_tile_pre(tile_manager.tile_start_col, tile_manager.tile_start_row);
+    }
+
+    __device__ __forceinline__ void init_bsm_addr() {
+        bsm_a_ptr = (PackTypeInt4*)smem_base;           //use 8k bytes, will load at most 32x128*sizeof(half), either m32k128 or m128k32
+        remaining_bsm_ptr = (float*)(smem_base + 0x2000 + 0x1000); //3k bytes
+        bsm_a_ptr += tv.slot_tid * (PAD_SLICE_K / (sizeof(PackTypeInt4) / sizeof(scalar_t))) + tv.slot_idx;
+        bsm_scales_ptr = (scalar_t*)(smem_base + 0x2000);      //use 128xsizeof(float)*2 = 1k
+        bsm_scales_ptr += (tv.wave_idx * SLOT + tv.slot_tid) * N_ITERS;
+        if constexpr(w_type_id == vllm::kU4.id()) {
+            bsm_zeros_ptr = (float*)(smem_base + 0x3000);
+            bsm_zeros_ptr += (tv.wave_idx * SLOT + tv.slot_tid) * N_ITERS;
+        }
+    }
+
+    __device__ __forceinline__ void write_c(int offset, const float& v) {
+        if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+            #ifdef BF16_HIGH_PRECISION
+            atomicAdd(C_temp + offset, v);
+            #else
+            atomicAdd(C+offset, (atomic_type)v);
+            #endif
+        } else {
+            atomicAdd(C + offset, (scalar_t)v);
+        }
+    }
+
+    //atomic write to c
+    __device__ __forceinline__ void write_c_pre() {
+        int k_broad = tv.slot_idx * 4;
+        int n_broad = (tv.wave_idx * SLOT + tv.slot_tid) * N_ITERS + tile_manager.tile_start_col_cache * TILE_N;
+        #pragma unroll
+        for (int miter = 0; miter < BLOCKS_M; miter++) {
+            #pragma unroll
+            for (int niter = 0; niter < N_ITERS; niter++) {
+                int store_n = n_broad + niter;
+                #pragma unroll
+                for (int miter2 = 0; miter2 < 4; miter2++) {
+                    int store_m = k_broad + miter * SLICE_M + miter2;
+                    if constexpr(HAS_M_PRED && HAS_NK_PRED) {
+                        if (store_m < m && store_n < n) {
+                            write_c(store_m * n + store_n, output[miter][niter][miter2]);
+			            }
+                    } else if constexpr(HAS_M_PRED) {
+                        if (store_m < m) {
+                            write_c(store_m * n + store_n, output[miter][niter][miter2]);
+			            }
+                    } else if constexpr(HAS_NK_PRED) {
+                        if (store_n < n) {
+                            write_c(store_m * n + store_n, output[miter][niter][miter2]);
+			            }
+                    } else {
+                        write_c(store_m * n + store_n, output[miter][niter][miter2]);
+		            }
+                }
+            }
+        }
+    }
+
+    template<bool KTAIL>
+    __device__ __forceinline__ void on_dequant_niters1(int k_idx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            swap_b_cache(0);
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+            if constexpr(!KTAIL) {
+                next_k_pre(); // preload gvm a/b
+                ldg_b(k_idx + 1); //preload b for next k
+                ldg_a(k_idx + 1); //preload a for next k
+            }
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0); //dequant b0
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            swap_b_cache(0);
+            ldg_b(k_idx, 1);
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+            dequant(0);
+            swap_b_cache(0);
+            if constexpr(!KTAIL) {
+                next_k_pre(); // preload gvm a/b
+                ldg_b(k_idx + 1); //preload b for next k
+                ldg_a(k_idx + 1); //preload a for next k
+            }
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0); //dequant b0
+        }
+    }
+
+    template<bool KTAIL>
+    __device__ __forceinline__ void on_dequant_niters2(int k_idx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            swap_b_cache(0);
+            dequant(0); //dequant b0
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+            swap_b_cache(1);
+            if constexpr(!KTAIL) {
+                next_k_pre(); // preload gvm a/b
+                ldg_b(k_idx + 1); //preload b for next k
+                ldg_a(k_idx + 1); //preload a for next k
+            }
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(1);   //dequant b64
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            swap_b_cache(0);
+            swap_b_cache(1);
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+            ldg_b(k_idx, 1);
+            dequant(0); //dequant b0
+            dequant(1);
+            swap_b_cache(0);
+            swap_b_cache(1);
+            if constexpr(!KTAIL) {
+                next_k_pre(); // preload gvm a/b
+                ldg_b(k_idx + 1); //preload b for next k
+                ldg_a(k_idx + 1); //preload a for next k
+            }
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0, 1);
+            dequant(1, 1);   //dequant b64
+        }
+    }
+
+    template<bool KTAIL>
+    __device__ __forceinline__ void on_dequant_niters3(int k_idx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            swap_b_cache(0);
+            dequant(0); //dequant b0
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+            swap_b_cache(1);
+            dequant(1);
+            swap_b_cache(2);
+            if constexpr(!KTAIL) {
+                next_k_pre(); // preload gvm a/b
+                ldg_b(k_idx + 1); //preload b for next k
+                ldg_a(k_idx + 1); //preload a for next k
+            }
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(2); //dequant b0
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            swap_b_cache(0);
+            swap_b_cache(1);
+            swap_b_cache(2);
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+            ldg_b(k_idx, 1);
+            dequant(0); //dequant b0
+            dequant(1); //dequant b0
+            dequant(2); //dequant b0
+            swap_b_cache(0);
+            swap_b_cache(1);
+            swap_b_cache(2);
+            if constexpr(!KTAIL) {
+                next_k_pre(); // preload gvm a/b
+                ldg_b(k_idx + 1); //preload b for next k
+                ldg_a(k_idx + 1); //preload a for next k
+            }
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0, 1); //dequant b0
+            dequant(1, 1); //dequant b0
+            dequant(2, 1); //dequant b0
+        }
+    }
+
+    template<bool KTAIL>
+    __device__ __forceinline__ void on_dequant_niters4(int k_idx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            swap_b_cache(0);
+            dequant(0); //dequant b0
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+            swap_b_cache(1);
+            dequant(1); //dequant b1
+            swap_b_cache(2);
+            swap_b_cache(3);
+            dequant(2); //dequant b2
+            if constexpr(!KTAIL) {
+                next_k_pre(); // preload gvm a/b
+                ldg_b(k_idx + 1); //preload b for next k
+                ldg_a(k_idx + 1); //preload a for next k
+            }
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(3); //dequant b3
+        } else {
+            swap_b_cache(0);
+            swap_b_cache(1);
+            swap_b_cache(2);
+            swap_b_cache(3);
+            dequant(0); //dequant b0
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+            ldg_b(k_idx, 1);
+            dequant(1); //dequant b0
+            dequant(2); //dequant b0
+            dequant(3); //dequant b0
+            swap_b_cache(0);
+            swap_b_cache(1);
+            swap_b_cache(2);
+            swap_b_cache(3);
+            if constexpr(!KTAIL) {
+                next_k_pre(); // preload gvm a/b
+                ldg_b(k_idx + 1); //preload b for next k
+                ldg_a(k_idx + 1); //preload a for next k
+            }
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0, 1); //dequant b3
+            dequant(1, 1); //dequant b3
+            dequant(2, 1); //dequant b3
+            dequant(3, 1); //dequant b3
+        }
+    }
+
+    template<bool KTAIL>
+    __device__ __forceinline__ void on_dequant(int kdx) {
+        if constexpr(N_ITERS == 1) on_dequant_niters1<KTAIL>(kdx);
+        else if constexpr(N_ITERS == 2) on_dequant_niters2<KTAIL>(kdx);
+        else if constexpr(N_ITERS == 3) on_dequant_niters3<KTAIL>(kdx);
+        else if constexpr(N_ITERS == 4) on_dequant_niters4<KTAIL>(kdx);
+    }
+};
+
+template<typename scalar_t,
+    const vllm::ScalarTypeId w_type_id,
+    const int THREADS,          // number of threads in a threadblock
+    const int BLOCKS_M,         // number of 16x16 blocks in the m
+                                // dimension (batchsize) of the
+                                // threadblock
+    const int BLOCKS_N,         // same for n dimension (output)
+    const int BLOCKS_K,         // same for k dimension (reduction)
+    const bool HAS_ACT_ORDER,   // whether act_order is enabled
+    const bool HAS_ZP,          // whether zero-points are enabled
+    const bool HAS_M_PRED = true,  //If we should use predictors to load m from gvm
+    const bool HAS_NK_PRED = true  //If we should use predictors to load nk from gvm
+    >
+__global__ void hgemm_gptq(
+    const PackTypeInt4* __restrict__ A,  // fp16 input matrix of shape mxk
+    const PackTypeInt4* __restrict__ B,  // 4bit quantized weight matrix of shape kxn
+    PackTypeInt4* __restrict__ C,        // fp16 output buffer of shape mxn
+    PackTypeInt4* __restrict__ C_tmp,    // fp32 tmp output buffer (for reduce)
+    const PackTypeInt4* __restrict__ scales_ptr,  // fp16 quantization scales of shape
+                                          // (k/groupsize)xn
+    const PackTypeInt4* __restrict__ zp_ptr,      // 4bit packed zero-points of shape
+                                          // (k/groupsize)x(n/pack_factor)
+    const int* __restrict__ g_idx,        // int32 group indices of shape k
+    int prob_m,           // batch dimension m
+    int prob_n,           // output dimension n
+    int prob_k,           // reduction dimension k
+    int quant_group_power2, // quant group means how many quanted values share the same scale and zero, this value restricts to 2^x where x >= 5
+    int max_iters,        // max tile iterations for one block
+    int* locks,           // extra global storage for barrier synchronization
+    bool use_fp32_reduce  // whether to use fp32 global reduce
+) {
+    int bidx = blockIdx.x;
+    __shared__ uint8_t smem_base[0x4000]; //4x16x256 = 16Kbytes
+    using LoadingManagerType = LoadingManager<scalar_t, w_type_id, THREADS, BLOCKS_M, BLOCKS_N, BLOCKS_K, HAS_ACT_ORDER, HAS_ZP, HAS_M_PRED, HAS_NK_PRED>;
+    LoadingManagerType loading_manager;
+    A += blockIdx.y * (MAX_BLOCKS_M * SLICE_M) * prob_k / (sizeof(PackTypeInt4) / sizeof(scalar_t));
+    C += blockIdx.y * (MAX_BLOCKS_M * SLICE_M) * prob_n / (sizeof(PackTypeInt4) / sizeof(scalar_t));
+    #ifdef BF16_HIGH_PRECISION
+    if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+        C_tmp += blockIdx.y * (MAX_BLOCKS_M * SLICE_M) * prob_n / (sizeof(PackTypeInt4) / sizeof(float));
+    }
+    #endif
+    loading_manager.set_address(A, B, C, C_tmp, scales_ptr, zp_ptr);
+    //loading_manager.init_address(prob_m, prob_n, prob_k, bidx, max_iters, smem_base);
+    loading_manager.init_address_pre(std::min<int>(MAX_BLOCKS_M*SLICE_M, prob_m - blockIdx.y * (MAX_BLOCKS_M * SLICE_M)), prob_n, prob_k, quant_group_power2, bidx, max_iters, smem_base);
+    loading_manager.clear_c();
+
+    while (max_iters > 0) {
+        loading_manager.init_bsm_addr(); //reset all bsm address for current tile
+        loading_manager.ldg_scales(); //Load all scales to bsm
+        loading_manager.ldg_zp();
+        loading_manager.ldg_b(0);    //load b0 and b64, two gvm
+        loading_manager.ldg_a(0);    //Load first k0~31 and all m, one ldg_b128, heavy load
+        loading_manager.sts_scales();
+        loading_manager.sts_zeros();
+        barrier_bsm;
+        loading_manager.lds_scales(); //load scale0 and scale64
+        loading_manager.pack_scales(); //pack scales into two v2f structure
+
+        int k_idx = 0;
+        if constexpr(BLOCKS_K > 1) {
+            #pragma unroll BLOCKS_K - 1
+            for (; k_idx < BLOCKS_K - 1; k_idx++) {
+                int m_idx = 0;
+                loading_manager.template on_dequant<false>(k_idx);
+                //Loop for 3 times so that we can add some loading instructions before matmul
+                if constexpr(BLOCKS_M > 1) {
+                    #pragma unroll BLOCKS_M - 1
+                    for (; m_idx < BLOCKS_M - 1; m_idx++) {
+                        loading_manager.lds_a(m_idx+1); //pre load next 16 batches
+                        loading_manager.matmul(m_idx); //do matmul
+                    }
+                }
+                barrier_bsm; //wait for the last 16 batches ready, becasue of k_idx will be change
+                loading_manager.next_k(); //modify gvm/bsm address of a and b
+                loading_manager.matmul(m_idx); //do matmul
+            }
+        }
+        int m_idx = 0;
+        loading_manager.template on_dequant<true>(k_idx);
+        if constexpr(BLOCKS_M > 1) {
+            #pragma unroll BLOCKS_M - 1
+            for (; m_idx < BLOCKS_M - 1; m_idx++) {
+                loading_manager.lds_a(m_idx+1); //pre load next 16 batches
+                loading_manager.matmul(m_idx); //do matmul
+            }
+        }
+        loading_manager.next_tile_pre();
+
+        loading_manager.matmul(m_idx); //do matmul
+        max_iters--;
+
+        if (loading_manager.tile_manager.need_save_data_pre()) {
+            loading_manager.write_c_pre(); // reduce and write back
+            loading_manager.clear_c();
+        }
+
+        barrier_bsm;
+    }
+}
+
+} //end of namespace __hgemm_singular_blocks_k
+
+namespace __hgemm_even_blocks_k {
+template<typename scalar_t, const vllm::ScalarTypeId w_type_id, int THREADS, int BLOCKS_M, int BLOCKS_N, int BLOCKS_K, bool HAS_ACT, bool HAS_ZP, bool HAS_M_PRED, bool HAS_NK_PRED>
+struct LoadingManager {
+    constexpr static int FragACount = 4;
+    using FragA = PackTypeInt2;
+    constexpr static int FragBCount = 1;
+    using FragB = PackType;
+    constexpr static int FragCCount = 4;
+    #ifdef BF16_HIGH_PRECISION
+    using FragC = scalar_t;
+    #else
+    //Directly use half as the final atomic type:
+    //1. Half precision and data range satisfies need of deepseek gemm
+    //2. C500 has no atomic instructions for bfloat16, we cannot atomic a bfloat16 memory
+    //3. The perfect precision type of atomic should be fp32, but the cost is too high to allocate a temp memory for float atomic
+    using atomic_type = half;
+    using FragC = atomic_type;
+    #endif
+    const FragA* A;
+    const FragA* A_loading;
+    const FragB* B;
+    const FragB* B_loading;
+    FragC* C;
+    float* C_temp;
+    using FragScaleLoading = half2;
+    using FragZeroLoading = uint32_t;
+    const FragScaleLoading* scales;
+    const FragScaleLoading* scales_loading;
+    const FragZeroLoading* zeros;
+    const FragZeroLoading* zeros_loading;
+
+    constexpr static int DOUBLE_SLICE_K = SLICE_K * 2;
+    constexpr static int DOUBLE_PAD_SLICE_K = SLICE_K * 2 + sizeof(PackTypeInt4) / sizeof(scalar_t);
+
+    int m;
+    int n;
+    int k;
+    int quant_group_power2;
+    uint8_t* smem_base;
+    int bidx;
+
+    PackTypeInt4* bsm_a_ptr;
+    scalar_t* bsm_scales_ptr;
+    float* bsm_zeros_ptr;
+    //float* remaining_bsm_ptr;
+
+    PackTypeInt2 local_a[BLOCKS_M][2];
+    PackType local_b[N_ITERS];
+    PackType local_b_cache[N_ITERS];
+    scalar_t local_dequanted_b[N_ITERS][PACK_RATIO_4BITS];
+    scalar_t local_dequanted_b_8bits[N_ITERS][2][PACK_RATIO_8BITS];
+    v2f local_scales[N_ITERS];
+    v2f local_zeros[N_ITERS];
+    FragScaleLoading temp_scales;
+    PackType temp_zeros;
+    float output[BLOCKS_M][N_ITERS][4];
+    FragA temp_a[LOADING_A_LOOP];
+
+    TileManager<BLOCKS_M, BLOCKS_N, BLOCKS_K> tile_manager;
+    ThreadView tv;
+
+    __device__ __forceinline__ void set_address(const PackTypeInt4* a,
+        const PackTypeInt4* b,
+        PackTypeInt4* c,
+        PackTypeInt4* c_temp,
+        const PackTypeInt4* scale_ptr,
+        const PackTypeInt4* zp_ptr = nullptr) {
+            A = (const FragA*)a;
+            B = (const FragB*)b;
+            C = (FragC*)c;
+            C_temp = (float*)c_temp;
+            scales = (const FragScaleLoading*)scale_ptr;
+            if constexpr(w_type_id == vllm::kU4.id()) {
+                zeros = (const FragZeroLoading*)zp_ptr;
+            }
+    }
+
+    __device__ __forceinline__ bool debug() {
+        #ifdef DEBUG
+        bool do_print = tv.wave_idx == 1 && tv.slot_idx == 0 && tv.slot_tid == 0;
+        return do_print;
+        #else
+        return false;
+        #endif
+    }
+
+    __device__ __forceinline__ void next_k0() {
+        //reset bsm a to base
+        bsm_a_ptr = (PackTypeInt4*)smem_base;
+        bsm_a_ptr += tv.slot_tid * (DOUBLE_PAD_SLICE_K / (sizeof(PackTypeInt4) / sizeof(scalar_t))) + tv.slot_idx;
+    }
+
+    __device__ __forceinline__ void next_k1() {
+        //Update only bsm_a_ptr
+        bsm_a_ptr = (PackTypeInt4*)smem_base;
+        bsm_a_ptr += tv.slot_tid * (DOUBLE_PAD_SLICE_K / (sizeof(PackTypeInt4) / sizeof(scalar_t))) + tv.slot_idx + WAVE_SLOTS;
+        //load k32~k63
+        //bsm_a_ptr += 4;
+    }
+
+    __device__ __forceinline__ void next_k0_pre() {
+        //A_loading += SLICE_K / FragACount;
+        //B_loading += SLICE_K / PACK_RATIO_4BITS * n;
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            B_loading += SLICE_K / PACK_RATIO_4BITS * n;
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            B_loading += SLICE_K / PACK_RATIO_8BITS * n;
+        }
+    }
+
+    __device__ __forceinline__ void next_k1_pre() {
+        A_loading += DOUBLE_SLICE_K / FragACount;
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            B_loading += SLICE_K / PACK_RATIO_4BITS * n;
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            B_loading += SLICE_K / PACK_RATIO_8BITS * n;
+        }
+    }
+
+    __device__ __forceinline__ void ldg_a(int k_idx) {
+        //32x64/2/256 = 16 / 4 = 4
+        int t = tv.tid;
+        int k_broad = tile_manager.tile_start_row * TILE_K + k_idx * SLICE_K;
+        #pragma unroll LOADING_A_LOOP
+        for (int i = 0; i < LOADING_A_LOOP; i++)  {
+            int reading_m = t / (DOUBLE_SLICE_K / FragACount);
+            int reading_k = t % (DOUBLE_SLICE_K / FragACount);
+            int gvm_offset = reading_m * k / FragACount + reading_k;
+            FragA* gvm_addr = (FragA*)A_loading + gvm_offset;
+            //FIXME: we cannot do slice k pad as ldg_b32_bsm_async seems does not support padding
+            if constexpr(HAS_M_PRED && HAS_NK_PRED) {
+                bool pred = reading_m < m;
+                bool pred_k = k_broad + reading_k * FragACount < k;
+                pred = pred && pred_k && tile_manager.global_pred;
+                ldg_b64_reg_noasync(temp_a[i], gvm_addr, pred, true);
+            } else if constexpr(HAS_M_PRED) {
+                bool pred = reading_m < m && tile_manager.global_pred;
+                ldg_b64_reg_noasync(temp_a[i], gvm_addr, pred, true);
+            } else if constexpr(HAS_NK_PRED) {
+                bool pred_k = k_broad + reading_k * FragACount < k && tile_manager.global_pred;
+                ldg_b64_reg_noasync(temp_a[i], gvm_addr, pred_k, true);
+            } else {
+                ldg_b64_reg_noasync(temp_a[i], gvm_addr, tile_manager.global_pred, true);
+            }
+            t += THREADS;
+        }
+    }
+
+    __device__ __forceinline__ void sts_a() {
+        FragA* to_bsm_a_ptr = (FragA*)smem_base;
+        int t = tv.tid;
+        #pragma unroll LOADING_A_LOOP
+        for (int i = 0; i < LOADING_A_LOOP; i++)  {
+            int reading_m = t / (DOUBLE_SLICE_K / FragACount);
+            int reading_k = t % (DOUBLE_SLICE_K / FragACount);
+            int bsm_offset = reading_m * (DOUBLE_PAD_SLICE_K / FragACount) + reading_k;
+            *(to_bsm_a_ptr + bsm_offset) = temp_a[i];
+            t += THREADS;
+        }
+    }
+
+    __device__ __forceinline__ void lds_a(int midx) {
+        *((PackTypeInt4*)local_a[midx]) = *bsm_a_ptr;
+        bsm_a_ptr += SLOT * (DOUBLE_PAD_SLICE_K / (sizeof(PackTypeInt4) / sizeof(scalar_t)));
+    }
+
+    //TODO: implement when N_ITERS==1 or N_ITERS==3
+    //korder used gptq_8bits, ldg_b will load two times in one SLICE_K
+    //For example, t0 loads packed_k0, packed_k1, and packed_k0 represents a packed 4 ks in first line of B,
+    //and packed_k1 represents a packed 4 ks in second line of B
+    __device__ __forceinline__ void ldg_b(int k_idx, int korder = 0) {
+        if constexpr(HAS_NK_PRED) {
+            bool pred_k = tile_manager.tile_start_row * TILE_K + k_idx * SLICE_K + tv.slot_idx * PACK_RATIO_4BITS + korder < k;
+            bool pred_n = tile_manager.tile_start_col * TILE_N + (tv.wave_idx * SLOT + tv.slot_tid)*N_ITERS < n;
+            bool pred = pred_n && pred_k && tile_manager.global_pred;
+            FragB* addr =  (FragB*)B_loading + korder * n;
+            if constexpr(N_ITERS == 2) {
+                ldg_b64_reg_noasync(*((PackTypeInt2*)local_b_cache), ((PackTypeInt2*)addr), pred, true);
+            } else if constexpr(N_ITERS == 4) {
+                ldg_b128_reg_noasync(*((PackTypeInt4*)local_b_cache), ((PackTypeInt4*)addr), pred, true);
+            }
+        } else {
+            FragB* addr =  (FragB*)B_loading + korder * n;
+            if constexpr(N_ITERS == 2) {
+                ldg_b64_reg_noasync(*((PackTypeInt2*)local_b_cache), ((PackTypeInt2*)addr), tile_manager.global_pred, true);
+            } else if constexpr(N_ITERS == 4) {
+                ldg_b128_reg_noasync(*((PackTypeInt4*)local_b_cache), ((PackTypeInt4*)addr), tile_manager.global_pred, true);
+            }
+        }
+    }
+
+    __device__ __forceinline__ void swap_b_cache(int i) {
+        local_b[i] = local_b_cache[i];
+    }
+
+    __device__ __forceinline__ void ldg_scales() {
+        bool pred = tv.tid < TILE_N / (sizeof(FragScaleLoading) / sizeof(scalar_t)) && tile_manager.global_pred;
+        if constexpr(HAS_NK_PRED) {
+            pred = pred && tv.tid < (n - tile_manager.tile_start_col * TILE_N) / (sizeof(FragScaleLoading) / sizeof(scalar_t));
+        }
+        //FragScaleLoading *scale_bsm = (FragScaleLoading*)(smem_base + 0x2000) + tv.tid;
+        FragScaleLoading *gvm_addr = (FragScaleLoading*)scales_loading + tv.tid;
+        //ldg_b32_bsm_async(scale_bsm, gvm_addr, pred, false);
+        ldg_b32_reg_noasync(*((PackType*)&temp_scales), gvm_addr, pred, true);
+    }
+
+    __device__ __forceinline__ void ldg_zp() {
+        if constexpr(w_type_id == vllm::kU4.id()) {
+            bool pred = (tv.tid < TILE_N / PACK_RATIO_4BITS) && tile_manager.global_pred;
+            if constexpr(HAS_NK_PRED) {
+                pred = pred && tv.tid < ((n - tile_manager.tile_start_col * TILE_N) / PACK_RATIO_4BITS);
+            }
+            FragZeroLoading *gvm_addr = (FragZeroLoading*)zeros_loading + tv.tid;
+            ldg_b32_reg_noasync(*((PackType*)&temp_zeros), gvm_addr, pred, true);
+        }
+    }
+
+    __device__ __forceinline__ void sts_scales() {
+        FragScaleLoading *scale_bsm = (FragScaleLoading*)(smem_base + 0x3000) + tv.tid;
+        *scale_bsm = temp_scales;
+    }
+
+    __device__ __forceinline__ void sts_zeros() {
+        if constexpr(w_type_id == vllm::kU4.id()) {
+            bool pred = (tv.tid < TILE_N / PACK_RATIO_4BITS) && tile_manager.global_pred;
+            if (pred) {
+                float temp[PACK_RATIO_4BITS];
+                decompress_zero_4bits(temp_zeros, temp);
+                float *zeros_bsm = (float*)(smem_base + 0x3400) + tv.tid * PACK_RATIO_4BITS;
+                for (int i = 0; i < PACK_RATIO_4BITS; i++) {
+                    *(zeros_bsm + i) = temp[i];
+                }
+            }
+        }
+    }
+
+    //TODO: implement when N_ITERS==1 or N_ITERS==3
+    __device__ __forceinline__ void lds_scales() {
+        if constexpr(N_ITERS==2) {
+            *((half2*)local_dequanted_b[0]) = *((half2*)bsm_scales_ptr);
+        } else if constexpr(N_ITERS==4) {
+            *((PackTypeInt2*)local_dequanted_b[0]) = *((PackTypeInt2*)bsm_scales_ptr);
+        }
+    }
+
+    __device__ __forceinline__ void pack_scales() {
+        if constexpr(w_type_id == vllm::kU4B8.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                float s = local_dequanted_b[0][i];
+                float z = -8 * s;
+                local_scales[i] = {s, s};
+                local_zeros[i] = {z, z};
+            }
+        } else if constexpr(w_type_id == vllm::kU4.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                float s;
+                if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+                    s = __bfloat162float(local_dequanted_b[0][i]);
+                } else {
+                    s = local_dequanted_b[0][i];
+                }
+                float z = *(bsm_zeros_ptr + i);
+                z = z * s;
+                local_scales[i] = {s, s};
+                local_zeros[i] = {z, z};
+            }
+        } else if constexpr(w_type_id == vllm::kU8B128.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                float s = local_dequanted_b[0][i];
+                float z = -128 * s;
+                local_scales[i] = {s, s};
+                local_zeros[i] = {z, z};
+            }
+        } else if constexpr(w_type_id == vllm::kU8.id()) {
+            // should apply zeros
+        }
+    }
+
+    __device__ __forceinline__ void dequant(int kdx, int korder = 0) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            dequant_gptq_4bits<scalar_t>(local_b[kdx], local_dequanted_b[kdx], local_scales[kdx], local_zeros[kdx]);
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            dequant_gptq_8bits<scalar_t>(local_b[kdx], local_dequanted_b_8bits[kdx][korder], local_scales[kdx], local_zeros[kdx]);
+        }
+    }
+
+    __device__ __forceinline__ void matmul(int mdx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                mma_16x16x16<scalar_t>(local_a[mdx][0], *((PackTypeInt2*)local_dequanted_b[i]), *((PackTypeInt4*)output[mdx][i]));
+                mma_16x16x16<scalar_t>(local_a[mdx][1], *((PackTypeInt2*)local_dequanted_b[i] + 1), *((PackTypeInt4*)output[mdx][i]));
+            }
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            #pragma unroll
+            for (int i = 0; i < N_ITERS; i++) {
+                mma_16x16x16<scalar_t>(local_a[mdx][0], *((PackTypeInt2*)local_dequanted_b_8bits[i][0]), *((PackTypeInt4*)output[mdx][i]));
+                mma_16x16x16<scalar_t>(local_a[mdx][1], *((PackTypeInt2*)local_dequanted_b_8bits[i][1]), *((PackTypeInt4*)output[mdx][i]));
+            }
+        }
+    }
+
+    __device__ __forceinline__ void clear_c() {
+        #pragma unroll
+        for (int miter = 0; miter < BLOCKS_M; miter++) {
+            #pragma unroll
+            for (int niter = 0; niter < N_ITERS; niter++) {
+                #pragma unroll
+                for (int miter2 = 0; miter2 < 4; miter2++) {
+                    output[miter][niter][miter2] = 0;
+                }
+            }
+        }
+    }
+
+    //functions for preloading next tile data
+    __device__ __forceinline__ void init_address_pre(int _m, int _n, int _k, int _quant_group_power2, int _bidx, int _iters, uint8_t *_smem_base) {
+        tv.init();
+        m = _m;
+        n = _n;
+        k = _k;
+        quant_group_power2 = _quant_group_power2;
+        bidx = _bidx;
+        smem_base = _smem_base;
+        tile_manager.init(m, n, k, bidx, _iters);
+        init_tile_pre(tile_manager.tile_start_col, tile_manager.tile_start_row);
+    }
+
+     __device__ __forceinline__ void init_tile_pre(int col, int row) {
+        //Initialize start slice address and set them to A_loading and B_loading
+        int offset_n = col * TILE_N;
+        int offset_k = row * TILE_K;
+        //A_loading address will always be valid
+        A_loading = A + offset_k / (FragACount);
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            B_loading = B + (offset_k / PACK_RATIO_4BITS * n + offset_n) / (FragBCount) + tv.slot_idx * n + (tv.wave_idx * SLOT + tv.slot_tid)*N_ITERS;
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            B_loading = B + (offset_k / PACK_RATIO_8BITS * n + offset_n) / (FragBCount) + tv.slot_idx * n * 2 + (tv.wave_idx * SLOT + tv.slot_tid)*N_ITERS;
+        }
+        scales_loading = scales + ((offset_k >> quant_group_power2) * n + offset_n) / (sizeof(FragScaleLoading)/sizeof(scalar_t));
+        if constexpr(w_type_id == vllm::kU4.id()) {
+            zeros_loading = zeros + ((offset_k >> quant_group_power2) * n + offset_n) / PACK_RATIO_4BITS;
+        }
+    }
+
+    __device__ __forceinline__ void next_tile_pre() {
+        tile_manager.next_tile_pre();
+        init_tile_pre(tile_manager.tile_start_col, tile_manager.tile_start_row);
+    }
+
+    __device__ __forceinline__ void init_bsm_addr() {
+        bsm_a_ptr = (PackTypeInt4*)smem_base;           //use 8k bytes, will load at most 32x128*sizeof(half), either m32k128 or m128k32
+        //remaining_bsm_ptr = (float*)(smem_base + 0x2000 + 0x1000); //3k bytes
+        bsm_a_ptr += tv.slot_tid * (DOUBLE_PAD_SLICE_K / (sizeof(PackTypeInt4) / sizeof(scalar_t))) + tv.slot_idx;
+        bsm_scales_ptr = (scalar_t*)(smem_base + 0x3000);      //use 128xsizeof(float)*2 = 1k
+        bsm_scales_ptr += (tv.wave_idx * SLOT + tv.slot_tid) * N_ITERS;
+        if constexpr(w_type_id == vllm::kU4.id()) {
+            bsm_zeros_ptr = (float*)(smem_base + 0x3400);      //use 128xsizeof(float)*2 = 1k
+            bsm_zeros_ptr += (tv.wave_idx * SLOT + tv.slot_tid) * N_ITERS;
+        }
+    }
+
+    __device__ __forceinline__ void write_c(int offset, const float& v) {
+        if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+            #ifdef BF16_HIGH_PRECISION
+	        atomicAdd(C_temp + offset, v);
+            #else
+            atomicAdd(C + offset, (atomic_type)v);
+            #endif
+        } else {
+            atomicAdd(C + offset, (scalar_t)v);
+        }
+    }
+
+    //atomic write to c
+    __device__ __forceinline__ void write_c_pre() {
+        int k_broad = tv.slot_idx * 4;
+        int n_broad = (tv.wave_idx * SLOT + tv.slot_tid) * N_ITERS + tile_manager.tile_start_col_cache * TILE_N;
+        #pragma unroll
+        for (int miter = 0; miter < BLOCKS_M; miter++) {
+            #pragma unroll
+            for (int niter = 0; niter < N_ITERS; niter++) {
+                int store_n = n_broad + niter;
+                #pragma unroll
+                for (int miter2 = 0; miter2 < 4; miter2++) {
+                    int store_m = k_broad + miter * SLICE_M + miter2;
+                    if constexpr(HAS_M_PRED && HAS_NK_PRED) {
+                        if (store_m < m && store_n < n) {
+			                write_c(store_m * n + store_n, output[miter][niter][miter2]);
+                        }
+                    } else if constexpr(HAS_M_PRED) {
+                        if (store_m < m) {
+			                write_c(store_m * n + store_n, output[miter][niter][miter2]);
+                        }
+                    } else if constexpr(HAS_NK_PRED) {
+                        if (store_n < n) {
+			                write_c(store_m * n + store_n, output[miter][niter][miter2]);
+                        }
+                    } else {
+			            write_c(store_m * n + store_n, output[miter][niter][miter2]);
+                    }
+                }
+            }
+        }
+    }
+
+    template<int K, bool KTAIL>
+    __device__ __forceinline__ void on_sts_a() {
+        if constexpr(K == 0) {
+            sts_a(); // reorder data_a from reg(gvm resouce) to bsm
+        }
+    }
+
+    template<int K, bool KTAIL>
+    __device__ __forceinline__ void on_preload(int k_idx) {
+        if constexpr(!KTAIL) {
+            if constexpr(K == 0) {
+                next_k0_pre(); // preload gvm a/b
+            } else {
+                next_k1_pre(); // preload gvm a/b
+            }
+            ldg_b(k_idx + K + 1); //preload b for next k
+            if constexpr(K == 1) {
+                ldg_a(k_idx + K + 1); //preload a for next k
+            }
+        } else {
+            next_tile_pre();
+            ldg_scales();
+            ldg_zp();
+            ldg_b(0);
+            ldg_a(0);
+        }
+    }
+
+    template<int K, bool KTAIL>
+    __device__ __forceinline__ void on_dequant_niters1(int k_idx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            swap_b_cache(0);
+            on_sts_a<K,KTAIL>(); // reorder data_a from reg(gvm resouce) to bsm
+            on_preload<K,KTAIL>(k_idx);
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0); //dequant b0
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            swap_b_cache(0);
+            ldg_b(0, 1);
+            dequant(0, 0);
+            on_sts_a<K,KTAIL>(); // reorder data_a from reg(gvm resouce) to bsm
+            swap_b_cache(0);
+            on_preload<K,KTAIL>(k_idx);
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0, 1); //dequant b0
+        }
+    }
+
+    template<int K, bool KTAIL>
+    __device__ __forceinline__ void on_dequant_niters2(int k_idx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            swap_b_cache(0);
+            dequant(0); //dequant b0
+            on_sts_a<K,KTAIL>();
+            swap_b_cache(1);
+            on_preload<K,KTAIL>(k_idx);
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(1);   //dequant b64
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            swap_b_cache(0);
+            swap_b_cache(1);
+            ldg_b(0, 1);
+            dequant(0); //dequant b0
+            on_sts_a<K,KTAIL>();
+            dequant(1);
+            swap_b_cache(0);
+            swap_b_cache(1);
+            on_preload<K,KTAIL>(k_idx);
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0, 1);   //dequant b64
+            dequant(1, 1);   //dequant b64
+        }
+    }
+
+    template<int K, bool KTAIL>
+    __device__ __forceinline__ void on_dequant_niters3(int k_idx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            swap_b_cache(0);
+            dequant(0); //dequant b0
+            on_sts_a<K,KTAIL>();
+            swap_b_cache(1);
+            dequant(1);
+            swap_b_cache(2);
+            on_preload<K,KTAIL>(k_idx);
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(2); //dequant b0
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            swap_b_cache(0);
+            swap_b_cache(1);
+            swap_b_cache(2);
+            ldg_b(0, 1);
+            dequant(0); //dequant b0
+            on_sts_a<K,KTAIL>();
+            swap_b_cache(1);
+            dequant(1);
+            dequant(2);
+            swap_b_cache(0);
+            swap_b_cache(1);
+            swap_b_cache(2);
+            on_preload<K,KTAIL>(k_idx);
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0, 1); //dequant b0
+            dequant(1, 1); //dequant b0
+            dequant(2, 1); //dequant b0
+        }
+    }
+
+    template<int K, bool KTAIL>
+    __device__ __forceinline__ void on_dequant_niters4(int k_idx) {
+        if constexpr(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU4B8.id()) {
+            swap_b_cache(0);
+            dequant(0); //dequant b0
+            on_sts_a<K,KTAIL>();
+            swap_b_cache(1);
+            dequant(1); //dequant b1
+            swap_b_cache(2);
+            swap_b_cache(3);
+            dequant(2); //dequant b2
+            on_preload<K,KTAIL>(k_idx);
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(3); //dequant b3
+        } else if constexpr(w_type_id == vllm::kU8.id() || w_type_id == vllm::kU8B128.id()) {
+            swap_b_cache(0);
+            swap_b_cache(1);
+            swap_b_cache(2);
+            swap_b_cache(3);
+            ldg_b(0, 1);
+            dequant(0); //dequant b0
+            on_sts_a<K,KTAIL>();
+            dequant(1); //dequant b1
+            dequant(2); //dequant b2
+            dequant(3); //dequant b2
+            swap_b_cache(0);
+            swap_b_cache(1);
+            swap_b_cache(2);
+            swap_b_cache(3);
+            on_preload<K,KTAIL>(k_idx);
+            barrier_bsm; //sync threads for sts_a
+            lds_a(0); //load first 16 batches into registers
+            dequant(0, 1); //dequant b0
+            dequant(1, 1); //dequant b1
+            dequant(2, 1); //dequant b2
+            dequant(3, 1); //dequant b3
+        }
+    }
+
+    template<int K, bool KTAIL>
+    __device__ __forceinline__ void on_dequant(int kdx) {
+        if constexpr(N_ITERS == 1) on_dequant_niters1<K,KTAIL>(kdx);
+        else if constexpr(N_ITERS == 2) on_dequant_niters2<K,KTAIL>(kdx);
+        else if constexpr(N_ITERS == 3) on_dequant_niters3<K,KTAIL>(kdx);
+        else if constexpr(N_ITERS == 4) on_dequant_niters4<K,KTAIL>(kdx);
+    }
+};
+
+template<typename scalar_t,
+    const vllm::ScalarTypeId w_type_id,
+    const int THREADS,          // number of threads in a threadblock
+    const int BLOCKS_M,         // number of 16x16 blocks in the m
+                                // dimension (batchsize) of the
+                                // threadblock
+    const int BLOCKS_N,         // same for n dimension (output)
+    const int BLOCKS_K,         // same for k dimension (reduction)
+    const bool HAS_ACT_ORDER,   // whether act_order is enabled
+    const bool HAS_ZP,          // whether zero-points are enabled
+    const bool HAS_M_PRED = true,  //If we should use predictors to load m from gvm
+    const bool HAS_NK_PRED = true  //If we should use predictors to load nk from gvm
+    >
+__global__ void hgemm_gptq(
+    const PackTypeInt4* __restrict__ A,  // fp16 input matrix of shape mxk
+    const PackTypeInt4* __restrict__ B,  // 4bit quantized weight matrix of shape kxn
+    PackTypeInt4* __restrict__ C,        // fp16 output buffer of shape mxn
+    PackTypeInt4* __restrict__ C_tmp,    // fp32 tmp output buffer (for reduce)
+    const PackTypeInt4* __restrict__ scales_ptr,  // fp16 quantization scales of shape
+                                          // (k/groupsize)xn
+    const PackTypeInt4* __restrict__ zp_ptr,      // 4bit packed zero-points of shape
+                                          // (k/groupsize)x(n/pack_factor)
+    const int* __restrict__ g_idx,        // int32 group indices of shape k
+    int prob_m,           // batch dimension m
+    int prob_n,           // output dimension n
+    int prob_k,           // reduction dimension k
+    int quant_group_power2,
+    int max_iters,        // max tile iterations for one block
+    int* locks,           // extra global storage for barrier synchronization
+    bool use_fp32_reduce  // whether to use fp32 global reduce
+) {
+    int bidx = blockIdx.x;
+    __shared__ uint8_t smem_base[0x4000]; //4x16x256 = 16Kbytes
+    using LoadingManagerType = LoadingManager<scalar_t, w_type_id, THREADS, BLOCKS_M, BLOCKS_N, BLOCKS_K, HAS_ACT_ORDER, HAS_ZP, HAS_M_PRED, HAS_NK_PRED>;
+    LoadingManagerType loading_manager;
+    A += blockIdx.y * (MAX_BLOCKS_M * SLICE_M) * prob_k / (sizeof(PackTypeInt4) / sizeof(scalar_t));
+    C += blockIdx.y * (MAX_BLOCKS_M * SLICE_M) * prob_n / (sizeof(PackTypeInt4) / sizeof(scalar_t));
+    #ifdef BF16_HIGH_PRECISION
+    if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+        C_tmp += blockIdx.y * (MAX_BLOCKS_M * SLICE_M) * prob_n / (sizeof(PackTypeInt4) / sizeof(float));
+    }
+    #endif
+    loading_manager.set_address(A, B, C, C_tmp, scales_ptr, zp_ptr);
+    //loading_manager.init_address(prob_m, prob_n, prob_k, bidx, max_iters, smem_base);
+    loading_manager.init_address_pre(std::min<int>(MAX_BLOCKS_M*SLICE_M, prob_m - blockIdx.y * (MAX_BLOCKS_M * SLICE_M)), prob_n, prob_k, quant_group_power2, bidx, max_iters, smem_base);
+
+    loading_manager.ldg_scales(); //Load all scales to bsm
+    loading_manager.ldg_zp();
+    loading_manager.ldg_b(0);    //load b in k0~31
+    loading_manager.ldg_a(0);    //Load first k0~63 and all m
+    loading_manager.clear_c();
+
+    while (max_iters > 0) {
+        loading_manager.init_bsm_addr(); //reset all bsm address for current tile
+        loading_manager.sts_scales();
+        loading_manager.sts_zeros();
+        barrier_bsm;
+        loading_manager.lds_scales(); //load scale0 and scale64
+        loading_manager.pack_scales(); //pack scales into two v2f structure
+
+        int k_idx = 0;
+        if constexpr(BLOCKS_K / 2 - 1 > 0) {
+            #pragma unroll BLOCKS_K / 2 - 1
+            for (int kloop = 0; kloop < BLOCKS_K / 2 - 1; kloop++) {
+                int m_idx = 0;
+                loading_manager.template on_dequant<0, false>(k_idx);
+                if constexpr(BLOCKS_M > 1) {
+                    #pragma unroll BLOCKS_M - 1
+                    for (; m_idx < BLOCKS_M - 1; m_idx++) {
+                        loading_manager.lds_a(m_idx+1); //pre load next 16 batches
+                        loading_manager.matmul(m_idx); //do matmul
+                    }
+                }
+                barrier_bsm; //wait for the last 16 batches ready, becasue of k_idx will be change
+                loading_manager.next_k1(); //modify gvm/bsm address of a and b
+                loading_manager.matmul(m_idx); //do matmul
+                m_idx = 0;
+                loading_manager.template on_dequant<1, false>(k_idx);
+                if constexpr(BLOCKS_M > 1) {
+                    #pragma unroll BLOCKS_M - 1
+                    for (; m_idx < BLOCKS_M - 1; m_idx++) {
+                        loading_manager.lds_a(m_idx+1); //pre load next 16 batches
+                        loading_manager.matmul(m_idx); //do matmul
+                    }
+                }
+                barrier_bsm; //wait for the last 16 batches ready, becasue of k_idx will be change
+                loading_manager.next_k0(); //modify gvm/bsm address of a and b
+                loading_manager.matmul(m_idx); //do matmul
+                k_idx += 2;
+            }
+        }
+
+        int m_idx = 0;
+        loading_manager.template on_dequant<0, false>(k_idx);
+        if constexpr(BLOCKS_M > 1) {
+            #pragma unroll BLOCKS_M - 1
+            for (; m_idx < BLOCKS_M - 1; m_idx++) {
+                loading_manager.lds_a(m_idx+1); //pre load next 16 batches
+                loading_manager.matmul(m_idx); //do matmul
+            }
+        }
+        barrier_bsm; //wait for the last 16 batches ready, becasue of k_idx will be change
+        loading_manager.next_k1(); //modify gvm/bsm address of a and b
+        loading_manager.matmul(m_idx); //do matmul
+        m_idx = 0;
+        loading_manager.template on_dequant<1, true>(k_idx);
+        if constexpr(BLOCKS_M > 1) {
+            #pragma unroll BLOCKS_M - 1
+            for (; m_idx < BLOCKS_M - 1; m_idx++) {
+                loading_manager.lds_a(m_idx+1); //pre load next 16 batches
+                loading_manager.matmul(m_idx); //do matmul
+            }
+        }
+        barrier_bsm; //wait for the last 16 batches ready, becasue of k_idx will be change
+        //loading_manager.next_tile_pre(); should move into on_dequant ?
+        loading_manager.matmul(m_idx); //do matmul
+
+        max_iters--;
+
+        if (loading_manager.tile_manager.need_save_data_pre()) {
+            loading_manager.write_c_pre(); // reduce and write back
+            loading_manager.clear_c();
+        }
+
+        barrier_bsm;
+    }
+}
+} //end of namespace __hgemm_even_blocks_k
+
+template<typename scalar_t,
+    const vllm::ScalarTypeId w_type_id,
+    int THREADS,
+    int BLOCKS_M,
+    int BLOCKS_N,
+    int BLOCKS_K,
+    bool HAS_ACT_ORDER,
+    bool HAS_ZP,
+    bool HAS_M_PRED,
+    bool HAS_NK_PRED>
+bool launch_gemm_gptq_kernel(const PackTypeInt4* A,
+    const PackTypeInt4* B,
+    PackTypeInt4* C,
+    PackTypeInt4* C_temp,
+    const PackTypeInt4* scales,
+    const PackTypeInt4* zeros,
+    int* g_idx, int m, int n, int k, int quant_group, int chunks, cudaStream_t stream = nullptr) {
+    int tiles_m = div_ceil(m, TILE_M);
+    int tiles_n = div_ceil(n, TILE_N);
+    int tiles_k = div_ceil(k, TILE_K);
+    if (TILE_K > quant_group && TILE_K % quant_group != 0) {
+        printf("Invalid TILE_K %d that can not be dived by QUANT_GROUP %d\n", TILE_K, quant_group);
+        return false;
+    }
+
+    int total_tiles = tiles_n * tiles_k;
+    int blocks = PEUS;
+    int iters = div_ceil(total_tiles, PEUS);
+    if (total_tiles < PEUS) {
+        if (TILE_K < quant_group) {
+            iters = quant_group / TILE_K;
+            blocks = div_ceil(total_tiles, iters);
+        } else {
+            iters = 1;
+            blocks = total_tiles;
+        }
+    } else {
+        if (TILE_K < quant_group) {
+            iters = div_ceil(iters, quant_group / TILE_K) * quant_group / TILE_K;
+            blocks = div_ceil(total_tiles, iters);
+        }
+    }
+    while (iters * blocks - total_tiles >= iters) {
+        blocks -= 1;
+    }
+
+    if (total_tiles < blocks) {
+        printf("total slice %d < blocks %d, Invalid configure\n", total_tiles, blocks);
+        return false;
+    }
+    // printf("Launching hgemm_gptq_4bits THREADS=%d, BLOCKS_M=%d, BLOCKS_N=%d, BLOCKS_K=%d, HAS_M_PRED=%d, HAS_NK_PRED=%d, tiles_n=%d, tiles_k = %d, total_tiles = %d, iters = %d, blocks = %d, chunks = %d, TILE=m%dn%dk%d\n",
+    // THREADS, BLOCKS_M, BLOCKS_N, BLOCKS_K, HAS_M_PRED, HAS_NK_PRED, tiles_n, tiles_k, total_tiles, iters, blocks, chunks, BLOCKS_M*SLICE_M, BLOCKS_N*SLICE_N, BLOCKS_K*SLICE_K
+    // );
+
+    if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+        size_t num_elem = size_t(m) * size_t(n);
+        size_t clean_blocks = std::max(size_t(1), (num_elem + clean_kernel_thread_num * clean_kernel_pack_num - 1)/ (clean_kernel_thread_num * clean_kernel_pack_num));
+        clean_zero<clean_kernel_thread_num, clean_kernel_pack_num><<<clean_blocks, clean_kernel_thread_num, 0, stream>>>((float*)C_temp, num_elem);
+    }
+
+    //It is better to do perm before launch kernel
+    if constexpr(BLOCKS_K % 2 == 1) {
+        __hgemm_singular_blocks_k::hgemm_gptq<scalar_t,
+            w_type_id,
+            THREADS,
+            BLOCKS_M, BLOCKS_N, BLOCKS_K,
+            HAS_ACT_ORDER, HAS_ZP, HAS_M_PRED, HAS_NK_PRED
+            ><<<dim3(std::max(blocks,1), std::max(chunks,1), 1), THREADS, 0, stream>>>(A, B, C, C_temp, scales, zeros, g_idx, m, n, k, get_power2(quant_group), iters, nullptr, false);
+    } else {
+        __hgemm_even_blocks_k::hgemm_gptq<scalar_t,
+            w_type_id,
+            THREADS,
+            BLOCKS_M, BLOCKS_N, BLOCKS_K,
+            HAS_ACT_ORDER, HAS_ZP, HAS_M_PRED, HAS_NK_PRED
+            ><<<dim3(std::max(blocks,1), std::max(chunks,1), 1), THREADS, 0, stream>>>(A, B, C, C_temp, scales, zeros, g_idx, m, n, k, get_power2(quant_group), iters, nullptr, false);
+    }
+
+    if constexpr(std::is_same_v<scalar_t, __maca_bfloat16>) {
+        size_t num_elem = size_t(m) * size_t(n);
+        size_t reduce_blocks = std::max(size_t(1), (num_elem  + reduce_kernel_thread_num * reduce_kernel_pack_num - 1) / (reduce_kernel_thread_num * reduce_kernel_pack_num));
+        all_reduce<reduce_kernel_thread_num, reduce_kernel_pack_num, false><<<reduce_blocks, reduce_kernel_thread_num, 0, stream>>>((float*)C_temp, (maca_bfloat16*)C, num_elem);
+    }
+
+    return true;
+}
+
+}
+
diff --git a/csrc/quantization/gptq/hgemv_nn_splitk_gptq.hpp b/csrc/quantization/gptq/hgemv_nn_splitk_gptq.hpp
new file mode 100644
index 000000000..6e0fdb0e3
--- /dev/null
+++ b/csrc/quantization/gptq/hgemv_nn_splitk_gptq.hpp
@@ -0,0 +1,633 @@
+// 2024 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
+#pragma once
+
+#include <mc_runtime.h>
+#include <maca_fp16.h>
+#include "Hgemm_common.cuh"
+#include "dequant.cuh"
+#define quant_packed_type uint32_t
+typedef __NATIVE_VECTOR__(2, float) v2f;
+template <int ThreadBlock, int BlockDimX, int BATCH>
+__global__ __launch_bounds__(256) void hgemv_nn_splitk_gptq(const half* __restrict__ srcB,
+                                                        const quant_packed_type* __restrict__ srcA,
+                                                        const quant_packed_type* __restrict__ zeros,
+                                                        const half* __restrict__ scales,
+                                                        half *dst,
+                                                        int m,
+                                                        int n,
+                                                        int k,
+                                                        int srcAStride,
+                                                        int dstStride,
+                                                        int k_div_sk,
+                                                        const int* __restrict__ b_perm = nullptr) {
+    constexpr int QBITS = 4;
+    constexpr int PACK_RATIO = (32 / QBITS);
+    constexpr int QUANT_GROUP = 128;
+    constexpr int LoopNum = QUANT_GROUP / (ThreadBlock / BlockDimX);
+    constexpr int N = BATCH;
+    const int k_stride = k;
+    const int k_block = (k_div_sk + QUANT_GROUP - 1) / QUANT_GROUP * QUANT_GROUP;
+    const int splitKOffset = blockIdx.y * k_block;
+    if (splitKOffset + k_block > k) k = k - splitKOffset;
+    else k = k_block;
+    srcA += splitKOffset * srcAStride / PACK_RATIO;
+    //srcB += splitKOffset;
+
+    constexpr int quant_groups = 1;
+    constexpr int m_per_thread = sizeof(float4)/sizeof(quant_packed_type);
+    constexpr int thread_groups = ThreadBlock / BlockDimX;
+    constexpr int group_elements = BlockDimX * m_per_thread;
+    constexpr int reduce_size = ThreadBlock * m_per_thread;
+    constexpr int data_cache_size = LoopNum * thread_groups * N + group_elements * 2;
+    float *bsm_b_ptr, *bsm_zeros_ptr, *bsm_scales_ptr, *smem;
+    if constexpr(reduce_size > data_cache_size) {
+        __shared__ float bsm_ptr[reduce_size];
+        bsm_b_ptr = bsm_ptr;
+        bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+        bsm_scales_ptr = bsm_zeros_ptr + group_elements;
+        smem = bsm_ptr;
+    } else {
+        __shared__ float bsm_ptr[data_cache_size];
+        bsm_b_ptr = bsm_ptr;
+        bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+        bsm_scales_ptr = bsm_zeros_ptr + group_elements;
+        smem = bsm_ptr;
+    }
+
+    const int zeros_stride = srcAStride / PACK_RATIO;
+    const int scales_stride = srcAStride;
+    zeros += splitKOffset * zeros_stride / QUANT_GROUP;
+    scales += splitKOffset * scales_stride / QUANT_GROUP;
+
+    dst += group_elements * blockIdx.x;
+    const int m_offset_a = group_elements * blockIdx.x;
+    const int m_offset_zeros = group_elements / PACK_RATIO * blockIdx.x;
+    const int m_offset_scales = group_elements * blockIdx.x;
+
+    //store splited fma results
+    v2f c_splited[N][m_per_thread];
+    for (int i = 0; i < N; i++) {
+        for (int j = 0; j < m_per_thread; j++) c_splited[i][j] = {0, 0};
+    }
+
+    int tid = threadIdx.x;
+    int tidCol = tid / BlockDimX;
+    int tidRow = tid % BlockDimX;
+
+    for (int i = 0; i < k; i += LoopNum * thread_groups) {
+        int quant_group = i / QUANT_GROUP;
+        constexpr int loading_pack = 2;
+        constexpr int loading_count = group_elements / loading_pack;
+        //Load needed zeros, scales
+        for (int x = tid; x < loading_count; x += BlockDimX) {
+            uint8_t temp_zeros = 0;
+            temp_zeros = ((uint8_t*)(zeros + quant_group * zeros_stride + m_offset_zeros))[x];
+            half temp_scales[2];
+            *((float*)temp_scales) = ((float*)(scales + quant_group * scales_stride + m_offset_scales))[x];
+            uint32_t z = temp_zeros;
+            uint32_t z0 = __builtin_mxc_ubfe(z, 0, QBITS) + 1;
+            uint32_t z1 = __builtin_mxc_ubfe(z, 4, QBITS) + 1;
+            float s1 = (float)(temp_scales[0]);
+            float s2 = (float)(temp_scales[1]);
+            //Store to shared memory
+            bsm_zeros_ptr[x*2] = (float)z0 * s1 * -1.0f;
+            bsm_zeros_ptr[x*2+1] = (float)z1 * s2 * -1.0f;
+            bsm_scales_ptr[x*2] = s1;
+            bsm_scales_ptr[x*2+1] = s2;
+        }
+
+        int loop_index = 0;
+        quant_packed_type A[m_per_thread];
+        //Load A
+        *((float4*)A) = *(float4*)(srcA + (loop_index + i + tidCol * LoopNum) / PACK_RATIO * srcAStride + m_offset_a + m_per_thread * tidRow);
+
+        //Load B and transform to float
+        if (b_perm != nullptr) {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[b_perm[splitKOffset + i + tidCol * LoopNum + x] + y * k_stride];
+                }
+            }
+        } else {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[splitKOffset + i + tidCol * LoopNum + x + y * k_stride];
+                }
+            }
+        }
+
+        __syncthreads();
+
+        //Load zero and scale from bsm
+        int m_index = tidRow * m_per_thread;
+
+        v2f local_scales[m_per_thread];
+        v2f local_scales_d16[m_per_thread];
+        for (int c = 0; c < m_per_thread; c++) {
+            float s = bsm_scales_ptr[m_index + c];
+            float s_d16 = s / 16;
+            local_scales[c] = {s, s};
+            local_scales_d16[c] = {s_d16, s_d16};
+        }
+        v2f local_zeros[m_per_thread];
+        for (int c = 0; c < m_per_thread; c++) local_zeros[c] = {bsm_zeros_ptr[m_index + c],bsm_zeros_ptr[m_index + c]};
+
+        const int shuffled_dequant_index[PACK_RATIO] = {0, 4, 1, 5, 2, 6, 3, 7};
+        for (; loop_index < LoopNum; loop_index += PACK_RATIO) {
+            if constexpr (N <= 4) {
+                v2f local_b[PACK_RATIO/2*N];
+                for (int y = 0; y < N; y++) {
+                    for (int x = 0; x < PACK_RATIO / 2; x++) {
+                        local_b[x+PACK_RATIO/2*y].x = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+y*thread_groups*LoopNum];
+                        local_b[x+PACK_RATIO/2*y].y = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+1+y*thread_groups*LoopNum];
+                    }
+                }
+                #pragma unroll m_per_thread
+                for (int c = 0; c < m_per_thread; c++) {
+                    uint32_t p0 = A[c] & 0x0f0f0f0f;
+                    uint32_t p1 = A[c] & 0xf0f0f0f0;
+                    float o1,o2,o3,o4,o5,o6,o7,o8;
+                    asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o1):"r"(p0));
+                    asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o2):"r"(p0));
+                    asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o3):"r"(p0));
+                    asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o4):"r"(p0));
+                    asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o5):"r"(p1));
+                    asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o6):"r"(p1));
+                    asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o7):"r"(p1));
+                    asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o8):"r"(p1));
+                    v2f a0 = {o1, o3};
+                    v2f a1 = {o5, o7};
+                    v2f a2 = {o2, o4};
+                    v2f a3 = {o6, o8};
+
+                    a0 = __builtin_mxc_pk_fma_f32(a0, local_scales[c], local_zeros[c]);
+                    a1 = __builtin_mxc_pk_fma_f32(a1, local_scales_d16[c], local_zeros[c]);
+                    a2 = __builtin_mxc_pk_fma_f32(a2, local_scales[c], local_zeros[c]);
+                    a3 = __builtin_mxc_pk_fma_f32(a3, local_scales_d16[c], local_zeros[c]);
+
+                    #pragma unroll N
+                    for (int y = 0; y < N; y++) {
+                        c_splited[y][c] = __builtin_mxc_pk_fma_f32(a0, local_b[0+PACK_RATIO/2*y], c_splited[y][c]);
+                        c_splited[y][c] = __builtin_mxc_pk_fma_f32(a1, local_b[1+PACK_RATIO/2*y], c_splited[y][c]);
+                        c_splited[y][c] = __builtin_mxc_pk_fma_f32(a2, local_b[2+PACK_RATIO/2*y], c_splited[y][c]);
+                        c_splited[y][c] = __builtin_mxc_pk_fma_f32(a3, local_b[3+PACK_RATIO/2*y], c_splited[y][c]);
+                    }
+                }
+            } else {
+                v2f local_b;
+                for (int x = 0; x < PACK_RATIO / 2; x++) {
+                    for (int y = 0; y < N; y++) {
+                        local_b.x = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+y*thread_groups*LoopNum];
+                        local_b.y = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+1+y*thread_groups*LoopNum];
+                        for (int c = 0; c < m_per_thread; c++) {
+                            float a1 = (float)__builtin_mxc_ubfe(A[c], QBITS * shuffled_dequant_index[x*2], QBITS);
+                            float a2 = (float)__builtin_mxc_ubfe(A[c], QBITS * shuffled_dequant_index[x*2+1], QBITS);
+                            v2f a = {a1, a2};
+                            a = __builtin_mxc_pk_fma_f32(a, local_scales[c], local_zeros[c]);
+                            c_splited[y][c] = __builtin_mxc_pk_fma_f32(a, local_b, c_splited[y][c]);
+                        }
+                    }
+                }
+            }
+
+            if constexpr(thread_groups != QUANT_GROUP / PACK_RATIO) {
+                if (loop_index + PACK_RATIO < LoopNum) {
+                    *((float4*)A) = *(float4*)(srcA + (loop_index + PACK_RATIO + i + tidCol * LoopNum) / PACK_RATIO * srcAStride + m_offset_a + m_per_thread * tidRow);
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    #pragma unroll N
+    for (int y = 0; y < N; y++) {
+        for (int i = 0; i < m_per_thread; i++) {
+            smem[tidCol + (tidRow * m_per_thread + i) * ThreadBlock / BlockDimX] = c_splited[y][i].x + c_splited[y][i].y;
+        }
+        __syncthreads();
+        constexpr int stride = ThreadBlock / BlockDimX;
+        int data_size = ThreadBlock * m_per_thread;
+        #pragma unroll
+        for (int i = ThreadBlock / BlockDimX / 2; i > 0; i /= 2) {
+            for (int j = tid; j < data_size / 2; j += ThreadBlock) {
+                int reduce_group = j / i;
+                int reduce_index = j % i;
+                smem[reduce_index + reduce_group * stride] += smem[reduce_index + reduce_group * stride + i];
+            }
+            __syncthreads();
+            data_size /= 2;
+        }
+        for (int i = tid; i < m_per_thread * BlockDimX; i += ThreadBlock) {
+            atomicAdd(dst + i + y * dstStride, (half)smem[i*stride]);
+        }
+        if constexpr(N > 1) {
+            if (y + 1 < N) {
+                __syncthreads();
+            }
+        }
+    }
+}
+
+
+template<int BATCH>
+__global__ __launch_bounds__(256) void hgemv_nn_splitk_gptq_tb256_bx256_kb128(const half* __restrict__ srcB,
+                                                        const quant_packed_type* __restrict__ srcA,
+                                                        const quant_packed_type* __restrict__ zeros,
+                                                        const half* __restrict__ scales,
+                                                        half *dst,
+                                                        int m,
+                                                        int n,
+                                                        int k,
+                                                        int srcAStride,
+                                                        int dstStride,
+                                                        int k_div_sk,
+                                                        const int* __restrict__ b_perm = nullptr) {
+    constexpr int QBITS = 4;
+    constexpr int PACK_RATIO = (32 / QBITS);
+    constexpr int QUANT_GROUP = 128;
+    constexpr int ThreadBlock = 256;
+    constexpr int BlockDimX = 256;
+    constexpr int LoopNum = QUANT_GROUP / (ThreadBlock / BlockDimX);
+    constexpr int N = BATCH;
+    const int k_stride = k;
+    constexpr int k_block = QUANT_GROUP;
+    const int splitKOffset = blockIdx.y * k_block;
+    k = k_block;
+    srcA += splitKOffset * srcAStride / PACK_RATIO;
+    //srcB += splitKOffset;
+
+    constexpr int quant_groups = 1;
+    constexpr int m_per_thread = sizeof(float4)/sizeof(quant_packed_type);
+    constexpr int thread_groups = ThreadBlock / BlockDimX;
+    constexpr int group_elements = BlockDimX * m_per_thread;
+    constexpr int reduce_size = ThreadBlock * m_per_thread;
+    constexpr int data_cache_size = LoopNum * thread_groups * N + group_elements + group_elements / 2;
+    float *bsm_b_ptr, *bsm_zeros_ptr;
+    half *bsm_scales_ptr;
+
+    __shared__ float bsm_ptr[2048];  //128*N+256*4+256*4/2 = 最大8K，每个block 256线程，占据半个PEU，一个AP可以跑满8个block
+    bsm_b_ptr = bsm_ptr;
+    bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+    bsm_scales_ptr = (half*)(bsm_zeros_ptr + group_elements);
+    //smem = bsm_ptr;
+
+    const int zeros_stride = srcAStride / PACK_RATIO;
+    const int scales_stride = srcAStride;
+    zeros += splitKOffset * zeros_stride / QUANT_GROUP;
+    scales += splitKOffset * scales_stride / QUANT_GROUP;
+
+    dst += group_elements * blockIdx.x;
+    const int m_offset_a = group_elements * blockIdx.x;
+    const int m_offset_zeros = group_elements / PACK_RATIO * blockIdx.x;
+    const int m_offset_scales = group_elements * blockIdx.x;
+
+    //store splited fma results
+    v2f c_splited[N][m_per_thread];
+    for (int i = 0; i < N; i++) {
+        for (int j = 0; j < m_per_thread; j++) c_splited[i][j] = {0, 0};
+    }
+
+    int tid = threadIdx.x;
+    int tidCol = tid / BlockDimX;
+    int tidRow = tid % BlockDimX;
+
+    for (int i = 0; i < k; i += LoopNum * thread_groups) {
+        int quant_group = i / QUANT_GROUP;
+        constexpr int loading_pack = 2;
+        constexpr int loading_count = group_elements / loading_pack;
+        //Load needed zeros, scales
+        for (int x = tid; x < loading_count; x += ThreadBlock) {
+            uint8_t temp_zeros = 0;
+            temp_zeros = ((uint8_t*)(zeros + quant_group * zeros_stride + m_offset_zeros))[x];
+            half temp_scales[2];
+            *((float*)temp_scales) = ((float*)(scales + quant_group * scales_stride + m_offset_scales))[x];
+            uint32_t z = temp_zeros;
+            uint32_t z0 = __builtin_mxc_ubfe(z, 0, QBITS) + 1;
+            uint32_t z1 = __builtin_mxc_ubfe(z, 4, QBITS) + 1;
+            float s1 = (float)(temp_scales[0]);
+            float s2 = (float)(temp_scales[1]);
+            //Store to shared memory
+            bsm_zeros_ptr[x*2] = (float)z0 * s1 * -1.0f;
+            bsm_zeros_ptr[x*2+1] = (float)z1 * s2 * -1.0f;
+            bsm_scales_ptr[x*2] = temp_scales[0];
+            bsm_scales_ptr[x*2+1] = temp_scales[1];
+        }
+
+        int loop_index = 0;
+        quant_packed_type A[m_per_thread];
+        //Load A
+        *((float4*)A) = *(float4*)(srcA + (loop_index + i + tidCol * LoopNum) / PACK_RATIO * srcAStride + m_offset_a + m_per_thread * tidRow);
+
+        //Load B and transform to float
+        if (b_perm != nullptr) {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[b_perm[splitKOffset + i + tidCol * LoopNum + x] + y * k_stride];
+                }
+            }
+        } else {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[splitKOffset + i + tidCol * LoopNum + x + y * k_stride];
+                }
+            }
+        }
+
+        __syncthreads();
+
+        //Load zero and scale from bsm
+        int m_index = tidRow * m_per_thread;
+
+        v2f local_scales[m_per_thread];
+        v2f local_scales_d16[m_per_thread];
+        for (int c = 0; c < m_per_thread; c++) {
+            float s = (float)bsm_scales_ptr[m_index + c];
+            float s_d16 = s / 16;
+            local_scales[c] = {s, s};
+            local_scales_d16[c] = {s_d16, s_d16};
+        }
+        v2f local_zeros[m_per_thread];
+        for (int c = 0; c < m_per_thread; c++) local_zeros[c] = {bsm_zeros_ptr[m_index + c],bsm_zeros_ptr[m_index + c]};
+
+        const int shuffled_dequant_index[PACK_RATIO] = {0, 4, 1, 5, 2, 6, 3, 7};
+
+        #pragma unroll 16
+        for (; loop_index < LoopNum; loop_index += PACK_RATIO) {
+            v2f local_b[PACK_RATIO/2*N];
+            for (int y = 0; y < N; y++) {
+                for (int x = 0; x < PACK_RATIO / 2; x++) {
+                    local_b[x+PACK_RATIO/2*y].x = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+y*thread_groups*LoopNum];
+                    local_b[x+PACK_RATIO/2*y].y = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+1+y*thread_groups*LoopNum];
+                }
+            }
+            #pragma unroll m_per_thread
+            for (int c = 0; c < m_per_thread; c++) {
+                uint32_t p0 = A[c] & 0x0f0f0f0f;
+                uint32_t p1 = A[c] & 0xf0f0f0f0;
+                float o1,o2,o3,o4,o5,o6,o7,o8;
+                asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o1):"r"(p0));
+                asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o2):"r"(p0));
+                asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o3):"r"(p0));
+                asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o4):"r"(p0));
+                asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o5):"r"(p1));
+                asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o6):"r"(p1));
+                asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o7):"r"(p1));
+                asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o8):"r"(p1));
+                v2f a0 = {o1, o3};
+                v2f a1 = {o5, o7};
+                v2f a2 = {o2, o4};
+                v2f a3 = {o6, o8};
+
+                a0 = __builtin_mxc_pk_fma_f32(a0, local_scales[c], local_zeros[c]);
+                a1 = __builtin_mxc_pk_fma_f32(a1, local_scales_d16[c], local_zeros[c]);
+                a2 = __builtin_mxc_pk_fma_f32(a2, local_scales[c], local_zeros[c]);
+                a3 = __builtin_mxc_pk_fma_f32(a3, local_scales_d16[c], local_zeros[c]);
+
+                #pragma unroll N
+                for (int y = 0; y < N; y++) {
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a0, local_b[0+PACK_RATIO/2*y], c_splited[y][c]);
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a1, local_b[1+PACK_RATIO/2*y], c_splited[y][c]);
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a2, local_b[2+PACK_RATIO/2*y], c_splited[y][c]);
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a3, local_b[3+PACK_RATIO/2*y], c_splited[y][c]);
+                }
+            }
+
+            if constexpr(thread_groups != QUANT_GROUP / PACK_RATIO) {
+                if (loop_index + PACK_RATIO < LoopNum) {
+                    *((float4*)A) = *(float4*)(srcA + (loop_index + PACK_RATIO + i + tidCol * LoopNum) / PACK_RATIO * srcAStride + m_offset_a + m_per_thread * tidRow);
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    // directly do atomic add, may cause partial write?
+    for (int y = 0; y < N; y++) {
+        for (int c = 0; c < m_per_thread; c++) {
+            atomicAdd(dst + tidRow * m_per_thread + c + y * dstStride, (half)(c_splited[y][c].x + c_splited[y][c].y));
+        }
+    }
+}
+
+
+template<>
+__global__ __launch_bounds__(256) void hgemv_nn_splitk_gptq_tb256_bx256_kb128<4>(const half* __restrict__ srcB,
+                                                        const quant_packed_type* __restrict__ srcA,
+                                                        const quant_packed_type* __restrict__ zeros,
+                                                        const half* __restrict__ scales,
+                                                        half *dst,
+                                                        int m,
+                                                        int n,
+                                                        int k,
+                                                        int srcAStride,
+                                                        int dstStride,
+                                                        int k_div_sk,
+                                                        const int* __restrict__ b_perm) {
+    constexpr int QBITS = 4;
+    constexpr int PACK_RATIO = (32 / QBITS);
+    constexpr int QUANT_GROUP = 128;
+    constexpr int ThreadBlock = 256;
+    constexpr int BlockDimX = 256;
+    constexpr int LoopNum = QUANT_GROUP / (ThreadBlock / BlockDimX);
+    constexpr int N = 4;
+    const int k_stride = k;
+    constexpr int k_block = QUANT_GROUP;
+    const int splitKOffset = blockIdx.y * k_block;
+    k = k_block;
+    srcA += splitKOffset * srcAStride / PACK_RATIO;
+    //srcB += splitKOffset;
+
+    constexpr int quant_groups = 1;
+    constexpr int m_per_thread = sizeof(float4)/sizeof(quant_packed_type);
+    constexpr int thread_groups = ThreadBlock / BlockDimX;
+    constexpr int group_elements = BlockDimX * m_per_thread;
+    constexpr int reduce_size = ThreadBlock * m_per_thread;
+    constexpr int data_cache_size = LoopNum * thread_groups * N + group_elements + group_elements / 2;
+    float *bsm_b_ptr, *bsm_zeros_ptr;
+    half *bsm_scales_ptr;
+
+    __shared__ float bsm_ptr[2048];  //128*N+256*4+256*4/2 = 最大8K，每个block 256线程，占据半个PEU，一个AP可以跑满8个block
+    bsm_b_ptr = bsm_ptr;
+    bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+    bsm_scales_ptr = (half*)(bsm_zeros_ptr + group_elements);
+    //smem = bsm_ptr;
+
+    const int zeros_stride = srcAStride / PACK_RATIO;
+    const int scales_stride = srcAStride;
+    zeros += splitKOffset * zeros_stride / QUANT_GROUP;
+    scales += splitKOffset * scales_stride / QUANT_GROUP;
+
+    dst += group_elements * blockIdx.x;
+    const int m_offset_a = group_elements * blockIdx.x;
+    const int m_offset_zeros = group_elements / PACK_RATIO * blockIdx.x;
+    const int m_offset_scales = group_elements * blockIdx.x;
+
+    //store splited fma results
+    v2f c_splited[N][m_per_thread];
+    for (int i = 0; i < N; i++) {
+        for (int j = 0; j < m_per_thread; j++) c_splited[i][j] = {0, 0};
+    }
+
+    int tid = threadIdx.x;
+    int tidCol = tid / BlockDimX;
+    int tidRow = tid % BlockDimX;
+
+    for (int i = 0; i < k; i += LoopNum * thread_groups) {
+        int quant_group = i / QUANT_GROUP;
+        constexpr int loading_pack = 2;
+        constexpr int loading_count = group_elements / loading_pack;
+        //Load needed zeros, scales
+        for (int x = tid; x < loading_count; x += ThreadBlock) {
+            uint8_t temp_zeros = 0;
+            temp_zeros = ((uint8_t*)(zeros + quant_group * zeros_stride + m_offset_zeros))[x];
+            half temp_scales[2];
+            *((float*)temp_scales) = ((float*)(scales + quant_group * scales_stride + m_offset_scales))[x];
+            uint32_t z = temp_zeros;
+            uint32_t z0 = __builtin_mxc_ubfe(z, 0, QBITS) + 1;
+            uint32_t z1 = __builtin_mxc_ubfe(z, 4, QBITS) + 1;
+            float s1 = (float)(temp_scales[0]);
+            float s2 = (float)(temp_scales[1]);
+            //Store to shared memory
+            bsm_zeros_ptr[x*2] = (float)z0 * s1 * -1.0f;
+            bsm_zeros_ptr[x*2+1] = (float)z1 * s2 * -1.0f;
+            bsm_scales_ptr[x*2] = temp_scales[0];
+            bsm_scales_ptr[x*2+1] = temp_scales[1];
+        }
+
+        int loop_index = 0;
+        quant_packed_type A[m_per_thread];
+        //Load A
+        *((float4*)A) = *(float4*)(srcA + (loop_index + i + tidCol * LoopNum) / PACK_RATIO * srcAStride + m_offset_a + m_per_thread * tidRow);
+
+        //Load B and transform to float
+        if (b_perm != nullptr) {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[b_perm[splitKOffset + i + tidCol * LoopNum + x] + y * k_stride];
+                }
+            }
+        } else {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[splitKOffset + i + tidCol * LoopNum + x + y * k_stride];
+                }
+            }
+        }
+
+        __syncthreads();
+
+        //Load zero and scale from bsm
+        int m_index = tidRow * m_per_thread;
+
+        v2f local_scales[m_per_thread];
+        for (int c = 0; c < m_per_thread; c++) {
+            float s = (float)bsm_scales_ptr[m_index + c];
+            local_scales[c] = {s,s};
+        }
+        v2f local_zeros[m_per_thread];
+        for (int c = 0; c < m_per_thread; c++) local_zeros[c] = {bsm_zeros_ptr[m_index + c], bsm_zeros_ptr[m_index + c]};
+
+        //const int shuffled_dequant_index[PACK_RATIO] = {0, 4, 1, 5, 2, 6, 3, 7};
+
+        #pragma unroll 16
+        for (; loop_index < LoopNum; loop_index += PACK_RATIO) {
+            //Split dequant and w*b into 4 parts so we can reduce registers usage from 114 to 76, and each peu will run 6 waves
+            v2f local_b[N];
+            uint32_t Aq[m_per_thread];
+            for (int y = 0; y < N; y++) {
+                int x = 0;
+                local_b[y].x = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+y*thread_groups*LoopNum];
+                local_b[y].y = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+1+y*thread_groups*LoopNum];
+            }
+            #pragma unroll m_per_thread
+            for (int c = 0; c < m_per_thread; c++) {
+                Aq[c] = A[c] & 0x0f0f0f0f;
+                float o1,o3;
+                asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o1):"r"(Aq[c]));
+                asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o3):"r"(Aq[c]));
+                v2f a0 = {o1, o3};
+
+                a0 = __builtin_mxc_pk_fma_f32(a0, local_scales[c], local_zeros[c]);
+
+                #pragma unroll N
+                for (int y = 0; y < N; y++) {
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a0, local_b[y], c_splited[y][c]);
+                }
+            }
+
+            for (int y = 0; y < N; y++) {
+                int x = 2;
+                local_b[y].x = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+y*thread_groups*LoopNum];
+                local_b[y].y = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+1+y*thread_groups*LoopNum];
+            }
+
+            #pragma unroll m_per_thread
+            for (int c = 0; c < m_per_thread; c++) {
+                float o2,o4;
+                asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o2):"r"(Aq[c]));
+                asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o4):"r"(Aq[c]));
+                v2f a2 = {o2, o4};
+
+                a2 = __builtin_mxc_pk_fma_f32(a2, local_scales[c], local_zeros[c]);
+
+                #pragma unroll N
+                for (int y = 0; y < N; y++) {
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a2, local_b[y], c_splited[y][c]);
+                }
+            }
+
+            for (int y = 0; y < N; y++) {
+                int x = 1;
+                local_b[y].x = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+y*thread_groups*LoopNum];
+                local_b[y].y = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+1+y*thread_groups*LoopNum];
+            }
+
+            for (int c = 0; c < m_per_thread; c++) {
+                Aq[c] = (A[c] >> 4) & 0x0f0f0f0f;
+                float o5,o7;
+                asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o5):"r"(Aq[c]));
+                asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o7):"r"(Aq[c]));
+                v2f a1 = {o5, o7};
+                a1 = __builtin_mxc_pk_fma_f32(a1, local_scales[c], local_zeros[c]);
+
+                #pragma unroll N
+                for (int y = 0; y < N; y++) {
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a1, local_b[y], c_splited[y][c]);
+                }
+            }
+
+            for (int y = 0; y < N; y++) {
+                int x = 3;
+                local_b[y].x = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+y*thread_groups*LoopNum];
+                local_b[y].y = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+1+y*thread_groups*LoopNum];
+            }
+
+            for (int c = 0; c < m_per_thread; c++) {
+                float o6,o8;
+                asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o6):"r"(Aq[c]));
+                asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o8):"r"(Aq[c]));
+                v2f a3 = {o6, o8};
+                a3 = __builtin_mxc_pk_fma_f32(a3, local_scales[c], local_zeros[c]);
+
+                #pragma unroll N
+                for (int y = 0; y < N; y++) {
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a3, local_b[y], c_splited[y][c]);
+                }
+            }
+
+            if constexpr(thread_groups != QUANT_GROUP / PACK_RATIO) {
+                if (loop_index + PACK_RATIO < LoopNum) {
+                    *((float4*)A) = *(float4*)(srcA + (loop_index + PACK_RATIO + i + tidCol * LoopNum) / PACK_RATIO * srcAStride + m_offset_a + m_per_thread * tidRow);
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    // directly do atomic add, may cause partial write?
+    for (int y = 0; y < N; y++) {
+        for (int c = 0; c < m_per_thread; c++) {
+            atomicAdd(dst + tidRow * m_per_thread + c + y * dstStride, (half)(c_splited[y][c].x + c_splited[y][c].y));
+        }
+    }
+}
diff --git a/csrc/quantization/gptq/hgemv_nn_splitk_gptq_int8.hpp b/csrc/quantization/gptq/hgemv_nn_splitk_gptq_int8.hpp
new file mode 100644
index 000000000..2a178d362
--- /dev/null
+++ b/csrc/quantization/gptq/hgemv_nn_splitk_gptq_int8.hpp
@@ -0,0 +1,215 @@
+#pragma once
+
+#include <mc_runtime.h>
+#include <maca_fp16.h>
+#include "Hgemm_common.cuh"
+#include "dequant.cuh"
+#define quant_packed_type uint32_t
+typedef __NATIVE_VECTOR__(2, float) v2f;
+template <int ThreadBlock, int BlockDimX, int BATCH>
+__global__ __launch_bounds__(256) void hgemv_nn_splitk_gptq_int8(const half* __restrict__ srcB,
+                                                        const quant_packed_type* __restrict__ srcA,
+                                                        const quant_packed_type* __restrict__ zeros,
+                                                        const half* __restrict__ scales,
+                                                        half *dst,
+                                                        int m,
+                                                        int n,
+                                                        int k,
+                                                        int srcAStride,
+                                                        int dstStride,
+                                                        int k_div_sk,
+                                                        const int* __restrict__ b_perm = nullptr) {
+    constexpr int QBITS = 8;
+    constexpr int PACK_RATIO = (32 / QBITS);
+    constexpr int QUANT_GROUP = 128;
+    constexpr int LoopNum = QUANT_GROUP / (ThreadBlock / BlockDimX);
+    constexpr int N = BATCH;
+    const int k_stride = k;
+    const int k_block = (k_div_sk + QUANT_GROUP - 1) / QUANT_GROUP * QUANT_GROUP;
+    const int splitKOffset = blockIdx.y * k_block;
+    if (splitKOffset + k_block > k) k = k - splitKOffset;
+    else k = k_block;
+    srcA += splitKOffset * srcAStride / PACK_RATIO;
+    //srcB += splitKOffset;
+
+    constexpr int quant_groups = 1;
+    constexpr int m_per_thread = sizeof(float4)/sizeof(quant_packed_type);
+    constexpr int thread_groups = ThreadBlock / BlockDimX;
+    constexpr int group_elements = BlockDimX * m_per_thread;
+    constexpr int reduce_size = ThreadBlock * m_per_thread;
+    constexpr int data_cache_size = LoopNum * thread_groups * N + group_elements * 2;
+    float *bsm_b_ptr, *bsm_zeros_ptr, *bsm_scales_ptr, *smem;
+    if constexpr(reduce_size > data_cache_size) {
+        __shared__ float bsm_ptr[reduce_size];
+        bsm_b_ptr = bsm_ptr;
+        bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+        bsm_scales_ptr = bsm_zeros_ptr + group_elements;
+        smem = bsm_ptr;
+    } else {
+        __shared__ float bsm_ptr[data_cache_size];
+        bsm_b_ptr = bsm_ptr;
+        bsm_zeros_ptr = bsm_ptr + LoopNum * thread_groups * N;
+        bsm_scales_ptr = bsm_zeros_ptr + group_elements;
+        smem = bsm_ptr;
+    }
+
+    const int zeros_stride = srcAStride / PACK_RATIO;
+    const int scales_stride = srcAStride;
+    zeros += splitKOffset * zeros_stride / QUANT_GROUP;
+    scales += splitKOffset * scales_stride / QUANT_GROUP;
+
+    dst += group_elements * blockIdx.x;
+    const int m_offset_a = group_elements * blockIdx.x;
+    const int m_offset_zeros = group_elements / PACK_RATIO * blockIdx.x;
+    const int m_offset_scales = group_elements * blockIdx.x;
+
+    //store splited fma results
+    v2f c_splited[N][m_per_thread];
+    for (int i = 0; i < N; i++) {
+        for (int j = 0; j < m_per_thread; j++) c_splited[i][j] = {0, 0};
+    }
+    
+    int tid = threadIdx.x;
+    int tidCol = tid / BlockDimX;
+    int tidRow = tid % BlockDimX;
+
+    for (int i = 0; i < k; i += LoopNum * thread_groups) {
+        int quant_group = i / QUANT_GROUP;
+        constexpr int loading_pack = 1;
+        constexpr int loading_count = group_elements / loading_pack;
+        //Load needed zeros, scales
+        for (int x = tid; x < loading_count; x += ThreadBlock) {
+            uint8_t temp_zeros = 0;
+            temp_zeros = ((uint8_t*)(zeros + quant_group * zeros_stride + m_offset_zeros))[x];
+            half temp_scales;
+            temp_scales = (scales + quant_group * scales_stride + m_offset_scales)[x];
+            uint32_t z = temp_zeros;
+            float z0;
+            asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(z0):"r"(z));
+            float s1 = (float)(temp_scales);
+            //Store to shared memory
+            //bsm_zeros_ptr[x] = (float)z0 * s1 * -1.0f;    //modify 11.19
+	    bsm_zeros_ptr[x] = (float)(z0+1) * s1 * -1.0f;
+            bsm_scales_ptr[x] = s1;
+            // if (i == 0 && blockIdx.x == 0) {
+            //     printf("tid %d, x = %d, temp_zero=%u, temp_scale=%f,  z0 = %f, s1 = %f\n", tid, x, (uint32_t)temp_zeros, (float)temp_scales, z0, s1);
+            // }
+        }
+
+        int loop_index = 0;
+        quant_packed_type A[m_per_thread];
+        //Load A
+        *((float4*)A) = *(float4*)(srcA + (loop_index + i + tidCol * LoopNum) / PACK_RATIO * srcAStride + m_offset_a + m_per_thread * tidRow);
+
+        //Load B and transform to float
+        if (b_perm != nullptr) {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[b_perm[splitKOffset + i + tidCol * LoopNum + x] + y * k_stride];
+                }
+            }
+        } else {
+            for (int y = 0; y < N; y++) {
+                for (int x = tidRow; x < LoopNum; x += BlockDimX) {
+                    bsm_b_ptr[x + tidCol * LoopNum + y * LoopNum * thread_groups] = srcB[splitKOffset + i + tidCol * LoopNum + x + y * k_stride];
+                }
+            }
+        }
+
+        __syncthreads();
+
+        //Load zero and scale from bsm
+        int m_index = tidRow * m_per_thread;
+
+        v2f local_scales[m_per_thread];
+        for (int c = 0; c < m_per_thread; c++) {
+            float s = bsm_scales_ptr[m_index + c];
+            local_scales[c] = {s, s};
+            // if (tid == 0 && c == 0 && blockIdx.x == 0 && blockIdx.y == 0) {
+            //     printf("get scale %f from index %d\n", s, m_index+c);
+            // }
+        }
+        v2f local_zeros[m_per_thread];
+        for (int c = 0; c < m_per_thread; c++) {
+            local_zeros[c] = {bsm_zeros_ptr[m_index + c],bsm_zeros_ptr[m_index + c]};
+            // if (tid == 0 && c == 0 && blockIdx.x == 0 && blockIdx.y == 0) {
+            //     printf("get zero %f from index %d\n", local_zeros[c].x, m_index+c);
+            // }
+        }
+
+        for (; loop_index < LoopNum; loop_index += PACK_RATIO) {
+            v2f local_b[PACK_RATIO/2*N];
+            for (int y = 0; y < N; y++) {
+                for (int x = 0; x < PACK_RATIO / 2; x++) {
+                    local_b[x+PACK_RATIO/2*y].x = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+y*thread_groups*LoopNum];
+                    local_b[x+PACK_RATIO/2*y].y = bsm_b_ptr[tidCol*LoopNum+loop_index+x*2+1+y*thread_groups*LoopNum];
+                }
+            }
+            #pragma unroll m_per_thread
+            for (int c = 0; c < m_per_thread; c++) {
+                float o1,o2,o3,o4;
+                asm volatile("cvt_b0tof32 %0,%1;\n":"=r"(o1):"r"(A[c]));
+                asm volatile("cvt_b1tof32 %0,%1;\n":"=r"(o2):"r"(A[c]));
+                asm volatile("cvt_b2tof32 %0,%1;\n":"=r"(o3):"r"(A[c]));
+                asm volatile("cvt_b3tof32 %0,%1;\n":"=r"(o4):"r"(A[c]));
+                v2f a0 = {o1, o2};
+                v2f a1 = {o3, o4};
+
+                // if (tid == 0 && c == 0 && blockIdx.x == 0 && blockIdx.y == 0) {
+                //     printf("GPU int8 a0=%f,%f, a1=%f,%f,scale=%f,zero=%f\n",
+                //     a0.x, a0.y, a1.x, a1.y, local_scales[c].x, local_zeros[c].x);
+                // }
+
+                a0 = __builtin_mxc_pk_fma_f32(a0, local_scales[c], local_zeros[c]);
+                a1 = __builtin_mxc_pk_fma_f32(a1, local_scales[c], local_zeros[c]);
+
+                // if (tid == 0 && c == 0 && blockIdx.x == 0 && blockIdx.y == 0) {
+                //     printf("GPU a=%x, a0=%f,%f,a1=%f,%f,b=%f,%f,%f,%f\n",
+                //         A[c], a0.x, a0.y, a1.x, a1.y, local_b[0].x, local_b[0].y, local_b[1].x, local_b[1].y
+                //     );
+                // }
+
+                #pragma unroll N
+                for (int y = 0; y < N; y++) {
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a0, local_b[0+PACK_RATIO/2*y], c_splited[y][c]);
+                    c_splited[y][c] = __builtin_mxc_pk_fma_f32(a1, local_b[1+PACK_RATIO/2*y], c_splited[y][c]);
+                }
+            }
+
+            if constexpr(thread_groups != QUANT_GROUP / PACK_RATIO) {
+                if (loop_index + PACK_RATIO < LoopNum) {
+                    *((float4*)A) = *(float4*)(srcA + (loop_index + PACK_RATIO + i + tidCol * LoopNum) / PACK_RATIO * srcAStride + m_offset_a + m_per_thread * tidRow);
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    #pragma unroll N
+    for (int y = 0; y < N; y++) {
+        for (int i = 0; i < m_per_thread; i++) {
+            smem[tidCol + (tidRow * m_per_thread + i) * ThreadBlock / BlockDimX] = c_splited[y][i].x + c_splited[y][i].y;
+        }
+        __syncthreads();
+        constexpr int stride = ThreadBlock / BlockDimX;
+        int data_size = ThreadBlock * m_per_thread;
+        #pragma unroll
+        for (int i = ThreadBlock / BlockDimX / 2; i > 0; i /= 2) {
+            for (int j = tid; j < data_size / 2; j += ThreadBlock) {
+                int reduce_group = j / i;
+                int reduce_index = j % i;
+                smem[reduce_index + reduce_group * stride] += smem[reduce_index + reduce_group * stride + i];
+            }
+            __syncthreads();
+            data_size /= 2;
+        }
+        for (int i = tid; i < m_per_thread * BlockDimX; i += ThreadBlock) {
+            atomicAdd(dst + i + y * dstStride, (half)smem[i*stride]);
+        }
+        if constexpr(N > 1) {
+            if (y + 1 < N) {
+                __syncthreads();
+            }
+        }
+    }
+}
diff --git a/csrc/quantization/gptq/hgemv_selector.hpp b/csrc/quantization/gptq/hgemv_selector.hpp
new file mode 100644
index 000000000..b9d12a4e1
--- /dev/null
+++ b/csrc/quantization/gptq/hgemv_selector.hpp
@@ -0,0 +1,287 @@
+#include <memory>
+#include <vector>
+#include <algorithm>
+#include "mc_runtime.h"
+#include "maca_fp16.h"
+
+struct KernelEventRecorder {
+    mcEvent_t _start;
+    mcEvent_t _stop;
+    float _eventMs = -1.f;
+
+    KernelEventRecorder() {
+        mcEventCreate(&_start);
+        mcEventCreate(&_stop);
+    }
+
+    ~KernelEventRecorder() {
+        mcEventDestroy(_start);
+        mcEventDestroy(_stop);
+    }
+
+    void start() {
+        mcEventRecord(_start, NULL);
+    }
+
+    float stop() {
+        mcEventRecord(_stop, NULL);
+        mcEventSynchronize(_stop);
+        mcEventElapsedTime(&_eventMs, _start, _stop);
+        return _eventMs;
+    }
+};
+namespace hgemv_selector {
+template<int quant_group=128, int pack_ratio=8, int m_per_thread=4>
+struct GemvParamAutoSelector {
+    int m;
+    int k;
+    int best_block_x = 0;
+    int best_split_k = 0;
+
+    std::pair<int,int> block_x_range;
+    std::pair<int,int> split_k_range;
+    bool _valid = false;
+
+private:
+    std::vector<std::pair<int, int>> param_candidates;
+    int warmup_iters = 0;
+    int current_block_x = 0;
+    int current_split_k = 0;
+    int current_perf_iter = 0;
+    std::vector<float> perf_times;
+    float kernel_best_time_ms_ave = 99999999.0f;
+    float best_band_width;
+    float data_size_gb;
+    std::shared_ptr<KernelEventRecorder> _r;
+    bool _selected = false;
+    const static int MAX_PERF_COUNT = 20;
+
+public:
+    GemvParamAutoSelector(int _m, int _n, int _k, bool allow_imcomplete_bx = false) : m(_m), k(_k)
+    {
+        if (!prepare_parameters(_m, _n, _k, block_x_range.second, split_k_range.second, allow_imcomplete_bx)) return;
+        block_x_range.first = block_x_range.second;
+        split_k_range.first = split_k_range.second;
+        for (int i = 0; i < 3; i++) if (block_x_range.first / 2 >= 16) block_x_range.first /= 2;
+        for (int i = 0; i < 3; i++) if (split_k_range.first % 2 == 0 && split_k_range.first / 2 > 0) split_k_range.first /= 2;
+        for (int i = block_x_range.first; i <= block_x_range.second; i *= 2) {
+            for (int j = split_k_range.first; j <= split_k_range.second; j *= 2) {
+                param_candidates.emplace_back(i, j);
+            }
+        }
+        if (split_k_range.second * quant_group != k) {
+            int max_split_k = k / quant_group;
+            if (max_split_k < 256) {
+                for (int i = block_x_range.first; i <= block_x_range.second; i *= 2) {
+                    param_candidates.emplace_back(i, max_split_k);
+                }
+            }
+        }
+
+        current_block_x = block_x_range.second;
+        current_split_k = split_k_range.second;
+        uint64_t data_size_in_bytes = _n * _k * sizeof(half);
+        data_size_in_bytes += (uint64_t)_m * _k / pack_ratio * sizeof(uint32_t);
+        data_size_in_bytes += (uint64_t)_m * _k / quant_group * sizeof(half);
+        data_size_in_bytes += (uint64_t)_m * _k / quant_group / pack_ratio * sizeof(uint32_t);
+        data_size_in_bytes += (uint64_t)_n * _m * sizeof(half);
+        data_size_gb = float(data_size_in_bytes) / 1024 / 1024 / 1024;
+        warmup_iters = 4;
+        _valid = true;
+    }
+
+    void select_in_warmup(const std::function<bool(int,int)> &f) {
+        if (!valid()) {
+            printf("Cannot run this selector!\n");
+            return;
+        }
+        if (_selected) {
+            f(best_block_x, best_split_k);
+            return;
+        };
+        //Warmup
+        for (auto iter = param_candidates.begin(); iter != param_candidates.end(); iter++) {
+            for (int i = 0; i < 5; i++) {
+                auto &p = *iter;
+                f(p.first, p.second);
+            }
+        }
+        _r.reset(new KernelEventRecorder());
+        kernel_best_time_ms_ave = 9999999.0f;
+        mcDeviceSynchronize();
+
+        for (auto iter = param_candidates.begin(); iter != param_candidates.end(); iter++) {
+            auto &p = *iter;
+            auto &bx = p.first;
+            auto &sk = p.second;
+            mcDeviceSynchronize();
+            _r->start();
+            bool launched = false;
+            for (int i = 0; i < MAX_PERF_COUNT; i++) {
+                launched = f(bx, sk);
+            }
+            auto ms = _r->stop();
+            if (launched && ms / MAX_PERF_COUNT < kernel_best_time_ms_ave) {
+                best_block_x = bx;
+                best_split_k = sk;
+                kernel_best_time_ms_ave = ms / MAX_PERF_COUNT;
+            }
+        }
+
+        _r.reset();
+        _selected = true;
+        warmup_iters = 0;
+    }
+
+    void run(const std::function<bool(int,int)> &f) {
+        if (!valid()) {
+            printf("Cannot run this selector!\n");
+            return;
+        }
+        if (warmup_iters > 0) {
+            f(current_block_x, current_split_k);
+            //printf("Calling warmup current_block_x = %d, current_split_k = %d\n", current_block_x, current_split_k);
+            if (current_block_x > block_x_range.first) current_block_x /= 2;
+            else if (current_split_k > split_k_range.first) current_split_k /= 2;
+            if (current_block_x == block_x_range.first && current_split_k >= split_k_range.first) {
+                current_block_x = block_x_range.second;
+                current_split_k = split_k_range.second;
+            }
+            warmup_iters--;
+            mcDeviceSynchronize();
+            return;
+        }
+
+        if (_selected) {
+            f(best_block_x, best_split_k);
+        } else {
+            if (!_r) {
+                _r = std::shared_ptr<KernelEventRecorder>(new KernelEventRecorder());
+                current_block_x = block_x_range.second;
+                current_split_k = split_k_range.second;
+                current_perf_iter = MAX_PERF_COUNT;
+            }
+            _r->start();
+            auto launched = f(current_block_x, current_split_k);
+            auto ms = _r->stop();
+            //printf("Launched with bx=%d,sk=%d,iter=%d,ms=%fms\n", current_block_x, current_split_k, current_perf_iter, ms);
+            if (!launched) {
+                printf("ERROR: Cannot run gemv with bx = %d, sk = %d\n", current_block_x, current_split_k);
+                return;
+            }
+            if (current_perf_iter-- > 0) {
+                perf_times.emplace_back(ms);
+                return;
+            }
+
+            std::sort(perf_times.begin(), perf_times.end());
+            float total_tm = 0;
+            for (size_t i = 1; i < perf_times.size() - 1; i++) {
+                total_tm += perf_times[i];
+            }
+
+            ms = total_tm /= (MAX_PERF_COUNT - 2);
+            perf_times.clear();
+            current_perf_iter = MAX_PERF_COUNT;
+            //printf("get ave time %fms\n", ms);
+
+            if (ms < kernel_best_time_ms_ave) {
+                best_block_x = current_block_x;
+                best_split_k = current_split_k;
+                kernel_best_time_ms_ave = ms;
+            }
+
+            if (current_split_k > split_k_range.first) {
+                current_split_k /= 2;
+            } else if (current_block_x > block_x_range.first){
+                current_split_k = split_k_range.second;
+                current_block_x /= 2;
+            } else if (current_block_x == block_x_range.first && current_split_k == split_k_range.first) {
+                _selected = true;
+                _r.reset();
+            } else {
+                printf("Error: all parameters are tried! current_block_x = %d, current_split_k = %d, block_x_range = %d,%d, split_k_range = %d,%d, best_block_x = %d, best_split_k = %d\n",
+                    current_block_x, current_split_k,
+                    block_x_range.first, block_x_range.second,
+                    split_k_range.first, split_k_range.second,
+                    best_block_x, best_split_k
+                );
+            }
+        }
+    }
+
+    bool valid() const { return _valid; }
+    bool selected() const {return _selected; }
+
+    float gemv_ave_time_us_cost() {
+        return kernel_best_time_ms_ave;
+    }
+
+    float gemv_bandwidth() {
+        return data_size_gb / (kernel_best_time_ms_ave / 1000.0f);
+    }
+
+    static bool prepare_parameters(int m, int n, int k, int &bx, int &sk, bool allow_imcomplete_bx = false) {
+        if (n > 4) return false;
+        if (k % quant_group != 0) return false;
+        if (m < 16 * m_per_thread) return false;
+        int max_split_k = k / quant_group;
+        int proper_splitk = 1;
+        while (proper_splitk*2 <= 256 && proper_splitk*2 <= max_split_k) proper_splitk *= 2;
+
+        int thread_block = 256;  //A thread block of 512 will wait for more warps when doing syncthreads
+
+        int proper_bx = 16;
+        if (m % (proper_bx * m_per_thread) != 0) return false;
+        //proper_bx * m_per_threads corresponds to group_elements, this value should not be larger than 2048 as limit of shared memory
+        while (proper_bx * 2 * m_per_thread && (m % (proper_bx*2*m_per_thread) == 0) && proper_bx * 2 <= 256 && proper_bx * 2 * m_per_thread <= 2048) proper_bx*=2;
+        if (allow_imcomplete_bx) {
+            int may_proper_bx = proper_bx * 2;
+            if (may_proper_bx <= 256 && may_proper_bx * m_per_thread <= 2048 && (m % (may_proper_bx*m_per_thread)) > may_proper_bx * m_per_thread / 4) {
+                proper_bx = may_proper_bx;
+            }
+        }
+
+        bx = proper_bx;
+        sk = proper_splitk;
+        //printf("m=%d,n=%d,k=%d,get bx=%d,sk=%d,allow_imcomplete_bx=%d\n", m, n, k, bx, sk, allow_imcomplete_bx);
+        return true;
+    }
+};
+
+template<int quant_group=128, int pack_ratio=8, int m_per_thread=4>
+class GemvSelectorHolder {
+private:
+    std::vector<GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread>> _selectors;
+    static std::shared_ptr<GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>> _holder;
+    static GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread> _invalid_selector;
+
+public:
+    static GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread>& selector(int m, int n, int k, bool allow_imcomplete_bx = false) {
+        if (!GemvSelectorHolder::_holder) {
+            GemvSelectorHolder::_holder.reset(new GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>());
+        }
+        int bx, sk;
+        if (!GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread>::prepare_parameters(m, n, k, bx, sk, allow_imcomplete_bx)) {
+            return _invalid_selector;
+        }
+        auto iter = std::find_if(_holder->_selectors.begin(), _holder->_selectors.end(),
+            [&](const GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread>& p) -> bool {
+                return p.m == m && p.k == k;
+            });
+        if (iter != _holder->_selectors.end()) return *iter;
+        GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread> sl(m, n, k, allow_imcomplete_bx);
+        if (!sl.valid()) {
+            return _invalid_selector;
+        }
+        _holder->_selectors.emplace_back(sl);
+        return _holder->_selectors.back();
+    }
+};
+
+template<int quant_group, int pack_ratio, int m_per_thread>
+std::shared_ptr<GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>> GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>::_holder;
+
+template<int quant_group, int pack_ratio, int m_per_thread>
+GemvParamAutoSelector<quant_group,pack_ratio,m_per_thread> GemvSelectorHolder<quant_group,pack_ratio,m_per_thread>::_invalid_selector(0,8,0);
+}
\ No newline at end of file
diff --git a/csrc/quantization/gptq/q_gemm.cu b/csrc/quantization/gptq/q_gemm.cu
index 6fad16e19..7a2c5eb6a 100644
--- a/csrc/quantization/gptq/q_gemm.cu
+++ b/csrc/quantization/gptq/q_gemm.cu
@@ -19,6 +19,17 @@ https://github.com/qwopqwop200/GPTQ-for-LLaMa
 #include "qdq_4.cuh"
 #include "qdq_8.cuh"
 
+#ifdef USE_MACA
+#include "hgemm_gptq.h"
+#include "scalar_type.hpp"
+
+#include "hgemv_nn_splitk_gptq.hpp"
+#include "hgemv_nn_splitk_gptq_int8.hpp"
+#include "hgemv_selector.hpp"
+#include "Hgemm_nn_128x32x128_8m1n8k_gptq-4bits.hpp"
+#include "Hgemm_nn_128x32x128_8m1n8k_gptq-8bits.hpp"
+#endif // USE_MACA
+
 namespace vllm {
 namespace gptq {
 
@@ -31,6 +42,7 @@ namespace gptq {
 #define THREADS_X 32
 #define THREADS_Y 32
 #define DIVIDE(x, size) (((x) + (size) - 1) / (size))
+#define QUANT_GROUP 128
 
 #if defined(USE_ROCM)
   #include <hipblas/hipblas.h>
@@ -734,26 +746,309 @@ fp_gemm_half_q_half_gptq_kernel pick_gemm_half_q_half_gptq_kernel(
   return NULL;
 }
 
+template <typename T>
+__global__ void blasMemset(T *data, size_t cnt, T init) {
+    size_t threads = gridDim.x * blockDim.x;
+    size_t itemsPerThread = (cnt + threads - 1) / threads;
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (size_t loop = 0; loop < itemsPerThread && loop * threads + tid < cnt; ++loop) {
+        data[loop * threads + tid] = init;
+    }
+}
+
+template <typename dstT, typename srcT, typename scalarT>
+__global__ void blasMemcpy(dstT *dst, const srcT *src, size_t cnt, scalarT beta) {
+    size_t threads = gridDim.x * blockDim.x;
+    size_t itemsPerThread = (cnt + threads - 1) / threads;
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (size_t loop = 0; loop < itemsPerThread && loop * threads + tid < cnt; ++loop) {
+        dst[loop * threads + tid] =
+            static_cast<double>(beta) * static_cast<double>(src[loop * threads + tid]);
+    }
+}
+
+template <typename reducT, typename outputT, typename scalarT>
+__global__ void blasReduc(outputT *dC_out, outputT *dC_in, reducT *d_acc, int count, int segs, scalarT beta)
+{
+    using accT = float;
+    size_t threads = gridDim.x * blockDim.x;
+    size_t itemsPerThread = (count + threads - 1) / threads;
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (size_t loop = 0; loop < itemsPerThread && (loop * threads + tid) < count; ++loop) {
+        accT acc = static_cast<accT>(beta) * static_cast<accT>(dC_in[loop * threads + tid]);
+        for (size_t SEG=0; SEG < segs; ++SEG)
+        {
+            acc += static_cast<accT>(d_acc[SEG * count + loop * threads + tid]);
+        }
+        dC_out[loop * threads + tid] = static_cast<outputT>(acc);
+    }
+}
+
+template<typename T_ACC, typename T_ACC_PACK, typename T, typename T_PACK>
+__global__ void split_reduce(const T_ACC* src, const int row, const int splitk, T* dest) {
+    constexpr int ELEMS = sizeof(T_ACC_PACK)/sizeof(T_ACC);
+    for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < row*sizeof(T_ACC)/sizeof(T_ACC_PACK); i += blockDim.x*gridDim.x) {
+        T_ACC_PACK p0 = ((T_ACC_PACK*)src)[i];
+        T_ACC p0_a[ELEMS];
+        for (int j = 0; j < ELEMS; j++) p0_a[j] = ((T_ACC*)&p0)[j];
+        for (int k = 1; k < splitk; k++) {
+            p0 = ((T_ACC_PACK*)src)[i + row / ELEMS * k];
+            for (int j = 0; j < ELEMS; j++) {
+                p0_a[j] += ((T_ACC*)&p0)[j];
+            }
+        }
+        T dest_pack[ELEMS];
+        for (int j = 0; j < ELEMS; j++) dest_pack[j] = p0_a[j];
+        ((T_PACK*)dest)[i] = *(T_PACK*)dest_pack;
+    }
+}
+
+template <int tileK, int tileN>
+__global__ void perm_b(half *output, const half *input, const int *idx, int k, int n, int ldb) {
+    int tid = threadIdx.x;
+    int row = blockIdx.x * tileK + tid;
+    if (row < k) {
+        int index = idx[row];
+        int col_offset = blockIdx.y * tileN;
+#pragma unroll 1
+        for (int i = 0; (i < tileN) && ((col_offset + i) < n); ++i) {
+            int col = col_offset + i;
+            output[row + ldb * col] = input[index + ldb * col];
+        }
+    }
+}
+
+#define SWITCH_CASE_BATCH(BlockDimX, SplitK, BATCH) \
+    case BATCH: {                                   \
+        CALL_GEMM(BlockDimX, SplitK, BATCH)         \
+        break;                                      \
+    }
+
+#define APPLY_HGEMM_BATCH(BlockDimX, SplitK, BATCH) \
+    switch(BATCH) {                                 \
+        SWITCH_CASE_BATCH(BlockDimX, SplitK, 1)           \
+        SWITCH_CASE_BATCH(BlockDimX, SplitK, 2)           \
+        SWITCH_CASE_BATCH(BlockDimX, SplitK, 3)           \
+        SWITCH_CASE_BATCH(BlockDimX, SplitK, 4)           \
+        default: {                                          \
+            launched = false;                               \
+            printf("ERROR: Unsupported BATCH %d\n", BATCH); \
+            break;                                          \
+        }                                                   \
+    }
+
+#define SWITCH_CASE_BlockDimX(BlockDimX, SplitK, BATCH) \
+    case BlockDimX: {                                   \
+        APPLY_HGEMM_BATCH(BlockDimX, SplitK, BATCH)    \
+        break;                                          \
+    }
+
+#define APPLY_HGEMM(BlockDimX, SplitK, BATCH)           \
+    switch (BlockDimX) {                                \
+        SWITCH_CASE_BlockDimX(16, SplitK, BATCH)        \
+        SWITCH_CASE_BlockDimX(32, SplitK, BATCH)        \
+        SWITCH_CASE_BlockDimX(64, SplitK, BATCH)        \
+        SWITCH_CASE_BlockDimX(128, SplitK, BATCH)       \
+        SWITCH_CASE_BlockDimX(256, SplitK, BATCH)       \
+        default: {                                                  \
+            launched = false;                                       \
+            printf("ERROR: Unsupported BlockDimX %d\n", BlockDimX); \
+            break;                                                  \
+        }                                                           \
+    }
+
+bool call_kernel(const half *srcB,
+    const quant_packed_type *srcA,
+    quant_packed_type *zeros, half *scales,
+    half* dst_D,
+    int m, int n, int k, int srcStride, int dstStride,
+    int block_x, int split_k, int bit,
+    const int* b_perm_D = nullptr) {
+    constexpr int ThreadBlock = 256;
+    const dim3 threadBlock = {static_cast<unsigned int>(ThreadBlock)};
+    const dim3 gridBlock = {static_cast<unsigned int>(m / (block_x * sizeof(float4) / sizeof(quant_packed_type))), static_cast<unsigned int>(split_k)};
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    if (split_k * QUANT_GROUP > k || k % QUANT_GROUP != 0) return false;
+    if (block_x < 16 || n > 4) return false;
+    bool launched = true;
+    #define CALL_GEMM(BX, SK, N) \
+    if (bit == 4) { \
+        if (QUANT_GROUP*SK == k && BX == 256) { \
+            hgemv_nn_splitk_gptq_tb256_bx256_kb128<N><<<gridBlock, threadBlock, 0, stream>>>( \
+                srcB, srcA, zeros, scales, dst_D, m, n, k, srcStride, dstStride, k/SK, b_perm_D); \
+        } else {    \
+            hgemv_nn_splitk_gptq<256, BX, N><<<gridBlock, threadBlock, 0, stream>>>( \
+                srcB, srcA, zeros, scales, dst_D, m, n, k, srcStride, dstStride, k/SK, b_perm_D);  \
+        } \
+    } \
+    if (bit == 8) { \
+        hgemv_nn_splitk_gptq_int8<256, BX, N><<<gridBlock, threadBlock, 0, stream>>>( \
+            srcB, srcA, zeros, scales, dst_D, m, n, k, srcStride, dstStride, k/SK, b_perm_D); \
+    }
+    APPLY_HGEMM(block_x, split_k, n);
+    return launched;
+}
+
 void gemm_half_q_half_cuda_part(const half* a, const uint32_t* b_q_weight,
                                 const uint32_t* b_gptq_qzeros,
                                 const half* b_gptq_scales, const int* b_q_perm,
                                 half* c, int size_m, int size_n, int size_k,
-                                int m_count, int groups, int bit) {
-  dim3 blockDim, gridDim;
-  blockDim.x = BLOCK_KN_SIZE;
-  blockDim.y = 1;
-  blockDim.z = 1;
-  gridDim.x = DIVIDE(size_n, BLOCK_KN_SIZE * 4);
-  gridDim.y = DIVIDE(size_m, m_count);
-  gridDim.z = DIVIDE(size_k, BLOCK_KN_SIZE);
+                                int m_count, int groups, int bit, bool m_sign, bool v_sign) {
+  if ((bit == 4 || bit == 8) && m_sign && !v_sign){
+        const int threads_n = 256;
+        const int tileM = 128;
+        const int tileN = 32;
+        const int tileK = 128;
+        int lda = size_n;
+        int ldb = size_k;
+        int ldc = size_n;
+
+        int splitk_iters = 3;
+        bool isSplitk = splitk_iters > 1;
+
+        uint32_t gridx = (size_n - 1) / tileM + 1;
+        uint32_t gridy = (size_m - 1) / tileN + 1;
+        uint32_t gridz = splitk_iters;
+
+        uint32_t* zeros = const_cast<uint32_t*>(b_gptq_qzeros);
+        half* scales = const_cast<half*>(b_gptq_scales);
+
+        dim3 dimBlock(threads_n, 1, 1);
+        dim3 dimGrid(gridx, gridy, gridz);
+        float alpha = 1.0, beta = 0.0;
+        const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+        bool isBetaZero = (beta == 0.0);
+        Operation_t trans_a = Operation_t(0);
+        Operation_t trans_b = Operation_t(0);
+
+        if (trans_a == OP_N && trans_b == OP_N && size_n % 8 == 0 && size_k % 8 == 0) {
+             half *dB_perm;
+             if (b_q_perm != nullptr) {
+                 mcMallocAsync((void **)&dB_perm, ldb * size_m * sizeof(input_type), stream);
+                 const int threads_n1 = 128;
+                 const int tileK1 = 128;
+                 const int tileN1 = 8;
+                 uint32_t gridx1 = (size_k - 1) / tileK1 + 1;
+                 uint32_t gridy1 = (size_m - 1) / tileN1 + 1;
+                 dim3 dimBlock1(threads_n1, 1, 1);
+                 dim3 dimGrid1(gridx1, gridy1, 1);
+                 perm_b<tileK1, tileN1><<<dimGrid1, dimBlock1, 0, stream>>>(dB_perm, a, b_q_perm, size_k, size_m, ldb);
+             }
+             const half *dB_actual = (b_q_perm != nullptr ? dB_perm : a);
+	     if (bit == 4) {
+                 if (!isSplitk) {
+                     if (isBetaZero) {
+                         Hgemm_nn_128x32x128_8m1n8k_gptq_4bit<OP_N, OP_N, true, tileM, tileN, tileK, true, false>
+                             <<<dimGrid, dimBlock, 0, stream>>>(size_n, size_m, size_k, alpha, beta, b_q_weight, lda, dB_actual, ldb, c,
+                                                                c, ldc, zeros, scales);
+                     } else {
+                         Hgemm_nn_128x32x128_8m1n8k_gptq_4bit<OP_N, OP_N, false, tileM, tileN, tileK, true, false>
+                             <<<dimGrid, dimBlock, 0, stream>>>(size_n, size_m, size_k, alpha, beta, b_q_weight, lda, dB_actual, ldb, c,
+                                                                c, ldc, zeros, scales);
+                     }
+                 } else {
+                     if (isBetaZero) {
+                         Hgemm_nn_128x32x128_8m1n8k_gptq_4bit<OP_N, OP_N, true, tileM, tileN, tileK, true, true>
+                             <<<dimGrid, dimBlock, 0, stream>>>(size_n, size_m, size_k, alpha, beta, b_q_weight, lda, dB_actual, ldb, c, c,
+                                                                ldc, zeros, scales, splitk_iters, c);
+                     }
+                     else {
+                         acc_type *d_acc;
+                         mcMalloc(reinterpret_cast<void **>(&d_acc), size_n * size_m * sizeof(acc_type));
+                         blasMemcpy<<<104, 512, 0, stream>>>(d_acc, c, size_n * size_m, beta);
+                         Hgemm_nn_128x32x128_8m1n8k_gptq_4bit<OP_N, OP_N, false, tileM, tileN, tileK, true, true>
+                             <<<dimGrid, dimBlock, 0, stream>>>(size_n, size_m, size_k, alpha, beta, b_q_weight, lda, dB_actual, ldb, c, c,
+                                                                ldc, zeros, scales, splitk_iters, d_acc);
+                         blasMemcpy<<<104, 512, 0, stream>>>(c, d_acc, size_n * size_m, 1);
+                         mcFree(d_acc);
+                     }
+                 }
+             }
+	     else if (bit == 8){
+                 if (!isSplitk) {
+                     if (isBetaZero) {
+                         Hgemm_nn_128x32x128_8m1n8k_gptq_8bit<OP_N, OP_N, true, tileM, tileN, tileK, true, false>
+                             <<<dimGrid, dimBlock, 0, stream>>>(size_n, size_m, size_k, alpha, beta, b_q_weight, lda, dB_actual, ldb, c,
+                                                                c, ldc, zeros, scales, 1, nullptr, nullptr);
+                     } else {
+                         Hgemm_nn_128x32x128_8m1n8k_gptq_8bit<OP_N, OP_N, false, tileM, tileN, tileK, true, false>
+                             <<<dimGrid, dimBlock, 0, stream>>>(size_n, size_m, size_k, alpha, beta, b_q_weight, lda, dB_actual, ldb, c,
+                                                                c, ldc, zeros, scales);
+                     }
+                 } else {
+                     if (isBetaZero) {
+                         Hgemm_nn_128x32x128_8m1n8k_gptq_8bit<OP_N, OP_N, true, tileM, tileN, tileK, true, true>
+                             <<<dimGrid, dimBlock, 0, stream>>>(size_n, size_m, size_k, alpha, beta, b_q_weight, lda, dB_actual, ldb, c, c,
+                                                                ldc, zeros, scales, splitk_iters, c);
+                     } else {
+                         acc_type *d_acc;
+                         mcMallocAsync(reinterpret_cast<void **>(&d_acc), size_n * size_m * sizeof(acc_type), stream);
+                         blasMemcpy<<<104, 512, 0, stream>>>(d_acc, c, size_n * size_m, beta);
+                         Hgemm_nn_128x32x128_8m1n8k_gptq_8bit<OP_N, OP_N, false, tileM, tileN, tileK, true, true>
+                             <<<dimGrid, dimBlock, 0, stream>>>(size_n, size_m, size_k, alpha, beta, b_q_weight, lda, dB_actual, ldb, c, c,
+                                                                ldc, zeros, scales, splitk_iters, d_acc);
+                         blasMemcpy<<<104, 512, 0, stream>>>(c, d_acc, size_n * size_m, 1);
+                         mcFreeAsync(d_acc, stream);
+                     }
+                 }
+             }
+             if (b_q_perm != nullptr) {
+                 mcFreeAsync(dB_perm, stream);
+             }
+        } else {
+            printf("Parameters not supported!\n");
+            return;
+        }
+  }
+  else if((bit == 4 || bit == 8) && v_sign){
+         constexpr int m_per_thread = 4;
+         uint32_t* zeros = const_cast<uint32_t*>(b_gptq_qzeros);
+         half* scales = const_cast<half*>(b_gptq_scales);
+         auto kernel_testing = [&](int bx, int sk) -> bool {
+             return call_kernel(a, b_q_weight, zeros, scales, c, size_n, size_m, size_k, size_n, size_n, bx, sk, bit, b_q_perm);
+         };
+         if (bit == 4){
+             auto& sl_warmup = hgemv_selector::GemvSelectorHolder<QUANT_GROUP,8,m_per_thread>::selector(size_n, size_m, size_k);
+             if (sl_warmup.valid()) {
+                 if (!sl_warmup.selected()) {
+                     sl_warmup.select_in_warmup(kernel_testing);
+                     mcMemset(c, 0, size_n * size_m * sizeof(half));
+                     sl_warmup.run(kernel_testing);
+                 } else {
+                     sl_warmup.run(kernel_testing);
+                 }
+             }
+         }
+         else {
+             auto& sl_warmup = hgemv_selector::GemvSelectorHolder<QUANT_GROUP,4,m_per_thread>::selector(size_n, size_m, size_k);
+             if (sl_warmup.valid()) {
+                 if (!sl_warmup.selected()) {
+                     sl_warmup.select_in_warmup(kernel_testing);
+                     mcMemset(c, 0, size_n * size_m * sizeof(half));
+                     sl_warmup.run(kernel_testing);
+                 } else {
+                     sl_warmup.run(kernel_testing);
+                 }
+             }
+         }
+  }
+  else {
+    dim3 blockDim, gridDim;
+    blockDim.x = BLOCK_KN_SIZE;
+    blockDim.y = 1;
+    blockDim.z = 1;
+    gridDim.x = DIVIDE(size_n, BLOCK_KN_SIZE * 4);
+    gridDim.y = DIVIDE(size_m, m_count);
+    gridDim.z = DIVIDE(size_k, BLOCK_KN_SIZE);
 
-  fp_gemm_half_q_half_gptq_kernel kernel =
-      pick_gemm_half_q_half_gptq_kernel(true, m_count, bit);
+    fp_gemm_half_q_half_gptq_kernel kernel =
+        pick_gemm_half_q_half_gptq_kernel(true, m_count, bit);
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  kernel<<<gridDim, blockDim, 0, stream>>>(a, b_q_weight, b_gptq_qzeros,
-                                           b_gptq_scales, c, size_m, size_n,
-                                           size_k, groups, b_q_perm);
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    kernel<<<gridDim, blockDim, 0, stream>>>(a, b_q_weight, b_gptq_qzeros,
+                                            b_gptq_scales, c, size_m, size_n,
+                                            size_k, groups, b_q_perm);
+  }
 }
 
 __global__ void reconstruct_exllama_8bit_kernel(
@@ -1487,56 +1782,339 @@ void reconstruct_gptq(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros,
                                            width, groups, out);
 }
 
+
+template <int tileK, int tileM, typename dtype>
+__global__ void perm_a(dtype *output, const dtype *input, const int *idx, int k, int m, int lda) {
+    int tid = threadIdx.x;
+    int row = blockIdx.x * tileK + tid;
+    int col_st = blockIdx.y * tileM;
+    if (row < k) {
+        int index = idx[row];
+        #pragma unroll tileM
+        for (int i = 0; i < tileM; ++i) {
+            int col = col_st + i;
+            if (col < m) {
+                output[row + lda * col] = input[index + lda * col];
+            }
+        }
+    }
+}
+
+template <typename input_tp, const vllm::ScalarTypeId w_type_id, typename output_tp, typename quant_packed_tp>
+bool launch_gemm_gptq(int m,
+                      int n,
+                      int k,
+                      int quant_group,
+                      const input_tp *dA,
+                      int lda,
+                      const quant_packed_tp *dB,
+                      int ldb,
+                      output_tp *dC,
+		                  float *dC_temp,
+                      int ldc,
+                      quant_packed_tp *d_zeros,
+                      input_tp *d_scales,
+                      const cudaStream_t stream,
+                      int chunks = 1) {
+    using namespace hgemm_marlin_gptq;
+    if(n % 16 != 0) {
+        printf("n %% 16 != 0, n = %d\n", n);
+        return false;
+    }
+    if(k % 32 != 0) {
+        printf("k %% 32 != 0, k = %d\n", k);
+        return false;
+    }
+    //const vllm::ScalarTypeId w_type_id = vllm::kU4B8.id();
+    const int THREADS = 256;
+    int BLOCKS_M = div_ceil(m, SLICE_M);
+    if(BLOCKS_M >= MAX_BLOCKS_M && BLOCKS_M % MAX_BLOCKS_M != 0) {
+        printf("Error: input m is error, m = %d, blocks_m = %d\n", m, BLOCKS_M);
+        return false;
+    }
+    if (BLOCKS_M > MAX_BLOCKS_M) BLOCKS_M = MAX_BLOCKS_M;
+    int BLOCKS_N = 8;
+    //int BLOCKS_K = 4;
+    //It is better let TILE_K = quant_group
+    //But if quant_group is too large, a quant_group can be divided into two parts
+    int BLOCKS_K = quant_group / SLICE_K;
+    if (quant_group > 128) BLOCKS_K = 128 / SLICE_K;
+
+    if (BLOCKS_M == 1 || BLOCKS_M == 2) {
+        BLOCKS_N = 16;
+    }
+    const bool HAS_ACT_ORDER = false;
+    //const bool HAS_ZP = false;
+    const bool HAS_ZP = (w_type_id == vllm::kU4.id()) || (w_type_id == vllm::kU8.id());
+    int *g_idx = nullptr;
+    bool HAS_NK_PRED = true;
+    bool HAS_M_PRED = true;
+    //int TILE_N = BLOCKS_N * SLICE_N;
+    //int TILE_K = BLOCKS_K * SLICE_K;
+    //int TILE_M = BLOCKS_M * SLICE_M;
+    if (n % TILE_N == 0 && k % TILE_K == 0) {
+        HAS_NK_PRED = false;
+    }
+    if (m % TILE_M == 0) {
+        HAS_M_PRED = false;
+    }
+
+#define LAUNCH_GPTQ(threads, bm, bn, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    else if (THREADS == threads && BLOCKS_M == bm && BLOCKS_N == bn \
+        && BLOCKS_K == bk  && HAS_ACT_ORDER == has_act_order \
+        && HAS_ZP == has_zp \
+        && HAS_M_PRED == has_m_pred && HAS_NK_PRED == has_nk_pred) { \
+            launch_gemm_gptq_kernel<input_tp, w_type_id, \
+                    threads, bm, bn, bk, has_act_order, has_zp, has_m_pred, has_nk_pred>( \
+                    (const PackTypeInt4*)dA, \
+                    (const PackTypeInt4*)dB, \
+                    (PackTypeInt4*)dC, \
+                    (PackTypeInt4*)dC_temp, \
+                    (const PackTypeInt4*)d_scales, \
+                    (const PackTypeInt4*)d_zeros, \
+                    nullptr, m, n, k, quant_group, chunks,\
+                    stream); \
+    }
+
+#define LAUNCH_GPTQ_K(bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ(256, 1, 16, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ(256, 2, 16, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ(256, 3, 8, bk, has_act_order, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ(256, 4, 8, bk, has_act_order, has_zp, has_nk_pred, has_m_pred)
+
+#define LAUNCH_GPTQ_ZP(has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ_K(1, false, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ_K(2, false, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ_K(4, false, has_zp, has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ_K(8, false, has_zp, has_nk_pred, has_m_pred)
+
+#define LAUNCH_GPTQ_PRED(has_nk_pred, has_m_pred) \
+    LAUNCH_GPTQ_ZP(false, has_nk_pred, has_m_pred) 
+    //LAUNCH_GPTQ_ZP(true, has_nk_pred, has_m_pred)
+
+    if (false) {
+
+    }
+    LAUNCH_GPTQ_PRED(true, true)
+    LAUNCH_GPTQ_PRED(true, false)
+    LAUNCH_GPTQ_PRED(false, true)
+    LAUNCH_GPTQ_PRED(false, false)
+    else {
+        printf("BLOCKS_M=%d, BLOCKS_N=%d, BLOCKS_k=%d, THREADS=%d, HAS_ACT_ORDER=%d, HAS_ZP=%d, quant_group=%d, HAS_M_PRED=%d, HAS_NK_PRED=%d is not supported\n",
+        BLOCKS_M, BLOCKS_N, BLOCKS_K, THREADS, HAS_ACT_ORDER, HAS_ZP, quant_group, HAS_M_PRED, HAS_NK_PRED);
+        return false;
+    }
+
+    return true;
+}
+
+#ifdef BF16_HIGH_PRECISION
+__global__ void vectorized_elementwise_fp32tobf16(float* input, __maca_bfloat16* output, int N) {
+    uint64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
+    if (tid < N) {
+        // printf("tid = %d, input = %f, output = %f\n", tid, input[tid], (float)(__maca_bfloat16)input[tid]);
+        *(__maca_bfloat16*)(output+tid) = (__maca_bfloat16)input[tid];
+    }
+}
+#else
+__global__ void vectorized_elementwise_fp16tobf16(__maca_bfloat16* input, int N) {
+    uint64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
+    if (tid < N) {
+        input[tid] = (float)(*(half*)(input+tid));
+    }
+}
+#endif
+
+
+template <typename input_tp, const vllm::ScalarTypeId w_type_id, typename output_tp, typename quant_packed_tp>
+bool launch_gemm(int m,
+                int n,
+                int k,
+                int quant_group,
+                const input_tp *dA,
+                int lda,
+                const quant_packed_tp *dB,
+                int ldb,
+                output_tp *dC,
+                float *dC_temp,
+                int ldc,
+                quant_packed_tp *d_zeros,
+                input_tp *d_scales,
+                const int* g_idx,
+                input_tp *perm_space,
+                bool is_gptq = true) {
+    using namespace hgemm_marlin_gptq;
+    //constexpr int max_blocks_m = 4;
+    int total_m_blocks = div_ceil(m, SLICE_M);
+    int chunks = total_m_blocks / MAX_BLOCKS_M;
+    int rest_blocks_m = total_m_blocks % MAX_BLOCKS_M;
+    // printf("m=%d,n=%d,k=%d,lda=%d,ldb=%d,ldc=%d,total_m_blocks=%d,chunks=%d,rest_blocks_m=%d\n",
+    //     m, n, k, lda, ldb, ldc, total_m_blocks, chunks, rest_blocks_m
+    // );
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    //input_tp *dA_perm;
+    if (g_idx != nullptr) {
+        //mcMalloc(reinterpret_cast<void **>(&dA_perm), k * m * sizeof(input_tp));
+        //mcMallocAsync(reinterpret_cast<void **>(&dA_perm), k * m * sizeof(input_tp), stream);
+        const int threads = 256;
+        const int tileK1 = 256;
+        const int tileM1 = 16;
+        uint32_t gridx1 = (k + tileK1 - 1) / tileK1;
+        uint32_t gridy1 = (m + tileM1 - 1) / tileM1;
+        dim3 dimBlock1(threads, 1, 1);
+        dim3 dimGrid1(gridx1, gridy1, 1);
+        perm_a<tileK1, tileM1, input_tp><<<dimGrid1, dimBlock1, 0, stream>>>(perm_space, dA, g_idx, k, m, k);
+    }
+    const input_tp *dA_actual = (g_idx != nullptr ? perm_space : dA);
+    bool ret = true;
+    if (chunks > 0) {
+        int real_m = m > chunks * MAX_BLOCKS_M * SLICE_M ? chunks * MAX_BLOCKS_M * SLICE_M : m;
+        if (is_gptq) {
+            ret = launch_gemm_gptq<input_tp, w_type_id, output_tp, quant_packed_tp>(
+                    real_m, n, k, quant_group, 
+                    dA_actual, lda, 
+                    dB, ldb, 
+                    dC, dC_temp, ldc, 
+                    d_zeros, d_scales, stream, chunks);
+        }
+    }
+    if (rest_blocks_m > 0) {
+        int m_offset = chunks * MAX_BLOCKS_M * SLICE_M;
+        if (is_gptq) {
+            ret = ret && launch_gemm_gptq<input_tp, w_type_id, output_tp, quant_packed_tp>(
+                            m - m_offset, n, k, quant_group, 
+                            dA_actual + lda * m_offset, lda, 
+                            dB, ldb, 
+                            dC + ldc * m_offset, dC_temp, ldc, 
+                            d_zeros, d_scales, stream, 1);
+        }
+    }
+
+//#if 0
+    if constexpr(std::is_same_v<input_tp, __maca_bfloat16>) {
+        uint64_t size = m * n;
+        uint64_t block = 512;
+        uint64_t grid = div_ceil(size, block);
+	    vectorized_elementwise_fp32tobf16<<<grid, block, 0, stream>>>((float*)dC_temp, (input_tp*)dC, size);
+    }
+#if 0
+    #ifdef BF16_HIGH_PRECISION
+	  vectorized_elementwise_fp32tobf16<<<grid, block, 0, stream>>>((float*)dC_temp, (input_tp*)dC, size);
+    #else
+          vectorized_elementwise_fp16tobf16<<<grid, block, 0, stream>>>((input_tp*)dC, size);
+    #endif
+#endif
+
+    return ret;
+}
+
+void gemm_bf16_q_bf16_cuda(const __maca_bfloat16* a,
+		           const uint32_t* b_q_weight,
+			   const uint32_t* b_gptq_qzeros,
+			   const __maca_bfloat16* b_gptq_scales, const int* b_g_idx,
+			   __maca_bfloat16* c, float* temp_space, int size_m, int size_n, int size_k,
+			   int bit, int group_size, __maca_bfloat16* perm_space) {
+  bool opt = ((group_size == 128) || (group_size == 64));
+  using scalar_t = __maca_bfloat16;
+  if ((bit == 4) && opt) {
+	  uint32_t* zeros = const_cast<uint32_t*>(b_gptq_qzeros);
+	  scalar_t* scales = const_cast<scalar_t*>(b_gptq_scales);
+	  launch_gemm<scalar_t, vllm::kU4B8.id(), scalar_t, quant_packed_type>(size_m, size_n, size_k,
+			  group_size, a, size_k, b_q_weight, size_n, c, temp_space, size_n, zeros, scales,
+			  b_g_idx, perm_space, true);
+  } else if ((bit == 8) && opt) {
+          uint32_t* zeros = const_cast<uint32_t*>(b_gptq_qzeros);
+          scalar_t* scales = const_cast<scalar_t*>(b_gptq_scales);
+          launch_gemm<scalar_t, vllm::kU8B128.id(), scalar_t, quant_packed_type>(size_m, size_n, size_k,
+                          group_size, a, size_k, b_q_weight, size_n, c, temp_space, size_n, zeros, scales,
+                          b_g_idx, perm_space, true);
+  } else {
+	  printf("Only supported bit-4 , bit-8 of block_size 128 or 64 now!\n");
+  }
+
+}
 void gemm_half_q_half_cuda(cublasHandle_t cublas_handle, const half* a,
                            const uint32_t* b_q_weight,
                            const uint32_t* b_gptq_qzeros,
                            const half* b_gptq_scales, const int* b_g_idx,
                            half* c, half* temp_dq, int size_m, int size_n,
-                           int size_k, int groups, bool use_exllama, int bit) {
+                           int size_k, int groups, bool use_exllama, int bit,
+                           int group_size, half* perm_space) {
   bool use_reconstruct;
-  if (use_exllama) {
-    use_reconstruct = ((bit == 8 && size_m > MAX_Q_GEMM_ROWS_8BIT) ||
-                       (bit != 8 && size_m > MAX_Q_GEMM_ROWS));
+  bool opt = ((group_size == 128) || (group_size == 64));
+  if ((bit == 4) && opt) {
+          if ((size_m <= 2) && (group_size == 128)) {
+                  gemm_half_q_half_cuda_part(a, b_q_weight, b_gptq_qzeros, b_gptq_scales,
+                                             b_g_idx, c, size_m, size_n, size_k,
+                                             BLOCK_M_SIZE_MAX, groups, bit, true, true);
+          } else {
+                  uint32_t* zeros = const_cast<uint32_t*>(b_gptq_qzeros);
+                  half* scales = const_cast<half*>(b_gptq_scales);
+                  launch_gemm<input_type, vllm::kU4B8.id(), output_type, quant_packed_type>(size_m, size_n, size_k,
+                                  group_size, a, size_k, b_q_weight, size_n, c, nullptr, size_n, zeros, scales,
+                                  b_g_idx, perm_space, true);
+          }
+  } else if ((bit == 8) && opt) {
+          uint32_t* zeros = const_cast<uint32_t*>(b_gptq_qzeros);
+          half* scales = const_cast<half*>(b_gptq_scales);
+          launch_gemm<input_type, vllm::kU8B128.id(), output_type, quant_packed_type>(size_m, size_n, size_k,
+                          group_size, a, size_k, b_q_weight, size_n, c, nullptr, size_n, zeros, scales,
+                          b_g_idx, perm_space, true);
   } else {
-    // The 2/3-bit kernels are somehow slower than dequant + gemm baseline, so
-    // we disabled them for now.
-    use_reconstruct = (bit < 4 || size_m > MAX_ALT_GEMM_ROWS);
-  }
-  if (use_reconstruct) {
-    // Reconstruct FP16 matrix, then cuBLAS
     if (use_exllama) {
-      reconstruct_exllama(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
-                          temp_dq, size_k, size_n, groups, bit);
+      use_reconstruct = ((bit == 8 && size_m > MAX_Q_GEMM_ROWS_8BIT) ||
+                        (bit != 8 && size_m > MAX_Q_GEMM_ROWS));
     } else {
-      reconstruct_gptq(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
-                       temp_dq, size_k, size_n, groups, bit);
+      // The 2/3-bit kernels are somehow slower than dequant + gemm baseline, so
+      // we disabled them for now.
+      use_reconstruct = (bit < 4 || size_m > 0);
     }
+    if (use_reconstruct) {
+      // Reconstruct FP16 matrix, then cuBLAS
+      if (use_exllama) {
+        reconstruct_exllama(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
+                            temp_dq, size_k, size_n, groups, bit);
+      } else {
+        reconstruct_gptq(b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
+                        temp_dq, size_k, size_n, groups, bit);
+      }
 
-    const half alpha = __float2half(1.0f);
-    const half beta = __float2half(0.0f);
-    cublasHgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, size_n, size_m, size_k,
-                &alpha, temp_dq, size_n, a, size_k, &beta, c, size_n);
-  } else if (use_exllama) {
-    // Quantized matmul
-    int max_chunks = size_m / BLOCK_M_SIZE_MAX;
-    int last_chunk = max_chunks * BLOCK_M_SIZE_MAX;
-    int last_chunk_size = size_m - last_chunk;
-
-    if (max_chunks) {
-      gemm_half_q_half_cuda_part(a, b_q_weight, b_gptq_qzeros, b_gptq_scales,
-                                 b_g_idx, c, last_chunk, size_n, size_k,
-                                 BLOCK_M_SIZE_MAX, groups, bit);
-    }
+      const half alpha = __float2half(1.0f);
+      const half beta = __float2half(0.0f);
+      cublasHgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, size_n, size_m, size_k,
+                  &alpha, temp_dq, size_n, a, size_k, &beta, c, size_n);
+    } else if (use_exllama) {
+      // Quantized matmul
+      int max_chunks = size_m / BLOCK_M_SIZE_MAX;
+      int last_chunk = max_chunks * BLOCK_M_SIZE_MAX;
+      int last_chunk_size = size_m - last_chunk;
+
+      bool m_sign;
+            bool v_sign;
+            if (group_size == 128) {
+                    m_sign = size_m <= 50;
+                    v_sign = size_m <= 4;
+            } else {
+                    m_sign = false;
+                    v_sign = false;
+            }
+
+      if (max_chunks) {
+        gemm_half_q_half_cuda_part(a, b_q_weight, b_gptq_qzeros, b_gptq_scales,
+                                  b_g_idx, c, last_chunk, size_n, size_k,
+                                  BLOCK_M_SIZE_MAX, groups, bit, m_sign, v_sign);
+      }
 
-    if (last_chunk_size) {
-      gemm_half_q_half_cuda_part(a + last_chunk * size_k, b_q_weight,
-                                 b_gptq_qzeros, b_gptq_scales, b_g_idx,
-                                 c + last_chunk * size_n, last_chunk_size,
-                                 size_n, size_k, last_chunk_size, groups, bit);
+      if (last_chunk_size) {
+        gemm_half_q_half_cuda_part(a + last_chunk * size_k, b_q_weight,
+                                  b_gptq_qzeros, b_gptq_scales, b_g_idx,
+                                  c + last_chunk * size_n, last_chunk_size,
+                                  size_n, size_k, last_chunk_size, groups, bit,  m_sign, v_sign);
+      }
+    } else {
+      gemm_half_q_half_alt(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
+                          c, size_m, size_n, size_k, bit);
     }
-  } else {
-    gemm_half_q_half_alt(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx,
-                         c, size_m, size_n, size_k, bit);
   }
 }
 
@@ -1823,25 +2401,45 @@ void shuffle_exllama_weight(uint32_t* q_weight, int* q_perm, int height,
 torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
                         torch::Tensor b_gptq_qzeros,
                         torch::Tensor b_gptq_scales, torch::Tensor b_g_idx,
-                        bool use_exllama, int64_t bit) {
+                        bool use_exllama, int64_t bit, int64_t group_size,
+                        torch::Tensor perm_space, torch::Tensor temp_space,
+                        bool dtype_bf16) {
   const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
   auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
-  at::Tensor c = torch::empty({a.size(0), b_q_weight.size(1)}, options);
+  at::Tensor c = torch::zeros({a.size(0), b_q_weight.size(1)}, options);
   at::Tensor temp_dq = torch::empty(
       {b_q_weight.size(0) * 32 / bit, b_q_weight.size(1)}, options);
 
-  vllm::gptq::gemm_half_q_half_cuda(
-      at::cuda::getCurrentCUDABlasHandle(), (const half*)a.data_ptr(),
-      (const uint32_t*)b_q_weight.data_ptr(),
-      (const uint32_t*)b_gptq_qzeros.data_ptr(),
-      (const half*)b_gptq_scales.data_ptr(),
-      b_g_idx.device().is_meta() ? NULL : (const int*)b_g_idx.data_ptr(),
-      (half*)c.data_ptr(), (half*)temp_dq.data_ptr(),
-      c.size(0),              // m
-      c.size(1),              // n
-      a.size(1),              // k
-      b_gptq_qzeros.size(0),  // group number
-      use_exllama, bit);
+  if (dtype_bf16) {
+      vllm::gptq::gemm_bf16_q_bf16_cuda(
+        (const __maca_bfloat16*)a.data_ptr(),
+        (const uint32_t*)b_q_weight.data_ptr(),
+        (const uint32_t*)b_gptq_qzeros.data_ptr(),
+        (const __maca_bfloat16*)b_gptq_scales.data_ptr(),
+        b_g_idx.device().is_meta() ? NULL : (const int*)b_g_idx.data_ptr(),
+        (__maca_bfloat16*)c.data_ptr(),
+	(float*)temp_space.data_ptr(),
+        c.size(0),              // m
+        c.size(1),              // n
+        a.size(1),              // k
+        bit, group_size,
+        (__maca_bfloat16*)perm_space.data_ptr());
+  } else {
+    vllm::gptq::gemm_half_q_half_cuda(
+        at::cuda::getCurrentCUDABlasHandle(), (const half*)a.data_ptr(),
+        (const uint32_t*)b_q_weight.data_ptr(),
+        (const uint32_t*)b_gptq_qzeros.data_ptr(),
+        (const half*)b_gptq_scales.data_ptr(),
+        b_g_idx.device().is_meta() ? NULL : (const int*)b_g_idx.data_ptr(),
+        (half*)c.data_ptr(), (half*)temp_dq.data_ptr(),
+        c.size(0),              // m
+        c.size(1),              // n
+        a.size(1),              // k
+        b_gptq_qzeros.size(0),  // group number
+        use_exllama, bit, group_size,
+        (half*)perm_space.data_ptr());
+  }
+
   return c;
 }
 
diff --git a/csrc/quantization/gptq/qdq_4.cuh b/csrc/quantization/gptq/qdq_4.cuh
index 7f65d2d28..6462584bc 100644
--- a/csrc/quantization/gptq/qdq_4.cuh
+++ b/csrc/quantization/gptq/qdq_4.cuh
@@ -85,7 +85,7 @@ __forceinline__ __device__ void dequant_4bit_8_prep_zero(const uint32_t zero,
   y1y16[0] = __half2half2(y1);
   y1y16[1] = __half2half2(y16);
 }
-
+typedef __NATIVE_VECTOR__(2, float) v2f;
 __forceinline__ __device__ void dequant_4bit_8_gptq(const uint32_t q_0,
                                                     half2 (&dq)[4],
                                                     half2 (&z1z16)[2],
@@ -112,12 +112,50 @@ __forceinline__ __device__ void dequant_4bit_8_gptq(const uint32_t q_0,
     dq[2] = __hfma2(q2.as_half2, y1y16[0], z1z16[0]);
     dq[3] = __hfma2(q3.as_half2, y1y16[1], z1z16[1]);
   } else {
+#ifndef USE_MACA
     dq[0] = __hadd2(q0.as_half2, z1z16[0]);  // half2( q[0] - z, q[1] - z )
     dq[1] = __hfma2(q1.as_half2, y1y16[1],
                     z1z16[1]);               // half2( q[2] - z, q[3] - z )
     dq[2] = __hadd2(q2.as_half2, z1z16[0]);  // half2( q[4] - z, q[5] - z )
     dq[3] = __hfma2(q3.as_half2, y1y16[1],
                     z1z16[1]);  // half2( q[6] - z, q[7] - z )
+#endif // USE_MACA
+#if 1
+    dq[0] = __hadd2(q0.as_half2, z1z16[0]);
+    dq[2] = __hadd2(q2.as_half2, z1z16[0]);
+    v2f q1_float2, q3_float2, y1y16_1_float2, z1z16_1_float2;
+    q1_float2[0] = __half2float(__low2half(q1.as_half2));
+    q1_float2[1] = __half2float(__high2half(q1.as_half2));
+
+    q3_float2[0] = __half2float(__low2half(q3.as_half2));
+    q3_float2[1] = __half2float(__high2half(q3.as_half2));
+
+    y1y16_1_float2[0] = __half2float(__low2half(y1y16[1]));
+    y1y16_1_float2[1] = __half2float(__high2half(y1y16[1]));
+
+    z1z16_1_float2[0] = __half2float(__low2half(z1z16[1]));
+    z1z16_1_float2[1] = __half2float(__high2half(z1z16[1]));
+    v2f result1, result3;
+    result1 = __builtin_mxc_pk_fma_f32(q1_float2, y1y16_1_float2, z1z16_1_float2);
+    result3 = __builtin_mxc_pk_fma_f32(q3_float2, y1y16_1_float2, z1z16_1_float2);
+    dq[1] = __floats2half2_rn(result1[0], result1[1]);
+    dq[3] = __floats2half2_rn(result3[0], result3[1]);
+#else
+    dq[0] = __hadd2(q0.as_half2, z1z16[0]);
+    dq[2] = __hadd2(q2.as_half2, z1z16[0]);
+
+    v2f q1_float2, q3_float2, y1y16_1_float2, z1z16_1_float2;
+    q1_float2 = __builtin_mxc_cvt_pk_f16tof32(*reinterpret_cast<const float*>(&q1.as_half2));
+    q3_float2 = __builtin_mxc_cvt_pk_f16tof32(*reinterpret_cast<const float*>(&q3.as_half2));
+    y1y16_1_float2 = __builtin_mxc_cvt_pk_f16tof32(*reinterpret_cast<const float*>(&y1y16[1]));
+    z1z16_1_float2 = __builtin_mxc_cvt_pk_f16tof32(*reinterpret_cast<const float*>(&z1z16[1]));
+
+    v2f result1 = __builtin_mxc_pk_fma_f32(q1_float2, y1y16_1_float2, z1z16_1_float2);
+    v2f result3 = __builtin_mxc_pk_fma_f32(q3_float2, y1y16_1_float2, z1z16_1_float2);
+
+    dq[1] = __floats2half2_rn(result1[0], result1[1]);
+    dq[3] = __floats2half2_rn(result3[0], result3[1]);
+#endif
   }
 }
 }  // namespace gptq
diff --git a/csrc/quantization/gptq/scalar_type.hpp b/csrc/quantization/gptq/scalar_type.hpp
new file mode 100644
index 000000000..46596ff49
--- /dev/null
+++ b/csrc/quantization/gptq/scalar_type.hpp
@@ -0,0 +1,551 @@
+#pragma once
+
+#include <variant>
+#include <tuple>
+//#include <torch/custom_class.h>
+
+namespace vllm {
+
+//
+//  ScalarType can represent a wide range of floating point and integer types,
+//  in particular it can be used to represent sub-byte data types (something
+//  that torch.dtype currently does not support).
+//
+//  ScalarTypeTorch is a subclass of ScalarType that is compatible with
+//  TORCH_LIBRARY, making it accessible from Python as well meaning this class
+//  can be used as a argument for custom operators, helping to simplify these
+//  interfaces.
+//
+//  The type definitions on the Python side can be found in: vllm/_core_ext.pyi
+//  these type definitions should be kept up to date with any Python API changes
+//  here.
+//
+class ScalarType {
+ public:
+  enum NanRepr : uint8_t {
+    NAN_NONE = 0,                // nans are not supported
+    NAN_IEEE_754 = 1,            // nans are: exp all 1s, mantissa not all 0s
+    NAN_EXTD_RANGE_MAX_MIN = 2,  // nans are: exp all 1s, mantissa all 1s
+
+    NAN_REPR_ID_MAX
+  };
+
+  constexpr ScalarType(uint8_t exponent, uint8_t mantissa, bool signed_,
+                       int32_t bias, bool finite_values_only = false,
+                       NanRepr nan_repr = NAN_IEEE_754)
+      : exponent(exponent),
+        mantissa(mantissa),
+        signed_(signed_),
+        bias(bias),
+        finite_values_only(finite_values_only),
+        nan_repr(nan_repr){};
+
+  static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) {
+    return ScalarType(0, size_bits - 1, true, bias);
+  }
+
+  static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) {
+    return ScalarType(0, size_bits, false, bias);
+  }
+
+  // IEEE 754 compliant floating point type
+  static constexpr ScalarType float_IEEE754(uint8_t exponent,
+                                            uint8_t mantissa) {
+    //TORCH_CHECK(mantissa > 0 && exponent > 0);
+    return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754);
+  }
+
+  // IEEE 754 non-compliant floating point type
+  static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
+                                     bool finite_values_only,
+                                     NanRepr nan_repr) {
+    //TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
+    //TORCH_CHECK(mantissa > 0 && exponent > 0);
+    //TORCH_CHECK(nan_repr != NAN_IEEE_754,
+    //            "use `float_IEEE754` constructor for floating point types that "
+    //            "follow IEEE 754 conventions");
+    return ScalarType(exponent, mantissa, true, 0, finite_values_only,
+                      nan_repr);
+  }
+
+  uint8_t const exponent;  // size of the exponent field (0 for integer types)
+  uint8_t const mantissa;  // size of the mantissa field (size of the integer
+                           // excluding the sign bit for integer types)
+  bool const signed_;  // flag if the type supports negative numbers (i.e. has a
+                       // sign bit)
+  int32_t const bias;  // stored values equal value + bias,
+                       // used for quantized type
+
+  // Extra Floating point info
+  bool const finite_values_only;  // i.e. no +/-inf if true
+  NanRepr const nan_repr;         // how NaNs are represented
+                                  // (not applicable for integer types)
+
+  using Id = int64_t;
+
+ private:
+  // Field size in id
+  template <typename T_>
+  static constexpr size_t member_id_field_width() {
+    using T = std::decay_t<T_>;
+    return std::is_same_v<T, bool> ? 1 : sizeof(T) * 8;
+  }
+
+  template <typename Fn, typename Init, typename Member, typename... Rest>
+  static constexpr auto reduce_members_helper(Fn f, Init val, Member member,
+                                              Rest... rest) {
+    auto new_val = f(val, member);
+    if constexpr (sizeof...(rest) > 0) {
+      return reduce_members_helper(f, new_val, rest...);
+    } else {
+      return new_val;
+    };
+  }
+
+  template <typename Fn, typename Init>
+  constexpr auto reduce_members(Fn f, Init init) const {
+    // Should be in constructor order for `from_id`
+    return reduce_members_helper(f, init, exponent, mantissa, signed_, bias,
+                                 finite_values_only, nan_repr);
+  };
+
+  template <typename Fn, typename Init>
+  static constexpr auto reduce_member_types(Fn f, Init init) {
+    constexpr auto dummy_type = ScalarType(0, 0, false, 0, false, NAN_NONE);
+    return dummy_type.reduce_members(f, init);
+  };
+
+  static constexpr auto id_size_bits() {
+    return reduce_member_types(
+        [](int acc, auto member) -> int {
+          return acc + member_id_field_width<decltype(member)>();
+        },
+        0);
+  }
+
+ public:
+  // unique id for this scalar type that can be computed at compile time for
+  //  c++17 template specialization this is not needed once we migrate to
+  //  c++20 and can pass literal classes as template parameters
+  constexpr Id id() const {
+    static_assert(id_size_bits() <= sizeof(Id) * 8,
+                  "ScalarType id is too large to be stored");
+
+    auto or_and_advance = [](std::pair<Id, uint32_t> result,
+                             auto member) -> std::pair<Id, uint32_t> {
+      auto [id, bit_offset] = result;
+      auto constexpr bits = member_id_field_width<decltype(member)>();
+      return {id | (int64_t(member) & ((uint64_t(1) << bits) - 1))
+                       << bit_offset,
+              bit_offset + bits};
+    };
+    return reduce_members(or_and_advance, std::pair<Id, uint32_t>{}).first;
+  }
+
+  // create a ScalarType from an id, for c++17 template specialization,
+  //  this is not needed once we migrate to c++20 and can pass literal
+  //  classes as template parameters
+  static constexpr ScalarType from_id(Id id) {
+    auto extract_and_advance = [id](auto result, auto member) {
+      using T = decltype(member);
+      auto [tuple, bit_offset] = result;
+      auto constexpr bits = member_id_field_width<T>();
+      auto extracted_val = static_cast<T>((int64_t(id) >> bit_offset) &
+                                          ((uint64_t(1) << bits) - 1));
+      auto new_tuple = std::tuple_cat(tuple, std::make_tuple(extracted_val));
+      return std::pair<decltype(new_tuple), int>{new_tuple, bit_offset + bits};
+    };
+
+    auto [tuple_args, _] = reduce_member_types(extract_and_advance,
+                                               std::pair<std::tuple<>, int>{});
+    return std::apply([](auto... args) { return ScalarType(args...); },
+                      tuple_args);
+  }
+
+  constexpr int64_t size_bits() const {
+    return mantissa + exponent + is_signed();
+  }
+  constexpr bool is_signed() const { return signed_; }
+  constexpr bool is_integer() const { return exponent == 0; }
+  constexpr bool is_floating_point() const { return exponent > 0; }
+  constexpr bool is_ieee_754() const {
+    return is_floating_point() && finite_values_only == false &&
+           nan_repr == NAN_IEEE_754;
+  }
+  constexpr bool has_nans() const {
+    return is_floating_point() && nan_repr != NAN_NONE;
+  }
+  constexpr bool has_infs() const {
+    return is_floating_point() && finite_values_only == false;
+  }
+  constexpr bool has_bias() const { return bias != 0; }
+
+ private:
+  double _floating_point_max() const {
+    //TORCH_CHECK(mantissa <= 52 && exponent <= 11,
+    //            "Cannot represent max/min as a double for type ", str());
+
+    uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
+    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
+      max_mantissa -= 1;
+    }
+
+    uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
+    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
+      //TORCH_CHECK(exponent < 11,
+      //            "Cannot represent max/min as a double for type ", str());
+      max_exponent += 1;
+    }
+
+    // adjust the exponent to match that of a double
+    //  for now we assume the exponent bias is the standard 2^(e-1) -1, (where e
+    //  is the exponent bits), there is some precedent for non-standard biases,
+    //  example `float8_e4m3b11fnuz` here: https://github.com/jax-ml/ml_dtypes
+    //  but to avoid premature over complication we are just assuming the
+    //  standard exponent bias until there is a need to support non-standard
+    //  biases
+    uint64_t exponent_bias = (uint64_t(1) << (exponent - 1)) - 1;
+    uint64_t exponent_bias_double = (uint64_t(1) << 10) - 1;  // double e = 11
+
+    uint64_t max_exponent_double =
+        max_exponent - exponent_bias + exponent_bias_double;
+
+    // shift the mantissa into the position for a double and
+    // the exponent
+    uint64_t double_raw =
+        (max_mantissa << (52 - mantissa)) | (max_exponent_double << 52);
+
+    return *reinterpret_cast<double*>(&double_raw);
+  }
+
+  constexpr std::variant<int64_t, double> _raw_max() const {
+    if (is_floating_point()) {
+      return {_floating_point_max()};
+    } else {
+      //TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
+      //            "Cannot represent max as a int64_t");
+      return {(int64_t(1) << mantissa) - 1};
+    }
+  }
+
+  constexpr std::variant<int64_t, double> _raw_min() const {
+    if (is_floating_point()) {
+      //TORCH_CHECK(is_signed(),
+      //            "We currently assume all floating point types are signed");
+      constexpr uint64_t sign_bit_double = (uint64_t(1) << 63);
+
+      double max = _floating_point_max();
+      uint64_t max_raw = *reinterpret_cast<uint64_t*>(&max);
+      uint64_t min_raw = max_raw | sign_bit_double;
+      return {*reinterpret_cast<double*>(&min_raw)};
+    } else {
+      //TORCH_CHECK(!is_signed() || size_bits() <= 64,
+      //            "Cannot represent min as a int64_t");
+      if (is_signed()) {
+        // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0
+        // then perform an arithmetic shift right to set all the bits above
+        // (size_bits() - 1) to 1
+        return {INT64_MIN >> (64 - size_bits())};
+      } else {
+        return {int64_t(0)};
+      }
+    }
+  }
+
+ public:
+  // Max representable value for this scalar type.
+  // (accounting for bias if there is one)
+  constexpr std::variant<int64_t, double> max() const {
+    return std::visit(
+        [this](auto x) -> std::variant<int64_t, double> { return {x - bias}; },
+        _raw_max());
+  }
+
+  // Min representable value for this scalar type.
+  // (accounting for bias if there is one)
+  constexpr std::variant<int64_t, double> min() const {
+    return std::visit(
+        [this](auto x) -> std::variant<int64_t, double> { return {x - bias}; },
+        _raw_min());
+  }
+
+  std::string str() const {
+    /* naming generally follows: https://github.com/jax-ml/ml_dtypes
+     * for floating point types (leading f) the scheme is:
+     *  `float<size_bits>_e<exponent_bits>m<mantissa_bits>[flags]`
+     *  flags:
+     *  - no-flags: means it follows IEEE 754 conventions
+     *  - f: means finite values only (no infinities)
+     *  - n: means nans are supported (non-standard encoding)
+     * for integer types the scheme is:
+     *  `[u]int<size_bits>[b<bias>]`
+     *  - if bias is not present it means its zero
+     */
+    if (is_floating_point()) {
+      auto ret = "float" + std::to_string(size_bits()) + "_e" +
+                 std::to_string(exponent) + "m" + std::to_string(mantissa);
+      if (!is_ieee_754()) {
+        if (finite_values_only) {
+          ret += "f";
+        }
+        if (nan_repr != NAN_NONE) {
+          ret += "n";
+        }
+      }
+      return ret;
+    } else {
+      auto ret = ((is_signed()) ? "int" : "uint") + std::to_string(size_bits());
+      if (has_bias()) {
+        ret += "b" + std::to_string(bias);
+      }
+      return ret;
+    }
+  }
+
+  constexpr bool operator==(ScalarType const& other) const {
+    return mantissa == other.mantissa && exponent == other.exponent &&
+           bias == other.bias && signed_ == other.signed_ &&
+           finite_values_only == other.finite_values_only &&
+           nan_repr == other.nan_repr;
+  }
+};
+#if 0
+// Create a TORCH_LIBRARY compatible version of ScalarType (i.e. inherit from
+//  torch::CustomClassHolder), we use multiple inheritance here since we cannot
+//  have ScalarType inherit from torch::CustomClassHolder and have a constexpr
+//  constructor at the same time (torch::CustomClassHolder does not have a
+//  constexpr destructor)
+// See also:
+// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
+class ScalarTypeTorch : public torch::CustomClassHolder, public ScalarType {
+ public:
+  ScalarTypeTorch(int64_t exponent, int64_t mantissa, int64_t bias,
+                  bool _signed)
+      : ScalarType(exponent, mantissa, bias, _signed){};
+
+  ScalarTypeTorch(ScalarType type) : ScalarType(type){};
+
+  using Base = ScalarType;
+  using Self = ScalarTypeTorch;
+  using SelfPtr = c10::intrusive_ptr<Self>;
+
+  static void check_size_bits(int64_t size_bits, bool signed_) {
+    TORCH_CHECK(
+        size_bits <=
+            std::numeric_limits<decltype(std::declval<Self>().mantissa)>::max(),
+        "size_bits bit width is too large to be represented");
+  }
+
+  static void check_bias(int64_t bias) {
+    using Bias = decltype(std::declval<Self>().bias);
+    TORCH_CHECK(bias <= std::numeric_limits<Bias>::max() &&
+                    bias >= std::numeric_limits<Bias>::min(),
+                "bias too large or small to be represented");
+  }
+
+  static void check_exponent(int64_t exponent) {
+    TORCH_CHECK(
+        exponent <=
+            std::numeric_limits<decltype(std::declval<Self>().exponent)>::max(),
+        "exponent bit width is too large to be represented");
+  }
+
+  static void check_mantissa(int64_t mantissa) {
+    TORCH_CHECK(
+        mantissa <=
+            std::numeric_limits<decltype(std::declval<Self>().mantissa)>::max(),
+        "mantissa bit width is too large to be represented");
+  }
+
+  static SelfPtr int_(int64_t size_bits, c10::optional<int64_t> bias) {
+    check_size_bits(size_bits, true);
+    check_bias(bias.value_or(0));
+    return c10::make_intrusive<Self>(
+        ScalarType::int_(size_bits, bias.value_or(0)));
+  }
+
+  static SelfPtr uint(int64_t size_bits, c10::optional<int64_t> bias) {
+    check_size_bits(size_bits, true);
+    check_bias(bias.value_or(0));
+    return c10::make_intrusive<Self>(
+        ScalarType::uint(size_bits, bias.value_or(0)));
+  }
+
+  static SelfPtr float_IEEE754(int64_t exponent, int64_t mantissa) {
+    check_mantissa(mantissa);
+    check_exponent(exponent);
+    return c10::make_intrusive<Self>(
+        ScalarType::float_IEEE754(exponent, mantissa));
+  }
+
+  static SelfPtr float_(int64_t exponent, int64_t mantissa,
+                        bool finite_values_only, int64_t nan_repr) {
+    check_mantissa(mantissa);
+    check_exponent(exponent);
+    return c10::make_intrusive<Self>(ScalarType::float_(
+        exponent, mantissa, finite_values_only, NanRepr(nan_repr)));
+  }
+
+  // This needs to be implemented and throw a TypeError in order for
+  // PyTorch's opcheck to work on ops that use ScalarTypes.
+  int64_t len() const {
+    throw c10::TypeError({__func__, __FILE__, static_cast<uint32_t>(__LINE__)},
+                         "__len__ not implemented");
+    return 0;
+  }
+
+  // Serialize a ScalarType into a tuple of pairs.  Where each pair
+  // is a (fieldname, value).
+  // For simplicity, we are just going to convert to a ScalarTypeId.
+  std::tuple<std::tuple<std::string, int64_t>> obj_flatten() const {
+    return {{"ScalarType", id()}};
+  }
+
+  // Deserialize a scalar type that has been serialized by obj_flatten,
+  // ostensibly from a tuple of (member name, value) pairs, but in reality
+  // just a ScalarTypeId.
+  static SelfPtr obj_unflatten(
+      std::tuple<std::tuple<std::string, int64_t>> const& flat_type) {
+    return c10::make_intrusive<Self>(
+        from_id(std::get<1>(std::get<0>(flat_type))));
+  }
+
+  template <typename T>
+  static void bind_readonly_property(torch::class_<Self>& cls,
+                                     std::string const& name, T Base::*field) {
+    auto getter_func_helper = [field = std::move(field)](SelfPtr const& self) {
+      if constexpr (std::is_member_function_pointer_v<decltype(field)>) {
+        return (self.get()->*field)();
+      } else {
+        return self.get()->*field;
+      }
+    };
+
+    auto getter_func = [field = std::move(field),
+                        getter_func_helper = std::move(getter_func_helper)](
+                           SelfPtr const& self) {
+      auto val = getter_func_helper(self);
+      // upconvert uint8_t, int32_t etc. to int64_t for python
+      if constexpr (std::is_integral_v<T>) {
+        return static_cast<int64_t>(val);
+      } else {
+        return val;
+      }
+    };
+
+    cls.def_property(name, getter_func);
+  }
+
+  template <typename MemberFunc, typename Cls>
+  static void bind_function(torch::class_<Self>& cls, const std::string& name,
+                            MemberFunc Cls::*member) {
+    cls.def(name, [member = std::move(member)](SelfPtr const& self) {
+      return (self.get()->*member)();
+    });
+  }
+
+  template <typename Func>
+  static void bind_function(torch::class_<Self>& cls, const std::string& name,
+                            Func func) {
+    cls.def(name, func);
+  }
+
+  template <typename Func>
+  static void bind_static_function(torch::class_<Self>& cls,
+                                   const std::string& name, Func func) {
+    cls.def_static(name, func);
+  }
+
+  static void bind_class(torch::Library& lib) {
+    auto cls = lib.class_<ScalarTypeTorch>("ScalarType")
+                   .def(torch::init<int64_t, int64_t, int64_t, bool>());
+
+    // Bind Properties
+    bind_readonly_property(cls, "mantissa", &Base::mantissa);
+    bind_readonly_property(cls, "exponent", &Base::exponent);
+    bind_readonly_property(cls, "bias", &Base::bias);
+    bind_readonly_property(cls, "signed", &Base::is_signed);
+    bind_readonly_property(cls, "size_bits", &Base::size_bits);
+
+    // Bind member functions
+    bind_function(cls, "is_signed", &Base::is_signed);
+    bind_function(cls, "is_integer", &Base::is_integer);
+    bind_function(cls, "is_floating_point", &Base::is_floating_point);
+    bind_function(cls, "is_ieee_754", &Base::is_ieee_754);
+    bind_function(cls, "has_nans", &Base::has_nans);
+    bind_function(cls, "has_infs", &Base::has_infs);
+    bind_function(cls, "has_bias", &Base::has_bias);
+
+    bind_function(cls, "max", [](SelfPtr const& self) {
+      return std::visit([](auto arg) { return c10::IValue(arg); },
+                        self.get()->max());
+    });
+    bind_function(cls, "min", [](SelfPtr const& self) {
+      return std::visit([](auto arg) { return c10::IValue(arg); },
+                        self.get()->min());
+    });
+
+    bind_function(cls, "__len__", &ScalarTypeTorch::len);
+    bind_function(cls, "__str__", &Base::str);
+    bind_function(cls, "__eq__", [](SelfPtr const& self, SelfPtr const& other) {
+      return *self == *other;
+    });
+    bind_function(cls, "__repr__", [](SelfPtr const& self) {
+      return "ScalarType." + self.get()->str();
+    });
+
+    bind_function(cls, "__obj_flatten__", &ScalarTypeTorch::obj_flatten);
+    bind_static_function(cls, "__obj_unflatten__",
+                         &ScalarTypeTorch::obj_unflatten);
+
+    // Bind static functions (convenience constructors)
+    bind_static_function(cls, "int_", &ScalarTypeTorch::int_);
+    bind_static_function(cls, "uint", &ScalarTypeTorch::uint);
+    bind_static_function(cls, "float_IEEE754", &ScalarTypeTorch::float_IEEE754);
+    bind_static_function(cls, "float_", &ScalarTypeTorch::float_);
+  }
+};
+#endif
+using ScalarTypeId = int64_t;
+//using ScalarTypeTorchPtr = c10::intrusive_ptr<ScalarTypeTorch>;
+
+// "rust style" names generally following:
+//   https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L60-L70
+static inline constexpr auto kS4 = ScalarType::int_(4);
+static inline constexpr auto kU4 = ScalarType::uint(4);
+static inline constexpr auto kU4B8 = ScalarType::uint(4, 8);
+static inline constexpr auto kS8 = ScalarType::int_(8);
+static inline constexpr auto kU8 = ScalarType::uint(8);
+static inline constexpr auto kU8B128 = ScalarType::uint(8, 128);
+
+static inline constexpr auto kFE3M2f =
+    ScalarType::float_(3, 2, true, ScalarType::NAN_NONE);
+static inline constexpr auto kFE4M3fn =
+    ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2);
+static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7);
+static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10);
+#if 0
+// Fixed width style names, generally following:
+//  https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57
+static inline constexpr auto kInt4 = kS4;
+static inline constexpr auto kUint4 = kU4;
+static inline constexpr auto kUint4b8 = kU4B8;
+static inline constexpr auto kInt8 = kS8;
+static inline constexpr auto kUint8 = kU8;
+static inline constexpr auto kUint8b128 = kU8B128;
+
+static inline constexpr auto kFloat6_e3m2f = kFE3M2f;
+static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn;
+static inline constexpr auto kFloat8_e5m2 = kFE5M2;
+static inline constexpr auto kFloat16_e8m7 = kFE8M7;
+static inline constexpr auto kFloat16_e5m10 = kFE5M10;
+
+// colloquial names
+static inline constexpr auto kHalf = kFE5M10;
+static inline constexpr auto kFloat16 = kHalf;
+static inline constexpr auto kBFloat16 = kFE8M7;
+
+static inline constexpr auto kFloat16Id = kFloat16.id();
+#endif
+};  // namespace vllm
diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp
index 1a1896b4c..cf326b922 100644
--- a/csrc/torch_bindings.cpp
+++ b/csrc/torch_bindings.cpp
@@ -222,6 +222,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
   // Quantization ops
 #ifndef USE_ROCM
+#ifndef USE_MACA
   // Quantized GEMM for AQLM.
   ops.def(
       "aqlm_gemm(Tensor input, Tensor codes, Tensor codebooks, "
@@ -236,21 +237,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int[] codebook_partition_sizes) -> Tensor",
       {stride_tag});
   ops.impl("aqlm_dequant", torch::kCUDA, &aqlm_dequant);
+#endif // USE_MACA
 
   // Quantized GEMM for AWQ.
   ops.def(
       "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
-      "Tensor _zeros, SymInt split_k_iters) -> Tensor",
-      {stride_tag});
+      "Tensor _zeros, SymInt split_k_iters, Tensor _temp_space, bool dtype_bf16) -> Tensor");
   ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);
 
   // Dequantization for AWQ.
   ops.def(
       "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
-      "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor",
-      {stride_tag});
+      "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
   ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
 
+  // Convert AWQ to GPTQ
+  ops.def(
+      "awq_to_gptq_4bit(Tensor qweight) -> Tensor");
+  ops.impl("awq_to_gptq_4bit", torch::kCUDA, &awq_to_gptq_4bit);
+
+#ifndef USE_MACA
   // Note about marlin kernel 'workspace' arguments:
   // Technically these should be mutable since they are modified by the kernel.
   // But since they are set back to zero once the kernel is finished we can
@@ -281,7 +287,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor",
       {stride_tag});
   //  conditionally compiled so impl in source file
-
+#endif // USE_MACA
   // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
   ops.def(
       "machete_supported_schedules("
@@ -319,6 +325,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
   ops.impl("permute_cols", torch::kCUDA, &permute_cols);
 
+#ifndef USE_MACA
   // gptq_marlin Optimized Quantized GEMM for GPTQ.
   ops.def(
       "gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
@@ -374,8 +381,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("ggml_moe_a8_vec", torch::kCUDA, &ggml_moe_a8_vec);
 
   ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
+#endif // USE_MACA
 
 #ifndef USE_ROCM
+#ifndef USE_MACA
   // marlin_qqq_gemm for QQQ.
   ops.def(
       "marlin_qqq_gemm(Tensor a, Tensor b_q_weight, "
@@ -400,6 +409,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()",
       {stride_tag});
   ops.impl("cutlass_fp4_group_mm", torch::kCUDA, &cutlass_fp4_group_mm);
+#endif // USE_MACA
 
   // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
   // quantization, as well as bias
@@ -420,6 +430,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       {stride_tag});
   ops.impl("cutlass_scaled_mm_azp", torch::kCUDA, &cutlass_scaled_mm_azp);
 
+#ifndef USE_MACA
   // Check if cutlass scaled_mm is supported for CUDA devices of the given
   // capability
   ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool");
@@ -498,6 +509,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // CUTLASS sparse matrix compressor
   ops.def("cutlass_sparse_compress(Tensor a) -> Tensor[]");
   ops.impl("cutlass_sparse_compress", &cutlass_sparse_compress);
+#endif // USE_MACA
 
   // CUTLASS MLA decode
   ops.def(
@@ -506,6 +518,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "                   Tensor page_table, float scale) -> ()");
   ops.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode);
 
+#ifndef USE_MACA
   // Compute NVFP4 block quantized tensor.
   ops.def(
       "scaled_fp4_quant(Tensor! output, Tensor input,"
@@ -523,6 +536,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // of the given capability
   ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
   ops.impl("cutlass_scaled_mm_supports_fp4", &cutlass_scaled_mm_supports_fp4);
+#endif // USE_MACA
 #endif
 
   // Quantized GEMM for GPTQ.
@@ -530,9 +544,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // to prevent the meta function registry.
   ops.def(
       "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
-      "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, int bit) "
-      "-> Tensor",
-      {stride_tag});
+      "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, int bit, int group_size, Tensor perm_space, "
+      "Tensor temp_space, bool dtype_bf16)-> Tensor");
   ops.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);
 
   // Post processing for GPTQ.
@@ -572,6 +585,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "Tensor!? azp) -> ()");
   ops.impl("dynamic_scaled_int8_quant", torch::kCUDA,
            &dynamic_scaled_int8_quant);
+  ops.def(
+      "dynamic_scaled_int8_mask_quant(Tensor! result, Tensor input, Tensor mask, Tensor! scale, "
+      "Tensor!? azp) -> ()");
+  ops.impl("dynamic_scaled_int8_mask_quant", torch::kCUDA,
+           &dynamic_scaled_int8_mask_quant);
+  ops.def(
+    "fused_silu_mul_dq_mask_quant_pack(Tensor! out, "
+    "Tensor input,"
+    "Tensor mask) -> ()");
+  ops.impl("fused_silu_mul_dq_mask_quant_pack", torch::kCUDA, &fused_silu_mul_dq_mask_quant_pack);
 
   // Mamba selective scan kernel
   ops.def(
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index a8a6f3417..0560f7aa5 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -533,7 +533,7 @@ Specified using `--task generate`.
 | `LlavaNextVideoForConditionalGeneration`     | LLaVA-NeXT-Video                                                         | T + V                                                                 | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc.                                                                                                                 |                       | ✅︎                          | ✅︎                      |
 | `LlavaOnevisionForConditionalGeneration`     | LLaVA-Onevision                                                          | T + I<sup>+</sup> + V<sup>+</sup>                                     | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc.                                                            |                       | ✅︎                          | ✅︎                      |
 | `MiniCPMO`                                   | MiniCPM-O                                                                | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup>                  | `openbmb/MiniCPM-o-2_6`, etc.                                                                                                                           | ✅︎                     | ✅︎                          | ✅︎                    |
-| `MiniCPMV`                                   | MiniCPM-V                                                                | T + I<sup>E+</sup> + V<sup>E+</sup>                                   | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, etc.                                                         | ✅︎                     |                            | ✅︎                    |
+| `MiniCPMV`                                   | MiniCPM-V                                                                | T + I<sup>E+</sup> + V<sup>E+</sup>                                   | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, etc.                                                         | ✅︎                     |                            | ✅︎                    |
 | `MiniMaxVL01ForConditionalGeneration`        | MiniMax-VL                                                               | T + I<sup>E+</sup>                                                    | `MiniMaxAI/MiniMax-VL-01`, etc.                                                                                                                         |                       | ✅︎                          |                       |
 | `Mistral3ForConditionalGeneration`           | Mistral3                                                                 | T + I<sup>+</sup>                                                     | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc.                                                                                                   | ✅︎                     | ✅︎                          | ✅︎                    |
 | `MllamaForConditionalGeneration`             | Llama 3.2                                                                | T + I<sup>+</sup>                                                     | `meta-llama/Llama-3.2-90B-Vision-Instruct`, `meta-llama/Llama-3.2-11B-Vision`, etc.                                                                     |                        |                             |                        |
diff --git a/env.sh b/env.sh
new file mode 100644
index 000000000..e95190521
--- /dev/null
+++ b/env.sh
@@ -0,0 +1,10 @@
+DEFAULT_DIR="/opt/maca"
+export MACA_PATH=${1:-$DEFAULT_DIR}
+export CUDA_PATH=/usr/local/cuda
+export CUCC_PATH=${MACA_PATH}/tools/cu-bridge
+export PATH=${CUDA_PATH}/bin:${MACA_PATH}/mxgpu_llvm/bin:${MACA_PATH}/bin:${CUCC_PATH}/tools:${CUCC_PATH}/bin:${PATH}
+export LD_LIBRARY_PATH=${MACA_PATH}/lib:${MACA_PATH}/ompi/lib:${MACA_PATH}/mxgpu_llvm/lib:${LD_LIBRARY_PATH}
+
+export VLLM_INSTALL_PUNICA_KERNELS=1
+
+echo "MACA PATH: ${MACA_PATH} Compile Code"
diff --git a/pyproject.toml b/pyproject.toml
index 307878f7e..685922ac3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,8 +15,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "vllm"
 authors = [{name = "vLLM Team"}]
-license = "Apache-2.0"
-license-files = ["LICENSE"]
+license = { "file"= "LICENSE" }
 readme = "README.md"
 description = "A high-throughput and memory-efficient inference and serving engine for LLMs"
 classifiers = [
diff --git a/requirements/build.txt b/requirements/build.txt
index 528cd3b53..cb528a965 100644
--- a/requirements/build.txt
+++ b/requirements/build.txt
@@ -4,7 +4,7 @@ ninja
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8
-torch==2.7.0
+torch==2.6.0
 wheel
 jinja2>=3.1.6
 regex
diff --git a/requirements/common.txt b/requirements/common.txt
index 871df3d21..ed3941e43 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -2,12 +2,14 @@ regex # Replace re for higher-performance regex matching
 cachetools
 psutil
 sentencepiece  # Required for LLaMA tokenizer.
-numpy
+numpy < 2.0
 requests >= 2.26.0
 tqdm
 blake3
 py-cpuinfo
-transformers >= 4.51.1
+# Fix transformers version to avoid mixtral models loading error (config["head_dim"] == null)
+# transformers >= 4.51.1
+transformers == 4.51.3
 huggingface-hub[hf_xet] >= 0.32.0  # Required for Xet downloads.
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
 protobuf # Required by LlamaTokenizer.
diff --git a/requirements/cuda.txt b/requirements/cuda.txt
index a71d9728f..bf32479b2 100644
--- a/requirements/cuda.txt
+++ b/requirements/cuda.txt
@@ -5,10 +5,10 @@ numba == 0.60.0; python_version == '3.9' # v0.61 doesn't support Python 3.9. Req
 numba == 0.61.2; python_version > '3.9'
 
 # Dependencies for NVIDIA GPUs
-ray[cgraph]>=2.43.0, !=2.44.* # Ray Compiled Graph, required for pipeline parallelism in V1.
-torch==2.7.0
-torchaudio==2.7.0
-# These must be updated alongside torch
-torchvision==0.22.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
-# https://github.com/facebookresearch/xformers/releases/tag/v0.0.30
-xformers==0.0.30; platform_system == 'Linux' and platform_machine == 'x86_64'  # Requires PyTorch >= 2.7
+# ray[cgraph]>=2.43.0, !=2.44.* # Ray Compiled Graph, required for pipeline parallelism in V1.
+torch==2.6.0
+# torchaudio==2.7.0
+# # These must be updated alongside torch
+# torchvision==0.22.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
+# # https://github.com/facebookresearch/xformers/releases/tag/v0.0.30
+# xformers==0.0.30; platform_system == 'Linux' and platform_machine == 'x86_64'  # Requires PyTorch >= 2.7
diff --git a/setup.py b/setup.py
index ea7cd0169..786206d86 100644
--- a/setup.py
+++ b/setup.py
@@ -16,9 +16,11 @@ import torch
 from packaging.version import Version, parse
 from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
-from setuptools_scm import get_version
+# from setuptools_scm import get_version
 from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
 
+USE_MACA = True
+CMAKE_EXECUTABLE = 'cmake' if not USE_MACA else 'cmake_maca'
 
 def load_module_from_path(module_name, path):
     spec = importlib.util.spec_from_file_location(module_name, path)
@@ -145,10 +147,19 @@ class cmake_build_ext(build_ext):
         default_cfg = "Debug" if self.debug else "RelWithDebInfo"
         cfg = envs.CMAKE_BUILD_TYPE or default_cfg
 
+        maca_version = get_maca_version_list()
+
         cmake_args = [
             '-DCMAKE_BUILD_TYPE={}'.format(cfg),
             '-DVLLM_TARGET_DEVICE={}'.format(VLLM_TARGET_DEVICE),
         ]
+        if USE_MACA:
+            maca_args_ext = ['-DUSE_MACA=ON',
+                '-DMACA_VERSION_MAJOR={}'.format(maca_version[0]),
+                '-DMACA_VERSION_MINOR={}'.format(maca_version[1]),
+                '-DMACA_VERSION_PATCH={}'.format(maca_version[2]),
+                '-DMACA_VERSION_BUILD={}'.format(maca_version[3]),]
+            cmake_args.extend(maca_args_ext)
 
         verbose = envs.VERBOSE
         if verbose:
@@ -203,16 +214,16 @@ class cmake_build_ext(build_ext):
             # Default build tool to whatever cmake picks.
             build_tool = []
         # Make sure we use the nvcc from CUDA_HOME
-        if _is_cuda():
-            cmake_args += [f'-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc']
+        # if _is_cuda():
+        #     cmake_args += [f'-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc']
         subprocess.check_call(
-            ['cmake', ext.cmake_lists_dir, *build_tool, *cmake_args],
+            [CMAKE_EXECUTABLE, ext.cmake_lists_dir, *build_tool, *cmake_args],
             cwd=self.build_temp)
 
     def build_extensions(self) -> None:
         # Ensure that CMake is present and working
         try:
-            subprocess.check_output(['cmake', '--version'])
+            subprocess.check_output([CMAKE_EXECUTABLE, '--version'])
         except OSError as e:
             raise RuntimeError('Cannot find CMake executable') from e
 
@@ -239,7 +250,7 @@ class cmake_build_ext(build_ext):
             *[f"--target={name}" for name in targets],
         ]
 
-        subprocess.check_call(["cmake", *build_args], cwd=self.build_temp)
+        subprocess.check_call([CMAKE_EXECUTABLE, *build_args], cwd=self.build_temp)
 
         # Install the libraries
         for ext in self.extensions:
@@ -258,7 +269,7 @@ class cmake_build_ext(build_ext):
 
             # prefix here should actually be the same for all components
             install_args = [
-                "cmake", "--install", ".", "--prefix", prefix, "--component",
+                CMAKE_EXECUTABLE, "--install", ".", "--prefix", prefix, "--component",
                 target_name(ext.name)
             ]
             subprocess.check_call(install_args, cwd=self.build_temp)
@@ -374,9 +385,9 @@ class repackage_wheel(build_ext):
             files_to_copy = [
                 "vllm/_C.abi3.so",
                 "vllm/_moe_C.abi3.so",
-                "vllm/_flashmla_C.abi3.so",
-                "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
-                "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
+                # "vllm/_flashmla_C.abi3.so",
+                # "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
+                # "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
                 "vllm/cumem_allocator.abi3.so",
                 # "vllm/_version.py", # not available in nightly wheels yet
             ]
@@ -544,9 +555,52 @@ def get_gaudi_sw_version():
             " ", "").split(":")[1][:-1].split("-")[0]
     return "0.0.0"  # when hl-smi is not available
 
+def get_maca_version():
+    """
+    Returns the MACA SDK Version
+    """
+    maca_path = str(os.getenv('MACA_PATH'))
+    if not os.path.exists(maca_path):
+        return None
+    file_full_path = os.path.join(maca_path, 'Version.txt')
+    if not os.path.isfile(file_full_path):
+        return None
+    
+    with open(file_full_path, 'r', encoding='utf-8') as file:
+        first_line = file.readline().strip()
+    return first_line.split(":")[-1]
+
+def get_maca_version_list():
+    version_str = get_maca_version()
+    version_list = list(map(int, (version_str or "0.0.0.0").split('.')))
+    version_list.extend([0] * (4 - len(version_list)))
+    return version_list
+    
+def get_git_commit():
+    curdir = os.path.dirname(__file__)
+    default_gitdir = os.path.normpath(os.path.join(curdir, ".git"))
+    print(default_gitdir)
+    try:
+        subprocess.check_output(["git", "--git-dir", default_gitdir, "config", "--global", "--add", "safe.directory", '*'])
+        commit_id = subprocess.check_output(["git", "--git-dir", default_gitdir, "rev-parse", "HEAD"]).decode("utf-8").strip()
+        return commit_id
+    except Exception as e:
+        print(f"Error: {e}")
+        return "git error"
+
+def write_to_file(file_path, content):
+    try:
+        with open(file_path, "w") as file:
+            file.write(content)
+        print(f"Content written to {file_path} successfully.")
+    except Exception as e:
+        print(f"Error writing to file: {e}")
 
 def get_vllm_version() -> str:
-    version = get_version(write_to="vllm/_version.py")
+    # version = get_version(write_to="vllm/_version.py")
+    commit_id = get_git_commit()
+    write_to_file("vllm/_release_info.txt", commit_id)
+    version = "0.9.1"
     sep = "+" if "+" not in version else "."  # dev versions might contain +
 
     if _no_device():
@@ -556,12 +610,18 @@ def get_vllm_version() -> str:
         if envs.VLLM_USE_PRECOMPILED:
             version += f"{sep}precompiled"
         else:
-            cuda_version = str(get_nvcc_cuda_version())
-            if cuda_version != MAIN_CUDA_VERSION:
-                cuda_version_str = cuda_version.replace(".", "")[:3]
-                # skip this for source tarball, required for pypi
-                if "sdist" not in sys.argv:
-                    version += f"{sep}cu{cuda_version_str}"
+            if not USE_MACA:
+                cuda_version = str(get_nvcc_cuda_version())
+                if cuda_version != MAIN_CUDA_VERSION:
+                    cuda_version_str = cuda_version.replace(".", "")[:3]
+                    # skip this for source tarball, required for pypi
+                    if "sdist" not in sys.argv:
+                        version += f"{sep}cu{cuda_version_str}"
+            else:
+                maca_version_str = get_maca_version()
+                torch_version = torch.__version__
+                major_minor_version = ".".join(torch_version.split(".")[:2])
+                version += f"{sep}maca{maca_version_str}torch{major_minor_version}"
     elif _is_hip():
         # Get the Rocm Version
         rocm_version = get_rocm_version() or torch.version.hip
@@ -648,7 +708,7 @@ if _is_cuda() or _is_hip():
 if _is_hip():
     ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
 
-if _is_cuda():
+if _is_cuda() and False:
     ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C"))
     if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.3"):
         # FA3 requires CUDA 12.3 or later
@@ -658,9 +718,11 @@ if _is_cuda():
         # not targeting a hopper system
         ext_modules.append(
             CMakeExtension(name="vllm._flashmla_C", optional=True))
+
+if _is_cuda():
     ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
 
-if _build_custom_ops():
+if _build_custom_ops() or True:
     ext_modules.append(CMakeExtension(name="vllm._C"))
 
 package_data = {
@@ -668,6 +730,8 @@ package_data = {
         "py.typed",
         "model_executor/layers/fused_moe/configs/*.json",
         "model_executor/layers/quantization/utils/configs/*.json",
+        "attention/backends/configs/*.json",
+        "_release_info.txt",
     ]
 }
 
diff --git a/tests/models/registry.py b/tests/models/registry.py
index e6543c197..eaa08a819 100644
--- a/tests/models/registry.py
+++ b/tests/models/registry.py
@@ -156,6 +156,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
                                          trust_remote_code=True),
     "DeepseekV3ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V3",  # noqa: E501
                                          trust_remote_code=True),
+    "Ernie4_5_ForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-0.3B-PT",
+                                        trust_remote_code=True),
+    "Ernie4_5_MoeForCausalLM": _HfExamplesInfo("baidu/ERNIE-4.5-21B-A3B-PT",
+                                        trust_remote_code=True),
     "ExaoneForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct"),  # noqa: E501
     "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"),  # noqa: E501
     "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"),
@@ -351,7 +355,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "MiniCPMO": _HfExamplesInfo("openbmb/MiniCPM-o-2_6",
                                 trust_remote_code=True),
     "MiniCPMV": _HfExamplesInfo("openbmb/MiniCPM-Llama3-V-2_5",
-                                extras={"2.6": "openbmb/MiniCPM-V-2_6"},  # noqa: E501
+                                extras={"2.6": "openbmb/MiniCPM-V-2_6", "4.0": "openbmb/MiniCPM-V-4"},  # noqa: E501
                                 trust_remote_code=True),
     "MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo("MiniMaxAI/MiniMax-VL-01", # noqa: E501
                                               trust_remote_code=True,
diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
index 92de1f5ef..6ae546dc2 100644
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
@@ -268,7 +268,6 @@ def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,
                                           cos_sin_cache, is_neox, rot_dim,
                                           cos_sin_cache_offsets)
 
-
 # layer norm ops
 def rms_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor,
              epsilon: float) -> None:
@@ -387,21 +386,30 @@ def awq_dequantize(qweight: torch.Tensor, scales: torch.Tensor,
 
 
 def awq_gemm(input: torch.Tensor, qweight: torch.Tensor, qzeros: torch.Tensor,
-             scales: torch.Tensor, split_k_iters: int) -> torch.Tensor:
+             scales: torch.Tensor, split_k_iters: int, temp_space: torch.Tensor, 
+             dtype_bf16: bool) -> torch.Tensor:
     if envs.VLLM_USE_TRITON_AWQ:
         from vllm.model_executor.layers.quantization.awq_triton import (
             awq_gemm_triton)
-        return awq_gemm_triton(input, qweight, qzeros, scales, split_k_iters)
-    return torch.ops._C.awq_gemm(input, qweight, qzeros, scales, split_k_iters)
+        return awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters)
+    return torch.ops._C.awq_gemm(input, qweight, scales, qzeros, split_k_iters,
+                                temp_space, dtype_bf16)
 
+# awq to gptq 4bit conversion
+def awq_to_gptq_4bit(qweight: torch.Tensor) -> torch.Tensor:
+    if envs.VLLM_USE_TRITON_AWQ:
+        return qweight
+    return torch.ops._C.awq_to_gptq_4bit(qweight)
 
 # gptq
 def gptq_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
               b_gptq_qzeros: torch.Tensor, b_gptq_scales: torch.Tensor,
               b_g_idx: torch.Tensor, use_exllama: bool,
-              bit: int) -> torch.Tensor:
+              bit: int, group_size: int, perm_space: torch.Tensor,
+              temp_space: torch.Tensor, dtype_bf16: bool) -> torch.Tensor:
     return torch.ops._C.gptq_gemm(a, b_q_weight, b_gptq_qzeros, b_gptq_scales,
-                                  b_g_idx, use_exllama, bit)
+                                  b_g_idx, use_exllama, bit, group_size, 
+                                  perm_space, temp_space, dtype_bf16)
 
 
 if hasattr(torch.ops._C, "gptq_gemm"):
@@ -410,7 +418,9 @@ if hasattr(torch.ops._C, "gptq_gemm"):
     def _gptq_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor,
                         b_gptq_qzeros: torch.Tensor,
                         b_gptq_scales: torch.Tensor, b_g_idx: torch.Tensor,
-                        use_exllama: bool, bit: int) -> torch.Tensor:
+                        use_exllama: bool, bit: int, 
+                        group_size: int, perm_space: torch.Tensor,
+                        temp_space: torch.Tensor, dtype_bf16: bool) -> torch.Tensor:
         return torch.empty((a.size(0), b_q_weight.size(1)),
                            dtype=a.dtype,
                            device=a.device)
@@ -503,7 +513,8 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
     @register_fake("_C::awq_gemm")
     def _awq_gemm_fake(input: torch.Tensor, qweight: torch.Tensor,
                        qzeros: torch.Tensor, scales: torch.Tensor,
-                       split_k_iters: torch.SymInt) -> torch.Tensor:
+                       split_k_iters: torch.SymInt, temp_space: torch.Tensor, 
+                       dtype_bf16: bool) -> torch.Tensor:
         num_in_feats = input.size(0)
         return torch.empty((split_k_iters, num_in_feats, qweight.size(1) * 8),
                            dtype=input.dtype,
@@ -645,8 +656,8 @@ if hasattr(torch.ops._C, "ggml_moe_a8_vec"):
 
 # cutlass
 def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool:
-    return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability)
-
+    # return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability)
+    return False
 
 def cutlass_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,
                           block_scale_a: torch.Tensor,
@@ -659,15 +670,24 @@ def cutlass_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,
                                        alpha)
     return out
 
-
 def cutlass_scaled_mm_supports_fp8(cuda_device_capability: int) -> bool:
-    return torch.ops._C.cutlass_scaled_mm_supports_fp8(cuda_device_capability)
+    # return torch.ops._C.cutlass_scaled_mm_supports_fp8(cuda_device_capability)
+    return True
 
 
 def cutlass_scaled_mm_supports_block_fp8(cuda_device_capability: int) -> bool:
-    return torch.ops._C.cutlass_scaled_mm_supports_block_fp8(
-        cuda_device_capability)
-
+    # return torch.ops._C.cutlass_scaled_mm_supports_block_fp8(
+    #     cuda_device_capability)
+    return True
+
+# Batch gemm in vllm, support w8a8 int8 quantization
+def cutlass_scaled_batch_mm(a: torch.Tensor, b: torch.Tensor,
+                            scale_a: torch.Tensor, scale_b: torch.Tensor,
+                            out_dtype: torch.dtype, bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+    assert (a.shape[0] == b.shape[0] and a.shape[2] == b.shape[1])
+    out = torch.empty((a.shape[0], a.shape[1], b.shape[2]), device = a.device, dtype = out_dtype)
+    torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias)
+    return out
 
 def cutlass_scaled_mm(a: torch.Tensor,
                       b: torch.Tensor,
@@ -1380,6 +1400,62 @@ def scaled_int8_quant(
                                            input_azp)
     return output, input_scales, input_azp
 
+def scaled_int8_quant_mask(
+    input: torch.Tensor,
+    mask:  torch.Tensor,
+    scale: Optional[torch.Tensor] = None,
+    azp: Optional[torch.Tensor] = None,
+    symmetric: bool = True
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """
+    Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
+
+    Args:
+        input: The input tensor to be quantized to int8.
+        scale: Optional scaling factor for the int8 quantization.
+            When not provided, we invoke dynamic-per-token quantization.
+        azp: Optional zero-point for the int8 quantization.
+            Must be provided for asymmetric quantization if `scale` is provided.
+        mask: mask
+        symmetric: Whether to use symmetric quantization (scale only, azp ignored).
+
+    Returns:
+      Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] : Output int8 tensor, scales, and optionally azp.
+    """
+    output = torch.empty_like(input, dtype=torch.int8)
+    if scale is not None:
+        # static-per-tensor quantization.
+        assert symmetric == (
+            azp
+            is None), "azp must only be provided for asymmetric quantization."
+        torch.ops._C.static_scaled_int8_quant(output, input, scale, azp)
+        return output, scale, azp
+
+    # dynamic-per-token quantization.
+    input_scales = torch.empty((input.numel() // input.shape[-1], 1),
+                               device=input.device,
+                               dtype=torch.float32)
+    input_azp = None if symmetric else torch.empty_like(input_scales,
+                                                        dtype=torch.int32)
+    torch.ops._C.dynamic_scaled_int8_mask_quant(output, input, mask, input_scales,
+                                           input_azp)
+    return output, input_scales, input_azp
+
+def fused_silu_mul_dq_mask_quant(
+    input: torch.Tensor,
+    mask:  torch.Tensor
+) -> torch.Tensor:
+    """
+    input shape [expert_num, token_num_padded, hidden_dim]
+    output shape [expert_num, token_num_padded, hidden_dim // 2], dtype bf16
+    masked_m shape [expert_num], indicates valid tokens per expert
+
+    implement silu_and_mul + quant + package
+    """
+    out_stride = (input.shape[-1] // 4 + 257) // 256 * 256
+    output = torch.empty((input.shape[0], input.shape[1], out_stride), device=input.device, dtype=input.dtype)
+    torch.ops._C.fused_silu_mul_dq_mask_quant_pack(output, input, mask)
+    return output
 
 # qqq ops
 def marlin_qqq_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
@@ -1555,6 +1631,14 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
     torch.ops._moe_C.topk_softmax(topk_weights, topk_ids,
                                   token_expert_indicies, gating_output)
 
+def fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
+                    topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+                    sorted_token_ids: torch.Tensor, expert_ids: torch.Tensor,
+                    num_tokens_post_padded: torch.Tensor, mul_routed_weight: bool, top_k: int, tileConfig: int) -> None:
+    torch.ops._moe_C.fused_moe_kernel(A, B, C,
+                    topk_weights, topk_ids,
+                    sorted_token_ids, expert_ids,
+                    num_tokens_post_padded, mul_routed_weight, top_k, tileConfig)
 
 def moe_wna16_marlin_gemm(input: torch.Tensor, output: Optional[torch.Tensor],
                           b_qweight: torch.Tensor, b_scales: torch.Tensor,
diff --git a/vllm/attention/backends/configs/tp8_merge.json b/vllm/attention/backends/configs/tp8_merge.json
new file mode 100644
index 000000000..773051b85
--- /dev/null
+++ b/vllm/attention/backends/configs/tp8_merge.json
@@ -0,0 +1,986 @@
+[
+    {
+        "BS": 1,
+        "L": 2,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 4,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 8,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 16,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 32,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 64,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 128,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 256,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 512,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 1024,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 2048,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 4096,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 8192,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 16384,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 32768,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 1,
+        "L": 65536,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 2,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 4,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 8,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 16,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 32,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 64,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 128,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 256,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 512,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 1024,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 2048,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 4096,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 8192,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 16384,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 32768,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 2,
+        "L": 65536,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 2,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 4,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 8,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 16,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 32,
+        "num_kv_splits": 13,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 64,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 128,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 256,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 512,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 1024,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 2048,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 4096,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 8192,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 16384,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 4,
+        "L": 32768,
+        "num_kv_splits": 16,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 2,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 8,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 16,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 32,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 64,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 128,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 256,
+        "num_kv_splits": 8,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 512,
+        "num_kv_splits": 13,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 1024,
+        "num_kv_splits": 13,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 2048,
+        "num_kv_splits": 13,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 4096,
+        "num_kv_splits": 13,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 8192,
+        "num_kv_splits": 13,
+        "num_stages": 2
+    },
+    {
+        "BS": 8,
+        "L": 16384,
+        "num_kv_splits": 13,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 32,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 64,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 128,
+        "num_kv_splits": 4,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 256,
+        "num_kv_splits": 6,
+        "num_stages": 2
+    },
+    {
+        "BS": 16,
+        "L": 512,
+        "num_kv_splits": 13,
+        "num_stages": 1
+    },
+    {
+        "BS": 16,
+        "L": 1024,
+        "num_kv_splits": 13,
+        "num_stages": 1
+    },
+    {
+        "BS": 16,
+        "L": 2048,
+        "num_kv_splits": 13,
+        "num_stages": 1
+    },
+    {
+        "BS": 16,
+        "L": 4096,
+        "num_kv_splits": 13,
+        "num_stages": 1
+    },
+    {
+        "BS": 16,
+        "L": 8192,
+        "num_kv_splits": 13,
+        "num_stages": 1
+    },
+    {
+        "BS": 32,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 32,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 32,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 32,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 32,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 32,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 2
+    },
+    {
+        "BS": 32,
+        "L": 128,
+        "num_kv_splits": 3,
+        "num_stages": 2
+    },
+    {
+        "BS": 32,
+        "L": 256,
+        "num_kv_splits": 6,
+        "num_stages": 1
+    },
+    {
+        "BS": 32,
+        "L": 512,
+        "num_kv_splits": 6,
+        "num_stages": 1
+    },
+    {
+        "BS": 32,
+        "L": 1024,
+        "num_kv_splits": 6,
+        "num_stages": 1
+    },
+    {
+        "BS": 32,
+        "L": 2048,
+        "num_kv_splits": 6,
+        "num_stages": 1
+    },
+    {
+        "BS": 32,
+        "L": 4096,
+        "num_kv_splits": 13,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 2,
+        "num_kv_splits": 4,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 128,
+        "num_kv_splits": 3,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 256,
+        "num_kv_splits": 3,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 512,
+        "num_kv_splits": 3,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 1024,
+        "num_kv_splits": 3,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 2048,
+        "num_kv_splits": 3,
+        "num_stages": 1
+    },
+    {
+        "BS": 64,
+        "L": 2048,
+        "num_kv_splits": 8,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 128,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 256,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 512,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 96,
+        "L": 1024,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 128,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 256,
+        "num_kv_splits": 3,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 512,
+        "num_kv_splits": 3,
+        "num_stages": 1
+    },
+    {
+        "BS": 128,
+        "L": 1024,
+        "num_kv_splits": 3,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 128,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 256,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 256,
+        "L": 512,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 512,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 512,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 512,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 512,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 512,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 512,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 512,
+        "L": 128,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 512,
+        "L": 256,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1024,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1024,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1024,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1024,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1024,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1024,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1024,
+        "L": 128,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1536,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1536,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1536,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1536,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1536,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 1536,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 2048,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 2048,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 2048,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 2048,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 2048,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 2048,
+        "L": 64,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 3072,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 3072,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 3072,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 3072,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 3072,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 4096,
+        "L": 2,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 4096,
+        "L": 4,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 4096,
+        "L": 8,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 4096,
+        "L": 16,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    },
+    {
+        "BS": 4096,
+        "L": 32,
+        "num_kv_splits": 2,
+        "num_stages": 1
+    }
+]
\ No newline at end of file
diff --git a/vllm/attention/backends/flash_attn.py b/vllm/attention/backends/flash_attn.py
index 73e377268..a46c0a8f4 100755
--- a/vllm/attention/backends/flash_attn.py
+++ b/vllm/attention/backends/flash_attn.py
@@ -23,13 +23,15 @@ from vllm.attention.backends.utils import (
     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.attention.utils.fa_utils import (flash_attn_supports_fp8,
-                                           get_flash_attn_version)
+# from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,
+#                                            get_flash_attn_version)
 from vllm.logger import init_logger
 from vllm.multimodal import MultiModalPlaceholderMap
 from vllm.utils import async_tensor_h2d, make_tensor_with_pad
-from vllm.vllm_flash_attn import (flash_attn_varlen_func,
-                                  flash_attn_with_kvcache)
+from flash_attn import (flash_attn_varlen_func,
+                        flash_attn_with_kvcache)
+def flash_attn_supports_fp8() -> bool:
+    return False
 
 if TYPE_CHECKING:
     from vllm.worker.model_runner import (ModelInputForGPUBuilder,
@@ -640,8 +642,8 @@ class FlashAttentionImpl(AttentionImpl):
         self.sliding_window = ((sliding_window - 1,
                                 0) if sliding_window is not None else (-1, -1))
         self.kv_cache_dtype = kv_cache_dtype
-        self.vllm_flash_attn_version = get_flash_attn_version(
-            requires_alibi=self.alibi_slopes is not None)
+        # self.vllm_flash_attn_version = get_flash_attn_version(
+        #     requires_alibi=self.alibi_slopes is not None)
         if is_quantized_kv_cache(self.kv_cache_dtype) and (
                 not self.kv_cache_dtype.startswith("fp8")
                 or not flash_attn_supports_fp8()):
@@ -809,7 +811,7 @@ class FlashAttentionImpl(AttentionImpl):
                         (num_kv_tokens, num_kv_heads, head_size))
 
                 descale_shape = (q_seq_start_loc.shape[0] - 1, key.shape[1])
-                flash_attn_varlen_func(
+                output[:num_prefill_query_tokens] = flash_attn_varlen_func(
                     q=query,
                     k=key,
                     v=value,
@@ -822,11 +824,11 @@ class FlashAttentionImpl(AttentionImpl):
                     window_size=window_size,
                     alibi_slopes=alibi_slopes,
                     softcap=logits_soft_cap,
-                    out=prefill_output,
-                    fa_version=self.vllm_flash_attn_version,
-                    q_descale=layer._q_scale.expand(descale_shape),
-                    k_descale=layer._k_scale.expand(descale_shape),
-                    v_descale=layer._v_scale.expand(descale_shape),
+                    # out=prefill_output,
+                    # fa_version=self.vllm_flash_attn_version,
+                    # q_descale=layer._q_scale.expand(descale_shape),
+                    # k_descale=layer._k_scale.expand(descale_shape),
+                    # v_descale=layer._v_scale.expand(descale_shape),
                 )
             else:
                 # prefix-enabled attention
@@ -837,13 +839,13 @@ class FlashAttentionImpl(AttentionImpl):
                 max_seq_len = max(prefill_meta.seq_lens)
                 descale_shape = (prefill_meta.query_start_loc.shape[0] - 1,
                                  key.shape[1])
-                flash_attn_varlen_func(  # noqa
+                output[:num_prefill_query_tokens] = flash_attn_varlen_func(  # noqa
                     q=query,
                     k=key_cache,
                     v=value_cache,
                     cu_seqlens_q=prefill_meta.query_start_loc,
                     max_seqlen_q=prefill_meta.max_query_len,
-                    seqused_k=prefill_meta.seq_lens_tensor,
+                    cu_seqlens_k=prefill_meta.seq_start_loc,
                     max_seqlen_k=max_seq_len,
                     softmax_scale=softmax_scale,
                     causal=True,
@@ -851,11 +853,11 @@ class FlashAttentionImpl(AttentionImpl):
                     alibi_slopes=alibi_slopes,
                     block_table=prefill_meta.block_tables,
                     softcap=logits_soft_cap,
-                    out=prefill_output,
-                    fa_version=self.vllm_flash_attn_version,
-                    q_descale=layer._q_scale.expand(descale_shape),
-                    k_descale=layer._k_scale.expand(descale_shape),
-                    v_descale=layer._v_scale.expand(descale_shape),
+                    # out=prefill_output,
+                    # fa_version=self.vllm_flash_attn_version,
+                    # q_descale=layer._q_scale.expand(descale_shape),
+                    # k_descale=layer._k_scale.expand(descale_shape),
+                    # v_descale=layer._v_scale.expand(descale_shape),
                 )
 
         if decode_meta := attn_metadata.decode_metadata:
@@ -872,13 +874,13 @@ class FlashAttentionImpl(AttentionImpl):
                 assert decode_meta.query_start_loc is not None
                 descale_shape = (decode_meta.query_start_loc.shape[0] - 1,
                                  key.shape[1])
-                flash_attn_varlen_func(
+                output[num_prefill_query_tokens:] = flash_attn_varlen_func(
                     q=decode_query,
                     k=key_cache,
                     v=value_cache,
                     cu_seqlens_q=decode_meta.query_start_loc,
                     max_seqlen_q=decode_meta.max_decode_query_len,
-                    seqused_k=decode_meta.seq_lens_tensor,
+                    cu_seqlens_k=decode_meta.seq_start_loc,
                     max_seqlen_k=decode_meta.max_decode_seq_len,
                     softmax_scale=softmax_scale,
                     causal=True,
@@ -886,11 +888,11 @@ class FlashAttentionImpl(AttentionImpl):
                     alibi_slopes=alibi_slopes,
                     softcap=logits_soft_cap,
                     block_table=decode_meta.block_tables,
-                    out=decode_output,
-                    fa_version=self.vllm_flash_attn_version,
-                    q_descale=layer._q_scale.expand(descale_shape),
-                    k_descale=layer._k_scale.expand(descale_shape),
-                    v_descale=layer._v_scale.expand(descale_shape),
+                    # out=decode_output,
+                    # fa_version=self.vllm_flash_attn_version,
+                    # q_descale=layer._q_scale.expand(descale_shape),
+                    # k_descale=layer._k_scale.expand(descale_shape),
+                    # v_descale=layer._v_scale.expand(descale_shape),
                 )
             else:
                 # Use flash_attn_with_kvcache for normal decoding.
@@ -900,7 +902,7 @@ class FlashAttentionImpl(AttentionImpl):
                     block_tables_arg,
                 ) = get_seq_len_block_table_args(decode_meta, False, attn_type)
                 descale_shape = (seq_lens_arg.shape[0], key_cache.shape[-2])
-                flash_attn_with_kvcache(
+                output[num_prefill_query_tokens:] = flash_attn_with_kvcache(
                     q=decode_query.unsqueeze(1),
                     k_cache=key_cache,
                     v_cache=value_cache,
@@ -911,12 +913,12 @@ class FlashAttentionImpl(AttentionImpl):
                     window_size=window_size,
                     alibi_slopes=alibi_slopes,
                     softcap=logits_soft_cap,
-                    out=decode_output.unsqueeze(1),
-                    fa_version=self.vllm_flash_attn_version,
-                    q_descale=layer._q_scale.expand(descale_shape),
-                    k_descale=layer._k_scale.expand(descale_shape),
-                    v_descale=layer._v_scale.expand(descale_shape),
-                )
+                    # out=decode_output.unsqueeze(1),
+                    # fa_version=self.vllm_flash_attn_version,
+                    # q_descale=layer._q_scale.expand(descale_shape),
+                    # k_descale=layer._k_scale.expand(descale_shape),
+                    # v_descale=layer._v_scale.expand(descale_shape),
+                ).squeeze(1)
         return output
 
 
diff --git a/vllm/attention/backends/flashinfer.py b/vllm/attention/backends/flashinfer.py
index a3937760f..5e8036b86 100644
--- a/vllm/attention/backends/flashinfer.py
+++ b/vllm/attention/backends/flashinfer.py
@@ -15,7 +15,8 @@ try:
     from flashinfer.decode import CUDAGraphBatchDecodeWithPagedKVCacheWrapper
     from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper
 
-    from vllm.vllm_flash_attn import flash_attn_varlen_func
+    # from vllm.vllm_flash_attn import flash_attn_varlen_func
+    from flash_attn import flash_attn_varlen_func
     FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
 except ImportError:
     # Avoid turning these types into variables during type checking
diff --git a/vllm/attention/backends/mla/common.py b/vllm/attention/backends/mla/common.py
index 78cf95288..103fdce57 100644
--- a/vllm/attention/backends/mla/common.py
+++ b/vllm/attention/backends/mla/common.py
@@ -208,7 +208,7 @@ from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
                                            compute_slot_mapping_start_idx,
                                            is_block_tables_empty)
 from vllm.attention.ops.merge_attn_states import merge_attn_states
-from vllm.attention.utils.fa_utils import get_flash_attn_version
+# from vllm.attention.utils.fa_utils import get_flash_attn_version
 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                LinearBase,
                                                UnquantizedLinearMethod)
@@ -239,6 +239,8 @@ if TYPE_CHECKING:
 
 is_hip = current_platform.is_rocm()
 
+def get_flash_attn_version():
+    return None
 
 class MLACommonBackend(AttentionBackend):
 
@@ -1135,7 +1137,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
                 del eye
                 # standardize to (output, input)
                 return dequant_weights.T
-            return layer.weight
+            return layer.weight if not envs.MACA_VLLM_USE_TN_2_NN else layer.weight.T
 
         # we currently do not have quantized bmm's which are needed for
         # `W_UV` and `W_UK_T`, we we just store fp16/bf16 copies and perform
diff --git a/vllm/attention/backends/triton_mla.py b/vllm/attention/backends/triton_mla.py
index e06f7d54e..5cf761817 100644
--- a/vllm/attention/backends/triton_mla.py
+++ b/vllm/attention/backends/triton_mla.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Type
 
 import torch
@@ -13,6 +14,44 @@ from vllm.attention.backends.mla.common import (MLACommonBackend,
 from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
 
 
+import json
+import os
+
+# TODO: Configure environment variables temporarily. New versions do not need to be configured
+os.environ['TRITON_ENABLE_MACA_OPT_MOVE_DOT_OPERANDS_OUT_LOOP'] = '1'
+os.environ['TRITON_ENABLE_MACA_CHAIN_DOT_OPT'] = '1'
+
+def load_config():
+    # Load JSON data from the file
+    json_path = config_file_path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "configs", "tp8_merge.json")
+    with open(json_path, 'r') as file:
+        data = json.load(file)
+    return data
+
+JSON_DATA = load_config()
+
+def find_best_mla_para(json_data, batch_size, input_len, tp_size):
+    best_match = None
+    best_batch_size_diff = float('inf')
+    best_input_len_diff = float('inf')
+    
+    for entry in json_data:
+        if entry["BS"] == batch_size and entry["L"] == input_len:
+            return entry["num_kv_splits"], entry['num_stages']
+        batch_size_diff = abs(entry["BS"] - batch_size)
+        input_len_diff = abs(entry["L"] - input_len)
+        
+        # Check if this is a better match than the current best match
+        if batch_size_diff < best_batch_size_diff or (batch_size_diff == best_batch_size_diff and input_len_diff < best_input_len_diff):
+            best_match = entry
+            best_batch_size_diff = batch_size_diff
+            best_input_len_diff = input_len_diff
+    
+    # If a match was found, return the best_kv_splits, otherwise return None
+    return best_match["num_kv_splits"],best_match["num_stages"]
+
+
 class TritonMLABackend(MLACommonBackend):
 
     @staticmethod
@@ -23,8 +62,37 @@ class TritonMLABackend(MLACommonBackend):
     def get_impl_cls() -> Type["TritonMLAImpl"]:
         return TritonMLAImpl
 
-
-class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
+    @staticmethod
+    def get_metadata_cls() -> Type["TritonMLAMetadata"]:
+        return TritonMLAMetadata
+
+@dataclass
+class TritonMLAMetadata(MLACommonMetadata):
+    num_kv_splits: int = 4  # TODO: heuristic
+    num_stages: int = 1
+
+    @property
+    def decode_metadata(self):
+        if self.num_decode_tokens == 0:
+            return None
+        if self._cached_decode_metadata is not None:
+            return self._cached_decode_metadata
+        
+        decode_metadata = super().decode_metadata
+        
+        if decode_metadata is not None:
+            if decode_metadata.seq_lens_tensor is not None:
+                batch = decode_metadata.seq_lens_tensor.shape[0]
+                max_seq_len = int(decode_metadata.seq_lens_tensor.max())
+                num_kv_splits, num_stages = find_best_mla_para(JSON_DATA, batch, max_seq_len, 8)
+            else:
+                num_kv_splits = self.num_kv_splits
+                num_stages = self.num_stages
+            decode_metadata.num_kv_splits = num_kv_splits
+            decode_metadata.num_stages = num_stages
+        return decode_metadata
+
+class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
 
     def __init__(
             self,
@@ -70,7 +138,7 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
         q_nope: torch.Tensor,
         q_pe: torch.Tensor,
         kv_c_and_k_pe_cache: torch.Tensor,
-        attn_metadata: MLACommonMetadata,
+        attn_metadata: TritonMLAMetadata,
     ) -> torch.Tensor:
         assert kv_c_and_k_pe_cache.numel() > 0
 
@@ -85,14 +153,12 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
                         dtype=q.dtype,
                         device=q.device)
 
-        num_kv_splits = 4  # TODO: heuristic
-
         # TODO(lucas) Allocate ahead of time
         attn_logits = torch.empty(
             (
                 B,
                 self.num_heads,
-                num_kv_splits,
+                decode_meta.num_kv_splits,
                 # NOTE(lucas) idk why the +1 is here but sglang has it so we
                 # just mirror that
                 self.kv_lora_rank + 1,
@@ -110,6 +176,6 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
         decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
                              decode_meta.block_tables,
                              decode_meta.seq_lens_tensor, attn_logits,
-                             num_kv_splits, self.scale, PAGE_SIZE)
+                             decode_meta.num_kv_splits, decode_meta.num_stages, self.scale, PAGE_SIZE)
 
         return self._v_up_proj(o)
diff --git a/vllm/attention/ops/flashmla.py b/vllm/attention/ops/flashmla.py
index b85f27ac4..0eaf50b29 100644
--- a/vllm/attention/ops/flashmla.py
+++ b/vllm/attention/ops/flashmla.py
@@ -10,29 +10,37 @@ from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
 
-if current_platform.is_cuda():
-    try:
-        import vllm._flashmla_C  # noqa: F401
-        _flashmla_C_AVAILABLE = True
-    except ImportError:
-        _flashmla_C_AVAILABLE = False
-else:
-    _flashmla_C_AVAILABLE = False
+# if current_platform.is_cuda():
+#     try:
+#         import vllm._flashmla_C  # noqa: F401
+#         _flashmla_C_AVAILABLE = True
+#     except ImportError:
+#         _flashmla_C_AVAILABLE = False
+# else:
+#     _flashmla_C_AVAILABLE = False
+try :
+    import flash_mla
+    _flashmla_AVAILABLE = True
+except ImportError as e:
+    logger.warning("Failed to import from flash_mla with %r on MACA Platform", e)
+    _flashmla_AVAILABLE = False
 
 
 def is_flashmla_supported() -> Tuple[bool, Optional[str]]:
     """
     Return: is_supported_flag, unsupported_reason (optional).
     """
-    if not current_platform.is_cuda():
-        return False, "FlashMLA is only supported on CUDA devices."
-    if current_platform.get_device_capability()[0] != 9:
-        return False, "FlashMLA is only supported on Hopper devices."
-    if not _flashmla_C_AVAILABLE:
-        return False, "vllm._flashmla_C is not available, likely was not "\
-            "compiled due to insufficient nvcc version or a supported arch "\
-            "(only sm90a currently) was not in the list of target arches to "\
-            "compile for."
+    # if not current_platform.is_cuda():
+    #     return False, "FlashMLA is only supported on CUDA devices."
+    # if current_platform.get_device_capability()[0] != 9:
+    #     return False, "FlashMLA is only supported on Hopper devices."
+    # if not _flashmla_C_AVAILABLE:
+    #     return False, "vllm._flashmla_C is not available, likely was not "\
+    #         "compiled due to insufficient nvcc version or a supported arch "\
+    #         "(only sm90a currently) was not in the list of target arches to "\
+    #         "compile for."
+    if not _flashmla_AVAILABLE:
+        return False, "flash_mla is not available"
     return True, None
 
 
@@ -52,7 +60,10 @@ def get_mla_metadata(
                                  dtype torch.int32.
         num_splits: (batch_size + 1), dtype torch.int32.
     """
-    return torch.ops._flashmla_C.get_mla_metadata(cache_seqlens,
+    # return torch.ops._flashmla_C.get_mla_metadata(cache_seqlens,
+    #                                               num_heads_per_head_k,
+    #                                               num_heads_k)
+    return flash_mla.flash_mla_interface.get_mla_metadata(cache_seqlens,
                                                   num_heads_per_head_k,
                                                   num_heads_k)
 
@@ -86,19 +97,30 @@ def flash_mla_with_kvcache(
         out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
         softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
     """
-    if softmax_scale is None:
-        softmax_scale = q.shape[-1]**(-0.5)
-    out, softmax_lse = torch.ops._flashmla_C.fwd_kvcache_mla(
+    # if softmax_scale is None:
+    #     softmax_scale = q.shape[-1]**(-0.5)
+    # out, softmax_lse = torch.ops._flashmla_C.fwd_kvcache_mla(
+    #     q,
+    #     k_cache,
+    #     None,
+    #     head_dim_v,
+    #     cache_seqlens,
+    #     block_table,
+    #     softmax_scale,
+    #     causal,
+    #     tile_scheduler_metadata,
+    #     num_splits,
+    # )
+    out, softmax_lse = flash_mla.flash_mla_interface.flash_mla_with_kvcache(
         q,
         k_cache,
-        None,
-        head_dim_v,
-        cache_seqlens,
         block_table,
-        softmax_scale,
-        causal,
+        cache_seqlens,
+        head_dim_v,
         tile_scheduler_metadata,
         num_splits,
+        softmax_scale,
+        causal,
     )
     return out, softmax_lse
 
diff --git a/vllm/attention/ops/triton_decode_attention.py b/vllm/attention/ops/triton_decode_attention.py
index c27b377ae..38bc59bbf 100644
--- a/vllm/attention/ops/triton_decode_attention.py
+++ b/vllm/attention/ops/triton_decode_attention.py
@@ -34,7 +34,7 @@ import logging
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
-is_hip_ = current_platform.is_rocm()
+# is_hip_ = current_platform.is_rocm()
 
 logger = logging.getLogger(__name__)
 
@@ -178,7 +178,8 @@ def _decode_att_m_fwd(
     page_size,
     logit_cap,
 ):
-    BLOCK = 64 if not is_hip_ else 8
+    # BLOCK = 64 if not is_hip_ else 8
+    BLOCK = 8
 
     NUM_KV_SPLITS = num_kv_splits
     Lk = k_buffer.shape[-1]
@@ -191,7 +192,8 @@ def _decode_att_m_fwd(
 
     num_warps = 4
     if kv_group_num != 1:
-        num_warps = 1 if is_hip_ else 2
+        # num_warps = 1 if is_hip_ else 2
+        num_warps = 1
 
     BLOCK_DMODEL = triton.next_power_of_2(Lk)
     BLOCK_DV = triton.next_power_of_2(Lv)
@@ -386,17 +388,18 @@ def _decode_grouped_att_m_fwd(
     Req_to_tokens,
     B_Seqlen,
     num_kv_splits,
+    num_stages,
     sm_scale,
     page_size,
     logit_cap,
 ):
-    BLOCK = 32
+    BLOCK = 16
     Lk = k_buffer.shape[-1]
     Lv = v_buffer.shape[-1]
 
     # [TODO] work around shmem limit on MI3xx
-    if is_hip_ and Lk >= 576:
-        BLOCK = 16
+    # if is_hip_ and Lk >= 576:
+    #     BLOCK = 16
 
     if Lk == 576:
         BLOCK_DMODEL = 512
@@ -420,17 +423,21 @@ def _decode_grouped_att_m_fwd(
         NUM_KV_SPLITS,
     )
 
-    extra_kargs = {}
-    num_stages = 2
-    if is_hip_:
-        # https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization/workload.html#mi300x-triton-kernel-performance-optimization
-        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
-        extra_kargs = {
-            "waves_per_eu": 1,
-            "matrix_instr_nonkdim": 16,
-            "kpack": 2
-        }
-        num_stages = 1
+    if num_stages == 1:
+        extra_kargs = {"scenario":"mla"}
+    elif num_stages == 2:
+        extra_kargs = {"scenario" : "mla", "pipeline" : "cpasync"}
+    else:
+        KeyError("num_stages should be 1 or 2") 
+    # if is_hip_:
+    #     # https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization/workload.html#mi300x-triton-kernel-performance-optimization
+    #     # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+    #     extra_kargs = {
+    #         "waves_per_eu": 1,
+    #         "matrix_instr_nonkdim": 16,
+    #         "kpack": 2
+    #     }
+    #     num_stages = 1
 
     _fwd_grouped_kernel_stage1[grid](
         q,
@@ -540,14 +547,14 @@ def _decode_softmax_reducev_fwd(
     NUM_KV_SPLITS = num_kv_splits
 
     extra_kargs = {}
-    if is_hip_:
-        # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
-        # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
-        extra_kargs = {
-            "waves_per_eu": 4,
-            "matrix_instr_nonkdim": 16,
-            "kpack": 2
-        }
+    # if is_hip_:
+    #     # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
+    #     # https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
+    #     extra_kargs = {
+    #         "waves_per_eu": 4,
+    #         "matrix_instr_nonkdim": 16,
+    #         "kpack": 2
+    #     }
 
     grid = (batch, head_num)
     _fwd_kernel_stage2[grid](
@@ -606,6 +613,7 @@ def decode_attention_fwd_grouped(
     b_seq_len,
     attn_logits,
     num_kv_splits,
+    num_stages,
     sm_scale,
     page_size,
     logit_cap=0.0,
@@ -618,6 +626,7 @@ def decode_attention_fwd_grouped(
         req_to_token,
         b_seq_len,
         num_kv_splits,
+        num_stages,
         sm_scale,
         page_size,
         logit_cap,
@@ -635,6 +644,7 @@ def decode_attention_fwd(
     b_seq_len,
     attn_logits,
     num_kv_splits,
+    num_stages,
     sm_scale,
     page_size=1,
     logit_cap=0.0,
@@ -668,6 +678,7 @@ def decode_attention_fwd(
             b_seq_len,
             attn_logits,
             num_kv_splits,
+            num_stages,
             sm_scale,
             page_size,
             logit_cap,
diff --git a/vllm/compilation/cuda_piecewise_backend.py b/vllm/compilation/cuda_piecewise_backend.py
index 993def49a..8c49ea6cc 100644
--- a/vllm/compilation/cuda_piecewise_backend.py
+++ b/vllm/compilation/cuda_piecewise_backend.py
@@ -14,6 +14,7 @@ from vllm.compilation.backends import VllmBackend
 from vllm.compilation.counter import compilation_counter
 from vllm.compilation.monitor import end_monitoring_torch_compile
 from vllm.config import VllmConfig
+from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.utils import weak_ref_tensors
 
@@ -137,7 +138,10 @@ class CUDAPiecewiseBackend:
             if self.is_last_graph and not self.to_be_compiled_sizes:
                 self.check_for_ending_compilation()
 
-        if not entry.use_cudagraph:
+        # Skip CUDA graphs if this entry doesn't use them OR
+        # if we're supposed to skip them globally
+        skip_cuda_graphs = get_forward_context().skip_cuda_graphs
+        if not entry.use_cudagraph or skip_cuda_graphs:
             return entry.runnable(*args)
 
         if entry.cudagraph is None:
diff --git a/vllm/config.py b/vllm/config.py
index 3fbb6015f..aa13f229e 100644
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -331,7 +331,7 @@ class ModelConfig:
     """Whether to disable sliding window. If True, we will disable the sliding
     window functionality of the model, capping to sliding window size. If the
     model does not support sliding window, this argument is ignored."""
-    disable_cascade_attn: bool = False
+    disable_cascade_attn: bool = True
     """Disable cascade attention for V1. While cascade attention does not
     change the mathematical correctness, disabling it could be useful for
     preventing potential numerical issues. Note that even if this is set to
@@ -899,13 +899,13 @@ class ModelConfig:
                     # Raise error if the override is not custom (custom would
                     # be in QUANTIZATION_METHODS but not QuantizationMethods)
                     # and hasn't been added to the overrides list.
-                    if (name in get_args(QuantizationMethods)
-                            and name not in overrides):
-                        raise ValueError(
-                            f"Quantization method {name} is an override but "
-                            "is has not been added to the `overrides` list "
-                            "above. This is necessary to ensure that the "
-                            "overrides are checked in order of preference.")
+                    # if (name in get_args(QuantizationMethods)
+                    #         and name not in overrides):
+                    #     raise ValueError(
+                    #         f"Quantization method {name} is an override but "
+                    #         "is has not been added to the `overrides` list "
+                    #         "above. This is necessary to ensure that the "
+                    #         "overrides are checked in order of preference.")
                     quant_method = quantization_override
                     self.quantization = quantization_override
                     break
@@ -1429,10 +1429,17 @@ class ModelConfig:
             spec_target_max_model_len=self.spec_target_max_model_len,
             encoder_config=self.encoder_config)
 
-        tokenizer_config = try_get_tokenizer_config(
-            self.tokenizer,
-            trust_remote_code=self.trust_remote_code,
-            revision=self.tokenizer_revision)
+        # For pooling models, the tokenizer's `model_max_length` is often a
+        # reliable source for the maximum sequence length. However, for
+        # generative models, this can be incorrect and unduly limit the
+        # context window (e.g., DeepSeek-R1). Therefore, we only consider
+        # tokenizer_config for pooling models.
+        tokenizer_config = None
+        if self.runner_type == "pooling":
+            tokenizer_config = try_get_tokenizer_config(
+                self.tokenizer,
+                trust_remote_code=self.trust_remote_code,
+                revision=self.tokenizer_revision)
 
         if tokenizer_config is None:
             return max_model_len
@@ -1765,7 +1772,7 @@ class ParallelConfig:
     sequentially in multiple batches. To avoid RAM OOM when using tensor
     parallel and large models."""
 
-    disable_custom_all_reduce: bool = False
+    disable_custom_all_reduce: bool = True
     """Disable the custom all-reduce kernel and fall back to NCCL."""
 
     tokenizer_pool_config: Optional[TokenizerPoolConfig] = None
@@ -2040,8 +2047,8 @@ class SchedulerConfig:
     NOTE: This will be replaced by speculative config in the future; it is
     present to enable correctness tests until then."""
 
-    cuda_graph_sizes: list[int] = field(default_factory=lambda: [512])
-    """Cuda graph capture sizes, default is 512.
+    cuda_graph_sizes: list[int] = field(default_factory=lambda: [256])
+    """Cuda graph capture sizes, default is 256.
     1. if one value is provided, then the capture list would follow the
     pattern: [1, 2, 4] + [i for i in range(8, cuda_graph_sizes + 1, 8)]
     2. more than one value (e.g. 1 2 128) is provided, then the capture list
@@ -4256,6 +4263,7 @@ class VllmConfig:
         from vllm import __version__
         vllm_factors.append(__version__)
         vllm_factors.append(envs.VLLM_USE_V1)
+        vllm_factors.append(envs.MACA_VLLM_USE_TN_2_NN)
         if self.model_config:
             vllm_factors.append(self.model_config.compute_hash())
         else:
@@ -4353,12 +4361,14 @@ class VllmConfig:
 
             if capability_tuple is not None:
                 capability = capability_tuple.to_int()
+                """
                 if capability < quant_config.get_min_capability():
                     raise ValueError(
                         f"The quantization method {model_config.quantization} "
                         "is not supported for the current GPU. Minimum "
                         f"capability: {quant_config.get_min_capability()}. "
                         f"Current capability: {capability}.")
+                """
             supported_dtypes = quant_config.get_supported_act_dtypes()
             if model_config.dtype not in supported_dtypes:
                 raise ValueError(
diff --git a/vllm/distributed/device_communicators/cuda_wrapper.py b/vllm/distributed/device_communicators/cuda_wrapper.py
index 2c38e8ed2..0052ba082 100644
--- a/vllm/distributed/device_communicators/cuda_wrapper.py
+++ b/vllm/distributed/device_communicators/cuda_wrapper.py
@@ -65,33 +65,33 @@ def find_loaded_library(lib_name) -> Optional[str]:
 class CudaRTLibrary:
     exported_functions = [
         # ​cudaError_t cudaSetDevice ( int  device )
-        Function("cudaSetDevice", cudaError_t, [ctypes.c_int]),
+        Function("mcSetDevice", cudaError_t, [ctypes.c_int]),
         # cudaError_t 	cudaDeviceSynchronize ( void )
-        Function("cudaDeviceSynchronize", cudaError_t, []),
+        Function("mcDeviceSynchronize", cudaError_t, []),
         # ​cudaError_t cudaDeviceReset ( void )
-        Function("cudaDeviceReset", cudaError_t, []),
+        Function("mcDeviceReset", cudaError_t, []),
 
         # const char* 	cudaGetErrorString ( cudaError_t error )
-        Function("cudaGetErrorString", ctypes.c_char_p, [cudaError_t]),
+        Function("mcGetErrorString", ctypes.c_char_p, [cudaError_t]),
 
         # ​cudaError_t 	cudaMalloc ( void** devPtr, size_t size )
-        Function("cudaMalloc", cudaError_t,
+        Function("mcMalloc", cudaError_t,
                  [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]),
         # ​cudaError_t 	cudaFree ( void* devPtr )
-        Function("cudaFree", cudaError_t, [ctypes.c_void_p]),
+        Function("mcFree", cudaError_t, [ctypes.c_void_p]),
         # ​cudaError_t cudaMemset ( void* devPtr, int  value, size_t count )
-        Function("cudaMemset", cudaError_t,
+        Function("mcMemset", cudaError_t,
                  [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]),
         # ​cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa
-        Function("cudaMemcpy", cudaError_t, [
+        Function("mcMemcpy", cudaError_t, [
             ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, cudaMemcpyKind
         ]),
 
         # cudaError_t cudaIpcGetMemHandle ( cudaIpcMemHandle_t* handle, void* devPtr ) # noqa
-        Function("cudaIpcGetMemHandle", cudaError_t,
+        Function("mcIpcGetMemHandle", cudaError_t,
                  [ctypes.POINTER(cudaIpcMemHandle_t), ctypes.c_void_p]),
         # ​cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int  flags ) # noqa
-        Function("cudaIpcOpenMemHandle", cudaError_t, [
+        Function("mcIpcOpenMemHandle", cudaError_t, [
             ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint
         ]),
     ]
@@ -106,7 +106,7 @@ class CudaRTLibrary:
 
     def __init__(self, so_file: Optional[str] = None):
         if so_file is None:
-            so_file = find_loaded_library("libcudart")
+            so_file = find_loaded_library("libmcruntime")
             if so_file is None:
                 so_file = envs.VLLM_CUDART_SO_PATH  # fallback to env var
             assert so_file is not None, \
@@ -135,39 +135,39 @@ class CudaRTLibrary:
             raise RuntimeError(f"CUDART error: {error_str}")
 
     def cudaGetErrorString(self, error: cudaError_t) -> str:
-        return self.funcs["cudaGetErrorString"](error).decode("utf-8")
+        return self.funcs["mcGetErrorString"](error).decode("utf-8")
 
     def cudaSetDevice(self, device: int) -> None:
-        self.CUDART_CHECK(self.funcs["cudaSetDevice"](device))
+        self.CUDART_CHECK(self.funcs["mcSetDevice"](device))
 
     def cudaDeviceSynchronize(self) -> None:
-        self.CUDART_CHECK(self.funcs["cudaDeviceSynchronize"]())
+        self.CUDART_CHECK(self.funcs["mcDeviceSynchronize"]())
 
     def cudaDeviceReset(self) -> None:
-        self.CUDART_CHECK(self.funcs["cudaDeviceReset"]())
+        self.CUDART_CHECK(self.funcs["mcDeviceReset"]())
 
     def cudaMalloc(self, size: int) -> ctypes.c_void_p:
         devPtr = ctypes.c_void_p()
-        self.CUDART_CHECK(self.funcs["cudaMalloc"](ctypes.byref(devPtr), size))
+        self.CUDART_CHECK(self.funcs["mcMalloc"](ctypes.byref(devPtr), size))
         return devPtr
 
     def cudaFree(self, devPtr: ctypes.c_void_p) -> None:
-        self.CUDART_CHECK(self.funcs["cudaFree"](devPtr))
+        self.CUDART_CHECK(self.funcs["mcFree"](devPtr))
 
     def cudaMemset(self, devPtr: ctypes.c_void_p, value: int,
                    count: int) -> None:
-        self.CUDART_CHECK(self.funcs["cudaMemset"](devPtr, value, count))
+        self.CUDART_CHECK(self.funcs["mcMemset"](devPtr, value, count))
 
     def cudaMemcpy(self, dst: ctypes.c_void_p, src: ctypes.c_void_p,
                    count: int) -> None:
         cudaMemcpyDefault = 4
         kind = cudaMemcpyDefault
-        self.CUDART_CHECK(self.funcs["cudaMemcpy"](dst, src, count, kind))
+        self.CUDART_CHECK(self.funcs["mcMemcpy"](dst, src, count, kind))
 
     def cudaIpcGetMemHandle(self,
                             devPtr: ctypes.c_void_p) -> cudaIpcMemHandle_t:
         handle = cudaIpcMemHandle_t()
-        self.CUDART_CHECK(self.funcs["cudaIpcGetMemHandle"](
+        self.CUDART_CHECK(self.funcs["mcIpcGetMemHandle"](
             ctypes.byref(handle), devPtr))
         return handle
 
@@ -175,6 +175,6 @@ class CudaRTLibrary:
                              handle: cudaIpcMemHandle_t) -> ctypes.c_void_p:
         cudaIpcMemLazyEnablePeerAccess = 1
         devPtr = ctypes.c_void_p()
-        self.CUDART_CHECK(self.funcs["cudaIpcOpenMemHandle"](
+        self.CUDART_CHECK(self.funcs["mcIpcOpenMemHandle"](
             ctypes.byref(devPtr), handle, cudaIpcMemLazyEnablePeerAccess))
         return devPtr
diff --git a/vllm/distributed/device_communicators/pynccl_wrapper.py b/vllm/distributed/device_communicators/pynccl_wrapper.py
index 04a4d0147..718d7ea82 100644
--- a/vllm/distributed/device_communicators/pynccl_wrapper.py
+++ b/vllm/distributed/device_communicators/pynccl_wrapper.py
@@ -128,18 +128,18 @@ class Function:
 class NCCLLibrary:
     exported_functions = [
         # const char* ncclGetErrorString(ncclResult_t result)
-        Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
+        Function("mcclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
         # ncclResult_t  ncclGetVersion(int *version);
-        Function("ncclGetVersion", ncclResult_t,
+        Function("mcclGetVersion", ncclResult_t,
                  [ctypes.POINTER(ctypes.c_int)]),
         # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
-        Function("ncclGetUniqueId", ncclResult_t,
+        Function("mcclGetUniqueId", ncclResult_t,
                  [ctypes.POINTER(ncclUniqueId)]),
         # ncclResult_t  ncclCommInitRank(
         #   ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
         # note that ncclComm_t is a pointer type, so the first argument
         # is a pointer to a pointer
-        Function("ncclCommInitRank", ncclResult_t, [
+        Function("mcclCommInitRank", ncclResult_t, [
             ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId,
             ctypes.c_int
         ]),
@@ -149,7 +149,7 @@ class NCCLLibrary:
         #   cudaStream_t stream);
         # note that cudaStream_t is a pointer type, so the last argument
         # is a pointer
-        Function("ncclAllReduce", ncclResult_t, [
+        Function("mcclAllReduce", ncclResult_t, [
             buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
             ncclRedOp_t, ncclComm_t, cudaStream_t
         ]),
@@ -160,7 +160,7 @@ class NCCLLibrary:
         #   cudaStream_t stream);
         # note that cudaStream_t is a pointer type, so the last argument
         # is a pointer
-        Function("ncclAllGather", ncclResult_t, [
+        Function("mcclAllGather", ncclResult_t, [
             buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
             ncclComm_t, cudaStream_t
         ]),
@@ -171,7 +171,7 @@ class NCCLLibrary:
         #   cudaStream_t stream);
         # note that cudaStream_t is a pointer type, so the last argument
         # is a pointer
-        Function("ncclReduceScatter", ncclResult_t, [
+        Function("mcclReduceScatter", ncclResult_t, [
             buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
             ncclRedOp_t, ncclComm_t, cudaStream_t
         ]),
@@ -179,7 +179,7 @@ class NCCLLibrary:
         # ncclResult_t  ncclSend(
         #   const void* sendbuff, size_t count, ncclDataType_t datatype,
         #   int dest, ncclComm_t comm, cudaStream_t stream);
-        Function("ncclSend", ncclResult_t, [
+        Function("mcclSend", ncclResult_t, [
             buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int,
             ncclComm_t, cudaStream_t
         ]),
@@ -187,7 +187,7 @@ class NCCLLibrary:
         # ncclResult_t  ncclRecv(
         #   void* recvbuff, size_t count, ncclDataType_t datatype,
         #   int src, ncclComm_t comm, cudaStream_t stream);
-        Function("ncclRecv", ncclResult_t, [
+        Function("mcclRecv", ncclResult_t, [
             buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int,
             ncclComm_t, cudaStream_t
         ]),
@@ -196,7 +196,7 @@ class NCCLLibrary:
         #   const void* sendbuff, void* recvbuff, size_t count,
         #   ncclDataType_t datatype, int root, ncclComm_t comm,
         #   cudaStream_t stream);
-        Function("ncclBroadcast", ncclResult_t, [
+        Function("mcclBroadcast", ncclResult_t, [
             buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
             ctypes.c_int, ncclComm_t, cudaStream_t
         ]),
@@ -206,7 +206,7 @@ class NCCLLibrary:
         # because Python object destruction can happen in random order,
         # it is better not to call it at all.
         # ncclResult_t  ncclCommDestroy(ncclComm_t comm);
-        Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
+        Function("mcclCommDestroy", ncclResult_t, [ncclComm_t]),
     ]
 
     # class attribute to store the mapping from the path to the library
@@ -249,16 +249,16 @@ class NCCLLibrary:
         self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
 
     def ncclGetErrorString(self, result: ncclResult_t) -> str:
-        return self._funcs["ncclGetErrorString"](result).decode("utf-8")
+        return self._funcs["mcclGetErrorString"](result).decode("utf-8")
 
     def NCCL_CHECK(self, result: ncclResult_t) -> None:
         if result != 0:
             error_str = self.ncclGetErrorString(result)
-            raise RuntimeError(f"NCCL error: {error_str}")
+            raise RuntimeError(f"MCCL error: {error_str}")
 
     def ncclGetVersion(self) -> str:
         version = ctypes.c_int()
-        self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
+        self.NCCL_CHECK(self._funcs["mcclGetVersion"](ctypes.byref(version)))
         version_str = str(version.value)
         # something like 21903 --> "2.19.3"
         major = version_str[0].lstrip("0")
@@ -268,14 +268,14 @@ class NCCLLibrary:
 
     def ncclGetUniqueId(self) -> ncclUniqueId:
         unique_id = ncclUniqueId()
-        self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](
+        self.NCCL_CHECK(self._funcs["mcclGetUniqueId"](
             ctypes.byref(unique_id)))
         return unique_id
 
     def ncclCommInitRank(self, world_size: int, unique_id: ncclUniqueId,
                          rank: int) -> ncclComm_t:
         comm = ncclComm_t()
-        self.NCCL_CHECK(self._funcs["ncclCommInitRank"](ctypes.byref(comm),
+        self.NCCL_CHECK(self._funcs["mcclCommInitRank"](ctypes.byref(comm),
                                                         world_size, unique_id,
                                                         rank))
         return comm
@@ -288,7 +288,7 @@ class NCCLLibrary:
         # both are aliases of `ctypes.c_int`
         # when we pass int to a function, it will be converted to `ctypes.c_int`
         # by ctypes automatically
-        self.NCCL_CHECK(self._funcs["ncclAllReduce"](sendbuff, recvbuff, count,
+        self.NCCL_CHECK(self._funcs["mcclAllReduce"](sendbuff, recvbuff, count,
                                                      datatype, op, comm,
                                                      stream))
 
@@ -300,7 +300,7 @@ class NCCLLibrary:
         # both are aliases of `ctypes.c_int`
         # when we pass int to a function, it will be converted to `ctypes.c_int`
         # by ctypes automatically
-        self.NCCL_CHECK(self._funcs["ncclReduceScatter"](sendbuff, recvbuff,
+        self.NCCL_CHECK(self._funcs["mcclReduceScatter"](sendbuff, recvbuff,
                                                          count, datatype, op,
                                                          comm, stream))
 
@@ -311,28 +311,28 @@ class NCCLLibrary:
         # which is an aliases of `ctypes.c_int`
         # when we pass int to a function, it will be converted to `ctypes.c_int`
         # by ctypes automatically
-        self.NCCL_CHECK(self._funcs["ncclAllGather"](sendbuff, recvbuff, count,
+        self.NCCL_CHECK(self._funcs["mcclAllGather"](sendbuff, recvbuff, count,
                                                      datatype, comm, stream))
 
     def ncclSend(self, sendbuff: buffer_type, count: int, datatype: int,
                  dest: int, comm: ncclComm_t, stream: cudaStream_t) -> None:
-        self.NCCL_CHECK(self._funcs["ncclSend"](sendbuff, count, datatype,
+        self.NCCL_CHECK(self._funcs["mcclSend"](sendbuff, count, datatype,
                                                 dest, comm, stream))
 
     def ncclRecv(self, recvbuff: buffer_type, count: int, datatype: int,
                  src: int, comm: ncclComm_t, stream: cudaStream_t) -> None:
-        self.NCCL_CHECK(self._funcs["ncclRecv"](recvbuff, count, datatype, src,
+        self.NCCL_CHECK(self._funcs["mcclRecv"](recvbuff, count, datatype, src,
                                                 comm, stream))
 
     def ncclBroadcast(self, sendbuff: buffer_type, recvbuff: buffer_type,
                       count: int, datatype: int, root: int, comm: ncclComm_t,
                       stream: cudaStream_t) -> None:
-        self.NCCL_CHECK(self._funcs["ncclBroadcast"](sendbuff, recvbuff, count,
+        self.NCCL_CHECK(self._funcs["mcclBroadcast"](sendbuff, recvbuff, count,
                                                      datatype, root, comm,
                                                      stream))
 
     def ncclCommDestroy(self, comm: ncclComm_t) -> None:
-        self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
+        self.NCCL_CHECK(self._funcs["mcclCommDestroy"](comm))
 
 
 __all__ = [
diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
index 10f87c49b..c596f8875 100644
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -1203,13 +1203,14 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     empty_cache = current_platform.empty_cache
     if empty_cache is not None:
         empty_cache()
+    """
     try:
         if not current_platform.is_cpu():
             torch._C._host_emptyCache()
     except AttributeError:
         logger.warning(
             "torch._C._host_emptyCache() only available in Pytorch >=2.5")
-
+    """
 
 def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
                         source_rank: int = 0) -> list[bool]:
diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py
index 6e3cb18fc..a5a8a812c 100644
--- a/vllm/entrypoints/llm.py
+++ b/vllm/entrypoints/llm.py
@@ -173,13 +173,14 @@ class LLM:
         cpu_offload_gb: float = 0,
         enforce_eager: bool = False,
         max_seq_len_to_capture: int = 8192,
-        disable_custom_all_reduce: bool = False,
+        disable_custom_all_reduce: bool = True,
         disable_async_output_proc: bool = False,
         hf_token: Optional[Union[bool, str]] = None,
         hf_overrides: Optional[HfOverrides] = None,
         mm_processor_kwargs: Optional[dict[str, Any]] = None,
         override_pooler_config: Optional[PoolerConfig] = None,
-        compilation_config: Optional[Union[int, dict[str, Any]]] = None,
+        compilation_config: Optional[Union[int, dict[str, Any],
+                                            CompilationConfig]] = None,
         **kwargs,
     ) -> None:
         """LLM constructor."""
diff --git a/vllm/env_override.py b/vllm/env_override.py
index 2bede4963..a127a5587 100644
--- a/vllm/env_override.py
+++ b/vllm/env_override.py
@@ -38,4 +38,4 @@ os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1'
 # see https://github.com/vllm-project/vllm/issues/10480
 os.environ['TORCHINDUCTOR_COMPILE_THREADS'] = '1'
 # see https://github.com/vllm-project/vllm/issues/10619
-torch._inductor.config.compile_threads = 1
+# torch._inductor.config.compile_threads = 1
diff --git a/vllm/envs.py b/vllm/envs.py
index 80c5f289b..1d351238b 100644
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -127,7 +127,7 @@ if TYPE_CHECKING:
     VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
     VLLM_SLEEP_WHEN_IDLE: bool = False
     VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
-
+    MACA_VLLM_USE_TN_2_NN: bool = True
 
 def get_default_cache_root():
     return os.getenv(
@@ -870,6 +870,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # processes via zmq.
     "VLLM_MQ_MAX_CHUNK_BYTES_MB":
     lambda: int(os.getenv("VLLM_MQ_MAX_CHUNK_BYTES_MB", "16")),
+    
+    # if set, enable loading weight by transpose
+    "MACA_VLLM_USE_TN_2_NN":
+    lambda: os.environ.get("MACA_VLLM_USE_TN_2_NN", "1") == "1",
 }
 
 # --8<-- [end:env-vars-definition]
diff --git a/vllm/executor/ray_distributed_executor.py b/vllm/executor/ray_distributed_executor.py
index bdc2b1f4c..174fbc26d 100644
--- a/vllm/executor/ray_distributed_executor.py
+++ b/vllm/executor/ray_distributed_executor.py
@@ -557,8 +557,16 @@ class RayDistributedExecutor(DistributedExecutorBase):
     def _compiled_ray_dag(self, enable_asyncio: bool):
         assert self.parallel_config.use_ray
         self._check_ray_cgraph_installation()
+        # Enlarge the default value of "RAY_CGRAPH_get_timeout" to 300 seconds
+        # (it is 10 seconds by default). This is a Ray environment variable to
+        # control the timeout of getting result from a compiled graph execution,
+        # i.e., the distributed execution that includes model forward runs and
+        # intermediate tensor communications, in the case of vllm.
+        os.environ.setdefault("RAY_CGRAPH_get_timeout", "300")  # noqa: SIM112
         from ray.dag import InputNode, MultiOutputNode
 
+        logger.info("RAY_CGRAPH_get_timeout is set to %s",
+                    os.environ["RAY_CGRAPH_get_timeout"])  # noqa: SIM112
         logger.info("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE = %s",
                     envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE)
         logger.info("VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM = %s",
@@ -570,14 +578,6 @@ class RayDistributedExecutor(DistributedExecutorBase):
                 "Invalid value for VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: "
                 f"{channel_type}. Valid values are: 'auto', 'nccl', or 'shm'.")
 
-        # Enlarge the default value of "RAY_CGRAPH_get_timeout" to 300 seconds
-        # (it is 10 seconds by default). This is a Ray environment variable to
-        # control the timeout of getting result from a compiled graph execution,
-        # i.e., the distributed execution that includes model forward runs and
-        # intermediate tensor communications, in the case of vllm.
-        os.environ.setdefault("RAY_CGRAPH_get_timeout", "300")  # noqa: SIM112
-        logger.info("RAY_CGRAPH_get_timeout is set to %s",
-                    os.environ["RAY_CGRAPH_get_timeout"])  # noqa: SIM112
 
         with InputNode() as input_data:
             # Example DAG: PP=2, TP=4
diff --git a/vllm/forward_context.py b/vllm/forward_context.py
index f3b0518a4..dd55b19fe 100644
--- a/vllm/forward_context.py
+++ b/vllm/forward_context.py
@@ -94,6 +94,7 @@ class ForwardContext:
     virtual_engine: int  # set dynamically for each forward pass
     # set dynamically for each forward pass
     dp_metadata: Optional[DPMetadata] = None
+    skip_cuda_graphs: bool = False
 
 
 _forward_context: Optional[ForwardContext] = None
@@ -108,11 +109,14 @@ def get_forward_context() -> ForwardContext:
 
 
 @contextmanager
-def set_forward_context(attn_metadata: Any,
-                        vllm_config: VllmConfig,
-                        virtual_engine: int = 0,
-                        num_tokens: Optional[int] = None,
-                        num_tokens_across_dp: Optional[torch.Tensor] = None):
+def set_forward_context(
+    attn_metadata: Any,
+    vllm_config: VllmConfig,
+    virtual_engine: int = 0,
+    num_tokens: Optional[int] = None,
+    num_tokens_across_dp: Optional[torch.Tensor] = None,
+    skip_cuda_graphs: bool = False,
+):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
     Here we can inject common logic for every model forward pass.
@@ -135,7 +139,9 @@ def set_forward_context(attn_metadata: Any,
         static_forward_context,
         virtual_engine=virtual_engine,
         attn_metadata=attn_metadata,
-        dp_metadata=dp_metadata)
+        dp_metadata=dp_metadata,
+        skip_cuda_graphs=skip_cuda_graphs,
+    )
 
     try:
         yield
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index ba1498e65..9d74c6857 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -4,9 +4,10 @@
 import functools
 import json
 import os
-from typing import Any, Callable, Optional
+from typing import Any, List, Callable, Optional
 
 import torch
+import math
 
 import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -82,6 +83,8 @@ def fused_moe_kernel_gptq_awq(
         BLOCK_SIZE_N: tl.constexpr,
         BLOCK_SIZE_K: tl.constexpr,
         GROUP_SIZE_M: tl.constexpr,
+        SPLIT_K: tl.constexpr,
+        ACCF32: tl.constexpr,
         MUL_ROUTED_WEIGHT: tl.constexpr,
         top_k: tl.constexpr,
         compute_type: tl.constexpr,
@@ -251,6 +254,21 @@ def fused_moe_kernel_gptq_awq(
     tl.store(c_ptrs, accumulator, mask=c_mask)
 
 
+@triton.heuristics(
+    {
+        "UPGRADE": lambda args: math.ceil((args["EM"] * args["N"]) / (args["BLOCK_SIZE_M"] * args["BLOCK_SIZE_N"])).bit_length() > 31,
+    }
+)
+@triton.heuristics(
+    {
+        "UPGRADE_A_OFFS": lambda args: (args["num_valid_tokens"] // args["top_k"] * args["stride_am"] + args["BLOCK_SIZE_K"] * args["stride_ak"]).bit_length() > 31,
+    }
+)
+@triton.heuristics(
+    {
+        "UPGRADE_B_OFFS": lambda args: ((args["E"]-1) * args["stride_be"] + (args["N"]-1) * args["stride_bn"]+(args["K"]-1) * args["stride_bk"]).bit_length() > 31,
+    }
+)
 @triton.jit
 def fused_moe_kernel(
     # Pointers to matrices
@@ -264,6 +282,7 @@ def fused_moe_kernel(
     expert_ids_ptr,
     num_tokens_post_padded_ptr,
     # Matrix dimensions
+    E,          # B.shape[0]
     N,
     K,
     EM,
@@ -292,13 +311,20 @@ def fused_moe_kernel(
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
     GROUP_SIZE_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    ACCF32: tl.constexpr,
     MUL_ROUTED_WEIGHT: tl.constexpr,
     top_k: tl.constexpr,
+    # experts_num: tl.constexpr,
     compute_type: tl.constexpr,
     use_fp8_w8a8: tl.constexpr,
     use_int8_w8a8: tl.constexpr,
     use_int8_w8a16: tl.constexpr,
     per_channel_quant: tl.constexpr,
+    UPGRADE: tl.constexpr,
+    UPGRADE_A_OFFS: tl.constexpr,
+    UPGRADE_B_OFFS: tl.constexpr,
+    FAST_F32_TO_BF16: tl.constexpr
 ):
     """
     Implements the fused computation for a Mixture of Experts (MOE) using
@@ -329,7 +355,12 @@ def fused_moe_kernel(
     # -----------------------------------------------------------
     # Map program ids `pid` to the block of C it should compute.
     # This is done in a grouped ordering to promote L2 data reuse.
-    pid = tl.program_id(axis=0)
+    if UPGRADE:
+        pid = tl.program_id(axis=0).to(tl.int64)
+        pid_z = tl.program_id(axis=1).to(tl.int64)
+    else:
+        pid = tl.program_id(axis=0)
+        pid_z = tl.program_id(axis=1)
     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
@@ -353,7 +384,14 @@ def fused_moe_kernel(
     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 UPGRADE_B_OFFS:
+        off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+    else:
+        off_experts = tl.load(expert_ids_ptr + pid_m)
+        
+    if UPGRADE_A_OFFS:
+        offs_token = offs_token.to(tl.int64)
+
     if off_experts == -1:
         # -----------------------------------------------------------
         # Write back zeros to the output when the expert is not
@@ -363,9 +401,16 @@ def fused_moe_kernel(
                               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)
+    if UPGRADE_B_OFFS:
+        offs_bn = (pid_n * BLOCK_SIZE_N +
+                   tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    else:
+        offs_bn = (pid_n * BLOCK_SIZE_N +
+                   tl.arange(0, BLOCK_SIZE_N)) % N
+        
+    # offs_k = tl.arange(0, BLOCK_SIZE_K)
+    offs_k = pid_z * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+
     a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am +
                       offs_k[None, :] * stride_ak)
 
@@ -376,7 +421,13 @@ def fused_moe_kernel(
             None, :] * stride_bsn
         b_scale = tl.load(b_scale_ptrs)
 
-    if use_fp8_w8a8 or use_int8_w8a8:
+    if use_int8_w8a8:  
+        a_scale = tl.load(a_scale_ptr+(offs_token[:, None] // top_k * stride_asm),mask=token_mask[:, None],other=0.0)
+        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:
         # block-wise
         if group_k > 0 and group_n > 0:
             a_scale_ptrs = a_scale_ptr + (offs_token // top_k) * stride_asm
@@ -402,23 +453,28 @@ def fused_moe_kernel(
     # 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)):
+    # accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.int32 if use_int8_w8a8 else tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_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),
+                    (offs_k[None, :] < K - k * BLOCK_SIZE_K * SPLIT_K),
                     other=0.0)
         b = tl.load(b_ptrs,
-                    mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
+                    mask=offs_k[:, None] < K - k * BLOCK_SIZE_K * SPLIT_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 or use_int8_w8a8:
+        elif use_int8_w8a8:
+            a = a.to(tl.int8)
+            accumulator += tl.dot(a, b,out_dtype=accumulator.dtype)
+        elif use_fp8_w8a8:
             if group_k > 0 and group_n > 0:
-                k_start = k * BLOCK_SIZE_K
+                k_start = k * BLOCK_SIZE_K * SPLIT_K
                 offs_ks = k_start // group_k
                 a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask,
                                   mask=token_mask,
@@ -436,8 +492,8 @@ def fused_moe_kernel(
         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
+        a_ptrs += BLOCK_SIZE_K * stride_ak * SPLIT_K
+        b_ptrs += BLOCK_SIZE_K * stride_bk * SPLIT_K
 
     if MUL_ROUTED_WEIGHT:
         moe_weight = tl.load(topk_weights_ptr + offs_token,
@@ -446,7 +502,15 @@ def fused_moe_kernel(
         accumulator = accumulator * moe_weight[:, None]
     if use_int8_w8a16:
         accumulator = (accumulator * b_scale).to(compute_type)
-    elif use_fp8_w8a8 or use_int8_w8a8:
+    elif use_int8_w8a8:
+        accumulator = accumulator.to(tl.float32)
+        accumulator = (accumulator * a_scale * b_scale)
+        if not ACCF32:
+            if FAST_F32_TO_BF16:
+                accumulator = accumulator.to(compute_type, "rtne_no_nan")
+            else:
+                accumulator = accumulator.to(compute_type)
+    elif use_fp8_w8a8:
         if group_k > 0 and group_n > 0:
             accumulator = accumulator.to(compute_type)
         else:
@@ -459,8 +523,10 @@ def fused_moe_kernel(
     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)
-
+    if SPLIT_K == 1:
+        tl.store(c_ptrs, accumulator, mask=c_mask)
+    else:
+        tl.atomic_add(c_ptrs, accumulator, mask=c_mask)
 
 def invoke_fused_moe_kernel(A: torch.Tensor,
                             B: torch.Tensor,
@@ -480,6 +546,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
                             use_int8_w8a8: bool,
                             use_int8_w8a16: bool,
                             use_int4_w4a16: bool,
+                            orig_acc_dtype: torch.dtype,
                             per_channel_quant: bool,
                             block_shape: Optional[list[int]] = None) -> None:
     assert topk_weights is not None or not mul_routed_weight
@@ -512,7 +579,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
         EM = min(sorted_token_ids.shape[0],
                  A.shape[0] * top_k * config['BLOCK_SIZE_M'])
     grid = lambda META: (triton.cdiv(EM, META['BLOCK_SIZE_M']) * triton.cdiv(
-        B.shape[1], META['BLOCK_SIZE_N']), )
+        B.shape[1], META['BLOCK_SIZE_N']), META['SPLIT_K'])
 
     if (use_int8_w8a16 or use_int4_w4a16) and \
             block_shape is not None and block_shape[1] > 0:
@@ -524,26 +591,27 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
             group_size=block_shape[1],
             num_experts=B.shape[0],
             bit=4 if use_int4_w4a16 else 8)
-        config = config.copy()
-        config.update(
-            get_moe_wna16_block_config(config=config,
-                                       use_moe_wna16_cuda=use_moe_wna16_cuda,
-                                       num_valid_tokens=num_tokens,
-                                       size_k=A.shape[1],
-                                       size_n=B.shape[1],
-                                       num_experts=B.shape[1],
-                                       group_size=block_shape[1],
-                                       real_top_k=top_k,
-                                       block_size_m=config["BLOCK_SIZE_M"]))
-
-        if use_moe_wna16_cuda:
+        # TODO: missing config for BLOCK_SIZE_K
+        # config = config.copy()
+        # config.update(
+        #     get_moe_wna16_block_config(config=config,
+        #                                use_moe_wna16_cuda=use_moe_wna16_cuda,
+        #                                num_valid_tokens=num_tokens,
+        #                                size_k=A.shape[1],
+        #                                size_n=B.shape[1],
+        #                                num_experts=B.shape[1],
+        #                                group_size=block_shape[1],
+        #                                real_top_k=top_k,
+        #                                block_size_m=config["BLOCK_SIZE_M"]))
+
+        if False and use_moe_wna16_cuda:
             bit = 4 if use_int4_w4a16 else 8
             ops.moe_wna16_gemm(A, C, B, B_scale, B_zp,
                                topk_weights if mul_routed_weight else None,
                                sorted_token_ids, expert_ids,
                                num_tokens_post_padded, top_k,
                                config["BLOCK_SIZE_M"], config["BLOCK_SIZE_N"],
-                               config["BLOCK_SIZE_K"], bit)
+                               config["BLOCK_SIZE_K"] * config["SPLIT_K"], bit)
             return
 
         fused_moe_kernel_gptq_awq[grid](
@@ -573,7 +641,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
             B_zp.stride(0) if B_zp is not None else 0,
             B_zp.stride(2) if B_zp is not None else 0,
             B_zp.stride(1) if B_zp is not None else 0,
-            block_k_diviable=A.shape[1] % config["BLOCK_SIZE_K"] == 0,
+            block_k_diviable=A.shape[1] % config["BLOCK_SIZE_K"] * config["SPLIT_K"] == 0,
             group_size=block_shape[1],
             MUL_ROUTED_WEIGHT=mul_routed_weight,
             top_k=top_k,
@@ -599,6 +667,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
             sorted_token_ids,
             expert_ids,
             num_tokens_post_padded,
+            B.shape[0],
             B.shape[1],
             B.shape[2],
             EM,
@@ -624,14 +693,18 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
             0 if block_shape is None else block_shape[1],
             MUL_ROUTED_WEIGHT=mul_routed_weight,
             top_k=top_k,
+            # experts_num=expert_ids.shape[0],
             compute_type=compute_type,
             use_fp8_w8a8=use_fp8_w8a8,
             use_int8_w8a8=use_int8_w8a8,
             use_int8_w8a16=use_int8_w8a16,
             per_channel_quant=per_channel_quant,
             BLOCK_SIZE_K=BLOCK_SIZE_K,
+            FAST_F32_TO_BF16 = True,
             **config,
         )
+    if config["ACCF32"]:
+       C = C.to(orig_acc_dtype)
 
 
 # Adapted from: https://github.com/sgl-project/sglang/pull/2628
@@ -639,7 +712,8 @@ def get_config_file_name(E: int,
                          N: int,
                          dtype: Optional[str],
                          block_shape: Optional[list[int]] = None) -> str:
-    device_name = current_platform.get_device_name().replace(" ", "_")
+    # device_name = current_platform.get_device_name().replace(" ", "_")
+    device_name = "Device_4000"
     dtype_selector = "" if not dtype else f",dtype={dtype}"
     block_shape_selector = ("" if not block_shape or not all(block_shape) else
                             f",block_shape={block_shape}").replace(" ", "")
@@ -654,6 +728,7 @@ def get_moe_configs(
     dtype: Optional[str],
     block_n: Optional[int] = None,
     block_k: Optional[int] = None,
+    H: int = 0,
 ) -> Optional[dict[int, Any]]:
     """
     Return optimized configurations for the fused MoE kernel.
@@ -668,9 +743,17 @@ def get_moe_configs(
     # directory
     block_shape = [block_n, block_k] if block_n and block_k else None
     json_file_name = get_config_file_name(E, N, dtype, block_shape)
+    json_file_name_new = f"H={H},{json_file_name}"
 
     config_file_path = os.path.join(
         os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
+    config_file_path_new = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name_new)
+
+    # First find H, E, N config file
+    if os.path.exists(config_file_path_new):
+        config_file_path = config_file_path_new
+
     if os.path.exists(config_file_path):
         with open(config_file_path) as f:
             logger.info("Using configuration from %s for MoE layer.",
@@ -773,21 +856,22 @@ def get_default_config(
             "num_warps": 4,
             "num_stages": 3 if not current_platform.is_rocm() else 2,
         }
-    elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None:
-        # moe wna16 kernels
-        # only set BLOCK_SIZE_M
-        # BLOCK_SIZE_N and BLOCK_SIZE_K would be set later
-        bit = 4 if dtype == "int4_w4a16" else 8
-        use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk,
-                                                       block_shape[1], E, bit)
-        if use_moe_wna16_cuda:
-            config = {"BLOCK_SIZE_M": min(16, M)}
-        elif M <= 20:
-            config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1}
-        elif M <= 40:
-            config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1}
-        else:
-            config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1}
+    # TODO: missing config for BLOCK_SIZE_K
+    # elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None:
+    #     # moe wna16 kernels
+    #     # only set BLOCK_SIZE_M
+    #     # BLOCK_SIZE_N and BLOCK_SIZE_K would be set later
+    #     bit = 4 if dtype == "int4_w4a16" else 8
+    #     use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk,
+    #                                                    block_shape[1], E, bit)
+    #     if use_moe_wna16_cuda:
+    #         config = {"BLOCK_SIZE_M": min(16, M)}
+    #     elif M <= 20:
+    #         config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1}
+    #     elif M <= 40:
+    #         config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1}
+    #     else:
+    #         config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1}
     elif is_marlin:
         for block_size_m in [8, 16, 32, 48, 64]:
             if M * topk / E / block_size_m < 0.9:
@@ -818,6 +902,7 @@ def try_get_optimal_moe_config(
     M: int,
     is_marlin: bool = False,
     block_shape: Optional[list[int]] = None,
+    H: int = 0,
 ):
     from vllm.model_executor.layers.fused_moe import get_config
     override_config = get_config()
@@ -826,11 +911,12 @@ def try_get_optimal_moe_config(
     else:
         # First try to load optimal config from the file
         E, _, N = w2_shape
-        if dtype == "int4_w4a16":
-            N = N * 2
+        # TODO: why we need N * 2
+        # if dtype == "int4_w4a16":
+        #     N = N * 2
         block_n = block_shape[0] if block_shape else 0
         block_k = block_shape[1] if block_shape else 0
-        configs = get_moe_configs(E, N, dtype, block_n, block_k)
+        configs = get_moe_configs(E, N, dtype, block_n, block_k, H)
 
         if configs:
             # If an optimal configuration map has been found, look up the
@@ -965,10 +1051,13 @@ def grouped_topk(
 def get_config_dtype_str(
         dtype: torch.dtype,
         use_int4_w4a16: Optional[bool] = False,
+        use_int8_w8a8: Optional[bool] = False,
         use_int8_w8a16: Optional[bool] = False,
         use_fp8_w8a8: Optional[bool] = False) -> Optional[str]:
     if use_fp8_w8a8:
         return "fp8_w8a8"
+    elif use_int8_w8a8:
+        return "int8_w8a8"
     elif use_int8_w8a16:
         return "int8_w8a16"
     elif use_int4_w4a16:
@@ -1014,7 +1103,7 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
                           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:
+                          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_w8a8, use_int8_w8a16, use_int4_w4a16,
@@ -1044,7 +1133,7 @@ def inplace_fused_experts_fake(
         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:
+        block_shape: Optional[List[int]] = None) -> None:
     pass
 
 
@@ -1078,7 +1167,7 @@ def outplace_fused_experts(
         w2_zp: Optional[torch.Tensor] = None,
         a1_scale: Optional[torch.Tensor] = None,
         a2_scale: Optional[torch.Tensor] = None,
-        block_shape: Optional[list[int]] = None) -> torch.Tensor:
+        block_shape: Optional[List[int]] = None) -> torch.Tensor:
     return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,
                               False, activation, apply_router_weight_on_input,
                               use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,
@@ -1108,7 +1197,7 @@ def outplace_fused_experts_fake(
         w2_zp: Optional[torch.Tensor] = None,
         a1_scale: Optional[torch.Tensor] = None,
         a2_scale: Optional[torch.Tensor] = None,
-        block_shape: Optional[list[int]] = None) -> torch.Tensor:
+        block_shape: Optional[List[int]] = None) -> torch.Tensor:
     return torch.empty_like(hidden_states)
 
 
@@ -1247,6 +1336,7 @@ def fused_experts_impl(
         torch.float32, torch.float16, torch.bfloat16
     ]
 
+    H = hidden_states.shape[-1]
     num_tokens = hidden_states.shape[0]
     E, N, _ = w1.shape
     K = w2.shape[1]
@@ -1258,6 +1348,7 @@ def fused_experts_impl(
     CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
     M = min(num_tokens, CHUNK_SIZE)
     config_dtype = get_config_dtype_str(use_fp8_w8a8=use_fp8_w8a8,
+                                        use_int8_w8a8=use_int8_w8a8,
                                         use_int8_w8a16=use_int8_w8a16,
                                         use_int4_w4a16=use_int4_w4a16,
                                         dtype=hidden_states.dtype)
@@ -1274,17 +1365,52 @@ def fused_experts_impl(
         top_k_num,
         config_dtype,
         block_shape=block_shape,
+        H=H,
     )
 
     config = get_config_func(M)
 
     # 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)
+    stage1_config = config["stage1"] if "stage1" in config else config
+    stage2_config = config["stage2"] if "stage2" in config else config
+    
+    if 'ACCF32' not in stage1_config:
+        stage1_config['ACCF32'] = False
+    if 'ACCF32' not in stage2_config:
+        stage2_config['ACCF32'] = False
+    if 'SPLIT_K' not in stage1_config:
+        stage1_config['SPLIT_K'] = 1
+    if 'SPLIT_K' not in stage2_config:
+        stage2_config['SPLIT_K'] = 1    
+
+    if stage1_config['ACCF32']:
+       acc_type1 = torch.float32
+    else:
+       acc_type1 = hidden_states.dtype
+    if stage2_config['ACCF32']:
+       acc_type2 = torch.float32
+    else:
+       acc_type2 = hidden_states.dtype
+       
+
+    if stage1_config['SPLIT_K'] > 1:
+        intermediate_cache1 = torch.zeros((M, topk_ids.shape[1], N),
+                                          device=hidden_states.device,
+                                          dtype=acc_type1)
+    else:
+        intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N),
+                                          device=hidden_states.device,
+                                          dtype=hidden_states.dtype)
+        
+    if stage2_config['SPLIT_K'] > 1:
+        intermediate_cache3 = torch.zeros((M, topk_ids.shape[1], w2.shape[1]),
+                                          device=hidden_states.device,
+                                          dtype=acc_type2)
+    else:
+        intermediate_cache3 = torch.empty((M, topk_ids.shape[1], w2.shape[1]),
+                                          device=hidden_states.device,
+                                          dtype=hidden_states.dtype)
 
     # This needs separate memory since it's used concurrently with cache1
     intermediate_cache2 = torch.empty((M * top_k_num, N // 2),
@@ -1329,18 +1455,30 @@ def fused_experts_impl(
         curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
         curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
 
-        qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
-            A=curr_hidden_states,
-            A_scale=a1_scale,
-            qtype=qtype,
-            per_channel_quant=per_channel_quant,
-            block_shape=block_shape)
-
         sorted_token_ids, expert_ids, num_tokens_post_padded = (
-            moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'],
+            moe_align_block_size(curr_topk_ids, stage1_config['BLOCK_SIZE_M'],
                                  global_num_experts, expert_map))
 
-        invoke_fused_moe_kernel(qcurr_hidden_states,
+        if (stage1_config['BLOCK_SIZE_M'] == 128 and use_int8_w8a8==False and (topk_ids.shape[1] == 1 or topk_ids.shape[1] == 2) and
+            (curr_hidden_states.dtype == torch.bfloat16 or curr_hidden_states.dtype == torch.float16) and
+            w1.shape[1] % 4 == 0 and w1.shape[2] % 8 == 0):
+            ops.fused_moe_kernel(curr_hidden_states, w1, intermediate_cache1,
+                                curr_topk_weights, curr_topk_ids, sorted_token_ids,
+                                expert_ids, num_tokens_post_padded, False,
+                                topk_ids.shape[1], 0)
+        else:
+            if stage2_config['BLOCK_SIZE_M'] != stage1_config['BLOCK_SIZE_M']:
+                sorted_token_ids, expert_ids, num_tokens_post_padded = (
+                moe_align_block_size(curr_topk_ids, stage2_config['BLOCK_SIZE_M'], global_num_experts, expert_map))
+
+            qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
+                A=curr_hidden_states,
+                A_scale=a1_scale,
+                qtype=qtype,
+                per_channel_quant=per_channel_quant,
+                block_shape=block_shape)
+            
+            invoke_fused_moe_kernel(qcurr_hidden_states,
                                 w1,
                                 intermediate_cache1,
                                 a1q_scale,
@@ -1352,12 +1490,13 @@ def fused_experts_impl(
                                 num_tokens_post_padded,
                                 apply_router_weight_on_input,
                                 top_k_num,
-                                config,
+                                stage1_config,
                                 compute_type=compute_type,
                                 use_fp8_w8a8=use_fp8_w8a8,
                                 use_int8_w8a8=use_int8_w8a8,
                                 use_int8_w8a16=use_int8_w8a16,
                                 use_int4_w4a16=use_int4_w4a16,
+                                orig_acc_dtype=hidden_states.dtype,
                                 per_channel_quant=per_channel_quant,
                                 block_shape=block_shape)
 
@@ -1370,14 +1509,21 @@ def fused_experts_impl(
         else:
             raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
-        qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
-            A=intermediate_cache2,
-            A_scale=a2_scale,
-            qtype=qtype,
-            per_channel_quant=per_channel_quant,
-            block_shape=block_shape)
-
-        invoke_fused_moe_kernel(qintermediate_cache2,
+        
+        if (stage2_config['BLOCK_SIZE_M'] == 128 and use_int8_w8a8==False and w2.shape[1] % 4 == 0 and w2.shape[2] % 8 == 0 and
+            (hidden_states.dtype == torch.bfloat16 or hidden_states.dtype == torch.float16)):
+            ops.fused_moe_kernel(intermediate_cache2, w2, intermediate_cache3,
+                                curr_topk_weights, curr_topk_ids, sorted_token_ids,
+                                expert_ids, num_tokens_post_padded, True, 1, 0)
+        else:
+            qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
+                A=intermediate_cache2,
+                A_scale=a2_scale,
+                qtype=qtype,
+                per_channel_quant=per_channel_quant,
+                block_shape=block_shape)
+
+            invoke_fused_moe_kernel(qintermediate_cache2,
                                 w2,
                                 intermediate_cache3,
                                 a2q_scale,
@@ -1389,12 +1535,13 @@ def fused_experts_impl(
                                 num_tokens_post_padded,
                                 not apply_router_weight_on_input,
                                 1,
-                                config,
+                                stage2_config,
                                 compute_type=compute_type,
                                 use_fp8_w8a8=use_fp8_w8a8,
                                 use_int8_w8a8=use_int8_w8a8,
                                 use_int8_w8a16=use_int8_w8a16,
                                 use_int4_w4a16=use_int4_w4a16,
+                                orig_acc_dtype=hidden_states.dtype,
                                 per_channel_quant=per_channel_quant,
                                 block_shape=block_shape)
 
diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py
index cf8e4ee65..074e690e7 100644
--- a/vllm/model_executor/layers/fused_moe/layer.py
+++ b/vllm/model_executor/layers/fused_moe/layer.py
@@ -35,6 +35,12 @@ from vllm.utils import direct_register_custom_op
 has_pplx = importlib.util.find_spec("pplx_kernels") is not None
 has_deepep = importlib.util.find_spec("deep_ep") is not None
 
+if has_deepep:
+    try:
+        import deep_ep
+    except ImportError:
+        has_deepep = False
+
 if current_platform.is_cuda_alike():
     from .fused_batched_moe import BatchedTritonExperts
     from .fused_moe import TritonExperts, fused_experts
diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py
index 692482c2e..cb49594f0 100644
--- a/vllm/model_executor/layers/fused_moe/utils.py
+++ b/vllm/model_executor/layers/fused_moe/utils.py
@@ -62,7 +62,8 @@ def _int8_quantize(
     if block_shape is None:
         assert per_act_token, \
             "int8 quantization only supports block or channel-wise"
-        A, A_scale = per_token_quant_int8(A)
+        # A, A_scale = per_token_quant_int8(A)
+        A, A_scale, _ = ops.scaled_int8_quant(A, A_scale)
     else:
         assert len(block_shape) == 2
         _, block_k = block_shape[0], block_shape[1]
diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
index 588aa8deb..cdb863ac4 100644
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -28,6 +28,8 @@ from vllm.model_executor.parameter import (BasevLLMParameter,
 # yapf: enable
 from vllm.model_executor.utils import set_weight_attrs
 
+import vllm.envs as envs
+
 logger = init_logger(__name__)
 
 WEIGHT_LOADER_V2_SUPPORTED = [
@@ -105,7 +107,11 @@ def adjust_scalar_to_fused_array(param, loaded_weight, shard_id):
         assert loaded_weight.shape[0] == 1
         loaded_weight = loaded_weight[0]
 
-    return param[shard_id], loaded_weight
+    # Support gemm_tn->gemm_nn
+    if envs.MACA_VLLM_USE_TN_2_NN:
+        return param[shard_id], loaded_weight.t()
+    else:
+        return param[shard_id], loaded_weight
 
 
 # TODO(Isotr0py): We might need a more flexible structure to handle
@@ -187,10 +193,17 @@ class UnquantizedLinearMethod(LinearMethodBase):
                        output_partition_sizes: list[int], input_size: int,
                        output_size: int, params_dtype: torch.dtype,
                        **extra_weight_attrs):
-        weight = Parameter(torch.empty(sum(output_partition_sizes),
-                                       input_size_per_partition,
+        # Support gemm_tn->gemm_nn here
+        if envs.MACA_VLLM_USE_TN_2_NN:
+            weight = Parameter(torch.empty(input_size_per_partition,
+                                       sum(output_partition_sizes),
                                        dtype=params_dtype),
                            requires_grad=False)
+        else:
+            weight = Parameter(torch.empty(sum(output_partition_sizes),
+                                           input_size_per_partition,
+                                           dtype=params_dtype),
+                           requires_grad=False)
         set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
         layer.register_parameter("weight", weight)
         set_weight_attrs(weight, extra_weight_attrs)
@@ -199,8 +212,11 @@ class UnquantizedLinearMethod(LinearMethodBase):
               layer: torch.nn.Module,
               x: torch.Tensor,
               bias: Optional[torch.Tensor] = None) -> torch.Tensor:
-
-        return dispatch_unquantized_gemm()(x, layer.weight, bias)
+        # Support gemm_tn->gemm_nn here
+        if envs.MACA_VLLM_USE_TN_2_NN and x.shape[-1] == layer.weight.shape[0]:
+            return dispatch_unquantized_gemm()(x, layer.weight.t(), bias)
+        else:
+            return dispatch_unquantized_gemm()(x, layer.weight, bias)
 
 
 class LinearBase(torch.nn.Module):
@@ -321,6 +337,11 @@ class ReplicatedLinear(LinearBase):
         if len(loaded_weight.shape) == 0:
             loaded_weight = loaded_weight.reshape(1)
 
+        # Support gemm_tn->gemm_nn here
+        is_quantization = not isinstance(self.quant_method, UnquantizedLinearMethod)
+        if envs.MACA_VLLM_USE_TN_2_NN and not is_quantization:
+            loaded_weight = loaded_weight.t()
+
         assert param.size() == loaded_weight.size(), (
             f"Tried to load weights of size {loaded_weight.size()}"
             f"to a parameter of size {param.size()}")
@@ -438,6 +459,8 @@ class ColumnParallelLinear(LinearBase):
         # bitsandbytes loads the weights of the specific portion
         # no need to narrow
         is_sharded_weight = is_sharded_weight or use_bitsandbytes_4bit
+        # Support gemm_tn->gemm_nn here
+        is_quantization = not isinstance(self.quant_method, UnquantizedLinearMethod)
 
         # Special case for GGUF
         is_gguf_weight = getattr(param, "is_gguf_weight", False)
@@ -456,7 +479,13 @@ class ColumnParallelLinear(LinearBase):
 
         param_data = param.data
         if output_dim is not None and not is_sharded_weight:
-            shard_size = param_data.shape[output_dim]
+            
+            # Support gemm_tn->gemm_nn here
+            if not envs.MACA_VLLM_USE_TN_2_NN or len(param_data.shape)==1 or is_quantization:
+                shard_size = param_data.shape[output_dim] 
+            else:
+                shard_size = param_data.shape[int(not(output_dim))]
+                
             start_idx = tp_rank * shard_size
             loaded_weight = loaded_weight.narrow(output_dim, start_idx,
                                                  shard_size)
@@ -466,6 +495,10 @@ class ColumnParallelLinear(LinearBase):
         if len(loaded_weight.shape) == 0:
             loaded_weight = loaded_weight.reshape(1)
 
+        # Support gemm_tn->gemm_nn here
+        if envs.MACA_VLLM_USE_TN_2_NN and not is_quantization:
+            loaded_weight = loaded_weight.t()
+
         assert param_data.shape == loaded_weight.shape
         param_data.copy_(loaded_weight)
 
@@ -596,7 +629,10 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         is_metadata = getattr(param, "is_metadata", False)
         # Special case for per-tensor scale to load scalar into fused array.
         needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
-
+        
+        # Support gemm_tn->gemm_nn here
+        is_quantization = not isinstance(self.quant_method, UnquantizedLinearMethod)
+        
         if loaded_shard_id is None:
             # Loaded weight is already fused on disk (mlp).
             # (e.g., Phi-3's gate_up_proj).
@@ -676,8 +712,12 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
                 shard_offset = loaded_weight.shape[output_dim] * \
                     loaded_shard_id
 
-            param_data = param_data.narrow(output_dim, shard_offset,
-                                           shard_size)
+            # Support gemm_tn->gemm_nn here
+            if not envs.MACA_VLLM_USE_TN_2_NN or is_quantization:
+                param_data = param_data.narrow(output_dim, shard_offset,shard_size)
+            else:
+                param_data = param_data.narrow(int(not(output_dim)), shard_offset,shard_size)
+                
             start_idx = tp_rank * shard_size
             if not is_sharded_weight:
                 loaded_weight = loaded_weight.narrow(output_dim, start_idx,
@@ -701,7 +741,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
                     "Loading a weight without `output_dim` attribute in "
                     "MergedColumnParallelLinear, assume the weight is "
                     "the same for all partitions.")
-
+        # Support gemm_tn->gemm_nn here
+        if envs.MACA_VLLM_USE_TN_2_NN and not is_quantization:
+            loaded_weight = loaded_weight.t()
         assert param_data.shape == loaded_weight.shape
         param_data.copy_(loaded_weight)
 
@@ -990,7 +1032,10 @@ class QKVParallelLinear(ColumnParallelLinear):
 
         # Special case for per-tensor scales in fused case.
         needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
-
+        
+        # Support gemm_tn->gemm_nn here
+        is_quantization = not isinstance(self.quant_method, UnquantizedLinearMethod)
+        
         if loaded_shard_id is None:
             # Loaded weight is already fused on disk (qkv).
             # (e.g., Phi-3's qkv_proj).
@@ -1097,8 +1142,12 @@ class QKVParallelLinear(ColumnParallelLinear):
                 shard_size, shard_offset = adjust_bitsandbytes_4bit_shard(
                     param, orig_qkv_offsets, loaded_shard_id)
 
-            param_data = param_data.narrow(output_dim, shard_offset,
-                                           shard_size)
+            # Support gemm_tn->gemm_nn here
+            if not envs.MACA_VLLM_USE_TN_2_NN or len(param_data.shape)==1 or is_quantization:
+                param_data = param_data.narrow(output_dim, shard_offset,shard_size)
+            else:
+                param_data = param_data.narrow(int(not(output_dim)), shard_offset,shard_size)
+            
             if loaded_shard_id == "q":
                 shard_id = tp_rank
             else:
@@ -1128,6 +1177,10 @@ class QKVParallelLinear(ColumnParallelLinear):
                     "QKVParallelLinear, assume the weight is the same "
                     "for all partitions.")
 
+        # Support gemm_tn->gemm_nn here
+        if envs.MACA_VLLM_USE_TN_2_NN and not is_quantization:
+            loaded_weight = loaded_weight.t()
+
         assert param_data.shape == loaded_weight.shape
         param_data.copy_(loaded_weight)
 
@@ -1245,8 +1298,15 @@ class RowParallelLinear(LinearBase):
             param.materialize(tuple(weight_shape), dtype=loaded_weight.dtype)
 
         param_data = param.data
+        # Support gemm_tn->gemm_nn here
+        is_quantization = not isinstance(self.quant_method, UnquantizedLinearMethod)
+        
         if input_dim is not None and not is_sharded_weight:
-            shard_size = param_data.shape[input_dim]
+            # Support gemm_tn->gemm_nn here
+            if not envs.MACA_VLLM_USE_TN_2_NN or is_quantization:
+                shard_size = param_data.shape[input_dim]
+            else:
+                shard_size = param_data.shape[int(not(input_dim))]
             start_idx = tp_rank * shard_size
             loaded_weight = loaded_weight.narrow(input_dim, start_idx,
                                                  shard_size)
@@ -1256,6 +1316,9 @@ class RowParallelLinear(LinearBase):
         if len(loaded_weight.shape) == 0:
             loaded_weight = loaded_weight.reshape(1)
 
+        # Support gemm_tn->gemm_nn here
+        if envs.MACA_VLLM_USE_TN_2_NN and not is_quantization:
+            loaded_weight = loaded_weight.t()
         assert param_data.shape == loaded_weight.shape
         param_data.copy_(loaded_weight)
 
diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py
index 1cb23e7a1..c3c184a08 100644
--- a/vllm/model_executor/layers/quantization/__init__.py
+++ b/vllm/model_executor/layers/quantization/__init__.py
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.quantization.base_config import (
 QuantizationMethods = Literal[
     "aqlm",
     "awq",
+    "gptq",
     "deepspeedfp",
     "tpu_int8",
     "fp8",
@@ -23,7 +24,7 @@ QuantizationMethods = Literal[
     "gptq_marlin",
     "gptq_bitblas",
     "awq_marlin",
-    "gptq",
+    # "gptq",
     "compressed-tensors",
     "bitsandbytes",
     "qqq",
@@ -126,9 +127,9 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
         "bitblas": BitBLASConfig,
         "gguf": GGUFConfig,
         "gptq_marlin_24": GPTQMarlin24Config,
-        "gptq_marlin": GPTQMarlinConfig,
+        "gptq_marlin": GPTQConfig,
         "gptq_bitblas": GPTQBitBLASConfig,
-        "awq_marlin": AWQMarlinConfig,
+        "awq_marlin": AWQConfig,
         "gptq": GPTQConfig,
         "compressed-tensors": CompressedTensorsConfig,
         "bitsandbytes": BitsAndBytesConfig,
diff --git a/vllm/model_executor/layers/quantization/awq.py b/vllm/model_executor/layers/quantization/awq.py
index f8bc3ab5e..4503e466e 100644
--- a/vllm/model_executor/layers/quantization/awq.py
+++ b/vllm/model_executor/layers/quantization/awq.py
@@ -13,7 +13,7 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig)
 from vllm.model_executor.parameter import (GroupQuantScaleParameter,
                                            PackedvLLMParameter)
-
+from vllm.utils import direct_register_custom_op
 
 class AWQConfig(QuantizationConfig):
     """Config class for AWQ.
@@ -50,7 +50,7 @@ class AWQConfig(QuantizationConfig):
         return "awq"
 
     def get_supported_act_dtypes(self) -> list[torch.dtype]:
-        return [torch.half]
+        return [torch.half, torch.bfloat16]
 
     @classmethod
     def get_min_capability(cls) -> int:
@@ -86,6 +86,58 @@ class AWQConfig(QuantizationConfig):
 def is_layer_skipped_awq(prefix: str, modules_to_not_convert: list[str]):
     return any(module_name in prefix for module_name in modules_to_not_convert)
 
+def _apply_awq_fake(x: torch.Tensor,
+                    qweight: torch.Tensor,
+                    scales: torch.Tensor,
+                    qzeros: torch.Tensor,
+                    bias: torch.Tensor,
+                    pack_factor: int,
+                    group_size: int) -> torch.Tensor:
+    out_shape = ()
+    if group_size % 32:
+        out_shape = (x.shape[:-1] + (qweight.shape[-1] * pack_factor, ))
+    else:
+        out_shape = (x.shape[:-1] + (qweight.shape[0], ))
+    return torch.empty(out_shape, dtype=x.dtype, device=x.device)
+
+def _apply_awq(x: torch.Tensor,
+               qweight: torch.Tensor,
+               scales: torch.Tensor,
+               qzeros: torch.Tensor,
+               bias: torch.Tensor,
+               pack_factor: int,
+               group_size: int) -> torch.Tensor:
+    out_shape = ()
+    reshaped_x = x.reshape(-1, x.shape[-1])
+    out = torch.empty(0)          
+    # num_tokens >= threshold
+    FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
+    # if (FP16_MATMUL_HEURISTIC_CONDITION and reshaped_x.dtype == torch.half) or self.quant_config.group_size != 128:
+    if group_size % 32:
+        out_shape = (x.shape[:-1] + (qweight.shape[-1] * pack_factor, ))
+        out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
+        out = torch.matmul(reshaped_x, out)
+    else:
+        num_out_channel = qweight.shape[0]
+        out_shape = (x.shape[:-1] + (num_out_channel, ))
+        temp_space = torch.empty(0, dtype=torch.float32, device=x.device)
+        if reshaped_x.dtype == torch.bfloat16:
+            temp_space = torch.zeros(reshaped_x.shape[0], num_out_channel,
+                                        dtype=torch.float32, device=x.device)
+        out = ops.awq_gemm(reshaped_x, qweight, qzeros, scales,
+                            pack_factor, temp_space,
+                            True if reshaped_x.dtype == torch.bfloat16 else False)
+    if bias is not None:
+        out.add_(bias)
+    return out.reshape(out_shape)
+
+direct_register_custom_op(
+    op_name="_apply_awq",
+    op_func=_apply_awq,
+    mutates_args=[],
+    fake_impl=_apply_awq_fake,
+    tags=(torch.Tag.needs_fixed_stride_order, ),
+)
 
 class AWQLinearMethod(LinearMethodBase):
     """Linear method for AWQ.
@@ -168,6 +220,12 @@ class AWQLinearMethod(LinearMethodBase):
                                           requires_grad=False)
         layer.scales = torch.nn.Parameter(layer.scales.data,
                                           requires_grad=False)
+        # warmup
+        if self.quant_config.group_size % 32:
+            pass
+        else:
+            qweight = ops.awq_to_gptq_4bit(layer.qweight)
+            layer.qweight = torch.nn.Parameter(qweight, requires_grad=False)
 
     def apply(self,
               layer: torch.nn.Module,
@@ -177,18 +235,7 @@ class AWQLinearMethod(LinearMethodBase):
         scales = layer.scales
         qzeros = layer.qzeros
         pack_factor = self.quant_config.pack_factor
-        out_shape = (x.shape[:-1] + (qweight.shape[-1] * pack_factor, ))
-        reshaped_x = x.reshape(-1, x.shape[-1])
-
-        # num_tokens >= threshold
-        FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
-
-        if FP16_MATMUL_HEURISTIC_CONDITION:
-            out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
-            out = torch.matmul(reshaped_x, out)
-        else:
-            out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros,
-                               pack_factor)
-        if bias is not None:
-            out.add_(bias)
-        return out.reshape(out_shape)
+        group_size = self.quant_config.group_size
+        
+        return torch.ops.vllm._apply_awq(x, qweight, scales, qzeros, 
+                                        bias, pack_factor, group_size)
diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py
index 78c5c75c0..0b50cef26 100644
--- a/vllm/model_executor/layers/quantization/base_config.py
+++ b/vllm/model_executor/layers/quantization/base_config.py
@@ -113,7 +113,10 @@ class QuantizationConfig(ABC):
            this method should only be overwritten by subclasses in exceptional 
            circumstances
         """
-        return None
+        if(user_quant != None):
+            return user_quant
+        else:
+            return hf_quant_cfg["quant_method"]
 
     @staticmethod
     def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any:
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
index 28c62fc5e..1538c371e 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
@@ -197,7 +197,7 @@ class CompressedTensorsConfig(QuantizationConfig):
                                 error: bool = True,
                                 match_exact: bool = False) -> bool:
         capability_tuple = current_platform.get_device_capability()
-
+        """
         if capability_tuple is not None:
             capability = capability_tuple.to_int()
             if match_exact:
@@ -217,6 +217,8 @@ class CompressedTensorsConfig(QuantizationConfig):
             return supported
         else:
             return False
+        """
+        return False
 
     def _is_fp4a4_nvfp4(self, weight_quant: BaseModel, input_quant: BaseModel):
 
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
index bc9d399cf..4beb287f4 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
@@ -96,7 +96,6 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
             raise RuntimeError(
                 f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}")
 
-
 class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
 
     def __init__(
diff --git a/vllm/model_executor/layers/quantization/gptq.py b/vllm/model_executor/layers/quantization/gptq.py
index d3ab1be3b..492a9c931 100644
--- a/vllm/model_executor/layers/quantization/gptq.py
+++ b/vllm/model_executor/layers/quantization/gptq.py
@@ -85,8 +85,9 @@ class GPTQConfig(QuantizationConfig):
         return "gptq"
 
     @classmethod
+
     def get_supported_act_dtypes(cls) -> list[torch.dtype]:
-        return [torch.half]
+        return [torch.half, torch.bfloat16]
 
     @classmethod
     # Need to figure it out
@@ -251,7 +252,7 @@ class GPTQLinearMethod(LinearMethodBase):
 
         # exllama needs to shuffle the weight after the weight is loaded
         # here we do the shuffle on first forward pass
-        if layer.exllama_state == ExllamaState.UNINITIALIZED:
+        if self.quant_config.group_size == 128 or self.quant_config.group_size == 64:
             if self.quant_config.desc_act:
                 layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int)
             else:
@@ -261,6 +262,63 @@ class GPTQLinearMethod(LinearMethodBase):
             layer.exllama_state = ExllamaState.READY
             ops.gptq_shuffle(layer.qweight, layer.g_idx,
                              self.quant_config.weight_bits)
+            
+            if layer.scales.dtype != torch.bfloat16:
+                perm_space = torch.empty(0)
+                temp_space = torch.empty(0)
+                if self.quant_config.weight_bits == 4:
+                    # warmup
+                    reshaped_x = torch.randn(1, layer.qweight.shape[0]*8, dtype=layer.scales.dtype, device="cuda")
+                    _ = ops.gptq_gemm(reshaped_x, layer.qweight, layer.qzeros,
+                                   layer.scales, layer.g_idx,
+                                   layer.exllama_state == ExllamaState.READY,
+                                   self.quant_config.weight_bits,
+                                   self.quant_config.group_size,
+                                   perm_space, temp_space,
+                                   False)
+                if self.quant_config.weight_bits == 8:
+                    # warmup
+                    reshaped_x = torch.randn(1, layer.qweight.shape[0]*4, dtype=layer.scales.dtype, device="cuda")
+                    _ = ops.gptq_gemm(reshaped_x, layer.qweight, layer.qzeros,
+                                   layer.scales, layer.g_idx,
+                                   layer.exllama_state == ExllamaState.READY,
+                                   self.quant_config.weight_bits,
+                                   self.quant_config.group_size,
+                                   perm_space, temp_space,
+                                   False)
+        else:
+            if layer.exllama_state == ExllamaState.UNINITIALIZED:
+                if self.quant_config.desc_act:
+                    layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int)
+                else:
+                    layer.g_idx.data = torch.empty((0, ),
+                                                   dtype=torch.int,
+                                                   device=layer.g_idx.device)
+                layer.exllama_state = ExllamaState.READY
+                ops.gptq_shuffle(layer.qweight, layer.g_idx,
+                                 self.quant_config.weight_bits)
+
+                """
+                perm_space = torch.empty(0)
+                if self.quant_config.weight_bits == 4:
+                    # warmup
+                    reshaped_x = torch.randn(1, layer.qweight.shape[0]*8, dtype=layer.scales.dtype, device="cuda")
+                    _ = ops.gptq_gemm(reshaped_x, layer.qweight, layer.qzeros,
+                                   layer.scales, layer.g_idx,
+                                   layer.exllama_state == ExllamaState.READY,
+                                   self.quant_config.weight_bits,
+                                   self.quant_config.group_size,
+                                   perm_space)
+                if self.quant_config.weight_bits == 8:
+                    # warmup
+                    reshaped_x = torch.randn(1, layer.qweight.shape[0]*4, dtype=layer.scales.dtype, device="cuda")
+                    _ = ops.gptq_gemm(reshaped_x, layer.qweight, layer.qzeros,
+                                   layer.scales, layer.g_idx,
+                                   layer.exllama_state == ExllamaState.READY,
+                                   self.quant_config.weight_bits,
+                                   self.quant_config.group_size,
+                                   perm_space)
+                """
 
     def apply(self,
               layer: torch.nn.Module,
@@ -269,10 +327,25 @@ class GPTQLinearMethod(LinearMethodBase):
         out_shape = x.shape[:-1] + (layer.qweight.shape[-1], )
         reshaped_x = x.reshape(-1, x.shape[-1])
 
+        perm_space = torch.empty(0)
+        temp_space = torch.empty(0)
+        if self.quant_config.weight_bits == 4 or self.quant_config.weight_bits == 8:
+            if self.quant_config.group_size == 128 or self.quant_config.group_size == 64:
+                if self.quant_config.desc_act:
+                    perm_space = torch.empty(reshaped_x.shape[0], reshaped_x.shape[1],
+                                             dtype=torch.float16, device="cuda")
+                    
+                if reshaped_x.dtype == torch.bfloat16:
+                    temp_space = torch.zeros(reshaped_x.shape[0], layer.qweight.shape[1],
+                                             dtype=torch.float32, device="cuda")
+
         output = ops.gptq_gemm(reshaped_x, layer.qweight, layer.qzeros,
                                layer.scales, layer.g_idx,
                                layer.exllama_state == ExllamaState.READY,
-                               self.quant_config.weight_bits)
+                               self.quant_config.weight_bits,
+                               self.quant_config.group_size,
+                               perm_space, temp_space,
+                               True if reshaped_x.dtype == torch.bfloat16 else False)
         if bias is not None:
             output.add_(bias)
         return output.reshape(out_shape)
diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py
index 3aa23f068..daae03c73 100644
--- a/vllm/model_executor/layers/quantization/moe_wna16.py
+++ b/vllm/model_executor/layers/quantization/moe_wna16.py
@@ -41,6 +41,7 @@ class MoeWNA16Config(QuantizationConfig):
             AWQMarlinConfig)
         from vllm.model_executor.layers.quantization.gptq_marlin import (
             GPTQMarlinConfig)
+        """
         if self.linear_quant_method == "gptq":
             self.use_marlin = GPTQMarlinConfig.is_gptq_marlin_compatible(
                 full_config)
@@ -59,7 +60,7 @@ class MoeWNA16Config(QuantizationConfig):
                 full_config)
         else:
             raise ValueError("moe_wna16 only support gptq and awq.")
-
+        """
         if modules_to_not_convert is None:
             self.modules_to_not_convert = []
         else:
@@ -390,13 +391,22 @@ class MoeWNA16Method(FusedMoEMethodBase):
                                     loaded_weight: torch.Tensor,
                                     weight_name: str, shard_id: str,
                                     expert_id: int):
+            if layer.ep_size > 1:
+                global_expert_id = expert_id
+                expert_id = layer._map_global_expert_id_to_local_expert_id(expert_id)
+                if expert_id == -1:
+                    return
+
             if "g_idx" in weight_name:
                 return
             if not layer.quant_config.has_zp and "qzeros" in weight_name:
                 return
 
             device = get_tp_group().device
-            tp_rank = get_tensor_model_parallel_rank()
+            if layer.ep_size > 1:
+                tp_rank = 0
+            else:
+                tp_rank = get_tensor_model_parallel_rank()
             loaded_weight = loaded_weight.to(device)
             shard_size = layer.intermediate_size_per_partition
 
@@ -433,16 +443,26 @@ class MoeWNA16Method(FusedMoEMethodBase):
                     layer.group_size_div_factor, 1)
 
             if "w13_qzeros" in weight_name:
-                tensor = loaded_weight.view(layer.tp_size, -1,
-                                            loaded_weight.size(1))[tp_rank]
+                if layer.ep_size > 1 :
+                    tensor = loaded_weight.view(-1, param.data[expert_id].shape[0] // 2,
+                                                loaded_weight.size(1))[tp_rank]
+                else:
+                    tensor = loaded_weight.view(layer.tp_size, -1,
+                                                loaded_weight.size(1))[tp_rank]
                 if shard_id == "w1":
                     param.data[expert_id, :shard_size // 2] = tensor
                 else:
                     param.data[expert_id, shard_size // 2:] = tensor
             elif "w2_qzeros" in weight_name:
-                param.data[expert_id] = loaded_weight.view(
-                    loaded_weight.size(0), layer.tp_size, -1)[:, tp_rank]
+                if layer.ep_size > 1 :
+                    param.data[expert_id] = loaded_weight.view(
+                        loaded_weight.size(0), -1, param.data[expert_id].shape[1])[:, tp_rank]
+                else:
+                    param.data[expert_id] = loaded_weight.view(
+                        loaded_weight.size(0), layer.tp_size, -1)[:, tp_rank]
             else:
+                if layer.ep_size > 1:
+                    expert_id = global_expert_id
                 weight_loader(param, loaded_weight, weight_name, shard_id,
                               expert_id)
 
diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py
index 08dc99e07..aaff93307 100644
--- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py
@@ -5,7 +5,7 @@
 import functools
 import json
 import os
-from typing import Any, Callable, Optional, Union
+from typing import Any, List, Callable, Optional, Union
 
 import torch
 
@@ -33,7 +33,7 @@ def cutlass_scaled_mm(
     B: torch.Tensor,
     As: torch.Tensor,
     Bs: torch.Tensor,
-    block_size: list[int],
+    block_size: List[int],
     output_dtype: torch.dtype = torch.float16,
 ) -> torch.Tensor:
     return ops.cutlass_scaled_mm(A,
@@ -48,7 +48,7 @@ def rocm_aiter_gemm_w8a8_blockscale_impl(
     B: torch.Tensor,
     As: torch.Tensor,
     Bs: torch.Tensor,
-    block_size: list[int],
+    block_size: List[int],
     output_dtype: torch.dtype = torch.float16,
 ) -> torch.Tensor:
     import aiter as rocm_aiter
@@ -61,7 +61,7 @@ def rocm_aiter_gemm_w8a8_blockscale_fake(
     B: torch.Tensor,
     As: torch.Tensor,
     Bs: torch.Tensor,
-    block_size: list[int],
+    block_size: List[int],
     output_dtype: torch.dtype = torch.float16,
 ) -> torch.Tensor:
 
@@ -103,7 +103,7 @@ def dispatch_w8a8_blockscale_func(
 def apply_w8a8_block_fp8_linear(
     input: torch.Tensor,
     weight: torch.Tensor,
-    block_size: list[int],
+    block_size: List[int],
     weight_scale: torch.Tensor,
     input_scale: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
@@ -156,7 +156,7 @@ def apply_w8a8_block_fp8_linear(
 def apply_w8a8_block_fp8_linear_fake(
     input: torch.Tensor,
     weight: torch.Tensor,
-    block_size: list[int],
+    block_size: List[int],
     weight_scale: torch.Tensor,
     input_scale: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
@@ -536,7 +536,7 @@ def w8a8_block_fp8_matmul(
     B: torch.Tensor,
     As: torch.Tensor,
     Bs: torch.Tensor,
-    block_size: list[int],
+    block_size: List[int],
     output_dtype: torch.dtype = torch.float16,
 ) -> torch.Tensor:
     """This function performs matrix multiplication with block-wise
diff --git a/vllm/model_executor/layers/rotary_embedding.py b/vllm/model_executor/layers/rotary_embedding.py
index 9de233896..53678ebeb 100644
--- a/vllm/model_executor/layers/rotary_embedding.py
+++ b/vllm/model_executor/layers/rotary_embedding.py
@@ -34,7 +34,7 @@ from vllm.model_executor.custom_op import CustomOp
 from vllm.platforms import current_platform
 
 if current_platform.is_cuda():
-    from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+    from flash_attn.layers.rotary import apply_rotary_emb
 
 
 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py
index 0f636d83a..489b621c9 100644
--- a/vllm/model_executor/layers/vocab_parallel_embedding.py
+++ b/vllm/model_executor/layers/vocab_parallel_embedding.py
@@ -6,6 +6,7 @@ from dataclasses import dataclass
 from typing import Optional
 
 import torch
+from vllm import envs
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter, UninitializedParameter
 
@@ -31,7 +32,13 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
                        output_size: int, params_dtype: torch.dtype,
                        **extra_weight_attrs):
         """Create weights for embedding layer."""
-        weight = Parameter(torch.empty(sum(output_partition_sizes),
+        if envs.MACA_VLLM_USE_TN_2_NN:
+            weight = Parameter(torch.empty(input_size_per_partition,
+                                       sum(output_partition_sizes),
+                                       dtype=params_dtype),
+                           requires_grad=False)
+        else:
+            weight = Parameter(torch.empty(sum(output_partition_sizes),
                                        input_size_per_partition,
                                        dtype=params_dtype),
                            requires_grad=False)
@@ -43,11 +50,17 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
               layer: torch.nn.Module,
               x: torch.Tensor,
               bias: Optional[torch.Tensor] = None) -> torch.Tensor:
-        return dispatch_unquantized_gemm()(x, layer.weight, bias)
+        if envs.MACA_VLLM_USE_TN_2_NN and x.shape[-1] == layer.weight.shape[0]:
+            return dispatch_unquantized_gemm()(x, layer.weight.t(), bias)
+        else:
+            return dispatch_unquantized_gemm()(x, layer.weight, bias)
 
     def embedding(self, layer: torch.nn.Module,
                   input_: torch.Tensor) -> torch.Tensor:
-        return F.embedding(input_, layer.weight)
+        if envs.MACA_VLLM_USE_TN_2_NN:
+            return F.embedding(input_, layer.weight.t())
+        else:
+            return F.embedding(input_, layer.weight)
 
 
 def pad_vocab_size(vocab_size: int,
@@ -138,13 +151,15 @@ class VocabParallelEmbeddingShardIndices:
         assert self.num_added_elements <= self.num_added_elements_padded
 
 
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+# @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+@torch.jit.script
 def get_masked_input_and_mask(
         input_: torch.Tensor, org_vocab_start_index: int,
         org_vocab_end_index: int, num_org_vocab_padding: int,
         added_vocab_start_index: int,
         added_vocab_end_index: int) -> tuple[torch.Tensor, torch.Tensor]:
     # torch.compile will fuse all of the pointwise ops below
+    # torch.jit.script will fuse all of the pointwise ops below
     # into a single kernel, making it very fast
     org_vocab_mask = (input_ >= org_vocab_start_index) & (
         input_ < org_vocab_end_index)
@@ -264,12 +279,12 @@ class VocabParallelEmbedding(torch.nn.Module):
             self.shard_indices.added_vocab_start_index)
 
         self.quant_method.create_weights(self,
-                                         self.embedding_dim,
-                                         [self.num_embeddings_per_partition],
-                                         self.embedding_dim,
-                                         self.num_embeddings_padded,
-                                         params_dtype=params_dtype,
-                                         weight_loader=self.weight_loader)
+                                        self.embedding_dim,
+                                        [self.num_embeddings_per_partition],
+                                        self.embedding_dim,
+                                        self.num_embeddings_padded,
+                                        params_dtype=params_dtype,
+                                        weight_loader=self.weight_loader)
 
     @classmethod
     def _get_indices(cls, vocab_size_padded: int, org_vocab_size_padded: int,
@@ -389,6 +404,13 @@ class VocabParallelEmbedding(torch.nn.Module):
         # Copy the data. Select chunk corresponding to current shard.
         loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
 
+        if envs.MACA_VLLM_USE_TN_2_NN:
+            loaded_weight = loaded_weight.t()
+            # we should padding last dimension after weight transpose
+            padding_needed = max(self.num_embeddings_per_partition - loaded_weight.size(-1), 0)
+            if padding_needed:
+                loaded_weight = torch.nn.functional.pad(loaded_weight, (0, padding_needed), value=0)
+
         if current_platform.is_hpu():
             # FIXME(kzawora): Weight copy with slicing bugs out on Gaudi here,
             # so we're using a workaround. Remove this when fixed in
diff --git a/vllm/model_executor/models/baichuan.py b/vllm/model_executor/models/baichuan.py
index 0de5de5e8..c8dfa32a0 100644
--- a/vllm/model_executor/models/baichuan.py
+++ b/vllm/model_executor/models/baichuan.py
@@ -394,7 +394,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 = self.lm_head_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
         self.logits_processor = LogitsProcessor(config.vocab_size)
diff --git a/vllm/model_executor/models/deepseek.py b/vllm/model_executor/models/deepseek.py
index 2f0202f1e..04d843f31 100644
--- a/vllm/model_executor/models/deepseek.py
+++ b/vllm/model_executor/models/deepseek.py
@@ -57,6 +57,7 @@ from .utils import (AutoWeightsLoader, extract_layer_index,
                     make_empty_intermediate_tensors_factory, make_layers,
                     maybe_prefix)
 
+import vllm.envs as envs
 
 class DeepseekMLP(nn.Module):
 
@@ -152,8 +153,17 @@ class DeepseekMoE(nn.Module):
         w2s = torch._utils._unflatten_dense_tensors(self.w2, w2)
         for data, param in zip(w2s, w2):
             param.data = data
-
         self.w2 = self.w2.view(len(w2), *w2s[0].shape)
+            
+        if envs.MACA_VLLM_USE_TN_2_NN:
+            self.w1 = self.w1.permute(0,2,1).contiguous()   
+            for expert, w in zip(self.experts, self.w1):
+                expert.gate_up_proj.weight.data = w.permute(1,0)
+                
+            self.w2 = self.w2.permute(0, 2, 1).contiguous()
+            for expert, w in zip(self.experts, self.w2):
+                expert.down_proj.weight.data = w.permute(1, 0)
+        
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py
index 6e6e74b0d..9dd8c7930 100644
--- a/vllm/model_executor/models/deepseek_mtp.py
+++ b/vllm/model_executor/models/deepseek_mtp.py
@@ -52,10 +52,6 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
         quant_config: Optional[QuantizationConfig] = None,
     ) -> None:
         super().__init__()
-        self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size,
-            config.hidden_size,
-        )
 
         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -74,8 +70,6 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
         inputs_embeds: Optional[torch.Tensor] = None,
         spec_step_index: int = 0,
     ) -> torch.Tensor:
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
         assert inputs_embeds is not None
         # masking inputs at position 0, as not needed by MTP
         inputs_embeds[positions == 0] = 0
@@ -112,6 +106,10 @@ class DeepSeekMultiTokenPredictor(nn.Module):
             for idx in range(self.mtp_start_layer_idx,
                              self.mtp_start_layer_idx + self.num_mtp_layers)
         })
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+        )
 
         self.logits_processor = LogitsProcessor(config.vocab_size)
 
@@ -123,6 +121,8 @@ class DeepSeekMultiTokenPredictor(nn.Module):
         inputs_embeds: Optional[torch.Tensor] = None,
         spec_step_idx: int = 0,
     ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
         current_step_idx = (spec_step_idx % self.num_mtp_layers)
         return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
             input_ids,
@@ -217,6 +217,11 @@ class DeepSeekMTP(nn.Module, SupportsPP):
                 # Skip loading extra bias for GPTQ models.
                 if name.endswith(".bias") and name not in params_dict:
                     continue
+                # According to DeepSeek-V3 Technical Report, MTP modules
+                # shares embedding layer. We only load the first weights.
+                if (spec_layer != self.model.mtp_start_layer_idx
+                        and ".layers" not in name):
+                    continue
 
                 param = params_dict[name]
                 weight_loader = param.weight_loader
@@ -253,17 +258,25 @@ class DeepSeekMTP(nn.Module, SupportsPP):
         """
         Rewrite the weight name to match the format of the original model.
         Add .mtp_block for modules in transformer layer block for spec layer
+        and rename shared layer weights to be top level.
         """
         spec_layer_weight_names = [
             "embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head"
         ]
+        shared_weight_names = ["embed_tokens"]
         spec_layer_weight = False
+        shared_weight = False
         for weight_name in spec_layer_weight_names:
             if weight_name in name:
                 spec_layer_weight = True
+                if weight_name in shared_weight_names:
+                    shared_weight = True
                 break
         if not spec_layer_weight:
             # treat rest weights as weights for transformer layer block
             name = name.replace(f"model.layers.{spec_layer}.",
                                 f"model.layers.{spec_layer}.mtp_block.")
+        elif shared_weight:
+            # treat shared weights as top level weights
+            name = name.replace(f"model.layers.{spec_layer}.", "model.")
         return name
diff --git a/vllm/model_executor/models/ernie45.py b/vllm/model_executor/models/ernie45.py
new file mode 100644
index 000000000..fcc7a1f17
--- /dev/null
+++ b/vllm/model_executor/models/ernie45.py
@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2025 The Baidu 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 Erine model compatible with HuggingFace weights."""
+from vllm.config import VllmConfig
+from vllm.model_executor.models.llama import LlamaForCausalLM
+
+from .utils import PPMissingLayer
+
+
+class Ernie4_5_ForCausalLM(LlamaForCausalLM):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        # Hack Llama model to fit HF format Ernie4.5 dense implementation
+        # Attention difference between Ernie and Llama:
+        # 1. rotary_dim and no Neox style.
+        # 2. There is no bias for o_proj in attention
+        for layer in self.model.layers:
+            if not isinstance(layer, PPMissingLayer):
+                layer.self_attn.rotary_emb.is_neox_style = False
+                layer.self_attn.o_proj.bias = None
+                layer.self_attn.o_proj.skip_bias_add = True
\ No newline at end of file
diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py
new file mode 100644
index 000000000..fb6390275
--- /dev/null
+++ b/vllm/model_executor/models/ernie45_moe.py
@@ -0,0 +1,585 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2025 The Baidu 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 ErineMoE model compatible with HuggingFace weights."""
+from collections.abc import Iterable
+from typing import Any, Optional, 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
+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.vocab_parallel_embedding import (
+    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 .interfaces import SupportsPP
+from .utils import (PPMissingLayer, extract_layer_index,
+                    is_pp_missing_parameter,
+                    make_empty_intermediate_tensors_factory, make_layers,
+                    maybe_prefix)
+
+logger = init_logger(__name__)
+
+
+class Ernie4_5_MoeMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        use_bias: bool = False,
+        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=use_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj")
+        self.down_proj = RowParallelLinear(intermediate_size,
+                                           hidden_size,
+                                           bias=use_bias,
+                                           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 Ernie4_5_MoeMoE(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        layer_idx = extract_layer_index(prefix)
+        self.layer_idx = layer_idx
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.moe_num_shared_experts = getattr(config, "moe_num_shared_experts",
+                                              None)
+
+        if self.tp_size > config.moe_num_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than "
+                f"the number of experts {config.moe_num_experts}.")
+
+        self.gate = ReplicatedLinear(config.hidden_size,
+                                     config.moe_num_experts,
+                                     bias=False,
+                                     quant_config=None,
+                                     prefix=f"{prefix}.gate")
+
+        self.experts = FusedMoE(num_experts=config.moe_num_experts,
+                                top_k=config.moe_k,
+                                hidden_size=config.hidden_size,
+                                intermediate_size=config.moe_intermediate_size,
+                                reduce_results=False,
+                                renormalize=True,
+                                quant_config=quant_config,
+                                prefix=f"{prefix}.experts")
+
+        if self.moe_num_shared_experts is not None:
+            intermediate_size = (config.moe_intermediate_size *
+                                 config.moe_num_shared_experts)
+            self.shared_experts = Ernie4_5_MoeMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.shared_experts",
+                reduce_results=self.experts.must_reduce_shared_expert_outputs(
+                ))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        orig_shape = hidden_states.shape
+        hidden_dim = hidden_states.shape[-1]
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        if self.moe_num_shared_experts is not None:
+            shared_output = self.shared_experts(hidden_states)
+
+        router_logits, _ = self.gate(hidden_states)
+
+        final_hidden_states = self.experts(hidden_states=hidden_states,
+                                           router_logits=router_logits)
+
+        if self.moe_num_shared_experts is not None and \
+              shared_output is not None:
+            final_hidden_states = final_hidden_states + shared_output
+
+        if self.tp_size > 1:
+            final_hidden_states = (
+                self.experts.maybe_all_reduce_tensor_model_parallel(
+                    final_hidden_states))
+
+        return final_hidden_states.view(orig_shape)
+
+
+class Ernie4_5_MoeAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: Optional[int] = None,
+        rope_theta: float = 500000,
+        rope_scaling: Optional[dict[str, Any]] = None,
+        max_position_embeddings: int = 131072,
+        rms_norm_eps: float = 1e-05,
+        qkv_bias: bool = False,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        layer_idx = extract_layer_index(prefix) if len(prefix) > 0 else 0
+        self.layer_idx = layer_idx
+        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,
+            is_neox_style=False,
+            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")
+
+    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)
+        q, k = self.rotary_emb(positions, q, k)
+
+        # Attention
+        attn_output = self.attn(q, k, v)
+        # Output projection
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Ernie4_5_MoeDecoderLayer(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", 500000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          131072)
+        self.self_attn = Ernie4_5_MoeAttention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            head_dim=getattr(config, 'head_dim', None),
+            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, 'use_bias', False),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+
+        layer_idx = extract_layer_index(prefix)
+        self.layer_idx = layer_idx
+
+        # MoE
+        moe_num_experts = getattr(config, "moe_num_experts", 0)
+        moe_layer_start_index = getattr(config, "moe_layer_start_index", 0)
+        moe_layer_end_index = getattr(config, "moe_layer_end_index",
+                                      config.num_hidden_layers - 1)
+        moe_layer_interval = getattr(config, "moe_layer_interval", 1)
+        use_moe = getattr(config, "use_moe", moe_num_experts > 0)
+
+        if (use_moe and ((layer_idx + 1) % moe_layer_interval == 0)
+                and layer_idx >= moe_layer_start_index
+                and layer_idx <= moe_layer_end_index):
+            self.mlp = Ernie4_5_MoeMoE(config=config,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.mlp")
+        else:
+            self.mlp = Ernie4_5_MoeMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                use_bias=getattr(config, 'use_bias', False),
+                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 Ernie4_5_MoeModel(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
+
+        if get_pp_group().is_first_rank or (config.tie_word_embeddings
+                                            and get_pp_group().is_last_rank):
+            self.embed_tokens = VocabParallelEmbedding(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=f"{prefix}.embed_tokens")
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: Ernie4_5_MoeDecoderLayer(config=config,
+                                                    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()
+
+        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
+
+
+class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": [
+            "gate_proj",
+            "up_proj",
+        ],
+    }
+
+    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 = Ernie4_5_MoeModel(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)
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        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 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.moe_num_experts)
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if self.config.tie_word_embeddings and name.endswith(
+                    "lm_head.weight"):
+                continue
+            # MTP will be supported soon.
+            if "mtp" in name:
+                continue
+
+            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
+
+                if (("mlp.experts." in name) and name not in params_dict):
+                    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
+
+                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.
+                    name = maybe_remap_kv_scale_name(name, params_dict)
+                    if name is None:
+                        continue
+
+                    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
\ No newline at end of file
diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py
index 4100fee0e..ca1df5c7d 100644
--- a/vllm/model_executor/models/minicpmv.py
+++ b/vllm/model_executor/models/minicpmv.py
@@ -38,6 +38,8 @@ from typing_extensions import TypeVar
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.awq import AWQConfig
+from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
 from vllm.model_executor.layers.resampler import (BaseResampler, Resampler2,
                                                   get_2d_sincos_pos_embed)
 from vllm.model_executor.model_loader.utils import set_default_torch_dtype
@@ -341,7 +343,9 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo):
 
     def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
         mm_limits = {"image": None}
-        if self.get_model_version() == (2, 6):
+        if self.get_model_version() == (2,
+                                        6) or self.get_model_version() == (4,
+                                                                           0):
             mm_limits["video"] = None
 
         return mm_limits
@@ -616,7 +620,8 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]):
         out_keys: set[str],
     ) -> dict[str, NestedTensors]:
         # This processor supports zipping prompt and mm_data together
-        if self.info.get_model_version() == (2, 6):
+        if self.info.get_model_version() == (
+                2, 6) or self.info.get_model_version() == (4, 0):
             inputs = super()._call_hf_processor(
                 prompt=prompts,  # type: ignore
                 mm_data=mm_data,
@@ -671,10 +676,18 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]):
         hf_processor_mm_kwargs: Mapping[str, object],
         out_mm_kwargs: MultiModalKwargs,
     ) -> Sequence[PromptUpdate]:
-        placeholder = {
-            "image": self.info.image_pattern,
-            "video": self.info.video_pattern,
-        }
+        placeholders = [("image", self.info.image_pattern),
+                        ("video", self.info.video_pattern)]
+
+        # hard code for inconsistency of encode-decode image_pattern
+        additional_placeholders = []
+        tokenizer = self.info.get_tokenizer()
+        for modality, pattern in placeholders:
+            sub_pattern = tokenizer.decode(
+                tokenizer.encode(pattern, add_special_tokens=False))
+            if sub_pattern != pattern:
+                additional_placeholders.append((modality, sub_pattern))
+        placeholders += additional_placeholders
 
         def get_image_replacement(item_idx: int):
             images = mm_items.get_items(
@@ -706,9 +719,9 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]):
 
         return [
             PromptReplacement(modality=modality,
-                              target=placeholder[modality],
+                              target=pattern,
                               replacement=get_replacement[modality])
-            for modality in ("image", "video")
+            for modality, pattern in placeholders
         ]
 
     def _get_mm_fields_config(
@@ -1244,11 +1257,124 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):
 
         return self.resampler(vision_embedding, tgt_sizes)
 
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self,
+                                   skip_prefixes=["apm.", "audio", "tts"])
+        return loader.load_weights(weights)
+
+
+class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
+    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__(vllm_config=vllm_config, prefix=prefix)
+        assert self.version == (4, 0)
+
+    def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):
+        if isinstance(quant_config, (AWQConfig, AWQMarlinConfig)):
+            return None
+        return quant_config
+
+    def init_llm(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> nn.Module:
+        return LlamaForCausalLM(vllm_config=vllm_config, prefix=prefix)
+
+    def init_vision_module(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> nn.Module:
+        quant_config = self._maybe_ignore_quant_config(quant_config)
+        model = Idefics2VisionTransformer(config.vision_config,
+                                          quant_config=quant_config,
+                                          prefix=prefix)
+        if self.config.drop_vision_last_layer:
+            model.encoder.layers = model.encoder.layers[:-1]
+        return model
+
+    def init_resampler(
+        self,
+        embed_dim: int,
+        vision_dim: int,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> nn.Module:
+        quant_config = self._maybe_ignore_quant_config(quant_config)
+        with set_default_torch_dtype(torch.float16):
+            # The resampler in 4.0 remains consistent with the one in 2.5/2.6.
+            resampler = Resampler2_5(num_queries=self.config.query_num,
+                                     embed_dim=embed_dim,
+                                     num_heads=embed_dim // 128,
+                                     kv_dim=vision_dim,
+                                     quant_config=quant_config,
+                                     prefix=prefix)
+
+        return resampler.to(device=current_platform.device_type,
+                            dtype=torch.get_default_dtype())
+
+    def get_vision_hidden_states(
+            self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
+        pixel_values = data["pixel_values"]
+        tgt_sizes = data["tgt_sizes"]
+
+        B = len(pixel_values)
+        P = pixel_values[0].shape[-2]
+        L = max(item.shape[-1] for item in pixel_values)
+        device = pixel_values[0].device
+        dtype = pixel_values[0].dtype
+
+        all_pixel_values = torch.zeros((B, 3, P, L),
+                                       dtype=dtype,
+                                       device=device)
+        for i, pixel_values_item in enumerate(pixel_values):
+            L_item = pixel_values_item.shape[-1]
+            all_pixel_values[i, ..., :L_item] = pixel_values_item
+
+        num_patches = tgt_sizes.prod(-1)
+        max_patches = num_patches.max().item()
+        assert isinstance(max_patches, int)
+
+        patch_attn_mask = torch.zeros((B, max_patches),
+                                      dtype=torch.bool,
+                                      device=device)
+        for i, num_patches_item in enumerate(num_patches):
+            patch_attn_mask[i, :num_patches_item] = True
+
+        vision_embedding = self.vpm(
+            all_pixel_values,
+            patch_attention_mask=patch_attn_mask.unsqueeze(1),
+            tgt_sizes=tgt_sizes,
+        )
+
+        return self.resampler(vision_embedding, tgt_sizes)
+
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self,
+                                   skip_prefixes=["apm.", "audio", "tts"])
+        return loader.load_weights(weights)
+
 
 _SUPPORT_VERSION = {
     (2, 0): MiniCPMV2_0,
     (2, 5): MiniCPMV2_5,
     (2, 6): MiniCPMV2_6,
+    (4, 0): MiniCPMV4_0,
 }
 
 
@@ -1276,8 +1402,10 @@ class MiniCPMV(MiniCPMVBaseModel, SupportsMultiModal, SupportsLoRA):
         # Dispatch class based on version
         instance_cls = _SUPPORT_VERSION.get(version)
         if instance_cls is None:
-            raise ValueError(
-                "Currently, MiniCPMV only supports versions 2.0, 2.5, and 2.6")
+            supported_versions = ", ".join(
+                [f"{v[0]}.{v[1]}" for v in sorted(_SUPPORT_VERSION.keys())])
+            raise ValueError(f"Currently, MiniCPMV only supports versions "
+                             f"{supported_versions}. Got version: {version}")
 
         # quant_config references base class members,
         # so update values before init is called
diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py
index a4f8a361e..5414e2370 100644
--- a/vllm/model_executor/models/qwen2_vl.py
+++ b/vllm/model_executor/models/qwen2_vl.py
@@ -236,7 +236,7 @@ def apply_rotary_pos_emb_vision(t: torch.Tensor,
     sin = freqs.sin()
     apply_rotary_emb = apply_rotary_emb_torch
     if current_platform.is_cuda():
-        from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+        from flash_attn.layers.rotary import apply_rotary_emb
     output = apply_rotary_emb(t_, cos, sin).type_as(t)
     return output
 
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index e82e36638..6703c3539 100644
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -51,6 +51,8 @@ _TEXT_GENERATION_MODELS = {
     "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
     "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
     "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
+    "Ernie4_5_ForCausalLM": ("ernie45", "Ernie4_5_ForCausalLM"),
+    "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
     "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"),
     "FalconForCausalLM": ("falcon", "FalconForCausalLM"),
     "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py
index 48d1aacba..dee2bb7a5 100644
--- a/vllm/platforms/cuda.py
+++ b/vllm/platforms/cuda.py
@@ -291,7 +291,7 @@ class CudaPlatformBase(Platform):
         # installed.
         if target_backend == _Backend.FLASH_ATTN:
             try:
-                import vllm.vllm_flash_attn  # noqa: F401
+                import flash_attn  # noqa: F401
                 from vllm.attention.backends.flash_attn import (  # noqa: F401
                     FlashAttentionBackend, flash_attn_supports_fp8)
 
diff --git a/vllm/third_party/pymcml.py b/vllm/third_party/pymcml.py
new file mode 100644
index 000000000..097007ea0
--- /dev/null
+++ b/vllm/third_party/pymcml.py
@@ -0,0 +1,5439 @@
+# 2025 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
+#####
+# Copyright (c) 2011-2023, NVIDIA Corporation.  All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#    * Redistributions of source code must retain the above copyright notice,
+#      this list of conditions and the following disclaimer.
+#    * Redistributions in binary form must reproduce the above copyright
+#      notice, this list of conditions and the following disclaimer in the
+#      documentation and/or other materials provided with the distribution.
+#    * Neither the name of the NVIDIA Corporation nor the names of its
+#      contributors may be used to endorse or promote products derived from
+#      this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+# THE POSSIBILITY OF SUCH DAMAGE.
+#####
+
+##
+# Python bindings for the MXSMLEX library
+##
+from ctypes import *
+from ctypes.util import find_library
+from functools import wraps
+import sys
+import os
+import threading
+import string
+
+## C Type mappings ##
+## Enums
+_nvmlEnableState_t = c_uint
+MXSMLEX_FEATURE_DISABLED    = 0
+MXSMLEX_FEATURE_ENABLED     = 1
+
+_nvmlBrandType_t = c_uint
+MXSMLEX_BRAND_UNKNOWN             = 0
+MXSMLEX_BRAND_QUADRO              = 1
+MXSMLEX_BRAND_TESLA               = 2
+MXSMLEX_BRAND_NVS                 = 3
+MXSMLEX_BRAND_GRID                = 4   # Deprecated from API reporting. Keeping definition for backward compatibility.
+MXSMLEX_BRAND_GEFORCE             = 5
+MXSMLEX_BRAND_TITAN               = 6
+MXSMLEX_BRAND_NVIDIA_VAPPS        = 7   # NVIDIA Virtual Applications
+MXSMLEX_BRAND_NVIDIA_VPC          = 8   # NVIDIA Virtual PC
+MXSMLEX_BRAND_NVIDIA_VCS          = 9   # NVIDIA Virtual Compute Server
+MXSMLEX_BRAND_NVIDIA_VWS          = 10  # NVIDIA RTX Virtual Workstation
+MXSMLEX_BRAND_NVIDIA_CLOUD_GAMING = 11  # NVIDIA Cloud Gaming
+MXSMLEX_BRAND_NVIDIA_VGAMING      = MXSMLEX_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility.
+MXSMLEX_BRAND_QUADRO_RTX          = 12
+MXSMLEX_BRAND_NVIDIA_RTX          = 13
+MXSMLEX_BRAND_NVIDIA              = 14
+MXSMLEX_BRAND_GEFORCE_RTX         = 15  # Unused
+MXSMLEX_BRAND_TITAN_RTX           = 16  # Unused
+MXSMLEX_BRAND_COUNT               = 17
+
+_nvmlTemperatureThresholds_t = c_uint
+MXSMLEX_TEMPERATURE_THRESHOLD_SHUTDOWN      = 0
+MXSMLEX_TEMPERATURE_THRESHOLD_SLOWDOWN      = 1
+MXSMLEX_TEMPERATURE_THRESHOLD_MEM_MAX       = 2
+MXSMLEX_TEMPERATURE_THRESHOLD_GPU_MAX       = 3
+MXSMLEX_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN  = 4
+MXSMLEX_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5
+MXSMLEX_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX  = 6
+MXSMLEX_TEMPERATURE_THRESHOLD_COUNT         = 7
+
+_nvmlTemperatureSensors_t = c_uint
+MXSMLEX_TEMPERATURE_GPU     = 0
+MXSMLEX_TEMPERATURE_COUNT   = 1
+
+
+_nvmlComputeMode_t = c_uint
+MXSMLEX_COMPUTEMODE_DEFAULT           = 0
+MXSMLEX_COMPUTEMODE_EXCLUSIVE_THREAD  = 1  ## Support Removed
+MXSMLEX_COMPUTEMODE_PROHIBITED        = 2
+MXSMLEX_COMPUTEMODE_EXCLUSIVE_PROCESS = 3
+MXSMLEX_COMPUTEMODE_COUNT             = 4
+
+_nvmlMemoryLocation_t = c_uint
+MXSMLEX_MEMORY_LOCATION_L1_CACHE = 0
+MXSMLEX_MEMORY_LOCATION_L2_CACHE = 1
+MXSMLEX_MEMORY_LOCATION_DEVICE_MEMORY = 2
+MXSMLEX_MEMORY_LOCATION_DRAM = 2
+MXSMLEX_MEMORY_LOCATION_REGISTER_FILE = 3
+MXSMLEX_MEMORY_LOCATION_TEXTURE_MEMORY = 4
+MXSMLEX_MEMORY_LOCATION_TEXTURE_SHM = 5
+MXSMLEX_MEMORY_LOCATION_CBU = 6
+MXSMLEX_MEMORY_LOCATION_SRAM = 7
+MXSMLEX_MEMORY_LOCATION_COUNT = 8
+
+MXSMLEX_NVLINK_MAX_LINKS = 18
+
+# For backwards compatibility, maintain the incorrectly-named "LANES" define
+MXSMLEX_NVLINK_MAX_LANES = MXSMLEX_NVLINK_MAX_LINKS
+
+_nvmlNvLinkErrorCounter_t = c_uint
+MXSMLEX_NVLINK_ERROR_DL_REPLAY = 0
+MXSMLEX_NVLINK_ERROR_DL_RECOVERY = 1
+MXSMLEX_NVLINK_ERROR_DL_CRC_FLIT = 2
+MXSMLEX_NVLINK_ERROR_DL_CRC_DATA = 3
+MXSMLEX_NVLINK_ERROR_DL_ECC_DATA = 4
+MXSMLEX_NVLINK_ERROR_COUNT = 5
+
+_nvmlNvLinkEccLaneErrorCounter_t = c_uint
+MXSMLEX_NVLINK_ERROR_DL_ECC_LANE0 = 0
+MXSMLEX_NVLINK_ERROR_DL_ECC_LANE1 = 1
+MXSMLEX_NVLINK_ERROR_DL_ECC_LANE2 = 2
+MXSMLEX_NVLINK_ERROR_DL_ECC_LANE3 = 3
+MXSMLEX_NVLINK_ERROR_DL_ECC_COUNT = 5
+
+_nvmlNvLinkCapability_t = c_uint
+MXSMLEX_NVLINK_CAP_P2P_SUPPORTED = 0
+MXSMLEX_NVLINK_CAP_SYSMEM_ACCESS = 1
+MXSMLEX_NVLINK_CAP_P2P_ATOMICS   = 2
+MXSMLEX_NVLINK_CAP_SYSMEM_ATOMICS= 3
+MXSMLEX_NVLINK_CAP_SLI_BRIDGE    = 4
+MXSMLEX_NVLINK_CAP_VALID         = 5
+MXSMLEX_NVLINK_CAP_COUNT         = 6
+
+_nvmlNvLinkUtilizationCountPktTypes_t = c_uint
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_NOP        = 0x1
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_READ       = 0x2
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_WRITE      = 0x4
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_RATOM      = 0x8
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_NRATOM     = 0x10
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_FLUSH      = 0x20
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_RESPDATA   = 0x40
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80
+MXSMLEX_NVLINK_COUNTER_PKTFILTER_ALL        = 0xFF
+
+_nvmlNvLinkUtilizationCountUnits_t = c_uint
+MXSMLEX_NVLINK_COUNTER_UNIT_CYCLES   = 0
+MXSMLEX_NVLINK_COUNTER_UNIT_PACKETS  = 1
+MXSMLEX_NVLINK_COUNTER_UNIT_BYTES    = 2
+MXSMLEX_NVLINK_COUNTER_UNIT_RESERVED = 3
+MXSMLEX_NVLINK_COUNTER_UNIT_COUNT    = 4
+
+_nvmlNvLinkDeviceType_t = c_uint
+MXSMLEX_NVLINK_DEVICE_TYPE_GPU     = 0x00
+MXSMLEX_NVLINK_DEVICE_TYPE_IBMNPU  = 0x01
+MXSMLEX_NVLINK_DEVICE_TYPE_SWITCH  = 0x02
+MXSMLEX_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF
+
+# These are deprecated, instead use _nvmlMemoryErrorType_t
+_nvmlEccBitType_t = c_uint
+MXSMLEX_SINGLE_BIT_ECC    = 0
+MXSMLEX_DOUBLE_BIT_ECC    = 1
+MXSMLEX_ECC_ERROR_TYPE_COUNT = 2
+
+_nvmlEccCounterType_t = c_uint
+MXSMLEX_VOLATILE_ECC      = 0
+MXSMLEX_AGGREGATE_ECC     = 1
+MXSMLEX_ECC_COUNTER_TYPE_COUNT = 2
+
+_nvmlMemoryErrorType_t = c_uint
+MXSMLEX_MEMORY_ERROR_TYPE_CORRECTED   = 0
+MXSMLEX_MEMORY_ERROR_TYPE_UNCORRECTED = 1
+MXSMLEX_MEMORY_ERROR_TYPE_COUNT       = 2
+
+_nvmlClockType_t = c_uint
+MXSMLEX_CLOCK_GRAPHICS  = 0
+MXSMLEX_CLOCK_SM        = 1
+MXSMLEX_CLOCK_MEM       = 2
+MXSMLEX_CLOCK_VIDEO     = 3
+MXSMLEX_CLOCK_COUNT     = 4
+
+_nvmlClockId_t = c_uint
+MXSMLEX_CLOCK_ID_CURRENT            = 0
+MXSMLEX_CLOCK_ID_APP_CLOCK_TARGET   = 1
+MXSMLEX_CLOCK_ID_APP_CLOCK_DEFAULT  = 2
+MXSMLEX_CLOCK_ID_CUSTOMER_BOOST_MAX = 3
+MXSMLEX_CLOCK_ID_COUNT              = 4
+
+_nvmlDriverModel_t = c_uint
+MXSMLEX_DRIVER_WDDM       = 0
+MXSMLEX_DRIVER_WDM        = 1
+MXSMLEX_DRIVER_MCDM       = 2
+
+MXSMLEX_MAX_GPU_PERF_PSTATES = 16
+
+_nvmlPstates_t = c_uint
+MXSMLEX_PSTATE_0               = 0
+MXSMLEX_PSTATE_1               = 1
+MXSMLEX_PSTATE_2               = 2
+MXSMLEX_PSTATE_3               = 3
+MXSMLEX_PSTATE_4               = 4
+MXSMLEX_PSTATE_5               = 5
+MXSMLEX_PSTATE_6               = 6
+MXSMLEX_PSTATE_7               = 7
+MXSMLEX_PSTATE_8               = 8
+MXSMLEX_PSTATE_9               = 9
+MXSMLEX_PSTATE_10              = 10
+MXSMLEX_PSTATE_11              = 11
+MXSMLEX_PSTATE_12              = 12
+MXSMLEX_PSTATE_13              = 13
+MXSMLEX_PSTATE_14              = 14
+MXSMLEX_PSTATE_15              = 15
+MXSMLEX_PSTATE_UNKNOWN         = 32
+
+_nvmlInforomObject_t = c_uint
+MXSMLEX_INFOROM_OEM            = 0
+MXSMLEX_INFOROM_ECC            = 1
+MXSMLEX_INFOROM_POWER          = 2
+MXSMLEX_INFOROM_COUNT          = 3
+
+_nvmlReturn_t = c_uint
+MXSMLEX_SUCCESS                         = 0
+MXSMLEX_ERROR_UNINITIALIZED             = 1
+MXSMLEX_ERROR_INVALID_ARGUMENT          = 2
+MXSMLEX_ERROR_NOT_SUPPORTED             = 3
+MXSMLEX_ERROR_NO_PERMISSION             = 4
+MXSMLEX_ERROR_ALREADY_INITIALIZED       = 5
+MXSMLEX_ERROR_NOT_FOUND                 = 6
+MXSMLEX_ERROR_INSUFFICIENT_SIZE         = 7
+MXSMLEX_ERROR_INSUFFICIENT_POWER        = 8
+MXSMLEX_ERROR_DRIVER_NOT_LOADED         = 9
+MXSMLEX_ERROR_TIMEOUT                   = 10
+MXSMLEX_ERROR_IRQ_ISSUE                 = 11
+MXSMLEX_ERROR_LIBRARY_NOT_FOUND         = 12
+MXSMLEX_ERROR_FUNCTION_NOT_FOUND        = 13
+MXSMLEX_ERROR_CORRUPTED_INFOROM         = 14
+MXSMLEX_ERROR_GPU_IS_LOST               = 15
+MXSMLEX_ERROR_RESET_REQUIRED            = 16
+MXSMLEX_ERROR_OPERATING_SYSTEM          = 17
+MXSMLEX_ERROR_LIB_RM_VERSION_MISMATCH   = 18
+MXSMLEX_ERROR_IN_USE                    = 19
+MXSMLEX_ERROR_MEMORY                    = 20
+MXSMLEX_ERROR_NO_DATA                   = 21
+MXSMLEX_ERROR_VGPU_ECC_NOT_SUPPORTED    = 22
+MXSMLEX_ERROR_INSUFFICIENT_RESOURCES    = 23
+MXSMLEX_ERROR_FREQ_NOT_SUPPORTED        = 24
+MXSMLEX_ERROR_ARGUMENT_VERSION_MISMATCH = 25
+MXSMLEX_ERROR_DEPRECATED                = 26
+MXSMLEX_ERROR_NOT_READY                 = 27
+MXSMLEX_ERROR_GPU_NOT_FOUND             = 28
+MXSMLEX_ERROR_INVALID_STATE             = 29
+MXSMLEX_ERROR_UNKNOWN                   = 999
+
+_nvmlFanState_t = c_uint
+MXSMLEX_FAN_NORMAL             = 0
+MXSMLEX_FAN_FAILED             = 1
+
+_nvmlFanControlPolicy_t = c_uint
+MXSMLEX_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0
+MXSMLEX_FAN_POLICY_MANUAL                   = 1
+
+_nvmlLedColor_t = c_uint
+MXSMLEX_LED_COLOR_GREEN        = 0
+MXSMLEX_LED_COLOR_AMBER        = 1
+
+_nvmlGpuOperationMode_t = c_uint
+MXSMLEX_GOM_ALL_ON                 = 0
+MXSMLEX_GOM_COMPUTE                = 1
+MXSMLEX_GOM_LOW_DP                 = 2
+
+_nvmlPageRetirementCause_t = c_uint
+MXSMLEX_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0
+MXSMLEX_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR           = 1
+MXSMLEX_PAGE_RETIREMENT_CAUSE_COUNT                          = 2
+
+_nvmlRestrictedAPI_t = c_uint
+MXSMLEX_RESTRICTED_API_SET_APPLICATION_CLOCKS                = 0
+MXSMLEX_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS               = 1
+MXSMLEX_RESTRICTED_API_COUNT                                 = 2
+
+_nvmlBridgeChipType_t = c_uint
+MXSMLEX_BRIDGE_CHIP_PLX = 0
+MXSMLEX_BRIDGE_CHIP_BRO4 = 1
+MXSMLEX_MAX_PHYSICAL_BRIDGE = 128
+
+_nvmlValueType_t = c_uint
+MXSMLEX_VALUE_TYPE_DOUBLE = 0
+MXSMLEX_VALUE_TYPE_UNSIGNED_INT = 1
+MXSMLEX_VALUE_TYPE_UNSIGNED_LONG = 2
+MXSMLEX_VALUE_TYPE_UNSIGNED_LONG_LONG = 3
+MXSMLEX_VALUE_TYPE_SIGNED_LONG_LONG = 4
+MXSMLEX_VALUE_TYPE_SIGNED_INT = 5
+MXSMLEX_VALUE_TYPE_COUNT = 6
+
+_nvmlPerfPolicyType_t = c_uint
+MXSMLEX_PERF_POLICY_POWER = 0
+MXSMLEX_PERF_POLICY_THERMAL = 1
+MXSMLEX_PERF_POLICY_SYNC_BOOST = 2
+MXSMLEX_PERF_POLICY_BOARD_LIMIT = 3
+MXSMLEX_PERF_POLICY_LOW_UTILIZATION = 4
+MXSMLEX_PERF_POLICY_RELIABILITY = 5
+MXSMLEX_PERF_POLICY_TOTAL_APP_CLOCKS = 10
+MXSMLEX_PERF_POLICY_TOTAL_BASE_CLOCKS = 11
+MXSMLEX_PERF_POLICY_COUNT = 12
+
+_nvmlEncoderQueryType_t = c_uint
+MXSMLEX_ENCODER_QUERY_H264 = 0
+MXSMLEX_ENCODER_QUERY_HEVC = 1
+MXSMLEX_ENCODER_QUERY_AV1 = 2
+MXSMLEX_ENCODER_QUERY_UNKNOWN = 255
+
+_nvmlFBCSessionType_t = c_uint
+MXSMLEX_FBC_SESSION_TYPE_UNKNOWN = 0
+MXSMLEX_FBC_SESSION_TYPE_TOSYS = 1
+MXSMLEX_FBC_SESSION_TYPE_CUDA = 2
+MXSMLEX_FBC_SESSION_TYPE_VID = 3
+MXSMLEX_FBC_SESSION_TYPE_HWENC = 4
+
+_nvmlDetachGpuState_t = c_uint
+MXSMLEX_DETACH_GPU_KEEP = 0
+MXSMLEX_DETACH_GPU_REMOVE = 1
+
+_nvmlPcieLinkState_t = c_uint
+MXSMLEX_PCIE_LINK_KEEP = 0
+MXSMLEX_PCIE_LINK_SHUT_DOWN = 1
+
+_nvmlSamplingType_t = c_uint
+MXSMLEX_TOTAL_POWER_SAMPLES = 0
+MXSMLEX_GPU_UTILIZATION_SAMPLES = 1
+MXSMLEX_MEMORY_UTILIZATION_SAMPLES = 2
+MXSMLEX_ENC_UTILIZATION_SAMPLES = 3
+MXSMLEX_DEC_UTILIZATION_SAMPLES = 4
+MXSMLEX_PROCESSOR_CLK_SAMPLES = 5
+MXSMLEX_MEMORY_CLK_SAMPLES = 6
+MXSMLEX_MODULE_POWER_SAMPLES = 7
+MXSMLEX_JPG_UTILIZATION_SAMPLES = 8
+MXSMLEX_OFA_UTILIZATION_SAMPLES = 9
+MXSMLEX_SAMPLINGTYPE_COUNT = 10
+
+_nvmlPcieUtilCounter_t = c_uint
+MXSMLEX_PCIE_UTIL_TX_BYTES = 0
+MXSMLEX_PCIE_UTIL_RX_BYTES = 1
+MXSMLEX_PCIE_UTIL_COUNT = 2
+
+_nvmlGpuTopologyLevel_t = c_uint
+MXSMLEX_TOPOLOGY_INTERNAL = 0
+MXSMLEX_TOPOLOGY_SINGLE = 10
+MXSMLEX_TOPOLOGY_MULTIPLE = 20
+MXSMLEX_TOPOLOGY_HOSTBRIDGE = 30
+MXSMLEX_TOPOLOGY_NODE = 40
+MXSMLEX_TOPOLOGY_CPU = MXSMLEX_TOPOLOGY_NODE
+MXSMLEX_TOPOLOGY_SYSTEM = 50
+
+_nvmlGpuP2PCapsIndex_t = c_uint
+MXSMLEX_P2P_CAPS_INDEX_READ = 0,
+MXSMLEX_P2P_CAPS_INDEX_WRITE = 1
+#MXSMLEX_P2P_CAPS_INDEX_MXLINK =2
+NVML_P2P_CAPS_INDEX_NVLINK = 2
+MXSMLEX_P2P_CAPS_INDEX_ATOMICS = 3
+#
+# MXSMLEX_P2P_CAPS_INDEX_PROP is deprecated.
+# Use MXSMLEX_P2P_CAPS_INDEX_PCI instead.
+#
+MXSMLEX_P2P_CAPS_INDEX_PROP = 4
+MXSMLEX_P2P_CAPS_INDEX_PCI = 4
+MXSMLEX_P2P_CAPS_INDEX_UNKNOWN = 5
+
+_nvmlGpuP2PStatus_t = c_uint
+#MXSMLEX_P2P_STATUS_OK     = 0
+NVML_P2P_STATUS_OK = 0
+MXSMLEX_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1
+MXSMLEX_P2P_STATUS_CHIPSET_NOT_SUPPORTED = MXSMLEX_P2P_STATUS_CHIPSET_NOT_SUPPORED
+MXSMLEX_P2P_STATUS_GPU_NOT_SUPPORTED = 2
+MXSMLEX_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED =3
+MXSMLEX_P2P_STATUS_DISABLED_BY_REGKEY =4
+MXSMLEX_P2P_STATUS_NOT_SUPPORTED =5
+MXSMLEX_P2P_STATUS_UNKNOWN =6
+
+_nvmlDeviceArchitecture_t = c_uint
+MXSMLEX_DEVICE_ARCH_KEPLER   = 2
+MXSMLEX_DEVICE_ARCH_MAXWELL  = 3
+MXSMLEX_DEVICE_ARCH_PASCAL   = 4
+MXSMLEX_DEVICE_ARCH_VOLTA    = 5
+MXSMLEX_DEVICE_ARCH_TURING   = 6
+MXSMLEX_DEVICE_ARCH_AMPERE   = 7
+MXSMLEX_DEVICE_ARCH_ADA      = 8
+MXSMLEX_DEVICE_ARCH_HOPPER   = 9
+MXSMLEX_DEVICE_ARCH_UNKNOWN  = 0xffffffff
+
+# PCI bus Types
+_nvmlBusType_t = c_uint
+MXSMLEX_BUS_TYPE_UNKNOWN = 0
+MXSMLEX_BUS_TYPE_PCI     = 1
+MXSMLEX_BUS_TYPE_PCIE    = 2
+MXSMLEX_BUS_TYPE_FPCI    = 3
+MXSMLEX_BUS_TYPE_AGP     = 4
+
+_nvmlPowerSource_t = c_uint
+MXSMLEX_POWER_SOURCE_AC         = 0x00000000
+MXSMLEX_POWER_SOURCE_BATTERY    = 0x00000001
+MXSMLEX_POWER_SOURCE_UNDERSIZED = 0x00000002
+
+_nvmlAdaptiveClockInfoStatus_t = c_uint
+MXSMLEX_ADAPTIVE_CLOCKING_INFO_STATUS_DISABLED = 0x00000000
+MXSMLEX_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001
+
+_nvmlClockLimitId_t = c_uint
+MXSMLEX_CLOCK_LIMIT_ID_RANGE_START = 0xffffff00
+MXSMLEX_CLOCK_LIMIT_ID_TDP         = 0xffffff01
+MXSMLEX_CLOCK_LIMIT_ID_UNLIMITED   = 0xffffff02
+
+_nvmlPcieLinkMaxSpeed_t = c_uint
+MXSMLEX_PCIE_LINK_MAX_SPEED_INVALID   = 0x00000000
+MXSMLEX_PCIE_LINK_MAX_SPEED_2500MBPS  = 0x00000001
+MXSMLEX_PCIE_LINK_MAX_SPEED_5000MBPS  = 0x00000002
+MXSMLEX_PCIE_LINK_MAX_SPEED_8000MBPS  = 0x00000003
+MXSMLEX_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004
+MXSMLEX_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005
+MXSMLEX_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006
+
+_nvmlAffinityScope_t = c_uint
+MXSMLEX_AFFINITY_SCOPE_NODE   = 0
+MXSMLEX_AFFINITY_SCOPE_SOCKET = 1
+
+# C preprocessor defined values
+nvmlFlagDefault             = 0
+nvmlFlagForce               = 1
+MXSMLEX_INIT_FLAG_NO_GPUS      = 1
+MXSMLEX_INIT_FLAG_NO_ATTACH    = 2
+
+MXSMLEX_MAX_GPC_COUNT          = 32
+
+# buffer size
+MXSMLEX_DEVICE_INFOROM_VERSION_BUFFER_SIZE      = 16
+MXSMLEX_DEVICE_UUID_BUFFER_SIZE                 = 80
+MXSMLEX_DEVICE_UUID_V2_BUFFER_SIZE              = 96
+MXSMLEX_SYSTEM_DRIVER_VERSION_BUFFER_SIZE       = 80
+MXSMLEX_SYSTEM_MXSMLEX_VERSION_BUFFER_SIZE         = 80
+MXSMLEX_DEVICE_NAME_BUFFER_SIZE                 = 64
+MXSMLEX_DEVICE_NAME_V2_BUFFER_SIZE              = 96
+MXSMLEX_DEVICE_SERIAL_BUFFER_SIZE               = 30
+MXSMLEX_DEVICE_PART_NUMBER_BUFFER_SIZE          = 80
+MXSMLEX_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE      = 80
+MXSMLEX_DEVICE_VBIOS_VERSION_BUFFER_SIZE        = 32
+MXSMLEX_DEVICE_PCI_BUS_ID_BUFFER_SIZE           = 32
+MXSMLEX_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE        = 16
+MXSMLEX_GRID_LICENSE_BUFFER_SIZE                = 128
+MXSMLEX_VGPU_NAME_BUFFER_SIZE                   = 64
+MXSMLEX_GRID_LICENSE_FEATURE_MAX_COUNT          = 3
+MXSMLEX_VGPU_METADATA_OPAQUE_DATA_SIZE          = sizeof(c_uint) + 256
+MXSMLEX_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE     = 256
+MXSMLEX_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE  = 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH
+
+# Format strings
+MXSMLEX_DEVICE_PCI_BUS_ID_LEGACY_FMT   = "%04X:%02X:%02X.0"
+MXSMLEX_DEVICE_PCI_BUS_ID_FMT          = "%08X:%02X:%02X.0"
+
+MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1)
+MXSMLEX_VALUE_NOT_AVAILABLE_uint = c_uint(-1)
+
+'''
+ Field Identifiers.
+
+ All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change.
+'''
+MXSMLEX_FI_DEV_ECC_CURRENT          = 1   # Current ECC mode. 1=Active. 0=Inactive
+MXSMLEX_FI_DEV_ECC_PENDING          = 2   # Pending ECC mode. 1=Active. 0=Inactive
+
+#ECC Count Totals
+MXSMLEX_FI_DEV_ECC_SBE_VOL_TOTAL    = 3   # Total single bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_VOL_TOTAL    = 4   # Total double bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_AGG_TOTAL    = 5   # Total single bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_AGG_TOTAL    = 6   # Total double bit aggregate (persistent) ECC errors
+#Individual ECC locations
+MXSMLEX_FI_DEV_ECC_SBE_VOL_L1       = 7   # L1 cache single bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_VOL_L1       = 8   # L1 cache double bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_VOL_L2       = 9   # L2 cache single bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_VOL_L2       = 10  # L2 cache double bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_VOL_DEV      = 11  # Device memory single bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_VOL_DEV      = 12  # Device memory double bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_VOL_REG      = 13  # Register file single bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_VOL_REG      = 14  # Register file double bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_VOL_TEX      = 15  # Texture memory single bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_VOL_TEX      = 16  # Texture memory double bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_VOL_CBU      = 17  # CBU double bit volatile ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_AGG_L1       = 18  # L1 cache single bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_AGG_L1       = 19  # L1 cache double bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_AGG_L2       = 20  # L2 cache single bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_AGG_L2       = 21  # L2 cache double bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_AGG_DEV      = 22  # Device memory single bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_AGG_DEV      = 23  # Device memory double bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_AGG_REG      = 24  # Register File single bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_AGG_REG      = 25  # Register File double bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_SBE_AGG_TEX      = 26  # Texture memory single bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_AGG_TEX      = 27  # Texture memory double bit aggregate (persistent) ECC errors
+MXSMLEX_FI_DEV_ECC_DBE_AGG_CBU      = 28  # CBU double bit aggregate ECC errors
+
+# Page Retirement
+MXSMLEX_FI_DEV_RETIRED_SBE          = 29  # Number of retired pages because of single bit errors
+MXSMLEX_FI_DEV_RETIRED_DBE          = 30  # Number of retired pages because of double bit errors
+MXSMLEX_FI_DEV_RETIRED_PENDING      = 31  # If any pages are pending retirement. 1=yes. 0=no.
+
+# NvLink Flit Error Counters
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0   = 32 # NVLink flow control CRC  Error Counter for Lane 0
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1   = 33 # NVLink flow control CRC  Error Counter for Lane 1
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2   = 34 # NVLink flow control CRC  Error Counter for Lane 2
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3   = 35 # NVLink flow control CRC  Error Counter for Lane 3
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4   = 36 # NVLink flow control CRC  Error Counter for Lane 4
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5   = 37 # NVLink flow control CRC  Error Counter for Lane 5
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = 38 # NVLink flow control CRC  Error Counter total for all Lanes
+
+# NvLink CRC Data Error Counters
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0   = 39 # NVLink data CRC Error Counter for Lane 0
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1   = 40 # NVLink data CRC Error Counter for Lane 1
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2   = 41 # NVLink data CRC Error Counter for Lane 2
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3   = 42 # NVLink data CRC Error Counter for Lane 3
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4   = 43 # NVLink data CRC Error Counter for Lane 4
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5   = 44 # NVLink data CRC Error Counter for Lane 5
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = 45 # NvLink data CRC Error Counter total for all Lanes
+
+# NvLink Replay Error Counters
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0     = 46 # NVLink Replay Error Counter for Lane 0
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1     = 47 # NVLink Replay Error Counter for Lane 1
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2     = 48 # NVLink Replay Error Counter for Lane 2
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3     = 49 # NVLink Replay Error Counter for Lane 3
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4     = 50 # NVLink Replay Error Counter for Lane 4
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5     = 51 # NVLink Replay Error Counter for Lane 5
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL  = 52 # NVLink Replay Error Counter total for all Lanes
+
+# NvLink Recovery Error Counters
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0   = 53 # NVLink Recovery Error Counter for Lane 0
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1   = 54 # NVLink Recovery Error Counter for Lane 1
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2   = 55 # NVLink Recovery Error Counter for Lane 2
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3   = 56 # NVLink Recovery Error Counter for Lane 3
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4   = 57 # NVLink Recovery Error Counter for Lane 4
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5   = 58 # NVLink Recovery Error Counter for Lane 5
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = 59 # NVLink Recovery Error Counter total for all Lanes
+
+# NvLink Bandwidth Counters
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L0    = 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L1    = 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L2    = 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L3    = 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L4    = 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L5    = 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes
+
+# NvLink Bandwidth Counters
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L0    = 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L1    = 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L2    = 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L3    = 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L4    = 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L5    = 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes
+
+# Perf Policy Counters
+MXSMLEX_FI_DEV_PERF_POLICY_POWER             = 74   # Perf Policy Counter for Power Policy
+MXSMLEX_FI_DEV_PERF_POLICY_THERMAL           = 75   # Perf Policy Counter for Thermal Policy
+MXSMLEX_FI_DEV_PERF_POLICY_SYNC_BOOST        = 76   # Perf Policy Counter for Sync boost Policy
+MXSMLEX_FI_DEV_PERF_POLICY_BOARD_LIMIT       = 77   # Perf Policy Counter for Board Limit
+MXSMLEX_FI_DEV_PERF_POLICY_LOW_UTILIZATION   = 78   # Perf Policy Counter for Low GPU Utilization Policy
+MXSMLEX_FI_DEV_PERF_POLICY_RELIABILITY       = 79   # Perf Policy Counter for Reliability Policy
+MXSMLEX_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS  = 80   # Perf Policy Counter for Total App Clock Policy
+MXSMLEX_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = 81   # Perf Policy Counter for Total Base Clocks Policy
+
+# Memory temperatures
+MXSMLEX_FI_DEV_MEMORY_TEMP  = 82 # Memory temperature for the device
+
+# Energy Counter
+MXSMLEX_FI_DEV_TOTAL_ENERGY_CONSUMPTION = 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded
+
+# NVLink Speed
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L0     = 84
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L1     = 85
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L2     = 86
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L3     = 87
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L4     = 88
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L5     = 89
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90
+
+# NVLink Link Count
+MXSMLEX_FI_DEV_NVLINK_LINK_COUNT = 91
+
+# Page Retirement pending fields
+MXSMLEX_FI_DEV_RETIRED_PENDING_SBE = 92
+MXSMLEX_FI_DEV_RETIRED_PENDING_DBE = 93
+
+# PCIe replay and replay rollover counters
+MXSMLEX_FI_DEV_PCIE_REPLAY_COUNTER = 94
+MXSMLEX_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95
+
+# NvLink Flit Error Counters
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6   = 96 # NVLink flow control CRC  Error Counter for Lane 6
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7   = 97 # NVLink flow control CRC  Error Counter for Lane 7
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8   = 98 # NVLink flow control CRC  Error Counter for Lane 8
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9   = 99 # NVLink flow control CRC  Error Counter for Lane 9
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10  = 100 # NVLink flow control CRC  Error Counter for Lane 10
+MXSMLEX_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11  = 101 # NVLink flow control CRC  Error Counter for Lane 11
+
+# NvLink CRC Data Error Counters
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6   = 102 # NVLink data CRC Error Counter for Lane 6
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7   = 103 # NVLink data CRC Error Counter for Lane 7
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8   = 104 # NVLink data CRC Error Counter for Lane 8
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9   = 105 # NVLink data CRC Error Counter for Lane 9
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10  = 106 # NVLink data CRC Error Counter for Lane 10
+MXSMLEX_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11  = 107 # NVLink data CRC Error Counter for Lane 11
+
+# NvLink Replay Error Counters
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6     = 108 # NVLink Replay Error Counter for Lane 6
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7     = 109 # NVLink Replay Error Counter for Lane 7
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8     = 110 # NVLink Replay Error Counter for Lane 8
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9     = 111 # NVLink Replay Error Counter for Lane 9
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10    = 112 # NVLink Replay Error Counter for Lane 10
+MXSMLEX_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11    = 113 # NVLink Replay Error Counter for Lane 11
+
+# NvLink Recovery Error Counters
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6   = 114 # NVLink Recovery Error Counter for Lane 6
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7   = 115 # NVLink Recovery Error Counter for Lane 7
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8   = 116 # NVLink Recovery Error Counter for Lane 8
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9   = 117 # NVLink Recovery Error Counter for Lane 9
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10  = 118 # NVLink Recovery Error Counter for Lane 10
+MXSMLEX_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11  = 119 # NVLink Recovery Error Counter for Lane 11
+
+# NvLink Bandwidth Counters
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L6    = 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L7    = 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L8    = 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L9    = 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L10   = 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C0_L11   = 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11
+
+# NvLink Bandwidth Counters
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L6    = 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L7    = 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L8    = 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L9    = 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L10   = 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10
+MXSMLEX_FI_DEV_NVLINK_BANDWIDTH_C1_L11   = 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11
+
+# NVLink Speed
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L6     = 132
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L7     = 133
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L8     = 134
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L9     = 135
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L10    = 136
+MXSMLEX_FI_DEV_NVLINK_SPEED_MBPS_L11    = 137
+
+# NVLink Throughput Counters
+MXSMLEX_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB
+MXSMLEX_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB
+MXSMLEX_FI_DEV_NVLINK_THROUGHPUT_RAW_TX  = 140 # NVLink TX Data + protocol overhead in KiB
+MXSMLEX_FI_DEV_NVLINK_THROUGHPUT_RAW_RX  = 141 # NVLink RX Data + protocol overhead in KiB
+
+# Row Remapper
+MXSMLEX_FI_DEV_REMAPPED_COR        = 142
+MXSMLEX_FI_DEV_REMAPPED_UNC        = 143
+MXSMLEX_FI_DEV_REMAPPED_PENDING    = 144
+MXSMLEX_FI_DEV_REMAPPED_FAILURE    = 145
+
+#Remote device NVLink ID
+MXSMLEX_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146
+
+# Number of NVLinks connected to NVSwitch
+MXSMLEX_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147
+
+# NvLink ECC Data Error Counters
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0    = 148 #< NVLink data ECC Error Counter for Link 0
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1    = 149 #< NVLink data ECC Error Counter for Link 1
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2    = 150 #< NVLink data ECC Error Counter for Link 2
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3    = 151 #< NVLink data ECC Error Counter for Link 3
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4    = 152 #< NVLink data ECC Error Counter for Link 4
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5    = 153 #< NVLink data ECC Error Counter for Link 5
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6    = 154 #< NVLink data ECC Error Counter for Link 6
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7    = 155 #< NVLink data ECC Error Counter for Link 7
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8    = 156 #< NVLink data ECC Error Counter for Link 8
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9    = 157 #< NVLink data ECC Error Counter for Link 9
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10   = 158 #< NVLink data ECC Error Counter for Link 10
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11   = 159 #< NVLink data ECC Error Counter for Link 11
+MXSMLEX_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = 160 #< NvLink data ECC Error Counter total for all Links
+
+MXSMLEX_FI_DEV_NVLINK_ERROR_DL_REPLAY            = 161
+MXSMLEX_FI_DEV_NVLINK_ERROR_DL_RECOVERY          = 162
+MXSMLEX_FI_DEV_NVLINK_ERROR_DL_CRC               = 163
+MXSMLEX_FI_DEV_NVLINK_GET_SPEED                  = 164
+MXSMLEX_FI_DEV_NVLINK_GET_STATE                  = 165
+MXSMLEX_FI_DEV_NVLINK_GET_VERSION                = 166
+
+MXSMLEX_FI_DEV_NVLINK_GET_POWER_STATE            = 167
+MXSMLEX_FI_DEV_NVLINK_GET_POWER_THRESHOLD        = 168
+
+MXSMLEX_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER       = 169
+
+MXSMLEX_FI_DEV_C2C_LINK_COUNT                    = 170
+MXSMLEX_FI_DEV_C2C_LINK_GET_STATUS               = 171
+MXSMLEX_FI_DEV_C2C_LINK_GET_MAX_BW               = 172
+
+MXSMLEX_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS     = 173
+MXSMLEX_FI_DEV_PCIE_COUNT_NAKS_RECEIVED          = 174
+MXSMLEX_FI_DEV_PCIE_COUNT_RECEIVER_ERROR         = 175
+MXSMLEX_FI_DEV_PCIE_COUNT_BAD_TLP                = 176
+MXSMLEX_FI_DEV_PCIE_COUNT_NAKS_SENT              = 177
+MXSMLEX_FI_DEV_PCIE_COUNT_BAD_DLLP               = 178
+MXSMLEX_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR        = 179
+MXSMLEX_FI_DEV_PCIE_COUNT_FATAL_ERROR            = 180
+MXSMLEX_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ        = 181
+MXSMLEX_FI_DEV_PCIE_COUNT_LCRC_ERROR             = 182
+MXSMLEX_FI_DEV_PCIE_COUNT_LANE_ERROR             = 183
+
+MXSMLEX_FI_DEV_IS_RESETLESS_MIG_SUPPORTED        = 184
+
+MXSMLEX_FI_DEV_POWER_AVERAGE                     = 185
+MXSMLEX_FI_DEV_POWER_INSTANT                     = 186
+MXSMLEX_FI_DEV_POWER_MIN_LIMIT                   = 187
+MXSMLEX_FI_DEV_POWER_MAX_LIMIT                   = 188
+MXSMLEX_FI_DEV_POWER_DEFAULT_LIMIT               = 189
+MXSMLEX_FI_DEV_POWER_CURRENT_LIMIT               = 190
+MXSMLEX_FI_DEV_ENERGY                            = 191
+MXSMLEX_FI_DEV_POWER_REQUESTED_LIMIT             = 192
+
+MXSMLEX_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT       = 193
+MXSMLEX_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT       = 194
+MXSMLEX_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT        = 195
+MXSMLEX_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT        = 196
+
+MXSMLEX_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE   = 199
+
+MXSMLEX_FI_MAX = 200 # One greater than the largest field ID defined above
+
+
+## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode
+MXSMLEX_GPU_VIRTUALIZATION_MODE_NONE        = 0  # Represents Bare Metal GPU
+MXSMLEX_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = 1  # Device is associated with GPU-Passthorugh
+MXSMLEX_GPU_VIRTUALIZATION_MODE_VGPU        = 2  # Device is associated with vGPU inside virtual machine.
+MXSMLEX_GPU_VIRTUALIZATION_MODE_HOST_VGPU   = 3  # Device is associated with VGX hypervisor in vGPU mode
+MXSMLEX_GPU_VIRTUALIZATION_MODE_HOST_VSGA   = 4  # Device is associated with VGX hypervisor in vSGA mode
+
+## Lib loading ##
+nvmlLib = None
+libLoadLock = threading.Lock()
+_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown
+
+## vGPU Management
+_nvmlVgpuTypeId_t   = c_uint
+_nvmlVgpuInstance_t = c_uint
+
+_nvmlVgpuVmIdType_t = c_uint
+MXSMLEX_VGPU_VM_ID_DOMAIN_ID    = 0
+MXSMLEX_VGPU_VM_ID_UUID         = 1
+
+_nvmlGridLicenseFeatureCode_t = c_uint
+MXSMLEX_GRID_LICENSE_FEATURE_CODE_UNKNOWN      = 0
+MXSMLEX_GRID_LICENSE_FEATURE_CODE_VGPU         = 1
+MXSMLEX_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX   = 2
+MXSMLEX_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = 2 # deprecated, use MXSMLEX_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX.
+MXSMLEX_GRID_LICENSE_FEATURE_CODE_GAMING       = 3
+MXSMLEX_GRID_LICENSE_FEATURE_CODE_COMPUTE      = 4
+
+_nvmlGridLicenseExpiryStatus_t = c_uint8
+MXSMLEX_GRID_LICENSE_EXPIRY_NOT_AVAILABLE    = 0,   # Expiry information not available
+MXSMLEX_GRID_LICENSE_EXPIRY_INVALID          = 1,   # Invalid expiry or error fetching expiry
+MXSMLEX_GRID_LICENSE_EXPIRY_VALID            = 2,   # Valid expiry
+MXSMLEX_GRID_LICENSE_EXPIRY_NOT_APPLICABLE   = 3,   # Expiry not applicable
+MXSMLEX_GRID_LICENSE_EXPIRY_PERMANENT        = 4,   # Permanent expiry
+
+_nvmlVgpuCapability_t = c_uint
+MXSMLEX_VGPU_CAP_NVLINK_P2P                    = 0  # vGPU P2P over NVLink is supported
+MXSMLEX_VGPU_CAP_GPUDIRECT                     = 1  # GPUDirect capability is supported
+MXSMLEX_VGPU_CAP_MULTI_VGPU_EXCLUSIVE          = 2  # vGPU profile cannot be mixed with other vGPU profiles in same VM
+MXSMLEX_VGPU_CAP_EXCLUSIVE_TYPE                = 3  # vGPU profile cannot run on a GPU alongside other profiles of different type
+MXSMLEX_VGPU_CAP_EXCLUSIVE_SIZE                = 4  # vGPU profile cannot run on a GPU alongside other profiles of different size
+MXSMLEX_VGPU_CAP_COUNT                         = 5
+
+_nvmlVgpuDriverCapability_t = c_uint
+MXSMLEX_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU          = 0  # Supports mixing of different vGPU profiles within one guest VM
+MXSMLEX_VGPU_DRIVER_CAP_COUNT                             = 1
+
+_nvmlDeviceVgpuCapability_t = c_uint
+MXSMLEX_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU             = 0  # Query if the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations
+MXSMLEX_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES  = 1  # Query if the GPU supports concurrent execution of timesliced vGPU profiles of differing types
+MXSMLEX_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES     = 2  # Query if the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes
+MXSMLEX_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW             = 3  # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second
+MXSMLEX_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW            = 4  # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second
+MXSMLEX_DEVICE_VGPU_CAP_DEVICE_STREAMING                  = 5  # Query if vGPU profiles on the GPU supports migration data streaming
+MXSMLEX_DEVICE_VGPU_CAP_MINI_QUARTER_GPU                  = 6  # Set/Get support of mini-quarter vGPU profiles
+MXSMLEX_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU          = 7  # Set/Get support for compute media engine vGPU profiles
+MXSMLEX_DEVICE_VGPU_CAP_COUNT                             = 8
+
+_nvmlVgpuGuestInfoState_t = c_uint
+MXSMLEX_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0
+MXSMLEX_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED   = 1
+
+_nvmlVgpuVmCompatibility_t = c_uint
+MXSMLEX_VGPU_VM_COMPATIBILITY_NONE         = 0x0
+MXSMLEX_VGPU_VM_COMPATIBILITY_COLD         = 0x1
+MXSMLEX_VGPU_VM_COMPATIBILITY_HIBERNATE    = 0x2
+MXSMLEX_VGPU_VM_COMPATIBILITY_SLEEP        = 0x4
+MXSMLEX_VGPU_VM_COMPATIBILITY_LIVE         = 0x8
+
+_nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint
+MXSMLEX_VGPU_COMPATIBILITY_LIMIT_NONE          = 0x0
+MXSMLEX_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER   = 0x1
+MXSMLEX_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER  = 0x2
+MXSMLEX_VGPU_COMPATIBILITY_LIMIT_GPU           = 0x4
+MXSMLEX_VGPU_COMPATIBILITY_LIMIT_OTHER         = 0x80000000
+
+_nvmlHostVgpuMode_t = c_uint
+MXSMLEX_HOST_VGPU_MODE_NON_SRIOV   = 0
+MXSMLEX_HOST_VGPU_MODE_SRIOV       = 1
+
+_nvmlConfComputeGpusReadyState_t = c_uint
+MXSMLEX_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0
+MXSMLEX_CC_ACCEPTING_CLIENT_REQUESTS_TRUE = 1
+
+_nvmlConfComputeGpuCaps_t = c_uint
+MXSMLEX_CC_SYSTEM_GPUS_CC_NOT_CAPABLE = 0
+MXSMLEX_CC_SYSTEM_GPUS_CC_CAPABLE = 1
+
+_nvmlConfComputeCpuCaps_t = c_uint
+MXSMLEX_CC_SYSTEM_CPU_CAPS_NONE = 0
+MXSMLEX_CC_SYSTEM_CPU_CAPS_AMD_SEV = 1
+MXSMLEX_CC_SYSTEM_CPU_CAPS_INTEL_TDX = 2
+
+_nvmlConfComputeDevToolsMode_t = c_uint
+MXSMLEX_CC_SYSTEM_DEVTOOLS_MODE_OFF = 0
+MXSMLEX_CC_SYSTEM_DEVTOOLS_MODE_ON = 1
+ 
+MXSMLEX_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0
+MXSMLEX_CC_SYSTEM_ENVIRONMENT_SIM = 1
+MXSMLEX_CC_SYSTEM_ENVIRONMENT_PROD = 2
+ 
+_nvmlConfComputeCcFeature_t = c_uint
+MXSMLEX_CC_SYSTEM_FEATURE_DISABLED = 0
+MXSMLEX_CC_SYSTEM_FEATURE_ENABLED = 1
+
+_nvmlConfComputeCcKeyRotationThreshAttackerAdv_t = c_uint
+MXSMLEX_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MIN = 50
+MXSMLEX_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MAX = 75
+
+# GSP firmware
+MXSMLEX_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40
+
+class MXSMLEXLibraryMismatchError(Exception):
+    pass
+
+## Error Checking ##
+class NVMLError(Exception):
+    _valClassMapping = dict()
+    # List of currently known error codes
+    _errcode_to_string = {
+        MXSMLEX_ERROR_UNINITIALIZED:       "Uninitialized",
+        MXSMLEX_ERROR_INVALID_ARGUMENT:    "Invalid Argument",
+        MXSMLEX_ERROR_NOT_SUPPORTED:       "Not Supported",
+        MXSMLEX_ERROR_NO_PERMISSION:       "Insufficient Permissions",
+        MXSMLEX_ERROR_ALREADY_INITIALIZED: "Already Initialized",
+        MXSMLEX_ERROR_NOT_FOUND:           "Not Found",
+        MXSMLEX_ERROR_INSUFFICIENT_SIZE:   "Insufficient Size",
+        MXSMLEX_ERROR_INSUFFICIENT_POWER:  "Insufficient External Power",
+        MXSMLEX_ERROR_DRIVER_NOT_LOADED:   "Driver Not Loaded",
+        MXSMLEX_ERROR_TIMEOUT:             "Timeout",
+        MXSMLEX_ERROR_IRQ_ISSUE:           "Interrupt Request Issue",
+        MXSMLEX_ERROR_LIBRARY_NOT_FOUND:   "MXSMLEX Shared Library Not Found",
+        MXSMLEX_ERROR_FUNCTION_NOT_FOUND:  "Function Not Found",
+        MXSMLEX_ERROR_CORRUPTED_INFOROM:   "Corrupted infoROM",
+        MXSMLEX_ERROR_GPU_IS_LOST:         "GPU is lost",
+        MXSMLEX_ERROR_RESET_REQUIRED:      "GPU requires restart",
+        MXSMLEX_ERROR_OPERATING_SYSTEM:    "The operating system has blocked the request.",
+        MXSMLEX_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an MXSMLEX/RM version mismatch.",
+        MXSMLEX_ERROR_MEMORY:              "Insufficient Memory",
+        MXSMLEX_ERROR_UNKNOWN:             "Unknown Error",
+        }
+    def __new__(typ, value):
+        '''
+        Maps value to a proper subclass of NVMLError.
+        See _extractNVMLErrorsAsClasses function for more details
+        '''
+        if typ == NVMLError:
+            typ = NVMLError._valClassMapping.get(value, typ)
+        obj = Exception.__new__(typ)
+        obj.value = value
+        return obj
+    def __str__(self):
+        try:
+            if self.value not in NVMLError._errcode_to_string:
+                NVMLError._errcode_to_string[self.value] = str(nvmlErrorString(self.value))
+            return NVMLError._errcode_to_string[self.value]
+        except NVMLError:
+            return "MXSMLEX Error with code %d" % self.value
+    def __eq__(self, other):
+        return self.value == other.value
+
+def nvmlExceptionClass(nvmlErrorCode):
+    if nvmlErrorCode not in NVMLError._valClassMapping:
+        raise ValueError('nvmlErrorCode %s is not valid' % nvmlErrorCode)
+    return NVMLError._valClassMapping[nvmlErrorCode]
+
+def _extractNVMLErrorsAsClasses():
+    '''
+    Generates a hierarchy of classes on top of NVMLError class.
+
+    Each MXSMLEX Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate
+    exceptions more easily.
+
+    NVMLError is a parent class. Each MXSMLEX_ERROR_* gets it's own subclass.
+    e.g. MXSMLEX_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized
+    '''
+    this_module = sys.modules[__name__]
+    nvmlErrorsNames = [x for x in dir(this_module) if x.startswith("MXSMLEX_ERROR_")]
+    for err_name in nvmlErrorsNames:
+        # e.g. Turn MXSMLEX_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized
+        class_name = "NVMLError_" + string.capwords(err_name.replace("MXSMLEX_ERROR_", ""), "_").replace("_", "")
+        err_val = getattr(this_module, err_name)
+        def gen_new(val):
+            def new(typ):
+                obj = NVMLError.__new__(typ, val)
+                return obj
+            return new
+        new_error_class = type(class_name, (NVMLError,), {'__new__': gen_new(err_val)})
+        new_error_class.__module__ = __name__
+        setattr(this_module, class_name, new_error_class)
+        NVMLError._valClassMapping[err_val] = new_error_class
+_extractNVMLErrorsAsClasses()
+
+def _nvmlCheckReturn(ret):
+    if (ret != MXSMLEX_SUCCESS):
+        raise NVMLError(ret)
+    return ret
+
+## Function access ##
+_nvmlGetFunctionPointer_cache = dict() # function pointers are cached to prevent unnecessary libLoadLock locking
+def _nvmlGetFunctionPointer(name):
+    global nvmlLib
+
+    if name in _nvmlGetFunctionPointer_cache:
+        return _nvmlGetFunctionPointer_cache[name]
+
+    libLoadLock.acquire()
+    try:
+        # ensure library was loaded
+        if (nvmlLib == None):
+            raise NVMLError(MXSMLEX_ERROR_UNINITIALIZED)
+        try:
+            _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name)
+            return _nvmlGetFunctionPointer_cache[name]
+        except AttributeError:
+            print("nvml error")
+            #raise NVMLError(MXSMLEX_ERROR_FUNCTION_NOT_FOUND)
+    finally:
+        # lock is always freed
+        libLoadLock.release()
+
+## Alternative object
+# Allows the object to be printed
+# Allows mismatched types to be assigned
+#  - like None when the Structure variant requires c_uint
+class nvmlFriendlyObject(object):
+    def __init__(self, dictionary):
+        for x in dictionary:
+            setattr(self, x, dictionary[x])
+    def __str__(self):
+        return self.__dict__.__str__()
+
+def nvmlStructToFriendlyObject(struct):
+    d = {}
+    for x in struct._fields_:
+        key = x[0]
+        value = getattr(struct, key)
+        # only need to convert from bytes if bytes, no need to check python version.
+        d[key] = value.decode() if isinstance(value, bytes) else value
+    obj = nvmlFriendlyObject(d)
+    return obj
+
+# pack the object so it can be passed to the MXSMLEX library
+def nvmlFriendlyObjectToStruct(obj, model):
+    for x in model._fields_:
+        key = x[0]
+        value = obj.__dict__[key]
+        # any c_char_p in python3 needs to be bytes, default encoding works fine.
+        if sys.version_info >= (3,):
+            setattr(model, key, value.encode())
+        else:
+            setattr(model, key, value)
+    return model
+
+## Unit structures
+class struct_c_nvmlUnit_t(Structure):
+    pass # opaque handle
+c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t)
+
+class _PrintableStructure(Structure):
+    """
+    Abstract class that produces nicer __str__ output than ctypes.Structure.
+    e.g. instead of:
+      >>> print str(obj)
+      <class_name object at 0x7fdf82fef9e0>
+    this class will print
+      class_name(field_name: formatted_value, field_name: formatted_value)
+
+    _fmt_ dictionary of <str _field_ name> -> <str format>
+    e.g. class that has _field_ 'hex_value', c_uint could be formatted with
+      _fmt_ = {"hex_value" : "%08X"}
+    to produce nicer output.
+    Default fomratting string for all fields can be set with key "<default>" like:
+      _fmt_ = {"<default>" : "%d MHz"} # e.g all values are numbers in MHz.
+    If not set it's assumed to be just "%s"
+
+    Exact format of returned str from this class is subject to change in the future.
+    """
+    _fmt_ = {}
+    def __str__(self):
+        result = []
+        for x in self._fields_:
+            key = x[0]
+            value = getattr(self, key)
+            fmt = "%s"
+            if key in self._fmt_:
+                fmt = self._fmt_[key]
+            elif "<default>" in self._fmt_:
+                fmt = self._fmt_["<default>"]
+            result.append(("%s: " + fmt) % (key, value))
+        return self.__class__.__name__ + "(" +  ", ".join(result) + ")"
+
+    def __getattribute__(self, name):
+        res = super(_PrintableStructure, self).__getattribute__(name)
+        # need to convert bytes to unicode for python3 don't need to for python2
+        # Python 2 strings are of both str and bytes
+        # Python 3 strings are not of type bytes
+        # ctypes should convert everything to the correct values otherwise
+        if isinstance(res, bytes):
+            if isinstance(res, str):
+                return res
+            return res.decode()
+        return res
+
+    def __setattr__(self, name, value):
+        if isinstance(value, str):
+            # encoding a python2 string returns the same value, since python2 strings are bytes already
+            # bytes passed in python3 will be ignored.
+            value = value.encode()
+        super(_PrintableStructure, self).__setattr__(name, value)
+
+class c_nvmlUnitInfo_t(_PrintableStructure):
+    _fields_ = [
+        ('name', c_char * 96),
+        ('id', c_char * 96),
+        ('serial', c_char * 96),
+        ('firmwareVersion', c_char * 96),
+    ]
+
+class c_nvmlC2cModeInfo_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('isC2cEnabled', c_uint)
+    ]
+
+nvmlC2cModeInfo_v1 = 0x1000008;
+
+class c_nvmlLedState_t(_PrintableStructure):
+    _fields_ = [
+        ('cause', c_char * 256),
+        ('color', _nvmlLedColor_t),
+    ]
+
+class c_nvmlPSUInfo_t(_PrintableStructure):
+    _fields_ = [
+        ('state', c_char * 256),
+        ('current', c_uint),
+        ('voltage', c_uint),
+        ('power', c_uint),
+    ]
+
+class c_nvmlUnitFanInfo_t(_PrintableStructure):
+    _fields_ = [
+        ('speed', c_uint),
+        ('state', _nvmlFanState_t),
+    ]
+
+class c_nvmlUnitFanSpeeds_t(_PrintableStructure):
+    _fields_ = [
+        ('fans', c_nvmlUnitFanInfo_t * 24),
+        ('count', c_uint)
+    ]
+
+## Device structures
+class struct_c_nvmlDevice_t(Structure):
+    pass # opaque handle
+c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t)
+
+class nvmlPciInfoExt_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('domain', c_uint),
+        ('bus', c_uint),
+        ('device', c_uint),
+        ('pciDeviceId', c_uint),
+        ('pciSubSystemId', c_uint),
+        ('baseClass', c_uint),
+        ('subClass', c_uint),
+        ('busId', c_char * MXSMLEX_DEVICE_PCI_BUS_ID_BUFFER_SIZE),
+    ]
+    _fmt_ = {
+            'version'        : "0x%04X",
+            'domain'         : "0x%04X",
+            'bus'            : "0x%02X",
+            'device'         : "0x%02X",
+            'pciDeviceId'    : "0x%08X",
+            'pciSubSystemId' : "0x%08X",
+            'baseClass'      : "0x%01X",
+            'subClass'       : "0x%01X",
+            }
+
+nvmlPciInfoExt_v1 = 0x1000040
+
+# Legacy pciInfo used for _v1 and _v2
+class nvmlPciInfo_v2_t(_PrintableStructure):
+    _fields_ = [
+        ('busId', c_char * MXSMLEX_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE),
+        ('domain', c_uint),
+        ('bus', c_uint),
+        ('device', c_uint),
+        ('pciDeviceId', c_uint),
+
+        # Added in 2.285
+        ('pciSubSystemId', c_uint),
+        ('reserved0', c_uint),
+        ('reserved1', c_uint),
+        ('reserved2', c_uint),
+        ('reserved3', c_uint),
+    ]
+    _fmt_ = {
+            'domain'         : "0x%04X",
+            'bus'            : "0x%02X",
+            'device'         : "0x%02X",
+            'pciDeviceId'    : "0x%08X",
+            'pciSubSystemId' : "0x%08X",
+            }
+
+class nvmlPciInfo_t(_PrintableStructure):
+    _fields_ = [
+        # Moved to the new busId location below
+        ('busIdLegacy', c_char * MXSMLEX_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE),
+        ('domain', c_uint),
+        ('bus', c_uint),
+        ('device', c_uint),
+        ('pciDeviceId', c_uint),
+
+        # Added in 2.285
+        ('pciSubSystemId', c_uint),
+        # New busId replaced the long deprecated and reserved fields with a
+        # field of the same size in 9.0
+        ('busId', c_char * MXSMLEX_DEVICE_PCI_BUS_ID_BUFFER_SIZE),
+    ]
+    _fmt_ = {
+            'domain'         : "0x%08X",
+            'bus'            : "0x%02X",
+            'device'         : "0x%02X",
+            'pciDeviceId'    : "0x%08X",
+            'pciSubSystemId' : "0x%08X",
+            }
+
+class c_nvmlExcludedDeviceInfo_t(_PrintableStructure):
+    _fields_ = [
+        ('pci', nvmlPciInfo_t),
+        ('uuid', c_char * MXSMLEX_DEVICE_UUID_BUFFER_SIZE)
+    ]
+
+class nvmlNvLinkUtilizationControl_t(_PrintableStructure):
+    _fields_ = [
+        ('units', _nvmlNvLinkUtilizationCountUnits_t),
+        ('pktfilter', _nvmlNvLinkUtilizationCountPktTypes_t),
+    ]
+
+class c_nvmlMemory_t(_PrintableStructure):
+    _fields_ = [
+        ('total', c_ulonglong),
+        ('free', c_ulonglong),
+        ('used', c_ulonglong),
+    ]
+    _fmt_ = {'<default>': "%d B"}
+
+class c_nvmlMemory_v2_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('total', c_ulonglong),
+        ('reserved', c_ulonglong),
+        ('free', c_ulonglong),
+        ('used', c_ulonglong),
+    ]
+    _fmt_ = {'<default>': "%d B"}
+
+nvmlMemory_v2 = 0x02000028
+
+class c_nvmlBAR1Memory_t(_PrintableStructure):
+    _fields_ = [
+        ('bar1Total', c_ulonglong),
+        ('bar1Free', c_ulonglong),
+        ('bar1Used', c_ulonglong),
+    ]
+    _fmt_ = {'<default>': "%d B"}
+
+class nvmlClkMonFaultInfo_t(Structure):
+    _fields_ = [("clkApiDomain", c_uint),
+                ("clkDomainFaultMask", c_uint)
+    ]
+
+class nvmlClkMonStatus_t(Structure):
+    _fields_ = [("bGlobalStatus", c_uint),
+                ("clkMonListSize", c_uint),
+                ("clkMonList", nvmlClkMonFaultInfo_t)
+    ]
+
+# On Windows with the WDDM driver, usedGpuMemory is reported as None
+# Code that processes this structure should check for None, I.E.
+#
+# if (info.usedGpuMemory == None):
+#     # TODO handle the error
+#     pass
+# else:
+#    print("Using %d MiB of memory" % (info.usedGpuMemory / 1024 / 1024))
+# endif
+#
+# See MXSMLEX documentation for more information
+class c_nvmlProcessInfo_v2_t(_PrintableStructure):
+    _fields_ = [
+        ('pid', c_uint),
+        ('usedGpuMemory', c_ulonglong),
+        ('gpuInstanceId', c_uint),
+        ('computeInstanceId', c_uint),
+    ]
+    _fmt_ = {'usedGpuMemory': "%d B"}
+
+c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t
+
+c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t
+
+_nvmlProcessMode_t = c_uint
+MXSMLEX_PROCESS_MODE_COMPUTE  = 0
+MXSMLEX_PROCESS_MODE_GRAPHICS = 1
+MXSMLEX_PROCESS_MODE_MPS      = 2
+
+class c_nvmlProcessDetail_v1_t(Structure):
+    _fields_ = [
+        ('pid', c_uint),
+        ('usedGpuMemory', c_ulonglong),
+        ('gpuInstanceId', c_uint),
+        ('computeInstanceId', c_uint),
+        ('usedGpuCcProtectedMemory', c_ulonglong),
+    ]
+
+class c_nvmlProcessDetailList_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('mode', _nvmlProcessMode_t),
+        ('numProcArrayEntries', c_uint),
+        ('procArray', POINTER(c_nvmlProcessDetail_v1_t)),
+    ]
+    _fmt_ = {'numProcArrayEntries': "%d B"}
+
+c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t
+
+nvmlProcessDetailList_v1 = 0x1000018
+
+class c_nvmlBridgeChipInfo_t(_PrintableStructure):
+    _fields_ = [
+        ('type', _nvmlBridgeChipType_t),
+        ('fwVersion', c_uint),
+    ]
+
+class c_nvmlBridgeChipHierarchy_t(_PrintableStructure):
+    _fields_ = [
+        ('bridgeCount', c_uint),
+        ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128),
+    ]
+
+class c_nvmlEccErrorCounts_t(_PrintableStructure):
+    _fields_ = [
+        ('l1Cache', c_ulonglong),
+        ('l2Cache', c_ulonglong),
+        ('deviceMemory', c_ulonglong),
+        ('registerFile', c_ulonglong),
+    ]
+
+class c_nvmlUtilization_t(_PrintableStructure):
+    _fields_ = [
+        ('gpu', c_uint),
+        ('memory', c_uint),
+    ]
+    _fmt_ = {'<default>': "%d %%"}
+
+# Added in 2.285
+class c_nvmlHwbcEntry_t(_PrintableStructure):
+    _fields_ = [
+        ('hwbcId', c_uint),
+        ('firmwareVersion', c_char * 32),
+    ]
+
+class c_nvmlValue_t(Union):
+    _fields_ = [
+        ('dVal', c_double),
+        ('uiVal', c_uint),
+        ('ulVal', c_ulong),
+        ('ullVal', c_ulonglong),
+        ('sllVal', c_longlong),
+        ('siVal', c_int),
+    ]
+
+class c_nvmlSample_t(_PrintableStructure):
+    _fields_ = [
+        ('timeStamp', c_ulonglong),
+        ('sampleValue', c_nvmlValue_t),
+    ]
+
+class c_nvmlViolationTime_t(_PrintableStructure):
+    _fields_ = [
+        ('referenceTime', c_ulonglong),
+        ('violationTime', c_ulonglong),
+    ]
+
+class c_nvmlFieldValue_t(_PrintableStructure):
+    _fields_ = [
+        ('fieldId', c_uint32),
+        ('scopeId', c_uint32),
+        ('timestamp', c_int64),
+        ('latencyUsec', c_int64),
+        ('valueType', _nvmlValueType_t),
+        ('nvmlReturn', _nvmlReturn_t),
+        ('value', c_nvmlValue_t)
+    ]
+
+class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('mode', c_uint),
+    ]
+
+VgpuHeterogeneousMode_v1 = 0x1000008
+
+class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('placementId', c_uint),
+    ]
+
+VgpuPlacementId_v1 = 0x1000008
+
+class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('count', c_uint),
+        ('placementSize', c_uint),
+        ('placementIds', POINTER(c_uint)),
+    ]
+
+VgpuPlacementList_v1 = 0x1000018
+
+class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure):
+    _fields_ = [
+        ('vgpuInstance', _nvmlVgpuInstance_t),
+        ('timeStamp', c_ulonglong),
+        ('smUtil', c_nvmlValue_t),
+        ('memUtil', c_nvmlValue_t),
+        ('encUtil', c_nvmlValue_t),
+        ('decUtil', c_nvmlValue_t),
+    ]
+
+class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('timeStamp', c_ulonglong),
+        ('vgpuInstance', _nvmlVgpuInstance_t),
+        ('smUtil', c_nvmlValue_t),
+        ('memUtil', c_nvmlValue_t),
+        ('encUtil', c_nvmlValue_t),
+        ('decUtil', c_nvmlValue_t),
+        ('jpgUtil', c_nvmlValue_t),
+        ('ofaUtil', c_nvmlValue_t),
+    ]
+
+class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('sampleValType', _nvmlValueType_t),
+        ('vgpuInstanceCount', c_uint),
+        ('lastSeenTimeStamp', c_ulonglong),
+        ('vgpuUtilArray', POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)),
+    ]
+
+VgpuInstancesUtilizationInfo_v1 = 0x01000020
+
+class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure):
+    _fields_ = [
+        ('vgpuInstance', _nvmlVgpuInstance_t),
+        ('pid', c_uint),
+        ('processName', c_char * MXSMLEX_VGPU_NAME_BUFFER_SIZE),
+        ('timeStamp', c_ulonglong),
+        ('smUtil', c_uint),
+        ('memUtil', c_uint),
+        ('encUtil', c_uint),
+        ('decUtil', c_uint),
+    ]
+
+class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('processName', c_char * MXSMLEX_VGPU_NAME_BUFFER_SIZE),
+        ('timeStamp', c_ulonglong),
+        ('vgpuInstance', _nvmlVgpuInstance_t),
+        ('pid', c_uint),
+        ('smUtil', c_uint),
+        ('memUtil', c_uint),
+        ('encUtil', c_uint),
+        ('decUtil', c_uint),
+        ('jpgUtil', c_uint),
+        ('ofaUtil', c_uint),
+    ]
+
+class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('vgpuProcessCount', c_uint),
+        ('lastSeenTimeStamp', c_ulonglong),
+        ('vgpuProcUtilArray', POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)),
+    ]
+
+VgpuProcessesUtilizationInfo_v1 = 0x01000018
+
+class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure):
+    _fields_ = [
+        ('year',    c_uint32),
+        ('month',   c_uint16),
+        ('day',     c_uint16),
+        ('hour',    c_uint16),
+        ('min',     c_uint16),
+        ('sec',     c_uint16),
+        ('status',  c_uint8),
+    ]
+
+MXSMLEX_GRID_LICENSE_STATE_UNKNOWN                 = 0
+MXSMLEX_GRID_LICENSE_STATE_UNINITIALIZED           = 1
+MXSMLEX_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2
+MXSMLEX_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED   = 3
+MXSMLEX_GRID_LICENSE_STATE_UNLICENSED              = 4
+MXSMLEX_GRID_LICENSE_STATE_LICENSED                = 5
+
+class c_nvmlVgpuLicenseInfo_t(_PrintableStructure):
+    _fields_ = [
+        ('isLicensed',      c_uint8),
+        ('licenseExpiry',   c_nvmlVgpuLicenseExpiry_t),
+        ('currentState',    c_uint),
+    ]
+
+class c_nvmlEncoderSession_t(_PrintableStructure):
+    _fields_ = [
+        ('sessionId', c_uint),
+        ('pid', c_uint),
+        ('vgpuInstance', _nvmlVgpuInstance_t),
+        ('codecType', c_uint),
+        ('hResolution', c_uint),
+        ('vResolution', c_uint),
+        ('averageFps', c_uint),
+        ('encodeLatency', c_uint),
+    ]
+
+class c_nvmlProcessUtilizationSample_t(_PrintableStructure):
+    _fields_ = [
+        ('pid', c_uint),
+        ('timeStamp', c_ulonglong),
+        ('smUtil', c_uint),
+        ('memUtil', c_uint),
+        ('encUtil', c_uint),
+        ('decUtil', c_uint),
+    ]
+
+class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('timeStamp', c_ulonglong),
+        ('pid', c_uint),
+        ('smUtil', c_uint),
+        ('memUtil', c_uint),
+        ('encUtil', c_uint),
+        ('decUtil', c_uint),
+        ('jpgUtil', c_uint),
+        ('ofaUtil', c_uint),
+    ]
+
+class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('processSamplesCount', c_uint),
+        ('lastSeenTimeStamp', c_ulonglong),
+        ('procUtilArray', POINTER(c_nvmlProcessUtilizationInfo_v1_t)),
+    ]
+
+ProcessesUtilizationInfo_v1 = 0x01000018
+
+class c_nvmlGridLicenseExpiry_t(_PrintableStructure):
+    _fields_ = [
+        ('year',    c_uint32),
+        ('month',   c_uint16),
+        ('day',     c_uint16),
+        ('hour',    c_uint16),
+        ('min',     c_uint16),
+        ('sec',     c_uint16),
+        ('status',  c_uint8),
+    ]
+
+class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure):
+    _fields_ = [
+        ('featureCode',    _nvmlGridLicenseFeatureCode_t),
+        ('featureState',   c_uint),
+        ('licenseInfo',    c_char * MXSMLEX_GRID_LICENSE_BUFFER_SIZE),
+        ('productName',    c_char * MXSMLEX_GRID_LICENSE_BUFFER_SIZE),
+        ('featureEnabled', c_uint),
+        ('licenseExpiry',  c_nvmlGridLicenseExpiry_t),
+    ]
+
+class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure):
+    _fields_ = [
+        ('isGridLicenseSupported',  c_int),
+        ('licensableFeaturesCount', c_uint),
+        ('gridLicensableFeatures',  c_nvmlGridLicensableFeature_v4_t * MXSMLEX_GRID_LICENSE_FEATURE_MAX_COUNT),
+    ]
+
+class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure):
+    _fields_ = [
+        ('featureCode', _nvmlGridLicenseFeatureCode_t),
+        ('featureState', c_uint),
+        ('licenseInfo', c_char * MXSMLEX_GRID_LICENSE_BUFFER_SIZE),
+        ('productName', c_char * MXSMLEX_GRID_LICENSE_BUFFER_SIZE),
+        ('featureEnabled', c_uint),
+    ]
+
+class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure):
+    _fields_ = [
+        ('isGridLicenseSupported', c_int),
+        ('licensableFeaturesCount', c_uint),
+        ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v3_t * MXSMLEX_GRID_LICENSE_FEATURE_MAX_COUNT),
+    ]
+
+class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure):
+    _fields_ = [
+        ('featureCode', _nvmlGridLicenseFeatureCode_t),
+        ('featureState', c_uint),
+        ('licenseInfo', c_char * MXSMLEX_GRID_LICENSE_BUFFER_SIZE),
+        ('productName', c_char * MXSMLEX_GRID_LICENSE_BUFFER_SIZE),
+    ]
+
+class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure):
+    _fields_ = [
+        ('isGridLicenseSupported', c_int),
+        ('licensableFeaturesCount', c_uint),
+        ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v2_t * MXSMLEX_GRID_LICENSE_FEATURE_MAX_COUNT),
+    ]
+
+class c_nvmlGridLicensableFeature_t(_PrintableStructure):
+    _fields_ = [
+        ('featureCode', _nvmlGridLicenseFeatureCode_t),
+        ('featureState', c_uint),
+        ('licenseInfo', c_char * MXSMLEX_GRID_LICENSE_BUFFER_SIZE),
+    ]
+
+class c_nvmlGridLicensableFeatures_t(_PrintableStructure):
+    _fields_ = [
+        ('isGridLicenseSupported', c_int),
+        ('licensableFeaturesCount', c_uint),
+        ('gridLicensableFeatures', c_nvmlGridLicensableFeature_t * MXSMLEX_GRID_LICENSE_FEATURE_MAX_COUNT),
+    ]
+
+## Event structures
+class struct_c_nvmlEventSet_t(Structure):
+    pass # opaque handle
+c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t)
+
+nvmlEventTypeSingleBitEccError     = 0x0000000000000001
+nvmlEventTypeDoubleBitEccError     = 0x0000000000000002
+nvmlEventTypePState                = 0x0000000000000004
+nvmlEventTypeXidCriticalError      = 0x0000000000000008
+nvmlEventTypeClock                 = 0x0000000000000010
+nvmlEventTypePowerSourceChange     = 0x0000000000000080
+nvmlEventMigConfigChange           = 0x0000000000000100
+nvmlEventTypeNone                  = 0x0000000000000000
+nvmlEventTypeAll                   = (
+                                        nvmlEventTypeNone
+                                        | nvmlEventTypeSingleBitEccError
+                                        | nvmlEventTypeDoubleBitEccError
+                                        | nvmlEventTypePState
+                                        | nvmlEventTypeClock
+                                        | nvmlEventTypePowerSourceChange
+                                        | nvmlEventTypeXidCriticalError
+                                        | nvmlEventMigConfigChange
+                                     )
+
+## Clock Event Reasons defines
+nvmlClocksEventReasonGpuIdle              = 0x0000000000000001
+nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002
+nvmlClocksEventReasonUserDefinedClocks         = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting
+nvmlClocksEventReasonSwPowerCap           = 0x0000000000000004
+nvmlClocksEventReasonHwSlowdown           = 0x0000000000000008
+nvmlClocksEventReasonSyncBoost            = 0x0000000000000010
+nvmlClocksEventReasonSwThermalSlowdown    = 0x0000000000000020
+nvmlClocksEventReasonHwThermalSlowdown    = 0x0000000000000040
+nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080
+nvmlClocksEventReasonDisplayClockSetting  = 0x0000000000000100
+nvmlClocksEventReasonNone                 = 0x0000000000000000
+nvmlClocksEventReasonAll                  = (
+                                                  nvmlClocksEventReasonNone |
+                                                  nvmlClocksEventReasonGpuIdle |
+                                                  nvmlClocksEventReasonApplicationsClocksSetting |
+                                                  nvmlClocksEventReasonSwPowerCap |
+                                                  nvmlClocksEventReasonHwSlowdown |
+                                                  nvmlClocksEventReasonSyncBoost |
+                                                  nvmlClocksEventReasonSwThermalSlowdown |
+                                                  nvmlClocksEventReasonHwThermalSlowdown |
+                                                  nvmlClocksEventReasonHwPowerBrakeSlowdown |
+                                                  nvmlClocksEventReasonDisplayClockSetting
+                                               )
+
+## Following have been deprecated
+nvmlClocksThrottleReasonGpuIdle              = 0x0000000000000001
+nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002
+nvmlClocksThrottleReasonUserDefinedClocks         = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting
+nvmlClocksThrottleReasonSwPowerCap           = 0x0000000000000004
+nvmlClocksThrottleReasonHwSlowdown           = 0x0000000000000008
+nvmlClocksThrottleReasonSyncBoost            = 0x0000000000000010
+nvmlClocksThrottleReasonSwThermalSlowdown    = 0x0000000000000020
+nvmlClocksThrottleReasonHwThermalSlowdown    = 0x0000000000000040
+nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080
+nvmlClocksThrottleReasonDisplayClockSetting  = 0x0000000000000100
+nvmlClocksThrottleReasonNone                 = 0x0000000000000000
+nvmlClocksThrottleReasonAll                  = (
+                                                  nvmlClocksThrottleReasonNone |
+                                                  nvmlClocksThrottleReasonGpuIdle |
+                                                  nvmlClocksThrottleReasonApplicationsClocksSetting |
+                                                  nvmlClocksThrottleReasonSwPowerCap |
+                                                  nvmlClocksThrottleReasonHwSlowdown |
+                                                  nvmlClocksThrottleReasonSyncBoost |
+                                                  nvmlClocksThrottleReasonSwThermalSlowdown |
+                                                  nvmlClocksThrottleReasonHwThermalSlowdown |
+                                                  nvmlClocksThrottleReasonHwPowerBrakeSlowdown |
+                                                  nvmlClocksThrottleReasonDisplayClockSetting
+                                               )
+
+class c_nvmlEventData_t(_PrintableStructure):
+    _fields_ = [
+        ('device', c_nvmlDevice_t),
+        ('eventType', c_ulonglong),
+        ('eventData', c_ulonglong),
+        ('gpuInstanceId', c_uint),
+        ('computeInstanceId', c_uint)
+    ]
+    _fmt_ = {'eventType': "0x%08X"}
+
+class c_nvmlAccountingStats_t(_PrintableStructure):
+    _fields_ = [
+        ('gpuUtilization', c_uint),
+        ('memoryUtilization', c_uint),
+        ('maxMemoryUsage', c_ulonglong),
+        ('time', c_ulonglong),
+        ('startTime', c_ulonglong),
+        ('isRunning', c_uint),
+        ('reserved', c_uint * 5)
+    ]
+
+class c_nvmlVgpuVersion_t(Structure):
+    _fields_ = [("minVersion", c_uint),
+                ("maxVersion", c_uint)
+               ]
+
+class c_nvmlVgpuMetadata_t(_PrintableStructure):
+    _fields_ = [("version", c_uint),
+                ("revision", c_uint),
+                ("guestInfoState", _nvmlVgpuGuestInfoState_t),
+                ("guestDriverVersion", c_char * MXSMLEX_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
+                ("hostDriverVersion", c_char * MXSMLEX_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
+                ("reserved", c_uint * 6),
+                ("vgpuVirtualizationCaps", c_uint),
+                ("guestVgpuVersion", c_uint),
+                ("opaqueDataSize", c_uint),
+                ("opaqueData", c_char * MXSMLEX_VGPU_METADATA_OPAQUE_DATA_SIZE)
+               ]
+
+class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure):
+    _fields_ = [("version", c_uint),
+                ("revision", c_uint),
+                ("hostDriverVersion", c_char * MXSMLEX_SYSTEM_DRIVER_VERSION_BUFFER_SIZE),
+                ("pgpuVirtualizationCaps", c_uint),
+                ("reserved", c_uint * 5),
+                ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t),
+                ("opaqueDataSize", c_uint),
+                ("opaqueData", c_char * MXSMLEX_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE)
+               ]
+
+class c_nvmlVgpuPgpuCompatibility_t(Structure):
+    _fields_ = [("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t),
+                ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t)
+               ]
+
+## vGPU scheduler policy defines
+MXSMLEX_VGPU_SCHEDULER_POLICY_UNKNOWN      = 0
+MXSMLEX_VGPU_SCHEDULER_POLICY_BEST_EFFORT  = 1
+MXSMLEX_VGPU_SCHEDULER_POLICY_EQUAL_SHARE  = 2
+MXSMLEX_VGPU_SCHEDULER_POLICY_FIXED_SHARE  = 3
+
+## Supported vGPU scheduler policy count
+MXSMLEX_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT  = 3
+
+MXSMLEX_SCHEDULER_SW_MAX_LOG_ENTRIES           = 200
+
+MXSMLEX_VGPU_SCHEDULER_ARR_DEFAULT   = 0
+MXSMLEX_VGPU_SCHEDULER_ARR_DISABLE   = 1
+MXSMLEX_VGPU_SCHEDULER_ARR_ENABLE    = 2
+
+class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure):
+    _fields_ = [
+        ('avgFactor',   c_uint),
+        ('timeslice',   c_uint),
+    ]
+
+class c_nvmlVgpuSchedData_t(_PrintableStructure):
+    _fields_ = [
+        ('timeslice',   c_uint),
+    ]
+
+class c_nvmlVgpuSchedulerParams_t(Union):
+    _fields_ = [
+        ('vgpuSchedDataWithARR', c_nvmlVgpuSchedDataWithARR_t),
+        ('vgpuSchedData',        c_nvmlVgpuSchedData_t),
+    ]
+
+class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure):
+    _fields_ = [
+        ('timestamp',                   c_ulonglong),
+        ('timeRunTotal',                c_ulonglong),
+        ('timeRun',                     c_ulonglong),
+        ('swRunlistId',                 c_uint),
+        ('targetTimeSlice',             c_ulonglong),
+        ('cumulativePreemptionTime',    c_ulonglong),
+    ]
+
+class c_nvmlVgpuSchedulerLog_t(_PrintableStructure):
+    _fields_ = [
+        ('engineId',        c_uint),
+        ('schedulerPolicy', c_uint),
+        ('arrMode',         c_uint),
+        ('schedulerParams', c_nvmlVgpuSchedulerParams_t),
+        ('entriesCount',    c_uint),
+        ('logEntries',      c_nvmlVgpuSchedulerLogEntry_t * MXSMLEX_SCHEDULER_SW_MAX_LOG_ENTRIES),
+    ]
+
+class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure):
+    _fields_ = [
+        ('schedulerPolicy', c_uint),
+        ('arrMode',         c_uint),
+        ('schedulerParams', c_nvmlVgpuSchedulerParams_t),
+    ]
+
+class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure):
+    _fields_ = [
+        ('avgFactor',   c_uint),
+        ('frequency',   c_uint),
+    ]
+
+class c_nvmlVgpuSchedSetData_t(_PrintableStructure):
+    _fields_ = [
+        ('timeslice',   c_uint),
+    ]
+
+class c_nvmlVgpuSchedulerSetParams_t(Union):
+    _fields_ = [
+        ('vgpuSchedDataWithARR', c_nvmlVgpuSchedSetDataWithARR_t),
+        ('vgpuSchedData',        c_nvmlVgpuSchedSetData_t),
+    ]
+
+class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure):
+    _fields_ = [
+        ('schedulerPolicy', c_uint),
+        ('enableARRMode',   c_uint),
+        ('schedulerParams', c_nvmlVgpuSchedulerSetParams_t),
+    ]
+
+class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure):
+    _fields_ = [
+        ('supportedSchedulers', c_uint * MXSMLEX_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT),
+        ('maxTimeslice',        c_uint),
+        ('minTimeslice',        c_uint),
+        ('isArrModeSupported',  c_uint),
+        ('maxFrequencyForARR',  c_uint),
+        ('minFrequencyForARR',  c_uint),
+        ('maxAvgFactorForARR',  c_uint),
+        ('minAvgFactorForARR',  c_uint),
+    ]
+
+class c_nvmlFBCStats_t(Structure):
+    _fields_ = [("sessionsCount", c_uint),
+                ("averageFPS", c_uint),
+                ("averageLatency", c_uint)
+               ]
+
+class c_nvmlFBCSession_t(_PrintableStructure):
+    _fields_ = [
+        ('sessionId', c_uint),
+        ('pid', c_uint),
+        ('vgpuInstance', _nvmlVgpuInstance_t),
+        ('displayOrdinal', c_uint),
+        ('sessionType', c_uint),
+        ('sessionFlags', c_uint),
+        ('hMaxResolution', c_uint),
+        ('vMaxResolution', c_uint),
+        ('hResolution', c_uint),
+        ('vResolution', c_uint),
+        ('averageFPS', c_uint),
+        ('averageLatency', c_uint),
+    ]
+
+MXSMLEX_DEVICE_MIG_DISABLE = 0x0
+MXSMLEX_DEVICE_MIG_ENABLE  = 0x1
+
+MXSMLEX_GPU_INSTANCE_PROFILE_1_SLICE      = 0x0
+MXSMLEX_GPU_INSTANCE_PROFILE_2_SLICE      = 0x1
+MXSMLEX_GPU_INSTANCE_PROFILE_3_SLICE      = 0x2
+MXSMLEX_GPU_INSTANCE_PROFILE_4_SLICE      = 0x3
+MXSMLEX_GPU_INSTANCE_PROFILE_7_SLICE      = 0x4
+MXSMLEX_GPU_INSTANCE_PROFILE_8_SLICE      = 0x5
+MXSMLEX_GPU_INSTANCE_PROFILE_6_SLICE      = 0x6
+MXSMLEX_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7
+MXSMLEX_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8
+MXSMLEX_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9
+MXSMLEX_GPU_INSTANCE_PROFILE_COUNT        = 0xA
+
+class c_nvmlGpuInstancePlacement_t(Structure):
+    _fields_ = [("start", c_uint),
+                ("size", c_uint)
+               ]
+
+class c_nvmlGpuInstanceProfileInfo_t(Structure):
+    _fields_ = [("id", c_uint),
+                ("isP2pSupported", c_uint),
+                ("sliceCount", c_uint),
+                ("instanceCount", c_uint),
+                ("multiprocessorCount", c_uint),
+                ("copyEngineCount", c_uint),
+                ("decoderCount", c_uint),
+                ("encoderCount", c_uint),
+                ("jpegCount", c_uint),
+                ("ofaCount", c_uint),
+                ("memorySizeMB", c_ulonglong),
+               ]
+
+nvmlGpuInstanceProfileInfo_v2 = 0x02000098
+
+class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure):
+    _fields_ = [("version", c_uint),
+                ("id", c_uint),
+                ("isP2pSupported", c_uint),
+                ("sliceCount", c_uint),
+                ("instanceCount", c_uint),
+                ("multiprocessorCount", c_uint),
+                ("copyEngineCount", c_uint),
+                ("decoderCount", c_uint),
+                ("encoderCount", c_uint),
+                ("jpegCount", c_uint),
+                ("ofaCount", c_uint),
+                ("memorySizeMB", c_ulonglong),
+                ("name", c_char * MXSMLEX_DEVICE_NAME_V2_BUFFER_SIZE)
+               ]
+    
+    def __init__(self):
+        super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__(version=nvmlGpuInstanceProfileInfo_v2)
+
+class c_nvmlGpuInstanceInfo_t(Structure):
+    _fields_ = [("device", c_nvmlDevice_t),
+                ("id", c_uint),
+                ("profileId", c_uint),
+                ("placement", c_nvmlGpuInstancePlacement_t)
+               ]
+
+class struct_c_nvmlGpuInstance_t(Structure):
+    pass # opaque handle
+c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t)
+
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_1_SLICE      = 0x0
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_2_SLICE      = 0x1
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_3_SLICE      = 0x2
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_4_SLICE      = 0x3
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_7_SLICE      = 0x4
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_8_SLICE      = 0x5
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_6_SLICE      = 0x6
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7
+MXSMLEX_COMPUTE_INSTANCE_PROFILE_COUNT        = 0x8
+
+MXSMLEX_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0
+MXSMLEX_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1
+
+class c_nvmlComputeInstancePlacement_t(Structure):
+    _fields_ = [("start", c_uint),
+                ("size", c_uint)
+               ]
+
+class c_nvmlComputeInstanceProfileInfo_t(Structure):
+    _fields_ = [("id", c_uint),
+                ("sliceCount", c_uint),
+                ("instanceCount", c_uint),
+                ("multiprocessorCount", c_uint),
+                ("sharedCopyEngineCount", c_uint),
+                ("sharedDecoderCount", c_uint),
+                ("sharedEncoderCount", c_uint),
+                ("sharedJpegCount", c_uint),
+                ("sharedOfaCount", c_uint)
+               ]
+
+nvmlComputeInstanceProfileInfo_v2 = 0x02000088
+
+class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure):
+    _fields_ = [("version", c_uint),
+                ("id", c_uint),
+                ("sliceCount", c_uint),
+                ("instanceCount", c_uint),
+                ("multiprocessorCount", c_uint),
+                ("sharedCopyEngineCount", c_uint),
+                ("sharedDecoderCount", c_uint),
+                ("sharedEncoderCount", c_uint),
+                ("sharedJpegCount", c_uint),
+                ("sharedOfaCount", c_uint),
+                ("name", c_char * MXSMLEX_DEVICE_NAME_V2_BUFFER_SIZE)
+               ]
+
+    def __init__(self):
+        super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__(version=nvmlComputeInstanceProfileInfo_v2)
+
+class c_nvmlComputeInstanceInfo_t(Structure):
+    _fields_ = [("device", c_nvmlDevice_t),
+                ("gpuInstance", c_nvmlGpuInstance_t),
+                ("id", c_uint),
+                ("profileId", c_uint),
+                ("placement", c_nvmlComputeInstancePlacement_t)
+               ]
+
+MXSMLEX_MAX_GPU_UTILIZATIONS = 8
+MXSMLEX_GPU_UTILIZATION_DOMAIN_GPU    = 0
+MXSMLEX_GPU_UTILIZATION_DOMAIN_FB     = 1
+MXSMLEX_GPU_UTILIZATION_DOMAIN_VID    = 2
+MXSMLEX_GPU_UTILIZATION_DOMAIN_BUS    = 3
+class c_nvmlGpuDynamicPstatesUtilization_t(Structure):
+    _fields_ = [("bIsPresent", c_uint, 1),
+                ("percentage", c_uint),
+                ("incThreshold", c_uint),
+                ("decThreshold", c_uint)]
+class c_nvmlGpuDynamicPstatesInfo_t(Structure):
+    _fields_ = [("flags", c_uint),
+                ("utilization", c_nvmlGpuDynamicPstatesUtilization_t * MXSMLEX_MAX_GPU_UTILIZATIONS)]
+
+MXSMLEX_MAX_THERMAL_SENSORS_PER_GPU = 3
+
+MXSMLEX_THERMAL_TARGET_NONE          = 0
+MXSMLEX_THERMAL_TARGET_GPU           = 1
+MXSMLEX_THERMAL_TARGET_MEMORY        = 2
+MXSMLEX_THERMAL_TARGET_POWER_SUPPLY  = 4
+MXSMLEX_THERMAL_TARGET_BOARD         = 8
+MXSMLEX_THERMAL_TARGET_VCD_BOARD     = 9
+MXSMLEX_THERMAL_TARGET_VCD_INLET     = 10
+MXSMLEX_THERMAL_TARGET_VCD_OUTLET    = 11
+MXSMLEX_THERMAL_TARGET_ALL           = 15
+MXSMLEX_THERMAL_TARGET_UNKNOWN       = -1
+
+MXSMLEX_THERMAL_CONTROLLER_NONE            = 0
+MXSMLEX_THERMAL_CONTROLLER_GPU_INTERNAL    = 1
+MXSMLEX_THERMAL_CONTROLLER_ADM1032         = 2
+MXSMLEX_THERMAL_CONTROLLER_ADT7461         = 3
+MXSMLEX_THERMAL_CONTROLLER_MAX6649         = 4
+MXSMLEX_THERMAL_CONTROLLER_MAX1617         = 5
+MXSMLEX_THERMAL_CONTROLLER_LM99            = 6
+MXSMLEX_THERMAL_CONTROLLER_LM89            = 7
+MXSMLEX_THERMAL_CONTROLLER_LM64            = 8
+MXSMLEX_THERMAL_CONTROLLER_G781            = 9
+MXSMLEX_THERMAL_CONTROLLER_ADT7473         = 10
+MXSMLEX_THERMAL_CONTROLLER_SBMAX6649       = 11
+MXSMLEX_THERMAL_CONTROLLER_VBIOSEVT        = 12
+MXSMLEX_THERMAL_CONTROLLER_OS              = 13
+MXSMLEX_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14
+MXSMLEX_THERMAL_CONTROLLER_NVSYSCON_E551   = 15
+MXSMLEX_THERMAL_CONTROLLER_MAX6649R        = 16
+MXSMLEX_THERMAL_CONTROLLER_ADT7473S        = 17
+MXSMLEX_THERMAL_CONTROLLER_UNKNOWN         = -1
+
+class c_nvmlGpuThermalSensor_t(Structure):
+    _fields_ = [("controller", c_int),
+                ("defaultMinTemp", c_int),
+                ("defaultMaxTemp", c_int),
+                ("currentTemp", c_int),
+                ("target", c_int)]
+class c_nvmlGpuThermalSettings_t(Structure):
+    _fields_ = [("count", c_uint),
+                ("sensor", c_nvmlGpuThermalSensor_t * MXSMLEX_MAX_THERMAL_SENSORS_PER_GPU)]
+
+class struct_c_nvmlComputeInstance_t(Structure):
+    pass # opaque handle
+c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t)
+
+class c_nvmlDeviceAttributes(Structure):
+    _fields_ = [("multiprocessorCount", c_uint),
+                ("sharedCopyEngineCount", c_uint),
+                ("sharedDecoderCount", c_uint),
+                ("sharedEncoderCount", c_uint),
+                ("sharedJpegCount", c_uint),
+                ("sharedOfaCount", c_uint),
+                ("gpuInstanceSliceCount", c_uint),
+                ("computeInstanceSliceCount", c_uint),
+                ("memorySizeMB", c_ulonglong),
+               ]
+
+class c_nvmlRowRemapperHistogramValues(Structure):
+    _fields_ = [("max", c_uint),
+                ("high", c_uint),
+                ("partial", c_uint),
+                ("low", c_uint),
+                ("none", c_uint)
+               ]
+
+MXSMLEX_GPU_CERT_CHAIN_SIZE                = 0x1000
+MXSMLEX_GPU_ATTESTATION_CERT_CHAIN_SIZE    = 0x1400
+MXSMLEX_CC_GPU_CEC_NONCE_SIZE              = 0x20
+MXSMLEX_CC_GPU_ATTESTATION_REPORT_SIZE     = 0x2000
+MXSMLEX_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000
+MXSMLEX_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0
+MXSMLEX_CC_CEC_ATTESTATION_REPORT_PRESENT     = 1
+
+class c_nvmlConfComputeSystemState_t(Structure):
+    _fields_ = [('environment', c_uint),
+                ('ccFeature', c_uint),
+                ('devToolsMode', c_uint),
+               ]
+
+nvmlSystemConfComputeSettings_v1 = 0x1000014
+
+class c_nvmlSystemConfComputeSettings_v1_t(Structure):
+    _fields_ = [('version', c_uint),
+                ('environment', c_uint),
+                ('ccFeature', c_uint),
+                ('devToolsMode', c_uint),
+                ('multiGpuMode', c_uint),
+               ]
+    def __init__(self):
+        super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__(version=nvmlSystemConfComputeSettings_v1)
+
+class c_nvmlConfComputeSystemCaps_t(Structure):
+    _fields_ = [('cpuCaps', c_uint),
+                ('gpusCaps', c_uint),
+               ]
+
+class c_nvmlConfComputeMemSizeInfo_t(Structure):
+    _fields_ = [('protectedMemSizeKib', c_ulonglong),
+                ('unprotectedMemSizeKib', c_ulonglong),
+               ]
+
+class c_nvmlConfComputeGpuCertificate_t(Structure):
+    _fields_ = [('certChainSize', c_uint),
+                ('attestationCertChainSize', c_uint),
+                ('certChain', c_uint8 * MXSMLEX_GPU_CERT_CHAIN_SIZE),
+                ('attestationCertChain', c_uint8 * MXSMLEX_GPU_ATTESTATION_CERT_CHAIN_SIZE),
+               ]
+
+class c_nvmlConfComputeGpuAttestationReport_t(Structure):
+    _fields_ = [('isCecAttestationReportPresent', c_uint),
+                ('attestationReportSize', c_uint),
+                ('cecAttestationReportSize', c_uint),
+                ('nonce', c_uint8 * MXSMLEX_CC_GPU_CEC_NONCE_SIZE),
+                ('attestationReport', c_uint8 * MXSMLEX_CC_GPU_ATTESTATION_REPORT_SIZE),
+                ('cecAttestationReport', c_uint8 * MXSMLEX_CC_GPU_CEC_ATTESTATION_REPORT_SIZE),
+               ]
+
+class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure):
+    _fields_ = [('version', c_uint),
+                ('maxAttackerAdvantage', c_ulong),
+               ]
+ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010
+
+class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure):
+    _fields_ = [('version', c_uint),
+                ('attackerAdvantage', c_ulong),
+               ]
+ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010
+
+
+## string/bytes conversion for ease of use
+def convertStrBytes(func):
+    '''
+    In python 3, strings are unicode instead of bytes, and need to be converted for ctypes
+    Args from caller: (1, 'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF>)
+    Args passed to function: (1, b'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF)>
+    ----
+    Returned from function: b'returned string'
+    Returned to caller: 'returned string'
+    '''
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        # encoding a str returns bytes in python 2 and 3
+        args = [arg.encode() if isinstance(arg, str) else arg for arg in args]
+        res = func(*args, **kwargs)
+        # In python 2, str and bytes are the same
+        # In python 3, str is unicode and should be decoded.
+        # Ctypes handles most conversions, this only effects c_char and char arrays.
+        if isinstance(res, bytes):
+            if isinstance(res, str):
+                return res
+            return res.decode()
+        return res
+
+    if sys.version_info >= (3,):
+        return wrapper
+    return func
+
+def throwOnVersionMismatch(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except NVMLError_FunctionNotFound:
+            raise MXSMLEXLibraryMismatchError("Unversioned function called and the "
+                                           "pyMXSMLEX version does not match the MXSMLEX lib version. "
+                                           "Either use matching pyMXSMLEX and MXSMLEX lib versions or "
+                                           "use a versioned function such as " + func.__name__ + "_v2")
+    return wrapper
+
+## C function wrappers ##
+def nvmlInitWithFlags(flags):
+    _LoadNvmlLibrary()
+
+    #
+    # Initialize the library
+    #
+    fn = _nvmlGetFunctionPointer("mxSmlExInit")
+    ret = fn(flags)
+    _nvmlCheckReturn(ret)
+
+    # Atomically update refcount
+    global _nvmlLib_refcount
+    libLoadLock.acquire()
+    _nvmlLib_refcount += 1
+    libLoadLock.release()
+    return None
+
+def nvmlInit():
+    nvmlInitWithFlags(0)
+    return None
+
+def _LoadNvmlLibrary():
+    '''
+    Load the library if it isn't loaded already
+    '''
+    global nvmlLib
+
+    if (nvmlLib == None):
+        # lock to ensure only one caller loads the library
+        libLoadLock.acquire()
+
+        try:
+            # ensure the library still isn't loaded
+            if (nvmlLib == None):
+                try:
+                    if (sys.platform[:3] == "win"):
+                        # cdecl calling convention
+                        try:
+                            # Check for nvml.dll in System32 first for DCH drivers
+                            nvmlLib = CDLL(os.path.join(os.getenv("WINDIR", "C:/Windows"), "System32/nvml.dll"))
+                        except OSError as ose:
+                            # If nvml.dll is not found in System32, it should be in ProgramFiles
+                            # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll
+                            nvmlLib = CDLL(os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll"))
+                    else:
+                        # assume linux
+                        nvmlLib = CDLL("libmxsml.so")
+                except OSError as ose:
+                    _nvmlCheckReturn(MXSMLEX_ERROR_LIBRARY_NOT_FOUND)
+                if (nvmlLib == None):
+                    _nvmlCheckReturn(MXSMLEX_ERROR_LIBRARY_NOT_FOUND)
+        finally:
+            # lock is always freed
+            libLoadLock.release()
+
+def nvmlShutdown():
+    #
+    # Leave the library loaded, but shutdown the interface
+    #
+    fn = _nvmlGetFunctionPointer("mxSmlExShutdown")
+    ret = fn()
+    _nvmlCheckReturn(ret)
+
+    # Atomically update refcount
+    global _nvmlLib_refcount
+    libLoadLock.acquire()
+    if (0 < _nvmlLib_refcount):
+        _nvmlLib_refcount -= 1
+    libLoadLock.release()
+    return None
+
+# Added in 2.285
+@convertStrBytes
+def nvmlErrorString(result):
+    fn = _nvmlGetFunctionPointer("mxSmlExErrorString")
+    fn.restype = c_char_p # otherwise return is an int
+    ret = fn(result)
+    return ret
+
+# Added in 2.285
+@convertStrBytes
+def nvmlSystemGetMXSMLEXVersion():
+    c_version = create_string_buffer(MXSMLEX_SYSTEM_MXSMLEX_VERSION_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetMXSMLEXVersion")
+    ret = fn(c_version, c_uint(MXSMLEX_SYSTEM_MXSMLEX_VERSION_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_version.value
+
+def nvmlSystemGetCudaDriverVersion():
+    c_cuda_version = c_int()
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetCudaDriverVersion")
+    ret = fn(byref(c_cuda_version))
+    _nvmlCheckReturn(ret)
+    return c_cuda_version.value
+
+def nvmlSystemGetCudaDriverVersion_v2():
+    c_cuda_version = c_int()
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetCudaDriverVersion_v2")
+    ret = fn(byref(c_cuda_version))
+    _nvmlCheckReturn(ret)
+    return c_cuda_version.value
+
+# Added in 2.285
+@convertStrBytes
+def nvmlSystemGetProcessName(pid):
+    c_name = create_string_buffer(1024)
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetProcessName")
+    ret = fn(c_uint(pid), c_name, c_uint(1024))
+    _nvmlCheckReturn(ret)
+    return c_name.value
+
+@convertStrBytes
+def nvmlSystemGetDriverVersion():
+    c_version = create_string_buffer(MXSMLEX_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetDriverVersion")
+    ret = fn(c_version, c_uint(MXSMLEX_SYSTEM_DRIVER_VERSION_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_version.value
+
+# Added in 2.285
+def nvmlSystemGetHicVersion():
+    c_count = c_uint(0)
+    hics = None
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetHicVersion")
+
+    # get the count
+    ret = fn(byref(c_count), None)
+
+    # this should only fail with insufficient size
+    if ((ret != MXSMLEX_SUCCESS) and
+        (ret != MXSMLEX_ERROR_INSUFFICIENT_SIZE)):
+        raise NVMLError(ret)
+
+    # If there are no hics
+    if (c_count.value == 0):
+        return []
+
+    hic_array = c_nvmlHwbcEntry_t * c_count.value
+    hics = hic_array()
+    ret = fn(byref(c_count), hics)
+    _nvmlCheckReturn(ret)
+    return hics
+
+## Unit get functions
+def nvmlUnitGetCount():
+    c_count = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetCount")
+    ret = fn(byref(c_count))
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+def nvmlUnitGetHandleByIndex(index):
+    c_index = c_uint(index)
+    unit = c_nvmlUnit_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetHandleByIndex")
+    ret = fn(c_index, byref(unit))
+    _nvmlCheckReturn(ret)
+    return unit
+
+def nvmlUnitGetUnitInfo(unit):
+    c_info = c_nvmlUnitInfo_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetUnitInfo")
+    ret = fn(unit, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return c_info
+
+def nvmlUnitGetLedState(unit):
+    c_state =  c_nvmlLedState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetLedState")
+    ret = fn(unit, byref(c_state))
+    _nvmlCheckReturn(ret)
+    return c_state
+
+def nvmlUnitGetPsuInfo(unit):
+    c_info = c_nvmlPSUInfo_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetPsuInfo")
+    ret = fn(unit, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return c_info
+
+def nvmlUnitGetTemperature(unit, type):
+    c_temp = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetTemperature")
+    ret = fn(unit, c_uint(type), byref(c_temp))
+    _nvmlCheckReturn(ret)
+    return c_temp.value
+
+def nvmlUnitGetFanSpeedInfo(unit):
+    c_speeds = c_nvmlUnitFanSpeeds_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetFanSpeedInfo")
+    ret = fn(unit, byref(c_speeds))
+    _nvmlCheckReturn(ret)
+    return c_speeds
+
+# added to API
+def nvmlUnitGetDeviceCount(unit):
+    c_count = c_uint(0)
+    # query the unit to determine device count
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetDevices")
+    ret = fn(unit, byref(c_count), None)
+    if (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        ret = MXSMLEX_SUCCESS
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+def nvmlUnitGetDevices(unit):
+    c_count = c_uint(nvmlUnitGetDeviceCount(unit))
+    device_array = c_nvmlDevice_t * c_count.value
+    c_devices = device_array()
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitGetDevices")
+    ret = fn(unit, byref(c_count), c_devices)
+    _nvmlCheckReturn(ret)
+    return c_devices
+
+## Device get functions
+def nvmlDeviceGetCount():
+    c_count = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetCount")
+    ret = fn(byref(c_count))
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+def nvmlDeviceGetHandleByIndex(index):
+    c_index = c_uint(index)
+    device = c_nvmlDevice_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExGetDeviceHandleByIndex")
+    ret = fn(c_index, byref(device))
+    _nvmlCheckReturn(ret)
+    return device
+
+@convertStrBytes
+def nvmlDeviceGetHandleBySerial(serial):
+    c_serial = c_char_p(serial)
+    device = c_nvmlDevice_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetHandleBySerial")
+    ret = fn(c_serial, byref(device))
+    _nvmlCheckReturn(ret)
+    return device
+
+@convertStrBytes
+def nvmlDeviceGetHandleByUUID(uuid):
+    c_uuid = c_char_p(uuid)
+    device = c_nvmlDevice_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetHandleByUUID")
+    ret = fn(c_uuid, byref(device))
+    _nvmlCheckReturn(ret)
+    return device
+
+@convertStrBytes
+def nvmlDeviceGetHandleByPciBusId(pciBusId):
+    c_busId = c_char_p(pciBusId)
+    device = c_nvmlDevice_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetHandleByPciBusId_v2")
+    ret = fn(c_busId, byref(device))
+    _nvmlCheckReturn(ret)
+    return device
+
+@convertStrBytes
+def nvmlDeviceGetName(handle):
+    c_name = create_string_buffer(MXSMLEX_DEVICE_NAME_V2_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetName")
+    ret = fn(handle, c_name, c_uint(MXSMLEX_DEVICE_NAME_V2_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_name.value
+
+def nvmlDeviceGetBoardId(handle):
+    c_id = c_uint();
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetBoardId")
+    ret = fn(handle, byref(c_id))
+    _nvmlCheckReturn(ret)
+    return c_id.value
+
+def nvmlDeviceGetMultiGpuBoard(handle):
+    c_multiGpu = c_uint();
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMultiGpuBoard")
+    ret = fn(handle, byref(c_multiGpu))
+    _nvmlCheckReturn(ret)
+    return c_multiGpu.value
+
+def nvmlDeviceGetBrand(handle):
+    c_type = _nvmlBrandType_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetBrand")
+    ret = fn(handle, byref(c_type))
+    _nvmlCheckReturn(ret)
+    return c_type.value
+
+def nvmlDeviceGetC2cModeInfoV1(handle):
+    c_info = c_nvmlC2cModeInfo_v1_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetC2cModeInfoV")
+    ret = fn(handle, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return c_info
+
+def nvmlDeviceGetC2cModeInfoV(handle):
+    return nvmlDeviceGetC2cModeInfoV1(handle)
+
+@convertStrBytes
+def nvmlDeviceGetBoardPartNumber(handle):
+    c_part_number = create_string_buffer(MXSMLEX_DEVICE_PART_NUMBER_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetBoardPartNumber")
+    ret = fn(handle, c_part_number, c_uint(MXSMLEX_DEVICE_PART_NUMBER_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_part_number.value
+
+@convertStrBytes
+def nvmlDeviceGetSerial(handle):
+    c_serial = create_string_buffer(MXSMLEX_DEVICE_SERIAL_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSerial")
+    ret = fn(handle, c_serial, c_uint(MXSMLEX_DEVICE_SERIAL_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_serial.value
+
+def nvmlDeviceGetModuleId(handle, moduleId):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetModuleId")
+    ret = fn(handle, moduleId)
+    return ret
+
+def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope):
+    affinity_array = c_ulonglong * nodeSetSize
+    c_affinity = affinity_array()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMemoryAffinity")
+    ret = fn(handle, nodeSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope))
+    _nvmlCheckReturn(ret)
+    return c_affinity
+
+def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope):
+    affinity_array = c_ulonglong * cpuSetSize
+    c_affinity = affinity_array()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetCpuAffinityWithinScope")
+    ret = fn(handle, cpuSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope))
+    _nvmlCheckReturn(ret)
+    return c_affinity
+
+def nvmlDeviceGetCpuAffinity(handle, cpuSetSize):
+    affinity_array = c_ulonglong * cpuSetSize
+    c_affinity = affinity_array()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetCpuAffinity")
+    ret = fn(handle, cpuSetSize, byref(c_affinity))
+    _nvmlCheckReturn(ret)
+    return c_affinity
+
+def nvmlDeviceSetCpuAffinity(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetCpuAffinity")
+    ret = fn(handle)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceClearCpuAffinity(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceClearCpuAffinity")
+    ret = fn(handle)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetNumaNodeId(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNumaNodeId")
+    node = c_int()
+    ret = fn(handle, byref(node))
+    _nvmlCheckReturn(ret)
+    return node.value
+
+def nvmlDeviceGetMinorNumber(handle):
+    c_minor_number = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMinorNumber")
+    ret = fn(handle, byref(c_minor_number))
+    _nvmlCheckReturn(ret)
+    return c_minor_number.value
+
+@convertStrBytes
+def nvmlDeviceGetUUID(handle):
+    c_uuid = create_string_buffer(MXSMLEX_DEVICE_UUID_V2_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetUUID")
+    ret = fn(handle, c_uuid, c_uint(MXSMLEX_DEVICE_UUID_V2_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_uuid.value
+
+@convertStrBytes
+def nvmlDeviceGetInforomVersion(handle, infoRomObject):
+    c_version = create_string_buffer(MXSMLEX_DEVICE_INFOROM_VERSION_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetInforomVersion")
+    ret = fn(handle, _nvmlInforomObject_t(infoRomObject),
+                 c_version, c_uint(MXSMLEX_DEVICE_INFOROM_VERSION_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_version.value
+
+# Added in 4.304
+@convertStrBytes
+def nvmlDeviceGetInforomImageVersion(handle):
+    c_version = create_string_buffer(MXSMLEX_DEVICE_INFOROM_VERSION_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetInforomImageVersion")
+    ret = fn(handle, c_version, c_uint(MXSMLEX_DEVICE_INFOROM_VERSION_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_version.value
+
+# Added in 4.304
+def nvmlDeviceGetInforomConfigurationChecksum(handle):
+    c_checksum = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetInforomConfigurationChecksum")
+    ret = fn(handle, byref(c_checksum))
+    _nvmlCheckReturn(ret)
+    return c_checksum.value
+
+# Added in 4.304
+def nvmlDeviceValidateInforom(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceValidateInforom")
+    ret = fn(handle)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetLastBBXFlushTime(handle):
+    c_timestamp = c_ulonglong()
+    c_durationUs = c_ulong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetLastBBXFlushTime")
+    ret = fn(handle, byref(c_timestamp), byref(c_durationUs))
+    _nvmlCheckReturn(ret)
+    return [c_timestamp.value, c_durationUs.value]
+
+def nvmlDeviceGetDisplayMode(handle):
+    c_mode = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDisplayMode")
+    ret = fn(handle, byref(c_mode))
+    _nvmlCheckReturn(ret)
+    return c_mode.value
+
+def nvmlDeviceGetDisplayActive(handle):
+    c_mode = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDisplayActive")
+    ret = fn(handle, byref(c_mode))
+    _nvmlCheckReturn(ret)
+    return c_mode.value
+
+
+def nvmlDeviceGetPersistenceMode(handle):
+    c_state = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPersistenceMode")
+    ret = fn(handle, byref(c_state))
+    _nvmlCheckReturn(ret)
+    return c_state.value
+
+def nvmlDeviceGetPciInfoExt(handle, c_info):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPciInfoExt")
+    ret = fn(handle, c_info)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetPciInfo_v3(handle):
+    c_info = nvmlPciInfo_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPciInfo_v3")
+    ret = fn(handle, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return c_info
+
+def nvmlDeviceGetPciInfo(handle):
+    return nvmlDeviceGetPciInfo_v3(handle)
+
+def nvmlDeviceGetClockInfo(handle, type):
+    c_clock = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetClockInfo")
+    ret = fn(handle, _nvmlClockType_t(type), byref(c_clock))
+    _nvmlCheckReturn(ret)
+    return c_clock.value
+
+# Added in 2.285
+def nvmlDeviceGetMaxClockInfo(handle, type):
+    c_clock = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMaxClockInfo")
+    ret = fn(handle, _nvmlClockType_t(type), byref(c_clock))
+    _nvmlCheckReturn(ret)
+    return c_clock.value
+
+# Added in 4.304
+def nvmlDeviceGetApplicationsClock(handle, type):
+    c_clock = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetApplicationsClock")
+    ret = fn(handle, _nvmlClockType_t(type), byref(c_clock))
+    _nvmlCheckReturn(ret)
+    return c_clock.value
+
+def nvmlDeviceGetMaxCustomerBoostClock(handle, type):
+    c_clock = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMaxCustomerBoostClock")
+    ret = fn(handle, _nvmlClockType_t(type), byref(c_clock))
+    _nvmlCheckReturn(ret)
+    return c_clock.value
+
+def nvmlDeviceGetClock(handle, type, id):
+    c_clock = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetClock")
+    ret = fn(handle, _nvmlClockType_t(type), _nvmlClockId_t(id), byref(c_clock))
+    _nvmlCheckReturn(ret)
+    return c_clock.value
+
+# Added in 5.319
+def nvmlDeviceGetDefaultApplicationsClock(handle, type):
+    c_clock = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDefaultApplicationsClock")
+    ret = fn(handle, _nvmlClockType_t(type), byref(c_clock))
+    _nvmlCheckReturn(ret)
+    return c_clock.value
+
+# Added in 4.304
+def nvmlDeviceGetSupportedMemoryClocks(handle):
+    # first call to get the size
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSupportedMemoryClocks")
+    ret = fn(handle, byref(c_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no clocks
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        clocks_array = c_uint * c_count.value
+        c_clocks = clocks_array()
+
+        # make the call again
+        ret = fn(handle, byref(c_count), c_clocks)
+        _nvmlCheckReturn(ret)
+
+        procs = []
+        for i in range(c_count.value):
+            procs.append(c_clocks[i])
+
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+# Added in 4.304
+def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz):
+    # first call to get the size
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSupportedGraphicsClocks")
+    ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no clocks
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        clocks_array = c_uint * c_count.value
+        c_clocks = clocks_array()
+
+        # make the call again
+        ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks)
+        _nvmlCheckReturn(ret)
+
+        procs = []
+        for i in range(c_count.value):
+            procs.append(c_clocks[i])
+
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetFanSpeed(handle):
+    c_speed = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetFanSpeed")
+    ret = fn(handle, byref(c_speed))
+    _nvmlCheckReturn(ret)
+    return c_speed.value
+
+def nvmlDeviceGetFanSpeed_v2(handle, fan):
+    c_speed = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetFanSpeed_v2")
+    ret = fn(handle, fan, byref(c_speed))
+    _nvmlCheckReturn(ret)
+    return c_speed.value
+
+def nvmlDeviceGetTargetFanSpeed(handle, fan):
+    c_speed = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetTargetFanSpeed")
+    ret = fn(handle, fan, byref(c_speed))
+    _nvmlCheckReturn(ret)
+    return c_speed.value
+
+def nvmlDeviceGetNumFans(device):
+    c_numFans = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNumFans")
+    ret = fn(device, byref(c_numFans))
+    _nvmlCheckReturn(ret)
+    return c_numFans.value
+
+def nvmlDeviceSetDefaultFanSpeed_v2(handle, index):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetDefaultFanSpeed_v2");
+    ret = fn(handle, index)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed, maxSpeed):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMinMaxFanSpeed")
+    ret = fn(handle, minSpeed, maxSpeed)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetFanControlPolicy_v2")
+    ret = fn(handle, fan, fanControlPolicy)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetFanControlPolicy")
+    ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy))
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetTemperature(handle, sensor):
+    c_temp = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetTemperature")
+    ret = fn(handle, _nvmlTemperatureSensors_t(sensor), byref(c_temp))
+    _nvmlCheckReturn(ret)
+    return c_temp.value
+
+def nvmlDeviceGetTemperatureThreshold(handle, threshold):
+    c_temp = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetTemperatureThreshold")
+    ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp))
+    _nvmlCheckReturn(ret)
+    return c_temp.value
+
+def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp):
+    c_temp = c_uint()
+    c_temp.value = temp
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetTemperatureThreshold")
+    ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp))
+    _nvmlCheckReturn(ret)
+    return None
+
+# DEPRECATED use nvmlDeviceGetPerformanceState
+def nvmlDeviceGetPowerState(handle):
+    c_pstate = _nvmlPstates_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPowerState")
+    ret = fn(handle, byref(c_pstate))
+    _nvmlCheckReturn(ret)
+    return c_pstate.value
+
+def nvmlDeviceGetPerformanceState(handle):
+    c_pstate = _nvmlPstates_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPerformanceState")
+    ret = fn(handle, byref(c_pstate))
+    _nvmlCheckReturn(ret)
+    return c_pstate.value
+
+def nvmlDeviceGetPowerManagementMode(handle):
+    c_pcapMode = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPowerManagementMode")
+    ret = fn(handle, byref(c_pcapMode))
+    _nvmlCheckReturn(ret)
+    return c_pcapMode.value
+
+def nvmlDeviceGetPowerManagementLimit(handle):
+    c_limit = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPowerManagementLimit")
+    ret = fn(handle, byref(c_limit))
+    _nvmlCheckReturn(ret)
+    return c_limit.value
+
+# Added in 4.304
+def nvmlDeviceGetPowerManagementLimitConstraints(handle):
+    c_minLimit = c_uint()
+    c_maxLimit = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPowerManagementLimitConstraints")
+    ret = fn(handle, byref(c_minLimit), byref(c_maxLimit))
+    _nvmlCheckReturn(ret)
+    return [c_minLimit.value, c_maxLimit.value]
+
+# Added in 4.304
+def nvmlDeviceGetPowerManagementDefaultLimit(handle):
+    c_limit = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPowerManagementDefaultLimit")
+    ret = fn(handle, byref(c_limit))
+    _nvmlCheckReturn(ret)
+    return c_limit.value
+
+
+# Added in 331
+def nvmlDeviceGetEnforcedPowerLimit(handle):
+    c_limit = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetEnforcedPowerLimit")
+    ret = fn(handle, byref(c_limit))
+    _nvmlCheckReturn(ret)
+    return c_limit.value
+
+def nvmlDeviceGetPowerUsage(handle):
+    c_watts = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPowerUsage")
+    ret = fn(handle, byref(c_watts))
+    _nvmlCheckReturn(ret)
+    return c_watts.value
+
+def nvmlDeviceGetTotalEnergyConsumption(handle):
+    c_millijoules = c_uint64()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetTotalEnergyConsumption")
+    ret = fn(handle, byref(c_millijoules))
+    _nvmlCheckReturn(ret)
+    return c_millijoules.value
+
+# Added in 4.304
+def nvmlDeviceGetGpuOperationMode(handle):
+    c_currState = _nvmlGpuOperationMode_t()
+    c_pendingState = _nvmlGpuOperationMode_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuOperationMode")
+    ret = fn(handle, byref(c_currState), byref(c_pendingState))
+    _nvmlCheckReturn(ret)
+    return [c_currState.value, c_pendingState.value]
+
+# Added in 4.304
+def nvmlDeviceGetCurrentGpuOperationMode(handle):
+    return nvmlDeviceGetGpuOperationMode(handle)[0]
+
+# Added in 4.304
+def nvmlDeviceGetPendingGpuOperationMode(handle):
+    return nvmlDeviceGetGpuOperationMode(handle)[1]
+
+def nvmlDeviceGetMemoryInfo(handle, version=None):
+    if not version:
+        c_memory = c_nvmlMemory_t()
+        fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMemoryInfo")
+    else:
+        c_memory = c_nvmlMemory_v2_t()
+        c_memory.version = version
+        fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMemoryInfo_v2")
+    ret = fn(handle, byref(c_memory))
+    _nvmlCheckReturn(ret)
+    return c_memory
+
+def nvmlDeviceGetBAR1MemoryInfo(handle):
+    c_bar1_memory = c_nvmlBAR1Memory_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetBAR1MemoryInfo")
+    ret = fn(handle, byref(c_bar1_memory))
+    _nvmlCheckReturn(ret)
+    return c_bar1_memory
+
+def nvmlDeviceGetComputeMode(handle):
+    c_mode = _nvmlComputeMode_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetComputeMode")
+    ret = fn(handle, byref(c_mode))
+    _nvmlCheckReturn(ret)
+    return c_mode.value
+
+def nvmlDeviceGetCudaComputeCapability(handle):
+    c_major = c_int()
+    c_minor = c_int()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetComputeCapability")
+    ret = fn(handle, byref(c_major), byref(c_minor))
+    _nvmlCheckReturn(ret)
+    return (c_major.value, c_minor.value)
+
+def nvmlDeviceGetEccMode(handle):
+    c_currState = _nvmlEnableState_t()
+    c_pendingState = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetEccMode")
+    ret = fn(handle, byref(c_currState), byref(c_pendingState))
+    _nvmlCheckReturn(ret)
+    return [c_currState.value, c_pendingState.value]
+
+# added to API
+def nvmlDeviceGetCurrentEccMode(handle):
+    return nvmlDeviceGetEccMode(handle)[0]
+
+# added to API
+def nvmlDeviceGetPendingEccMode(handle):
+    return nvmlDeviceGetEccMode(handle)[1]
+
+def nvmlDeviceGetDefaultEccMode(handle):
+    c_defaultState = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDefaultEccMode")
+    ret = fn(handle, byref(c_defaultState))
+    _nvmlCheckReturn(ret)
+    return [c_defaultState.value]
+
+def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType):
+    c_count = c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetTotalEccErrors")
+    ret = fn(handle, _nvmlMemoryErrorType_t(errorType),
+                 _nvmlEccCounterType_t(counterType), byref(c_count))
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+# This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter
+def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType):
+    c_counts = c_nvmlEccErrorCounts_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDetailedEccErrors")
+    ret = fn(handle, _nvmlMemoryErrorType_t(errorType),
+                 _nvmlEccCounterType_t(counterType), byref(c_counts))
+    _nvmlCheckReturn(ret)
+    return c_counts
+
+# Added in 4.304
+def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType):
+    c_count = c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMemoryErrorCounter")
+    ret = fn(handle,
+             _nvmlMemoryErrorType_t(errorType),
+             _nvmlEccCounterType_t(counterType),
+             _nvmlMemoryLocation_t(locationType),
+             byref(c_count))
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+def nvmlDeviceGetUtilizationRates(handle):
+    c_util = c_nvmlUtilization_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetUtilizationRates")
+    ret = fn(handle, byref(c_util))
+    _nvmlCheckReturn(ret)
+    return c_util
+
+def nvmlDeviceGetEncoderUtilization(handle):
+    c_util = c_uint()
+    c_samplingPeriod = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetEncoderUtilization")
+    ret = fn(handle, byref(c_util), byref(c_samplingPeriod))
+    _nvmlCheckReturn(ret)
+    return [c_util.value, c_samplingPeriod.value]
+
+def nvmlDeviceGetDecoderUtilization(handle):
+    c_util = c_uint()
+    c_samplingPeriod = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDecoderUtilization")
+    ret = fn(handle, byref(c_util), byref(c_samplingPeriod))
+    _nvmlCheckReturn(ret)
+    return [c_util.value, c_samplingPeriod.value]
+
+def nvmlDeviceGetJpgUtilization(handle):
+    c_util = c_uint()
+    c_samplingPeriod = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetJpgUtilization")
+    ret = fn(handle, byref(c_util), byref(c_samplingPeriod))
+    _nvmlCheckReturn(ret)
+    return [c_util.value, c_samplingPeriod.value]
+
+def nvmlDeviceGetOfaUtilization(handle):
+    c_util = c_uint()
+    c_samplingPeriod = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetOfaUtilization")
+    ret = fn(handle, byref(c_util), byref(c_samplingPeriod))
+    _nvmlCheckReturn(ret)
+    return [c_util.value, c_samplingPeriod.value]
+
+def nvmlDeviceGetPcieReplayCounter(handle):
+    c_replay = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPcieReplayCounter")
+    ret = fn(handle, byref(c_replay))
+    _nvmlCheckReturn(ret)
+    return c_replay.value
+
+def nvmlDeviceGetDriverModel(handle):
+    c_currModel = _nvmlDriverModel_t()
+    c_pendingModel = _nvmlDriverModel_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDriverModel")
+    ret = fn(handle, byref(c_currModel), byref(c_pendingModel))
+    _nvmlCheckReturn(ret)
+    return [c_currModel.value, c_pendingModel.value]
+
+# added to API
+def nvmlDeviceGetCurrentDriverModel(handle):
+    return nvmlDeviceGetDriverModel(handle)[0]
+
+# added to API
+def nvmlDeviceGetPendingDriverModel(handle):
+    return nvmlDeviceGetDriverModel(handle)[1]
+
+# Added in 2.285
+@convertStrBytes
+def nvmlDeviceGetVbiosVersion(handle):
+    c_version = create_string_buffer(MXSMLEX_DEVICE_VBIOS_VERSION_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetVbiosVersion")
+    ret = fn(handle, c_version, c_uint(MXSMLEX_DEVICE_VBIOS_VERSION_BUFFER_SIZE))
+    _nvmlCheckReturn(ret)
+    return c_version.value
+
+# Added in 2.285
+def nvmlDeviceGetComputeRunningProcesses_v2(handle):
+    # first call to get the size
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetComputeRunningProcesses_v2")
+    ret = fn(handle, byref(c_count), None)
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no running processes
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        # oversize the array incase more processes are created
+        c_count.value = c_count.value * 2 + 5
+        proc_array = c_nvmlProcessInfo_v2_t * c_count.value
+        c_procs = proc_array()
+        # make the call again
+        ret = fn(handle, byref(c_count), c_procs)
+        _nvmlCheckReturn(ret)
+        procs = []
+        for i in range(c_count.value):
+            # use an alternative struct for this object
+            obj = nvmlStructToFriendlyObject(c_procs[i])
+            if (obj.usedGpuMemory == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+                # special case for WDDM on Windows, see comment above
+                obj.usedGpuMemory = None
+            procs.append(obj)
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+# Added in 2.285
+def nvmlDeviceGetComputeRunningProcesses_v3(handle):
+    # first call to get the size
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetComputeRunningProcesses_v3")
+    ret = fn(handle, byref(c_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no running processes
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        # oversize the array incase more processes are created
+        c_count.value = c_count.value * 2 + 5
+        proc_array = c_nvmlProcessInfo_v3_t * c_count.value
+        c_procs = proc_array()
+
+        # make the call again
+        ret = fn(handle, byref(c_count), c_procs)
+        _nvmlCheckReturn(ret)
+
+        procs = []
+        for i in range(c_count.value):
+            # use an alternative struct for this object
+            obj = nvmlStructToFriendlyObject(c_procs[i])
+            if (obj.usedGpuMemory == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+                # special case for WDDM on Windows, see comment above
+                obj.usedGpuMemory = None
+            procs.append(obj)
+
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+@throwOnVersionMismatch
+def nvmlDeviceGetComputeRunningProcesses(handle):
+    return nvmlDeviceGetComputeRunningProcesses_v3(handle)
+
+def nvmlDeviceGetGraphicsRunningProcesses_v2(handle):
+    # first call to get the size
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGraphicsRunningProcesses_v2")
+    ret = fn(handle, byref(c_count), None)
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no running processes
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        # oversize the array incase more processes are created
+        c_count.value = c_count.value * 2 + 5
+        proc_array = c_nvmlProcessInfo_v2_t * c_count.value
+        c_procs = proc_array()
+        # make the call again
+        ret = fn(handle, byref(c_count), c_procs)
+        _nvmlCheckReturn(ret)
+        procs = []
+        for i in range(c_count.value):
+            # use an alternative struct for this object
+            obj = nvmlStructToFriendlyObject(c_procs[i])
+            if (obj.usedGpuMemory == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+                # special case for WDDM on Windows, see comment above
+                obj.usedGpuMemory = None
+            procs.append(obj)
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetGraphicsRunningProcesses_v3(handle):
+    # first call to get the size
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGraphicsRunningProcesses_v3")
+    ret = fn(handle, byref(c_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no running processes
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        # oversize the array incase more processes are created
+        c_count.value = c_count.value * 2 + 5
+        proc_array = c_nvmlProcessInfo_v3_t * c_count.value
+        c_procs = proc_array()
+
+        # make the call again
+        ret = fn(handle, byref(c_count), c_procs)
+        _nvmlCheckReturn(ret)
+
+        procs = []
+        for i in range(c_count.value):
+            # use an alternative struct for this object
+            obj = nvmlStructToFriendlyObject(c_procs[i])
+            if (obj.usedGpuMemory == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+                # special case for WDDM on Windows, see comment above
+                obj.usedGpuMemory = None
+            procs.append(obj)
+
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+@throwOnVersionMismatch
+def nvmlDeviceGetGraphicsRunningProcesses(handle):
+    return nvmlDeviceGetGraphicsRunningProcesses_v3(handle)
+
+@throwOnVersionMismatch
+def nvmlDeviceGetMPSComputeRunningProcesses(handle):
+    return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle)
+
+def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle):
+    # first call to get the size
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMPSComputeRunningProcesses_v2")
+    ret = fn(handle, byref(c_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no running processes
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        # oversize the array incase more processes are created
+        c_count.value = c_count.value * 2 + 5
+        proc_array = c_nvmlProcessInfo_v2_t * c_count.value
+        c_procs = proc_array()
+
+        # make the call again
+        ret = fn(handle, byref(c_count), c_procs)
+        _nvmlCheckReturn(ret)
+
+        procs = []
+        for i in range(c_count.value):
+            # use an alternative struct for this object
+            obj = nvmlStructToFriendlyObject(c_procs[i])
+            if (obj.usedGpuMemory == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+                # special case for WDDM on Windows, see comment above
+                obj.usedGpuMemory = None
+            procs.append(obj)
+
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle):
+    # first call to get the size
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMPSComputeRunningProcesses_v3")
+    ret = fn(handle, byref(c_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no running processes
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        # oversize the array incase more processes are created
+        c_count.value = c_count.value * 2 + 5
+        proc_array = c_nvmlProcessInfo_v3_t * c_count.value
+        c_procs = proc_array()
+
+        # make the call again
+        ret = fn(handle, byref(c_count), c_procs)
+        _nvmlCheckReturn(ret)
+
+        procs = []
+        for i in range(c_count.value):
+            # use an alternative struct for this object
+            obj = nvmlStructToFriendlyObject(c_procs[i])
+            if (obj.usedGpuMemory == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+                # special case for WDDM on Windows, see comment above
+                obj.usedGpuMemory = None
+            procs.append(obj)
+
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetRunningProcessDetailList(handle, version, mode):
+    c_processDetailList = c_nvmlProcessDetailList_t()
+    c_processDetailList.version = version
+    c_processDetailList.mode = mode
+
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetRunningProcessDetailList")
+
+    # first call to get the size
+    ret = fn(handle, byref(c_processDetailList))
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no running processes
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        c_procs = c_nvmlProcessDetail_v1_t * c_processDetailList.numProcArrayEntries
+        c_processDetailList.procArray = cast((c_procs)(), POINTER(c_nvmlProcessDetail_v1_t))
+
+        # make the call again
+        ret = fn(handle, byref(c_processDetailList))
+        _nvmlCheckReturn(ret)
+
+        procs = []
+        for i in range(c_processDetailList.numProcArrayEntries):
+            # use an alternative struct for this object
+            obj = c_processDetailList.procArray[i]
+            if (obj.usedGpuMemory == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+                obj.usedGpuMemory = None
+            if (obj.usedGpuCcProtectedMemory == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+                obj.usedGpuCcProtectedMemory = None
+            procs.append(obj)
+
+        return procs
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetAutoBoostedClocksEnabled(handle):
+    c_isEnabled = _nvmlEnableState_t()
+    c_defaultIsEnabled = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetAutoBoostedClocksEnabled")
+    ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled))
+    _nvmlCheckReturn(ret)
+    return [c_isEnabled.value, c_defaultIsEnabled.value]
+    #Throws MXSMLEX_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+
+## Set functions
+def nvmlUnitSetLedState(unit, color):
+    fn = _nvmlGetFunctionPointer("mxSmlExUnitSetLedState")
+    ret = fn(unit, _nvmlLedColor_t(color))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceSetPersistenceMode(handle, mode):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetPersistenceMode")
+    ret = fn(handle, _nvmlEnableState_t(mode))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceSetComputeMode(handle, mode):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetComputeMode")
+    ret = fn(handle, _nvmlComputeMode_t(mode))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceSetEccMode(handle, mode):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetEccMode")
+    ret = fn(handle, _nvmlEnableState_t(mode))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceClearEccErrorCounts(handle, counterType):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceClearEccErrorCounts")
+    ret = fn(handle, _nvmlEccCounterType_t(counterType))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceSetDriverModel(handle, model):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetDriverModel")
+    ret = fn(handle, _nvmlDriverModel_t(model))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetAutoBoostedClocksEnabled")
+    ret = fn(handle, _nvmlEnableState_t(enabled))
+    _nvmlCheckReturn(ret)
+    return None
+    #Throws MXSMLEX_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+
+def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetDefaultAutoBoostedClocksEnabled")
+    ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags))
+    _nvmlCheckReturn(ret)
+    return None
+    #Throws MXSMLEX_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
+
+def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetGpuLockedClocks")
+    ret = fn(handle, c_uint(minGpuClockMHz), c_uint(maxGpuClockMHz))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceResetGpuLockedClocks(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceResetGpuLockedClocks")
+    ret = fn(handle)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetMemoryLockedClocks")
+    ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceResetMemoryLockedClocks(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceResetMemoryLockedClocks")
+    ret = fn(handle)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetClkMonStatus")
+    ret = fn(handle, c_clkMonInfo)
+    return ret
+
+# Added in 4.304
+def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetApplicationsClocks")
+    ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz))
+    _nvmlCheckReturn(ret)
+    return None
+
+# Added in 4.304
+def nvmlDeviceResetApplicationsClocks(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceResetApplicationsClocks")
+    ret = fn(handle)
+    _nvmlCheckReturn(ret)
+    return None
+
+# Added in 4.304
+def nvmlDeviceSetPowerManagementLimit(handle, limit):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetPowerManagementLimit")
+    ret = fn(handle, c_uint(limit))
+    _nvmlCheckReturn(ret)
+    return None
+
+# Added in 4.304
+def nvmlDeviceSetGpuOperationMode(handle, mode):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetGpuOperationMode")
+    ret = fn(handle, _nvmlGpuOperationMode_t(mode))
+    _nvmlCheckReturn(ret)
+    return None
+
+# Added in 2.285
+def nvmlEventSetCreate():
+    fn = _nvmlGetFunctionPointer("mxSmlExEventSetCreate")
+    eventSet = c_nvmlEventSet_t()
+    ret = fn(byref(eventSet))
+    _nvmlCheckReturn(ret)
+    return eventSet
+
+# Added in 2.285
+def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceRegisterEvents")
+    ret = fn(handle, c_ulonglong(eventTypes), eventSet)
+    _nvmlCheckReturn(ret)
+    return None
+
+# Added in 2.285
+def nvmlDeviceGetSupportedEventTypes(handle):
+    c_eventTypes = c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSupportedEventTypes")
+    ret = fn(handle, byref(c_eventTypes))
+    _nvmlCheckReturn(ret)
+    return c_eventTypes.value
+
+# raises MXSMLEX_ERROR_TIMEOUT exception on timeout
+def nvmlEventSetWait_v2(eventSet, timeoutms):
+    fn = _nvmlGetFunctionPointer("mxSmlExEventSetWait_v2")
+    data = c_nvmlEventData_t()
+    ret = fn(eventSet, byref(data), c_uint(timeoutms))
+    _nvmlCheckReturn(ret)
+    return data
+
+def nvmlEventSetWait(eventSet, timeoutms):
+    return nvmlEventSetWait_v2(eventSet, timeoutms)
+
+# Added in 2.285
+def nvmlEventSetFree(eventSet):
+    fn = _nvmlGetFunctionPointer("mxSmlExEventSetFree")
+    ret = fn(eventSet)
+    _nvmlCheckReturn(ret)
+    return None
+
+# Added in 3.295
+def nvmlDeviceOnSameBoard(handle1, handle2):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceOnSameBoard")
+    onSameBoard = c_int()
+    ret = fn(handle1, handle2, byref(onSameBoard))
+    _nvmlCheckReturn(ret)
+    return (onSameBoard.value != 0)
+
+# Added in 3.295
+def nvmlDeviceGetCurrPcieLinkGeneration(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetCurrPcieLinkGeneration")
+    gen = c_uint()
+    ret = fn(handle, byref(gen))
+    _nvmlCheckReturn(ret)
+    return gen.value
+
+# Added in 3.295
+def nvmlDeviceGetMaxPcieLinkGeneration(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMaxPcieLinkGeneration")
+    gen = c_uint()
+    ret = fn(handle, byref(gen))
+    _nvmlCheckReturn(ret)
+    return gen.value
+
+# Added in 3.295
+def nvmlDeviceGetCurrPcieLinkWidth(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetCurrPcieLinkWidth")
+    width = c_uint()
+    ret = fn(handle, byref(width))
+    _nvmlCheckReturn(ret)
+    return width.value
+
+# Added in 3.295
+def nvmlDeviceGetMaxPcieLinkWidth(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMaxPcieLinkWidth")
+    width = c_uint()
+    ret = fn(handle, byref(width))
+    _nvmlCheckReturn(ret)
+    return width.value
+
+def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuMaxPcieLinkGeneration")
+    gen = c_uint()
+    ret = fn(handle, byref(gen))
+    _nvmlCheckReturn(ret)
+    return gen.value
+
+# Added in 4.304
+def nvmlDeviceGetSupportedClocksThrottleReasons(handle):
+    c_reasons= c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSupportedClocksThrottleReasons")
+    ret = fn(handle, byref(c_reasons))
+    _nvmlCheckReturn(ret)
+    return c_reasons.value
+
+def nvmlDeviceGetSupportedClocksEventReasons(handle):
+    c_reasons= c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSupportedClocksEventReasons")
+    ret = fn(handle, byref(c_reasons))
+    _nvmlCheckReturn(ret)
+    return c_reasons.value
+
+# Added in 4.304
+def nvmlDeviceGetCurrentClocksThrottleReasons(handle):
+    c_reasons= c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetCurrentClocksThrottleReasons")
+    ret = fn(handle, byref(c_reasons))
+    _nvmlCheckReturn(ret)
+    return c_reasons.value
+
+def nvmlDeviceGetCurrentClocksEventReasons(handle):
+    c_reasons= c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetCurrentClocksEventReasons")
+    ret = fn(handle, byref(c_reasons))
+    _nvmlCheckReturn(ret)
+    return c_reasons.value
+
+# Added in 5.319
+def nvmlDeviceGetIndex(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetIndex")
+    c_index = c_uint()
+    ret = fn(handle, byref(c_index))
+    _nvmlCheckReturn(ret)
+    return c_index.value
+
+# Added in 5.319
+def nvmlDeviceGetAccountingMode(handle):
+    c_mode = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetAccountingMode")
+    ret = fn(handle, byref(c_mode))
+    _nvmlCheckReturn(ret)
+    return c_mode.value
+
+def nvmlDeviceSetAccountingMode(handle, mode):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetAccountingMode")
+    ret = fn(handle, _nvmlEnableState_t(mode))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceClearAccountingPids(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceClearAccountingPids")
+    ret = fn(handle)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetAccountingStats(handle, pid):
+    stats = c_nvmlAccountingStats_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetAccountingStats")
+    ret = fn(handle, c_uint(pid), byref(stats))
+    _nvmlCheckReturn(ret)
+    if (stats.maxMemoryUsage == MXSMLEX_VALUE_NOT_AVAILABLE_ulonglong.value):
+        # special case for WDDM on Windows, see comment above
+        stats.maxMemoryUsage = None
+    return stats
+
+def nvmlDeviceGetAccountingPids(handle):
+    count = c_uint(nvmlDeviceGetAccountingBufferSize(handle))
+    pids = (c_uint * count.value)()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetAccountingPids")
+    ret = fn(handle, byref(count), pids)
+    _nvmlCheckReturn(ret)
+    return list(map(int, pids[0:count.value]))
+
+def nvmlDeviceGetAccountingBufferSize(handle):
+    bufferSize = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetAccountingBufferSize")
+    ret = fn(handle, byref(bufferSize))
+    _nvmlCheckReturn(ret)
+    return int(bufferSize.value)
+
+def nvmlDeviceGetRetiredPages(device, sourceFilter):
+    c_source = _nvmlPageRetirementCause_t(sourceFilter)
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetRetiredPages")
+
+    # First call will get the size
+    ret = fn(device, c_source, byref(c_count), None)
+
+    # this should only fail with insufficient size
+    if ((ret != MXSMLEX_SUCCESS) and
+        (ret != MXSMLEX_ERROR_INSUFFICIENT_SIZE)):
+        raise NVMLError(ret)
+
+    # call again with a buffer
+    # oversize the array for the rare cases where additional pages
+    # are retired between MXSMLEX calls
+    c_count.value = c_count.value * 2 + 5
+    page_array = c_ulonglong * c_count.value
+    c_pages = page_array()
+    ret = fn(device, c_source, byref(c_count), c_pages)
+    _nvmlCheckReturn(ret)
+    return list(map(int, c_pages[0:c_count.value]))
+
+def nvmlDeviceGetRetiredPages_v2(device, sourceFilter):
+    c_source = _nvmlPageRetirementCause_t(sourceFilter)
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetRetiredPages_v2")
+
+    # First call will get the size
+    ret = fn(device, c_source, byref(c_count), None)
+
+    # this should only fail with insufficient size
+    if ((ret != MXSMLEX_SUCCESS) and
+        (ret != MXSMLEX_ERROR_INSUFFICIENT_SIZE)):
+        raise NVMLError(ret)
+
+    # call again with a buffer
+    # oversize the array for the rare cases where additional pages
+    # are retired between MXSMLEX calls
+    c_count.value = c_count.value * 2 + 5
+    page_array = c_ulonglong * c_count.value
+    c_pages = page_array()
+    times_array = c_ulonglong * c_count.value
+    c_times = times_array()
+    ret = fn(device, c_source, byref(c_count), c_pages, c_times)
+    _nvmlCheckReturn(ret)
+    return [ { 'address': int(c_pages[i]), 'timestamp': int(c_times[i]) } for i in range(c_count.value) ];
+
+def nvmlDeviceGetRetiredPagesPendingStatus(device):
+    c_pending = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetRetiredPagesPendingStatus")
+    ret = fn(device, byref(c_pending))
+    _nvmlCheckReturn(ret)
+    return int(c_pending.value)
+
+def nvmlDeviceGetAPIRestriction(device, apiType):
+    c_permission = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetAPIRestriction")
+    ret = fn(device, _nvmlRestrictedAPI_t(apiType), byref(c_permission))
+    _nvmlCheckReturn(ret)
+    return int(c_permission.value)
+
+def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetAPIRestriction")
+    ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetBridgeChipInfo(handle):
+    bridgeHierarchy = c_nvmlBridgeChipHierarchy_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetBridgeChipInfo")
+    ret = fn(handle, byref(bridgeHierarchy))
+    _nvmlCheckReturn(ret)
+    return bridgeHierarchy
+
+def nvmlDeviceGetSamples(device, sampling_type, timeStamp):
+    c_sampling_type = _nvmlSamplingType_t(sampling_type)
+    c_time_stamp = c_ulonglong(timeStamp)
+    c_sample_count = c_uint(0)
+    c_sample_value_type = _nvmlValueType_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSamples")
+
+    ## First Call gets the size
+    ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), None)
+
+    # Stop if this fails
+    if (ret != MXSMLEX_SUCCESS):
+        raise NVMLError(ret)
+
+    sampleArray = c_sample_count.value * c_nvmlSample_t
+    c_samples = sampleArray()
+    ret = fn(device, c_sampling_type, c_time_stamp,  byref(c_sample_value_type), byref(c_sample_count), c_samples)
+    _nvmlCheckReturn(ret)
+    return (c_sample_value_type.value, c_samples[0:c_sample_count.value])
+
+def nvmlDeviceGetViolationStatus(device, perfPolicyType):
+    c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType)
+    c_violTime = c_nvmlViolationTime_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetViolationStatus")
+
+    ## Invoke the method to get violation time
+    ret = fn(device, c_perfPolicy_type, byref(c_violTime))
+    _nvmlCheckReturn(ret)
+    return c_violTime
+
+def nvmlDeviceGetPcieThroughput(device, counter):
+    c_util = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPcieThroughput")
+    ret = fn(device, _nvmlPcieUtilCounter_t(counter), byref(c_util))
+    _nvmlCheckReturn(ret)
+    return c_util.value
+
+def nvmlSystemGetTopologyGpuSet(cpuNumber):
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetTopologyGpuSet")
+
+    # First call will get the size
+    ret = fn(cpuNumber, byref(c_count), None)
+
+    if ret != MXSMLEX_SUCCESS:
+        raise NVMLError(ret)
+    # call again with a buffer
+    device_array = c_nvmlDevice_t * c_count.value
+    c_devices = device_array()
+    ret = fn(cpuNumber, byref(c_count), c_devices)
+    _nvmlCheckReturn(ret)
+    return list(c_devices[0:c_count.value])
+
+def nvmlDeviceGetTopologyNearestGpus(device, level):
+    c_count = c_uint(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetTopologyNearestGpus")
+
+    # First call will get the size
+    ret = fn(device, level, byref(c_count), None)
+
+    if ret != MXSMLEX_SUCCESS:
+        raise NVMLError(ret)
+
+    # call again with a buffer
+    device_array = c_nvmlDevice_t * c_count.value
+    c_devices = device_array()
+    ret = fn(device, level, byref(c_count), c_devices)
+    _nvmlCheckReturn(ret)
+    return list(c_devices[0:c_count.value])
+
+def nvmlDeviceGetTopologyCommonAncestor(device1, device2):
+    c_level = _nvmlGpuTopologyLevel_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetTopologyCommonAncestor")
+    ret = fn(device1, device2, byref(c_level))
+    _nvmlCheckReturn(ret)
+    return c_level.value
+
+def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter):
+    c_rxcounter = c_ulonglong()
+    c_txcounter = c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNvLinkUtilizationCounter")
+    ret = fn(device, link, counter, byref(c_rxcounter), byref(c_txcounter))
+    _nvmlCheckReturn(ret)
+    return (c_rxcounter.value, c_txcounter.value)
+
+def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceFreezeNvLinkUtilizationCounter")
+    ret = fn(device, link, counter, freeze)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceResetNvLinkUtilizationCounter")
+    ret = fn(device, link, counter)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceSetNvLinkUtilizationControl(device, link, counter, control, reset):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetNvLinkUtilizationControl")
+    ret = fn(device, link, counter, byref(control), reset)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter):
+    c_control = nvmlNvLinkUtilizationControl_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNvLinkUtilizationControl")
+    ret = fn(device, link, counter, byref(c_control))
+    _nvmlCheckReturn(ret)
+    return c_control
+
+def nvmlDeviceGetNvLinkCapability(device, link, capability):
+    c_capResult = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNvLinkCapability")
+    ret = fn(device, link, capability, byref(c_capResult))
+    _nvmlCheckReturn(ret)
+    return c_capResult.value
+
+def nvmlDeviceGetNvLinkErrorCounter(device, link, counter):
+    c_result = c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNvLinkErrorCounter")
+    ret = fn(device, link, counter, byref(c_result))
+    _nvmlCheckReturn(ret)
+    return c_result.value
+
+def nvmlDeviceResetNvLinkErrorCounters(device, link):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceResetNvLinkErrorCounters")
+    ret = fn(device, link)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetNvLinkRemotePciInfo(device, link):
+    c_pci = nvmlPciInfo_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNvLinkRemotePciInfo_v2")
+    ret = fn(device, link, byref(c_pci))
+    _nvmlCheckReturn(ret)
+    return c_pci
+
+def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link):
+    c_type = _nvmlNvLinkDeviceType_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNvLinkRemoteDeviceType")
+    ret = fn(handle, link, byref(c_type))
+    _nvmlCheckReturn(ret)
+    return c_type.value
+
+def nvmlDeviceGetNvLinkState(device, link):
+    c_isActive = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNvLinkState")
+    ret = fn(device, link, byref(c_isActive))
+    _nvmlCheckReturn(ret)
+    return c_isActive.value
+
+def nvmlDeviceGetNvLinkVersion(device, link):
+    c_version = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNvLinkVersion")
+    ret = fn(device, link, byref(c_version))
+    _nvmlCheckReturn(ret)
+    return c_version.value
+
+def nvmlDeviceModifyDrainState(pciInfo, newState):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceModifyDrainState")
+    ret = fn(pointer(pciInfo), newState)
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceQueryDrainState(pciInfo):
+    c_newState = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceQueryDrainState")
+    ret = fn(pointer(pciInfo), byref(c_newState))
+    _nvmlCheckReturn(ret)
+    return c_newState.value
+
+def nvmlDeviceRemoveGpu(pciInfo):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceRemoveGpu")
+    ret = fn(pointer(pciInfo))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceDiscoverGpus(pciInfo):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceDiscoverGpus")
+    ret = fn(pointer(pciInfo))
+    _nvmlCheckReturn(ret)
+    return None
+
+def nvmlDeviceGetFieldValues(handle, fieldIds):
+    values_arr = c_nvmlFieldValue_t * len(fieldIds)
+    values = values_arr()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetFieldValues")
+
+    for i, fieldId in enumerate(fieldIds):
+        try:
+            (values[i].fieldId, values[i].scopeId) = fieldId
+        except TypeError:
+            values[i].fieldId = fieldId
+
+    ret = fn(handle, c_int32(len(fieldIds)), byref(values))
+    _nvmlCheckReturn(ret)
+    return values
+
+def nvmlDeviceClearFieldValues(handle, fieldIds):
+    values_arr = c_nvmlFieldValue_t * len(fieldIds)
+    values = values_arr()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceClearFieldValues")
+
+    for i, fieldId in enumerate(fieldIds):
+        try:
+            (values[i].fieldId, values[i].scopeId) = fieldId
+        except TypeError:
+            values[i].fieldId = fieldId
+
+    ret = fn(handle, c_int32(len(fieldIds)), byref(values))
+    _nvmlCheckReturn(ret)
+    return values
+
+def nvmlDeviceGetVirtualizationMode(handle):
+    c_virtualization_mode = c_ulonglong()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetVirtualizationMode")
+    ret = fn(handle, byref(c_virtualization_mode))
+    _nvmlCheckReturn(ret)
+    return c_virtualization_mode.value
+
+def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetVirtualizationMode")
+    return fn(handle, virtualization_mode)
+
+def nvmlDeviceGetVgpuHeterogeneousMode(handle):
+    c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0)
+    c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuHeterogeneousMode")
+    ret = fn(handle, byref(c_vgpuHeterogeneousMode))
+    _nvmlCheckReturn(ret)
+    return c_vgpuHeterogeneousMode.mode
+
+def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode):
+    c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0)
+    c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1
+    c_vgpuHeterogeneousMode.mode = heterogeneous_mode
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetVgpuHeterogeneousMode")
+    ret = fn(handle, byref(c_vgpuHeterogeneousMode))
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlVgpuInstanceGetPlacementId(vgpuInstance):
+    c_placement = c_nvmlVgpuPlacementId_v1_t(0)
+    c_placement.version = VgpuPlacementId_v1
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetPlacementId")
+    ret = fn(vgpuInstance, byref(c_placement))
+    _nvmlCheckReturn(ret)
+    return c_placement.placementId
+
+def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId):
+    c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t()
+    c_max_instances = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetMaxInstances")
+    ret = fn(handle, vgpuTypeId, byref(c_max_instances))
+    _nvmlCheckReturn(ret)
+    c_placements = c_uint * c_max_instances.value
+    c_vgpu_placements.version = VgpuPlacementList_v1
+    c_vgpu_placements.placementIds = c_placements()
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuTypeSupportedPlacements")
+    ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements))
+    _nvmlCheckReturn(ret)
+    return c_vgpu_placements
+
+def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId):
+    c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t()
+    c_max_instances = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetMaxInstances")
+    ret = fn(handle, vgpuTypeId, byref(c_max_instances))
+    _nvmlCheckReturn(ret)
+    c_placements = c_uint * c_max_instances.value
+    c_vgpu_placements.version = VgpuPlacementList_v1
+    c_vgpu_placements.placementIds = c_placements()
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuTypeCreatablePlacements")
+    ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements))
+    _nvmlCheckReturn(ret)
+    return c_vgpu_placements
+
+def nvmlGetVgpuDriverCapabilities(capability):
+    c_capResult = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExGetVgpuDriverCapabilities")
+    ret = fn(_nvmlVgpuDriverCapability_t(capability), byref(c_capResult))
+    _nvmlCheckReturn(ret)
+    return c_capResult.value
+
+def nvmlDeviceGetVgpuCapabilities(handle, capability):
+    c_capResult = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuCapabilities")
+    ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), byref(c_capResult))
+    _nvmlCheckReturn(ret)
+    return c_capResult.value
+
+def nvmlDeviceSetVgpuCapabilities(handle, capability, state):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetVgpuCapabilities")
+    ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetSupportedVgpus(handle):
+    # first call to get the size
+    c_vgpu_count = c_uint(0)
+
+    fn =  _nvmlGetFunctionPointer("mxSmlExDeviceGetSupportedVgpus")
+    ret = fn(handle, byref(c_vgpu_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no supported vGPUs
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value
+        c_vgpu_type_ids = vgpu_type_ids_array()
+
+        # make the call again
+        ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids)
+        _nvmlCheckReturn(ret)
+        vgpus = []
+        for i in range(c_vgpu_count.value):
+            vgpus.append(c_vgpu_type_ids[i])
+        return vgpus
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetCreatableVgpus(handle):
+    # first call to get the size
+    c_vgpu_count = c_uint(0)
+
+    fn =  _nvmlGetFunctionPointer("mxSmlExDeviceGetCreatableVgpus")
+    ret = fn(handle, byref(c_vgpu_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no supported vGPUs
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value
+        c_vgpu_type_ids = vgpu_type_ids_array()
+
+        # make the call again
+        ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids)
+        _nvmlCheckReturn(ret)
+        vgpus = []
+        for i in range(c_vgpu_count.value):
+            vgpus.append(c_vgpu_type_ids[i])
+        return vgpus
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId):
+    c_profile_id = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetGpuInstanceProfileId")
+    ret = fn(vgpuTypeId, byref(c_profile_id))
+    _nvmlCheckReturn(ret)
+    return (c_profile_id.value)
+
+@convertStrBytes
+def nvmlVgpuTypeGetClass(vgpuTypeId):
+    c_class = create_string_buffer(MXSMLEX_DEVICE_NAME_BUFFER_SIZE)
+    c_buffer_size = c_uint(MXSMLEX_DEVICE_NAME_BUFFER_SIZE)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetClass")
+    ret = fn(vgpuTypeId, c_class, byref(c_buffer_size))
+    _nvmlCheckReturn(ret)
+    return c_class.value
+
+@convertStrBytes
+def nvmlVgpuTypeGetName(vgpuTypeId):
+    c_name = create_string_buffer(MXSMLEX_DEVICE_NAME_BUFFER_SIZE)
+    c_buffer_size = c_uint(MXSMLEX_DEVICE_NAME_BUFFER_SIZE)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetName")
+    ret = fn(vgpuTypeId, c_name, byref(c_buffer_size))
+    _nvmlCheckReturn(ret)
+    return c_name.value
+
+def nvmlVgpuTypeGetDeviceID(vgpuTypeId):
+    c_device_id    = c_ulonglong(0)
+    c_subsystem_id = c_ulonglong(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetDeviceID")
+    ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id))
+    _nvmlCheckReturn(ret)
+    return (c_device_id.value, c_subsystem_id.value)
+
+def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId):
+    c_fb_size = c_ulonglong(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetFramebufferSize")
+    ret = fn(vgpuTypeId, byref(c_fb_size))
+    _nvmlCheckReturn(ret)
+    return c_fb_size.value
+
+def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId):
+    c_num_heads = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetNumDisplayHeads")
+    ret = fn(vgpuTypeId, byref(c_num_heads))
+    _nvmlCheckReturn(ret)
+    return c_num_heads.value
+
+def nvmlVgpuTypeGetResolution(vgpuTypeId):
+    c_xdim = c_uint(0)
+    c_ydim = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetResolution")
+    ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim))
+    _nvmlCheckReturn(ret)
+    return (c_xdim.value, c_ydim.value)
+
+@convertStrBytes
+def nvmlVgpuTypeGetLicense(vgpuTypeId):
+    c_license = create_string_buffer(MXSMLEX_GRID_LICENSE_BUFFER_SIZE)
+    c_buffer_size = c_uint(MXSMLEX_GRID_LICENSE_BUFFER_SIZE)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetLicense")
+    ret = fn(vgpuTypeId, c_license, c_buffer_size)
+    _nvmlCheckReturn(ret)
+    return c_license.value
+
+def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId):
+    c_frl_config = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetFrameRateLimit")
+    ret = fn(vgpuTypeId, byref(c_frl_config))
+    _nvmlCheckReturn(ret)
+    return c_frl_config.value
+
+def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId):
+    c_gsp_heap = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetGspHeapSize")
+    ret = fn(vgpuTypeId, byref(c_gsp_heap))
+    _nvmlCheckReturn(ret)
+    return c_gsp_heap.value
+
+def nvmlVgpuTypeGetFbReservation(vgpuTypeId):
+    c_fb_reservation = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetFbReservation")
+    ret = fn(vgpuTypeId, byref(c_fb_reservation))
+    _nvmlCheckReturn(ret)
+    return c_fb_reservation.value
+
+def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId):
+    c_max_instances = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetMaxInstances")
+    ret = fn(handle, vgpuTypeId, byref(c_max_instances))
+    _nvmlCheckReturn(ret)
+    return c_max_instances.value
+
+def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId):
+    c_max_instances_per_vm = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetMaxInstancesPerVm")
+    ret = fn(vgpuTypeId, byref(c_max_instances_per_vm))
+    _nvmlCheckReturn(ret)
+    return c_max_instances_per_vm.value
+
+def nvmlDeviceGetActiveVgpus(handle):
+    # first call to get the size
+    c_vgpu_count = c_uint(0)
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetActiveVgpus")
+    ret = fn(handle, byref(c_vgpu_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no active vGPUs
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_count.value
+        c_vgpu_instances = vgpu_instance_array()
+
+        # make the call again
+        ret = fn(handle, byref(c_vgpu_count), c_vgpu_instances)
+        _nvmlCheckReturn(ret)
+        vgpus = []
+        for i in range(c_vgpu_count.value):
+            vgpus.append(c_vgpu_instances[i])
+        return vgpus
+    else:
+        # error case
+        raise NVMLError(ret)
+
+@convertStrBytes
+def nvmlVgpuInstanceGetVmID(vgpuInstance):
+    c_vm_id = create_string_buffer(MXSMLEX_DEVICE_UUID_BUFFER_SIZE)
+    c_buffer_size = c_uint(MXSMLEX_GRID_LICENSE_BUFFER_SIZE)
+    c_vm_id_type  = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetVmID")
+    ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type))
+    _nvmlCheckReturn(ret)
+    return (c_vm_id.value, c_vm_id_type.value)
+
+@convertStrBytes
+def nvmlVgpuInstanceGetUUID(vgpuInstance):
+    c_uuid = create_string_buffer(MXSMLEX_DEVICE_UUID_BUFFER_SIZE)
+    c_buffer_size = c_uint(MXSMLEX_DEVICE_UUID_BUFFER_SIZE)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetUUID")
+    ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size)
+    _nvmlCheckReturn(ret)
+    return c_uuid.value
+
+@convertStrBytes
+def nvmlVgpuInstanceGetMdevUUID(vgpuInstance):
+    c_uuid = create_string_buffer(MXSMLEX_DEVICE_UUID_BUFFER_SIZE)
+    c_buffer_size = c_uint(MXSMLEX_DEVICE_UUID_BUFFER_SIZE)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetMdevUUID")
+    ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size)
+    _nvmlCheckReturn(ret)
+    return c_uuid.value
+
+@convertStrBytes
+def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance):
+    c_driver_version = create_string_buffer(MXSMLEX_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)
+    c_buffer_size = c_uint(MXSMLEX_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetVmDriverVersion")
+    ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size)
+    _nvmlCheckReturn(ret)
+    return c_driver_version.value
+
+def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance):
+    c_license_status = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetLicenseStatus")
+    ret = fn(vgpuInstance, byref(c_license_status))
+    _nvmlCheckReturn(ret)
+    return c_license_status.value
+
+def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance):
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetLicenseInfo_v2")
+    c_license_info = c_nvmlVgpuLicenseInfo_t()
+    ret = fn(vgpuInstance, byref(c_license_info))
+    _nvmlCheckReturn(ret)
+    return c_license_info
+
+def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance):
+    return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance)
+
+def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance):
+    c_frl = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetFrameRateLimit")
+    ret = fn(vgpuInstance, byref(c_frl))
+    _nvmlCheckReturn(ret)
+    return c_frl.value
+
+def nvmlVgpuInstanceGetEccMode(vgpuInstance):
+    c_mode = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetEccMode")
+    ret = fn(vgpuInstance, byref(c_mode))
+    _nvmlCheckReturn(ret)
+    return c_mode.value
+
+def nvmlVgpuInstanceGetType(vgpuInstance):
+    c_vgpu_type = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetType")
+    ret = fn(vgpuInstance, byref(c_vgpu_type))
+    _nvmlCheckReturn(ret)
+    return c_vgpu_type.value
+
+def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance):
+    c_encoder_capacity = c_ulonglong(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetEncoderCapacity")
+    ret = fn(vgpuInstance, byref(c_encoder_capacity))
+    _nvmlCheckReturn(ret)
+    return c_encoder_capacity.value
+
+def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity):
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceSetEncoderCapacity")
+    return fn(vgpuInstance, encoder_capacity)
+
+def nvmlVgpuInstanceGetFbUsage(vgpuInstance):
+    c_fb_usage = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetFbUsage")
+    ret = fn(vgpuInstance, byref(c_fb_usage))
+    _nvmlCheckReturn(ret)
+    return c_fb_usage.value
+
+def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability):
+    c_cap_result = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuTypeGetCapabilities")
+    ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result))
+    _nvmlCheckReturn(ret)
+    return (c_cap_result.value)
+
+def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance):
+    c_id = c_uint(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetGpuInstanceId")
+    ret = fn(vgpuInstance, byref(c_id))
+    _nvmlCheckReturn(ret)
+    return (c_id.value)
+
+@convertStrBytes
+def nvmlVgpuInstanceGetGpuPciId(vgpuInstance):
+    c_vgpuPciId = create_string_buffer(MXSMLEX_DEVICE_PCI_BUS_ID_BUFFER_SIZE)
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetGpuPciId")
+    ret = fn(vgpuInstance, c_vgpuPciId, byref(c_uint(MXSMLEX_DEVICE_PCI_BUS_ID_BUFFER_SIZE)))
+    _nvmlCheckReturn(ret)
+    return c_vgpuPciId.value
+
+def nvmlDeviceGetVgpuUtilization(handle, timeStamp):
+    # first call to get the size
+    c_vgpu_count = c_uint(0)
+    c_time_stamp = c_ulonglong(timeStamp)
+    c_sample_value_type = _nvmlValueType_t()
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuUtilization")
+    ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no active vGPUs
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        sampleArray = c_vgpu_count.value * c_nvmlVgpuInstanceUtilizationSample_t
+        c_samples = sampleArray()
+
+        # make the call again
+        ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), c_samples)
+        _nvmlCheckReturn(ret)
+
+        return c_samples[0:c_vgpu_count.value]
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp):
+    # first call to get the size
+    c_time_stamp = c_ulonglong(timeStamp)
+    c_vgpuUtilInfo = c_nvmlVgpuInstancesUtilizationInfo_v1_t(0)
+    c_vgpuUtilInfo.version = VgpuInstancesUtilizationInfo_v1
+    c_vgpuUtilInfo.sampleValType = _nvmlValueType_t()
+    c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0)
+    c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuInstancesUtilizationInfo")
+    ret = fn(handle, byref(c_vgpuUtilInfo))
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no active vGPUs
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        sampleArray = c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t
+        c_samples = sampleArray()
+        c_vgpuUtilInfo.vgpuUtilArray = c_samples
+
+        # make the call again
+        ret = fn(handle, byref(c_vgpuUtilInfo))
+        _nvmlCheckReturn(ret)
+
+        return c_samples[0:c_vgpuUtilInfo.vgpuInstanceCount]
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex):
+    c_p2pstatus = _nvmlGpuP2PStatus_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetP2PStatus")
+    ret = fn(device1, device2,p2pIndex, byref(c_p2pstatus))
+    _nvmlCheckReturn(ret)
+    return c_p2pstatus.value
+
+def nvmlDeviceGetGridLicensableFeatures_v4(handle):
+    c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGridLicensableFeatures_v4")
+    ret = fn(handle, byref(c_get_grid_licensable_features))
+    _nvmlCheckReturn(ret)
+
+    return (c_get_grid_licensable_features)
+
+def nvmlDeviceGetGridLicensableFeatures(handle):
+    return nvmlDeviceGetGridLicensableFeatures_v4(handle)
+
+def nvmlDeviceGetGspFirmwareVersion(handle, version):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGspFirmwareVersion")
+    ret = fn(handle, version)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetGspFirmwareMode(handle, isEnabled, defaultMode):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGspFirmwareMode")
+    ret = fn(handle, isEnabled, defaultMode)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType):
+    c_encoder_capacity = c_ulonglong(0)
+    c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType)
+
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetEncoderCapacity")
+    ret = fn(handle, c_encoderQuery_type, byref(c_encoder_capacity))
+    _nvmlCheckReturn(ret)
+    return c_encoder_capacity.value
+
+def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp):
+    # first call to get the size
+    c_vgpu_count = c_uint(0)
+    c_time_stamp = c_ulonglong(timeStamp)
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuProcessUtilization")
+    ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no active vGPUs
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        sampleArray = c_vgpu_count.value * c_nvmlVgpuProcessUtilizationSample_t
+        c_samples = sampleArray()
+
+        # make the call again
+        ret = fn(handle, c_time_stamp, byref(c_vgpu_count), c_samples)
+        _nvmlCheckReturn(ret)
+
+        return c_samples[0:c_vgpu_count.value]
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp):
+    # first call to get the size
+    c_time_stamp = c_ulonglong(timeStamp)
+    c_vgpuProcUtilInfo = c_nvmlVgpuProcessesUtilizationInfo_v1_t(0)
+    c_vgpuProcUtilInfo.version = VgpuProcessesUtilizationInfo_v1
+    c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0)
+    c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuProcessesUtilizationInfo")
+    ret = fn(handle, byref(c_vgpuProcUtilInfo))
+
+    if (ret == MXSMLEX_SUCCESS):
+        # special case, no active vGPUs
+        return []
+    elif (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        sampleArray = c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t
+        c_samples = sampleArray()
+        c_vgpuProcUtilInfo.vgpuProcUtilArray = c_samples
+
+        # make the call again
+        ret = fn(handle, byref(c_vgpuProcUtilInfo))
+        _nvmlCheckReturn(ret)
+
+        return c_samples[0:c_vgpuProcUtilInfo.vgpuProcessCount]
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetEncoderStats(handle):
+    c_encoderCount = c_ulonglong(0)
+    c_encodeFps = c_ulonglong(0)
+    c_encoderLatency = c_ulonglong(0)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetEncoderStats")
+    ret = fn(handle, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency))
+    _nvmlCheckReturn(ret)
+    return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value)
+
+def nvmlDeviceGetEncoderSessions(handle):
+    # first call to get the size
+    c_session_count = c_uint(0)
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetEncoderSessions")
+    ret = fn(handle, byref(c_session_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        if (c_session_count.value != 0):
+            # typical case
+            session_array = c_nvmlEncoderSession_t * c_session_count.value
+            c_sessions = session_array()
+
+            # make the call again
+            ret = fn(handle, byref(c_session_count), c_sessions)
+            _nvmlCheckReturn(ret)
+            sessions = []
+            for i in range(c_session_count.value):
+                sessions.append(c_sessions[i])
+            return sessions
+        else:
+            return []  # no active sessions
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetFBCStats(handle):
+    c_fbcStats = c_nvmlFBCStats_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetFBCStats")
+    ret = fn(handle, byref(c_fbcStats))
+    _nvmlCheckReturn(ret)
+    return c_fbcStats
+
+def nvmlDeviceGetFBCSessions(handle):
+    # first call to get the size
+    c_session_count = c_uint(0)
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetFBCSessions")
+    ret = fn(handle, byref(c_session_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        if (c_session_count.value != 0):
+            # typical case
+            session_array = c_nvmlFBCSession_t * c_session_count.value
+            c_sessions = session_array()
+
+            # make the call again
+            ret = fn(handle, byref(c_session_count), c_sessions)
+            _nvmlCheckReturn(ret)
+            sessions = []
+            for i in range(c_session_count.value):
+                sessions.append(c_sessions[i])
+            return sessions
+        else:
+            return []  # no active sessions
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlVgpuInstanceGetEncoderStats(vgpuInstance):
+    c_encoderCount    = c_ulonglong(0)
+    c_encodeFps       = c_ulonglong(0)
+    c_encoderLatency  = c_ulonglong(0)
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetEncoderStats")
+    ret = fn(vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency))
+    _nvmlCheckReturn(ret)
+    return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value)
+
+def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance):
+    # first call to get the size
+    c_session_count = c_uint(0)
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetEncoderSessions")
+    ret = fn(vgpuInstance, byref(c_session_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        if (c_session_count.value != 0):
+            # typical case
+            session_array = c_nvmlEncoderSession_t * c_session_count.value
+            c_sessions = session_array()
+
+            # make the call again
+            ret = fn(vgpuInstance, byref(c_session_count), c_sessions)
+            _nvmlCheckReturn(ret)
+            sessions = []
+            for i in range(c_session_count.value):
+                sessions.append(c_sessions[i])
+            return sessions
+        else:
+            return []  # no active sessions
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlVgpuInstanceGetFBCStats(vgpuInstance):
+    c_fbcStats = c_nvmlFBCStats_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetFBCStats")
+    ret = fn(vgpuInstance, byref(c_fbcStats))
+    _nvmlCheckReturn(ret)
+    return c_fbcStats
+
+def nvmlVgpuInstanceGetFBCSessions(vgpuInstance):
+    # first call to get the size
+    c_session_count = c_uint(0)
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetFBCSessions")
+    ret = fn(vgpuInstance, byref(c_session_count), None)
+
+    if (ret == MXSMLEX_SUCCESS):
+        if (c_session_count.value != 0):
+            # typical case
+            session_array = c_nvmlFBCSession_t * c_session_count.value
+            c_sessions = session_array()
+
+            # make the call again
+            ret = fn(vgpuInstance, byref(c_session_count), c_sessions)
+            _nvmlCheckReturn(ret)
+            sessions = []
+            for i in range(c_session_count.value):
+                sessions.append(c_sessions[i])
+            return sessions
+        else:
+            return []  # no active sessions
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetProcessUtilization(handle, timeStamp):
+    # first call to get the size
+    c_count = c_uint(0)
+    c_time_stamp = c_ulonglong(timeStamp)
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetProcessUtilization")
+    ret = fn(handle, None, byref(c_count), c_time_stamp)
+
+    if (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        sampleArray = c_count.value * c_nvmlProcessUtilizationSample_t
+        c_samples = sampleArray()
+
+        # make the call again
+        ret = fn(handle, c_samples, byref(c_count), c_time_stamp)
+        _nvmlCheckReturn(ret)
+
+        return c_samples[0:c_count.value]
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp):
+    # first call to get the size
+    c_time_stamp = c_ulonglong(timeStamp)
+    c_processesUtilInfo = c_nvmlProcessesUtilizationInfo_v1_t(0)
+    c_processesUtilInfo.version = ProcessesUtilizationInfo_v1
+    c_processesUtilInfo.processSamplesCount = c_uint(0)
+    c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp
+
+    fn  = _nvmlGetFunctionPointer("mxSmlExDeviceGetProcessesUtilizationInfo")
+    ret = fn(handle, byref(c_processesUtilInfo))
+
+    if (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        # typical case
+        sampleArray = c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t
+        c_samples = sampleArray()
+        c_processesUtilInfo.procUtilArray = c_samples
+
+        # make the call again
+        ret = fn(handle, byref(c_processesUtilInfo))
+        _nvmlCheckReturn(ret)
+
+        return c_samples[0:c_processesUtilInfo.processSamplesCount]
+    else:
+        # error case
+        raise NVMLError(ret)
+
+def nvmlVgpuInstanceGetMetadata(vgpuInstance):
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetMetadata")
+    c_vgpuMetadata = c_nvmlVgpuMetadata_t()
+    c_bufferSize = c_uint(0)
+    # Make the first MXSMLEX API call to get the c_bufferSize value.
+    # We have already allocated required buffer above.
+    ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize))
+    if (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize))
+        _nvmlCheckReturn(ret)
+    else:
+        raise NVMLError(ret)
+    return c_vgpuMetadata
+
+def nvmlDeviceGetVgpuMetadata(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuMetadata")
+    c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t()
+    c_bufferSize = c_uint(0)
+    # Make the first MXSMLEX API call to get the c_bufferSize value.
+    # We have already allocated required buffer above.
+    ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize))
+    if (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize))
+        _nvmlCheckReturn(ret)
+    else:
+        raise NVMLError(ret)
+    return c_vgpuPgpuMetadata
+
+def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata):
+    fn = _nvmlGetFunctionPointer("mxSmlExGetVgpuCompatibility")
+    c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t()
+    ret = fn(byref(vgpuMetadata), byref(pgpuMetadata), byref(c_vgpuPgpuCompatibility))
+    _nvmlCheckReturn(ret)
+    return c_vgpuPgpuCompatibility
+
+@convertStrBytes
+def nvmlDeviceGetPgpuMetadataString(handle):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPgpuMetadataString")
+    c_pgpuMetadata = create_string_buffer(MXSMLEX_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE)
+    c_bufferSize = c_uint(0)
+    # Make the first MXSMLEX API call to get the c_bufferSize value.
+    # We have already allocated required buffer above.
+    ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize))
+    if (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize))
+        _nvmlCheckReturn(ret)
+    else:
+        raise NVMLError(ret)
+    return (c_pgpuMetadata.value, c_bufferSize.value)
+
+def nvmlDeviceGetVgpuSchedulerLog(handle):
+    c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuSchedulerLog")
+    ret = fn(handle, byref(c_vgpu_sched_log))
+    _nvmlCheckReturn(ret)
+    return c_vgpu_sched_log
+
+def nvmlDeviceGetVgpuSchedulerState(handle):
+    c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuSchedulerState")
+    ret = fn(handle, byref(c_vgpu_sched_state))
+    _nvmlCheckReturn(ret)
+    return c_vgpu_sched_state
+
+def nvmlDeviceGetVgpuSchedulerCapabilities(handle):
+    c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetVgpuSchedulerCapabilities")
+    ret = fn(handle, byref(c_vgpu_sched_caps))
+    _nvmlCheckReturn(ret)
+    return c_vgpu_sched_caps
+
+def nvmlDeviceSetVgpuSchedulerState(handle, sched_state):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetVgpuSchedulerState")
+    ret = fn(handle, byref(sched_state))
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlSetVgpuVersion(vgpuVersion):
+    fn = _nvmlGetFunctionPointer("mxSmlExSetVgpuVersion")
+    ret = fn(byref(vgpuVersion))
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlGetVgpuVersion(supported, current):
+    fn = _nvmlGetFunctionPointer("mxSmlExGetVgpuVersion")
+    ret = fn(byref(supported), byref(current))
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlVgpuInstanceGetAccountingMode(vgpuInstance):
+    c_mode = _nvmlEnableState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetAccountingMode")
+    ret = fn(vgpuInstance, byref(c_mode))
+    _nvmlCheckReturn(ret)
+    return c_mode.value
+
+def nvmlVgpuInstanceGetAccountingPids(vgpuInstance):
+    c_pidCount = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetAccountingPids")
+    ret = fn(vgpuInstance, byref(c_pidCount), None)
+    if (ret == MXSMLEX_ERROR_INSUFFICIENT_SIZE):
+        sampleArray = c_pidCount.value * c_uint
+        c_pidArray = sampleArray()
+        ret = fn(vgpuInstance, byref(c_pidCount), byref(c_pidArray))
+        _nvmlCheckReturn(ret)
+    else:
+        raise NVMLError(ret)
+    return (c_pidCount, c_pidArray)
+
+def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid):
+    c_accountingStats = c_nvmlAccountingStats_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceGetAccountingStats")
+    ret = fn(vgpuInstance, pid, byref(c_accountingStats))
+    _nvmlCheckReturn(ret)
+    return c_accountingStats
+
+def nvmlVgpuInstanceClearAccountingPids(vgpuInstance):
+    fn = _nvmlGetFunctionPointer("mxSmlExVgpuInstanceClearAccountingPids")
+    ret = fn(vgpuInstance)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlGetExcludedDeviceCount():
+    c_count = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExGetExcludedDeviceCount")
+    ret = fn(byref(c_count))
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+def nvmlGetExcludedDeviceInfoByIndex(index):
+    c_index = c_uint(index)
+    info = c_nvmlExcludedDeviceInfo_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExGetExcludedDeviceInfoByIndex")
+    ret = fn(c_index, byref(info))
+    _nvmlCheckReturn(ret)
+    return info
+
+def nvmlDeviceGetHostVgpuMode(handle):
+    c_host_vgpu_mode = _nvmlHostVgpuMode_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetHostVgpuMode")
+    ret = fn(handle, byref(c_host_vgpu_mode))
+    _nvmlCheckReturn(ret)
+    return c_host_vgpu_mode.value
+
+def nvmlDeviceSetMigMode(device, mode):
+    c_activationStatus = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetMigMode")
+    ret = fn(device, mode, byref(c_activationStatus))
+    _nvmlCheckReturn(ret)
+    return c_activationStatus.value
+
+def nvmlDeviceGetMigMode(device):
+    c_currentMode = c_uint()
+    c_pendingMode = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMigMode")
+    ret = fn(device, byref(c_currentMode), byref(c_pendingMode))
+    _nvmlCheckReturn(ret)
+    return [c_currentMode.value, c_pendingMode.value]
+
+def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2):
+    if version == 2:
+        c_info = c_nvmlGpuInstanceProfileInfo_v2_t()
+        fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuInstanceProfileInfoV")
+    elif version == 1:
+        c_info = c_nvmlGpuInstanceProfileInfo_t()
+        fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuInstanceProfileInfo")
+    else:
+        raise NVMLError(MXSMLEX_ERROR_FUNCTION_NOT_FOUND)
+    ret = fn(device, profile, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return c_info
+
+# Define function alias for the API exposed by MXSMLEX
+nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo
+
+def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId):
+    c_count = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuInstanceRemainingCapacity")
+    ret = fn(device, profileId, byref(c_count))
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+def nvmlDeviceGetGpuInstancePossiblePlacements(device, profileId, placementsRef, countRef):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuInstancePossiblePlacements_v2")
+    ret = fn(device, profileId, placementsRef, countRef)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceCreateGpuInstance(device, profileId):
+    c_instance = c_nvmlGpuInstance_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceCreateGpuInstance")
+    ret = fn(device, profileId, byref(c_instance))
+    _nvmlCheckReturn(ret)
+    return c_instance
+
+def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement):
+    c_instance = c_nvmlGpuInstance_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceCreateGpuInstanceWithPlacement")
+    ret = fn(device, profileId, placement, byref(c_instance))
+    _nvmlCheckReturn(ret)
+    return c_instance
+
+def nvmlGpuInstanceDestroy(gpuInstance):
+    fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceDestroy")
+    ret = fn(gpuInstance)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuInstances")
+    ret = fn(device, profileId, gpuInstancesRef, countRef)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId):
+    c_instance = c_nvmlGpuInstance_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuInstanceById")
+    ret = fn(device, gpuInstanceId, byref(c_instance))
+    _nvmlCheckReturn(ret)
+    return c_instance
+
+def nvmlGpuInstanceGetInfo(gpuInstance):
+    c_info = c_nvmlGpuInstanceInfo_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceGetInfo")
+    ret = fn(gpuInstance, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return c_info
+
+def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, version=2):
+    if version == 2:
+        c_info = c_nvmlComputeInstanceProfileInfo_v2_t()
+        fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceGetComputeInstanceProfileInfoV")
+    elif version == 1:
+        c_info = c_nvmlComputeInstanceProfileInfo_t()
+        fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceGetComputeInstanceProfileInfo")
+    else:
+        raise NVMLError(MXSMLEX_ERROR_FUNCTION_NOT_FOUND) 
+    ret = fn(device, profile, engProfile, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return c_info
+
+# Define function alias for the API exposed by MXSMLEX
+nvmlGpuInstanceGetComputeInstanceProfileInfoV = nvmlGpuInstanceGetComputeInstanceProfileInfo
+
+def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId):
+    c_count = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceGetComputeInstanceRemainingCapacity")
+    ret = fn(gpuInstance, profileId, byref(c_count))
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+def nvmlGpuInstanceGetComputeInstancePossiblePlacements(gpuInstance, profileId, placementsRef, countRef):
+    fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceGetComputeInstancePossiblePlacements")
+    ret = fn(gpuInstance, profileId, placementsRef, countRef)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId):
+    c_instance = c_nvmlComputeInstance_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceCreateComputeInstance")
+    ret = fn(gpuInstance, profileId, byref(c_instance))
+    _nvmlCheckReturn(ret)
+    return c_instance
+
+def nvmlGpuInstanceCreateComputeInstanceWithPlacement(gpuInstance, profileId, placement):
+    c_instance = c_nvmlComputeInstance_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceCreateComputeInstanceWithPlacement")
+    ret = fn(gpuInstance, profileId, placement, byref(c_instance))
+    _nvmlCheckReturn(ret)
+    return c_instance
+
+def nvmlComputeInstanceDestroy(computeInstance):
+    fn = _nvmlGetFunctionPointer("mxSmlExComputeInstanceDestroy")
+    ret = fn(computeInstance)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlGpuInstanceGetComputeInstances(gpuInstance, profileId, computeInstancesRef, countRef):
+    fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceGetComputeInstances")
+    ret = fn(gpuInstance, profileId, computeInstancesRef, countRef)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId):
+    c_instance = c_nvmlComputeInstance_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExGpuInstanceGetComputeInstanceById")
+    ret = fn(gpuInstance, computeInstanceId, byref(c_instance))
+    _nvmlCheckReturn(ret)
+    return c_instance
+
+def nvmlComputeInstanceGetInfo_v2(computeInstance):
+    c_info = c_nvmlComputeInstanceInfo_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExComputeInstanceGetInfo_v2")
+    ret = fn(computeInstance, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return c_info
+
+def nvmlComputeInstanceGetInfo(computeInstance):
+    return nvmlComputeInstanceGetInfo_v2(computeInstance)
+
+def nvmlDeviceIsMigDeviceHandle(device):
+    c_isMigDevice = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceIsMigDeviceHandle")
+    ret = fn(device, byref(c_isMigDevice))
+    _nvmlCheckReturn(ret)
+    return c_isMigDevice
+
+def nvmlDeviceGetGpuInstanceId(device):
+    c_gpuInstanceId = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuInstanceId")
+    ret = fn(device, byref(c_gpuInstanceId))
+    _nvmlCheckReturn(ret)
+    return c_gpuInstanceId.value
+
+def nvmlDeviceGetComputeInstanceId(device):
+    c_computeInstanceId = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetComputeInstanceId")
+    ret = fn(device, byref(c_computeInstanceId))
+    _nvmlCheckReturn(ret)
+    return c_computeInstanceId.value
+
+def nvmlDeviceGetMaxMigDeviceCount(device):
+    c_count = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMaxMigDeviceCount")
+    ret = fn(device, byref(c_count))
+    _nvmlCheckReturn(ret)
+    return c_count.value
+
+def nvmlDeviceGetMigDeviceHandleByIndex(device, index):
+    c_index = c_uint(index)
+    migDevice = c_nvmlDevice_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMigDeviceHandleByIndex")
+    ret = fn(device, c_index, byref(migDevice))
+    _nvmlCheckReturn(ret)
+    return migDevice
+
+def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice):
+    device = c_nvmlDevice_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDeviceHandleFromMigDeviceHandle")
+    ret = fn(migDevice, byref(device))
+    _nvmlCheckReturn(ret)
+    return device
+
+def nvmlDeviceGetAttributes_v2(device):
+    c_attrs = c_nvmlDeviceAttributes()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetAttributes_v2")
+    ret = fn(device, byref(c_attrs))
+    _nvmlCheckReturn(ret)
+    return c_attrs
+
+def nvmlDeviceGetAttributes(device):
+    return nvmlDeviceGetAttributes_v2(device)
+
+def nvmlDeviceGetRemappedRows(device):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetRemappedRows")
+    c_corr = c_uint()
+    c_unc = c_uint()
+    c_bpending = c_uint()
+    c_bfailure = c_uint()
+    ret = fn(device, byref(c_corr), byref(c_unc), byref(c_bpending), byref(c_bfailure))
+    _nvmlCheckReturn(ret)
+    return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value)
+
+def nvmlDeviceGetRowRemapperHistogram(device):
+    c_vals = c_nvmlRowRemapperHistogramValues()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetRowRemapperHistogram")
+    ret = fn(device, byref(c_vals))
+    _nvmlCheckReturn(ret)
+    return c_vals
+
+def nvmlDeviceGetArchitecture(device):
+    arch = _nvmlDeviceArchitecture_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetArchitecture")
+    ret = fn(device, byref(arch))
+    _nvmlCheckReturn(ret)
+    return arch.value
+
+def nvmlDeviceGetBusType(device):
+    c_busType = _nvmlBusType_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetBusType")
+    ret = fn(device, byref(c_busType))
+    _nvmlCheckReturn(ret)
+    return c_busType.value
+
+def nvmlDeviceGetIrqNum(device):
+    c_irqNum = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetIrqNum")
+    ret = fn(device, byref(c_irqNum))
+    _nvmlCheckReturn(ret)
+    return c_irqNum.value
+
+def nvmlDeviceGetNumGpuCores(device):
+    c_numCores = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetNumGpuCores")
+    ret = fn(device, byref(c_numCores))
+    _nvmlCheckReturn(ret)
+    return c_numCores.value
+
+def nvmlDeviceGetPowerSource(device):
+    c_powerSource = _nvmlPowerSource_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPowerSource")
+    ret = fn(device, byref(c_powerSource))
+    _nvmlCheckReturn(ret)
+    return c_powerSource.value
+
+def nvmlDeviceGetMemoryBusWidth(device):
+    c_memBusWidth = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMemoryBusWidth")
+    ret = fn(device, byref(c_memBusWidth))
+    _nvmlCheckReturn(ret)
+    return c_memBusWidth.value
+
+def nvmlDeviceGetPcieLinkMaxSpeed(device):
+    c_speed = _nvmlPcieLinkMaxSpeed_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPcieLinkMaxSpeed")
+    ret = fn(device, byref(c_speed))
+    _nvmlCheckReturn(ret)
+    return c_speed.value
+
+def nvmlDeviceGetAdaptiveClockInfoStatus(device):
+    c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetAdaptiveClockInfoStatus")
+    ret = fn(device, byref(c_adaptiveClockInfoStatus))
+    _nvmlCheckReturn(ret)
+    return c_adaptiveClockInfoStatus.value
+
+def nvmlDeviceGetPcieSpeed(device):
+    c_speed = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetPcieSpeed")
+    ret = fn(device, byref(c_speed))
+    _nvmlCheckReturn(ret)
+    return c_speed.value
+
+def nvmlDeviceGetDynamicPstatesInfo(device, c_dynamicpstatesinfo):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetDynamicPstatesInfo");
+    ret = fn(device, c_dynamicpstatesinfo)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceSetFanSpeed_v2(handle, index, speed):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetFanSpeed_v2");
+    ret = fn(handle, index, speed)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetThermalSettings(device, sensorindex, c_thermalsettings):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetThermalSettings");
+    ret = fn(device, sensorindex, c_thermalsettings)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetMinMaxClockOfPState(device, type, pstate, minClockMHz, maxClockMHz):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMinMaxClockOfPState");
+    ret = fn(device, _nvmlClockType_t(type), _nvmlClockType_t(pstate), minClockMHz, maxClockMHz)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetSupportedPerformanceStates(device):
+    pstates = []
+    c_count = c_uint(MXSMLEX_MAX_GPU_PERF_PSTATES)
+    c_size = sizeof(c_uint)*c_count.value
+
+    # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration.
+    pstates_array = _nvmlPstates_t * c_count.value
+    c_pstates = pstates_array()
+
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSupportedPerformanceStates")
+    ret = fn(device, c_pstates, c_size)
+    _nvmlCheckReturn(ret)
+
+    for value in c_pstates:
+        if value != MXSMLEX_PSTATE_UNKNOWN:
+            pstates.append(value)
+
+    return pstates
+
+def nvmlDeviceGetGpcClkVfOffset(device):
+    offset = c_int32()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpcClkVfOffset")
+    ret = fn(device, byref(offset))
+    _nvmlCheckReturn(ret)
+    return offset.value
+
+def nvmlDeviceSetGpcClkVfOffset(device, offset):
+    c_offset = c_int32(offset)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetGpcClkVfOffset")
+    ret = fn(device, c_offset)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset, maxOffset):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpcClkMinMaxVfOffset")
+    ret = fn(device, minOffset, maxOffset)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetMemClkVfOffset(device):
+    offset = c_int32()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMemClkVfOffset")
+    ret = fn(device, byref(offset))
+    _nvmlCheckReturn(ret)
+    return offset.value
+
+def nvmlDeviceSetMemClkVfOffset(device, offset):
+    c_offset = c_int32(offset)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetMemClkVfOffset")
+    ret = fn(device, c_offset)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset, maxOffset):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetMemClkMinMaxVfOffset")
+    ret = fn(device, minOffset, maxOffset)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlSystemSetConfComputeGpusReadyState(state):
+    c_state = c_uint(state)
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemSetConfComputeGpusReadyState")
+    ret = fn(c_state)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlSystemGetConfComputeGpusReadyState():
+    c_state = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetConfComputeGpusReadyState")
+    ret = fn(byref(c_state))
+    _nvmlCheckReturn(ret)
+    return c_state.value
+
+def nvmlSystemGetConfComputeCapabilities():
+    c_ccSysCaps = c_nvmlConfComputeSystemCaps_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetConfComputeCapabilities")
+    ret = fn(byref(c_ccSysCaps))
+    _nvmlCheckReturn(ret)
+    return c_ccSysCaps
+
+def nvmlSystemGetConfComputeState():
+    c_state = c_nvmlConfComputeSystemState_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetConfComputeState")
+    ret = fn(byref(c_state))
+    _nvmlCheckReturn(ret)
+    return c_state
+
+def nvmlSystemGetConfComputeSettings(settings):
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetConfComputeSettings")
+    return fn(settings)
+
+def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetConfComputeUnprotectedMemSize")
+    ret = fn(device, c_ccMemSize)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetConfComputeMemSizeInfo(device):
+    c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetConfComputeMemSizeInfo")
+    ret = fn(device, byref(c_ccMemSize))
+    _nvmlCheckReturn(ret)
+    return c_ccMemSize
+
+def nvmlDeviceGetConfComputeProtectedMemoryUsage(device):
+    c_memory = c_nvmlMemory_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetConfComputeProtectedMemoryUsage")
+    ret = fn(device, byref(c_memory))
+    _nvmlCheckReturn(ret)
+    return c_memory
+
+def nvmlDeviceGetConfComputeGpuCertificate(device):
+    c_cert = c_nvmlConfComputeGpuCertificate_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetConfComputeGpuCertificate")
+    ret = fn(device, byref(c_cert))
+    _nvmlCheckReturn(ret)
+    return c_cert
+
+def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce):
+    c_attestReport = c_nvmlConfComputeGpuAttestationReport_t()
+    c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce))
+    setattr(c_attestReport, 'nonce', c_nonce_arr)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetConfComputeGpuAttestationReport")
+    ret = fn(device, byref(c_attestReport))
+    _nvmlCheckReturn(ret)
+    return c_attestReport
+
+def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv):
+    c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0)
+    c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1
+    c_keyRotationThrInfo.maxAttackerAdvantage = max_atk_adv
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemSetConfComputeKeyRotationThresholdInfo")
+    ret = fn(byref(c_keyRotationThrInfo))
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlSystemGetConfComputeKeyRotationThresholdInfo():
+    c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0)
+    c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetConfComputeKeyRotationThresholdInfo")
+    ret = fn(byref(c_keyRotationThrInfo))
+    _nvmlCheckReturn(ret)
+    return c_keyRotationThrInfo
+
+## GPM ##
+#########
+
+## Enums/defines
+
+#### GPM Metric Identifiers
+MXSMLEX_GPM_METRIC_GRAPHICS_UTIL           = 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_SM_UTIL                 = 2 # Percentage of SMs that were busy. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_SM_OCCUPANCY            = 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_INTEGER_UTIL            = 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_ANY_TENSOR_UTIL         = 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_DFMA_TENSOR_UTIL        = 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_HMMA_TENSOR_UTIL        = 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_IMMA_TENSOR_UTIL        = 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_DRAM_BW_UTIL            = 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_FP64_UTIL               = 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_FP32_UTIL               = 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_FP16_UTIL               = 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_PCIE_TX_PER_SEC         = 20 # PCIe traffic from this GPU in MiB/sec
+MXSMLEX_GPM_METRIC_PCIE_RX_PER_SEC         = 21 # PCIe traffic to this GPU in MiB/sec
+MXSMLEX_GPM_METRIC_NVDEC_0_UTIL            = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVDEC_1_UTIL            = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVDEC_2_UTIL            = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVDEC_3_UTIL            = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVDEC_4_UTIL            = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVDEC_5_UTIL            = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVDEC_6_UTIL            = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVDEC_7_UTIL            = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVJPG_0_UTIL            = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVJPG_1_UTIL            = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVJPG_2_UTIL            = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVJPG_3_UTIL            = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVJPG_4_UTIL            = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVJPG_5_UTIL            = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVJPG_6_UTIL            = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVJPG_7_UTIL            = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVOFA_0_UTIL            = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0
+MXSMLEX_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = 60 # NvLink read bandwidth for all links in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = 61 # NvLink write bandwidth for all links in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L0_RX_PER_SEC    = 62 # NvLink read bandwidth for link 0 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L0_TX_PER_SEC    = 63 # NvLink write bandwidth for link 0 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L1_RX_PER_SEC    = 64 # NvLink read bandwidth for link 1 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L1_TX_PER_SEC    = 65 # NvLink write bandwidth for link 1 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L2_RX_PER_SEC    = 66 # NvLink read bandwidth for link 2 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L2_TX_PER_SEC    = 67 # NvLink write bandwidth for link 2 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L3_RX_PER_SEC    = 68 # NvLink read bandwidth for link 3 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L3_TX_PER_SEC    = 69 # NvLink write bandwidth for link 3 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L4_RX_PER_SEC    = 70 # NvLink read bandwidth for link 4 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L4_TX_PER_SEC    = 71 # NvLink write bandwidth for link 4 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L5_RX_PER_SEC    = 72 # NvLink read bandwidth for link 5 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L5_TX_PER_SEC    = 73 # NvLink write bandwidth for link 5 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L6_RX_PER_SEC    = 74 # NvLink read bandwidth for link 6 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L6_TX_PER_SEC    = 75 # NvLink write bandwidth for link 6 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L7_RX_PER_SEC    = 76 # NvLink read bandwidth for link 7 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L7_TX_PER_SEC    = 77 # NvLink write bandwidth for link 7 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L8_RX_PER_SEC    = 78 # NvLink read bandwidth for link 8 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L8_TX_PER_SEC    = 79 # NvLink write bandwidth for link 8 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L9_RX_PER_SEC    = 80 # NvLink read bandwidth for link 9 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L9_TX_PER_SEC    = 81 # NvLink write bandwidth for link 9 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L10_RX_PER_SEC   = 82 # NvLink read bandwidth for link 10 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L10_TX_PER_SEC   = 83 # NvLink write bandwidth for link 10 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L11_RX_PER_SEC   = 84 # NvLink read bandwidth for link 11 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L11_TX_PER_SEC   = 85 # NvLink write bandwidth for link 11 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L12_RX_PER_SEC   = 86 # NvLink read bandwidth for link 12 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L12_TX_PER_SEC   = 87 # NvLink write bandwidth for link 12 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L13_RX_PER_SEC   = 88 # NvLink read bandwidth for link 13 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L13_TX_PER_SEC   = 89 # NvLink write bandwidth for link 13 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L14_RX_PER_SEC   = 90 # NvLink read bandwidth for link 14 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L14_TX_PER_SEC   = 91 # NvLink write bandwidth for link 14 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L15_RX_PER_SEC   = 92 # NvLink read bandwidth for link 15 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L15_TX_PER_SEC   = 93 # NvLink write bandwidth for link 15 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L16_RX_PER_SEC   = 94 # NvLink read bandwidth for link 16 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L16_TX_PER_SEC   = 95 # NvLink write bandwidth for link 16 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L17_RX_PER_SEC   = 96 # NvLink read bandwidth for link 17 in MiB/sec
+MXSMLEX_GPM_METRIC_NVLINK_L17_TX_PER_SEC   = 97 # NvLink write bandwidth for link 17 in MiB/sec
+MXSMLEX_GPM_METRIC_MAX                     = 98
+
+## Structs
+
+class c_nvmlUnitInfo_t(_PrintableStructure):
+    _fields_ = [
+        ('name', c_char * 96),
+        ('id', c_char * 96),
+        ('serial', c_char * 96),
+        ('firmwareVersion', c_char * 96),
+    ]
+
+class struct_c_nvmlGpmSample_t(Structure):
+    pass # opaque handle
+c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t)
+
+class c_metricInfo_t(Structure):
+    _fields_ = [
+        ("shortName", c_char_p),
+        ("longName", c_char_p),
+        ("unit", c_char_p),
+    ]
+
+class c_nvmlGpmMetric_t(_PrintableStructure):
+    _fields_ = [
+        ('metricId', c_uint),
+        ('nvmlReturn', _nvmlReturn_t),
+        ('value', c_double),
+        ('metricInfo', c_metricInfo_t)
+    ]
+
+class c_nvmlGpmMetricsGet_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('numMetrics', c_uint),
+        ('sample1', c_nvmlGpmSample_t),
+        ('sample2', c_nvmlGpmSample_t),
+        ('metrics', c_nvmlGpmMetric_t * MXSMLEX_GPM_METRIC_MAX)
+    ]
+
+MXSMLEX_GPM_METRICS_GET_VERSION = 1
+
+class c_nvmlGpmSupport_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('isSupportedDevice', c_uint),
+    ]
+
+MXSMLEX_GPM_SUPPORT_VERSION = 1
+
+## Functions
+
+def nvmlGpmMetricsGet(metricsGet):
+    fn = _nvmlGetFunctionPointer("mxSmlExGpmMetricsGet")
+    ret = fn(byref(metricsGet))
+    _nvmlCheckReturn(ret)
+    return metricsGet
+
+def nvmlGpmSampleFree(gpmSample):
+    fn = _nvmlGetFunctionPointer("mxSmlExGpmSampleFree")
+    ret = fn(gpmSample)
+    _nvmlCheckReturn(ret)
+    return
+
+def nvmlGpmSampleAlloc():
+    gpmSample = c_nvmlGpmSample_t()
+    fn = _nvmlGetFunctionPointer("mxSmlExGpmSampleAlloc")
+    ret = fn(byref(gpmSample))
+    _nvmlCheckReturn(ret)
+    return gpmSample
+
+def nvmlGpmSampleGet(device, gpmSample):
+    fn = _nvmlGetFunctionPointer("mxSmlExGpmSampleGet")
+    ret = fn(device, gpmSample)
+    _nvmlCheckReturn(ret)
+    return gpmSample
+
+def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample):
+    fn = _nvmlGetFunctionPointer("mxSmlExGpmMigSampleGet")
+    ret = fn(device, gpuInstanceId, gpmSample)
+    _nvmlCheckReturn(ret)
+    return gpmSample
+
+def nvmlGpmQueryDeviceSupport(device):
+    gpmSupport = c_nvmlGpmSupport_t()
+    gpmSupport.version = MXSMLEX_GPM_SUPPORT_VERSION
+    fn = _nvmlGetFunctionPointer("mxSmlExGpmQueryDeviceSupport")
+    ret = fn(device, byref(gpmSupport))
+    _nvmlCheckReturn(ret)
+    return gpmSupport
+
+def nvmlGpmSetStreamingEnabled(device, state):
+    c_state = c_uint(state)
+    fn = _nvmlGetFunctionPointer("mxSmlExGpmSetStreamingEnabled")
+    ret = fn(device, c_state)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlGpmQueryIfStreamingEnabled(device):
+    c_state = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExGpmQueryIfStreamingEnabled")
+    ret = fn(device, byref(c_state))
+    _nvmlCheckReturn(ret)
+    return c_state.value
+
+# Low Power Structure and Function
+
+MXSMLEX_NVLINK_POWER_STATE_HIGH_SPEED    = 0x0
+MXSMLEX_NVLINK_POWER_STATE_LOW           = 0x1
+
+MXSMLEX_NVLINK_LOW_POWER_THRESHOLD_MIN   = 0x1
+MXSMLEX_NVLINK_LOW_POWER_THRESHOLD_MAX   = 0x1FFF
+MXSMLEX_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF
+
+class c_nvmlNvLinkPowerThres_t(Structure):
+    _fields_ = [
+        ("lowPwrThreshold", c_uint),
+    ]
+
+def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold):
+    c_info = c_nvmlNvLinkPowerThres_t()
+    c_info.lowPwrThreshold = l1threshold
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetNvLinkDeviceLowPowerThreshold")
+    ret = fn(device, byref(c_info))
+    _nvmlCheckReturn(ret)
+    return ret 
+
+_nvmlGpuFabricState_t = c_uint
+MXSMLEX_GPU_FABRIC_STATE_NOT_SUPPORTED = 0
+MXSMLEX_GPU_FABRIC_STATE_NOT_STARTED   = 1
+MXSMLEX_GPU_FABRIC_STATE_IN_PROGRESS   = 2
+MXSMLEX_GPU_FABRIC_STATE_COMPLETED     = 3
+
+class c_nvmlGpuFabricInfo_t(_PrintableStructure):
+    _fields_ = [
+        ("clusterUuid", c_char * MXSMLEX_DEVICE_UUID_BUFFER_SIZE),
+        ("status", _nvmlReturn_t),
+        ("cliqueId", c_uint32),
+        ("state", _nvmlGpuFabricState_t)
+    ]
+
+nvmlGpuFabricInfo_v2 = 0x02000024
+
+class c_nvmlGpuFabricInfoV_t(_PrintableStructure):
+    _fields_ = [
+        ("version", c_uint),
+        ("clusterUuid", c_char * MXSMLEX_DEVICE_UUID_BUFFER_SIZE),
+        ("status", _nvmlReturn_t),
+        ("cliqueId", c_uint32),
+        ("state", _nvmlGpuFabricState_t),
+        ("healthMask", c_uint32)
+    ]
+
+    def __init__(self):
+        super(c_nvmlGpuFabricInfoV_t, self).__init__(version=nvmlGpuFabricInfo_v2)
+
+def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuFabricInfo");
+    ret = fn(device, gpuFabricInfo)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetGpuFabricInfoV");
+    ret = fn(device, gpuFabricInfo)
+    _nvmlCheckReturn(ret)
+    return ret
+
+######################
+## Enums/defines
+#### MXSMLEX GPU NVLINK BW MODE
+MXSMLEX_GPU_NVLINK_BW_MODE_FULL      = 0x0
+MXSMLEX_GPU_NVLINK_BW_MODE_OFF       = 0x1
+MXSMLEX_GPU_NVLINK_BW_MODE_MIN       = 0x2
+MXSMLEX_GPU_NVLINK_BW_MODE_HALF      = 0x3
+MXSMLEX_GPU_NVLINK_BW_MODE_3QUARTER  = 0x4
+MXSMLEX_GPU_NVLINK_BW_MODE_COUNT     = 0x5
+
+def nvmlSystemSetNvlinkBwMode(mode):
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemSetNvlinkBwMode")
+    ret = fn(mode)
+    _nvmlCheckReturn(ret)
+    return ret
+
+def nvmlSystemGetNvlinkBwMode():
+    mode = c_uint()
+    fn = _nvmlGetFunctionPointer("mxSmlExSystemGetNvlinkBwMode")
+    ret = fn(byref(mode))
+    _nvmlCheckReturn(ret)
+    return mode.value
+
+_nvmlPowerScopeType_t = c_uint
+MXSMLEX_POWER_SCOPE_GPU     = 0
+MXSMLEX_POWER_SCOPE_MODULE  = 1
+MXSMLEX_POWER_SCOPE_MEMORY  = 2
+
+class c_nvmlPowerValue_v2_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('powerScope', _nvmlPowerScopeType_t),
+        ('powerValueMw', c_uint),
+    ]
+    _fmt_ = {'<default>': "%d B"}
+
+nvmlPowerValue_v2 = 0x0200000C
+
+def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version=nvmlPowerValue_v2):
+    c_powerScope = _nvmlPowerScopeType_t(powerScope)
+    c_powerValue = c_nvmlPowerValue_v2_t()
+    c_powerValue.version = c_uint(version)
+    c_powerValue.powerScope = c_powerScope
+    c_powerValue.powerValueMw = c_uint(powerLimit)
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceSetPowerManagementLimit_v2")
+    ret = fn(device, byref(c_powerValue))
+    return ret
+
+class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure):
+    _fields_ = [
+        ('version', c_uint),
+        ('aggregateUncParity', c_ulonglong),
+        ('aggregateUncSecDed', c_ulonglong),
+        ('aggregateCor', c_ulonglong),
+        ('volatileUncParity', c_ulonglong),
+        ('volatileUncSecDed', c_ulonglong),
+        ('volatileCor', c_ulonglong),
+        ('aggregateUncBucketL2', c_ulonglong),
+        ('aggregateUncBucketSm', c_ulonglong),
+        ('aggregateUncBucketPcie', c_ulonglong),
+        ('aggregateUncBucketMcu', c_ulonglong),
+        ('aggregateUncBucketOther', c_ulonglong),
+        ('bThresholdExceeded', c_uint)
+    ]
+
+    def __init__(self):
+        super(c_nvmlEccSramErrorStatus_v1_t, self).__init__(version=nvmlEccSramErrorStatus_v1)
+
+nvmlEccSramErrorStatus_v1 = 0x1000068
+
+def nvmlDeviceGetSramEccErrorStatus(device, status):
+    fn = _nvmlGetFunctionPointer("mxSmlExDeviceGetSramEccErrorStatus")
+    ret = fn(device, status)
+    _nvmlCheckReturn(ret)
+    return ret
+
diff --git a/vllm/utils.py b/vllm/utils.py
index d8dd5f284..12590b707 100644
--- a/vllm/utils.py
+++ b/vllm/utils.py
@@ -1118,7 +1118,7 @@ def find_nccl_library() -> str:
             so_file)
     else:
         if torch.version.cuda is not None:
-            so_file = "libnccl.so.2"
+            so_file = "libmccl.so"
         elif torch.version.hip is not None:
             so_file = "librccl.so.1"
         else:
@@ -2698,7 +2698,7 @@ def import_pynvml():
     After all the troubles, we decide to copy the official `pynvml`
     module to our codebase, and use it directly.
     """
-    import vllm.third_party.pynvml as pynvml
+    import vllm.third_party.pymcml as pynvml
     return pynvml
 
 
diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py
index d7a580c28..1c4604cc2 100644
--- a/vllm/v1/attention/backends/cpu_attn.py
+++ b/vllm/v1/attention/backends/cpu_attn.py
@@ -7,7 +7,8 @@ from vllm.attention.backends.torch_sdpa import (TorchSDPABackendImpl,
                                                 TorchSDPAMetadata)
 from vllm.attention.backends.utils import CommonAttentionState
 from vllm.attention.ops.ipex_attn import PagedAttention
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
+                                              CommonAttentionMetadata)
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.v1.worker.block_table import BlockTable
@@ -53,7 +54,7 @@ class TorchSDPABackend:
         return False
 
 
-class TorchSDPAMetadataBuilderV1:
+class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
 
     def __init__(self, runner: CPUModelRunner, kv_cache_spec: AttentionSpec,
                  block_table: BlockTable) -> None:
@@ -118,9 +119,12 @@ class TorchSDPAMetadataBuilderV1:
 
         return True
 
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata):
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
         runner = self.runner
         block_table = self.block_table
         seq_lens_np = runner.seq_lens_np[:num_reqs]
diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py
index 91a7c43cd..22d361c9b 100755
--- a/vllm/v1/attention/backends/flash_attn.py
+++ b/vllm/v1/attention/backends/flash_attn.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Attention layer with FlashAttention."""
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Optional
 
 import numpy as np
 import torch
@@ -13,29 +13,32 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               is_quantized_kv_cache)
 from vllm.attention.layer import Attention
 from vllm.attention.ops.merge_attn_states import merge_attn_states
-from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,
-                                           get_flash_attn_version)
+# from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,
+#                                            get_flash_attn_version)
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     get_kv_connector_cache_layout)
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils import cdiv
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
+                                              CommonAttentionMetadata)
 from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.v1.worker.block_table import BlockTable
 
 if TYPE_CHECKING:
-    from vllm.v1.core.sched.output import SchedulerOutput
-    from vllm.v1.worker.gpu_input_batch import InputBatch
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 if current_platform.is_cuda():
-    from vllm.vllm_flash_attn import (flash_attn_varlen_func,
-                                      get_scheduler_metadata)
+    from flash_attn import (flash_attn_varlen_func, flash_attn_with_kvcache)
 
 logger = init_logger(__name__)
 
+def flash_attn_supports_fp8() -> bool:
+    return False
+
+def get_flash_attn_version():
+    return None
 
 class FlashAttentionBackend(AttentionBackend):
 
@@ -43,7 +46,7 @@ class FlashAttentionBackend(AttentionBackend):
 
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
-        return [32, 64, 96, 128, 160, 192, 224, 256]
+        return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256]
 
     @staticmethod
     def get_name() -> str:
@@ -102,6 +105,22 @@ class FlashAttentionMetadata:
     query_start_loc: torch.Tensor
     max_seq_len: int
     seq_lens: torch.Tensor
+
+    # For handling prefill decode split
+    num_decodes: int
+    num_decode_tokens: int
+    decode_query_start_loc: torch.Tensor
+    decode_max_seq_len: int
+    decode_seq_lens: torch.Tensor
+    decode_block_table: torch.Tensor
+
+    num_prefills: int
+    num_prefill_tokens: int
+    prefill_query_start_loc: torch.Tensor
+    prefill_max_seq_len: int
+    prefill_seq_lens: torch.Tensor
+    prefill_block_table: torch.Tensor
+
     block_table: torch.Tensor
     slot_mapping: torch.Tensor
 
@@ -306,7 +325,9 @@ def _get_sliding_window_configs(
     return sliding_window_configs
 
 
-class FlashAttentionMetadataBuilder:
+class FlashAttentionMetadataBuilder(
+        AttentionMetadataBuilder[FlashAttentionMetadata]):
+    full_cudagraph_supported: ClassVar[bool] = True  # Decode-only
 
     def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
                  block_table: BlockTable):
@@ -325,9 +346,9 @@ class FlashAttentionMetadataBuilder:
 
         self.aot_schedule = (get_flash_attn_version() == 3)
         self.use_full_cuda_graph = compilation_config.full_cuda_graph
-        if self.use_full_cuda_graph and not self.aot_schedule:
-            raise ValueError("Full CUDA graph mode requires AOT scheduling, "
-                             "which requires FlashAttention 3.")
+        # if self.use_full_cuda_graph and not self.aot_schedule:
+        #     raise ValueError("Full CUDA graph mode requires AOT scheduling, "
+        #                      "which requires FlashAttention 3.")
         self.scheduler_metadata = torch.zeros(self.runner.max_num_reqs + 1,
                                               dtype=torch.int32,
                                               device=self.runner.device)
@@ -338,11 +359,95 @@ class FlashAttentionMetadataBuilder:
 
     def reorder_batch(self, input_batch: "InputBatch",
                       scheduler_output: "SchedulerOutput") -> bool:
-        return False
+        # We now want to reorder the batch so that the "decode" requests are and
+        # the front and the "prefill" requests are at the using the least amount
+        # swaps possible. (NOTE for now we loosely use "decode" to mean requests
+        # where attention is likely memory-bound and "prefill" to mean requests
+        # where attention is likely compute-bound, TODO(lucas): figure out a
+        # better naming here)
+        decodes = []
+        prefills = []
+        num_decode_tokens = 0
+        num_prefill_tokens = 0
+
+        for i, req_id in enumerate(input_batch.req_ids):
+            num_tokens = scheduler_output.num_scheduled_tokens[req_id]
+            # for now treat 1 scheduled token as "decode" even if its not,
+            # we should update this to something like < 8 in the future but
+            # currently the decode run only supports num_tokens = 1
+            if num_tokens == 1:
+                decodes.append(i)
+                num_decode_tokens += num_tokens
+            else:
+                prefills.append(i)
+                num_prefill_tokens += num_tokens
+
+        # We hope that this is fairly minimal since decodes
+        # should be around for a number of iterations so hopefully they are
+        # relatively stationary (and new request are generally appended to the
+        # persistent batch so already should be at the back)
+        # To achieve this we loop over the decodes in descending order and
+        # the prefills in ascending order. We swap decodes from the  "back"
+        # i.e. past where the last decode should be in the reodorered with
+        # prefills from the front of the batch.
+        # `decodes` and `prefills` are already in ascending order just based on
+        # the above loop
+        num_decodes = len(decodes)
+        num_prefills = len(prefills)
+        modified_batch = False
+
+        for i in range(1, min(num_decodes, num_prefills) + 1):
+            # If the decode is at the "back" of the batch, i, we can swap it
+            # with the prefill closest to the front of the batch
+            decode_idx = decodes[num_decodes - i]
+            if decode_idx < num_decodes:
+                break
+
+            input_batch.swap_states(prefills[i - 1], decode_idx)
+            modified_batch = True
+
+        # Save for next `build` call
+        # TODO(lucas): this is a bit of a hack, we should probably have a
+        # better way of doing this
+        self._num_decodes = num_decodes
+        self._num_prefills = num_prefills
+        self._num_decode_tokens = num_decode_tokens
+        self._num_prefill_tokens = num_prefill_tokens
+
+        return modified_batch
+
+    def build_for_cudagraph_capture(
+            self, common_attn_metadata: CommonAttentionMetadata) -> FlashAttentionMetadata:
+        """
+        This method builds the metadata for full cudagraph capture.
+        Currently, only decode is supported for full cudagraphs with MLA.
+        """
+        m = common_attn_metadata
+        assert m.num_reqs == m.num_actual_tokens, \
+            "MLA only supports decode-only full CUDAGraph capture. " \
+            "Make sure all cudagraph capture sizes <= max_num_seq."
+
+        m.max_query_len = 1  # decode-only
+
+        # Update state usually set in reorder_batch.
+        self._num_decodes = m.num_reqs
+        self._num_decode_tokens = m.num_actual_tokens
+        self._num_prefills = 0
+        self._num_prefill_tokens = 0
+        return self.build(0, m)
+
+    def build(
+        self, common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata
+    ) -> FlashAttentionMetadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
+        assert self._num_decodes + self._num_prefills == num_reqs
+        assert (self._num_decode_tokens +
+                self._num_prefill_tokens == num_actual_tokens)
 
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
-              common_attn_metadata: CommonAttentionMetadata):
         max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
         query_start_loc = common_attn_metadata.query_start_loc
         seq_lens = common_attn_metadata.seq_lens
@@ -358,6 +463,30 @@ class FlashAttentionMetadataBuilder:
 
         slot_mapping = block_table.slot_mapping[:num_actual_tokens]
 
+        # For handling prefill decode split
+        if self._num_decodes > 0:
+            decode_max_seq_len = int(self.runner.seq_lens_np[:self._num_decodes].max())
+            decode_query_start_loc = common_attn_metadata.query_start_loc[:self._num_decodes + 1]
+            decode_seq_lens = common_attn_metadata.seq_lens[:self._num_decodes]
+            decode_block_table_tensor = block_table.get_device_tensor()[:self._num_decodes]
+        else:
+            decode_max_seq_len = 0
+            decode_query_start_loc = None
+            decode_seq_lens = None
+            decode_block_table_tensor = None
+
+        if self._num_prefills > 0:
+            prefill_max_seq_len = int(self.runner.seq_lens_np[self._num_decodes:num_reqs].max())
+            prefill_query_start_loc = (common_attn_metadata.query_start_loc[self._num_decodes:num_reqs + 1] -
+                                       common_attn_metadata.query_start_loc[self._num_decodes])
+            prefill_seq_lens = common_attn_metadata.seq_lens[self._num_decodes:num_reqs]
+            prefill_block_table_tensor = block_table.get_device_tensor()[self._num_decodes:num_reqs]
+        else:
+            prefill_max_seq_len = 0
+            prefill_query_start_loc = None
+            prefill_seq_lens = None
+            prefill_block_table_tensor = None
+        
         if self.aot_sliding_window is None:
             self.aot_sliding_window = (-1, -1)
             # For the AOT scheduler we need the sliding window value to be
@@ -376,20 +505,20 @@ class FlashAttentionMetadataBuilder:
 
         def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
                      max_seq_len, causal):
-            if self.aot_schedule:
-                return get_scheduler_metadata(
-                    batch_size=batch_size,
-                    max_seqlen_q=max_query_len,
-                    max_seqlen_k=max_seq_len,
-                    cache_seqlens=seqlens,
-                    num_heads_q=self.num_heads_q,
-                    num_heads_kv=self.num_heads_kv,
-                    headdim=self.headdim,
-                    page_size=self.block_size,
-                    cu_seqlens_q=cu_query_lens,
-                    causal=causal,
-                    window_size=self.aot_sliding_window,
-                )
+            # if self.aot_schedule:
+            #     return get_scheduler_metadata(
+            #         batch_size=batch_size,
+            #         max_seqlen_q=max_query_len,
+            #         max_seqlen_k=max_seq_len,
+            #         cache_seqlens=seqlens,
+            #         num_heads_q=self.num_heads_q,
+            #         num_heads_kv=self.num_heads_kv,
+            #         headdim=self.headdim,
+            #         page_size=self.block_size,
+            #         cu_seqlens_q=cu_query_lens,
+            #         causal=causal,
+            #         window_size=self.aot_sliding_window,
+            #     )
             return None
 
         # for local attention
@@ -465,17 +594,17 @@ class FlashAttentionMetadataBuilder:
                                           max_seq_len=max_seq_len,
                                           causal=True)
 
-        if self.use_full_cuda_graph:
-            assert scheduler_metadata is not None
-            n = scheduler_metadata.shape[0]
-            self.scheduler_metadata[:n].copy_(scheduler_metadata,
-                                              non_blocking=True)
-            # NOTE(woosuk): We should zero out the rest of the scheduler
-            # metadata to guarantee the correctness. Otherwise, some thread
-            # blocks may use the invalid scheduler metadata and overwrite the
-            # output buffer.
-            self.scheduler_metadata[n:] = 0
-            scheduler_metadata = self.scheduler_metadata[:n]
+        # if self.use_full_cuda_graph:
+        #     assert scheduler_metadata is not None
+        #     n = scheduler_metadata.shape[0]
+        #     self.scheduler_metadata[:n].copy_(scheduler_metadata,
+        #                                       non_blocking=True)
+        #     # NOTE(woosuk): We should zero out the rest of the scheduler
+        #     # metadata to guarantee the correctness. Otherwise, some thread
+        #     # blocks may use the invalid scheduler metadata and overwrite the
+        #     # output buffer.
+        #     self.scheduler_metadata[n:] = 0
+        #     scheduler_metadata = self.scheduler_metadata[:n]
 
         attn_metadata = FlashAttentionMetadata(
             num_actual_tokens=num_actual_tokens,
@@ -483,6 +612,19 @@ class FlashAttentionMetadataBuilder:
             query_start_loc=query_start_loc,
             max_seq_len=max_seq_len,
             seq_lens=seq_lens,
+            # For handling prefill decode split
+            num_decodes=self._num_decodes,
+            num_decode_tokens=self._num_decode_tokens,
+            decode_query_start_loc=decode_query_start_loc,
+            decode_max_seq_len=decode_max_seq_len,
+            decode_seq_lens=decode_seq_lens,
+            decode_block_table=decode_block_table_tensor,
+            num_prefills=self._num_prefills,
+            num_prefill_tokens=self._num_prefill_tokens,
+            prefill_query_start_loc=prefill_query_start_loc,
+            prefill_max_seq_len=prefill_max_seq_len,
+            prefill_seq_lens=prefill_seq_lens,
+            prefill_block_table=prefill_block_table_tensor,
             block_table=block_table_tensor,
             slot_mapping=slot_mapping,
             use_cascade=use_cascade,
@@ -496,6 +638,10 @@ class FlashAttentionMetadataBuilder:
         )
         return attn_metadata
 
+    def can_run_in_cudagraph(
+            self, common_attn_metadata: CommonAttentionMetadata) -> bool:
+        return common_attn_metadata.max_query_len == 1
+
     def use_cascade_attention(self, *args, **kwargs) -> bool:
         return use_cascade_attention(*args, **kwargs)
 
@@ -634,6 +780,43 @@ class FlashAttentionImpl(AttentionImpl):
         # Compute attention and update output up to `num_actual_tokens`.
         use_local_attn = \
             (self.use_irope and attn_metadata.local_attn_metadata is not None)
+        
+        # For handling prefill decode split
+        if not attn_metadata.use_cascade and not use_local_attn:
+            num_decode_tokens = attn_metadata.num_decode_tokens
+            if attn_metadata.num_prefills > 0:
+                cu_prefix_kv_lens = torch.tensor([0] + attn_metadata.prefill_seq_lens.tolist(), device=attn_metadata.prefill_seq_lens.device, dtype=torch.int32).cumsum(dim=0, dtype=torch.int32)
+                output[num_decode_tokens:num_actual_tokens] = flash_attn_varlen_func(
+                    q=query[num_decode_tokens:num_actual_tokens],
+                    k=key_cache,
+                    v=value_cache,
+                    block_table=attn_metadata.prefill_block_table,
+                    cu_seqlens_q=attn_metadata.prefill_query_start_loc,
+                    cu_seqlens_k=cu_prefix_kv_lens,
+                    max_seqlen_q=attn_metadata.max_query_len,
+                    max_seqlen_k=attn_metadata.prefill_max_seq_len,
+                    softmax_scale=self.scale,
+                    causal=True,
+                    window_size=self.sliding_window,
+                    alibi_slopes=self.alibi_slopes,
+                    softcap=self.logits_soft_cap,
+                )
+            if attn_metadata.num_decodes > 0:
+                # Use flash_attn_with_kvcache for normal decoding.
+                decode_query = query[:num_decode_tokens]
+                output[:num_decode_tokens] = flash_attn_with_kvcache(
+                    q=decode_query.unsqueeze(1),
+                    k_cache=key_cache,
+                    v_cache=value_cache,
+                    block_table=attn_metadata.decode_block_table,
+                    cache_seqlens=attn_metadata.decode_seq_lens,
+                    softmax_scale=self.scale,
+                    causal=True,
+                    window_size=self.sliding_window,
+                    alibi_slopes=self.alibi_slopes,
+                    softcap=self.logits_soft_cap,
+                ).squeeze(1)
+            return output
 
         if not attn_metadata.use_cascade or use_local_attn:
             if use_local_attn:
@@ -653,16 +836,18 @@ class FlashAttentionImpl(AttentionImpl):
                 block_table = attn_metadata.block_table
                 scheduler_metadata = attn_metadata.scheduler_metadata
 
-            descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
-
-            flash_attn_varlen_func(
+            # descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
+            cu_prefix_kv_lens = torch.tensor([0] + seqused_k.tolist(), seqused_k.device, 
+                                             dtype=torch.int32).cumsum(dim=0, dtype=torch.int32)
+            output[:num_actual_tokens] = flash_attn_varlen_func(
                 q=query[:num_actual_tokens],
                 k=key_cache,
                 v=value_cache,
-                out=output[:num_actual_tokens],
+                # out=output[:num_actual_tokens],
                 cu_seqlens_q=cu_seqlens_q,
                 max_seqlen_q=max_seqlen_q,
-                seqused_k=seqused_k,
+                # seqused_k=seqused_k,
+                cu_seqlens_k=cu_prefix_kv_lens,
                 max_seqlen_k=max_seqlen_k,
                 softmax_scale=self.scale,
                 causal=True,
@@ -670,11 +855,11 @@ class FlashAttentionImpl(AttentionImpl):
                 window_size=self.sliding_window,
                 block_table=block_table,
                 softcap=self.logits_soft_cap,
-                scheduler_metadata=scheduler_metadata,
-                fa_version=self.vllm_flash_attn_version,
-                q_descale=layer._q_scale.expand(descale_shape),
-                k_descale=layer._k_scale.expand(descale_shape),
-                v_descale=layer._v_scale.expand(descale_shape),
+                # scheduler_metadata=scheduler_metadata,
+                # fa_version=self.vllm_flash_attn_version,
+                # q_descale=layer._q_scale.expand(descale_shape),
+                # k_descale=layer._k_scale.expand(descale_shape),
+                # v_descale=layer._v_scale.expand(descale_shape),
             )
             return output
 
@@ -813,12 +998,15 @@ def cascade_attention(
     descale_shape = (cu_prefix_query_lens.shape[0] - 1, key_cache.shape[-2])
 
     # Process shared prefix.
+    cu_prefix_kv_lens = torch.tensor([0] + prefix_kv_lens.tolist(), 
+                                     device=prefix_kv_lens.device, 
+                                     dtype=torch.int32).cumsum(dim=0, dtype=torch.int32)
     prefix_output, prefix_lse = flash_attn_varlen_func(
         q=query,
         k=key_cache,
         v=value_cache,
         cu_seqlens_q=cu_prefix_query_lens,
-        seqused_k=prefix_kv_lens,
+        cu_seqlens_k=cu_prefix_kv_lens,
         max_seqlen_q=num_tokens,
         max_seqlen_k=common_prefix_len,
         softmax_scale=softmax_scale,
@@ -827,25 +1015,28 @@ def cascade_attention(
         block_table=block_table[:1],
         softcap=logits_soft_cap,
         return_softmax_lse=True,
-        scheduler_metadata=prefix_scheduler_metadata,
-        fa_version=fa_version,
-        q_descale=q_descale.expand(descale_shape)
-        if q_descale is not None else None,
-        k_descale=k_descale.expand(descale_shape)
-        if k_descale is not None else None,
-        v_descale=v_descale.expand(descale_shape)
-        if v_descale is not None else None,
+        # scheduler_metadata=prefix_scheduler_metadata,
+        # fa_version=fa_version,
+        # q_descale=q_descale.expand(descale_shape)
+        # if q_descale is not None else None,
+        # k_descale=k_descale.expand(descale_shape)
+        # if k_descale is not None else None,
+        # v_descale=v_descale.expand(descale_shape)
+        # if v_descale is not None else None,
     )
 
     descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])
 
     # Process suffix per query.
+    cu_suffix_kv_lens = torch.tensor([0] + suffix_kv_lens.tolist(), 
+                                     device=suffix_kv_lens.device, 
+                                     dtype=torch.int32).cumsum(dim=0, dtype=torch.int32)
     suffix_output, suffix_lse = flash_attn_varlen_func(
         q=query,
         k=key_cache,
         v=value_cache,
         cu_seqlens_q=cu_query_lens,
-        seqused_k=suffix_kv_lens,
+        cu_seqlens_k=cu_suffix_kv_lens,
         max_seqlen_q=max_query_len,
         max_seqlen_k=max_kv_len - common_prefix_len,
         softmax_scale=softmax_scale,
@@ -854,14 +1045,14 @@ def cascade_attention(
         block_table=block_table[:, num_common_kv_blocks:],
         softcap=logits_soft_cap,
         return_softmax_lse=True,
-        scheduler_metadata=suffix_scheduler_metadata,
-        fa_version=fa_version,
-        q_descale=q_descale.expand(descale_shape)
-        if q_descale is not None else None,
-        k_descale=k_descale.expand(descale_shape)
-        if k_descale is not None else None,
-        v_descale=v_descale.expand(descale_shape)
-        if v_descale is not None else None,
+        # scheduler_metadata=suffix_scheduler_metadata,
+        # fa_version=fa_version,
+        # q_descale=q_descale.expand(descale_shape)
+        # if q_descale is not None else None,
+        # k_descale=k_descale.expand(descale_shape)
+        # if k_descale is not None else None,
+        # v_descale=v_descale.expand(descale_shape)
+        # if v_descale is not None else None,
     )
 
     # Merge prefix and suffix outputs, and store the result in output.
diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py
index b15bb4b31..f8455b5dc 100755
--- a/vllm/v1/attention/backends/flashinfer.py
+++ b/vllm/v1/attention/backends/flashinfer.py
@@ -8,8 +8,8 @@ from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
-                        BatchPrefillWithPagedKVCacheWrapper,
-                        MultiLevelCascadeAttentionWrapper)
+                        BatchPrefillWithPagedKVCacheWrapper)
+                        #MultiLevelCascadeAttentionWrapper
 
 import vllm.envs as envs
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
@@ -18,7 +18,8 @@ from vllm.attention.layer import Attention
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.logger import init_logger
 from vllm.v1.attention.backends.flash_attn import use_cascade_attention
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
+                                              CommonAttentionMetadata)
 from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.v1.worker.block_table import BlockTable
 
@@ -184,7 +185,7 @@ class FlashInferMetadata:
 
     prefill_wrapper: Optional[BatchPrefillWithPagedKVCacheWrapper] = None
     decode_wrapper: Optional[BatchDecodeWithPagedKVCacheWrapper] = None
-    cascade_wrapper: Optional[MultiLevelCascadeAttentionWrapper] = None
+    # cascade_wrapper: Optional[MultiLevelCascadeAttentionWrapper] = None
 
     @property
     def query_start_loc(self):
@@ -202,7 +203,7 @@ class FlashInferMetadata:
                 f" received {self.head_dim}.")
 
 
-class FlashInferMetadataBuilder:
+class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
 
     def __init__(self, runner: GPUModelRunner, kv_cache_spec: AttentionSpec,
                  block_table: BlockTable):
@@ -306,17 +307,17 @@ class FlashInferMetadataBuilder:
                 use_tensor_cores=use_tensor_cores)
         return self._decode_wrapper
 
-    def _get_cascade_wrapper(self):
-        if self._cascade_wrapper is None:
-            self._cascade_wrapper = MultiLevelCascadeAttentionWrapper(
-                2, self._get_workspace_buffer(), "NHD")
-        return self._cascade_wrapper
+    # def _get_cascade_wrapper(self):
+    #     if self._cascade_wrapper is None:
+    #         self._cascade_wrapper = MultiLevelCascadeAttentionWrapper(
+    #             2, self._get_workspace_buffer(), "NHD")
+    #     return self._cascade_wrapper
 
     def _plan(self, attn_metadata: FlashInferMetadata):
         if self.global_hyperparameters is None:
             self.global_hyperparameters = infer_global_hyperparameters(
                 get_per_layer_parameters(self.vllm_config))
-        if attn_metadata.use_cascade:
+        if attn_metadata.use_cascade and False: # not supported
             attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
             attn_metadata.cascade_wrapper.plan(
                 [attn_metadata.shared_qo_indptr, attn_metadata.qo_indptr],
@@ -399,9 +400,11 @@ class FlashInferMetadataBuilder:
                     kv_data_type=attn_metadata.data_type,
                 )
 
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata):
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+
         assert self._num_decodes + self._num_prefills == num_reqs
         assert (self._num_decode_tokens +
                 self._num_prefill_tokens == num_actual_tokens)
@@ -486,6 +489,9 @@ class FlashInferMetadataBuilder:
         return attn_metadata
 
     def use_cascade_attention(self, *args, **kwargs) -> bool:
+        logger.warning_once(
+                "Using cascade attention in FlashInfer is not supported yet")
+        return False
         if self.kv_cache_spec.dtype != self.runner.model_config.dtype:
             # TODO: The cascade wrapper currently does not support setting
             # kv cache dtype to something different from query dtype.
@@ -603,11 +609,11 @@ class FlashInferImpl(AttentionImpl):
         output_padded = output
         output = output[:num_actual_tokens]
 
-        if attn_metadata.use_cascade:
-            # Cascade attention (rare case).
-            assert attn_metadata.cascade_wrapper is not None
-            output.copy_(attn_metadata.cascade_wrapper.run(query, kv_cache))
-            return output
+        # if attn_metadata.use_cascade:
+        #     # Cascade attention (rare case).
+        #     assert attn_metadata.cascade_wrapper is not None
+        #     output.copy_(attn_metadata.cascade_wrapper.run(query, kv_cache))
+        #     return output
 
         num_decode_tokens = attn_metadata.num_decode_tokens
         num_prefill_tokens = attn_metadata.num_prefill_tokens
diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py
index 5b473b146..c3efb938e 100644
--- a/vllm/v1/attention/backends/flex_attention.py
+++ b/vllm/v1/attention/backends/flex_attention.py
@@ -15,7 +15,8 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               is_quantized_kv_cache)
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
+                                              CommonAttentionMetadata)
 from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.v1.worker.block_table import BlockTable
 
@@ -25,8 +26,6 @@ if current_platform.is_cuda():
 logger = init_logger(__name__)
 
 if TYPE_CHECKING:
-    from vllm.v1.core.sched.output import SchedulerOutput
-    from vllm.v1.worker.gpu_input_batch import InputBatch
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 create_block_mask_compiled = torch.compile(create_block_mask,
@@ -256,7 +255,8 @@ class FlexAttentionMetadata:
         self.block_mask = self.build_block_mask()
 
 
-class FlexAttentionMetadataBuilder:
+class FlexAttentionMetadataBuilder(
+        AttentionMetadataBuilder[FlexAttentionMetadata]):
 
     def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
                  block_table: BlockTable):
@@ -272,13 +272,12 @@ class FlexAttentionMetadataBuilder:
         self.kv_cache_spec = kv_cache_spec
         self.block_table = block_table
 
-    def reorder_batch(self, input_batch: "InputBatch",
-                      scheduler_output: "SchedulerOutput") -> bool:
-        return False
-
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata):
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
         max_seq_len = self.runner.seq_lens_np[:num_reqs].max()
         query_start_loc = common_attn_metadata.query_start_loc
         seq_lens = common_attn_metadata.seq_lens
@@ -332,9 +331,6 @@ class FlexAttentionMetadataBuilder:
         )
         return out
 
-    def use_cascade_attention(self, *args, **kwargs) -> bool:
-        return False
-
 
 class FlexAttentionImpl(AttentionImpl):
     sliding_window: Optional[tuple[int, int]]
diff --git a/vllm/v1/attention/backends/mla/common.py b/vllm/v1/attention/backends/mla/common.py
index e6b4f6404..6f28c0f47 100644
--- a/vllm/v1/attention/backends/mla/common.py
+++ b/vllm/v1/attention/backends/mla/common.py
@@ -200,14 +200,15 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
                                               MLAAttentionImpl)
 from vllm.attention.backends.utils import get_mla_dims
 from vllm.attention.ops.merge_attn_states import merge_attn_states
-from vllm.attention.utils.fa_utils import get_flash_attn_version
+# from vllm.attention.utils.fa_utils import get_flash_attn_version
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                LinearBase,
                                                UnquantizedLinearMethod)
 from vllm.platforms import current_platform
 from vllm.utils import cdiv, round_down
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
+                                              CommonAttentionMetadata)
 from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.v1.worker.block_table import BlockTable
 
@@ -223,9 +224,13 @@ if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
     from vllm.v1.worker.gpu_input_batch import InputBatch
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+    
+from vllm import envs
 
 logger = init_logger(__name__)
 
+def get_flash_attn_version():
+    return None
 
 class MLACommonBackend(AttentionBackend):
 
@@ -329,7 +334,7 @@ class MLACommonMetadata(Generic[D]):
 M = TypeVar("M", bound=MLACommonMetadata)
 
 
-class MLACommonMetadataBuilder(Generic[M]):
+class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
     """
     NOTE: Please read the comment at the top of the file before trying to
     understand this class
@@ -350,7 +355,7 @@ class MLACommonMetadataBuilder(Generic[M]):
         self.num_heads = model_config.get_num_attention_heads(
             runner.parallel_config)
         self.mla_dims = get_mla_dims(model_config)
-        self.aot_schedule = current_platform.is_cuda()
+        self.aot_schedule = False # current_platform.is_cuda()
         self.kv_cache_spec = kv_cache_spec
 
         # Dont try to access the runner on AMD
@@ -450,9 +455,32 @@ class MLACommonMetadataBuilder(Generic[M]):
             seq_lens=seq_lens,
         )
 
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build_for_cudagraph_capture(
+            self, common_attn_metadata: CommonAttentionMetadata) -> M:
+        """
+        This method builds the metadata for full cudagraph capture.
+        Currently, only decode is supported for full cudagraphs with MLA.
+        """
+        m = common_attn_metadata
+        assert m.num_reqs == m.num_actual_tokens, \
+            "MLA only supports decode-only full CUDAGraph capture. " \
+            "Make sure all cudagraph capture sizes <= max_num_seq."
+
+        m.max_query_len = 1  # decode-only
+
+        # Update state usually set in reorder_batch.
+        self._num_decodes = m.num_reqs
+        self._num_decode_tokens = m.num_actual_tokens
+        self._num_prefills = 0
+        self._num_prefill_tokens = 0
+        return self.build(0, m)
+
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata) -> M:
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
         assert self._num_decodes + self._num_prefills == num_reqs
 
         # Note(simon): be careful about the CPU <> GPU memory movement in this
@@ -461,8 +489,11 @@ class MLACommonMetadataBuilder(Generic[M]):
         device = self.runner.device
         block_table = self.block_table
         block_table_tensor = block_table.get_device_tensor()[:num_reqs]
-        slot_mapping = block_table.slot_mapping_cpu[:num_actual_tokens].to(
-            device, non_blocking=True).long()
+        block_table.slot_mapping[:num_actual_tokens].copy_(
+            block_table.slot_mapping_cpu[:num_actual_tokens],
+            non_blocking=True)
+        block_table.slot_mapping[num_actual_tokens:].fill_(-1)
+        slot_mapping = block_table.slot_mapping[:num_actual_tokens]
 
         query_start_loc = common_attn_metadata.query_start_loc
         seq_lens = common_attn_metadata.seq_lens
@@ -564,8 +595,9 @@ class MLACommonMetadataBuilder(Generic[M]):
             decode=decode_metadata,
         )
 
-    def use_cascade_attention(self, *args, **kwargs) -> bool:
-        return False
+    def can_run_in_cudagraph(
+            self, common_attn_metadata: CommonAttentionMetadata) -> bool:
+        return common_attn_metadata.max_query_len == 1
 
 
 class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
@@ -644,14 +676,25 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             maybe_padded_v = torch.nn.functional.pad(
                 v, [0, q.shape[-1] - v.shape[-1]], value=0)
 
-        attn_out = self.flash_attn_varlen_func(
-            q=q,
-            k=k,
-            v=maybe_padded_v,
-            return_softmax_lse=return_softmax_lse,
-            softmax_scale=softmax_scale,
-            **kwargs,
-        )
+        if is_vllm_fa:
+            attn_out = self.flash_attn_varlen_func(
+                q=q,
+                k=k,
+                v=maybe_padded_v,
+                return_softmax_lse=return_softmax_lse,
+                softmax_scale=softmax_scale,
+                **kwargs,
+            )
+        else:
+            # Use return_attn_probs instead of return_softmax_lse for RoCM
+            attn_out = self.flash_attn_varlen_func(
+                q=q,
+                k=k,
+                v=maybe_padded_v,
+                return_attn_probs=return_softmax_lse,
+                softmax_scale=softmax_scale,
+                **kwargs,
+            )
 
         # Unpack the output if there is multiple results
         lse = None
@@ -695,7 +738,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
                 del eye
                 # standardize to (output, input)
                 return dequant_weights.T
-            return layer.weight
+            return layer.weight if not envs.MACA_VLLM_USE_TN_2_NN else layer.weight.T
 
         # we currently do not have quantized bmm's which are needed for
         # `W_UV` and `W_UK_T`, we we just store fp16/bf16 copies and perform
@@ -805,7 +848,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
     ) -> torch.Tensor:
         assert attn_metadata.prefill is not None
 
-        has_context = attn_metadata.prefill.chunked_context is not None
+        has_context = False # attn_metadata.prefill.chunked_context is not None
         kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\
             -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
         k_nope, v = kv_nope\
@@ -823,7 +866,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             max_seqlen_k=attn_metadata.prefill.max_query_len,
             softmax_scale=self.scale,
             causal=True,
-            return_softmax_lse=has_context,
+            # return_softmax_lse=has_context,
         )
 
         if has_context:
@@ -842,9 +885,10 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
 
         # unpad if necessary
         if self._pad_v:
-            output = output[..., :v.shape[-1]]
+            output = output.view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]] \
+                .reshape(-1, self.num_heads * v.shape[-1])
 
-        return output.flatten(start_dim=-2)
+        return output
 
     @abstractmethod
     def _forward_decode(
diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py
index 318b8ede1..87733c54f 100644
--- a/vllm/v1/attention/backends/mla/flashmla.py
+++ b/vllm/v1/attention/backends/mla/flashmla.py
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any, ClassVar, Optional
 
 import torch
 
@@ -44,7 +44,7 @@ class FlashMLABackend(MLACommonBackend):
 
 @dataclass
 class FlashMLADecodeMetadata(MLACommonDecodeMetadata):
-    tile_scheduler_metadata: tuple[torch.Tensor, torch.Tensor]
+    tile_scheduler_metadata: torch.Tensor
     num_splits: torch.Tensor
 
 
@@ -54,13 +54,17 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):
 
 
 class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
+    full_cudagraph_supported: ClassVar[bool] = True  # Decode-only
 
     def __init__(self, runner, kv_cache_spec: AttentionSpec,
                  block_table: BlockTable):
-        super().__init__(runner, kv_cache_spec, block_table)
+        super().__init__(runner, kv_cache_spec, block_table, FlashMLAMetadata)
 
         self.num_q_heads = self.runner.model_config.get_num_attention_heads(
             self.runner.parallel_config)
+        
+        self.cg_buf_tile_scheduler_metadata = None
+        self.cg_buf_num_splits = None
 
     def _build_decode(self, block_table_tensor: torch.Tensor,
                       seq_lens: torch.Tensor) -> FlashMLADecodeMetadata:
@@ -70,6 +74,30 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
             self.num_q_heads,
             1, # MQA for the decode path
         )
+        
+        if self.runner.full_cuda_graph:
+            # First time around (CUDAGraph capture), allocate the static buffer
+            if self.cg_buf_tile_scheduler_metadata is None:
+                self.cg_buf_tile_scheduler_metadata = tile_scheduler_metadata
+                self.cg_buf_num_splits = num_splits
+            else:
+                assert self.cg_buf_num_splits is not None
+
+                # Metadata per-SM, fixed size (#SMs, TileMetadataSize)
+                assert (self.cg_buf_tile_scheduler_metadata.size() ==
+                        tile_scheduler_metadata.size())
+                self.cg_buf_tile_scheduler_metadata.\
+                    copy_(tile_scheduler_metadata)
+                tile_scheduler_metadata = self.cg_buf_tile_scheduler_metadata
+
+                # Num splits is per-batch, varying size (batch_size,)
+                n = num_splits.size(0)
+                # make sure static buffer is large enough
+                assert n <= self.cg_buf_num_splits.size(0)
+                num_splits_view = self.cg_buf_num_splits[:n]
+                num_splits_view.copy_(num_splits)
+                self.cg_buf_num_splits[n:].fill_(0)  # fill the rest with 0s
+                num_splits = num_splits_view
 
         return FlashMLADecodeMetadata(
             block_table=block_table_tensor,
diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
index 1f0406a7a..9fbca2e95 100644
--- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py
@@ -66,7 +66,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
 
     def __init__(self, runner, kv_cache_spec: AttentionSpec,
                  block_table: BlockTable):
-        super().__init__(runner, kv_cache_spec, block_table)
+        super().__init__(runner, kv_cache_spec, block_table, AiterMLAMetadata)
         assert self.kv_cache_spec.block_size == 1, "AITER MLA" \
             "only supports block size 1."
 
diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py
index e26d79091..523f17263 100644
--- a/vllm/v1/attention/backends/mla/triton_mla.py
+++ b/vllm/v1/attention/backends/mla/triton_mla.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass
 from typing import Any, Optional
 
 import torch
@@ -10,11 +11,20 @@ from vllm.attention.backends.abstract import (AttentionType,
 from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
 from vllm.logger import init_logger
 from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
+                                                   MLACommonDecodeMetadata,
                                                    MLACommonImpl,
-                                                   MLACommonMetadata)
-
+                                                   MLACommonMetadata,
+                                                   MLACommonMetadataBuilder)
+from vllm.attention.backends.triton_mla import (load_config,
+                                                find_best_mla_para)
 logger = init_logger(__name__)
 
+import os
+# TODO: Configure environment variables temporarily. New versions do not need to be configured
+os.environ['TRITON_ENABLE_MACA_OPT_MOVE_DOT_OPERANDS_OUT_LOOP'] = '1'
+os.environ['TRITON_ENABLE_MACA_CHAIN_DOT_OPT'] = '1'
+
+JSON_DATA = load_config()
 
 class TritonMLABackend(MLACommonBackend):
 
@@ -22,10 +32,42 @@ class TritonMLABackend(MLACommonBackend):
     def get_name() -> str:
         return "TRITON_MLA_VLLM_V1"
 
+    @staticmethod
+    def get_metadata_cls() -> type["TritonMLAMetadata"]:
+        return TritonMLAMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["TritonMLAMetadataBuilder"]:
+        return TritonMLAMetadataBuilder
+
     @staticmethod
     def get_impl_cls() -> type["TritonMLAImpl"]:
         return TritonMLAImpl
-
+@dataclass
+class TritonMLADecodeMetadata(MLACommonDecodeMetadata):
+    num_kv_splits: int
+    num_stages: int
+
+@dataclass
+class TritonMLAMetadata(MLACommonMetadata[TritonMLADecodeMetadata]):
+    pass
+
+class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[TritonMLAMetadata]):
+    def _build_decode(self, block_table_tensor: torch.Tensor,
+                      seq_lens: torch.Tensor) -> TritonMLADecodeMetadata:
+        if seq_lens is not None:
+            batch = seq_lens.shape[0]
+            max_seq_len = int(seq_lens.max())
+            num_kv_splits, num_stages = find_best_mla_para(JSON_DATA, batch, max_seq_len, 8)
+        else:
+            num_kv_splits = 4
+            num_stages = 1
+        return TritonMLADecodeMetadata(
+            block_table=block_table_tensor,
+            seq_lens=seq_lens,
+            num_kv_splits=num_kv_splits,
+            num_stages=num_stages,
+        )
 
 class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
 
@@ -90,14 +132,12 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
                         dtype=q.dtype,
                         device=q.device)
 
-        num_kv_splits = 4  # TODO: heuristic
-
         # TODO(lucas) Allocate ahead of time
         attn_logits = torch.empty(
             (
                 B,
                 self.num_heads,
-                num_kv_splits,
+                attn_metadata.decode.num_kv_splits,
                 # NOTE(lucas) idk why the +1 is here but sglang has it so we
                 # just mirror that
                 self.kv_lora_rank + 1,
@@ -115,6 +155,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
         decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
                              attn_metadata.decode.block_table,
                              attn_metadata.decode.seq_lens, attn_logits,
-                             num_kv_splits, self.scale, PAGE_SIZE)
+                             attn_metadata.decode.num_kv_splits,
+                             attn_metadata.decode.num_stages,
+                             self.scale, PAGE_SIZE)
 
         return self._v_up_proj(o)
diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py
index 72c764353..c53674f25 100644
--- a/vllm/v1/attention/backends/utils.py
+++ b/vllm/v1/attention/backends/utils.py
@@ -1,15 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import abc
+from abc import abstractmethod
 from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar
 
+import numpy as np
 import torch
 
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.worker.gpu_input_batch import InputBatch
 
 @dataclass
 class CommonAttentionMetadata:
     """
-    Attention metadata attributes that can be shared by layers in different KV
-    cache groups and thus having different block table.
+    Per-batch attention metadata, shared across layers and backends.
+    AttentionMetadataBuilder instances use it to construct per-layer metadata.
     """
 
     query_start_loc: torch.Tensor
@@ -18,6 +25,65 @@ class CommonAttentionMetadata:
     """(batch_size,), the length of each request including both computed tokens
     and newly scheduled tokens"""
 
+    num_reqs: int
+    """Number of requests"""
+    num_actual_tokens: int
+    """Total number of tokens in batch"""
+    max_query_len: int
+    """Longest query in batch"""
+
+M = TypeVar("M")
+
+
+class AttentionMetadataBuilder(abc.ABC, Generic[M]):
+    # Does this backend/builder support CUDA Graphs for attention.
+    full_cudagraph_supported: ClassVar[bool] = False
+
+    @abstractmethod
+    def build(self, common_prefix_len: int,
+              common_attn_metadata: CommonAttentionMetadata) -> M:
+        """
+        Central method that builds attention metadata.
+        Some builders (MLA) require reorder_batch to be called prior to build.
+        """
+        raise NotImplementedError
+
+    def can_run_in_cudagraph(
+            self, common_attn_metadata: CommonAttentionMetadata) -> bool:
+        """
+        Can this batch (with given metadata) use CUDA Graphs for attention.
+        """
+        return False
+
+    def build_for_cudagraph_capture(
+            self, common_attn_metadata: CommonAttentionMetadata) -> M:
+        """
+        Build attention metadata for CUDA graph capture. Uses build by default.
+        Subclasses that override this method should call self.build or
+        super().build_for_cudagraph_capture.
+        """
+        return self.build(common_prefix_len=0,
+                          common_attn_metadata=common_attn_metadata)
+
+    def use_cascade_attention(
+        self,
+        common_prefix_len: int,
+        query_lens: np.ndarray,
+        num_query_heads: int,
+        num_kv_heads: int,
+        use_alibi: bool,
+        use_sliding_window: bool,
+        num_sms: int,
+    ) -> bool:
+        return False
+
+    def reorder_batch(self, input_batch: "InputBatch",
+                      scheduler_output: "SchedulerOutput") -> bool:
+        """
+        This method can reorder the batch if desired by the backend.
+        :return: Has the batch been reordered (default False).
+        """
+        return False
 
 def validate_kv_sharing_target(current_layer_name, target_layer_name,
                                static_forward_context):
diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
index f36a491a1..dd8730aed 100644
--- a/vllm/v1/engine/core.py
+++ b/vllm/v1/engine/core.py
@@ -799,6 +799,7 @@ class DPEngineCoreProc(EngineCoreProc):
             str(current_platform.device_id_to_physical_device_id(i))
             for i in range(local_dp_rank * world_size, (local_dp_rank + 1) *
                            world_size))
+        os.environ["MACA_VISIBLE_DEVICES"] = os.environ[device_control_env_var]
 
         self.dp_rank = dp_rank
         self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()
@@ -917,7 +918,7 @@ class DPEngineCoreActor(DPEngineCoreProc):
         # Ray sets CUDA_VISIBLE_DEVICES to empty string,
         # we clean this up to be able to properly initialize
         # data parallel groups.
-        del os.environ['CUDA_VISIBLE_DEVICES']
+        # del os.environ['CUDA_VISIBLE_DEVICES']
 
         super().__init__(vllm_config, on_head_node, "", executor_class,
                          log_stats)
diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py
index c6fe2d339..9ec2f5ce2 100644
--- a/vllm/v1/engine/detokenizer.py
+++ b/vllm/v1/engine/detokenizer.py
@@ -17,6 +17,13 @@ from vllm.v1.engine import EngineCoreRequest
 
 logger = init_logger(__name__)
 
+# Only tokenizers >= 0.21.1 supports DecodeStream used for
+# FastIncrementalDetokenizer.
+USE_FAST_DETOKENIZER = version.parse(
+    tokenizers.__version__) >= version.parse("0.21.1")
+
+# Error string from https://github.com/huggingface/tokenizers/blob/909fdde2a4ffedd9295206f705eb612be2a91b12/tokenizers/src/tokenizer/mod.rs#L1042
+INVALID_PREFIX_ERR_MSG = "Invalid prefix encountered"
 
 class IncrementalDetokenizer:
 
@@ -46,10 +53,9 @@ class IncrementalDetokenizer:
             # No tokenizer => skipping detokenization.
             return IncrementalDetokenizer()
 
-        if (isinstance(tokenizer, PreTrainedTokenizerFast) and version.parse(
-                tokenizers.__version__) >= version.parse("0.21.1")):
+        if USE_FAST_DETOKENIZER and isinstance(tokenizer,
+                                               PreTrainedTokenizerFast):
             # Fast tokenizer => use tokenizers library DecodeStream.
-            # And only tokenizers >= 0.21.1 supports Fast Detokenizer.
             return FastIncrementalDetokenizer(tokenizer, request)
 
         # Fall back to slow python-based incremental detokenization.
@@ -157,8 +163,11 @@ class FastIncrementalDetokenizer(BaseIncrementalDetokenizer):
         super().__init__(request)
 
         sampling_params = request.sampling_params
+
+        self.request_id = request.request_id
+        self.skip_special_tokens = sampling_params.skip_special_tokens
         self.stream = DecodeStream(
-            skip_special_tokens=sampling_params.skip_special_tokens)
+            skip_special_tokens=self.skip_special_tokens)
 
         self.tokenizer: Tokenizer = tokenizer._tokenizer
 
@@ -174,7 +183,7 @@ class FastIncrementalDetokenizer(BaseIncrementalDetokenizer):
 
         # Prime the stream.
         for tid in prompt_suffix:
-            self.stream.step(self.tokenizer, tid)
+            self._protected_step(tid)
 
         self.spaces_between_special_tokens = (
             sampling_params.skip_special_tokens
@@ -199,7 +208,7 @@ class FastIncrementalDetokenizer(BaseIncrementalDetokenizer):
                 self.spaces_between_special_tokens = True
 
     def decode_next(self, next_token_id: int) -> str:
-        token = self.stream.step(self.tokenizer, next_token_id)
+        token = self._protected_step(next_token_id)
 
         if not self.spaces_between_special_tokens:
             special_token = self.added_token_ids.get(next_token_id)
@@ -211,6 +220,22 @@ class FastIncrementalDetokenizer(BaseIncrementalDetokenizer):
 
         return token or ""
 
+    def _protected_step(self, next_token_id: int) -> Optional[str]:
+        try:
+            token = self.stream.step(self.tokenizer, next_token_id)
+        except Exception as e:
+            if str(e) != INVALID_PREFIX_ERR_MSG:
+                raise e
+            # Recover from edge case where tokenizer can produce non-monotonic,
+            # invalid UTF-8 output, which breaks the internal state of
+            # tokenizers' DecodeStream.
+            # See https://github.com/vllm-project/vllm/issues/17448.
+            logger.warning(
+                "Encountered invalid prefix detokenization error"
+                " for request %s, resetting decode stream.", self.request_id)
+            self.stream = DecodeStream(self.skip_special_tokens)
+            token = self.stream.step(self.tokenizer, next_token_id)
+        return token
 
 class SlowIncrementalDetokenizer(BaseIncrementalDetokenizer):
 
diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py
index b2354c533..efaa85145 100644
--- a/vllm/v1/sample/rejection_sampler.py
+++ b/vllm/v1/sample/rejection_sampler.py
@@ -446,7 +446,7 @@ def rejection_greedy_sample_kernel(
         is_greedy = True
     else:
         is_greedy = tl.load(is_greedy_ptr + req_idx)
-    if not is_greedy:
+    if is_greedy is None:
         # Early exit for non-greedy sampling requests.
         return
 
@@ -494,7 +494,7 @@ def rejection_random_sample_kernel(
 ):
     req_idx = tl.program_id(0)
     is_greedy = tl.load(is_greedy_ptr + req_idx)
-    if is_greedy:
+    if is_greedy is not None:
         # Early exit for greedy sampling requests.
         return
 
diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py
index 4b5c9b7ec..47b0fbf54 100644
--- a/vllm/v1/spec_decode/eagle.py
+++ b/vllm/v1/spec_decode/eagle.py
@@ -138,15 +138,17 @@ class EagleProposer:
             max_query_len = query_lens.max().item()
 
             common_attn_metadata = CommonAttentionMetadata(
-                query_start_loc=cu_num_tokens, seq_lens=seq_lens)
+                query_start_loc=cu_num_tokens,
+                seq_lens=seq_lens,
+                num_reqs=batch_size,
+                num_actual_tokens=num_tokens,
+                max_query_len=max_query_len,
+            )
 
             assert self.runner is not None
 
             # FIXME: need to consider multiple kv_cache_groups
-            attn_metadata = self.runner.attn_metadata_builder.build(
-                num_reqs=batch_size,
-                num_actual_tokens=num_tokens,
-                max_query_len=max_query_len,
+            attn_metadata = self.runner.attn_metadata_builders[0].build(
                 common_prefix_len=0,
                 common_attn_metadata=common_attn_metadata,
             )
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index b1bc727e1..88d1d35ba 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -15,10 +15,8 @@ import torch.nn as nn
 
 import vllm.envs as envs
 from vllm.attention import AttentionType, get_attn_backend
-from vllm.attention.backends.abstract import (AttentionBackend,
-                                              AttentionMetadataBuilder)
+from vllm.attention.backends.abstract import AttentionBackend
 from vllm.attention.layer import Attention
-from vllm.attention.utils.fa_utils import get_flash_attn_version
 from vllm.config import (CompilationLevel, VllmConfig,
                          get_layers_from_vllm_config)
 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
@@ -40,7 +38,8 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
                         GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,
                         check_use_alibi, is_pin_memory_available)
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
+                                              CommonAttentionMetadata)
 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
 from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
                                         KVCacheConfig, KVCacheSpec,
@@ -74,6 +73,24 @@ else:
 logger = init_logger(__name__)
 
 
+# --- 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:
+        print("Try to using FLAGGEMS...")
+        import flag_gems
+        flag_gems.enable(record=True, unused=["exponential_", "softmax", "masked_fill_", "index"], path="/tmp/gems_oplist.log.txt")
+        logger.info("Successfully enabled flag_gems as default ops implementation.")
+     except ImportError as e:
+        # Throw an exception directly if failure occurs
+        raise ImportError("Failed to import 'flag_gems'. Please install flag_gems or set USE_FLAGGEMS=false to disable it.") from e
+     except Exception as e:
+        # Throw an exception directly if failure occurs
+        raise RuntimeError(f"Failed to enable 'flag_gems': {e}. Please check your flag_gems installation or set USE_FLAGGEMS=false to disable it.") from e
+# --- FLAGSCALE MODIFICATION END ---
+
+
 class GPUModelRunner(LoRAModelRunnerMixin):
 
     def __init__(
@@ -84,6 +101,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.vllm_config = vllm_config
         self.model_config = vllm_config.model_config
         self.cache_config = vllm_config.cache_config
+        self.compilation_config = vllm_config.compilation_config
         self.lora_config = vllm_config.lora_config
         self.load_config = vllm_config.load_config
         self.parallel_config = vllm_config.parallel_config
@@ -192,7 +210,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             block_sizes=[self.cache_config.block_size],
         )
 
-        self.use_cuda_graph = (self.vllm_config.compilation_config.level
+        self.use_cuda_graph = (self.compilation_config.level
                                == CompilationLevel.PIECEWISE
                                and not self.model_config.enforce_eager)
         # TODO(woosuk): Provide an option to tune the max cudagraph batch size.
@@ -200,8 +218,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # self.cudagraph_batch_sizes sorts in ascending order.
         # The batch sizes in the config are in descending order.
         self.cudagraph_batch_sizes = list(
-            reversed(
-                self.vllm_config.compilation_config.cudagraph_capture_sizes))
+           reversed(self.compilation_config.cudagraph_capture_sizes))
+
+        self.full_cuda_graph = self.compilation_config.full_cuda_graph
 
         # Cache the device properties.
         self._init_device_properties()
@@ -552,7 +571,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
     def _prepare_inputs(
         self,
         scheduler_output: "SchedulerOutput",
-    ) -> tuple[dict[str, Any], torch.Tensor, Optional[SpecDecodeMetadata]]:
+     ) -> tuple[dict[str, Any], bool, torch.Tensor,
+               Optional[SpecDecodeMetadata]]:
+        """
+        :return: tuple[
+            attn_metadata: layer-to-attention_metadata mapping,
+            attention_cuda_graphs: whether attention can run in cudagraph
+            logits_indices, spec_decode_metadata
+        ]
+        """
         total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         assert total_num_scheduled_tokens > 0
         num_reqs = self.input_batch.num_reqs
@@ -666,7 +693,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         seq_lens = self.seq_lens[:num_reqs]
 
         common_attn_metadata = CommonAttentionMetadata(
-            query_start_loc=query_start_loc, seq_lens=seq_lens)
+            query_start_loc=query_start_loc,
+            seq_lens=seq_lens,
+            num_reqs=num_reqs,
+            num_actual_tokens=total_num_scheduled_tokens,
+            max_query_len=max_num_scheduled_tokens,
+        )
 
         attn_metadata: dict[str, Any] = {}
         # Prepare the attention metadata for each KV cache group and make layers
@@ -676,25 +708,28 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
             # Prepare for cascade attention if enabled & beneficial.
             common_prefix_len = 0
+            builder = self.attn_metadata_builders[kv_cache_group_id]
             if self.cascade_attn_enabled:
                 common_prefix_len = self._compute_cascade_attn_prefix_len(
                     num_scheduled_tokens,
                     scheduler_output.
                     num_common_prefix_blocks[kv_cache_group_id],
                     kv_cache_group_spec.kv_cache_spec,
-                    self.attn_metadata_builders[kv_cache_group_id],
+                    builder,
                 )
 
-            attn_metadata_i = (
-                self.attn_metadata_builders[kv_cache_group_id].build(
-                    num_reqs=num_reqs,
-                    num_actual_tokens=total_num_scheduled_tokens,
-                    max_query_len=max_num_scheduled_tokens,
-                    common_prefix_len=common_prefix_len,
-                    common_attn_metadata=common_attn_metadata))
+            attn_metadata_i = (builder.build(
+                common_prefix_len=common_prefix_len,
+                common_attn_metadata=common_attn_metadata,
+            ))
+
             for layer_name in kv_cache_group_spec.layer_names:
                 attn_metadata[layer_name] = attn_metadata_i
 
+        attention_cuda_graphs = all(
+            b.can_run_in_cudagraph(common_attn_metadata)
+            for b in self.attn_metadata_builders)
+
         use_spec_decode = len(
             scheduler_output.scheduled_spec_decode_tokens) > 0
         if not use_spec_decode:
@@ -723,7 +758,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         if self.lora_config:
             self.set_active_loras(self.input_batch, num_scheduled_tokens)
 
-        return attn_metadata, logits_indices, spec_decode_metadata
+        return (attn_metadata, attention_cuda_graphs, logits_indices,
+                spec_decode_metadata)
 
     def _compute_cascade_attn_prefix_len(
         self,
@@ -1115,7 +1151,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         assert self.intermediate_tensors is not None
 
         tp = self.vllm_config.parallel_config.tensor_parallel_size
-        enabled_sp = self.vllm_config.compilation_config.pass_config. \
+        enabled_sp = self.compilation_config.pass_config. \
             enable_sequence_parallelism
         if enabled_sp:
             # When sequence parallelism is enabled, we always pad num_tokens
@@ -1183,8 +1219,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             return self.kv_connector_no_forward(scheduler_output)
 
         # Prepare the decoder inputs.
-        attn_metadata, logits_indices, spec_decode_metadata = (
-            self._prepare_inputs(scheduler_output))
+        (attn_metadata, attention_cuda_graphs, logits_indices,
+         spec_decode_metadata) = (self._prepare_inputs(scheduler_output))
         num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         if (self.use_cuda_graph
                 and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
@@ -1197,7 +1233,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             # Pad tokens to multiple of tensor_parallel_size when
             # enabled collective fusion for SP
             tp_size = self.vllm_config.parallel_config.tensor_parallel_size
-            if self.vllm_config.compilation_config.pass_config. \
+            if self.compilation_config.pass_config. \
                 enable_sequence_parallelism and tp_size > 1:
                 from vllm.utils import round_up
                 num_input_tokens = round_up(num_scheduled_tokens, tp_size)
@@ -1249,12 +1285,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             intermediate_tensors = self.sync_and_slice_intermediate_tensors(
                 num_input_tokens, intermediate_tensors, True)
 
+        # Some attention backends only support CUDA Graphs in pure decode.
+        # If attention doesn't support CUDA Graphs for this batch, but we
+        # compiled with full CUDA graphs, we have to skip them entirely.
+        skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
+
         # Run the decoder.
         # Use persistent buffers for CUDA graphs.
-        with set_forward_context(attn_metadata,
-                                 self.vllm_config,
-                                 num_tokens=num_input_tokens,
-                                 num_tokens_across_dp=num_tokens_across_dp):
+        with set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                skip_cuda_graphs=skip_cuda_graphs,
+            ):
             self.maybe_setup_kv_connector(scheduler_output)
 
             model_output = self.model(
@@ -1763,7 +1807,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
     def _dummy_run(
         self,
         num_tokens: int,
-        skip_attn: bool = True,
+        capture_attn_cudagraph: bool = False,
     ) -> torch.Tensor:
 
         # Padding for DP
@@ -1784,9 +1828,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         num_scheduled_tokens = np.array(num_scheduled_tokens_list,
                                         dtype=np.int32)
 
-        if skip_attn:
-            attn_metadata: Optional[dict[str, Any]] = None
-        else:
+        attn_metadata: Optional[dict[str, Any]] = None
+        if capture_attn_cudagraph:
+            attn_metadata = {}
             query_start_loc = self.query_start_loc[:num_reqs + 1]
             # Make sure max_model_len is used at the graph capture time.
             self.seq_lens_np[:num_reqs] = self.max_model_len
@@ -1796,19 +1840,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             seq_lens = self.seq_lens[:num_reqs]
 
             common_attn_metadata = CommonAttentionMetadata(
-                query_start_loc=query_start_loc, seq_lens=seq_lens)
+                query_start_loc=query_start_loc,
+                seq_lens=seq_lens,
+                num_reqs=num_reqs,
+                num_actual_tokens=num_tokens,
+                max_query_len=num_tokens,
+            )
 
-            attn_metadata = {}
             for kv_cache_group_id, kv_cache_group_spec in enumerate(
                     self.kv_cache_config.kv_cache_groups):
-                attn_metadata_i = (
-                    self.attn_metadata_builders[kv_cache_group_id].build(
-                        num_reqs=num_reqs,
-                        num_actual_tokens=num_tokens,
-                        max_query_len=num_tokens,
-                        common_prefix_len=0,
-                        common_attn_metadata=common_attn_metadata,
-                    ))
+                attn_metadata_i = self.attn_metadata_builders[
+                    kv_cache_group_id].build_for_cudagraph_capture(
+                        common_attn_metadata)
                 for layer_name in kv_cache_group_spec.layer_names:
                     attn_metadata[layer_name] = attn_metadata_i
 
@@ -1880,7 +1923,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         dummy_metadata = SamplingMetadata(
             temperature=dummy_tensors(0.5),
-            all_greedy=False,
+            all_greedy=True,
             all_random=False,
             top_p=dummy_tensors(0.9),
             top_k=dummy_tensors(logits.size(1) - 1),
@@ -2033,12 +2076,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # Capture the large shapes first so that the smaller shapes
         # can reuse the memory pool allocated for the large shapes.
         with graph_capture(device=self.device):
-            skip_attn = not self.vllm_config.compilation_config.full_cuda_graph
+            full_cg = self.full_cuda_graph
             for num_tokens in reversed(self.cudagraph_batch_sizes):
-                for _ in range(self.vllm_config.compilation_config.
-                               cudagraph_num_of_warmups):
-                    self._dummy_run(num_tokens, skip_attn=skip_attn)
-                self._dummy_run(num_tokens, skip_attn=skip_attn)
+                for _ in range(
+                        self.compilation_config.cudagraph_num_of_warmups):
+                    self._dummy_run(num_tokens, capture_attn_cudagraph=full_cg)
+                self._dummy_run(num_tokens, capture_attn_cudagraph=full_cg)
 
         end_time = time.perf_counter()
         end_free_gpu_memory = torch.cuda.mem_get_info()[0]
@@ -2081,20 +2124,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     "Non-Attention backend is not supported by V1 "
                     "GPUModelRunner.")
 
-            if self.vllm_config.compilation_config.full_cuda_graph:
-                attn_backend_name = attn_backend_i.__name__
-                flash_attn_version = get_flash_attn_version()
-                if attn_backend_name != "FlashAttentionBackend" or \
-                    flash_attn_version != 3:
-                    raise ValueError(
-                        f"full_cuda_graph is only supported with "
-                        f"FA3. Current attention backend is "
-                        f"{attn_backend_name}, FlashAttention version is "
-                        f"{flash_attn_version}.")
-
             block_table_i = self.input_batch.block_table[i]
             attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
-                weakref.proxy(self), kv_cache_spec, block_table_i)
+                weakref.proxy(self),
+                kv_cache_spec,
+                block_table_i,
+            )
+
+            if (self.full_cuda_graph
+                and not attn_metadata_builder_i.full_cudagraph_supported):
+                raise ValueError(
+                    f"Full CUDAGraph not supported for "
+                    f"{attn_backend_i.__name__}. Turn off CompilationConfig."
+                    f"full_cuda_graph or use a different attention backend.")
+
             self.attn_backends.append(attn_backend_i)
             self.attn_metadata_builders.append(attn_metadata_builder_i)
 
@@ -2237,10 +2280,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 kv_caches,
             )
 
-        bind_kv_cache(
-            kv_caches,
-            self.vllm_config.compilation_config.static_forward_context,
-            self.kv_caches)
+        bind_kv_cache(kv_caches,
+                      self.compilation_config.static_forward_context,
+                      self.kv_caches)
         return kv_caches
 
     def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
diff --git a/vllm/version.py b/vllm/version.py
index 6c88b1b5a..11e21ccd5 100644
--- a/vllm/version.py
+++ b/vllm/version.py
@@ -6,11 +6,11 @@ try:
 except Exception as e:
     import warnings
 
-    warnings.warn(f"Failed to read commit hash:\n{e}",
-                  RuntimeWarning,
-                  stacklevel=2)
+    # warnings.warn(f"Failed to read commit hash:\n{e}",
+    #               RuntimeWarning,
+    #               stacklevel=2)
 
-    __version__ = "dev"
+    __version__ = "0.9.1"
     __version_tuple__ = (0, 0, __version__)
 
 
diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
index 82db6617b..dcae1a098 100644
--- a/vllm/worker/model_runner.py
+++ b/vllm/worker/model_runner.py
@@ -78,6 +78,24 @@ 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:
+        print("Try to using FLAGGEMS...")
+        import flag_gems
+        flag_gems.enable(record=True, unused=["exponential_", "softmax", "masked_fill_", "index"], path="/tmp/gems_oplist.log.txt")
+        logger.info("Successfully enabled flag_gems as default ops implementation.")
+     except ImportError as e:
+        # Throw an exception directly if failure occurs
+        raise ImportError("Failed to import 'flag_gems'. Please install flag_gems or set USE_FLAGGEMS=false to disable it.") from e
+     except Exception as e:
+        # Throw an exception directly if failure occurs
+        raise RuntimeError(f"Failed to enable 'flag_gems': {e}. Please check your flag_gems installation or set USE_FLAGGEMS=false to disable it.") from e
+# --- FLAGSCALE MODIFICATION END ---
+
+
 @dataclass(frozen=True)
 class ModelInputForGPU(ModelRunnerInputBase):
     """
diff --git a/vllm/worker/multi_step_model_runner.py b/vllm/worker/multi_step_model_runner.py
index cc0cc855e..927475bbe 100644
--- a/vllm/worker/multi_step_model_runner.py
+++ b/vllm/worker/multi_step_model_runner.py
@@ -34,7 +34,7 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 MULTI_STEP_ATTENTION_BACKENDS = [
-    "FLASH_ATTN", "ROCM_FLASH", "FLASHINFER", "NO_ATTENTION"
+    "FLASH_ATTN", "ROCM_FLASH", "FLASHINFER", "NO_ATTENTION", "TRITON_MLA"
 ]
 MULTI_STEP_CHUNKED_PREFILL_ATTENTION_BACKENDS = ["FLASH_ATTN", "FLASHINFER"]
 
diff --git a/vllm/worker/worker_base.py b/vllm/worker/worker_base.py
index 0b37caa71..7db5f1496 100644
--- a/vllm/worker/worker_base.py
+++ b/vllm/worker/worker_base.py
@@ -540,6 +540,8 @@ class WorkerWrapperBase:
             # suppress the warning in `update_environment_variables`
             del os.environ[key]
         update_environment_variables(envs)
+        if key in os.environ:
+            os.environ["MACA_VISIBLE_DEVICES"] = os.environ[key]
 
     def init_worker(self, all_kwargs: List[Dict[str, Any]]) -> None:
         """
